From dbc941ff8a571fad156ac8817bfed6e95e227bc2 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 22:22:46 +0200 Subject: [PATCH 1/3] [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, + ) From 3f2cc63a50a681873450250303b727667a052407 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 13 Aug 2026 13:40:17 +0200 Subject: [PATCH 2/3] [PyTorch] Declare the remaining basic ops and forward linear+bias as custom ops Extends the compute-halves pattern from Bias and the activations to every remaining basic operation, plus the first fused one: - custom ops: ConstantScale, Dropout, Quantize, L2Normalization, LayerNorm, RMSNorm, SwiGLU, ClampedSwiGLU, BasicLinear, AllReduce, AllGather, ReduceScatter, and ForwardLinearBiasActivation (fused; its backward walks the basic ops' own custom ops, so forward and backward groupings differ); - compile_inline for pure-PyTorch ops (Identity, Reshape, AddExtraInput, MakeExtraOutput): Dynamo traces their eager passes directly; the in-place variants are gated out since they mutate enclosing-scope tensors; - a forward compute half may return None for "the input, unchanged", since a custom op may not return one of its own inputs (Quantize outside autocast, single-rank communication ops, evaluation-mode Dropout); - custom-op registration moves to FusibleOperation so fused ops share it, with a name-collision guard: re-registering under an existing name makes torch.library destroy the first registration, leaving the base class's compile_ops holding dangling operator handles (module.rmsnorm.RMSNorm subclasses ops.basic.RMSNorm under the same class name), which surfaced as heap-layout-dependent vector::reserve / bad_alloc errors; - BackwardAddRMSNorm rebuilds the kernel's 2D view from the raw saved input, since the converted RMSNorm forward no longer saves the dequantized view. Still gated out: accumulate_into_main_grad (mutates the param's main_grad), grouped operations, userbuffers, backward fused ops, delayed scaling. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_ops_custom_ops.py | 76 ++- tests/pytorch/test_torch_compile.py | 93 +++- .../pytorch/ops/basic/activation.py | 2 +- .../pytorch/ops/basic/add_extra_input.py | 9 + .../pytorch/ops/basic/all_gather.py | 130 ++++-- .../pytorch/ops/basic/all_reduce.py | 97 +++- .../pytorch/ops/basic/basic_linear.py | 358 +++++++++++--- .../pytorch/ops/basic/constant_scale.py | 83 +++- .../pytorch/ops/basic/dropout.py | 190 +++++--- .../pytorch/ops/basic/identity.py | 4 + .../pytorch/ops/basic/l2normalization.py | 135 ++++-- .../pytorch/ops/basic/layer_norm.py | 241 ++++++---- .../pytorch/ops/basic/make_extra_output.py | 10 + .../pytorch/ops/basic/quantize.py | 131 ++++-- .../pytorch/ops/basic/reduce_scatter.py | 126 +++-- .../pytorch/ops/basic/reshape.py | 4 + .../pytorch/ops/basic/rmsnorm.py | 226 ++++++--- .../pytorch/ops/basic/swiglu.py | 435 +++++++++++------- .../pytorch/ops/fused/backward_add_rmsnorm.py | 7 +- .../fused/forward_linear_bias_activation.py | 166 ++++++- transformer_engine/pytorch/ops/fuser.py | 44 +- transformer_engine/pytorch/ops/op.py | 225 ++++++--- 22 files changed, 2048 insertions(+), 744 deletions(-) diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py index 2ad8464c71..7f4cfcac75 100644 --- a/tests/pytorch/test_ops_custom_ops.py +++ b/tests/pytorch/test_ops_custom_ops.py @@ -74,6 +74,35 @@ def _build_bias(): return op +def _build_layer_norm(): + return te.ops.LayerNorm(_HIDDEN, device=_device, dtype=_dtype) + + +def _build_rmsnorm(): + return te.ops.RMSNorm(_HIDDEN, device=_device, dtype=_dtype) + + +def _build_basic_linear(): + torch.manual_seed(0) + return te.ops.BasicLinear(_HIDDEN, _HIDDEN, device=_device, dtype=_dtype) + + +def _ensure_process_group() -> None: + """Single-rank default process group, for the communication operations.""" + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group( + "gloo", init_method="tcp://localhost:29765", rank=0, world_size=1 + ) + + +def _build_comm(cls): + def build(): + _ensure_process_group() + return cls() + + return build + + @dataclass class OpCase: """One operation and how to build it.""" @@ -91,6 +120,23 @@ class OpCase: OpCase(name=cls.__name__, build=cls) for cls in (te.ops.GELU, te.ops.ReLU, te.ops.SiLU, te.ops.GEGLU, te.ops.ReGLU) ), + OpCase(name="ConstantScale", build=lambda: te.ops.ConstantScale(0.5), quantizes_output=False), + # p=0.0 keeps every element, so the fused RNG kernel is deterministic and + # the eager / compiled outputs are comparable. + OpCase(name="Dropout", build=lambda: te.ops.Dropout(0.0), quantizes_output=False), + OpCase(name="L2Normalization", build=te.ops.L2Normalization, quantizes_output=False), + OpCase(name="LayerNorm", build=_build_layer_norm, num_grads=3), + OpCase(name="RMSNorm", build=_build_rmsnorm, num_grads=2), + OpCase(name="SwiGLU", build=te.ops.SwiGLU), + OpCase(name="ClampedSwiGLU", build=te.ops.ClampedSwiGLU), + # Outside autocast Quantize is an identity, so this exercises the + # None-output ("the input, unchanged") path across the op boundary. + OpCase(name="Quantize", build=te.ops.Quantize, quantizes_output=False), + OpCase(name="BasicLinear", build=_build_basic_linear, quantizes_output=False, num_grads=2), + # Single rank, so the communication operations take their trivial paths. + OpCase(name="AllReduce", build=_build_comm(te.ops.AllReduce), quantizes_output=False), + OpCase(name="AllGather", build=_build_comm(te.ops.AllGather), quantizes_output=False), + OpCase(name="ReduceScatter", build=_build_comm(te.ops.ReduceScatter), quantizes_output=False), ] _CASE_IDS = [case.name for case in _OP_CASES] _QUANTIZING = [case for case in _OP_CASES if case.quantizes_output] @@ -190,7 +236,9 @@ def test_op_fake_matches_real(case: OpCase, fp8_output: bool) -> None: 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) + # A None output means "the input, unchanged". + out_shape = real_out.shape if real_out is not None else x.shape + dy = torch.randn(*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)) @@ -215,6 +263,9 @@ def test_op_matches_eager(case: OpCase) -> None: dx_ref = x.grad.clone() y, _saved, _attrs = forward_fn(_resolve(op, x.detach(), fp8_output=False)) + if y is None: + # "The input, unchanged": the caller substitutes, as compiled_op_forward does. + y = x.detach() _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 @@ -242,6 +293,10 @@ def test_op_compiles_fullgraph(case: OpCase, fp8_output: bool) -> None: def fwd(x_): out, saved, _attrs = forward_fn(_resolve(op, x_, fp8_output=fp8_output)) + if out is None: + # "The input, unchanged": the caller substitutes. The branch is + # decided by the fake at trace time, so it is Dynamo-safe. + out = x_ # 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): @@ -299,12 +354,29 @@ def test_op_is_compile_supported(case: OpCase) -> 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() + op = te.ops.ScaledSReLU() assert op.compile_ops is None reason = op.compile_unsupported_reason() assert reason is not None and "compute halves" in reason +@_cuda +def test_inline_ops_are_compile_supported() -> None: + """Pure-PyTorch operations opt in via compile_inline instead of custom ops.""" + for op in ( + te.ops.Identity(), + te.ops.Reshape((-1,)), + te.ops.AddExtraInput(), + te.ops.MakeExtraOutput(), + ): + assert op.compile_ops is None, type(op).__name__ + assert op.compile_unsupported_reason() is None, type(op).__name__ + + # The in-place variants mutate tensors from an enclosing scope. + assert te.ops.AddExtraInput(in_place=True).compile_unsupported_reason() is not None + assert te.ops.MakeExtraOutput(in_place=True).compile_unsupported_reason() is not None + + @_cuda def test_non_value_opaque_quantizer_is_gated() -> None: """A quantizer torch.compile cannot specialize on must be refused. diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 8c6ca357bb..14c4b63317 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1766,6 +1766,24 @@ def test_te_ops_single_op_group_compiles(): _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp()), base) +class _EagerOnlyScaleOp(BasicOperation): + """Test-only operation that never declared the compute halves.""" + + def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> None: + super().__init__() + self.scale = torch.nn.Parameter(torch.full((), 2.0, device=device, dtype=dtype)) + + def op_forward(self, ctx, input_, *, prev_op_grad_output_quantizer, next_op_input_quantizer): + del prev_op_grad_output_quantizer, next_op_input_quantizer + if ctx.requires_grad: + ctx.save_for_backward(input_) + return input_ * self.scale + + def op_backward(self, ctx, grad_output): + (x,) = ctx.saved_tensors + return grad_output * self.scale, ((grad_output * x).sum(),) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_te_ops_unsupported_group_still_compiles_eagerly(): """An operation without the compute halves runs its eager implementation. @@ -1776,10 +1794,81 @@ def test_te_ops_unsupported_group_still_compiles_eagerly(): mutation of anything from an enclosing scope -- have to hold on both paths. """ torch._dynamo.reset() - assert te.ops.Identity().compile_unsupported_reason() is not None + assert _EagerOnlyScaleOp().compile_unsupported_reason() is not None + + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + _assert_sequential_matches_eager(lambda: te.ops.Sequential(_EagerOnlyScaleOp()), base) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_te_ops_inline_op_group_compiles(): + """A ``compile_inline`` operation's eager pass is traced directly by Dynamo.""" + torch._dynamo.reset() + assert te.ops.Identity().compile_unsupported_reason() is None + assert te.ops.Reshape((-1,)).compile_unsupported_reason() is None base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") - _assert_sequential_matches_eager(lambda: te.ops.Sequential(te.ops.Identity()), base) + _assert_sequential_matches_eager( + lambda: te.ops.Sequential(te.ops.Reshape((64, 32)), _ScaleOp(), te.ops.Reshape((32, 64))), + base, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_te_ops_converted_pipeline_compiles(): + """A pipeline of real converted operations runs compiled, end to end. + + The operations are chosen so no fusion pattern matches: each stays its own + single-op group, which is what the compiled path supports so far. + """ + torch._dynamo.reset() + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + + def make_model(): + torch.manual_seed(0) + return te.ops.Sequential( + te.ops.LayerNorm(64, device="cuda", dtype=torch.bfloat16), + te.ops.GELU(), + te.ops.Dropout(0.0), # p=0.0 keeps the fused RNG kernel deterministic + te.ops.RMSNorm(64, device="cuda", dtype=torch.bfloat16), + te.ops.ConstantScale(0.5), + te.ops.SwiGLU(), + ) + + _assert_sequential_matches_eager(make_model, base) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_te_ops_basic_linear_compiles(): + """``BasicLinear`` runs through its custom op under ``fullgraph=True``.""" + torch._dynamo.reset() + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + + def make_model(): + torch.manual_seed(0) + return te.ops.Sequential(te.ops.BasicLinear(64, 64, device="cuda", dtype=torch.bfloat16)) + + _assert_sequential_matches_eager(make_model, base) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_te_ops_fused_linear_bias_compiles(): + """Linear+bias fuses into ``ForwardLinearBiasActivation`` in the forward + grouping; the fused op runs through its own custom op while the backward + walks the basic operations' custom ops -- the two groupings differ, which is + exactly what pipeline-level autograd permits.""" + torch._dynamo.reset() + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + + def make_model(): + torch.manual_seed(0) + return te.ops.Sequential( + te.ops.BasicLinear(64, 64, device="cuda", dtype=torch.bfloat16), + te.ops.Bias(64, device="cuda", dtype=torch.bfloat16), + te.ops.GELU(), + ) + + _assert_sequential_matches_eager(make_model, base) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index d4baa48d93..4f2707364c 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -8,7 +8,7 @@ import abc from collections.abc import Iterable from dataclasses import dataclass -from typing import Any, Callable, Dict, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple, Union import torch diff --git a/transformer_engine/pytorch/ops/basic/add_extra_input.py b/transformer_engine/pytorch/ops/basic/add_extra_input.py index fc3ca9cade..f863e86c14 100644 --- a/transformer_engine/pytorch/ops/basic/add_extra_input.py +++ b/transformer_engine/pytorch/ops/basic/add_extra_input.py @@ -38,10 +38,19 @@ class AddExtraInput(BasicOperation): # Operation expects buffer for output tensor num_extra_inputs: int = 1 + # A plain add that Dynamo traces directly. The in-place variant mutates the + # extra input -- a tensor from the enclosing scope -- so it is gated out. + compile_inline = True + def __init__(self, *, in_place: bool = False): super().__init__() self._in_place = in_place + def compile_unsupported_reason(self) -> Optional[str]: + if self._in_place: + return f"{self.__class__.__name__}(in_place=True) mutates its extra input" + return super().compile_unsupported_reason() + def op_forward(self, *args, **kwargs) -> None: raise RuntimeError( "{self.__class__.__name__} operation has " diff --git a/transformer_engine/pytorch/ops/basic/all_gather.py b/transformer_engine/pytorch/ops/basic/all_gather.py index 4e5c192876..219de287ba 100644 --- a/transformer_engine/pytorch/ops/basic/all_gather.py +++ b/transformer_engine/pytorch/ops/basic/all_gather.py @@ -5,7 +5,8 @@ """Fusible operation for all-gather.""" from __future__ import annotations -from typing import Optional +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Union import torch @@ -13,6 +14,28 @@ from .._common import maybe_dequantize from ..op import BasicOperation, OperationContext from ...tensor import Quantizer +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec + +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class AllGatherFwdArgs: + """Flat, ``self``-free inputs to the all-gather forward.""" + + input_: TensorOrQuantized + process_group: Optional[torch.distributed.ProcessGroup] + process_group_size: int + + +@dataclass(slots=True) +class AllGatherBwdArgs: + """Flat inputs to the all-gather backward.""" + + grad_output: Optional[torch.Tensor] = None + process_group: Optional[torch.distributed.ProcessGroup] = None + process_group_size: Optional[int] = None class AllGather(BasicOperation): @@ -28,6 +51,10 @@ class AllGather(BasicOperation): """ + fwd_args_type = AllGatherFwdArgs + bwd_args_type = AllGatherBwdArgs + num_grad_inputs = 1 + def __init__( self, process_group: Optional[torch.distributed.ProcessGroup] = None, @@ -36,49 +63,80 @@ def __init__( self.process_group: Optional[torch.distributed.ProcessGroup] = process_group self.process_group_size: int = torch.distributed.get_world_size(process_group) - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - out: torch.Tensor - if self.process_group_size == 1: - out = input_.detach() - else: - out, _ = gather_along_first_dim(input_, self.process_group) - return out - - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - - # Trivial case - if self.process_group_size == 1: - return grad_output, () + @classmethod + def forward_compute( + cls, args: AllGatherFwdArgs + ) -> Tuple[Optional[torch.Tensor], Tuple[()], Dict[str, Any]]: + # Trivial case: the input, unchanged (a custom op may not return it). + if args.process_group_size == 1: + return None, (), {} + out, _ = gather_along_first_dim(args.input_, args.process_group) + return out, (), {} + + @classmethod + def forward_fake( + cls, args: AllGatherFwdArgs + ) -> Tuple[Optional[TensorSpec], Tuple[()], Dict[str, Any]]: + if args.process_group_size == 1: + return None, (), {} + x = args.input_ + shape = tuple(x.shape) + out_shape = (shape[0] * args.process_group_size, *shape[1:]) + return TensorSpec(shape=out_shape, dtype=x.dtype, device=x.device), (), {} + + @classmethod + def backward_compute(cls, args: AllGatherBwdArgs) -> Tuple[Optional[torch.Tensor]]: + # Trivial case: the incoming gradient, unchanged. + if args.process_group_size == 1: + return (None,) # Tensor dimensions - output_dims = grad_output.size() - if not output_dims or output_dims[0] % self.process_group_size != 0: + dy = args.grad_output + output_dims = dy.size() + if not output_dims or output_dims[0] % args.process_group_size != 0: raise RuntimeError( "Attempted to reduce-scatter a tensor " f"with shape={list(output_dims)} " - f"over {self.process_group_size} processes" + f"over {args.process_group_size} processes" ) input_dims = list(output_dims) - input_dims[0] //= self.process_group_size - - # Check output gradient tensor - dy = maybe_dequantize(grad_output.contiguous()) + input_dims[0] //= args.process_group_size # Perform reduce-scatter + dy = maybe_dequantize(dy.contiguous()) dx = torch.empty(input_dims, dtype=dy.dtype, device=dy.device) - torch.distributed.reduce_scatter_tensor( - dx, - dy, - group=self.process_group, + torch.distributed.reduce_scatter_tensor(dx, dy, group=args.process_group) + return (dx,) + + @classmethod + def backward_fake(cls, args: AllGatherBwdArgs) -> Tuple[Optional[TensorSpec]]: + if args.process_group_size == 1: + return (None,) + dy = args.grad_output + shape = tuple(dy.shape) + dx_shape = (shape[0] // args.process_group_size, *shape[1:]) + return (TensorSpec(shape=dx_shape, dtype=dy.dtype, device=dy.device),) + + 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, + ) -> AllGatherFwdArgs: + del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer + return AllGatherFwdArgs( + input_=input_, + process_group=self.process_group, + process_group_size=self.process_group_size, + ) + + def resolve_bwd_args( + self, ctx: OperationContext, grad_output: torch.Tensor + ) -> AllGatherBwdArgs: + return AllGatherBwdArgs( + grad_output=grad_output, + process_group=self.process_group, + process_group_size=self.process_group_size, ) - return dx, () diff --git a/transformer_engine/pytorch/ops/basic/all_reduce.py b/transformer_engine/pytorch/ops/basic/all_reduce.py index f2e4b2481d..6e4d962c7b 100644 --- a/transformer_engine/pytorch/ops/basic/all_reduce.py +++ b/transformer_engine/pytorch/ops/basic/all_reduce.py @@ -5,13 +5,34 @@ """Fusible operation for all-reduce.""" from __future__ import annotations -from typing import Optional +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Union import torch from .._common import maybe_dequantize from ..op import BasicOperation, OperationContext from ...tensor import Quantizer +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec + +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class AllReduceFwdArgs: + """Flat, ``self``-free inputs to the all-reduce forward.""" + + input_: TensorOrQuantized + process_group: Optional[torch.distributed.ProcessGroup] + process_group_size: int + + +@dataclass(slots=True) +class AllReduceBwdArgs: + """Flat inputs to the all-reduce backward.""" + + grad_output: Optional[torch.Tensor] = None class AllReduce(BasicOperation): @@ -29,6 +50,10 @@ class AllReduce(BasicOperation): """ + fwd_args_type = AllReduceFwdArgs + bwd_args_type = AllReduceBwdArgs + num_grad_inputs = 1 + def __init__( self, process_group: Optional[torch.distributed.ProcessGroup] = None, @@ -38,26 +63,58 @@ def __init__( self.process_group: Optional[torch.distributed.ProcessGroup] = process_group self._reduce_in_backward: bool = reduce_in_backward - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: + @classmethod + def forward_compute( + cls, args: AllReduceFwdArgs + ) -> Tuple[Optional[torch.Tensor], Tuple[()], Dict[str, Any]]: + # Trivial case: the input, unchanged (a custom op may not return it). + if args.process_group_size == 1: + return None, (), {} + + # All-reduce into a fresh buffer: the eager implementation reduced the + # dequantized tensor in place, which for a plain contiguous input is the + # input itself -- a custom op may neither mutate nor return its inputs. + x = maybe_dequantize(args.input_.contiguous()) + if x is args.input_: + x = x.clone() + torch.distributed.all_reduce(x, group=args.process_group) + return x, (), {} - # Trivial case - if torch.distributed.get_world_size(self.process_group) == 1: - return input_ + @classmethod + def forward_fake( + cls, args: AllReduceFwdArgs + ) -> Tuple[Optional[TensorSpec], Tuple[()], Dict[str, Any]]: + if args.process_group_size == 1: + return None, (), {} + x = args.input_ + return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), (), {} - # Perform all-reduce - x = maybe_dequantize(input_.contiguous()) - torch.distributed.all_reduce(x, group=self.process_group) - return x + @classmethod + def backward_compute(cls, args: AllReduceBwdArgs) -> Tuple[Optional[torch.Tensor]]: + del args + return (None,) - def op_backward( + @classmethod + def backward_fake(cls, args: AllReduceBwdArgs) -> Tuple[Optional[TensorSpec]]: + del args + return (None,) + + def resolve_fwd_args( self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - return grad_output, () + input_: torch.Tensor, + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> AllReduceFwdArgs: + del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer + return AllReduceFwdArgs( + input_=input_, + process_group=self.process_group, + process_group_size=torch.distributed.get_world_size(self.process_group), + ) + + def resolve_bwd_args( + self, ctx: OperationContext, grad_output: torch.Tensor + ) -> AllReduceBwdArgs: + return AllReduceBwdArgs(grad_output=grad_output) diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index cb429055a4..7927b308d9 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -7,13 +7,13 @@ from __future__ import annotations from collections.abc import Callable, Iterable import contextlib +from dataclasses import dataclass import math -from typing import Any, Optional +from typing import Any, Optional, Union import torch from ...cpp_extensions import general_gemm -from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...distributed import ( CudaRNGStatesTracker, gather_along_first_dim, @@ -32,7 +32,6 @@ from ...utils import ( canonicalize_device, canonicalize_dtype, - clear_tensor_data, devices_match, ) from ..op import BasicOperation, OperationContext @@ -43,6 +42,10 @@ is_quantized_tensor, maybe_dequantize, ) +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec + +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] def _wait_async(handle: Optional[Any]) -> None: @@ -51,6 +54,80 @@ def _wait_async(handle: Optional[Any]) -> None: handle.wait() +def _value_is_quantized(x: Any) -> bool: + """Quantized-ness of a value that may be a tensor or its ``TensorSpec``.""" + if isinstance(x, TensorSpec): + return x.quantizer is not None + return is_quantized_tensor(x) + + +@dataclass(slots=True) +class BasicLinearFwdArgs: + """Flat, ``self``-free inputs to the linear forward.""" + + input_: TensorOrQuantized + weight: TensorOrQuantized + dtype: torch.dtype + tensor_parallel_mode: Optional[str] + tensor_parallel_group: Optional[torch.distributed.ProcessGroup] + sequence_parallel: bool + with_quantized_compute: bool + backward_override: Optional[str] + input_quantizer: Optional[Quantizer] + weight_quantizer: Optional[Quantizer] + output_quantizer: Optional[Quantizer] + grad_output_quantizer: Optional[Quantizer] + grad_input_quantizer: Optional[Quantizer] + input_requires_grad: bool + weight_requires_grad: bool + + +@dataclass(slots=True) +class BasicLinearBwdArgs: + """Flat inputs to the linear backward.""" + + grad_output: Optional[torch.Tensor] = None + saved_input: Optional[TensorOrQuantized] = None + saved_weight: Optional[TensorOrQuantized] = None + input_requires_grad: bool = True + weight_requires_grad: bool = True + dtype: Optional[torch.dtype] = None + grad_weight: Optional[torch.Tensor] = None + accumulate_into_grad_weight: bool = False + tensor_parallel_mode: Optional[str] = None + tensor_parallel_group: Optional[torch.distributed.ProcessGroup] = None + sequence_parallel: bool = False + with_quantized_compute: bool = False + input_quantizer: Optional[Quantizer] = None + weight_quantizer: Optional[Quantizer] = None + grad_output_quantizer: Optional[Quantizer] = None + grad_input_quantizer: Optional[Quantizer] = None + + +def _saved_input_is_fresh(args: BasicLinearFwdArgs, input_is_quantized: bool) -> bool: + """Whether the forward's saved input is created inside the op. + + Mirrors ``_functional_forward``: the input crosses back out only when it was + freshly quantized inside (no gather). In every other case the saved value + would *be* the op's input -- which a custom op may not return -- so the + backward rebuilds it from the operation's input instead. Static in the args, + so the fake can apply the same rule. + """ + if not args.weight_requires_grad or args.backward_override == "high_precision": + return False + if not args.with_quantized_compute or input_is_quantized: + return False + with_x_all_gather = args.tensor_parallel_mode == "column" and args.sequence_parallel + return not with_x_all_gather + + +def _saved_weight_is_fresh(args: BasicLinearFwdArgs, weight_is_quantized: bool) -> bool: + """Whether the forward's saved weight is created inside the op.""" + if not args.input_requires_grad or args.backward_override == "high_precision": + return False + return args.with_quantized_compute and not weight_is_quantized + + class BasicLinear(BasicOperation): """Apply linear transformation: :math:`y = x A^T` @@ -990,39 +1067,196 @@ def _functional_backward( _wait_async(dx_async) return dx, dw - def op_forward( + fwd_args_type = BasicLinearFwdArgs + bwd_args_type = BasicLinearBwdArgs + num_grad_inputs = 2 # grad input, grad weight + + @classmethod + def forward_compute(cls, args: BasicLinearFwdArgs): + output, x_local, w = BasicLinear._functional_forward( + input=args.input_, + weight=args.weight, + dtype=args.dtype, + tensor_parallel_mode=args.tensor_parallel_mode, + tensor_parallel_group=args.tensor_parallel_group, + sequence_parallel=args.sequence_parallel, + with_quantized_compute=args.with_quantized_compute, + backward_override=args.backward_override, + input_quantizer=args.input_quantizer, + weight_quantizer=args.weight_quantizer, + output_quantizer=args.output_quantizer, + input_requires_grad=args.input_requires_grad, + weight_requires_grad=args.weight_requires_grad, + ) + + # Only tensors created inside the op may cross back out as saved + # tensors; in every other case the backward rebuilds them from the + # operation's input and weight (see ``_saved_input_is_fresh``). + saved_input = ( + x_local if _saved_input_is_fresh(args, is_quantized_tensor(args.input_)) else None + ) + saved_weight = w if _saved_weight_is_fresh(args, is_quantized_tensor(args.weight)) else None + ctx_attrs = { + "with_quantized_compute": ( + args.with_quantized_compute and args.backward_override is None + ), + "backward_override": args.backward_override, + "input_quantizer": args.input_quantizer, + "weight_quantizer": args.weight_quantizer, + "grad_output_quantizer": args.grad_output_quantizer, + "grad_input_quantizer": args.grad_input_quantizer, + "dtype": args.dtype, + "input_requires_grad": args.input_requires_grad, + "weight_requires_grad": args.weight_requires_grad, + } + return output, (saved_input, saved_weight), ctx_attrs + + @classmethod + def forward_fake(cls, args: BasicLinearFwdArgs): + x = args.input_ + w = args.weight + device = x.device + input_dims = tuple(x.shape) + out_features = tuple(w.shape)[0] + + # Output geometry, mirroring the gather / reduce-scatter in the real impl + out_rows = input_dims[0] if input_dims else 1 + if args.tensor_parallel_mode == "column" and args.sequence_parallel: + out_rows *= torch.distributed.get_world_size(args.tensor_parallel_group) + elif args.tensor_parallel_mode == "row" and args.sequence_parallel: + out_rows //= torch.distributed.get_world_size(args.tensor_parallel_group) + out_shape = (out_rows, *input_dims[1:-1], out_features) + + output_quantizer = args.output_quantizer + if not args.with_quantized_compute or args.tensor_parallel_mode == "row": + output_quantizer = None + out = TensorSpec( + shape=out_shape, dtype=args.dtype, quantizer=output_quantizer, device=device + ) + + saved_input = None + if _saved_input_is_fresh(args, _value_is_quantized(x)): + saved_input = TensorSpec( + shape=input_dims, dtype=args.dtype, quantizer=args.input_quantizer, device=device + ) + saved_input.update_usage(rowwise_usage=False, columnwise_usage=True) + saved_weight = None + if _saved_weight_is_fresh(args, _value_is_quantized(w)): + saved_weight = TensorSpec( + shape=tuple(w.shape), + dtype=args.dtype, + quantizer=args.weight_quantizer, + device=device, + ) + if args.input_requires_grad and args.backward_override is None: + saved_weight.update_usage(rowwise_usage=False, columnwise_usage=True) + ctx_attrs = { + "with_quantized_compute": ( + args.with_quantized_compute and args.backward_override is None + ), + "backward_override": args.backward_override, + "input_quantizer": args.input_quantizer, + "weight_quantizer": args.weight_quantizer, + "grad_output_quantizer": args.grad_output_quantizer, + "grad_input_quantizer": args.grad_input_quantizer, + "dtype": args.dtype, + "input_requires_grad": args.input_requires_grad, + "weight_requires_grad": args.weight_requires_grad, + } + return out, (saved_input, saved_weight), ctx_attrs + + @classmethod + def backward_compute(cls, args: BasicLinearBwdArgs): + grad_input, grad_weight = BasicLinear._functional_backward( + grad_output=args.grad_output, + input=args.saved_input, + weight=args.saved_weight, + input_requires_grad=args.input_requires_grad, + weight_requires_grad=args.weight_requires_grad, + dtype=args.dtype, + grad_weight=args.grad_weight, + accumulate_into_grad_weight=args.accumulate_into_grad_weight, + tensor_parallel_mode=args.tensor_parallel_mode, + tensor_parallel_group=args.tensor_parallel_group, + sequence_parallel=args.sequence_parallel, + with_quantized_compute=args.with_quantized_compute, + input_quantizer=args.input_quantizer, + weight_quantizer=args.weight_quantizer, + grad_output_quantizer=args.grad_output_quantizer, + grad_input_quantizer=args.grad_input_quantizer, + ) + return grad_input, grad_weight + + @classmethod + def backward_fake(cls, args: BasicLinearBwdArgs): + dy = args.grad_output + device = dy.device + dy_dims = tuple(dy.shape) + in_features = tuple(args.saved_weight.shape)[-1] if args.saved_weight is not None else None + + grad_input = None + if args.input_requires_grad: + dx_rows = dy_dims[0] if dy_dims else 1 + if args.tensor_parallel_mode == "row" and args.sequence_parallel: + dx_rows *= torch.distributed.get_world_size(args.tensor_parallel_group) + elif args.tensor_parallel_mode == "column" and args.sequence_parallel: + dx_rows //= torch.distributed.get_world_size(args.tensor_parallel_group) + grad_input_quantizer = args.grad_input_quantizer + if not args.with_quantized_compute or args.tensor_parallel_mode == "column": + grad_input_quantizer = None + grad_input = TensorSpec( + shape=(dx_rows, *dy_dims[1:-1], in_features), + dtype=args.dtype, + quantizer=grad_input_quantizer, + device=device, + ) + + grad_weight = None + if args.weight_requires_grad: + dw_dtype = args.dtype if args.grad_weight is None else args.grad_weight.dtype + dw_shape = ( + tuple(args.grad_weight.shape) + if args.grad_weight is not None + else (dy_dims[-1], tuple(args.saved_input.shape)[-1]) + ) + grad_weight = TensorSpec(shape=dw_shape, dtype=dw_dtype, device=device) + return grad_input, grad_weight + + def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple: + # A None slot means the tensor was not created inside the op; the + # backward rebuilds it from the operation's input / weight. When the + # corresponding grad is not required the slot simply goes unread. + saved_input, saved_weight = saved + if saved_input is None: + saved_input = input_ + if saved_weight is None: + saved_weight = self.weight + return (saved_input, saved_weight) + + 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: - - # Check which grads are required - input_requires_grad = ctx.requires_grad - weight_requires_grad = ctx.requires_grad and self.weight.requires_grad + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> BasicLinearFwdArgs: + input_requires_grad = requires_grad + weight_requires_grad = requires_grad and self.weight.requires_grad - # Quantizers - input_quantizer = self.get_quantizer("forward", 0) - weight_quantizer = self.get_quantizer("forward", 1) - output_quantizer = next_op_input_quantizer - grad_output_quantizer = self.get_quantizer("backward", 0) - grad_input_quantizer = prev_op_grad_output_quantizer with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() if with_quantized_compute: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override else: backward_override = None - # Get autocast dtype if needed if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") else: dtype = self.weight.dtype - # Linear forward - output, x_local, w = BasicLinear._functional_forward( - input=input_, + return BasicLinearFwdArgs( + input_=input_, weight=self.weight, dtype=dtype, tensor_parallel_mode=self.tensor_parallel_mode, @@ -1030,51 +1264,22 @@ def op_forward( sequence_parallel=self.sequence_parallel, with_quantized_compute=with_quantized_compute, backward_override=backward_override, - input_quantizer=input_quantizer, - weight_quantizer=weight_quantizer, - output_quantizer=output_quantizer, + input_quantizer=self.get_quantizer("forward", 0), + weight_quantizer=self.get_quantizer("forward", 1), + output_quantizer=next_op_input_quantizer, + grad_output_quantizer=self.get_quantizer("backward", 0), + grad_input_quantizer=prev_op_grad_output_quantizer, input_requires_grad=input_requires_grad, weight_requires_grad=weight_requires_grad, ) - # Save state for backward pass - if ctx.requires_grad: - if backward_override == "high_precision": - saved_input = input_ if weight_requires_grad else None - saved_weight = self.weight if input_requires_grad else None - else: - saved_input = x_local - saved_weight = w - if is_cpu_offload_enabled(): - # No special CPU offloading logic is needed for weights. saved_weight is - # either self.weight (nn.Parameter, auto-excluded from offload) or a - # workspace freshly created each forward pass. - mark_activation_offload(saved_input) - ctx.save_for_backward(saved_input, saved_weight) - ctx.with_quantized_compute = with_quantized_compute and backward_override is None - ctx.backward_override = backward_override - ctx.input_quantizer = input_quantizer - ctx.weight_quantizer = weight_quantizer - ctx.grad_output_quantizer = grad_output_quantizer - ctx.grad_input_quantizer = grad_input_quantizer - ctx.dtype = dtype - ctx.input_requires_grad = input_requires_grad - ctx.weight_requires_grad = weight_requires_grad - - return output + def resolve_bwd_args( + self, ctx: OperationContext, grad_output: torch.Tensor + ) -> BasicLinearBwdArgs: + saved_input, saved_weight = ctx.saved_tensors - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: - - # Saved tensors from forward pass - (x_local, w) = ctx.saved_tensors - - # Megatron-LM wgrad fusion - # Note: Get grad tensor from param so we can accumulate - # directly into it. + # Megatron-LM wgrad fusion: accumulate directly into the param's + # main_grad. Gated out under compile (see compile_unsupported_reason). accumulate_into_main_grad = self._accumulate_into_main_grad grad_weight = None if ctx.weight_requires_grad and accumulate_into_main_grad: @@ -1085,11 +1290,10 @@ def op_backward( else: accumulate_into_main_grad = False - # Linear backward pass - grad_input, grad_weight = BasicLinear._functional_backward( + return BasicLinearBwdArgs( grad_output=grad_output, - input=x_local, - weight=w, + saved_input=saved_input, + saved_weight=saved_weight, input_requires_grad=ctx.input_requires_grad, weight_requires_grad=ctx.weight_requires_grad, dtype=ctx.dtype, @@ -1105,12 +1309,24 @@ def op_backward( grad_input_quantizer=ctx.grad_input_quantizer, ) - # Clear input tensor if possible - clear_tensor_data(x_local) + def compile_unsupported_reason(self) -> Optional[str]: + if self._accumulate_into_main_grad: + # The wgrad GEMM writes into the param's main_grad, a tensor from an + # enclosing scope, which a custom op may not mutate. + return "accumulate_into_main_grad" + return super().compile_unsupported_reason() + + def op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: + args = self.resolve_bwd_args(ctx, grad_output) + grad_input, grad_weight = self.backward_compute(args) - # Megatron-LM wgrad fusion - # Note: Return dummy tensor for grad weight if needed. - if accumulate_into_main_grad: + # Megatron-LM wgrad fusion: the real wgrad went into main_grad, so hand + # autograd a dummy tensor. + if args.accumulate_into_grad_weight: grad_weight = get_dummy_wgrads_for_params([self.weight])[0] return grad_input, [grad_weight] diff --git a/transformer_engine/pytorch/ops/basic/constant_scale.py b/transformer_engine/pytorch/ops/basic/constant_scale.py index d4b3660acf..08af965bda 100644 --- a/transformer_engine/pytorch/ops/basic/constant_scale.py +++ b/transformer_engine/pytorch/ops/basic/constant_scale.py @@ -5,36 +5,85 @@ """Fusible operation for constant scaling.""" from __future__ import annotations -from typing import Optional +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Union import torch -from transformer_engine.pytorch.ops.op import ( - BasicOperation, - OperationContext, -) +from ..op import BasicOperation, OperationContext +from .._common import maybe_dequantize from ...tensor import Quantizer +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec + +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class ConstantScaleFwdArgs: + """Flat, ``self``-free inputs to the constant-scale forward.""" + + input_: TensorOrQuantized + scale: float + + +@dataclass(slots=True) +class ConstantScaleBwdArgs: + """Flat inputs to the constant-scale backward.""" + + grad_output: Optional[torch.Tensor] = None + scale: Optional[float] = None class ConstantScale(BasicOperation): """Multiply by a constant""" + fwd_args_type = ConstantScaleFwdArgs + bwd_args_type = ConstantScaleBwdArgs + num_grad_inputs = 1 + def __init__(self, scale: float) -> None: super().__init__() self.scale = scale - def op_forward( + @classmethod + def forward_compute( + cls, args: ConstantScaleFwdArgs + ) -> Tuple[torch.Tensor, Tuple[()], Dict[str, Any]]: + x = maybe_dequantize(args.input_) + return x * args.scale, (), {} + + @classmethod + def forward_fake( + cls, args: ConstantScaleFwdArgs + ) -> Tuple[TensorSpec, Tuple[()], Dict[str, Any]]: + x = args.input_ + out = TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device) + return out, (), {} + + @classmethod + def backward_compute(cls, args: ConstantScaleBwdArgs) -> Tuple[torch.Tensor]: + dy = maybe_dequantize(args.grad_output) + return (dy * args.scale,) + + @classmethod + def backward_fake(cls, args: ConstantScaleBwdArgs) -> Tuple[TensorSpec]: + dy = args.grad_output + return (TensorSpec(shape=tuple(dy.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: - return input_ * self.scale + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> ConstantScaleFwdArgs: + del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer + return ConstantScaleFwdArgs(input_=input_, scale=self.scale) - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - return grad_output * self.scale, () + def resolve_bwd_args( + self, ctx: OperationContext, grad_output: torch.Tensor + ) -> ConstantScaleBwdArgs: + del ctx + return ConstantScaleBwdArgs(grad_output=grad_output, scale=self.scale) diff --git a/transformer_engine/pytorch/ops/basic/dropout.py b/transformer_engine/pytorch/ops/basic/dropout.py index 8850604aad..3845ecb5af 100644 --- a/transformer_engine/pytorch/ops/basic/dropout.py +++ b/transformer_engine/pytorch/ops/basic/dropout.py @@ -5,16 +5,42 @@ """Fusible operation for dropout.""" from __future__ import annotations -from typing import Optional +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Union import torch import transformer_engine_torch as tex -from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...tensor import Quantizer from ...tensor.storage.float8_tensor_storage import Float8TensorStorage +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec from .._common import maybe_autocast_dtype, maybe_dequantize from ..op import BasicOperation, OperationContext +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class DropoutFwdArgs: + """Flat, ``self``-free inputs to the dropout forward.""" + + input_: TensorOrQuantized + dtype: torch.dtype + impl: str + dropout_probability: float + requires_grad: bool + + +@dataclass(slots=True) +class DropoutBwdArgs: + """Flat inputs to the dropout backward.""" + + grad_output: Optional[torch.Tensor] = None + mask: Optional[torch.Tensor] = None + dtype: Optional[torch.dtype] = None + impl: Optional[str] = None + dropout_probability: Optional[float] = None + class Dropout(BasicOperation): """Randomly zero out tensor entries during training @@ -25,81 +51,115 @@ class Dropout(BasicOperation): """ + fwd_args_type = DropoutFwdArgs + bwd_args_type = DropoutBwdArgs + num_grad_inputs = 1 + def __init__(self, p: float) -> None: super().__init__() self.dropout_probability: float = p - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - - # Output dtype - dtype = maybe_autocast_dtype(default_dtype=input_.dtype) - - # Choose implementation - impl = None - if not self.training: - impl = "evaluation" - elif input_.numel() % 16 == 0 and dtype in (torch.float16, torch.bfloat16): - impl = "fused" - else: - impl = "unfused" - - # Perform dropout - out: torch.Tensor - mask: Optional[torch.Tensor] = None + @classmethod + def forward_compute( + cls, args: DropoutFwdArgs + ) -> Tuple[Optional[torch.Tensor], Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + impl = args.impl + ctx_attrs = { + "impl": impl, + "dtype": args.dtype, + "dropout_probability": args.dropout_probability, + } if impl == "evaluation": - out = input_ - elif impl == "fused": - x = input_ + # The input, unchanged; a custom op may not return its own input. + return None, (None,) if args.requires_grad else (), ctx_attrs + mask: torch.Tensor + if impl == "fused": + x = args.input_ if not isinstance(x, Float8TensorStorage): - x = maybe_dequantize(x, dtype=dtype) - out, mask = tex.dropout_fwd(x, self.dropout_probability) + x = maybe_dequantize(x, dtype=args.dtype) + out, mask = tex.dropout_fwd(x, args.dropout_probability) elif impl == "unfused": - x = maybe_dequantize(input_, dtype=dtype) - keep_prob = 1 - self.dropout_probability + x = maybe_dequantize(args.input_, dtype=args.dtype) + keep_prob = 1 - args.dropout_probability mask = torch.empty_like(x) mask.bernoulli_(keep_prob) mask *= 1 / keep_prob out = x * mask else: raise ValueError(f"Unsupported forward implementation {impl}") - - # Save context for backward - if ctx.requires_grad: - if is_cpu_offload_enabled(): - mark_activation_offload(mask) - ctx.save_for_backward(mask) - ctx.impl = impl - ctx.dropout_probability = self.dropout_probability - ctx.dtype = dtype - - return out - - def op_backward( + return out, (mask,) if args.requires_grad else (), ctx_attrs + + @classmethod + def forward_fake( + cls, args: DropoutFwdArgs + ) -> Tuple[Optional[TensorSpec], Tuple[Optional[TensorSpec], ...], Dict[str, Any]]: + x = args.input_ + ctx_attrs = { + "impl": args.impl, + "dtype": args.dtype, + "dropout_probability": args.dropout_probability, + } + if args.impl == "evaluation": + return None, (None,) if args.requires_grad else (), ctx_attrs + shape = tuple(x.shape) + numel = 1 + for d in shape: + numel *= d + out = TensorSpec(shape=shape, dtype=args.dtype, device=x.device) + if args.impl == "fused": + mask = TensorSpec(shape=(numel // 8,), dtype=torch.uint8, device=x.device) + else: + mask = TensorSpec(shape=shape, dtype=args.dtype, device=x.device) + return out, (mask,) if args.requires_grad else (), ctx_attrs + + @classmethod + def backward_compute(cls, args: DropoutBwdArgs) -> Tuple[Optional[torch.Tensor]]: + if args.impl == "evaluation": + return (None,) + dy = maybe_dequantize(args.grad_output, dtype=args.dtype) + if args.impl == "fused": + return (tex.dropout_bwd(dy, args.mask, args.dropout_probability),) + if args.impl == "unfused": + return (dy * args.mask,) + raise ValueError(f"Unsupported backward implementation {args.impl}") + + @classmethod + def backward_fake(cls, args: DropoutBwdArgs) -> Tuple[Optional[TensorSpec]]: + if args.impl == "evaluation": + return (None,) + dy = args.grad_output + return (TensorSpec(shape=tuple(dy.shape), dtype=args.dtype, device=dy.device),) + + def resolve_fwd_args( self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - - # Saved tensors from forward pass - (mask,) = ctx.saved_tensors - - # Perform dropout backward pass - grad_input: torch.Tensor - if ctx.impl == "evaluation": - grad_input = grad_output - elif ctx.impl == "fused": - dy = maybe_dequantize(grad_output, dtype=ctx.dtype) - grad_input = tex.dropout_bwd(dy, mask, ctx.dropout_probability) - elif ctx.impl == "unfused": - dy = maybe_dequantize(grad_output, dtype=ctx.dtype) - grad_input = dy * mask + input_: torch.Tensor, + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> DropoutFwdArgs: + del prev_op_grad_output_quantizer, next_op_input_quantizer + dtype = maybe_autocast_dtype(default_dtype=input_.dtype) + if not self.training: + impl = "evaluation" + elif input_.numel() % 16 == 0 and dtype in (torch.float16, torch.bfloat16): + impl = "fused" else: - raise ValueError(f"Unsupported backward implementation {ctx.impl}") - - return grad_input, () + impl = "unfused" + return DropoutFwdArgs( + input_=input_, + dtype=dtype, + impl=impl, + dropout_probability=self.dropout_probability, + requires_grad=requires_grad, + ) + + def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> DropoutBwdArgs: + (mask,) = ctx.saved_tensors + return DropoutBwdArgs( + grad_output=grad_output, + mask=mask, + dtype=ctx.dtype, + impl=ctx.impl, + dropout_probability=ctx.dropout_probability, + ) diff --git a/transformer_engine/pytorch/ops/basic/identity.py b/transformer_engine/pytorch/ops/basic/identity.py index 9e90bd98c0..ee8e424780 100644 --- a/transformer_engine/pytorch/ops/basic/identity.py +++ b/transformer_engine/pytorch/ops/basic/identity.py @@ -19,6 +19,10 @@ class Identity(BasicOperation): """Return input tensor""" + # Pure passthrough: a custom op may not return its own input, but Dynamo + # can trace the eager pass directly. + compile_inline = True + def op_forward( self, ctx: OperationContext, diff --git a/transformer_engine/pytorch/ops/basic/l2normalization.py b/transformer_engine/pytorch/ops/basic/l2normalization.py index be155c9356..5898938a7a 100644 --- a/transformer_engine/pytorch/ops/basic/l2normalization.py +++ b/transformer_engine/pytorch/ops/basic/l2normalization.py @@ -5,13 +5,13 @@ """Fusable operation for L2 Normalization.""" from __future__ import annotations -from typing import Optional +from dataclasses import dataclass +from typing import Optional, Union import os import torch from ...torch_version import torch_version -from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...jit import ( l2normalization_fused, l2normalization_fwd_fused, @@ -20,10 +20,32 @@ warmup_jit_l2normalization_all_dtypes, ) from ...tensor import Quantizer -from ...utils import clear_tensor_data +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec from ..op import BasicOperation, OperationContext from .._common import maybe_dequantize +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class L2NormalizationFwdArgs: + """Flat, ``self``-free inputs to the L2-normalization forward.""" + + input_: TensorOrQuantized + eps: float + requires_grad: bool + + +@dataclass(slots=True) +class L2NormalizationBwdArgs: + """Flat inputs to the L2-normalization backward.""" + + grad_output: Optional[torch.Tensor] = None + saved_input: Optional[TensorOrQuantized] = None + rsqrt_norm: Optional[torch.Tensor] = None + eps: Optional[float] = None + class L2Normalization(BasicOperation): r"""L2 Normalization @@ -77,54 +99,67 @@ def __init__( for hidden_size in common_hidden_sizes: warmup_jit_l2normalization_all_dtypes(hidden_size, seq_length, micro_batch_size) - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - # Use input directly - torch.compile can handle multi-dimensional tensors - x = maybe_dequantize(input_) + fwd_args_type = L2NormalizationFwdArgs + bwd_args_type = L2NormalizationBwdArgs + num_grad_inputs = 1 - # Check if backward pass is needed - requires_grad = ctx.requires_grad - - # Compute L2 normalization using fused implementation + @classmethod + def forward_compute(cls, args: L2NormalizationFwdArgs): # L2 norm: x / sqrt(sum(x^2) + eps) = x * rsqrt(sum(x^2) + eps) - if requires_grad: - # Training: use version that returns output and intermediate values for backward pass - y, rsqrt_norm = l2normalization_fwd_fused(x, self.eps) - else: - # Inference: use lightweight version that only returns output - y = l2normalization_fused(x, self.eps) - rsqrt_norm = None # Not needed for inference - - # Save state for backward pass - if requires_grad: - if is_cpu_offload_enabled(): - mark_activation_offload(x, rsqrt_norm) - ctx.save_for_backward(x, rsqrt_norm) - - return y - - def op_backward( + x = maybe_dequantize(args.input_) + if args.requires_grad: + y, rsqrt_norm = l2normalization_fwd_fused(x, args.eps) + # x is derived from the input by a possibly-no-op dequantize, so it + # may *be* the input, which a custom op may not return; the backward + # rebuilds it from the operation's input instead. + return y, (rsqrt_norm,), {"eps": args.eps} + y = l2normalization_fused(x, args.eps) + return y, (), {"eps": args.eps} + + @classmethod + def forward_fake(cls, args: L2NormalizationFwdArgs): + x = args.input_ + shape = tuple(x.shape) + y = TensorSpec(shape=shape, dtype=x.dtype, device=x.device) + saved = () + if args.requires_grad: + rsqrt_norm = TensorSpec(shape=shape[:-1] + (1,), dtype=torch.float32, device=x.device) + saved = (rsqrt_norm,) + return y, saved, {"eps": args.eps} + + @classmethod + def backward_compute(cls, args: L2NormalizationBwdArgs): + x = maybe_dequantize(args.saved_input) + dy = maybe_dequantize(args.grad_output) + # Recalculates l2_norm_squared_eps from x + return (l2normalization_backward_fused(dy, x, args.rsqrt_norm, args.eps),) + + @classmethod + def backward_fake(cls, args: L2NormalizationBwdArgs): + x = args.saved_input + return (TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device),) + + def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple: + return (input_, *saved) + + def resolve_fwd_args( self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - - # Saved tensors from forward pass + input_: torch.Tensor, + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> L2NormalizationFwdArgs: + del prev_op_grad_output_quantizer, next_op_input_quantizer + return L2NormalizationFwdArgs(input_=input_, eps=self.eps, requires_grad=requires_grad) + + def resolve_bwd_args( + self, ctx: OperationContext, grad_output: torch.Tensor + ) -> L2NormalizationBwdArgs: x, rsqrt_norm = ctx.saved_tensors - - dy = maybe_dequantize(grad_output) - - # Compute L2 norm backward pass using fused implementation - recalculates l2_norm_squared_eps - dx = l2normalization_backward_fused(dy, x, rsqrt_norm, self.eps) - - # Clear saved tensors if possible - clear_tensor_data(x) - clear_tensor_data(rsqrt_norm) - - # No parameters, so empty tuple for param grads - return dx, () + return L2NormalizationBwdArgs( + grad_output=grad_output, + saved_input=x, + rsqrt_norm=rsqrt_norm, + eps=ctx.eps, + ) diff --git a/transformer_engine/pytorch/ops/basic/layer_norm.py b/transformer_engine/pytorch/ops/basic/layer_norm.py index 3372952add..d10fd51874 100644 --- a/transformer_engine/pytorch/ops/basic/layer_norm.py +++ b/transformer_engine/pytorch/ops/basic/layer_norm.py @@ -6,21 +6,22 @@ from __future__ import annotations from collections.abc import Iterable +from dataclasses import dataclass import math import os -from typing import Optional +from typing import Any, Optional, Union import torch from transformer_engine_torch import layernorm_bwd, layernorm_fwd from ...constants import TE_DType -from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...export import is_in_onnx_export_mode from ...tensor import Quantizer +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec from ...utils import ( canonicalize_device, canonicalize_dtype, - clear_tensor_data, devices_match, ) from ..op import BasicOperation, OperationContext @@ -30,6 +31,37 @@ maybe_dequantize, ) +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class LayerNormFwdArgs: + """Flat, ``self``-free inputs to the layer-norm forward.""" + + input_: TensorOrQuantized + weight: torch.Tensor + bias: torch.Tensor + eps: float + zero_centered_gamma: bool + sm_margin: int + dtype: torch.dtype + output_quantizer: Optional[Quantizer] + requires_grad: bool + + +@dataclass(slots=True) +class LayerNormBwdArgs: + """Flat inputs to the layer-norm backward.""" + + grad_output: Optional[torch.Tensor] = None + saved_input: Optional[TensorOrQuantized] = None + means: Optional[torch.Tensor] = None + rstdevs: Optional[torch.Tensor] = None + weight: Optional[torch.Tensor] = None + zero_centered_gamma: Optional[bool] = None + sm_margin: Optional[int] = None + dtype: Optional[torch.dtype] = None + class LayerNorm(BasicOperation): r"""Layer Normalization @@ -177,100 +209,155 @@ def pre_first_fuser_forward(self) -> None: if self.weight.device.type == "meta" or self.bias.device.type == "meta": self.reset_parameters() - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - if is_in_onnx_export_mode(): - return self.op_onnx_forward(input_) - - # Fall back to a high-precision output when fused quantization is unsupported. - next_op_input_quantizer = get_fused_normalization_quantizer(next_op_input_quantizer) + fwd_args_type = LayerNormFwdArgs + bwd_args_type = LayerNormBwdArgs + num_grad_inputs = 3 # grad input, grad weight, grad bias - # Check tensor dims - weight = self.weight - weight_dims = tuple(weight.size()) - input_dims = tuple(input_.size()) - if len(input_dims) < len(weight_dims) or input_dims[-len(weight_dims) :] != weight_dims: - raise ValueError( - f"Input tensor (shape={input_dims}) " - f"and weight tensor (shape={weight_dims}) are not compatible" - ) - - # Check input tensors + @classmethod + def forward_compute(cls, args: LayerNormFwdArgs): + weight_dims = tuple(args.weight.size()) + input_dims = tuple(args.input_.size()) inner_dim = math.prod(weight_dims) - dtype = maybe_autocast_dtype(default_dtype=weight.dtype) - x = maybe_dequantize(input_.contiguous(), dtype).view((-1, inner_dim)) - w = maybe_dequantize(self.weight, dtype).view((inner_dim,)) - b = maybe_dequantize(self.bias, dtype).view((inner_dim,)) + dtype = args.dtype + x = maybe_dequantize(args.input_.contiguous(), dtype).view((-1, inner_dim)) + w = maybe_dequantize(args.weight, dtype).view((inner_dim,)) + b = maybe_dequantize(args.bias, dtype).view((inner_dim,)) - # Compute layer norm - sm_margin = self._sm_margins["forward" if ctx.requires_grad else "inference"] y, means, rstdevs = layernorm_fwd( x, w, b, - self.eps, + args.eps, None, - next_op_input_quantizer, + args.output_quantizer, TE_DType[dtype], - sm_margin, - self.zero_centered_gamma, + args.sm_margin, + args.zero_centered_gamma, ) - - # Save state for backward pass - if ctx.requires_grad: - if is_cpu_offload_enabled(): - mark_activation_offload(x, means, rstdevs) - ctx.save_for_backward(x, means, rstdevs) - ctx.dtype = dtype - - # Reshape output tensor out = y.view(input_dims) - return out - - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - - # Saved tensors from forward pass - x, means, rstdevs = ctx.saved_tensors - - # Tensor dims - weight_dims = self.weight.size() + # x is a view of the (possibly no-op) dequantized input, which a custom + # op may not return; the backward rebuilds it from the operation's input. + saved = (means, rstdevs) if args.requires_grad else () + return out, saved, {"dtype": dtype} + + @classmethod + def forward_fake(cls, args: LayerNormFwdArgs): + input_dims = tuple(args.input_.shape) + inner_dim = math.prod(tuple(args.weight.shape)) + outer_dim = math.prod(input_dims) // inner_dim + device = args.input_.device + out = TensorSpec( + shape=input_dims, dtype=args.dtype, quantizer=args.output_quantizer, device=device + ) + saved = () + if args.requires_grad: + stats = TensorSpec(shape=(outer_dim,), dtype=torch.float32, device=device) + saved = (stats, stats) + return out, saved, {"dtype": args.dtype} + + @classmethod + def backward_compute(cls, args: LayerNormBwdArgs): + weight_dims = tuple(args.weight.size()) inner_dim = math.prod(weight_dims) + dtype = args.dtype + x = maybe_dequantize(args.saved_input.contiguous(), dtype).view((-1, inner_dim)) + dy = maybe_dequantize(args.grad_output.contiguous(), dtype).view(x.size()) + w = maybe_dequantize(args.weight, dtype).view((inner_dim,)) - # Check input tensors - dtype = ctx.dtype - dy = maybe_dequantize(grad_output.contiguous(), dtype).view(x.size()) - w = maybe_dequantize(self.weight, dtype).view((inner_dim,)) - - # Compute layer norm backward pass dx, dw, db = layernorm_bwd( dy, x, - means, - rstdevs, + args.means, + args.rstdevs, w, - self._sm_margins["backward"], - self.zero_centered_gamma, + args.sm_margin, + args.zero_centered_gamma, ) - - # Clear saved tensors if possible - clear_tensor_data(x) - clear_tensor_data(means) - clear_tensor_data(rstdevs) - - # Reshape results - grad_input = dx.view(grad_output.size()) + grad_input = dx.view(args.grad_output.size()) grad_weight = dw.view(weight_dims) grad_bias = db.view(weight_dims) - return grad_input, (grad_weight, grad_bias) + return grad_input, grad_weight, grad_bias + + @classmethod + def backward_fake(cls, args: LayerNormBwdArgs): + device = args.grad_output.device + weight_dims = tuple(args.weight.shape) + grad_input = TensorSpec( + shape=tuple(args.grad_output.shape), dtype=args.dtype, device=device + ) + grad_param = TensorSpec(shape=weight_dims, dtype=args.dtype, device=device) + return grad_input, grad_param, grad_param + + def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple: + return (input_, *saved) + + 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, + ) -> LayerNormFwdArgs: + del prev_op_grad_output_quantizer + + # Fall back to a high-precision output when fused quantization is unsupported. + output_quantizer = get_fused_normalization_quantizer(next_op_input_quantizer) + + # Check tensor dims + weight_dims = tuple(self.weight.size()) + input_dims = tuple(input_.size()) + if len(input_dims) < len(weight_dims) or input_dims[-len(weight_dims) :] != weight_dims: + raise ValueError( + f"Input tensor (shape={input_dims}) " + f"and weight tensor (shape={weight_dims}) are not compatible" + ) + + return LayerNormFwdArgs( + input_=input_, + weight=self.weight, + bias=self.bias, + eps=self.eps, + zero_centered_gamma=self.zero_centered_gamma, + sm_margin=self._sm_margins["forward" if requires_grad else "inference"], + dtype=maybe_autocast_dtype(default_dtype=self.weight.dtype), + output_quantizer=output_quantizer, + requires_grad=requires_grad, + ) + + def resolve_bwd_args( + self, ctx: OperationContext, grad_output: torch.Tensor + ) -> LayerNormBwdArgs: + saved_input, means, rstdevs = ctx.saved_tensors + return LayerNormBwdArgs( + grad_output=grad_output, + saved_input=saved_input, + means=means, + rstdevs=rstdevs, + weight=self.weight, + zero_centered_gamma=self.zero_centered_gamma, + sm_margin=self._sm_margins["backward"], + dtype=ctx.dtype, + ) + + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + **kwargs: Any, + ) -> torch.Tensor: + if is_in_onnx_export_mode(): + return self.op_onnx_forward(input_) + return super().op_forward( + ctx, + input_, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + **kwargs, + ) def op_onnx_forward( self, diff --git a/transformer_engine/pytorch/ops/basic/make_extra_output.py b/transformer_engine/pytorch/ops/basic/make_extra_output.py index 0d9c870262..13ed14b2d9 100644 --- a/transformer_engine/pytorch/ops/basic/make_extra_output.py +++ b/transformer_engine/pytorch/ops/basic/make_extra_output.py @@ -43,10 +43,20 @@ class MakeExtraOutput(BasicOperation): # Operation expects buffer for output tensor num_extra_outputs: int = 1 + # A passthrough plus an add in backward, traced directly by Dynamo. The + # in-place variant mutates the extra output's gradient, which may come from + # the enclosing scope, so it is gated out. + compile_inline = True + def __init__(self, *, in_place: bool = False): super().__init__() self._in_place: bool = in_place + def compile_unsupported_reason(self) -> Optional[str]: + if self._in_place: + return f"{self.__class__.__name__}(in_place=True) mutates its extra output's gradient" + return super().compile_unsupported_reason() + def op_forward(self, *args, **kwargs) -> None: raise RuntimeError( "{self.__class__.__name__} operation has " diff --git a/transformer_engine/pytorch/ops/basic/quantize.py b/transformer_engine/pytorch/ops/basic/quantize.py index d0c1137d91..8504d841db 100644 --- a/transformer_engine/pytorch/ops/basic/quantize.py +++ b/transformer_engine/pytorch/ops/basic/quantize.py @@ -5,7 +5,8 @@ """Fusible operation for quantization.""" from __future__ import annotations -from typing import Optional +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Union import torch @@ -13,6 +14,35 @@ from .._common import is_quantized_tensor from ..op import BasicOperation, OperationContext from ...tensor import Quantizer +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec + +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +def _is_quantized(x: Any) -> bool: + """Quantized-ness of a value that may be a tensor or its ``TensorSpec``.""" + if isinstance(x, TensorSpec): + return x.quantizer is not None + return is_quantized_tensor(x) + + +@dataclass(slots=True) +class QuantizeFwdArgs: + """Flat, ``self``-free inputs to the quantize forward.""" + + input_: TensorOrQuantized + forward_quantizer: Optional[Quantizer] + quantize_backward: bool + backward_quantizer: Optional[Quantizer] + + +@dataclass(slots=True) +class QuantizeBwdArgs: + """Flat inputs to the quantize backward.""" + + grad_output: Optional[torch.Tensor] = None + backward_quantizer: Optional[Quantizer] = None class Quantize(BasicOperation): @@ -30,6 +60,10 @@ class Quantize(BasicOperation): """ + fwd_args_type = QuantizeFwdArgs + bwd_args_type = QuantizeBwdArgs + num_grad_inputs = 1 + def __init__( self, forward: bool = True, @@ -46,39 +80,78 @@ def num_quantizers(self, mode: str) -> int: return 1 return 0 - def op_forward( + @classmethod + def forward_compute( + cls, args: QuantizeFwdArgs + ) -> Tuple[Optional[TensorOrQuantized], Tuple[()], Dict[str, Any]]: + ctx_attrs = { + "backward_quantizer": args.backward_quantizer if args.quantize_backward else None + } + quantizer = args.forward_quantizer + if quantizer is not None and not is_quantized_tensor(args.input_): + return quantizer(args.input_), (), ctx_attrs + # The input, unchanged; a custom op may not return its own input. + return None, (), ctx_attrs + + @classmethod + def forward_fake( + cls, args: QuantizeFwdArgs + ) -> Tuple[Optional[TensorSpec], Tuple[()], Dict[str, Any]]: + ctx_attrs = { + "backward_quantizer": args.backward_quantizer if args.quantize_backward else None + } + x = args.input_ + quantizer = args.forward_quantizer + if quantizer is not None and not _is_quantized(x): + out = TensorSpec( + shape=tuple(x.shape), dtype=x.dtype, quantizer=quantizer, device=x.device + ) + return out, (), ctx_attrs + return None, (), ctx_attrs + + @classmethod + def backward_compute(cls, args: QuantizeBwdArgs) -> Tuple[Optional[torch.Tensor]]: + quantizer = args.backward_quantizer + if quantizer is not None and not is_quantized_tensor(args.grad_output): + return (quantizer(args.grad_output),) + return (None,) + + @classmethod + def backward_fake(cls, args: QuantizeBwdArgs) -> Tuple[Optional[TensorSpec]]: + dy = args.grad_output + quantizer = args.backward_quantizer + if quantizer is not None and not _is_quantized(dy): + return ( + TensorSpec( + shape=tuple(dy.shape), dtype=dy.dtype, quantizer=quantizer, device=dy.device + ), + ) + return (None,) + + 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: - - # Check if FP8 is enabled + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> QuantizeFwdArgs: + del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer fp8_enabled = FP8GlobalStateManager.is_fp8_enabled() quantize_forward = fp8_enabled and self._quantize_forward quantize_backward = fp8_enabled and self._quantize_backward - - # Backward quantization is controlled by recipe backward override. if fp8_enabled: recipe = FP8GlobalStateManager.get_fp8_recipe() quantize_backward = quantize_backward and recipe.backward_override is None - - # Quantize if needed - out = input_ - if quantize_forward and not is_quantized_tensor(out): - out = self.get_quantizer("forward", 0)(out) - - if ctx.requires_grad: - ctx.quantize_backward = quantize_backward - return out - - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - grad_input = grad_output - if ctx.quantize_backward and not is_quantized_tensor(grad_input): - grad_input = self.get_quantizer("backward", 0)(grad_input) - return grad_input, () + return QuantizeFwdArgs( + input_=input_, + forward_quantizer=self.get_quantizer("forward", 0) if quantize_forward else None, + quantize_backward=quantize_backward, + backward_quantizer=self.get_quantizer("backward", 0) if quantize_backward else None, + ) + + def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> QuantizeBwdArgs: + return QuantizeBwdArgs( + grad_output=grad_output, + backward_quantizer=ctx.backward_quantizer, + ) diff --git a/transformer_engine/pytorch/ops/basic/reduce_scatter.py b/transformer_engine/pytorch/ops/basic/reduce_scatter.py index 0169da2490..ff299e116a 100644 --- a/transformer_engine/pytorch/ops/basic/reduce_scatter.py +++ b/transformer_engine/pytorch/ops/basic/reduce_scatter.py @@ -5,7 +5,8 @@ """Fusible operation for reduce-scatter.""" from __future__ import annotations -from typing import Optional +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Union import torch @@ -13,6 +14,28 @@ from .._common import maybe_dequantize from ..op import BasicOperation, OperationContext from ...tensor import Quantizer +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec + +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class ReduceScatterFwdArgs: + """Flat, ``self``-free inputs to the reduce-scatter forward.""" + + input_: TensorOrQuantized + process_group: Optional[torch.distributed.ProcessGroup] + process_group_size: int + + +@dataclass(slots=True) +class ReduceScatterBwdArgs: + """Flat inputs to the reduce-scatter backward.""" + + grad_output: Optional[torch.Tensor] = None + process_group: Optional[torch.distributed.ProcessGroup] = None + process_group_size: Optional[int] = None class ReduceScatter(BasicOperation): @@ -28,6 +51,10 @@ class ReduceScatter(BasicOperation): """ + fwd_args_type = ReduceScatterFwdArgs + bwd_args_type = ReduceScatterBwdArgs + num_grad_inputs = 1 + def __init__( self, process_group: Optional[torch.distributed.ProcessGroup] = None, @@ -36,45 +63,80 @@ def __init__( self.process_group: Optional[torch.distributed.ProcessGroup] = process_group self.process_group_size: int = torch.distributed.get_world_size(process_group) - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - - # Trivial case - if self.process_group_size == 1: - return input_.detach() + @classmethod + def forward_compute( + cls, args: ReduceScatterFwdArgs + ) -> Tuple[Optional[torch.Tensor], Tuple[()], Dict[str, Any]]: + # Trivial case: the input, unchanged (a custom op may not return it). + if args.process_group_size == 1: + return None, (), {} # Tensor dimensions - input_dims = input_.size() - if not input_dims or input_dims[0] % self.process_group_size != 0: + x = args.input_ + input_dims = x.size() + if not input_dims or input_dims[0] % args.process_group_size != 0: raise RuntimeError( "Attempted to reduce-scatter a tensor " f"with shape={list(input_dims)} " - f"over {self.process_group_size} processes" + f"over {args.process_group_size} processes" ) output_dims = list(input_dims) - output_dims[0] //= self.process_group_size - - # Check input tensor - x = maybe_dequantize(input_.contiguous()) + output_dims[0] //= args.process_group_size # Perform reduce-scatter + x = maybe_dequantize(x.contiguous()) y = torch.empty(output_dims, dtype=x.dtype, device=x.device) - torch.distributed.reduce_scatter_tensor(y, x, group=self.process_group) - return y - - def op_backward( + torch.distributed.reduce_scatter_tensor(y, x, group=args.process_group) + return y, (), {} + + @classmethod + def forward_fake( + cls, args: ReduceScatterFwdArgs + ) -> Tuple[Optional[TensorSpec], Tuple[()], Dict[str, Any]]: + if args.process_group_size == 1: + return None, (), {} + x = args.input_ + shape = tuple(x.shape) + out_shape = (shape[0] // args.process_group_size, *shape[1:]) + return TensorSpec(shape=out_shape, dtype=x.dtype, device=x.device), (), {} + + @classmethod + def backward_compute(cls, args: ReduceScatterBwdArgs) -> Tuple[Optional[torch.Tensor]]: + # Trivial case: the incoming gradient, unchanged. + if args.process_group_size == 1: + return (None,) + dx, _ = gather_along_first_dim(args.grad_output, args.process_group) + return (dx,) + + @classmethod + def backward_fake(cls, args: ReduceScatterBwdArgs) -> Tuple[Optional[TensorSpec]]: + if args.process_group_size == 1: + return (None,) + dy = args.grad_output + shape = tuple(dy.shape) + dx_shape = (shape[0] * args.process_group_size, *shape[1:]) + return (TensorSpec(shape=dx_shape, dtype=dy.dtype, device=dy.device),) + + def resolve_fwd_args( self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - grad_input: torch.Tensor - if self.process_group_size == 1: - grad_input = grad_output.detach() - else: - grad_input, _ = gather_along_first_dim(grad_output, self.process_group) - return grad_input, () + input_: torch.Tensor, + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> ReduceScatterFwdArgs: + del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer + return ReduceScatterFwdArgs( + input_=input_, + process_group=self.process_group, + process_group_size=self.process_group_size, + ) + + def resolve_bwd_args( + self, ctx: OperationContext, grad_output: torch.Tensor + ) -> ReduceScatterBwdArgs: + return ReduceScatterBwdArgs( + grad_output=grad_output, + process_group=self.process_group, + process_group_size=self.process_group_size, + ) diff --git a/transformer_engine/pytorch/ops/basic/reshape.py b/transformer_engine/pytorch/ops/basic/reshape.py index 4a171c294b..f64ab25295 100644 --- a/transformer_engine/pytorch/ops/basic/reshape.py +++ b/transformer_engine/pytorch/ops/basic/reshape.py @@ -30,6 +30,10 @@ class Reshape(BasicOperation): """ + # Both passes are a single view op: a custom op may return neither its + # input nor a view of it, but Dynamo traces reshape natively. + compile_inline = True + def __init__(self, shape: Iterable[int]) -> None: super().__init__() self._shape = tuple(shape) diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm.py b/transformer_engine/pytorch/ops/basic/rmsnorm.py index 96851c8db4..7120d30084 100644 --- a/transformer_engine/pytorch/ops/basic/rmsnorm.py +++ b/transformer_engine/pytorch/ops/basic/rmsnorm.py @@ -6,21 +6,22 @@ from __future__ import annotations from collections.abc import Iterable +from dataclasses import dataclass import math import os -from typing import Optional +from typing import Any, Optional, Union import torch from transformer_engine_torch import rmsnorm_bwd, rmsnorm_fwd from ...constants import TE_DType -from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...export import is_in_onnx_export_mode from ...tensor import Quantizer +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec from ...utils import ( canonicalize_device, canonicalize_dtype, - clear_tensor_data, devices_match, ) from ..op import BasicOperation, OperationContext @@ -30,6 +31,35 @@ maybe_dequantize, ) +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class RMSNormFwdArgs: + """Flat, ``self``-free inputs to the RMSNorm forward.""" + + input_: TensorOrQuantized + weight: torch.Tensor + eps: float + zero_centered_gamma: bool + sm_margin: int + dtype: torch.dtype + output_quantizer: Optional[Quantizer] + requires_grad: bool + + +@dataclass(slots=True) +class RMSNormBwdArgs: + """Flat inputs to the RMSNorm backward.""" + + grad_output: Optional[torch.Tensor] = None + saved_input: Optional[TensorOrQuantized] = None + rstdevs: Optional[torch.Tensor] = None + weight: Optional[torch.Tensor] = None + zero_centered_gamma: Optional[bool] = None + sm_margin: Optional[int] = None + dtype: Optional[torch.dtype] = None + class RMSNorm(BasicOperation): r"""Root Mean Square Layer Normalization @@ -160,95 +190,145 @@ def pre_first_fuser_forward(self) -> None: if self.weight.device.type == "meta": self.reset_parameters() - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - if is_in_onnx_export_mode(): - return self.op_onnx_forward(input_) - - # Fall back to a high-precision output when fused quantization is unsupported. - next_op_input_quantizer = get_fused_normalization_quantizer(next_op_input_quantizer) + fwd_args_type = RMSNormFwdArgs + bwd_args_type = RMSNormBwdArgs + num_grad_inputs = 2 # grad input, grad weight - # Check tensor dims - weight = self.weight - weight_dims = tuple(weight.size()) - input_dims = tuple(input_.size()) - if len(input_dims) < len(weight_dims) or input_dims[-len(weight_dims) :] != weight_dims: - raise ValueError( - f"Input tensor (shape={input_dims}) " - f"and weight tensor (shape={weight_dims}) are not compatible" - ) - - # Check input tensors + @classmethod + def forward_compute(cls, args: RMSNormFwdArgs): + weight_dims = tuple(args.weight.size()) + input_dims = tuple(args.input_.size()) inner_dim = math.prod(weight_dims) - dtype = maybe_autocast_dtype(default_dtype=weight.dtype) - x = maybe_dequantize(input_.contiguous(), dtype).view((-1, inner_dim)) - w = maybe_dequantize(self.weight, dtype).view((inner_dim,)) + dtype = args.dtype + x = maybe_dequantize(args.input_.contiguous(), dtype).view((-1, inner_dim)) + w = maybe_dequantize(args.weight, dtype).view((inner_dim,)) - # Compute RMSNorm - sm_margin = self._sm_margins["forward" if ctx.requires_grad else "inference"] y, _, rstdevs = rmsnorm_fwd( x, w, - self.eps, + args.eps, None, - next_op_input_quantizer, + args.output_quantizer, TE_DType[dtype], - sm_margin, - self.zero_centered_gamma, + args.sm_margin, + args.zero_centered_gamma, ) - - # Save state for backward pass - if ctx.requires_grad: - if is_cpu_offload_enabled(): - mark_activation_offload(x, rstdevs) - ctx.save_for_backward(x, rstdevs) - ctx.dtype = dtype - - # Reshape output tensor out = y.view(input_dims) - return out - - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - - # Saved tensors from forward pass - x, rstdevs = ctx.saved_tensors - - # Tensor dims - weight_dims = self.weight.size() + # x is a view of the (possibly no-op) dequantized input, which a custom + # op may not return; the backward rebuilds it from the operation's input. + saved = (rstdevs,) if args.requires_grad else () + return out, saved, {"dtype": dtype} + + @classmethod + def forward_fake(cls, args: RMSNormFwdArgs): + input_dims = tuple(args.input_.shape) + inner_dim = math.prod(tuple(args.weight.shape)) + outer_dim = math.prod(input_dims) // inner_dim + device = args.input_.device + out = TensorSpec( + shape=input_dims, dtype=args.dtype, quantizer=args.output_quantizer, device=device + ) + saved = () + if args.requires_grad: + saved = (TensorSpec(shape=(outer_dim,), dtype=torch.float32, device=device),) + return out, saved, {"dtype": args.dtype} + + @classmethod + def backward_compute(cls, args: RMSNormBwdArgs): + weight_dims = tuple(args.weight.size()) inner_dim = math.prod(weight_dims) + dtype = args.dtype + x = maybe_dequantize(args.saved_input.contiguous(), dtype).view((-1, inner_dim)) + dy = maybe_dequantize(args.grad_output.contiguous(), dtype).view(x.size()) + w = maybe_dequantize(args.weight, dtype).view((inner_dim,)) - # Check input tensors - dtype = ctx.dtype - dy = maybe_dequantize(grad_output.contiguous(), dtype).view(x.size()) - w = maybe_dequantize(self.weight, dtype).view((inner_dim,)) - - # Compute RMSNorm backward pass dx, dw = rmsnorm_bwd( dy, x, - rstdevs, + args.rstdevs, w, - self._sm_margins["backward"], - self.zero_centered_gamma, + args.sm_margin, + args.zero_centered_gamma, ) + grad_input = dx.view(args.grad_output.size()) + grad_weight = dw.view(weight_dims) + return grad_input, grad_weight - # Clear saved tensors if possible - clear_tensor_data(x) - clear_tensor_data(rstdevs) + @classmethod + def backward_fake(cls, args: RMSNormBwdArgs): + device = args.grad_output.device + grad_input = TensorSpec( + shape=tuple(args.grad_output.shape), dtype=args.dtype, device=device + ) + grad_weight = TensorSpec(shape=tuple(args.weight.shape), dtype=args.dtype, device=device) + return grad_input, grad_weight - # Reshape results - grad_input = dx.view(grad_output.size()) - grad_weight = dw.view(weight_dims) - return grad_input, (grad_weight,) + def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple: + return (input_, *saved) + + 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, + ) -> RMSNormFwdArgs: + del prev_op_grad_output_quantizer + + # Fall back to a high-precision output when fused quantization is unsupported. + output_quantizer = get_fused_normalization_quantizer(next_op_input_quantizer) + + # Check tensor dims + weight_dims = tuple(self.weight.size()) + input_dims = tuple(input_.size()) + if len(input_dims) < len(weight_dims) or input_dims[-len(weight_dims) :] != weight_dims: + raise ValueError( + f"Input tensor (shape={input_dims}) " + f"and weight tensor (shape={weight_dims}) are not compatible" + ) + + return RMSNormFwdArgs( + input_=input_, + weight=self.weight, + eps=self.eps, + zero_centered_gamma=self.zero_centered_gamma, + sm_margin=self._sm_margins["forward" if requires_grad else "inference"], + dtype=maybe_autocast_dtype(default_dtype=self.weight.dtype), + output_quantizer=output_quantizer, + requires_grad=requires_grad, + ) + + def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> RMSNormBwdArgs: + saved_input, rstdevs = ctx.saved_tensors + return RMSNormBwdArgs( + grad_output=grad_output, + saved_input=saved_input, + rstdevs=rstdevs, + weight=self.weight, + zero_centered_gamma=self.zero_centered_gamma, + sm_margin=self._sm_margins["backward"], + dtype=ctx.dtype, + ) + + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + **kwargs: Any, + ) -> torch.Tensor: + if is_in_onnx_export_mode(): + return self.op_onnx_forward(input_) + return super().op_forward( + ctx, + input_, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + **kwargs, + ) def op_onnx_forward( self, diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 02f330ede3..c77c2b20da 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -6,7 +6,8 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any, Optional +from dataclasses import dataclass +from typing import Any, Dict, Optional, Union import torch @@ -14,12 +15,74 @@ from ...constants import DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...tensor import Float8CurrentScalingQuantizer, Quantizer +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec from ...utils import clear_tensor_data from ..op import BasicOperation, OperationContext from .._common import maybe_dequantize __all__ = ["SwiGLU", "ClampedSwiGLU", "ScaledSwiGLU", "ScaledClampedQGeGLU"] +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +def _deinterleave_glu(t: torch.Tensor, interleave_size: int) -> torch.Tensor: + """Convert block-interleaved gates/units into the concatenated layout.""" + shape = t.size() + t = t.reshape(-1, shape[-1] // (2 * interleave_size), 2, interleave_size) + t = t.transpose(1, 2).contiguous() + return t.view(shape) + + +def _interleave_glu(t: torch.Tensor, interleave_size: int) -> torch.Tensor: + """Inverse of :func:`_deinterleave_glu`.""" + shape = t.size() + t = t.reshape(-1, 2, shape[-1] // (2 * interleave_size), interleave_size) + t = t.transpose(1, 2).contiguous() + return t.view(shape) + + +@dataclass(slots=True) +class SwiGLUFwdArgs: + """Flat, ``self``-free inputs to the SwiGLU forward.""" + + input_: TensorOrQuantized + dtype: torch.dtype + output_quantizer: Optional[Quantizer] + cache_quantized_input: bool + glu_interleave_size: Optional[int] + requires_grad: bool + prev_op_grad_output_quantizer: Optional[Quantizer] + + +@dataclass(slots=True) +class SwiGLUBwdArgs: + """Flat inputs to the SwiGLU backward.""" + + grad_output: Optional[torch.Tensor] = None + saved_input: Optional[TensorOrQuantized] = None + dtype: Optional[torch.dtype] = None + glu_interleave_size: Optional[int] = None + grad_input_quantizer: Optional[Quantizer] = None + + +@dataclass(slots=True) +class ClampedSwiGLUFwdArgs(SwiGLUFwdArgs): + """Flat, ``self``-free inputs to the clamped-SwiGLU forward.""" + + limit: float = 7.0 + alpha: float = 1.702 + glu_linear_offset: float = 1.0 + + +@dataclass(slots=True) +class ClampedSwiGLUBwdArgs(SwiGLUBwdArgs): + """Flat inputs to the clamped-SwiGLU backward.""" + + limit: Optional[float] = None + alpha: Optional[float] = None + glu_linear_offset: Optional[float] = None + class SwiGLU(BasicOperation): r"""Swish gated linear unit @@ -80,112 +143,124 @@ def __init__( self.cache_quantized_input: bool = cache_quantized_input self.glu_interleave_size: Optional[int] = glu_interleave_size - 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})") + fwd_args_type = SwiGLUFwdArgs + bwd_args_type = SwiGLUBwdArgs + num_grad_inputs = 1 - # Check input tensor - input_ = maybe_dequantize(input_.contiguous(), dtype) + @classmethod + def forward_compute(cls, args: SwiGLUFwdArgs): + x = maybe_dequantize(args.input_.contiguous(), args.dtype) - # Remove interleaving if needed - swiglu_in = input_ - if self.glu_interleave_size is not None: - shape = swiglu_in.size() - swiglu_in = swiglu_in.reshape( - -1, - shape[-1] // (2 * self.glu_interleave_size), - 2, - self.glu_interleave_size, - ) - swiglu_in = swiglu_in.transpose(1, 2).contiguous() - swiglu_in = swiglu_in.view(shape) + swiglu_in = x + if args.glu_interleave_size is not None: + swiglu_in = _deinterleave_glu(swiglu_in, args.glu_interleave_size) - # Launch kernel - out = tex.swiglu(swiglu_in, next_op_input_quantizer) + out = tex.swiglu(swiglu_in, args.output_quantizer) - # Quantize input to FP8 before caching if needed - if self.cache_quantized_input: - input_quantizer = Float8CurrentScalingQuantizer( - DType.kFloat8E4M3, - input_.device, - ) + if args.cache_quantized_input: + input_quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, x.device) input_quantizer.set_usage(rowwise=True, columnwise=False) - input_ = input_quantizer(input_) - - # Save state for backward pass - if ctx.requires_grad: - if is_cpu_offload_enabled(): - mark_activation_offload(input_) - ctx.save_for_backward(input_) - ctx.dtype = dtype - ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer - - return out - - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: + x = input_quantizer(x) + # Only the re-quantized input is handed back; otherwise x may *be* the + # input (dequantize + contiguous are no-ops for a plain contiguous + # tensor), which a custom op may not return. The backward rebuilds it + # from the operation's input instead. + saved = (x,) if (args.requires_grad and args.cache_quantized_input) else () + return out, saved, cls._ctx_attrs(args) + + @classmethod + def forward_fake(cls, args: SwiGLUFwdArgs): + x = args.input_ + shape = tuple(x.shape) + out = TensorSpec( + shape=(*shape[:-1], shape[-1] // 2), + 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 out, saved, cls._ctx_attrs(args) - # Saved tensors from forward pass - (input_,) = ctx.saved_tensors + @staticmethod + def _ctx_attrs(args: SwiGLUFwdArgs) -> Dict[str, Any]: + return { + "dtype": args.dtype, + "prev_op_grad_output_quantizer": args.prev_op_grad_output_quantizer, + } - # Make sure tensors have correct dtypes - x = maybe_dequantize(input_.contiguous(), ctx.dtype) - dy = maybe_dequantize(grad_output.contiguous(), ctx.dtype) + @classmethod + def backward_compute(cls, args: SwiGLUBwdArgs): + x = maybe_dequantize(args.saved_input.contiguous(), args.dtype) + dy = maybe_dequantize(args.grad_output.contiguous(), args.dtype) - # Remove interleaving if needed swiglu_in = x - if self.glu_interleave_size is not None: - shape = swiglu_in.size() - swiglu_in = swiglu_in.reshape( - -1, - shape[-1] // (2 * self.glu_interleave_size), - 2, - self.glu_interleave_size, - ) - swiglu_in = swiglu_in.transpose(1, 2).contiguous() - swiglu_in = swiglu_in.view(shape) - - # Quantizer for grad input - quantizer = ctx.prev_op_grad_output_quantizer - if self.glu_interleave_size is not None: + quantizer = args.grad_input_quantizer + if args.glu_interleave_size is not None: + swiglu_in = _deinterleave_glu(swiglu_in, args.glu_interleave_size) quantizer = None - # Launch kernel - grad_swiglu_in = tex.dswiglu(dy, swiglu_in, quantizer) + dx = tex.dswiglu(dy, swiglu_in, quantizer) + if args.glu_interleave_size is not None: + dx = _interleave_glu(dx, args.glu_interleave_size) + return (dx,) - # Apply interleaving if needed - dx = grad_swiglu_in - if self.glu_interleave_size is not None: - shape = dx.size() - dx = dx.reshape( - -1, - 2, - shape[-1] // (2 * self.glu_interleave_size), - self.glu_interleave_size, - ) - dx = dx.transpose(1, 2).contiguous() - dx = dx.view(shape) + @classmethod + def backward_fake(cls, args: SwiGLUBwdArgs): + x = args.saved_input + quantizer = args.grad_input_quantizer + if args.glu_interleave_size is not None: + quantizer = None + return ( + TensorSpec( + shape=tuple(x.shape), dtype=args.dtype, quantizer=quantizer, device=x.device + ), + ) - # Clear input tensor if possible - clear_tensor_data(input_) + def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple: + return saved or (input_,) - 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, + ) -> SwiGLUFwdArgs: + 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 SwiGLUFwdArgs( + input_=input_, + dtype=dtype, + output_quantizer=next_op_input_quantizer, + cache_quantized_input=self.cache_quantized_input, + glu_interleave_size=self.glu_interleave_size, + requires_grad=requires_grad, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + ) + + def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> SwiGLUBwdArgs: + (saved_input,) = ctx.saved_tensors + return SwiGLUBwdArgs( + grad_output=grad_output, + saved_input=saved_input, + dtype=ctx.dtype, + glu_interleave_size=self.glu_interleave_size, + grad_input_quantizer=ctx.prev_op_grad_output_quantizer, + ) class ClampedSwiGLU(BasicOperation): @@ -267,109 +342,111 @@ def _tex_clamped_dswiglu( self.glu_linear_offset, ) - 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})") + fwd_args_type = ClampedSwiGLUFwdArgs + bwd_args_type = ClampedSwiGLUBwdArgs + num_grad_inputs = 1 - # Check input tensor - x = maybe_dequantize(input_.contiguous(), dtype) + @classmethod + def forward_compute(cls, args: ClampedSwiGLUFwdArgs): + x = maybe_dequantize(args.input_.contiguous(), args.dtype) - # Remove interleaving if needed swiglu_in = x - if self.glu_interleave_size is not None: - shape = swiglu_in.size() - swiglu_in = swiglu_in.reshape( - -1, - shape[-1] // (2 * self.glu_interleave_size), - 2, - self.glu_interleave_size, - ) - swiglu_in = swiglu_in.transpose(1, 2).contiguous() - swiglu_in = swiglu_in.view(shape) + if args.glu_interleave_size is not None: + swiglu_in = _deinterleave_glu(swiglu_in, args.glu_interleave_size) - # Launch kernel - out = self._tex_clamped_swiglu_forward(swiglu_in, next_op_input_quantizer) + out = tex.clamped_swiglu( + swiglu_in, + args.output_quantizer, + args.limit, + args.alpha, + args.glu_linear_offset, + ) - # Quantize input to FP8 before caching if needed - if self.cache_quantized_input: + if args.cache_quantized_input: input_quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, x.device) input_quantizer.set_usage(rowwise=True, columnwise=False) x = input_quantizer(x) + # Same saved-tensor rule as SwiGLU: only a freshly quantized input may + # be handed back; otherwise the backward rebuilds it from the op input. + saved = (x,) if (args.requires_grad and args.cache_quantized_input) else () + return out, saved, SwiGLU._ctx_attrs(args) - # 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 out + @classmethod + def forward_fake(cls, args: ClampedSwiGLUFwdArgs): + return SwiGLU.forward_fake(args) - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - - # Saved tensors from forward pass - (input_,) = ctx.saved_tensors + @classmethod + def backward_compute(cls, args: ClampedSwiGLUBwdArgs): + x = maybe_dequantize(args.saved_input.contiguous(), args.dtype) + dy = maybe_dequantize(args.grad_output.contiguous(), args.dtype) - # Make sure tensors have correct dtypes - x = maybe_dequantize(input_.contiguous(), ctx.dtype) - dy = maybe_dequantize(grad_output.contiguous(), ctx.dtype) - - # Remove interleaving if needed swiglu_in = x - if self.glu_interleave_size is not None: - shape = swiglu_in.size() - swiglu_in = swiglu_in.reshape( - -1, - shape[-1] // (2 * self.glu_interleave_size), - 2, - self.glu_interleave_size, - ) - swiglu_in = swiglu_in.transpose(1, 2).contiguous() - swiglu_in = swiglu_in.view(shape) - - # Quantizer for grad input - quantizer = ctx.prev_op_grad_output_quantizer - if self.glu_interleave_size is not None: + quantizer = args.grad_input_quantizer + if args.glu_interleave_size is not None: + swiglu_in = _deinterleave_glu(swiglu_in, args.glu_interleave_size) quantizer = None - # Launch kernel - grad_swiglu_in = self._tex_clamped_dswiglu(dy, swiglu_in, quantizer) + dx = tex.clamped_dswiglu( + dy, + swiglu_in, + quantizer, + args.limit, + args.alpha, + args.glu_linear_offset, + ) + if args.glu_interleave_size is not None: + dx = _interleave_glu(dx, args.glu_interleave_size) + return (dx,) + + @classmethod + def backward_fake(cls, args: ClampedSwiGLUBwdArgs): + return SwiGLU.backward_fake(args) - # Apply interleaving if needed - dx = grad_swiglu_in - if self.glu_interleave_size is not None: - shape = dx.size() - dx = dx.reshape( - -1, - 2, - shape[-1] // (2 * self.glu_interleave_size), - self.glu_interleave_size, - ) - dx = dx.transpose(1, 2).contiguous() - dx = dx.view(shape) + def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple: + return saved or (input_,) - # Clear input tensor if possible - clear_tensor_data(input_) + 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, + ) -> ClampedSwiGLUFwdArgs: + 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 ClampedSwiGLUFwdArgs( + input_=input_, + dtype=dtype, + output_quantizer=next_op_input_quantizer, + cache_quantized_input=self.cache_quantized_input, + glu_interleave_size=self.glu_interleave_size, + requires_grad=requires_grad, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + limit=self.limit, + alpha=self.alpha, + glu_linear_offset=self.glu_linear_offset, + ) - return dx, () + def resolve_bwd_args( + self, ctx: OperationContext, grad_output: torch.Tensor + ) -> ClampedSwiGLUBwdArgs: + (saved_input,) = ctx.saved_tensors + return ClampedSwiGLUBwdArgs( + grad_output=grad_output, + saved_input=saved_input, + dtype=ctx.dtype, + glu_interleave_size=self.glu_interleave_size, + grad_input_quantizer=ctx.prev_op_grad_output_quantizer, + limit=self.limit, + alpha=self.alpha, + glu_linear_offset=self.glu_linear_offset, + ) class _ScaledGLU(BasicOperation): diff --git a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py index a3c81e60c8..cfcf5b4980 100644 --- a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py +++ b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py @@ -44,8 +44,10 @@ def fuser_backward( rmsnorm_op = self.basic_ops[1] rmsnorm_op_ctx = basic_op_ctxs[1] - # Saved tensors from forward pass - x, rstdevs = rmsnorm_op_ctx.saved_tensors + # Saved tensors from forward pass. The forward saves the operation's raw + # input (a custom op may not return one of its own inputs), so rebuild + # the 2D view the kernel expects here. + saved_input, rstdevs = rmsnorm_op_ctx.saved_tensors # Tensor dims weight_dims = rmsnorm_op.weight.size() @@ -53,6 +55,7 @@ def fuser_backward( # Check input tensors dtype = rmsnorm_op_ctx.dtype + x = maybe_dequantize(saved_input.contiguous(), dtype).view((-1, inner_dim)) extra_grad = basic_op_grad_extra_outputs[0][0] dy = maybe_dequantize(grad_output.contiguous(), dtype).view(x.size()) w = maybe_dequantize(rmsnorm_op.weight, dtype).view((inner_dim,)) diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py index 8df929f799..4f459d01e5 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py @@ -6,7 +6,8 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any, Optional +from dataclasses import dataclass +from typing import Any, Dict, Optional import torch @@ -14,9 +15,34 @@ from ...quantization import FP8GlobalStateManager from ...tensor import Quantizer from ..basic import BasicLinear, Bias +from ..basic.basic_linear import ( + BasicLinearFwdArgs, + _saved_input_is_fresh, + _saved_weight_is_fresh, +) +from .._common import is_quantized_tensor from ..op import FusedOperation, FusibleOperation, OperationContext +@dataclass(slots=True) +class ForwardLinearBiasActivationFwdArgs(BasicLinearFwdArgs): + """Flat inputs to the fused linear+bias forward.""" + + bias: Optional[torch.Tensor] = None + + +@dataclass(slots=True) +class ForwardLinearBiasActivationBwdArgs: + """Placeholder: this fused op has no backward of its own. + + The backward pass walks the basic operations' own custom ops -- that is the + point of pipeline-level autograd -- so this container exists only to + satisfy the registration's schema contract. + """ + + grad_output: Optional[torch.Tensor] = None + + class ForwardLinearBiasActivation(FusedOperation): """Fused forward GEMM + bias + activation @@ -50,6 +76,144 @@ def __init__( # Index of each basic operations self._op_idxs: dict[str, Optional[int]] = op_idxs + fwd_args_type = ForwardLinearBiasActivationFwdArgs + bwd_args_type = ForwardLinearBiasActivationBwdArgs + num_grad_inputs = 1 # never used; the backward walks the basic ops + + @classmethod + def forward_compute(cls, args: ForwardLinearBiasActivationFwdArgs): + output, x_local, w = BasicLinear._functional_forward( + input=args.input_, + weight=args.weight, + bias=args.bias, + dtype=args.dtype, + tensor_parallel_mode=args.tensor_parallel_mode, + tensor_parallel_group=args.tensor_parallel_group, + sequence_parallel=args.sequence_parallel, + with_quantized_compute=args.with_quantized_compute, + backward_override=args.backward_override, + input_quantizer=args.input_quantizer, + weight_quantizer=args.weight_quantizer, + output_quantizer=args.output_quantizer, + input_requires_grad=args.input_requires_grad, + weight_requires_grad=args.weight_requires_grad, + ) + saved_input = ( + x_local if _saved_input_is_fresh(args, is_quantized_tensor(args.input_)) else None + ) + saved_weight = w if _saved_weight_is_fresh(args, is_quantized_tensor(args.weight)) else None + return output, (saved_input, saved_weight), cls._nest_ctx_attrs(args) + + @classmethod + def forward_fake(cls, args: ForwardLinearBiasActivationFwdArgs): + out, saved, _ = BasicLinear.forward_fake(args) + return out, saved, cls._nest_ctx_attrs(args) + + @staticmethod + def _nest_ctx_attrs(args: ForwardLinearBiasActivationFwdArgs) -> Dict[str, Any]: + """Per-basic-op context attrs; ``scatter_ctx`` unpacks them.""" + bias_grad_input_quantizer = args.grad_output_quantizer + if args.backward_override is not None: + bias_grad_input_quantizer = None + return { + "linear": { + "with_quantized_compute": ( + args.with_quantized_compute and args.backward_override is None + ), + "backward_override": args.backward_override, + "input_quantizer": args.input_quantizer, + "weight_quantizer": args.weight_quantizer, + "grad_output_quantizer": args.grad_output_quantizer, + "grad_input_quantizer": args.grad_input_quantizer, + "dtype": args.dtype, + "input_requires_grad": args.input_requires_grad, + "weight_requires_grad": args.weight_requires_grad, + }, + "bias": {"grad_input_quantizer": bias_grad_input_quantizer}, + } + + @classmethod + def backward_compute(cls, args): + raise RuntimeError( + "ForwardLinearBiasActivation has no backward of its own; " + "the backward pass walks the basic operations." + ) + + @classmethod + def backward_fake(cls, args): + raise RuntimeError( + "ForwardLinearBiasActivation has no backward of its own; " + "the backward pass walks the basic operations." + ) + + def resolve_fuser_fwd_args( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> ForwardLinearBiasActivationFwdArgs: + del basic_op_kwargs # the gate rejected any; Bias takes none + linear_op = self.basic_ops[self._op_idxs["linear"]] + linear_op_ctx = basic_op_ctxs[self._op_idxs["linear"]] + bias = None + if self._op_idxs["bias"] is not None: + bias = self.basic_ops[self._op_idxs["bias"]].bias + + input_requires_grad = linear_op_ctx.requires_grad + weight_requires_grad = linear_op_ctx.requires_grad and linear_op.weight.requires_grad + + with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() + if with_quantized_compute: + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None + + if torch.is_autocast_enabled(): + dtype = torch.get_autocast_dtype("cuda") + else: + dtype = linear_op.weight.dtype + + return ForwardLinearBiasActivationFwdArgs( + input_=input_, + weight=linear_op.weight, + dtype=dtype, + tensor_parallel_mode=linear_op.tensor_parallel_mode, + tensor_parallel_group=linear_op.tensor_parallel_group, + sequence_parallel=linear_op.sequence_parallel, + with_quantized_compute=with_quantized_compute, + backward_override=backward_override, + input_quantizer=linear_op.get_quantizer("forward", 0), + weight_quantizer=linear_op.get_quantizer("forward", 1), + output_quantizer=next_op_input_quantizer, + grad_output_quantizer=linear_op.get_quantizer("backward", 0), + grad_input_quantizer=prev_op_grad_output_quantizer, + input_requires_grad=input_requires_grad, + weight_requires_grad=weight_requires_grad, + bias=bias, + ) + + def scatter_ctx( + self, + basic_op_ctxs: list[OperationContext], + saved: tuple, + ctx_attrs: dict[str, Any], + input_: torch.Tensor, + ) -> None: + linear_op = self.basic_ops[self._op_idxs["linear"]] + linear_op_ctx = basic_op_ctxs[self._op_idxs["linear"]] + if linear_op_ctx.requires_grad: + linear_op_ctx.save_for_backward(*linear_op.saved_for_backward(saved, input_)) + for name, value in ctx_attrs["linear"].items(): + setattr(linear_op_ctx, name, value) + if self._op_idxs["bias"] is not None: + bias_op_ctx = basic_op_ctxs[self._op_idxs["bias"]] + if bias_op_ctx.requires_grad: + for name, value in ctx_attrs["bias"].items(): + setattr(bias_op_ctx, name, value) + def fuser_forward( self, basic_op_ctxs: list[OperationContext], diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index a1d8d260d9..ff896a1bf9 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -142,15 +142,27 @@ def forward( if next_op is not None: next_op_input_quantizer = next_op.get_input_quantizer() - if use_compiled: - x = op.compiled_op_forward( - basic_op_ctxs[basic_op_idxs[0]], - x, - prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, - next_op_input_quantizer=next_op_input_quantizer, - **basic_op_kwargs[basic_op_idxs[0]], - ) - fused_op_extra_outputs = [()] + # An inline op's eager pass is pure PyTorch, so Dynamo traces it + # directly; only ops with registered custom ops are routed there. + if use_compiled and not op.compile_inline: + if op.is_fused_op: + x = op.compiled_fuser_forward( + [basic_op_ctxs[idx] for idx in basic_op_idxs], + x, + 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], + ) + fused_op_extra_outputs = [() for _ in basic_op_idxs] + else: + x = op.compiled_op_forward( + basic_op_ctxs[basic_op_idxs[0]], + x, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + **basic_op_kwargs[basic_op_idxs[0]], + ) + fused_op_extra_outputs = [()] else: x, fused_op_extra_outputs = op.fuser_forward( [basic_op_ctxs[idx] for idx in basic_op_idxs], @@ -283,7 +295,7 @@ def backward( # Backward op grad_extra_outputs = [basic_op_grad_extra_outputs[idx] for idx in basic_op_idxs] - if func_ctx.use_compiled: + if func_ctx.use_compiled and not op.compile_inline: dx, grad_params_one = op.compiled_op_backward(basic_op_ctxs[basic_op_idxs[0]], dx) fused_op_grad_params = [grad_params_one] fused_op_grad_extra_inputs = [()] @@ -547,9 +559,15 @@ def maybe_fuse_ops( def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) -> Optional[str]: """Why this group may not run through its operations' custom ops.""" - if len(self._forward_ops) != self._num_basic_ops: - # A fused op covers several basic ops; only single-op groups so far. - return "fused operations are not supported yet" + for op, _ in self._forward_ops: + # A forward fused op compiles like a basic one, through its own + # registered custom op; one that has not declared the compute + # halves sends the group to eager. + if op.is_fused_op and op.compile_ops is None: + return f"{type(op).__name__} does not implement the compute halves" + for op, _ in self._backward_ops: + if op.is_fused_op: + return "backward fused operations are not supported yet" for op, kwargs in zip(self._basic_ops, basic_op_kwargs, strict=True): # A kwarg an operation declares is resolved into its args container # like any other config. Anything else -- notably the preallocated diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index c3af65a831..df25ad3740 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -60,6 +60,95 @@ def save_for_backward(self, *tensors: Optional[torch.Tensor]) -> None: class FusibleOperation(torch.nn.Module, metaclass=abc.ABCMeta): """Tensor operation supported by the operation fuser""" + # torch.compile support, shared by basic and fused operations. An operation + # opts in by declaring the two arg containers and implementing the compute + # classmethods below; this base class then registers its custom ops. + fwd_args_type: Optional[type] = None + bwd_args_type: Optional[type] = None + # Gradients returned by backward_compute: the input's, then any parameters'. + num_grad_inputs: int = 1 + # (forward_fn, backward_fn) pair, or None if the operation cannot be compiled. + compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None + # An operation whose eager passes are pure PyTorch (no TE kernels, no global + # state) opts in with this instead of the compute halves: Dynamo traces its + # eager forward/backward directly, so no custom op is registered. + compile_inline: bool = False + + # Registered op names so far. Two classes may share a name (e.g. the + # ``module.rmsnorm.RMSNorm`` wrapper subclasses ``ops.basic.RMSNorm``); + # re-registering under the same name makes ``torch.library`` destroy the + # first registration, leaving the base class's compile_ops holding dangling + # operator handles -- which surfaces later as heap-layout-dependent + # ``vector::reserve`` / ``bad_alloc`` errors from schema parsing. + _compile_op_name_counts: dict[str, int] = {} + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + if cls.fwd_args_type is None or cls.bwd_args_type is None: + return + if getattr(cls.forward_compute, "__isabstractmethod__", False): + return + for name, arg_type in ( + ("fwd_args_type", cls.fwd_args_type), + ("bwd_args_type", cls.bwd_args_type), + ): + # The op schema is built from the container's fields, so this is the + # framework's actual requirement -- check it where it is declared. + if not dataclasses.is_dataclass(arg_type): + raise TypeError(f"{cls.__name__}.{name} must be a dataclass") + # One registration per class. The compute halves are bound here, so a + # subclass that only swaps kernels (the activations) still gets its own + # op without repeating any of this. Import order is deterministic, so + # the deduplicated name is identical across ranks. + base_name = cls.__name__.lower() + count = FusibleOperation._compile_op_name_counts.get(base_name, 0) + FusibleOperation._compile_op_name_counts[base_name] = count + 1 + cls.compile_ops = register_custom_op( + op_name=base_name if count == 0 else f"{base_name}_{count + 1}", + fwd_arg_type=cls.fwd_args_type, + fwd_impl=cls.forward_compute, + fwd_fake_impl=cls.forward_fake, + bwd_arg_type=cls.bwd_args_type, + bwd_impl=cls.backward_compute, + bwd_fake_impl=cls.backward_fake, + num_grad_inputs=cls.num_grad_inputs, + ) + + # ------------------------------------------------------------------ # + # Compute halves. Classmethods, not free functions: they belong to the + # operation, and binding to the class is what lets a family of operations + # share one implementation while dispatching to per-class kernels. + # ------------------------------------------------------------------ # + + @classmethod + def forward_compute(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]: + """Pure forward: ``(output, tensors_to_save, ctx_attrs)``. + + Takes everything through ``args``; must not read ``self`` or global + state, both of which are invisible to the compiler at this point. + """ + raise NotImplementedError + + @classmethod + def forward_fake(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]: + """Allocation-free twin of :meth:`forward_compute` over ``TensorSpec``. + + Runs as a meta kernel, outside the traced frame, and more than once per + compile, so it must be a pure function of ``args`` -- a read of global + state here is unguarded and can silently disagree with the real impl. + """ + raise NotImplementedError + + @classmethod + def backward_compute(cls, args: Any) -> tuple: + """Pure backward: ``num_grad_inputs`` gradients.""" + raise NotImplementedError + + @classmethod + def backward_fake(cls, args: Any) -> tuple: + """Allocation-free twin of :meth:`backward_compute`.""" + raise NotImplementedError + @property @abc.abstractmethod def is_fused_op(self) -> bool: @@ -185,47 +274,9 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta): # Number of extra tensor outputs num_extra_outputs: int = 0 - # torch.compile support. An operation opts in by declaring the two arg - # containers and implementing the four compute classmethods below; the base - # class then registers its custom ops and drives them from op_forward / - # op_backward, so no operation writes that plumbing itself. - fwd_args_type: Optional[type] = None - bwd_args_type: Optional[type] = None - # Gradients returned by backward_compute: the input's, then any parameters'. - num_grad_inputs: int = 1 # Forward kwargs this operation accepts, resolved into fwd_args_type like # any other config. A kwarg carries no gradient and must not be mutated. fwd_kwarg_names: tuple[str, ...] = () - # (forward_fn, backward_fn) pair, or None if the operation cannot be compiled. - compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None - - def __init_subclass__(cls, **kwargs) -> None: - super().__init_subclass__(**kwargs) - if cls.fwd_args_type is None or cls.bwd_args_type is None: - return - if getattr(cls.forward_compute, "__isabstractmethod__", False): - return - for name, arg_type in ( - ("fwd_args_type", cls.fwd_args_type), - ("bwd_args_type", cls.bwd_args_type), - ): - # The op schema is built from the container's fields, so this is the - # framework's actual requirement -- check it where it is declared. - if not dataclasses.is_dataclass(arg_type): - raise TypeError(f"{cls.__name__}.{name} must be a dataclass") - # One registration per class. The compute halves are bound here, so a - # subclass that only swaps kernels (the activations) still gets its own - # op without repeating any of this. - cls.compile_ops = register_custom_op( - op_name=cls.__name__.lower(), - fwd_arg_type=cls.fwd_args_type, - fwd_impl=cls.forward_compute, - fwd_fake_impl=cls.forward_fake, - bwd_arg_type=cls.bwd_args_type, - bwd_impl=cls.backward_compute, - bwd_fake_impl=cls.backward_fake, - num_grad_inputs=cls.num_grad_inputs, - ) def __init__(self) -> None: super().__init__() @@ -234,41 +285,6 @@ def __init__(self) -> None: self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None - # ------------------------------------------------------------------ # - # Compute halves. Classmethods, not free functions: they belong to the - # operation, and binding to the class is what lets a family of operations - # share one implementation while dispatching to per-class kernels. - # ------------------------------------------------------------------ # - - @classmethod - def forward_compute(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]: - """Pure forward: ``(output, tensors_to_save, ctx_attrs)``. - - Takes everything through ``args``; must not read ``self`` or global - state, both of which are invisible to the compiler at this point. - """ - raise NotImplementedError - - @classmethod - def forward_fake(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]: - """Allocation-free twin of :meth:`forward_compute` over ``TensorSpec``. - - Runs as a meta kernel, outside the traced frame, and more than once per - compile, so it must be a pure function of ``args`` -- a read of global - state here is unguarded and can silently disagree with the real impl. - """ - raise NotImplementedError - - @classmethod - def backward_compute(cls, args: Any) -> tuple: - """Pure backward: ``num_grad_inputs`` gradients.""" - raise NotImplementedError - - @classmethod - def backward_fake(cls, args: Any) -> tuple: - """Allocation-free twin of :meth:`backward_compute`.""" - raise NotImplementedError - def compile_unsupported_reason(self) -> Optional[str]: """Why this operation cannot go through its custom op, or ``None``. @@ -277,7 +293,7 @@ def compile_unsupported_reason(self) -> Optional[str]: Recipe-level limits are not checked here -- they belong to whoever reads the recipe, which is the fuser. """ - if self.compile_ops is None: + if self.compile_ops is None and not self.compile_inline: return f"{self.__class__.__name__} does not implement the compute halves" for mode in ("forward", "backward"): for index in range(self.num_quantizers(mode)): @@ -607,6 +623,10 @@ def op_forward( **kwargs, ) output, saved, ctx_attrs = self.forward_compute(args) + if output is None: + # "The input, unchanged": a custom op may not return one of its own + # inputs, so the compute half hands back None instead. + output = input_ if ctx.requires_grad: ctx.save_for_backward(*self.saved_for_backward(saved, input_)) for name, value in ctx_attrs.items(): @@ -637,6 +657,8 @@ def compiled_op_forward( **kwargs, ) output, saved, ctx_attrs = self.compile_ops[0](args) + if output is None: + output = input_ if ctx.requires_grad: ctx.save_for_backward(*self.saved_for_backward(saved, input_)) for name, value in ctx_attrs.items(): @@ -971,6 +993,61 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None: for op in self.basic_ops: op.pre_fuser_forward(requires_grad=requires_grad) + def resolve_fuser_fwd_args( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> Any: + """Gather the fused forward's inputs into a flat, ``self``-free container. + + The fused counterpart of ``BasicOperation.resolve_fwd_args``; reads the + basic ops' state and global state in the traced region. + """ + raise NotImplementedError + + def scatter_ctx( + self, + basic_op_ctxs: list[OperationContext], + saved: tuple, + ctx_attrs: dict[str, Any], + input_: torch.Tensor, + ) -> None: + """Distribute the forward's saved tensors and attrs onto the basic ops' + contexts, applying any saved-tensor substitutions (a custom op may not + return one of its own inputs).""" + raise NotImplementedError + + def compiled_fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> torch.Tensor: + """:meth:`fuser_forward` routed through this fused operation's custom op. + + Unlike a basic op, a fused op spans several contexts; the op-specific + ``scatter_ctx`` writes each one. + """ + args = self.resolve_fuser_fwd_args( + basic_op_ctxs, + input_, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + basic_op_kwargs=basic_op_kwargs, + ) + output, saved, ctx_attrs = self.compile_ops[0](args) + if output is None: + output = input_ + self.scatter_ctx(basic_op_ctxs, saved, ctx_attrs, input_) + return output + def forward( self, input: torch.Tensor, # pylint: disable=redefined-builtin From a6a91b9627038f032eb1a40636252ff80aff2754 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 13 Aug 2026 14:41:33 +0200 Subject: [PATCH 3/3] [PyTorch] Cache the compile gate's static half per fusion state The gate ran in full on every fuser call, including pure eager: fused-op and per-op checks (which walk each operation's quantizers) were recomputed although they only change when maybe_fuse_ops resets the fusion layout or recipe state. Cache that half there; per call only the kwargs are checked, and calls without kwargs skip even that. The once-per-process warning is unaffected. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/ops/fuser.py | 36 ++++++++++++++++++++----- 1 file changed, 30 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index ff896a1bf9..5a0aba2c96 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -410,6 +410,11 @@ def __init__( self.backward_override = None self._last_amax_history_len = 0 + # Why the group cannot run through its custom ops, ignoring per-call + # kwargs. Depends only on the fusion layout and the operations' state, + # so it is recomputed with them in maybe_fuse_ops rather than per call. + self._static_compile_unsupported_reason: Optional[str] = None + # Flatten list of parameters self._basic_op_params = [list(op.parameters()) for op in self._basic_ops] self._basic_op_num_params = list(map(len, self._basic_op_params)) @@ -551,14 +556,18 @@ def maybe_fuse_ops( # Save current fusion params self.recipe_type, self.first_op_requiring_backward, self.backward_override = fusion_params + # The gate's static half follows the fusion layout and recipe state, + # both of which were just rebuilt. + self._static_compile_unsupported_reason = self._compute_static_unsupported_reason() + # Save amax history length if isinstance(recipe, DelayedScaling): self._last_amax_history_len = recipe.amax_history_len else: self._last_amax_history_len = 0 - def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) -> Optional[str]: - """Why this group may not run through its operations' custom ops.""" + def _compute_static_unsupported_reason(self) -> Optional[str]: + """The gate's per-configuration half; cached by ``maybe_fuse_ops``.""" for op, _ in self._forward_ops: # A forward fused op compiles like a basic one, through its own # registered custom op; one that has not declared the compute @@ -568,7 +577,26 @@ def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) -> for op, _ in self._backward_ops: if op.is_fused_op: return "backward fused operations are not supported yet" + for op in self._basic_ops: + reason = op.compile_unsupported_reason() + if reason is not None: + return reason + return None + + def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) -> Optional[str]: + """Why this group may not run through its operations' custom ops. + + The configuration-dependent half is cached (recomputed only when + ``maybe_fuse_ops`` resets the fusion layout or recipe state); only the + per-call kwargs are checked here, and calls without kwargs skip even + that. + """ + reason = self._static_compile_unsupported_reason + if reason is not None: + return reason for op, kwargs in zip(self._basic_ops, basic_op_kwargs, strict=True): + if not kwargs: + continue # A kwarg an operation declares is resolved into its args container # like any other config. Anything else -- notably the preallocated # buffers of the grouped operations -- is written to by the op, and a @@ -584,10 +612,6 @@ def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) -> values = sorted(name for name, v in kwargs.items() if not isinstance(v, torch.Tensor)) if values: return f"{type(op).__name__} keyword arguments {values} are not tensors" - for op in self._basic_ops: - reason = op.compile_unsupported_reason() - if reason is not None: - return reason return None def _use_compiled(self, basic_op_kwargs: list[dict[str, Any]]) -> bool: