From b1574f0a9791d8142ea58a51eab04d5f2028d28e Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 28 Jul 2026 23:05:14 +0000 Subject: [PATCH 1/7] produce/consume extra output Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 42 ++++++ transformer_engine/pytorch/ops/fuser.py | 174 ++++++++++++++++++++---- transformer_engine/pytorch/ops/op.py | 38 ++++++ 3 files changed, 224 insertions(+), 30 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 66857d8125..fc262b30cf 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -436,6 +436,48 @@ def test_extra_tensors(self, size: int = 16) -> None: torch.testing.assert_close(x3, x3_orig + x2 + b) torch.testing.assert_close(x4, x4_orig + x3) + def test_internal_extra_tensor_channel_fanout(self, size: int = 16) -> None: + """An internal extra output can feed multiple later consumers.""" + producer = te_ops.MakeExtraOutput() + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer1.set_extra_input_channel(0, "route") + consumer2.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer1, consumer2) + + x = torch.rand((size,), requires_grad=True) + y = model(x) + + # Main path: x -> x + route -> x + route + route. + torch.testing.assert_close(y, 3 * x) + y.sum().backward() + # The channel fan-out contributes two independent gradient paths. + torch.testing.assert_close(x.grad, torch.full_like(x, 3)) + + # Internal slots are unavailable before forward, so grad discovery + # must tolerate them when no public input requires gradients. + x_no_grad = x.detach() + torch.testing.assert_close(model(x_no_grad), 3 * x_no_grad) + + def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None: + """Unbound slots remain public when other slots use internal channels.""" + producer = te_ops.MakeExtraOutput() + internal_consumer = te_ops.AddExtraInput() + external_consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + internal_consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, internal_consumer, external_consumer) + + x = torch.rand((size,), requires_grad=True) + extra = torch.rand((size,), requires_grad=True) + y = model(x, extra) + + torch.testing.assert_close(y, 2 * x + extra) + y.sum().backward() + torch.testing.assert_close(x.grad, torch.full_like(x, 2)) + torch.testing.assert_close(extra.grad, torch.ones_like(extra)) + class TestFuser: """Tests for operation fusion infrastructure""" diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 09ffb004dd..33742137db 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -102,12 +102,16 @@ def forward( for tensor in (input_,) + params_and_extra_inputs: tensor._do_not_clear = True - # Unflatten list of parameters and extra tensor inputs - extra_inputs = params_and_extra_inputs[-fuser.num_extra_inputs :] - basic_op_extra_inputs = [] - for op in fuser._basic_ops: - xs, extra_inputs = _split_tuple(extra_inputs, op.num_extra_inputs) - basic_op_extra_inputs.append(xs) + # Place public extra inputs into their basic-op slots. Slots bound to + # internal channels are filled lazily as their producers execute. + extra_inputs = params_and_extra_inputs[len(fuser._flat_basic_op_params) :] + basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ + [None] * op.num_extra_inputs for op in fuser._basic_ops + ] + for tensor, (op_idx, input_idx) in zip( + extra_inputs, fuser._external_extra_input_slots + ): + basic_op_extra_inputs[op_idx][input_idx] = tensor # Apply forward ops x = input_ @@ -118,8 +122,31 @@ def forward( for idx in basic_op_idxs: basic_op_ctxs[idx].requires_grad = idx >= fuser.first_op_requiring_backward - # Forward op - extra_inputs = [basic_op_extra_inputs[idx] for idx in basic_op_idxs] + # Forward op. Resolve internal channel inputs from outputs of + # earlier basic ops. A fusion may consume an earlier channel, but + # may not contain both its producer and consumer. + for idx in basic_op_idxs: + for input_idx, source in enumerate( + fuser._basic_op_extra_input_sources[idx] + ): + if source is None: + continue + producer_idx, output_idx = source + if producer_idx in basic_op_idxs: + raise RuntimeError( + "An operation fusion contains both producer and consumer " + f"of extra tensor channel " + f"{fuser._basic_op_extra_output_channels[producer_idx][output_idx]!r}" + ) + producer_outputs = extra_outputs[producer_idx] + if producer_outputs is None: + raise RuntimeError( + f"Extra tensor channel producer op {producer_idx} has not run" + ) + basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx] + op_extra_inputs = [ + tuple(basic_op_extra_inputs[idx]) for idx in basic_op_idxs + ] prev_op_idx = basic_op_idxs[0] - 1 prev_op = fuser._basic_ops[prev_op_idx] if prev_op_idx >= 0 else None prev_op_grad_output_quantizer = None @@ -134,18 +161,21 @@ def forward( x, fused_op_extra_outputs = op.fuser_forward( [basic_op_ctxs[idx] for idx in basic_op_idxs], x, - basic_op_extra_inputs=extra_inputs, + basic_op_extra_inputs=op_extra_inputs, prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, next_op_input_quantizer=next_op_input_quantizer, basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], ) for idx, ys in zip(basic_op_idxs, fused_op_extra_outputs): for y in ys: - if set_output_requires_grad: + if set_output_requires_grad and ( + y.is_floating_point() or y.is_complex() + ): y.requires_grad_(idx >= fuser.first_op_requiring_backward) extra_outputs[idx] = ys - # Flatten list of extra outputs + # Validate extra outputs and flatten only public slots. Outputs bound + # to channels stay internal to the fuser. extra_outputs_flat = [] for idx, ys in enumerate(extra_outputs): ys = list(ys) @@ -156,7 +186,9 @@ def forward( "{num_extra_outputs} extra inputs, " f"but got {len(ys)}" ) - extra_outputs_flat.extend(ys) + for output_idx, y in enumerate(ys): + if fuser._basic_op_extra_output_channels[idx][output_idx] is None: + extra_outputs_flat.append(y) # Save context for backward pass if func_ctx is not None: @@ -188,10 +220,12 @@ def forward( func_ctx.basic_op_num_params = fuser._basic_op_num_params func_ctx.num_extra_inputs = fuser.num_extra_inputs func_ctx.num_extra_outputs = len(extra_outputs_flat) + func_ctx.fuser = fuser func_ctx.is_first_module = is_first_module # Mark output tensors as not deletable in backward - for tensor in [x] + extra_outputs_flat: + all_extra_outputs = [y for ys in extra_outputs for y in ys] + for tensor in [x] + all_extra_outputs: tensor._do_not_clear = True if set_output_requires_grad: @@ -224,21 +258,28 @@ def backward( ctx.saved_tensors = saved_tensors[slice(*ctx._saved_tensors_range)] ctx._saved_tensors_range = None - # Unflatten list of extra tensor output grads + fuser = func_ctx.fuser + + # Place public extra-output grads into their basic-op slots. Internal + # output grads are accumulated from channel consumers during backward. if len(grad_extra_outputs) != func_ctx.num_extra_outputs: raise ValueError( f"Expected grads for {func_ctx.num_extra_outputs} extra tensor outputs, " f"but got {len(grad_extra_outputs)}" ) - basic_op_grad_extra_outputs = [] - for op in basic_ops: - dys, grad_extra_outputs = _split_tuple(grad_extra_outputs, op.num_extra_outputs) - basic_op_grad_extra_outputs.append(dys) + basic_op_grad_extra_outputs: list[list[Optional[torch.Tensor]]] = [ + [None] * op.num_extra_outputs for op in basic_ops + ] + for grad, (op_idx, output_idx) in zip( + grad_extra_outputs, fuser._external_extra_output_slots + ): + basic_op_grad_extra_outputs[op_idx][output_idx] = grad # Apply backward ops dx = grad_output grad_params = [None for _ in range(len(basic_ops))] grad_extra_inputs = [None for _ in range(len(basic_ops))] + channel_grads: dict[str, torch.Tensor] = {} for op, basic_op_idxs in reversed(backward_ops): # Stop if no more gradients are required @@ -246,18 +287,39 @@ def backward( dx = None break - # Backward op - grad_extra_outputs = [basic_op_grad_extra_outputs[idx] for idx in basic_op_idxs] + # Backward op. Supply gradients accumulated from every consumer of + # each internal channel. + for idx in basic_op_idxs: + for output_idx, channel in enumerate( + fuser._basic_op_extra_output_channels[idx] + ): + if channel is not None: + basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get( + channel + ) + op_grad_extra_outputs = [ + tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs + ] dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( [basic_op_ctxs[idx] for idx in basic_op_idxs], dx, - basic_op_grad_extra_outputs=grad_extra_outputs, + basic_op_grad_extra_outputs=op_grad_extra_outputs, ) for idx, dparams in zip(basic_op_idxs, fused_op_grad_params): grad_params[idx] = dparams basic_op_ctxs[idx].saved_tensors = None for idx, dxs in zip(basic_op_idxs, fused_op_grad_extra_inputs): grad_extra_inputs[idx] = dxs + for input_idx, grad in enumerate(dxs): + source = fuser._basic_op_extra_input_sources[idx][input_idx] + if source is None or grad is None: + continue + producer_idx, output_idx = source + channel = fuser._basic_op_extra_output_channels[producer_idx][output_idx] + previous_grad = channel_grads.get(channel) + channel_grads[channel] = ( + grad if previous_grad is None else previous_grad + grad + ) # Flatten list of parameter gradients grad_params_flat = [] @@ -288,7 +350,9 @@ def backward( f"for {num_extra_inputs} extra inputs, " f"but got {len(dxs)}" ) - grad_extra_inputs_flat.extend(dxs) + for input_idx, grad in enumerate(dxs): + if fuser._basic_op_extra_input_sources[idx][input_idx] is None: + grad_extra_inputs_flat.append(grad) # Update FP8 scaling factors if func_ctx.is_first_module and not _is_graph_capturing(): @@ -342,7 +406,52 @@ def __init__( # Number of extra tensor inputs self._basic_op_num_extra_inputs: list[int] = list(op.num_extra_inputs for op in basic_ops) - self.num_extra_inputs: int = sum(self._basic_op_num_extra_inputs) + self._basic_op_extra_input_sources: list[list[Optional[tuple[int, int]]]] = [ + [None] * op.num_extra_inputs for op in basic_ops + ] + self._basic_op_extra_output_channels: list[list[Optional[str]]] = [ + list(op._extra_output_channels) for op in basic_ops + ] + self._external_extra_input_slots: list[tuple[int, int]] = [] + self._external_extra_output_slots: list[tuple[int, int]] = [] + + # Resolve named channels in pipeline order. Channels deliberately only + # connect an output to later inputs, which keeps execution acyclic. + channel_producers: dict[str, tuple[int, int]] = {} + consumed_channels: set[str] = set() + for op_idx, op in enumerate(basic_ops): + for input_idx, channel in enumerate(op._extra_input_channels): + if channel is None: + self._external_extra_input_slots.append((op_idx, input_idx)) + continue + if channel not in channel_producers: + raise ValueError( + f"Extra tensor channel {channel!r} consumed by op {op_idx} " + f"({type(op).__name__}) has no earlier producer" + ) + self._basic_op_extra_input_sources[op_idx][input_idx] = channel_producers[ + channel + ] + consumed_channels.add(channel) + for output_idx, channel in enumerate(op._extra_output_channels): + if channel is None: + self._external_extra_output_slots.append((op_idx, output_idx)) + continue + if channel in channel_producers: + producer_idx, _ = channel_producers[channel] + raise ValueError( + f"Extra tensor channel {channel!r} has multiple producers " + f"(ops {producer_idx} and {op_idx})" + ) + channel_producers[channel] = (op_idx, output_idx) + + unused_channels = channel_producers.keys() - consumed_channels + if unused_channels: + channels = ", ".join(repr(channel) for channel in sorted(unused_channels)) + raise ValueError(f"Extra tensor channels have no consumers: {channels}") + + self.num_extra_inputs = len(self._external_extra_input_slots) + self.num_extra_outputs = len(self._external_extra_output_slots) # Ops for forward and backward pass, will be populated in maybe_fuse_ops self._forward_ops: list[tuple[FusibleOperation, list[int]]] @@ -432,7 +541,9 @@ def maybe_fuse_ops( first_op_requiring_backward = self._num_basic_ops for op_idx in range(self._num_basic_ops): op_inputs = itertools.chain(self._basic_op_params[op_idx], extra_inputs[op_idx]) - if any(tensor.requires_grad for tensor in op_inputs): + if any( + tensor is not None and tensor.requires_grad for tensor in op_inputs + ): first_op_requiring_backward = op_idx break @@ -517,12 +628,15 @@ def __call__( if basic_op_kwargs is None: basic_op_kwargs = [{}] * self._num_basic_ops - # Unflatten list of extra tensor inputs - extra_inputs_copy = list(extra_inputs) - basic_op_extra_inputs = [] - for op in self._basic_ops: - xs, extra_inputs_copy = _split_tuple(extra_inputs_copy, op.num_extra_inputs) - basic_op_extra_inputs.append(xs) + # Place public extra inputs into their basic-op slots. Internal slots + # are not available until forward executes their producers. + basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ + [None] * op.num_extra_inputs for op in self._basic_ops + ] + for tensor, (op_idx, input_idx) in zip( + extra_inputs, self._external_extra_input_slots + ): + basic_op_extra_inputs[op_idx][input_idx] = tensor # Get environment state recipe = None diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 5106ec9e0a..3159ae0f0d 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -187,10 +187,48 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta): def __init__(self) -> None: super().__init__() + # Optional names for extra-tensor channels internal to an OperationFuser. + # Unbound slots remain public inputs/outputs, preserving the original API. + self._extra_input_channels: list[Optional[str]] = [None] * self.num_extra_inputs + self._extra_output_channels: list[Optional[str]] = [None] * self.num_extra_outputs + # Objects for quantization self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None + def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOperation: + """Bind an extra input slot to an internal fuser channel. + + A bound slot receives the matching extra output from an earlier + operation in the same fuser instead of consuming a public extra input. + Passing ``None`` removes the binding. + """ + if not 0 <= index < self.num_extra_inputs: + raise IndexError( + f"Extra input index {index} is out of range for " + f"{type(self).__name__} with {self.num_extra_inputs} extra inputs" + ) + if channel is not None and (not isinstance(channel, str) or not channel): + raise ValueError("Extra input channel must be a non-empty string or None") + self._extra_input_channels[index] = channel + return self + + def set_extra_output_channel(self, index: int, channel: Optional[str]) -> BasicOperation: + """Bind an extra output slot to an internal fuser channel. + + A bound slot can feed one or more later operations and is not returned + as a public extra output. Passing ``None`` removes the binding. + """ + if not 0 <= index < self.num_extra_outputs: + raise IndexError( + f"Extra output index {index} is out of range for " + f"{type(self).__name__} with {self.num_extra_outputs} extra outputs" + ) + if channel is not None and (not isinstance(channel, str) or not channel): + raise ValueError("Extra output channel must be a non-empty string or None") + self._extra_output_channels[index] = channel + return self + @property def is_fused_op(self) -> bool: return False From 63192abdcaf135d1929f8e2ab73d25d69b9f5870 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 4 Aug 2026 04:28:38 +0000 Subject: [PATCH 2/7] allow for fusions with producer/consumer being part of same fuser with error handling tests Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 468 ++++++++++++++++++ .../pytorch/ops/basic/grouped_linear.py | 2 + transformer_engine/pytorch/ops/fuser.py | 74 ++- 3 files changed, 534 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index fc262b30cf..67358493b0 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -23,6 +23,7 @@ OUTPUT_BUFFER_KEY, GRAD_INPUT_BUFFER_KEY, ) +from transformer_engine.pytorch.ops.fuser import OperationFuser from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV from transformer_engine.pytorch.ops.fused import ( @@ -478,6 +479,473 @@ def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None torch.testing.assert_close(x.grad, torch.full_like(x, 2)) torch.testing.assert_close(extra.grad, torch.ones_like(extra)) + def test_moe_style_dispatch_combine_extra_channels( + self, + *, + group_size: int = 4, + hidden_size: int = 32, + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + ) -> None: + """Wire Dispatch-style extras through GroupedLinear / ScaledActivation / Combine. + + Stand-in for ``te.Sequential(Dispatch, GroupedLinear, ScaledActivation, + GroupedLinear, Combine)`` once real Dispatch/Combine ops land. ``m_splits`` + and ``probs`` are produced once and fan out to later consumers via named + channels so the public call is ``model(x, m_splits, probs)``. + """ + + class FakeDispatch(te_ops.BasicOperation): + """Stand-in MoE dispatch: passthrough hidden states, emit routing extras. + + Real Dispatch would permute tokens. Extra inputs are the externally + provided ``m_splits`` and ``probs``; matching extra outputs are bound + to internal channels for later consumers. A third extra output is a + stub ``routing_map`` for Combine. + """ + + num_extra_inputs = 2 + num_extra_outputs = 3 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("FakeDispatch uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("FakeDispatch uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs, + input_, + *, + basic_op_extra_inputs, + **unused, + ): + m_splits, probs = basic_op_extra_inputs[0] + # Stub row-id map: real Dispatch would emit permute indices. + routing_map = torch.arange( + input_.size(0), device=input_.device, dtype=torch.int64 + ) + return input_, [(m_splits, probs, routing_map)] + + def fuser_backward( + self, + basic_op_ctxs, + grad_output, + *, + basic_op_grad_extra_outputs, + ): + # Channel grads from GroupedLinear / ScaledActivation land here. + return ( + grad_output, + [()], + [tuple(basic_op_grad_extra_outputs[0][:2])], + ) + + class FakeCombine(te_ops.BasicOperation): + """Stand-in MoE combine: consumes Dispatch ``routing_map``, identity path. + + Real Combine would unpermute with the routing map. + """ + + num_extra_inputs = 1 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("FakeCombine uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("FakeCombine uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs, + input_, + *, + basic_op_extra_inputs, + **unused, + ): + routing_map = basic_op_extra_inputs[0][0] + if routing_map is None: + raise RuntimeError("FakeCombine expected routing_map channel") + if int(routing_map.numel()) != int(input_.size(0)): + raise RuntimeError( + "FakeCombine routing_map length does not match tokens" + ) + return input_, [()] + + def fuser_backward( + self, + basic_op_ctxs, + grad_output, + *, + basic_op_grad_extra_outputs, + ): + del basic_op_grad_extra_outputs + return grad_output, [()], [(None,)] + + split_sizes = torch.tensor([8, 16, 8, 8], dtype=torch.int64, device=device)[ + :group_size + ] + num_tokens = int(split_sizes.sum()) + in_shape = (num_tokens, hidden_size) + + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + min=-0.25, + max=0.25, + test_dtype=dtype, + test_device=device, + ) + probs_ref, probs_test = make_reference_and_test_tensors( + (num_tokens,), + min=-0.25, + max=0.25, + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + in_shape, + min=-0.25, + max=0.25, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Reference: GroupedLinear + ScaledSReLU + GroupedLinear (no dispatch permute). + # Run the PyTorch reference in the same dtype/device as TE. Cross-device + # float32 GEMMs (CPU vs CUDA) differ enough that probs grads — a + # reduction over hidden — can miss dtype_tols even when channel wiring + # is correct. + fc1_w_refs, fc1_w_tests = [], [] + fc2_w_refs, fc2_w_tests = [], [] + for _ in range(group_size): + w1_ref, w1_test = make_reference_and_test_tensors( + (hidden_size, hidden_size), + min=-0.25, + max=0.25, + test_dtype=dtype, + test_device=device, + ) + w2_ref, w2_test = make_reference_and_test_tensors( + (hidden_size, hidden_size), + min=-0.25, + max=0.25, + test_dtype=dtype, + test_device=device, + ) + fc1_w_refs.append(w1_test.detach().clone()) + fc1_w_tests.append(w1_test) + fc2_w_refs.append(w2_test.detach().clone()) + fc2_w_tests.append(w2_test) + x_ref = x_test.detach().clone().requires_grad_(True) + probs_ref = probs_test.detach().clone().requires_grad_(True) + dy_ref = dy_test.detach().clone() + xs = torch.split(x_ref, split_sizes.tolist()) + probs = torch.split(probs_ref, split_sizes.tolist()) + ys = [] + for group_idx in range(group_size): + fc1_out = torch.nn.functional.linear(xs[group_idx], fc1_w_refs[group_idx]) + act_out = torch.nn.functional.relu(fc1_out).square() + fc2_in = act_out * probs[group_idx].unsqueeze(-1) + ys.append(torch.nn.functional.linear(fc2_in, fc2_w_refs[group_idx])) + y_ref = torch.cat(ys) + y_ref.backward(dy_ref) + + dispatch = FakeDispatch() + fc1 = te_ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ) + activation = te_ops.ScaledSReLU() + fc2 = te_ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ) + combine = FakeCombine() + + # Bind channels: Dispatch fans out, consumers bind by name. + dispatch.set_extra_output_channel(0, "m_splits") + dispatch.set_extra_output_channel(1, "probs") + dispatch.set_extra_output_channel(2, "routing_map") + fc1.set_extra_input_channel(0, "m_splits") + activation.set_extra_input_channel(0, "probs") + fc2.set_extra_input_channel(0, "m_splits") + combine.set_extra_input_channel(0, "routing_map") + + model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) + with torch.no_grad(): + for group_idx in range(group_size): + getattr(fc1, f"weight{group_idx}").copy_(fc1_w_tests[group_idx]) + getattr(fc2, f"weight{group_idx}").copy_(fc2_w_tests[group_idx]) + del fc1_w_tests, fc2_w_tests + + # Only Dispatch's extras remain public: model(x, m_splits, probs). + y_test = model(x_test, split_sizes, probs_test) + y_test.backward(dy_test) + + tols = dtype_tols(dtype) + assert_close(y_test, y_ref, **tols) + assert_close_grads(x_test, x_ref, **tols) + assert_close_grads(probs_test, probs_ref, **tols) + + def test_fused_op_with_internal_producer_consumer(self, size: int = 16) -> None: + """A fused op may contain both a channel producer and its consumer. + + This is the MegaMoE path: Dispatch + MLP + Combine collapse into one + fused kernel that wires channels internally instead of via the fuser. + """ + + class FakeDispatch(te_ops.BasicOperation): + num_extra_inputs = 1 + num_extra_outputs = 1 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("FakeDispatch uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("FakeDispatch uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs, + input_, + *, + basic_op_extra_inputs, + **unused, + ): + (route,) = basic_op_extra_inputs[0] + return input_, [(route,)] + + def fuser_backward( + self, + basic_op_ctxs, + grad_output, + *, + basic_op_grad_extra_outputs, + ): + return ( + grad_output, + [()], + [tuple(basic_op_grad_extra_outputs[0])], + ) + + class MegaMoELike(te_ops.FusedOperation): + """Fused stub that owns both producer and consumer of ``route``.""" + + _enabled = True + + def __init__(self, dispatch, consumer) -> None: + super().__init__((dispatch, consumer)) + + def fuser_forward( + self, + basic_op_ctxs, + input_, + *, + basic_op_extra_inputs, + **unused, + ): + # Consumer slot is intentionally unset: producer is in this fusion. + route = basic_op_extra_inputs[0][0] + assert basic_op_extra_inputs[1][0] is None + out = input_ + route + return out, [(route,), ()] + + def fuse_mega_moe_like(ops, **unused): + if not MegaMoELike._enabled: + return ops + MegaMoELike._enabled = False + out = [] + window, ops = ops[:2], ops[2:] + while len(window) == 2: + if isinstance(window[0], FakeDispatch) and isinstance( + window[1], te_ops.AddExtraInput + ): + window = [MegaMoELike(*window)] + else: + out.append(window[0]) + window = window[1:] + window, ops = window + ops[:1], ops[1:] + out.extend(window + ops) + return out + + dispatch = FakeDispatch() + consumer = te_ops.AddExtraInput() + dispatch.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(dispatch, consumer) + + te_ops.register_forward_fusion(fuse_mega_moe_like) + x = torch.rand((size,), requires_grad=True) + route = torch.rand((size,), requires_grad=True) + y = model(x, route) + torch.testing.assert_close(y, x + route) + y.sum().backward() + torch.testing.assert_close(x.grad, torch.ones_like(x)) + torch.testing.assert_close(route.grad, torch.ones_like(route)) + + +class TestExtraTensorChannels: + """Error handling and grad coverage for named extra-tensor channels.""" + + def test_consumer_channel_without_producer(self) -> None: + """Extra input bound to a channel that no earlier op produces.""" + consumer = te_ops.AddExtraInput() + consumer.set_extra_input_channel(0, "missing") + with pytest.raises(ValueError, match="has no earlier producer"): + OperationFuser([consumer]) + + def test_consumer_before_producer(self) -> None: + """Channels only connect forward; a later producer does not satisfy an earlier consumer.""" + consumer = te_ops.AddExtraInput() + producer = te_ops.MakeExtraOutput() + consumer.set_extra_input_channel(0, "route") + producer.set_extra_output_channel(0, "route") + with pytest.raises(ValueError, match="has no earlier producer"): + OperationFuser([consumer, producer]) + + def test_set_extra_channel_rejects_invalid_index(self) -> None: + """Slot indices must be in range; negatives and OOB are rejected at bind time.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + + with pytest.raises(IndexError, match="out of range"): + producer.set_extra_output_channel(-1, "route") + with pytest.raises(IndexError, match="out of range"): + producer.set_extra_output_channel(1, "route") + with pytest.raises(IndexError, match="out of range"): + consumer.set_extra_input_channel(-1, "route") + with pytest.raises(IndexError, match="out of range"): + consumer.set_extra_input_channel(1, "route") + + def test_set_extra_channel_rejects_invalid_name(self) -> None: + """Channel names must be non-empty strings.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + + with pytest.raises(ValueError, match="non-empty string"): + producer.set_extra_output_channel(0, "") + with pytest.raises(ValueError, match="non-empty string"): + consumer.set_extra_input_channel(0, "") + with pytest.raises(ValueError, match="non-empty string"): + producer.set_extra_output_channel(0, 123) # type: ignore[arg-type] + + def test_duplicate_extra_output_channel_names(self) -> None: + """Two extra outputs may not publish the same channel name.""" + producer1 = te_ops.MakeExtraOutput() + producer2 = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer1.set_extra_output_channel(0, "route") + producer2.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + with pytest.raises(ValueError, match="multiple producers"): + OperationFuser([producer1, producer2, consumer]) + + def test_duplicate_extra_output_channels_on_same_op(self) -> None: + """A single op with multiple extras still cannot reuse a channel name.""" + + class DualExtraOutput(te_ops.BasicOperation): + num_extra_outputs = 2 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("DualExtraOutput uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("DualExtraOutput uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + return input_, [(input_, input_)] + + def fuser_backward( + self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs + ): + g0, g1 = basic_op_grad_extra_outputs[0] + grad_extra = torch.zeros_like(grad_output) + if g0 is not None: + grad_extra = grad_extra + g0 + if g1 is not None: + grad_extra = grad_extra + g1 + return grad_output + grad_extra, [()], [()] + + producer = DualExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + producer.set_extra_output_channel(1, "route") + consumer.set_extra_input_channel(0, "route") + with pytest.raises(ValueError, match="multiple producers"): + OperationFuser([producer, consumer]) + + def test_unused_extra_output_channel(self) -> None: + """Every produced channel must have at least one consumer.""" + producer = te_ops.MakeExtraOutput() + producer.set_extra_output_channel(0, "orphan") + with pytest.raises(ValueError, match="have no consumers"): + OperationFuser([producer]) + + def test_one_extra_input_has_single_source(self) -> None: + """Each extra-input slot binds to one channel / one producer source. + + Rebinding replaces the previous name; the abandoned producer channel + then fails as unused rather than attaching two sources to one input. + """ + producer_a = te_ops.MakeExtraOutput() + producer_b = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer_a.set_extra_output_channel(0, "a") + producer_b.set_extra_output_channel(0, "b") + consumer.set_extra_input_channel(0, "a") + consumer.set_extra_input_channel(0, "b") + with pytest.raises(ValueError, match="have no consumers"): + OperationFuser([producer_a, producer_b, consumer]) + + # Valid single binding: consumer input 0 is fed only by producer_a. + consumer.set_extra_input_channel(0, "a") + producer_b.set_extra_output_channel(0, None) + fuser = OperationFuser([producer_a, producer_b, consumer]) + assert fuser._basic_op_extra_input_sources[2] == [(0, 0)] + assert fuser.num_extra_inputs == 0 + assert fuser.num_extra_outputs == 1 # producer_b's unbound extra output + + def test_channel_fanout_accumulates_grads(self, size: int = 16) -> None: + """Grads from every consumer of a channel are accumulated into the producer.""" + producer = te_ops.MakeExtraOutput() + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer1.set_extra_input_channel(0, "route") + consumer2.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer1, consumer2) + + x = torch.rand((size,), requires_grad=True) + y = model(x) + # Forward: x -> x+x -> x+x+x + torch.testing.assert_close(y, 3 * x) + + dy = torch.rand((size,)) + y.backward(dy) + # Main path contributes dy; each AddExtraInput also routes dy back + # through the channel into MakeExtraOutput's extra-output grad, which + # is added again into dx. Total: dy (main) + dy + dy (two consumers). + torch.testing.assert_close(x.grad, 3 * dy) + + def test_mixed_internal_external_grad(self, size: int = 16) -> None: + """Internal channel grads and public extra-input grads both flow correctly.""" + producer = te_ops.MakeExtraOutput() + internal_consumer = te_ops.AddExtraInput() + external_consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + internal_consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, internal_consumer, external_consumer) + + x = torch.rand((size,), requires_grad=True) + extra = torch.rand((size,), requires_grad=True) + y = model(x, extra) + torch.testing.assert_close(y, 2 * x + extra) + + dy = torch.rand((size,)) + y.backward(dy) + torch.testing.assert_close(x.grad, 2 * dy) + torch.testing.assert_close(extra.grad, dy) + class TestFuser: """Tests for operation fusion infrastructure""" diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index be931829ea..3f4ca2d4b2 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -152,6 +152,8 @@ def __init__( self._scale_bias: bool = scale_bias and bias if self._scale_bias: self.num_extra_inputs = 2 + # BasicOperation.__init__ sized channel lists from the class default (1). + self._extra_input_channels = [None] * self.num_extra_inputs self.wgrad_store = WeightGradStore(delay_wgrad_compute) self.wgrad_accumulation_and_reduce_hooks: list = [] diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 33742137db..2655ea4319 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -123,8 +123,9 @@ def forward( basic_op_ctxs[idx].requires_grad = idx >= fuser.first_op_requiring_backward # Forward op. Resolve internal channel inputs from outputs of - # earlier basic ops. A fusion may consume an earlier channel, but - # may not contain both its producer and consumer. + # earlier basic ops. When a fusion contains both producer and + # consumer, leave the consumer slot unset so the fused op can + # wire the channel itself for idx in basic_op_idxs: for input_idx, source in enumerate( fuser._basic_op_extra_input_sources[idx] @@ -133,16 +134,25 @@ def forward( continue producer_idx, output_idx = source if producer_idx in basic_op_idxs: - raise RuntimeError( - "An operation fusion contains both producer and consumer " - f"of extra tensor channel " - f"{fuser._basic_op_extra_output_channels[producer_idx][output_idx]!r}" - ) + # fused op will wire the channel itself internally + continue producer_outputs = extra_outputs[producer_idx] if producer_outputs is None: raise RuntimeError( f"Extra tensor channel producer op {producer_idx} has not run" ) + if ( + output_idx >= len(producer_outputs) + or producer_outputs[output_idx] is None + ): + raise RuntimeError( + f"Extra tensor channel producer op {producer_idx} " + f"({type(fuser._basic_ops[producer_idx]).__name__}) " + f"did not emit extra output {output_idx} for " + f"consumer op {idx} " + f"({type(fuser._basic_ops[idx]).__name__}) " + f"input {input_idx}" + ) basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx] op_extra_inputs = [ tuple(basic_op_extra_inputs[idx]) for idx in basic_op_idxs @@ -167,7 +177,12 @@ def forward( basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], ) for idx, ys in zip(basic_op_idxs, fused_op_extra_outputs): - for y in ys: + for output_idx, y in enumerate(ys): + if y is None: + raise RuntimeError( + f"Op {idx} ({type(fuser._basic_ops[idx]).__name__}) " + f"did not emit extra output {output_idx}" + ) if set_output_requires_grad and ( y.is_floating_point() or y.is_complex() ): @@ -315,6 +330,10 @@ def backward( if source is None or grad is None: continue producer_idx, output_idx = source + # Producer already ran inside this fusion; the fused op + # must apply these grads itself rather than via channel_grads. + if producer_idx in basic_op_idxs: + continue channel = fuser._basic_op_extra_output_channels[producer_idx][output_idx] previous_grad = channel_grads.get(channel) channel_grads[channel] = ( @@ -420,7 +439,8 @@ def __init__( channel_producers: dict[str, tuple[int, int]] = {} consumed_channels: set[str] = set() for op_idx, op in enumerate(basic_ops): - for input_idx, channel in enumerate(op._extra_input_channels): + for input_idx in range(op.num_extra_inputs): + channel = op._extra_input_channels[input_idx] if channel is None: self._external_extra_input_slots.append((op_idx, input_idx)) continue @@ -433,7 +453,9 @@ def __init__( channel ] consumed_channels.add(channel) - for output_idx, channel in enumerate(op._extra_output_channels): + for output_idx, channel in enumerate( + self._basic_op_extra_output_channels[op_idx] + ): if channel is None: self._external_extra_output_slots.append((op_idx, output_idx)) continue @@ -450,6 +472,38 @@ def __init__( channels = ", ".join(repr(channel) for channel in sorted(unused_channels)) raise ValueError(f"Extra tensor channels have no consumers: {channels}") + # Every channel-bound extra input must be wired to a matching producer + # extra output. External slots remain unbound (source is None). + for op_idx, sources in enumerate(self._basic_op_extra_input_sources): + op = basic_ops[op_idx] + for input_idx, source in enumerate(sources): + channel = op._extra_input_channels[input_idx] + if channel is None: + if source is not None: + raise RuntimeError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is external but has a " + f"producer source {source}" + ) + continue + if source is None: + raise ValueError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is bound to channel {channel!r} " + f"but has no producer" + ) + producer_idx, output_idx = source + producer_channel = self._basic_op_extra_output_channels[producer_idx][ + output_idx + ] + if producer_channel != channel: + raise ValueError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is bound to channel {channel!r}, " + f"but producer op {producer_idx} extra output {output_idx} " + f"is bound to {producer_channel!r}" + ) + self.num_extra_inputs = len(self._external_extra_input_slots) self.num_extra_outputs = len(self._external_extra_output_slots) From 3b4b523937b8741784f3f3c1e27433326dfdde0c Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 4 Aug 2026 18:18:43 +0000 Subject: [PATCH 3/7] cleanup Signed-off-by: Varun Thumbe --- .../pytorch/ops/basic/grouped_linear.py | 6 +-- transformer_engine/pytorch/ops/fuser.py | 37 +++++++++++-------- 2 files changed, 24 insertions(+), 19 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 3f4ca2d4b2..7faad6536b 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -147,13 +147,11 @@ def __init__( delay_wgrad_compute: bool = False, scale_bias: bool = False, ) -> None: - super().__init__() - + # Decide before BasicOperation.__init__ sizes _extra_input_channels. self._scale_bias: bool = scale_bias and bias if self._scale_bias: self.num_extra_inputs = 2 - # BasicOperation.__init__ sized channel lists from the class default (1). - self._extra_input_channels = [None] * self.num_extra_inputs + super().__init__() self.wgrad_store = WeightGradStore(delay_wgrad_compute) self.wgrad_accumulation_and_reduce_hooks: list = [] diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 2655ea4319..f0f5f36fe9 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -102,7 +102,7 @@ def forward( for tensor in (input_,) + params_and_extra_inputs: tensor._do_not_clear = True - # Place public extra inputs into their basic-op slots. Slots bound to + # Place user provided extra inputs into their basic-op slots. Slots bound to # internal channels are filled lazily as their producers execute. extra_inputs = params_and_extra_inputs[len(fuser._flat_basic_op_params) :] basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ @@ -183,14 +183,10 @@ def forward( f"Op {idx} ({type(fuser._basic_ops[idx]).__name__}) " f"did not emit extra output {output_idx}" ) - if set_output_requires_grad and ( - y.is_floating_point() or y.is_complex() - ): - y.requires_grad_(idx >= fuser.first_op_requiring_backward) extra_outputs[idx] = ys # Validate extra outputs and flatten only public slots. Outputs bound - # to channels stay internal to the fuser. + # to channels stay internal to the fuser and are not marked. extra_outputs_flat = [] for idx, ys in enumerate(extra_outputs): ys = list(ys) @@ -202,7 +198,11 @@ def forward( f"but got {len(ys)}" ) for output_idx, y in enumerate(ys): + # Output is not bound to a channel consumed by another op, + # so it is a public output. if fuser._basic_op_extra_output_channels[idx][output_idx] is None: + if set_output_requires_grad: + y.requires_grad_(idx >= fuser.first_op_requiring_backward) extra_outputs_flat.append(y) # Save context for backward pass @@ -228,14 +228,16 @@ def forward( if fuser.first_op_requiring_backward < fuser._num_basic_ops: is_first_module = FP8GlobalStateManager.is_first_fp8_module() - # Other context + # Other context. Save only the wiring metadata needed by + # backward instead of the whole OperationFuser. func_ctx.backward_ops = fuser._backward_ops func_ctx.basic_ops = fuser._basic_ops func_ctx.basic_op_ctxs = basic_op_ctxs func_ctx.basic_op_num_params = fuser._basic_op_num_params - func_ctx.num_extra_inputs = fuser.num_extra_inputs func_ctx.num_extra_outputs = len(extra_outputs_flat) - func_ctx.fuser = fuser + func_ctx.external_extra_output_slots = fuser._external_extra_output_slots + func_ctx.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels + func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources func_ctx.is_first_module = is_first_module # Mark output tensors as not deletable in backward @@ -273,7 +275,10 @@ def backward( ctx.saved_tensors = saved_tensors[slice(*ctx._saved_tensors_range)] ctx._saved_tensors_range = None - fuser = func_ctx.fuser + # Channel wiring saved from forward + external_extra_output_slots = func_ctx.external_extra_output_slots + basic_op_extra_output_channels = func_ctx.basic_op_extra_output_channels + basic_op_extra_input_sources = func_ctx.basic_op_extra_input_sources # Place public extra-output grads into their basic-op slots. Internal # output grads are accumulated from channel consumers during backward. @@ -286,7 +291,7 @@ def backward( [None] * op.num_extra_outputs for op in basic_ops ] for grad, (op_idx, output_idx) in zip( - grad_extra_outputs, fuser._external_extra_output_slots + grad_extra_outputs, external_extra_output_slots ): basic_op_grad_extra_outputs[op_idx][output_idx] = grad @@ -306,7 +311,7 @@ def backward( # each internal channel. for idx in basic_op_idxs: for output_idx, channel in enumerate( - fuser._basic_op_extra_output_channels[idx] + basic_op_extra_output_channels[idx] ): if channel is not None: basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get( @@ -326,7 +331,7 @@ def backward( for idx, dxs in zip(basic_op_idxs, fused_op_grad_extra_inputs): grad_extra_inputs[idx] = dxs for input_idx, grad in enumerate(dxs): - source = fuser._basic_op_extra_input_sources[idx][input_idx] + source = basic_op_extra_input_sources[idx][input_idx] if source is None or grad is None: continue producer_idx, output_idx = source @@ -334,7 +339,7 @@ def backward( # must apply these grads itself rather than via channel_grads. if producer_idx in basic_op_idxs: continue - channel = fuser._basic_op_extra_output_channels[producer_idx][output_idx] + channel = basic_op_extra_output_channels[producer_idx][output_idx] previous_grad = channel_grads.get(channel) channel_grads[channel] = ( grad if previous_grad is None else previous_grad + grad @@ -370,7 +375,9 @@ def backward( f"but got {len(dxs)}" ) for input_idx, grad in enumerate(dxs): - if fuser._basic_op_extra_input_sources[idx][input_idx] is None: + # Only append public grad extra inputs to the list + # to be returned to the user. + if basic_op_extra_input_sources[idx][input_idx] is None: grad_extra_inputs_flat.append(grad) # Update FP8 scaling factors From de38ed836a4c54ea84a2c2db2e8519c6e48ba197 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 4 Aug 2026 19:03:26 +0000 Subject: [PATCH 4/7] minor cleanup Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 3 ++- transformer_engine/pytorch/ops/fuser.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 67358493b0..9cecd4531b 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -903,7 +903,8 @@ def test_one_extra_input_has_single_source(self) -> None: fuser = OperationFuser([producer_a, producer_b, consumer]) assert fuser._basic_op_extra_input_sources[2] == [(0, 0)] assert fuser.num_extra_inputs == 0 - assert fuser.num_extra_outputs == 1 # producer_b's unbound extra output + # producer_b's unbound extra output remains public + assert fuser._external_extra_output_slots == [(1, 0)] def test_channel_fanout_accumulates_grads(self, size: int = 16) -> None: """Grads from every consumer of a channel are accumulated into the producer.""" diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index f0f5f36fe9..cc706ab91f 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -510,9 +510,9 @@ def __init__( f"but producer op {producer_idx} extra output {output_idx} " f"is bound to {producer_channel!r}" ) - + # Used by Sequential to determine the number of extra inputs + # needed for each OperationFuser module in the sequence. self.num_extra_inputs = len(self._external_extra_input_slots) - self.num_extra_outputs = len(self._external_extra_output_slots) # Ops for forward and backward pass, will be populated in maybe_fuse_ops self._forward_ops: list[tuple[FusibleOperation, list[int]]] From 385b0d51b3b074d00108e1ba0373bf6ba6b57678 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 4 Aug 2026 23:10:17 +0000 Subject: [PATCH 5/7] dispatch combine impl Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/moe_ep_reference.py | 792 ++++++++++++++++++ tests/pytorch/distributed/run_ep.py | 160 ++++ .../pytorch/ops/basic/__init__.py | 2 + transformer_engine/pytorch/ops/combine.py | 153 ++++ transformer_engine/pytorch/ops/dispatch.py | 214 +++++ 5 files changed, 1321 insertions(+) create mode 100644 tests/pytorch/distributed/moe_ep_reference.py create mode 100644 transformer_engine/pytorch/ops/combine.py create mode 100644 transformer_engine/pytorch/ops/dispatch.py diff --git a/tests/pytorch/distributed/moe_ep_reference.py b/tests/pytorch/distributed/moe_ep_reference.py new file mode 100644 index 0000000000..3fab88a50b --- /dev/null +++ b/tests/pytorch/distributed/moe_ep_reference.py @@ -0,0 +1,792 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Pure PyTorch semantic reference for a SwiGLU MoE with expert parallelism. + +The implementation deliberately favors readable semantics over performance. It +supports a one-rank execution path and a variable-size ``all_to_all_single`` EP +path, plus BF16, MXFP8, and NVFP4 block-scaled public outputs. + +The quantized tensor layouts are logical (unswizzled) layouts. A production +kernel may reorder scale factors internally without changing this API contract. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Optional, Sequence, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +class MoeFormat(str, Enum): + """Public and communication formats supported by the reference.""" + + BF16 = "bf16" + MXFP8 = "mxfp8" + NVFP4 = "nvfp4" + + +def _parse_format(value: Union[MoeFormat, str]) -> MoeFormat: + if isinstance(value, MoeFormat): + return value + try: + return MoeFormat(value.lower()) + except (AttributeError, ValueError) as exc: + choices = ", ".join(item.value for item in MoeFormat) + raise ValueError(f"unsupported format {value!r}; expected one of: {choices}") from exc + + +def _require_torch_dtype(name: str) -> torch.dtype: + dtype = getattr(torch, name, None) + if dtype is None: + raise RuntimeError(f"this PyTorch build does not provide torch.{name}") + return dtype + + +def _normalize_axis(axis: int, ndim: int) -> int: + normalized = axis + ndim if axis < 0 else axis + if normalized < 0 or normalized >= ndim: + raise IndexError(f"axis {axis} is out of range for a {ndim}-D tensor") + return normalized + + +def _shape_with_axis(shape: Sequence[int], axis: int, value: int) -> Tuple[int, ...]: + result = list(shape) + result[axis] = value + return tuple(result) + + +def _ceil_div(numerator: int, denominator: int) -> int: + return (numerator + denominator - 1) // denominator + + +@dataclass(frozen=True) +class BlockScaledTensor: + """Portable data-plus-scale representation for MXFP8 or NVFP4. + + ``logical_shape`` describes the dequantized tensor. For MXFP8, ``data`` + has that shape and uses E4M3. For NVFP4, ``data`` is a uint8 tensor with + two E2M1 values per byte along ``axis`` (low nibble first). ``scale`` + replaces that axis by one scale per block. + """ + + data: torch.Tensor + scale: torch.Tensor + format: Union[MoeFormat, str] + logical_shape: Tuple[int, ...] + axis: int = -1 + + def __post_init__(self) -> None: + fmt = _parse_format(self.format) + if fmt is MoeFormat.BF16: + raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") + shape = tuple(int(dim) for dim in self.logical_shape) + if not shape or any(dim < 0 for dim in shape): + raise ValueError(f"logical_shape must contain non-negative dimensions, got {shape}") + axis = _normalize_axis(self.axis, len(shape)) + object.__setattr__(self, "format", fmt) + object.__setattr__(self, "logical_shape", shape) + object.__setattr__(self, "axis", axis) + self._validate_storage() + + @property + def block_size(self) -> int: + return 32 if self.format is MoeFormat.MXFP8 else 16 + + @property + def shape(self) -> Tuple[int, ...]: + return self.logical_shape + + @property + def device(self) -> torch.device: + return self.data.device + + def _validate_storage(self) -> None: + if self.data.device != self.scale.device: + raise ValueError("block-scaled data and scale must be on the same device") + + logical_extent = self.logical_shape[self.axis] + scale_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, self.block_size), + ) + if tuple(self.scale.shape) != scale_shape: + raise ValueError(f"scale shape must be {scale_shape}, got {tuple(self.scale.shape)}") + + if self.format is MoeFormat.MXFP8: + expected_dtype = _require_torch_dtype("float8_e4m3fn") + expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") + data_shape = self.logical_shape + if self.data.dtype != expected_dtype: + raise TypeError(f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}") + else: + expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + if self.data.dtype != torch.uint8 and self.data.dtype != fp4_dtype: + raise TypeError("nvfp4 data must be packed uint8 or torch.float4_e2m1fn_x2") + data_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, 2), + ) + + if tuple(self.data.shape) != data_shape: + raise ValueError(f"data shape must be {data_shape}, got {tuple(self.data.shape)}") + if self.scale.dtype != expected_scale_dtype: + raise TypeError(f"scale must have dtype {expected_scale_dtype}, got {self.scale.dtype}") + + def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Return the logical tensor with block scales applied.""" + + logical_extent = self.logical_shape[self.axis] + scale = self.scale.movedim(self.axis, -1).float() + expanded_scale = scale.repeat_interleave(self.block_size, dim=-1)[..., :logical_extent] + + if self.format is MoeFormat.MXFP8: + values = self.data.movedim(self.axis, -1).float() + else: + packed = self.data + if packed.dtype != torch.uint8: + packed = packed.view(torch.uint8) + packed = packed.movedim(self.axis, -1) + low = packed & 0x0F + high = packed >> 4 + codes = torch.stack((low, high), dim=-1).flatten(-2)[..., :logical_extent] + table = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float32, + device=packed.device, + ) + values = table[codes.long()] + + return (values * expanded_scale).movedim(-1, self.axis).to(dtype) + + +def _nearest_e2m1_codes(values: torch.Tensor) -> torch.Tensor: + """Quantize to E2M1 nibble codes with round-to-nearest, ties-to-even.""" + + levels = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float32, + device=values.device, + ) + magnitudes = values.abs().unsqueeze(-1) + distances = (magnitudes - levels).abs() + minimum = distances.amin(dim=-1, keepdim=True) + candidates = distances == minimum + codes = torch.arange(8, dtype=torch.int64, device=values.device) + any_code = torch.where(candidates, codes, 8).amin(dim=-1) + even_code = torch.where(candidates & ((codes & 1) == 0), codes, 8).amin(dim=-1) + magnitude_code = torch.where(even_code < 8, even_code, any_code) + sign_code = torch.signbit(values).to(torch.int64) << 3 + return magnitude_code | sign_code + + +def quantize_blockwise( + tensor: torch.Tensor, + format: Union[MoeFormat, str], + *, + axis: int = -1, +) -> BlockScaledTensor: + """Quantize a floating tensor into logical MXFP8 or NVFP4 blocks. + + MXFP8 uses 32-value blocks, E4M3 payloads, and E8M0 scales rounded toward + positive infinity. NVFP4 uses 16-value blocks, packed E2M1 payloads, and + E4M3 scales rounded to nearest. + """ + + fmt = _parse_format(format) + if fmt is MoeFormat.BF16: + raise ValueError("quantize_blockwise requires mxfp8 or nvfp4") + if not tensor.is_floating_point(): + raise TypeError(f"tensor must be floating point, got {tensor.dtype}") + + axis = _normalize_axis(axis, tensor.ndim) + logical_shape = tuple(tensor.shape) + moved = tensor.float().movedim(axis, -1) + logical_extent = moved.shape[-1] + block_size = 32 if fmt is MoeFormat.MXFP8 else 16 + block_count = _ceil_div(logical_extent, block_size) + padded_extent = block_count * block_size + if padded_extent != logical_extent: + moved = F.pad(moved, (0, padded_extent - logical_extent)) + blocks = moved.reshape(*moved.shape[:-1], block_count, block_size) + + value_limit = 448.0 if fmt is MoeFormat.MXFP8 else 6.0 + scale_float = blocks.abs().amax(dim=-1) / value_limit + if fmt is MoeFormat.MXFP8: + safe_scale = torch.where(scale_float > 0, scale_float, 1.0) + scale_float = torch.where( + scale_float > 0, + torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), + torch.zeros_like(scale_float), + ) + scale_dtype = _require_torch_dtype("float8_e8m0fnu") + else: + scale_dtype = _require_torch_dtype("float8_e4m3fn") + + scale = scale_float.to(scale_dtype) + scale_for_math = scale.float() + reciprocal = torch.where(scale_for_math > 0, scale_for_math.reciprocal(), 0.0) + normalized = (blocks * reciprocal.unsqueeze(-1)).clamp(-value_limit, value_limit) + + if fmt is MoeFormat.MXFP8: + data_dtype = _require_torch_dtype("float8_e4m3fn") + data = normalized.to(data_dtype).reshape(*moved.shape)[..., :logical_extent] + else: + codes = _nearest_e2m1_codes(normalized).reshape(*moved.shape) + low = codes[..., 0::2] + high = codes[..., 1::2] + data = (low | (high << 4)).to(torch.uint8)[..., : _ceil_div(logical_extent, 2)] + + return BlockScaledTensor( + data=data.movedim(-1, axis).contiguous(), + scale=scale.movedim(-1, axis).contiguous(), + format=fmt, + logical_shape=logical_shape, + axis=axis, + ) + + +MoeTensor = Union[torch.Tensor, BlockScaledTensor] + + +@dataclass(frozen=True) +class _DispatchPlan: + """Send-side routing derived from ``topk_idx``; identical in fwd and bwd.""" + + send_expert: torch.Tensor # local expert id per sent route + send_weight: torch.Tensor # router weight per sent route + send_token_idx: torch.Tensor # source token per sent route + send_slot_idx: torch.Tensor # source top-k slot per sent route + send_counts: Tuple[int, ...] # routes sent to each rank + recv_counts: Tuple[int, ...] # routes received from each rank + + +def _tensor_device(tensor: MoeTensor) -> torch.device: + return tensor.device + + +def _decode_tensor( + tensor: MoeTensor, + *, + name: str, + expected_shape: Tuple[int, ...], + quantized_axis: int, +) -> torch.Tensor: + if isinstance(tensor, BlockScaledTensor): + if tensor.logical_shape != expected_shape: + raise ValueError(f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}") + if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): + raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") + return tensor.dequantize() + + if tuple(tensor.shape) != expected_shape: + raise ValueError(f"{name} shape must be {expected_shape}, got {tuple(tensor.shape)}") + if not tensor.is_floating_point(): + raise TypeError(f"{name} must be floating point or BlockScaledTensor, got {tensor.dtype}") + return tensor.float() + + +def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: + if format is MoeFormat.BF16: + return tensor.to(torch.bfloat16).float() + return quantize_blockwise(tensor, format, axis=-1).dequantize() + + +class MoeEpReference: + """Reference implementation of routed SwiGLU experts plus EP dispatch. + + Global experts are assigned contiguously: rank ``r`` owns + ``[r * experts_per_rank, (r + 1) * experts_per_rank)``. Pass an explicit + initialized process group for multi-rank execution; ``None`` means a + one-rank reference even if the default distributed group is initialized. + """ + + def __init__( + self, + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + ep_group: Optional[dist.ProcessGroup] = None, + max_tokens_per_rank: Optional[int] = None, + output_format: Union[MoeFormat, str] = MoeFormat.BF16, + combine_format: Union[MoeFormat, str] = MoeFormat.BF16, + apply_topk_in_fc1: bool = True, + gate_up_clamp: Optional[float] = None, + generate_c: bool = False, + ) -> None: + for name, value in ( + ("num_experts", num_experts), + ("hidden_size", hidden_size), + ("intermediate_size", intermediate_size), + ("top_k", top_k), + ): + if not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + if top_k > num_experts: + raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") + if max_tokens_per_rank is not None and max_tokens_per_rank < 0: + raise ValueError("max_tokens_per_rank must be non-negative") + + if ep_group is None: + ep_size, ep_rank = 1, 0 + else: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("ep_group requires an initialized torch.distributed process group") + ep_size = dist.get_world_size(ep_group) + ep_rank = dist.get_rank(ep_group) + if num_experts % ep_size != 0: + raise ValueError(f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})") + + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.top_k = top_k + self.ep_group = ep_group + self.ep_size = ep_size + self.ep_rank = ep_rank + self.experts_per_rank = num_experts // ep_size + self.max_tokens_per_rank = max_tokens_per_rank + self.output_format = _parse_format(output_format) + self.combine_format = _parse_format(combine_format) + self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) + self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) + self.generate_c = bool(generate_c) + + for name, fmt in (("output_format", self.output_format), ("combine_format", self.combine_format)): + required_multiple = 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + if hidden_size % required_multiple != 0: + raise ValueError(f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for {name}={fmt.value}") + + def __repr__(self) -> str: + return ( + f"{type(self).__name__}(" + f"experts={self.num_experts}, local_experts={self.experts_per_rank}, " + f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " + f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " + f"output={self.output_format.value}, combine={self.combine_format.value})" + ) + + def _collective_device(self, device: torch.device) -> torch.device: + """Device the process group can run ``all_to_all_single`` on. + + Gloo only implements all-to-all for CPU tensors, so CUDA tensors are + staged through host memory; NCCL groups communicate in place. + """ + if device.type != "cpu" and dist.get_backend(self.ep_group) == "gloo": + return torch.device("cpu") + return device + + def _exchange_counts(self, send_counts: torch.Tensor) -> torch.Tensor: + if self.ep_size == 1: + return send_counts.clone() + comm_device = self._collective_device(send_counts.device) + staged = send_counts.to(comm_device) + recv_counts = torch.empty_like(staged) + dist.all_to_all_single(recv_counts, staged, group=self.ep_group) + return recv_counts.to(send_counts.device) + + def _all_to_all( + self, + send: torch.Tensor, + send_counts: Sequence[int], + recv_counts: Sequence[int], + ) -> torch.Tensor: + if self.ep_size == 1: + return send.clone() + comm_device = self._collective_device(send.device) + staged = send.contiguous().to(comm_device) + output_shape = (sum(recv_counts), *send.shape[1:]) + recv = torch.empty(output_shape, dtype=send.dtype, device=comm_device) + dist.all_to_all_single( + recv, + staged, + output_split_sizes=list(recv_counts), + input_split_sizes=list(send_counts), + group=self.ep_group, + ) + return recv.to(send.device) + + def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> _DispatchPlan: + """Route valid ``topk_idx`` entries to destination ranks, stably by rank. + + Backward reuses this so gradient re-dispatch reproduces the exact + forward route order. + """ + + device = topk_idx.device + token_count = topk_idx.shape[0] + flat_expert = topk_idx.reshape(-1).to(torch.int64) + flat_weight = topk_weights.reshape(-1).float() + valid = flat_expert != -1 + invalid_negative = flat_expert < -1 + invalid_high = flat_expert >= self.num_experts + if bool((invalid_negative | invalid_high).any().item()): + bad = flat_expert[invalid_negative | invalid_high][0].item() + raise ValueError(f"topk_idx contains out-of-range expert id {bad}") + + flat_token = torch.arange(token_count, device=device).repeat_interleave(self.top_k) + flat_slot = torch.arange(self.top_k, device=device).repeat(token_count) + expert = flat_expert[valid] + destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") + order = torch.argsort(destination, stable=True) + + send_counts_tensor = torch.bincount(destination.index_select(0, order), minlength=self.ep_size).to(torch.int64) + recv_counts_tensor = self._exchange_counts(send_counts_tensor) + return _DispatchPlan( + send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), + send_weight=flat_weight[valid].index_select(0, order), + send_token_idx=flat_token[valid].index_select(0, order), + send_slot_idx=flat_slot[valid].index_select(0, order), + send_counts=tuple(int(v) for v in send_counts_tensor.cpu().tolist()), + recv_counts=tuple(int(v) for v in recv_counts_tensor.cpu().tolist()), + ) + + def _run_local_experts( + self, + tokens: torch.Tensor, + local_expert_idx: torch.Tensor, + route_weight: torch.Tensor, + fc1_weight: torch.Tensor, + fc2_weight: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + output = torch.empty( + (tokens.shape[0], self.hidden_size), + dtype=torch.float32, + device=tokens.device, + ) + fc1_c_rows = [] if self.generate_c else None + for expert in range(self.experts_per_rank): + positions = torch.nonzero(local_expert_idx == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + expert_tokens = tokens.index_select(0, positions) + gate_up = expert_tokens @ fc1_weight[expert] + if fc1_c_rows is not None: + # Raw pre-SwiGLU accumulator: before clamp, no router weight. + fc1_c_rows.append(gate_up.to(torch.bfloat16)) + gate, up = gate_up.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + gate = gate.clamp(max=self.gate_up_clamp) + up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + intermediate = F.silu(gate) * up + weights = route_weight.index_select(0, positions).unsqueeze(-1) + if self.apply_topk_in_fc1: + intermediate = intermediate * weights + expert_output = intermediate @ fc2_weight[expert] + if not self.apply_topk_in_fc1: + expert_output = expert_output * weights + expert_output = _format_round_trip(expert_output, self.combine_format) + output.index_copy_(0, positions, expert_output) + fc1_c = None + if fc1_c_rows is not None: + fc1_c = torch.cat(fc1_c_rows) if fc1_c_rows else torch.empty((0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device) + return output, fc1_c + + def __call__( + self, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ) -> Union[MoeTensor, Tuple[MoeTensor, torch.Tensor, torch.Tensor]]: + """Run dispatch, local experts, return routing, top-k reduce, and encode. + + Shapes: + activation: ``(T, H)`` + fc1_weight: ``(E_local, H, 2 * I)`` + fc2_weight: ``(E_local, I, H)`` + topk_idx/topk_weights: ``(T, K)`` + + Returns the ``(T, H)`` result, or ``(result, fc1_c, route_metadata)`` + when constructed with ``generate_c=True``. ``fc1_c`` is the BF16 + pre-SwiGLU FC1 accumulator of every route this rank's experts + processed, ``(local_routes, 2 * I)``, grouped by local expert and + ordered within each expert by (source rank, source token-major route + order); captured before the gate/up clamp, without the router weight. + ``route_metadata`` is Int32 ``(local_routes, 4)`` with columns + ``(local_expert, src_rank, src_token, src_slot)``; row ``i`` identifies + the route behind ``fc1_c`` row ``i`` for the backward gradient + re-dispatch. + """ + + if topk_idx.ndim != 2: + raise ValueError(f"topk_idx must be 2-D, got shape {tuple(topk_idx.shape)}") + token_count = topk_idx.shape[0] + route_shape = (token_count, self.top_k) + if tuple(topk_idx.shape) != route_shape: + raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") + if tuple(topk_weights.shape) != route_shape: + raise ValueError(f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}") + if topk_idx.dtype not in (torch.int32, torch.int64): + raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") + if not topk_weights.is_floating_point(): + raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") + if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: + raise ValueError(f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}") + + device = _tensor_device(activation) + inputs = { + "fc1_weight": _tensor_device(fc1_weight), + "fc2_weight": _tensor_device(fc2_weight), + "topk_idx": topk_idx.device, + "topk_weights": topk_weights.device, + } + for name, input_device in inputs.items(): + if input_device != device: + raise ValueError(f"{name} must be on {device}, got {input_device}") + + activation_float = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + ) + + plan = self._dispatch_plan(topk_idx, topk_weights) + send_token_idx = plan.send_token_idx + send_slot_idx = plan.send_slot_idx + send_counts, recv_counts = plan.send_counts, plan.recv_counts + send_tokens = activation_float.index_select(0, send_token_idx) + + recv_tokens = self._all_to_all(send_tokens, send_counts, recv_counts) + recv_expert = self._all_to_all(plan.send_expert, send_counts, recv_counts) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + + route_metadata = None + if self.generate_c: + recv_src_rank = torch.repeat_interleave( + torch.arange(self.ep_size, device=device), + torch.tensor(recv_counts, device=device), + ) + recv_token = self._all_to_all(send_token_idx, send_counts, recv_counts) + recv_slot = self._all_to_all(send_slot_idx, send_counts, recv_counts) + # Stable sort by local expert reproduces the fc1_c row order + # (grouped by expert; source order preserved within each group). + fc1_c_order = torch.argsort(recv_expert, stable=True) + route_metadata = torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1).index_select(0, fc1_c_order).to(torch.int32) + # recv rows are ordered by source rank, then that source's token-major + # route order, so the per-expert position grouping below realizes the + # documented fc1_c ordering. + recv_output, fc1_c = self._run_local_experts( + recv_tokens, + recv_expert, + recv_weight, + fc1_float, + fc2_float, + ) + + returned = self._all_to_all(recv_output, recv_counts, send_counts) + combine_plane = torch.zeros( + (token_count * self.top_k, self.hidden_size), + dtype=torch.float32, + device=device, + ) + send_flat_slot = send_token_idx * self.top_k + send_slot_idx + combine_plane.index_copy_(0, send_flat_slot, returned) + reduced = combine_plane.view(token_count, self.top_k, self.hidden_size).sum(dim=1) + + if self.output_format is MoeFormat.BF16: + output = reduced.to(torch.bfloat16) + else: + output = quantize_blockwise(reduced, self.output_format, axis=-1) + if self.generate_c: + return output, fc1_c, route_metadata + return output + + def backward( + self, + grad_output: torch.Tensor, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + fc1_c: torch.Tensor, + route_metadata: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Backward pass consuming the ``generate_c=True`` stash. + + ``fc1_c`` is the recompute source: gate/up, the clamp masks, SwiGLU, + and the FC2 input are all rebuilt from it, so no post-SwiGLU forward + intermediate needs to be saved. ``route_metadata`` alone reconstructs + the mapping between re-dispatched rows and ``fc1_c`` rows and drives + the gradient return scatter. + + Quantization round-trips (input decode, ``combine_format``, + ``output_format``) are treated as straight-through identities; + ``grad_output`` is the ``(T, H)`` gradient of the dequantized output. + + Returns ``(grad_activation, grad_fc1_weight, grad_fc2_weight, + grad_topk_weights)`` in float32. + """ + + if not self.generate_c: + raise RuntimeError("backward requires the operator to be constructed with generate_c=True") + token_count = topk_idx.shape[0] + if tuple(grad_output.shape) != (token_count, self.hidden_size): + raise ValueError(f"grad_output shape must be {(token_count, self.hidden_size)}, got {tuple(grad_output.shape)}") + if not grad_output.is_floating_point(): + raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") + + device = _tensor_device(activation) + two_i = 2 * self.intermediate_size + activation_float = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, two_i), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + ) + if fc1_c.shape != (int(route_metadata.shape[0]), two_i): + raise ValueError(f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got {tuple(fc1_c.shape)}") + + # Re-dispatch the FC1 inputs, router weights, and output gradients + # along the identical forward routes. + plan = self._dispatch_plan(topk_idx, topk_weights) + send_counts, recv_counts = plan.send_counts, plan.recv_counts + grad_output_float = grad_output.float() + recv_tokens = self._all_to_all(activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + recv_grad = self._all_to_all(grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + + # route_metadata rows are in fc1_c order; sorting them by + # (src_rank, src_token, src_slot) reproduces the receive order, giving + # the permutation between re-dispatched rows and fc1_c rows. + metadata = route_metadata.to(device=device, dtype=torch.int64) + local_routes = metadata.shape[0] + if local_routes > 0: + token_span = int(metadata[:, 2].max().item()) + 1 + recv_key = (metadata[:, 1] * token_span + metadata[:, 2]) * self.top_k + metadata[:, 3] + perm = torch.argsort(recv_key) # perm[j] = fc1_c row at receive position j + else: + perm = torch.empty((0,), dtype=torch.int64, device=device) + x_rows = torch.empty_like(recv_tokens) + x_rows.index_copy_(0, perm, recv_tokens) + w_rows = torch.empty_like(recv_weight) + w_rows.index_copy_(0, perm, recv_weight) + dy_rows = torch.empty_like(recv_grad) + dy_rows.index_copy_(0, perm, recv_grad) + + c_rows = fc1_c.float() + expert_rows = metadata[:, 0] + d_x_rows = torch.zeros((local_routes, self.hidden_size), dtype=torch.float32, device=device) + d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) + grad_fc1 = torch.zeros_like(fc1_float) + grad_fc2 = torch.zeros_like(fc2_float) + for expert in range(self.experts_per_rank): + positions = torch.nonzero(expert_rows == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + c = c_rows.index_select(0, positions) + x = x_rows.index_select(0, positions) + w = w_rows.index_select(0, positions).unsqueeze(-1) + d_y = dy_rows.index_select(0, positions) + + gate, up = c.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + g = gate.clamp(max=self.gate_up_clamp) + u = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + else: + g, u = gate, up + sig = torch.sigmoid(g) + s = g * sig + h = s * u + + if self.apply_topk_in_fc1: + h_fc2 = h * w + d_y_pre = d_y + else: + h_fc2 = h + d_y_pre = d_y * w + grad_fc2[expert] = h_fc2.transpose(0, 1) @ d_y_pre + d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) + if self.apply_topk_in_fc1: + d_h = d_h_fc2 * w + d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1) + else: + d_h = d_h_fc2 + d_w_rows[positions] = (d_y * (h @ fc2_float[expert])).sum(dim=-1) + + d_g = d_h * u * (sig * (1 + g * (1 - sig))) + d_u = d_h * s + if self.gate_up_clamp is not None: + d_gate = d_g * (gate <= self.gate_up_clamp) + d_up = d_u * ((up >= -self.gate_up_clamp) & (up <= self.gate_up_clamp)) + else: + d_gate, d_up = d_g, d_u + d_c = torch.cat((d_gate, d_up), dim=-1) + grad_fc1[expert] = x.transpose(0, 1) @ d_c + d_x_rows.index_copy_(0, positions, d_c @ fc1_float[expert].transpose(0, 1)) + + # Return the route gradients to their source ranks and scatter-add. + returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) + returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) + grad_activation = torch.zeros((token_count, self.hidden_size), dtype=torch.float32, device=device) + grad_activation.index_add_(0, plan.send_token_idx, returned_dx) + grad_topk_weights = torch.zeros((token_count * self.top_k,), dtype=torch.float32, device=device) + grad_topk_weights.index_copy_(0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw) + return ( + grad_activation, + grad_fc1, + grad_fc2, + grad_topk_weights.view(token_count, self.top_k), + ) + + +__all__ = [ + "BlockScaledTensor", + "MoeEpReference", + "MoeFormat", + "MoeTensor", + "quantize_blockwise", +] diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 4778498b7d..5b965e10bc 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -11,6 +11,8 @@ import torch import torch.distributed as dist +from moe_ep_reference import MoeEpReference +from transformer_engine.pytorch import ops as te_ops from transformer_engine.pytorch.ep import ( EpBuffer, ep_bootstrap, @@ -121,6 +123,30 @@ def _make_identity_inputs(rank, ep_size, device="cuda"): ) +def _make_moe_inputs(rank, ep_size, device="cuda"): + """Random activations and top-k routing representative of a router output.""" + generator = torch.Generator(device=device) + generator.manual_seed(2026 + rank) + num_experts = ep_size * NUM_LOCAL_EXPERTS + tokens = torch.randn( + TOKENS_PER_RANK, + HIDDEN_DIM, + generator=generator, + dtype=torch.float32, + device=device, + ).mul_(0.25) + router_logits = torch.randn( + TOKENS_PER_RANK, + num_experts, + generator=generator, + dtype=torch.float32, + device=device, + ) + topk_logits, topk_idx = torch.topk(router_logits, TOP_K, dim=-1) + topk_weights = torch.softmax(topk_logits, dim=-1) + return topk_idx, tokens.to(torch.bfloat16), topk_weights + + class _Cfg: rank: int world_size: int @@ -554,6 +580,140 @@ def zero_grads(): rtol=5e-2, ) + @_eager_test_include + def test_fusible_dispatch_combine_moe(self): + """Fusible NCCL EP MoE matches the PyTorch all-to-all reference.""" + if not EAGER: + self.skipTest( + "variable grouped-linear splits require eager EP output sizing" + ) + + topk_idx, tokens, topk_weights = _make_moe_inputs( + self.cfg.rank, self.cfg.ep_size + ) + num_experts = NUM_LOCAL_EXPERTS + intermediate_dim = HIDDEN_DIM + + dispatch_buffer = self._make_buffer() + dispatch = te_ops.Dispatch(dispatch_buffer) + fc1 = te_ops.GroupedLinear( + num_experts, + HIDDEN_DIM, + 2 * intermediate_dim, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear( + num_experts, + intermediate_dim, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + combine = te_ops.Combine(dispatch_buffer, num_local_tokens=TOKENS_PER_RANK) + + dispatch.set_extra_output_channel(0, "m_splits") + dispatch.set_extra_output_channel(1, "token_probs") + fc1.set_extra_input_channel(0, "m_splits") + activation.set_extra_input_channel(0, "token_probs") + fc2.set_extra_input_channel(0, "m_splits") + model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) + + generator = torch.Generator(device=self.cfg.device) + generator.manual_seed(1234 + self.cfg.rank) + with torch.no_grad(): + for expert_idx in range(num_experts): + getattr(fc1, f"weight{expert_idx}").uniform_( + -0.1, 0.1, generator=generator + ) + getattr(fc2, f"weight{expert_idx}").uniform_( + -0.1, 0.1, generator=generator + ) + + ref_fc1_weights = torch.stack( + [ + getattr(fc1, f"weight{idx}").detach().transpose(0, 1) + for idx in range(num_experts) + ] + ) + ref_fc2_weights = torch.stack( + [ + getattr(fc2, f"weight{idx}").detach().transpose(0, 1) + for idx in range(num_experts) + ] + ) + reference = MoeEpReference( + num_experts=self.cfg.num_experts, + hidden_size=HIDDEN_DIM, + intermediate_size=intermediate_dim, + top_k=TOP_K, + ep_group=self.ep_group, + max_tokens_per_rank=TOKENS_PER_RANK, + apply_topk_in_fc1=True, + generate_c=True, + ) + ref_output, fc1_c, route_metadata = reference( + tokens, + ref_fc1_weights, + ref_fc2_weights, + topk_idx, + topk_weights, + ) + + test_tokens = tokens.detach().clone().requires_grad_(True) + test_topk_weights = topk_weights.detach().clone().requires_grad_(True) + test_output = model(test_tokens, topk_idx, test_topk_weights) + + grad_output = torch.linspace( + -0.2, + 0.2, + TOKENS_PER_RANK * HIDDEN_DIM, + device=self.cfg.device, + dtype=torch.float32, + ).reshape(TOKENS_PER_RANK, HIDDEN_DIM) + grad_output = grad_output.to(torch.bfloat16) + ref_grads = reference.backward( + grad_output, + tokens, + ref_fc1_weights, + ref_fc2_weights, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + test_output.backward(grad_output) + torch.cuda.synchronize() + + torch.testing.assert_close( + test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2 + ) + torch.testing.assert_close( + test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2 + ) + torch.testing.assert_close( + test_topk_weights.grad, + ref_grads[3], + atol=5e-2, + rtol=5e-2, + ) + for expert_idx in range(num_experts): + torch.testing.assert_close( + getattr(fc1, f"weight{expert_idx}").grad.float(), + ref_grads[1][expert_idx].transpose(0, 1), + atol=5e-2, + rtol=5e-2, + ) + torch.testing.assert_close( + getattr(fc2, f"weight{expert_idx}").grad.float(), + ref_grads[2][expert_idx].transpose(0, 1), + atol=5e-2, + rtol=5e-2, + ) + @_zero_copy_test_include @_eager_test_include def test_combine_autograd(self): diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 6def36ffc7..4ea116793f 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -23,6 +23,8 @@ from .basic_linear import BasicLinear from .bias import Bias from .constant_scale import ConstantScale +from .combine import Combine +from .dispatch import Dispatch from .dropout import Dropout from .grouped_linear import GroupedLinear from .identity import Identity diff --git a/transformer_engine/pytorch/ops/combine.py b/transformer_engine/pytorch/ops/combine.py new file mode 100644 index 0000000000..19f78f7356 --- /dev/null +++ b/transformer_engine/pytorch/ops/combine.py @@ -0,0 +1,153 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible NCCL expert-parallel combine operation.""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from ...ep import EpBuffer, _alloc_io, is_symm_backed +from ...tensor import Quantizer +from ..op import BasicOperation, OperationContext + + +def _validate_grad_buffer( + tensor: Optional[torch.Tensor], + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + if tuple(tensor.shape) != shape: + raise ValueError( + f"grad_out shape {tuple(tensor.shape)} does not match {shape}." + ) + if tensor.dtype is not dtype: + raise TypeError(f"grad_out must have dtype {dtype}, got {tensor.dtype}.") + if tensor.device != device: + raise ValueError(f"grad_out must be on {device}, got {tensor.device}.") + if not tensor.is_contiguous(): + raise ValueError("grad_out must be contiguous.") + if tensor.requires_grad: + raise ValueError("grad_out must not require gradients.") + return tensor + + +class Combine(BasicOperation): + """Combine pre-weighted local expert outputs with NCCL EP. + + The operation uses routing state produced by a :class:`Dispatch` with the + same :class:`EpBuffer`. + """ + + def __init__( + self, buffer: EpBuffer, *, num_local_tokens: Optional[int] = None + ) -> None: + super().__init__() + self.buffer = buffer + self.num_local_tokens = ( + buffer.max_tokens_per_rank + if num_local_tokens is None + else int(num_local_tokens) + ) + if self.num_local_tokens < 0: + raise ValueError("num_local_tokens must be non-negative.") + + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + **kwargs: Any, + ) -> torch.Tensor: + del prev_op_grad_output_quantizer, next_op_input_quantizer + if input_.dtype is not torch.bfloat16: + raise NotImplementedError( + f"NCCL EP requires BF16 combine input, got {input_.dtype}." + ) + if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: + raise ValueError( + f"Combine input must have shape (R, {self.buffer.hidden_dim}), " + f"got {tuple(input_.shape)}." + ) + + expert_out = input_ + if self.buffer.zero_copy: + expert_out = _alloc_io( + tuple(input_.shape), + input_.dtype, + input_.device, + True, + ) + expert_out.copy_(input_) + + result = torch.empty( + self.num_local_tokens, + self.buffer.hidden_dim, + dtype=input_.dtype, + device=input_.device, + ) + torch.ops.transformer_engine_ep.combine( + self.buffer.handle_mem, + expert_out, + result, + ) + + if ctx.requires_grad: + grad_out = kwargs.get("grad_out") + if self.buffer.eager and grad_out is not None: + raise ValueError( + "eager mode sizes combine gradients per step and cannot use " + "a caller-supplied grad_out" + ) + grad_out = _validate_grad_buffer( + grad_out, + shape=tuple(input_.shape), + dtype=input_.dtype, + device=input_.device, + ) + if ( + self.buffer.zero_copy + and grad_out is not None + and not is_symm_backed(grad_out) + ): + raise ValueError( + "zero-copy Combine grad_out must be symmetric-memory-backed." + ) + ctx.grad_out = grad_out + ctx.input_shape = tuple(input_.shape) + ctx.input_dtype = input_.dtype + ctx.save_for_backward(self.buffer.handle_mem) + + return result + + def op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, tuple[()]]: + (handle_mem,) = ctx.saved_tensors + grad_output = grad_output.contiguous() + grad_input = ctx.grad_out + if grad_input is None: + grad_input = _alloc_io( + ctx.input_shape, + ctx.input_dtype, + grad_output.device, + self.buffer.zero_copy, + ) + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + grad_output, + grad_input, + ) + return grad_input, () + diff --git a/transformer_engine/pytorch/ops/dispatch.py b/transformer_engine/pytorch/ops/dispatch.py new file mode 100644 index 0000000000..ab89f6e130 --- /dev/null +++ b/transformer_engine/pytorch/ops/dispatch.py @@ -0,0 +1,214 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible NCCL expert-parallel dispatch operation.""" + +from __future__ import annotations + +from typing import Any, Iterable, Optional + +import torch + +from ...ep import EpBuffer, _alloc_io, ep_prepare +from ...tensor import Quantizer +from ..op import BasicOperation, OperationContext + + +def _validate_output_buffer( + name: str, + tensor: Optional[torch.Tensor], + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} shape {tuple(tensor.shape)} does not match {shape}.") + if tensor.dtype is not dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}.") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}.") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous.") + if tensor.requires_grad: + raise ValueError(f"{name} must not require gradients.") + return tensor + + +class Dispatch(BasicOperation): + """Dispatch BF16 tokens to local experts with NCCL EP. + + The extra inputs are routing indices and FP32 routing weights. The extra + outputs are local tokens-per-expert and received routing weights. + """ + + num_extra_inputs: int = 2 + num_extra_outputs: int = 2 + + def __init__(self, buffer: EpBuffer) -> None: + super().__init__() + self.buffer = buffer + + def op_forward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Dispatch uses fuser_forward") + + def op_backward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Dispatch uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + del prev_op_grad_output_quantizer, next_op_input_quantizer + topk_idx, topk_weights = basic_op_extra_inputs[0] + kwargs = basic_op_kwargs[0] + + if input_.dtype is not torch.bfloat16: + raise NotImplementedError( + f"NCCL EP requires BF16 dispatch input, got {input_.dtype}." + ) + if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: + raise ValueError( + f"Dispatch input must have shape (T, {self.buffer.hidden_dim}), " + f"got {tuple(input_.shape)}." + ) + if topk_idx.dtype not in (torch.int32, torch.int64): + raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}.") + expected_route_shape = (input_.shape[0], self.buffer.top_k) + if tuple(topk_idx.shape) != expected_route_shape: + raise ValueError( + f"topk_idx shape must be {expected_route_shape}, got {tuple(topk_idx.shape)}." + ) + if tuple(topk_weights.shape) != expected_route_shape: + raise ValueError( + f"topk_weights shape must be {expected_route_shape}, " + f"got {tuple(topk_weights.shape)}." + ) + if topk_weights.dtype is not torch.float32: + raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") + for name, tensor in (("topk_idx", topk_idx), ("topk_weights", topk_weights)): + if tensor.device != input_.device: + raise ValueError( + f"{name} must be on {input_.device}, got {tensor.device}." + ) + + recv_tokens = kwargs.get("recv_tokens") + recv_topk_weights = kwargs.get("recv_topk_weights") + if self.buffer.eager and ( + recv_tokens is not None or recv_topk_weights is not None + ): + raise ValueError( + "eager mode sizes dispatch outputs per step and cannot use " + "caller-supplied receive buffers" + ) + + tokens_per_expert = ep_prepare(self.buffer, topk_idx) + rows = ( + self.buffer._host_total_recv_tokens + if self.buffer.eager + else self.buffer.recv_capacity_per_rank + ) + if rows is None: + raise RuntimeError("NCCL EP dispatch receive size is unavailable.") + rows = int(rows) + recv_shape = (rows, self.buffer.hidden_dim) + recv_tokens = _validate_output_buffer( + "recv_tokens", + recv_tokens, + shape=recv_shape, + dtype=self.buffer.payload_dtype, + device=self.buffer.device, + ) + recv_topk_weights = _validate_output_buffer( + "recv_topk_weights", + recv_topk_weights, + shape=(rows,), + dtype=torch.float32, + device=self.buffer.device, + ) + if recv_tokens is None: + recv_tokens = _alloc_io( + recv_shape, + self.buffer.payload_dtype, + self.buffer.device, + self.buffer.zero_copy, + ) + if recv_topk_weights is None: + recv_topk_weights = _alloc_io( + (rows,), + torch.float32, + self.buffer.device, + self.buffer.zero_copy, + ) + + torch.ops.transformer_engine_ep.dispatch( + self.buffer.handle_mem, + topk_idx, + input_, + topk_weights, + recv_tokens, + recv_topk_weights, + ) + + ctx = basic_op_ctxs[0] + if ctx.requires_grad: + ctx.input_shape = tuple(input_.shape) + ctx.input_dtype = input_.dtype + ctx.topk_weights_shape = tuple(topk_weights.shape) + ctx.save_for_backward(self.buffer.handle_mem) + + return recv_tokens, [(tokens_per_expert, recv_topk_weights)] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + ctx = basic_op_ctxs[0] + (handle_mem,) = ctx.saved_tensors + grad_output = grad_output.contiguous() + + grad_recv_weights = basic_op_grad_extra_outputs[0][1] + if grad_recv_weights is None: + grad_recv_weights = torch.zeros( + grad_output.shape[0], + dtype=torch.float32, + device=grad_output.device, + ) + else: + grad_recv_weights = grad_recv_weights.to(dtype=torch.float32).contiguous() + + grad_input = torch.empty( + ctx.input_shape, + dtype=ctx.input_dtype, + device=grad_output.device, + ) + grad_topk_weights = torch.empty( + ctx.topk_weights_shape, + dtype=torch.float32, + device=grad_output.device, + ) + torch.ops.transformer_engine_ep.dispatch_bwd( + handle_mem, + grad_output, + grad_recv_weights, + grad_input, + grad_topk_weights, + ) + return grad_input, [()], [(None, grad_topk_weights)] + From ad3b044f08bd5821850930c4debd39fb3b775c4c Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 5 Aug 2026 01:15:06 +0000 Subject: [PATCH 6/7] fusible ops test Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/run_ep.py | 134 ++++++++++++++++++++++++++++ 1 file changed, 134 insertions(+) diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 5b965e10bc..2c81782a0e 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -714,6 +714,140 @@ def test_fusible_dispatch_combine_moe(self): rtol=5e-2, ) + @_eager_test_include + def test_fusible_dispatch_combine_moe(self): + """Fusible NCCL EP MoE matches the PyTorch all-to-all reference.""" + if not EAGER: + self.skipTest( + "variable grouped-linear splits require eager EP output sizing" + ) + + topk_idx, tokens, topk_weights = _make_moe_inputs( + self.cfg.rank, self.cfg.ep_size + ) + num_experts = NUM_LOCAL_EXPERTS + intermediate_dim = HIDDEN_DIM + + dispatch_buffer = self._make_buffer() + dispatch = te_ops.Dispatch(dispatch_buffer) + fc1 = te_ops.GroupedLinear( + num_experts, + HIDDEN_DIM, + 2 * intermediate_dim, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear( + num_experts, + intermediate_dim, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + combine = te_ops.Combine(dispatch_buffer, num_local_tokens=TOKENS_PER_RANK) + + dispatch.set_extra_output_channel(0, "m_splits") + dispatch.set_extra_output_channel(1, "token_probs") + fc1.set_extra_input_channel(0, "m_splits") + activation.set_extra_input_channel(0, "token_probs") + fc2.set_extra_input_channel(0, "m_splits") + model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) + + generator = torch.Generator(device=self.cfg.device) + generator.manual_seed(1234 + self.cfg.rank) + with torch.no_grad(): + for expert_idx in range(num_experts): + getattr(fc1, f"weight{expert_idx}").uniform_( + -0.1, 0.1, generator=generator + ) + getattr(fc2, f"weight{expert_idx}").uniform_( + -0.1, 0.1, generator=generator + ) + + ref_fc1_weights = torch.stack( + [ + getattr(fc1, f"weight{idx}").detach().transpose(0, 1) + for idx in range(num_experts) + ] + ) + ref_fc2_weights = torch.stack( + [ + getattr(fc2, f"weight{idx}").detach().transpose(0, 1) + for idx in range(num_experts) + ] + ) + reference = MoeEpReference( + num_experts=self.cfg.num_experts, + hidden_size=HIDDEN_DIM, + intermediate_size=intermediate_dim, + top_k=TOP_K, + ep_group=self.ep_group, + max_tokens_per_rank=TOKENS_PER_RANK, + apply_topk_in_fc1=True, + generate_c=True, + ) + ref_output, fc1_c, route_metadata = reference( + tokens, + ref_fc1_weights, + ref_fc2_weights, + topk_idx, + topk_weights, + ) + + test_tokens = tokens.detach().clone().requires_grad_(True) + test_topk_weights = topk_weights.detach().clone().requires_grad_(True) + test_output = model(test_tokens, topk_idx, test_topk_weights) + + grad_output = torch.linspace( + -0.2, + 0.2, + TOKENS_PER_RANK * HIDDEN_DIM, + device=self.cfg.device, + dtype=torch.float32, + ).reshape(TOKENS_PER_RANK, HIDDEN_DIM) + grad_output = grad_output.to(torch.bfloat16) + ref_grads = reference.backward( + grad_output, + tokens, + ref_fc1_weights, + ref_fc2_weights, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + test_output.backward(grad_output) + torch.cuda.synchronize() + + torch.testing.assert_close( + test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2 + ) + torch.testing.assert_close( + test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2 + ) + torch.testing.assert_close( + test_topk_weights.grad, + ref_grads[3], + atol=5e-2, + rtol=5e-2, + ) + for expert_idx in range(num_experts): + torch.testing.assert_close( + getattr(fc1, f"weight{expert_idx}").grad.float(), + ref_grads[1][expert_idx].transpose(0, 1), + atol=5e-2, + rtol=5e-2, + ) + torch.testing.assert_close( + getattr(fc2, f"weight{expert_idx}").grad.float(), + ref_grads[2][expert_idx].transpose(0, 1), + atol=5e-2, + rtol=5e-2, + ) + @_zero_copy_test_include @_eager_test_include def test_combine_autograd(self): From 5fb0d3af53e90bf8cf421c705e018307c71676df Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:46:15 +0000 Subject: [PATCH 7/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/moe_ep_reference.py | 90 ++++++++++++++----- tests/pytorch/distributed/run_ep.py | 68 ++++---------- tests/pytorch/test_fusible_ops.py | 16 +--- transformer_engine/pytorch/ops/combine.py | 27 ++---- transformer_engine/pytorch/ops/dispatch.py | 13 +-- transformer_engine/pytorch/ops/fuser.py | 55 +++--------- 6 files changed, 113 insertions(+), 156 deletions(-) diff --git a/tests/pytorch/distributed/moe_ep_reference.py b/tests/pytorch/distributed/moe_ep_reference.py index 3fab88a50b..280c4f0d21 100644 --- a/tests/pytorch/distributed/moe_ep_reference.py +++ b/tests/pytorch/distributed/moe_ep_reference.py @@ -123,7 +123,9 @@ def _validate_storage(self) -> None: expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") data_shape = self.logical_shape if self.data.dtype != expected_dtype: - raise TypeError(f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}") + raise TypeError( + f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}" + ) else: expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) @@ -298,7 +300,9 @@ def _decode_tensor( ) -> torch.Tensor: if isinstance(tensor, BlockScaledTensor): if tensor.logical_shape != expected_shape: - raise ValueError(f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}") + raise ValueError( + f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}" + ) if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") return tensor.dequantize() @@ -357,11 +361,15 @@ def __init__( ep_size, ep_rank = 1, 0 else: if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError("ep_group requires an initialized torch.distributed process group") + raise RuntimeError( + "ep_group requires an initialized torch.distributed process group" + ) ep_size = dist.get_world_size(ep_group) ep_rank = dist.get_rank(ep_group) if num_experts % ep_size != 0: - raise ValueError(f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})") + raise ValueError( + f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})" + ) self.num_experts = num_experts self.hidden_size = hidden_size @@ -378,10 +386,18 @@ def __init__( self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) self.generate_c = bool(generate_c) - for name, fmt in (("output_format", self.output_format), ("combine_format", self.combine_format)): - required_multiple = 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + for name, fmt in ( + ("output_format", self.output_format), + ("combine_format", self.combine_format), + ): + required_multiple = ( + 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + ) if hidden_size % required_multiple != 0: - raise ValueError(f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for {name}={fmt.value}") + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for" + f" {name}={fmt.value}" + ) def __repr__(self) -> str: return ( @@ -456,7 +472,9 @@ def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") order = torch.argsort(destination, stable=True) - send_counts_tensor = torch.bincount(destination.index_select(0, order), minlength=self.ep_size).to(torch.int64) + send_counts_tensor = torch.bincount( + destination.index_select(0, order), minlength=self.ep_size + ).to(torch.int64) recv_counts_tensor = self._exchange_counts(send_counts_tensor) return _DispatchPlan( send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), @@ -505,7 +523,13 @@ def _run_local_experts( output.index_copy_(0, positions, expert_output) fc1_c = None if fc1_c_rows is not None: - fc1_c = torch.cat(fc1_c_rows) if fc1_c_rows else torch.empty((0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device) + fc1_c = ( + torch.cat(fc1_c_rows) + if fc1_c_rows + else torch.empty( + (0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device + ) + ) return output, fc1_c def __call__( @@ -543,13 +567,17 @@ def __call__( if tuple(topk_idx.shape) != route_shape: raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") if tuple(topk_weights.shape) != route_shape: - raise ValueError(f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}") + raise ValueError( + f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}" + ) if topk_idx.dtype not in (torch.int32, torch.int64): raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") if not topk_weights.is_floating_point(): raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: - raise ValueError(f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}") + raise ValueError( + f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}" + ) device = _tensor_device(activation) inputs = { @@ -602,7 +630,11 @@ def __call__( # Stable sort by local expert reproduces the fc1_c row order # (grouped by expert; source order preserved within each group). fc1_c_order = torch.argsort(recv_expert, stable=True) - route_metadata = torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1).index_select(0, fc1_c_order).to(torch.int32) + route_metadata = ( + torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) + .index_select(0, fc1_c_order) + .to(torch.int32) + ) # recv rows are ordered by source rank, then that source's token-major # route order, so the per-expert position grouping below realizes the # documented fc1_c ordering. @@ -660,10 +692,15 @@ def backward( """ if not self.generate_c: - raise RuntimeError("backward requires the operator to be constructed with generate_c=True") + raise RuntimeError( + "backward requires the operator to be constructed with generate_c=True" + ) token_count = topk_idx.shape[0] if tuple(grad_output.shape) != (token_count, self.hidden_size): - raise ValueError(f"grad_output shape must be {(token_count, self.hidden_size)}, got {tuple(grad_output.shape)}") + raise ValueError( + f"grad_output shape must be {(token_count, self.hidden_size)}, got" + f" {tuple(grad_output.shape)}" + ) if not grad_output.is_floating_point(): raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") @@ -688,16 +725,23 @@ def backward( quantized_axis=1, ) if fc1_c.shape != (int(route_metadata.shape[0]), two_i): - raise ValueError(f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got {tuple(fc1_c.shape)}") + raise ValueError( + f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got" + f" {tuple(fc1_c.shape)}" + ) # Re-dispatch the FC1 inputs, router weights, and output gradients # along the identical forward routes. plan = self._dispatch_plan(topk_idx, topk_weights) send_counts, recv_counts = plan.send_counts, plan.recv_counts grad_output_float = grad_output.float() - recv_tokens = self._all_to_all(activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + recv_tokens = self._all_to_all( + activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) - recv_grad = self._all_to_all(grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts) + recv_grad = self._all_to_all( + grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) # route_metadata rows are in fc1_c order; sorting them by # (src_rank, src_token, src_slot) reproduces the receive order, giving @@ -771,10 +815,16 @@ def backward( # Return the route gradients to their source ranks and scatter-add. returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) - grad_activation = torch.zeros((token_count, self.hidden_size), dtype=torch.float32, device=device) + grad_activation = torch.zeros( + (token_count, self.hidden_size), dtype=torch.float32, device=device + ) grad_activation.index_add_(0, plan.send_token_idx, returned_dx) - grad_topk_weights = torch.zeros((token_count * self.top_k,), dtype=torch.float32, device=device) - grad_topk_weights.index_copy_(0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw) + grad_topk_weights = torch.zeros( + (token_count * self.top_k,), dtype=torch.float32, device=device + ) + grad_topk_weights.index_copy_( + 0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw + ) return ( grad_activation, grad_fc1, diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 2c81782a0e..a69418088c 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -584,13 +584,9 @@ def zero_grads(): def test_fusible_dispatch_combine_moe(self): """Fusible NCCL EP MoE matches the PyTorch all-to-all reference.""" if not EAGER: - self.skipTest( - "variable grouped-linear splits require eager EP output sizing" - ) + self.skipTest("variable grouped-linear splits require eager EP output sizing") - topk_idx, tokens, topk_weights = _make_moe_inputs( - self.cfg.rank, self.cfg.ep_size - ) + topk_idx, tokens, topk_weights = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) num_experts = NUM_LOCAL_EXPERTS intermediate_dim = HIDDEN_DIM @@ -626,24 +622,14 @@ def test_fusible_dispatch_combine_moe(self): generator.manual_seed(1234 + self.cfg.rank) with torch.no_grad(): for expert_idx in range(num_experts): - getattr(fc1, f"weight{expert_idx}").uniform_( - -0.1, 0.1, generator=generator - ) - getattr(fc2, f"weight{expert_idx}").uniform_( - -0.1, 0.1, generator=generator - ) + getattr(fc1, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) + getattr(fc2, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) ref_fc1_weights = torch.stack( - [ - getattr(fc1, f"weight{idx}").detach().transpose(0, 1) - for idx in range(num_experts) - ] + [getattr(fc1, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] ) ref_fc2_weights = torch.stack( - [ - getattr(fc2, f"weight{idx}").detach().transpose(0, 1) - for idx in range(num_experts) - ] + [getattr(fc2, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] ) reference = MoeEpReference( num_experts=self.cfg.num_experts, @@ -688,12 +674,8 @@ def test_fusible_dispatch_combine_moe(self): test_output.backward(grad_output) torch.cuda.synchronize() - torch.testing.assert_close( - test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2 - ) - torch.testing.assert_close( - test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2 - ) + torch.testing.assert_close(test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2) torch.testing.assert_close( test_topk_weights.grad, ref_grads[3], @@ -718,13 +700,9 @@ def test_fusible_dispatch_combine_moe(self): def test_fusible_dispatch_combine_moe(self): """Fusible NCCL EP MoE matches the PyTorch all-to-all reference.""" if not EAGER: - self.skipTest( - "variable grouped-linear splits require eager EP output sizing" - ) + self.skipTest("variable grouped-linear splits require eager EP output sizing") - topk_idx, tokens, topk_weights = _make_moe_inputs( - self.cfg.rank, self.cfg.ep_size - ) + topk_idx, tokens, topk_weights = _make_moe_inputs(self.cfg.rank, self.cfg.ep_size) num_experts = NUM_LOCAL_EXPERTS intermediate_dim = HIDDEN_DIM @@ -760,24 +738,14 @@ def test_fusible_dispatch_combine_moe(self): generator.manual_seed(1234 + self.cfg.rank) with torch.no_grad(): for expert_idx in range(num_experts): - getattr(fc1, f"weight{expert_idx}").uniform_( - -0.1, 0.1, generator=generator - ) - getattr(fc2, f"weight{expert_idx}").uniform_( - -0.1, 0.1, generator=generator - ) + getattr(fc1, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) + getattr(fc2, f"weight{expert_idx}").uniform_(-0.1, 0.1, generator=generator) ref_fc1_weights = torch.stack( - [ - getattr(fc1, f"weight{idx}").detach().transpose(0, 1) - for idx in range(num_experts) - ] + [getattr(fc1, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] ) ref_fc2_weights = torch.stack( - [ - getattr(fc2, f"weight{idx}").detach().transpose(0, 1) - for idx in range(num_experts) - ] + [getattr(fc2, f"weight{idx}").detach().transpose(0, 1) for idx in range(num_experts)] ) reference = MoeEpReference( num_experts=self.cfg.num_experts, @@ -822,12 +790,8 @@ def test_fusible_dispatch_combine_moe(self): test_output.backward(grad_output) torch.cuda.synchronize() - torch.testing.assert_close( - test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2 - ) - torch.testing.assert_close( - test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2 - ) + torch.testing.assert_close(test_output.float(), ref_output.float(), atol=5e-2, rtol=5e-2) + torch.testing.assert_close(test_tokens.grad.float(), ref_grads[0], atol=5e-2, rtol=5e-2) torch.testing.assert_close( test_topk_weights.grad, ref_grads[3], diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 9cecd4531b..dba1cab263 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -523,9 +523,7 @@ def fuser_forward( ): m_splits, probs = basic_op_extra_inputs[0] # Stub row-id map: real Dispatch would emit permute indices. - routing_map = torch.arange( - input_.size(0), device=input_.device, dtype=torch.int64 - ) + routing_map = torch.arange(input_.size(0), device=input_.device, dtype=torch.int64) return input_, [(m_splits, probs, routing_map)] def fuser_backward( @@ -568,9 +566,7 @@ def fuser_forward( if routing_map is None: raise RuntimeError("FakeCombine expected routing_map channel") if int(routing_map.numel()) != int(input_.size(0)): - raise RuntimeError( - "FakeCombine routing_map length does not match tokens" - ) + raise RuntimeError("FakeCombine routing_map length does not match tokens") return input_, [()] def fuser_backward( @@ -583,9 +579,7 @@ def fuser_backward( del basic_op_grad_extra_outputs return grad_output, [()], [(None,)] - split_sizes = torch.tensor([8, 16, 8, 8], dtype=torch.int64, device=device)[ - :group_size - ] + split_sizes = torch.tensor([8, 16, 8, 8], dtype=torch.int64, device=device)[:group_size] num_tokens = int(split_sizes.sum()) in_shape = (num_tokens, hidden_size) @@ -855,9 +849,7 @@ def op_backward(self, *args, **kwargs): def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): return input_, [(input_, input_)] - def fuser_backward( - self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs - ): + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): g0, g1 = basic_op_grad_extra_outputs[0] grad_extra = torch.zeros_like(grad_output) if g0 is not None: diff --git a/transformer_engine/pytorch/ops/combine.py b/transformer_engine/pytorch/ops/combine.py index 19f78f7356..4d4740c517 100644 --- a/transformer_engine/pytorch/ops/combine.py +++ b/transformer_engine/pytorch/ops/combine.py @@ -25,9 +25,7 @@ def _validate_grad_buffer( if tensor is None: return None if tuple(tensor.shape) != shape: - raise ValueError( - f"grad_out shape {tuple(tensor.shape)} does not match {shape}." - ) + raise ValueError(f"grad_out shape {tuple(tensor.shape)} does not match {shape}.") if tensor.dtype is not dtype: raise TypeError(f"grad_out must have dtype {dtype}, got {tensor.dtype}.") if tensor.device != device: @@ -46,15 +44,11 @@ class Combine(BasicOperation): same :class:`EpBuffer`. """ - def __init__( - self, buffer: EpBuffer, *, num_local_tokens: Optional[int] = None - ) -> None: + def __init__(self, buffer: EpBuffer, *, num_local_tokens: Optional[int] = None) -> None: super().__init__() self.buffer = buffer self.num_local_tokens = ( - buffer.max_tokens_per_rank - if num_local_tokens is None - else int(num_local_tokens) + buffer.max_tokens_per_rank if num_local_tokens is None else int(num_local_tokens) ) if self.num_local_tokens < 0: raise ValueError("num_local_tokens must be non-negative.") @@ -70,9 +64,7 @@ def op_forward( ) -> torch.Tensor: del prev_op_grad_output_quantizer, next_op_input_quantizer if input_.dtype is not torch.bfloat16: - raise NotImplementedError( - f"NCCL EP requires BF16 combine input, got {input_.dtype}." - ) + raise NotImplementedError(f"NCCL EP requires BF16 combine input, got {input_.dtype}.") if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: raise ValueError( f"Combine input must have shape (R, {self.buffer.hidden_dim}), " @@ -114,14 +106,8 @@ def op_forward( dtype=input_.dtype, device=input_.device, ) - if ( - self.buffer.zero_copy - and grad_out is not None - and not is_symm_backed(grad_out) - ): - raise ValueError( - "zero-copy Combine grad_out must be symmetric-memory-backed." - ) + if self.buffer.zero_copy and grad_out is not None and not is_symm_backed(grad_out): + raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") ctx.grad_out = grad_out ctx.input_shape = tuple(input_.shape) ctx.input_dtype = input_.dtype @@ -150,4 +136,3 @@ def op_backward( grad_input, ) return grad_input, () - diff --git a/transformer_engine/pytorch/ops/dispatch.py b/transformer_engine/pytorch/ops/dispatch.py index ab89f6e130..9762ed76f7 100644 --- a/transformer_engine/pytorch/ops/dispatch.py +++ b/transformer_engine/pytorch/ops/dispatch.py @@ -73,9 +73,7 @@ def fuser_forward( kwargs = basic_op_kwargs[0] if input_.dtype is not torch.bfloat16: - raise NotImplementedError( - f"NCCL EP requires BF16 dispatch input, got {input_.dtype}." - ) + raise NotImplementedError(f"NCCL EP requires BF16 dispatch input, got {input_.dtype}.") if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: raise ValueError( f"Dispatch input must have shape (T, {self.buffer.hidden_dim}), " @@ -97,15 +95,11 @@ def fuser_forward( raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") for name, tensor in (("topk_idx", topk_idx), ("topk_weights", topk_weights)): if tensor.device != input_.device: - raise ValueError( - f"{name} must be on {input_.device}, got {tensor.device}." - ) + raise ValueError(f"{name} must be on {input_.device}, got {tensor.device}.") recv_tokens = kwargs.get("recv_tokens") recv_topk_weights = kwargs.get("recv_topk_weights") - if self.buffer.eager and ( - recv_tokens is not None or recv_topk_weights is not None - ): + if self.buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): raise ValueError( "eager mode sizes dispatch outputs per step and cannot use " "caller-supplied receive buffers" @@ -211,4 +205,3 @@ def fuser_backward( grad_topk_weights, ) return grad_input, [()], [(None, grad_topk_weights)] - diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index cc706ab91f..63d3f4a6be 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -108,9 +108,7 @@ def forward( basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ [None] * op.num_extra_inputs for op in fuser._basic_ops ] - for tensor, (op_idx, input_idx) in zip( - extra_inputs, fuser._external_extra_input_slots - ): + for tensor, (op_idx, input_idx) in zip(extra_inputs, fuser._external_extra_input_slots): basic_op_extra_inputs[op_idx][input_idx] = tensor # Apply forward ops @@ -127,9 +125,7 @@ def forward( # consumer, leave the consumer slot unset so the fused op can # wire the channel itself for idx in basic_op_idxs: - for input_idx, source in enumerate( - fuser._basic_op_extra_input_sources[idx] - ): + for input_idx, source in enumerate(fuser._basic_op_extra_input_sources[idx]): if source is None: continue producer_idx, output_idx = source @@ -141,10 +137,7 @@ def forward( raise RuntimeError( f"Extra tensor channel producer op {producer_idx} has not run" ) - if ( - output_idx >= len(producer_outputs) - or producer_outputs[output_idx] is None - ): + if output_idx >= len(producer_outputs) or producer_outputs[output_idx] is None: raise RuntimeError( f"Extra tensor channel producer op {producer_idx} " f"({type(fuser._basic_ops[producer_idx]).__name__}) " @@ -154,9 +147,7 @@ def forward( f"input {input_idx}" ) basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx] - op_extra_inputs = [ - tuple(basic_op_extra_inputs[idx]) for idx in basic_op_idxs - ] + op_extra_inputs = [tuple(basic_op_extra_inputs[idx]) for idx in basic_op_idxs] prev_op_idx = basic_op_idxs[0] - 1 prev_op = fuser._basic_ops[prev_op_idx] if prev_op_idx >= 0 else None prev_op_grad_output_quantizer = None @@ -290,9 +281,7 @@ def backward( basic_op_grad_extra_outputs: list[list[Optional[torch.Tensor]]] = [ [None] * op.num_extra_outputs for op in basic_ops ] - for grad, (op_idx, output_idx) in zip( - grad_extra_outputs, external_extra_output_slots - ): + for grad, (op_idx, output_idx) in zip(grad_extra_outputs, external_extra_output_slots): basic_op_grad_extra_outputs[op_idx][output_idx] = grad # Apply backward ops @@ -310,13 +299,9 @@ def backward( # Backward op. Supply gradients accumulated from every consumer of # each internal channel. for idx in basic_op_idxs: - for output_idx, channel in enumerate( - basic_op_extra_output_channels[idx] - ): + for output_idx, channel in enumerate(basic_op_extra_output_channels[idx]): if channel is not None: - basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get( - channel - ) + basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get(channel) op_grad_extra_outputs = [ tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs ] @@ -341,9 +326,7 @@ def backward( continue channel = basic_op_extra_output_channels[producer_idx][output_idx] previous_grad = channel_grads.get(channel) - channel_grads[channel] = ( - grad if previous_grad is None else previous_grad + grad - ) + channel_grads[channel] = grad if previous_grad is None else previous_grad + grad # Flatten list of parameter gradients grad_params_flat = [] @@ -456,13 +439,9 @@ def __init__( f"Extra tensor channel {channel!r} consumed by op {op_idx} " f"({type(op).__name__}) has no earlier producer" ) - self._basic_op_extra_input_sources[op_idx][input_idx] = channel_producers[ - channel - ] + self._basic_op_extra_input_sources[op_idx][input_idx] = channel_producers[channel] consumed_channels.add(channel) - for output_idx, channel in enumerate( - self._basic_op_extra_output_channels[op_idx] - ): + for output_idx, channel in enumerate(self._basic_op_extra_output_channels[op_idx]): if channel is None: self._external_extra_output_slots.append((op_idx, output_idx)) continue @@ -497,12 +476,10 @@ def __init__( raise ValueError( f"Extra input {input_idx} of op {op_idx} " f"({type(op).__name__}) is bound to channel {channel!r} " - f"but has no producer" + "but has no producer" ) producer_idx, output_idx = source - producer_channel = self._basic_op_extra_output_channels[producer_idx][ - output_idx - ] + producer_channel = self._basic_op_extra_output_channels[producer_idx][output_idx] if producer_channel != channel: raise ValueError( f"Extra input {input_idx} of op {op_idx} " @@ -602,9 +579,7 @@ def maybe_fuse_ops( first_op_requiring_backward = self._num_basic_ops for op_idx in range(self._num_basic_ops): op_inputs = itertools.chain(self._basic_op_params[op_idx], extra_inputs[op_idx]) - if any( - tensor is not None and tensor.requires_grad for tensor in op_inputs - ): + if any(tensor is not None and tensor.requires_grad for tensor in op_inputs): first_op_requiring_backward = op_idx break @@ -694,9 +669,7 @@ def __call__( basic_op_extra_inputs: list[list[Optional[torch.Tensor]]] = [ [None] * op.num_extra_inputs for op in self._basic_ops ] - for tensor, (op_idx, input_idx) in zip( - extra_inputs, self._external_extra_input_slots - ): + for tensor, (op_idx, input_idx) in zip(extra_inputs, self._external_extra_input_slots): basic_op_extra_inputs[op_idx][input_idx] = tensor # Get environment state