Skip to content

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

Merged
apbose merged 2 commits into
mainfrom
abose/torchTRT_MD_tests_fix
Jul 20, 2026
Merged

fix(distributed): subgroup NCCL regression + nccl_utils TRT 11 cleanup#4407
apbose merged 2 commits 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,

@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>
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 5a37370 to b07e1a0 Compare July 20, 2026 16:56

@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/py/torch_tensorrt/dynamo/lowering/passes/eliminate_sym_min_int64_max.py	2026-07-20 16:58:14.770814+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/eliminate_sym_min_int64_max.py	2026-07-20 16:58:38.866999+00:00
@@ -2,11 +2,10 @@

import torch
from torch.fx import GraphModule, Node

from .pass_utils import clean_up_graph_after_modifications
-

_INT64_MAX = 2**63 - 1
_SYM_MIN = getattr(torch, "sym_min", None)

@apbose
apbose force-pushed the abose/torchTRT_MD_tests_fix branch from 0be982b to 7bba49f Compare July 20, 2026 17:23
@github-actions github-actions Bot added component: lowering Issues re: The lowering / preprocessing passes component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Jul 20, 2026
@apbose
apbose merged commit 50cced9 into main Jul 20, 2026
135 of 148 checks passed
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: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant