diff --git a/.github/workflows/_test-linux.yml b/.github/workflows/_test-linux.yml index 9f4a2c56b9..063f50ac39 100644 --- a/.github/workflows/_test-linux.yml +++ b/.github/workflows/_test-linux.yml @@ -142,6 +142,8 @@ jobs: build-matrix: ${{ needs.filter-matrix.outputs.matrix }} pre-script: packaging/pre_build_script.sh use-rtx: ${{ inputs.use-rtx }} + # Per-suite runner from the manifest; "" falls back to matrix.validation_runner. + runner: ${{ matrix.runner }} fail-on-empty: true script: | set -euo pipefail diff --git a/core/runtime/TRTEngine.h b/core/runtime/TRTEngine.h index b6a3badbe0..099c28f5f8 100644 --- a/core/runtime/TRTEngine.h +++ b/core/runtime/TRTEngine.h @@ -22,14 +22,29 @@ #include "core/runtime/TensorRTBindingNames.h" #include "core/util/prelude.h" -// TensorRT 10.16+ has native NCCL collective support via IExecutionContext::setCommunicator() -#if NV_TENSORRT_MAJOR > 10 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 16) +// Native NCCL collective support is exposed via IExecutionContext::setCommunicator() +// together with IDistCollectiveLayer. Two independent release lines ship that API: +// +// * TensorRT-RTX 1.5+ -- NvInferVersion.h defines TRT_MAJOR_RTX/TRT_MINOR_RTX and +// then aliases NV_TENSORRT_MAJOR/MINOR to them, so NV_TENSORRT_MAJOR is 1 on RTX +// and the mainline ">= 10.16" comparison below can never match. Detect the RTX +// package first and version-check against its own numbering. +// * TensorRT 10.16+ -- mainline. +// +// Do not collapse these into a single NV_TENSORRT_MAJOR/MINOR test: the two lines use +// incompatible numbering schemes. See is_tensorrt_version_supported() in +// py/torch_tensorrt/_utils.py for the Python-side equivalent of the same problem. +#if defined(TRT_MAJOR_RTX) +#if TRT_MAJOR_RTX > 1 || (TRT_MAJOR_RTX == 1 && TRT_MINOR_RTX >= 5) +#define TRT_HAS_NATIVE_NCCL 1 +#endif +#elif NV_TENSORRT_MAJOR > 10 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 16) #define TRT_HAS_NATIVE_NCCL 1 #endif // Full TRT NCCL collectives support requires both: // 1. PyTorch built with NCCL (USE_C10D_NCCL defined via Bazel) -// 2. TensorRT 10.16+ (TRT_HAS_NATIVE_NCCL defined above) +// 2. A TensorRT exposing the native collectives API (TRT_HAS_NATIVE_NCCL above) #if defined(USE_C10D_NCCL) && defined(TRT_HAS_NATIVE_NCCL) #define ENABLE_TRT_NCCL_COLLECTIVES 1 #endif diff --git a/tests/ci/runner.py b/tests/ci/runner.py index 70974e87a9..793b152fcb 100644 --- a/tests/ci/runner.py +++ b/tests/ci/runner.py @@ -318,6 +318,10 @@ def matrix(**filters: str | None) -> list[dict[str, str]]: "variant": var, "tier": s.tier, "cwd": s.for_variant(var)["cwd"], + # "" means "no override" -- linux-test.yml falls back to + # matrix.validation_runner. Set on suites that need specific + # hardware (e.g. multi-GPU for distributed). + "runner": s.for_variant(var)["runner"] or "", } for s, var in select(**filters) ] diff --git a/tests/ci/suites.py b/tests/ci/suites.py index 800803b7e7..804537cc93 100644 --- a/tests/ci/suites.py +++ b/tests/ci/suites.py @@ -81,6 +81,7 @@ class Suite: setup: tuple[str, ...] = () # named pre-steps: hub|executorch|cuda-core|mpi follow: tuple[tuple[str, ...], ...] = () # extra argv to run AFTER pytest env: dict[str, str] = field(default_factory=dict) + runner: str | None = None # GHA runner label; None = matrix.validation_runner overrides: dict[str, dict[str, Any]] = field(default_factory=dict) # per-variant def for_variant(self, variant: Variant) -> dict[str, Any]: @@ -101,6 +102,7 @@ def for_variant(self, variant: Variant) -> dict[str, Any]: "setup", "follow", "env", + "runner", ) } base.update(self.overrides.get(variant, {})) @@ -310,10 +312,30 @@ def for_variant(self, variant: Variant) -> dict[str, Any]: jobs="auto", verbose=True, reruns=False, - variants=("standard",), + variants=("standard", "rtx"), platforms=("linux-x86_64",), setup=("mpi",), env={"USE_HOST_DEPS": "1", "CI_BUILD": "1", "USE_TRTLLM_PLUGINS": "1"}, + # The --multirank follow-ups need 2 GPUs, so this suite cannot run on + # the default single-GPU validation_runner. + runner="linux.g4dn.12xlarge.nvidia.gpu", + # TensorRT-RTX has no TensorRT-LLM plugin path, so multi-device runs + # entirely on the native TRT DistCollective API. Drop test_nccl_ops.py + # (every test in it is gated on ENABLED_FEATURES.trtllm_for_nccl and + # would no-op) and USE_TRTLLM_PLUGINS along with it. + overrides={ + "rtx": { + "paths": ( + "distributed/test_native_nccl.py", + "distributed/test_export_save_load.py", + ), + "env": {"USE_HOST_DEPS": "1", "CI_BUILD": "1"}, + # Multi-GPU box: the --multirank follow-ups need 2 devices. + # g5 is A10G (SM 8.6) rather than g4dn's T4 (SM 7.5), since the + # TensorRT docs describe DistCollective as needing Ampere+. + "runner": "linux.g5.12xlarge.nvidia.gpu", + } + }, follow=( ( "-m", diff --git a/tests/py/dynamo/distributed/test_native_nccl.py b/tests/py/dynamo/distributed/test_native_nccl.py index 4e879ee37e..92e3831fd7 100644 --- a/tests/py/dynamo/distributed/test_native_nccl.py +++ b/tests/py/dynamo/distributed/test_native_nccl.py @@ -92,6 +92,16 @@ def has_nccl_collectives() -> bool: return False +def is_tensorrt_rtx() -> bool: + """Check if this is a TensorRT-RTX build.""" + try: + from torch_tensorrt._features import ENABLED_FEATURES + + return bool(ENABLED_FEATURES.tensorrt_rtx) + except Exception: + return False + + def _cpp_runtime_available() -> bool: """Return True when the C++ Torch-TensorRT runtime extension is loaded. @@ -2188,6 +2198,125 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: print(f"[Rank {rank}] PASS _multirank_pg_migration", flush=True) +def _find_py_engine(mod: nn.Module) -> Any: + """Return the Python-runtime ``TRTEngine`` inside a compiled module, or None.""" + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule + from torch_tensorrt.dynamo.runtime._TRTEngine import TRTEngine + + if isinstance(mod, TorchTensorRTModule) and isinstance(mod.engine, TRTEngine): + return mod.engine + for child in mod.children() if isinstance(mod, nn.Module) else []: + found = _find_py_engine(child) + if found is not None: + return found + return None + + +def _multirank_rtx_cudagraph_keeps_nccl_comm( + rank: int, world_size: int, device: torch.device +) -> None: + """TRT-RTX: enabling cudagraphs must not cost an MD engine its communicator. + + On TensorRT-RTX, manual ``torch.cuda.CUDAGraph`` capture is unsafe (lazy + kernel specialization), so ``_TRTEngine._execute_standard`` switches the + engine to RTX-native CUDA graphs on the first execute when cudagraphs are + enabled. That switch runs ``update_runtime_settings()`` -> + ``invalidate_context()``. + + The NCCL communicator is bound to the *IExecutionContext* by + ``setup_nccl_comm()`` (``set_communicator()``). Invalidation drops that + context but leaves ``_nccl_comm`` set, and the re-bind guard in + ``execute()`` is ``if self._nccl_comm is None`` -- so the replacement + context never receives the communicator and cannot be repaired. + + Unlike standard TensorRT, where a context invalidation requires the user to + deliberately change runtime settings, on RTX this fires with no user action + beyond turning cudagraphs on. + + The ``_nccl_comm`` assertion is deliberate: it fails fast rather than + hanging inside a collective, which matters in CI. + """ + import torch_tensorrt + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor.parallel import ( + ColwiseParallel, + RowwiseParallel, + parallelize_module, + ) + from torch_tensorrt._features import ENABLED_FEATURES + from torch_tensorrt.distributed._distributed import distributed_context + from torch_tensorrt.distributed._nccl_utils import setup_nccl_for_torch_tensorrt + + if not ENABLED_FEATURES.tensorrt_rtx: + print( + f"[Rank {rank}] SKIP _multirank_rtx_cudagraph_keeps_nccl_comm " + "(not a TensorRT-RTX build)", + flush=True, + ) + return + + setup_nccl_for_torch_tensorrt() + device_mesh = init_device_mesh("cuda", (world_size,)) + + class TinyMLP(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = nn.Linear(16, 32, bias=False) + self.relu = nn.ReLU() + self.fc2 = nn.Linear(32, 16, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc2(self.relu(self.fc1(x))) + + torch.manual_seed(42) + model = TinyMLP().to(device) + parallelize_module( + model, device_mesh, {"fc1": ColwiseParallel(), "fc2": RowwiseParallel()} + ) + + torch.manual_seed(0) + inp = torch.randn(4, 16, device=device) + + # Python runtime: set_communicator() lives in _TRTEngine, not the C++ engine. + with distributed_context(dist.group.WORLD): + trt_model = torch.compile( + model, + backend="torch_tensorrt", + dynamic=False, + options={ + "min_block_size": 1, + "use_distributed_mode_trace": True, + "use_python_runtime": True, + }, + ) + with torch.no_grad(): + expected = trt_model(inp) + + engine = _find_py_engine(trt_model) + if engine is None: + raise AssertionError("Could not locate a Python TRTEngine in compiled model") + assert engine._has_nccl_ops, "Engine is not a multi-device engine" + assert engine._nccl_comm is not None, "Communicator was never bound" + + # Turning cudagraphs on is enough: the first execute auto-switches to + # RTX-native graphs and invalidates the context. + with torch_tensorrt.runtime.enable_cudagraphs(trt_model) as cg_model: + with torch.no_grad(): + out = cg_model(inp) + + assert engine._nccl_comm is not None, ( + "the RTX cudagraph switch invalidated the IExecutionContext and the " + "communicator was not re-bound to the replacement context" + ) + + _check_close(out, expected, f"output after RTX cudagraph switch rank={rank}") + + print( + f"[Rank {rank}] PASS _multirank_rtx_cudagraph_keeps_nccl_comm", + flush=True, + ) + + # ============================================================================ # Section 8 — Multi-rank pytest tests (MultiProcessTestCase, requires 2 GPUs) # ============================================================================ @@ -2319,6 +2448,15 @@ def test_pg_migration(self) -> None: device = self._init_dist() _multirank_pg_migration(self.rank, self.world_size, device) + @unittest.skipIf(not is_tensorrt_rtx(), "TensorRT-RTX only") + @unittest.skipIf(not has_nccl_collectives(), "No NCCL collective support available") + @requires_nccl() + @skip_if_lt_x_gpu(2) + def test_rtx_cudagraph_keeps_nccl_comm(self) -> None: + """RTX cudagraph auto-switch invalidates the context; the comm must survive.""" + device = self._init_dist() + _multirank_rtx_cudagraph_keeps_nccl_comm(self.rank, self.world_size, device) + # ============================================================================ # Section 9 — torchrun / mpirun entry point (legacy multi-rank runner) @@ -2346,6 +2484,8 @@ def run_multirank_tests() -> None: _multirank_cpp_runtime_bind_nccl, _multirank_distributed_mode_context_switch, _multirank_pg_migration, + # Self-skips on non-RTX builds. + _multirank_rtx_cudagraph_keeps_nccl_comm, ] failed = []