From b554e722aa4cee313e46563696c2224b0b206ec1 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 4 Jun 2026 00:06:54 +0000 Subject: [PATCH 1/2] [PyTorch] Add joint forward-backward op fusion pass Introduce a third operation fusion pass for joint forward-backward fusions, applied before the forward-only and backward-only passes. A joint fused op implements both fuser_forward and fuser_backward, so the two halves can cooperate (e.g. the forward saving reduced state that only its own backward knows how to recompute) and need not be individually interchangeable with the unfused ops. Add register_forward_backward_fusion, split the fusion application and basic-op reconciliation in OperationFuser so the forward/backward passes build on the joint grouping, and spell out the interchangeability contracts in the register_*_fusion docstrings. Add a custom joint fusion unit test and a user guide section. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Tim Moon --- docs/api/pytorch.rst | 2 + docs/examples/op_fuser/op_fuser.rst | 69 +++++++++- tests/pytorch/test_fusible_ops.py | 145 +++++++++++++++++++++ transformer_engine/pytorch/ops/__init__.py | 6 +- transformer_engine/pytorch/ops/fuser.py | 138 +++++++++++++++++--- transformer_engine/pytorch/ops/op.py | 23 +++- 6 files changed, 354 insertions(+), 29 deletions(-) diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 99d850d04d..14289095a0 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -166,6 +166,8 @@ Operation fuser .. autoapifunction:: transformer_engine.pytorch.ops.register_backward_fusion +.. autoapifunction:: transformer_engine.pytorch.ops.register_forward_backward_fusion + .. autoapiclass:: transformer_engine.pytorch.ops.Linear .. autoapiclass:: transformer_engine.pytorch.ops.AddExtraInput diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index 9613ba74b3..dd17191e58 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -317,9 +317,10 @@ of context objects for all the corresponding ``BasicOperation`` s. .. warning:: - Remember the contract that the fused operation must produce outputs - that are interchangeable with the corresponding basic operation - outputs. + Forward-only and backward-only fused operations must produce outputs + that are interchangeable with the corresponding basic operations, + since the opposite pass is fused independently. Joint forward-backward + fusions (described below) relax this contract. In order to make these fused operations useful, they should be registered with the operation fuser. To do this, first implement a @@ -351,3 +352,65 @@ and then register it with the ``register_forward_fusion`` or # Register fusion with operation fuser te.ops.register_forward_fusion(fuse_axpy_ops) + +Joint forward-backward fusions +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The fusions above replace operations in either the forward or the +backward pass, and they must remain interchangeable with the unfused +operations because the opposite pass is fused independently. Some +optimizations, however, benefit from co-designing the two passes. For +example, a fused operation's forward pass might skip saving a tensor if +it knows that its backward pass can recompute it. + +A *joint* forward-backward fusion expresses this coupling. It is a single +``FusedOperation`` that implements both ``fuser_forward`` and +``fuser_backward``, registered with ``register_forward_backward_fusion``. +Joint fusions are applied before the forward-only and backward-only +fusion passes, so the same fused operation is seen by both passes. Unlike +forward-only or backward-only fusions, the two halves need not be +individually interchangeable with the unfused operations; only the +forward/backward pair must be jointly equivalent. + +.. code-block:: python + + class LinearSiLU(te.ops.FusedOperation): + + def __init__(self, linear: te.ops.Linear, silu: te.ops.SiLU) -> None: + super().__init__((linear, silu)) # Equivalent basic ops + + def fuser_forward(self, basic_op_ctxs, input_, **unused): + weight = self.basic_ops[0].weight + out = torch.nn.functional.silu(torch.matmul(input_, weight.T)) + # Save reduced state: the backward recomputes the SiLU input + # rather than saving it. + basic_op_ctxs[0].save_for_backward(input_, weight) + return out, [(), ()] + + def fuser_backward(self, basic_op_ctxs, grad_output, **unused): + x, w = basic_op_ctxs[0].saved_tensors + y = torch.matmul(x, w.T) # Recompute SiLU input + s = torch.sigmoid(y) + dy = grad_output * s * (1 + y * (1 - s)) # SiLU backward + return ( + torch.matmul(dy, w), # Grad input + [(torch.matmul(dy.T, x),), ()], # Grad params for each basic op + [(), ()], # Grad extra inputs for each basic op + ) + + def fuse_linear_silu(ops, **unused): + """Sliding window scan to perform LinearSiLU fusion""" + out = [] + window, ops = ops[:2], ops[2:] + while len(window) == 2: + if isinstance(window[0], te.ops.Linear) and isinstance(window[1], te.ops.SiLU): + window = [LinearSiLU(window[0], window[1])] + else: + out.append(window[0]) + window = window[1:] + window, ops = window + ops[:1], ops[1:] + out.extend(window + ops) + return out + + # Register joint fusion with operation fuser + te.ops.register_forward_backward_fusion(fuse_linear_silu) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 57316e019f..2e2fc8a19f 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -5117,6 +5117,151 @@ def fuse_ops( torch.testing.assert_close(dx_test, x_ref.grad, **tols) torch.testing.assert_close(dw_test, w_ref.grad, **tols) + def test_custom_forward_backward_fused_op( + self, + *, + shape: Iterable[int] = (7, 11), + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + ): + """Custom joint forward-backward fused op + + A single fused op implements both ``fuser_forward`` and + ``fuser_backward``. Because the same op owns both passes, the + forward saves reduced state (just the linear input and weight) + and lets its own backward recompute the SiLU input rather than + saving it. + + """ + + class CustomLinearSiLU(te.ops.FusedOperation): + """Custom joint fused op for GEMM + SiLU""" + + _enabled = True + + def __init__(self, *, linear, silu) -> None: + super().__init__((linear, silu)) + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + **unused, + ) -> torch.Tensor: + weight = self.basic_ops[0].weight + dtype = weight.dtype + + # Forward compute + y = torch.matmul(input_, weight.T) + out = torch.nn.functional.silu(y) + + # Save reduced state for the joint backward. Note that we + # do not save the SiLU input ``y``; the backward recomputes + # it from the linear inputs. + linear_op_ctx = basic_op_ctxs[0] + linear_op_ctx.save_for_backward(input_, weight) + linear_op_ctx.dtype = dtype + + return out, [(), ()] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + **unused, + ) -> torch.Tensor: + + # Load reduced state from the joint forward + linear_op_ctx = basic_op_ctxs[0] + x, w = linear_op_ctx.saved_tensors + dtype = linear_op_ctx.dtype + + # Recompute SiLU input and its gradient in FP64 + x = x.double() + w = w.double() + dout = grad_output.double() + y = torch.matmul(x, w.T) + s = torch.sigmoid(y) + dsilu = s * (1 + y * (1 - s)) + dy = dout * dsilu + + # Linear backward + dx = torch.matmul(dy, w).to(dtype=dtype) + dw = torch.matmul(dy.T, x).to(dtype=dtype) + + # grad_input, grad params per basic op, grad extra inputs per basic op + return dx, [(dw,), ()], [(), ()] + + @staticmethod + def fuse_ops( + ops: list[FusibleOperation], + **unused, + ) -> list[FusibleOperation]: + """Apply fusion the first time this function is called""" + if CustomLinearSiLU._enabled: + CustomLinearSiLU._enabled = False + op = CustomLinearSiLU(linear=ops[0], silu=ops[1]) + return [op] + ops[2:] + return ops + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + shape, + test_dtype=dtype, + test_device=device, + ) + w_ref, w_test = make_reference_and_test_tensors( + (shape[-1], shape[-1]), + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Plain PyTorch implementation + y_ref = torch.nn.functional.linear(x_ref, w_ref) + y_ref = torch.nn.functional.silu(y_ref) + y_ref.backward(dy_ref) + + # Implementation with joint fusible operation + te.ops.register_forward_backward_fusion(CustomLinearSiLU.fuse_ops) + model = te.ops.Sequential( + te.ops.Linear(shape[-1], shape[-1], bias=False), + te.ops.SiLU(), + ) + with torch.no_grad(): + model[0].weight.copy_(w_test) + del w_test + y_test = model(x_test) + y_test.backward(dy_test) + + # Check that operations have been fused in both passes, using the + # same fused op object + forward_ops = model._module_groups[0]._forward_ops + backward_ops = model._module_groups[0]._backward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], CustomLinearSiLU) + assert len(backward_ops) == 1 + assert isinstance(backward_ops[0][0], CustomLinearSiLU) + assert forward_ops[0][0] is backward_ops[0][0] + + # Expected numerical error + tols = dtype_tols(dtype) + if dtype == torch.float32: + tols = dtype_tols(torch.float16) # TF32 GEMM + + # Check results + y_test = y_test.to(dtype=torch.float64, device="cpu") + dx_test = x_test.grad.to(dtype=torch.float64, device="cpu") + dw_test = model[0].weight.grad.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(y_test, y_ref, **tols) + torch.testing.assert_close(dx_test, x_ref.grad, **tols) + torch.testing.assert_close(dw_test, w_ref.grad, **tols) + class TestTrainingLoops: diff --git a/transformer_engine/pytorch/ops/__init__.py b/transformer_engine/pytorch/ops/__init__.py index 99f51a9c7a..5e21bb39e8 100644 --- a/transformer_engine/pytorch/ops/__init__.py +++ b/transformer_engine/pytorch/ops/__init__.py @@ -9,7 +9,11 @@ """ from .basic import * -from .fuser import register_backward_fusion, register_forward_fusion +from .fuser import ( + register_backward_fusion, + register_forward_backward_fusion, + register_forward_fusion, +) from .linear import Linear from .op import BasicOperation, FusedOperation, FusibleOperation from .sequential import Sequential diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 5283af8144..33d6f177a9 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -307,6 +307,12 @@ def backward( class OperationFuser: """Manages forward and backward passes for a pipeline of operations + Operations are fused with three passes (see ``register_*_fusion``): + + 1. Joint forward-backward fusions. + 2. Forward-only fusions. + 3. Backward-only fusions. + Parameters ---------- ops : list of FusibleOperation @@ -315,6 +321,7 @@ class OperationFuser: """ # Functions to perform operation fusion + forward_backward_fusion_functions: list[OperationFusionFunction] = [] forward_fusion_functions: list[OperationFusionFunction] = [] backward_fusion_functions: list[OperationFusionFunction] = [] @@ -352,19 +359,30 @@ def __init__( self._basic_op_num_params = list(map(len, self._basic_op_params)) self._flat_basic_op_params = sum(self._basic_op_params, []) - @classmethod - def _fuse_ops( - cls, - basic_ops: Sequence[BasicOperation], + @staticmethod + def _apply_fusions( + ops: Iterable[FusibleOperation], fusion_funcs: Iterable[OperationFusionFunction], recipe: Optional[Recipe], - ) -> list[tuple[FusibleOperation, list[int]]]: - """Apply operation fusions""" - - # Apply op fusions - fused_ops = list(basic_ops) + ) -> list[FusibleOperation]: + """Apply a sequence of fusion functions to a list of ops""" + fused_ops = list(ops) for func in fusion_funcs: fused_ops = func(fused_ops, recipe=recipe) + return fused_ops + + @staticmethod + def _match_basic_ops( + fused_ops: Sequence[FusibleOperation], + basic_ops: Sequence[BasicOperation], + ) -> list[tuple[FusibleOperation, list[int]]]: + """Map a fused op list back to basic op indices + + Verifies that the fused ops expand to exactly ``basic_ops`` in + order, and annotates each (possibly fused) op with the indices + of the basic ops it covers. + + """ def raise_mismatch_error() -> None: """Throw error indicating invalid op fusion""" @@ -381,13 +399,13 @@ def raise_mismatch_error() -> None: if isinstance(op, FusedOperation): idxs = [] for basic_op in op.basic_ops: - if basic_op is not basic_ops[idx]: + if idx >= len(basic_ops) or basic_op is not basic_ops[idx]: raise_mismatch_error() idxs.append(idx) idx += 1 out.append((op, idxs)) else: - if op is not basic_ops[idx]: + if idx >= len(basic_ops) or op is not basic_ops[idx]: raise_mismatch_error() out.append((op, [idx])) idx += 1 @@ -449,16 +467,29 @@ def maybe_fuse_ops( for op in self._basic_ops: op.pre_first_fuser_forward() - # Prepare basic op lists for fusions - self._forward_ops = OperationFuser._fuse_ops( + # Apply joint forward-backward fusions first + joint_ops = OperationFuser._apply_fusions( self._basic_ops, - OperationFuser.forward_fusion_functions, + OperationFuser.forward_backward_fusion_functions, recipe=recipe, ) - self._backward_ops = OperationFuser._fuse_ops( + + # Apply forward-only and backward-only fusions + self._forward_ops = OperationFuser._match_basic_ops( + OperationFuser._apply_fusions( + joint_ops, + OperationFuser.forward_fusion_functions, + recipe=recipe, + ), + self._basic_ops, + ) + self._backward_ops = OperationFuser._match_basic_ops( + OperationFuser._apply_fusions( + joint_ops, + OperationFuser.backward_fusion_functions, + recipe=recipe, + ), self._basic_ops, - OperationFuser.backward_fusion_functions, - recipe=recipe, ) # Save current fusion params @@ -525,11 +556,64 @@ def __call__( return _OperationFuserAutogradFunction.apply(*args) +def register_forward_backward_fusion( + op_fusion_func: OperationFusionFunction, + prepend: bool = False, +) -> None: + """Register a joint forward-backward operation fusion. + + A joint fusion replaces a run of basic ops with a single fused op + that implements *both* ``fuser_forward`` and ``fuser_backward``. + Unlike forward-only or backward-only fusions (see + ``register_forward_fusion`` / ``register_backward_fusion``), the two + halves need not be individually interchangeable with the unfused + ops; only the forward/backward pair must be jointly equivalent. This + lets the forward pass cooperate with its own backward, e.g. saving + state that only its backward knows how to handle. + + Joint fusions are applied before the forward-only and backward-only + fusion passes, so a joint fused op is seen by both passes. The + forward-only and backward-only passes then fuse the remaining ops + independently. + + The fusion function should have the following signature: + + .. code-block:: python + + func(ops, *, recipe) -> updated ops + + Parameters + ---------- + op_fusion_func: function + Function that takes a list of operations and may substitute + them with fused operations. + prepend: bool, default = ``False`` + Whether the operation fuser should apply this fusion function + first within the joint fusion pass. The default is to apply it + last. + + """ + if prepend: + OperationFuser.forward_backward_fusion_functions.insert(0, op_fusion_func) + else: + OperationFuser.forward_backward_fusion_functions.append(op_fusion_func) + + def register_forward_fusion( op_fusion_func: OperationFusionFunction, prepend: bool = False, ) -> None: - """Register function to perform operation fusion for forward pass. + """Register a forward-only operation fusion. + + A forward-only fusion replaces a run of basic ops with a single + fused op that implements ``fuser_forward``. Because the backward + pass is fused independently (see ``register_backward_fusion``), the + fused op's forward must be interchangeable with the corresponding + basic ops' forward: it must produce the same output and save state in + each basic op's context that the unfused backward can consume. If the + forward and backward need to cooperate (e.g. the forward saving + reduced state that only a matching backward can handle), use + ``register_forward_backward_fusion`` instead. The fusion function should have the following signature: @@ -544,7 +628,8 @@ def register_forward_fusion( them with fused operations. prepend: bool, default = ``False`` Whether the operation fuser should apply this fusion function - first. The default is to apply it last. + first within the forward fusion pass. The default is to apply it + last. """ if prepend: @@ -557,7 +642,17 @@ def register_backward_fusion( op_fusion_func: OperationFusionFunction, prepend: bool = False, ) -> None: - """Register function to perform operation fusion for backward pass. + """Register a backward-only operation fusion. + + A backward-only fusion replaces a run of basic ops with a single + fused op that implements ``fuser_backward``. Because the forward + pass is fused independently (see ``register_forward_fusion``), the + fused op's backward must be interchangeable with the corresponding + basic ops' backward: it must consume the state saved in each basic + op's context by the unfused forward and produce the same gradients. + If the forward and backward need to cooperate (e.g. the forward + saving reduced state that only a matching backward can handle), use + ``register_forward_backward_fusion`` instead. The fusion function should have the following signature: @@ -572,7 +667,8 @@ def register_backward_fusion( them with fused operations. prepend: bool, default = ``False`` Whether the operation fuser should apply this fusion function - first. The default is to apply it last. + first within the backward fusion pass. The default is to apply it + last. """ if prepend: diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 1687187230..86bd60ed9c 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -697,10 +697,25 @@ def _load_from_state_dict(self, *args, **kwargs) -> None: class FusedOperation(FusibleOperation): """Compound tensor operation supported by the operation fuser - If the forward or backward passes are defined, they must be - functionally equivalent to the forward/backward passes of the - corresponding basic ops. This class should hold no parameters or - other state, but should access them from the basic ops. + A fused op corresponds to a run of basic ops. Depending on which + fusion pass produces it (see ``fuser.py``), the equivalence contract + differs: + + - Forward-only or backward-only fused ops (from + ``register_forward_fusion`` / ``register_backward_fusion``): the + defined pass must be functionally equivalent to the corresponding + basic ops' pass, since the opposite pass may be fused + independently. + - Joint forward-backward fused ops (from + ``register_forward_backward_fusion``): the op implements both + ``fuser_forward`` and ``fuser_backward``, and only the pair must + be jointly equivalent to the basic ops' forward and backward. The + two halves need not be individually interchangeable, so the + forward may save state that only its own backward knows + how to consume. + + This class should hold no parameters or other state, but should + access them from the basic ops. Parameters ---------- From 54da51cb28628234040eab03dfa6d879e3f9af35 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Mon, 8 Jun 2026 15:45:03 -0700 Subject: [PATCH 2/2] Review suggestion from @vthumbe1503 Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- transformer_engine/pytorch/ops/fuser.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 33d6f177a9..09ffb004dd 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -372,7 +372,7 @@ def _apply_fusions( return fused_ops @staticmethod - def _match_basic_ops( + def _map_to_basic_ops( fused_ops: Sequence[FusibleOperation], basic_ops: Sequence[BasicOperation], ) -> list[tuple[FusibleOperation, list[int]]]: @@ -475,7 +475,7 @@ def maybe_fuse_ops( ) # Apply forward-only and backward-only fusions - self._forward_ops = OperationFuser._match_basic_ops( + self._forward_ops = OperationFuser._map_to_basic_ops( OperationFuser._apply_fusions( joint_ops, OperationFuser.forward_fusion_functions, @@ -483,7 +483,7 @@ def maybe_fuse_ops( ), self._basic_ops, ) - self._backward_ops = OperationFuser._match_basic_ops( + self._backward_ops = OperationFuser._map_to_basic_ops( OperationFuser._apply_fusions( joint_ops, OperationFuser.backward_fusion_functions,