diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index dd17191e58..b7019b6237 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -151,6 +151,153 @@ arguments and the extra outputs will be returned. the block has been split into two sections, each with one branching operation. +Extra tensor channels +""""""""""""""""""""" + +Branching operations can also route their extra inputs and outputs within +the same ``Sequential`` via named channels. Extra output tensors with a +specified channel can be consumed by later operations in the same ``Sequential`` +and may optionally be returned to the caller. Extra input tensors with a specified +channel are accessed internally instead of being provided as arguments +to ``Sequential``. + +With a channel, the residual block above can be expressed using one +``Sequential``: + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + make_residual = te.ops.MakeExtraOutput() + add_residual = te.ops.AddExtraInput() + make_residual.set_extra_output_channel( + 0, "residual", output_to_caller=False + ) + add_residual.set_extra_input_channel(0, "residual") + + block = te.ops.Sequential( + te.ops.LayerNorm(4096), + make_residual, + te.ops.Linear(4096, 28672), + te.ops.SwiGLU(), + te.ops.Linear(14336, 4096), + add_residual, + ) + + # The residual is routed internally and omitted from the public outputs. + x = torch.randn(16384, 4096, device="cuda") + y = block(x) + +Channels are also useful for mixture-of-experts blocks. The following +example assumes custom ``Dispatch`` and ``Combine`` basic operations. +``Dispatch`` has one public extra input containing router probabilities +and three extra outputs: split sizes, token probabilities, and a +routing map. ``Combine`` consumes the routing map. + +.. code-block:: python + + import transformer_engine.pytorch as te + from my_ops import Dispatch, Combine + + num_experts = 8 + hidden_size = 4096 + ffn_size = 14336 + + dispatch = Dispatch(num_experts) + fc1 = te.ops.GroupedLinear( + num_experts, hidden_size, 2 * ffn_size, bias=False + ) + activation = te.ops.ScaledSwiGLU() + fc2 = te.ops.GroupedLinear( + num_experts, ffn_size, hidden_size, bias=False + ) + combine = Combine(num_experts) + + # Dispatch extra outputs: + # 0: split sizes, 1: token probabilities, 2: routing map + dispatch.set_extra_output_channel( + 0, "m_splits", output_to_caller=False + ) + dispatch.set_extra_output_channel( + 1, "probs", output_to_caller=False + ) + dispatch.set_extra_output_channel( + 2, "routing_map", output_to_caller=False + ) + + 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") + + moe = te.ops.Sequential(dispatch, fc1, activation, fc2, combine) + + # Dispatch's extra input has no channel, so the caller passes router_probs. + # Channels supply all later extra inputs internally. The channel outputs + # are not returned because output_to_caller=False. + y = moe(x, router_probs) + +Channels cannot connect operations in different ``OperationFuser`` +instances. In particular, an ordinary PyTorch module inside a +``Sequential`` splits the fusible operations on either side into +separate fusers. The following channel connection is therefore not +supported: + +.. code-block:: python + + make_residual = te.ops.MakeExtraOutput() + add_residual = te.ops.AddExtraInput() + make_residual.set_extra_output_channel(0, "residual") + add_residual.set_extra_input_channel(0, "residual") + + block = te.ops.Sequential( + make_residual, + torch.nn.Identity(), # Splits the operations into separate fusers. + add_residual, + ) + +Use the public extra output and extra input interfaces, as in the +two-``Sequential`` example above, when the producer and consumer cannot +be placed in the same ``OperationFuser``. + +The following conditions apply to extra tensor channels: + +- Every named extra input must have a matching producer earlier in the + same fuser. Leave an extra input unnamed when the caller should provide it. +- An output channel name has at most one producer, but its output may + fan out to multiple consumers. +- A named output does not require a consumer. It is returned as a public + extra output by default. +- A channel is scoped to one ``OperationFuser``. In a ``Sequential``, + ordinary PyTorch modules split adjacent fusible operations into + separate fusers, and channels cannot cross that boundary. +- The caller passes unnamed extra inputs. Named, channel-connected extra + input slots do not appear in the ``Sequential`` arguments. +- ``set_extra_output_channel`` accepts ``output_to_caller`` (``True`` by + default). Public extra outputs are returned in their original + basic-operation and slot order. Gradients supplied for a returned output + are combined with gradients from its internal channel consumers. An + unused extra-output gradient may be ``None`` and is treated as zero. +- Set ``output_to_caller=False`` for a channel tensor that should remain + internal. Removing a channel binding with ``channel=None`` restores that + output as public. +- Channel bindings are captured when an ``OperationFuser`` (or the + fusers inside a ``Sequential``) is constructed. Rebinding a channel with + ``set_extra_input_channel`` / ``set_extra_output_channel`` afterward + invalidates any fuser that captured it, so a new ``OperationFuser`` or + ``Sequential`` must be constructed before the next forward pass. + +Channel-connected basic operations may still be replaced by registered +``FusedOperation`` implementations. If a fused operation contains both +the producer and consumer of a channel, its ``fuser_forward`` and +``fuser_backward`` implementations are responsible for routing the +tensor and its gradient between those basic operations. For a non-public +channel fully owned by one forward fusion, ``fuser_forward`` may return +``None`` in the corresponding basic-operation output slot. A tensor is +still required when the output is public or when a consumer is outside +that forward fusion. + Developer guide --------------- diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 067f75d725..d8aab26e11 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -11,6 +11,7 @@ import torch import torch.distributed as dist +from transformer_engine.pytorch import ops as te_ops from transformer_engine.pytorch.ep import ( EpBuffer, ep_bootstrap, @@ -23,6 +24,7 @@ _ep_combine_raw, _ep_dispatch_raw, ) +from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" @@ -37,6 +39,7 @@ HIDDEN_DIM = 32 TOP_K = 2 TOKENS_PER_RANK = 4 +INTERMEDIATE_DIM = 16 def _zero_copy_test_include(fn): @@ -121,6 +124,31 @@ def _make_identity_inputs(rank, ep_size, device="cuda"): ) +def _make_moe_inputs(rank, ep_size, device="cuda"): + """Deterministic BF16 activations and FP32 top-k router weights.""" + generator = torch.Generator(device=device) + generator.manual_seed(2026 + rank) + tokens = ( + torch.randn( + TOKENS_PER_RANK, + HIDDEN_DIM, + generator=generator, + dtype=torch.float32, + device=device, + ) + * 0.25 + ).to(torch.bfloat16) + router_logits = torch.randn( + TOKENS_PER_RANK, + ep_size * NUM_LOCAL_EXPERTS, + generator=generator, + dtype=torch.float32, + device=device, + ) + topk_logits, topk_idx = torch.topk(router_logits, TOP_K, dim=-1) + return topk_idx, tokens, torch.softmax(topk_logits, dim=-1) + + class _Cfg: rank: int world_size: int @@ -232,6 +260,42 @@ def _moe_step(self, buffer, topk_idx, tokens, w): expert_out = self._weighted(recv_t, recv_w_out) return ep_combine(buffer, expert_out) + def _make_moe_model(self, *, fuse_ops=True): + """Build a BF16 EP MoE Sequential. + + With ``fuse_ops=True``, dispatch routing extras stay internal + (``output_to_caller=False``) so :class:`FusedMoeEp` can claim the + sequence. With ``fuse_ops=False``, those extras are returned to the + caller, which blocks fusion. + """ + buffer = self._make_buffer() + dispatch = te_ops.Dispatch(buffer) + fc1 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + HIDDEN_DIM, + 2 * INTERMEDIATE_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + INTERMEDIATE_DIM, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) + + dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=not fuse_ops) + dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=not fuse_ops) + fc1.set_extra_input_channel(0, "tokens_per_expert") + activation.set_extra_input_channel(0, "routing_weights") + fc2.set_extra_input_channel(0, "tokens_per_expert") + return te_ops.Sequential(dispatch, fc1, activation, fc2, combine), fc1, fc2 + # Prepare @_eager_test_include @@ -419,6 +483,100 @@ def test_caller_provides_grad_expert_out(self): # the caller-owned buffer was used as the combine-bwd scatter target self.assertGreater(gbuf.abs().sum().item(), 0.0) + @_eager_test_include + def test_bf16_moe_sequential_fusion(self): + """Reference-backed fusion matches the unfused BF16 EP MoE sequence. + + ``fuse_ops=True`` keeps dispatch routing extras internal so fusion + fires; ``fuse_ops=False`` returns them to the caller and blocks it. + """ + if not EAGER: + self.skipTest("variable-size reference comparison requires eager EP mode") + + fused, fused_fc1, fused_fc2 = self._make_moe_model(fuse_ops=True) + unfused, unfused_fc1, unfused_fc2 = self._make_moe_model(fuse_ops=False) + generator = torch.Generator(device=self.cfg.device) + generator.manual_seed(3100 + self.cfg.rank) + with torch.no_grad(): + for fused_op, unfused_op in ((fused_fc1, unfused_fc1), (fused_fc2, unfused_fc2)): + for expert in range(NUM_LOCAL_EXPERTS): + weight = ( + torch.randn( + getattr(fused_op, f"weight{expert}").shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + getattr(fused_op, f"weight{expert}").copy_(weight) + getattr(unfused_op, f"weight{expert}").copy_(weight) + + topk_idx, tokens, topk_weights = _make_moe_inputs( + self.cfg.rank, + self.cfg.ep_size, + self.cfg.device, + ) + fused_tokens = tokens.detach().clone().requires_grad_(True) + unfused_tokens = tokens.detach().clone().requires_grad_(True) + fused_topk_weights = topk_weights.detach().clone().requires_grad_(True) + unfused_topk_weights = topk_weights.detach().clone().requires_grad_(True) + + fused_out = fused( + fused_tokens, + topk_idx, + fused_topk_weights, + ) + unfused_out, tokens_per_expert, recv_topk_weights = unfused( + unfused_tokens, + topk_idx, + unfused_topk_weights, + ) + + fused_forward_ops = fused._module_groups[0]._forward_ops + unfused_forward_ops = unfused._module_groups[0]._forward_ops + self.assertEqual(len(fused_forward_ops), 1) + self.assertIsInstance(fused_forward_ops[0][0], FusedMoeEp) + self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in unfused_forward_ops)) + self.assertIsInstance(fused_out, torch.Tensor) + self.assertEqual(fused_out.dtype, torch.bfloat16) + self.assertEqual(unfused_out.dtype, torch.bfloat16) + self.assertEqual(tokens_per_expert.shape, (NUM_LOCAL_EXPERTS,)) + self.assertEqual(tokens_per_expert.dtype, torch.int64) + self.assertEqual(recv_topk_weights.dtype, torch.float32) + + dy = ( + torch.randn( + fused_out.shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + fused_out.backward(dy) + unfused_out.backward(dy) + torch.cuda.synchronize() + + # The two BF16 paths use different grouped-GEMM and reduction orders. + # Observed forward abs error reaches ~3e-5 on near-zero values where + # rtol does not apply, so atol must stay above that (1e-5 is too tight). + tolerances = {"rtol": 1.6e-2, "atol": 5e-5} + torch.testing.assert_close(fused_out, unfused_out, **tolerances) + torch.testing.assert_close(fused_tokens.grad, unfused_tokens.grad, **tolerances) + torch.testing.assert_close( + fused_topk_weights.grad, + unfused_topk_weights.grad, + **tolerances, + ) + for fused_op, unfused_op in ((fused_fc1, unfused_fc1), (fused_fc2, unfused_fc2)): + for expert in range(NUM_LOCAL_EXPERTS): + fused_grad = getattr(fused_op, f"weight{expert}").grad + unfused_grad = getattr(unfused_op, f"weight{expert}").grad + self.assertEqual(fused_grad.dtype, torch.bfloat16) + self.assertEqual(unfused_grad.dtype, torch.bfloat16) + torch.testing.assert_close(fused_grad, unfused_grad, **tolerances) + @_zero_copy_test_include def test_zero_copy_pool_auto_alloc(self): """Zero-copy with recv/grad left None: ep_dispatch/ep_combine allocate their IO diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 66857d8125..36c02a3154 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 ( @@ -437,6 +438,535 @@ def test_extra_tensors(self, size: int = 16) -> None: torch.testing.assert_close(x4, x4_orig + x3) +class _DualExtraOutput(te_ops.BasicOperation): + """Test helper: one op with two scaled extra outputs.""" + + num_extra_outputs = 2 + + def __init__(self, scales: tuple[float, float] = (1.0, 1.0)) -> None: + super().__init__() + self._scales = scales + + 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): + del basic_op_ctxs, basic_op_extra_inputs + s0, s1 = self._scales + return input_, [(s0 * input_, s1 * input_)] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + del basic_op_ctxs + s0, s1 = self._scales + g0, g1 = basic_op_grad_extra_outputs[0] + grad_extra = torch.zeros_like(grad_output) + if g0 is not None: + grad_extra = grad_extra + s0 * g0 + if g1 is not None: + grad_extra = grad_extra + s1 * g1 + return grad_output + grad_extra, [()], [()] + + +class TestExtraTensorChannels: + """Error handling and grad coverage for named extra-tensor channels.""" + + @pytest.mark.parametrize("with_extra_grad", (True, False)) + def test_internal_residual_connection( + self, + with_extra_grad: bool, + size: int = 16, + ) -> None: + """A channel can keep a residual connection inside a Sequential.""" + residual = te_ops.MakeExtraOutput() + body = te_ops.Bias(size=size, device="cpu") + add_residual = te_ops.AddExtraInput() + residual.set_extra_output_channel(0, "residual") + add_residual.set_extra_input_channel(0, "residual") + + model = te_ops.Sequential(residual, body, add_residual) + x = torch.rand((size,), requires_grad=True) + y, residual_out = model(x) + + torch.testing.assert_close(y, 2 * x + body.bias) + torch.testing.assert_close(residual_out, x) + dy = torch.rand_like(y) + if with_extra_grad: + dresidual = torch.rand_like(residual_out) + torch.autograd.backward((y, residual_out), (dy, dresidual)) + expected_dx = 2 * dy + dresidual + else: + y.backward(dy) + expected_dx = 2 * dy + torch.testing.assert_close(x.grad, expected_dx) + torch.testing.assert_close(body.bias.grad, dy) + + @pytest.mark.parametrize("output_to_caller", (True, False)) + def test_unconsumed_extra_output( + self, + output_to_caller: bool, + size: int = 16, + ) -> None: + """Unused public or hidden-unconsumed MakeExtraOutput treats None grad as zero.""" + residual = te_ops.MakeExtraOutput() + body = te_ops.Bias(size=size, device="cpu") + if not output_to_caller: + residual.set_extra_output_channel(0, "unused", output_to_caller=False) + + model = te_ops.Sequential(residual, body) + x = torch.rand((size,), requires_grad=True) + outputs = model(x) + if output_to_caller: + y, residual_out = outputs + torch.testing.assert_close(residual_out, x) + else: + assert isinstance(outputs, torch.Tensor) + y = outputs + + torch.testing.assert_close(y, x + body.bias) + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, dy) + torch.testing.assert_close(body.bias.grad, dy) + + @pytest.mark.parametrize("fusion_kind", ("forward", "backward", "forward_backward")) + @pytest.mark.parametrize("output_to_caller", (True, False)) + def test_fused_internal_residual_connection( + self, + fusion_kind: str, + output_to_caller: bool, + size: int = 16, + ) -> None: + """Forward, backward, and joint fusions can own an internal channel.""" + + class FusedResidual(te_ops.FusedOperation): + """Fuse MakeExtraOutput, Bias, and AddExtraInput.""" + + _enabled = True + + def __init__(self, residual, body, add_residual) -> None: + super().__init__((residual, body, add_residual)) + + def fuser_forward( + self, + basic_op_ctxs, + input_, + *, + basic_op_extra_inputs, + **unused, + ): + del basic_op_ctxs + # The consumer slot is internal to this fusion, so the + # OperationFuser deliberately leaves it unset. + assert basic_op_extra_inputs[2][0] is None + residual_out = input_ if output_to_caller else None + return 2 * input_ + self.basic_ops[1].bias, [(residual_out,), (), ()] + + def fuser_backward( + self, + basic_op_ctxs, + grad_output, + *, + basic_op_grad_extra_outputs, + ): + del basic_op_ctxs + # The fusion owns the internal residual edge. The fuser also + # supplies a gradient when the residual is a public output. + grad_residual = basic_op_grad_extra_outputs[0][0] + return ( + 2 * grad_output + + (torch.zeros_like(grad_output) if grad_residual is None else grad_residual), + [(), (grad_output,), ()], + [(), (), (grad_output,)], + ) + + def fuse_residual(ops, **unused): + if not FusedResidual._enabled: + return ops + if ( + len(ops) == 3 + and isinstance(ops[0], te_ops.MakeExtraOutput) + and isinstance(ops[1], te_ops.Bias) + and isinstance(ops[2], te_ops.AddExtraInput) + ): + # We want to enable this fusion just for this test. + # Hence disable it after fusing it once in the test. + FusedResidual._enabled = False + return [FusedResidual(*ops)] + return ops + + residual = te_ops.MakeExtraOutput() + body = te_ops.Bias(size=size, device="cpu") + add_residual = te_ops.AddExtraInput() + residual.set_extra_output_channel( + 0, + "residual", + output_to_caller=output_to_caller, + ) + add_residual.set_extra_input_channel(0, "residual") + model = te_ops.Sequential(residual, body, add_residual) + + if fusion_kind == "forward": + te_ops.register_forward_fusion(fuse_residual, prepend=True) + elif fusion_kind == "backward": + te_ops.register_backward_fusion(fuse_residual, prepend=True) + else: + te_ops.register_forward_backward_fusion(fuse_residual, prepend=True) + x = torch.rand((size,), requires_grad=True) + outputs = model(x) + if output_to_caller: + y, residual_out = outputs + else: + assert isinstance(outputs, torch.Tensor) + y = outputs + + forward_ops = model._module_groups[0]._forward_ops + backward_ops = model._module_groups[0]._backward_ops + if fusion_kind in ("forward", "forward_backward"): + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], FusedResidual) + else: + assert len(forward_ops) == 3 + if fusion_kind in ("backward", "forward_backward"): + assert len(backward_ops) == 1 + assert isinstance(backward_ops[0][0], FusedResidual) + else: + assert len(backward_ops) == 3 + if fusion_kind == "forward_backward": + assert backward_ops[0][0] is forward_ops[0][0] + torch.testing.assert_close(y, 2 * x + body.bias) + dy = torch.rand_like(y) + if output_to_caller: + dresidual = torch.rand_like(residual_out) + torch.autograd.backward((y, residual_out), (dy, dresidual)) + expected_dx = 2 * dy + dresidual + else: + y.backward(dy) + expected_dx = 2 * dy + torch.testing.assert_close(x.grad, expected_dx) + torch.testing.assert_close(body.bias.grad, dy) + + 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, route = model(x) + + # Main path: x -> x + route -> x + route + route. + torch.testing.assert_close(y, 3 * x) + torch.testing.assert_close(route, x) + dy = torch.rand_like(y) + droute = torch.rand_like(route) + torch.autograd.backward((y, route), (dy, droute)) + # The channel fan-out contributes two independent gradient paths. + torch.testing.assert_close(x.grad, 3 * dy + droute) + + # Internal slots are unavailable before forward, so grad discovery + # must tolerate them when no public input requires gradients. + x_no_grad = x.detach() + y_no_grad, route_no_grad = model(x_no_grad) + torch.testing.assert_close(y_no_grad, 3 * x_no_grad) + torch.testing.assert_close(route_no_grad, x_no_grad) + + def test_internal_extra_tensor_channel_can_be_hidden(self, size: int = 16) -> None: + """A non-public channel still propagates forward and backward.""" + producer = te_ops.MakeExtraOutput() + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route", output_to_caller=False) + 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) + assert isinstance(y, torch.Tensor) + torch.testing.assert_close(y, 3 * x) + + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, 3 * dy) + + 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, route = model(x, extra) + + torch.testing.assert_close(y, 2 * x + extra) + torch.testing.assert_close(route, x) + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, 2 * dy) + torch.testing.assert_close(extra.grad, dy) + + def test_named_extra_input_requires_producer(self) -> None: + """A named input cannot fall back to a caller-provided tensor.""" + consumer = te_ops.AddExtraInput() + consumer.set_extra_input_channel(0, "missing") + with pytest.raises(ValueError, match="has no 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] + with pytest.raises(TypeError, match="output_to_caller must be a bool"): + producer.set_extra_output_channel( + 0, + "route", + output_to_caller=1, # type: ignore[arg-type] + ) + + def test_extra_channel_change_requires_new_sequential(self, size: int = 16) -> None: + """Rebinding a channel invalidates the fuser that captured it.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer) + + x = torch.rand((size,)) + y, route = model(x) + torch.testing.assert_close(y, 2 * x) + torch.testing.assert_close(route, x) + + consumer.set_extra_input_channel(0, None) + extra = torch.rand_like(x) + with pytest.raises(RuntimeError, match="Construct a new OperationFuser"): + model(x, extra) + + model = te_ops.Sequential(producer, consumer) + y, route = model(x, extra) + torch.testing.assert_close(y, x + extra) + torch.testing.assert_close(route, x) + + producer.set_extra_output_channel(0, "route", output_to_caller=False) + with pytest.raises(RuntimeError, match="Construct a new OperationFuser"): + model(x, extra) + + @pytest.mark.parametrize("layout", ("two_ops", "same_op")) + def test_duplicate_extra_output_channel_names(self, layout: str) -> None: + """A channel name may have at most one producer, across ops or slots.""" + consumer = te_ops.AddExtraInput() + consumer.set_extra_input_channel(0, "route") + if layout == "two_ops": + producer1 = te_ops.MakeExtraOutput() + producer2 = te_ops.MakeExtraOutput() + producer1.set_extra_output_channel(0, "route") + producer2.set_extra_output_channel(0, "route") + ops = [producer1, producer2, consumer] + else: + producer = _DualExtraOutput() + producer.set_extra_output_channel(0, "route") + producer.set_extra_output_channel(1, "route") + ops = [producer, consumer] + with pytest.raises(ValueError, match="multiple producers"): + OperationFuser(ops) + + def test_one_extra_input_has_single_source(self, size: int = 16) -> None: + """Rebinding selects one source and leaves the other output public.""" + 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") + fuser = OperationFuser([producer_a, producer_b, consumer]) + assert fuser._basic_op_extra_input_sources[2] == [(1, 0)] + assert fuser.num_extra_inputs == 0 + + x = torch.rand((size,)) + y, output_a, output_b = fuser(x) + torch.testing.assert_close(y, 2 * x) + torch.testing.assert_close(output_a, x) + torch.testing.assert_close(output_b, x) + + def test_mixed_public_and_hidden_channel_outputs(self, size: int = 16) -> None: + """Only configured public outputs are returned, in slot order.""" + producer = _DualExtraOutput(scales=(2.0, 3.0)) + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "internal", output_to_caller=False) + producer.set_extra_output_channel(1, "public") + consumer.set_extra_input_channel(0, "internal") + model = te_ops.Sequential(producer, consumer) + + x = torch.rand((size,), requires_grad=True) + y, public = model(x) + torch.testing.assert_close(y, 3 * x) + torch.testing.assert_close(public, 3 * x) + + dy = torch.rand_like(y) + dpublic = torch.rand_like(public) + torch.autograd.backward((y, public), (dy, dpublic)) + torch.testing.assert_close(x.grad, 3 * dy + 3 * dpublic) + + @pytest.mark.parametrize("output_to_caller", (True, False)) + def test_fused_op_cannot_omit_required_channel_output( + self, + output_to_caller: bool, + size: int = 16, + ) -> None: + """A fusion must materialize public outputs and cross-fusion channels.""" + + class FusedProducer(te_ops.FusedOperation): + _enabled = True + + def __init__(self, producer) -> None: + super().__init__((producer,)) + + def fuser_forward(self, basic_op_ctxs, input_, **unused): + del basic_op_ctxs + return input_, [(None,)] + + def fuse_producer(ops, **unused): + if ( + FusedProducer._enabled + and len(ops) == 2 + and isinstance(ops[0], te_ops.MakeExtraOutput) + and isinstance(ops[1], te_ops.AddExtraInput) + ): + FusedProducer._enabled = False + return [FusedProducer(ops[0]), ops[1]] + return ops + + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel( + 0, + "route", + output_to_caller=output_to_caller, + ) + consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer) + te_ops.register_forward_fusion(fuse_producer, prepend=True) + + x = torch.rand((size,), requires_grad=True) + error = "is public" if output_to_caller else "outside its forward fusion" + with pytest.raises(RuntimeError, match=error): + model(x) + + def test_fresh_internal_output_preserves_grad_requirement(self) -> None: + """A freshly computed internal channel still receives a consumer gradient.""" + + # A BasicOperation with one extra output that is freshly computed instead of + # retrieved from a previous op's tensor. + class MakeScale(te_ops.BasicOperation): + num_extra_outputs = 1 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("MakeScale uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("MakeScale uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + del basic_op_extra_inputs + basic_op_ctxs[0].save_for_backward(input_) + return input_, [(input_.square().mean(dim=-1),)] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + (input_,) = basic_op_ctxs[0].saved_tensors + grad_scale = basic_op_grad_extra_outputs[0][0] + assert grad_scale is not None + grad_input = grad_output + grad_scale.unsqueeze(-1) * 2 * input_ / input_.size(-1) + return grad_input, [()], [()] + + class ScaleByExtra(te_ops.BasicOperation): + num_extra_inputs = 1 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("ScaleByExtra uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("ScaleByExtra uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + scale = basic_op_extra_inputs[0][0] + ctx = basic_op_ctxs[0] + # Match scaled activations: only compute scale grads when the + # fuser marked this fresh internal channel as requiring grad. + ctx.extra_input_requires_grad = scale.requires_grad + ctx.save_for_backward(input_, scale) + return input_ * scale.unsqueeze(-1), [()] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + del basic_op_grad_extra_outputs + ctx = basic_op_ctxs[0] + input_, scale = ctx.saved_tensors + grad_input = grad_output * scale.unsqueeze(-1) + grad_scale = ( + (grad_output * input_).sum(dim=-1) if ctx.extra_input_requires_grad else None + ) + return grad_input, [()], [(grad_scale,)] + + producer = MakeScale() + consumer = ScaleByExtra() + producer.set_extra_output_channel(0, "scale", output_to_caller=False) + consumer.set_extra_input_channel(0, "scale") + model = te_ops.Sequential(producer, consumer) + + x_ref = torch.randn((5, 8), requires_grad=True) + x_test = x_ref.detach().clone().requires_grad_(True) + scale_ref = x_ref.square().mean(dim=-1) + y_ref = x_ref * scale_ref.unsqueeze(-1) + y_test = model(x_test) + assert isinstance(y_test, torch.Tensor) + torch.testing.assert_close(y_test, y_ref) + + dy = torch.rand_like(y_ref) + y_ref.backward(dy) + y_test.backward(dy) + torch.testing.assert_close(x_test.grad, x_ref.grad) + + class TestFuser: """Tests for operation fusion infrastructure""" @@ -2961,6 +3491,7 @@ def test_backward_activation_bias( @pytest.mark.parametrize("in_shape", ((-1,), (6, 16, -1))) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("zero_centered_gamma", (False, True)) + @pytest.mark.parametrize("with_extra_grad", (True, False)) def test_backward_add_rmsnorm( self, *, @@ -2970,6 +3501,7 @@ def test_backward_add_rmsnorm( device: torch.device = "cuda", eps: float = 0.3, zero_centered_gamma: bool, + with_extra_grad: bool, ) -> None: """Fused backward RMNorm + add""" @@ -3008,7 +3540,10 @@ def test_backward_add_rmsnorm( else: y1_ref = x_ref / torch.sqrt(eps + var_ref) * w_ref y2_ref = x_ref - (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + if with_extra_grad: + (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + else: + (y1_ref * dy1_ref).sum().backward() # Implementation with fusible operations model = te_ops.Sequential( @@ -3025,7 +3560,10 @@ def test_backward_add_rmsnorm( model[1].weight.copy_(w_test) del w_test y1_test, y2_test = model(x_test) - (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + if with_extra_grad: + (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + else: + (y1_test * dy1_test).sum().backward() # Check that backward operations have been fused backward_ops = model._module_groups[0]._backward_ops @@ -3047,6 +3585,7 @@ def test_backward_add_rmsnorm( @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("with_extra_grad", (True, False)) def test_backward_linear_add( self, *, @@ -3056,6 +3595,7 @@ def test_backward_linear_add( device: torch.device = "cuda", quantization: Optional[str], quantized_weight: bool = False, + with_extra_grad: bool, ) -> None: """Backward dgrad GEMM + add""" @@ -3103,7 +3643,10 @@ def test_backward_linear_add( # Plain PyTorch implementation y1_ref = torch.nn.functional.linear(x_ref, w_ref) y2_ref = x_ref - (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + if with_extra_grad: + (y1_ref * dy1_ref + y2_ref * dy2_ref).sum().backward() + else: + (y1_ref * dy1_ref).sum().backward() # Implementation with fusible operations recipe = make_recipe(quantization) @@ -3123,7 +3666,10 @@ def test_backward_linear_add( del w_test with te.autocast(enabled=quantized_compute, recipe=recipe): y1_test, y2_test = model(x_test) - (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + if with_extra_grad: + (y1_test * dy1_test + y2_test * dy2_test).sum().backward() + else: + (y1_test * dy1_test).sum().backward() # Check that backward operations have been fused backward_ops = model._module_groups[0]._backward_ops diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 42f605367a..28c4a49a13 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -21,6 +21,7 @@ __all__ = [ "EpBuffer", "ep_bootstrap", + "get_ep_group", "is_ep_bootstrapped", "ep_finalize", "ep_dispatch", @@ -175,6 +176,11 @@ def is_ep_bootstrapped() -> bool: return _BOOTSTRAPPED +def get_ep_group() -> Optional[dist.ProcessGroup]: + """Return the process group registered by :func:`ep_bootstrap`.""" + return _EP_GROUP + + def ep_finalize() -> None: """Optional explicit EP teardown; idempotent. diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py new file mode 100644 index 0000000000..8a8d58ad0f --- /dev/null +++ b/transformer_engine/pytorch/ep_reference.py @@ -0,0 +1,881 @@ +# 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, + dtype: torch.dtype, +) -> 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(dtype=dtype) + + 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.to(dtype=dtype) + + +def _format_round_trip( + tensor: torch.Tensor, + format: MoeFormat, + *, + dtype: torch.dtype, +) -> torch.Tensor: + if format is MoeFormat.BF16: + return tensor.to(torch.bfloat16).to(dtype) + return quantize_blockwise(tensor, format, axis=-1).dequantize(dtype=dtype) + + +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, + compute_dtype: torch.dtype = torch.float32, + ) -> 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 compute_dtype not in (torch.float32, torch.bfloat16): + raise ValueError( + f"compute_dtype must be torch.float32 or torch.bfloat16, got {compute_dtype}" + ) + + 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) + self.compute_dtype = compute_dtype + + 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" + f" {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}, " + f"compute_dtype={self.compute_dtype})" + ) + + 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=self.compute_dtype, + 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, + dtype=self.compute_dtype, + ) + 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, + dtype=self.compute_dtype, + ) + 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, + dtype=self.compute_dtype, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ) + + 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).to( + dtype=self.compute_dtype + ) + + 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=self.compute_dtype, + 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)``. Activation and expert weight gradients use + ``compute_dtype``. Router weights are cast to ``compute_dtype`` for the + activation scaling path; the returned + ``grad_topk_weights`` remain 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" + 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}") + + 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, + dtype=self.compute_dtype, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, two_i), + quantized_axis=1, + dtype=self.compute_dtype, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ) + 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" + 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.to(dtype=self.compute_dtype) + recv_tokens = self._all_to_all( + activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) + # Same FP32-on-wire / compute_dtype-for-activation cast as forward. + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts).to( + dtype=self.compute_dtype + ) + 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.to(dtype=self.compute_dtype) + expert_rows = metadata[:, 0] + d_x_rows = torch.zeros( + (local_routes, self.hidden_size), + dtype=self.compute_dtype, + 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).to(dtype=self.compute_dtype).float() + else: + d_h = d_h_fc2 + d_w_rows[positions] = ( + (d_y * (h @ fc2_float[expert])).sum(dim=-1).to(dtype=self.compute_dtype).float() + ) + + 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=self.compute_dtype, + 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/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 6def36ffc7..531dde233f 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -22,7 +22,9 @@ from .all_reduce import AllReduce from .basic_linear import BasicLinear from .bias import Bias +from .combine import Combine from .constant_scale import ConstantScale +from .dispatch import Dispatch from .dropout import Dropout from .grouped_linear import GroupedLinear from .identity import Identity diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 8137051322..a974d41ef9 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -6,7 +6,7 @@ from __future__ import annotations import abc -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -401,7 +401,7 @@ def fuser_forward( prev_op_grad_output_quantizer: Optional[Quantizer], # pylint: disable=unused-argument next_op_input_quantizer: Optional[Quantizer], # pylint: disable=unused-argument basic_op_kwargs: list[dict[str, Any]], # pylint: disable=unused-argument - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: if self.activation_recompute_in_mlp: raise RuntimeError( f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " diff --git a/transformer_engine/pytorch/ops/basic/add_extra_input.py b/transformer_engine/pytorch/ops/basic/add_extra_input.py index fc3ca9cade..9af399f2dc 100644 --- a/transformer_engine/pytorch/ops/basic/add_extra_input.py +++ b/transformer_engine/pytorch/ops/basic/add_extra_input.py @@ -5,7 +5,7 @@ """Fusible operation for adding extra input tensor.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -67,7 +67,7 @@ def fuser_forward( 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]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: extra_input = basic_op_extra_inputs[0][0] if self._in_place: extra_input = extra_input.detach() diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py new file mode 100644 index 0000000000..4d4740c517 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -0,0 +1,138 @@ +# 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/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py new file mode 100644 index 0000000000..9762ed76f7 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -0,0 +1,207 @@ +# 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)] diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index e1980d2943..27213fe7cf 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -147,11 +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 + super().__init__() self.wgrad_store = WeightGradStore(delay_wgrad_compute) self.wgrad_accumulation_and_reduce_hooks: list = [] @@ -1003,7 +1003,7 @@ def fuser_forward( 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]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: num_groups = self.num_groups weight_param = self.weight if self.single_grouped_weight else self.weight0 device = weight_param.device diff --git a/transformer_engine/pytorch/ops/basic/make_extra_output.py b/transformer_engine/pytorch/ops/basic/make_extra_output.py index 0d9c870262..47b7c6495d 100644 --- a/transformer_engine/pytorch/ops/basic/make_extra_output.py +++ b/transformer_engine/pytorch/ops/basic/make_extra_output.py @@ -5,7 +5,7 @@ """Make extra tensor output in operation fuser.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -72,7 +72,7 @@ def fuser_forward( 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]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: return input_, [(input_,)] def fuser_backward( @@ -80,14 +80,17 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[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]]], ]: grad_extra_output = basic_op_grad_extra_outputs[0][0] - if self._in_place: + if grad_extra_output is None: + # Extra output is not consumed, so its gradient is zero + grad_input = grad_output + elif self._in_place: grad_extra_output += grad_output grad_input = grad_extra_output else: diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index fb663c0480..598126e2cb 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -5,7 +5,7 @@ """Fusible operation for SwiGLU and variants.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any, Optional import torch @@ -429,7 +429,7 @@ def fuser_forward( 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]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: if self.activation_recompute_in_mlp: raise RuntimeError( f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " diff --git a/transformer_engine/pytorch/ops/fused/__init__.py b/transformer_engine/pytorch/ops/fused/__init__.py index dc9dcd6dc3..7aa45a89f5 100644 --- a/transformer_engine/pytorch/ops/fused/__init__.py +++ b/transformer_engine/pytorch/ops/fused/__init__.py @@ -35,3 +35,4 @@ GroupedMLP_CuTeGEMMGLU, GroupedMLP_CuTeGEMMUnary, ) +from .moe_ep import FusedMoeEp # pylint: disable=wrong-import-position diff --git a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py index a3c81e60c8..483749da47 100644 --- a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py +++ b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py @@ -33,7 +33,7 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], ) -> tuple[ torch.Tensor, list[tuple[Optional[torch.Tensor], ...]], @@ -56,7 +56,11 @@ def fuser_backward( extra_grad = basic_op_grad_extra_outputs[0][0] dy = maybe_dequantize(grad_output.contiguous(), dtype).view(x.size()) w = maybe_dequantize(rmsnorm_op.weight, dtype).view((inner_dim,)) - add = maybe_dequantize(extra_grad.contiguous(), dtype).view(x.size()) + if extra_grad is None: + # Extra output is not consumed, so its gradient is zero + add = torch.zeros_like(dy) + else: + add = maybe_dequantize(extra_grad.contiguous(), dtype).view(x.size()) # Compute RMSNorm backward pass dx, dw = tex.rmsnorm_bwd_add( diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_add.py b/transformer_engine/pytorch/ops/fused/backward_linear_add.py index 382fecfd07..60b1320b2e 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_add.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_add.py @@ -40,7 +40,7 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], ) -> tuple[ torch.Tensor, list[tuple[Optional[torch.Tensor], ...]], @@ -67,19 +67,21 @@ def fuser_backward( else: accumulate_into_main_grad = False - # Linear backward pass + # Linear backward pass. Skip in-place add when there is no + # extra-output gradient. grad_input = basic_op_grad_extra_outputs[0][0] + accumulate_into_grad_input = grad_input is not None grad_input, grad_weight = BasicLinear._functional_backward( grad_output=grad_output, input=x_local, weight=w, input_requires_grad=linear_op_ctx.input_requires_grad, weight_requires_grad=linear_op_ctx.weight_requires_grad, - dtype=grad_input.dtype, + dtype=None if grad_input is None else grad_input.dtype, grad_weight=grad_weight, accumulate_into_grad_weight=accumulate_into_main_grad, grad_input=grad_input, - accumulate_into_grad_input=True, + accumulate_into_grad_input=accumulate_into_grad_input, tensor_parallel_mode=linear_op.tensor_parallel_mode, tensor_parallel_group=linear_op.tensor_parallel_group, sequence_parallel=linear_op.sequence_parallel, diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py index 8df929f799..6fa63675b4 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + bias + activation.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Sequence from typing import Any, Optional import torch @@ -59,7 +59,7 @@ def fuser_forward( 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]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations idx = self._op_idxs["linear"] diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index 5376a7d264..28586360f5 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + bias + add.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Sequence from typing import Any, Optional import torch @@ -57,7 +57,7 @@ def fuser_forward( 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]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations idx = self._op_idxs["linear"] diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index abeb39adfa..277263e0ec 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -5,7 +5,7 @@ """Fused operation for forward GEMM + scale + add.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Sequence from typing import Any, Optional import torch @@ -47,7 +47,7 @@ def fuser_forward( 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]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations linear_op = self.basic_ops[0] diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 76d51673f0..e618ef89c3 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -6,7 +6,7 @@ from __future__ import annotations -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Sequence import functools import os from importlib.metadata import PackageNotFoundError, version as get_pkg_version @@ -975,7 +975,7 @@ def fuser_forward( 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]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations fc1_op, activation_op, fc2_op = self.basic_ops fc1_ctx, _activation_ctx, fc2_ctx = basic_op_ctxs diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py new file mode 100644 index 0000000000..dd0f46c7ab --- /dev/null +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -0,0 +1,299 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Reference-backed BF16 expert-parallel MoE fusion.""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from typing import Any, Optional + +import torch + +from ...ep import get_ep_group +from ...ep_reference import MoeEpReference +from ...quantization import Recipe +from ...tensor import Quantizer +from ..basic import Combine, Dispatch, GroupedLinear, ScaledSwiGLU +from ..fuser import register_forward_backward_fusion +from ..op import FusedOperation, FusibleOperation, OperationContext + + +def _weight_list(op: GroupedLinear) -> list[torch.Tensor]: + """Return per-expert weights in their registered order.""" + return [getattr(op, f"weight{idx}") for idx in range(op.num_groups)] + + +def _reference_weights(op: GroupedLinear) -> torch.Tensor: + """Pack ``(out, in)`` expert weights into reference ``(E, in, out)`` layout.""" + return torch.stack([weight.transpose(0, 1) for weight in _weight_list(op)]) + + +def _grouped_linear_supported(op: GroupedLinear) -> bool: + weights = _weight_list(op) if not op.single_grouped_weight else [] + return ( + not op.use_bias + and not op._scale_bias + and not op.single_grouped_weight + and not op.single_grouped_bias + and not op._accumulate_into_main_grad + and not op._is_distributed_weight() + and not op.wgrad_store.delay_wgrad_compute() + and bool(weights) + and all(weight.dtype is torch.bfloat16 for weight in weights) + ) + + +def _routing_extras_internal( + dispatch: Dispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, +) -> bool: + """Whether the dispatch routing extras stay inside the fusion. + + The fused op keeps tokens-per-expert and the received routing weights + internal to :class:`MoeEpReference`, so it can only replace the sequence + when those two outputs feed exactly these ops and are not returned to the + caller. + """ + tokens_per_expert, routing_weights = dispatch._extra_output_channels + if tokens_per_expert is None or routing_weights is None: + return False + if any(dispatch._extra_output_to_caller): + return False + return ( + fc1._extra_input_channels[0] == tokens_per_expert + and fc2._extra_input_channels[0] == tokens_per_expert + and activation._extra_input_channels[0] == routing_weights + ) + + +def _routing_extras_internal( + dispatch: Dispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, +) -> bool: + """Whether the dispatch routing extras stay inside the fusion. + + The fused op keeps tokens-per-expert and the received routing weights + internal to :class:`MoeEpReference`, so it can only replace the sequence + when those two outputs feed exactly these ops and are not returned to the + caller. + """ + tokens_per_expert, routing_weights = dispatch._extra_output_channels + if tokens_per_expert is None or routing_weights is None: + return False + if any(dispatch._extra_output_to_caller): + return False + return ( + fc1._extra_input_channels[0] == tokens_per_expert + and fc2._extra_input_channels[0] == tokens_per_expert + and activation._extra_input_channels[0] == routing_weights + ) + + +def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bool: + if recipe is not None or len(window) != 5: + return False + dispatch, fc1, activation, fc2, combine = window + if not ( + isinstance(dispatch, Dispatch) + and isinstance(fc1, GroupedLinear) + and isinstance(activation, ScaledSwiGLU) + and isinstance(fc2, GroupedLinear) + and isinstance(combine, Combine) + ): + return False + if dispatch.buffer is not combine.buffer: + return False + buffer = dispatch.buffer + if not buffer.eager or buffer.payload_dtype is not torch.bfloat16: + return False + if not (_grouped_linear_supported(fc1) and _grouped_linear_supported(fc2)): + return False + if activation.activation_recompute_in_mlp or activation.glu_interleave_size is not None: + return False + if not _routing_extras_internal(dispatch, fc1, activation, fc2): + return False + return ( + fc1.num_groups == buffer.num_local_experts + and fc2.num_groups == buffer.num_local_experts + and fc1.in_features == buffer.hidden_dim + and fc2.out_features == buffer.hidden_dim + and fc1.out_features == 2 * fc2.in_features + ) + + +class FusedMoeEp(FusedOperation): + """Joint BF16 fusion implemented with :class:`MoeEpReference`.""" + + def __init__( + self, + *, + dispatch: Dispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, + combine: Combine, + ) -> None: + super().__init__([dispatch, fc1, activation, fc2, combine]) + ep_group = get_ep_group() + ep_size = 1 if ep_group is None else ep_group.size() + self._reference = MoeEpReference( + num_experts=dispatch.buffer.num_local_experts * ep_size, + hidden_size=dispatch.buffer.hidden_dim, + intermediate_size=fc2.in_features, + top_k=dispatch.buffer.top_k, + ep_group=ep_group, + max_tokens_per_rank=dispatch.buffer.max_tokens_per_rank, + compute_dtype=torch.bfloat16, + apply_topk_in_fc1=True, + generate_c=True, + ) + + @property + def dispatch(self) -> Dispatch: + return self.basic_ops[0] + + @property + def fc1(self) -> GroupedLinear: + return self.basic_ops[1] + + @property + def fc2(self) -> GroupedLinear: + return self.basic_ops[3] + + 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, Sequence[Sequence[Optional[torch.Tensor]]]]: + del prev_op_grad_output_quantizer, next_op_input_quantizer + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"FusedMoeEp requires BF16 input, got {input_.dtype}.") + if any(kwargs for kwargs in basic_op_kwargs): + raise NotImplementedError("FusedMoeEp does not support per-operation output buffers.") + + topk_idx, topk_weights = basic_op_extra_inputs[0] + if topk_weights.dtype is not torch.float32: + raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") + fc1_weight = _reference_weights(self.fc1) + fc2_weight = _reference_weights(self.fc2) + output, fc1_c, route_metadata = self._reference( + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ) + + if any(ctx.requires_grad for ctx in basic_op_ctxs): + basic_op_ctxs[0].save_for_backward( + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + + # Dispatch extras are channel-bound with output_to_caller=False and are + # only consumed by ops inside this fusion, so they need not be materialized. + return output, [ + (None, None), + (), + (), + (), + (), + ] + + 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]]], + ]: + del basic_op_grad_extra_outputs + ( + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) = basic_op_ctxs[0].saved_tensors + grad_input, grad_fc1, grad_fc2, grad_topk_weights = self._reference.backward( + grad_output, + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + + fc1_param_grads = [ + grad_fc1[idx].transpose(0, 1) if weight.requires_grad else None + for idx, weight in enumerate(_weight_list(self.fc1)) + ] + fc2_param_grads = [ + grad_fc2[idx].transpose(0, 1) if weight.requires_grad else None + for idx, weight in enumerate(_weight_list(self.fc2)) + ] + return ( + grad_input, + [(), fc1_param_grads, (), fc2_param_grads, ()], + [(None, grad_topk_weights), (None,), (None,), (None,), ()], + ) + + +def fuse_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused: Any, +) -> list[FusibleOperation]: + """Fuse supported five-op BF16 EP MoE sequences.""" + del unused + out: list[FusibleOperation] = [] + idx = 0 + while idx < len(ops): + window = ops[idx : idx + 5] + if _matches(window, recipe): + dispatch, fc1, activation, fc2, combine = window + out.append( + FusedMoeEp( + dispatch=dispatch, + fc1=fc1, + activation=activation, + fc2=fc2, + combine=combine, + ) + ) + idx += 5 + else: + out.append(ops[idx]) + idx += 1 + return out + + +register_forward_backward_fusion(fuse_ops, prepend=True) + + +__all__ = ["FusedMoeEp"] diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 3a8ff5438d..cb5f8a16e6 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -5,7 +5,7 @@ """Linear layer forward with Userbuffers communication.""" from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Sequence from typing import Any, Optional import torch @@ -286,7 +286,7 @@ def fuser_forward( 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]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: # Get basic operations idx = self._op_idxs["linear"] diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 09ffb004dd..9f346bd4c8 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -102,23 +102,45 @@ 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 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]]] = [ + [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_ - extra_outputs = [None] * fuser._num_basic_ops + extra_outputs: list[Optional[Sequence[Optional[torch.Tensor]]]] = [ + None + ] * fuser._num_basic_ops for op, basic_op_idxs in fuser._forward_ops: # Set if backward op is required for idx in basic_op_idxs: basic_op_ctxs[idx].requires_grad = idx >= fuser.first_op_requiring_backward - # Forward op + # Resolve internal channel inputs from outputs of + # 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]): + if source is None: + continue + producer_idx, output_idx = source + if producer_idx in basic_op_idxs: + # fused op will wire the channel itself internally + continue + producer_outputs = extra_outputs[producer_idx] + basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx] + + # Prepare args for op forward extra_inputs = [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 @@ -139,24 +161,53 @@ def forward( next_op_input_quantizer=next_op_input_quantizer, basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], ) + if len(fused_op_extra_outputs) != len(basic_op_idxs): + raise RuntimeError( + f"Expected {type(op).__name__} to generate extra outputs for " + f"{len(basic_op_idxs)} basic operations, " + f"but got {len(fused_op_extra_outputs)}" + ) for idx, ys in zip(basic_op_idxs, fused_op_extra_outputs): - for y in ys: - if set_output_requires_grad: - y.requires_grad_(idx >= fuser.first_op_requiring_backward) + num_extra_outputs = fuser._basic_ops[idx].num_extra_outputs + if len(ys) != num_extra_outputs: + raise RuntimeError( + f"Expected op {idx} to generate {num_extra_outputs} extra outputs, " + f"but got {len(ys)}" + ) + for output_idx, y in enumerate(ys): + if y is None: + # Extra output can be None if it is not required by any operations outside the fusion + # and is not required to be outputted to the caller. + output_to_caller = fuser._basic_op_extra_output_to_caller[idx][output_idx] + consumers = fuser._basic_op_extra_output_consumers[idx][output_idx] + needed_outside_fusion = any( + consumer_idx not in basic_op_idxs for consumer_idx in consumers + ) + if output_to_caller: + raise RuntimeError( + f"Op {idx} extra output {output_idx} is public, " + f"but {type(op).__name__} returned None" + ) + if needed_outside_fusion: + raise RuntimeError( + f"Op {idx} extra output {output_idx} is required by an " + "operation outside its forward fusion, " + f"but {type(op).__name__} returned None" + ) + continue + if ( + set_output_requires_grad + and idx >= fuser.first_op_requiring_backward + and y.is_floating_point() + ): + y.requires_grad_(True) extra_outputs[idx] = ys - # Flatten list of extra outputs - extra_outputs_flat = [] - for idx, ys in enumerate(extra_outputs): - ys = list(ys) - num_extra_outputs = fuser._basic_ops[idx].num_extra_outputs - if len(ys) != num_extra_outputs: - raise RuntimeError( - f"Expected op {idx} to generate " - "{num_extra_outputs} extra inputs, " - f"but got {len(ys)}" - ) - extra_outputs_flat.extend(ys) + # Collect caller-visible extra outputs in basic-op and slot order. + extra_outputs_flat = [ + extra_outputs[op_idx][output_idx] + for op_idx, output_idx in fuser._public_extra_output_slots + ] # Save context for backward pass if func_ctx is not None: @@ -186,12 +237,19 @@ def forward( 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.external_extra_input_slots = fuser._external_extra_input_slots + func_ctx.public_extra_output_slots = fuser._public_extra_output_slots + func_ctx.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels + func_ctx.basic_op_extra_output_consumers = fuser._basic_op_extra_output_consumers + 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 - for tensor in [x] + extra_outputs_flat: + for tensor in itertools.chain( + (x,), + (y for ys in extra_outputs for y in ys if y is not None), + ): tensor._do_not_clear = True if set_output_requires_grad: @@ -224,21 +282,32 @@ def backward( ctx.saved_tensors = saved_tensors[slice(*ctx._saved_tensors_range)] ctx._saved_tensors_range = None - # Unflatten list of extra tensor output grads + # Channel wiring saved from forward + basic_op_extra_output_channels = func_ctx.basic_op_extra_output_channels + basic_op_extra_output_consumers = func_ctx.basic_op_extra_output_consumers + basic_op_extra_input_sources = func_ctx.basic_op_extra_input_sources + + # Place caller-provided extra-output grads into their basic-op slots. + # Gradients from internal channel consumers are added 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, + func_ctx.public_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,7 +315,17 @@ def backward( dx = None break - # Backward op + # 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]): + if basic_op_extra_output_consumers[idx][output_idx]: + channel_grad = channel_grads.get(channel) + if channel_grad is not None: + output_grad = basic_op_grad_extra_outputs[idx][output_idx] + basic_op_grad_extra_outputs[idx][output_idx] = ( + channel_grad if output_grad is None else output_grad + channel_grad + ) grad_extra_outputs = [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], @@ -258,6 +337,18 @@ def backward( 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 = basic_op_extra_input_sources[idx][input_idx] + 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 = 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 = [] @@ -275,20 +366,22 @@ def backward( grad_params_flat.extend(dparams) # Flatten list of parameter gradients - grad_extra_inputs_flat = [] for idx, dxs in enumerate(grad_extra_inputs): num_extra_inputs = basic_ops[idx].num_extra_inputs if dxs is None: - dxs = [None for _ in range(num_extra_inputs)] - else: - dxs = list(dxs) - if len(dxs) != num_extra_inputs: + grad_extra_inputs[idx] = (None,) * num_extra_inputs + elif len(dxs) != num_extra_inputs: raise RuntimeError( f"Expected op {idx} to generate grads " f"for {num_extra_inputs} extra inputs, " f"but got {len(dxs)}" ) - grad_extra_inputs_flat.extend(dxs) + + # Collect the gradient for each public extra input. + grad_extra_inputs_flat = [ + grad_extra_inputs[op_idx][input_idx] + for op_idx, input_idx in func_ctx.external_extra_input_slots + ] # Update FP8 scaling factors if func_ctx.is_first_module and not _is_graph_capturing(): @@ -339,10 +432,104 @@ def __init__( basic_ops.append(op) self._num_basic_ops: int = len(basic_ops) self._basic_ops: list[BasicOperation] = basic_ops + # Capture channel routing from each basic op. If any op later rebinds a + # channel it flips this flag, essentially invalidating this fuser's forward pass. + self._channels_stale: bool = False + for op in self._basic_ops: + op._capturing_op_fusers.add(self) # 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._basic_op_extra_output_to_caller: list[list[bool]] = [ + list(op._extra_output_to_caller) for op in basic_ops + ] + self._basic_op_extra_output_consumers: list[list[list[int]]] = [ + [[] for _ in range(op.num_extra_outputs)] for op in basic_ops + ] + self._external_extra_input_slots: list[tuple[int, int]] = [] + self._public_extra_output_slots: list[tuple[int, int]] = [] + + # Find channel producers and reject ambiguous names. + channel_producers: dict[str, tuple[int, int]] = {} + for op_idx, op in enumerate(basic_ops): + for output_idx, channel in enumerate(self._basic_op_extra_output_channels[op_idx]): + if channel is None: + 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) + + # Resolve inputs. Named inputs must have an earlier producer; + # unnamed inputs remain public. + 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 + producer = channel_producers.get(channel) + if producer is None: + raise ValueError( + f"Extra tensor channel {channel!r} consumed by op {op_idx} " + f"({type(op).__name__}) has no producer" + ) + producer_idx, _ = producer + if producer_idx >= op_idx: + 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] = producer + producer_idx, output_idx = producer + self._basic_op_extra_output_consumers[producer_idx][output_idx].append(op_idx) + + # Record caller-visible outputs in stable basic-op and slot order. + for op_idx, op in enumerate(basic_ops): + for output_idx in range(op.num_extra_outputs): + if self._basic_op_extra_output_to_caller[op_idx][output_idx]: + self._public_extra_output_slots.append((op_idx, output_idx)) + + # 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 RuntimeError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is bound to channel {channel!r} " + "without a producer source" + ) + 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}" + ) + # 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) # Ops for forward and backward pass, will be populated in maybe_fuse_ops self._forward_ops: list[tuple[FusibleOperation, list[int]]] @@ -432,7 +619,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.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 @@ -507,6 +694,12 @@ def __call__( *extra_inputs: torch.Tensor, basic_op_kwargs: Optional[list[dict[str, Any]]] = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: + if self._channels_stale: + raise RuntimeError( + "Extra tensor channels changed after this OperationFuser captured " + "its routing. Construct a new OperationFuser." + ) + # Verify extra input count if len(extra_inputs) != self.num_extra_inputs: raise ValueError( @@ -517,12 +710,13 @@ 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..bb02a5f0ee 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -6,10 +6,11 @@ from __future__ import annotations import abc -from collections.abc import Iterable +from collections.abc import Iterable, Sequence import dataclasses import pickle from typing import Any, Optional +import weakref import torch @@ -85,11 +86,11 @@ def fuser_forward( basic_op_ctxs: list[OperationContext], input_: torch.Tensor, *, - basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + basic_op_extra_inputs: Sequence[Sequence[Optional[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]]]: + ) -> tuple[torch.Tensor, Sequence[Sequence[Optional[torch.Tensor]]]]: """Forward pass This op is either a basic op or the fusion of basic ops, so @@ -104,8 +105,9 @@ def fuser_forward( Contexts for basic operations input_: torch.Tensor Input tensor - basic_op_extra_inputs: list of torch.Tensor - Extra tensor inputs to basic operations + basic_op_extra_inputs: sequence of sequences of torch.Tensor + Extra tensor inputs to basic operations. An internal input + owned by this fused operation may be ``None``. prev_op_grad_output_quantizer: Quantizer, optional The grad_output_quantizer of the preceeding operation next_op_input_quantizer: Quantizer, optional @@ -118,8 +120,9 @@ def fuser_forward( ------- torch.Tensor: Output tensor. - Iterable of torch.Tensor: - Extra tensor outputs from basic operations. + Sequence of sequences of torch.Tensor: + Extra tensor outputs from basic operations. A non-public + channel owned by this fused operation may be ``None``. """ raise NotImplementedError( @@ -131,11 +134,11 @@ def fuser_backward( basic_op_ctxs: list[OperationContext], grad_output: torch.Tensor, *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + basic_op_grad_extra_outputs: Sequence[Sequence[Optional[torch.Tensor]]], ) -> tuple[ torch.Tensor, - Iterable[Iterable[Optional[torch.Tensor]]], - Iterable[Iterable[Optional[torch.Tensor]]], + Sequence[Sequence[Optional[torch.Tensor]]], + Sequence[Sequence[Optional[torch.Tensor]]], ]: """Backward pass @@ -159,9 +162,9 @@ def fuser_backward( ------- torch.Tensor: Loss gradient w.r.t. operation input - Iterable of iterable of torch.Tensor: + Sequence of sequences of torch.Tensor: Loss gradients w.r.t. parameters for basic operations - Iterable of iterable of torch.Tensor: + Sequence of sequences of torch.Tensor: Loss gradients w.r.t. extra tensor inputs to basic operations @@ -187,10 +190,85 @@ 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 + self._extra_output_to_caller: list[bool] = [True] * self.num_extra_outputs + # OperationFusers that have captured this op's channel routing. Weak + # refs avoid cycles and drop automatically when a fuser is discarded. + self._capturing_op_fusers: weakref.WeakSet = weakref.WeakSet() + # Objects for quantization self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None + def _invalidate_capturing_op_fusers(self) -> None: + """Mark fusers that captured this op's channels as stale. + + Channel routing is cheap to change, but any OperationFuser that + already captured it must be rebuilt. Flipping a flag here keeps the + per-call check O(1) instead of rescanning versions every forward. + """ + for fuser in self._capturing_op_fusers: + fuser._channels_stale = True + self._capturing_op_fusers.clear() + + 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") + if self._extra_input_channels[index] == channel: + return self + self._extra_input_channels[index] = channel + self._invalidate_capturing_op_fusers() + return self + + def set_extra_output_channel( + self, + index: int, + channel: Optional[str], + *, + output_to_caller: bool = True, + ) -> BasicOperation: + """Bind an extra output slot to an internal fuser channel. + + A bound slot can feed one or more later operations. By default, the + output is also returned to the caller. Set ``output_to_caller=False`` + to keep it internal to the fuser. Passing ``channel=None`` removes the + binding and restores the output as public. + """ + 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") + if not isinstance(output_to_caller, bool): + raise TypeError("output_to_caller must be a bool") + if channel is None: + output_to_caller = True + if ( + self._extra_output_channels[index] == channel + and self._extra_output_to_caller[index] == output_to_caller + ): + return self + self._extra_output_channels[index] = channel + self._extra_output_to_caller[index] = output_to_caller + self._invalidate_capturing_op_fusers() + return self + @property def is_fused_op(self) -> bool: return False