diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py new file mode 100644 index 0000000000..2ad8464c71 --- /dev/null +++ b/tests/pytorch/test_ops_custom_ops.py @@ -0,0 +1,329 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Per-op custom ops for ``transformer_engine.pytorch.ops``. + +An operation opts into ``torch.compile`` by declaring its two argument +containers and implementing four compute halves; ``BasicOperation`` registers +the custom ops and drives them. The checks here are the same for every such +operation and are driven by ``_OP_CASES`` -- adding an operation means adding +one entry. + +Each operation is checked three ways: + +* the data-free fake agrees with the real impl, slot for slot -- the compiled + path slices a flat ``Tensor[]`` payload by what the fake said, so a + disagreement is a silently misassembled tensor rather than an error; +* the registered op reproduces the eager operation; +* both halves trace under ``fullgraph=True``, with and without an FP8 output. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Optional, Tuple + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch.constants import DType +from transformer_engine.pytorch.dynamo import TensorSpec +from transformer_engine.pytorch.dynamo.custom_op import _spec_slot_count, _value_to_flat_tensors +from transformer_engine.pytorch.ops.op import OperationContext +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage +from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer + +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +_device = "cuda" +_dtype = torch.bfloat16 +_HIDDEN = 64 + + +@pytest.fixture(autouse=True) +def _fresh_dynamo(): + """Compile each case from scratch. + + The compiled helpers below are closures over one operation, so every case + recompiles the same code object; without a reset the parametrized runs walk + into Dynamo's recompilation limit and fall back to eager, which is an error + under ``fullgraph=True``. + """ + torch._dynamo.reset() + yield + torch._dynamo.reset() + + +def _fp8_quantizer() -> Any: + """An FP8 quantizer, standing in for the next operation's input quantizer.""" + quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, torch.device(_device)) + quantizer.set_usage(rowwise=True, columnwise=False) + return quantizer + + +# --------------------------------------------------------------------------- # +# The operations under test +# --------------------------------------------------------------------------- # + + +def _build_bias(): + op = te.ops.Bias(_HIDDEN, device=_device, dtype=_dtype) + with torch.no_grad(): + op.bias.copy_(torch.randn_like(op.bias)) + return op + + +@dataclass +class OpCase: + """One operation and how to build it.""" + + name: str + build: Callable[[], Any] + quantizes_output: bool = True + num_grads: int = 1 + in_shape: Tuple[int, ...] = (16, _HIDDEN) + + +_OP_CASES = [ + OpCase(name="Bias", build=_build_bias, quantizes_output=False, num_grads=2), + *( + OpCase(name=cls.__name__, build=cls) + for cls in (te.ops.GELU, te.ops.ReLU, te.ops.SiLU, te.ops.GEGLU, te.ops.ReGLU) + ), +] +_CASE_IDS = [case.name for case in _OP_CASES] +_QUANTIZING = [case for case in _OP_CASES if case.quantizes_output] + + +def _make_input(case: OpCase, requires_grad: bool = False) -> torch.Tensor: + return torch.randn(*case.in_shape, device=_device, dtype=_dtype, requires_grad=requires_grad) + + +def _resolve(op, x, *, fp8_output: bool): + return op.resolve_fwd_args( + x, + requires_grad=True, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=_fp8_quantizer() if fp8_output else None, + ) + + +def _forward_through_ctx(op, x, *, fp8_output: bool): + """Run the eager forward and hand back a context the backward can read.""" + ctx = OperationContext() + ctx.requires_grad = True + y = op.op_forward( + ctx, + x, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=_fp8_quantizer() if fp8_output else None, + ) + ctx.saved_tensors = ctx.to_save + return y, ctx + + +# --------------------------------------------------------------------------- # +# Conformance: the fake must describe what the real impl produces +# --------------------------------------------------------------------------- # + + +def _geometry(value: Any) -> Optional[Tuple]: + """Shape / dtype / quantizer type of a real value, or ``None`` for a sentinel.""" + if value is None: + return None + quantizer = getattr(value, "_quantizer", None) + if not isinstance(value, QuantizedTensorStorage): + quantizer = None + return (tuple(value.shape), value.dtype, type(quantizer) if quantizer else None) + + +def _spec_geometry(spec: Optional[TensorSpec]) -> Optional[Tuple]: + if spec is None: + return None + return (tuple(spec.shape), spec.dtype, type(spec.quantizer) if spec.quantizer else None) + + +def _as_sequence(values: Any) -> Tuple: + if values is None: + return () + if not isinstance(values, (tuple, list)): + return (values,) + return tuple(values) + + +def assert_values_match_specs(real: Any, specs: Any, what: str) -> None: + """Require that the fake describes what the real impl produced. + + The invariant the compiled path relies on is the flat ``Tensor[]`` slot + layout, so that is what gets checked, using the framework's own helpers. + Geometry is compared on top. + """ + real_seq, spec_seq = _as_sequence(real), _as_sequence(specs) + assert len(real_seq) == len(spec_seq), f"{what}: count differs" + for i, (value, spec) in enumerate(zip(real_seq, spec_seq)): + real_slots = len(_value_to_flat_tensors(value)) + fake_slots = _spec_slot_count(spec) + assert real_slots == fake_slots, f"{what}[{i}]: {real_slots} slots vs fake {fake_slots}" + assert _geometry(value) == _spec_geometry(spec), f"{what}[{i}]: geometry differs" + + +# --------------------------------------------------------------------------- # +# Tests +# --------------------------------------------------------------------------- # + + +@_cuda +@pytest.mark.parametrize("case", _OP_CASES, ids=_CASE_IDS) +@pytest.mark.parametrize("fp8_output", [False, True]) +def test_op_fake_matches_real(case: OpCase, fp8_output: bool) -> None: + """The fake must predict output geometry, including whether it is quantized.""" + op = case.build() + cls = type(op) + x = _make_input(case) + args = _resolve(op, x, fp8_output=fp8_output) + + real_out, real_saved, real_attrs = cls.forward_compute(args) + fake_out, fake_saved, fake_attrs = cls.forward_fake(args) + assert_values_match_specs(real_out, fake_out, "forward output") + assert_values_match_specs(real_saved, fake_saved, "saved tensor") + assert set(real_attrs) == set(fake_attrs), "ctx_attrs keys disagree" + + _y, ctx = _forward_through_ctx(op, x, fp8_output=fp8_output) + dy = torch.randn(*real_out.shape, device=_device, dtype=_dtype) + bwd_args = op.resolve_bwd_args(ctx, dy) + real_grads = _as_sequence(cls.backward_compute(bwd_args)) + fake_grads = _as_sequence(cls.backward_fake(bwd_args)) + assert len(real_grads) == len(fake_grads) == case.num_grads + for i, (value, spec) in enumerate(zip(real_grads, fake_grads)): + # One slot per gradient regardless of quantization, so only geometry matters. + assert _geometry(value) == _spec_geometry(spec), f"gradient[{i}]: geometry differs" + + +@_cuda +@pytest.mark.parametrize("case", _OP_CASES, ids=_CASE_IDS) +def test_op_matches_eager(case: OpCase) -> None: + """The registered op must reproduce the eager operation.""" + op = case.build() + assert op.compile_ops is not None, f"{case.name}: custom ops failed to register" + forward_fn, backward_fn = op.compile_ops + + x = _make_input(case, requires_grad=True) + y_ref = op(x) + dy = torch.randn_like(y_ref) + y_ref.backward(dy) + dx_ref = x.grad.clone() + + y, _saved, _attrs = forward_fn(_resolve(op, x.detach(), fp8_output=False)) + _y_eager, ctx = _forward_through_ctx(op, x.detach(), fp8_output=False) + grads = backward_fn(op.resolve_bwd_args(ctx, dy)) + # A None gradient means "the incoming gradient, unchanged" -- a custom op may + # not return one of its own inputs. + dx = grads[0] if grads[0] is not None else dy + + torch.testing.assert_close(y, y_ref) + torch.testing.assert_close(dx, dx_ref) + + +@_cuda +@pytest.mark.parametrize("case", _OP_CASES, ids=_CASE_IDS) +@pytest.mark.parametrize("fp8_output", [False, True]) +def test_op_compiles_fullgraph(case: OpCase, fp8_output: bool) -> None: + """Every operation's halves must trace under ``fullgraph=True``. + + Run under ``no_grad``: the halves carry no autograd of their own (that is the + point of ``register_custom_op``), so a caller wires them into + ``autograd.Function`` -- see ``test_ops_hop_poc.py``. + """ + op = case.build() + assert op.compile_ops is not None, f"{case.name}: custom ops failed to register" + forward_fn, backward_fn = op.compile_ops + x = _make_input(case) + + def fwd(x_): + out, saved, _attrs = forward_fn(_resolve(op, x_, fp8_output=fp8_output)) + # An FP8 output crosses the boundary as its inner buffers and is rebuilt + # on the far side; compare the buffers, since dequantize is not traceable. + if isinstance(out, QuantizedTensorStorage): + return (out._data, out._scale_inv, *saved) + return (out, *saved) + + with torch.no_grad(): + expected = fwd(x) + got = torch.compile(fwd, fullgraph=True)(x) + assert len(got) == len(expected) + for a, b in zip(got, expected): + torch.testing.assert_close(a, b) + + y, ctx = _forward_through_ctx(op, x, fp8_output=fp8_output) + dy = torch.randn(*y.shape, device=_device, dtype=_dtype) + bwd_args = op.resolve_bwd_args(ctx, dy) + + def bwd(args): + return backward_fn(args) + + with torch.no_grad(): + expected_grads = bwd(bwd_args) + got_grads = torch.compile(bwd, fullgraph=True)(bwd_args) + for a, b in zip(got_grads, expected_grads): + if a is None or b is None: + assert a is b is None + continue + torch.testing.assert_close(a, b) + + +@_cuda +@pytest.mark.parametrize("case", _QUANTIZING, ids=[c.name for c in _QUANTIZING]) +def test_op_returns_fp8(case: OpCase) -> None: + """With a next-operation quantizer the op must hand back an FP8 tensor. + + This is what matters for a pipeline: an operation quantizes its output with + the *next* operation's input quantizer, so a quantized tensor crosses the + custom-op boundary -- as its inner buffers, rebuilt on the far side -- rather + than being dequantized at it. + """ + op = case.build() + forward_fn, _ = op.compile_ops + y, _saved, _attrs = forward_fn(_resolve(op, _make_input(case), fp8_output=True)) + assert isinstance(y, QuantizedTensorStorage), f"expected a quantized output, got {type(y)}" + + +@_cuda +@pytest.mark.parametrize("case", _OP_CASES, ids=_CASE_IDS) +def test_op_is_compile_supported(case: OpCase) -> None: + """A converted operation must not gate itself out in a plain configuration.""" + op = case.build() + assert op.compile_unsupported_reason() is None + + +@_cuda +def test_unconverted_op_reports_a_reason() -> None: + """An operation without the compute halves must say so rather than compile.""" + op = te.ops.Identity() + assert op.compile_ops is None + reason = op.compile_unsupported_reason() + assert reason is not None and "compute halves" in reason + + +@_cuda +def test_non_value_opaque_quantizer_is_gated() -> None: + """A quantizer torch.compile cannot specialize on must be refused. + + Delayed scaling is the case that matters: its quantizer holds live + scale/amax tensors, so baking it into the graph would silently freeze stale + scales. None of the operations converted so far hold quantizers, so the + check is driven directly until one that does lands. + """ + from transformer_engine.common.recipe import DelayedScaling + from transformer_engine.pytorch.quantization import RecipeState + + op = te.ops.GELU() + assert op.compile_unsupported_reason() is None + + state = RecipeState.create(DelayedScaling(), mode="forward", num_quantizers=1) + (quantizer,) = state.make_quantizers() + op.num_quantizers = lambda mode: 1 if mode == "forward" else 0 + op.get_quantizer = lambda mode, index: quantizer + + reason = op.compile_unsupported_reason() + assert reason is not None and "value-opaque" in reason diff --git a/tests/pytorch/test_ops_hop_poc.py b/tests/pytorch/test_ops_hop_poc.py new file mode 100644 index 0000000000..a917f4e32f --- /dev/null +++ b/tests/pytorch/test_ops_hop_poc.py @@ -0,0 +1,221 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Proof of concept: a pipeline-level ``autograd.Function`` traced as a HOP. + +This is the load-bearing assumption behind compiling ``ops.Sequential``, checked +end to end before the fuser is touched: + +1. Dynamo traces a ``torch.autograd.Function`` as the ``autograd_function_apply`` + higher-order op, so **both** its forward and its backward end up in the graph. +2. That lets the forward and the backward walk **different op groupings**, which + is what ``OperationFuser`` does -- forward fuses linear+bias while backward + fuses bias+activation, and the two partitions overlap without nesting. An op + whose autograd is registered per op could not do this: autograd would compose + the backward out of the forward's nodes. +3. ``OperationContext`` objects are created inside the forward and read in the + backward. Dynamo permits mutating objects created *within* the HOP scope, so + they never have to cross an op schema. +4. A quantized (FP8) tensor is produced by one op and consumed by the next + *inside* the traced region. + +The ops themselves are the real registered custom ops; only the pipeline driver +is written out longhand here, standing in for the fuser. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch.constants import DType +from transformer_engine.pytorch.ops.basic.activation import ActivationBwdArgs +from transformer_engine.pytorch.ops.basic.bias import BiasBwdArgs +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage +from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer + +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +_device = "cuda" + + +class _OpCtx: + """Stand-in for ``OperationContext``: created in forward, read in backward.""" + + def __init__(self) -> None: + self.saved: tuple = () + self.attrs: dict = {} + + +def _fp8_quantizer() -> Any: + quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, torch.device(_device)) + quantizer.set_usage(rowwise=True, columnwise=False) + return quantizer + + +class _Pipeline(torch.autograd.Function): + """Runs a fixed 3-op pipeline: Bias -> GELU -> Bias. + + The forward walks the ops one by one; the backward walks a *coarser* + grouping, handling (bias1, gelu) as a single step. That asymmetry is the + whole point -- it is only expressible because autograd is owned here, at the + pipeline level, rather than by each op. + """ + + @staticmethod + def forward(func_ctx, x, bias1, bias2, ops, quantize_middle): + bias_fwd, _ = ops["bias"] + act_fwd, _ = ops["act"] + + ctxs = [_OpCtx() for _ in range(3)] + + # op 0: bias + args0 = ops["bias_op1"].resolve_fwd_args( + x, requires_grad=True, prev_op_grad_output_quantizer=None + ) + y0, saved0, attrs0 = bias_fwd(args0) + ctxs[0].saved, ctxs[0].attrs = saved0, attrs0 + + # op 1: activation, quantizing its output with the next op's quantizer + args1 = ops["act_op"].resolve_fwd_args( + y0, + requires_grad=True, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=_fp8_quantizer() if quantize_middle else None, + ) + y1, saved1, attrs1 = act_fwd(args1) + ctxs[1].saved = ops["act_op"].saved_for_backward(saved1, y0) + ctxs[1].attrs = attrs1 + + # op 2: bias, consuming what op 1 produced (FP8 when quantize_middle) + args2 = ops["bias_op2"].resolve_fwd_args( + y1, requires_grad=True, prev_op_grad_output_quantizer=None + ) + y2, saved2, attrs2 = bias_fwd(args2) + ctxs[2].saved, ctxs[2].attrs = saved2, attrs2 + + func_ctx.ctxs = ctxs + func_ctx.ops = ops + func_ctx.save_for_backward(x, bias1, bias2) + return y2 + + @staticmethod + def backward(func_ctx, grad_output): + _, bias_bwd = func_ctx.ops["bias"] + _, act_bwd = func_ctx.ops["act"] + ctxs = func_ctx.ctxs + + # Backward group A: op 2 alone. + dy2, db2 = bias_bwd( + BiasBwdArgs( + grad_output=grad_output, + grad_input_quantizer=ctxs[2].attrs["grad_input_quantizer"], + ) + ) + if dy2 is None: + dy2 = grad_output + + # Backward group B: ops 1 and 0 handled together -- a coarser grouping + # than the forward used. A real fused backward op would replace these + # two calls; what matters here is that the grouping may differ at all. + (dy1,) = act_bwd( + ActivationBwdArgs( + grad_output=dy2, + saved_input=ctxs[1].saved[0], + dtype=ctxs[1].attrs["dtype"], + grad_input_quantizer=ctxs[1].attrs["prev_op_grad_output_quantizer"], + ) + ) + dy0, db1 = bias_bwd( + BiasBwdArgs( + grad_output=dy1, + grad_input_quantizer=ctxs[0].attrs["grad_input_quantizer"], + ) + ) + if dy0 is None: + dy0 = dy1 + + return dy0, db1, db2, None, None + + +def _build(dtype: torch.dtype): + bias_op1 = te.ops.Bias(64, device=_device, dtype=dtype) + bias_op2 = te.ops.Bias(64, device=_device, dtype=dtype) + act_op = te.ops.GELU() + with torch.no_grad(): + bias_op1.bias.copy_(torch.randn_like(bias_op1.bias)) + bias_op2.bias.copy_(torch.randn_like(bias_op2.bias)) + ops = { + "bias": bias_op1.compile_ops, + "act": act_op.compile_ops, + "bias_op1": bias_op1, + "bias_op2": bias_op2, + "act_op": act_op, + } + return ops, bias_op1, bias_op2 + + +def _run(x, ops, bias_op1, bias_op2, quantize_middle): + return _Pipeline.apply(x, bias_op1.bias, bias_op2.bias, ops, quantize_middle) + + +@_cuda +@pytest.mark.parametrize("quantize_middle", [False, True]) +def test_hop_pipeline_compiles_fullgraph(quantize_middle) -> None: + """The pipeline must trace as a HOP and match eager, forward and backward.""" + dtype = torch.bfloat16 + ops, bias_op1, bias_op2 = _build(dtype) + x = torch.randn(16, 64, device=_device, dtype=dtype, requires_grad=True) + dy = torch.randn(16, 64, device=_device, dtype=dtype) + + def step(x_): + return _run(x_, ops, bias_op1, bias_op2, quantize_middle) + + # Eager reference. + out_ref = step(x) + out_ref.backward(dy) + grads_ref = [x.grad.clone(), bias_op1.bias.grad.clone(), bias_op2.bias.grad.clone()] + + for t in (x, bias_op1.bias, bias_op2.bias): + t.grad = None + + compiled = torch.compile(step, fullgraph=True) + out = compiled(x) + out.backward(dy) + grads = [x.grad, bias_op1.bias.grad, bias_op2.bias.grad] + + torch.testing.assert_close(out, out_ref) + for got, expected in zip(grads, grads_ref): + torch.testing.assert_close(got, expected) + + +@_cuda +def test_hop_pipeline_carries_fp8_between_ops() -> None: + """The middle tensor really is FP8, and it is consumed by the next op.""" + dtype = torch.bfloat16 + ops, bias_op1, bias_op2 = _build(dtype) + x = torch.randn(16, 64, device=_device, dtype=dtype, requires_grad=True) + + bias_fwd, _ = ops["bias"] + act_fwd, _ = ops["act"] + args0 = ops["bias_op1"].resolve_fwd_args( + x.detach(), requires_grad=False, prev_op_grad_output_quantizer=None + ) + y0, _, _ = bias_fwd(args0) + args1 = ops["act_op"].resolve_fwd_args( + y0, + requires_grad=False, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=_fp8_quantizer(), + ) + y1, _, _ = act_fwd(args1) + assert isinstance(y1, QuantizedTensorStorage), f"expected FP8, got {type(y1)}" + + args2 = ops["bias_op2"].resolve_fwd_args( + y1, requires_grad=False, prev_op_grad_output_quantizer=None + ) + y2, _, _ = bias_fwd(args2) + assert y2 is not None diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 68b1174474..cb362f92f1 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -4,6 +4,7 @@ import abc import contextlib +import dataclasses import pytest import torch @@ -29,6 +30,7 @@ from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear +from transformer_engine.pytorch.ops.op import BasicOperation from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, Quantizer @@ -1520,3 +1522,141 @@ def fn(inp): "Unexpected recompilation(s) across different batch sizes: " f"{recompile_count_after - recompile_count_baseline} recompile(s) detected" ) + + +# --------------------------------------------------------------------------- # +# transformer_engine.pytorch.ops under torch.compile +# --------------------------------------------------------------------------- # + + +@dataclasses.dataclass(slots=True) +class _ScaleFwdArgs: + """Flat, ``self``-free inputs to the test operation's forward.""" + + input_: torch.Tensor + scale: torch.Tensor + + +@dataclasses.dataclass(slots=True) +class _ScaleBwdArgs: + """Flat inputs to the test operation's backward.""" + + grad_output: torch.Tensor = None + saved_input: torch.Tensor = None + scale: torch.Tensor = None + + +class _ScaleOp(BasicOperation): + """Test-only operation: multiply by a learnable scalar. + + Exists so the fuser's compiled path can be exercised without depending on + which real operations happen to declare their compute halves. It is the + smallest operation that still has a parameter gradient and a saved tensor. + """ + + fwd_args_type = _ScaleFwdArgs + bwd_args_type = _ScaleBwdArgs + num_grad_inputs = 2 # grad input, grad scale + + 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)) + + @classmethod + def forward_compute(cls, args): + return args.input_ * args.scale, (), {} + + @classmethod + def forward_fake(cls, args): + x = args.input_ + return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), (), {} + + @classmethod + def backward_compute(cls, args): + dy = args.grad_output + return dy * args.scale, (dy * args.saved_input).sum() + + @classmethod + def backward_fake(cls, args): + dy = args.grad_output + return ( + TensorSpec(shape=tuple(dy.shape), dtype=dy.dtype, device=dy.device), + TensorSpec(shape=(), dtype=dy.dtype, device=dy.device), + ) + + def saved_for_backward(self, saved, input_): + # The forward produces no distinct tensor for its input, and a custom op + # may not return one of its own inputs. + del saved + return (input_,) + + def resolve_fwd_args( + self, + input_, + *, + requires_grad, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=None, + ): + del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer + return _ScaleFwdArgs(input_=input_, scale=self.scale) + + def resolve_bwd_args(self, ctx, grad_output): + (x,) = ctx.saved_tensors + return _ScaleBwdArgs(grad_output=grad_output, saved_input=x, scale=self.scale) + + +def _assert_sequential_matches_eager(model, compiled, base): + """Run a Sequential eagerly and compiled on identical inputs; compare both + the output and every parameter gradient.""" + inp_eager = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_eager = model(inp_eager) + out_eager.sum().backward() + ref_out = out_eager.detach().clone() + ref_igrad = inp_eager.grad.detach().clone() + ref_pgrads = [p.grad.detach().clone() for p in model.parameters()] + + inp_compiled = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_compiled = compiled(inp_compiled).clone() + out_compiled.sum().backward() + + torch.testing.assert_close(out_compiled, ref_out) + torch.testing.assert_close(inp_compiled.grad, ref_igrad) + for got, expected in zip(model.parameters(), ref_pgrads): + torch.testing.assert_close(got.grad, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_te_ops_single_op_group_compiles(): + """``fullgraph=True`` over an ``OperationFuser`` group holding one operation. + + The pipeline-level ``autograd.Function`` is traced as a higher-order op and + calls the operation's custom ops inside, so forward and backward both end up + in the graph. + """ + torch._dynamo.reset() + model = te.ops.Sequential(_ScaleOp()) + compiled = torch.compile(model, fullgraph=True) + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + _assert_sequential_matches_eager(model, compiled, base) + + +@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. + + Note that this is not a fallback: under ``fullgraph=True`` there is no + leaving the graph, so the pipeline is traced either way and only the choice + of implementation changes. That is why the tracing constraints -- no + mutation of anything from an enclosing scope -- have to hold on both paths. + """ + torch._dynamo.reset() + op = te.ops.Identity() + assert op.compile_unsupported_reason() is not None + + model = te.ops.Sequential(op) + compiled = torch.compile(model, fullgraph=True) + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + _assert_sequential_matches_eager(model, compiled, base) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 3598e54daa..1382cc1c92 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,7 +6,7 @@ from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer from .tensor_spec import TensorSpec, to_tensor_spec -from .custom_op import register_custom_op +from .custom_op import register_custom_op, register_custom_op_with_autograd __all__ = [ "register_value_opaque_quantizer", @@ -14,4 +14,5 @@ "TensorSpec", "to_tensor_spec", "register_custom_op", + "register_custom_op_with_autograd", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index b529deb75a..5fb1c4612e 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -6,7 +6,7 @@ Turns a TE module's eager forward/backward into ``torch.library`` custom ops so ``torch.compile(fullgraph=True)`` traces them as single graph nodes -- no graph -break into the eager ``autograd.Function``. ``register_custom_op`` is the entry +break into the eager ``autograd.Function``. ``register_custom_op_with_autograd`` is the entry point (its docstring documents the per-callable contract); ``module/linear.py`` is the first user. Internal framework API -- exported from ``transformer_engine.pytorch.dynamo``, not re-exported at the top level. @@ -50,7 +50,7 @@ only when its value is trivial (``None`` / all-``None``) at call time. What runs where. Each op registers a data-free fake (``register_fake``) so it -traces under ``torch.compile`` without allocating. ``register_custom_op`` returns +traces under ``torch.compile`` without allocating. ``register_custom_op_with_autograd`` returns ``forward_fn`` -- the drop-in for the eager ``autograd.Function.apply``. A forward call through it: @@ -1036,7 +1036,267 @@ def _split_fwd_fake_result( # --------------------------------------------------------------------------- # -# Op registration +# Two-tier op pair: the registration both public entry points build on +# --------------------------------------------------------------------------- # + + +@dataclasses.dataclass +class _OpPair: + """One registered forward/backward pair, and what a caller needs to drive it.""" + + fwd_adapters: List[_Adapter] + bwd_adapters: List[_Adapter] + fwd_arg_names: List[str] + bwd_arg_names: List[str] + fwd_tensor_field_names: List[str] + bwd_tensor_field_names: List[str] + base_fwd_def: Any + base_fwd_op: Any + base_bwd_op: Any + wrapper_fwd_def: Any + wrapper_fwd_op: Any + wrapper_bwd_op: Any + + +def _register_two_tier_pair( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: int, +) -> _OpPair: + """Define an operation's forward and backward as two-tier custom ops. + + Everything that is common to :func:`register_custom_op_with_autograd` and + :func:`register_custom_op`: schemas from the argument + containers, the base kernels, the wrapper ops that flatten + ``QuantizedTensor`` subclass inputs, and the passthrough registrations. + Autograd is deliberately not touched here -- that is what the two entry + points differ on. + """ + subclass_list = _all_quantized_tensor_subclasses() + + fwd_adapters = _get_adapters(fwd_arg_type) + bwd_adapters = _get_adapters(bwd_arg_type) + fwd_tensor_field_names = _tensor_field_names(fwd_adapters) + bwd_tensor_field_names = _tensor_field_names(bwd_adapters) + + fwd_schema_args, fwd_arg_names = _build_schema(fwd_adapters) + bwd_schema_args, bwd_arg_names = _build_schema(bwd_adapters) + fwd_schema = f"{fwd_schema_args} -> Tensor[]" + bwd_schema = f"{bwd_schema_args} -> Tensor[]" + + wrapper_fwd_name = op_name + wrapper_bwd_name = f"{op_name}_backward" + base_fwd_name = f"{op_name}_base" + base_bwd_name = f"{wrapper_bwd_name}_base" + base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" + + base_fwd_def = _register_kernel( + op_name=base_fwd_name, + schema_str=fwd_schema, + arg_type=fwd_arg_type, + arg_names=fwd_arg_names, + adapters=fwd_adapters, + tensor_field_names=fwd_tensor_field_names, + impl=fwd_impl, + fake_impl=fwd_fake_impl, + format_result=_format_fwd_result, + ) + _register_kernel( + op_name=base_bwd_name, + schema_str=bwd_schema, + arg_type=bwd_arg_type, + arg_names=bwd_arg_names, + adapters=bwd_adapters, + tensor_field_names=bwd_tensor_field_names, + impl=bwd_impl, + fake_impl=bwd_fake_impl, + format_result=lambda g: _format_bwd_result(g, num_grad_inputs, base_bwd_qualname), + ) + + base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) + base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) + + wrapper_fwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_fwd_name, + schema_str=fwd_schema, + base_op=base_fwd_op, + adapters=fwd_adapters, + ) + wrapper_bwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_bwd_name, + schema_str=bwd_schema, + base_op=base_bwd_op, + adapters=bwd_adapters, + ) + wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) + wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) + + fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_adapters) + bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_adapters) + + def _fwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) + return base_fwd_op(*new_args) + + def _bwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) + return base_bwd_op(*new_args) + + for sub in subclass_list: + wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) + wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule) + + for op in (wrapper_fwd_op, wrapper_bwd_op, base_fwd_op, base_bwd_op): + _quantized_tensor_passthrough_ops.add(op.default) + + return _OpPair( + fwd_adapters=fwd_adapters, + bwd_adapters=bwd_adapters, + fwd_arg_names=fwd_arg_names, + bwd_arg_names=bwd_arg_names, + fwd_tensor_field_names=fwd_tensor_field_names, + bwd_tensor_field_names=bwd_tensor_field_names, + base_fwd_def=base_fwd_def, + base_fwd_op=base_fwd_op, + base_bwd_op=base_bwd_op, + wrapper_fwd_def=wrapper_fwd_def, + wrapper_fwd_op=wrapper_fwd_op, + wrapper_bwd_op=wrapper_bwd_op, + ) + + +def _reassemble(specs: List[Any], payload: List[torch.Tensor], cursor: int = 0): + """Rebuild the values described by ``specs`` from ``payload[cursor:]``.""" + out: List[Any] = [] + for spec in specs: + n = _spec_slot_count(spec) + chunk = [_decode_none(t) for t in payload[cursor : cursor + n]] + cursor += n + out.append(_spec_reassemble(spec, chunk)) + return out, cursor + + +# --------------------------------------------------------------------------- # +# Op registration: the forward/backward pair, and the autograd-wired variant +# --------------------------------------------------------------------------- # + + +def register_custom_op( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: int, +) -> Optional[Tuple[Callable[[Any], Any], Callable[[Any], Any]]]: + """Register an op's forward and backward as two independent custom ops. + + Autograd is the caller's: it decides how the two are wired, which is what + lets a pipeline-level ``torch.autograd.Function`` -- traced by Dynamo as a + higher-order op -- group the forward and backward passes differently, as + ``ops.OperationFuser`` does. :func:`register_custom_op_with_autograd` builds + on this and wires them the usual way instead. + + Both ops are two-tier, so ``QuantizedTensor`` subclass inputs pass through + without dequantization. + + Contracts, mirroring :func:`register_custom_op_with_autograd`: + + * ``fwd_impl(fwd_args) -> (*user_outputs, tensors_to_save, ctx_attrs)`` + * ``fwd_fake_impl`` -- its data-free twin over :class:`TensorSpec` + * ``bwd_impl(bwd_args) -> tuple`` of ``num_grad_inputs`` gradients + * ``bwd_fake_impl`` -- its data-free twin + + Returns ``(forward_fn, backward_fn)``: + + * ``forward_fn(fwd_args) -> (outputs, saved_tensors, ctx_attrs)`` -- + ``outputs`` is a single value or a tuple, mirroring ``fwd_impl``'s user + outputs; ``saved_tensors`` is the reassembled ``tensors_to_save`` tuple, + which the caller is expected to persist (e.g. ``ctx.save_for_backward``). + * ``backward_fn(bwd_args) -> tuple`` of gradients. + + Returns ``None`` if registration fails (warned once), so callers can fall + back to eager rather than breaking import. + """ + try: + return _register_custom_op_impl( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_arg_type=bwd_arg_type, + bwd_impl=bwd_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=num_grad_inputs, + ) + except (ImportError, AttributeError, RuntimeError, TypeError) as e: + warn_compile_unsupported( + f"could not register the autograd-free custom ops '{op_name}' ({type(e).__name__}: {e})" + ) + return None + + +def _register_custom_op_impl( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: int, +) -> Tuple[Callable[[Any], Any], Callable[[Any], Any]]: + """Body of :func:`register_custom_op`; see it for semantics.""" + pair = _register_two_tier_pair( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_arg_type=bwd_arg_type, + bwd_impl=bwd_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=num_grad_inputs, + ) + + def forward_fn(fwd_args): + spec_obj = _spec_view(fwd_args, pair.fwd_tensor_field_names) + user_specs, saved_specs, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) + kwargs = _args_to_slots(fwd_args, pair.fwd_adapters) + payload = pair.wrapper_fwd_op(*[kwargs[name] for name in pair.fwd_arg_names]) + + outputs, cursor = _reassemble(user_specs, payload) + saved, _ = _reassemble(saved_specs, payload, cursor) + return (outputs[0] if len(outputs) == 1 else tuple(outputs)), tuple(saved), ctx_attrs + + def backward_fn(bwd_args): + # Unlike the forward payload, each grad occupies exactly one slot + # (``_format_bwd_result`` materializes a TensorSpec grad), so there is + # nothing to reassemble. + kwargs = _args_to_slots(bwd_args, pair.bwd_adapters) + payload = pair.wrapper_bwd_op(*[kwargs[name] for name in pair.bwd_arg_names]) + return tuple(_decode_none(t) for t in payload) + + return forward_fn, backward_fn + + +# --------------------------------------------------------------------------- # +# Autograd-wired registration # --------------------------------------------------------------------------- # @@ -1257,7 +1517,10 @@ def _forward(*flat: Any) -> List[torch.Tensor]: return base_op(*new_args) op = torch.library.custom_op( - f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", _forward, mutates_args=(), schema=schema_str + f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", + _forward, + mutates_args=(), + schema=schema_str, ) op.register_fake(_forward) return op @@ -1277,7 +1540,7 @@ def _all_quantized_tensor_subclasses() -> List[type]: return found -def register_custom_op( +def register_custom_op_with_autograd( *, op_name: str, input_tensors_for_grad: List[str], @@ -1350,7 +1613,7 @@ def register_custom_op( ``torch.compile`` (a graph break) rather than breaking import. """ try: - return _register_custom_op_impl( + return _register_custom_op_with_autograd_impl( op_name=op_name, input_tensors_for_grad=input_tensors_for_grad, fwd_arg_type=fwd_arg_type, @@ -1368,7 +1631,7 @@ def register_custom_op( return None -def _register_custom_op_impl( +def _register_custom_op_with_autograd_impl( *, op_name: str, input_tensors_for_grad: List[str], @@ -1380,7 +1643,7 @@ def _register_custom_op_impl( fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], ) -> Callable[..., Any]: - """Body of :func:`register_custom_op`; see it for semantics.""" + """Body of :func:`register_custom_op_with_autograd`; see it for semantics.""" # Existence check at the API boundary: every ``input_tensors_for_grad`` name # must be an actual field of ``fwd_arg_type`` (differentiability -- whether # that field can carry a gradient -- is checked later in @@ -1390,124 +1653,42 @@ def _register_custom_op_impl( if missing: raise ValueError(f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}") - wrapper_fwd_name = op_name - wrapper_bwd_name = f"{op_name}_backward" - base_fwd_name = f"{op_name}_base" - base_bwd_name = f"{wrapper_bwd_name}_base" - subclass_list = _all_quantized_tensor_subclasses() - - fwd_adapters = _get_adapters(fwd_arg_type) - bwd_adapters = _get_adapters(backward_arg_type) - fwd_tensor_field_names = _tensor_field_names(fwd_adapters) - bwd_tensor_field_names = _tensor_field_names(bwd_adapters) - - fwd_schema_args, fwd_arg_names = _build_schema(fwd_adapters) - bwd_schema_args, bwd_arg_names = _build_schema(bwd_adapters) - - num_grad_inputs = len(input_tensors_for_grad) - slot_count, grad_targets = _resolve_grad_targets(fwd_adapters, input_tensors_for_grad) - - fwd_schema = f"{fwd_schema_args} -> Tensor[]" - bwd_schema = f"{bwd_schema_args} -> Tensor[]" - - base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" - - base_fwd_def = _register_kernel( - op_name=base_fwd_name, - schema_str=fwd_schema, - arg_type=fwd_arg_type, - arg_names=fwd_arg_names, - adapters=fwd_adapters, - tensor_field_names=fwd_tensor_field_names, - impl=fwd_impl, - fake_impl=fwd_fake_impl, - format_result=_format_fwd_result, - ) - _register_kernel( - op_name=base_bwd_name, - schema_str=bwd_schema, - arg_type=backward_arg_type, - arg_names=bwd_arg_names, - adapters=bwd_adapters, - tensor_field_names=bwd_tensor_field_names, - impl=backward_impl, - fake_impl=bwd_fake_impl, - format_result=lambda g: _format_bwd_result(g, num_grad_inputs, base_bwd_qualname), - ) - - base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) - base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) - - wrapper_fwd_def = _register_wrapper_op( - wrapper_op_name=wrapper_fwd_name, - schema_str=fwd_schema, - base_op=base_fwd_op, - adapters=fwd_adapters, - ) - wrapper_bwd_def = _register_wrapper_op( - wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op + pair = _register_two_tier_pair( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_arg_type=backward_arg_type, + bwd_impl=backward_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=len(input_tensors_for_grad), ) + slot_count, grad_targets = _resolve_grad_targets(pair.fwd_adapters, input_tensors_for_grad) autograd_common = { "fwd_arg_type": fwd_arg_type, - "fwd_arg_names": fwd_arg_names, - "fwd_adapters": fwd_adapters, - "fwd_tensor_field_names": fwd_tensor_field_names, - "bwd_arg_names": bwd_arg_names, - "bwd_adapters": bwd_adapters, + "fwd_arg_names": pair.fwd_arg_names, + "fwd_adapters": pair.fwd_adapters, + "fwd_tensor_field_names": pair.fwd_tensor_field_names, + "bwd_arg_names": pair.bwd_arg_names, + "bwd_adapters": pair.bwd_adapters, "slot_count": slot_count, "grad_targets": grad_targets, "setup_context_user": setup_context, "backward_obj_type": backward_arg_type, "fwd_fake_impl": fwd_fake_impl, } - wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) - wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) - - _register_autograd_for_op(fwd_op=base_fwd_def, bwd_op=base_bwd_op, **autograd_common) - _register_autograd_for_op(fwd_op=wrapper_fwd_def, bwd_op=wrapper_bwd_op, **autograd_common) - - fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_adapters) - bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_adapters) - - def _fwd_rule(mode, func, types, args, kwargs): - del mode, func, types, kwargs - new_args = list(args) - for sub in subclass_list: - _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) - return base_fwd_op(*new_args) - - def _bwd_rule(mode, func, types, args, kwargs): - del mode, func, types, kwargs - new_args = list(args) - for sub in subclass_list: - _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) - return base_bwd_op(*new_args) - - for sub in subclass_list: - wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) - wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule) - - _quantized_tensor_passthrough_ops.add(wrapper_fwd_op.default) - _quantized_tensor_passthrough_ops.add(wrapper_bwd_op.default) - _quantized_tensor_passthrough_ops.add(base_fwd_op.default) - _quantized_tensor_passthrough_ops.add(base_bwd_op.default) + _register_autograd_for_op(fwd_op=pair.base_fwd_def, bwd_op=pair.base_bwd_op, **autograd_common) + _register_autograd_for_op( + fwd_op=pair.wrapper_fwd_def, bwd_op=pair.wrapper_bwd_op, **autograd_common + ) def forward_fn(fwd_args): - spec_obj = _spec_view(fwd_args, fwd_tensor_field_names) - user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) - kwargs = _args_to_slots(fwd_args, fwd_adapters) - flat_in = [kwargs[name] for name in fwd_arg_names] - result = wrapper_fwd_op(*flat_in) - - cursor = 0 - outputs: List[Any] = [] - for spec in user_fakes: - n = _spec_slot_count(spec) - chunk = [_decode_none(t) for t in result[cursor : cursor + n]] - cursor += n - outputs.append(_spec_reassemble(spec, chunk)) - + spec_obj = _spec_view(fwd_args, pair.fwd_tensor_field_names) + user_specs, _saved_specs, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) + kwargs = _args_to_slots(fwd_args, pair.fwd_adapters) + payload = pair.wrapper_fwd_op(*[kwargs[name] for name in pair.fwd_arg_names]) + outputs, _ = _reassemble(user_specs, payload) if len(outputs) == 1: return outputs[0] return tuple(outputs) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 219263773d..0e4eb634e0 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -83,7 +83,7 @@ prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorSpec, register_custom_op, is_value_opaque_quantizer +from ..dynamo import TensorSpec, register_custom_op_with_autograd, is_value_opaque_quantizer from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -1753,7 +1753,7 @@ def _linear_backward_impl_fake( # Custom op used under ``torch.compile``. -_linear_op = register_custom_op( +_linear_op = register_custom_op_with_autograd( op_name="linear", input_tensors_for_grad=["weight", "inp", "bias"], fwd_arg_type=LinearFwdArgs, diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index f4beffe90c..d4baa48d93 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -7,13 +7,16 @@ from __future__ import annotations import abc from collections.abc import Iterable -from typing import Any, Optional +from dataclasses import dataclass +from typing import Any, Callable, Dict, Optional, Tuple, Union import torch import transformer_engine_torch as tex from ...constants import DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload +from ...dynamo import TensorSpec +from ...quantized_tensor import QuantizedTensorStorage from ...tensor.float8_tensor import Float8CurrentScalingQuantizer, Quantizer from ...utils import clear_tensor_data from ..op import BasicOperation, OperationContext @@ -33,6 +36,39 @@ "SiLU", ] +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class ActivationFwdArgs: + """Flat, ``self``-free inputs to an activation forward.""" + + input_: TensorOrQuantized + dtype: torch.dtype + output_quantizer: Optional[Quantizer] + cache_quantized_input: bool + requires_grad: bool + prev_op_grad_output_quantizer: Optional[Quantizer] + + +@dataclass(slots=True) +class ActivationBwdArgs: + """Flat inputs to an activation backward.""" + + grad_output: Optional[torch.Tensor] = None + saved_input: Optional[TensorOrQuantized] = None + dtype: Optional[torch.dtype] = None + grad_input_quantizer: Optional[Quantizer] = None + + +def _activation_output_shape( + input_shape: Tuple[int, ...], halves_last_dim: bool +) -> Tuple[int, ...]: + """Output shape of an activation: GLU variants consume pairs along the inner dim.""" + if not halves_last_dim: + return input_shape + return (*input_shape[:-1], input_shape[-1] // 2) + class _ActivationOperation(BasicOperation, metaclass=abc.ABCMeta): r"""Apply activation function @@ -67,83 +103,137 @@ def __init__(self, *, cache_quantized_input: bool = False): super().__init__() self.cache_quantized_input: bool = cache_quantized_input + @staticmethod @abc.abstractmethod - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: """Forward implementation Implementation from transformer_engine.pytorch.cpp_extensions. """ + @staticmethod @abc.abstractmethod - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: """Backward implementation Implementation from transformer_engine_torch. """ - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - - # Compute dtype - dtype: torch.dtype - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") - else: - dtype = input_.dtype - if dtype not in (torch.float32, torch.float16, torch.bfloat16): - raise RuntimeError(f"Unsupported dtype ({dtype})") - - # Check input tensor - x = maybe_dequantize(input_.contiguous(), dtype) + # GLU variants consume pairs along the inner dimension; set per subclass. + _output_halves_last_dim: bool = False - # Launch kernel - y = self._activation_forward_impl(x, next_op_input_quantizer) + fwd_args_type = ActivationFwdArgs + bwd_args_type = ActivationBwdArgs + num_grad_inputs = 1 - # Quantize input to FP8 before caching if needed - if self.cache_quantized_input: + @classmethod + def forward_compute(cls, args: ActivationFwdArgs): + x = maybe_dequantize(args.input_.contiguous(), args.dtype) + y = cls._activation_forward_impl(x, args.output_quantizer) + if args.cache_quantized_input: input_quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, x.device) input_quantizer.set_usage(rowwise=True, columnwise=False) x = input_quantizer(x) + # Only the re-quantized input is handed back. Otherwise x is derived from + # the input by dequantize + contiguous, both no-ops for an already-plain + # contiguous tensor, so x would *be* the input -- which a custom op may + # not return. Whether those calls are no-ops depends on strides, which + # the fake cannot see, so the rule has to be static: the backward + # rebuilds its input from the operation's input instead. + saved = (x,) if (args.requires_grad and args.cache_quantized_input) else () + return y, saved, cls._ctx_attrs(args) + + @classmethod + def forward_fake(cls, args: ActivationFwdArgs): + x = args.input_ + shape = tuple(x.shape) + y = TensorSpec( + shape=_activation_output_shape(shape, cls._output_halves_last_dim), + dtype=args.dtype, + quantizer=args.output_quantizer, + device=x.device, + ) + saved = () + if args.requires_grad and args.cache_quantized_input: + saved_quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, x.device) + saved_quantizer.set_usage(rowwise=True, columnwise=False) + saved = ( + TensorSpec( + shape=shape, dtype=args.dtype, quantizer=saved_quantizer, device=x.device + ), + ) + return y, saved, cls._ctx_attrs(args) + + @staticmethod + def _ctx_attrs(args: ActivationFwdArgs) -> Dict[str, Any]: + return { + "dtype": args.dtype, + "prev_op_grad_output_quantizer": args.prev_op_grad_output_quantizer, + } + + @classmethod + def backward_compute(cls, args: ActivationBwdArgs): + x = maybe_dequantize(args.saved_input.contiguous(), args.dtype) + dy = maybe_dequantize(args.grad_output.contiguous(), x.dtype) + return (cls._activation_backward_impl(dy, x, args.grad_input_quantizer),) + + @classmethod + def backward_fake(cls, args: ActivationBwdArgs): + return ( + TensorSpec( + shape=tuple(args.saved_input.shape), + dtype=args.dtype, + quantizer=args.grad_input_quantizer, + device=args.saved_input.device, + ), + ) - # Save state for backward pass - if ctx.requires_grad: - if is_cpu_offload_enabled(): - mark_activation_offload(x) - ctx.save_for_backward(x) - ctx.dtype = dtype - ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer - - return y - - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: + def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple: + # Without cache_quantized_input the forward keeps nothing, so the + # backward rebuilds its input from the operation's input. + return saved or (input_,) - # Saved tensors from forward pass + def resolve_bwd_args( + self, ctx: OperationContext, grad_output: torch.Tensor + ) -> ActivationBwdArgs: (x,) = ctx.saved_tensors + return ActivationBwdArgs( + grad_output=grad_output, + saved_input=x, + dtype=ctx.dtype, + grad_input_quantizer=ctx.prev_op_grad_output_quantizer, + ) - # Check input tensor - x = maybe_dequantize(x.contiguous(), ctx.dtype) - - # Check grad output tensor - dy = maybe_dequantize(grad_output.contiguous(), x.dtype) - - # Launch kernel - dx = self._activation_backward_impl(dy, x, ctx.prev_op_grad_output_quantizer) - - # Clear input tensor if possible - clear_tensor_data(x) - - return dx, () + def resolve_fwd_args( + self, + input_: torch.Tensor, + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> ActivationFwdArgs: + """Gather everything the forward needs into a flat, self-free container. + + Reads the autocast state, so it must run in the traced region (where + Dynamo guards that read), never inside the custom op. + """ + dtype: torch.dtype + if torch.is_autocast_enabled(): + dtype = torch.get_autocast_dtype("cuda") + else: + dtype = input_.dtype + if dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise RuntimeError(f"Unsupported dtype ({dtype})") + return ActivationFwdArgs( + input_=input_, + dtype=dtype, + output_quantizer=next_op_input_quantizer, + cache_quantized_input=self.cache_quantized_input, + requires_grad=requires_grad, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + ) class GELU(_ActivationOperation): @@ -159,10 +249,12 @@ class GELU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.gelu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dgelu(*args, **kwargs) @@ -191,10 +283,14 @@ class GLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + _output_halves_last_dim: bool = True + + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.glu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dglu(*args, **kwargs) @@ -226,10 +322,14 @@ class GEGLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + _output_halves_last_dim: bool = True + + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.geglu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dgeglu(*args, **kwargs) @@ -245,10 +345,12 @@ class QGELU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.qgelu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dqgelu(*args, **kwargs) @@ -278,10 +380,14 @@ class QGEGLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + _output_halves_last_dim: bool = True + + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.qgeglu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dqgeglu(*args, **kwargs) @@ -294,10 +400,12 @@ class ReLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.relu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.drelu(*args, **kwargs) @@ -323,10 +431,14 @@ class ReGLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + _output_halves_last_dim: bool = True + + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.reglu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dreglu(*args, **kwargs) @@ -341,10 +453,12 @@ class SReLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.srelu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dsrelu(*args, **kwargs) @@ -483,10 +597,14 @@ class SReGLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + _output_halves_last_dim: bool = True + + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.sreglu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dsreglu(*args, **kwargs) @@ -499,8 +617,10 @@ class SiLU(_ActivationOperation): """ - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_forward_impl(*args, **kwargs) -> torch.Tensor: return tex.silu(*args, **kwargs) - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + @staticmethod + def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: return tex.dsilu(*args, **kwargs) diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index 88f563b2c5..9be54a9e1d 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -5,7 +5,8 @@ """Fusible operation for bias.""" from __future__ import annotations -from typing import Optional +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Union import torch @@ -14,6 +15,28 @@ from ..op import BasicOperation, OperationContext from ...utils import canonicalize_device, canonicalize_dtype from ...tensor import Quantizer +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec + +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class BiasFwdArgs: + """Flat, ``self``-free inputs to the bias forward.""" + + input_: TensorOrQuantized + bias: torch.Tensor + local_size: int + grad_input_quantizer: Optional[Quantizer] + + +@dataclass(slots=True) +class BiasBwdArgs: + """Flat inputs to the bias backward.""" + + grad_output: Optional[torch.Tensor] = None + grad_input_quantizer: Optional[Quantizer] = None class Bias(BasicOperation): @@ -37,6 +60,10 @@ class Bias(BasicOperation): """ + fwd_args_type = BiasFwdArgs + bwd_args_type = BiasBwdArgs + num_grad_inputs = 2 # grad input, grad bias + def __init__( self, size: int, @@ -113,37 +140,84 @@ def pre_first_fuser_forward(self) -> None: if self.bias.device.type == "meta": self.reset_parameters() - def op_forward( + @classmethod + def forward_compute(cls, args: BiasFwdArgs) -> Tuple[torch.Tensor, Tuple[()], Dict[str, Any]]: + """Add the bias. Saves no tensors; the backward only needs the quantizer.""" + x = args.input_ + b = args.bias.view([1] * (x.dim() - 1) + [args.local_size]) + return x + b, (), {"grad_input_quantizer": args.grad_input_quantizer} + + @classmethod + def forward_fake(cls, args: BiasFwdArgs) -> Tuple[TensorSpec, Tuple[()], Dict[str, Any]]: + x = args.input_ + out = TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device) + return out, (), {"grad_input_quantizer": args.grad_input_quantizer} + + @classmethod + def backward_compute(cls, args: BiasBwdArgs) -> Tuple[Optional[torch.Tensor], torch.Tensor]: + """Reduce the gradient over every dimension but the inner one. + + Returns ``None`` for the grad input when it is ``grad_output`` + unchanged: a custom op may not return one of its own inputs, and cloning + would cost a full-size copy on the common (unquantized) path. + """ + dy = args.grad_output + if dy.dim() > 1: + quantizer = args.grad_input_quantizer + if quantizer is None: + return None, dy.sum(tuple(range(dy.dim() - 1))) + db, dy = tex.bgrad_quantize(dy, quantizer) + return dy, db + return None, dy + + @classmethod + def backward_fake(cls, args: BiasBwdArgs) -> Tuple[Optional[TensorSpec], TensorSpec]: + dy = args.grad_output + shape = tuple(dy.shape) + if len(shape) > 1: + grad_bias = TensorSpec(shape=(shape[-1],), dtype=dy.dtype, device=dy.device) + quantizer = args.grad_input_quantizer + if quantizer is None: + return None, grad_bias + grad_input = TensorSpec( + shape=shape, dtype=dy.dtype, quantizer=quantizer, device=dy.device + ) + return grad_input, grad_bias + return None, TensorSpec(shape=shape, dtype=dy.dtype, device=dy.device) + + def resolve_fwd_args( self, - ctx: OperationContext, input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - x = input_ - b = self.bias.view([1] * (x.dim() - 1) + [self.local_size]) - - if ctx.requires_grad: - ctx.grad_input_quantizer = prev_op_grad_output_quantizer + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> BiasFwdArgs: + """Gather everything the forward needs into a flat, self-free container. + + Reads module config and global FP8 state, so it must run in the traced + region (where Dynamo guards those reads), never inside the custom op. + + Bias never quantizes its output, so ``next_op_input_quantizer`` is + accepted (every operation resolves through the same signature) and + ignored. + """ + del next_op_input_quantizer + grad_input_quantizer = None + if requires_grad: + grad_input_quantizer = prev_op_grad_output_quantizer if FP8GlobalStateManager.is_fp8_enabled(): - fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() - if fp8_recipe.backward_override is not None: - ctx.grad_input_quantizer = None - - return x + b + if FP8GlobalStateManager.get_fp8_recipe().backward_override is not None: + grad_input_quantizer = None + return BiasFwdArgs( + input_=input_, + bias=self.bias, + local_size=self.local_size, + grad_input_quantizer=grad_input_quantizer, + ) - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - dy = grad_output - if dy.dim() > 1: - quantizer = ctx.grad_input_quantizer - if quantizer is None: - db = dy.sum(tuple(range(dy.dim() - 1))) - else: - db, dy = tex.bgrad_quantize(dy, quantizer) - else: - db = dy - return dy, (db,) + def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> BiasBwdArgs: + return BiasBwdArgs( + grad_output=grad_output, + grad_input_quantizer=ctx.grad_input_quantizer, + ) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 09ffb004dd..fe1eed9ad4 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -6,6 +6,7 @@ from __future__ import annotations from collections.abc import Callable, Iterable, Sequence +import copy import itertools from typing import Any, Optional, TypeAlias @@ -13,6 +14,7 @@ from ..quantization import FP8GlobalStateManager, Recipe, DelayedScaling from ..quantized_tensor import prepare_for_saving, restore_from_func_ctx +from ..dynamo.quantizer_opaque import warn_compile_unsupported from .op import ( BasicOperation, FusibleOperation, @@ -66,6 +68,7 @@ def forward( fuser: OperationFuser, basic_op_kwargs: list[dict[str, Any]], set_output_requires_grad: bool, + use_compiled: bool, *params_and_extra_inputs: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Forward pass @@ -82,6 +85,9 @@ def forward( Keyword arguments to BasicOperation set_output_requires_grad: bool Whether to set ``requires_grad`` flags on returned tensors + use_compiled: bool + Whether to call the operations' custom ops instead of their eager + implementations. Decided once per group by ``OperationFuser``. *params_and_extra_inputs: torch.Tensor Other tensor inputs to include in autograd graph. Consists of parameter tensors, followed by extra operation inputs. @@ -98,9 +104,14 @@ def forward( # Operation autograd contexts basic_op_ctxs = [OperationContext() for _ in range(fuser._num_basic_ops)] - # Mark input tensors as not deletable in backward - for tensor in (input_,) + params_and_extra_inputs: - tensor._do_not_clear = True + # Mark input tensors as not deletable in backward. Skipped whenever this + # is being traced -- not merely when the custom ops are used: these + # tensors are created outside this function, and a higher-order op may + # not mutate anything from an enclosing scope. Under fullgraph there is + # no falling back out of the graph, so the constraint holds either way. + if not torch.compiler.is_compiling(): + for tensor in (input_,) + params_and_extra_inputs: + tensor._do_not_clear = True # Unflatten list of parameters and extra tensor inputs extra_inputs = params_and_extra_inputs[-fuser.num_extra_inputs :] @@ -131,14 +142,23 @@ def forward( if next_op is not None: next_op_input_quantizer = next_op.get_input_quantizer() - x, fused_op_extra_outputs = op.fuser_forward( - [basic_op_ctxs[idx] for idx in basic_op_idxs], - x, - basic_op_extra_inputs=extra_inputs, - 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], - ) + 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, + ) + fused_op_extra_outputs = [()] + else: + x, fused_op_extra_outputs = op.fuser_forward( + [basic_op_ctxs[idx] for idx in basic_op_idxs], + x, + basic_op_extra_inputs=extra_inputs, + 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], + ) for idx, ys in zip(basic_op_idxs, fused_op_extra_outputs): for y in ys: if set_output_requires_grad: @@ -176,9 +196,13 @@ def forward( func_ctx.save_for_backward(*tensors_to_save) func_ctx.tensor_objects = tensor_objects - # Whether to perform recipe update in backward pass + # Whether to perform recipe update in backward pass. Skipped under + # compile: this reads and flips global FP8 state, and delayed + # scaling -- the only recipe it serves -- is gated out anyway. is_first_module = False - if fuser.first_op_requiring_backward < fuser._num_basic_ops: + if not torch.compiler.is_compiling() and ( + fuser.first_op_requiring_backward < fuser._num_basic_ops + ): is_first_module = FP8GlobalStateManager.is_first_fp8_module() # Other context @@ -189,12 +213,17 @@ def forward( func_ctx.num_extra_inputs = fuser.num_extra_inputs func_ctx.num_extra_outputs = len(extra_outputs_flat) func_ctx.is_first_module = is_first_module + func_ctx.use_compiled = use_compiled - # Mark output tensors as not deletable in backward - for tensor in [x] + extra_outputs_flat: - tensor._do_not_clear = True + # Mark output tensors as not deletable in backward (eager only; see above) + if not torch.compiler.is_compiling(): + for tensor in [x] + extra_outputs_flat: + tensor._do_not_clear = True - if set_output_requires_grad: + # Autograd marks the outputs of an ``apply`` itself, so this is only + # needed on the eager path -- and AOTAutograd's functionalization drops + # a requires_grad_() applied to a graph output anyway. + if set_output_requires_grad and not torch.compiler.is_compiling(): x.requires_grad_(fuser.first_op_requiring_backward < fuser._num_basic_ops) if extra_outputs_flat: @@ -219,7 +248,12 @@ def backward( # Restore saved tensors saved_tensors = restore_from_func_ctx(func_ctx) - # Unflatten list of saved tensors + # Unflatten list of saved tensors. Under compile the contexts were + # created in the forward, which is a different subgraph, so writing to + # them here would be a side effect on an enclosing scope; copy them into + # this one instead. The copy carries the attributes the forward set. + if torch.compiler.is_compiling(): + basic_op_ctxs = [copy.copy(ctx) for ctx in basic_op_ctxs] for ctx in basic_op_ctxs: ctx.saved_tensors = saved_tensors[slice(*ctx._saved_tensors_range)] ctx._saved_tensors_range = None @@ -248,14 +282,22 @@ def backward( # Backward op grad_extra_outputs = [basic_op_grad_extra_outputs[idx] for idx in basic_op_idxs] - dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( - [basic_op_ctxs[idx] for idx in basic_op_idxs], - dx, - basic_op_grad_extra_outputs=grad_extra_outputs, - ) + if func_ctx.use_compiled: + 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 = [()] + else: + dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( + [basic_op_ctxs[idx] for idx in basic_op_idxs], + dx, + basic_op_grad_extra_outputs=grad_extra_outputs, + ) for idx, dparams in zip(basic_op_idxs, fused_op_grad_params): grad_params[idx] = dparams - basic_op_ctxs[idx].saved_tensors = None + # Dropping the reference frees the activation early; on the + # compiled path the graph owns that lifetime instead. + if not torch.compiler.is_compiling(): + basic_op_ctxs[idx].saved_tensors = None for idx, dxs in zip(basic_op_idxs, fused_op_grad_extra_inputs): grad_extra_inputs[idx] = dxs @@ -299,6 +341,7 @@ def backward( None, # fuser None, # basic_op_kwargs None, # set_output_requires_grad + None, # use_compiled *grad_params_flat, *grad_extra_inputs_flat, ) @@ -501,6 +544,37 @@ def maybe_fuse_ops( else: self._last_amax_history_len = 0 + def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) -> Optional[str]: + """Why this group may not run through its operations' custom ops.""" + 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" + if any(kwargs for kwargs in basic_op_kwargs): + return "operation keyword arguments are not supported" + 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: + """Whether this group runs through its operations' custom ops. + + Decided once for the whole group: a pipeline compiles as a whole, so one + unsupported operation sends all of them to eager. + + The reason is reported from the eager path only -- ``warnings.warn`` is + not traceable, so warning from inside the traced region would itself + break the graph. A configuration that is never run eagerly therefore + falls back silently. + """ + reason = self._compile_unsupported_reason(basic_op_kwargs) + if reason is None: + return torch.compiler.is_compiling() + if not torch.compiler.is_compiling(): + warn_compile_unsupported(f"running {type(self).__name__} eagerly: {reason}") + return False + def __call__( self, input: torch.Tensor, # pylint: disable=redefined-builtin @@ -541,11 +615,14 @@ def __call__( # Note: We call forward directly when is_grad_enabled=False, # which can expose non-leaf tensors to the inner ops. Avoid # problems in this case by passing set_output_requires_grad=False. + use_compiled = self._use_compiled(basic_op_kwargs) + args = ( input, self, basic_op_kwargs, is_grad_enabled, # set_output_requires_grad + use_compiled, *self._flat_basic_op_params, *extra_inputs, ) diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 5106ec9e0a..2cf28e1e39 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -9,7 +9,7 @@ from collections.abc import Iterable import dataclasses import pickle -from typing import Any, Optional +from typing import Any, Callable, Optional import torch @@ -22,6 +22,7 @@ autocast, ) from ..tensor import Quantizer +from ..dynamo import is_value_opaque_quantizer, register_custom_op @dataclasses.dataclass @@ -184,6 +185,45 @@ 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_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__() @@ -191,6 +231,93 @@ 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``. + + Asked per operation, but acted on per fuser group: a pipeline compiles + as a whole, so one unsupported operation sends the whole group to eager. + Recipe-level limits are not checked here -- they belong to whoever reads + the recipe, which is the fuser. + """ + if self.compile_ops is None: + return f"{self.__class__.__name__} does not implement the compute halves" + for mode in ("forward", "backward"): + for index in range(self.num_quantizers(mode)): + quantizer = self.get_quantizer(mode, index) + if quantizer is not None and not is_value_opaque_quantizer(quantizer): + # Delayed scaling holds live scale/amax tensors, so its + # quantizer cannot be specialized on and would be baked into + # the graph as a stale constant. + return ( + f"{type(quantizer).__name__} is not a torch.compile value-opaque quantizer" + ) + return None + + 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, + ) -> Any: + """Gather the forward's inputs into a flat, ``self``-free container. + + This is where module config and global state are read, so it belongs in + the traced region where Dynamo guards those reads -- never inside the + custom op. + """ + raise NotImplementedError + + def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> Any: + """Rebuild the backward's inputs from the forward's saved state.""" + raise NotImplementedError + + def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple: + """Tensors to persist, given what the forward handed back. + + An operation whose backward needs its input but whose forward does not + produce a distinct tensor for it overrides this; a custom op may not + return one of its own inputs. + """ + del input_ + return saved + @property def is_fused_op(self) -> bool: return False @@ -425,7 +552,6 @@ def _load_fp8_metas(self, fp8_metas: Optional[dict[str, Any]]) -> None: self._fp8_metas[mode][fp8_meta_key].scale.copy_(scale) self._fp8_metas[mode][fp8_meta_key].amax_history.copy_(amax_history) - @abc.abstractmethod def op_forward( self, ctx: OperationContext, @@ -437,6 +563,10 @@ def op_forward( ) -> torch.Tensor: """Forward pass + Operations that declare the compute halves inherit this: it resolves the + arguments, runs the forward, and records what the backward will need. The + rest override it. + Parameters ---------- ctx: OperationContext @@ -454,8 +584,63 @@ def op_forward( Output tensor """ + if self.fwd_args_type is None: + raise NotImplementedError( + f"{self.__class__.__name__} implements neither op_forward nor the compute halves" + ) + if kwargs: + raise ValueError(f"{self.__class__.__name__} forward does not expect keyword arguments") + args = self.resolve_fwd_args( + input_, + requires_grad=ctx.requires_grad, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + ) + output, saved, ctx_attrs = self.forward_compute(args) + if ctx.requires_grad: + ctx.save_for_backward(*self.saved_for_backward(saved, input_)) + for name, value in ctx_attrs.items(): + setattr(ctx, name, value) + return output + + def compiled_op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + ) -> torch.Tensor: + """:meth:`op_forward` routed through this operation's custom op. + + Same bookkeeping, but the computation crosses an op boundary so Dynamo + sees one graph node instead of tracing into the kernels. + """ + args = self.resolve_fwd_args( + input_, + requires_grad=ctx.requires_grad, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + ) + output, saved, ctx_attrs = self.compile_ops[0](args) + if ctx.requires_grad: + ctx.save_for_backward(*self.saved_for_backward(saved, input_)) + for name, value in ctx_attrs.items(): + setattr(ctx, name, value) + return output + + def compiled_op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: + """:meth:`op_backward` routed through this operation's custom op.""" + grads = self.compile_ops[1](self.resolve_bwd_args(ctx, grad_output)) + grad_input = grads[0] + if grad_input is None: + grad_input = grad_output + return grad_input, tuple(grads[1:]) - @abc.abstractmethod def op_backward( self, ctx: OperationContext, @@ -463,6 +648,8 @@ def op_backward( ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: """Backward pass + Counterpart to the inherited :meth:`op_forward`. + Parameters ---------- ctx: OperationContext @@ -478,6 +665,17 @@ def op_backward( Loss gradients w.r.t. parameters """ + if self.bwd_args_type is None: + raise NotImplementedError( + f"{self.__class__.__name__} implements neither op_backward nor the compute halves" + ) + grads = self.backward_compute(self.resolve_bwd_args(ctx, grad_output)) + grad_input = grads[0] + if grad_input is None: + # "The incoming gradient, unchanged": a custom op may not return one + # of its own inputs, so the compute half hands back None instead. + grad_input = grad_output + return grad_input, tuple(grads[1:]) def fuser_forward( self, diff --git a/transformer_engine/pytorch/ops/sequential.py b/transformer_engine/pytorch/ops/sequential.py index cb5dfecb9f..b8724ca460 100644 --- a/transformer_engine/pytorch/ops/sequential.py +++ b/transformer_engine/pytorch/ops/sequential.py @@ -179,9 +179,7 @@ def forward( or grouped MLP. """ - # Create module groups if needed - if self._module_groups is None: - self._module_groups = self._make_module_groups(self._modules.values()) + module_groups = self._get_module_groups() # Route op kwargs to each module group's basic ops group_op_kwargs = self._resolve_op_kwargs(op_kwargs) @@ -189,7 +187,7 @@ def forward( # Forward pass for each module group x = input extra_outputs: list[torch.Tensor] = [] - for group_idx, module_group in enumerate(self._module_groups): + for group_idx, module_group in enumerate(module_groups): if isinstance(module_group, OperationFuser): xs, extra_inputs = ( (x,) + extra_inputs[: module_group.num_extra_inputs], @@ -208,6 +206,17 @@ def forward( return (x,) + tuple(extra_outputs) return x + def _get_module_groups(self) -> list[OperationFuser | torch.nn.Module]: + """Module groups, built once. + + Kept out of the forward pass: building them constructs ``OperationFuser`` + and fused-operation objects, and an ``nn.Module`` cannot be constructed + inside a traced region. + """ + if self._module_groups is None: + self._module_groups = self._make_module_groups(self._modules.values()) + return self._module_groups + def _resolve_op_kwargs( self, op_kwargs: Optional[dict[torch.nn.Module | int, dict[str, Any]]],