diff --git a/tests/pytorch/test_ops_custom_ops.py b/tests/pytorch/test_ops_custom_ops.py new file mode 100644 index 0000000000..b262e4d904 --- /dev/null +++ b/tests/pytorch/test_ops_custom_ops.py @@ -0,0 +1,187 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Per-op custom ops for ``transformer_engine.pytorch.ops``. + +Each fusible operation registers its forward and backward as independent custom +ops (``register_op_halves``) so a compiled pipeline can call them directly. The +tests here check each half in isolation -- numerics against the eager op, and +that the data-free fake agrees with the real impl slot for slot. + +The fake is what the compiler believes; if it disagrees with the real impl the +result is a silently misassembled tensor rather than an error, so the +conformance check runs for every op. +""" + +from __future__ import annotations + +from typing import Any, List, Optional, Tuple + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch.dynamo import TensorSpec +from transformer_engine.pytorch.ops.basic.bias import ( + BiasBwdArgs, + BiasFwdArgs, + _bias_backward_impl, + _bias_backward_impl_fake, + _bias_forward_impl, + _bias_forward_impl_fake, + _bias_ops, +) +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage + +# Test setup +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +_device = "cuda" + + +# --------------------------------------------------------------------------- # +# Conformance: the fake must describe what the real impl produces +# --------------------------------------------------------------------------- # + + +def _describe(value: Any) -> Optional[Tuple]: + """Structural fingerprint of a real tensor or of the spec describing it.""" + if value is None: + return None + if isinstance(value, TensorSpec): + quantizer = value.quantizer + return (tuple(value.shape), value.dtype, type(quantizer) if quantizer else None) + quantizer = getattr(value, "_quantizer", None) + if not isinstance(value, QuantizedTensorStorage): + quantizer = None + return (tuple(value.shape), value.dtype, type(quantizer) if quantizer else None) + + +def _describe_all(values: Any) -> List[Optional[Tuple]]: + if values is None: + return [] + if not isinstance(values, (tuple, list)): + values = (values,) + return [_describe(v) for v in values] + + +def assert_fwd_fake_matches_real(args: Any, impl, fake_impl) -> None: + """Run a forward impl and its fake on the same args; require agreement. + + Compares user outputs, saved tensors (count included -- the compiled path + slices a flat payload by the fake's saved-tensor list) and ``ctx_attrs`` + keys. + """ + real_out, real_saved, real_attrs = impl(args) + fake_out, fake_saved, fake_attrs = fake_impl(args) + + assert _describe_all(real_out) == _describe_all(fake_out), "forward outputs disagree" + assert _describe_all(real_saved) == _describe_all(fake_saved), "saved tensors disagree" + assert set(real_attrs) == set(fake_attrs), "ctx_attrs keys disagree" + + +def assert_bwd_fake_matches_real(args: Any, impl, fake_impl) -> None: + """Run a backward impl and its fake on the same args; require agreement.""" + real = impl(args) + fake = fake_impl(args) + assert _describe_all(real) == _describe_all(fake), "gradients disagree" + + +# --------------------------------------------------------------------------- # +# Bias +# --------------------------------------------------------------------------- # + + +def _bias_op(size: int, dtype: torch.dtype) -> te.ops.Bias: + op = te.ops.Bias(size, device=_device, dtype=dtype) + with torch.no_grad(): + op.bias.copy_(torch.randn_like(op.bias)) + return op + + +def _bias_fwd_args(shape, size: int, dtype: torch.dtype) -> BiasFwdArgs: + op = _bias_op(size, dtype) + x = torch.randn(*shape, device=_device, dtype=dtype) + return op.resolve_fwd_args(x, requires_grad=True, prev_op_grad_output_quantizer=None) + + +@_cuda +@pytest.mark.parametrize("shape", [(16, 32), (2, 8, 32), (32,)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_bias_fake_conformance(shape, dtype) -> None: + args = _bias_fwd_args(shape, shape[-1], dtype) + assert_fwd_fake_matches_real(args, _bias_forward_impl, _bias_forward_impl_fake) + + dy = torch.randn(*shape, device=_device, dtype=dtype) + bwd_args = BiasBwdArgs(grad_output=dy, grad_input_quantizer=None) + assert_bwd_fake_matches_real(bwd_args, _bias_backward_impl, _bias_backward_impl_fake) + + +@_cuda +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_bias_custom_op_matches_eager(dtype) -> None: + """The registered op must reproduce the eager operation, forward and backward.""" + assert _bias_ops is not None, "bias custom ops failed to register" + forward_fn, backward_fn = _bias_ops + + shape, size = (16, 32), 32 + op = _bias_op(size, dtype) + x = torch.randn(*shape, device=_device, dtype=dtype, requires_grad=True) + dy = torch.randn(*shape, device=_device, dtype=dtype) + + # Reference: the op as used today. + y_ref = op(x) + y_ref.backward(dy) + dx_ref, db_ref = x.grad.clone(), op.bias.grad.clone() + + # Same computation through the custom ops. + args = op.resolve_fwd_args(x.detach(), requires_grad=True, prev_op_grad_output_quantizer=None) + y, saved, ctx_attrs = forward_fn(args) + assert saved == (), "bias saves no tensors" + dx, db = backward_fn( + BiasBwdArgs(grad_output=dy, grad_input_quantizer=ctx_attrs["grad_input_quantizer"]) + ) + # A None grad input means "grad_output unchanged" -- a custom op may not + # return one of its own inputs. + if dx is None: + dx = dy + + torch.testing.assert_close(y, y_ref) + torch.testing.assert_close(dx, dx_ref) + torch.testing.assert_close(db, db_ref) + + +@_cuda +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_bias_custom_op_compiles_fullgraph(dtype) -> None: + """Both halves must trace without a graph break. + + Run under ``no_grad``: the halves carry no autograd of their own (that is + the point of ``register_op_halves``), so a caller is expected to wire them + into its own ``autograd.Function``. + """ + assert _bias_ops is not None, "bias custom ops failed to register" + forward_fn, backward_fn = _bias_ops + + shape, size = (16, 32), 32 + op = _bias_op(size, dtype) + x = torch.randn(*shape, device=_device, dtype=dtype) + dy = torch.randn(*shape, device=_device, dtype=dtype) + + def fwd(x_): + args = op.resolve_fwd_args(x_, requires_grad=True, prev_op_grad_output_quantizer=None) + out, _saved, _attrs = forward_fn(args) + return out + + def bwd(dy_): + return backward_fn(BiasBwdArgs(grad_output=dy_, grad_input_quantizer=None)) + + with torch.no_grad(): + torch.testing.assert_close(torch.compile(fwd, fullgraph=True)(x), fwd(x)) + compiled_grads = torch.compile(bwd, fullgraph=True)(dy) + expected_grads = bwd(dy) + for got, expected in zip(compiled_grads, expected_grads): + if got is None or expected is None: + assert got is expected is None + continue + torch.testing.assert_close(got, expected) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 3598e54daa..a580ea95e8 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,7 +6,7 @@ from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer from .tensor_spec import TensorSpec, to_tensor_spec -from .custom_op import register_custom_op +from .custom_op import register_custom_op, register_op_halves __all__ = [ "register_value_opaque_quantizer", @@ -14,4 +14,5 @@ "TensorSpec", "to_tensor_spec", "register_custom_op", + "register_op_halves", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 2a4f64d9a8..a35269aa07 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -183,9 +183,7 @@ def is_simple_value(cls, value: Any) -> bool: if _is_opaque_value_type(type(value)): return True if isinstance(value, dict): - return all( - isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items() - ) + return all(isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items()) if isinstance(value, (list, tuple)): return all(cls.is_simple_value(v) for v in value) return False @@ -309,7 +307,9 @@ def _collect(value: Any) -> None: register_opaque_type(OpaqueValueBundle, typ="value") _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) -except Exception as e: # pylint: disable=broad-exception-caught # pragma: no cover - older torch without opaque_object +except ( + Exception +) as e: # pylint: disable=broad-exception-caught # pragma: no cover - older torch without opaque_object warn_compile_unsupported(f"could not register OpaqueValueBundle as an opaque type ({e})") _is_opaque_value_type = None _is_opaque_reference_type = None @@ -1166,9 +1166,7 @@ def _setup_context(ctx, inputs, output): ctx_attrs, tuple(saved_list), ) - tensors_to_save, tensor_objects = prepare_for_saving( - *(tensors_to_save_from_setup or ()) - ) + tensors_to_save, tensor_objects = prepare_for_saving(*(tensors_to_save_from_setup or ())) ctx.tensor_objects = tensor_objects ctx.save_for_backward(*tensors_to_save) ctx.bwd_obj = bwd_obj @@ -1247,7 +1245,9 @@ def _register_wrapper_op( """ subclass_list = _all_quantized_tensor_subclasses() input_flatten_enabled = bool(subclass_list) and adapters is not None - slot_offsets = _collect_tensor_or_quantized_slot_offsets(adapters) if input_flatten_enabled else [] + slot_offsets = ( + _collect_tensor_or_quantized_slot_offsets(adapters) if input_flatten_enabled else [] + ) def _forward(*flat: Any) -> List[torch.Tensor]: if not input_flatten_enabled: @@ -1258,7 +1258,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 @@ -1389,9 +1392,7 @@ def _register_custom_op_impl( fwd_field_names = {f.name for f in dataclasses.fields(fwd_arg_type)} missing = [n for n in input_tensors_for_grad if n not in fwd_field_names] if missing: - raise ValueError( - f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {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" @@ -1516,3 +1517,194 @@ def forward_fn(fwd_args): return tuple(outputs) return forward_fn + + +# --------------------------------------------------------------------------- # +# Split registration: forward and backward as independent ops (no autograd) +# --------------------------------------------------------------------------- # + + +def register_op_halves( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: int, +) -> Optional[Tuple[Callable[[Any], Any], Callable[[Any], Any]]]: + """Register an op's forward and backward as two *independent* custom ops. + + Unlike :func:`register_custom_op`, no autograd is registered. The caller + wires forward to backward itself -- e.g. a pipeline-level + ``torch.autograd.Function`` that Dynamo traces as a higher-order op, so the + forward and backward passes can be grouped differently (which is what + ``ops.OperationFuser`` does). Both halves are still two-tier, so + ``QuantizedTensor`` subclass inputs pass through without dequantization. + + Contracts, mirroring :func:`register_custom_op`: + + * ``fwd_impl(fwd_args) -> (*user_outputs, tensors_to_save, ctx_attrs)`` + * ``fwd_fake_impl`` -- its data-free twin over :class:`TensorSpec` + * ``bwd_impl(bwd_args) -> tuple`` of ``num_grad_inputs`` gradients + * ``bwd_fake_impl`` -- its data-free twin + + Returns ``(forward_fn, backward_fn)``: + + * ``forward_fn(fwd_args) -> (outputs, saved_tensors, ctx_attrs)`` -- + ``outputs`` is a single value or a tuple, mirroring ``fwd_impl``'s user + outputs; ``saved_tensors`` is the reassembled ``tensors_to_save`` tuple, + which the caller is expected to persist (e.g. ``ctx.save_for_backward``). + * ``backward_fn(bwd_args) -> tuple`` of gradients. + + Returns ``None`` if registration fails (warned once), so callers can fall + back to eager rather than breaking import. + """ + try: + return _register_op_halves_impl( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_arg_type=bwd_arg_type, + bwd_impl=bwd_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=num_grad_inputs, + ) + except (ImportError, AttributeError, RuntimeError, TypeError) as e: + warn_compile_unsupported( + f"could not register the custom op halves '{op_name}' ({type(e).__name__}: {e})" + ) + return None + + +def _register_op_halves_impl( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: int, +) -> Tuple[Callable[[Any], Any], Callable[[Any], Any]]: + """Body of :func:`register_op_halves`; see it for semantics.""" + subclass_list = _all_quantized_tensor_subclasses() + + fwd_adapters = _get_adapters(fwd_arg_type) + bwd_adapters = _get_adapters(bwd_arg_type) + fwd_tensor_field_names = _tensor_field_names(fwd_adapters) + bwd_tensor_field_names = _tensor_field_names(bwd_adapters) + + fwd_schema_args, fwd_arg_names = _build_schema(fwd_adapters) + bwd_schema_args, bwd_arg_names = _build_schema(bwd_adapters) + fwd_schema = f"{fwd_schema_args} -> Tensor[]" + bwd_schema = f"{bwd_schema_args} -> Tensor[]" + + wrapper_fwd_name = op_name + wrapper_bwd_name = f"{op_name}_backward" + base_fwd_name = f"{op_name}_base" + base_bwd_name = f"{wrapper_bwd_name}_base" + base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" + + _register_kernel( + op_name=base_fwd_name, + schema_str=fwd_schema, + arg_type=fwd_arg_type, + arg_names=fwd_arg_names, + adapters=fwd_adapters, + tensor_field_names=fwd_tensor_field_names, + impl=fwd_impl, + fake_impl=fwd_fake_impl, + format_result=_format_fwd_result, + ) + _register_kernel( + op_name=base_bwd_name, + schema_str=bwd_schema, + arg_type=bwd_arg_type, + arg_names=bwd_arg_names, + adapters=bwd_adapters, + tensor_field_names=bwd_tensor_field_names, + impl=bwd_impl, + fake_impl=bwd_fake_impl, + format_result=lambda g: _format_bwd_result(g, num_grad_inputs, base_bwd_qualname), + ) + + base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) + base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) + + wrapper_fwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_fwd_name, + schema_str=fwd_schema, + base_op=base_fwd_op, + adapters=fwd_adapters, + ) + wrapper_bwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_bwd_name, + schema_str=bwd_schema, + base_op=base_bwd_op, + adapters=bwd_adapters, + ) + wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) + wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) + + fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_adapters) + bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_adapters) + + def _fwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) + return base_fwd_op(*new_args) + + def _bwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) + return base_bwd_op(*new_args) + + for sub in subclass_list: + wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) + wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule) + + for op in (wrapper_fwd_op, wrapper_bwd_op, base_fwd_op, base_bwd_op): + _quantized_tensor_passthrough_ops.add(op.default) + + def _reassemble(specs: List[Any], payload: List[torch.Tensor], cursor: int) -> Tuple[List, int]: + """Rebuild values described by ``specs`` from ``payload[cursor:]``.""" + out: List[Any] = [] + for spec in specs: + n = _spec_slot_count(spec) + chunk = [_decode_none(t) for t in payload[cursor : cursor + n]] + cursor += n + out.append(_spec_reassemble(spec, chunk)) + return out, cursor + + def forward_fn(fwd_args): + spec_obj = _spec_view(fwd_args, fwd_tensor_field_names) + user_specs, saved_specs, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(spec_obj)) + kwargs = _args_to_slots(fwd_args, fwd_adapters) + payload = wrapper_fwd_op(*[kwargs[name] for name in fwd_arg_names]) + + outputs, cursor = _reassemble(user_specs, payload, 0) + saved, _ = _reassemble(saved_specs, payload, cursor) + return ( + (outputs[0] if len(outputs) == 1 else tuple(outputs)), + tuple(saved), + ctx_attrs, + ) + + def backward_fn(bwd_args): + # Unlike the forward payload, each grad occupies exactly one slot + # (``_format_bwd_result`` materializes a TensorSpec grad), so there is + # nothing to reassemble. + kwargs = _args_to_slots(bwd_args, bwd_adapters) + payload = wrapper_bwd_op(*[kwargs[name] for name in bwd_arg_names]) + return tuple(_decode_none(t) for t in payload) + + return forward_fn, backward_fn diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index 88f563b2c5..e427267ccd 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -5,7 +5,8 @@ """Fusible operation for bias.""" from __future__ import annotations -from typing import Optional +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Union import torch @@ -14,6 +15,94 @@ from ..op import BasicOperation, OperationContext from ...utils import canonicalize_device, canonicalize_dtype from ...tensor import Quantizer +from ...quantized_tensor import QuantizedTensorStorage +from ...dynamo import TensorSpec, register_op_halves + +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class BiasFwdArgs: + """Flat, ``self``-free inputs to the bias forward.""" + + input_: TensorOrQuantized + bias: torch.Tensor + local_size: int + grad_input_quantizer: Optional[Quantizer] + + +@dataclass(slots=True) +class BiasBwdArgs: + """Flat inputs to the bias backward.""" + + grad_output: Optional[torch.Tensor] = None + grad_input_quantizer: Optional[Quantizer] = None + + +def _bias_forward_impl( + args: BiasFwdArgs, +) -> Tuple[torch.Tensor, Tuple[()], Dict[str, Any]]: + """Bias forward. Saves no tensors; backward only needs the quantizer.""" + x = args.input_ + b = args.bias.view([1] * (x.dim() - 1) + [args.local_size]) + return x + b, (), {"grad_input_quantizer": args.grad_input_quantizer} + + +def _bias_forward_impl_fake( + args: BiasFwdArgs, +) -> Tuple[TensorSpec, Tuple[()], Dict[str, Any]]: + """Allocation-free fake of :func:`_bias_forward_impl`.""" + x = args.input_ + out = TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device) + return out, (), {"grad_input_quantizer": args.grad_input_quantizer} + + +def _bias_backward_impl( + args: BiasBwdArgs, +) -> Tuple[Optional[torch.Tensor], torch.Tensor]: + """Bias backward: reduce the grad over all but the inner dimension. + + Returns ``None`` for the grad input when it is ``grad_output`` unchanged. A + custom op may not return one of its own inputs, and cloning would cost a + full-size copy on the common (unquantized) path; the caller substitutes + ``grad_output`` instead. + """ + dy = args.grad_output + if dy.dim() > 1: + quantizer = args.grad_input_quantizer + if quantizer is None: + return None, dy.sum(tuple(range(dy.dim() - 1))) + db, dy = tex.bgrad_quantize(dy, quantizer) + return dy, db + return None, dy + + +def _bias_backward_impl_fake( + args: BiasBwdArgs, +) -> Tuple[Optional[TensorSpec], TensorSpec]: + """Allocation-free fake of :func:`_bias_backward_impl`.""" + dy = args.grad_output + shape = tuple(dy.shape) + if len(shape) > 1: + grad_bias = TensorSpec(shape=(shape[-1],), dtype=dy.dtype, device=dy.device) + quantizer = args.grad_input_quantizer + if quantizer is None: + return None, grad_bias + grad_input = TensorSpec(shape=shape, dtype=dy.dtype, quantizer=quantizer, device=dy.device) + return grad_input, grad_bias + return None, TensorSpec(shape=shape, dtype=dy.dtype, device=dy.device) + + +_bias_ops = register_op_halves( + op_name="bias", + fwd_arg_type=BiasFwdArgs, + fwd_impl=_bias_forward_impl, + fwd_fake_impl=_bias_forward_impl_fake, + bwd_arg_type=BiasBwdArgs, + bwd_impl=_bias_backward_impl, + bwd_fake_impl=_bias_backward_impl_fake, + num_grad_inputs=2, +) class Bias(BasicOperation): @@ -113,6 +202,31 @@ def pre_first_fuser_forward(self) -> None: if self.bias.device.type == "meta": self.reset_parameters() + def resolve_fwd_args( + self, + input_: torch.Tensor, + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer], + ) -> BiasFwdArgs: + """Gather everything the forward needs into a flat, self-free container. + + Reads module config and global FP8 state, so it must run in the traced + region (where Dynamo guards those reads), never inside the custom op. + """ + grad_input_quantizer = None + if requires_grad: + grad_input_quantizer = prev_op_grad_output_quantizer + if FP8GlobalStateManager.is_fp8_enabled(): + if FP8GlobalStateManager.get_fp8_recipe().backward_override is not None: + grad_input_quantizer = None + return BiasFwdArgs( + input_=input_, + bias=self.bias, + local_size=self.local_size, + grad_input_quantizer=grad_input_quantizer, + ) + def op_forward( self, ctx: OperationContext, @@ -120,30 +234,30 @@ def op_forward( prev_op_grad_output_quantizer: Optional[Quantizer], next_op_input_quantizer: Optional[Quantizer], ) -> torch.Tensor: - x = input_ - b = self.bias.view([1] * (x.dim() - 1) + [self.local_size]) - + del next_op_input_quantizer # Bias never quantizes its output + args = self.resolve_fwd_args( + input_, + requires_grad=ctx.requires_grad, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + ) + out, saved, ctx_attrs = _bias_forward_impl(args) if ctx.requires_grad: - ctx.grad_input_quantizer = prev_op_grad_output_quantizer - if FP8GlobalStateManager.is_fp8_enabled(): - fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() - if fp8_recipe.backward_override is not None: - ctx.grad_input_quantizer = None - - return x + b + ctx.save_for_backward(*saved) + for name, value in ctx_attrs.items(): + setattr(ctx, name, value) + return out def op_backward( self, ctx: OperationContext, grad_output: torch.Tensor, ) -> tuple[torch.Tensor, tuple[()]]: - dy = grad_output - if dy.dim() > 1: - quantizer = ctx.grad_input_quantizer - if quantizer is None: - db = dy.sum(tuple(range(dy.dim() - 1))) - else: - db, dy = tex.bgrad_quantize(dy, quantizer) - else: - db = dy - return dy, (db,) + grad_input, grad_bias = _bias_backward_impl( + BiasBwdArgs( + grad_output=grad_output, + grad_input_quantizer=ctx.grad_input_quantizer, + ) + ) + if grad_input is None: + grad_input = grad_output + return grad_input, (grad_bias,)