diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 1286492a6e..0c8f34e5e2 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -6,6 +6,7 @@ import pytest import torch +from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode try: from torch._opaque_base import OpaqueBaseMeta @@ -24,14 +25,20 @@ from transformer_engine.common import recipe from transformer_engine.pytorch.constants import FP8FwdTensorIdx, FP8BwdTensorIdx 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.tensor.float8_tensor import Float8CurrentScalingQuantizer -from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer +from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, _STORAGE_REGISTRY +from transformer_engine.pytorch.dynamo import TensorProto, to_tensor_proto from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, is_fp8_block_scaling_available, is_nvfp4_available, + Float8Quantizer, + Float8BlockQuantizer, + MXFP8Quantizer, ) from utils import recipe_id @@ -384,3 +391,428 @@ def fn(inp): out = compiled(inp) out.sum().backward() + + +# --------------------------------------------------------------------------- +# Value-opaque quantizers +# --------------------------------------------------------------------------- + + +def _mxfp8(dtype=tex.DType.kFloat8E4M3): + return MXFP8Quantizer(fp8_dtype=dtype) + + +def _blockwise(force_pow_2_scales=True): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=force_pow_2_scales, + ) + + +def _current_scaling(amax_epsilon=0.0): + return Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device=torch.device("cpu"), + amax_epsilon=amax_epsilon, + ) + + +def _nvfp4(with_rht=True): + # Default with_rht=True so the quantize round-trip below exercises the + # derived ``rht_matrix`` tensor (the field most likely to be dropped on + # value-key reconstruction). + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_rht=with_rht, + ) + + +def _hw_available(quantizer): + """Whether this HW can actually run the quantize kernel for *quantizer*.""" + if isinstance(quantizer, MXFP8Quantizer): + return mxfp8_available + if isinstance(quantizer, NVFP4Quantizer): + return nvfp4_available + if isinstance(quantizer, Float8BlockQuantizer): + return fp8_block_scaling_available + return fp8_available # Float8CurrentScalingQuantizer + + +# (factory, kwargs producing a different-but-valid config) +_VALUE_QUANTIZERS = [ + pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), + pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"), + pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), + pytest.param( + _nvfp4, + {"with_rht": False}, + id="nvfp4", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), + reason="NVFP4Quantizer requires CUDA to construct", + ), + ), +] + + +@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) +def test_quantizer_value_object(factory, other_kwargs): + """Value semantics + ``__fx_repr__`` round-trip via the production FX path.""" + a, b = factory(), factory() + # Same config -> equal, same hash, interchangeable as a dict/set key. + assert a is not b + assert a == b + assert hash(a) == hash(b) + assert {a: "x"}[b] == "x" + # Different config -> not equal. + assert a != factory(**other_kwargs) + + # ``__fx_repr__`` (used by torch.compile codegen) rebuilds an equal object. + repr_str, globals_ = a.__fx_repr__() + rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used + assert rebuilt == a and rebuilt is not a + assert hash(rebuilt) == hash(a) + # The deprecated amax-reduction group is never part of the value. + assert getattr(rebuilt, "amax_reduction_group", None) is None + + # The rebuilt quantizer must also *behave* identically, not just compare + # equal: equality only looks at the value key, so a field the kernel needs + # but that is absent from the key (e.g. NVFP4's derived ``rht_matrix``) would + # slip through the checks above and only blow up at quantize time. Run the + # real quantize kernel on both and require bit-exact results. + if torch.cuda.is_available() and _hw_available(a): + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0) + + +def test_value_quantizer_rejects_process_group(): + """A value quantizer holding a live ProcessGroup must refuse to be turned + into a value key / FX constant (raise), not silently drop the group.""" + import torch.distributed as dist # pylint: disable=import-outside-toplevel + + created = not dist.is_initialized() + if created: + dist.init_process_group(backend="gloo", store=dist.HashStore(), rank=0, world_size=1) + try: + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + q.amax_reduction_group = dist.group.WORLD + # Every value-materialization path must reject it (hash, eq, __fx_repr__). + with pytest.raises(TypeError): + hash(q) + with pytest.raises(TypeError): + q.__fx_repr__() + finally: + if created: + dist.destroy_process_group() + + +if _opaque_available: + # A minimal custom op taking a tensor and a value-opaque quantizer that + # quantizes + dequantizes inside it, one per production quantizer class. + # ``test_quantizer_value_object_fullgraph`` drives this under + # ``torch.compile(fullgraph=True)`` so the quantizer is used *inside* the + # graph -- proving the opaque-type registration took effect (a graph break + # would make ``fullgraph=True`` raise). + _qdq_lib = torch.library.Library("test_te_qdq", "DEF") + _QDQ_OPS = {} + for _qcls in ( + MXFP8Quantizer, + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + NVFP4Quantizer, + ): + _op = f"qdq_{_qcls.__name__}" + _qdq_lib.define(f"{_op}(Tensor x, {get_opaque_type_name(_qcls)} q) -> Tensor") + + @torch.library.impl(f"test_te_qdq::{_op}", "CompositeExplicitAutograd", lib=_qdq_lib) + def _qdq_impl(x, q): + return q(x).dequantize() + + @torch.library.register_fake(f"test_te_qdq::{_op}", lib=_qdq_lib) + def _qdq_fake(x, q): + return torch.empty_like(x) + + _QDQ_OPS[_qcls] = getattr(torch.ops.test_te_qdq, _op) + + +@pytest.mark.skipif( + not _opaque_available, + reason="torch.compile opaque-object support requires PyTorch >= 2.11", +) +@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) +def test_quantizer_value_object_fullgraph(factory, other_kwargs): + """Quantizer is usable *inside* a torch.compile(fullgraph=True) graph. + + A custom op quantizes+dequantizes with the (opaque value) quantizer; the + compiled result must match eager. ``fullgraph=True`` raises on any graph + break, so this proves the opaque-type registration actually took effect -- + unlike merely passing the quantizer through. + """ + q = factory() + if not (torch.cuda.is_available() and _hw_available(q)): + pytest.skip("format not supported on this HW") + + op = _QDQ_OPS[type(q)] + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def fn(inp): + return op(inp, q) + + ref = fn(x) + torch._dynamo.reset() + out = torch.compile(fn, fullgraph=True)(x) + torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0) + + +# --------------------------------------------------------------------------- +# torch.compile-traceable allocation primitives + TensorProto +# --------------------------------------------------------------------------- + + +# (factory, logical shape) -- shapes respect MXFP8 (mult. of 32) / blockwise (128) +# / NVFP4 (mult. of 16) constraints. +_PROTO_QUANTIZERS = [ + pytest.param(_current_scaling, (4, 8), id="fp8_current_scaling"), + pytest.param(_mxfp8, (64, 128), id="mxfp8"), + pytest.param(_blockwise, (128, 256), id="fp8_blockwise"), + pytest.param( + _nvfp4, + (64, 128), + id="nvfp4", + marks=pytest.mark.skipif( + not nvfp4_available, + reason="NVFP4 is not available", + ), + ), +] + + +def _build_from_primitives(quantizer, shape, dtype, device="cpu"): + """Assemble a quantized tensor straight from the quantizer primitives: + ``alloc_tensors`` (buffers) + ``create_metadata`` (ctx) + the storage's + ``__tensor_unflatten__`` -- i.e. exactly what ``TensorProto.create_tensor`` + does, but without going through :class:`TensorProto`. + """ + names = tuple(quantizer._describe_buffers(shape)) # pylint: disable=protected-access + ctx = quantizer.create_metadata(shape, dtype=dtype) + buffers = quantizer.alloc_tensors(shape, device=device) + inner = {name: buffers[name] for name in names} + storage_cls = _STORAGE_REGISTRY[ctx["cls"]] + # Row-major (contiguous) outer stride for ``__tensor_unflatten__``; ``meta`` + # device computes it without allocating storage. + outer_stride = torch.empty(tuple(shape), device="meta").stride() + return storage_cls.__tensor_unflatten__(inner, ctx, tuple(shape), outer_stride) + + +def _signature(tensor, names): + """Comparable shape/dtype fingerprint of a tensor and its inner buffers.""" + sig = {"__shape__": tuple(tensor.shape), "__dtype__": tensor.dtype} + for name in names: + buf = getattr(tensor, name) + sig[name] = (tuple(buf.shape), buf.dtype) + return sig + + +def _skip_if_dequantize_unsupported(q): + """Skip when this HW can't run ``dequantize()`` for the quantizer's format. + + ``dequantize()`` runs the real kernel on CUDA, so each format has its own + availability gate (mirrors the ``is_*_available`` checks in test_numerics). + """ + if isinstance(q, MXFP8Quantizer): + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + elif isinstance(q, NVFP4Quantizer): + if not nvfp4_available: + pytest.skip("NVFP4 is not available") + elif isinstance(q, Float8BlockQuantizer): + if not fp8_block_scaling_available: + pytest.skip("FP8 block scaling is not available") + elif not fp8_available: # Float8 current scaling + pytest.skip(reason_for_no_fp8) + + +# ----- Quantizer primitives ----- + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_primitives_unflatten_compiles(factory, shape): + """create_metadata + alloc_tensors + __tensor_unflatten__ compose and trace + under ``fullgraph=True`` (CPU), without TensorProto.""" + q = factory() + names = tuple(q._describe_buffers(shape)) # pylint: disable=protected-access + + def fn(x): + t = _build_from_primitives(q, shape, x.dtype, device=x.device) + # Read every buffer into the result so the alloc + unflatten can't be + # eliminated as dead code -- forces the whole build path into the graph. + acc = x.new_zeros(()) + for name in names: + acc = acc + getattr(t, name).float().sum() + return acc + + x = torch.zeros(*shape, dtype=torch.bfloat16) + torch._dynamo.reset() + out = torch.compile(fn, fullgraph=True)(x) + assert out.shape == () + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_alloc_tensors_fake(factory, shape): + """``alloc_tensors`` produces FakeTensors with the described shapes/dtypes.""" + q = factory() + bufs = q._describe_buffers(shape) # pylint: disable=protected-access + with FakeTensorMode(): + alloc = q.alloc_tensors(shape, device="cpu") + assert set(alloc) == set(bufs) + for name, (buf_shape, buf_dtype) in bufs.items(): + assert isinstance(alloc[name], FakeTensor) + assert tuple(alloc[name].shape) == tuple(buf_shape) + assert alloc[name].dtype == buf_dtype + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_storage_flatten_unflatten_roundtrip(factory, shape): + """Storage ``__tensor_flatten__`` / ``__tensor_unflatten__`` round-trips. + + Build a tensor from ``alloc_tensors`` + ``create_metadata``, flatten it, then + unflatten and verify shape/dtype and every inner buffer match before vs after. + """ + q = factory() + _skip_if_dequantize_unsupported(q) + + tensor = _build_from_primitives(q, shape, torch.bfloat16) + names = tuple(q._describe_buffers(shape)) # pylint: disable=protected-access + # Fill buffers with deterministic data (empty() may contain NaNs) so the + # round-trip can be checked by value via dequantize(). + for name in names: + buf = getattr(tensor, name) + buf.copy_(torch.arange(buf.numel(), device=buf.device).reshape(buf.shape)) + before = _signature(tensor, names) + expected = tensor.dequantize() + + flat_names, flat_ctx = tensor.__tensor_flatten__() + assert set(flat_names) == set(names) + inner = {name: getattr(tensor, name) for name in flat_names} + rebuilt = type(tensor).__tensor_unflatten__( + inner, flat_ctx, tuple(tensor.shape), tensor.stride() + ) + + assert isinstance(rebuilt, QuantizedTensor) + assert _signature(rebuilt, flat_names) == before + # The reconstructed tensor dequantizes to the same values. + torch.testing.assert_close(rebuilt.dequantize(), expected, atol=0, rtol=0, equal_nan=True) + + +# ----- TensorProto ----- + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_tensor_proto_matches_primitives(factory, shape): + """TensorProto is a thin wrapper: its ``create_metadata`` / + ``create_inner_tensors`` / ``create_tensor`` match building everything + directly from the quantizer primitives.""" + q = factory() + proto = TensorProto(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu")) + assert proto.is_quantized + + # Metadata matches the quantizer's. + assert proto.create_metadata() == q.create_metadata(shape, dtype=torch.bfloat16) + + # inner_names + create_inner_tensors match _describe_buffers. + bufs = q._describe_buffers(shape) # pylint: disable=protected-access + names = tuple(bufs) + assert proto.inner_names() == names + inner = proto.create_inner_tensors() + assert len(inner) == len(names) + for name, buf in zip(names, inner): + exp_shape, exp_dtype = bufs[name] + assert tuple(buf.shape) == tuple(exp_shape) + assert buf.dtype == exp_dtype + + # The assembled tensor matches one built directly from the primitives. + direct = _build_from_primitives(q, shape, torch.bfloat16) + assert _signature(proto.create_tensor(), names) == _signature(direct, names) + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_tensor_proto_create_tensor_eager(factory, shape): + """``create_tensor`` (no fake) yields a real quantized tensor.""" + q = factory() + proto = TensorProto(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu")) + out = proto.create_tensor() + assert isinstance(out, QuantizedTensor) + assert tuple(out.shape) == tuple(shape) + assert out.dtype == torch.bfloat16 + for name in proto.inner_names(): + assert not isinstance(getattr(out, name), FakeTensor) + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_tensor_proto_create_tensor_fake(factory, shape): + """``create_tensor`` under ``FakeTensorMode`` yields a fake-backed quantized + tensor with the right shape/dtype and fake inner buffers.""" + q = factory() + proto = TensorProto(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu")) + with FakeTensorMode(): + out = proto.create_tensor() + assert isinstance(out, QuantizedTensor) + assert tuple(out.shape) == tuple(shape) + assert out.dtype == torch.bfloat16 + for name in proto.inner_names(): + assert isinstance(getattr(out, name), FakeTensor) + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_tensor_proto_create_tensor_compiles(factory, shape): + """``TensorProto.create_tensor`` traces under ``fullgraph=True`` (CPU).""" + q = factory() + + def fn(x): + proto = TensorProto(shape=tuple(x.shape), dtype=x.dtype, quantizer=q, device=x.device) + t = proto.create_tensor() + acc = x.new_zeros(()) + for name in proto.inner_names(): + acc = acc + getattr(t, name).float().sum() + return acc + + x = torch.zeros(*shape, dtype=torch.bfloat16) + torch._dynamo.reset() + out = torch.compile(fn, fullgraph=True)(x) + assert out.shape == () + + +def test_to_tensor_proto_plain(): + """``to_tensor_proto`` describes a plain tensor.""" + t = torch.empty(2, 3, dtype=torch.float32) + proto = to_tensor_proto(t) + assert not proto.is_quantized + assert proto.shape == (2, 3) + assert proto.dtype == torch.float32 + assert proto.inner_names() == ("data",) + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +def test_to_tensor_proto_quantized(factory, shape): + """``to_tensor_proto`` round-trips a quantized tensor back into a proto.""" + q = factory() + tensor = TensorProto( + shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu") + ).create_tensor() + + proto = to_tensor_proto(tensor) + assert proto.is_quantized + assert proto.shape == tuple(shape) + assert proto.dtype == torch.bfloat16 + # Same buffer layout as the original tensor. + assert proto.inner_names() == tuple( + q._describe_buffers(shape) + ) # pylint: disable=protected-access + # Rebuilding from the derived proto matches the original tensor's structure. + assert _signature(proto.create_tensor(), proto.inner_names()) == _signature( + tensor, proto.inner_names() + ) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py new file mode 100644 index 0000000000..f932a7d9c3 --- /dev/null +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""torch.compile glue for Transformer Engine.""" + +from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer +from .tensor_proto import TensorProto, to_tensor_proto + +__all__ = [ + "register_value_opaque_quantizer", + "is_value_opaque_quantizer", + "TensorProto", + "to_tensor_proto", +] diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py new file mode 100644 index 0000000000..8b8b3caa69 --- /dev/null +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -0,0 +1,116 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Value-opaque quantizers for torch.compile.""" + +from __future__ import annotations +from typing import Any, Dict, Tuple + +from ..constants import DType + + +# Registration marks the class with this attribute rather than recording it in a +# module-level set. It looks odd but is a deliberate workaround: the check must +# stay traceable when it runs inside a torch.compile graph -- Dynamo can bake a +# ``getattr`` on the opaque quantizer into a constant, but cannot evaluate +# ``type(q) in some_set`` (no equality/hash rules for the opaque class object), +# which would graph-break under ``fullgraph=True``. +_VALUE_OPAQUE_FLAG = "_te_compile_value_opaque" + + +def is_value_opaque_quantizer(quantizer: Any) -> bool: + """Whether *quantizer*'s class is registered as a torch.compile value-opaque + type.""" + return getattr(quantizer, _VALUE_OPAQUE_FLAG, False) + + +def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: + """Rebuild a tensorless quantizer of type *cls* from its value items. + + Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the + generated FX code calls this to materialize the quantizer constant. + """ + # Bypass ``__init__`` and restore the value attributes directly: the value + # items already capture every value-defining field (including derived ones), + # and the constructors have heterogeneous signatures / side effects. + obj = cls.__new__(cls) + field_names = set() + for name, value in items: + if name == "dtype": + value = DType.cast(value) + object.__setattr__(obj, name, value) + field_names.add(name) + # The deprecated amax-reduction group is not a value field; initialize it to + # None so attribute access keeps working on the rebuilt quantizer. + if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"): + object.__setattr__(obj, "amax_reduction_group", None) + # Restore non-value derived state that ``__init__`` would normally build but + # that cannot live in the value key (e.g. NVFP4's ``rht_matrix`` tensor). + finalize = getattr(obj, "_rebuild_derived_state", None) + if finalize is not None: + finalize() + return obj + + +def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]: + """``__fx_repr__`` for value-opaque quantizers (attached at registration). + + Returns an evaluable expression that rebuilds the quantizer via + :func:`_rebuild_quantizer`, capturing both the helper and the quantizer + class itself in the FX globals so codegen can resolve them with no global + registry and no qualname collisions. + + Raises ``TypeError`` (via :meth:`Quantizer._value_key`) if the quantizer + stores a process group (e.g. a non-``None`` deprecated + ``amax_reduction_group``): live distributed state must never be baked into + the graph as a constant. Pass the reduction group per quantize call instead + of storing it on the quantizer. + """ + cls = type(self) + items = self._value_key()[1] + return ( + f"_rebuild_quantizer({cls.__name__}, {items!r})", + {"_rebuild_quantizer": _rebuild_quantizer, cls.__name__: cls}, + ) + + +def register_value_opaque_quantizer(cls: type) -> None: + """Register a tensorless quantizer class as a torch.compile value opaque type. + + Attaches ``__fx_repr__`` and registers the class with + ``torch._library.opaque_object``. Safe to call on any PyTorch build: on + versions without the opaque-object API it only attaches ``__fx_repr__`` + (harmless), so Transformer Engine keeps importing and running in eager mode. + + The quantizer class must already provide value ``__eq__`` / ``__hash__`` and + a non-``None`` ``_value_fields`` (see + :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`). + """ + # Stamp the class so it can be recognized as value-opaque in dynamo-traced + # code (used to fall back to eager for unregistered quantizers). + setattr(cls, _VALUE_OPAQUE_FLAG, True) + + # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the + # class, so attach it before registering. + if "__fx_repr__" not in cls.__dict__: + cls.__fx_repr__ = _quantizer_fx_repr + + try: + from torch._library.opaque_object import ( # pylint: disable=import-outside-toplevel + register_opaque_type, + is_opaque_value_type, + ) + except (ImportError, AttributeError): + # Older PyTorch without the opaque-object API: eager value semantics + # still work; torch.compile specialization on the quantizer does not. + return + + try: + if not is_opaque_value_type(cls): + register_opaque_type(cls, typ="value") + except (RuntimeError, TypeError): + # Keep TE importable: neither the opaque-type query nor the registration + # must crash the import, e.g. on PyTorch versions with only partial / + # experimental opaque-object support. + pass diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py new file mode 100644 index 0000000000..911b4151bf --- /dev/null +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -0,0 +1,177 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""TensorProto: a data-free description of a tensor / quantized tensor.""" + +from __future__ import annotations +import copy as _copy +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +import torch + + +def _contiguous_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]: + """Row-major (contiguous) stride for ``shape``.""" + stride: list = [] + acc = 1 + for dim in reversed(shape): + stride.append(acc) + acc *= dim + return tuple(reversed(stride)) + + +@dataclass +class TensorProto: + """A data-free *prototype* of a tensor or quantized tensor. + + Captures ``shape`` / ``dtype`` and, for quantized tensors, the + (value-opaque) ``quantizer`` -- enough to rebuild a tensor without holding + storage. The common abstraction over plain ``torch.Tensor``, + ``QuantizedTensorStorage`` and ``QuantizedTensor``, used for custom-op fake + impls and for reassembling a quantized tensor from bare buffers. + """ + + shape: Tuple[int, ...] + dtype: torch.dtype + quantizer: Optional[Any] = None + requires_grad: bool = False + device: Optional[torch.device] = field(default=None) + + def __post_init__(self) -> None: + # Own a private copy of the quantizer so usage changes (update_usage) + # never touch the shared, value-opaque quantizer. The copy inherits the + # quantizer's current row-/column-wise usage as this proto's layout. + if self.quantizer is not None: + q = self.quantizer + self.quantizer = q.copy() if hasattr(q, "copy") else _copy.copy(q) + + @property + def is_quantized(self) -> bool: + """Whether this proto describes a quantized tensor.""" + return self.quantizer is not None + + def update_usage( + self, + *, + rowwise_usage: Optional[bool] = None, + columnwise_usage: Optional[bool] = None, + ) -> None: + """Mirror ``QuantizedTensor.update_usage`` on the proto's buffer layout. + + Applied to the proto's own quantizer copy, so the shared (value-opaque) + quantizer is never mutated. No-op for plain (non-quantized) protos. + """ + if self.quantizer is None: + return + self.quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) + + def inner_names(self) -> Tuple[str, ...]: + """Names of the flat tensor buffers backing this proto, in order. + + The real op flattens a quantized output via the storage's + ``__tensor_flatten__`` -- i.e. ``_FLATTEN_TENSOR_BUFFERS`` order, keeping + only the present buffers. ``_describe_buffers`` may emit the same buffers + in a different (per-usage) order (e.g. NVFP4 groups each amax right after + its scale), so reorder to the canonical flatten order here to keep the + fake layout aligned with the real one slot-for-slot. + """ + if self.quantizer is None: + return ("data",) + # pylint: disable=protected-access + described = list(self.quantizer._describe_buffers(tuple(self.shape)).keys()) + storage_cls = self.quantizer._storage_metadata(self.dtype)["cls"] + flatten_order = [attr for attr, _ in storage_cls._FLATTEN_TENSOR_BUFFERS] + extra = [name for name in described if name not in flatten_order] + if extra: + raise RuntimeError( + f"{storage_cls.__name__} describes buffer(s) {extra} absent from its " + f"_FLATTEN_TENSOR_BUFFERS {flatten_order}; the fake layout cannot be " + "aligned with the real one slot-for-slot." + ) + return tuple(name for name in flatten_order if name in described) + + def create_metadata(self) -> Dict[str, Any]: + """Data-free ``__tensor_unflatten__`` context describing this tensor.""" + if self.quantizer is None: + return { + "is_tensor": True, + "is_quantized": False, + "dtype": self.dtype, + "requires_grad": self.requires_grad, + } + return self.quantizer.create_metadata( + tuple(self.shape), dtype=self.dtype, requires_grad=self.requires_grad + ) + + def create_inner_tensors(self) -> List[torch.Tensor]: + """Materialize the flat inner buffers (in :meth:`inner_names` order). + + Under ``register_fake`` the ``torch.empty`` calls produce ``FakeTensor``s; + ``requires_grad`` is left default (managed by ``register_autograd``). + """ + device = self.device if self.device is not None else torch.device("cuda") + if self.quantizer is None: + return [torch.empty(tuple(self.shape), dtype=self.dtype, device=device)] + inner = self.quantizer.alloc_tensors(tuple(self.shape), device=device) + return [inner[name] for name in self.inner_names()] + + def create_tensor(self) -> torch.Tensor: + """Materialize an (uninitialized) tensor matching this proto (traceable). + + Quantized protos reassemble the :meth:`create_inner_tensors` buffers via + the storage's ``__tensor_unflatten__``. + """ + if self.quantizer is None: + device = self.device if self.device is not None else torch.device("cuda") + return torch.empty( + tuple(self.shape), + dtype=self.dtype, + device=device, + requires_grad=self.requires_grad, + ) + from ..quantized_tensor import ( # pylint: disable=import-outside-toplevel + _STORAGE_REGISTRY, + ) + + shape = tuple(self.shape) + ctx = self.create_metadata() + inner = dict(zip(self.inner_names(), self.create_inner_tensors())) + storage_cls = _STORAGE_REGISTRY[ctx["cls"]] + return storage_cls.__tensor_unflatten__(inner, ctx, shape, _contiguous_stride(shape)) + + +def to_tensor_proto(tensor: Any) -> TensorProto: + """Build a :class:`TensorProto` describing ``tensor``. + + Works for plain ``torch.Tensor`` and for ``QuantizedTensorStorage`` / + ``QuantizedTensor``. A *bare* storage exposes its shape via ``.size()`` and + its (fake) dtype via ``_dtype`` rather than ``.shape`` / ``.dtype``. + """ + from ..quantized_tensor import ( # pylint: disable=import-outside-toplevel + QuantizedTensorStorage, + ) + + requires_grad = bool(getattr(tensor, "requires_grad", False)) + if isinstance(tensor, QuantizedTensorStorage): + shape = getattr(tensor, "shape", None) + if shape is None: + shape = tensor.size() + dtype = getattr(tensor, "dtype", None) + if dtype is None: + dtype = getattr(tensor, "_dtype", None) + return TensorProto( + shape=tuple(shape), + dtype=dtype, + quantizer=getattr(tensor, "_quantizer", None), + requires_grad=requires_grad, + device=tensor.device, + ) + return TensorProto( + shape=tuple(tensor.shape), + dtype=tensor.dtype, + quantizer=None, + requires_grad=requires_grad, + device=tensor.device, + ) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index fed367bce5..473483faff 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -68,6 +68,7 @@ prepare_for_saving, restore_from_func_ctx, ) +from ..dynamo import TensorProto from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -92,7 +93,7 @@ class LinearFwdArgs: # --- Differentiable tensors (also passed positionally to autograd) --- weight: TensorOrQuantized - inp: torch.Tensor + inp: TensorOrQuantized bias: Optional[torch.Tensor] # --- Non-differentiable cached tensors --- @@ -301,7 +302,15 @@ def _linear_forward_impl( # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) - backward_needs_input = is_grad_enabled and weight.requires_grad + # Use the requires-grad flags captured into ``args`` at op-call time rather + # than the live tensors': the fake impl (``_linear_forward_impl_fake``) keys + # the number of FP8 inner buffers it emits off ``args.*_requires_grad``, so + # the real impl must agree to keep the custom-op output arity stable. Under + # ``torch.compile`` with CUDA-graph trees (``mode="reduce-overhead"``) the + # static graph inputs are detached during capture, so live + # ``weight.requires_grad`` / ``inp.requires_grad`` flip to False mid-capture + # and would otherwise diverge from the fake (schema/arity mismatch). + backward_needs_input = is_grad_enabled and args.weight_requires_grad with_input_all_gather_nccl = ( parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop ) @@ -418,7 +427,7 @@ def _linear_forward_impl( # No need to set the quantizer states if weight is already quantized # for debug mode we create quantizer every iteration, thus we need to set the quantizer states if weight_quantizer is not None and (not isinstance(weight, QuantizedTensor) or debug): - columnwise_usage = is_grad_enabled and inp.requires_grad and not is_fsdp2 + columnwise_usage = is_grad_enabled and args.input_requires_grad and not is_fsdp2 if backward_override is not None: columnwise_usage = False if not columnwise_usage: @@ -598,12 +607,24 @@ def _linear_forward_impl( if is_fsdp2 and weightmat is not weight: wt_save = None - # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` - # rebuilds the refs from ``inp`` / ``weight`` / ``bias``. - # Needed for torch.compile to work correctly. + # Dedup save slots that alias forward inputs or other op returns; + # ``_linear_setup_ctx`` rebuilds the refs. Needed because a custom op may + # not return a tensor that aliases an input or another return: the cached + # FP8 weight is the same tensor as ``new_weight_workspace`` (a return, on a + # cache miss) or ``weight_workspace`` (an input, on a cache hit). + if wt_save is None: + wt_alias = None + elif wt_save is weight: + wt_alias = "weight" + elif new_weight_workspace is not None and wt_save is new_weight_workspace: + wt_alias = "new_workspace" + elif args.weight_workspace is not None and wt_save is args.weight_workspace: + wt_alias = "weight_workspace" + else: + wt_alias = None saved_tensor_aliases = ( "inp" if saved_inputmat is inp else None, - "weight" if wt_save is weight else None, + wt_alias, "weight", # ``saved_weight`` slot is always the weight parameter "bias" if bias is not None else None, ) @@ -622,10 +643,214 @@ def _linear_forward_impl( return out, new_weight_workspace, tensors_to_save_from_forward, None, ctx_attrs +def _linear_forward_impl_fake( + args: LinearFwdArgs, +) -> Tuple[TensorProto, Optional[TensorProto], Optional[Tuple[Any, ...]], None, Optional[Dict]]: + """Shape/metadata-only twin of :func:`_linear_forward_impl` for torch.compile, + returning ``TensorProto`` descriptors for the outputs and saved tensors instead + of allocating real data.""" + if args.fsdp_group is not None and args.is_grad_enabled: + raise NotImplementedError( + "Compile-time Linear forward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + + weight = args.weight + inp = args.inp + bias = args.bias + input_quantizer = args.input_quantizer + weight_quantizer = args.weight_quantizer + output_quantizer = args.output_quantizer + fp8 = args.fp8 + debug = args.debug + fp8_or_debug = fp8 or debug + is_grad_enabled = args.is_grad_enabled + activation_dtype = args.activation_dtype + save_original_input = args.save_original_input + if args.backward_override == "high_precision": + save_original_input = True + + out_features, _ = weight.shape + backward_needs_input = is_grad_enabled and args.weight_requires_grad + + own_quantized_input = False + inputmat_is_storage = False + inputmat_aliases_inp = False + if fp8_or_debug: + if inp.is_quantized: + # Primary-quantized input reused as-is. + inputmat_is_storage = True + inputmat_aliases_inp = True + else: + if input_quantizer is None: + raise ValueError("Missing quantizer for input tensor") + input_quantizer.set_usage( + rowwise=True, + columnwise=( + backward_needs_input + and not save_original_input + and args.backward_override is None + ), + ) + own_quantized_input = True + inputmat_is_storage = True + else: + inputmat_aliases_inp = inp.dtype == activation_dtype + + if save_original_input: + inputmat_aliases_inp = True + inputmat_is_storage = False + + # ------------------------------------------------------ + # Weight pipeline -- mirror ``quantize_weight`` / ``cast_if_needed``. + # ``new_weight_workspace`` is a fresh fake storage only on the + # cache-miss + ``cache_weight`` path, else ``None``. + # ------------------------------------------------------ + new_weight_workspace = None + weightmat = None + weightmat_is_storage = False + weightmat_aliases_weight = False + if fp8_or_debug: + if weight_quantizer is not None and (not weight.is_quantized or debug): + columnwise_usage = is_grad_enabled and args.input_requires_grad and not args.is_fsdp2 + if args.backward_override is not None: + columnwise_usage = False + if not columnwise_usage: + columnwise_usage = ( + is_fp8_activation_recompute_enabled() + and not in_fp8_activation_recompute_phase() + ) + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + elif weight.is_quantized: + weight_quantizer = weight.quantizer + + if weight.is_quantized: + # Primary-quantized weight: the impl reuses it as ``weightmat``. + weightmat = weight + weightmat_is_storage = True + weightmat_aliases_weight = True + else: + weightmat = TensorProto( + shape=tuple(weight.shape), + dtype=activation_dtype, + quantizer=weight_quantizer, + device=weight.device, + ) + weightmat_is_storage = True + update_ws = args.is_first_microbatch is None or args.is_first_microbatch + if args.cache_weight and update_ws and args.weight_workspace is None: + new_weight_workspace = TensorProto( + shape=tuple(weight.shape), + dtype=activation_dtype, + quantizer=weight_quantizer, + device=weight.device, + ) + else: + weightmat_aliases_weight = weight.dtype == activation_dtype + weightmat = TensorProto( + shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device + ) + + if output_quantizer is not None: + output_quantizer.set_usage(rowwise=True, columnwise=False) + + # ------------------------------------------------------ + # Output tensor: y = x @ w^T (quantized iff an output quantizer is set). + # ------------------------------------------------------ + out_leading = inp.shape[0] + if args.parallel_mode == "column" and args.sequence_parallel: + out_leading = out_leading * args.tp_size + elif args.parallel_mode == "row" and args.sequence_parallel: + out_leading = out_leading // args.tp_size + out = TensorProto( + shape=(out_leading, *tuple(inp.shape[1:-1]), out_features), + dtype=activation_dtype, + quantizer=output_quantizer, + requires_grad=is_grad_enabled and (args.input_requires_grad or args.weight_requires_grad), + device=inp.device, + ) + + # ------------------------------------------------------ + # Backward state -- saved-tensor layout + # (saved_inputmat, wt_save, saved_weight, bias) with name-based aliasing. + # ------------------------------------------------------ + tensors_to_save_from_forward = None + ctx_attrs = None + if is_grad_enabled: + # Slot 0 -- ``saved_inputmat``. + inputmat_alias = None + saved_inputmat = None + if backward_needs_input: + if inputmat_aliases_inp: + inputmat_alias = "inp" + elif inputmat_is_storage: + saved_inputmat = TensorProto( + shape=tuple(inp.shape), + dtype=activation_dtype, + quantizer=input_quantizer, + device=inp.device, + ) + # Mirror ``_linear_forward_impl``'s post-quantization + # ``inputmat.update_usage(...)`` so the saved input's buffer layout + # matches -- driven by the same conditions as the real impl. + if own_quantized_input and not save_original_input: + if args.backward_override is not None: + saved_inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) + elif ( + args.backward_input_needs_gather + and weight_quantizer is not None + and weight_quantizer.supports_only_rowwise_all_gather() + ): + saved_inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) + else: + saved_inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) + else: + saved_inputmat = TensorProto( + shape=tuple(inp.shape), dtype=activation_dtype, device=inp.device + ) + + # Slot 1 -- ``wt_save``. Mirror the real impl's alias dedup: the cached + # FP8 weight is shared with ``new_weight_workspace`` (a return, on a cache + # miss) or the ``weight_workspace`` input (on a cache hit), so it is + # reconstructed in ``_linear_setup_ctx`` rather than saved twice. + wt_alias = None + wt_save = None + if weightmat_aliases_weight: + wt_alias = "weight" + elif args.is_fsdp2: + pass # FSDP2 re-quantizes from the gathered weight in backward. + elif weightmat_is_storage and new_weight_workspace is not None: + wt_alias = "new_workspace" + elif weightmat_is_storage and args.weight_workspace is not None: + wt_alias = "weight_workspace" + elif weightmat_is_storage: + wt_save = weightmat + else: + wt_save = TensorProto( + shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device + ) + + # Slot 2 -- ``saved_weight`` (always aliased to ``weight``). + # Slot 3 -- ``bias`` (aliased to ``bias`` when present, else absent). + saved_tensor_aliases = ( + inputmat_alias, + wt_alias, + "weight", + "bias" if bias is not None else None, + ) + tensors_to_save_from_forward = (saved_inputmat, wt_save, None, None) + ctx_attrs = { + "fsdp_shapes": [], + "saved_tensor_aliases": saved_tensor_aliases, + } + + return out, new_weight_workspace, tensors_to_save_from_forward, None, ctx_attrs + + def _linear_setup_ctx( bwd_args: LinearBwdArgs, fwd_args: LinearFwdArgs, - out: torch.Tensor, + fwd_outputs: Tuple[Any, ...], ctx_attrs: Dict, tensors_to_save_from_forward: Tuple[Any, ...], ) -> Tuple[Any, ...]: @@ -638,7 +863,8 @@ def _linear_setup_ctx( for FSDP2 re-quantization) without having to mutate the structured metadata returned by ``prepare_for_saving``. """ - del out # No-op; kept for symmetry with the compile-time helper signature. + # ``fwd_outputs`` are the op's user outputs ``(out, new_weight_workspace)``; + # only ``new_weight_workspace`` is needed here, to rebuild the deduped weight. inp = fwd_args.inp weight = fwd_args.weight @@ -728,6 +954,10 @@ def _linear_setup_ctx( saved_inputmat = inp if wt_save_alias == "weight": wt_save = weight + elif wt_save_alias == "new_workspace": + wt_save = fwd_outputs[1] + elif wt_save_alias == "weight_workspace": + wt_save = fwd_args.weight_workspace if saved_weight_alias == "weight": saved_weight = weight if bias_alias == "bias": @@ -1311,6 +1541,68 @@ def wgrad_gemm( ) +def _linear_backward_impl_fake( + args: LinearBwdArgs, +) -> Tuple[Optional[TensorProto], Optional[TensorProto], Optional[TensorProto]]: + """Allocation-free fake of :func:`_linear_backward` on ``TensorProto``. + + The saved-tensor fields of ``args`` carry + :class:`~transformer_engine.pytorch.dynamo.TensorProto` instances. Returns + ``(wgrad, dgrad, grad_bias)`` protos describing the nature of the gradients, + mirroring the real backward's return contract without allocating storage. + + Tensor-/sequence-parallel gather/scatter happens inside the eager backward + custom op and is opaque to ``torch.compile``: ``dgrad`` always carries the + rank-local input shape and ``wgrad`` the local weight shape, so no extra + shape modeling is needed here. + """ + if args.fsdp_group is not None: + raise NotImplementedError( + "Fake Linear backward does not support manual TE FSDP " + "(fsdp_group is not None); use FSDP2 or MCore FSDP." + ) + + weight = args.saved_weight if args.saved_weight is not None else args.weight_fp8 + out_dtype = args.activation_dtype + out_features, in_features = weight.shape + + # Mirror ``_linear_backward``: ``set_usage`` on ``grad_input_quantizer`` + # influences ``dgrad``'s buffer layout. + if args.grad_input_quantizer is not None: + args.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) + + dgrad = None + if args.requires_dgrad: + # dgrad has the logical input shape and may be quantized for the next op. + dgrad = TensorProto( + shape=tuple(args.inp_shape), + dtype=out_dtype, + quantizer=args.grad_input_quantizer, + device=args.grad_output.device, + ) + + wgrad = None + if args.requires_wgrad and not args.fuse_wgrad_accumulation: + # wgrad has the weight's shape; quantized iff an fp8 wgrad output is + # requested (mirrors ``quantization_params=grad_weight_quantizer``), + # otherwise high precision. Under fuse_wgrad_accumulation the grad is + # written into ``main_grad`` in place and no wgrad tensor is returned. + wgrad = TensorProto( + shape=(out_features, in_features), + dtype=out_dtype, + quantizer=args.grad_weight_quantizer, + device=weight.device, + ) + + grad_bias = None + if args.use_bias and args.requires_wgrad: + grad_bias = TensorProto( + shape=(out_features,), dtype=out_dtype, device=args.grad_output.device + ) + + return wgrad, dgrad, grad_bias + + class _Linear(torch.autograd.Function): """Linear semi-top level module Calls custom cuda extensions. @@ -1350,7 +1642,7 @@ def forward( tensors_to_save_from_setup = _linear_setup_ctx( bwd_args, fwd_args, - out, + (out, new_weight_workspace), ctx_attrs, tensors_to_save_from_forward, ) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index cfe488aae5..86c4c94f7c 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -16,6 +16,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe +from transformer_engine.pytorch.constants import dist_group_type from transformer_engine.pytorch.tensor._quantization_helpers import ( _QuantizeFunc, _IdentityFunc, @@ -23,12 +24,29 @@ ) +def _contains_process_group(value: Any) -> bool: + """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``. + + Checks the value directly and one level of ``tuple``/``list`` nesting, which + covers the shapes a quantizer value field could plausibly take. + """ + if isinstance(value, dist_group_type): + return True + if isinstance(value, (tuple, list)): + return any(_contains_process_group(item) for item in value) + return False + + # Custom ops that should pass through __torch_dispatch__ without unwrapping # QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that # handle quantized tensors internally. _quantized_tensor_passthrough_ops: set = set() +#: Maps storage / wrapper class qualname -> class object, for ``__tensor_unflatten__``. +_STORAGE_REGISTRY: Dict[str, type] = {} + + class QuantizedTensorStorage: r"""Base class for all TensorStorage classes. @@ -132,6 +150,66 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: f"{self.__class__.__name__} class does not implement copy_from_storage function" ) + # ----- PyTorch subclass flatten protocol (torch.compile / TensorProto) ----- + + # Subclasses declare their tensor buffers once, as ``(attribute_name, + # constructor_kwarg)`` pairs in flatten order; everything else returned by + # :meth:`get_metadata` is treated as non-tensor context. + _FLATTEN_TENSOR_BUFFERS: Tuple[Tuple[str, str], ...] = () + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + # Register every storage / wrapper class so ``__tensor_unflatten__`` can + # resolve the concrete class from its qualname inside an FX graph. + _STORAGE_REGISTRY[cls.__qualname__] = cls + + def _flatten_nontensor_kwargs(self) -> Dict[str, Any]: + """Non-tensor constructor kwargs (scalars, dtype, quantizer).""" + tensor_kwargs = {kwarg for _, kwarg in self._FLATTEN_TENSOR_BUFFERS} + return {k: v for k, v in self.get_metadata().items() if k not in tensor_kwargs} + + def __tensor_flatten__(self) -> Tuple[list, Dict[str, Any]]: + """Return ``(inner_tensor_attr_names, context)``; see class comment.""" + present = [ + attr for attr, _ in self._FLATTEN_TENSOR_BUFFERS if getattr(self, attr) is not None + ] + ctx = { + "cls": type(self).__qualname__, + "is_tensor": isinstance(self, QuantizedTensor), + "requires_grad": ( + bool(self.requires_grad) if isinstance(self, QuantizedTensor) else False + ), + "nontensor_kwargs": self._flatten_nontensor_kwargs(), + } + return present, ctx + + @staticmethod + def __tensor_unflatten__( + inner_tensors: Dict[str, torch.Tensor], + ctx: Dict[str, Any], + outer_size: Iterable[int], + outer_stride: Optional[Iterable[int]], + ) -> QuantizedTensorStorage: + """Rebuild a storage / wrapper from flat tensors + context.""" + cls = _STORAGE_REGISTRY[ctx["cls"]] + kwargs: Dict[str, Any] = dict(ctx["nontensor_kwargs"]) + # Map each declared buffer back to its constructor kwarg (absent -> None). + for attr, kwarg in cls._FLATTEN_TENSOR_BUFFERS: + kwargs[kwarg] = inner_tensors.get(attr) + if not ctx["is_tensor"]: + return cls(**kwargs) + # Wrapper subclass: it also needs outer shape / dtype / device / stride. + fake_dtype = kwargs.get("fake_dtype") + device = next((t.device for t in inner_tensors.values() if t is not None), None) + return cls( + shape=tuple(outer_size), + dtype=fake_dtype, + requires_grad=ctx["requires_grad"], + device=device, + stride=tuple(outer_stride) if outer_stride is not None else None, + **kwargs, + ) + def prepare_for_saving( *tensors: Union[torch.Tensor, QuantizedTensorStorage], @@ -349,6 +427,71 @@ def make_empty( result.requires_grad_(True) return result + # ----- Data-free buffer/metadata primitives backing TensorProto ----- + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + """Return ``{attr_name: (buffer_shape, buffer_dtype)}`` for the buffers + this quantizer would allocate for a logical tensor of ``shape``. + + Keys must match the buffer attribute names declared in the storage's + ``_FLATTEN_TENSOR_BUFFERS`` and respect the quantizer's usage flags. + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not implement _describe_buffers; " + "it cannot be used with TensorProto / pure-Python allocation" + ) + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + """Non-tensor context for the produced storage. + + Returns ``{"cls": , "nontensor_kwargs": {...}}`` where ``cls`` is + the concrete class to instantiate (wrapper subclass for user-visible + tensors, bare storage class for ``internal`` quantizers) and + ``nontensor_kwargs`` are its non-tensor constructor kwargs (e.g. + ``fp8_dtype``, ``quantizer``, ``fake_dtype``). + """ + raise NotImplementedError( + f"{self.__class__.__name__} does not implement _storage_metadata; " + "it cannot be used with TensorProto / pure-Python allocation" + ) + + def alloc_tensors( + self, + shape: Iterable[int], + *, + device: Optional[Union[torch.device, str]] = None, + ) -> Dict[str, torch.Tensor]: + """Allocate (uninitialized) the flat buffers for ``shape``. + + Returns ``{attr_name: torch.Tensor}`` suitable as the ``inner_tensors`` + argument of the storage's ``__tensor_unflatten__``. + """ + device = torch.device(device if device is not None else "cuda") + return { + attr: torch.empty(buf_shape, dtype=buf_dtype, device=device) + for attr, (buf_shape, buf_dtype) in self._describe_buffers(tuple(shape)).items() + } + + def create_metadata( + self, + _shape: Iterable[int], + *, + dtype: torch.dtype, + requires_grad: bool = False, + ) -> Dict[str, Any]: + """Build the data-free ``__tensor_unflatten__`` context describing the + quantized tensor this quantizer would produce for ``shape`` / ``dtype``. + """ + meta = self._storage_metadata(dtype) + return { + "cls": meta["cls"].__qualname__, + "is_tensor": not self.internal, + "requires_grad": requires_grad, + "nontensor_kwargs": meta["nontensor_kwargs"], + } + def calibrate(self, tensor: torch.Tensor) -> None: """Calibrate quantizer state @@ -408,6 +551,78 @@ def get_usages(self) -> Dict[str, bool]: "columnwise": self.columnwise_usage, } + #: Attributes shared by every quantizer that take part in value identity. + _BASE_VALUE_FIELDS: Tuple[str, ...] = ( + "rowwise_usage", + "columnwise_usage", + "internal", + "optimize_for_gemm", + ) + + def _value_fields(self) -> Optional[Tuple[str, ...]]: + """Subclass-specific value-defining attribute names, or ``None``. + + Returning ``None`` (the default) means the quantizer cannot be represented as + a value opaque object and keeps identity-based equality/hashing. + This also means that passing such a quantizer as an argument to a custom op + causes a graph break under torch.compile, since it cannot be baked into the + FX graph as a constant. + """ + return None + + def _check_value_has_no_process_group(self) -> None: + # A value quantizer is baked into the FX graph as a constant via its + # value key, which cannot carry live distributed state. Enforced here -- + # the single point every value-materialization path (``__eq__`` / + # ``__hash__`` / ``__fx_repr__``) goes through -- so a custom + # ``__fx_repr__`` cannot bypass it. Reject any field holding a + # ProcessGroup (e.g. the deprecated ``amax_reduction_group``) rather than + # silently dropping it; pass the reduction group per quantize call. + for name, value in vars(self).items(): + if _contains_process_group(value): + raise TypeError( + f"{type(self).__name__} cannot be used as a torch.compile value " + f"object: attribute {name!r} holds a torch.distributed.ProcessGroup, " + "which is live distributed state and must not be baked into an FX " + "graph. Pass the amax reduction group per quantize call instead of " + "storing it on the quantizer." + ) + + def _value_key(self) -> Tuple[Any, ...]: + """Hashable, reproducible key identifying this quantizer's value. + + Only valid for value quantizers (``_value_fields()`` is not ``None``). + """ + fields = self._value_fields() # pylint: disable=assignment-from-none + assert fields is not None, f"{type(self).__name__} is not a value quantizer" + self._check_value_has_no_process_group() + items = [] + for name in self._BASE_VALUE_FIELDS + tuple(fields): + value = getattr(self, name) + if name == "dtype": + # ``DType`` is an ``IntEnum``; store the int so the key stays + # plain: hashable and ``repr``-reproducible for FX codegen. + value = int(value) + items.append((name, value)) + return (type(self).__qualname__, tuple(items)) + + def __eq__(self, other: object) -> Any: + # Value quantizers compare by configuration; everything else keeps the + # default identity semantics (returning ``NotImplemented`` makes Python + # fall back to identity). ``_value_key`` rejects a stored ProcessGroup. + if self is other: + return True + if self._value_fields() is None or type(self) is not type(other): + return NotImplemented + if other._value_fields() is None: + return NotImplemented + return self._value_key() == other._value_key() + + def __hash__(self) -> int: + if self._value_fields() is None: + return object.__hash__(self) + return hash(self._value_key()) + class QuantizedTensor(torch.Tensor): """Abstract base class for tensor with quantized data diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index ba46508d74..c816b4fb04 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -7,13 +7,14 @@ from collections.abc import Iterable import math import warnings -from typing import Any, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple, Union import torch import transformer_engine_torch as tex from transformer_engine.common.recipe import Float8BlockScaling, Recipe from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc from ..constants import DType from ..utils import devices_match, round_up_to_nearest_multiple @@ -69,6 +70,42 @@ def copy(self) -> Float8BlockQuantizer: return quantizer + def _value_fields(self) -> Tuple[str, ...]: + return ("dtype", "block_len", "amax_epsilon", "force_pow_2_scales", "block_scaling_dim") + + # ----- TensorProto / pure-Python allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": Float8BlockwiseQTensorStorage if self.internal else Float8BlockwiseQTensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "is_2D_scaled": self.block_scaling_dim == 2, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # Blockwise FP8 scales are FP32; columnwise data is stored transposed. + if self.rowwise_usage: + buffers["_rowwise_data"] = (shape, torch.uint8) + buffers["_rowwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=False)), + torch.float32, + ) + if self.columnwise_usage: + buffers["_columnwise_data"] = (tuple(self.get_columnwise_shape(shape)), torch.uint8) + buffers["_columnwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=True)), + torch.float32, + ) + return buffers + def update_quantized( self, src: torch.Tensor, @@ -211,6 +248,9 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return Float8BlockScaling +register_value_opaque_quantizer(Float8BlockQuantizer) + + class Float8BlockwiseQTensor(Float8BlockwiseQTensorStorage, QuantizedTensor): """Tensor class with FP8 data quantized via NxN blocks or 1xN blocks. diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index e26abf7df0..6d3a53b3d7 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -4,7 +4,7 @@ """Tensor class with FP8 data""" from __future__ import annotations -from typing import Any, Optional, Tuple, Iterable, Union +from typing import Any, Dict, Optional, Tuple, Iterable, Union import warnings import torch from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState @@ -15,9 +15,10 @@ Float8CurrentScaling, Recipe, ) -from ..utils import canonicalize_process_group, devices_match +from ..utils import canonicalize_process_group, devices_match, is_non_tn_fp8_gemm_supported from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc from ..constants import dist_group_type, DType @@ -386,6 +387,45 @@ def supports_only_rowwise_all_gather(self) -> bool: """ return True + def _value_fields(self) -> Tuple[str, ...]: + # ``amax_reduction_group`` is intentionally excluded: it is a deprecated + # process group (not a value). If one is actually stored, ``__fx_repr__`` + # raises so it can never be baked into a torch.compile graph. + return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction") + + # ----- TensorProto / pure-Python allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": Float8TensorStorage if self.internal else Float8Tensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # Mirror the C++ quantizer allocation (csrc/quantizer.cpp): on non-TN-capable + # archs (Blackwell+) a single ``_data`` buffer backs both row- and column-wise + # usage and no separate transpose is materialized. This must match what the + # real kernel produces so the torch.compile fake layout lines up slot-for-slot. + non_tn = is_non_tn_fp8_gemm_supported() + if self.rowwise_usage or non_tn: + buffers["_data"] = (shape, torch.uint8) + if self.columnwise_usage and not non_tn: + buffers["_transpose"] = ((shape[-1], *shape[:-1]), torch.uint8) + # Per-tensor scale-inv is always present for current scaling. + buffers["_scale_inv"] = ((1,), torch.float32) + return buffers + + +register_value_opaque_quantizer(Float8CurrentScalingQuantizer) + class Float8Tensor(Float8TensorStorage, QuantizedTensor): """Experimental tensor class with FP8 data diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index d759aaf5c4..f804a96f24 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -6,7 +6,7 @@ from __future__ import annotations from collections.abc import Iterable import math -from typing import Optional, Tuple, Union, Any +from typing import Optional, Tuple, Union, Any, Dict import warnings import torch @@ -18,6 +18,7 @@ from ..utils import devices_match, round_up_to_nearest_multiple from .storage.mxfp8_tensor_storage import MXFP8TensorStorage, _FromMXFP8Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc aten = torch.ops.aten @@ -57,6 +58,41 @@ def copy(self) -> MXFP8Quantizer: return quantizer + def _value_fields(self) -> Tuple[str, ...]: + return ("dtype",) + + # ----- TensorProto / pure-Python allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": MXFP8TensorStorage if self.internal else MXFP8Tensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "with_gemm_swizzled_scales": self.optimize_for_gemm, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + if self.rowwise_usage: + buffers["_rowwise_data"] = (shape, torch.uint8) + buffers["_rowwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=False)), + torch.uint8, + ) + if self.columnwise_usage: + buffers["_columnwise_data"] = (shape, torch.uint8) + buffers["_columnwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=True)), + torch.uint8, + ) + return buffers + def update_quantized( self, src: torch.Tensor, @@ -1058,3 +1094,6 @@ def backward( ) return dgrad, None return grad.view(ctx.shape), None + + +register_value_opaque_quantizer(MXFP8Quantizer) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index aa92be004f..8d51480f51 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -7,7 +7,7 @@ from collections.abc import Iterable import math import warnings -from typing import Dict, Optional, Tuple, Union +from typing import Any, Dict, Optional, Tuple, Union import functools import torch @@ -23,6 +23,7 @@ from .storage.nvfp4_tensor_storage import NVFP4TensorStorage, _FromNVFP4Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc aten = torch.ops.aten @@ -173,6 +174,7 @@ def __init__( self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") + self._with_random_sign_mask = with_random_sign_mask self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( with_random_sign_mask, torch.cuda.current_device() ) @@ -184,6 +186,16 @@ def __getstate__(self): state["amax_reduction_group"] = None return state + def _rebuild_derived_state(self) -> None: + """Restore the derived ``rht_matrix`` after value-key reconstruction. + + ``rht_matrix`` is a ``torch.Tensor`` built from ``_with_random_sign_mask`` + and the device, so it cannot be part of the (hashable) value key. + ``_rebuild_quantizer`` calls this hook to rebuild it; the ``lru_cache`` on + :func:`get_rht_matrix` makes an already-seen (flag, device) a cheap hit. + """ + self.rht_matrix = get_rht_matrix(self._with_random_sign_mask, torch.cuda.current_device()) + def update_quantized( self, src: torch.Tensor, @@ -333,6 +345,79 @@ def _canonicalized_amax_reduction_group(self) -> dist_group_type: def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return NVFP4BlockScaling + def _value_fields(self) -> Tuple[str, ...]: + # ``amax_reduction_group`` is intentionally excluded: it is a deprecated + # process group, not a value (``_value_key`` rejects a stored group). + # ``rht_matrix_random_sign_mask_t`` is a device-independent int derived + # from ``_with_random_sign_mask``; kept in the key so the rebuilt + # quantizer carries it without recomputation. + return ( + "dtype", + "with_rht", + "with_post_rht_amax", + "with_2d_quantization", + "stochastic_rounding", + "row_scaled_nvfp4", + "nvfp4_use_4over6", + "nvfp4_e4m3_max", + "nvfp4_4over6_err_mode", + "_with_random_sign_mask", + "rht_matrix_random_sign_mask_t", + "with_amax_reduction", + ) + + # ----- TensorProto / pure-Python allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": NVFP4TensorStorage if self.internal else NVFP4Tensor, + "nontensor_kwargs": { + "fp4_dtype": self.dtype, + "quantizer": self, + "with_gemm_swizzled_scales": self.optimize_for_gemm, + "row_scaled_nvfp4": self.row_scaled_nvfp4, + "nvfp4_use_4over6": self.nvfp4_use_4over6, + "nvfp4_e4m3_max": self.nvfp4_e4m3_max, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # FP4 data packs 2 values per byte (uint8); block scales are E4M3 stored + # as uint8; amax buffers are FP32 (per-row when row-scaled, else scalar). + # Order matches NVFP4TensorStorage._FLATTEN_TENSOR_BUFFERS (the canonical + # __tensor_flatten__ order): data + scale_inv per usage first, amax last. + # Workaround: call @staticmethods via the class, not the instance -- + # instance access breaks torch.compile guard generation (pytorch #182741). + if self.rowwise_usage: + buffers["_rowwise_data"] = (type(self).convert_shape_for_fp4(shape), torch.uint8) + buffers["_rowwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=False)), + torch.uint8, + ) + if self.columnwise_usage: + buffers["_columnwise_data"] = ( + type(self).convert_shape_for_fp4(type(self).get_columnwise_shape(shape)), + torch.uint8, + ) + buffers["_columnwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=True)), + torch.uint8, + ) + if self.rowwise_usage: + amax_rowwise_shape = (math.prod(shape[:-1]),) if self.row_scaled_nvfp4 else (1,) + buffers["_amax_rowwise"] = (amax_rowwise_shape, torch.float32) + if self.columnwise_usage: + buffers["_amax_columnwise"] = ((1,), torch.float32) + return buffers + + +register_value_opaque_quantizer(NVFP4Quantizer) + class NVFP4Tensor(NVFP4TensorStorage, QuantizedTensor): """Quantized tensor class with FP4 data diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index f7a3dae70b..ec2c40ef9a 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -35,6 +35,15 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): _columnwise_scale_inv: Optional[torch.Tensor] _is_2D_scaled: bool + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_rowwise_data", "rowwise_data"), + ("_rowwise_scale_inv", "rowwise_scale_inv"), + ("_columnwise_data", "columnwise_data"), + ("_columnwise_scale_inv", "columnwise_scale_inv"), + ) + def __new__( cls, rowwise_data: Optional[torch.Tensor], diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index a97162f91c..1419a02559 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -75,6 +75,14 @@ class Float8TensorStorage(QuantizedTensorStorage): _transpose: Optional[torch.Tensor] _transpose_invalid: bool + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_data", "data"), + ("_transpose", "data_transpose"), + ("_scale_inv", "fp8_scale_inv"), + ) + def __new__( cls, *args, diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index ea592cd989..a2a3bf2f4c 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -82,6 +82,15 @@ class MXFP8TensorStorage(QuantizedTensorStorage): # GEMM _with_gemm_swizzled_scales: bool + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_rowwise_data", "rowwise_data"), + ("_rowwise_scale_inv", "rowwise_scale_inv"), + ("_columnwise_data", "columnwise_data"), + ("_columnwise_scale_inv", "columnwise_scale_inv"), + ) + def __new__( cls, rowwise_data: Optional[torch.Tensor], diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 53bb5e7c11..5ed6d1d641 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -108,6 +108,17 @@ class NVFP4TensorStorage(QuantizedTensorStorage): # Global E4M3 scale bound used by this NVFP4 tensor _nvfp4_e4m3_max: int + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_rowwise_data", "rowwise_data"), + ("_rowwise_scale_inv", "rowwise_scale_inv"), + ("_columnwise_data", "columnwise_data"), + ("_columnwise_scale_inv", "columnwise_scale_inv"), + ("_amax_rowwise", "amax_rowwise"), + ("_amax_columnwise", "amax_columnwise"), + ) + def __new__( cls, rowwise_data: Optional[torch.Tensor],