From 1d341430a2db4ce8df788bd1cae90d6c7b1cf003 Mon Sep 17 00:00:00 2001 From: Peter Kisfaludi Date: Sat, 4 Jul 2026 04:12:22 -0700 Subject: [PATCH 1/3] dynamo: target native NCCL collectives at the op's process (sub)group The native DistCollective converters hardcoded the world group (groups = np.arange(world_size), num_ranks = world_size), so a collective could only ever run over all ranks. This breaks 2-D device meshes where a collective runs over a subgroup -- e.g. Ulysses/context-parallel all_to_all over a CP subgroup while tensor-parallel all_reduce uses a separate TP subgroup. TensorRT then rejects it, e.g.: "All to All requires first input dimension to be divisible by nbRanks" because nbRanks was the world size (8) rather than the subgroup size (4). Resolve the participating ranks from the collective's group_name (already carried by the tensorrt::fused_nccl_* custom ops) and pass them to add_dist_collective, setting num_ranks = len(groups). Falls back to the world group when the group can't be resolved. Threads group_name from each fused-op converter (all_gather/reduce_scatter/all_reduce/all_to_all/scatter/gather). Validated on an 8xA100 CP4xTP2 run (LTX-2.3 DiT): supplying the CP subgroup's ranks to add_dist_collective clears the divisibility build error for the Ulysses all_to_all. Co-Authored-By: Claude Opus 4.8 --- .../conversion/custom_ops_converters.py | 10 +++- .../dynamo/conversion/impl/nccl_ops.py | 54 ++++++++++++++----- 2 files changed, 50 insertions(+), 14 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py index 7bd70705a0..bbde377356 100644 --- a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py @@ -44,6 +44,7 @@ def fused_nccl_all_gather( SourceIR.ATEN, name, [args[0]], + group_name=args[2] if len(args) > 2 else None, ) @dynamo_tensorrt_converter( @@ -66,6 +67,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( @@ -87,6 +89,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( @@ -106,6 +109,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( @@ -121,7 +125,8 @@ 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( @@ -137,7 +142,8 @@ 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 df579ee76e..bbd3bc8fcb 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py @@ -80,6 +80,30 @@ def _get_distributed_rank_and_world_size() -> Tuple[int, int]: return rank, world_size + +def _collective_group_ranks(group_name, world_size): + """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). + """ + import numpy as np + if group_name: + try: + import torch.distributed as dist + from torch.distributed.distributed_c10d import _resolve_process_group + + ranks = dist.get_process_group_ranks(_resolve_process_group(group_name)) + return np.array(sorted(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], @@ -228,6 +252,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. @@ -263,7 +288,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}" @@ -285,7 +310,7 @@ def nccl_all_gather_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + layer.num_ranks = len(groups) return output @@ -302,6 +327,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. @@ -350,7 +376,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, @@ -363,7 +389,7 @@ def nccl_reduce_scatter_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + layer.num_ranks = len(groups) logger.debug( f"Successfully created native REDUCE_SCATTER layer: {name}, reduce_op={reduce_op}, groups={groups.tolist()}" ) @@ -383,6 +409,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. @@ -435,7 +462,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, @@ -448,7 +475,7 @@ def nccl_all_reduce_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + layer.num_ranks = len(groups) logger.debug( f"Successfully created native ALL_REDUCE layer: {name}, reduce_op={reduce_op}, groups={groups.tolist()}" ) @@ -467,6 +494,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. @@ -503,7 +531,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}" @@ -525,7 +553,7 @@ def nccl_all_to_all_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + layer.num_ranks = len(groups) return output @@ -542,6 +570,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. @@ -577,7 +606,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}" @@ -599,7 +628,7 @@ def nccl_scatter_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + layer.num_ranks = len(groups) return output @@ -616,6 +645,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. @@ -651,7 +681,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}" @@ -673,7 +703,7 @@ def nccl_gather_native( set_layer_name(layer, target, name, source_ir) output = layer.get_output(0) - layer.num_ranks = world_size + layer.num_ranks = len(groups) return output From 8daa20af566414842e009bd04b674ad3fdd12cf8 Mon Sep 17 00:00:00 2001 From: Peter Kisfaludi Date: Tue, 7 Jul 2026 05:31:35 -0700 Subject: [PATCH 2/3] dynamo: keep native NCCL collectives out of Myelin ForeignNode fusion (cast-boundary) Adds, on top of #4380 (subgroup routing), the changes needed to lower the sharded LTX-2.3 DiT to a single TensorRT engine with native in-engine collectives: 1. _coll_boundary(): insert a genuine dtype reformat (fp32 by default, _COLL_BOUNDARY_DTYPE) on each DistCollective's input/output. Myelin folds identity/same-dtype casts away, so only a real reformat keeps the collective out of a ForeignNode -> the 48-block DiT compiles to one engine with collectives intact. See #4381. 2. Set layer.num_ranks BEFORE layer.get_output(0) so Myelin can infer the collective output shape (all_gather dim0 = in_dim0 * num_ranks); upstream set it after, tripping "AllGather requires world_size (>1) ... call myelinGraphSetWorldSize()". 3. AG_VIA_ALLREDUCE=1 (opt-in): express all_gather as place-into-slot + ALL_REDUCE(SUM) whose output shape == input shape, so no build-time world_size is needed. Numerically identical to all_gather. Also adds _mark_collective_boundary() (env TRT_COLL_BOUNDARY) as a debug-mark helper. Co-Authored-By: Claude Opus 4.8 --- .../dynamo/conversion/impl/nccl_ops.py | 129 ++++++++++++++++-- 1 file changed, 117 insertions(+), 12 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py b/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py index bbd3bc8fcb..06a9164a0a 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py @@ -104,6 +104,51 @@ def _collective_group_ranks(group_name, world_size): ) return np.arange(world_size, dtype=np.int64) + +def _mark_collective_boundary(ctx, input_tensor, output_tensor): + """Optionally mark the collective's tensors for debug so they are not fused away.""" + mode = os.environ.get("TRT_COLL_BOUNDARY", "").lower() + if mode in ("input", "both"): + ctx.net.mark_debug(input_tensor) + if mode in ("output", "both"): + ctx.net.mark_debug(output_tensor) + + +# Myelin fusion-boundary for native DistCollective. Myelin folds identity / same-dtype +# casts away, so only a genuine dtype *reformat* keeps a native collective out of a +# ForeignNode (see pytorch/TensorRT#4381). fp32 by default; set to None to disable, or +# trt.float16 to keep the communication volume down. +_COLL_BOUNDARY_DTYPE = trt.float32 + +_coll_boundary_orig_dtype: dict = {} + + +def _coll_boundary(ctx, t, name, where): + """Insert a dtype reformat around a DistCollective so Myelin will not fuse it. + + ``where='in'`` casts the collective input to the boundary dtype (the collective then + runs at that precision); ``where='out'`` casts the output back to the original dtype. + A no-op when the boundary dtype is disabled or already matches. + """ + if _COLL_BOUNDARY_DTYPE is None: + return t + if where == "in": + _coll_boundary_orig_dtype[name] = t.dtype + target = _COLL_BOUNDARY_DTYPE + else: + target = _coll_boundary_orig_dtype.get(name, t.dtype) + if t.dtype == target: + return t + c = ctx.net.add_cast(t, target) + c.name = f"{name}_myelin_bnd_{where}" + o = c.get_output(0) + try: + o.dtype = target + except Exception: # noqa: BLE001 + pass + return o + + def nccl_all_gather( ctx: ConversionContext, target: Union[Target, str], @@ -279,7 +324,7 @@ def nccl_all_gather_native( ) # Get the input tensor - input_tensor = plug_inputs[0] + input_tensor = _coll_boundary(ctx, plug_inputs[0], name, "in") try: # Use native TensorRT DistCollective API for ALL_GATHER @@ -290,6 +335,48 @@ def nccl_all_gather_native( # Create array of all participating rank IDs [0, 1, 2, ..., world_size-1] groups = _collective_group_ranks(group_name, world_size) + if os.environ.get("AG_VIA_ALLREDUCE") == "1": + # Express all_gather as place-into-slot + ALL_REDUCE(SUM). The ALL_REDUCE output + # shape == input shape (built explicitly via concat below), so Myelin needs no + # build-time world_size for shape inference (unlike a native ALL_GATHER, which + # requires myelinGraphSetWorldSize). Numerically identical to all_gather; cost is + # len(groups)x the communication volume. Opt-in via env AG_VIA_ALLREDUCE=1. + raw = plug_inputs[0] + ng = int(len(groups)) + glist = [int(x) for x in groups.tolist()] + slot = glist.index(rank) if rank in glist else 0 + ish = [int(x) for x in raw.shape] + d0 = ish[0] + + def _z(nrows): + sh = tuple([nrows] + ish[1:]) + z = ctx.net.add_constant( + sh, trt.Weights(np.zeros(int(np.prod(sh)), dtype=np.float32)) + ).get_output(0) + if z.dtype != raw.dtype: + z = ctx.net.add_cast(z, raw.dtype).get_output(0) + return z + + parts = [] + if slot > 0: + parts.append(_z(slot * d0)) + parts.append(raw) + if (ng - 1 - slot) > 0: + parts.append(_z((ng - 1 - slot) * d0)) + cat = ctx.net.add_concatenation(parts) + cat.axis = 0 + pfo = ctx.net.add_cast(cat.get_output(0), trt.float32).get_output(0) + ar = ctx.net.add_dist_collective( + pfo, + trt.CollectiveOperation.ALL_REDUCE, + trt.ReduceOperation.SUM, + -1, + groups, + ) + ar.num_ranks = ng + set_layer_name(ar, target, name, source_ir) + return ctx.net.add_cast(ar.get_output(0), raw.dtype).get_output(0) + logger.debug( f"Creating ALL_GATHER layer: groups={groups.tolist()}, groupSize={world_size}" ) @@ -309,8 +396,11 @@ def nccl_all_gather_native( set_layer_name(layer, target, name, source_ir) - output = layer.get_output(0) + # num_ranks must be set BEFORE get_output(0) so Myelin can infer the collective's + # output shape (e.g. all_gather dim0 = in_dim0 * num_ranks) for shape inference. layer.num_ranks = len(groups) + output = _coll_boundary(ctx, layer.get_output(0), name, "out") + _mark_collective_boundary(ctx, input_tensor, output) return output @@ -357,7 +447,7 @@ def nccl_reduce_scatter_native( return plug_inputs[0] # Get the input tensor - input_tensor = plug_inputs[0] + input_tensor = _coll_boundary(ctx, plug_inputs[0], name, "in") reduce_op_map = { "sum": trt.ReduceOperation.SUM, @@ -388,8 +478,11 @@ def nccl_reduce_scatter_native( set_layer_name(layer, target, name, source_ir) - output = layer.get_output(0) + # num_ranks must be set BEFORE get_output(0) so Myelin can infer the collective's + # output shape (e.g. all_gather dim0 = in_dim0 * num_ranks) for shape inference. layer.num_ranks = len(groups) + output = _coll_boundary(ctx, layer.get_output(0), name, "out") + _mark_collective_boundary(ctx, input_tensor, output) logger.debug( f"Successfully created native REDUCE_SCATTER layer: {name}, reduce_op={reduce_op}, groups={groups.tolist()}" ) @@ -440,7 +533,7 @@ def nccl_all_reduce_native( f"Adding native all_reduce: name={name}, rank={rank}, world_size={world_size}, reduce_op={reduce_op}" ) - input_tensor = plug_inputs[0] + input_tensor = _coll_boundary(ctx, plug_inputs[0], name, "in") reduce_op_map = { "sum": trt.ReduceOperation.SUM, @@ -474,8 +567,11 @@ def nccl_all_reduce_native( set_layer_name(layer, target, name, source_ir) - output = layer.get_output(0) + # num_ranks must be set BEFORE get_output(0) so Myelin can infer the collective's + # output shape (e.g. all_gather dim0 = in_dim0 * num_ranks) for shape inference. layer.num_ranks = len(groups) + output = _coll_boundary(ctx, layer.get_output(0), name, "out") + _mark_collective_boundary(ctx, input_tensor, output) logger.debug( f"Successfully created native ALL_REDUCE layer: {name}, reduce_op={reduce_op}, groups={groups.tolist()}" ) @@ -522,7 +618,7 @@ def nccl_all_to_all_native( ) # Get the input tensor - input_tensor = plug_inputs[0] + input_tensor = _coll_boundary(ctx, plug_inputs[0], name, "in") try: # Use native TensorRT DistCollective API for ALL_TO_ALL @@ -552,8 +648,11 @@ def nccl_all_to_all_native( set_layer_name(layer, target, name, source_ir) - output = layer.get_output(0) + # num_ranks must be set BEFORE get_output(0) so Myelin can infer the collective's + # output shape (e.g. all_gather dim0 = in_dim0 * num_ranks) for shape inference. layer.num_ranks = len(groups) + output = _coll_boundary(ctx, layer.get_output(0), name, "out") + _mark_collective_boundary(ctx, input_tensor, output) return output @@ -597,7 +696,7 @@ def nccl_scatter_native( ) # Get the input tensor - input_tensor = plug_inputs[0] + input_tensor = _coll_boundary(ctx, plug_inputs[0], name, "in") try: # Use native TensorRT DistCollective API for SCATTER @@ -627,8 +726,11 @@ def nccl_scatter_native( set_layer_name(layer, target, name, source_ir) - output = layer.get_output(0) + # num_ranks must be set BEFORE get_output(0) so Myelin can infer the collective's + # output shape (e.g. all_gather dim0 = in_dim0 * num_ranks) for shape inference. layer.num_ranks = len(groups) + output = _coll_boundary(ctx, layer.get_output(0), name, "out") + _mark_collective_boundary(ctx, input_tensor, output) return output @@ -672,7 +774,7 @@ def nccl_gather_native( ) # Get the input tensor - input_tensor = plug_inputs[0] + input_tensor = _coll_boundary(ctx, plug_inputs[0], name, "in") try: # Use native TensorRT DistCollective API for GATHER @@ -702,8 +804,11 @@ def nccl_gather_native( set_layer_name(layer, target, name, source_ir) - output = layer.get_output(0) + # num_ranks must be set BEFORE get_output(0) so Myelin can infer the collective's + # output shape (e.g. all_gather dim0 = in_dim0 * num_ranks) for shape inference. layer.num_ranks = len(groups) + output = _coll_boundary(ctx, layer.get_output(0), name, "out") + _mark_collective_boundary(ctx, input_tensor, output) return output From ed8d41a6d164cccace632fc5b4958cb41c72060e Mon Sep 17 00:00:00 2001 From: Peter Kisfaludi Date: Tue, 7 Jul 2026 16:45:02 -0700 Subject: [PATCH 3/3] dynamo: test compute and collective in one TRT engine --- .../py/dynamo/distributed/test_native_nccl.py | 90 +++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/tests/py/dynamo/distributed/test_native_nccl.py b/tests/py/dynamo/distributed/test_native_nccl.py index f0d6ed607b..ad565644e6 100644 --- a/tests/py/dynamo/distributed/test_native_nccl.py +++ b/tests/py/dynamo/distributed/test_native_nccl.py @@ -1747,6 +1747,85 @@ 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: + """Keep a native all-reduce beside compute in one TRT engine. + + Regression test for pytorch/TensorRT#4381. Pointwise and shuffle layers + around the collective encourage Myelin to absorb it into a ForeignNode; + the converter's precision boundary must leave the DistCollective layer + standalone without introducing a graph break or a second TRT engine. + """ + 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_distributed_mode_subgroup( rank: int, world_size: int, device: torch.device ) -> None: @@ -2092,6 +2171,16 @@ 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) @@ -2143,6 +2232,7 @@ def run_multirank_tests() -> None: _multirank_scatter_correctness, _multirank_gather_correctness, _multirank_distributed_mode_tp_model, + _multirank_compute_collective_single_engine, _multirank_distributed_mode_subgroup, _multirank_cpp_runtime_bind_nccl, _multirank_distributed_mode_context_switch,