From dbc941ff8a571fad156ac8817bfed6e95e227bc2 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 22:22:46 +0200 Subject: [PATCH] [PyTorch] Declare Bias and the activations as custom ops Each declares its two argument containers and four compute classmethods, so BasicOperation registers its custom ops and the fuser can run it under torch.compile. Nothing in the eager path changes: op_forward / op_backward drive the same implementations. The activations share one implementation and dispatch to their per-class kernel through cls, so all ten subclasses get their own registered op without a factory. Their dispatch methods become staticmethods, which is what 'the compute half must not depend on an instance' means in practice. They are also the first operations here that hand back a quantized tensor: with a next-operation input quantizer the kernel writes FP8 directly, and the tensor crosses the op boundary as its inner buffers, rebuilt on the far side from the fake's TensorSpec. Two contracts the implementations had to respect. A custom op may not return one of its own inputs, so a backward that passes its gradient through returns None for that slot and the caller substitutes; cloning would cost a full-size copy on the common path. The same trick does not work for saved tensors -- the fake cannot see strides, so it cannot predict whether contiguous() will be a no-op, and the mismatch surfaces as an inductor assertion on the sentinel's rank. Hence the activations keep their input only when cache_quantized_input is set. test_ops_custom_ops.py is driven by a single list of operations: adding one means adding one entry. test_ops_hop_poc.py runs a three-operation pipeline whose backward walks a coarser grouping than its forward, which an operation with per-op autograd could not express. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_ops_custom_ops.py | 329 ++++++++++++++++++ tests/pytorch/test_ops_hop_poc.py | 221 ++++++++++++ .../pytorch/ops/basic/activation.py | 270 ++++++++++---- transformer_engine/pytorch/ops/basic/bias.py | 136 ++++++-- 4 files changed, 850 insertions(+), 106 deletions(-) create mode 100644 tests/pytorch/test_ops_custom_ops.py create mode 100644 tests/pytorch/test_ops_hop_poc.py diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py new file mode 100644 index 0000000000..2ad8464c71 --- /dev/null +++ b/tests/pytorch/test_ops_custom_ops.py @@ -0,0 +1,329 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Per-op custom ops for ``transformer_engine.pytorch.ops``. + +An operation opts into ``torch.compile`` by declaring its two argument +containers and implementing four compute halves; ``BasicOperation`` registers +the custom ops and drives them. The checks here are the same for every such +operation and are driven by ``_OP_CASES`` -- adding an operation means adding +one entry. + +Each operation is checked three ways: + +* the data-free fake agrees with the real impl, slot for slot -- the compiled + path slices a flat ``Tensor[]`` payload by what the fake said, so a + disagreement is a silently misassembled tensor rather than an error; +* the registered op reproduces the eager operation; +* both halves trace under ``fullgraph=True``, with and without an FP8 output. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Optional, Tuple + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch.constants import DType +from transformer_engine.pytorch.dynamo import TensorSpec +from transformer_engine.pytorch.dynamo.custom_op import _spec_slot_count, _value_to_flat_tensors +from transformer_engine.pytorch.ops.op import OperationContext +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage +from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer + +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +_device = "cuda" +_dtype = torch.bfloat16 +_HIDDEN = 64 + + +@pytest.fixture(autouse=True) +def _fresh_dynamo(): + """Compile each case from scratch. + + The compiled helpers below are closures over one operation, so every case + recompiles the same code object; without a reset the parametrized runs walk + into Dynamo's recompilation limit and fall back to eager, which is an error + under ``fullgraph=True``. + """ + torch._dynamo.reset() + yield + torch._dynamo.reset() + + +def _fp8_quantizer() -> Any: + """An FP8 quantizer, standing in for the next operation's input quantizer.""" + quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, torch.device(_device)) + quantizer.set_usage(rowwise=True, columnwise=False) + return quantizer + + +# --------------------------------------------------------------------------- # +# The operations under test +# --------------------------------------------------------------------------- # + + +def _build_bias(): + op = te.ops.Bias(_HIDDEN, device=_device, dtype=_dtype) + with torch.no_grad(): + op.bias.copy_(torch.randn_like(op.bias)) + return op + + +@dataclass +class OpCase: + """One operation and how to build it.""" + + name: str + build: Callable[[], Any] + quantizes_output: bool = True + num_grads: int = 1 + in_shape: Tuple[int, ...] = (16, _HIDDEN) + + +_OP_CASES = [ + OpCase(name="Bias", build=_build_bias, quantizes_output=False, num_grads=2), + *( + OpCase(name=cls.__name__, build=cls) + for cls in (te.ops.GELU, te.ops.ReLU, te.ops.SiLU, te.ops.GEGLU, te.ops.ReGLU) + ), +] +_CASE_IDS = [case.name for case in _OP_CASES] +_QUANTIZING = [case for case in _OP_CASES if case.quantizes_output] + + +def _make_input(case: OpCase, requires_grad: bool = False) -> torch.Tensor: + return torch.randn(*case.in_shape, device=_device, dtype=_dtype, requires_grad=requires_grad) + + +def _resolve(op, x, *, fp8_output: bool): + return op.resolve_fwd_args( + x, + requires_grad=True, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=_fp8_quantizer() if fp8_output else None, + ) + + +def _forward_through_ctx(op, x, *, fp8_output: bool): + """Run the eager forward and hand back a context the backward can read.""" + ctx = OperationContext() + ctx.requires_grad = True + y = op.op_forward( + ctx, + x, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=_fp8_quantizer() if fp8_output else None, + ) + ctx.saved_tensors = ctx.to_save + return y, ctx + + +# --------------------------------------------------------------------------- # +# Conformance: the fake must describe what the real impl produces +# --------------------------------------------------------------------------- # + + +def _geometry(value: Any) -> Optional[Tuple]: + """Shape / dtype / quantizer type of a real value, or ``None`` for a sentinel.""" + if value is None: + return None + quantizer = getattr(value, "_quantizer", None) + if not isinstance(value, QuantizedTensorStorage): + quantizer = None + return (tuple(value.shape), value.dtype, type(quantizer) if quantizer else None) + + +def _spec_geometry(spec: Optional[TensorSpec]) -> Optional[Tuple]: + if spec is None: + return None + return (tuple(spec.shape), spec.dtype, type(spec.quantizer) if spec.quantizer else None) + + +def _as_sequence(values: Any) -> Tuple: + if values is None: + return () + if not isinstance(values, (tuple, list)): + return (values,) + return tuple(values) + + +def assert_values_match_specs(real: Any, specs: Any, what: str) -> None: + """Require that the fake describes what the real impl produced. + + The invariant the compiled path relies on is the flat ``Tensor[]`` slot + layout, so that is what gets checked, using the framework's own helpers. + Geometry is compared on top. + """ + real_seq, spec_seq = _as_sequence(real), _as_sequence(specs) + assert len(real_seq) == len(spec_seq), f"{what}: count differs" + for i, (value, spec) in enumerate(zip(real_seq, spec_seq)): + real_slots = len(_value_to_flat_tensors(value)) + fake_slots = _spec_slot_count(spec) + assert real_slots == fake_slots, f"{what}[{i}]: {real_slots} slots vs fake {fake_slots}" + assert _geometry(value) == _spec_geometry(spec), f"{what}[{i}]: geometry differs" + + +# --------------------------------------------------------------------------- # +# Tests +# --------------------------------------------------------------------------- # + + +@_cuda +@pytest.mark.parametrize("case", _OP_CASES, ids=_CASE_IDS) +@pytest.mark.parametrize("fp8_output", [False, True]) +def test_op_fake_matches_real(case: OpCase, fp8_output: bool) -> None: + """The fake must predict output geometry, including whether it is quantized.""" + op = case.build() + cls = type(op) + x = _make_input(case) + args = _resolve(op, x, fp8_output=fp8_output) + + real_out, real_saved, real_attrs = cls.forward_compute(args) + fake_out, fake_saved, fake_attrs = cls.forward_fake(args) + assert_values_match_specs(real_out, fake_out, "forward output") + assert_values_match_specs(real_saved, fake_saved, "saved tensor") + assert set(real_attrs) == set(fake_attrs), "ctx_attrs keys disagree" + + _y, ctx = _forward_through_ctx(op, x, fp8_output=fp8_output) + dy = torch.randn(*real_out.shape, device=_device, dtype=_dtype) + bwd_args = op.resolve_bwd_args(ctx, dy) + real_grads = _as_sequence(cls.backward_compute(bwd_args)) + fake_grads = _as_sequence(cls.backward_fake(bwd_args)) + assert len(real_grads) == len(fake_grads) == case.num_grads + for i, (value, spec) in enumerate(zip(real_grads, fake_grads)): + # One slot per gradient regardless of quantization, so only geometry matters. + assert _geometry(value) == _spec_geometry(spec), f"gradient[{i}]: geometry differs" + + +@_cuda +@pytest.mark.parametrize("case", _OP_CASES, ids=_CASE_IDS) +def test_op_matches_eager(case: OpCase) -> None: + """The registered op must reproduce the eager operation.""" + op = case.build() + assert op.compile_ops is not None, f"{case.name}: custom ops failed to register" + forward_fn, backward_fn = op.compile_ops + + x = _make_input(case, requires_grad=True) + y_ref = op(x) + dy = torch.randn_like(y_ref) + y_ref.backward(dy) + dx_ref = x.grad.clone() + + y, _saved, _attrs = forward_fn(_resolve(op, x.detach(), fp8_output=False)) + _y_eager, ctx = _forward_through_ctx(op, x.detach(), fp8_output=False) + grads = backward_fn(op.resolve_bwd_args(ctx, dy)) + # A None gradient means "the incoming gradient, unchanged" -- a custom op may + # not return one of its own inputs. + dx = grads[0] if grads[0] is not None else dy + + torch.testing.assert_close(y, y_ref) + torch.testing.assert_close(dx, dx_ref) + + +@_cuda +@pytest.mark.parametrize("case", _OP_CASES, ids=_CASE_IDS) +@pytest.mark.parametrize("fp8_output", [False, True]) +def test_op_compiles_fullgraph(case: OpCase, fp8_output: bool) -> None: + """Every operation's halves must trace under ``fullgraph=True``. + + Run under ``no_grad``: the halves carry no autograd of their own (that is the + point of ``register_custom_op``), so a caller wires them into + ``autograd.Function`` -- see ``test_ops_hop_poc.py``. + """ + op = case.build() + assert op.compile_ops is not None, f"{case.name}: custom ops failed to register" + forward_fn, backward_fn = op.compile_ops + x = _make_input(case) + + def fwd(x_): + out, saved, _attrs = forward_fn(_resolve(op, x_, fp8_output=fp8_output)) + # An FP8 output crosses the boundary as its inner buffers and is rebuilt + # on the far side; compare the buffers, since dequantize is not traceable. + if isinstance(out, QuantizedTensorStorage): + return (out._data, out._scale_inv, *saved) + return (out, *saved) + + with torch.no_grad(): + expected = fwd(x) + got = torch.compile(fwd, fullgraph=True)(x) + assert len(got) == len(expected) + for a, b in zip(got, expected): + torch.testing.assert_close(a, b) + + y, ctx = _forward_through_ctx(op, x, fp8_output=fp8_output) + dy = torch.randn(*y.shape, device=_device, dtype=_dtype) + bwd_args = op.resolve_bwd_args(ctx, dy) + + def bwd(args): + return backward_fn(args) + + with torch.no_grad(): + expected_grads = bwd(bwd_args) + got_grads = torch.compile(bwd, fullgraph=True)(bwd_args) + for a, b in zip(got_grads, expected_grads): + if a is None or b is None: + assert a is b is None + continue + torch.testing.assert_close(a, b) + + +@_cuda +@pytest.mark.parametrize("case", _QUANTIZING, ids=[c.name for c in _QUANTIZING]) +def test_op_returns_fp8(case: OpCase) -> None: + """With a next-operation quantizer the op must hand back an FP8 tensor. + + This is what matters for a pipeline: an operation quantizes its output with + the *next* operation's input quantizer, so a quantized tensor crosses the + custom-op boundary -- as its inner buffers, rebuilt on the far side -- rather + than being dequantized at it. + """ + op = case.build() + forward_fn, _ = op.compile_ops + y, _saved, _attrs = forward_fn(_resolve(op, _make_input(case), fp8_output=True)) + assert isinstance(y, QuantizedTensorStorage), f"expected a quantized output, got {type(y)}" + + +@_cuda +@pytest.mark.parametrize("case", _OP_CASES, ids=_CASE_IDS) +def test_op_is_compile_supported(case: OpCase) -> None: + """A converted operation must not gate itself out in a plain configuration.""" + op = case.build() + assert op.compile_unsupported_reason() is None + + +@_cuda +def test_unconverted_op_reports_a_reason() -> None: + """An operation without the compute halves must say so rather than compile.""" + op = te.ops.Identity() + assert op.compile_ops is None + reason = op.compile_unsupported_reason() + assert reason is not None and "compute halves" in reason + + +@_cuda +def test_non_value_opaque_quantizer_is_gated() -> None: + """A quantizer torch.compile cannot specialize on must be refused. + + Delayed scaling is the case that matters: its quantizer holds live + scale/amax tensors, so baking it into the graph would silently freeze stale + scales. None of the operations converted so far hold quantizers, so the + check is driven directly until one that does lands. + """ + from transformer_engine.common.recipe import DelayedScaling + from transformer_engine.pytorch.quantization import RecipeState + + op = te.ops.GELU() + assert op.compile_unsupported_reason() is None + + state = RecipeState.create(DelayedScaling(), mode="forward", num_quantizers=1) + (quantizer,) = state.make_quantizers() + op.num_quantizers = lambda mode: 1 if mode == "forward" else 0 + op.get_quantizer = lambda mode, index: quantizer + + reason = op.compile_unsupported_reason() + assert reason is not None and "value-opaque" in reason diff --git a/tests/pytorch/test_ops_hop_poc.py b/tests/pytorch/test_ops_hop_poc.py new file mode 100644 index 0000000000..a917f4e32f --- /dev/null +++ b/tests/pytorch/test_ops_hop_poc.py @@ -0,0 +1,221 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Proof of concept: a pipeline-level ``autograd.Function`` traced as a HOP. + +This is the load-bearing assumption behind compiling ``ops.Sequential``, checked +end to end before the fuser is touched: + +1. Dynamo traces a ``torch.autograd.Function`` as the ``autograd_function_apply`` + higher-order op, so **both** its forward and its backward end up in the graph. +2. That lets the forward and the backward walk **different op groupings**, which + is what ``OperationFuser`` does -- forward fuses linear+bias while backward + fuses bias+activation, and the two partitions overlap without nesting. An op + whose autograd is registered per op could not do this: autograd would compose + the backward out of the forward's nodes. +3. ``OperationContext`` objects are created inside the forward and read in the + backward. Dynamo permits mutating objects created *within* the HOP scope, so + they never have to cross an op schema. +4. A quantized (FP8) tensor is produced by one op and consumed by the next + *inside* the traced region. + +The ops themselves are the real registered custom ops; only the pipeline driver +is written out longhand here, standing in for the fuser. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch.constants import DType +from transformer_engine.pytorch.ops.basic.activation import ActivationBwdArgs +from transformer_engine.pytorch.ops.basic.bias import BiasBwdArgs +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage +from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer + +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +_device = "cuda" + + +class _OpCtx: + """Stand-in for ``OperationContext``: created in forward, read in backward.""" + + def __init__(self) -> None: + self.saved: tuple = () + self.attrs: dict = {} + + +def _fp8_quantizer() -> Any: + quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, torch.device(_device)) + quantizer.set_usage(rowwise=True, columnwise=False) + return quantizer + + +class _Pipeline(torch.autograd.Function): + """Runs a fixed 3-op pipeline: Bias -> GELU -> Bias. + + The forward walks the ops one by one; the backward walks a *coarser* + grouping, handling (bias1, gelu) as a single step. That asymmetry is the + whole point -- it is only expressible because autograd is owned here, at the + pipeline level, rather than by each op. + """ + + @staticmethod + def forward(func_ctx, x, bias1, bias2, ops, quantize_middle): + bias_fwd, _ = ops["bias"] + act_fwd, _ = ops["act"] + + ctxs = [_OpCtx() for _ in range(3)] + + # op 0: bias + args0 = ops["bias_op1"].resolve_fwd_args( + x, requires_grad=True, prev_op_grad_output_quantizer=None + ) + y0, saved0, attrs0 = bias_fwd(args0) + ctxs[0].saved, ctxs[0].attrs = saved0, attrs0 + + # op 1: activation, quantizing its output with the next op's quantizer + args1 = ops["act_op"].resolve_fwd_args( + y0, + requires_grad=True, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=_fp8_quantizer() if quantize_middle else None, + ) + y1, saved1, attrs1 = act_fwd(args1) + ctxs[1].saved = ops["act_op"].saved_for_backward(saved1, y0) + ctxs[1].attrs = attrs1 + + # op 2: bias, consuming what op 1 produced (FP8 when quantize_middle) + args2 = ops["bias_op2"].resolve_fwd_args( + y1, requires_grad=True, prev_op_grad_output_quantizer=None + ) + y2, saved2, attrs2 = bias_fwd(args2) + ctxs[2].saved, ctxs[2].attrs = saved2, attrs2 + + func_ctx.ctxs = ctxs + func_ctx.ops = ops + func_ctx.save_for_backward(x, bias1, bias2) + return y2 + + @staticmethod + def backward(func_ctx, grad_output): + _, bias_bwd = func_ctx.ops["bias"] + _, act_bwd = func_ctx.ops["act"] + ctxs = func_ctx.ctxs + + # Backward group A: op 2 alone. + dy2, db2 = bias_bwd( + BiasBwdArgs( + grad_output=grad_output, + grad_input_quantizer=ctxs[2].attrs["grad_input_quantizer"], + ) + ) + if dy2 is None: + dy2 = grad_output + + # Backward group B: ops 1 and 0 handled together -- a coarser grouping + # than the forward used. A real fused backward op would replace these + # two calls; what matters here is that the grouping may differ at all. + (dy1,) = act_bwd( + ActivationBwdArgs( + grad_output=dy2, + saved_input=ctxs[1].saved[0], + dtype=ctxs[1].attrs["dtype"], + grad_input_quantizer=ctxs[1].attrs["prev_op_grad_output_quantizer"], + ) + ) + dy0, db1 = bias_bwd( + BiasBwdArgs( + grad_output=dy1, + grad_input_quantizer=ctxs[0].attrs["grad_input_quantizer"], + ) + ) + if dy0 is None: + dy0 = dy1 + + return dy0, db1, db2, None, None + + +def _build(dtype: torch.dtype): + bias_op1 = te.ops.Bias(64, device=_device, dtype=dtype) + bias_op2 = te.ops.Bias(64, device=_device, dtype=dtype) + act_op = te.ops.GELU() + with torch.no_grad(): + bias_op1.bias.copy_(torch.randn_like(bias_op1.bias)) + bias_op2.bias.copy_(torch.randn_like(bias_op2.bias)) + ops = { + "bias": bias_op1.compile_ops, + "act": act_op.compile_ops, + "bias_op1": bias_op1, + "bias_op2": bias_op2, + "act_op": act_op, + } + return ops, bias_op1, bias_op2 + + +def _run(x, ops, bias_op1, bias_op2, quantize_middle): + return _Pipeline.apply(x, bias_op1.bias, bias_op2.bias, ops, quantize_middle) + + +@_cuda +@pytest.mark.parametrize("quantize_middle", [False, True]) +def test_hop_pipeline_compiles_fullgraph(quantize_middle) -> None: + """The pipeline must trace as a HOP and match eager, forward and backward.""" + dtype = torch.bfloat16 + ops, bias_op1, bias_op2 = _build(dtype) + x = torch.randn(16, 64, device=_device, dtype=dtype, requires_grad=True) + dy = torch.randn(16, 64, device=_device, dtype=dtype) + + def step(x_): + return _run(x_, ops, bias_op1, bias_op2, quantize_middle) + + # Eager reference. + out_ref = step(x) + out_ref.backward(dy) + grads_ref = [x.grad.clone(), bias_op1.bias.grad.clone(), bias_op2.bias.grad.clone()] + + for t in (x, bias_op1.bias, bias_op2.bias): + t.grad = None + + compiled = torch.compile(step, fullgraph=True) + out = compiled(x) + out.backward(dy) + grads = [x.grad, bias_op1.bias.grad, bias_op2.bias.grad] + + torch.testing.assert_close(out, out_ref) + for got, expected in zip(grads, grads_ref): + torch.testing.assert_close(got, expected) + + +@_cuda +def test_hop_pipeline_carries_fp8_between_ops() -> None: + """The middle tensor really is FP8, and it is consumed by the next op.""" + dtype = torch.bfloat16 + ops, bias_op1, bias_op2 = _build(dtype) + x = torch.randn(16, 64, device=_device, dtype=dtype, requires_grad=True) + + bias_fwd, _ = ops["bias"] + act_fwd, _ = ops["act"] + args0 = ops["bias_op1"].resolve_fwd_args( + x.detach(), requires_grad=False, prev_op_grad_output_quantizer=None + ) + y0, _, _ = bias_fwd(args0) + args1 = ops["act_op"].resolve_fwd_args( + y0, + requires_grad=False, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=_fp8_quantizer(), + ) + y1, _, _ = act_fwd(args1) + assert isinstance(y1, QuantizedTensorStorage), f"expected FP8, got {type(y1)}" + + args2 = ops["bias_op2"].resolve_fwd_args( + y1, requires_grad=False, prev_op_grad_output_quantizer=None + ) + y2, _, _ = bias_fwd(args2) + assert y2 is not None diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index f4beffe90c..d4baa48d93 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -7,13 +7,16 @@ from __future__ import annotations import abc from collections.abc import Iterable -from typing import Any, Optional +from dataclasses import dataclass +from typing import Any, Callable, Dict, Optional, Tuple, Union import torch import transformer_engine_torch as tex from ...constants import DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload +from ...dynamo import TensorSpec +from ...quantized_tensor import QuantizedTensorStorage from ...tensor.float8_tensor import Float8CurrentScalingQuantizer, Quantizer from ...utils import clear_tensor_data from ..op import BasicOperation, OperationContext @@ -33,6 +36,39 @@ "SiLU", ] +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class ActivationFwdArgs: + """Flat, ``self``-free inputs to an activation forward.""" + + input_: TensorOrQuantized + dtype: torch.dtype + output_quantizer: Optional[Quantizer] + cache_quantized_input: bool + requires_grad: bool + prev_op_grad_output_quantizer: Optional[Quantizer] + + +@dataclass(slots=True) +class ActivationBwdArgs: + """Flat inputs to an activation backward.""" + + grad_output: Optional[torch.Tensor] = None + saved_input: Optional[TensorOrQuantized] = None + dtype: Optional[torch.dtype] = None + grad_input_quantizer: Optional[Quantizer] = None + + +def _activation_output_shape( + input_shape: Tuple[int, ...], halves_last_dim: bool +) -> Tuple[int, ...]: + """Output shape of an activation: GLU variants consume pairs along the inner dim.""" + if not halves_last_dim: + return input_shape + return (*input_shape[:-1], input_shape[-1] // 2) + class _ActivationOperation(BasicOperation, metaclass=abc.ABCMeta): r"""Apply activation function @@ -67,83 +103,137 @@ def __init__(self, *, cache_quantized_input: bool = False): super().__init__() self.cache_quantized_input: bool = cache_quantized_input + @staticmethod @abc.abstractmethod - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: """Forward implementation Implementation from transformer_engine.pytorch.cpp_extensions. """ + @staticmethod @abc.abstractmethod - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: """Backward implementation Implementation from transformer_engine_torch. """ - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - - # Compute dtype - dtype: torch.dtype - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") - else: - dtype = input_.dtype - if dtype not in (torch.float32, torch.float16, torch.bfloat16): - raise RuntimeError(f"Unsupported dtype ({dtype})") - - # Check input tensor - x = maybe_dequantize(input_.contiguous(), dtype) + # GLU variants consume pairs along the inner dimension; set per subclass. + _output_halves_last_dim: bool = False - # Launch kernel - y = self._activation_forward_impl(x, next_op_input_quantizer) + fwd_args_type = ActivationFwdArgs + bwd_args_type = ActivationBwdArgs + num_grad_inputs = 1 - # Quantize input to FP8 before caching if needed - if self.cache_quantized_input: + @classmethod + def forward_compute(cls, args: ActivationFwdArgs): + x = maybe_dequantize(args.input_.contiguous(), args.dtype) + y = cls._activation_forward_impl(x, args.output_quantizer) + if args.cache_quantized_input: input_quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, x.device) input_quantizer.set_usage(rowwise=True, columnwise=False) x = input_quantizer(x) + # Only the re-quantized input is handed back. Otherwise x is derived from + # the input by dequantize + contiguous, both no-ops for an already-plain + # contiguous tensor, so x would *be* the input -- which a custom op may + # not return. Whether those calls are no-ops depends on strides, which + # the fake cannot see, so the rule has to be static: the backward + # rebuilds its input from the operation's input instead. + saved = (x,) if (args.requires_grad and args.cache_quantized_input) else () + return y, saved, cls._ctx_attrs(args) + + @classmethod + def forward_fake(cls, args: ActivationFwdArgs): + x = args.input_ + shape = tuple(x.shape) + y = TensorSpec( + shape=_activation_output_shape(shape, cls._output_halves_last_dim), + dtype=args.dtype, + quantizer=args.output_quantizer, + device=x.device, + ) + saved = () + if args.requires_grad and args.cache_quantized_input: + saved_quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, x.device) + saved_quantizer.set_usage(rowwise=True, columnwise=False) + saved = ( + TensorSpec( + shape=shape, dtype=args.dtype, quantizer=saved_quantizer, device=x.device + ), + ) + return y, saved, cls._ctx_attrs(args) + + @staticmethod + def _ctx_attrs(args: ActivationFwdArgs) -> Dict[str, Any]: + return { + "dtype": args.dtype, + "prev_op_grad_output_quantizer": args.prev_op_grad_output_quantizer, + } + + @classmethod + def backward_compute(cls, args: ActivationBwdArgs): + x = maybe_dequantize(args.saved_input.contiguous(), args.dtype) + dy = maybe_dequantize(args.grad_output.contiguous(), x.dtype) + return (cls._activation_backward_impl(dy, x, args.grad_input_quantizer),) + + @classmethod + def backward_fake(cls, args: ActivationBwdArgs): + return ( + TensorSpec( + shape=tuple(args.saved_input.shape), + dtype=args.dtype, + quantizer=args.grad_input_quantizer, + device=args.saved_input.device, + ), + ) - # Save state for backward pass - if ctx.requires_grad: - if is_cpu_offload_enabled(): - mark_activation_offload(x) - ctx.save_for_backward(x) - ctx.dtype = dtype - ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer - - return y - - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: + def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple: + # Without cache_quantized_input the forward keeps nothing, so the + # backward rebuilds its input from the operation's input. + return saved or (input_,) - # Saved tensors from forward pass + def resolve_bwd_args( + self, ctx: OperationContext, grad_output: torch.Tensor + ) -> ActivationBwdArgs: (x,) = ctx.saved_tensors + return ActivationBwdArgs( + grad_output=grad_output, + saved_input=x, + dtype=ctx.dtype, + grad_input_quantizer=ctx.prev_op_grad_output_quantizer, + ) - # Check input tensor - x = maybe_dequantize(x.contiguous(), ctx.dtype) - - # Check grad output tensor - dy = maybe_dequantize(grad_output.contiguous(), x.dtype) - - # Launch kernel - dx = self._activation_backward_impl(dy, x, ctx.prev_op_grad_output_quantizer) - - # Clear input tensor if possible - clear_tensor_data(x) - - return dx, () + def resolve_fwd_args( + self, + input_: torch.Tensor, + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> ActivationFwdArgs: + """Gather everything the forward needs into a flat, self-free container. + + Reads the autocast state, so it must run in the traced region (where + Dynamo guards that read), never inside the custom op. + """ + dtype: torch.dtype + if torch.is_autocast_enabled(): + dtype = torch.get_autocast_dtype("cuda") + else: + dtype = input_.dtype + if dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise RuntimeError(f"Unsupported dtype ({dtype})") + return ActivationFwdArgs( + input_=input_, + dtype=dtype, + output_quantizer=next_op_input_quantizer, + cache_quantized_input=self.cache_quantized_input, + requires_grad=requires_grad, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + ) class GELU(_ActivationOperation): @@ -159,10 +249,12 @@ class GELU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.gelu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dgelu(*args, **kwargs) @@ -191,10 +283,14 @@ class GLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + _output_halves_last_dim: bool = True + + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.glu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dglu(*args, **kwargs) @@ -226,10 +322,14 @@ class GEGLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + _output_halves_last_dim: bool = True + + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.geglu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dgeglu(*args, **kwargs) @@ -245,10 +345,12 @@ class QGELU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.qgelu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dqgelu(*args, **kwargs) @@ -278,10 +380,14 @@ class QGEGLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + _output_halves_last_dim: bool = True + + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.qgeglu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dqgeglu(*args, **kwargs) @@ -294,10 +400,12 @@ class ReLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.relu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.drelu(*args, **kwargs) @@ -323,10 +431,14 @@ class ReGLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + _output_halves_last_dim: bool = True + + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.reglu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dreglu(*args, **kwargs) @@ -341,10 +453,12 @@ class SReLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.srelu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dsrelu(*args, **kwargs) @@ -483,10 +597,14 @@ class SReGLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + _output_halves_last_dim: bool = True + + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.sreglu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dsreglu(*args, **kwargs) @@ -499,8 +617,10 @@ class SiLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.silu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dsilu(*args, **kwargs) diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index 88f563b2c5..9be54a9e1d 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -5,7 +5,8 @@ """Fusible operation for bias.""" from __future__ import annotations -from typing import Optional +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Union import torch @@ -14,6 +15,28 @@ from ..op import BasicOperation, OperationContext from ...utils import canonicalize_device, canonicalize_dtype from ...tensor import Quantizer +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec + +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class BiasFwdArgs: + """Flat, ``self``-free inputs to the bias forward.""" + + input_: TensorOrQuantized + bias: torch.Tensor + local_size: int + grad_input_quantizer: Optional[Quantizer] + + +@dataclass(slots=True) +class BiasBwdArgs: + """Flat inputs to the bias backward.""" + + grad_output: Optional[torch.Tensor] = None + grad_input_quantizer: Optional[Quantizer] = None class Bias(BasicOperation): @@ -37,6 +60,10 @@ class Bias(BasicOperation): """ + fwd_args_type = BiasFwdArgs + bwd_args_type = BiasBwdArgs + num_grad_inputs = 2 # grad input, grad bias + def __init__( self, size: int, @@ -113,37 +140,84 @@ def pre_first_fuser_forward(self) -> None: if self.bias.device.type == "meta": self.reset_parameters() - def op_forward( + @classmethod + def forward_compute(cls, args: BiasFwdArgs) -> Tuple[torch.Tensor, Tuple[()], Dict[str, Any]]: + """Add the bias. Saves no tensors; the backward only needs the quantizer.""" + x = args.input_ + b = args.bias.view([1] * (x.dim() - 1) + [args.local_size]) + return x + b, (), {"grad_input_quantizer": args.grad_input_quantizer} + + @classmethod + def forward_fake(cls, args: BiasFwdArgs) -> Tuple[TensorSpec, Tuple[()], Dict[str, Any]]: + x = args.input_ + out = TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device) + return out, (), {"grad_input_quantizer": args.grad_input_quantizer} + + @classmethod + def backward_compute(cls, args: BiasBwdArgs) -> Tuple[Optional[torch.Tensor], torch.Tensor]: + """Reduce the gradient over every dimension but the inner one. + + Returns ``None`` for the grad input when it is ``grad_output`` + unchanged: a custom op may not return one of its own inputs, and cloning + would cost a full-size copy on the common (unquantized) path. + """ + dy = args.grad_output + if dy.dim() > 1: + quantizer = args.grad_input_quantizer + if quantizer is None: + return None, dy.sum(tuple(range(dy.dim() - 1))) + db, dy = tex.bgrad_quantize(dy, quantizer) + return dy, db + return None, dy + + @classmethod + def backward_fake(cls, args: BiasBwdArgs) -> Tuple[Optional[TensorSpec], TensorSpec]: + dy = args.grad_output + shape = tuple(dy.shape) + if len(shape) > 1: + grad_bias = TensorSpec(shape=(shape[-1],), dtype=dy.dtype, device=dy.device) + quantizer = args.grad_input_quantizer + if quantizer is None: + return None, grad_bias + grad_input = TensorSpec( + shape=shape, dtype=dy.dtype, quantizer=quantizer, device=dy.device + ) + return grad_input, grad_bias + return None, TensorSpec(shape=shape, dtype=dy.dtype, device=dy.device) + + def resolve_fwd_args( self, - ctx: OperationContext, input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - x = input_ - b = self.bias.view([1] * (x.dim() - 1) + [self.local_size]) - - if ctx.requires_grad: - ctx.grad_input_quantizer = prev_op_grad_output_quantizer + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> BiasFwdArgs: + """Gather everything the forward needs into a flat, self-free container. + + Reads module config and global FP8 state, so it must run in the traced + region (where Dynamo guards those reads), never inside the custom op. + + Bias never quantizes its output, so ``next_op_input_quantizer`` is + accepted (every operation resolves through the same signature) and + ignored. + """ + del next_op_input_quantizer + grad_input_quantizer = None + if requires_grad: + grad_input_quantizer = prev_op_grad_output_quantizer if FP8GlobalStateManager.is_fp8_enabled(): - fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() - if fp8_recipe.backward_override is not None: - ctx.grad_input_quantizer = None - - return x + b + if FP8GlobalStateManager.get_fp8_recipe().backward_override is not None: + grad_input_quantizer = None + return BiasFwdArgs( + input_=input_, + bias=self.bias, + local_size=self.local_size, + grad_input_quantizer=grad_input_quantizer, + ) - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - dy = grad_output - if dy.dim() > 1: - quantizer = ctx.grad_input_quantizer - if quantizer is None: - db = dy.sum(tuple(range(dy.dim() - 1))) - else: - db, dy = tex.bgrad_quantize(dy, quantizer) - else: - db = dy - return dy, (db,) + def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> BiasBwdArgs: + return BiasBwdArgs( + grad_output=grad_output, + grad_input_quantizer=ctx.grad_input_quantizer, + )