From 08cef59890c4c63fd467970a949ab5d625ba6352 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 16:53:55 +0200 Subject: [PATCH 01/11] [PyTorch] Add register_op_halves: fwd/bwd custom ops without autograd glue Registers an op's forward and backward as two independent two-tier custom ops and returns callables for both, leaving autograd to the caller. This lets a pipeline-level autograd.Function (which Dynamo traces as a higher-order op) group the forward and backward passes differently, as ops.OperationFuser does. Reuses the existing schema/adapter/TensorSpec machinery; register_custom_op is untouched. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/__init__.py | 3 +- .../pytorch/dynamo/custom_op.py | 187 ++++++++++++++++++ 2 files changed, 189 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 3598e54daa..a580ea95e8 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_op_halves __all__ = [ "register_value_opaque_quantizer", @@ -14,4 +14,5 @@ "TensorSpec", "to_tensor_spec", "register_custom_op", + "register_op_halves", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index b529deb75a..91dbb9ee5e 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1513,3 +1513,190 @@ def forward_fn(fwd_args): return tuple(outputs) return forward_fn + + +# --------------------------------------------------------------------------- # +# Split registration: forward and backward as independent ops (no autograd) +# --------------------------------------------------------------------------- # + + +def register_op_halves( + *, + 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. + + Unlike :func:`register_custom_op`, no autograd is registered. The caller + wires forward to backward itself -- e.g. a pipeline-level + ``torch.autograd.Function`` that Dynamo traces as a higher-order op, so the + forward and backward passes can be grouped differently (which is what + ``ops.OperationFuser`` does). Both halves are still two-tier, so + ``QuantizedTensor`` subclass inputs pass through without dequantization. + + Contracts, mirroring :func:`register_custom_op`: + + * ``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_op_halves_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 custom op halves '{op_name}' ({type(e).__name__}: {e})" + ) + return None + + +def _register_op_halves_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_op_halves`; see it for semantics.""" + 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}" + + _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) + + def _reassemble(specs: List[Any], payload: List[torch.Tensor], cursor: int) -> Tuple[List, int]: + """Rebuild 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 + + def forward_fn(fwd_args): + spec_obj = _spec_view(fwd_args, 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, fwd_adapters) + payload = wrapper_fwd_op(*[kwargs[name] for name in fwd_arg_names]) + + outputs, cursor = _reassemble(user_specs, payload, 0) + 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, bwd_adapters) + payload = wrapper_bwd_op(*[kwargs[name] for name in bwd_arg_names]) + return tuple(_decode_none(t) for t in payload) + + return forward_fn, backward_fn From bb3f9307a3b2896bd809ed622adc799b9e562ffb Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 16:56:51 +0200 Subject: [PATCH 02/11] [PyTorch] Declare the Bias operation as a custom op Splits Bias into config resolution, pure compute and ctx saving, then registers the compute halves via register_op_halves. resolve_fwd_args reads module config and global FP8 state, so it stays in the traced region where Dynamo guards those reads; the impls take everything as arguments and never touch self. op_forward/op_backward keep their signatures and drive the same impls, so the eager path is unchanged. Adds tests/pytorch/test_ops_custom_ops.py with the fake-vs-real conformance harness every subsequent op will reuse. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_ops_custom_ops.py | 177 +++++++++++++++++++ transformer_engine/pytorch/ops/basic/bias.py | 154 +++++++++++++--- 2 files changed, 310 insertions(+), 21 deletions(-) create mode 100644 tests/pytorch/test_ops_custom_ops.py diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py new file mode 100644 index 0000000000..01b6abe579 --- /dev/null +++ b/tests/pytorch/test_ops_custom_ops.py @@ -0,0 +1,177 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Per-op custom ops for ``transformer_engine.pytorch.ops``. + +Each fusible operation registers its forward and backward as independent custom +ops (``register_op_halves``) so a compiled pipeline can call them directly. The +tests here check each half in isolation -- numerics against the eager op, and +that the data-free fake agrees with the real impl slot for slot. + +The fake is what the compiler believes; if it disagrees with the real impl the +result is a silently misassembled tensor rather than an error, so the +conformance check runs for every op. +""" + +from __future__ import annotations + +from typing import Any, List, Optional, Tuple + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch.dynamo import TensorSpec +from transformer_engine.pytorch.ops.basic.bias import ( + BiasBwdArgs, + BiasFwdArgs, + _bias_backward_impl, + _bias_backward_impl_fake, + _bias_forward_impl, + _bias_forward_impl_fake, + _bias_ops, +) +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage + +# Test setup +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +_device = "cuda" + + +# --------------------------------------------------------------------------- # +# Conformance: the fake must describe what the real impl produces +# --------------------------------------------------------------------------- # + + +def _describe(value: Any) -> Optional[Tuple]: + """Structural fingerprint of a real tensor or of the spec describing it.""" + if value is None: + return None + if isinstance(value, TensorSpec): + quantizer = value.quantizer + return (tuple(value.shape), value.dtype, type(quantizer) if quantizer else 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 _describe_all(values: Any) -> List[Optional[Tuple]]: + if values is None: + return [] + if not isinstance(values, (tuple, list)): + values = (values,) + return [_describe(v) for v in values] + + +def assert_fwd_fake_matches_real(args: Any, impl, fake_impl) -> None: + """Run a forward impl and its fake on the same args; require agreement. + + Compares user outputs, saved tensors (count included -- the compiled path + slices a flat payload by the fake's saved-tensor list) and ``ctx_attrs`` + keys. + """ + real_out, real_saved, real_attrs = impl(args) + fake_out, fake_saved, fake_attrs = fake_impl(args) + + assert _describe_all(real_out) == _describe_all(fake_out), "forward outputs disagree" + assert _describe_all(real_saved) == _describe_all(fake_saved), "saved tensors disagree" + assert set(real_attrs) == set(fake_attrs), "ctx_attrs keys disagree" + + +def assert_bwd_fake_matches_real(args: Any, impl, fake_impl) -> None: + """Run a backward impl and its fake on the same args; require agreement.""" + real = impl(args) + fake = fake_impl(args) + assert _describe_all(real) == _describe_all(fake), "gradients disagree" + + +# --------------------------------------------------------------------------- # +# Bias +# --------------------------------------------------------------------------- # + + +def _bias_op(size: int, dtype: torch.dtype) -> te.ops.Bias: + op = te.ops.Bias(size, device=_device, dtype=dtype) + with torch.no_grad(): + op.bias.copy_(torch.randn_like(op.bias)) + return op + + +def _bias_fwd_args(shape, size: int, dtype: torch.dtype) -> BiasFwdArgs: + op = _bias_op(size, dtype) + x = torch.randn(*shape, device=_device, dtype=dtype) + return op.resolve_fwd_args(x, requires_grad=True, prev_op_grad_output_quantizer=None) + + +@_cuda +@pytest.mark.parametrize("shape", [(16, 32), (2, 8, 32), (32,)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_bias_fake_conformance(shape, dtype) -> None: + args = _bias_fwd_args(shape, shape[-1], dtype) + assert_fwd_fake_matches_real(args, _bias_forward_impl, _bias_forward_impl_fake) + + dy = torch.randn(*shape, device=_device, dtype=dtype) + bwd_args = BiasBwdArgs(grad_output=dy, grad_input_quantizer=None) + assert_bwd_fake_matches_real(bwd_args, _bias_backward_impl, _bias_backward_impl_fake) + + +@_cuda +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_bias_custom_op_matches_eager(dtype) -> None: + """The registered op must reproduce the eager operation, forward and backward.""" + assert _bias_ops is not None, "bias custom ops failed to register" + forward_fn, backward_fn = _bias_ops + + shape, size = (16, 32), 32 + op = _bias_op(size, dtype) + x = torch.randn(*shape, device=_device, dtype=dtype, requires_grad=True) + dy = torch.randn(*shape, device=_device, dtype=dtype) + + # Reference: the op as used today. + y_ref = op(x) + y_ref.backward(dy) + dx_ref, db_ref = x.grad.clone(), op.bias.grad.clone() + + # Same computation through the custom ops. + args = op.resolve_fwd_args( + x.detach(), requires_grad=True, prev_op_grad_output_quantizer=None + ) + y, saved, ctx_attrs = forward_fn(args) + assert saved == (), "bias saves no tensors" + dx, db = backward_fn( + BiasBwdArgs(grad_output=dy, grad_input_quantizer=ctx_attrs["grad_input_quantizer"]) + ) + + torch.testing.assert_close(y, y_ref) + torch.testing.assert_close(dx, dx_ref) + torch.testing.assert_close(db, db_ref) + + +@_cuda +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_bias_custom_op_compiles_fullgraph(dtype) -> None: + """Both halves must trace without a graph break.""" + assert _bias_ops is not None, "bias custom ops failed to register" + forward_fn, backward_fn = _bias_ops + + shape, size = (16, 32), 32 + op = _bias_op(size, dtype) + x = torch.randn(*shape, device=_device, dtype=dtype) + dy = torch.randn(*shape, device=_device, dtype=dtype) + + def fwd(x_): + args = op.resolve_fwd_args( + x_, requires_grad=True, prev_op_grad_output_quantizer=None + ) + out, _saved, _attrs = forward_fn(args) + return out + + def bwd(dy_): + return backward_fn(BiasBwdArgs(grad_output=dy_, grad_input_quantizer=None)) + + torch.testing.assert_close(torch.compile(fwd, fullgraph=True)(x), fwd(x)) + compiled_grads = torch.compile(bwd, fullgraph=True)(dy) + for got, expected in zip(compiled_grads, bwd(dy)): + torch.testing.assert_close(got, expected) diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index 88f563b2c5..c026c1365c 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,94 @@ 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, register_op_halves + + +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 + + +def _bias_forward_impl( + args: BiasFwdArgs, +) -> Tuple[torch.Tensor, Tuple[()], Dict[str, Any]]: + """Bias forward. Saves no tensors; 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} + + +def _bias_forward_impl_fake( + args: BiasFwdArgs, +) -> Tuple[TensorSpec, Tuple[()], Dict[str, Any]]: + """Allocation-free fake of :func:`_bias_forward_impl`.""" + x = args.input_ + out = TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device) + return out, (), {"grad_input_quantizer": args.grad_input_quantizer} + + +def _bias_backward_impl( + args: BiasBwdArgs, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Bias backward: reduce the grad over all but the inner dimension.""" + dy = args.grad_output + if dy.dim() > 1: + quantizer = args.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 _bias_backward_impl_fake( + args: BiasBwdArgs, +) -> Tuple[TensorSpec, TensorSpec]: + """Allocation-free fake of :func:`_bias_backward_impl`. + + Mirrors its branching: with a quantizer the grad input is quantized in place + of the reduction, otherwise both grads stay in high precision. + """ + dy = args.grad_output + shape = tuple(dy.shape) + quantizer = args.grad_input_quantizer if len(shape) > 1 else None + grad_bias_shape = (shape[-1],) if len(shape) > 1 else shape + grad_input = TensorSpec( + shape=shape, dtype=dy.dtype, quantizer=quantizer, device=dy.device + ) + grad_bias = TensorSpec(shape=grad_bias_shape, dtype=dy.dtype, device=dy.device) + return grad_input, grad_bias + + +_bias_ops = register_op_halves( + op_name="bias", + fwd_arg_type=BiasFwdArgs, + fwd_impl=_bias_forward_impl, + fwd_fake_impl=_bias_forward_impl_fake, + bwd_arg_type=BiasBwdArgs, + bwd_impl=_bias_backward_impl, + bwd_fake_impl=_bias_backward_impl_fake, + num_grad_inputs=2, +) class Bias(BasicOperation): @@ -113,6 +202,31 @@ def pre_first_fuser_forward(self) -> None: if self.bias.device.type == "meta": self.reset_parameters() + def resolve_fwd_args( + self, + input_: torch.Tensor, + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer], + ) -> 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. + """ + grad_input_quantizer = None + if requires_grad: + grad_input_quantizer = prev_op_grad_output_quantizer + if FP8GlobalStateManager.is_fp8_enabled(): + 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_forward( self, ctx: OperationContext, @@ -120,30 +234,28 @@ def op_forward( 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]) - + del next_op_input_quantizer # Bias never quantizes its output + args = self.resolve_fwd_args( + input_, + requires_grad=ctx.requires_grad, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + ) + out, saved, ctx_attrs = _bias_forward_impl(args) if ctx.requires_grad: - ctx.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 + ctx.save_for_backward(*saved) + for name, value in ctx_attrs.items(): + setattr(ctx, name, value) + return out 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,) + grad_input, grad_bias = _bias_backward_impl( + BiasBwdArgs( + grad_output=grad_output, + grad_input_quantizer=ctx.grad_input_quantizer, + ) + ) + return grad_input, (grad_bias,) From c5693bd9f5b8c0bbcf19324d51e76ec7f027d07b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 17:08:32 +0200 Subject: [PATCH 03/11] [PyTorch] Return None from a backward that passes its grad through A custom op may not return one of its own inputs, which Bias did whenever the grad input is grad_output unchanged. Cloning would cost a full-size copy on the common unquantized path, so the impl returns None for that slot and the caller substitutes grad_output. Pass-through grads are common (Identity, Reshape, ConstantScale, Quantize), so this is the convention those ops will follow too. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_ops_custom_ops.py | 30 ++++++++---- .../pytorch/dynamo/custom_op.py | 11 ++++- transformer_engine/pytorch/ops/basic/bias.py | 46 ++++++++++--------- 3 files changed, 53 insertions(+), 34 deletions(-) diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py index 01b6abe579..b262e4d904 100644 --- a/tests/pytorch/test_ops_custom_ops.py +++ b/tests/pytorch/test_ops_custom_ops.py @@ -135,14 +135,16 @@ def test_bias_custom_op_matches_eager(dtype) -> None: dx_ref, db_ref = x.grad.clone(), op.bias.grad.clone() # Same computation through the custom ops. - args = op.resolve_fwd_args( - x.detach(), requires_grad=True, prev_op_grad_output_quantizer=None - ) + args = op.resolve_fwd_args(x.detach(), requires_grad=True, prev_op_grad_output_quantizer=None) y, saved, ctx_attrs = forward_fn(args) assert saved == (), "bias saves no tensors" dx, db = backward_fn( BiasBwdArgs(grad_output=dy, grad_input_quantizer=ctx_attrs["grad_input_quantizer"]) ) + # A None grad input means "grad_output unchanged" -- a custom op may not + # return one of its own inputs. + if dx is None: + dx = dy torch.testing.assert_close(y, y_ref) torch.testing.assert_close(dx, dx_ref) @@ -152,7 +154,12 @@ def test_bias_custom_op_matches_eager(dtype) -> None: @_cuda @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_bias_custom_op_compiles_fullgraph(dtype) -> None: - """Both halves must trace without a graph break.""" + """Both halves must trace without a graph break. + + Run under ``no_grad``: the halves carry no autograd of their own (that is + the point of ``register_op_halves``), so a caller is expected to wire them + into its own ``autograd.Function``. + """ assert _bias_ops is not None, "bias custom ops failed to register" forward_fn, backward_fn = _bias_ops @@ -162,16 +169,19 @@ def test_bias_custom_op_compiles_fullgraph(dtype) -> None: dy = torch.randn(*shape, device=_device, dtype=dtype) def fwd(x_): - args = op.resolve_fwd_args( - x_, requires_grad=True, prev_op_grad_output_quantizer=None - ) + args = op.resolve_fwd_args(x_, requires_grad=True, prev_op_grad_output_quantizer=None) out, _saved, _attrs = forward_fn(args) return out def bwd(dy_): return backward_fn(BiasBwdArgs(grad_output=dy_, grad_input_quantizer=None)) - torch.testing.assert_close(torch.compile(fwd, fullgraph=True)(x), fwd(x)) - compiled_grads = torch.compile(bwd, fullgraph=True)(dy) - for got, expected in zip(compiled_grads, bwd(dy)): + with torch.no_grad(): + torch.testing.assert_close(torch.compile(fwd, fullgraph=True)(x), fwd(x)) + compiled_grads = torch.compile(bwd, fullgraph=True)(dy) + expected_grads = bwd(dy) + for got, expected in zip(compiled_grads, expected_grads): + if got is None or expected is None: + assert got is expected is None + continue torch.testing.assert_close(got, expected) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 91dbb9ee5e..eb9754bc36 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1257,7 +1257,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 @@ -1689,7 +1692,11 @@ def forward_fn(fwd_args): outputs, cursor = _reassemble(user_specs, payload, 0) saved, _ = _reassemble(saved_specs, payload, cursor) - return (outputs[0] if len(outputs) == 1 else tuple(outputs)), tuple(saved), ctx_attrs + 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 diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index c026c1365c..e427267ccd 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -18,7 +18,6 @@ from ...quantized_tensor import QuantizedTensorStorage from ...dynamo import TensorSpec, register_op_halves - TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] @@ -60,37 +59,38 @@ def _bias_forward_impl_fake( def _bias_backward_impl( args: BiasBwdArgs, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Bias backward: reduce the grad over all but the inner dimension.""" +) -> Tuple[Optional[torch.Tensor], torch.Tensor]: + """Bias backward: reduce the grad over all but the inner dimension. + + 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; the caller substitutes + ``grad_output`` instead. + """ dy = args.grad_output if dy.dim() > 1: quantizer = args.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 + return None, dy.sum(tuple(range(dy.dim() - 1))) + db, dy = tex.bgrad_quantize(dy, quantizer) + return dy, db + return None, dy def _bias_backward_impl_fake( args: BiasBwdArgs, -) -> Tuple[TensorSpec, TensorSpec]: - """Allocation-free fake of :func:`_bias_backward_impl`. - - Mirrors its branching: with a quantizer the grad input is quantized in place - of the reduction, otherwise both grads stay in high precision. - """ +) -> Tuple[Optional[TensorSpec], TensorSpec]: + """Allocation-free fake of :func:`_bias_backward_impl`.""" dy = args.grad_output shape = tuple(dy.shape) - quantizer = args.grad_input_quantizer if len(shape) > 1 else None - grad_bias_shape = (shape[-1],) if len(shape) > 1 else shape - grad_input = TensorSpec( - shape=shape, dtype=dy.dtype, quantizer=quantizer, device=dy.device - ) - grad_bias = TensorSpec(shape=grad_bias_shape, dtype=dy.dtype, device=dy.device) - return grad_input, grad_bias + 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) _bias_ops = register_op_halves( @@ -258,4 +258,6 @@ def op_backward( grad_input_quantizer=ctx.grad_input_quantizer, ) ) + if grad_input is None: + grad_input = grad_output return grad_input, (grad_bias,) From a9e00307e93ad876b12c0bda11b65340adc5a8de Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 18:56:23 +0200 Subject: [PATCH 04/11] [PyTorch] Declare the activation operations as custom ops Each activation subclass registers its own pair of compute halves, keyed by the class: the kernel pair is fixed by the class and a callable has no place in an op schema, so per-class registration keeps the op registry bounded by the op zoo rather than by the model. The subclass dispatch methods become staticmethods, which is what 'the compute half must not depend on an instance' means in practice; existing self._activation_forward_impl(...) call sites are unaffected. Activations are the first ops here that hand back a quantized tensor: with a next-operation input quantizer the kernel writes FP8 directly, so an FP8 tensor crosses the custom-op boundary instead of being dequantized at it. The forward's saved tensor is the input itself whenever dequantize and contiguous are both no-ops, so it follows the same None convention as the Bias backward. The conformance harness now checks the invariant the compiled path actually relies on -- the flat Tensor[] slot layout, via the framework's own helpers -- rather than approximating it. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_ops_custom_ops.py | 197 +++++++++-- .../pytorch/ops/basic/activation.py | 325 ++++++++++++++---- 2 files changed, 443 insertions(+), 79 deletions(-) diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py index b262e4d904..2427f56a58 100644 --- a/tests/pytorch/test_ops_custom_ops.py +++ b/tests/pytorch/test_ops_custom_ops.py @@ -23,6 +23,8 @@ import transformer_engine.pytorch as te 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.basic.activation import ActivationBwdArgs from transformer_engine.pytorch.ops.basic.bias import ( BiasBwdArgs, BiasFwdArgs, @@ -44,47 +46,70 @@ # --------------------------------------------------------------------------- # -def _describe(value: Any) -> Optional[Tuple]: - """Structural fingerprint of a real tensor or of the spec describing it.""" +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 - if isinstance(value, TensorSpec): - quantizer = value.quantizer - return (tuple(value.shape), value.dtype, type(quantizer) if quantizer else 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 _describe_all(values: Any) -> List[Optional[Tuple]]: +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 [] + return () if not isinstance(values, (tuple, list)): - values = (values,) - return [_describe(v) for v in values] + return (values,) + return tuple(values) -def assert_fwd_fake_matches_real(args: Any, impl, fake_impl) -> None: - """Run a forward impl and its fake on the same args; require agreement. +def assert_values_match_specs(real: Any, specs: Any, what: str) -> None: + """Require that the fake describes what the real impl produced. - Compares user outputs, saved tensors (count included -- the compiled path - slices a flat payload by the fake's saved-tensor list) and ``ctx_attrs`` - keys. + The invariant the compiled path actually 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, but only where the real impl produced + a value: ``None`` is a legal payload meaning "an input, unchanged" (a custom + op may not return one of its own inputs), and it occupies the same single + slot as an unquantized tensor. """ + 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}" + if value is not None: + assert _geometry(value) == _spec_geometry(spec), f"{what}[{i}]: geometry differs" + + +def assert_fwd_fake_matches_real(args: Any, impl, fake_impl) -> None: + """Run a forward impl and its fake on the same args; require agreement.""" real_out, real_saved, real_attrs = impl(args) fake_out, fake_saved, fake_attrs = fake_impl(args) - assert _describe_all(real_out) == _describe_all(fake_out), "forward outputs disagree" - assert _describe_all(real_saved) == _describe_all(fake_saved), "saved tensors disagree" + 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" def assert_bwd_fake_matches_real(args: Any, impl, fake_impl) -> None: - """Run a backward impl and its fake on the same args; require agreement.""" - real = impl(args) - fake = fake_impl(args) - assert _describe_all(real) == _describe_all(fake), "gradients disagree" + """Run a backward impl and its fake on the same args; require agreement. + + Gradients are not slot-counted: the backward payload holds exactly one slot + per gradient regardless of quantization, so only geometry is compared. + """ + real_seq, spec_seq = _as_sequence(impl(args)), _as_sequence(fake_impl(args)) + assert len(real_seq) == len(spec_seq), "gradient count differs" + for i, (value, spec) in enumerate(zip(real_seq, spec_seq)): + assert _geometry(value) == _spec_geometry(spec), f"gradient[{i}]: geometry differs" # --------------------------------------------------------------------------- # @@ -185,3 +210,135 @@ def bwd(dy_): assert got is expected is None continue torch.testing.assert_close(got, expected) + + +# --------------------------------------------------------------------------- # +# Activations +# --------------------------------------------------------------------------- # + + +def _fp8_quantizer(dtype: torch.dtype) -> Any: + """A current-scaling FP8 quantizer, as a next op would hand down.""" + from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer + from transformer_engine.pytorch.constants import DType + + del dtype + quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, torch.device(_device)) + quantizer.set_usage(rowwise=True, columnwise=False) + return quantizer + + +_ACTIVATIONS = [te.ops.GELU, te.ops.ReLU, te.ops.SiLU, te.ops.GEGLU, te.ops.ReGLU] + + +@_cuda +@pytest.mark.parametrize("cls", _ACTIVATIONS) +@pytest.mark.parametrize("quantize_output", [False, True]) +def test_activation_fake_conformance(cls, quantize_output) -> None: + """The fake must predict output geometry, including whether it is quantized.""" + dtype = torch.bfloat16 + op = cls() + x = torch.randn(16, 64, device=_device, dtype=dtype) + args = op.resolve_fwd_args( + x, + requires_grad=True, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=_fp8_quantizer(dtype) if quantize_output else None, + ) + assert_fwd_fake_matches_real(args, op._impls.forward, op._impls.forward_fake) + + out, saved, ctx_attrs = op._impls.forward(args) + dy = torch.randn_like(out if not quantize_output else out.dequantize()) + bwd_args = ActivationBwdArgs( + grad_output=dy, + saved_input=x if saved[0] is None else saved[0], + dtype=ctx_attrs["dtype"], + grad_input_quantizer=None, + ) + assert_bwd_fake_matches_real(bwd_args, op._impls.backward, op._impls.backward_fake) + + +@_cuda +@pytest.mark.parametrize("cls", _ACTIVATIONS) +def test_activation_custom_op_returns_fp8(cls) -> None: + """With a next-op quantizer the registered op must hand back an FP8 tensor. + + This is the case that 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 rather than being dequantized at it. + """ + assert cls._impls.ops is not None, f"{cls.__name__} custom ops failed to register" + forward_fn, _ = cls._impls.ops + + dtype = torch.bfloat16 + op = cls() + x = torch.randn(16, 64, device=_device, dtype=dtype) + args = op.resolve_fwd_args( + x, + requires_grad=True, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=_fp8_quantizer(dtype), + ) + y, saved, _ctx_attrs = forward_fn(args) + + assert isinstance(y, QuantizedTensorStorage), f"expected a quantized output, got {type(y)}" + y_ref, _saved_ref, _ = op._impls.forward(args) + torch.testing.assert_close(y.dequantize(), y_ref.dequantize()) + assert len(saved) == 1 + + +@_cuda +@pytest.mark.parametrize("cls", _ACTIVATIONS) +def test_activation_custom_op_matches_eager(cls) -> None: + dtype = torch.bfloat16 + op = cls() + forward_fn, backward_fn = cls._impls.ops + + x = torch.randn(16, 64, device=_device, dtype=dtype, requires_grad=True) + y_ref = op(x) + dy = torch.randn_like(y_ref) + y_ref.backward(dy) + dx_ref = x.grad.clone() + + args = op.resolve_fwd_args( + x.detach(), + requires_grad=True, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=None, + ) + y, saved, ctx_attrs = forward_fn(args) + # None means "the input, unchanged" -- a custom op may not return its own input. + saved_input = x.detach() if saved[0] is None else saved[0] + (dx,) = backward_fn( + ActivationBwdArgs( + grad_output=dy, + saved_input=saved_input, + dtype=ctx_attrs["dtype"], + grad_input_quantizer=ctx_attrs["prev_op_grad_output_quantizer"], + ) + ) + + torch.testing.assert_close(y, y_ref) + torch.testing.assert_close(dx, dx_ref) + + +@_cuda +@pytest.mark.parametrize("cls", [te.ops.GELU, te.ops.GEGLU]) +def test_activation_custom_op_compiles_fullgraph(cls) -> None: + dtype = torch.bfloat16 + op = cls() + forward_fn, _ = cls._impls.ops + x = torch.randn(16, 64, device=_device, dtype=dtype) + + def fwd(x_): + args = op.resolve_fwd_args( + x_, + requires_grad=False, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=None, + ) + out, _saved, _attrs = forward_fn(args) + return out + + with torch.no_grad(): + torch.testing.assert_close(torch.compile(fwd, fullgraph=True)(x), fwd(x)) diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index f4beffe90c..d8bfe01369 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, register_op_halves +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,151 @@ "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) + + +@dataclass(slots=True) +class _ActivationImpls: + """One activation class's compute halves, plus their registered custom ops. + + Held in a container rather than as class attributes so the functions stay + plain functions instead of being bound as methods on attribute lookup. + """ + + forward: Callable[[ActivationFwdArgs], Any] + forward_fake: Callable[[ActivationFwdArgs], Any] + backward: Callable[[ActivationBwdArgs], Any] + backward_fake: Callable[[ActivationBwdArgs], Any] + ops: Optional[Tuple[Callable[..., Any], Callable[..., Any]]] + + +def _make_activation_ops( + op_name: str, + forward_kernel: Callable[..., torch.Tensor], + backward_kernel: Callable[..., torch.Tensor], + halves_last_dim: bool, +) -> _ActivationImpls: + """Build and register the compute halves for one activation class. + + The kernel pair is baked in here rather than passed as an argument: it is + fixed by the class, and a callable has no place in an op schema. + """ + + def forward_impl(args: ActivationFwdArgs): + x = maybe_dequantize(args.input_.contiguous(), args.dtype) + y = forward_kernel(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) + saved = () + if args.requires_grad: + # Both dequantize and contiguous are no-ops for an input that is + # already plain and contiguous, leaving x as the input itself. A + # custom op may not return one of its own inputs, so hand back None + # and let the caller substitute; the slot count is unchanged because + # a None value and an unquantized spec both occupy one slot. + saved = (None if x is args.input_ else x,) + ctx_attrs = { + "dtype": args.dtype, + "prev_op_grad_output_quantizer": args.prev_op_grad_output_quantizer, + } + return y, saved, ctx_attrs + + def forward_fake_impl(args: ActivationFwdArgs): + x = args.input_ + shape = tuple(x.shape) + y = TensorSpec( + shape=_activation_output_shape(shape, halves_last_dim), + dtype=args.dtype, + quantizer=args.output_quantizer, + device=x.device, + ) + saved = () + if args.requires_grad: + # ``cache_quantized_input`` re-quantizes the dequantized input with a + # fresh current-scaling quantizer, so the saved tensor is quantized + # exactly when that flag is set -- never because the input was. + saved_quantizer = None + if 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 + ), + ) + ctx_attrs = { + "dtype": args.dtype, + "prev_op_grad_output_quantizer": args.prev_op_grad_output_quantizer, + } + return y, saved, ctx_attrs + + def backward_impl(args: ActivationBwdArgs): + x = maybe_dequantize(args.saved_input.contiguous(), args.dtype) + dy = maybe_dequantize(args.grad_output.contiguous(), x.dtype) + dx = backward_kernel(dy, x, args.grad_input_quantizer) + return (dx,) + + def backward_fake_impl(args: ActivationBwdArgs): + shape = tuple(args.saved_input.shape) + return ( + TensorSpec( + shape=shape, + dtype=args.dtype, + quantizer=args.grad_input_quantizer, + device=args.saved_input.device, + ), + ) + + return _ActivationImpls( + forward=forward_impl, + forward_fake=forward_fake_impl, + backward=backward_impl, + backward_fake=backward_fake_impl, + ops=register_op_halves( + op_name=op_name, + fwd_arg_type=ActivationFwdArgs, + fwd_impl=forward_impl, + fwd_fake_impl=forward_fake_impl, + bwd_arg_type=ActivationBwdArgs, + bwd_impl=backward_impl, + bwd_fake_impl=backward_fake_impl, + num_grad_inputs=1, + ), + ) + class _ActivationOperation(BasicOperation, metaclass=abc.ABCMeta): r"""Apply activation function @@ -67,31 +215,55 @@ 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( + # GLU variants consume pairs along the inner dimension; set per subclass. + _output_halves_last_dim: bool = False + _impls: Optional[_ActivationImpls] = None + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + # The kernel pair is fixed by the class, so each subclass gets its own + # registration; the op registry stays bounded by the op zoo, not by the + # model. + if getattr(cls._activation_forward_impl, "__isabstractmethod__", False): + return + cls._impls = _make_activation_ops( + op_name=f"activation_{cls.__name__.lower()}", + forward_kernel=cls._activation_forward_impl, + backward_kernel=cls._activation_backward_impl, + halves_last_dim=cls._output_halves_last_dim, + ) + + def resolve_fwd_args( self, - ctx: OperationContext, input_: torch.Tensor, + *, + requires_grad: bool, prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: + ) -> ActivationFwdArgs: + """Gather everything the forward needs into a flat, self-free container. - # Compute dtype + 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") @@ -99,27 +271,36 @@ def op_forward( 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, + ) - # Check input tensor - x = maybe_dequantize(input_.contiguous(), dtype) - - # Launch kernel - y = self._activation_forward_impl(x, next_op_input_quantizer) - - # Quantize input to FP8 before caching if needed - if self.cache_quantized_input: - input_quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, x.device) - input_quantizer.set_usage(rowwise=True, columnwise=False) - x = input_quantizer(x) - - # Save state for backward pass + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + ) -> torch.Tensor: + 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, + ) + y, saved, ctx_attrs = self._impls.forward(args) if ctx.requires_grad: + saved = tuple(input_ if t is None else t for t in saved) 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 - + mark_activation_offload(*saved) + ctx.save_for_backward(*saved) + for name, value in ctx_attrs.items(): + setattr(ctx, name, value) return y def op_backward( @@ -127,22 +308,18 @@ def op_backward( ctx: OperationContext, grad_output: torch.Tensor, ) -> tuple[torch.Tensor, tuple[()]]: - - # Saved tensors from forward pass (x,) = ctx.saved_tensors - - # 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 + (dx,) = self._impls.backward( + ActivationBwdArgs( + grad_output=grad_output, + saved_input=x, + dtype=ctx.dtype, + grad_input_quantizer=ctx.prev_op_grad_output_quantizer, + ) + ) + # Eager only: the compiled path leaves saved tensors to the graph's own + # memory planning, and clearing an op input would lie to functionalization. clear_tensor_data(x) - return dx, () @@ -159,10 +336,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 +370,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 +409,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 +432,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 +467,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 +487,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 +518,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 +540,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 +684,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 +704,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) From cd2655c019c6df26f84c079041a1a88e8ea6a54d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 19:00:44 +0200 Subject: [PATCH 05/11] [PyTorch] Add a proof of concept for the pipeline-level HOP Checks the assumption the whole approach rests on, before the fuser is touched: Dynamo traces a pipeline-level autograd.Function as the autograd_function_apply higher-order op, so the forward and the backward can walk different op groupings -- which is what OperationFuser does, and what an op with per-op autograd could not express. The pipeline runs Bias -> GELU -> Bias through the real registered custom ops, creates context objects inside the forward and reads them in the backward, and carries an FP8 tensor from one op to the next inside the traced region. Forward and backward match eager under fullgraph=True. Also makes the activation's saved tensor a static choice. It used to be returned as None when dequantize and contiguous were both no-ops, but the fake cannot see strides and so cannot predict that; the resulting metadata mismatch surfaced as an inductor assertion on the sentinel's rank. The op now keeps the input only when cache_quantized_input is set, and the backward rebuilds its input from the operation's input otherwise. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_ops_custom_ops.py | 6 +- tests/pytorch/test_ops_hop_poc.py | 221 ++++++++++++++++++ .../pytorch/ops/basic/activation.py | 30 ++- 3 files changed, 237 insertions(+), 20 deletions(-) create mode 100644 tests/pytorch/test_ops_hop_poc.py diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py index 2427f56a58..0766c94393 100644 --- a/tests/pytorch/test_ops_custom_ops.py +++ b/tests/pytorch/test_ops_custom_ops.py @@ -251,7 +251,7 @@ def test_activation_fake_conformance(cls, quantize_output) -> None: dy = torch.randn_like(out if not quantize_output else out.dequantize()) bwd_args = ActivationBwdArgs( grad_output=dy, - saved_input=x if saved[0] is None else saved[0], + saved_input=saved[0] if saved else x, dtype=ctx_attrs["dtype"], grad_input_quantizer=None, ) @@ -284,7 +284,7 @@ def test_activation_custom_op_returns_fp8(cls) -> None: assert isinstance(y, QuantizedTensorStorage), f"expected a quantized output, got {type(y)}" y_ref, _saved_ref, _ = op._impls.forward(args) torch.testing.assert_close(y.dequantize(), y_ref.dequantize()) - assert len(saved) == 1 + assert saved == (), "without cache_quantized_input the op keeps nothing" @_cuda @@ -308,7 +308,7 @@ def test_activation_custom_op_matches_eager(cls) -> None: ) y, saved, ctx_attrs = forward_fn(args) # None means "the input, unchanged" -- a custom op may not return its own input. - saved_input = x.detach() if saved[0] is None else saved[0] + saved_input = saved[0] if saved else x.detach() (dx,) = backward_fn( ActivationBwdArgs( grad_output=dy, diff --git a/tests/pytorch/test_ops_hop_poc.py b/tests/pytorch/test_ops_hop_poc.py new file mode 100644 index 0000000000..98dd3a1020 --- /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 = saved1 or (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": te.ops.basic.bias._bias_ops, + "act": type(act_op)._impls.ops, + "bias_op1": bias_op1, + "bias_op2": bias_op2, + "act_op": act_op, + } + return ops, bias_op1, bias_op2 + + +def _run(x, ops, bias_op1, bias_op2, quantize_middle): + return _Pipeline.apply(x, bias_op1.bias, bias_op2.bias, ops, quantize_middle) + + +@_cuda +@pytest.mark.parametrize("quantize_middle", [False, True]) +def test_hop_pipeline_compiles_fullgraph(quantize_middle) -> None: + """The pipeline must trace as a HOP and match eager, forward and backward.""" + dtype = torch.bfloat16 + ops, bias_op1, bias_op2 = _build(dtype) + x = torch.randn(16, 64, device=_device, dtype=dtype, requires_grad=True) + dy = torch.randn(16, 64, device=_device, dtype=dtype) + + def step(x_): + return _run(x_, ops, bias_op1, bias_op2, quantize_middle) + + # Eager reference. + out_ref = step(x) + out_ref.backward(dy) + grads_ref = [x.grad.clone(), bias_op1.bias.grad.clone(), bias_op2.bias.grad.clone()] + + for t in (x, bias_op1.bias, bias_op2.bias): + t.grad = None + + compiled = torch.compile(step, fullgraph=True) + out = compiled(x) + out.backward(dy) + grads = [x.grad, bias_op1.bias.grad, bias_op2.bias.grad] + + torch.testing.assert_close(out, out_ref) + for got, expected in zip(grads, grads_ref): + torch.testing.assert_close(got, expected) + + +@_cuda +def test_hop_pipeline_carries_fp8_between_ops() -> None: + """The middle tensor really is FP8, and it is consumed by the next op.""" + dtype = torch.bfloat16 + ops, bias_op1, bias_op2 = _build(dtype) + x = torch.randn(16, 64, device=_device, dtype=dtype, requires_grad=True) + + bias_fwd, _ = ops["bias"] + act_fwd, _ = ops["act"] + args0 = ops["bias_op1"].resolve_fwd_args( + x.detach(), requires_grad=False, prev_op_grad_output_quantizer=None + ) + y0, _, _ = bias_fwd(args0) + args1 = ops["act_op"].resolve_fwd_args( + y0, + requires_grad=False, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=_fp8_quantizer(), + ) + y1, _, _ = act_fwd(args1) + assert isinstance(y1, QuantizedTensorStorage), f"expected FP8, got {type(y1)}" + + args2 = ops["bias_op2"].resolve_fwd_args( + y1, requires_grad=False, prev_op_grad_output_quantizer=None + ) + y2, _, _ = bias_fwd(args2) + assert y2 is not None diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index d8bfe01369..5d607edf8c 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -104,14 +104,13 @@ def forward_impl(args: ActivationFwdArgs): input_quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, x.device) input_quantizer.set_usage(rowwise=True, columnwise=False) x = input_quantizer(x) - saved = () - if args.requires_grad: - # Both dequantize and contiguous are no-ops for an input that is - # already plain and contiguous, leaving x as the input itself. A - # custom op may not return one of its own inputs, so hand back None - # and let the caller substitute; the slot count is unchanged because - # a None value and an unquantized spec both occupy one slot. - saved = (None if x is args.input_ else x,) + # Only the re-quantized input is handed back. Otherwise x is derived from + # the input by dequantize + contiguous, both of which are 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 a static one: + # the caller passes the input to the backward, which dequantizes it there. + saved = (x,) if (args.requires_grad and args.cache_quantized_input) else () ctx_attrs = { "dtype": args.dtype, "prev_op_grad_output_quantizer": args.prev_op_grad_output_quantizer, @@ -128,14 +127,9 @@ def forward_fake_impl(args: ActivationFwdArgs): device=x.device, ) saved = () - if args.requires_grad: - # ``cache_quantized_input`` re-quantizes the dequantized input with a - # fresh current-scaling quantizer, so the saved tensor is quantized - # exactly when that flag is set -- never because the input was. - saved_quantizer = None - if args.cache_quantized_input: - saved_quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, x.device) - saved_quantizer.set_usage(rowwise=True, columnwise=False) + 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 @@ -295,7 +289,9 @@ def op_forward( ) y, saved, ctx_attrs = self._impls.forward(args) if ctx.requires_grad: - saved = tuple(input_ if t is None else t for t in saved) + # Without ``cache_quantized_input`` the op keeps nothing: the backward + # rebuilds its input from the operation's input tensor. + saved = saved or (input_,) if is_cpu_offload_enabled(): mark_activation_offload(*saved) ctx.save_for_backward(*saved) From e385a0ba8ac57d678e30259d31c441b176ff7291 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 21:03:12 +0200 Subject: [PATCH 06/11] [PyTorch] Make the per-op tests table-driven The checks were written per operation, so each one restated the same three questions in its own shape. They are now driven by a single _OP_CASES list: adding an operation means adding one entry, and every operation gets the same fake-conformance, matches-eager and fullgraph checks, each run with and without an FP8 output. resolve_fwd_args takes the same arguments on every operation to make that possible; Bias accepts next_op_input_quantizer and ignores it, as its op_forward already did. Compile cases reset Dynamo between runs: the compiled helpers are closures over one operation, so the parametrized runs otherwise walk into the recompilation limit and fall back to eager, which fullgraph=True rejects. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_ops_custom_ops.py | 463 +++++++++--------- .../pytorch/ops/basic/activation.py | 4 +- transformer_engine/pytorch/ops/basic/bias.py | 8 +- 3 files changed, 230 insertions(+), 245 deletions(-) diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py index 0766c94393..5a3d210bce 100644 --- a/tests/pytorch/test_ops_custom_ops.py +++ b/tests/pytorch/test_ops_custom_ops.py @@ -4,41 +4,157 @@ """Per-op custom ops for ``transformer_engine.pytorch.ops``. -Each fusible operation registers its forward and backward as independent custom +Every fusible operation registers its forward and backward as independent custom ops (``register_op_halves``) so a compiled pipeline can call them directly. The -tests here check each half in isolation -- numerics against the eager op, and -that the data-free fake agrees with the real impl slot for slot. +checks here are the same for every operation and are driven by ``_OP_CASES``: +adding an operation means adding one entry to that list. -The fake is what the compiler believes; if it disagrees with the real impl the -result is a silently misassembled tensor rather than an error, so the -conformance check runs for every op. +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 typing import Any, List, Optional, Tuple +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.basic import bias as _bias_mod from transformer_engine.pytorch.ops.basic.activation import ActivationBwdArgs -from transformer_engine.pytorch.ops.basic.bias import ( - BiasBwdArgs, - BiasFwdArgs, - _bias_backward_impl, - _bias_backward_impl_fake, - _bias_forward_impl, - _bias_forward_impl_fake, - _bias_ops, -) +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 -# Test setup _cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@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() + + _device = "cuda" +_dtype = torch.bfloat16 +_HIDDEN = 64 + + +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 +# --------------------------------------------------------------------------- # + + +@dataclass +class OpCase: + """One operation and how to drive it. + + Everything except ``bwd_args`` is uniform; that one callable exists because + each operation's backward container names its own fields. + """ + + name: str + build: Callable[[], Any] + bwd_args: Callable[..., Any] + impls: Callable[[Any], Any] + quantizes_output: bool = True + num_grads: int = 1 + in_shape: Tuple[int, ...] = (16, _HIDDEN) + + +def _bias_bwd_args(grad_output, saved, ctx_attrs, input_): + del saved, input_ + return BiasBwdArgs( + grad_output=grad_output, + grad_input_quantizer=ctx_attrs["grad_input_quantizer"], + ) + + +def _activation_bwd_args(grad_output, saved, ctx_attrs, input_): + return ActivationBwdArgs( + grad_output=grad_output, + saved_input=saved[0] if saved else input_, + dtype=ctx_attrs["dtype"], + grad_input_quantizer=ctx_attrs["prev_op_grad_output_quantizer"], + ) + + +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 + + +class _BiasImpls: + """The Bias module's halves, in the same shape the activations expose.""" + + forward = staticmethod(_bias_mod._bias_forward_impl) + forward_fake = staticmethod(_bias_mod._bias_forward_impl_fake) + backward = staticmethod(_bias_mod._bias_backward_impl) + backward_fake = staticmethod(_bias_mod._bias_backward_impl_fake) + ops = _bias_mod._bias_ops + + +_OP_CASES = [ + OpCase( + name="Bias", + build=_build_bias, + bwd_args=_bias_bwd_args, + impls=lambda op: _BiasImpls, + quantizes_output=False, + num_grads=2, + ), + *( + OpCase( + name=cls.__name__, + build=cls, + bwd_args=_activation_bwd_args, + impls=lambda op: type(op)._impls, + ) + 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, + ) # --------------------------------------------------------------------------- # @@ -73,12 +189,9 @@ def _as_sequence(values: Any) -> Tuple: 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 actually 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, but only where the real impl produced - a value: ``None`` is a legal payload meaning "an input, unchanged" (a custom - op may not return one of its own inputs), and it occupies the same single - slot as an unquantized tensor. + 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" @@ -86,259 +199,125 @@ def assert_values_match_specs(real: Any, specs: Any, what: str) -> None: 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}" - if value is not None: - assert _geometry(value) == _spec_geometry(spec), f"{what}[{i}]: geometry differs" - - -def assert_fwd_fake_matches_real(args: Any, impl, fake_impl) -> None: - """Run a forward impl and its fake on the same args; require agreement.""" - real_out, real_saved, real_attrs = impl(args) - fake_out, fake_saved, fake_attrs = fake_impl(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" - - -def assert_bwd_fake_matches_real(args: Any, impl, fake_impl) -> None: - """Run a backward impl and its fake on the same args; require agreement. - - Gradients are not slot-counted: the backward payload holds exactly one slot - per gradient regardless of quantization, so only geometry is compared. - """ - real_seq, spec_seq = _as_sequence(impl(args)), _as_sequence(fake_impl(args)) - assert len(real_seq) == len(spec_seq), "gradient count differs" - for i, (value, spec) in enumerate(zip(real_seq, spec_seq)): - assert _geometry(value) == _spec_geometry(spec), f"gradient[{i}]: geometry differs" + assert _geometry(value) == _spec_geometry(spec), f"{what}[{i}]: geometry differs" # --------------------------------------------------------------------------- # -# Bias +# Tests # --------------------------------------------------------------------------- # -def _bias_op(size: int, dtype: torch.dtype) -> te.ops.Bias: - op = te.ops.Bias(size, device=_device, dtype=dtype) - with torch.no_grad(): - op.bias.copy_(torch.randn_like(op.bias)) - return op - - -def _bias_fwd_args(shape, size: int, dtype: torch.dtype) -> BiasFwdArgs: - op = _bias_op(size, dtype) - x = torch.randn(*shape, device=_device, dtype=dtype) - return op.resolve_fwd_args(x, requires_grad=True, prev_op_grad_output_quantizer=None) - - @_cuda -@pytest.mark.parametrize("shape", [(16, 32), (2, 8, 32), (32,)]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_bias_fake_conformance(shape, dtype) -> None: - args = _bias_fwd_args(shape, shape[-1], dtype) - assert_fwd_fake_matches_real(args, _bias_forward_impl, _bias_forward_impl_fake) +@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() + impls = case.impls(op) + x = _make_input(case) + args = _resolve(op, x, fp8_output=fp8_output) - dy = torch.randn(*shape, device=_device, dtype=dtype) - bwd_args = BiasBwdArgs(grad_output=dy, grad_input_quantizer=None) - assert_bwd_fake_matches_real(bwd_args, _bias_backward_impl, _bias_backward_impl_fake) + real_out, real_saved, real_attrs = impls.forward(args) + fake_out, fake_saved, fake_attrs = impls.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" + + dy = torch.randn(*real_out.shape, device=_device, dtype=_dtype) + bwd_args = case.bwd_args(dy, real_saved, real_attrs, x) + real_grads = _as_sequence(impls.backward(bwd_args)) + fake_grads = _as_sequence(impls.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("dtype", [torch.float32, torch.bfloat16]) -def test_bias_custom_op_matches_eager(dtype) -> None: +@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, forward and backward.""" - assert _bias_ops is not None, "bias custom ops failed to register" - forward_fn, backward_fn = _bias_ops + op = case.build() + ops = case.impls(op).ops + assert ops is not None, f"{case.name}: custom ops failed to register" + forward_fn, backward_fn = ops - shape, size = (16, 32), 32 - op = _bias_op(size, dtype) - x = torch.randn(*shape, device=_device, dtype=dtype, requires_grad=True) - dy = torch.randn(*shape, device=_device, dtype=dtype) - - # Reference: the op as used today. + x = _make_input(case, requires_grad=True) y_ref = op(x) + dy = torch.randn_like(y_ref) y_ref.backward(dy) - dx_ref, db_ref = x.grad.clone(), op.bias.grad.clone() + dx_ref = x.grad.clone() - # Same computation through the custom ops. - args = op.resolve_fwd_args(x.detach(), requires_grad=True, prev_op_grad_output_quantizer=None) + args = _resolve(op, x.detach(), fp8_output=False) y, saved, ctx_attrs = forward_fn(args) - assert saved == (), "bias saves no tensors" - dx, db = backward_fn( - BiasBwdArgs(grad_output=dy, grad_input_quantizer=ctx_attrs["grad_input_quantizer"]) - ) - # A None grad input means "grad_output unchanged" -- a custom op may not - # return one of its own inputs. - if dx is None: - dx = dy + grads = backward_fn(case.bwd_args(dy, saved, ctx_attrs, x.detach())) + # 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) - torch.testing.assert_close(db, db_ref) @_cuda -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_bias_custom_op_compiles_fullgraph(dtype) -> None: - """Both halves must trace without a graph break. - - Run under ``no_grad``: the halves carry no autograd of their own (that is - the point of ``register_op_halves``), so a caller is expected to wire them - into its own ``autograd.Function``. +@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_op_halves``), so a caller wires them into its own + ``autograd.Function`` -- see ``test_ops_hop_poc.py``. """ - assert _bias_ops is not None, "bias custom ops failed to register" - forward_fn, backward_fn = _bias_ops + op = case.build() + ops = case.impls(op).ops + assert ops is not None, f"{case.name}: custom ops failed to register" + forward_fn, backward_fn = ops - shape, size = (16, 32), 32 - op = _bias_op(size, dtype) - x = torch.randn(*shape, device=_device, dtype=dtype) - dy = torch.randn(*shape, device=_device, dtype=dtype) + x = _make_input(case) def fwd(x_): - args = op.resolve_fwd_args(x_, requires_grad=True, prev_op_grad_output_quantizer=None) - out, _saved, _attrs = forward_fn(args) - return out - - def bwd(dy_): - return backward_fn(BiasBwdArgs(grad_output=dy_, grad_input_quantizer=None)) + 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(): - torch.testing.assert_close(torch.compile(fwd, fullgraph=True)(x), fwd(x)) - compiled_grads = torch.compile(bwd, fullgraph=True)(dy) - expected_grads = bwd(dy) - for got, expected in zip(compiled_grads, expected_grads): - if got is None or expected is None: - assert got is expected is None - continue - torch.testing.assert_close(got, expected) - - -# --------------------------------------------------------------------------- # -# Activations -# --------------------------------------------------------------------------- # + 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) + _out, saved, attrs = forward_fn(_resolve(op, x, fp8_output=fp8_output)) + dy = torch.randn(*_as_sequence(_out)[0].shape, device=_device, dtype=_dtype) -def _fp8_quantizer(dtype: torch.dtype) -> Any: - """A current-scaling FP8 quantizer, as a next op would hand down.""" - from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer - from transformer_engine.pytorch.constants import DType + def bwd(dy_, input_): + return backward_fn(case.bwd_args(dy_, saved, attrs, input_)) - del dtype - quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, torch.device(_device)) - quantizer.set_usage(rowwise=True, columnwise=False) - return quantizer - - -_ACTIVATIONS = [te.ops.GELU, te.ops.ReLU, te.ops.SiLU, te.ops.GEGLU, te.ops.ReGLU] - - -@_cuda -@pytest.mark.parametrize("cls", _ACTIVATIONS) -@pytest.mark.parametrize("quantize_output", [False, True]) -def test_activation_fake_conformance(cls, quantize_output) -> None: - """The fake must predict output geometry, including whether it is quantized.""" - dtype = torch.bfloat16 - op = cls() - x = torch.randn(16, 64, device=_device, dtype=dtype) - args = op.resolve_fwd_args( - x, - requires_grad=True, - prev_op_grad_output_quantizer=None, - next_op_input_quantizer=_fp8_quantizer(dtype) if quantize_output else None, - ) - assert_fwd_fake_matches_real(args, op._impls.forward, op._impls.forward_fake) - - out, saved, ctx_attrs = op._impls.forward(args) - dy = torch.randn_like(out if not quantize_output else out.dequantize()) - bwd_args = ActivationBwdArgs( - grad_output=dy, - saved_input=saved[0] if saved else x, - dtype=ctx_attrs["dtype"], - grad_input_quantizer=None, - ) - assert_bwd_fake_matches_real(bwd_args, op._impls.backward, op._impls.backward_fake) + with torch.no_grad(): + expected_grads = bwd(dy, x) + got_grads = torch.compile(bwd, fullgraph=True)(dy, x) + 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("cls", _ACTIVATIONS) -def test_activation_custom_op_returns_fp8(cls) -> None: - """With a next-op quantizer the registered op must hand back an FP8 tensor. - - This is the case that 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 rather than being dequantized at it. +@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. """ - assert cls._impls.ops is not None, f"{cls.__name__} custom ops failed to register" - forward_fn, _ = cls._impls.ops - - dtype = torch.bfloat16 - op = cls() - x = torch.randn(16, 64, device=_device, dtype=dtype) - args = op.resolve_fwd_args( - x, - requires_grad=True, - prev_op_grad_output_quantizer=None, - next_op_input_quantizer=_fp8_quantizer(dtype), - ) - y, saved, _ctx_attrs = forward_fn(args) - + op = case.build() + forward_fn, _ = case.impls(op).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)}" - y_ref, _saved_ref, _ = op._impls.forward(args) - torch.testing.assert_close(y.dequantize(), y_ref.dequantize()) - assert saved == (), "without cache_quantized_input the op keeps nothing" - - -@_cuda -@pytest.mark.parametrize("cls", _ACTIVATIONS) -def test_activation_custom_op_matches_eager(cls) -> None: - dtype = torch.bfloat16 - op = cls() - forward_fn, backward_fn = cls._impls.ops - - x = torch.randn(16, 64, device=_device, dtype=dtype, requires_grad=True) - y_ref = op(x) - dy = torch.randn_like(y_ref) - y_ref.backward(dy) - dx_ref = x.grad.clone() - - args = op.resolve_fwd_args( - x.detach(), - requires_grad=True, - prev_op_grad_output_quantizer=None, - next_op_input_quantizer=None, - ) - y, saved, ctx_attrs = forward_fn(args) - # None means "the input, unchanged" -- a custom op may not return its own input. - saved_input = saved[0] if saved else x.detach() - (dx,) = backward_fn( - ActivationBwdArgs( - grad_output=dy, - saved_input=saved_input, - dtype=ctx_attrs["dtype"], - grad_input_quantizer=ctx_attrs["prev_op_grad_output_quantizer"], - ) - ) - - torch.testing.assert_close(y, y_ref) - torch.testing.assert_close(dx, dx_ref) - - -@_cuda -@pytest.mark.parametrize("cls", [te.ops.GELU, te.ops.GEGLU]) -def test_activation_custom_op_compiles_fullgraph(cls) -> None: - dtype = torch.bfloat16 - op = cls() - forward_fn, _ = cls._impls.ops - x = torch.randn(16, 64, device=_device, dtype=dtype) - - def fwd(x_): - args = op.resolve_fwd_args( - x_, - requires_grad=False, - prev_op_grad_output_quantizer=None, - next_op_input_quantizer=None, - ) - out, _saved, _attrs = forward_fn(args) - return out - - with torch.no_grad(): - torch.testing.assert_close(torch.compile(fwd, fullgraph=True)(x), fwd(x)) diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 5d607edf8c..6c499bb36e 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -250,8 +250,8 @@ def resolve_fwd_args( input_: torch.Tensor, *, requires_grad: bool, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], + 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. diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index e427267ccd..a4ca6df9eb 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -207,13 +207,19 @@ def resolve_fwd_args( input_: torch.Tensor, *, requires_grad: bool, - prev_op_grad_output_quantizer: Optional[Quantizer], + 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 From e84de33b24188074dc04a446f00a0f3deec5b0b9 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 21:11:04 +0200 Subject: [PATCH 07/11] [PyTorch] Move the compute halves into the operations themselves The impls and fakes lived outside the class they belong to -- module-level functions for Bias, closures from a factory for the activations -- and each operation then repeated the same op_forward/op_backward plumbing. They are now classmethods on the operation, and BasicOperation does the rest: __init_subclass__ registers the custom ops for any operation that declares fwd_args_type/bwd_args_type, and op_forward/op_backward are written once in the base. Declaring an operation means two dataclasses, four compute halves, resolve_fwd_args and resolve_bwd_args -- no registration or plumbing. Classmethods rather than static functions because the binding is load bearing: the whole activation family shares one implementation and dispatches to its per-class kernel through cls, so it needs no factory. Two hooks cover what genuinely differs: saved_for_backward, for an operation whose backward needs its input but whose forward produces no distinct tensor for it, and resolve_bwd_args, since each backward container names its own fields. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_ops_custom_ops.py | 147 ++++------- tests/pytorch/test_ops_hop_poc.py | 6 +- .../pytorch/ops/basic/activation.py | 247 ++++++------------ transformer_engine/pytorch/ops/basic/bias.py | 154 ++++------- transformer_engine/pytorch/ops/op.py | 136 +++++++++- 5 files changed, 328 insertions(+), 362 deletions(-) diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py index 5a3d210bce..562a1bd447 100644 --- a/tests/pytorch/test_ops_custom_ops.py +++ b/tests/pytorch/test_ops_custom_ops.py @@ -4,10 +4,11 @@ """Per-op custom ops for ``transformer_engine.pytorch.ops``. -Every fusible operation registers its forward and backward as independent custom -ops (``register_op_halves``) so a compiled pipeline can call them directly. The -checks here are the same for every operation and are driven by ``_OP_CASES``: -adding an operation means adding one entry to that list. +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: @@ -30,13 +31,14 @@ 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.basic import bias as _bias_mod -from transformer_engine.pytorch.ops.basic.activation import ActivationBwdArgs -from transformer_engine.pytorch.ops.basic.bias import BiasBwdArgs +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) @@ -53,11 +55,6 @@ def _fresh_dynamo(): torch._dynamo.reset() -_device = "cuda" -_dtype = torch.bfloat16 -_HIDDEN = 64 - - def _fp8_quantizer() -> Any: """An FP8 quantizer, standing in for the next operation's input quantizer.""" quantizer = Float8CurrentScalingQuantizer(DType.kFloat8E4M3, torch.device(_device)) @@ -70,40 +67,6 @@ def _fp8_quantizer() -> Any: # --------------------------------------------------------------------------- # -@dataclass -class OpCase: - """One operation and how to drive it. - - Everything except ``bwd_args`` is uniform; that one callable exists because - each operation's backward container names its own fields. - """ - - name: str - build: Callable[[], Any] - bwd_args: Callable[..., Any] - impls: Callable[[Any], Any] - quantizes_output: bool = True - num_grads: int = 1 - in_shape: Tuple[int, ...] = (16, _HIDDEN) - - -def _bias_bwd_args(grad_output, saved, ctx_attrs, input_): - del saved, input_ - return BiasBwdArgs( - grad_output=grad_output, - grad_input_quantizer=ctx_attrs["grad_input_quantizer"], - ) - - -def _activation_bwd_args(grad_output, saved, ctx_attrs, input_): - return ActivationBwdArgs( - grad_output=grad_output, - saved_input=saved[0] if saved else input_, - dtype=ctx_attrs["dtype"], - grad_input_quantizer=ctx_attrs["prev_op_grad_output_quantizer"], - ) - - def _build_bias(): op = te.ops.Bias(_HIDDEN, device=_device, dtype=_dtype) with torch.no_grad(): @@ -111,32 +74,21 @@ def _build_bias(): return op -class _BiasImpls: - """The Bias module's halves, in the same shape the activations expose.""" +@dataclass +class OpCase: + """One operation and how to build it.""" - forward = staticmethod(_bias_mod._bias_forward_impl) - forward_fake = staticmethod(_bias_mod._bias_forward_impl_fake) - backward = staticmethod(_bias_mod._bias_backward_impl) - backward_fake = staticmethod(_bias_mod._bias_backward_impl_fake) - ops = _bias_mod._bias_ops + 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, - bwd_args=_bias_bwd_args, - impls=lambda op: _BiasImpls, - quantizes_output=False, - num_grads=2, - ), + OpCase(name="Bias", build=_build_bias, quantizes_output=False, num_grads=2), *( - OpCase( - name=cls.__name__, - build=cls, - bwd_args=_activation_bwd_args, - impls=lambda op: type(op)._impls, - ) + OpCase(name=cls.__name__, build=cls) for cls in (te.ops.GELU, te.ops.ReLU, te.ops.SiLU, te.ops.GEGLU, te.ops.ReGLU) ), ] @@ -157,6 +109,20 @@ def _resolve(op, x, *, fp8_output: bool): ) +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 # --------------------------------------------------------------------------- # @@ -213,20 +179,21 @@ def assert_values_match_specs(real: Any, specs: Any, what: str) -> None: 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() - impls = case.impls(op) + cls = type(op) x = _make_input(case) args = _resolve(op, x, fp8_output=fp8_output) - real_out, real_saved, real_attrs = impls.forward(args) - fake_out, fake_saved, fake_attrs = impls.forward_fake(args) + 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 = case.bwd_args(dy, real_saved, real_attrs, x) - real_grads = _as_sequence(impls.backward(bwd_args)) - fake_grads = _as_sequence(impls.backward_fake(bwd_args)) + 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. @@ -236,11 +203,10 @@ def test_op_fake_matches_real(case: OpCase, fp8_output: bool) -> None: @_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, forward and backward.""" + """The registered op must reproduce the eager operation.""" op = case.build() - ops = case.impls(op).ops - assert ops is not None, f"{case.name}: custom ops failed to register" - forward_fn, backward_fn = ops + 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) @@ -248,9 +214,9 @@ def test_op_matches_eager(case: OpCase) -> None: y_ref.backward(dy) dx_ref = x.grad.clone() - args = _resolve(op, x.detach(), fp8_output=False) - y, saved, ctx_attrs = forward_fn(args) - grads = backward_fn(case.bwd_args(dy, saved, ctx_attrs, x.detach())) + 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 @@ -270,10 +236,8 @@ def test_op_compiles_fullgraph(case: OpCase, fp8_output: bool) -> None: ``autograd.Function`` -- see ``test_ops_hop_poc.py``. """ op = case.build() - ops = case.impls(op).ops - assert ops is not None, f"{case.name}: custom ops failed to register" - forward_fn, backward_fn = ops - + 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_): @@ -291,15 +255,16 @@ def fwd(x_): for a, b in zip(got, expected): torch.testing.assert_close(a, b) - _out, saved, attrs = forward_fn(_resolve(op, x, fp8_output=fp8_output)) - dy = torch.randn(*_as_sequence(_out)[0].shape, device=_device, dtype=_dtype) + 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(dy_, input_): - return backward_fn(case.bwd_args(dy_, saved, attrs, input_)) + def bwd(args): + return backward_fn(args) with torch.no_grad(): - expected_grads = bwd(dy, x) - got_grads = torch.compile(bwd, fullgraph=True)(dy, x) + 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 @@ -318,6 +283,6 @@ def test_op_returns_fp8(case: OpCase) -> None: than being dequantized at it. """ op = case.build() - forward_fn, _ = case.impls(op).ops + 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)}" diff --git a/tests/pytorch/test_ops_hop_poc.py b/tests/pytorch/test_ops_hop_poc.py index 98dd3a1020..a917f4e32f 100644 --- a/tests/pytorch/test_ops_hop_poc.py +++ b/tests/pytorch/test_ops_hop_poc.py @@ -87,7 +87,7 @@ def forward(func_ctx, x, bias1, bias2, ops, quantize_middle): next_op_input_quantizer=_fp8_quantizer() if quantize_middle else None, ) y1, saved1, attrs1 = act_fwd(args1) - ctxs[1].saved = saved1 or (y0,) + 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) @@ -149,8 +149,8 @@ def _build(dtype: torch.dtype): bias_op1.bias.copy_(torch.randn_like(bias_op1.bias)) bias_op2.bias.copy_(torch.randn_like(bias_op2.bias)) ops = { - "bias": te.ops.basic.bias._bias_ops, - "act": type(act_op)._impls.ops, + "bias": bias_op1.compile_ops, + "act": act_op.compile_ops, "bias_op1": bias_op1, "bias_op2": bias_op2, "act_op": act_op, diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 6c499bb36e..d4baa48d93 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -15,7 +15,7 @@ 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, register_op_halves +from ...dynamo import TensorSpec from ...quantized_tensor import QuantizedTensorStorage from ...tensor.float8_tensor import Float8CurrentScalingQuantizer, Quantizer from ...utils import clear_tensor_data @@ -70,112 +70,6 @@ def _activation_output_shape( return (*input_shape[:-1], input_shape[-1] // 2) -@dataclass(slots=True) -class _ActivationImpls: - """One activation class's compute halves, plus their registered custom ops. - - Held in a container rather than as class attributes so the functions stay - plain functions instead of being bound as methods on attribute lookup. - """ - - forward: Callable[[ActivationFwdArgs], Any] - forward_fake: Callable[[ActivationFwdArgs], Any] - backward: Callable[[ActivationBwdArgs], Any] - backward_fake: Callable[[ActivationBwdArgs], Any] - ops: Optional[Tuple[Callable[..., Any], Callable[..., Any]]] - - -def _make_activation_ops( - op_name: str, - forward_kernel: Callable[..., torch.Tensor], - backward_kernel: Callable[..., torch.Tensor], - halves_last_dim: bool, -) -> _ActivationImpls: - """Build and register the compute halves for one activation class. - - The kernel pair is baked in here rather than passed as an argument: it is - fixed by the class, and a callable has no place in an op schema. - """ - - def forward_impl(args: ActivationFwdArgs): - x = maybe_dequantize(args.input_.contiguous(), args.dtype) - y = forward_kernel(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 of which are 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 a static one: - # the caller passes the input to the backward, which dequantizes it there. - saved = (x,) if (args.requires_grad and args.cache_quantized_input) else () - ctx_attrs = { - "dtype": args.dtype, - "prev_op_grad_output_quantizer": args.prev_op_grad_output_quantizer, - } - return y, saved, ctx_attrs - - def forward_fake_impl(args: ActivationFwdArgs): - x = args.input_ - shape = tuple(x.shape) - y = TensorSpec( - shape=_activation_output_shape(shape, 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 - ), - ) - ctx_attrs = { - "dtype": args.dtype, - "prev_op_grad_output_quantizer": args.prev_op_grad_output_quantizer, - } - return y, saved, ctx_attrs - - def backward_impl(args: ActivationBwdArgs): - x = maybe_dequantize(args.saved_input.contiguous(), args.dtype) - dy = maybe_dequantize(args.grad_output.contiguous(), x.dtype) - dx = backward_kernel(dy, x, args.grad_input_quantizer) - return (dx,) - - def backward_fake_impl(args: ActivationBwdArgs): - shape = tuple(args.saved_input.shape) - return ( - TensorSpec( - shape=shape, - dtype=args.dtype, - quantizer=args.grad_input_quantizer, - device=args.saved_input.device, - ), - ) - - return _ActivationImpls( - forward=forward_impl, - forward_fake=forward_fake_impl, - backward=backward_impl, - backward_fake=backward_fake_impl, - ops=register_op_halves( - op_name=op_name, - fwd_arg_type=ActivationFwdArgs, - fwd_impl=forward_impl, - fwd_fake_impl=forward_fake_impl, - bwd_arg_type=ActivationBwdArgs, - bwd_impl=backward_impl, - bwd_fake_impl=backward_fake_impl, - num_grad_inputs=1, - ), - ) - - class _ActivationOperation(BasicOperation, metaclass=abc.ABCMeta): r"""Apply activation function @@ -229,20 +123,87 @@ def _activation_backward_impl(*args, **kwargs) -> torch.Tensor: # GLU variants consume pairs along the inner dimension; set per subclass. _output_halves_last_dim: bool = False - _impls: Optional[_ActivationImpls] = None - - def __init_subclass__(cls, **kwargs) -> None: - super().__init_subclass__(**kwargs) - # The kernel pair is fixed by the class, so each subclass gets its own - # registration; the op registry stays bounded by the op zoo, not by the - # model. - if getattr(cls._activation_forward_impl, "__isabstractmethod__", False): - return - cls._impls = _make_activation_ops( - op_name=f"activation_{cls.__name__.lower()}", - forward_kernel=cls._activation_forward_impl, - backward_kernel=cls._activation_backward_impl, - halves_last_dim=cls._output_halves_last_dim, + + fwd_args_type = ActivationFwdArgs + bwd_args_type = ActivationBwdArgs + num_grad_inputs = 1 + + @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, + ), + ) + + 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_,) + + 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, ) def resolve_fwd_args( @@ -274,50 +235,6 @@ def resolve_fwd_args( prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, ) - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - 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, - ) - y, saved, ctx_attrs = self._impls.forward(args) - if ctx.requires_grad: - # Without ``cache_quantized_input`` the op keeps nothing: the backward - # rebuilds its input from the operation's input tensor. - saved = saved or (input_,) - if is_cpu_offload_enabled(): - mark_activation_offload(*saved) - ctx.save_for_backward(*saved) - for name, value in ctx_attrs.items(): - setattr(ctx, name, value) - return y - - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - (x,) = ctx.saved_tensors - (dx,) = self._impls.backward( - ActivationBwdArgs( - grad_output=grad_output, - saved_input=x, - dtype=ctx.dtype, - grad_input_quantizer=ctx.prev_op_grad_output_quantizer, - ) - ) - # Eager only: the compiled path leaves saved tensors to the graph's own - # memory planning, and clearing an op input would lie to functionalization. - clear_tensor_data(x) - return dx, () - class GELU(_ActivationOperation): r"""Gaussian Error Linear Unit diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index a4ca6df9eb..9be54a9e1d 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -16,7 +16,7 @@ from ...utils import canonicalize_device, canonicalize_dtype from ...tensor import Quantizer from ...quantized_tensor import QuantizedTensorStorage -from ...dynamo import TensorSpec, register_op_halves +from ...dynamo import TensorSpec TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] @@ -39,72 +39,6 @@ class BiasBwdArgs: grad_input_quantizer: Optional[Quantizer] = None -def _bias_forward_impl( - args: BiasFwdArgs, -) -> Tuple[torch.Tensor, Tuple[()], Dict[str, Any]]: - """Bias forward. Saves no tensors; 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} - - -def _bias_forward_impl_fake( - args: BiasFwdArgs, -) -> Tuple[TensorSpec, Tuple[()], Dict[str, Any]]: - """Allocation-free fake of :func:`_bias_forward_impl`.""" - x = args.input_ - out = TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device) - return out, (), {"grad_input_quantizer": args.grad_input_quantizer} - - -def _bias_backward_impl( - args: BiasBwdArgs, -) -> Tuple[Optional[torch.Tensor], torch.Tensor]: - """Bias backward: reduce the grad over all but the inner dimension. - - 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; the caller substitutes - ``grad_output`` instead. - """ - 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 - - -def _bias_backward_impl_fake( - args: BiasBwdArgs, -) -> Tuple[Optional[TensorSpec], TensorSpec]: - """Allocation-free fake of :func:`_bias_backward_impl`.""" - 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) - - -_bias_ops = register_op_halves( - op_name="bias", - fwd_arg_type=BiasFwdArgs, - fwd_impl=_bias_forward_impl, - fwd_fake_impl=_bias_forward_impl_fake, - bwd_arg_type=BiasBwdArgs, - bwd_impl=_bias_backward_impl, - bwd_fake_impl=_bias_backward_impl_fake, - num_grad_inputs=2, -) - - class Bias(BasicOperation): """Apply additive bias @@ -126,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, @@ -202,6 +140,51 @@ def pre_first_fuser_forward(self) -> None: if self.bias.device.type == "meta": self.reset_parameters() + @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, input_: torch.Tensor, @@ -233,37 +216,8 @@ def resolve_fwd_args( grad_input_quantizer=grad_input_quantizer, ) - def op_forward( - self, - ctx: OperationContext, - input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - ) -> torch.Tensor: - del next_op_input_quantizer # Bias never quantizes its output - args = self.resolve_fwd_args( - input_, - requires_grad=ctx.requires_grad, - prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, - ) - out, saved, ctx_attrs = _bias_forward_impl(args) - if ctx.requires_grad: - ctx.save_for_backward(*saved) - for name, value in ctx_attrs.items(): - setattr(ctx, name, value) - return out - - def op_backward( - self, - ctx: OperationContext, - grad_output: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[()]]: - grad_input, grad_bias = _bias_backward_impl( - BiasBwdArgs( - grad_output=grad_output, - grad_input_quantizer=ctx.grad_input_quantizer, - ) + 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, ) - if grad_input is None: - grad_input = grad_output - return grad_input, (grad_bias,) diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 5106ec9e0a..1973c52975 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 register_op_halves @dataclasses.dataclass @@ -184,6 +185,37 @@ 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) from register_op_halves, or None if unsupported. + 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 + # 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_op_halves( + 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 +223,71 @@ 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 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 +522,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 +533,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 +554,25 @@ 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 - @abc.abstractmethod def op_backward( self, ctx: OperationContext, @@ -463,6 +580,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 +597,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, From 6379e8082c002335cbb8f451ae323971ac325ac0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 21:19:34 +0200 Subject: [PATCH 08/11] [PyTorch] Let an operation say why it cannot be compiled Linear puts compile_unsupported_reason on its forward args, because the module is the compile boundary and the check sits next to the choice between the custom op and eager. In ops/ the boundary is the fuser group, not the operation, so the check belongs elsewhere: - a pipeline compiles as a whole, so the decision is one per group, aggregated over its operations; - the args only exist after resolve_fwd_args, while most reasons are known earlier and more cheaply from the operation itself; - recipe-level limits belong to whoever reads the recipe, which is the fuser. So the hook goes on BasicOperation. The default refuses an operation without the compute halves, and any quantizer torch.compile cannot specialize on -- delayed scaling holds live scale/amax tensors, so baking its quantizer into the graph would silently freeze stale scales. __init_subclass__ now also checks that the argument containers are dataclasses, which is what the framework actually requires of them: the op schema is built from their fields. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_ops_custom_ops.py | 41 ++++++++++++++++++++++++++++ transformer_engine/pytorch/ops/op.py | 32 +++++++++++++++++++++- 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py index 562a1bd447..0f3a07b7e3 100644 --- a/tests/pytorch/test_ops_custom_ops.py +++ b/tests/pytorch/test_ops_custom_ops.py @@ -286,3 +286,44 @@ def test_op_returns_fp8(case: OpCase) -> None: 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/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 1973c52975..b80999139d 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -22,7 +22,7 @@ autocast, ) from ..tensor import Quantizer -from ..dynamo import register_op_halves +from ..dynamo import is_value_opaque_quantizer, register_op_halves @dataclasses.dataclass @@ -202,6 +202,14 @@ def __init_subclass__(cls, **kwargs) -> 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. @@ -258,6 +266,28 @@ 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, From c90665a6df236118e5ad9fa6a88d0921b88a6f9b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 21:31:26 +0200 Subject: [PATCH 09/11] [PyTorch] Share the two-tier registration between both entry points register_op_halves had grown as a copy of register_custom_op's body minus the autograd wiring -- about two thirds of it was the same code. Both now build on _register_two_tier_pair, which owns everything they genuinely share: schemas from the argument containers, the base kernels, the wrapper ops that flatten QuantizedTensor subclass inputs, and the passthrough registrations. What is left in each entry point is only what they differ on. Renamed to register_custom_op_without_autograd. 'Halves' said nothing about why the function exists; the one thing a reader needs is that, unlike register_custom_op, it leaves autograd to the caller. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_ops_custom_ops.py | 2 +- transformer_engine/pytorch/dynamo/__init__.py | 4 +- .../pytorch/dynamo/custom_op.py | 411 +++++++++--------- transformer_engine/pytorch/ops/op.py | 6 +- 4 files changed, 204 insertions(+), 219 deletions(-) diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py index 0f3a07b7e3..ef0b871abc 100644 --- a/tests/pytorch/test_ops_custom_ops.py +++ b/tests/pytorch/test_ops_custom_ops.py @@ -232,7 +232,7 @@ 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_op_halves``), so a caller wires them into its own + point of ``register_custom_op_without_autograd``), so a caller wires them into ``autograd.Function`` -- see ``test_ops_hop_poc.py``. """ op = case.build() diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index a580ea95e8..85b444a914 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, register_op_halves +from .custom_op import register_custom_op, register_custom_op_without_autograd __all__ = [ "register_value_opaque_quantizer", @@ -14,5 +14,5 @@ "TensorSpec", "to_tensor_spec", "register_custom_op", - "register_op_halves", + "register_custom_op_without_autograd", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index eb9754bc36..c37338e712 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1035,6 +1035,159 @@ def _split_fwd_fake_result( return user_fakes, saved_fakes, ctx_attrs +# --------------------------------------------------------------------------- # +# 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` and + :func:`register_custom_op_without_autograd`: 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 # --------------------------------------------------------------------------- # @@ -1393,124 +1546,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) @@ -1523,7 +1594,7 @@ def forward_fn(fwd_args): # --------------------------------------------------------------------------- # -def register_op_halves( +def register_custom_op_without_autograd( *, op_name: str, fwd_arg_type: type, @@ -1562,7 +1633,7 @@ def register_op_halves( back to eager rather than breaking import. """ try: - return _register_op_halves_impl( + return _register_custom_op_without_autograd_impl( op_name=op_name, fwd_arg_type=fwd_arg_type, fwd_impl=fwd_impl, @@ -1574,12 +1645,12 @@ def register_op_halves( ) except (ImportError, AttributeError, RuntimeError, TypeError) as e: warn_compile_unsupported( - f"could not register the custom op halves '{op_name}' ({type(e).__name__}: {e})" + f"could not register the autograd-free custom ops '{op_name}' ({type(e).__name__}: {e})" ) return None -def _register_op_halves_impl( +def _register_custom_op_without_autograd_impl( *, op_name: str, fwd_arg_type: type, @@ -1590,120 +1661,34 @@ def _register_op_halves_impl( bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], num_grad_inputs: int, ) -> Tuple[Callable[[Any], Any], Callable[[Any], Any]]: - """Body of :func:`register_op_halves`; see it for semantics.""" - 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}" - - _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, + """Body of :func:`register_custom_op_without_autograd`; 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, ) - 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) - - def _reassemble(specs: List[Any], payload: List[torch.Tensor], cursor: int) -> Tuple[List, int]: - """Rebuild 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 def forward_fn(fwd_args): - spec_obj = _spec_view(fwd_args, fwd_tensor_field_names) + 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, fwd_adapters) - payload = wrapper_fwd_op(*[kwargs[name] for name in fwd_arg_names]) + 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, 0) + 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, - ) + 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, bwd_adapters) - payload = wrapper_bwd_op(*[kwargs[name] for name in bwd_arg_names]) + 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 diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index b80999139d..ee81b637fe 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -22,7 +22,7 @@ autocast, ) from ..tensor import Quantizer -from ..dynamo import is_value_opaque_quantizer, register_op_halves +from ..dynamo import is_value_opaque_quantizer, register_custom_op_without_autograd @dataclasses.dataclass @@ -193,7 +193,7 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta): 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) from register_op_halves, or None if unsupported. + # (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: @@ -213,7 +213,7 @@ def __init_subclass__(cls, **kwargs) -> None: # 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_op_halves( + cls.compile_ops = register_custom_op_without_autograd( op_name=cls.__name__.lower(), fwd_arg_type=cls.fwd_args_type, fwd_impl=cls.forward_compute, From 12427bb9c7e2d3eecb76de8a1c79dcf4791d994e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 21:41:45 +0200 Subject: [PATCH 10/11] [PyTorch] Name the op pair as the primitive, autograd as the addition register_custom_op_without_autograd named a thing by what it lacked, which had it backwards: after the two entry points were factored onto a shared registration, the forward/backward pair is the primitive and autograd is what the other one adds on top. So the pair takes the plain name and the wired variant becomes register_custom_op_with_autograd. The file follows the same order, primitive first, and the pair's docstring no longer defines itself by negation. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_ops_custom_ops.py | 2 +- transformer_engine/pytorch/dynamo/__init__.py | 4 +- .../pytorch/dynamo/custom_op.py | 230 +++++++++--------- transformer_engine/pytorch/module/linear.py | 4 +- transformer_engine/pytorch/ops/op.py | 4 +- 5 files changed, 123 insertions(+), 121 deletions(-) diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py index ef0b871abc..2ad8464c71 100644 --- a/tests/pytorch/test_ops_custom_ops.py +++ b/tests/pytorch/test_ops_custom_ops.py @@ -232,7 +232,7 @@ 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_without_autograd``), so a caller wires them into + point of ``register_custom_op``), so a caller wires them into ``autograd.Function`` -- see ``test_ops_hop_poc.py``. """ op = case.build() diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 85b444a914..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, register_custom_op_without_autograd +from .custom_op import register_custom_op, register_custom_op_with_autograd __all__ = [ "register_value_opaque_quantizer", @@ -14,5 +14,5 @@ "TensorSpec", "to_tensor_spec", "register_custom_op", - "register_custom_op_without_autograd", + "register_custom_op_with_autograd", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index c37338e712..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: @@ -1071,8 +1071,8 @@ def _register_two_tier_pair( ) -> _OpPair: """Define an operation's forward and backward as two-tier custom ops. - Everything that is common to :func:`register_custom_op` and - :func:`register_custom_op_without_autograd`: schemas from the argument + 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 @@ -1189,7 +1189,114 @@ def _reassemble(specs: List[Any], payload: List[torch.Tensor], cursor: int = 0): # --------------------------------------------------------------------------- # -# Op registration +# 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 # --------------------------------------------------------------------------- # @@ -1433,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], @@ -1506,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, @@ -1524,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], @@ -1536,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 @@ -1587,108 +1694,3 @@ def forward_fn(fwd_args): return tuple(outputs) return forward_fn - - -# --------------------------------------------------------------------------- # -# Split registration: forward and backward as independent ops (no autograd) -# --------------------------------------------------------------------------- # - - -def register_custom_op_without_autograd( - *, - 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. - - Unlike :func:`register_custom_op`, no autograd is registered. The caller - wires forward to backward itself -- e.g. a pipeline-level - ``torch.autograd.Function`` that Dynamo traces as a higher-order op, so the - forward and backward passes can be grouped differently (which is what - ``ops.OperationFuser`` does). Both halves are still two-tier, so - ``QuantizedTensor`` subclass inputs pass through without dequantization. - - Contracts, mirroring :func:`register_custom_op`: - - * ``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_without_autograd_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_without_autograd_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_without_autograd`; 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 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/op.py b/transformer_engine/pytorch/ops/op.py index ee81b637fe..d2b7febc4e 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -22,7 +22,7 @@ autocast, ) from ..tensor import Quantizer -from ..dynamo import is_value_opaque_quantizer, register_custom_op_without_autograd +from ..dynamo import is_value_opaque_quantizer, register_custom_op @dataclasses.dataclass @@ -213,7 +213,7 @@ def __init_subclass__(cls, **kwargs) -> None: # 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_without_autograd( + cls.compile_ops = register_custom_op( op_name=cls.__name__.lower(), fwd_arg_type=cls.fwd_args_type, fwd_impl=cls.forward_compute, From 96c35f1f328aaaad914f50db2fc27d4c427b44d5 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 22:09:00 +0200 Subject: [PATCH 11/11] [PyTorch] Compile an OperationFuser group holding one operation Wires the fuser's compiled path end to end, so a group whose operations declare their compute halves runs through their custom ops under torch.compile(fullgraph=True). The pipeline-level autograd.Function is traced as a higher-order op, which is what lets its forward and backward walk different op groupings later on. Four things blocked tracing, all of them side effects reaching outside the higher-order op's scope: - OperationContext objects are created in the forward, but the backward is a separate subgraph, so writing to them there mutates an enclosing scope; the backward copies them into its own scope instead; - requires_grad_ on an output, which AOTAutograd's functionalization drops anyway -- autograd marks the outputs of an apply() itself; - _do_not_clear on inputs and outputs; - warnings.warn from the gate, which is not traceable, so the reason is now reported from the eager path only. These are gated on being traced rather than on using the custom ops. Under fullgraph there is no leaving the graph, so an unsupported operation does not 'fall back' -- the pipeline is traced either way and only the choice of implementation changes, which means the tracing constraints hold on both paths. Sequential builds its module groups outside the forward pass, since that constructs nn.Modules. Tested with a test-only operation, so the fuser's path does not depend on which real operations happen to be converted. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 140 +++++++++++++++++++ transformer_engine/pytorch/ops/fuser.py | 125 +++++++++++++---- transformer_engine/pytorch/ops/op.py | 38 +++++ transformer_engine/pytorch/ops/sequential.py | 17 ++- 4 files changed, 292 insertions(+), 28 deletions(-) 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/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 d2b7febc4e..2cf28e1e39 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -603,6 +603,44 @@ def op_forward( 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:]) + def op_backward( self, ctx: OperationContext, 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]]],