[cherry-pick][release/2.13] fix(distributed): subgroup NCCL regression + 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>

Author: 3 people

py/torch_tensorrt/distributed/_nccl_utils.py

@@ -7,25 +7,31 @@
 
 Background:
 -----------
-TensorRT's dlopen("libnccl.so") may load a different NCCL library than PyTorch,
-causing crashes when sharing NCCL communicators.
+TensorRT 11.0+ (ncclWrapper.cpp) searches for NCCL in this order:
 
-- PyTorch loads NCCL via RPATH baked at compile time (libnccl.so.2)
-- TensorRT lazy-loads NCCL via dlopen("libnccl.so") at runtime
+    for (char const* name : {"libnccl.so.2", "libnccl.so"})
+        dlopen(name, ...)
 
-The mismatch occurs because:
-1. pip's nvidia-nccl-cu* package only ships libnccl.so.2 (no libnccl.so symlink)
-2. TRT specifically looks for libnccl.so, misses pip's copy, falls back to system NCCL
+Because it tries "libnccl.so.2" first, it matches the exact soname of PyTorch's
+already-loaded NCCL library. glibc returns the same handle, so TRT and PyTorch
+share one NCCL instance — no symlink required.
+
+The only remaining requirement is that the pip nccl directory is in
+LD_LIBRARY_PATH so a cold dlopen (library not yet loaded) can find the file.
+Pre-loading libnccl.so.2 with RTLD_GLOBAL before TRT's first dlopen call
+guarantees soname reuse regardless of LD_LIBRARY_PATH timing.
+
+Note on older TRT (< 11.0):
+    Pre-11.0 runtimes only called dlopen("libnccl.so"), which does not match
+    PyTorch's libnccl.so.2 soname. A libnccl.so symlink was required.
+    ensure_nccl_symlink() is kept for that case but is no longer called by
+    setup_nccl_for_torch_tensorrt() on TRT 11.0+.
 
 Environments:
 -------------
 - NGC containers: No action needed (both use system NCCL)
-- pip install torch: Requires symlink + LD_LIBRARY_PATH setup
-
-Future:
--------
-TensorRT 11.0 will support TRT_NCCL_LIBRARY env var, eliminating the need for
-symlink workarounds.
+- pip install torch + TRT 11.0+: LD_LIBRARY_PATH + libnccl.so.2 pre-load only
+- pip install torch + TRT < 11.0: call ensure_nccl_symlink() manually first
 """
 
 import ctypes
@@ -120,34 +126,36 @@ def check_nccl_library_path() -> bool:
 
 def setup_nccl_for_torch_tensorrt() -> None:
     """
-    Setup NCCL library for TensorRT distributed inference.
+    Setup NCCL library for TensorRT 11.0+ distributed inference.
+
+    TRT 11.0+ tries dlopen("libnccl.so.2") before "libnccl.so", so it matches
+    PyTorch's already-loaded library by soname and reuses the same instance.
+    This function ensures that reuse works by:
 
-    This function:
-    1. Detects if nvidia.nccl pip package is installed
-    2. Creates libnccl.so symlink if needed
-    3. Pre-loads libnccl.so via ctypes (helps Python runtime path)
-    4. Updates LD_LIBRARY_PATH for dynamic loaders
+    1. Prepending the pip nccl directory to LD_LIBRARY_PATH so a cold dlopen
+       can find libnccl.so.2 if it isn't yet in-process.
+    2. Pre-loading libnccl.so.2 with RTLD_GLOBAL so TRT's subsequent
+       dlopen("libnccl.so.2") gets the same handle via soname reuse rather
+       than opening a second copy.
 
-    Note: TRT's internal loader (libLoader.cpp) reads LD_LIBRARY_PATH at
-    process launch time, not when updated via os.environ. For the C++ TRT
-    runtime path, LD_LIBRARY_PATH must be set before the process starts:
+    No symlink is created. For NGC containers (system NCCL) this is a no-op.
 
-        NCCL_LIB=$(python -c "from torch_tensorrt.distributed._nccl_utils import get_nccl_library_path; print(get_nccl_library_path())")
-        LD_LIBRARY_PATH="$NCCL_LIB:$LD_LIBRARY_PATH" python script.py
+    Note: for the LD_LIBRARY_PATH update to be guaranteed effective on the
+    very first TRT dlopen, set it before process start:
 
