Skip to content

[PyTorch] Fix GIL/refcount abort in Comm+GEMM overlap and NCCL-EP init bindings - #3203

Merged
denera merged 1 commit into
NVIDIA:mainfrom
denera:pytorch/comm-overlap-pybind-gil-fix
Jul 10, 2026
Merged

[PyTorch] Fix GIL/refcount abort in Comm+GEMM overlap and NCCL-EP init bindings#3203
denera merged 1 commit into
NVIDIA:mainfrom
denera:pytorch/comm-overlap-pybind-gil-fix

Conversation

@denera

@denera denera commented Jul 10, 2026

Copy link
Copy Markdown
Collaborator

Description

The factory py::init bindings for CommOverlap/CommOverlapP2P and the ep_initialize binding were registered with an outer py::call_guard<py::gil_scoped_release>(), which releases the GIL for the entire call -- including regions where pybind11 performs Python reference-count operations that require the GIL to be held. This triggers a pybind11 GIL/ref-count hazard assert when TE is compiled with NVTE_BUILD_DEBUG=1.

PR fixes the bug by dropping the outer call guard and releasing the GIL narrowly, only around the blocking C++ native constructor/initializer. The GIL is reacquired before pybind11 casts the arguments, registers the holder and destroys the by-value object parameters.

Fixes # 3199

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@denera
denera requested a review from timmoon10 July 10, 2026 18:37
@denera denera self-assigned this Jul 10, 2026
@denera denera added the bug Something isn't working label Jul 10, 2026
@denera
denera requested a review from ksivaman as a code owner July 10, 2026 18:37
@greptile-apps

greptile-apps Bot commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR narrows where the PyTorch bindings release the GIL.

  • ep_initialize now keeps the GIL during pybind argument handling and releases it only around native EP initialization.
  • CommOverlap and CommOverlapP2P constructors now release the GIL inside the init lambda around native construction.
  • The outer pybind call guards were removed from the affected init bindings.

Confidence Score: 5/5

This looks safe to merge.

  • No blocking issues found in the changed code.

Important Files Changed

Filename Overview
transformer_engine/pytorch/csrc/extensions/ep.cpp The EP init binding now releases the GIL only around the native initializer while keeping pybind object handling under the GIL.
transformer_engine/pytorch/csrc/extensions/pybind.cpp The CommOverlap bindings now release the GIL inside the constructor lambdas after pybind has converted the Python arguments.

Reviews (1): Last reviewed commit: "Fix GIL/refcount abort in Comm+GEMM over..." | Re-trigger Greptile

@timmoon10 timmoon10 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@denera
denera merged commit 6377ca1 into NVIDIA:main Jul 10, 2026
11 of 15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants