From 96404d02eacfc7760d97c139de9a90716cec26f9 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 22:21:38 +0200 Subject: [PATCH 1/4] [PyTorch] Register an op's forward and backward without autograd glue register_custom_op now defines an operation's forward and backward as two independent two-tier custom ops and hands back both, leaving autograd to the caller. That is what lets a pipeline-level autograd.Function decide how the two are wired, and so group the forward and backward passes differently -- which is what ops.OperationFuser does. The variant that wires autograd itself keeps the old behaviour under register_custom_op_with_autograd, and is now built on the same registration: the pair is the primitive, autograd is what the other one adds. About two thirds of the two bodies were the same code before. BasicOperation gains the plumbing an operation needs to opt in: declare two argument containers and implement four compute classmethods, and __init_subclass__ registers the custom ops while op_forward / op_backward are written once in the base. compile_unsupported_reason lets an operation say why it cannot be compiled -- it sits here rather than on the args, as Linear has it, because in ops/ the compile boundary is the fuser group, not the operation. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/__init__.py | 3 +- .../pytorch/dynamo/custom_op.py | 409 +++++++++++++----- transformer_engine/pytorch/module/linear.py | 4 +- transformer_engine/pytorch/ops/op.py | 204 ++++++++- 4 files changed, 500 insertions(+), 120 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 3598e54daa..1382cc1c92 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,7 +6,7 @@ from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer from .tensor_spec import TensorSpec, to_tensor_spec -from .custom_op import register_custom_op +from .custom_op import register_custom_op, register_custom_op_with_autograd __all__ = [ "register_value_opaque_quantizer", @@ -14,4 +14,5 @@ "TensorSpec", "to_tensor_spec", "register_custom_op", + "register_custom_op_with_autograd", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index b529deb75a..5fb1c4612e 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -6,7 +6,7 @@ Turns a TE module's eager forward/backward into ``torch.library`` custom ops so ``torch.compile(fullgraph=True)`` traces them as single graph nodes -- no graph -break into the eager ``autograd.Function``. ``register_custom_op`` is the entry +break into the eager ``autograd.Function``. ``register_custom_op_with_autograd`` is the entry point (its docstring documents the per-callable contract); ``module/linear.py`` is the first user. Internal framework API -- exported from ``transformer_engine.pytorch.dynamo``, not re-exported at the top level. @@ -50,7 +50,7 @@ only when its value is trivial (``None`` / all-``None``) at call time. What runs where. Each op registers a data-free fake (``register_fake``) so it -traces under ``torch.compile`` without allocating. ``register_custom_op`` returns +traces under ``torch.compile`` without allocating. ``register_custom_op_with_autograd`` returns ``forward_fn`` -- the drop-in for the eager ``autograd.Function.apply``. A forward call through it: @@ -1036,7 +1036,267 @@ def _split_fwd_fake_result( # --------------------------------------------------------------------------- # -# Op registration +# Two-tier op pair: the registration both public entry points build on +# --------------------------------------------------------------------------- # + + +@dataclasses.dataclass +class _OpPair: + """One registered forward/backward pair, and what a caller needs to drive it.""" + + fwd_adapters: List[_Adapter] + bwd_adapters: List[_Adapter] + fwd_arg_names: List[str] + bwd_arg_names: List[str] + fwd_tensor_field_names: List[str] + bwd_tensor_field_names: List[str] + base_fwd_def: Any + base_fwd_op: Any + base_bwd_op: Any + wrapper_fwd_def: Any + wrapper_fwd_op: Any + wrapper_bwd_op: Any + + +def _register_two_tier_pair( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: int, +) -> _OpPair: + """Define an operation's forward and backward as two-tier custom ops. + + Everything that is common to :func:`register_custom_op_with_autograd` and + :func:`register_custom_op`: schemas from the argument + containers, the base kernels, the wrapper ops that flatten + ``QuantizedTensor`` subclass inputs, and the passthrough registrations. + Autograd is deliberately not touched here -- that is what the two entry + points differ on. + """ + subclass_list = _all_quantized_tensor_subclasses() + + fwd_adapters = _get_adapters(fwd_arg_type) + bwd_adapters = _get_adapters(bwd_arg_type) + fwd_tensor_field_names = _tensor_field_names(fwd_adapters) + bwd_tensor_field_names = _tensor_field_names(bwd_adapters) + + fwd_schema_args, fwd_arg_names = _build_schema(fwd_adapters) + bwd_schema_args, bwd_arg_names = _build_schema(bwd_adapters) + fwd_schema = f"{fwd_schema_args} -> Tensor[]" + bwd_schema = f"{bwd_schema_args} -> Tensor[]" + + wrapper_fwd_name = op_name + wrapper_bwd_name = f"{op_name}_backward" + base_fwd_name = f"{op_name}_base" + base_bwd_name = f"{wrapper_bwd_name}_base" + base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" + + base_fwd_def = _register_kernel( + op_name=base_fwd_name, + schema_str=fwd_schema, + arg_type=fwd_arg_type, + arg_names=fwd_arg_names, + adapters=fwd_adapters, + tensor_field_names=fwd_tensor_field_names, + impl=fwd_impl, + fake_impl=fwd_fake_impl, + format_result=_format_fwd_result, + ) + _register_kernel( + op_name=base_bwd_name, + schema_str=bwd_schema, + arg_type=bwd_arg_type, + arg_names=bwd_arg_names, + adapters=bwd_adapters, + tensor_field_names=bwd_tensor_field_names, + impl=bwd_impl, + fake_impl=bwd_fake_impl, + format_result=lambda g: _format_bwd_result(g, num_grad_inputs, base_bwd_qualname), + ) + + base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) + base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) + + wrapper_fwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_fwd_name, + schema_str=fwd_schema, + base_op=base_fwd_op, + adapters=fwd_adapters, + ) + wrapper_bwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_bwd_name, + schema_str=bwd_schema, + base_op=base_bwd_op, + adapters=bwd_adapters, + ) + wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) + wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) + + fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_adapters) + bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_adapters) + + def _fwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) + return base_fwd_op(*new_args) + + def _bwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) + return base_bwd_op(*new_args) + + for sub in subclass_list: + wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) + wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule) + + for op in (wrapper_fwd_op, wrapper_bwd_op, base_fwd_op, base_bwd_op): + _quantized_tensor_passthrough_ops.add(op.default) + + return _OpPair( + fwd_adapters=fwd_adapters, + bwd_adapters=bwd_adapters, + fwd_arg_names=fwd_arg_names, + bwd_arg_names=bwd_arg_names, + fwd_tensor_field_names=fwd_tensor_field_names, + bwd_tensor_field_names=bwd_tensor_field_names, + base_fwd_def=base_fwd_def, + base_fwd_op=base_fwd_op, + base_bwd_op=base_bwd_op, + wrapper_fwd_def=wrapper_fwd_def, + wrapper_fwd_op=wrapper_fwd_op, + wrapper_bwd_op=wrapper_bwd_op, + ) + + +def _reassemble(specs: List[Any], payload: List[torch.Tensor], cursor: int = 0): + """Rebuild the values described by ``specs`` from ``payload[cursor:]``.""" + out: List[Any] = [] + for spec in specs: + n = _spec_slot_count(spec) + chunk = [_decode_none(t) for t in payload[cursor : cursor + n]] + cursor += n + out.append(_spec_reassemble(spec, chunk)) + return out, cursor + + +# --------------------------------------------------------------------------- # +# Op registration: the forward/backward pair, and the autograd-wired variant +# --------------------------------------------------------------------------- # + + +def register_custom_op( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: int, +) -> Optional[Tuple[Callable[[Any], Any], Callable[[Any], Any]]]: + """Register an op's forward and backward as two independent custom ops. + + Autograd is the caller's: it decides how the two are wired, which is what + lets a pipeline-level ``torch.autograd.Function`` -- traced by Dynamo as a + higher-order op -- group the forward and backward passes differently, as + ``ops.OperationFuser`` does. :func:`register_custom_op_with_autograd` builds + on this and wires them the usual way instead. + + Both ops are two-tier, so ``QuantizedTensor`` subclass inputs pass through + without dequantization. + + Contracts, mirroring :func:`register_custom_op_with_autograd`: + + * ``fwd_impl(fwd_args) -> (*user_outputs, tensors_to_save, ctx_attrs)`` + * ``fwd_fake_impl`` -- its data-free twin over :class:`TensorSpec` + * ``bwd_impl(bwd_args) -> tuple`` of ``num_grad_inputs`` gradients + * ``bwd_fake_impl`` -- its data-free twin + + Returns ``(forward_fn, backward_fn)``: + + * ``forward_fn(fwd_args) -> (outputs, saved_tensors, ctx_attrs)`` -- + ``outputs`` is a single value or a tuple, mirroring ``fwd_impl``'s user + outputs; ``saved_tensors`` is the reassembled ``tensors_to_save`` tuple, + which the caller is expected to persist (e.g. ``ctx.save_for_backward``). + * ``backward_fn(bwd_args) -> tuple`` of gradients. + + Returns ``None`` if registration fails (warned once), so callers can fall + back to eager rather than breaking import. + """ + try: + return _register_custom_op_impl( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_arg_type=bwd_arg_type, + bwd_impl=bwd_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=num_grad_inputs, + ) + except (ImportError, AttributeError, RuntimeError, TypeError) as e: + warn_compile_unsupported( + f"could not register the autograd-free custom ops '{op_name}' ({type(e).__name__}: {e})" + ) + return None + + +def _register_custom_op_impl( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: int, +) -> Tuple[Callable[[Any], Any], Callable[[Any], Any]]: + """Body of :func:`register_custom_op`; see it for semantics.""" + pair = _register_two_tier_pair( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_arg_type=bwd_arg_type, + bwd_impl=bwd_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=num_grad_inputs, + ) + + def forward_fn(fwd_args): + spec_obj = _spec_view(fwd_args, pair.fwd_tensor_field_names) + user_specs, saved_specs, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) + kwargs = _args_to_slots(fwd_args, pair.fwd_adapters) + payload = pair.wrapper_fwd_op(*[kwargs[name] for name in pair.fwd_arg_names]) + + outputs, cursor = _reassemble(user_specs, payload) + saved, _ = _reassemble(saved_specs, payload, cursor) + return (outputs[0] if len(outputs) == 1 else tuple(outputs)), tuple(saved), ctx_attrs + + def backward_fn(bwd_args): + # Unlike the forward payload, each grad occupies exactly one slot + # (``_format_bwd_result`` materializes a TensorSpec grad), so there is + # nothing to reassemble. + kwargs = _args_to_slots(bwd_args, pair.bwd_adapters) + payload = pair.wrapper_bwd_op(*[kwargs[name] for name in pair.bwd_arg_names]) + return tuple(_decode_none(t) for t in payload) + + return forward_fn, backward_fn + + +# --------------------------------------------------------------------------- # +# Autograd-wired registration # --------------------------------------------------------------------------- # @@ -1257,7 +1517,10 @@ def _forward(*flat: Any) -> List[torch.Tensor]: return base_op(*new_args) op = torch.library.custom_op( - f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", _forward, mutates_args=(), schema=schema_str + f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", + _forward, + mutates_args=(), + schema=schema_str, ) op.register_fake(_forward) return op @@ -1277,7 +1540,7 @@ def _all_quantized_tensor_subclasses() -> List[type]: return found -def register_custom_op( +def register_custom_op_with_autograd( *, op_name: str, input_tensors_for_grad: List[str], @@ -1350,7 +1613,7 @@ def register_custom_op( ``torch.compile`` (a graph break) rather than breaking import. """ try: - return _register_custom_op_impl( + return _register_custom_op_with_autograd_impl( op_name=op_name, input_tensors_for_grad=input_tensors_for_grad, fwd_arg_type=fwd_arg_type, @@ -1368,7 +1631,7 @@ def register_custom_op( return None -def _register_custom_op_impl( +def _register_custom_op_with_autograd_impl( *, op_name: str, input_tensors_for_grad: List[str], @@ -1380,7 +1643,7 @@ def _register_custom_op_impl( fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], ) -> Callable[..., Any]: - """Body of :func:`register_custom_op`; see it for semantics.""" + """Body of :func:`register_custom_op_with_autograd`; see it for semantics.""" # Existence check at the API boundary: every ``input_tensors_for_grad`` name # must be an actual field of ``fwd_arg_type`` (differentiability -- whether # that field can carry a gradient -- is checked later in @@ -1390,124 +1653,42 @@ def _register_custom_op_impl( if missing: raise ValueError(f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}") - wrapper_fwd_name = op_name - wrapper_bwd_name = f"{op_name}_backward" - base_fwd_name = f"{op_name}_base" - base_bwd_name = f"{wrapper_bwd_name}_base" - subclass_list = _all_quantized_tensor_subclasses() - - fwd_adapters = _get_adapters(fwd_arg_type) - bwd_adapters = _get_adapters(backward_arg_type) - fwd_tensor_field_names = _tensor_field_names(fwd_adapters) - bwd_tensor_field_names = _tensor_field_names(bwd_adapters) - - fwd_schema_args, fwd_arg_names = _build_schema(fwd_adapters) - bwd_schema_args, bwd_arg_names = _build_schema(bwd_adapters) - - num_grad_inputs = len(input_tensors_for_grad) - slot_count, grad_targets = _resolve_grad_targets(fwd_adapters, input_tensors_for_grad) - - fwd_schema = f"{fwd_schema_args} -> Tensor[]" - bwd_schema = f"{bwd_schema_args} -> Tensor[]" - - base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" - - base_fwd_def = _register_kernel( - op_name=base_fwd_name, - schema_str=fwd_schema, - arg_type=fwd_arg_type, - arg_names=fwd_arg_names, - adapters=fwd_adapters, - tensor_field_names=fwd_tensor_field_names, - impl=fwd_impl, - fake_impl=fwd_fake_impl, - format_result=_format_fwd_result, - ) - _register_kernel( - op_name=base_bwd_name, - schema_str=bwd_schema, - arg_type=backward_arg_type, - arg_names=bwd_arg_names, - adapters=bwd_adapters, - tensor_field_names=bwd_tensor_field_names, - impl=backward_impl, - fake_impl=bwd_fake_impl, - format_result=lambda g: _format_bwd_result(g, num_grad_inputs, base_bwd_qualname), - ) - - base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) - base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) - - wrapper_fwd_def = _register_wrapper_op( - wrapper_op_name=wrapper_fwd_name, - schema_str=fwd_schema, - base_op=base_fwd_op, - adapters=fwd_adapters, - ) - wrapper_bwd_def = _register_wrapper_op( - wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op + pair = _register_two_tier_pair( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_arg_type=backward_arg_type, + bwd_impl=backward_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=len(input_tensors_for_grad), ) + slot_count, grad_targets = _resolve_grad_targets(pair.fwd_adapters, input_tensors_for_grad) autograd_common = { "fwd_arg_type": fwd_arg_type, - "fwd_arg_names": fwd_arg_names, - "fwd_adapters": fwd_adapters, - "fwd_tensor_field_names": fwd_tensor_field_names, - "bwd_arg_names": bwd_arg_names, - "bwd_adapters": bwd_adapters, + "fwd_arg_names": pair.fwd_arg_names, + "fwd_adapters": pair.fwd_adapters, + "fwd_tensor_field_names": pair.fwd_tensor_field_names, + "bwd_arg_names": pair.bwd_arg_names, + "bwd_adapters": pair.bwd_adapters, "slot_count": slot_count, "grad_targets": grad_targets, "setup_context_user": setup_context, "backward_obj_type": backward_arg_type, "fwd_fake_impl": fwd_fake_impl, } - wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) - wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) - - _register_autograd_for_op(fwd_op=base_fwd_def, bwd_op=base_bwd_op, **autograd_common) - _register_autograd_for_op(fwd_op=wrapper_fwd_def, bwd_op=wrapper_bwd_op, **autograd_common) - - fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_adapters) - bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_adapters) - - def _fwd_rule(mode, func, types, args, kwargs): - del mode, func, types, kwargs - new_args = list(args) - for sub in subclass_list: - _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) - return base_fwd_op(*new_args) - - def _bwd_rule(mode, func, types, args, kwargs): - del mode, func, types, kwargs - new_args = list(args) - for sub in subclass_list: - _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) - return base_bwd_op(*new_args) - - for sub in subclass_list: - wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) - wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule) - - _quantized_tensor_passthrough_ops.add(wrapper_fwd_op.default) - _quantized_tensor_passthrough_ops.add(wrapper_bwd_op.default) - _quantized_tensor_passthrough_ops.add(base_fwd_op.default) - _quantized_tensor_passthrough_ops.add(base_bwd_op.default) + _register_autograd_for_op(fwd_op=pair.base_fwd_def, bwd_op=pair.base_bwd_op, **autograd_common) + _register_autograd_for_op( + fwd_op=pair.wrapper_fwd_def, bwd_op=pair.wrapper_bwd_op, **autograd_common + ) def forward_fn(fwd_args): - spec_obj = _spec_view(fwd_args, fwd_tensor_field_names) - user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) - kwargs = _args_to_slots(fwd_args, fwd_adapters) - flat_in = [kwargs[name] for name in fwd_arg_names] - result = wrapper_fwd_op(*flat_in) - - cursor = 0 - outputs: List[Any] = [] - for spec in user_fakes: - n = _spec_slot_count(spec) - chunk = [_decode_none(t) for t in result[cursor : cursor + n]] - cursor += n - outputs.append(_spec_reassemble(spec, chunk)) - + spec_obj = _spec_view(fwd_args, pair.fwd_tensor_field_names) + user_specs, _saved_specs, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) + kwargs = _args_to_slots(fwd_args, pair.fwd_adapters) + payload = pair.wrapper_fwd_op(*[kwargs[name] for name in pair.fwd_arg_names]) + outputs, _ = _reassemble(user_specs, payload) if len(outputs) == 1: return outputs[0] return tuple(outputs) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 219263773d..0e4eb634e0 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -83,7 +83,7 @@ prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorSpec, register_custom_op, is_value_opaque_quantizer +from ..dynamo import TensorSpec, register_custom_op_with_autograd, is_value_opaque_quantizer from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -1753,7 +1753,7 @@ def _linear_backward_impl_fake( # Custom op used under ``torch.compile``. -_linear_op = register_custom_op( +_linear_op = register_custom_op_with_autograd( op_name="linear", input_tensors_for_grad=["weight", "inp", "bias"], fwd_arg_type=LinearFwdArgs, diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 5106ec9e0a..2cf28e1e39 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -9,7 +9,7 @@ from collections.abc import Iterable import dataclasses import pickle -from typing import Any, Optional +from typing import Any, Callable, Optional import torch @@ -22,6 +22,7 @@ autocast, ) from ..tensor import Quantizer +from ..dynamo import is_value_opaque_quantizer, register_custom_op @dataclasses.dataclass @@ -184,6 +185,45 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta): # Number of extra tensor outputs num_extra_outputs: int = 0 + # torch.compile support. An operation opts in by declaring the two arg + # containers and implementing the four compute classmethods below; the base + # class then registers its custom ops and drives them from op_forward / + # op_backward, so no operation writes that plumbing itself. + fwd_args_type: Optional[type] = None + bwd_args_type: Optional[type] = None + # Gradients returned by backward_compute: the input's, then any parameters'. + num_grad_inputs: int = 1 + # (forward_fn, backward_fn) pair, or None if the operation cannot be compiled. + compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + if cls.fwd_args_type is None or cls.bwd_args_type is None: + return + if getattr(cls.forward_compute, "__isabstractmethod__", False): + return + for name, arg_type in ( + ("fwd_args_type", cls.fwd_args_type), + ("bwd_args_type", cls.bwd_args_type), + ): + # The op schema is built from the container's fields, so this is the + # framework's actual requirement -- check it where it is declared. + if not dataclasses.is_dataclass(arg_type): + raise TypeError(f"{cls.__name__}.{name} must be a dataclass") + # One registration per class. The compute halves are bound here, so a + # subclass that only swaps kernels (the activations) still gets its own + # op without repeating any of this. + cls.compile_ops = register_custom_op( + op_name=cls.__name__.lower(), + fwd_arg_type=cls.fwd_args_type, + fwd_impl=cls.forward_compute, + fwd_fake_impl=cls.forward_fake, + bwd_arg_type=cls.bwd_args_type, + bwd_impl=cls.backward_compute, + bwd_fake_impl=cls.backward_fake, + num_grad_inputs=cls.num_grad_inputs, + ) + def __init__(self) -> None: super().__init__() @@ -191,6 +231,93 @@ def __init__(self) -> None: self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None + # ------------------------------------------------------------------ # + # Compute halves. Classmethods, not free functions: they belong to the + # operation, and binding to the class is what lets a family of operations + # share one implementation while dispatching to per-class kernels. + # ------------------------------------------------------------------ # + + @classmethod + def forward_compute(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]: + """Pure forward: ``(output, tensors_to_save, ctx_attrs)``. + + Takes everything through ``args``; must not read ``self`` or global + state, both of which are invisible to the compiler at this point. + """ + raise NotImplementedError + + @classmethod + def forward_fake(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]: + """Allocation-free twin of :meth:`forward_compute` over ``TensorSpec``. + + Runs as a meta kernel, outside the traced frame, and more than once per + compile, so it must be a pure function of ``args`` -- a read of global + state here is unguarded and can silently disagree with the real impl. + """ + raise NotImplementedError + + @classmethod + def backward_compute(cls, args: Any) -> tuple: + """Pure backward: ``num_grad_inputs`` gradients.""" + raise NotImplementedError + + @classmethod + def backward_fake(cls, args: Any) -> tuple: + """Allocation-free twin of :meth:`backward_compute`.""" + raise NotImplementedError + + def compile_unsupported_reason(self) -> Optional[str]: + """Why this operation cannot go through its custom op, or ``None``. + + Asked per operation, but acted on per fuser group: a pipeline compiles + as a whole, so one unsupported operation sends the whole group to eager. + Recipe-level limits are not checked here -- they belong to whoever reads + the recipe, which is the fuser. + """ + if self.compile_ops is None: + return f"{self.__class__.__name__} does not implement the compute halves" + for mode in ("forward", "backward"): + for index in range(self.num_quantizers(mode)): + quantizer = self.get_quantizer(mode, index) + if quantizer is not None and not is_value_opaque_quantizer(quantizer): + # Delayed scaling holds live scale/amax tensors, so its + # quantizer cannot be specialized on and would be baked into + # the graph as a stale constant. + return ( + f"{type(quantizer).__name__} is not a torch.compile value-opaque quantizer" + ) + return None + + def resolve_fwd_args( + self, + input_: torch.Tensor, + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> Any: + """Gather the forward's inputs into a flat, ``self``-free container. + + This is where module config and global state are read, so it belongs in + the traced region where Dynamo guards those reads -- never inside the + custom op. + """ + raise NotImplementedError + + def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> Any: + """Rebuild the backward's inputs from the forward's saved state.""" + raise NotImplementedError + + def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple: + """Tensors to persist, given what the forward handed back. + + An operation whose backward needs its input but whose forward does not + produce a distinct tensor for it overrides this; a custom op may not + return one of its own inputs. + """ + del input_ + return saved + @property def is_fused_op(self) -> bool: return False @@ -425,7 +552,6 @@ def _load_fp8_metas(self, fp8_metas: Optional[dict[str, Any]]) -> None: self._fp8_metas[mode][fp8_meta_key].scale.copy_(scale) self._fp8_metas[mode][fp8_meta_key].amax_history.copy_(amax_history) - @abc.abstractmethod def op_forward( self, ctx: OperationContext, @@ -437,6 +563,10 @@ def op_forward( ) -> torch.Tensor: """Forward pass + Operations that declare the compute halves inherit this: it resolves the + arguments, runs the forward, and records what the backward will need. The + rest override it. + Parameters ---------- ctx: OperationContext @@ -454,8 +584,63 @@ def op_forward( Output tensor """ + if self.fwd_args_type is None: + raise NotImplementedError( + f"{self.__class__.__name__} implements neither op_forward nor the compute halves" + ) + if kwargs: + raise ValueError(f"{self.__class__.__name__} forward does not expect keyword arguments") + args = self.resolve_fwd_args( + input_, + requires_grad=ctx.requires_grad, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + ) + output, saved, ctx_attrs = self.forward_compute(args) + if ctx.requires_grad: + ctx.save_for_backward(*self.saved_for_backward(saved, input_)) + for name, value in ctx_attrs.items(): + setattr(ctx, name, value) + return output + + def compiled_op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + ) -> torch.Tensor: + """:meth:`op_forward` routed through this operation's custom op. + + Same bookkeeping, but the computation crosses an op boundary so Dynamo + sees one graph node instead of tracing into the kernels. + """ + args = self.resolve_fwd_args( + input_, + requires_grad=ctx.requires_grad, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + ) + output, saved, ctx_attrs = self.compile_ops[0](args) + if ctx.requires_grad: + ctx.save_for_backward(*self.saved_for_backward(saved, input_)) + for name, value in ctx_attrs.items(): + setattr(ctx, name, value) + return output + + def compiled_op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: + """:meth:`op_backward` routed through this operation's custom op.""" + grads = self.compile_ops[1](self.resolve_bwd_args(ctx, grad_output)) + grad_input = grads[0] + if grad_input is None: + grad_input = grad_output + return grad_input, tuple(grads[1:]) - @abc.abstractmethod def op_backward( self, ctx: OperationContext, @@ -463,6 +648,8 @@ def op_backward( ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: """Backward pass + Counterpart to the inherited :meth:`op_forward`. + Parameters ---------- ctx: OperationContext @@ -478,6 +665,17 @@ def op_backward( Loss gradients w.r.t. parameters """ + if self.bwd_args_type is None: + raise NotImplementedError( + f"{self.__class__.__name__} implements neither op_backward nor the compute halves" + ) + grads = self.backward_compute(self.resolve_bwd_args(ctx, grad_output)) + grad_input = grads[0] + if grad_input is None: + # "The incoming gradient, unchanged": a custom op may not return one + # of its own inputs, so the compute half hands back None instead. + grad_input = grad_output + return grad_input, tuple(grads[1:]) def fuser_forward( self, From e25ec8dcfb71de58ba2a41a7781dd20fb986d8a7 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 5 Aug 2026 22:21:51 +0200 Subject: [PATCH 2/4] [PyTorch] Compile an OperationFuser group holding one operation A group whose operations declare their compute halves now 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 will later let its forward and backward walk different op groupings. Four side effects reached outside the higher-order op's scope and had to go: - 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 reported from the eager path only. They 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. 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 declare their halves. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 140 +++++++++++++++++++ transformer_engine/pytorch/ops/fuser.py | 125 +++++++++++++---- transformer_engine/pytorch/ops/sequential.py | 17 ++- 3 files changed, 254 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/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]]], From 0c2e100422644e0cc3fa40c67f2e18c64acc10d5 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 6 Aug 2026 19:41:39 +0200 Subject: [PATCH 3/4] [PyTorch] Accept an operation's declared forward kwargs under compile An operation lists the forward kwargs it takes in fwd_kwarg_names. They are resolved into its args container like any other config, in the traced Python where Dynamo guards them, so they reach the custom op through the existing schema -- a value is guarded, a tensor is lifted into the graph, and a quantized one crosses as its inner buffers. An undeclared kwarg still sends the whole group to eager. That is not a schema limitation, as the old message implied: the kwargs that remain are the grouped operations' preallocated buffers, which the op writes to, and a custom op may not mutate a tensor from an enclosing scope. A kwarg carries no gradient. This matches the eager path, where kwargs never entered the autograd graph either, and is why only read-only ones are accepted. The fuser test helper now builds a separate model for the eager and the compiled pass. Previously both shared one model and the eager pass ran first to produce the reference, so the compiled pass was always traced on a model whose module groups, fusions and pre_first_fuser_forward had already run. Those paths are now traced as well. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 199 ++++++++++++++++++++---- transformer_engine/pytorch/ops/fuser.py | 11 +- transformer_engine/pytorch/ops/op.py | 22 ++- 3 files changed, 200 insertions(+), 32 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index cb362f92f1..4aa10fba50 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -5,6 +5,7 @@ import abc import contextlib import dataclasses +from typing import Union import pytest import torch @@ -33,7 +34,11 @@ 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 +from transformer_engine.pytorch.quantized_tensor import ( + QuantizedTensor, + QuantizedTensorStorage, + Quantizer, +) from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec from transformer_engine.pytorch import ( is_fp8_available, @@ -1606,26 +1611,146 @@ def resolve_bwd_args(self, ctx, grad_output): return _ScaleBwdArgs(grad_output=grad_output, saved_input=x, scale=self.scale) -def _assert_sequential_matches_eager(model, compiled, base): +@dataclasses.dataclass(slots=True) +class _ScaleKwargsFwdArgs: + """Flat inputs to the kwarg-taking test operation's forward.""" + + input_: torch.Tensor + scale: torch.Tensor + extra_scale: float + offset: Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclasses.dataclass(slots=True) +class _ScaleKwargsBwdArgs: + """Flat inputs to the kwarg-taking test operation's backward.""" + + grad_output: torch.Tensor = None + saved_input: torch.Tensor = None + scale: torch.Tensor = None + extra_scale: float = 1.0 + + +class _ScaleWithKwargsOp(BasicOperation): + """Test-only operation taking forward kwargs: a value and a tensor. + + ``offset`` is declared as tensor-or-quantized, so a quantized kwarg crosses + the op boundary as its inner buffers. Neither kwarg carries a gradient -- + that is what "read-only" means here. + """ + + fwd_args_type = _ScaleKwargsFwdArgs + bwd_args_type = _ScaleKwargsBwdArgs + num_grad_inputs = 2 # grad input, grad scale + fwd_kwarg_names = ("extra_scale", "offset") + + 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): + offset = args.offset + if isinstance(offset, QuantizedTensor): + offset = offset.dequantize() + out = args.input_ * args.scale * args.extra_scale + offset + return out, (), {"extra_scale": args.extra_scale} + + @classmethod + def forward_fake(cls, args): + x = args.input_ + return ( + TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), + (), + {"extra_scale": args.extra_scale}, + ) + + @classmethod + def backward_compute(cls, args): + dy = args.grad_output + return ( + dy * args.scale * args.extra_scale, + (dy * args.saved_input).sum() * args.extra_scale, + ) + + @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_): + del saved + return (input_,) + + def resolve_fwd_args( + self, + input_, + *, + requires_grad, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=None, + extra_scale=1.0, + offset=None, + ): + del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer + if offset is None: + offset = torch.zeros((), device=input_.device, dtype=input_.dtype) + return _ScaleKwargsFwdArgs( + input_=input_, + scale=self.scale, + extra_scale=extra_scale, + offset=offset, + ) + + def resolve_bwd_args(self, ctx, grad_output): + (x,) = ctx.saved_tensors + return _ScaleKwargsBwdArgs( + grad_output=grad_output, + saved_input=x, + scale=self.scale, + extra_scale=ctx.extra_scale, + ) + + +def _assert_sequential_matches_eager(make_model, base, op_kwargs_seq=(None,)): """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()] + the output and every parameter gradient. - 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() + Each pass gets its own freshly built model, so the compiled one is traced on + a first run: nothing has built the module groups, resolved the fusions or run + ``pre_first_fuser_forward`` on it beforehand. ``make_model`` must therefore + build deterministically identical models. + + Several ``op_kwargs`` are run in order on the same pair of models, which is + what exercises Dynamo's guards on a kwarg value. + """ + eager_model = make_model() + compiled_model = make_model() + compiled = torch.compile(compiled_model, fullgraph=True) - 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) + for op_kwargs in op_kwargs_seq: + call_kwargs = {} if op_kwargs is None else {"op_kwargs": op_kwargs} + + inp_eager = base.detach().clone().requires_grad_(True) + eager_model.zero_grad(set_to_none=True) + out_eager = eager_model(inp_eager, **call_kwargs) + 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 eager_model.parameters()] + + inp_compiled = base.detach().clone().requires_grad_(True) + compiled_model.zero_grad(set_to_none=True) + out_compiled = compiled(inp_compiled, **call_kwargs).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(compiled_model.parameters(), ref_pgrads): + torch.testing.assert_close(got.grad, expected) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -1637,10 +1762,8 @@ def test_te_ops_single_op_group_compiles(): 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) + _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp()), base) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -1653,10 +1776,34 @@ def test_te_ops_unsupported_group_still_compiles_eagerly(): mutation of anything from an enclosing scope -- have to hold on both paths. """ torch._dynamo.reset() - op = te.ops.Identity() - assert op.compile_unsupported_reason() is not None + assert te.ops.Identity().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) + _assert_sequential_matches_eager(lambda: te.ops.Sequential(te.ops.Identity()), base) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_te_ops_forward_kwargs_compile(): + """Forward kwargs reach the operation through its custom op. + + Covers both kinds at once: a value, which Dynamo guards on -- hence the + second call with a different one -- and a tensor, quantized here, which + crosses the op boundary as its inner buffers. + """ + torch._dynamo.reset() + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device=torch.device("cuda"), + ) + offset = quantizer(torch.randn(64, dtype=torch.bfloat16, device="cuda")) + + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + _assert_sequential_matches_eager( + lambda: te.ops.Sequential(_ScaleWithKwargsOp()), + base, + op_kwargs_seq=( + {0: {"extra_scale": 3.0, "offset": offset}}, + {0: {"extra_scale": 5.0, "offset": offset}}, + ), + ) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index fe1eed9ad4..89d3d45744 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -148,6 +148,7 @@ def forward( x, prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, next_op_input_quantizer=next_op_input_quantizer, + **basic_op_kwargs[basic_op_idxs[0]], ) fused_op_extra_outputs = [()] else: @@ -549,8 +550,14 @@ def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) -> 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, kwargs in zip(self._basic_ops, basic_op_kwargs): + # A kwarg an operation declares is resolved into its args container + # like any other config. Anything else -- notably the preallocated + # buffers of the grouped operations -- is written to by the op, and a + # custom op may not mutate a tensor from an enclosing scope. + unsupported = sorted(name for name in kwargs if name not in op.fwd_kwarg_names) + if unsupported: + return f"{type(op).__name__} does not support keyword arguments {unsupported}" for op in self._basic_ops: reason = op.compile_unsupported_reason() if reason is not None: diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 2cf28e1e39..c3af65a831 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -193,6 +193,9 @@ 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 kwargs this operation accepts, resolved into fwd_args_type like + # any other config. A kwarg carries no gradient and must not be mutated. + fwd_kwarg_names: tuple[str, ...] = () # (forward_fn, backward_fn) pair, or None if the operation cannot be compiled. compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None @@ -295,12 +298,15 @@ def resolve_fwd_args( requires_grad: bool, prev_op_grad_output_quantizer: Optional[Quantizer] = None, next_op_input_quantizer: Optional[Quantizer] = None, + **kwargs: Any, ) -> 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. + custom op. ``kwargs`` are the caller's forward kwargs, restricted to + ``fwd_kwarg_names``; an operation declaring them supplies their defaults + here, since a kwarg may be absent. """ raise NotImplementedError @@ -588,13 +594,17 @@ def op_forward( 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") + unsupported = sorted(name for name in kwargs if name not in self.fwd_kwarg_names) + if unsupported: + raise ValueError( + f"{self.__class__.__name__} forward does not accept keyword arguments {unsupported}" + ) 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, + **kwargs, ) output, saved, ctx_attrs = self.forward_compute(args) if ctx.requires_grad: @@ -610,17 +620,21 @@ def compiled_op_forward( *, prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], + **kwargs: Any, ) -> 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. + sees one graph node instead of tracing into the kernels. ``kwargs`` are + not validated here -- the fuser's gate already rejected a group whose + kwargs an operation does not declare. """ 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, + **kwargs, ) output, saved, ctx_attrs = self.compile_ops[0](args) if ctx.requires_grad: From b389608d0ef35cc974833b1df6c0eb84529a86a2 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 7 Aug 2026 16:25:23 +0200 Subject: [PATCH 4/4] [PyTorch] Restrict compiled forward kwargs to tensors A value kwarg does not survive a second call. The other fields of an args container are read off the module and are constant across calls, so they are baked into the graph; a kwarg changes per call, and on the second value Dynamo hands over a symbolic scalar, which OpaqueValueBundle cannot carry -- it fails with AsPythonConstantNotImplementedError, not with a graph break. Measured on int and float alike; specialize_float=True cures only the float, and is global. The gate now takes tensor kwargs only, so a value sends the group to eager deterministically instead of failing on its second call. A 0-d tensor is the way to pass a scalar: it is a graph input, so it recompiles for no value at all. The test carries a quantized offset that changes on every call and confirms no recompilation, then adds a value kwarg to cover the gated path. That last call keeps its offset unquantized on purpose: a gated group runs the eager implementation, which is traced directly rather than hidden behind a custom op, and dequantize() graph-breaks there. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 21 ++++++++++++++------- transformer_engine/pytorch/ops/fuser.py | 16 ++++++++++++---- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 4aa10fba50..8c6ca357bb 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1785,25 +1785,32 @@ def test_te_ops_unsupported_group_still_compiles_eagerly(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) def test_te_ops_forward_kwargs_compile(): - """Forward kwargs reach the operation through its custom op. + """A tensor forward kwarg reaches the operation through its custom op. - Covers both kinds at once: a value, which Dynamo guards on -- hence the - second call with a different one -- and a tensor, quantized here, which - crosses the op boundary as its inner buffers. + The tensor is quantized, so it crosses the op boundary as its inner buffers, + and it changes between calls, which a graph input absorbs without a + recompilation. The last call adds a value kwarg: that one is gated onto the + eager implementation, since Dynamo turns a changed scalar into a symbol that + cannot be carried as opaque config. """ torch._dynamo.reset() quantizer = Float8CurrentScalingQuantizer( fp8_dtype=tex.DType.kFloat8E4M3, device=torch.device("cuda"), ) - offset = quantizer(torch.randn(64, dtype=torch.bfloat16, device="cuda")) + + def offset(value): + return quantizer(torch.full((64,), value, dtype=torch.bfloat16, device="cuda")) base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") _assert_sequential_matches_eager( lambda: te.ops.Sequential(_ScaleWithKwargsOp()), base, op_kwargs_seq=( - {0: {"extra_scale": 3.0, "offset": offset}}, - {0: {"extra_scale": 5.0, "offset": offset}}, + {0: {"offset": offset(0.5)}}, + {0: {"offset": offset(1.5)}}, + # No quantized offset here: this call runs the eager implementation, + # which is traced directly, and dequantize() is not traceable. + {0: {"extra_scale": 3.0}}, ), ) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 89d3d45744..a1d8d260d9 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -550,14 +550,22 @@ def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) -> if len(self._forward_ops) != self._num_basic_ops: # A fused op covers several basic ops; only single-op groups so far. return "fused operations are not supported yet" - for op, kwargs in zip(self._basic_ops, basic_op_kwargs): + for op, kwargs in zip(self._basic_ops, basic_op_kwargs, strict=True): # A kwarg an operation declares is resolved into its args container # like any other config. Anything else -- notably the preallocated # buffers of the grouped operations -- is written to by the op, and a # custom op may not mutate a tensor from an enclosing scope. - unsupported = sorted(name for name in kwargs if name not in op.fwd_kwarg_names) - if unsupported: - return f"{type(op).__name__} does not support keyword arguments {unsupported}" + undeclared = sorted(name for name in kwargs if name not in op.fwd_kwarg_names) + if undeclared: + return f"{type(op).__name__} does not support keyword arguments {undeclared}" + # Only tensors. The other fields of an args container are values read + # off the module, constant across calls and baked into the graph; a + # kwarg changes per call, and on the second value Dynamo hands over a + # symbolic scalar, which cannot go into an opaque value bundle. Pass a + # 0-d tensor instead -- it is a graph input, so it does not recompile. + values = sorted(name for name, v in kwargs.items() if not isinstance(v, torch.Tensor)) + if values: + return f"{type(op).__name__} keyword arguments {values} are not tensors" for op in self._basic_ops: reason = op.compile_unsupported_reason() if reason is not None: