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..06a9164a0a 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py @@ -80,6 +80,75 @@ 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 _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], @@ -228,6 +297,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. @@ -254,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 @@ -263,7 +333,49 @@ 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) + + 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}" @@ -284,8 +396,11 @@ 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 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 @@ -302,6 +417,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. @@ -331,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, @@ -350,7 +466,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, @@ -362,8 +478,11 @@ 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 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()}" ) @@ -383,6 +502,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. @@ -413,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, @@ -435,7 +555,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, @@ -447,8 +567,11 @@ 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 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()}" ) @@ -467,6 +590,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. @@ -494,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 @@ -503,7 +627,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}" @@ -524,8 +648,11 @@ 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 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 @@ -542,6 +669,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. @@ -568,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 @@ -577,7 +705,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}" @@ -598,8 +726,11 @@ def nccl_scatter_native( set_layer_name(layer, target, name, source_ir) - output = layer.get_output(0) - layer.num_ranks = world_size + # 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 @@ -616,6 +747,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. @@ -642,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 @@ -651,7 +783,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}" @@ -672,8 +804,11 @@ def nccl_gather_native( set_layer_name(layer, target, name, source_ir) - output = layer.get_output(0) - layer.num_ranks = world_size + # 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 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,