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..f9f8cdae9d 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,12 @@ 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. + 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 +1321,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 +1345,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."""