-    For NGC containers (system NCCL), this is a no-op.
+        NCCL_LIB=$(python -c "from torch_tensorrt.distributed._nccl_utils \\
+            import get_nccl_library_path; print(get_nccl_library_path())")
+        LD_LIBRARY_PATH="$NCCL_LIB:$LD_LIBRARY_PATH" torchrun ...
     """
     global _nccl_setup_checked
 
-    # Only check once per process
     if _nccl_setup_checked:
         return
     _nccl_setup_checked = True
 
     nccl_lib_dir = get_nccl_library_path()
 
     if nccl_lib_dir is None:
-        # NGC container or system NCCL - no action needed
         logger.debug(
             "nvidia.nccl package not found. "
             "Assuming system NCCL is used by both PyTorch and TensorRT."
@@ -156,12 +164,8 @@ def setup_nccl_for_torch_tensorrt() -> None:
 
     logger.debug(f"Found nvidia.nccl package at: {nccl_lib_dir}")
 
-    # Ensure symlink exists
-    symlink_ok = ensure_nccl_symlink(nccl_lib_dir)
-
-    # Ensure LD_LIBRARY_PATH includes the NCCL directory so TRT's dlopen("libnccl.so")
-    # finds the same library PyTorch already loaded.  dlopen() reads LD_LIBRARY_PATH
-    # dynamically, so updating os.environ here takes effect for subsequent loads.
+    # Prepend pip nccl dir to LD_LIBRARY_PATH so a cold dlopen("libnccl.so.2")
+    # finds the right file when the library isn't yet loaded in-process.
     ld_library_path = os.environ.get("LD_LIBRARY_PATH", "")
     if nccl_lib_dir not in ld_library_path:
         os.environ["LD_LIBRARY_PATH"] = (
@@ -171,18 +175,19 @@ def setup_nccl_for_torch_tensorrt() -> None:
     else:
         logger.debug(f"LD_LIBRARY_PATH already includes NCCL directory: {nccl_lib_dir}")
 
-    if symlink_ok:
-        # Pre-load libnccl.so into the process with RTLD_GLOBAL so that TRT's
-        # subsequent dlopen("libnccl.so") inside setCommunicator() finds the
-        # already-loaded library rather than searching LD_LIBRARY_PATH again.
-        nccl_so = os.path.join(nccl_lib_dir, "libnccl.so")
+    # Pre-load libnccl.so.2 with RTLD_GLOBAL so TRT's dlopen("libnccl.so.2")
+    # gets soname reuse (same handle) instead of loading a second instance.
+    nccl_so_2 = os.path.join(nccl_lib_dir, "libnccl.so.2")
+    if os.path.exists(nccl_so_2):
         try:
-            ctypes.CDLL(nccl_so, mode=ctypes.RTLD_GLOBAL)
-            logger.debug(f"Pre-loaded NCCL library: {nccl_so}")
+            ctypes.CDLL(nccl_so_2, mode=ctypes.RTLD_GLOBAL)
+            logger.debug(f"Pre-loaded NCCL library: {nccl_so_2}")
         except OSError as e:
-            logger.warning(f"Failed to pre-load NCCL library {nccl_so}: {e}")
+            logger.warning(f"Failed to pre-load NCCL library {nccl_so_2}: {e}")
+    else:
+        logger.warning(f"libnccl.so.2 not found at {nccl_so_2}")
 
-        logger.debug("NCCL library setup complete")
+    logger.debug("NCCL library setup complete")
 
 
 def initialize_nccl_comm(device: Optional[int] = None) -> None:

tests/py/dynamo/distributed/test_native_nccl.py

@@ -45,7 +45,6 @@
 import sys
 import threading
 import unittest
-from contextlib import nullcontext
 from typing import Any
 from unittest.mock import MagicMock, patch
 
@@ -93,6 +92,21 @@ def has_nccl_collectives() -> bool:
         return False
 
 
+def _cpp_runtime_available() -> bool:
+    """Return True when the C++ Torch-TensorRT runtime extension is loaded.
+
+    When True, _TRTEngine registers neither its opaque type nor the Python
+    tensorrt::execute_engine custom op (see _TRTEngine.py:1376), so
+    TestPythonRuntimePickle cannot run and must be skipped.
+    """
+    try:
+        from torch_tensorrt._features import ENABLED_FEATURES
+
+        return bool(ENABLED_FEATURES.torch_tensorrt_runtime)
+    except Exception:
+        return False
+
+
 # ---------------------------------------------------------------------------
 # Shared test fakes
 # ---------------------------------------------------------------------------
@@ -1204,7 +1218,6 @@ def _run(self, model: nn.Module, inputs: list[torch.Tensor]) -> None:
                 backend="torch_tensorrt",
                 dynamic=False,
                 options={
-                    "use_python_runtime": True,
                     "min_block_size": 1,
                     "use_distributed_mode_trace": True,
                 },
@@ -1323,6 +1336,10 @@ def test_group_name_survives_context_exit(self) -> None:
     not has_nccl_collectives(),
     "Skipped: No NCCL collective support (neither native TRT collectives nor TRT-LLM).",
 )
+@unittest.skipIf(
+    _cpp_runtime_available(),
+    "Skipped: Python runtime (_TRTEngine) op not registered when C++ extension is available.",
+)
 class TestPythonRuntimePickle(unittest.TestCase):
     """Verifies that _nccl_comm is stripped on pickle and reset on unpickle.
 
@@ -1344,8 +1361,12 @@ def tearDownClass(cls) -> None:
             dist.destroy_process_group()
 
     def _compile_small_model(self) -> Any:
-        """Return a compiled module backed by a Python ``TRTEngine``."""
-        import torch_tensorrt
+        """Return a compiled module backed by a Python ``TRTEngine``.
+
+        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:
@@ -1359,21 +1380,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                 backend="torch_tensorrt",
                 dynamic=False,
                 options={
-                    "use_python_runtime": True,
                     "min_block_size": 1,
                 },
             )
-            _ = trt_model(inp)  # trigger compilation
+            _ = trt_model(inp)  # trigger compilation → setup_engine() → _TRTEngine
         return trt_model
 
     @classmethod
     def _find_python_trt_engine(cls, obj: Any) -> Any:
         """Walk a compiled module and return the Python ``TRTEngine`` instance.
 
-        With the unified runtime, the wrapping ``TorchTensorRTModule`` is the
-        same regardless of backend; ``use_python_runtime=True`` causes its
-        ``.engine`` attribute to be a Python ``TRTEngine`` (vs the C++
-        ``torch.classes.tensorrt.Engine`` otherwise).
+        Only meaningful when C++ extension is absent (ENABLED_FEATURES.torch_tensorrt_runtime=False),
+        which is enforced by the class-level skipIf decorator.
         """
         from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import (
             TorchTensorRTModule,
@@ -1633,7 +1651,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             backend="torch_tensorrt",
             dynamic=False,
             options={
-                "use_python_runtime": True,
                 "min_block_size": 1,
                 "use_distributed_mode_trace": True,
             },
@@ -1736,7 +1753,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             backend="torch_tensorrt",
             dynamic=False,
             options={
-                "use_python_runtime": True,
                 "min_block_size": 1,
                 "use_distributed_mode_trace": True,
             },
@@ -1770,6 +1786,10 @@ def _multirank_distributed_mode_subgroup(
     # New subgroup with all ranks (different group object from WORLD)
     subgroup = dist.new_group(ranks=list(range(world_size)))
     sg_name = subgroup.group_name
+    # Force subgroup NCCL comm init — PyTorch creates ncclComm_t lazily on first
+    # collective; bind_nccl_comm() in execute_engine reads getCommPtr() which is 0
+    # until at least one collective has run on the group.
+    dist.barrier(group=subgroup)
 
     class AllReduceModel(nn.Module):
         def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -1791,7 +1811,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             backend="torch_tensorrt",
             dynamic=False,
             options={
-                "use_python_runtime": True,
                 "min_block_size": 1,
                 "use_distributed_mode_trace": True,
             },
@@ -1872,6 +1891,9 @@ def _multirank_distributed_mode_context_switch(
     sg2 = dist.new_group(ranks=list(range(world_size)))
     sg1_name = sg1.group_name
     sg2_name = sg2.group_name
+    # Force both subgroup NCCL comms to init before TRT tries to bind them.
+    dist.barrier(group=sg1)
+    dist.barrier(group=sg2)
 
     class AllReduceModel(nn.Module):
         def __init__(self, name: str) -> None:
@@ -1893,7 +1915,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                 backend="torch_tensorrt",
                 dynamic=False,
                 options={
-                    "use_python_runtime": True,
                     "min_block_size": 1,
                     "use_distributed_mode_trace": True,
                 },
@@ -1947,50 +1968,41 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     expected_sum = world_size * (world_size + 1) // 2
     expected = torch.full((1, 4), float(expected_sum), device=device)
 
-    for label, use_python_runtime in [("cpp", False), ("python", True)]:
-        # ---- Step 1: compile + run with default world group ----
-        model = AllReduceModel(world_name).to(device).eval()
+    # ---- Step 1: compile + run with default world group ----
+    model = AllReduceModel(world_name).to(device).eval()
 
-        with distributed_context(world_group):
-            trt_model = torch.compile(
-                model,
-                backend="torch_tensorrt",
-                dynamic=False,
-                options={
-                    "use_python_runtime": use_python_runtime,
-                    "min_block_size": 1,
-                    "use_distributed_mode_trace": True,
-                },
-            )
-            with torch.no_grad():
-                out_world = trt_model(inp)
+    with distributed_context(world_group):
+        trt_model = torch.compile(
+            model,
+            backend="torch_tensorrt",
+            dynamic=False,
+            options={
+                "min_block_size": 1,
+                "use_distributed_mode_trace": True,
+            },
+        )
+        with torch.no_grad():
+            out_world = trt_model(inp)
 
-        _check_close(out_world, expected, f"[{label}] world group rank={rank}")
+    _check_close(out_world, expected, f"world group rank={rank}")
 
-        # ---- Step 2: migrate to subgroup via distributed_context(subgroup, model) ----
-        # This calls set_distributed_mode() which resets nccl_initialized on the
-        # C++ engine, and keeps _state.pg = subgroup active for the Python runtime's
-        # lazy setup_nccl_comm() call.
-        with distributed_context(subgroup, trt_model) as migrated_model:
-            with torch.no_grad():
-                out_sub = migrated_model(inp)
+    # ---- Step 2: migrate to subgroup via distributed_context(subgroup, model) ----
+    # set_distributed_mode() resets nccl_initialized on the C++ engine so
+    # bind_nccl_comm() re-runs on the next forward with the new group name.
+    with distributed_context(subgroup, trt_model) as migrated_model:
+        with torch.no_grad():
+            out_sub = migrated_model(inp)
 
-        _check_close(out_sub, expected, f"[{label}] migrated to subgroup rank={rank}")
+    _check_close(out_sub, expected, f"migrated to subgroup rank={rank}")
 
-        # ---- Step 3: set_distributed_mode (persistent, outside context) ----
-        subgroup2 = dist.new_group(ranks=list(range(world_size)))
-        torch_tensorrt.distributed.set_distributed_mode(subgroup2, trt_model)
-        # _state.pg is NOT set here — Python runtime falls back to world group
-        # for lazy setup_nccl_comm; C++ runtime uses the pinned group name.
-        # For C++ runtime only (Python runtime needs _state.pg active):
-        if not use_python_runtime:
-            with torch.no_grad():
-                out_pin = trt_model(inp)
-            _check_close(
-                out_pin, expected, f"[{label}] set_distributed_mode rank={rank}"
-            )
+    # ---- Step 3: set_distributed_mode (persistent, outside context) ----
+    subgroup2 = dist.new_group(ranks=list(range(world_size)))
+    torch_tensorrt.distributed.set_distributed_mode(subgroup2, trt_model)
+    with torch.no_grad():
+        out_pin = trt_model(inp)
+    _check_close(out_pin, expected, f"set_distributed_mode rank={rank}")
 
-        print(f"[Rank {rank}] PASS _multirank_pg_migration [{label}]", flush=True)
+    print(f"[Rank {rank}] PASS _multirank_pg_migration", flush=True)
 
 
 # ============================================================================
@@ -2096,7 +2108,7 @@ def test_distributed_mode_tp_model(self) -> None:
     @requires_nccl()
     @skip_if_lt_x_gpu(2)
     def test_distributed_mode_subgroup(self) -> None:
-        """distributed_context() with a non-default TP subgroup routes NCCL correctly."""
+        """C++ runtime with a non-default TP subgroup routes NCCL correctly."""
         device = self._init_dist()
         _multirank_distributed_mode_subgroup(self.rank, self.world_size, device)
 
@@ -2112,7 +2124,7 @@ def test_cpp_runtime_bind_nccl(self) -> None:
     @requires_nccl()
     @skip_if_lt_x_gpu(2)
     def test_distributed_mode_context_switch(self) -> None:
-        """Switching distributed_context between two subgroups routes to correct communicator."""
+        """C++ runtime: switching distributed_context between two subgroups routes correctly."""
         device = self._init_dist()
         _multirank_distributed_mode_context_switch(self.rank, self.world_size, device)
 
@@ -2140,8 +2152,12 @@ def run_multirank_tests() -> None:
         _multirank_all_gather_correctness,
         _multirank_reduce_scatter_all_reduce_ops,
         _multirank_all_to_all_correctness,
-        _multirank_scatter_correctness,
-        _multirank_gather_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)
+        ],
         _multirank_distributed_mode_tp_model,
         _multirank_distributed_mode_subgroup,
         _multirank_cpp_runtime_bind_nccl,