From 2c6b32d240a16266f90fd45e9acd2e2df653bd36 Mon Sep 17 00:00:00 2001 From: Peter Kisfaludi Date: Fri, 28 Aug 2026 23:27:59 -0700 Subject: [PATCH 1/6] dynamo: route native NCCL collectives at the op's process (sub)group The native collective converters built every collective as if it spanned the whole world: the rank array handed to addDistCollective was always arange(world_size) and num_ranks was world_size. On a mesh with more than one dimension that is wrong -- a tensor-parallel all_reduce would reduce across every rank, mixing data from separate CP groups, with no error raised. Resolve the op's group_name instead and pass that group's global ranks, with num_ranks set to the size of that group. Preserve the group's ordering. get_process_group_ranks() returns global ranks indexed by group rank, and that mapping is the collective's semantics: all_gather concatenates by group rank, reduce_scatter sends chunk i to group rank i. Sorting silently renumbers any group not created in ascending order. Verified on 8x A100 driving TensorRT's native collectives directly: - all_reduce over two 4-rank groups, rank r holding r+1, correct result 10/10/10/10/26/26/26/26. Passing the group's ranks gives that; arange(world) reduces across everything. - all_gather over a group whose rank array is reversed returns 8,7,6,5,4,3,2,1; sorting the same array returns 1,2,3,4,5,6,7,8. The rank array and the bound communicator must share an index space -- a mismatch returns zeros rather than raising -- so num_ranks now tracks len(groups). Tests are CPU-only (no GPU, NCCL or process group): TestCollectiveGroupRanks covers group resolution and ordering, TestNativeCollectiveNumRanks covers the rank count for all six converters. Both were checked against the pre-fix code to confirm they fail rather than pass vacuously. test_two_dimensional_mesh_routing needs four ranks -- one axis for CP, one for TP -- so it gets its own class. TestMultirankNccl fixes world_size at 2, where the CP and TP groups collapse onto the same pair and the test cannot distinguish correct routing from the world-group behaviour; registered there it would have skipped and reported success. The shared spawn/process-group setup moves to a base class that holds no tests of its own. Group seeding uses a real collective rather than dist.barrier(), which for NCCL infers the device when device_ids is omitted and can hang the group. --- .../conversion/custom_ops_converters.py | 20 +- .../dynamo/conversion/impl/nccl_ops.py | 80 +++- .../py/dynamo/distributed/test_native_nccl.py | 367 +++++++++++++++++- 3 files changed, 445 insertions(+), 22 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py index 587e793deba..757f827d2ad 100644 --- a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py @@ -114,6 +114,7 @@ def fused_nccl_all_gather( SourceIR.ATEN, name, [args[0]], + group_name=args[2] if len(args) > 2 else None, ) @dynamo_tensorrt_converter( @@ -138,6 +139,7 @@ def fused_nccl_reduce_scatter( name, [args[0]], reduce_op=reduce_op, + group_name=args[3] if len(args) > 3 else None, ) @dynamo_tensorrt_converter( @@ -161,6 +163,7 @@ def fused_nccl_all_reduce( name, [args[0]], reduce_op=reduce_op, + group_name=args[2] if len(args) > 2 else None, ) @dynamo_tensorrt_converter( @@ -182,6 +185,7 @@ def fused_nccl_all_to_all( SourceIR.ATEN, name, [args[0]], + group_name=args[3] if len(args) > 3 else None, ) @dynamo_tensorrt_converter( @@ -199,7 +203,13 @@ def fused_nccl_scatter( """Scatter using native TensorRT DistCollective API.""" root = args[1] if len(args) > 1 else 0 return impl.nccl_ops.nccl_scatter_native( - ctx, target, SourceIR.ATEN, name, [args[0]], root=root + ctx, + target, + SourceIR.ATEN, + name, + [args[0]], + root=root, + group_name=args[2] if len(args) > 2 else None, ) @dynamo_tensorrt_converter( @@ -217,7 +227,13 @@ def fused_nccl_gather( """Gather using native TensorRT DistCollective API.""" root = args[1] if len(args) > 1 else 0 return impl.nccl_ops.nccl_gather_native( - ctx, target, SourceIR.ATEN, name, [args[0]], root=root + ctx, + target, + SourceIR.ATEN, + name, + [args[0]], + root=root, + group_name=args[2] if len(args) > 2 else None, ) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py b/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py index dfa8cb83e6d..3e1dd835f58 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py @@ -83,6 +83,38 @@ def _get_distributed_rank_and_world_size() -> Tuple[int, int]: return rank, world_size +def _collective_group_ranks(group_name: Optional[str], world_size: int) -> np.ndarray: + """Global ranks of the collective's process group. + + The native ``add_dist_collective`` layer needs the set of ranks that participate in + *this* collective. Resolving it from the op's ``group_name`` lets a collective target a + process **subgroup** (e.g. context/sequence-parallel over one subgroup while tensor-parallel + uses another -- a 2-D device mesh) instead of always the whole world. Falls back to the world + group when the group cannot be resolved (single-program / group not created in this process). + """ + if group_name: + try: + import torch.distributed as dist + from torch.distributed.distributed_c10d import _resolve_process_group + + # Preserve the group's own ordering: get_process_group_ranks() returns global + # ranks indexed by *group* rank, and that mapping is what layout-sensitive + # collectives are defined against. all_gather concatenates by group rank, + # reduce_scatter sends chunk i to group rank i, all_to_all permutes by it, and a + # scatter/gather root names a position in it. Sorting would silently renumber the + # group whenever it was not created in ascending order -- e.g. a device mesh + # yielding [5, 2] would be read as [2, 5], swapping which rank receives which + # slice. Sorting is not needed for agreement either: every member resolves the + # same group_name to the same group and therefore already sees the same order. + ranks = dist.get_process_group_ranks(_resolve_process_group(group_name)) + return np.array(ranks, dtype=np.int64) + except Exception as e: # noqa: BLE001 + logger.warning( + f"Could not resolve process group '{group_name}' ({e}); using world group" + ) + return np.arange(world_size, dtype=np.int64) + + def nccl_all_gather( ctx: ConversionContext, target: Union[Target, str], @@ -231,6 +263,7 @@ def nccl_all_gather_native( source_ir: Optional[SourceIR], name: str, plug_inputs: Tuple[Argument, ...], + group_name: Optional[str] = None, ) -> trt.ITensor: """ Implement all_gather using native TensorRT DistCollective API. @@ -266,7 +299,7 @@ def nccl_all_gather_native( import numpy as np # Create array of all participating rank IDs [0, 1, 2, ..., world_size-1] - groups = np.arange(world_size, dtype=np.int64) + groups = _collective_group_ranks(group_name, world_size) logger.debug( f"Creating ALL_GATHER layer: groups={groups.tolist()}, groupSize={world_size}" @@ -288,7 +321,10 @@ def nccl_all_gather_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + # num_ranks is the size of *this collective's* group, not the world size: TensorRT + # takes the participants from `groups` and their count from num_ranks, so a subgroup + # collective left at world_size describes a group it does not have. + layer.num_ranks = len(groups) return output @@ -305,6 +341,7 @@ def nccl_reduce_scatter_native( name: str, plug_inputs: Tuple[Argument, ...], reduce_op: str = "sum", + group_name: Optional[str] = None, ) -> trt.ITensor: """ Implement reduce_scatter using native TensorRT DistCollective API. @@ -353,7 +390,7 @@ def nccl_reduce_scatter_native( trt_reduce_op = reduce_op_map[reduce_op.lower()] try: - groups = np.arange(world_size, dtype=np.int64) + groups = _collective_group_ranks(group_name, world_size) layer = ctx.net.add_dist_collective( input_tensor, @@ -366,7 +403,10 @@ def nccl_reduce_scatter_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + # num_ranks is the size of *this collective's* group, not the world size: TensorRT + # takes the participants from `groups` and their count from num_ranks, so a subgroup + # collective left at world_size describes a group it does not have. + layer.num_ranks = len(groups) logger.debug( f"Successfully created native REDUCE_SCATTER layer: {name}, reduce_op={reduce_op}, groups={groups.tolist()}" ) @@ -386,6 +426,7 @@ def nccl_all_reduce_native( name: str, plug_inputs: Tuple[Argument, ...], reduce_op: str = "sum", + group_name: Optional[str] = None, ) -> trt.ITensor: """ Implement all_reduce using native TensorRT DistCollective API. @@ -438,7 +479,7 @@ def nccl_all_reduce_native( # Create array of all participating rank IDs [0, 1, ..., world_size-1] # Passing None for groups can be treated as a no-op by TRT; use an explicit # rank array (same as ALL_GATHER) to ensure the reduction is performed. - groups = np.arange(world_size, dtype=np.int64) + groups = _collective_group_ranks(group_name, world_size) layer = ctx.net.add_dist_collective( input_tensor, @@ -451,7 +492,10 @@ def nccl_all_reduce_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + # num_ranks is the size of *this collective's* group, not the world size: TensorRT + # takes the participants from `groups` and their count from num_ranks, so a subgroup + # collective left at world_size describes a group it does not have. + layer.num_ranks = len(groups) logger.debug( f"Successfully created native ALL_REDUCE layer: {name}, reduce_op={reduce_op}, groups={groups.tolist()}" ) @@ -470,6 +514,7 @@ def nccl_all_to_all_native( source_ir: Optional[SourceIR], name: str, plug_inputs: Tuple[Argument, ...], + group_name: Optional[str] = None, ) -> trt.ITensor: """ Implement all_to_all using native TensorRT DistCollective API. @@ -506,7 +551,7 @@ def nccl_all_to_all_native( import numpy as np # Create array of all participating rank IDs [0, 1, 2, ..., world_size-1] - groups = np.arange(world_size, dtype=np.int64) + groups = _collective_group_ranks(group_name, world_size) logger.debug( f"Creating ALL_TO_ALL layer: groups={groups.tolist()}, groupSize={world_size}" @@ -528,7 +573,10 @@ def nccl_all_to_all_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + # num_ranks is the size of *this collective's* group, not the world size: TensorRT + # takes the participants from `groups` and their count from num_ranks, so a subgroup + # collective left at world_size describes a group it does not have. + layer.num_ranks = len(groups) return output @@ -545,6 +593,7 @@ def nccl_scatter_native( name: str, plug_inputs: Tuple[Argument, ...], root: int = 0, + group_name: Optional[str] = None, ) -> trt.ITensor: """ Implement scatter using native TensorRT DistCollective API. @@ -580,7 +629,7 @@ def nccl_scatter_native( import numpy as np # Create array of all participating rank IDs [0, 1, 2, ..., world_size-1] - groups = np.arange(world_size, dtype=np.int64) + groups = _collective_group_ranks(group_name, world_size) logger.debug( f"Creating scatter layer: groups={groups.tolist()}, groupSize={world_size}" @@ -602,7 +651,10 @@ def nccl_scatter_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + # num_ranks is the size of *this collective's* group, not the world size: TensorRT + # takes the participants from `groups` and their count from num_ranks, so a subgroup + # collective left at world_size describes a group it does not have. + layer.num_ranks = len(groups) return output @@ -619,6 +671,7 @@ def nccl_gather_native( name: str, plug_inputs: Tuple[Argument, ...], root: int = 0, + group_name: Optional[str] = None, ) -> trt.ITensor: """ Implement gather using native TensorRT DistCollective API. @@ -654,7 +707,7 @@ def nccl_gather_native( import numpy as np # Create array of all participating rank IDs [0, 1, 2, ..., world_size-1] - groups = np.arange(world_size, dtype=np.int64) + groups = _collective_group_ranks(group_name, world_size) logger.debug( f"Creating gather layer: groups={groups.tolist()}, groupSize={world_size}" @@ -676,7 +729,10 @@ def nccl_gather_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + # num_ranks is the size of *this collective's* group, not the world size: TensorRT + # takes the participants from `groups` and their count from num_ranks, so a subgroup + # collective left at world_size describes a group it does not have. + layer.num_ranks = len(groups) return output diff --git a/tests/py/dynamo/distributed/test_native_nccl.py b/tests/py/dynamo/distributed/test_native_nccl.py index f6850143db4..44d30e4a12e 100644 --- a/tests/py/dynamo/distributed/test_native_nccl.py +++ b/tests/py/dynamo/distributed/test_native_nccl.py @@ -83,6 +83,16 @@ def is_trtllm_for_nccl() -> bool: return False +def has_native_collective_api() -> bool: + """Whether this TensorRT exposes the DistCollective API (TRT 10.16+).""" + try: + import tensorrt as trt + + return hasattr(trt, "CollectiveOperation") + except Exception: + return False + + def has_nccl_collectives() -> bool: """Check if any NCCL collective backend is available (native TRT or TRT-LLM).""" try: @@ -711,6 +721,133 @@ def __init__(self) -> None: # ============================================================================ +class TestCollectiveGroupRanks(unittest.TestCase): + """_collective_group_ranks must preserve group-rank order — no GPU / no dist required.""" + + def _resolve_with(self, ranks, group_name="pg_under_test"): + """Run _collective_group_ranks with the process-group lookup stubbed to *ranks*.""" + from torch_tensorrt.dynamo.conversion.impl import nccl_ops + + import torch.distributed as dist + from torch.distributed import distributed_c10d + + real_resolve = getattr(distributed_c10d, "_resolve_process_group", None) + real_get = dist.get_process_group_ranks + distributed_c10d._resolve_process_group = lambda name: _FakeGroup(name) + dist.get_process_group_ranks = lambda group: list(ranks) + try: + return nccl_ops._collective_group_ranks(group_name, world_size=8) + finally: + dist.get_process_group_ranks = real_get + if real_resolve is not None: + distributed_c10d._resolve_process_group = real_resolve + + def test_group_rank_order_is_preserved(self) -> None: + """A non-ascending group must not be renumbered. + + get_process_group_ranks() returns global ranks *ordered by group rank*, and that + mapping is what layout-sensitive collectives are defined against: all_gather + concatenates by group rank, reduce_scatter sends chunk i to group rank i, + all_to_all permutes by it, and a scatter/gather root names a position in it. + Sorting a group created as [5, 2] into [2, 5] swaps which rank receives which + slice, silently producing wrong results rather than an error. + """ + self.assertEqual(list(self._resolve_with([5, 2])), [5, 2]) + self.assertEqual(list(self._resolve_with([3, 1, 2, 0])), [3, 1, 2, 0]) + + def test_ascending_group_is_unchanged(self) -> None: + """The common ascending case is unaffected by preserving order.""" + self.assertEqual(list(self._resolve_with([0, 1, 2, 3])), [0, 1, 2, 3]) + + def test_unresolvable_group_falls_back_to_world(self) -> None: + """An unresolvable group name falls back to the world group rather than raising.""" + from torch_tensorrt.dynamo.conversion.impl import nccl_ops + + result = nccl_ops._collective_group_ranks("no_such_group", world_size=4) + self.assertEqual(list(result), [0, 1, 2, 3]) + + +class TestNativeCollectiveNumRanks(unittest.TestCase): + """``num_ranks`` must equal the number of ranks in the collective's rank array. + + TensorRT takes a collective's participants from the rank array passed to + ``add_dist_collective`` and their count from ``num_ranks``; the two describe the same + group and must agree. Leaving ``num_ranks`` at the world size while the array holds a + subgroup describes a group the collective does not have. + + On hardware the mismatch is not loud: an all_reduce over a 4-rank group whose rank array + did not line up with the bound communicator returned zeros rather than raising. + + This drives each converter against a stubbed network, so there is no GPU, no NCCL and no + process group -- the fake layer never reaches TensorRT. + """ + + class _RecordingLayer: + """Stands in for an ``IDistCollectiveLayer``, capturing what the converter sets.""" + + def __init__(self) -> None: + object.__setattr__(self, "num_ranks", None) + + def get_output(self, index: int) -> object: + return object() + + def _num_ranks_for(self, converter_name: str, group_size: int) -> object: + """Return the num_ranks a converter set, given a group of ``group_size`` ranks.""" + import numpy as np + from unittest import mock + + from torch_tensorrt import _features + from torch_tensorrt.dynamo.conversion.impl import nccl_ops + + layer = self._RecordingLayer() + ctx = mock.MagicMock() + ctx.net.add_dist_collective.return_value = layer + + converter = getattr(nccl_ops, converter_name) + # The converters are wrapped in @needs_native_collectives, which raises unless the + # feature is on. It reads the module global at call time, so replacing the namedtuple + # is enough -- this test is about the rank count, not about the runtime being built + # with NCCL. + enabled = _features.ENABLED_FEATURES._replace(native_trt_collectives=True) + with mock.patch.object( + _features, "ENABLED_FEATURES", enabled + ), mock.patch.object(nccl_ops, "set_layer_name"), mock.patch.object( + nccl_ops, + "_get_distributed_rank_and_world_size", + # world is deliberately larger than the group, so num_ranks=world_size + # is distinguishable from num_ranks=len(groups). + return_value=(0, group_size * 2), + ), mock.patch.object( + nccl_ops, + "_collective_group_ranks", + return_value=np.arange(group_size, dtype=np.int64), + ): + converter(ctx, "target", None, "collective", (mock.MagicMock(),)) + return layer.num_ranks + + @unittest.skipIf( + not has_native_collective_api(), + "TensorRT build has no CollectiveOperation (needs TRT 10.16+)", + ) + def test_num_ranks_is_group_size_not_world_size(self) -> None: + group_size = 4 + for converter_name in ( + "nccl_all_gather_native", + "nccl_reduce_scatter_native", + "nccl_all_reduce_native", + "nccl_all_to_all_native", + "nccl_scatter_native", + "nccl_gather_native", + ): + with self.subTest(converter=converter_name): + self.assertEqual( + self._num_ranks_for(converter_name, group_size), + group_size, + f"{converter_name} set num_ranks to the world size instead of the " + "size of the collective's own group", + ) + + class TestNcclUtils(unittest.TestCase): """Tests for _nccl_utils.py functions — no GPU / no dist required.""" @@ -1959,6 +2096,180 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: _check_close(pt_out, trt_out, f"TP MLP distributed_context rank={rank}") +def _multirank_compute_collective_single_engine( + rank: int, world_size: int, device: torch.device +) -> None: + """Compile compute and a native all-reduce into a single TRT engine. + + Pointwise and shuffle layers around the collective encourage Myelin to absorb it + into a ForeignNode. The whole module must still lower to one engine -- no graph + break, no second engine, no fallback to the PyTorch collective -- and produce the + same values as eager. + + This also covers the converter setting ``num_ranks`` before the output is + requested: without it the build fails in Myelin shape inference rather than + producing a wrong answer, so a green result here means that ordering held. + """ + import torch_tensorrt + from torch_tensorrt.distributed._distributed import distributed_context + from torch_tensorrt.distributed._nccl_utils import setup_nccl_for_torch_tensorrt + + setup_nccl_for_torch_tensorrt() + group = dist.group.WORLD + + hidden = 64 + batch = 2 + sequence = 8 + + class ComputeCollective(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc_in = nn.Linear(hidden, hidden) + self.fc_out = nn.Linear(hidden, hidden) + self.gain = nn.Parameter(torch.randn(hidden)) + self.bias = nn.Parameter(torch.randn(hidden)) + + def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + x = self.fc_in(x) + x = (x * self.gain + self.bias + residual).unsqueeze(0) + x = x.reshape(batch * sequence, hidden) + dist.all_reduce(x) + x = x.reshape(batch, sequence, hidden) + return self.fc_out(x * self.gain - self.bias + residual) + + torch.manual_seed(42) + model = ComputeCollective().to(device=device, dtype=torch.bfloat16).eval() + torch.manual_seed(1234 + rank) + inp = torch.randn(batch, sequence, hidden, device=device, dtype=torch.bfloat16) + residual = torch.randn_like(inp) + + with torch.no_grad(): + eager_out = model(inp, residual) + exported = torch.export.export(model, (inp, residual), strict=False) + with distributed_context(group): + trt_model = torch_tensorrt.dynamo.compile( + exported, + inputs=[inp, residual], + min_block_size=1, + use_distributed_mode_trace=True, + use_python_runtime=False, + ) + + # Compilation must not hide the collective in a PyTorch fallback. + trt_engines = [ + name for name, _ in trt_model.named_modules() if "_run_on_acc" in name + ] + assert ( + len(trt_engines) == 1 + ), f"expected one TRT engine, found {trt_engines}" + graph = str(trt_model.graph) if hasattr(trt_model, "graph") else "" + for token in ("all_reduce", "_c10d_functional", "wait_tensor"): + assert token not in graph, f"collective escaped TRT engine: {graph}" + + trt_out = trt_model(inp, residual) + + torch.testing.assert_close(trt_out, eager_out, atol=2e-2, rtol=2e-2) + print( + f"[Rank {rank}] PASS compute+collective single-engine regression", + flush=True, + ) + + +def _multirank_two_dimensional_mesh_routing( + rank: int, world_size: int, device: torch.device +) -> None: + """Each collective must reduce over its own subgroup, not the world group. + + Topology is the one TensorRT MR !49040 was validated on: four ranks arranged as a + 2-D mesh, with context parallelism over one axis and tensor parallelism over the + other. + + CP groups: [0, 2] and [1, 3] + TP groups: [0, 1] and [2, 3] + + This is what distinguishes correct subgroup routing from the old world-group + behaviour, and it is the only test that can: with rank r contributing r + 1, a + TP all-reduce owes {1+2=3, 3+4=7} and a CP all-reduce owes {1+3=4, 2+4=6}, while a + collective wrongly routed to the world group would return 10 everywhere. The + existing subgroup test cannot catch this -- it builds its group from *all* ranks, + so subgroup and world are numerically indistinguishable. + """ + if world_size != 4: + print( + f"[SKIP] _multirank_two_dimensional_mesh_routing requires world_size == 4, " + f"got {world_size}" + ) + return + import torch_tensorrt + from torch_tensorrt.distributed._distributed import distributed_context + from torch_tensorrt.distributed._nccl_utils import setup_nccl_for_torch_tensorrt + + setup_nccl_for_torch_tensorrt() + + # Every rank must create every group, in the same order, for the handles to match. + cp_groups = [dist.new_group(ranks=[0, 2]), dist.new_group(ranks=[1, 3])] + tp_groups = [dist.new_group(ranks=[0, 1]), dist.new_group(ranks=[2, 3])] + cp_group = cp_groups[rank % 2] + tp_group = tp_groups[rank // 2] + # PyTorch creates the ncclComm_t lazily; bind_nccl_comm() reads a null pointer until + # at least one collective has run on the group. A real collective is used rather than + # dist.barrier(), which for NCCL infers the device when device_ids is omitted and can + # pick the wrong one on a multi-GPU rank, hanging the group. + seed = torch.zeros(1, device=device) + for group in cp_groups + tp_groups: + if rank in dist.get_process_group_ranks(group): + dist.all_reduce(seed, group=group) + + class AllReduceOnGroup(nn.Module): + def __init__(self, group_name: str) -> None: + super().__init__() + self.group_name = group_name + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = torch.ops._c10d_functional.all_reduce.default( + x, "sum", self.group_name + ) + return torch.ops._c10d_functional.wait_tensor.default(out) + + inp = torch.full((1, 8), float(rank + 1), device=device) + world_total = float(sum(r + 1 for r in range(world_size))) + + for axis, group in (("TP", tp_group), ("CP", cp_group)): + members = dist.get_process_group_ranks(group) + expected_value = float(sum(r + 1 for r in members)) + model = AllReduceOnGroup(group.group_name).to(device).eval() + + with distributed_context(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 = trt_model(inp) + + expected = torch.full((1, 8), expected_value, device=device) + if not torch.allclose(out, expected): + got = out[0, 0].item() + # Name the specific regression when the value matches the world sum: that is + # what a collective routed to the world group instead of `members` returns. + hint = ( + f" -- that is the world sum, so the collective was routed to the world " + f"group instead of {members}" + if abs(got - world_total) < 1e-6 + else "" + ) + raise AssertionError( + f"rank {rank}: {axis} all-reduce over {members} returned {got}, " + f"expected {expected_value}{hint}" + ) + + print( + f"[rank {rank}] 2-D mesh routing OK (CP and TP each reduced over their own group)" + ) + + def _multirank_distributed_mode_subgroup( rank: int, world_size: int, device: torch.device ) -> None: @@ -2325,17 +2636,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # ============================================================================ -class TestMultirankNccl(MultiProcessTestCase): - """Multi-rank NCCL tests as pytest-compatible MultiProcessTestCase. +class MultirankNcclBase(MultiProcessTestCase): + """Shared harness for the multi-rank tests. Subclasses set ``world_size``. - Each test spawns 2 worker processes via torch.multiprocessing. Requires - exactly 2 CUDA GPUs. Run with: - - pytest distributed/test_native_nccl.py::TestMultirankNccl -v + Deliberately not named ``Test*`` so neither pytest nor unittest collects it on its + own -- it holds no tests, only the process spawn and the process-group setup. """ - world_size = 2 - def setUp(self) -> None: super().setUp() self._spawn_processes() @@ -2364,6 +2671,15 @@ def _init_dist(self) -> torch.device: dist.barrier() # seeds ncclComm_t before any TRT bind_nccl_comm() call return torch.device(f"cuda:{local}") + +class TestMultirankNccl(MultirankNcclBase): + """Multi-rank NCCL tests on 2 GPUs. + + pytest distributed/test_native_nccl.py::TestMultirankNccl -v + """ + + world_size = 2 + @requires_nccl() @skip_if_lt_x_gpu(2) def test_all_reduce_correctness(self) -> None: @@ -2419,6 +2735,14 @@ def test_distributed_mode_tp_model(self) -> None: device = self._init_dist() _multirank_distributed_mode_tp_model(self.rank, self.world_size, device) + @unittest.skipIf(not has_nccl_collectives(), "No NCCL collective support available") + @requires_nccl() + @skip_if_lt_x_gpu(2) + def test_compute_collective_single_engine(self) -> None: + """Compute and a native collective compile and run in one TRT engine.""" + device = self._init_dist() + _multirank_compute_collective_single_engine(self.rank, self.world_size, device) + @unittest.skipIf(not has_nccl_collectives(), "No NCCL collective support available") @requires_nccl() @skip_if_lt_x_gpu(2) @@ -2472,6 +2796,31 @@ def test_comm_survives_disable_profiling(self) -> None: ) +class TestMultirankNccl4GPU(MultirankNcclBase): + """The 2-D mesh test, which needs four ranks -- one axis for CP, one for TP. + + It lives in its own class because ``world_size`` is fixed per class and + ``TestMultirankNccl`` runs at 2. At 2 ranks the CP and TP groups collapse onto the + same pair, so the test could not tell correct subgroup routing from the old + world-group behaviour, and would report success without checking anything. + + pytest distributed/test_native_nccl.py::TestMultirankNccl4GPU -v + """ + + world_size = 4 + + @unittest.skipIf(not has_nccl_collectives(), "No NCCL collective support available") + @requires_nccl() + @skip_if_lt_x_gpu(4) + def test_two_dimensional_mesh_routing(self) -> None: + """CP and TP collectives each route to their own subgroup, not the world group.""" + self.assertEqual( + self.world_size, 4, "the 2-D mesh test is meaningless below 4 ranks" + ) + device = self._init_dist() + _multirank_two_dimensional_mesh_routing(self.rank, self.world_size, device) + + # ============================================================================ # Section 9 — torchrun / mpirun entry point (legacy multi-rank runner) # ============================================================================ @@ -2494,7 +2843,9 @@ def run_multirank_tests() -> None: _multirank_gather_correctness(i, r, ws, dev) for i in range(ws) ], _multirank_distributed_mode_tp_model, + _multirank_compute_collective_single_engine, _multirank_distributed_mode_subgroup, + _multirank_two_dimensional_mesh_routing, _multirank_cpp_runtime_bind_nccl, _multirank_distributed_mode_context_switch, _multirank_pg_migration, From a502786748c7f33ef231b640d91941d57e9e87a1 Mon Sep 17 00:00:00 2001 From: Peter Kisfaludi Date: Wed, 9 Sep 2026 10:26:15 -0700 Subject: [PATCH 2/6] =?UTF-8?q?dynamo:=20address=20review=20=E2=80=94=20fa?= =?UTF-8?q?il=20on=20unresolvable=20group,=20fix=20mesh=20test=20group=20h?= =?UTF-8?q?andling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three fixes from apbose's review: - _collective_group_ranks no longer falls back to the world group when the op named a group that cannot be resolved. Widening a subgroup collective to every rank is the silent wrong-results failure this resolution exists to prevent, so it now raises. A missing group_name still means the world group, which is correct by construction. - The 2-D mesh test read membership with get_process_group_ranks() on every group it created. new_group() returns GroupMember.NON_GROUP_MEMBER on ranks outside the group and get_process_group_ranks() raises on that sentinel, so ranks would have died in setup. Membership now comes from the rank lists the groups were built from. - The mesh test bound a subgroup. addDistCollective's rank array is documented as 'rank IDs in the communicator', selecting a subset of it, so the array and the bound communicator must share an index space; global ranks of a group like [2,3] are out of range in its own 2-rank communicator. It now binds the world group, which is also the only configuration that lets one engine carry collectives on two different subgroups. Adds that two-subgroup case: TP and CP all_reduces in a single model, scaled between them so the result (100) is distinguishable from two world-routed reduces (400). --- .../dynamo/conversion/impl/nccl_ops.py | 35 ++++-- .../py/dynamo/distributed/test_native_nccl.py | 113 +++++++++++++++--- 2 files changed, 120 insertions(+), 28 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py b/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py index 3e1dd835f58..d54bc8c1b43 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py @@ -84,13 +84,17 @@ def _get_distributed_rank_and_world_size() -> Tuple[int, int]: def _collective_group_ranks(group_name: Optional[str], world_size: int) -> np.ndarray: - """Global ranks of the collective's process group. + """Ranks that participate in this collective, as IDs in the bound communicator. - The native ``add_dist_collective`` layer needs the set of ranks that participate in - *this* collective. Resolving it from the op's ``group_name`` lets a collective target a - process **subgroup** (e.g. context/sequence-parallel over one subgroup while tensor-parallel - uses another -- a 2-D device mesh) instead of always the whole world. Falls back to the world - group when the group cannot be resolved (single-program / group not created in this process). + ``addDistCollective`` documents ``groups`` as "a flat array of rank IDs in the + communicator", selecting the subset that takes part and defining each one's group-local + rank by position. The runtime binds one communicator per engine, so these are that + communicator's rank IDs -- global ranks when the world communicator is bound, which is + what lets a single engine host collectives on several different subgroups (e.g. context + parallel on one mesh axis and tensor parallel on the other). + + Resolving the op's ``group_name`` is what lets a collective target a subgroup at all; + without it every collective is built over the whole world. """ if group_name: try: @@ -106,12 +110,19 @@ def _collective_group_ranks(group_name: Optional[str], world_size: int) -> np.nd # yielding [5, 2] would be read as [2, 5], swapping which rank receives which # slice. Sorting is not needed for agreement either: every member resolves the # same group_name to the same group and therefore already sees the same order. - ranks = dist.get_process_group_ranks(_resolve_process_group(group_name)) - return np.array(ranks, dtype=np.int64) - except Exception as e: # noqa: BLE001 - logger.warning( - f"Could not resolve process group '{group_name}' ({e}); using world group" - ) + pg = _resolve_process_group(group_name) + except Exception as e: + # Do not fall back to the world group here. The op named a group, so falling + # back would build a collective spanning every rank -- the exact silent + # wrong-results failure this function exists to prevent. Fail the build instead. + raise RuntimeError( + f"Collective names process group '{group_name}', which could not be " + f"resolved in this process ({e}). Refusing to fall back to the world " + f"group: that would reduce across every rank and silently return wrong " + f"results. Ensure the group is created before compiling." + ) from e + return np.array(dist.get_process_group_ranks(pg), dtype=np.int64) + # No group named: the collective is over the world group by construction. return np.arange(world_size, dtype=np.int64) diff --git a/tests/py/dynamo/distributed/test_native_nccl.py b/tests/py/dynamo/distributed/test_native_nccl.py index 44d30e4a12e..8ef687e00d4 100644 --- a/tests/py/dynamo/distributed/test_native_nccl.py +++ b/tests/py/dynamo/distributed/test_native_nccl.py @@ -759,12 +759,25 @@ def test_ascending_group_is_unchanged(self) -> None: """The common ascending case is unaffected by preserving order.""" self.assertEqual(list(self._resolve_with([0, 1, 2, 3])), [0, 1, 2, 3]) - def test_unresolvable_group_falls_back_to_world(self) -> None: - """An unresolvable group name falls back to the world group rather than raising.""" + def test_unresolvable_group_raises_rather_than_widening_to_world(self) -> None: + """A named-but-unresolvable group must fail the build, not silently go world-wide. + + Falling back would build a collective over every rank, which is the silent + wrong-results failure this resolution exists to prevent. + """ from torch_tensorrt.dynamo.conversion.impl import nccl_ops - result = nccl_ops._collective_group_ranks("no_such_group", world_size=4) - self.assertEqual(list(result), [0, 1, 2, 3]) + with self.assertRaises(RuntimeError) as cm: + nccl_ops._collective_group_ranks("no_such_group", world_size=4) + self.assertIn("no_such_group", str(cm.exception)) + + def test_no_group_name_uses_world(self) -> None: + """With no group named, the collective is over the world group by construction.""" + from torch_tensorrt.dynamo.conversion.impl import nccl_ops + + self.assertEqual( + list(nccl_ops._collective_group_ranks(None, world_size=4)), [0, 1, 2, 3] + ) class TestNativeCollectiveNumRanks(unittest.TestCase): @@ -2207,17 +2220,24 @@ def _multirank_two_dimensional_mesh_routing( setup_nccl_for_torch_tensorrt() # Every rank must create every group, in the same order, for the handles to match. - cp_groups = [dist.new_group(ranks=[0, 2]), dist.new_group(ranks=[1, 3])] - tp_groups = [dist.new_group(ranks=[0, 1]), dist.new_group(ranks=[2, 3])] - cp_group = cp_groups[rank % 2] - tp_group = tp_groups[rank // 2] - # PyTorch creates the ncclComm_t lazily; bind_nccl_comm() reads a null pointer until - # at least one collective has run on the group. A real collective is used rather than + cp_ranks = [[0, 2], [1, 3]] + tp_ranks = [[0, 1], [2, 3]] + cp_groups = [dist.new_group(ranks=r) for r in cp_ranks] + tp_groups = [dist.new_group(ranks=r) for r in tp_ranks] + cp_members, cp_group = cp_ranks[rank % 2], cp_groups[rank % 2] + tp_members, tp_group = tp_ranks[rank // 2], tp_groups[rank // 2] + + # PyTorch creates the ncclComm_t lazily; bind_nccl_comm() reads a null pointer until at + # least one collective has run on the group. A real collective is used rather than # dist.barrier(), which for NCCL infers the device when device_ids is omitted and can # pick the wrong one on a multi-GPU rank, hanging the group. + # + # Membership is read off the rank lists rather than queried: new_group() hands back + # GroupMember.NON_GROUP_MEMBER on ranks outside the group, and get_process_group_ranks() + # raises on that sentinel. seed = torch.zeros(1, device=device) - for group in cp_groups + tp_groups: - if rank in dist.get_process_group_ranks(group): + for members, group in zip(cp_ranks + tp_ranks, cp_groups + tp_groups): + if rank in members: dist.all_reduce(seed, group=group) class AllReduceOnGroup(nn.Module): @@ -2234,12 +2254,22 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: inp = torch.full((1, 8), float(rank + 1), device=device) world_total = float(sum(r + 1 for r in range(world_size))) - for axis, group in (("TP", tp_group), ("CP", cp_group)): - members = dist.get_process_group_ranks(group) + # Bind the WORLD communicator, not either subgroup. addDistCollective's rank array is + # "rank IDs in the communicator" and selects a subset of it, so the communicator has to + # be the one those IDs are numbered in. Binding a subgroup instead would make the + # global ranks of a group like [2, 3] out of range for its own 2-rank communicator. + # Binding the world group is also what lets one engine carry collectives on several + # different subgroups, since the runtime binds a single communicator per engine. + world_group = dist.group.WORLD + + for axis, group, members in ( + ("TP", tp_group, tp_members), + ("CP", cp_group, cp_members), + ): expected_value = float(sum(r + 1 for r in members)) model = AllReduceOnGroup(group.group_name).to(device).eval() - with distributed_context(group): + with distributed_context(world_group): trt_model = torch.compile( model, backend="torch_tensorrt", @@ -2265,8 +2295,59 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: f"expected {expected_value}{hint}" ) + # Both axes inside ONE model, i.e. one engine carrying collectives on two different + # subgroups. This is the case binding the world communicator is meant to support. + # + # A plain TP-then-CP sum would be indistinguishable from the bug: summing over both + # mesh axes covers all four ranks, which is exactly the world sum. Scaling between the + # two collectives breaks that symmetry -- correct routing gives (3 or 7) * 10 summed + # down the CP axis = 100 everywhere, while two world-routed all_reduces give 400. + class TwoAxisModel(nn.Module): + def __init__(self, tp_name: str, cp_name: str) -> None: + super().__init__() + self.tp_name = tp_name + self.cp_name = cp_name + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = torch.ops._c10d_functional.all_reduce.default(x, "sum", self.tp_name) + y = torch.ops._c10d_functional.wait_tensor.default(y) + y = y * 10.0 + z = torch.ops._c10d_functional.all_reduce.default(y, "sum", self.cp_name) + return torch.ops._c10d_functional.wait_tensor.default(z) + + two_axis = TwoAxisModel(tp_group.group_name, cp_group.group_name).to(device).eval() + with distributed_context(world_group): + trt_two_axis = torch.compile( + two_axis, + backend="torch_tensorrt", + dynamic=False, + options={"min_block_size": 1, "use_distributed_mode_trace": True}, + ) + with torch.no_grad(): + combined = trt_two_axis(inp) + + # Each CP peer contributes its own TP-group sum, scaled by 10. + def tp_sum_for(r: int) -> float: + return float(sum(m + 1 for m in tp_ranks[r // 2])) + + cp_expected = sum(tp_sum_for(peer) * 10.0 for peer in cp_members) + if not torch.allclose(combined, torch.full((1, 8), cp_expected, device=device)): + got = combined[0, 0].item() + both_world = world_total * 10.0 * world_size + hint = ( + " -- that is what two world-routed all_reduces return, so neither collective " + "was routed to its own subgroup" + if abs(got - both_world) < 1e-6 + else "" + ) + raise AssertionError( + f"rank {rank}: TP({tp_members}) then CP({cp_members}) in one model returned " + f"{got}, expected {cp_expected}{hint}" + ) + print( - f"[rank {rank}] 2-D mesh routing OK (CP and TP each reduced over their own group)" + f"[rank {rank}] 2-D mesh routing OK (CP and TP each reduced over their own group, " + f"separately and in a single engine)" ) From b32b1514866c5b09d77758db3065eb9fbe18e5b5 Mon Sep 17 00:00:00 2001 From: Peter Kisfaludi Date: Wed, 9 Sep 2026 14:07:55 -0700 Subject: [PATCH 3/6] test: gate the 2-D mesh test on TensorRT child-communicator support The mesh test needs one engine to carry collectives on two different groups. That is what addDistCollective's `groups` array is for: setCommunicator installs a single communicator that "must be uniform across all multi-device instances", and each layer names its group as a subset of it. So global rank IDs are the right thing to emit, and a CP x TP mesh works with one bound communicator. Measured on 8 GPUs, every rank binding the same 8-rank communicator: TRT 11.2.0.86-md-moe-ep subset [0,1,2,3] -> 10 to its members, [4,5,6,7] -> 26 to its members CP4 x TP2 in ONE engine -> 360 on every rank (TP sum, scaled, then CP sum) the 4-GPU mesh test passes, including the two-axis single-engine case TRT 11.2.1.2 a subset containing rank 0 is silently ignored and reduces over the whole communicator (36 instead of 10) a subset without rank 0 fails: "Did not properly set the child communicator" So the capability exists and the design is right, but it is broken on the TensorRT the CI pins. Detecting it needs a real multi-rank run, so the test opts in via TORCHTRT_TEST_COLLECTIVE_SUBGROUPS=1 rather than failing CI over a TensorRT limitation. The helper docstring carries the measurements so the reason is not folklore. --- .../py/dynamo/distributed/test_native_nccl.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/py/dynamo/distributed/test_native_nccl.py b/tests/py/dynamo/distributed/test_native_nccl.py index 8ef687e00d4..79a47a3e300 100644 --- a/tests/py/dynamo/distributed/test_native_nccl.py +++ b/tests/py/dynamo/distributed/test_native_nccl.py @@ -93,6 +93,26 @@ def has_native_collective_api() -> bool: return False +def trt_supports_collective_subgroups() -> bool: + """Whether this TensorRT can build a child communicator for a subset of the bound one. + + ``addDistCollective``'s ``groups`` array selects a subset of the communicator installed by + ``setCommunicator``, which is what lets one engine carry collectives on several different + groups -- a CP x TP mesh, for instance. Measured on 8 GPUs, every rank binding the same + 8-rank communicator and passing the same rank array: + + TRT 11.2.0.86-md-moe-ep subsets work: ``[0,1,2,3]`` -> 10, ``[4,5,6,7]`` -> 26, and + CP4 x TP2 in a single engine -> 360 on every rank. + TRT 11.2.1.2 a subset containing rank 0 is silently ignored and reduces over + the whole communicator (36 instead of 10); a subset without + rank 0 fails with "Did not properly set the child communicator". + + Detecting this needs a real multi-rank run, so it cannot be probed cheaply from a skipIf. + Opt in with ``TORCHTRT_TEST_COLLECTIVE_SUBGROUPS=1`` on a TensorRT that supports it. + """ + return os.environ.get("TORCHTRT_TEST_COLLECTIVE_SUBGROUPS") == "1" + + def has_nccl_collectives() -> bool: """Check if any NCCL collective backend is available (native TRT or TRT-LLM).""" try: @@ -2891,6 +2911,11 @@ class TestMultirankNccl4GPU(MultirankNcclBase): world_size = 4 @unittest.skipIf(not has_nccl_collectives(), "No NCCL collective support available") + @unittest.skipUnless( + trt_supports_collective_subgroups(), + "TensorRT cannot build child communicators for a subset of the bound communicator " + "(broken on 11.2.1.2); set TORCHTRT_TEST_COLLECTIVE_SUBGROUPS=1 on a build that can", + ) @requires_nccl() @skip_if_lt_x_gpu(4) def test_two_dimensional_mesh_routing(self) -> None: From dde0553a50665cc35974fcf5d6e32b149564d675 Mon Sep 17 00:00:00 2001 From: Peter Kisfaludi Date: Wed, 9 Sep 2026 14:29:15 -0700 Subject: [PATCH 4/6] test: name the TensorRT fix the mesh test depends on The gate said "broken on 11.2.1.2", which is true but not actionable. The routing was fixed by TensorRT MR !49040 ("Fix DistCollective subgroup communicator routing", 2026-08-20), which derives the ncclCommSplit colour/key from the parent communicator's rank instead of assuming rank 0. Present on TensorRT main and rel-11.4; absent from the 11.2.1.x releases and rel-11.3. So the requirement is TensorRT >= 11.4, not a mystery version. --- .../py/dynamo/distributed/test_native_nccl.py | 35 +++++++++++-------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/tests/py/dynamo/distributed/test_native_nccl.py b/tests/py/dynamo/distributed/test_native_nccl.py index 79a47a3e300..cea4d385f96 100644 --- a/tests/py/dynamo/distributed/test_native_nccl.py +++ b/tests/py/dynamo/distributed/test_native_nccl.py @@ -94,21 +94,28 @@ def has_native_collective_api() -> bool: def trt_supports_collective_subgroups() -> bool: - """Whether this TensorRT can build a child communicator for a subset of the bound one. + """Whether this TensorRT routes a collective to a subset of the bound communicator. - ``addDistCollective``'s ``groups`` array selects a subset of the communicator installed by - ``setCommunicator``, which is what lets one engine carry collectives on several different - groups -- a CP x TP mesh, for instance. Measured on 8 GPUs, every rank binding the same - 8-rank communicator and passing the same rank array: + ``setCommunicator`` installs one communicator that "must be uniform across all + multi-device instances", and each ``addDistCollective`` layer names its group as a subset + of it. TensorRT splits a child communicator per group, which is what lets one engine carry + collectives on several groups -- a CP x TP mesh. - TRT 11.2.0.86-md-moe-ep subsets work: ``[0,1,2,3]`` -> 10, ``[4,5,6,7]`` -> 26, and - CP4 x TP2 in a single engine -> 360 on every rank. - TRT 11.2.1.2 a subset containing rank 0 is silently ignored and reduces over - the whole communicator (36 instead of 10); a subset without - rank 0 fails with "Did not properly set the child communicator". + That routing was fixed by TensorRT **MR !49040** ("Fix DistCollective subgroup communicator + routing", 2026-08-20), which derives the split colour/key from the parent communicator's + rank instead of assuming rank 0. It is present on TensorRT ``main`` and ``rel-11.4``, and + absent from the 11.2.1.x releases and ``rel-11.3``. - Detecting this needs a real multi-rank run, so it cannot be probed cheaply from a skipIf. - Opt in with ``TORCHTRT_TEST_COLLECTIVE_SUBGROUPS=1`` on a TensorRT that supports it. + Measured on 8 GPUs, every rank binding the same 8-rank communicator: + + with the fix subset ``[0,1,2,3]`` -> 10 to its members, ``[4,5,6,7]`` -> 26 to its + members, and CP4 x TP2 in a single engine -> 360 on every rank. + without it a subset containing rank 0 is silently ignored and reduces over the + whole communicator (36 instead of 10); a subset without rank 0 fails + with "Did not properly set the child communicator". + + Detecting this needs a real multi-rank run, so it cannot be probed from a skipIf. Opt in + with ``TORCHTRT_TEST_COLLECTIVE_SUBGROUPS=1`` once building against TensorRT >= 11.4. """ return os.environ.get("TORCHTRT_TEST_COLLECTIVE_SUBGROUPS") == "1" @@ -2913,8 +2920,8 @@ class TestMultirankNccl4GPU(MultirankNcclBase): @unittest.skipIf(not has_nccl_collectives(), "No NCCL collective support available") @unittest.skipUnless( trt_supports_collective_subgroups(), - "TensorRT cannot build child communicators for a subset of the bound communicator " - "(broken on 11.2.1.2); set TORCHTRT_TEST_COLLECTIVE_SUBGROUPS=1 on a build that can", + "needs TensorRT subgroup routing (TRT MR !49040, in main and rel-11.4; absent from " + "11.2.1.x) — set TORCHTRT_TEST_COLLECTIVE_SUBGROUPS=1 when building against it", ) @requires_nccl() @skip_if_lt_x_gpu(4) From eac638ea607fa16ad814c9c45527faee399dbedf Mon Sep 17 00:00:00 2001 From: Peter Kisfaludi Date: Tue, 15 Sep 2026 14:37:48 -0700 Subject: [PATCH 5/6] test: build the single-engine collective test in fp32, not bf16 The distributed CI lane runs on linux.g4dn (T4, Turing), which has no BF16 support, so the bf16 network failed to build at all: IBuilder::buildEngineWithConfig: Error Code 9: API Usage Error (Networks with BF16 precision require hardware with BF16 support) FAILED test_compute_collective_single_engine - assert cuda_engine Nothing in the test is precision-specific: it checks that compute and a native collective lower to a single engine with no PyTorch fallback, and that the numbers match eager. fp32 tests that just as well and actually runs on the CI hardware, so the tolerance tightens from 2e-2 to 1e-3. Also drops a stale claim from the docstring. It said a green result proved num_ranks was set before the output was requested; that ordering was measured to be a no-op (engines are byte-identical either way) and the reordering is no longer part of this PR, so the test never covered it. --- tests/py/dynamo/distributed/test_native_nccl.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/py/dynamo/distributed/test_native_nccl.py b/tests/py/dynamo/distributed/test_native_nccl.py index cea4d385f96..4692d56adbc 100644 --- a/tests/py/dynamo/distributed/test_native_nccl.py +++ b/tests/py/dynamo/distributed/test_native_nccl.py @@ -2146,9 +2146,10 @@ def _multirank_compute_collective_single_engine( break, no second engine, no fallback to the PyTorch collective -- and produce the same values as eager. - This also covers the converter setting ``num_ranks`` before the output is - requested: without it the build fails in Myelin shape inference rather than - producing a wrong answer, so a green result here means that ordering held. + Built in fp32 rather than bf16: the distributed CI lane runs on T4 (``linux.g4dn``), + where a bf16 network fails to build outright with "Networks with BF16 precision + require hardware with BF16 support". Nothing here is precision-specific -- the point + is that the collective stays inside the single engine and the numbers match eager. """ import torch_tensorrt from torch_tensorrt.distributed._distributed import distributed_context @@ -2178,9 +2179,9 @@ def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: return self.fc_out(x * self.gain - self.bias + residual) torch.manual_seed(42) - model = ComputeCollective().to(device=device, dtype=torch.bfloat16).eval() + model = ComputeCollective().to(device=device, dtype=torch.float32).eval() torch.manual_seed(1234 + rank) - inp = torch.randn(batch, sequence, hidden, device=device, dtype=torch.bfloat16) + inp = torch.randn(batch, sequence, hidden, device=device, dtype=torch.float32) residual = torch.randn_like(inp) with torch.no_grad(): @@ -2208,7 +2209,7 @@ def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: trt_out = trt_model(inp, residual) - torch.testing.assert_close(trt_out, eager_out, atol=2e-2, rtol=2e-2) + torch.testing.assert_close(trt_out, eager_out, atol=1e-3, rtol=1e-3) print( f"[Rank {rank}] PASS compute+collective single-engine regression", flush=True, From df7925da4354e7666732c180d08f5305d36bdf04 Mon Sep 17 00:00:00 2001 From: apbose Date: Wed, 30 Sep 2026 22:43:20 -0700 Subject: [PATCH 6/6] fix: update NCCL changes for current lint --- .../dynamo/conversion/impl/nccl_ops.py | 20 ++++++------------- .../py/dynamo/distributed/test_native_nccl.py | 5 ++--- 2 files changed, 8 insertions(+), 17 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py b/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py index d54bc8c1b43..53a81f0b328 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py @@ -267,7 +267,7 @@ def nccl_reduce_scatter( return layer.get_output(0) -@needs_native_collectives +@needs_native_collectives # type: ignore[misc] def nccl_all_gather_native( ctx: ConversionContext, target: Union[Target, str], @@ -307,8 +307,6 @@ def nccl_all_gather_native( # Use native TensorRT DistCollective API for ALL_GATHER # For ALL_GATHER, the reduce operation and root rank parameters are ignored # The last parameter (group) can be None to include all ranks - import numpy as np - # Create array of all participating rank IDs [0, 1, 2, ..., world_size-1] groups = _collective_group_ranks(group_name, world_size) @@ -344,7 +342,7 @@ def nccl_all_gather_native( raise -@needs_native_collectives +@needs_native_collectives # type: ignore[misc] def nccl_reduce_scatter_native( ctx: ConversionContext, target: Union[Target, str], @@ -429,7 +427,7 @@ def nccl_reduce_scatter_native( raise -@needs_native_collectives +@needs_native_collectives # type: ignore[misc] def nccl_all_reduce_native( ctx: ConversionContext, target: Union[Target, str], @@ -518,7 +516,7 @@ def nccl_all_reduce_native( raise -@needs_native_collectives +@needs_native_collectives # type: ignore[misc] def nccl_all_to_all_native( ctx: ConversionContext, target: Union[Target, str], @@ -559,8 +557,6 @@ def nccl_all_to_all_native( # Use native TensorRT DistCollective API for ALL_TO_ALL # For ALL_TO_ALL, the reduce operation and root rank parameters are ignored # The last parameter (group) can be None to include all ranks - import numpy as np - # Create array of all participating rank IDs [0, 1, 2, ..., world_size-1] groups = _collective_group_ranks(group_name, world_size) @@ -596,7 +592,7 @@ def nccl_all_to_all_native( raise -@needs_native_collectives +@needs_native_collectives # type: ignore[misc] def nccl_scatter_native( ctx: ConversionContext, target: Union[Target, str], @@ -637,8 +633,6 @@ def nccl_scatter_native( # Use native TensorRT DistCollective API for SCATTER # For SCATTER, the reduce operation parameter is ignored # The last parameter (group) can be None to include all ranks - import numpy as np - # Create array of all participating rank IDs [0, 1, 2, ..., world_size-1] groups = _collective_group_ranks(group_name, world_size) @@ -674,7 +668,7 @@ def nccl_scatter_native( raise -@needs_native_collectives +@needs_native_collectives # type: ignore[misc] def nccl_gather_native( ctx: ConversionContext, target: Union[Target, str], @@ -715,8 +709,6 @@ def nccl_gather_native( # Use native TensorRT DistCollective API for GATHER # For GATHER, the reduce operation parameter is ignored # The last parameter (group) can be None to include all ranks - import numpy as np - # Create array of all participating rank IDs [0, 1, 2, ..., world_size-1] groups = _collective_group_ranks(group_name, world_size) diff --git a/tests/py/dynamo/distributed/test_native_nccl.py b/tests/py/dynamo/distributed/test_native_nccl.py index 4692d56adbc..88b56a65875 100644 --- a/tests/py/dynamo/distributed/test_native_nccl.py +++ b/tests/py/dynamo/distributed/test_native_nccl.py @@ -753,10 +753,9 @@ class TestCollectiveGroupRanks(unittest.TestCase): def _resolve_with(self, ranks, group_name="pg_under_test"): """Run _collective_group_ranks with the process-group lookup stubbed to *ranks*.""" - from torch_tensorrt.dynamo.conversion.impl import nccl_ops - import torch.distributed as dist from torch.distributed import distributed_c10d + from torch_tensorrt.dynamo.conversion.impl import nccl_ops real_resolve = getattr(distributed_c10d, "_resolve_process_group", None) real_get = dist.get_process_group_ranks @@ -833,9 +832,9 @@ def get_output(self, index: int) -> object: def _num_ranks_for(self, converter_name: str, group_size: int) -> object: """Return the num_ranks a converter set, given a group of ``group_size`` ranks.""" - import numpy as np from unittest import mock + import numpy as np from torch_tensorrt import _features from torch_tensorrt.dynamo.conversion.impl import nccl_ops