[PyTorch] Add joint forward-backward op fusion pass

docs/api/pytorch.rst

@@ -176,6 +176,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

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)

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:
 

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