[PyTorch] TE 2.17 Userbuffers initialization crashes in CommOverlapP2P with a GIL/refcount error

[cuichenx]

[PyTorch] TE 2.17 Userbuffers initialization crashes in CommOverlapP2P with a GIL/refcount error

Describe the bug

Upgrading only Transformer Engine from 2.16 to 2.17 causes the public PyTorch initialize_ub() path to abort while constructing its first CommOverlapP2P communicator. The failure occurs with the legacy Userbuffers backend (with_cublasmp=False) before training starts:

!!! [UBP2P] UBuf 1
pybind11::handle::inc_ref() is being called while the GIL is either not held or invalid.
The failing pybind11::handle::inc_ref() call was triggered on a tuple object.
munmap_chunk(): invalid pointer
Fatal Python error: Aborted

The second process subsequently reports:

Segfault encountered
cfree
pybind11::class_<CommOverlapP2P, std::shared_ptr<CommOverlapP2P>, ...>
    ::dealloc_without_manipulating_gil(...)

This is deterministic in two CI attempts with TE 2.17. The same test initializes all Userbuffers communicators and passes with TE 2.16.

Steps/Code to reproduce bug

The regression is reproduced by Megatron-Bridge PR #4696, which changes only the TE source pin and lockfile:

• Passing: 2.16.0+d64bc14d at d64bc14dc87eb658ab98839e4b7687595ee53e2d
• Failing: 2.17.0+2e559f06 at 2e559f062497bef768dfbe9d7e45548fadeca80a

From that PR checkout, run the two-rank recipe test:

uv run python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 \
  -m pytest -v -s -x \
  tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py

The first two TP=1 cases pass. The Nemotron 3 Nano case sets tensor_model_parallel_size=2, sequence_parallel=True, and tp_comm_overlap=True, then aborts in:

transformer_engine/pytorch/module/base.py:448 in add_ub
transformer_engine/pytorch/module/base.py:538 in initialize_ub

CI evidence:

• TE 2.17, attempt 1: failed job
• TE 2.17, attempt 2: failed job with the native CommOverlapP2P deallocation stack
• TE 2.16 control: passing job; all four pretrain cases pass

The failure uses the documented default with_cublasmp=False; it is not a cuBLASMp execution path.

Expected behavior

initialize_ub() should construct the TP-overlap communicators without touching Python reference counts while the GIL is released. Behavior should match TE 2.16, where all Userbuffers communicators initialize and training proceeds.

Environment overview (please complete the following information)

• Environment location: Docker on AWS GitHub Actions GPU runners (nemo-ci-aws-gpu-x2)
• Transformer Engine install: Git source pin resolved by uv
• Test launch: two local processes / two visible GPUs
• Userbuffers backend: with_cublasmp=False, bootstrapped through torch.distributed

Environment details

• OS: Linux container
• Python: 3.12.3
• PyTorch: 2.12.1
• Passing Transformer Engine: 2.16.0+d64bc14d
• Failing Transformer Engine: 2.17.0+2e559f06
• The full container and runner metadata is in the linked jobs.

Device details

• Two-GPU AWS CI runner; the launcher selects the H100 BF16 recipe configuration.

Additional context

The likely source-level regression is 815bf369e60f5d56ef20badd36b6aa5e359adb48 / PR #2443. That change replaced the direct py::init<...> bindings for both CommOverlap and CommOverlapP2P with factory lambdas returning std::shared_ptr, while retaining the outer py::call_guard<py::gil_scoped_release>():

• CommOverlap factory and call guard
• CommOverlapP2P factory and call guard

The CI failure is directly observed in CommOverlapP2P. CommOverlap has the same new factory/guard pattern, so it may have the same lifetime hazard even though this run reaches the P2P factory first.

My source-level suspicion is that the outer call guard leaves the GIL released while pybind11 converts/wraps the factory return. A candidate fix would retain the GIL across the factory return boundary and, if construction needs to run without it, use a narrower py::gil_scoped_release only around std::make_shared, allowing the GIL to be reacquired before returning to pybind11. This hypothesis still needs confirmation with a patched TE build and a two-rank regression test.

[balasaajay]

Adding another data point from Megatron-LM PR 5680 CI: https://github.com/NVIDIA/Megatron-LM/actions/runs/29059466380?pr=5680

Three real test jobs failed with the same pybind11 GIL/refcount abort signature:

• tests/unit_tests/transformer/moe/**/*.py - latest: https://github.com/NVIDIA/Megatron-LM/actions/runs/29059466380/job/86263669371
  • First failing test observed: tests/unit_tests/transformer/moe/test_token_dispatcher.py::TestFlexDispatcher::test_forward_backward[False-ncclep-False-1-8]
• tests/unit_tests/a2a_overlap/**/*.py - latest: https://github.com/NVIDIA/Megatron-LM/actions/runs/29059466380/job/86263669405
  • First failing test observed: tests/unit_tests/a2a_overlap/test_delay_wgrad_compute.py::TestDelayWgradCompute::test_overlap_dispatch_backward_with_experts_wgrad[fp8_flag0-flex-ncclep-None]
• tests/unit_tests/**/*.py - gb200 latest: https://github.com/NVIDIA/Megatron-LM/actions/runs/29059466380/job/86263671680
  • First failing test observed: tests/unit_tests/transformer/moe/test_paged_stashing.py::TestNcclEpPagedStashing::test_forward_backward_4_layers

Nemo_CICD_Test failed only as the rollup job because it detected the failed unit jobs.

Common signature across the failing jobs:

NCCL EP: using timeout=...
pybind11::handle::dec_ref() is being called while the GIL is either not held or invalid.
The failing pybind11::handle::dec_ref() call was triggered on a torch.dtype object.
terminate called after throwing an instance of 'std::runtime_error'
  what():  pybind11::handle::dec_ref() PyGILState_Check() failure.
Fatal Python error: Aborted
...
torch.distributed.elastic.multiprocessing.errors.ChildFailedError
exitcode  : -6  (SIGABRT)

The stack in these Megatron-LM failures goes through NCCL EP bootstrap rather than the Userbuffers CommOverlapP2P path from the original report:

transformer_engine/pytorch/ep.py:125 in ep_bootstrap
megatron/core/transformer/moe/fused_a2a.py:646 in ensure_nccl_ep_bootstrapped
megatron/core/transformer/moe/token_dispatcher.py:1535 in _ensure_bootstrap

So this looks like the same TE 2.17 pybind11 GIL/refcount regression family, but affecting the NCCL EP bootstrap path as well.

[ptrendx]

@denera please take a look

[cuichenx]

this failure is limited to debug/assertion-enabled builds; normal optimized builds with -DNDEBUG do not reproduce it.
Megatron-Bridge CI was not explicitly a debug build, but its setuptools build path dropped Python’s default -DNDEBUG, effectively enabling the same pybind11 assertions. This is fixed by #3203