diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 68b1174474..4aa10fba50 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -4,6 +4,8 @@ import abc import contextlib +import dataclasses +from typing import Union import pytest import torch @@ -29,9 +31,14 @@ from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear +from transformer_engine.pytorch.ops.op import BasicOperation from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer -from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, Quantizer +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, @@ -1520,3 +1527,283 @@ def fn(inp): "Unexpected recompilation(s) across different batch sizes: " f"{recompile_count_after - recompile_count_baseline} recompile(s) detected" ) + + +# --------------------------------------------------------------------------- # +# transformer_engine.pytorch.ops under torch.compile +# --------------------------------------------------------------------------- # + + +@dataclasses.dataclass(slots=True) +class _ScaleFwdArgs: + """Flat, ``self``-free inputs to the test operation's forward.""" + + input_: torch.Tensor + scale: torch.Tensor + + +@dataclasses.dataclass(slots=True) +class _ScaleBwdArgs: + """Flat inputs to the test operation's backward.""" + + grad_output: torch.Tensor = None + saved_input: torch.Tensor = None + scale: torch.Tensor = None + + +class _ScaleOp(BasicOperation): + """Test-only operation: multiply by a learnable scalar. + + Exists so the fuser's compiled path can be exercised without depending on + which real operations happen to declare their compute halves. It is the + smallest operation that still has a parameter gradient and a saved tensor. + """ + + fwd_args_type = _ScaleFwdArgs + bwd_args_type = _ScaleBwdArgs + num_grad_inputs = 2 # grad input, grad scale + + def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> None: + super().__init__() + self.scale = torch.nn.Parameter(torch.full((), 2.0, device=device, dtype=dtype)) + + @classmethod + def forward_compute(cls, args): + return args.input_ * args.scale, (), {} + + @classmethod + def forward_fake(cls, args): + x = args.input_ + return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), (), {} + + @classmethod + def backward_compute(cls, args): + dy = args.grad_output + return dy * args.scale, (dy * args.saved_input).sum() + + @classmethod + def backward_fake(cls, args): + dy = args.grad_output + return ( + TensorSpec(shape=tuple(dy.shape), dtype=dy.dtype, device=dy.device), + TensorSpec(shape=(), dtype=dy.dtype, device=dy.device), + ) + + def saved_for_backward(self, saved, input_): + # The forward produces no distinct tensor for its input, and a custom op + # may not return one of its own inputs. + del saved + return (input_,) + + def resolve_fwd_args( + self, + input_, + *, + requires_grad, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=None, + ): + del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer + return _ScaleFwdArgs(input_=input_, scale=self.scale) + + def resolve_bwd_args(self, ctx, grad_output): + (x,) = ctx.saved_tensors + return _ScaleBwdArgs(grad_output=grad_output, saved_input=x, scale=self.scale) + + +@dataclasses.dataclass(slots=True) +class _ScaleKwargsFwdArgs: + """Flat inputs to the kwarg-taking test operation's forward.""" + + input_: torch.Tensor + scale: torch.Tensor + extra_scale: float + offset: Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclasses.dataclass(slots=True) +class _ScaleKwargsBwdArgs: + """Flat inputs to the kwarg-taking test operation's backward.""" + + grad_output: torch.Tensor = None + saved_input: torch.Tensor = None + scale: torch.Tensor = None + extra_scale: float = 1.0 + + +class _ScaleWithKwargsOp(BasicOperation): + """Test-only operation taking forward kwargs: a value and a tensor. + + ``offset`` is declared as tensor-or-quantized, so a quantized kwarg crosses + the op boundary as its inner buffers. Neither kwarg carries a gradient -- + that is what "read-only" means here. + """ + + fwd_args_type = _ScaleKwargsFwdArgs + bwd_args_type = _ScaleKwargsBwdArgs + num_grad_inputs = 2 # grad input, grad scale + fwd_kwarg_names = ("extra_scale", "offset") + + def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> None: + super().__init__() + self.scale = torch.nn.Parameter(torch.full((), 2.0, device=device, dtype=dtype)) + + @classmethod + def forward_compute(cls, args): + offset = args.offset + if isinstance(offset, QuantizedTensor): + offset = offset.dequantize() + out = args.input_ * args.scale * args.extra_scale + offset + return out, (), {"extra_scale": args.extra_scale} + + @classmethod + def forward_fake(cls, args): + x = args.input_ + return ( + TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), + (), + {"extra_scale": args.extra_scale}, + ) + + @classmethod + def backward_compute(cls, args): + dy = args.grad_output + return ( + dy * args.scale * args.extra_scale, + (dy * args.saved_input).sum() * args.extra_scale, + ) + + @classmethod + def backward_fake(cls, args): + dy = args.grad_output + return ( + TensorSpec(shape=tuple(dy.shape), dtype=dy.dtype, device=dy.device), + TensorSpec(shape=(), dtype=dy.dtype, device=dy.device), + ) + + def saved_for_backward(self, saved, input_): + del saved + return (input_,) + + def resolve_fwd_args( + self, + input_, + *, + requires_grad, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=None, + extra_scale=1.0, + offset=None, + ): + del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer + if offset is None: + offset = torch.zeros((), device=input_.device, dtype=input_.dtype) + return _ScaleKwargsFwdArgs( + input_=input_, + scale=self.scale, + extra_scale=extra_scale, + offset=offset, + ) + + def resolve_bwd_args(self, ctx, grad_output): + (x,) = ctx.saved_tensors + return _ScaleKwargsBwdArgs( + grad_output=grad_output, + saved_input=x, + scale=self.scale, + extra_scale=ctx.extra_scale, + ) + + +def _assert_sequential_matches_eager(make_model, base, op_kwargs_seq=(None,)): + """Run a Sequential eagerly and compiled on identical inputs; compare both + the output and every parameter gradient. + + 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_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(): + """Forward kwargs reach the operation through its custom op. + + Covers both kinds at once: a value, which Dynamo guards on -- hence the + second call with a different one -- and a tensor, quantized here, which + crosses the op boundary as its inner buffers. + """ + torch._dynamo.reset() + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device=torch.device("cuda"), + ) + offset = quantizer(torch.randn(64, dtype=torch.bfloat16, device="cuda")) + + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + _assert_sequential_matches_eager( + lambda: te.ops.Sequential(_ScaleWithKwargsOp()), + base, + op_kwargs_seq=( + {0: {"extra_scale": 3.0, "offset": offset}}, + {0: {"extra_scale": 5.0, "offset": offset}}, + ), + ) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 3598e54daa..1382cc1c92 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,7 +6,7 @@ from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer from .tensor_spec import TensorSpec, to_tensor_spec -from .custom_op import register_custom_op +from .custom_op import register_custom_op, register_custom_op_with_autograd __all__ = [ "register_value_opaque_quantizer", @@ -14,4 +14,5 @@ "TensorSpec", "to_tensor_spec", "register_custom_op", + "register_custom_op_with_autograd", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index b529deb75a..5fb1c4612e 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -6,7 +6,7 @@ Turns a TE module's eager forward/backward into ``torch.library`` custom ops so ``torch.compile(fullgraph=True)`` traces them as single graph nodes -- no graph -break into the eager ``autograd.Function``. ``register_custom_op`` is the entry +break into the eager ``autograd.Function``. ``register_custom_op_with_autograd`` is the entry point (its docstring documents the per-callable contract); ``module/linear.py`` is the first user. Internal framework API -- exported from ``transformer_engine.pytorch.dynamo``, not re-exported at the top level. @@ -50,7 +50,7 @@ only when its value is trivial (``None`` / all-``None``) at call time. What runs where. Each op registers a data-free fake (``register_fake``) so it -traces under ``torch.compile`` without allocating. ``register_custom_op`` returns +traces under ``torch.compile`` without allocating. ``register_custom_op_with_autograd`` returns ``forward_fn`` -- the drop-in for the eager ``autograd.Function.apply``. A forward call through it: @@ -1036,7 +1036,267 @@ def _split_fwd_fake_result( # --------------------------------------------------------------------------- # -# Op registration +# Two-tier op pair: the registration both public entry points build on +# --------------------------------------------------------------------------- # + + +@dataclasses.dataclass +class _OpPair: + """One registered forward/backward pair, and what a caller needs to drive it.""" + + fwd_adapters: List[_Adapter] + bwd_adapters: List[_Adapter] + fwd_arg_names: List[str] + bwd_arg_names: List[str] + fwd_tensor_field_names: List[str] + bwd_tensor_field_names: List[str] + base_fwd_def: Any + base_fwd_op: Any + base_bwd_op: Any + wrapper_fwd_def: Any + wrapper_fwd_op: Any + wrapper_bwd_op: Any + + +def _register_two_tier_pair( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: int, +) -> _OpPair: + """Define an operation's forward and backward as two-tier custom ops. + + Everything that is common to :func:`register_custom_op_with_autograd` and + :func:`register_custom_op`: schemas from the argument + containers, the base kernels, the wrapper ops that flatten + ``QuantizedTensor`` subclass inputs, and the passthrough registrations. + Autograd is deliberately not touched here -- that is what the two entry + points differ on. + """ + subclass_list = _all_quantized_tensor_subclasses() + + fwd_adapters = _get_adapters(fwd_arg_type) + bwd_adapters = _get_adapters(bwd_arg_type) + fwd_tensor_field_names = _tensor_field_names(fwd_adapters) + bwd_tensor_field_names = _tensor_field_names(bwd_adapters) + + fwd_schema_args, fwd_arg_names = _build_schema(fwd_adapters) + bwd_schema_args, bwd_arg_names = _build_schema(bwd_adapters) + fwd_schema = f"{fwd_schema_args} -> Tensor[]" + bwd_schema = f"{bwd_schema_args} -> Tensor[]" + + wrapper_fwd_name = op_name + wrapper_bwd_name = f"{op_name}_backward" + base_fwd_name = f"{op_name}_base" + base_bwd_name = f"{wrapper_bwd_name}_base" + base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" + + base_fwd_def = _register_kernel( + op_name=base_fwd_name, + schema_str=fwd_schema, + arg_type=fwd_arg_type, + arg_names=fwd_arg_names, + adapters=fwd_adapters, + tensor_field_names=fwd_tensor_field_names, + impl=fwd_impl, + fake_impl=fwd_fake_impl, + format_result=_format_fwd_result, + ) + _register_kernel( + op_name=base_bwd_name, + schema_str=bwd_schema, + arg_type=bwd_arg_type, + arg_names=bwd_arg_names, + adapters=bwd_adapters, + tensor_field_names=bwd_tensor_field_names, + impl=bwd_impl, + fake_impl=bwd_fake_impl, + format_result=lambda g: _format_bwd_result(g, num_grad_inputs, base_bwd_qualname), + ) + + base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) + base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) + + wrapper_fwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_fwd_name, + schema_str=fwd_schema, + base_op=base_fwd_op, + adapters=fwd_adapters, + ) + wrapper_bwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_bwd_name, + schema_str=bwd_schema, + base_op=base_bwd_op, + adapters=bwd_adapters, + ) + wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) + wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) + + fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_adapters) + bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_adapters) + + def _fwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) + return base_fwd_op(*new_args) + + def _bwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) + return base_bwd_op(*new_args) + + for sub in subclass_list: + wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) + wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule) + + for op in (wrapper_fwd_op, wrapper_bwd_op, base_fwd_op, base_bwd_op): + _quantized_tensor_passthrough_ops.add(op.default) + + return _OpPair( + fwd_adapters=fwd_adapters, + bwd_adapters=bwd_adapters, + fwd_arg_names=fwd_arg_names, + bwd_arg_names=bwd_arg_names, + fwd_tensor_field_names=fwd_tensor_field_names, + bwd_tensor_field_names=bwd_tensor_field_names, + base_fwd_def=base_fwd_def, + base_fwd_op=base_fwd_op, + base_bwd_op=base_bwd_op, + wrapper_fwd_def=wrapper_fwd_def, + wrapper_fwd_op=wrapper_fwd_op, + wrapper_bwd_op=wrapper_bwd_op, + ) + + +def _reassemble(specs: List[Any], payload: List[torch.Tensor], cursor: int = 0): + """Rebuild the values described by ``specs`` from ``payload[cursor:]``.""" + out: List[Any] = [] + for spec in specs: + n = _spec_slot_count(spec) + chunk = [_decode_none(t) for t in payload[cursor : cursor + n]] + cursor += n + out.append(_spec_reassemble(spec, chunk)) + return out, cursor + + +# --------------------------------------------------------------------------- # +# Op registration: the forward/backward pair, and the autograd-wired variant +# --------------------------------------------------------------------------- # + + +def register_custom_op( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: int, +) -> Optional[Tuple[Callable[[Any], Any], Callable[[Any], Any]]]: + """Register an op's forward and backward as two independent custom ops. + + Autograd is the caller's: it decides how the two are wired, which is what + lets a pipeline-level ``torch.autograd.Function`` -- traced by Dynamo as a + higher-order op -- group the forward and backward passes differently, as + ``ops.OperationFuser`` does. :func:`register_custom_op_with_autograd` builds + on this and wires them the usual way instead. + + Both ops are two-tier, so ``QuantizedTensor`` subclass inputs pass through + without dequantization. + + Contracts, mirroring :func:`register_custom_op_with_autograd`: + + * ``fwd_impl(fwd_args) -> (*user_outputs, tensors_to_save, ctx_attrs)`` + * ``fwd_fake_impl`` -- its data-free twin over :class:`TensorSpec` + * ``bwd_impl(bwd_args) -> tuple`` of ``num_grad_inputs`` gradients + * ``bwd_fake_impl`` -- its data-free twin + + Returns ``(forward_fn, backward_fn)``: + + * ``forward_fn(fwd_args) -> (outputs, saved_tensors, ctx_attrs)`` -- + ``outputs`` is a single value or a tuple, mirroring ``fwd_impl``'s user + outputs; ``saved_tensors`` is the reassembled ``tensors_to_save`` tuple, + which the caller is expected to persist (e.g. ``ctx.save_for_backward``). + * ``backward_fn(bwd_args) -> tuple`` of gradients. + + Returns ``None`` if registration fails (warned once), so callers can fall + back to eager rather than breaking import. + """ + try: + return _register_custom_op_impl( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_arg_type=bwd_arg_type, + bwd_impl=bwd_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=num_grad_inputs, + ) + except (ImportError, AttributeError, RuntimeError, TypeError) as e: + warn_compile_unsupported( + f"could not register the autograd-free custom ops '{op_name}' ({type(e).__name__}: {e})" + ) + return None + + +def _register_custom_op_impl( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: int, +) -> Tuple[Callable[[Any], Any], Callable[[Any], Any]]: + """Body of :func:`register_custom_op`; see it for semantics.""" + pair = _register_two_tier_pair( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_arg_type=bwd_arg_type, + bwd_impl=bwd_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=num_grad_inputs, + ) + + def forward_fn(fwd_args): + spec_obj = _spec_view(fwd_args, pair.fwd_tensor_field_names) + user_specs, saved_specs, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) + kwargs = _args_to_slots(fwd_args, pair.fwd_adapters) + payload = pair.wrapper_fwd_op(*[kwargs[name] for name in pair.fwd_arg_names]) + + outputs, cursor = _reassemble(user_specs, payload) + saved, _ = _reassemble(saved_specs, payload, cursor) + return (outputs[0] if len(outputs) == 1 else tuple(outputs)), tuple(saved), ctx_attrs + + def backward_fn(bwd_args): + # Unlike the forward payload, each grad occupies exactly one slot + # (``_format_bwd_result`` materializes a TensorSpec grad), so there is + # nothing to reassemble. + kwargs = _args_to_slots(bwd_args, pair.bwd_adapters) + payload = pair.wrapper_bwd_op(*[kwargs[name] for name in pair.bwd_arg_names]) + return tuple(_decode_none(t) for t in payload) + + return forward_fn, backward_fn + + +# --------------------------------------------------------------------------- # +# Autograd-wired registration # --------------------------------------------------------------------------- # @@ -1257,7 +1517,10 @@ def _forward(*flat: Any) -> List[torch.Tensor]: return base_op(*new_args) op = torch.library.custom_op( - f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", _forward, mutates_args=(), schema=schema_str + f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", + _forward, + mutates_args=(), + schema=schema_str, ) op.register_fake(_forward) return op @@ -1277,7 +1540,7 @@ def _all_quantized_tensor_subclasses() -> List[type]: return found -def register_custom_op( +def register_custom_op_with_autograd( *, op_name: str, input_tensors_for_grad: List[str], @@ -1350,7 +1613,7 @@ def register_custom_op( ``torch.compile`` (a graph break) rather than breaking import. """ try: - return _register_custom_op_impl( + return _register_custom_op_with_autograd_impl( op_name=op_name, input_tensors_for_grad=input_tensors_for_grad, fwd_arg_type=fwd_arg_type, @@ -1368,7 +1631,7 @@ def register_custom_op( return None -def _register_custom_op_impl( +def _register_custom_op_with_autograd_impl( *, op_name: str, input_tensors_for_grad: List[str], @@ -1380,7 +1643,7 @@ def _register_custom_op_impl( fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], ) -> Callable[..., Any]: - """Body of :func:`register_custom_op`; see it for semantics.""" + """Body of :func:`register_custom_op_with_autograd`; see it for semantics.""" # Existence check at the API boundary: every ``input_tensors_for_grad`` name # must be an actual field of ``fwd_arg_type`` (differentiability -- whether # that field can carry a gradient -- is checked later in @@ -1390,124 +1653,42 @@ def _register_custom_op_impl( if missing: raise ValueError(f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}") - wrapper_fwd_name = op_name - wrapper_bwd_name = f"{op_name}_backward" - base_fwd_name = f"{op_name}_base" - base_bwd_name = f"{wrapper_bwd_name}_base" - subclass_list = _all_quantized_tensor_subclasses() - - fwd_adapters = _get_adapters(fwd_arg_type) - bwd_adapters = _get_adapters(backward_arg_type) - fwd_tensor_field_names = _tensor_field_names(fwd_adapters) - bwd_tensor_field_names = _tensor_field_names(bwd_adapters) - - fwd_schema_args, fwd_arg_names = _build_schema(fwd_adapters) - bwd_schema_args, bwd_arg_names = _build_schema(bwd_adapters) - - num_grad_inputs = len(input_tensors_for_grad) - slot_count, grad_targets = _resolve_grad_targets(fwd_adapters, input_tensors_for_grad) - - fwd_schema = f"{fwd_schema_args} -> Tensor[]" - bwd_schema = f"{bwd_schema_args} -> Tensor[]" - - base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" - - base_fwd_def = _register_kernel( - op_name=base_fwd_name, - schema_str=fwd_schema, - arg_type=fwd_arg_type, - arg_names=fwd_arg_names, - adapters=fwd_adapters, - tensor_field_names=fwd_tensor_field_names, - impl=fwd_impl, - fake_impl=fwd_fake_impl, - format_result=_format_fwd_result, - ) - _register_kernel( - op_name=base_bwd_name, - schema_str=bwd_schema, - arg_type=backward_arg_type, - arg_names=bwd_arg_names, - adapters=bwd_adapters, - tensor_field_names=bwd_tensor_field_names, - impl=backward_impl, - fake_impl=bwd_fake_impl, - format_result=lambda g: _format_bwd_result(g, num_grad_inputs, base_bwd_qualname), - ) - - base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) - base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) - - wrapper_fwd_def = _register_wrapper_op( - wrapper_op_name=wrapper_fwd_name, - schema_str=fwd_schema, - base_op=base_fwd_op, - adapters=fwd_adapters, - ) - wrapper_bwd_def = _register_wrapper_op( - wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op + pair = _register_two_tier_pair( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_arg_type=backward_arg_type, + bwd_impl=backward_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=len(input_tensors_for_grad), ) + slot_count, grad_targets = _resolve_grad_targets(pair.fwd_adapters, input_tensors_for_grad) autograd_common = { "fwd_arg_type": fwd_arg_type, - "fwd_arg_names": fwd_arg_names, - "fwd_adapters": fwd_adapters, - "fwd_tensor_field_names": fwd_tensor_field_names, - "bwd_arg_names": bwd_arg_names, - "bwd_adapters": bwd_adapters, + "fwd_arg_names": pair.fwd_arg_names, + "fwd_adapters": pair.fwd_adapters, + "fwd_tensor_field_names": pair.fwd_tensor_field_names, + "bwd_arg_names": pair.bwd_arg_names, + "bwd_adapters": pair.bwd_adapters, "slot_count": slot_count, "grad_targets": grad_targets, "setup_context_user": setup_context, "backward_obj_type": backward_arg_type, "fwd_fake_impl": fwd_fake_impl, } - wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) - wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) - - _register_autograd_for_op(fwd_op=base_fwd_def, bwd_op=base_bwd_op, **autograd_common) - _register_autograd_for_op(fwd_op=wrapper_fwd_def, bwd_op=wrapper_bwd_op, **autograd_common) - - fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_adapters) - bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_adapters) - - def _fwd_rule(mode, func, types, args, kwargs): - del mode, func, types, kwargs - new_args = list(args) - for sub in subclass_list: - _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) - return base_fwd_op(*new_args) - - def _bwd_rule(mode, func, types, args, kwargs): - del mode, func, types, kwargs - new_args = list(args) - for sub in subclass_list: - _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) - return base_bwd_op(*new_args) - - for sub in subclass_list: - wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) - wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule) - - _quantized_tensor_passthrough_ops.add(wrapper_fwd_op.default) - _quantized_tensor_passthrough_ops.add(wrapper_bwd_op.default) - _quantized_tensor_passthrough_ops.add(base_fwd_op.default) - _quantized_tensor_passthrough_ops.add(base_bwd_op.default) + _register_autograd_for_op(fwd_op=pair.base_fwd_def, bwd_op=pair.base_bwd_op, **autograd_common) + _register_autograd_for_op( + fwd_op=pair.wrapper_fwd_def, bwd_op=pair.wrapper_bwd_op, **autograd_common + ) def forward_fn(fwd_args): - spec_obj = _spec_view(fwd_args, fwd_tensor_field_names) - user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) - kwargs = _args_to_slots(fwd_args, fwd_adapters) - flat_in = [kwargs[name] for name in fwd_arg_names] - result = wrapper_fwd_op(*flat_in) - - cursor = 0 - outputs: List[Any] = [] - for spec in user_fakes: - n = _spec_slot_count(spec) - chunk = [_decode_none(t) for t in result[cursor : cursor + n]] - cursor += n - outputs.append(_spec_reassemble(spec, chunk)) - + spec_obj = _spec_view(fwd_args, pair.fwd_tensor_field_names) + user_specs, _saved_specs, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) + kwargs = _args_to_slots(fwd_args, pair.fwd_adapters) + payload = pair.wrapper_fwd_op(*[kwargs[name] for name in pair.fwd_arg_names]) + outputs, _ = _reassemble(user_specs, payload) if len(outputs) == 1: return outputs[0] return tuple(outputs) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 219263773d..0e4eb634e0 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -83,7 +83,7 @@ prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorSpec, register_custom_op, is_value_opaque_quantizer +from ..dynamo import TensorSpec, register_custom_op_with_autograd, is_value_opaque_quantizer from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -1753,7 +1753,7 @@ def _linear_backward_impl_fake( # Custom op used under ``torch.compile``. -_linear_op = register_custom_op( +_linear_op = register_custom_op_with_autograd( op_name="linear", input_tensors_for_grad=["weight", "inp", "bias"], fwd_arg_type=LinearFwdArgs, diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 09ffb004dd..89d3d45744 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -6,6 +6,7 @@ from __future__ import annotations from collections.abc import Callable, Iterable, Sequence +import copy import itertools from typing import Any, Optional, TypeAlias @@ -13,6 +14,7 @@ from ..quantization import FP8GlobalStateManager, Recipe, DelayedScaling from ..quantized_tensor import prepare_for_saving, restore_from_func_ctx +from ..dynamo.quantizer_opaque import warn_compile_unsupported from .op import ( BasicOperation, FusibleOperation, @@ -66,6 +68,7 @@ def forward( fuser: OperationFuser, basic_op_kwargs: list[dict[str, Any]], set_output_requires_grad: bool, + use_compiled: bool, *params_and_extra_inputs: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Forward pass @@ -82,6 +85,9 @@ def forward( Keyword arguments to BasicOperation set_output_requires_grad: bool Whether to set ``requires_grad`` flags on returned tensors + use_compiled: bool + Whether to call the operations' custom ops instead of their eager + implementations. Decided once per group by ``OperationFuser``. *params_and_extra_inputs: torch.Tensor Other tensor inputs to include in autograd graph. Consists of parameter tensors, followed by extra operation inputs. @@ -98,9 +104,14 @@ def forward( # Operation autograd contexts basic_op_ctxs = [OperationContext() for _ in range(fuser._num_basic_ops)] - # Mark input tensors as not deletable in backward - for tensor in (input_,) + params_and_extra_inputs: - tensor._do_not_clear = True + # Mark input tensors as not deletable in backward. Skipped whenever this + # is being traced -- not merely when the custom ops are used: these + # tensors are created outside this function, and a higher-order op may + # not mutate anything from an enclosing scope. Under fullgraph there is + # no falling back out of the graph, so the constraint holds either way. + if not torch.compiler.is_compiling(): + for tensor in (input_,) + params_and_extra_inputs: + tensor._do_not_clear = True # Unflatten list of parameters and extra tensor inputs extra_inputs = params_and_extra_inputs[-fuser.num_extra_inputs :] @@ -131,14 +142,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_compiled: + x = op.compiled_op_forward( + basic_op_ctxs[basic_op_idxs[0]], + x, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + **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], + ) for idx, ys in zip(basic_op_idxs, fused_op_extra_outputs): for y in ys: if set_output_requires_grad: @@ -176,9 +197,13 @@ def forward( func_ctx.save_for_backward(*tensors_to_save) func_ctx.tensor_objects = tensor_objects - # Whether to perform recipe update in backward pass + # Whether to perform recipe update in backward pass. Skipped under + # compile: this reads and flips global FP8 state, and delayed + # scaling -- the only recipe it serves -- is gated out anyway. is_first_module = False - if fuser.first_op_requiring_backward < fuser._num_basic_ops: + if not torch.compiler.is_compiling() and ( + fuser.first_op_requiring_backward < fuser._num_basic_ops + ): is_first_module = FP8GlobalStateManager.is_first_fp8_module() # Other context @@ -189,12 +214,17 @@ def forward( func_ctx.num_extra_inputs = fuser.num_extra_inputs func_ctx.num_extra_outputs = len(extra_outputs_flat) func_ctx.is_first_module = is_first_module + func_ctx.use_compiled = use_compiled - # Mark output tensors as not deletable in backward - for tensor in [x] + extra_outputs_flat: - tensor._do_not_clear = True + # Mark output tensors as not deletable in backward (eager only; see above) + if not torch.compiler.is_compiling(): + for tensor in [x] + extra_outputs_flat: + tensor._do_not_clear = True - if set_output_requires_grad: + # Autograd marks the outputs of an ``apply`` itself, so this is only + # needed on the eager path -- and AOTAutograd's functionalization drops + # a requires_grad_() applied to a graph output anyway. + if set_output_requires_grad and not torch.compiler.is_compiling(): x.requires_grad_(fuser.first_op_requiring_backward < fuser._num_basic_ops) if extra_outputs_flat: @@ -219,7 +249,12 @@ def backward( # Restore saved tensors saved_tensors = restore_from_func_ctx(func_ctx) - # Unflatten list of saved tensors + # Unflatten list of saved tensors. Under compile the contexts were + # created in the forward, which is a different subgraph, so writing to + # them here would be a side effect on an enclosing scope; copy them into + # this one instead. The copy carries the attributes the forward set. + if torch.compiler.is_compiling(): + basic_op_ctxs = [copy.copy(ctx) for ctx in basic_op_ctxs] for ctx in basic_op_ctxs: ctx.saved_tensors = saved_tensors[slice(*ctx._saved_tensors_range)] ctx._saved_tensors_range = None @@ -248,14 +283,22 @@ def backward( # Backward op grad_extra_outputs = [basic_op_grad_extra_outputs[idx] for idx in basic_op_idxs] - dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( - [basic_op_ctxs[idx] for idx in basic_op_idxs], - dx, - basic_op_grad_extra_outputs=grad_extra_outputs, - ) + if func_ctx.use_compiled: + dx, grad_params_one = op.compiled_op_backward(basic_op_ctxs[basic_op_idxs[0]], dx) + fused_op_grad_params = [grad_params_one] + fused_op_grad_extra_inputs = [()] + else: + dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( + [basic_op_ctxs[idx] for idx in basic_op_idxs], + dx, + basic_op_grad_extra_outputs=grad_extra_outputs, + ) for idx, dparams in zip(basic_op_idxs, fused_op_grad_params): grad_params[idx] = dparams - basic_op_ctxs[idx].saved_tensors = None + # Dropping the reference frees the activation early; on the + # compiled path the graph owns that lifetime instead. + if not torch.compiler.is_compiling(): + basic_op_ctxs[idx].saved_tensors = None for idx, dxs in zip(basic_op_idxs, fused_op_grad_extra_inputs): grad_extra_inputs[idx] = dxs @@ -299,6 +342,7 @@ def backward( None, # fuser None, # basic_op_kwargs None, # set_output_requires_grad + None, # use_compiled *grad_params_flat, *grad_extra_inputs_flat, ) @@ -501,6 +545,43 @@ def maybe_fuse_ops( else: self._last_amax_history_len = 0 + def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) -> Optional[str]: + """Why this group may not run through its operations' custom ops.""" + if len(self._forward_ops) != self._num_basic_ops: + # A fused op covers several basic ops; only single-op groups so far. + return "fused operations are not supported yet" + for op, kwargs in zip(self._basic_ops, basic_op_kwargs): + # A kwarg an operation declares is resolved into its args container + # like any other config. Anything else -- notably the preallocated + # buffers of the grouped operations -- is written to by the op, and a + # custom op may not mutate a tensor from an enclosing scope. + unsupported = sorted(name for name in kwargs if name not in op.fwd_kwarg_names) + if unsupported: + return f"{type(op).__name__} does not support keyword arguments {unsupported}" + for op in self._basic_ops: + reason = op.compile_unsupported_reason() + if reason is not None: + return reason + return None + + def _use_compiled(self, basic_op_kwargs: list[dict[str, Any]]) -> bool: + """Whether this group runs through its operations' custom ops. + + Decided once for the whole group: a pipeline compiles as a whole, so one + unsupported operation sends all of them to eager. + + The reason is reported from the eager path only -- ``warnings.warn`` is + not traceable, so warning from inside the traced region would itself + break the graph. A configuration that is never run eagerly therefore + falls back silently. + """ + reason = self._compile_unsupported_reason(basic_op_kwargs) + if reason is None: + return torch.compiler.is_compiling() + if not torch.compiler.is_compiling(): + warn_compile_unsupported(f"running {type(self).__name__} eagerly: {reason}") + return False + def __call__( self, input: torch.Tensor, # pylint: disable=redefined-builtin @@ -541,11 +622,14 @@ def __call__( # Note: We call forward directly when is_grad_enabled=False, # which can expose non-leaf tensors to the inner ops. Avoid # problems in this case by passing set_output_requires_grad=False. + use_compiled = self._use_compiled(basic_op_kwargs) + args = ( input, self, basic_op_kwargs, is_grad_enabled, # set_output_requires_grad + use_compiled, *self._flat_basic_op_params, *extra_inputs, ) diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 5106ec9e0a..c3af65a831 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -9,7 +9,7 @@ from collections.abc import Iterable import dataclasses import pickle -from typing import Any, Optional +from typing import Any, Callable, Optional import torch @@ -22,6 +22,7 @@ autocast, ) from ..tensor import Quantizer +from ..dynamo import is_value_opaque_quantizer, register_custom_op @dataclasses.dataclass @@ -184,6 +185,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__() @@ -191,6 +234,96 @@ def __init__(self) -> None: self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None + # ------------------------------------------------------------------ # + # Compute halves. Classmethods, not free functions: they belong to the + # operation, and binding to the class is what lets a family of operations + # share one implementation while dispatching to per-class kernels. + # ------------------------------------------------------------------ # + + @classmethod + def forward_compute(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]: + """Pure forward: ``(output, tensors_to_save, ctx_attrs)``. + + Takes everything through ``args``; must not read ``self`` or global + state, both of which are invisible to the compiler at this point. + """ + raise NotImplementedError + + @classmethod + def forward_fake(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]: + """Allocation-free twin of :meth:`forward_compute` over ``TensorSpec``. + + Runs as a meta kernel, outside the traced frame, and more than once per + compile, so it must be a pure function of ``args`` -- a read of global + state here is unguarded and can silently disagree with the real impl. + """ + raise NotImplementedError + + @classmethod + def backward_compute(cls, args: Any) -> tuple: + """Pure backward: ``num_grad_inputs`` gradients.""" + raise NotImplementedError + + @classmethod + def backward_fake(cls, args: Any) -> tuple: + """Allocation-free twin of :meth:`backward_compute`.""" + raise NotImplementedError + + def compile_unsupported_reason(self) -> Optional[str]: + """Why this operation cannot go through its custom op, or ``None``. + + Asked per operation, but acted on per fuser group: a pipeline compiles + as a whole, so one unsupported operation sends the whole group to eager. + Recipe-level limits are not checked here -- they belong to whoever reads + the recipe, which is the fuser. + """ + if self.compile_ops is None: + return f"{self.__class__.__name__} does not implement the compute halves" + for mode in ("forward", "backward"): + for index in range(self.num_quantizers(mode)): + quantizer = self.get_quantizer(mode, index) + if quantizer is not None and not is_value_opaque_quantizer(quantizer): + # Delayed scaling holds live scale/amax tensors, so its + # quantizer cannot be specialized on and would be baked into + # the graph as a stale constant. + return ( + f"{type(quantizer).__name__} is not a torch.compile value-opaque quantizer" + ) + return None + + def resolve_fwd_args( + self, + input_: torch.Tensor, + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + **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 saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple: + """Tensors to persist, given what the forward handed back. + + An operation whose backward needs its input but whose forward does not + produce a distinct tensor for it overrides this; a custom op may not + return one of its own inputs. + """ + del input_ + return saved + @property def is_fused_op(self) -> bool: return False @@ -425,7 +558,6 @@ def _load_fp8_metas(self, fp8_metas: Optional[dict[str, Any]]) -> None: self._fp8_metas[mode][fp8_meta_key].scale.copy_(scale) self._fp8_metas[mode][fp8_meta_key].amax_history.copy_(amax_history) - @abc.abstractmethod def op_forward( self, ctx: OperationContext, @@ -437,6 +569,10 @@ def op_forward( ) -> torch.Tensor: """Forward pass + Operations that declare the compute halves inherit this: it resolves the + arguments, runs the forward, and records what the backward will need. The + rest override it. + Parameters ---------- ctx: OperationContext @@ -454,8 +590,71 @@ 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, saved, ctx_attrs = self.forward_compute(args) + if ctx.requires_grad: + ctx.save_for_backward(*self.saved_for_backward(saved, input_)) + for name, value in ctx_attrs.items(): + setattr(ctx, name, value) + return output + + def compiled_op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + **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, saved, ctx_attrs = self.compile_ops[0](args) + if ctx.requires_grad: + ctx.save_for_backward(*self.saved_for_backward(saved, input_)) + for name, value in ctx_attrs.items(): + setattr(ctx, name, value) + return output + + def compiled_op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: + """:meth:`op_backward` routed through this operation's custom op.""" + grads = self.compile_ops[1](self.resolve_bwd_args(ctx, grad_output)) + grad_input = grads[0] + if grad_input is None: + grad_input = grad_output + return grad_input, tuple(grads[1:]) - @abc.abstractmethod def op_backward( self, ctx: OperationContext, @@ -463,6 +662,8 @@ def op_backward( ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: """Backward pass + Counterpart to the inherited :meth:`op_forward`. + Parameters ---------- ctx: OperationContext @@ -478,6 +679,17 @@ def op_backward( Loss gradients w.r.t. parameters """ + if self.bwd_args_type is None: + raise NotImplementedError( + f"{self.__class__.__name__} implements neither op_backward nor the compute halves" + ) + grads = self.backward_compute(self.resolve_bwd_args(ctx, grad_output)) + grad_input = grads[0] + if grad_input is None: + # "The incoming gradient, unchanged": a custom op may not return one + # of its own inputs, so the compute half hands back None instead. + grad_input = grad_output + return grad_input, tuple(grads[1:]) def fuser_forward( self, diff --git a/transformer_engine/pytorch/ops/sequential.py b/transformer_engine/pytorch/ops/sequential.py index cb5dfecb9f..b8724ca460 100644 --- a/transformer_engine/pytorch/ops/sequential.py +++ b/transformer_engine/pytorch/ops/sequential.py @@ -179,9 +179,7 @@ def forward( or grouped MLP. """ - # Create module groups if needed - if self._module_groups is None: - self._module_groups = self._make_module_groups(self._modules.values()) + module_groups = self._get_module_groups() # Route op kwargs to each module group's basic ops group_op_kwargs = self._resolve_op_kwargs(op_kwargs) @@ -189,7 +187,7 @@ def forward( # Forward pass for each module group x = input extra_outputs: list[torch.Tensor] = [] - for group_idx, module_group in enumerate(self._module_groups): + for group_idx, module_group in enumerate(module_groups): if isinstance(module_group, OperationFuser): xs, extra_inputs = ( (x,) + extra_inputs[: module_group.num_extra_inputs], @@ -208,6 +206,17 @@ def forward( return (x,) + tuple(extra_outputs) return x + def _get_module_groups(self) -> list[OperationFuser | torch.nn.Module]: + """Module groups, built once. + + Kept out of the forward pass: building them constructs ``OperationFuser`` + and fused-operation objects, and an ``nn.Module`` cannot be constructed + inside a traced region. + """ + if self._module_groups is None: + self._module_groups = self._make_module_groups(self._modules.values()) + return self._module_groups + def _resolve_op_kwargs( self, op_kwargs: Optional[dict[torch.nn.Module | int, dict[str, Any]]],