diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index f8e09d5ce1..c02863a1a0 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -4,10 +4,12 @@ import abc import contextlib +import dataclasses import os import re import sys import warnings +from typing import Union import pytest import torch @@ -37,12 +39,18 @@ from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear +from transformer_engine.pytorch.ops.fuser import OperationFuser +from transformer_engine.pytorch.ops.op import BasicOperation from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer -from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, Quantizer -from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec +from transformer_engine.pytorch.quantized_tensor import ( + QuantizedTensor, + QuantizedTensorStorage, + Quantizer, +) +from transformer_engine.pytorch.dynamo import ForwardResult, TensorSpec, to_tensor_spec from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, @@ -2339,3 +2347,330 @@ def fn(inp): "Unexpected recompilation(s) across different batch sizes: " f"{unique_graphs_after - unique_graphs_baseline} extra graph(s) compiled" ) + + +# --------------------------------------------------------------------------- # +# 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 ForwardResult(args.input_ * args.scale) + + @classmethod + def forward_fake(cls, args): + x = args.input_ + return ForwardResult(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 setup_context(self, ctx, args, aux): + del aux + ctx.save_for_backward(args.input_, args.scale) + + 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, scale = ctx.saved_tensors + return _ScaleBwdArgs(grad_output=grad_output, saved_input=x, scale=scale) + + +class _BackwardScalePair(te.ops.FusedOperation): + """Backward-only fusion for the compile gate test.""" + + def fuser_backward(self, basic_op_ctxs, grad_output, **unused): + dx, grad_params_1 = self.basic_ops[1].op_backward(basic_op_ctxs[1], grad_output) + dx, grad_params_0 = self.basic_ops[0].op_backward(basic_op_ctxs[0], dx) + return dx, [grad_params_0, grad_params_1], [(), ()] + + +def _fuse_backward_scale_pair(ops, **unused): + if len(ops) == 2 and all(isinstance(op, _ScaleOp) for op in ops): + return [_BackwardScalePair(ops)] + return ops + + +@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 ForwardResult(out) + + @classmethod + def forward_fake(cls, args): + x = args.input_ + return ForwardResult(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 * 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 setup_context(self, ctx, args, aux): + del aux + ctx.save_for_backward(args.input_, args.scale) + ctx.extra_scale = args.extra_scale + + 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, scale = ctx.saved_tensors + return _ScaleKwargsBwdArgs( + grad_output=grad_output, + saved_input=x, + scale=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. + + 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) + + 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") +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() + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp()), base) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_te_ops_backward_fusion_uses_eager_implementations(): + """A backward fusion prevents the group from using basic-op custom ops.""" + te.ops.register_backward_fusion(_fuse_backward_scale_pair, prepend=True) + try: + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + with pytest.warns(UserWarning, match="backward fusion"): + _assert_sequential_matches_eager( + lambda: te.ops.Sequential(_ScaleOp(), _ScaleOp()), base + ) + finally: + OperationFuser.backward_fusion_functions.remove(_fuse_backward_scale_pair) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("compile_model", [False, True], ids=["eager", "compiled"]) +def test_te_ops_setup_context_saves_parameter(compile_model): + """Backward observes mutation of a tensor used by the forward.""" + op = _ScaleOp() + model = te.ops.Sequential(op) + if compile_model: + model = torch.compile(model, fullgraph=True) + x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) + y = model(x) + with torch.no_grad(): + op.scale.add_(1) + with pytest.raises(RuntimeError, match="modified by an inplace operation"): + y.sum().backward() + + +@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() + assert te.ops.Identity().compile_unsupported_reason() is not None + + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + _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(): + """A tensor forward kwarg reaches the operation through its custom op. + + 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"), + ) + + 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: {"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/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index e42eb8f9f6..88a2d716a7 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,13 +6,20 @@ 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, TensorOrQuantized +from .custom_op import ( + ForwardResult, + register_custom_op, + register_custom_op_with_autograd, + TensorOrQuantized, +) __all__ = [ "register_value_opaque_quantizer", "is_value_opaque_quantizer", "TensorSpec", "to_tensor_spec", + "ForwardResult", "register_custom_op", + "register_custom_op_with_autograd", "TensorOrQuantized", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 00846d615a..1f5df798f6 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -6,11 +6,16 @@ Registers TE modules' eager forward/backward as ``torch.library`` custom ops so ``torch.compile(fullgraph=True)`` traces them as single graph nodes. -``register_custom_op`` is the entry point; ``module/linear.py`` is the first user. +``register_custom_op_with_autograd`` is the entry point for a module that +wires autograd on the op itself (``module/linear.py`` is the first user); +``register_custom_op`` hands back the forward and backward ops separately, for a +caller that drives autograd at a higher level (``ops/fuser.py``). A TE forward/backward implementation takes one dataclass argument (``fwd_arg_type`` / ``bwd_arg_type``, e.g. ``LinearFwdArgs``) whose fields mix tensors, quantized tensors, quantizers, process groups and plain Python values. +The autograd-free forward returns ``ForwardResult(output, aux)``; the +autograd-wired API keeps its saved-tensor and context-metadata contract. A ``torch.library`` custom op is narrower: it only accepts flat schema slots (tensors plus opaque objects) and returns a flat ``Tensor[]``. @@ -41,9 +46,9 @@ 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 -``forward_fn`` -- the drop-in for the eager ``autograd.Function.apply``. A forward -call through it: +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: * runs the fake ``fwd_fake_impl`` on ``TensorSpec`` descriptors (data-free; see ``tensor_spec.py``) and parses its result into an ``_OutputPlan`` -- the @@ -111,6 +116,15 @@ _TE_OP_NAMESPACE = "transformer_engine_compile" + +@dataclasses.dataclass(frozen=True, slots=True) +class ForwardResult: + """Output and fresh auxiliary tensors produced by an autograd-free forward.""" + + output: Any + aux: tuple = () + + # Annotation for an op arg field that may hold a plain tensor, a quantized # tensor subclass or a *bare* ``QuantizedTensorStorage`` (the internal-quantizer # optimization). Matched exactly by ``_TensorOrQuantizedAdapter``. @@ -930,7 +944,7 @@ def _slice_user_grads( # --------------------------------------------------------------------------- # -# Op registration +# Op registration: base and wrapper ops, autograd wiring # --------------------------------------------------------------------------- # @@ -1166,7 +1180,242 @@ def _all_quantized_tensor_subclasses() -> List[type]: return found +@dataclasses.dataclass(frozen=True) +class _OpPair: + """One registered forward/backward pair, and what a caller needs to drive it.""" + + fwd_plan: _ArgPlan + bwd_plan: _ArgPlan + base_fwd_def: Any + base_bwd_op: Any + wrapper_fwd_def: Any + wrapper_fwd_op: Any + wrapper_bwd_op: Any + + def call_forward( + self, fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], fwd_args: Any + ) -> Tuple[_OutputPlan, List[torch.Tensor]]: + """Run the forward op on ``fwd_args``: its output plan and flat payload.""" + spec_obj = _spec_view(fwd_args, self.fwd_plan.tensor_field_names()) + out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj)) + kwargs = self.fwd_plan.pack(fwd_args) + payload = self.wrapper_fwd_op(*[kwargs[name] for name in self.fwd_plan.slot_names]) + return out_plan, payload + + +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_with_autograd`: the arg plans, 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. + """ + 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_plan = _parse_arg_type(fwd_arg_type) + bwd_plan = _parse_arg_type(bwd_arg_type) + + fwd_schema = f"{fwd_plan.schema_str} -> Tensor[]" + bwd_schema = f"{bwd_plan.schema_str} -> Tensor[]" + + base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" + + base_fwd_def = _register_base_op( + op_name=base_fwd_name, + schema_str=fwd_schema, + plan=fwd_plan, + impl=fwd_impl, + fake_impl=fwd_fake_impl, + pack_result=_pack_fwd_result, + ) + _register_base_op( + op_name=base_bwd_name, + schema_str=bwd_schema, + plan=bwd_plan, + impl=bwd_impl, + fake_impl=bwd_fake_impl, + pack_result=lambda g: _pack_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) + + fwd_slot_offsets = fwd_plan.tensor_or_quantized_offsets() + bwd_slot_offsets = bwd_plan.tensor_or_quantized_offsets() + + wrapper_fwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_fwd_name, + schema_str=fwd_schema, + base_op=base_fwd_op, + slot_offsets=fwd_slot_offsets, + subclasses=subclass_list, + ) + # Pass-through: a subclass input reaches the base op through the dispatch + # rule below, never through the wrapper body. + wrapper_bwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op + ) + 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_rule = _make_dispatch_rule( + _make_slot_forwarder(base_fwd_op, fwd_slot_offsets, subclass_list) + ) + _bwd_rule = _make_dispatch_rule( + _make_slot_forwarder(base_bwd_op, bwd_slot_offsets, subclass_list) + ) + + 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_plan=fwd_plan, + bwd_plan=bwd_plan, + base_fwd_def=base_fwd_def, + base_bwd_op=base_bwd_op, + wrapper_fwd_def=wrapper_fwd_def, + wrapper_fwd_op=wrapper_fwd_op, + wrapper_bwd_op=wrapper_bwd_op, + ) + + +# --------------------------------------------------------------------------- # +# 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. + + Callable contracts: + + * ``fwd_impl(fwd_args) -> ForwardResult(output, aux)`` + * ``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) -> (output, aux)`` -- ``aux`` contains only fresh + tensors produced by the custom op. The caller decides which tensors and + metadata to persist for backward. + * ``backward_fn(bwd_args) -> tuple`` of gradients. + + Returns ``None`` if registration fails (recorded 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: + record_compile_disabled( + 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.""" + + def adapt_forward(impl): + def wrapped(args): + result = impl(args) + if not isinstance(result, ForwardResult): + raise TypeError( + f"autograd-free fwd impl must return ForwardResult, got {type(result).__name__}" + ) + return result.output, result.aux, None + + return wrapped + + adapted_fwd_impl = adapt_forward(fwd_impl) + adapted_fwd_fake_impl = adapt_forward(fwd_fake_impl) + pair = _register_two_tier_pair( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=adapted_fwd_impl, + fwd_fake_impl=adapted_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): + out_plan, payload = pair.call_forward(adapted_fwd_fake_impl, fwd_args) + outputs = out_plan.user_outputs(payload) + aux = out_plan.saved_tensors(payload) + return outputs[0], tuple(aux) + + def backward_fn(bwd_args): + # Unlike the forward payload, each grad occupies exactly one slot + # (``_pack_bwd_result`` materializes a TensorSpec grad), so there is + # nothing to reassemble. + kwargs = pair.bwd_plan.pack(bwd_args) + payload = pair.wrapper_bwd_op(*[kwargs[name] for name in pair.bwd_plan.slot_names]) + return tuple(_decode_none(t) for t in payload) + + return forward_fn, backward_fn + + +def register_custom_op_with_autograd( *, op_name: str, input_tensors_for_grad: List[str], @@ -1231,7 +1480,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, @@ -1249,7 +1498,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], @@ -1261,7 +1510,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 @@ -1271,96 +1520,32 @@ 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_plan = _parse_arg_type(fwd_arg_type) - bwd_plan = _parse_arg_type(bwd_arg_type) - - num_grad_inputs = len(input_tensors_for_grad) - grad_targets = fwd_plan.resolve_grad_targets(input_tensors_for_grad) - - fwd_schema = f"{fwd_plan.schema_str} -> Tensor[]" - bwd_schema = f"{bwd_plan.schema_str} -> Tensor[]" - - base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" - - base_fwd_def = _register_base_op( - op_name=base_fwd_name, - schema_str=fwd_schema, - plan=fwd_plan, - impl=fwd_impl, - fake_impl=fwd_fake_impl, - pack_result=_pack_fwd_result, - ) - _register_base_op( - op_name=base_bwd_name, - schema_str=bwd_schema, - plan=bwd_plan, - impl=bwd_impl, - fake_impl=bwd_fake_impl, - pack_result=lambda g: _pack_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) - - fwd_slot_offsets = fwd_plan.tensor_or_quantized_offsets() - bwd_slot_offsets = bwd_plan.tensor_or_quantized_offsets() - - wrapper_fwd_def = _register_wrapper_op( - wrapper_op_name=wrapper_fwd_name, - schema_str=fwd_schema, - base_op=base_fwd_op, - slot_offsets=fwd_slot_offsets, - subclasses=subclass_list, - ) - # Pass-through: a subclass input reaches the base op through the dispatch - # rule below, never through the wrapper body. - 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=bwd_arg_type, + bwd_impl=bwd_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=len(input_tensors_for_grad), ) autograd_common = { - "fwd_plan": fwd_plan, - "bwd_plan": bwd_plan, - "grad_targets": grad_targets, + "fwd_plan": pair.fwd_plan, + "bwd_plan": pair.bwd_plan, + "grad_targets": pair.fwd_plan.resolve_grad_targets(input_tensors_for_grad), "setup_context_user": setup_context, "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_rule = _make_dispatch_rule( - _make_slot_forwarder(base_fwd_op, fwd_slot_offsets, subclass_list) - ) - _bwd_rule = _make_dispatch_rule( - _make_slot_forwarder(base_bwd_op, bwd_slot_offsets, subclass_list) + _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 ) - 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) - def forward_fn(fwd_args): - spec_obj = _spec_view(fwd_args, fwd_plan.tensor_field_names()) - out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj)) - kwargs = fwd_plan.pack(fwd_args) - flat_in = [kwargs[name] for name in fwd_plan.slot_names] - result = wrapper_fwd_op(*flat_in) - - outputs = out_plan.user_outputs(result) + out_plan, payload = pair.call_forward(fwd_fake_impl, fwd_args) + outputs = out_plan.user_outputs(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 55fc69ef7f..3b3c99facd 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -88,7 +88,7 @@ from ..dynamo import ( TensorSpec, TensorOrQuantized, - register_custom_op, + register_custom_op_with_autograd, is_value_opaque_quantizer, ) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer @@ -1788,7 +1788,7 @@ def _linear_backward_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/fuser.py b/transformer_engine/pytorch/ops/fuser.py index fd66529ba8..87a3f18d96 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 ..utils import warn_compile_eager_fallback 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_custom_ops: 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_custom_ops: bool + Whether to call the operations' custom ops instead of tracing 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 # Place user provided extra inputs into their basic-op slots. Slots bound to # internal channels are filled lazily as their producers execute. @@ -153,14 +164,24 @@ 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_custom_ops: + x = op.compiled_op_forward( + basic_op_ctxs[basic_op_idxs[0]], + x, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + **basic_op_kwargs[basic_op_idxs[0]], + ) + fused_op_extra_outputs = [()] + else: + x, fused_op_extra_outputs = op.fuser_forward( + [basic_op_ctxs[idx] for idx in basic_op_idxs], + 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 len(fused_op_extra_outputs) != len(basic_op_idxs): raise RuntimeError( f"Expected {type(op).__name__} to generate extra outputs for " @@ -227,9 +248,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 @@ -244,15 +269,20 @@ def forward( func_ctx.basic_op_extra_output_consumers = fuser._basic_op_extra_output_consumers func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources func_ctx.is_first_module = is_first_module - - # Mark output tensors as not deletable in backward - for tensor in itertools.chain( - (x,), - (y for ys in extra_outputs for y in ys if y is not None), - ): - tensor._do_not_clear = True - - if set_output_requires_grad: + func_ctx.use_custom_ops = use_custom_ops + + # Mark output tensors as not deletable in backward (eager only; see above) + if not torch.compiler.is_compiling(): + for tensor in itertools.chain( + (x,), + (y for ys in extra_outputs for y in ys if y is not None), + ): + tensor._do_not_clear = True + + # 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: @@ -277,7 +307,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 @@ -327,14 +362,22 @@ def backward( channel_grad if output_grad is None else output_grad + channel_grad ) 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_custom_ops: + 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 for input_idx, grad in enumerate(dxs): @@ -392,6 +435,7 @@ def backward( None, # fuser None, # basic_op_kwargs None, # set_output_requires_grad + None, # use_custom_ops *grad_params_flat, *grad_extra_inputs_flat, ) @@ -691,6 +735,54 @@ def maybe_fuse_ops( else: self._last_amax_history_len = 0 + def _custom_ops_unsupported_reason( + self, basic_op_kwargs: list[dict[str, Any]] + ) -> Optional[str]: + """Why this group may not run through its operations' custom ops.""" + for mode, ops in (("forward", self._forward_ops), ("backward", self._backward_ops)): + if len(ops) != self._num_basic_ops or any( + op is not self._basic_ops[idx] or basic_op_idxs != [idx] + for idx, (op, basic_op_idxs) in enumerate(ops) + ): + return f"a {mode} fusion" + 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. + undeclared = sorted(name for name in kwargs if name not in op.fwd_kwarg_names) + if undeclared: + return f"{type(op).__name__} with undeclared 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__} with non-tensor keyword arguments {values}" + for op in self._basic_ops: + if op.num_extra_inputs or op.num_extra_outputs: + return f"{type(op).__name__} with extra tensor inputs or outputs" + reason = op.compile_unsupported_reason() + if reason is not None: + return reason + return None + + def _use_custom_ops(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. + """ + if not torch.compiler.is_compiling(): + return False + reason = self._custom_ops_unsupported_reason(basic_op_kwargs) + if reason is None: + return True + warn_compile_eager_fallback(reason) + return False + def __call__( self, input: torch.Tensor, # pylint: disable=redefined-builtin @@ -733,11 +825,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_custom_ops = self._use_custom_ops(basic_op_kwargs) + args = ( input, self, basic_op_kwargs, is_grad_enabled, # set_output_requires_grad + use_custom_ops, *self._flat_basic_op_params, *extra_inputs, ) diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index d057d46816..6c5b997ee2 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -9,7 +9,7 @@ from collections.abc import Iterable, Sequence 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 ForwardResult, is_value_opaque_quantizer, register_custom_op @dataclasses.dataclass @@ -186,6 +187,48 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta): # Number of extra tensor outputs num_extra_outputs: int = 0 + # torch.compile support. An operation opts in by declaring the two arg + # containers and implementing the four compute classmethods below; the base + # class then registers its custom ops and drives them from op_forward / + # op_backward, so no operation writes that plumbing itself. + fwd_args_type: Optional[type] = None + bwd_args_type: Optional[type] = None + # Gradients returned by backward_compute: the input's, then any parameters'. + num_grad_inputs: int = 1 + # Forward kwargs this operation accepts, resolved into fwd_args_type like + # any other config. A kwarg carries no gradient and must not be mutated. + fwd_kwarg_names: tuple[str, ...] = () + # (forward_fn, backward_fn) pair, or None if the operation cannot be compiled. + compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + if cls.fwd_args_type is None or cls.bwd_args_type is None: + return + if getattr(cls.forward_compute, "__isabstractmethod__", False): + return + for name, arg_type in ( + ("fwd_args_type", cls.fwd_args_type), + ("bwd_args_type", cls.bwd_args_type), + ): + # The op schema is built from the container's fields, so this is the + # framework's actual requirement -- check it where it is declared. + if not dataclasses.is_dataclass(arg_type): + raise TypeError(f"{cls.__name__}.{name} must be a dataclass") + # One registration per class. The compute halves are bound here, so a + # subclass that only swaps kernels (the activations) still gets its own + # op without repeating any of this. + cls.compile_ops = register_custom_op( + op_name=cls.__name__.lower(), + fwd_arg_type=cls.fwd_args_type, + fwd_impl=cls.forward_compute, + fwd_fake_impl=cls.forward_fake, + bwd_arg_type=cls.bwd_args_type, + bwd_impl=cls.backward_compute, + bwd_fake_impl=cls.backward_fake, + num_grad_inputs=cls.num_grad_inputs, + ) + def __init__(self) -> None: super().__init__() @@ -274,6 +317,91 @@ def set_extra_output_channel( self._extra_output_to_caller[index] = output_to_caller return self + # ------------------------------------------------------------------ # + # 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) -> ForwardResult: + """Forward computation over explicit arguments. + + 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) -> ForwardResult: + """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__} without 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__} (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, + **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. ``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 + + 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 setup_context(self, ctx: OperationContext, args: Any, aux: tuple) -> None: + """Prepare backward state from the original arguments and fresh auxiliary tensors.""" + del args + ctx.save_for_backward(*aux) + @property def is_fused_op(self) -> bool: return False @@ -508,7 +636,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, @@ -520,6 +647,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 @@ -537,8 +668,67 @@ 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" + ) + 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, + ) + result = self.forward_compute(args) + if ctx.requires_grad: + self.setup_context(ctx, args, result.aux) + return result.output + + def compiled_op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + **kwargs: Any, + ) -> torch.Tensor: + """: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. ``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, aux = self.compile_ops[0](args) + if ctx.requires_grad: + self.setup_context(ctx, args, aux) + 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, @@ -546,6 +736,8 @@ def op_backward( ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: """Backward pass + Counterpart to the inherited :meth:`op_forward`. + Parameters ---------- ctx: OperationContext @@ -561,6 +753,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,