Skip to content

fix(distributed): subgroup NCCL regression + nccl_utils TRT 11 cleanup#4407

Open
apbose wants to merge 1 commit into
mainfrom
abose/torchTRT_MD_tests_fix
Open

fix(distributed): subgroup NCCL regression + nccl_utils TRT 11 cleanup#4407
apbose wants to merge 1 commit into
mainfrom
abose/torchTRT_MD_tests_fix

Conversation

@apbose

@apbose apbose commented Jul 15, 2026

Copy link
Copy Markdown
Collaborator

test_native_nccl.py:

  • Fix test_distributed_mode_subgroup and test_distributed_mode_context_switch regressions introduced by PR Python runtime rework  #4222 (Python runtime rework): add dist.barrier(group=subgroup) before first TRT forward so PyTorch's lazy ncclComm_t is initialized and bind_nccl_comm() sees a non-zero getCommPtr().
  • Remove all use_python_runtime=True flags (deprecated no-op in 2.13).
  • Add _cpp_runtime_available() guard to TestPythonRuntimePickle so it skips in C++ builds where _TRTEngine's custom op is never registered.
  • Fix _multirank_scatter_correctness / _multirank_gather_correctness in run_multirank_tests(): wrap with lambdas that loop over root ranks to match their (root, rank, world_size, device) signature (broken since PR TRT 11 MD Ops #4321).
  • Remove _multirank_distributed_mode_subgroup_python and _multirank_distributed_mode_context_switch_python (cannot work when C++ extension is loaded); collapse _multirank_pg_migration to cpp-only.

_nccl_utils.py:

  • setup_nccl_for_torch_tensorrt(): drop ensure_nccl_symlink() call and pre-load libnccl.so.2 directly instead of via the unversioned symlink. TRT 11.0+ (ncclWrapper.cpp) tries dlopen("libnccl.so.2") first, matching PyTorch's soname exactly; glibc reuses the already-loaded handle so no symlink is needed. ensure_nccl_symlink() is kept for TRT < 11.0 users.
  • Update module docstring to reflect confirmed TRT 11.0+ behavior.

@meta-cla meta-cla Bot added the cla signed label Jul 15, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests component: api [Python] Issues re: Python API labels Jul 15, 2026
@github-actions github-actions Bot requested a review from lanluo-nvidia July 15, 2026 05:32

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/distributed/test_native_nccl.py	2026-07-15 05:32:11.252442+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/distributed/test_native_nccl.py	2026-07-15 05:32:39.841891+00:00
@@ -1365,10 +1365,11 @@

        This class is only reached when ENABLED_FEATURES.torch_tensorrt_runtime=False
        (C++ extension not installed), so setup_engine() naturally instantiates
        _TRTEngine without any patching needed.
        """
+
        class LinearModel(nn.Module):
            def forward(self, x: torch.Tensor) -> torch.Tensor:
                return x + 1

        model = LinearModel().cuda().eval()
@@ -2149,12 +2150,16 @@
    tests = [
        _multirank_all_reduce_correctness,
        _multirank_all_gather_correctness,
        _multirank_reduce_scatter_all_reduce_ops,
        _multirank_all_to_all_correctness,
-        lambda r, ws, dev: [_multirank_scatter_correctness(i, r, ws, dev) for i in range(ws)],
-        lambda r, ws, dev: [_multirank_gather_correctness(i, r, ws, dev) for i in range(ws)],
+        lambda r, ws, dev: [
+            _multirank_scatter_correctness(i, r, ws, dev) for i in range(ws)
+        ],
+        lambda r, ws, dev: [
+            _multirank_gather_correctness(i, r, ws, dev) for i in range(ws)
+        ],
        _multirank_distributed_mode_tp_model,
        _multirank_distributed_mode_subgroup,
        _multirank_cpp_runtime_bind_nccl,
        _multirank_distributed_mode_context_switch,
        _multirank_pg_migration,

test_native_nccl.py:
- Fix test_distributed_mode_subgroup and test_distributed_mode_context_switch
  regressions introduced by PR #4222 (Python runtime rework): add
  dist.barrier(group=subgroup) before first TRT forward so PyTorch's lazy
  ncclComm_t is initialized and bind_nccl_comm() sees a non-zero getCommPtr().
- Remove all use_python_runtime=True flags (deprecated no-op in 2.13).
- Add _cpp_runtime_available() guard to TestPythonRuntimePickle so it skips
  in C++ builds where _TRTEngine's custom op is never registered.
- Fix _multirank_scatter_correctness / _multirank_gather_correctness in
  run_multirank_tests(): wrap with lambdas that loop over root ranks to match
  their (root, rank, world_size, device) signature (broken since PR #4321).
- Remove _multirank_distributed_mode_subgroup_python and
  _multirank_distributed_mode_context_switch_python (cannot work when C++
  extension is loaded); collapse _multirank_pg_migration to cpp-only.

_nccl_utils.py:
- setup_nccl_for_torch_tensorrt(): drop ensure_nccl_symlink() call and
  pre-load libnccl.so.2 directly instead of via the unversioned symlink.
  TRT 11.0+ (ncclWrapper.cpp) tries dlopen("libnccl.so.2") first, matching
  PyTorch's soname exactly; glibc reuses the already-loaded handle so no
  symlink is needed. ensure_nccl_symlink() is kept for TRT < 11.0 users.
- Update module docstring to reflect confirmed TRT 11.0+ behavior.
@apbose apbose force-pushed the abose/torchTRT_MD_tests_fix branch from eaced4d to 5a37370 Compare July 15, 2026 05:40
lanluo-nvidia pushed a commit that referenced this pull request Jul 15, 2026
…n + nccl_utils TRT 11 cleanup (#4407) (#4408)

Co-authored-by: Apurba Bose <abose@umb-b200-138.cl1u1.colossus.nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant