diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index dd17191e58..5b9bd838ca 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -151,6 +151,151 @@ 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. +- 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 first constructed. Changing + ``set_extra_input_channel`` / ``set_extra_output_channel`` afterward + requires constructing a new ``OperationFuser`` or ``Sequential``. + +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/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 66857d8125..a06c6da06f 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,500 @@ 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("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: + """Sequential does not auto-rebuild after a channel configuration change.""" + 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] + basic_op_ctxs[0].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 + input_, scale = basic_op_ctxs[0].saved_tensors + grad_input = grad_output * scale.unsqueeze(-1) + grad_scale = (grad_output * input_).sum(dim=-1) + 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""" 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/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..5ad00e2fd1 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( 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/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 909e5a8a9b..694208b60f 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 @@ -979,7 +979,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/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..50c1c11ffa 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,102 @@ def __init__( basic_ops.append(op) self._num_basic_ops: int = len(basic_ops) self._basic_ops: list[BasicOperation] = basic_ops + self._basic_op_extra_channels_versions = [ + op._extra_channels_version for op in self._basic_ops + ] # 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]]] @@ -359,6 +544,16 @@ def __init__( self._basic_op_num_params = list(map(len, self._basic_op_params)) self._flat_basic_op_params = sum(self._basic_op_params, []) + def has_stale_op_channels(self) -> bool: + """Whether an operation's extra tensor channels have changed.""" + return any( + op._extra_channels_version != version + for op, version in zip( + self._basic_ops, + self._basic_op_extra_channels_versions, + ) + ) + @staticmethod def _apply_fusions( ops: Iterable[FusibleOperation], @@ -432,7 +627,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 +702,12 @@ def __call__( *extra_inputs: torch.Tensor, basic_op_kwargs: Optional[list[dict[str, Any]]] = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: + if self.has_stale_op_channels(): + 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 +718,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..2325002645 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -6,7 +6,7 @@ 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 @@ -85,11 +85,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 +104,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 +119,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,7 +133,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: Sequence[Sequence[Optional[torch.Tensor]]], ) -> tuple[ torch.Tensor, Iterable[Iterable[Optional[torch.Tensor]]], @@ -187,10 +189,72 @@ 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 + self._extra_channels_version = 0 + # Objects for quantization self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None + def set_extra_input_channel(self, index: int, channel: Optional[str]) -> BasicOperation: + """Bind an extra input slot to an internal fuser channel. + + A bound slot receives the matching extra output from an earlier + operation in the same fuser instead of consuming a public extra input. + Passing ``None`` removes the binding. + """ + if not 0 <= index < self.num_extra_inputs: + raise IndexError( + f"Extra input index {index} is out of range for " + f"{type(self).__name__} with {self.num_extra_inputs} extra inputs" + ) + if channel is not None and (not isinstance(channel, str) or not channel): + raise ValueError("Extra input channel must be a non-empty string or None") + if self._extra_input_channels[index] == channel: + return self + self._extra_input_channels[index] = channel + self._extra_channels_version += 1 + 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._extra_channels_version += 1 + return self + @property def is_fused_op(self) -> bool: return False