diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py new file mode 100644 index 0000000000..7f4cfcac75 --- /dev/null +++ b/tests/pytorch/test_ops_custom_ops.py @@ -0,0 +1,401 @@ +# 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 + + +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.""" + + 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) + ), + 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] + + +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) + # 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)) + 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)) + 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 + # 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)) + 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): + 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.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. + + 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/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 f4beffe90c..4f2707364c 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, 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/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/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, + ) 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..5a0aba2c96 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 = [()] @@ -398,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)) @@ -539,18 +556,47 @@ 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 _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 + # 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 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.""" - 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" + """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 @@ -566,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: 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