From 8eb0b2ed71f3d70a9ae92390c52f42200b9c8151 Mon Sep 17 00:00:00 2001 From: apbose Date: Tue, 4 Aug 2026 16:48:35 -0700 Subject: [PATCH 1/2] feat: enable native multi-device TensorRT on TensorRT-RTX builds TensorRT-RTX 1.5 ships the full multi-device API (IDistCollectiveLayer, addDistCollective, setNbRanks, IExecutionContext::setCommunicator), but Torch-TensorRT compiled all of it out of RTX builds. core/runtime/TRTEngine.h gated TRT_HAS_NATIVE_NCCL on NV_TENSORRT_MAJOR/MINOR >= 10.16. TensorRT-RTX 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 that comparison never matched. ENABLE_TRT_NCCL_COLLECTIVES was therefore never defined, and the whole MD runtime -- bind_nccl_comm, set_group_name, release_nccl_comm, the lazy bind in execute_engine.cpp, and NATIVE_TRT_COLLECTIVES_AVAIL -- was absent from RTX builds. The two release lines use incompatible version schemes, so detect the RTX package first and version-check against its own numbering, mirroring is_tensorrt_version_supported() on the Python side. Nothing else needed changing: the Bazel NCCL detection probes PyTorch (RTX-agnostic on Linux), USE_C10D_NCCL and nccl_headers already reach the RTX configs, the tensorrt_rtx bindings expose add_dist_collective / CollectiveOperation / ReduceOperation / set_communicator / num_ranks, and MD is not behind a PreviewFeature on RTX so the existing hasattr guard in _TRTInterpreter correctly skips it. CI: the `distributed` suite was variants=("standard",) and so never ran against RTX. Add the rtx variant, overriding away test_nccl_ops.py and USE_TRTLLM_PLUGINS -- both are TensorRT-LLM-only and would no-op on RTX. Also add a per-suite `runner` field to the manifest. The suite's --multirank follow-ups need 2 GPUs, but #4397 dropped the explicit multi-GPU runner the old build-test-linux-x86_64.yml set for this job, so it had been falling back to the single-GPU validation_runner. Restore g4dn.12xlarge for standard; rtx uses g5.12xlarge (A10G/SM 8.6) since the TensorRT docs describe DistCollective as requiring Ampere or newer. Not yet validated on hardware; needs a 2-GPU RTX Linux box. --- .github/workflows/_test-linux.yml | 2 ++ core/runtime/TRTEngine.h | 21 ++++++++++++++++++--- tests/ci/runner.py | 4 ++++ tests/ci/suites.py | 24 +++++++++++++++++++++++- 4 files changed, 47 insertions(+), 4 deletions(-) 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", From 2ba781a6a1764d1f310b2ad3de92cfa9bf3979d6 Mon Sep 17 00:00:00 2001 From: apbose Date: Wed, 5 Aug 2026 20:02:15 -0700 Subject: [PATCH 2/2] test: use project-standard RTOL/ATOL in the single-rank NCCL tests --- tests/py/dynamo/distributed/test_native_nccl.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/tests/py/dynamo/distributed/test_native_nccl.py b/tests/py/dynamo/distributed/test_native_nccl.py index 4e879ee37e..2cf4d8da35 100644 --- a/tests/py/dynamo/distributed/test_native_nccl.py +++ b/tests/py/dynamo/distributed/test_native_nccl.py @@ -1236,6 +1236,7 @@ def tearDownClass(cls) -> None: def _run(self, model: nn.Module, inputs: list[torch.Tensor]) -> None: """Compile with torch_tensorrt and verify output matches PyTorch.""" import torch_tensorrt + from torch_tensorrt.dynamo.utils import ATOL, RTOL model = model.cuda().eval() inputs_cuda = [t.cuda() for t in inputs] @@ -1254,7 +1255,15 @@ def _run(self, model: nn.Module, inputs: list[torch.Tensor]) -> None: ) out = trt_model(*inputs_cuda) - torch.testing.assert_close(ref, out, atol=1e-4, rtol=1e-4) + # Project-standard tolerance (torch_tensorrt.dynamo.utils), the same one + # every converter test uses via tests/py/dynamo/conversion/harness.py. + # TensorRT selects tensor-core kernels for batched matmuls, giving ~2^-11 + # (~4.9e-4) rounding versus PyTorch's FP32 reference -- well inside 5e-3 + # but over the 1e-4 this file previously hard-coded. Note disable_tf32 + # cannot be used to avoid it on TensorRT-RTX: it clears a builder flag, + # and RTX networks are strongly typed (see _TRTInterpreter), where + # builder precision flags do not apply. + torch.testing.assert_close(ref, out, atol=ATOL, rtol=RTOL) def test_all_reduce_single_rank(self) -> None: """all_reduce compiles and produces correct output on a single rank.""" @@ -1315,6 +1324,7 @@ def _run_dynamic( ) -> None: """Mark a dim dynamic (min/max), compile at the opt shape, verify at other shapes.""" import torch_tensorrt # noqa: F401 + from torch_tensorrt.dynamo.utils import ATOL, RTOL model = model.cuda().eval() opt_cuda = [t.cuda() for t in opt_inputs] @@ -1338,7 +1348,10 @@ def _run_dynamic( check_cuda = [t.cuda() for t in check] ref = model(*check_cuda) out = trt_model(*check_cuda) - torch.testing.assert_close(ref, out, atol=1e-4, rtol=1e-4) + # See _run: project-standard tolerance. These checks run at + # batch > 1, where TensorRT picks a tensor-core kernel, so the + # error is ~4.9e-4 rather than the ~1e-6 seen at batch 1. + torch.testing.assert_close(ref, out, atol=ATOL, rtol=RTOL) def test_all_reduce_single_rank_dynamic(self) -> None: """all_reduce compiles with a dynamic seq dim and is correct at other shapes."""