diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index dd17191e58..ffec6baa4e 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -113,33 +113,33 @@ quantized compute. Branching operations ^^^^^^^^^^^^^^^^^^^^ -The operation fuser supports very limited branching behavior. While -the operations must be in sequential order, some operations can accept +The operation fuser supports limited branching behavior. While the +operations must be in sequential order, some operations can accept extra inputs or produce extra outputs. For example, ``AddExtraInput`` -will add an extra input tensor to the intermediate tensor and -``MakeExtraOutput`` will return the intermediate tensor as an extra -output. When calling a ``Sequential`` that contains any of these -branching operations, the extra inputs should be passed in as -arguments and the extra outputs will be returned. +adds an extra input tensor to the intermediate tensor, and +``MakeExtraOutput`` returns the intermediate tensor as an extra output. +When calling a ``Sequential`` that contains any of these branching +operations, the extra inputs should be passed as arguments and the +extra outputs will be returned after the main output. .. code-block:: python import torch import transformer_engine.pytorch as te - # Construct MLP with residual connection + # Construct an MLP with a residual connection. fc1 = te.ops.Sequential( te.ops.LayerNorm(4096), - te.ops.MakeExtraOutput(), # Output residual + te.ops.MakeExtraOutput(), # Output the residual. te.ops.Linear(4096, 28672), te.ops.SwiGLU(), ) fc2 = te.ops.Sequential( te.ops.Linear(14336, 4096), - te.ops.AddExtraInput(), # Add residual + te.ops.AddExtraInput(), # Add the residual. ) - # Forward pass + # Pass the extra output from fc1 as the extra input to fc2. x = torch.randn(16384, 4096, device="cuda") y, residual = fc1(x) y = fc2(y, residual) @@ -147,9 +147,131 @@ arguments and the extra outputs will be returned. .. figure:: ./residual_layernorm_mlp.png :align: center - Operations for an MLP block with a residual connection. Note that - the block has been split into two sections, each with one branching - operation. + Operations for an MLP block with a residual connection. The block + is split into two sections so that the caller can pass the extra + output from the first section to the second. + +Extra tensor channels +""""""""""""""""""""" + +Extra inputs and outputs may optionally specify a channel. Assigning +the same channel name to an extra output and one or more later extra +inputs routes the tensor internally within the same +``OperationFuser``. Slots bound to channels are removed from the public +``Sequential`` interface. + +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, so the caller receives only y. + 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") + 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. + 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: + +- A producer must appear before all of its consumers. Backward edges + and cycles are not supported. +- A channel has exactly one producer, but its output may fan out to + multiple consumers. +- Every named output channel must have at least one consumer, and the + channel names on the producer and consumers must match. +- 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 have no channel assigned and + receives extra outputs that have no channel assigned. Slots assigned + to channels are internal and do not appear in the ``Sequential`` + arguments or return value. + +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..5362fac716 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,457 @@ def test_extra_tensors(self, size: int = 16) -> None: torch.testing.assert_close(x4, x4_orig + x3) +class TestExtraTensorChannels: + """Error handling and grad coverage for named extra-tensor channels.""" + + def test_internal_residual_connection(self, 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 = model(x) + + torch.testing.assert_close(y, 2 * x + body.bias) + y.sum().backward() + torch.testing.assert_close(x.grad, torch.full_like(x, 2)) + + @pytest.mark.parametrize("fusion_kind", ("forward", "backward", "forward_backward")) + def test_fused_internal_residual_connection( + self, + fusion_kind: str, + 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, including its + # contribution to the input gradient. + assert basic_op_grad_extra_outputs[0][0] is None + return ( + 2 * grad_output, + [(), (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) + ): + 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 = 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) + y.backward(dy) + torch.testing.assert_close(x.grad, 2 * dy) + 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 = model(x) + + # Main path: x -> x + route -> x + route + route. + torch.testing.assert_close(y, 3 * x) + dy = torch.rand_like(y) + y.backward(dy) + # The channel fan-out contributes two independent gradient paths. + torch.testing.assert_close(x.grad, 3 * dy) + + # Internal slots are unavailable before forward, so grad discovery + # must tolerate them when no public input requires gradients. + x_no_grad = x.detach() + torch.testing.assert_close(model(x_no_grad), 3 * x_no_grad) + + def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None: + """Unbound slots remain public when other slots use internal channels.""" + producer = te_ops.MakeExtraOutput() + internal_consumer = te_ops.AddExtraInput() + external_consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + internal_consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, internal_consumer, external_consumer) + + x = torch.rand((size,), requires_grad=True) + extra = torch.rand((size,), requires_grad=True) + y = model(x, extra) + + torch.testing.assert_close(y, 2 * x + extra) + 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_input_fanout(self, size: int = 16) -> None: + """One public tensor supplies every unmatched input with the same channel.""" + 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) + y = model(x, 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_set_extra_channel_rejects_mutation_after_fuser_construction(self) -> None: + """Channel routing is immutable after it has been captured by a fuser.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + fuser = OperationFuser([producer, consumer]) + assert fuser.num_extra_inputs == 0 + with pytest.raises(RuntimeError, match="cannot be changed"): + producer.set_extra_output_channel(0, None) + with pytest.raises(RuntimeError, match="cannot be changed"): + consumer.set_extra_input_channel(0, None) + + def test_duplicate_extra_output_channel_names(self) -> None: + """Two extra outputs may not publish the same channel name.""" + producer1 = te_ops.MakeExtraOutput() + producer2 = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer1.set_extra_output_channel(0, "route") + producer2.set_extra_output_channel(0, "route") + consumer.set_extra_input_channel(0, "route") + with pytest.raises(ValueError, match="multiple producers"): + OperationFuser([producer1, producer2, consumer]) + + def test_duplicate_extra_output_channels_on_same_op(self) -> None: + """A single op with multiple extras still cannot reuse a channel name.""" + + class DualExtraOutput(te_ops.BasicOperation): + num_extra_outputs = 2 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("DualExtraOutput uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("DualExtraOutput uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + return input_, [(input_, input_)] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + g0, g1 = basic_op_grad_extra_outputs[0] + grad_extra = torch.zeros_like(grad_output) + if g0 is not None: + grad_extra = grad_extra + g0 + if g1 is not None: + grad_extra = grad_extra + g1 + return grad_output + grad_extra, [()], [()] + + producer = DualExtraOutput() + consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + producer.set_extra_output_channel(1, "route") + consumer.set_extra_input_channel(0, "route") + with pytest.raises(ValueError, match="multiple producers"): + OperationFuser([producer, consumer]) + + def test_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) -> 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 + assert fuser._external_extra_output_slots == [(0, 0)] + + def test_mixed_channel_outputs_accept_generators(self, size: int = 16) -> None: + """Generator outputs support mixed internal and public channel slots.""" + + class DualExtraOutput(te_ops.BasicOperation): + num_extra_outputs = 2 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("DualExtraOutput uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("DualExtraOutput uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + del basic_op_ctxs, basic_op_extra_inputs + outputs = (2 * input_, 3 * input_) + return input_, (iter(outputs) for _ in range(1)) + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + del basic_op_ctxs + grad_internal, grad_public = basic_op_grad_extra_outputs[0] + return ( + grad_output + 2 * grad_internal + 3 * grad_public, + (iter(()) for _ in range(1)), + (iter(()) for _ in range(1)), + ) + + producer = DualExtraOutput() + 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, 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) + + def test_fresh_internal_output_preserves_grad_requirement(self) -> None: + """A fresh internal tensor requests its gradient from a scaled activation.""" + + 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 = 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) + + def test_grouped_linear_scale_bias_channels(self) -> None: + """Both GroupedLinear extra inputs can be supplied by channels.""" + + class RouteExtras(te_ops.BasicOperation): + num_extra_inputs = 2 + num_extra_outputs = 2 + + def op_forward(self, *args, **kwargs): + raise RuntimeError("RouteExtras uses fuser_forward") + + def op_backward(self, *args, **kwargs): + raise RuntimeError("RouteExtras uses fuser_backward") + + def fuser_forward(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused): + del basic_op_ctxs + return input_, [basic_op_extra_inputs[0]] + + def fuser_backward(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs): + del basic_op_ctxs + return grad_output, [()], [basic_op_grad_extra_outputs[0]] + + group_size, in_features, out_features = 2, 8, 6 + split_sizes = torch.tensor((3, 2), dtype=torch.int32, device="cuda") + num_tokens = int(split_sizes.sum()) + x = torch.randn((num_tokens, in_features), device="cuda", requires_grad=True) + scales = torch.randn((num_tokens,), device="cuda", requires_grad=True) + + producer = RouteExtras() + linear = te_ops.GroupedLinear( + group_size, + in_features, + out_features, + bias=True, + scale_bias=True, + device="cuda", + dtype=torch.float32, + ) + producer.set_extra_output_channel(0, "split_sizes") + producer.set_extra_output_channel(1, "bias_scales") + linear.set_extra_input_channel(0, "split_sizes") + linear.set_extra_input_channel(1, "bias_scales") + model = te_ops.Sequential(producer, linear) + + x_ref = x.detach().clone().requires_grad_(True) + scales_ref = scales.detach().clone().requires_grad_(True) + ys_ref = [] + for group_idx, (x_group, scale_group) in enumerate( + zip( + torch.split(x_ref, split_sizes.tolist()), + torch.split(scales_ref, split_sizes.tolist()), + ) + ): + weight = getattr(linear, f"weight{group_idx}") + bias = getattr(linear, f"bias{group_idx}") + ys_ref.append( + torch.nn.functional.linear(x_group, weight) + scale_group.unsqueeze(-1) * bias + ) + y_ref = torch.cat(ys_ref) + y_test = model(x, split_sizes, scales) + dy = torch.rand_like(y_test) + grads_ref = torch.autograd.grad( + y_ref, + (x_ref, scales_ref, *linear.parameters()), + dy, + ) + y_test.backward(dy) + + tols = dtype_tols(torch.float16) # Grouped GEMM uses TF32 for FP32 inputs. + torch.testing.assert_close(y_test, y_ref, **tols) + torch.testing.assert_close(x.grad, grads_ref[0], **tols) + torch.testing.assert_close(scales.grad, grads_ref[1], **tols) + for param, grad_ref in zip(linear.parameters(), grads_ref[2:]): + torch.testing.assert_close(param.grad, grad_ref, **tols) + + class TestFuser: """Tests for operation fusion infrastructure""" diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index be931829ea..7faad6536b 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 = [] diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 09ffb004dd..7c641ce7fa 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -102,12 +102,15 @@ 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, slots in zip(extra_inputs, fuser._external_extra_input_slots): + for op_idx, input_idx in slots: + basic_op_extra_inputs[op_idx][input_idx] = tensor # Apply forward ops x = input_ @@ -118,8 +121,34 @@ def forward( for idx in basic_op_idxs: basic_op_ctxs[idx].requires_grad = idx >= fuser.first_op_requiring_backward - # Forward op - extra_inputs = [basic_op_extra_inputs[idx] for idx in basic_op_idxs] + # Forward op. Resolve internal channel inputs from outputs of + # earlier basic ops. 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] + if producer_outputs is None: + raise RuntimeError( + f"Extra tensor channel producer op {producer_idx} has not run" + ) + if output_idx >= len(producer_outputs) or producer_outputs[output_idx] is None: + raise RuntimeError( + f"Extra tensor channel producer op {producer_idx} " + f"({type(fuser._basic_ops[producer_idx]).__name__}) " + f"did not emit extra output {output_idx} for " + f"consumer op {idx} " + f"({type(fuser._basic_ops[idx]).__name__}) " + f"input {input_idx}" + ) + basic_op_extra_inputs[idx][input_idx] = producer_outputs[output_idx] + op_extra_inputs = [tuple(basic_op_extra_inputs[idx]) for idx in basic_op_idxs] prev_op_idx = basic_op_idxs[0] - 1 prev_op = fuser._basic_ops[prev_op_idx] if prev_op_idx >= 0 else None prev_op_grad_output_quantizer = None @@ -134,29 +163,45 @@ def forward( x, fused_op_extra_outputs = op.fuser_forward( [basic_op_ctxs[idx] for idx in basic_op_idxs], x, - basic_op_extra_inputs=extra_inputs, + basic_op_extra_inputs=op_extra_inputs, prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, next_op_input_quantizer=next_op_input_quantizer, basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], ) + fused_op_extra_outputs = tuple(tuple(ys) for ys in fused_op_extra_outputs) + 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: + raise RuntimeError( + f"Op {idx} ({type(fuser._basic_ops[idx]).__name__}) " + f"did not emit extra output {output_idx}" + ) + if ( + set_output_requires_grad + and idx >= fuser.first_op_requiring_backward + and (y.is_floating_point() or y.is_complex()) + ): + 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) + # Flatten public extra outputs. Matched channels stay internal, while + # unnamed slots and named channels without consumers remain public. + extra_outputs_flat = [ + extra_outputs[op_idx][output_idx] + for op_idx, output_idx in fuser._external_extra_output_slots + ] # Save context for backward pass if func_ctx is not None: @@ -181,17 +226,23 @@ def forward( if fuser.first_op_requiring_backward < fuser._num_basic_ops: is_first_module = FP8GlobalStateManager.is_first_fp8_module() - # Other context + # Other context. Save only the wiring metadata needed by + # backward instead of the whole OperationFuser. func_ctx.backward_ops = fuser._backward_ops func_ctx.basic_ops = fuser._basic_ops func_ctx.basic_op_ctxs = basic_op_ctxs func_ctx.basic_op_num_params = fuser._basic_op_num_params - func_ctx.num_extra_inputs = fuser.num_extra_inputs func_ctx.num_extra_outputs = len(extra_outputs_flat) + func_ctx.external_extra_input_slots = fuser._external_extra_input_slots + func_ctx.external_extra_output_slots = fuser._external_extra_output_slots + func_ctx.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels + func_ctx.basic_op_extra_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 - for tensor in [x] + extra_outputs_flat: + all_extra_outputs = [y for ys in extra_outputs for y in ys] + for tensor in [x] + all_extra_outputs: tensor._do_not_clear = True if set_output_requires_grad: @@ -224,21 +275,30 @@ 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 + external_extra_output_slots = func_ctx.external_extra_output_slots + basic_op_extra_output_channels = func_ctx.basic_op_extra_output_channels + basic_op_extra_output_is_internal = func_ctx.basic_op_extra_output_is_internal + basic_op_extra_input_sources = func_ctx.basic_op_extra_input_sources + + # Place public extra-output grads into their basic-op slots. Internal + # output grads are accumulated from channel consumers during backward. 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, external_extra_output_slots): + basic_op_grad_extra_outputs[op_idx][output_idx] = grad # Apply backward ops dx = grad_output grad_params = [None for _ in range(len(basic_ops))] grad_extra_inputs = [None for _ in range(len(basic_ops))] + channel_grads: dict[str, torch.Tensor] = {} for op, basic_op_idxs in reversed(backward_ops): # Stop if no more gradients are required @@ -246,18 +306,51 @@ def backward( dx = None break - # Backward op - grad_extra_outputs = [basic_op_grad_extra_outputs[idx] for idx in basic_op_idxs] + # Backward op. Supply gradients accumulated from every consumer of + # each internal channel. + for idx in basic_op_idxs: + for output_idx, channel in enumerate(basic_op_extra_output_channels[idx]): + if basic_op_extra_output_is_internal[idx][output_idx]: + basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get(channel) + op_grad_extra_outputs = [ + tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs + ] dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( [basic_op_ctxs[idx] for idx in basic_op_idxs], dx, - basic_op_grad_extra_outputs=grad_extra_outputs, + basic_op_grad_extra_outputs=op_grad_extra_outputs, ) + fused_op_grad_params = tuple(tuple(grads) for grads in fused_op_grad_params) + fused_op_grad_extra_inputs = tuple(tuple(grads) for grads in fused_op_grad_extra_inputs) + if len(fused_op_grad_params) != len(basic_op_idxs): + raise RuntimeError( + f"Expected {type(op).__name__} to generate parameter grads for " + f"{len(basic_op_idxs)} basic operations, but got " + f"{len(fused_op_grad_params)}" + ) + if len(fused_op_grad_extra_inputs) != len(basic_op_idxs): + raise RuntimeError( + f"Expected {type(op).__name__} to generate extra-input grads for " + f"{len(basic_op_idxs)} basic operations, but got " + f"{len(fused_op_grad_extra_inputs)}" + ) for idx, dparams in zip(basic_op_idxs, fused_op_grad_params): grad_params[idx] = dparams basic_op_ctxs[idx].saved_tensors = None for idx, dxs in zip(basic_op_idxs, fused_op_grad_extra_inputs): grad_extra_inputs[idx] = dxs + for input_idx, grad in enumerate(dxs): + source = 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 +368,27 @@ 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) + + # One public tensor may fan out to several unmatched slots sharing a + # channel name, so sum their gradients before returning to autograd. + grad_extra_inputs_flat = [] + for slots in func_ctx.external_extra_input_slots: + grad = None + for op_idx, input_idx in slots: + slot_grad = grad_extra_inputs[op_idx][input_idx] + if slot_grad is not None: + grad = slot_grad if grad is None else grad + slot_grad + grad_extra_inputs_flat.append(grad) # Update FP8 scaling factors if func_ctx.is_first_module and not _is_graph_capturing(): @@ -328,6 +428,8 @@ class OperationFuser: def __init__( self, ops: list[FusibleOperation], + *, + lock_extra_channels: bool = True, ) -> None: # Get list of basic operations @@ -342,7 +444,96 @@ def __init__( # Number of extra tensor inputs self._basic_op_num_extra_inputs: list[int] = list(op.num_extra_inputs for op in basic_ops) - self.num_extra_inputs: int = sum(self._basic_op_num_extra_inputs) + self._basic_op_extra_input_sources: list[list[Optional[tuple[int, int]]]] = [ + [None] * op.num_extra_inputs for op in basic_ops + ] + self._basic_op_extra_output_channels: list[list[Optional[str]]] = [ + list(op._extra_output_channels) for op in basic_ops + ] + self._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[list[tuple[int, int]]] = [] + self._external_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. A channel with an earlier producer is internal. A + # channel without a producer is public, with one positional tensor + # fanning out to every public slot that shares the channel name. + external_input_channels: dict[str, int] = {} + consumed_channels: set[str] = set() + for op_idx, op in enumerate(basic_ops): + for input_idx, channel in enumerate(op._extra_input_channels): + if channel is None: + self._external_extra_input_slots.append([(op_idx, input_idx)]) + continue + producer = channel_producers.get(channel) + if producer is None: + group_idx = external_input_channels.get(channel) + if group_idx is None: + group_idx = len(self._external_extra_input_slots) + external_input_channels[channel] = group_idx + self._external_extra_input_slots.append([]) + self._external_extra_input_slots[group_idx].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) + + # Unnamed outputs and named outputs without local consumers are public. + 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 + else: + self._external_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: + 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 +550,11 @@ def __init__( self._basic_op_num_params = list(map(len, self._basic_op_params)) self._flat_basic_op_params = sum(self._basic_op_params, []) + # Persistent fusers capture channel routing as structural state. + if lock_extra_channels: + for op in self._basic_ops: + op._lock_extra_channels() + @staticmethod def _apply_fusions( ops: Iterable[FusibleOperation], @@ -432,7 +628,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 @@ -517,12 +713,14 @@ 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, slots in zip(extra_inputs, self._external_extra_input_slots): + for op_idx, input_idx in 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..d668a7c904 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -187,10 +187,73 @@ 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_locked = False + # 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: + """Assign a channel name to an extra input slot. + + The slot receives the matching extra output from an earlier operation + in the same fuser. If there is no producer in the fuser, the slot + remains public; one public tensor fans out to all input slots with the + same channel name. Passing ``None`` removes the name. Channels cannot + be changed after the operation has been attached to a persistent + ``OperationFuser``. + """ + 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._assert_extra_channels_mutable() + self._extra_input_channels[index] = channel + return self + + def set_extra_output_channel(self, index: int, channel: Optional[str]) -> BasicOperation: + """Assign a channel name to an extra output slot. + + The slot feeds matching extra inputs on later operations in the same + fuser. If there are no consumers in the fuser, the slot remains a + public extra output. Passing ``None`` removes the name. Channels cannot + be changed after the operation has been attached to a persistent + ``OperationFuser``. + """ + 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._assert_extra_channels_mutable() + self._extra_output_channels[index] = channel + return self + + def _assert_extra_channels_mutable(self) -> None: + """Check that channel routing has not been captured by a fuser.""" + if self._extra_channels_locked: + raise RuntimeError( + "Extra tensor channels cannot be changed after an operation has been " + "attached to an OperationFuser" + ) + + def _lock_extra_channels(self) -> None: + """Prevent changes after a fuser has captured the channel routing.""" + self._extra_channels_locked = True + @property def is_fused_op(self) -> bool: return False @@ -535,7 +598,7 @@ def forward( """Apply operation""" from .fuser import OperationFuser - return OperationFuser([self])( + return OperationFuser([self], lock_extra_channels=False)( input, *extra_inputs, basic_op_kwargs=[kwargs], @@ -770,7 +833,7 @@ def forward( basic_op_kwargs = [{} for _ in range(len(self.basic_ops))] from .fuser import OperationFuser - return OperationFuser([self])( + return OperationFuser([self], lock_extra_channels=False)( input, *extra_inputs, basic_op_kwargs=basic_op_kwargs,