diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index dd17191e58..1466e38121 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -151,6 +151,132 @@ 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 other operations, in addition to being returned from the ``Sequential``. Extra input tensors with a specified channel are accessed internally instead of being provided as arguments to ``Seqential``. + +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") + 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 is also returned to the caller. + x = torch.randn(16384, 4096, device="cuda") + y, residual = 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") + dispatch.set_extra_output_channel(1, "probs") + dispatch.set_extra_output_channel(2, "routing_map") + + fc1.set_extra_input_channel(0, "m_splits") + activation.set_extra_input_channel(0, "probs") + fc2.set_extra_input_channel(0, "m_splits") + combine.set_extra_input_channel(0, "routing_map") + + 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, while Dispatch's + # extra outputs are still returned in their original order. + y, m_splits, probs, routing_map = 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: + +- A producer must appear before all of its consumers. Backward edges + and cycles are not supported. +- 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 still returned as + a public extra output. +- 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 extra inputs that are not connected to an earlier + producer in the same fuser. Channel-connected extra input slots do + not appear in the ``Sequential`` arguments. +- The caller receives every extra output in the original basic-operation + and slot order. This includes channel-bound outputs that are also + consumed internally. Gradients supplied for a returned output are + combined with gradients from its internal channel consumers. +- 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. + Developer guide --------------- diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 66857d8125..b50c9c3fc1 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,417 @@ 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("with_extra_grad", (True, False)) + def test_fused_internal_residual_connection( + self, + fusion_kind: str, + with_extra_grad: 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 + return 2 * input_ + self.basic_ops[1].bias, [(input_,), (), ()] + + 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 the gradient from the public residual 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. + 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") + 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) + y, residual_out = model(x) + + 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 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) + + 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_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_external_named_extra_inputs_remain_separate(self, size: int = 16) -> None: + """Unmatched inputs with the same channel require separate public tensors.""" + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + consumer1.set_extra_input_channel(0, "external") + consumer2.set_extra_input_channel(0, "external") + model = te_ops.Sequential(consumer1, consumer2) + + x = torch.rand((size,), requires_grad=True) + extra = torch.rand((size,), requires_grad=True) + with pytest.raises(ValueError, match="Expected 2 extra inputs but got 1"): + model(x, extra) + y = model(x, extra, extra) + torch.testing.assert_close(y, x + 2 * extra) + + dy = torch.rand_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, dy) + torch.testing.assert_close(extra.grad, 2 * dy) + + def test_consumer_before_producer(self) -> None: + """Channels only connect forward; a later producer does not satisfy an earlier consumer.""" + consumer = te_ops.AddExtraInput() + producer = te_ops.MakeExtraOutput() + consumer.set_extra_input_channel(0, "route") + producer.set_extra_output_channel(0, "route") + with pytest.raises(ValueError, match="has no earlier producer"): + OperationFuser([consumer, producer]) + + def test_set_extra_channel_rejects_invalid_index(self) -> None: + """Slot indices must be in range; negatives and OOB are rejected at bind time.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + + with pytest.raises(IndexError, match="out of range"): + producer.set_extra_output_channel(-1, "route") + with pytest.raises(IndexError, match="out of range"): + producer.set_extra_output_channel(1, "route") + with pytest.raises(IndexError, match="out of range"): + consumer.set_extra_input_channel(-1, "route") + with pytest.raises(IndexError, match="out of range"): + consumer.set_extra_input_channel(1, "route") + + def test_set_extra_channel_rejects_invalid_name(self) -> None: + """Channel names must be non-empty strings.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + + with pytest.raises(ValueError, match="non-empty string"): + producer.set_extra_output_channel(0, "") + with pytest.raises(ValueError, match="non-empty string"): + consumer.set_extra_input_channel(0, "") + with pytest.raises(ValueError, match="non-empty string"): + producer.set_extra_output_channel(0, 123) # type: ignore[arg-type] + + def test_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) + + @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_named_extra_output_without_consumer_is_public(self, size: int = 16) -> None: + """A named output remains public when its fuser has no consumer.""" + producer = te_ops.MakeExtraOutput() + producer.set_extra_output_channel(0, "orphan") + x = torch.rand((size,), requires_grad=True) + y, extra = producer(x) + torch.testing.assert_close(y, x) + torch.testing.assert_close(extra, x) + + 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_channel_outputs_are_public(self, size: int = 16) -> None: + """Both internally consumed and unconsumed channel outputs are public.""" + producer = _DualExtraOutput(scales=(2.0, 3.0)) + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "internal") + 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, internal, public = model(x) + torch.testing.assert_close(y, 3 * x) + torch.testing.assert_close(internal, 2 * x) + torch.testing.assert_close(public, 3 * x) + + dy = torch.rand_like(y) + dinternal = torch.rand_like(internal) + dpublic = torch.rand_like(public) + torch.autograd.backward((y, internal, public), (dy, dinternal, dpublic)) + torch.testing.assert_close(x.grad, 3 * dy + 2 * dinternal + 3 * dpublic) + + def test_fresh_internal_output_preserves_grad_requirement(self) -> None: + """A fresh internal tensor requests its gradient from a scaled activation.""" + + # 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, [()], [()] + + producer = MakeScale() + activation = te_ops.ScaledSReLU() + producer.set_extra_output_channel(0, "scale") + activation.set_extra_input_channel(0, "scale") + model = te_ops.Sequential(producer, activation) + + x_ref = torch.randn((5, 8), device="cuda", requires_grad=True) + x_test = x_ref.detach().clone().requires_grad_(True) + scale_ref = x_ref.square().mean(dim=-1) + y_ref = torch.nn.functional.relu(x_ref).square() * scale_ref.unsqueeze(-1) + y_test, _scale_test = model(x_test) + 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 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/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..3c95845c80 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -102,12 +102,17 @@ 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_ @@ -118,7 +123,22 @@ def forward( 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 +159,30 @@ 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 ( + 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) + extra_outputs_flat = [y for ys in extra_outputs for y in ys] # Save context for backward pass if func_ctx is not None: @@ -186,8 +212,11 @@ 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.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels + func_ctx.basic_op_extra_output_is_internal = fuser._basic_op_extra_output_is_internal + 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 @@ -224,21 +253,31 @@ 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_is_internal = func_ctx.basic_op_extra_output_is_internal + 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 = [] + basic_op_grad_extra_outputs: list[list[Optional[torch.Tensor]]] = [] 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) + grads, grad_extra_outputs = _split_tuple( + grad_extra_outputs, + op.num_extra_outputs, + ) + basic_op_grad_extra_outputs.append(list(grads)) # 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 +285,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_is_internal[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 +307,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 +336,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 +402,93 @@ 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_is_internal: list[list[bool]] = [ + [False] * op.num_extra_outputs for op in basic_ops + ] + self._external_extra_input_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. A channel with an earlier producer is internal. + # Every input without an earlier producer is a separate public input. + consumed_channels: set[str] = set() + for op_idx, op in enumerate(basic_ops): + for input_idx, channel in enumerate(op._extra_input_channels): + if channel is None: + self._external_extra_input_slots.append((op_idx, input_idx)) + continue + producer = channel_producers.get(channel) + # If no producer for named channel, this is a public input. + if producer is None: + self._external_extra_input_slots.append((op_idx, input_idx)) + continue + 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 + consumed_channels.add(channel) + + # All extra outputs remain public, including outputs consumed internally. + 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 not None and channel in consumed_channels: + self._basic_op_extra_output_is_internal[op_idx][output_idx] = True + + # 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: + continue + 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 +505,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 +588,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 +663,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 +679,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..85e5641078 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 @@ -89,7 +89,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]]]: """Forward pass This op is either a basic op or the fusion of basic ops, so @@ -118,7 +118,7 @@ def fuser_forward( ------- torch.Tensor: Output tensor. - Iterable of torch.Tensor: + Sequence of torch.Tensor: Extra tensor outputs from basic operations. """ @@ -187,10 +187,55 @@ 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_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]) -> BasicOperation: + """Bind an extra output slot to an internal fuser channel. + + A bound slot can feed one or more later operations and is not returned + as a public extra output. Passing ``None`` removes the binding. + """ + if not 0 <= index < self.num_extra_outputs: + raise IndexError( + f"Extra output index {index} is out of range for " + f"{type(self).__name__} with {self.num_extra_outputs} extra outputs" + ) + if channel is not None and (not isinstance(channel, str) or not channel): + raise ValueError("Extra output channel must be a non-empty string or None") + if self._extra_output_channels[index] == channel: + return self + self._extra_output_channels[index] = channel + self._extra_channels_version += 1 + return self + @property def is_fused_op(self) -> bool: return False