diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index dd17191e58..7b003091e0 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -113,43 +113,115 @@ quantized compute. Branching operations ^^^^^^^^^^^^^^^^^^^^ -The operation fuser supports very 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. +The operation fuser supports limited branching behavior. While the +operations must be in sequential order, basic operations may declare +extra tensor inputs and outputs. By default, an extra tensor slot has +no channel assigned and is part of the public ``Sequential`` interface: +the caller provides extra inputs as arguments, and extra outputs are +returned after the main output. Assigning the same channel name to an +output slot and a later input slot connects them internally instead. .. code-block:: python import torch import transformer_engine.pytorch as te - # Construct MLP with residual connection - fc1 = te.ops.Sequential( + # Keep a residual connection inside one Sequential. + 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), - te.ops.MakeExtraOutput(), # Output residual + make_residual, te.ops.Linear(4096, 28672), te.ops.SwiGLU(), - ) - fc2 = te.ops.Sequential( te.ops.Linear(14336, 4096), - te.ops.AddExtraInput(), # Add residual + add_residual, ) - # Forward pass x = torch.randn(16384, 4096, device="cuda") - y, residual = fc1(x) - y = fc2(y, residual) + y = block(x) .. 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. + +Extra tensor channels +""""""""""""""""""""" + +An extra output and one or more later extra inputs can be assigned the +same channel name. This routes the tensor inside the +``OperationFuser`` and removes the bound slots from the public +``Sequential`` interface. In the residual example above, the caller +therefore receives only ``y`` and does not need to pass the residual +back into the block. + +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) + +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..1616d830ff 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,295 @@ 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)) + + def test_fused_internal_residual_connection(self, size: int = 16) -> None: + """A fused op can implement both ends of an internal channel.""" + + class FusedResidual(te_ops.FusedOperation): + """Fuse MakeExtraOutput, Bias, and AddExtraInput in forward.""" + + _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_, [(input_,), (), ()] + + 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.Identity) + and isinstance(ops[2], te_ops.AddExtraInput) + ): + FusedResidual._enabled = False + return [FusedResidual(*ops)] + return ops + + residual = te_ops.MakeExtraOutput() + body = te_ops.Identity() + 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) + + te_ops.register_forward_fusion(fuse_residual, prepend=True) + x = torch.rand((size,), requires_grad=True) + y = model(x) + + forward_ops = model._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], FusedResidual) + torch.testing.assert_close(y, 2 * x) + y.sum().backward() + torch.testing.assert_close(x.grad, torch.full_like(x, 2)) + + def test_internal_extra_tensor_channel_fanout(self, size: int = 16) -> None: + """An internal extra output can feed multiple later consumers.""" + producer = te_ops.MakeExtraOutput() + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer1.set_extra_input_channel(0, "route") + consumer2.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer1, consumer2) + + x = torch.rand((size,), requires_grad=True) + y = model(x) + + # Main path: x -> x + route -> x + route + route. + torch.testing.assert_close(y, 3 * x) + y.sum().backward() + # The channel fan-out contributes two independent gradient paths. + torch.testing.assert_close(x.grad, torch.full_like(x, 3)) + + # Internal slots are unavailable before forward, so grad discovery + # must tolerate them when no public input requires gradients. + x_no_grad = x.detach() + torch.testing.assert_close(model(x_no_grad), 3 * x_no_grad) + + def test_internal_and_external_extra_tensor_inputs(self, size: int = 16) -> None: + """Unbound slots remain public when other slots use internal channels.""" + producer = te_ops.MakeExtraOutput() + internal_consumer = te_ops.AddExtraInput() + external_consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + internal_consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, internal_consumer, external_consumer) + + x = torch.rand((size,), requires_grad=True) + extra = torch.rand((size,), requires_grad=True) + y = model(x, extra) + + torch.testing.assert_close(y, 2 * x + extra) + y.sum().backward() + torch.testing.assert_close(x.grad, torch.full_like(x, 2)) + torch.testing.assert_close(extra.grad, torch.ones_like(extra)) + + def test_consumer_channel_without_producer(self) -> None: + """Extra input bound to a channel that no earlier op produces.""" + consumer = te_ops.AddExtraInput() + consumer.set_extra_input_channel(0, "missing") + with pytest.raises(ValueError, match="has no earlier producer"): + OperationFuser([consumer]) + + def test_consumer_before_producer(self) -> None: + """Channels only connect forward; a later producer does not satisfy an earlier consumer.""" + consumer = te_ops.AddExtraInput() + producer = te_ops.MakeExtraOutput() + consumer.set_extra_input_channel(0, "route") + producer.set_extra_output_channel(0, "route") + with pytest.raises(ValueError, match="has no earlier producer"): + OperationFuser([consumer, producer]) + + def test_set_extra_channel_rejects_invalid_index(self) -> None: + """Slot indices must be in range; negatives and OOB are rejected at bind time.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + + with pytest.raises(IndexError, match="out of range"): + producer.set_extra_output_channel(-1, "route") + with pytest.raises(IndexError, match="out of range"): + producer.set_extra_output_channel(1, "route") + with pytest.raises(IndexError, match="out of range"): + consumer.set_extra_input_channel(-1, "route") + with pytest.raises(IndexError, match="out of range"): + consumer.set_extra_input_channel(1, "route") + + def test_set_extra_channel_rejects_invalid_name(self) -> None: + """Channel names must be non-empty strings.""" + producer = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + + with pytest.raises(ValueError, match="non-empty string"): + producer.set_extra_output_channel(0, "") + with pytest.raises(ValueError, match="non-empty string"): + consumer.set_extra_input_channel(0, "") + with pytest.raises(ValueError, match="non-empty string"): + producer.set_extra_output_channel(0, 123) # type: ignore[arg-type] + + def test_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_unused_extra_output_channel(self) -> None: + """Every produced channel must have at least one consumer.""" + producer = te_ops.MakeExtraOutput() + producer.set_extra_output_channel(0, "orphan") + with pytest.raises(ValueError, match="have no consumers"): + OperationFuser([producer]) + + def test_one_extra_input_has_single_source(self) -> None: + """Each extra-input slot binds to one channel / one producer source. + + Rebinding replaces the previous name; the abandoned producer channel + then fails as unused rather than attaching two sources to one input. + """ + producer_a = te_ops.MakeExtraOutput() + producer_b = te_ops.MakeExtraOutput() + consumer = te_ops.AddExtraInput() + producer_a.set_extra_output_channel(0, "a") + producer_b.set_extra_output_channel(0, "b") + consumer.set_extra_input_channel(0, "a") + consumer.set_extra_input_channel(0, "b") + with pytest.raises(ValueError, match="have no consumers"): + OperationFuser([producer_a, producer_b, consumer]) + + # Valid single binding: consumer input 0 is fed only by producer_a. + consumer.set_extra_input_channel(0, "a") + producer_b.set_extra_output_channel(0, None) + fuser = OperationFuser([producer_a, producer_b, consumer]) + assert fuser._basic_op_extra_input_sources[2] == [(0, 0)] + assert fuser.num_extra_inputs == 0 + # producer_b's unbound extra output remains public + assert fuser._external_extra_output_slots == [(1, 0)] + + def test_channel_fanout_accumulates_grads(self, size: int = 16) -> None: + """Grads from every consumer of a channel are accumulated into the producer.""" + producer = te_ops.MakeExtraOutput() + consumer1 = te_ops.AddExtraInput() + consumer2 = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + consumer1.set_extra_input_channel(0, "route") + consumer2.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, consumer1, consumer2) + + x = torch.rand((size,), requires_grad=True) + y = model(x) + # Forward: x -> x+x -> x+x+x + torch.testing.assert_close(y, 3 * x) + + dy = torch.rand((size,)) + y.backward(dy) + # Main path contributes dy; each AddExtraInput also routes dy back + # through the channel into MakeExtraOutput's extra-output grad, which + # is added again into dx. Total: dy (main) + dy + dy (two consumers). + torch.testing.assert_close(x.grad, 3 * dy) + + def test_mixed_internal_external_grad(self, size: int = 16) -> None: + """Internal channel grads and public extra-input grads both flow correctly.""" + producer = te_ops.MakeExtraOutput() + internal_consumer = te_ops.AddExtraInput() + external_consumer = te_ops.AddExtraInput() + producer.set_extra_output_channel(0, "route") + internal_consumer.set_extra_input_channel(0, "route") + model = te_ops.Sequential(producer, internal_consumer, external_consumer) + + x = torch.rand((size,), requires_grad=True) + extra = torch.rand((size,), requires_grad=True) + y = model(x, extra) + torch.testing.assert_close(y, 2 * x + extra) + + dy = torch.rand((size,)) + y.backward(dy) + torch.testing.assert_close(x.grad, 2 * dy) + torch.testing.assert_close(extra.grad, dy) + + class TestFuser: """Tests for operation fusion infrastructure""" diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index be931829ea..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..c4d7e1ec10 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -102,12 +102,14 @@ 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,8 +120,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,18 +162,22 @@ def forward( x, fused_op_extra_outputs = op.fuser_forward( [basic_op_ctxs[idx] for idx in basic_op_idxs], x, - basic_op_extra_inputs=extra_inputs, + basic_op_extra_inputs=op_extra_inputs, prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, next_op_input_quantizer=next_op_input_quantizer, basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], ) for idx, ys in zip(basic_op_idxs, fused_op_extra_outputs): - for y in ys: - if set_output_requires_grad: - y.requires_grad_(idx >= fuser.first_op_requiring_backward) + 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}" + ) extra_outputs[idx] = ys - # Flatten list of extra outputs + # Validate extra outputs and flatten only public slots. Outputs bound + # to channels stay internal to the fuser and are not marked. extra_outputs_flat = [] for idx, ys in enumerate(extra_outputs): ys = list(ys) @@ -156,7 +188,13 @@ def forward( "{num_extra_outputs} extra inputs, " f"but got {len(ys)}" ) - extra_outputs_flat.extend(ys) + for output_idx, y in enumerate(ys): + # Output is not bound to a channel consumed by another op, + # so it is a public output. + if fuser._basic_op_extra_output_channels[idx][output_idx] is None: + if set_output_requires_grad: + y.requires_grad_(idx >= fuser.first_op_requiring_backward) + extra_outputs_flat.append(y) # Save context for backward pass if func_ctx is not None: @@ -181,17 +219,21 @@ 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_output_slots = fuser._external_extra_output_slots + func_ctx.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels + func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources func_ctx.is_first_module = is_first_module # Mark output tensors as not deletable in backward - 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 +266,29 @@ 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_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 +296,37 @@ 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 channel is not None: + basic_op_grad_extra_outputs[idx][output_idx] = channel_grads.get(channel) + op_grad_extra_outputs = [ + tuple(basic_op_grad_extra_outputs[idx]) for idx in basic_op_idxs + ] dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( [basic_op_ctxs[idx] for idx in basic_op_idxs], dx, - basic_op_grad_extra_outputs=grad_extra_outputs, + basic_op_grad_extra_outputs=op_grad_extra_outputs, ) for idx, dparams in zip(basic_op_idxs, fused_op_grad_params): grad_params[idx] = dparams basic_op_ctxs[idx].saved_tensors = None for idx, dxs in zip(basic_op_idxs, fused_op_grad_extra_inputs): grad_extra_inputs[idx] = dxs + for input_idx, grad in enumerate(dxs): + source = 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 = [] @@ -288,7 +357,11 @@ def backward( f"for {num_extra_inputs} extra inputs, " f"but got {len(dxs)}" ) - grad_extra_inputs_flat.extend(dxs) + for input_idx, grad in enumerate(dxs): + # Only append public grad extra inputs to the list + # to be returned to the user. + if basic_op_extra_input_sources[idx][input_idx] is None: + grad_extra_inputs_flat.append(grad) # Update FP8 scaling factors if func_ctx.is_first_module and not _is_graph_capturing(): @@ -342,7 +415,81 @@ def __init__( # Number of extra tensor inputs self._basic_op_num_extra_inputs: list[int] = list(op.num_extra_inputs for op in basic_ops) - self.num_extra_inputs: int = sum(self._basic_op_num_extra_inputs) + self._basic_op_extra_input_sources: list[list[Optional[tuple[int, int]]]] = [ + [None] * op.num_extra_inputs for op in basic_ops + ] + self._basic_op_extra_output_channels: list[list[Optional[str]]] = [ + list(op._extra_output_channels) for op in basic_ops + ] + self._external_extra_input_slots: list[tuple[int, int]] = [] + self._external_extra_output_slots: list[tuple[int, int]] = [] + + # Resolve named channels in pipeline order. Channels deliberately only + # connect an output to later inputs, which keeps execution acyclic. + channel_producers: dict[str, tuple[int, int]] = {} + consumed_channels: set[str] = set() + for op_idx, op in enumerate(basic_ops): + for input_idx in range(op.num_extra_inputs): + channel = op._extra_input_channels[input_idx] + if channel is None: + self._external_extra_input_slots.append((op_idx, input_idx)) + continue + if channel not in channel_producers: + raise ValueError( + f"Extra tensor channel {channel!r} consumed by op {op_idx} " + f"({type(op).__name__}) has no earlier producer" + ) + self._basic_op_extra_input_sources[op_idx][input_idx] = channel_producers[channel] + consumed_channels.add(channel) + for output_idx, channel in enumerate(self._basic_op_extra_output_channels[op_idx]): + if channel is None: + self._external_extra_output_slots.append((op_idx, output_idx)) + continue + if channel in channel_producers: + producer_idx, _ = channel_producers[channel] + raise ValueError( + f"Extra tensor channel {channel!r} has multiple producers " + f"(ops {producer_idx} and {op_idx})" + ) + channel_producers[channel] = (op_idx, output_idx) + + unused_channels = channel_producers.keys() - consumed_channels + if unused_channels: + channels = ", ".join(repr(channel) for channel in sorted(unused_channels)) + raise ValueError(f"Extra tensor channels have no consumers: {channels}") + + # Every channel-bound extra input must be wired to a matching producer + # extra output. External slots remain unbound (source is None). + for op_idx, sources in enumerate(self._basic_op_extra_input_sources): + op = basic_ops[op_idx] + for input_idx, source in enumerate(sources): + channel = op._extra_input_channels[input_idx] + if channel is None: + if source is not None: + raise RuntimeError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is external but has a " + f"producer source {source}" + ) + continue + if source is None: + raise ValueError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is bound to channel {channel!r} " + "but has no producer" + ) + producer_idx, output_idx = source + producer_channel = self._basic_op_extra_output_channels[producer_idx][output_idx] + if producer_channel != channel: + raise ValueError( + f"Extra input {input_idx} of op {op_idx} " + f"({type(op).__name__}) is bound to channel {channel!r}, " + f"but producer op {producer_idx} extra output {output_idx} " + f"is bound to {producer_channel!r}" + ) + # 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 +506,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, []) + # Channel routing is structural fuser state. Prevent the basic ops from + # changing their bindings after this fuser captures them. + for op in self._basic_ops: + op._lock_extra_channels() + @staticmethod def _apply_fusions( ops: Iterable[FusibleOperation], @@ -432,7 +584,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 +669,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..d9f9624ab0 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -187,10 +187,70 @@ 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: + """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. Bindings cannot be changed after + the operation has been attached to an ``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: + """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. Bindings + cannot be changed after the operation has been attached to an + ``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