diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index f8e09d5ce1..4342e1a330 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,11 +39,17 @@ 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.quantized_tensor import ( + QuantizedTensor, + QuantizedTensorStorage, + Quantizer, +) from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec from transformer_engine.pytorch import ( is_fp8_available, @@ -2339,3 +2347,320 @@ 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 + + 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 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 + + 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, () + + @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 * 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_multi_op_group_uses_eager_implementations(): + """A group of several operations is gated onto the eager implementations.""" + torch._dynamo.reset() + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + with pytest.warns(UserWarning, match="several operations"): + _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp(), _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") +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..042c209c1f 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,7 +6,11 @@ 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 ( + register_custom_op, + register_custom_op_with_autograd, + TensorOrQuantized, +) __all__ = [ "register_value_opaque_quantizer", @@ -14,5 +18,6 @@ "TensorSpec", "to_tensor_spec", "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..4fa0299f2d 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 an ``(output, aux)`` tuple; 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,7 @@ _TE_OP_NAMESPACE = "transformer_engine_compile" + # 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``. @@ -824,14 +830,16 @@ def _pack_fwd_result(result: Any) -> List[torch.Tensor]: return flat -def _pack_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List[torch.Tensor]: +def _pack_bwd_result( + grads: Any, num_grad_inputs: Optional[int], op_qualname: str +) -> List[torch.Tensor]: """Pack a backward-impl return tuple into the op's ``Tensor[]`` payload. - Each grad occupies exactly one slot (validated against ``num_grad_inputs``); - a :class:`TensorSpec` grad is materialized into a single tensor. + Each grad occupies exactly one slot (validated against ``num_grad_inputs`` + when given); a :class:`TensorSpec` grad is materialized into a single tensor. """ grads = list(grads) - if len(grads) != num_grad_inputs: + if num_grad_inputs is not None and len(grads) != num_grad_inputs: raise RuntimeError( f"{op_qualname} expected bwd_impl to return {num_grad_inputs} grads " f"(one per input_tensors_for_grad entry), got {len(grads)}" @@ -930,7 +938,7 @@ def _slice_user_grads( # --------------------------------------------------------------------------- # -# Op registration +# Op registration: base and wrapper ops, autograd wiring, one op # --------------------------------------------------------------------------- # @@ -1166,7 +1174,238 @@ def _all_quantized_tensor_subclasses() -> List[type]: return found +@dataclasses.dataclass(frozen=True) +class _RegisteredOp: + """One registered custom op: its arg plan and the base / wrapper definitions.""" + + plan: _ArgPlan + base_def: Any + base_op: Any + wrapper_def: Any + wrapper_op: Any + + def __call__(self, args: Any) -> List[torch.Tensor]: + """Pack the args dataclass into slots and call the wrapper op.""" + kwargs = self.plan.pack(args) + return self.wrapper_op(*[kwargs[name] for name in self.plan.slot_names]) + + +def _register_op( + *, + name: str, + arg_type: type, + impl: Callable[[Any], Any], + fake_impl: Callable[[Any], Any], + pack_result: Callable[[Any], List[torch.Tensor]], + flatten_in_body: bool, +) -> _RegisteredOp: + """Define one two-tier custom op: the base kernel, the wrapper op that lets + ``QuantizedTensor`` subclasses be inputs, and the passthrough registrations. + + ``flatten_in_body`` also flattens subclass inputs inside the wrapper body, + not only through the ``register_torch_dispatch`` rules. + """ + plan = _parse_arg_type(arg_type) + schema = f"{plan.schema_str} -> Tensor[]" + subclasses = _all_quantized_tensor_subclasses() + slot_offsets = plan.tensor_or_quantized_offsets() + namespace = getattr(torch.ops, _TE_OP_NAMESPACE) + + base_def = _register_base_op( + op_name=f"{name}_base", + schema_str=schema, + plan=plan, + impl=impl, + fake_impl=fake_impl, + pack_result=pack_result, + ) + base_op = getattr(namespace, f"{name}_base") + wrapper_def = _register_wrapper_op( + wrapper_op_name=name, + schema_str=schema, + base_op=base_op, + slot_offsets=slot_offsets if flatten_in_body else (), + subclasses=subclasses if flatten_in_body else (), + ) + wrapper_op = getattr(namespace, name) + + rule = _make_dispatch_rule(_make_slot_forwarder(base_op, slot_offsets, subclasses)) + for sub in subclasses: + wrapper_def.register_torch_dispatch(sub, rule) + _quantized_tensor_passthrough_ops.update((base_op.default, wrapper_op.default)) + + return _RegisteredOp( + plan=plan, + base_def=base_def, + base_op=base_op, + wrapper_def=wrapper_def, + wrapper_op=wrapper_op, + ) + + +def _register_forward_op( + *, name: str, arg_type: type, impl: Callable[[Any], Any], fake_impl: Callable[[Any], Any] +) -> _RegisteredOp: + return _register_op( + name=name, + arg_type=arg_type, + impl=impl, + fake_impl=fake_impl, + pack_result=_pack_fwd_result, + flatten_in_body=True, + ) + + +def _register_backward_op( + *, + name: str, + arg_type: type, + impl: Callable[[Any], Any], + fake_impl: Callable[[Any], Any], + num_grad_inputs: Optional[int], +) -> _RegisteredOp: + # Pass-through body: a subclass input reaches the base op through the + # dispatch rule, never through the wrapper body. + qualname = f"{_TE_OP_NAMESPACE}::{name}_base" + return _register_op( + name=name, + arg_type=arg_type, + impl=impl, + fake_impl=fake_impl, + pack_result=lambda g: _pack_bwd_result(g, num_grad_inputs, qualname), + flatten_in_body=False, + ) + + +def _run_forward( + fwd_op: _RegisteredOp, 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, fwd_op.plan.tensor_field_names()) + out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj)) + return out_plan, fwd_op(fwd_args) + + +# --------------------------------------------------------------------------- # +# Op registration: the autograd-free 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: Optional[int] = None, +) -> 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) -> (output, aux)`` -- ``aux`` is a tuple of fresh tensors + * ``fwd_fake_impl`` -- its data-free twin over :class:`TensorSpec` + * ``bwd_impl(bwd_args) -> tuple`` of gradients (``num_grad_inputs`` of them, + if given) + * ``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: Optional[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, tuple) or len(result) != 2: + raise TypeError( + "autograd-free fwd impl must return an (output, aux) tuple, got" + f" {type(result).__name__}" + ) + output, aux = result + return output, tuple(aux), None + + return wrapped + + adapted_fwd_fake_impl = adapt_forward(fwd_fake_impl) + fwd_op = _register_forward_op( + name=op_name, + arg_type=fwd_arg_type, + impl=adapt_forward(fwd_impl), + fake_impl=adapted_fwd_fake_impl, + ) + bwd_op = _register_backward_op( + name=f"{op_name}_backward", + arg_type=bwd_arg_type, + impl=bwd_impl, + fake_impl=bwd_fake_impl, + num_grad_inputs=num_grad_inputs, + ) + + def forward_fn(fwd_args): + out_plan, payload = _run_forward(fwd_op, 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. + return tuple(_decode_none(t) for t in bwd_op(bwd_args)) + + return forward_fn, backward_fn + + +def register_custom_op_with_autograd( *, op_name: str, input_tensors_for_grad: List[str], @@ -1231,7 +1470,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 +1488,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 +1500,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 +1510,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, + fwd_op = _register_forward_op( + name=op_name, arg_type=fwd_arg_type, impl=fwd_impl, fake_impl=fwd_fake_impl ) - _register_base_op( - op_name=base_bwd_name, - schema_str=bwd_schema, - plan=bwd_plan, + bwd_op = _register_backward_op( + name=f"{op_name}_backward", + arg_type=bwd_arg_type, 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 + num_grad_inputs=len(input_tensors_for_grad), ) autograd_common = { - "fwd_plan": fwd_plan, - "bwd_plan": bwd_plan, - "grad_targets": grad_targets, + "fwd_plan": fwd_op.plan, + "bwd_plan": bwd_op.plan, + "grad_targets": fwd_op.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) + _register_autograd_for_op(fwd_op=fwd_op.base_def, bwd_op=bwd_op.base_op, **autograd_common) + _register_autograd_for_op( + fwd_op=fwd_op.wrapper_def, bwd_op=bwd_op.wrapper_op, **autograd_common ) - _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) - - _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 = _run_forward(fwd_op, 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..e51160a15c 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,11 +362,16 @@ 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 @@ -392,6 +432,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 +732,56 @@ 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" + if self._num_basic_ops != 1: + return "a group of several operations" + 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 +824,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..f2b9377cea 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -8,8 +8,9 @@ import abc from collections.abc import Iterable, Sequence import dataclasses +import inspect import pickle -from typing import Any, Optional +from typing import Any, Callable, Optional import torch @@ -22,6 +23,11 @@ autocast, ) from ..tensor import Quantizer +from ..dynamo import is_value_opaque_quantizer, register_custom_op + +_FIXED_RESOLVE_FWD_KWARGS = frozenset( + ("requires_grad", "prev_op_grad_output_quantizer", "next_op_input_quantizer") +) @dataclasses.dataclass @@ -186,6 +192,54 @@ 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 + # Forward kwargs this operation accepts: the keyword-only parameters of its + # resolve_fwd_args beyond the fixed ones. 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") + params = inspect.signature(cls.resolve_fwd_args).parameters + if any(p.kind is p.VAR_KEYWORD for p in params.values()): + raise TypeError(f"{cls.__name__}.resolve_fwd_args must name its keyword arguments") + cls.fwd_kwarg_names = tuple( + name + for name, p in params.items() + if p.kind is p.KEYWORD_ONLY and name not in _FIXED_RESOLVE_FWD_KWARGS + ) + # 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, + ) + def __init__(self) -> None: super().__init__() @@ -274,6 +328,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) -> tuple[Any, tuple]: + """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) -> tuple[Any, tuple]: + """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: the input's gradient, then the parameters'.""" + 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 +647,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 +658,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 +679,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, + ) + output, aux = self.forward_compute(args) + if ctx.requires_grad: + self.setup_context(ctx, args, aux) + 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], + **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 +747,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 +764,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,