diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 926724250f..5fac0a89a6 100644 --- a/docs/api/pytorch.rst +++ b/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 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 14a52249f3..7c75d11e3b 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..09ffb004dd 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 _map_to_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._map_to_basic_ops( + OperationFuser._apply_fusions( + joint_ops, + OperationFuser.forward_fusion_functions, + recipe=recipe, + ), + self._basic_ops, + ) + self._backward_ops = OperationFuser._map_to_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 ----------