From 6d3eecea21704f32d98cf80a9c57572e2b0e85c0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 16 Jun 2026 17:29:49 +0200 Subject: [PATCH 01/28] Add TensorProto mechanism for data-free quantized tensor allocation Squashed PR #8 (tensor_proto_mechanism) onto the rebased base. Adds TensorProto (pure-Python, torch.compile-traceable quantized-tensor allocation via Quantizer.alloc_tensors + storage __tensor_flatten__/__tensor_unflatten__), Linear fake fwd/bwd impls for the custom-op path, and tests. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 252 ++++++++++++++++ transformer_engine/pytorch/dynamo/__init__.py | 3 + .../pytorch/dynamo/tensor_proto.py | 172 +++++++++++ transformer_engine/pytorch/module/linear.py | 275 +++++++++++++++++- .../pytorch/quantized_tensor.py | 127 ++++++++ .../pytorch/tensor/float8_blockwise_tensor.py | 35 ++- .../pytorch/tensor/float8_tensor.py | 34 ++- .../pytorch/tensor/mxfp8_tensor.py | 36 ++- .../pytorch/tensor/nvfp4_tensor.py | 45 ++- .../float8_blockwise_tensor_storage.py | 9 + .../tensor/storage/float8_tensor_storage.py | 8 + .../tensor/storage/mxfp8_tensor_storage.py | 9 + .../tensor/storage/nvfp4_tensor_storage.py | 11 + 13 files changed, 1008 insertions(+), 8 deletions(-) create mode 100644 transformer_engine/pytorch/dynamo/tensor_proto.py diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 8ebce563ce..50757fcd75 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 @@ -27,6 +28,10 @@ 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.dynamo.tensor_proto import _contiguous_stride from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, @@ -466,6 +471,8 @@ def test_quantizer_value_object(factory): 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 @@ -554,3 +561,248 @@ def fn(inp): 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 torch.cuda.is_available(), + reason="NVFP4Quantizer requires CUDA to construct", + ), + ), +] + + +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"]] + return storage_cls.__tensor_unflatten__(inner, ctx, tuple(shape), _contiguous_stride(shape)) + + +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 index ee860c78e3..f932a7d9c3 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -5,8 +5,11 @@ """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/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py new file mode 100644 index 0000000000..b5248e3ee4 --- /dev/null +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -0,0 +1,172 @@ +# 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] + ordered = [name for name in flatten_order if name in described] + ordered += [name for name in described if name not in flatten_order] + return tuple(ordered) + + 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 78a4d31852..3b3ff0e0cb 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 --- @@ -303,7 +304,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 ) @@ -420,7 +429,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: @@ -624,6 +633,204 @@ 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``. + 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: + 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, @@ -1313,6 +1520,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. diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index c4194e7f52..e887924519 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -30,6 +30,10 @@ _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. @@ -133,6 +137,64 @@ 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], @@ -350,6 +412,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 diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 09fde86f17..7f43ac463e 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -7,7 +7,7 @@ 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 @@ -73,6 +73,39 @@ def copy(self) -> Float8BlockQuantizer: return quantizer + # ----- 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, diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 2e0491d49a..a4da60d7f3 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,7 +15,7 @@ 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 @@ -387,6 +387,36 @@ def supports_only_rowwise_all_gather(self) -> bool: """ return True + # ----- 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) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 3045662216..ded81faf0a 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 @@ -58,6 +58,40 @@ def copy(self) -> MXFP8Quantizer: return quantizer + # ----- 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]] = {} + # MXFP8 block scales are stored as uint8 (E8M0); data buffers keep the + # logical shape (rowwise) and its transpose (columnwise). + 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, diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index ccf06ac166..da5c07718b 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 @@ -345,6 +345,49 @@ def _canonicalized_amax_reduction_group(self) -> dist_group_type: def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return NVFP4BlockScaling + # ----- 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). + if self.rowwise_usage: + buffers["_rowwise_data"] = (self.convert_shape_for_fp4(shape), torch.uint8) + buffers["_rowwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=False)), + torch.uint8, + ) + 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["_columnwise_data"] = ( + self.convert_shape_for_fp4(self.get_columnwise_shape(shape)), + torch.uint8, + ) + buffers["_columnwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=True)), + torch.uint8, + ) + buffers["_amax_columnwise"] = ((1,), torch.float32) + return buffers + register_value_opaque_quantizer(NVFP4Quantizer) 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 b24a4e9144..3b9b0f11f0 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -126,6 +126,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 374d0e1e72..e1078674e7 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -76,6 +76,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 606ac9e74b..27a44e55fb 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -83,6 +83,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 09f040ba67..b43c5e93f0 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -109,6 +109,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], From 4bdaa60df268e9027cf8e3860144297a050d8623 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 22 Jun 2026 12:45:34 +0200 Subject: [PATCH 02/28] [PyTorch] torch.compile: dedup cached FP8 weight from saved-for-backward The cached FP8 weight is the same tensor returned as new_weight_workspace (cache miss) or passed in as weight_workspace (cache hit). A custom op may not return a tensor that aliases an input or another return, so mark those slots and reconstruct wt_save in _linear_setup_ctx instead of saving it twice. Mirrored in the fake impl so the saved-slot layout matches. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 40 ++++++++++++++++----- 1 file changed, 32 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 3b3ff0e0cb..b487d88556 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -609,12 +609,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, ) @@ -800,13 +812,20 @@ def _linear_forward_impl_fake( shape=tuple(inp.shape), dtype=activation_dtype, device=inp.device ) - # Slot 1 -- ``wt_save``. + # 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: @@ -834,7 +853,7 @@ def _linear_forward_impl_fake( 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, ...]: @@ -847,7 +866,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 @@ -937,6 +957,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": @@ -1621,7 +1645,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, ) From 3be6fc79cbeccc21a6150d05d45d899a8c462819 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 22 Jun 2026 13:17:19 +0200 Subject: [PATCH 03/28] [PyTorch] nvfp4: emit _describe_buffers in canonical flatten order NVFP4Quantizer._describe_buffers grouped each amax right after its scale (per-usage), diverging from NVFP4TensorStorage._FLATTEN_TENSOR_BUFFERS (amax buffers last). The order is functionally irrelevant (buffers are consumed by name in alloc_tensors and reordered in TensorProto.inner_names), but aligning it makes describe/flatten agree and fixes test_to_tensor_proto_quantized[nvfp4]. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index da5c07718b..8f65acc735 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -368,14 +368,14 @@ def _describe_buffers( 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. if self.rowwise_usage: buffers["_rowwise_data"] = (self.convert_shape_for_fp4(shape), torch.uint8) buffers["_rowwise_scale_inv"] = ( tuple(self.get_scale_shape(shape, columnwise=False)), torch.uint8, ) - 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["_columnwise_data"] = ( self.convert_shape_for_fp4(self.get_columnwise_shape(shape)), @@ -385,6 +385,10 @@ def _describe_buffers( 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 From 8a9d90c28bac4be90f8673b9b2f6a2c181a1cabd Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 08:00:58 +0200 Subject: [PATCH 04/28] Address review: error on undescribed buffers, gate nvfp4 test on HW support - TensorProto.inner_names now raises if the quantizer describes buffer(s) absent from the storage's _FLATTEN_TENSOR_BUFFERS, instead of silently appending them. - Gate the nvfp4 proto-quantizer param on nvfp4_available so it skips on hardware without NVFP4 support rather than failing. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 14 +++++++++----- transformer_engine/pytorch/dynamo/tensor_proto.py | 11 ++++++++--- transformer_engine/pytorch/module/linear.py | 3 +-- transformer_engine/pytorch/quantized_tensor.py | 4 +++- transformer_engine/pytorch/tensor/mxfp8_tensor.py | 2 -- 5 files changed, 21 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 50757fcd75..727567726e 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -31,7 +31,6 @@ 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.dynamo.tensor_proto import _contiguous_stride from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, @@ -579,8 +578,8 @@ def fn(inp): (64, 128), id="nvfp4", marks=pytest.mark.skipif( - not torch.cuda.is_available(), - reason="NVFP4Quantizer requires CUDA to construct", + not nvfp4_available, + reason="NVFP4 is not available", ), ), ] @@ -597,7 +596,10 @@ def _build_from_primitives(quantizer, shape, dtype, device="cpu"): buffers = quantizer.alloc_tensors(shape, device=device) inner = {name: buffers[name] for name in names} storage_cls = _STORAGE_REGISTRY[ctx["cls"]] - return storage_cls.__tensor_unflatten__(inner, ctx, tuple(shape), _contiguous_stride(shape)) + # 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): @@ -801,7 +803,9 @@ def test_to_tensor_proto_quantized(factory, shape): 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 + 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/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index b5248e3ee4..911b4151bf 100644 --- a/transformer_engine/pytorch/dynamo/tensor_proto.py +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -83,9 +83,14 @@ def inner_names(self) -> Tuple[str, ...]: 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] - ordered = [name for name in flatten_order if name in described] - ordered += [name for name in described if name not in flatten_order] - return tuple(ordered) + 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.""" diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index b487d88556..975c2992b7 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -768,8 +768,7 @@ def _linear_forward_impl_fake( 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), + requires_grad=is_grad_enabled and (args.input_requires_grad or args.weight_requires_grad), device=inp.device, ) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index e887924519..42b0143589 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -163,7 +163,9 @@ def __tensor_flatten__(self) -> Tuple[list, Dict[str, Any]]: ctx = { "cls": type(self).__qualname__, "is_tensor": isinstance(self, QuantizedTensor), - "requires_grad": bool(self.requires_grad) if isinstance(self, QuantizedTensor) else False, + "requires_grad": ( + bool(self.requires_grad) if isinstance(self, QuantizedTensor) else False + ), "nontensor_kwargs": self._flatten_nontensor_kwargs(), } return present, ctx diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index ded81faf0a..c639795a80 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -76,8 +76,6 @@ def _describe_buffers( ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: shape = tuple(shape) buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} - # MXFP8 block scales are stored as uint8 (E8M0); data buffers keep the - # logical shape (rowwise) and its transpose (columnwise). if self.rowwise_usage: buffers["_rowwise_data"] = (shape, torch.uint8) buffers["_rowwise_scale_inv"] = ( From e36cf6d871e0cf8f97a5e46c9b2a58b0edaced92 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 16:25:44 +0200 Subject: [PATCH 05/28] [PyTorch] Workaround torch.compile staticmethod guard bug in NVFP4 _describe_buffers Access NVFP4Quantizer @staticmethods (convert_shape_for_fp4, get_columnwise_shape) via the class instead of the instance. Under torch.compile, instance access of a @staticmethod on a value-opaque object crashes Dynamo guard generation with "'function' object has no attribute '__func__'" (pytorch/pytorch#182741). Temporary workaround until the PyTorch-side fix lands. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 8f65acc735..9e468ee458 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -370,15 +370,17 @@ def _describe_buffers( # 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"] = (self.convert_shape_for_fp4(shape), torch.uint8) + 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"] = ( - self.convert_shape_for_fp4(self.get_columnwise_shape(shape)), + type(self).convert_shape_for_fp4(type(self).get_columnwise_shape(shape)), torch.uint8, ) buffers["_columnwise_scale_inv"] = ( From 90fc49483ea47f880439f3d70f5d64ce90c2efff Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 13 Jul 2026 10:46:31 +0200 Subject: [PATCH 06/28] [PyTorch] Document TensorOrQuantized union (review feedback) The union is intentional: fields may carry bare QuantizedTensorStorage objects (internal-quantizer optimization), and the annotation is introspected in the follow-up custom-op PR to build the op schema with flatten/unflatten slots. Also note the size()/.shape asymmetry and how TensorProto handles it. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 975c2992b7..ecfb31a11d 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -84,6 +84,11 @@ __all__ = ["Linear"] +# Fields with this union may hold a *bare* ``QuantizedTensorStorage`` (see its +# docstring), not just a plain / subclass tensor. The annotation is also +# machine-read when the args bag becomes a torch.compile custom op (follow-up +# PR): such fields get flatten/unflatten slots so a bare storage can cross the +# op boundary -- a plain ``Tensor`` slot cannot carry it. TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] From a3bf5c1d4435db1a9acd9c545b322d26b639658e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 13 Jul 2026 10:56:07 +0200 Subject: [PATCH 07/28] [PyTorch] Add shape property to QuantizedTensorStorage Make .shape valid on bare storages (derived from size()), so Tensor, QuantizedTensor, bare storage and TensorProto all expose the same attribute. Wrapper subclasses defer to the native TensorBase.shape. Simplifies the shape fallback in to_tensor_proto. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/tensor_proto.py | 9 +++------ transformer_engine/pytorch/quantized_tensor.py | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index 911b4151bf..8881024104 100644 --- a/transformer_engine/pytorch/dynamo/tensor_proto.py +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -146,8 +146,8 @@ 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``. + ``QuantizedTensor``. A *bare* storage exposes its (fake) dtype via + ``_dtype`` rather than ``.dtype``. """ from ..quantized_tensor import ( # pylint: disable=import-outside-toplevel QuantizedTensorStorage, @@ -155,14 +155,11 @@ def to_tensor_proto(tensor: Any) -> TensorProto: 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), + shape=tuple(tensor.shape), dtype=dtype, quantizer=getattr(tensor, "_quantizer", None), requires_grad=requires_grad, diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 42b0143589..d9740c820e 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -53,6 +53,20 @@ class QuantizedTensorStorage: _dtype: torch.dtype _quantizer: Optional[Quantizer] + @property + def shape(self) -> torch.Size: + """Logical tensor shape, valid on bare storages and wrapper tensors alike. + + Wrapper subclasses (also ``torch.Tensor``) defer to the native tensor + shape (``size()`` on a bare storage may reconstruct the shape from the + columnwise buffer, which is not necessarily the outer shape); bare + storages derive it from ``size()``. + """ + if isinstance(self, torch.Tensor): + # pylint: disable=unnecessary-dunder-call + return torch._C.TensorBase.shape.__get__(self, type(self)) + return torch.Size(self.size()) + def update_usage( self, rowwise_usage: Optional[bool] = None, From 37b0d17ebc1fabb169c432f8ba045af2e9efdf22 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 13 Jul 2026 10:57:46 +0200 Subject: [PATCH 08/28] [PyTorch] Test parity of Python alloc vs C++ make_empty Address review: build the same quantized tensor via make_empty (C++, tex.create_empty_quantized_tensor) and via the Python primitives (_describe_buffers + create_metadata + alloc_tensors + __tensor_unflatten__) and check structural parity (class, buffer set, per-buffer shape/dtype/device, flatten context) and functional parity (the real quantize kernel writes bit-identical results into both, dequantize matches), across quantizer families x rowwise/columnwise x wrapper/internal. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 134 +++++++++++++++++++++++++++- 1 file changed, 133 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 727567726e..6b034abfc0 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -29,7 +29,11 @@ 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.quantized_tensor import ( + QuantizedTensor, + Quantizer, + _STORAGE_REGISTRY, +) from transformer_engine.pytorch.dynamo import TensorProto, to_tensor_proto from transformer_engine.pytorch import ( is_fp8_available, @@ -702,6 +706,134 @@ def test_storage_flatten_unflatten_roundtrip(factory, shape): torch.testing.assert_close(rebuilt.dequantize(), expected, atol=0, rtol=0, equal_nan=True) +_USAGE_COMBOS = [ + pytest.param(True, True, id="rowwise_columnwise"), + pytest.param(True, False, id="rowwise_only"), + pytest.param(False, True, id="columnwise_only"), +] + + +@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS) +@pytest.mark.parametrize("rowwise, columnwise", _USAGE_COMBOS) +@pytest.mark.parametrize("internal", [False, True], ids=["wrapper", "internal"]) +def test_python_alloc_matches_cpp_make_empty(factory, shape, rowwise, columnwise, internal): + """Pure-Python allocation is interchangeable with the C++ allocation. + + Builds the same quantized tensor twice: via ``Quantizer.make_empty`` + (``tex.create_empty_quantized_tensor``, the C++ path) and via the Python + primitives ``_describe_buffers`` + ``create_metadata`` + ``alloc_tensors`` + + ``__tensor_unflatten__`` (exactly what ``TensorProto.create_tensor`` + does). Checks: + + * structural parity -- same concrete class, buffer set, per-buffer + shape/dtype/device, logical shape/dtype and flatten context; + * functional parity -- the real C++ quantize kernel writes bit-identical + results into the Python-allocated buffers as into the C++-allocated + ones, proving the Python buffer description matches the layout + (padding/alignment) the kernels expect. + """ + + # Two independent, identically-configured quantizers so no state can leak + # between the two allocation paths. + def make_quantizer(): + q = factory() + q.set_usage(rowwise=rowwise, columnwise=columnwise) + q.internal = internal + return q + + q_ref = make_quantizer() + if not (torch.cuda.is_available() and _hw_available(q_ref)): + pytest.skip("format not supported on this HW") + q_py = make_quantizer() + + ref = q_ref.make_empty(shape, dtype=torch.bfloat16, device="cuda") + py = _build_from_primitives(q_py, shape, torch.bfloat16, device="cuda") + + # --- Structural parity --- + assert type(py) is type(ref) + ref_names, ref_ctx = ref.__tensor_flatten__() + py_names, py_ctx = py.__tensor_flatten__() + assert set(py_names) == set(ref_names) + for name in ref_names: + rbuf, pbuf = getattr(ref, name), getattr(py, name) + assert tuple(pbuf.shape) == tuple(rbuf.shape), name + assert pbuf.dtype == rbuf.dtype, name + assert pbuf.device == rbuf.device, name + + # Logical shape / dtype (bare storages are not torch.Tensors: they expose + # size() and _dtype instead of .shape / .dtype). + if isinstance(ref, QuantizedTensor): + assert tuple(py.shape) == tuple(ref.shape) == tuple(shape) + assert py.dtype == ref.dtype == torch.bfloat16 + else: + assert tuple(py.size()) == tuple(ref.size()) == tuple(shape) + # pylint: disable=protected-access + assert py._dtype == ref._dtype == torch.bfloat16 + + # Flatten context. The quantizer entry needs special handling: production + # quantizers get a value-based __eq__ from register_value_opaque_quantizer, + # but fall back to field-wise comparison for classes that don't define one + # (plain object.__eq__ is identity, which would spuriously fail). + assert set(py_ctx) == set(ref_ctx) + for key in ("cls", "is_tensor", "requires_grad"): + assert py_ctx[key] == ref_ctx[key], key + ref_kwargs, py_kwargs = ref_ctx["nontensor_kwargs"], py_ctx["nontensor_kwargs"] + assert set(py_kwargs) == set(ref_kwargs) + for key in ref_kwargs: + rv, pv = ref_kwargs[key], py_kwargs[key] + if isinstance(rv, Quantizer) or isinstance(pv, Quantizer): + assert type(pv) is type(rv), key + assert (pv.rowwise_usage, pv.columnwise_usage, pv.internal) == ( + rv.rowwise_usage, + rv.columnwise_usage, + rv.internal, + ), key + if type(rv).__eq__ is not object.__eq__: + assert pv == rv, key + else: + assert pv == rv, key + + # --- Functional parity: run the real C++ quantize kernel into both --- + x = torch.randn(*shape, dtype=torch.bfloat16, device="cuda") + + def _quantize_into(quantizer, dst): + if internal: + # update_quantized() only accepts the wrapper classes; internal + # (bare storage) tensors are filled through the same underlying + # kernel binding directly. + tex.quantize(x, quantizer, dst, None) + else: + quantizer.update_quantized(x, dst) + + # Some combos are rejected by the quantize kernel itself regardless of who + # allocated the tensor (e.g. FP8 current-scaling columnwise-only on + # TN-capable archs: there is no rowwise data buffer and + # nvte_compute_scale_from_amax asserts on it). Parity then means the + # Python-allocated tensor is rejected the same way -- not a silent skip. + try: + _quantize_into(q_ref, ref) + except RuntimeError: + with pytest.raises(RuntimeError): + _quantize_into(q_py, py) + return + _quantize_into(q_py, py) + for name in ref_names: + torch.testing.assert_close( + getattr(py, name), getattr(ref, name), rtol=0.0, atol=0.0, equal_nan=True + ) + + # Value check through dequantize(). Some layouts cannot dequantize at all + # (e.g. FP8 columnwise-only raises NotImplementedError) -- the C++-allocated + # reference defines what is supported, and when it raises, the bitwise + # buffer equality above already proves value parity. + try: + expected = ref.dequantize() + except NotImplementedError: + expected = None + if expected is not None: + torch.testing.assert_close(py.dequantize(), expected, rtol=0.0, atol=0.0) + + # ----- TensorProto ----- From 4feacdcd62bc848e52170b4a4cad86c7532a816e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 13 Jul 2026 13:57:12 +0200 Subject: [PATCH 09/28] [PyTorch] Simplify comment on captured requires_grad flags Address review: the change stands on its own as a correctness fix; drop the detailed (and imprecise) fake-impl/cudagraph justification. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index ecfb31a11d..80eceab15c 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -310,13 +310,9 @@ def _linear_forward_impl( # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) # 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). + # than the live tensors': under ``torch.compile`` the live flags are + # sometimes unreliable, and the fake impl keys its output arity off the + # same captured flags, so real and fake must agree. 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 From 5adb6e49c3e42d147a5b695960406796a5895562 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 13 Jul 2026 14:35:21 +0200 Subject: [PATCH 10/28] [PyTorch] Drop redundant None guards in wt_save alias dedup Address review: wt_save is known non-None past the first branch, so 'X is not None and wt_save is X' reduces to 'wt_save is X'. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 80eceab15c..ba144de03f 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -619,9 +619,9 @@ def _linear_forward_impl( 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: + elif wt_save is new_weight_workspace: wt_alias = "new_workspace" - elif args.weight_workspace is not None and wt_save is args.weight_workspace: + elif wt_save is args.weight_workspace: wt_alias = "weight_workspace" else: wt_alias = None From 40d2d7bccef99f3491851d0b230d3349eccca771 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 13 Jul 2026 14:43:33 +0200 Subject: [PATCH 11/28] [PyTorch] Use torch._prims_common.make_contiguous_strides_for Address review: replace the local _contiguous_stride helper with the torch one (stable at this path since v1.13); it also matches the ATen contiguous-stride convention for zero-size dims and handles SymInts. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/tensor_proto.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index 8881024104..69a4c5527a 100644 --- a/transformer_engine/pytorch/dynamo/tensor_proto.py +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -10,16 +10,7 @@ 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)) +from torch._prims_common import make_contiguous_strides_for @dataclass @@ -139,7 +130,9 @@ def create_tensor(self) -> torch.Tensor: 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)) + return storage_cls.__tensor_unflatten__( + inner, ctx, shape, make_contiguous_strides_for(shape) + ) def to_tensor_proto(tensor: Any) -> TensorProto: From 78e95a6ff43265ab8d0313035abe585cad97c4d8 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 13 Jul 2026 14:47:09 +0200 Subject: [PATCH 12/28] [PyTorch] Raise on update_usage of a non-quantized TensorProto Address review: a silent no-op diverges from the real object's behavior (plain torch.Tensor has no update_usage), which is exactly the class of fake/real mismatches the proto is meant to avoid. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/tensor_proto.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index 69a4c5527a..c2dfe201b0 100644 --- a/transformer_engine/pytorch/dynamo/tensor_proto.py +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -52,10 +52,11 @@ def update_usage( """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. + quantizer is never mutated. Raises on plain (non-quantized) protos -- + a real plain ``torch.Tensor`` has no ``update_usage`` either. """ if self.quantizer is None: - return + raise ValueError("update_usage called on a non-quantized TensorProto") self.quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) def inner_names(self) -> Tuple[str, ...]: From b3229cc3ee6bf7a48f095f3a621c3a3a7a931920 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 13 Jul 2026 15:07:43 +0200 Subject: [PATCH 13/28] [PyTorch] Collapse to_tensor_proto branches Address review: after the QuantizedTensorStorage.shape property the storage and plain-tensor paths differed only in getattr fallbacks (dtype/_dtype, _quantizer), which work uniformly for all input kinds; drop the isinstance branch and the local import it needed. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/tensor_proto.py | 22 +++++-------------- 1 file changed, 5 insertions(+), 17 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index c2dfe201b0..5cc5a19171 100644 --- a/transformer_engine/pytorch/dynamo/tensor_proto.py +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -143,26 +143,14 @@ def to_tensor_proto(tensor: Any) -> TensorProto: ``QuantizedTensor``. A *bare* storage exposes its (fake) dtype via ``_dtype`` rather than ``.dtype``. """ - from ..quantized_tensor import ( # pylint: disable=import-outside-toplevel - QuantizedTensorStorage, - ) - requires_grad = bool(getattr(tensor, "requires_grad", False)) - if isinstance(tensor, QuantizedTensorStorage): - dtype = getattr(tensor, "dtype", None) - if dtype is None: - dtype = getattr(tensor, "_dtype", None) - return TensorProto( - shape=tuple(tensor.shape), - dtype=dtype, - quantizer=getattr(tensor, "_quantizer", None), - requires_grad=requires_grad, - device=tensor.device, - ) + dtype = getattr(tensor, "dtype", None) + if dtype is None: + dtype = getattr(tensor, "_dtype", None) return TensorProto( shape=tuple(tensor.shape), - dtype=tensor.dtype, - quantizer=None, + dtype=dtype, + quantizer=getattr(tensor, "_quantizer", None), requires_grad=requires_grad, device=tensor.device, ) From 2b19368f07595761fcd7dee020523b0b4d2dcddb Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 13 Jul 2026 15:37:40 +0200 Subject: [PATCH 14/28] [PyTorch] Drop dead weight_fp8 fallback in fake backward Address review: the saved_weight slot is unconditionally aliased to the weight parameter in forward, so it is never None in backward. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index ba144de03f..dc21aa6c28 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1565,7 +1565,8 @@ def _linear_backward_impl_fake( "(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 + # The ``saved_weight`` slot is always aliased to the weight parameter. + weight = args.saved_weight out_dtype = args.activation_dtype out_features, in_features = weight.shape From c3dec8306affdcbf7015c587732154225bc1ff2b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 13 Jul 2026 15:42:55 +0200 Subject: [PATCH 15/28] [PyTorch] Drop redundant comment on saved_weight Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index dc21aa6c28..3f91a63b37 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1565,7 +1565,6 @@ def _linear_backward_impl_fake( "(fsdp_group is not None); use FSDP2 or MCore FSDP." ) - # The ``saved_weight`` slot is always aliased to the weight parameter. weight = args.saved_weight out_dtype = args.activation_dtype out_features, in_features = weight.shape From dd99206d797fc310b9e886d0ce0fb023a4964324 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 16:49:42 +0200 Subject: [PATCH 16/28] [PyTorch] Add TensorProto.assemble for rebuilding from ready-made buffers Factor the quantized-reassembly tail of create_tensor into assemble(), which rebuilds a tensor from already-materialized inner buffers (in inner_names() order). create_tensor becomes assemble(create_inner_tensors()). This is the reusable primitive the torch.compile custom-op boundary uses to rebuild an op's quantized outputs from its flat Tensor[] payload. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/tensor_proto.py | 36 ++++++++++++------- 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index 5cc5a19171..5a7c7e59c6 100644 --- a/transformer_engine/pytorch/dynamo/tensor_proto.py +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -109,32 +109,44 @@ def create_inner_tensors(self) -> List[torch.Tensor]: 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). + def assemble(self, inner_tensors: List[torch.Tensor]) -> torch.Tensor: + """Rebuild the tensor from ready-made ``inner_tensors`` (in :meth:`inner_names` + order). Shared by :meth:`create_tensor` (fresh buffers) and the custom-op + boundary (buffers arriving from an op's flat ``Tensor[]`` payload). - Quantized protos reassemble the :meth:`create_inner_tensors` buffers via - the storage's ``__tensor_unflatten__``. + Non-quantized protos are the single inner tensor as-is; quantized protos + are reassembled into the storage/wrapper via ``__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, - ) + return inner_tensors[0] 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())) + inner = dict(zip(self.inner_names(), inner_tensors)) storage_cls = _STORAGE_REGISTRY[ctx["cls"]] return storage_cls.__tensor_unflatten__( inner, ctx, shape, make_contiguous_strides_for(shape) ) + def create_tensor(self) -> torch.Tensor: + """Materialize an (uninitialized) tensor matching this proto (traceable). + + Quantized protos reassemble freshly-allocated :meth:`create_inner_tensors` + buffers via :meth:`assemble`. + """ + 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, + ) + return self.assemble(self.create_inner_tensors()) + def to_tensor_proto(tensor: Any) -> TensorProto: """Build a :class:`TensorProto` describing ``tensor``. From 7e7201393b9412cf3275d9fb3f719e8de90d9545 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 28 Jul 2026 14:28:01 +0200 Subject: [PATCH 17/28] Fix blockwise alloc parity test and drop dead CUDA skips test_python_alloc_matches_cpp_make_empty compared buffers the quantize kernel never writes: the scale-inv padding is allocated uninitialized by both paths, so the bit-exact comparison saw random bytes and failed on H100/B200 for fp8_blockwise. Zero every buffer before quantizing, so the comparison covers kernel output only. Also drop the param-level skips on the nvfp4 entries of _PROTO_QUANTIZERS and _VALUE_QUANTIZERS. is_fp8_available() and friends run at import time and go through torch.cuda.current_device(), so this module cannot be collected without CUDA at all and skipif(not torch.cuda.is_available()) never fires; the same goes for the torch.cuda.is_available() halves of the _hw_available() guards. Gating nvfp4 on nvfp4_available was also inconsistent with MXFP8 and blockwise, which are gated at runtime and only in the tests that run a kernel -- the allocation primitives themselves are pure Python and describe the layout on any HW. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 33 ++++++++++++----------------- 1 file changed, 13 insertions(+), 20 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 42119616ca..c95f3ae36b 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -719,14 +719,7 @@ def _hw_available(quantizer): pytest.param(_mxfp8, id="mxfp8"), pytest.param(_blockwise, id="float8_blockwise"), pytest.param(_current_scaling, id="float8_current_scaling"), - pytest.param( - _nvfp4, - id="nvfp4", - marks=pytest.mark.skipif( - not torch.cuda.is_available(), - reason="NVFP4Quantizer requires CUDA to construct", - ), - ), + pytest.param(_nvfp4, id="nvfp4"), ] @@ -748,7 +741,7 @@ def test_quantizer_value_object(factory): # 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): + if _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) @@ -817,7 +810,7 @@ def test_quantizer_value_object_fullgraph(factory): unlike merely passing the quantizer through. """ q = factory() - if not (torch.cuda.is_available() and _hw_available(q)): + if not _hw_available(q): pytest.skip("format not supported on this HW") op = _QDQ_OPS[type(q)] @@ -839,19 +832,13 @@ def fn(inp): # (factory, logical shape) -- shapes respect MXFP8 (mult. of 32) / blockwise (128) # / NVFP4 (mult. of 16) constraints. +# Format support is gated at runtime, in the tests that run a kernel; the rest is +# pure Python and works on any HW. _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", - ), - ), + pytest.param(_nvfp4, (64, 128), id="nvfp4"), ] @@ -1008,7 +995,7 @@ def make_quantizer(): return q q_ref = make_quantizer() - if not (torch.cuda.is_available() and _hw_available(q_ref)): + if not _hw_available(q_ref): pytest.skip("format not supported on this HW") q_py = make_quantizer() @@ -1071,6 +1058,12 @@ def _quantize_into(quantizer, dst): else: quantizer.update_quantized(x, dst) + # Scale-inv padding is never written by the kernel and both paths allocate it + # uninitialized; zero it so the comparison below covers only kernel output. + for name in ref_names: + getattr(ref, name).zero_() + getattr(py, name).zero_() + # Some combos are rejected by the quantize kernel itself regardless of who # allocated the tensor (e.g. FP8 current-scaling columnwise-only on # TN-capable archs: there is no rowwise data buffer and From 65cc3349c9c4c4869df2c7e3a397edbb6dfd8e81 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 28 Jul 2026 14:39:26 +0200 Subject: [PATCH 18/28] Mirror quantize_weight's cache semantics in the Linear fake forward _linear_forward_impl_fake diverged from quantize_weight on the weight workspace in three ways: - it produced a new workspace only when update_ws was true, but the real cache-miss path returns (out, out) whenever cache=True, regardless of update_workspace; a first call with is_first_microbatch=False therefore lost the workspace and the "new_workspace" saved-weight alias; - it treated any non-None cached workspace as a hit, while the real path runs _is_weight_workspace_valid() first and falls through to a miss when the cached buffer layout no longer matches the quantizer's usage; - it kept quantizer.internal, so the descriptor resolved to a bare storage class, while the real path quantizes persistent workspaces with internal=False and caches wrapper tensors. On a cache hit the weightmat is now the workspace descriptor itself, and on a miss with cache_weight it is the same proto object returned as the new workspace, matching quantize_weight's aliasing. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 26 +++++++++++++-------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 75692cce77..4aa788ee1f 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -25,6 +25,7 @@ is_ub_initialized, using_cublasmp_backend, quantize_weight, + _is_weight_workspace_valid, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -74,7 +75,7 @@ prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorProto +from ..dynamo import TensorProto, to_tensor_proto from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -751,21 +752,26 @@ def _linear_forward_impl_fake( 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( + workspace = args.weight_workspace + if workspace is not None and weight_quantizer is not None: + if not _is_weight_workspace_valid(workspace, weight_quantizer): + workspace = None + if workspace is not None: + weightmat = to_tensor_proto(workspace) + else: + weightmat = TensorProto( shape=tuple(weight.shape), dtype=activation_dtype, quantizer=weight_quantizer, device=weight.device, ) + if args.cache_weight: + # Persistent cache entries are wrappers, not bare storages. + if weightmat.quantizer is not None: + weightmat.quantizer.internal = False + new_weight_workspace = weightmat + weightmat.update_usage(rowwise_usage=True) else: weightmat_aliases_weight = weight.dtype == activation_dtype weightmat = TensorProto( From 3d102222485fc422992138f38f8a3177a5e79a20 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 28 Jul 2026 14:42:25 +0200 Subject: [PATCH 19/28] Honor the dequantized backward override in the Linear fake forward Eager forward forces save_original_input=False for backward_override="dequantized", but the fake only handled "high_precision". With save_original_input=True and that override, the fake aliased the original input into saved-tensor slot 0 while eager saved a quantized input with rowwise-only usage, so the saved payload layout and the compiled backward setup disagreed. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 4aa788ee1f..f82aa95b39 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -691,6 +691,8 @@ def _linear_forward_impl_fake( save_original_input = args.save_original_input if args.backward_override == "high_precision": save_original_input = True + elif args.backward_override == "dequantized": + save_original_input = False out_features, _ = weight.shape backward_needs_input = is_grad_enabled and args.weight_requires_grad From a257fa3482e628d63d4e3c73ffd9fd008369b819 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 28 Jul 2026 14:44:56 +0200 Subject: [PATCH 20/28] Include the bias in the Linear fake output's differentiability The output proto's requires_grad considered only the input and the weight, so a frozen input and weight with a trainable bias described the output as non-differentiable while eager _Linear.apply produces a differentiable one. bias_requires_grad is already False when there is no bias. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index f82aa95b39..64a4d3463c 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -795,7 +795,8 @@ def _linear_forward_impl_fake( 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), + requires_grad=is_grad_enabled + and (args.input_requires_grad or args.weight_requires_grad or args.bias_requires_grad), device=inp.device, ) From 0414586567d0db8189dd68baf0db3aab9f2e5c22 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 28 Jul 2026 15:41:30 +0200 Subject: [PATCH 21/28] Move the Linear fake impls out to the custom-op branch _linear_forward_impl_fake / _linear_backward_impl_fake, and the eager-side changes that existed only to support them (reading the requires_grad flags off LinearFwdArgs, the new_workspace/weight_workspace alias dedup and the _linear_setup_ctx signature carrying (out, new_weight_workspace)), have no caller in this PR: nothing registers them as a custom op's fake, so nothing exercises them here. They belong with the custom-op registration that consumes them. This PR is left as the TensorProto mechanism proper -- the proto, the storage flatten protocol and the pure-Python quantizer allocation hooks -- which the new tests do cover. linear.py returns to its upstream state. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 322 +------------------- 1 file changed, 10 insertions(+), 312 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 64a4d3463c..6b0f941bc4 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -25,7 +25,6 @@ is_ub_initialized, using_cublasmp_backend, quantize_weight, - _is_weight_workspace_valid, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -75,7 +74,6 @@ prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorProto, to_tensor_proto from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -91,11 +89,6 @@ __all__ = ["Linear"] -# Fields with this union may hold a *bare* ``QuantizedTensorStorage`` (see its -# docstring), not just a plain / subclass tensor. The annotation is also -# machine-read when the args bag becomes a torch.compile custom op (follow-up -# PR): such fields get flatten/unflatten slots so a bare storage can cross the -# op boundary -- a plain ``Tensor`` slot cannot carry it. TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] @@ -105,7 +98,7 @@ class LinearFwdArgs: # --- Differentiable tensors (also passed positionally to autograd) --- weight: TensorOrQuantized - inp: TensorOrQuantized + inp: torch.Tensor bias: Optional[torch.Tensor] # --- Non-differentiable cached tensors --- @@ -317,11 +310,7 @@ def _linear_forward_impl( # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) - # Use the requires-grad flags captured into ``args`` at op-call time rather - # than the live tensors': under ``torch.compile`` the live flags are - # sometimes unreliable, and the fake impl keys its output arity off the - # same captured flags, so real and fake must agree. - backward_needs_input = is_grad_enabled and args.weight_requires_grad + backward_needs_input = is_grad_enabled and weight.requires_grad with_input_all_gather_nccl = ( parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop ) @@ -446,7 +435,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 args.input_requires_grad and not is_fsdp2 + columnwise_usage = is_grad_enabled and inp.requires_grad and not is_fsdp2 if backward_override is not None: columnwise_usage = False if not columnwise_usage: @@ -629,24 +618,12 @@ def _linear_forward_impl( if is_dist_weight: wt_save = None - # 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 wt_save is new_weight_workspace: - wt_alias = "new_workspace" - elif wt_save is args.weight_workspace: - wt_alias = "weight_workspace" - else: - wt_alias = 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. saved_tensor_aliases = ( "inp" if saved_inputmat is inp else None, - wt_alias, + "weight" if wt_save is weight else None, "weight", # ``saved_weight`` slot is always the weight parameter "bias" if bias is not None else None, ) @@ -665,222 +642,10 @@ 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 - elif args.backward_override == "dequantized": - save_original_input = False - - 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_is_storage = True - workspace = args.weight_workspace - if workspace is not None and weight_quantizer is not None: - if not _is_weight_workspace_valid(workspace, weight_quantizer): - workspace = None - if workspace is not None: - weightmat = to_tensor_proto(workspace) - else: - weightmat = TensorProto( - shape=tuple(weight.shape), - dtype=activation_dtype, - quantizer=weight_quantizer, - device=weight.device, - ) - if args.cache_weight: - # Persistent cache entries are wrappers, not bare storages. - if weightmat.quantizer is not None: - weightmat.quantizer.internal = False - new_weight_workspace = weightmat - weightmat.update_usage(rowwise_usage=True) - 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 or args.bias_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, - fwd_outputs: Tuple[Any, ...], + out: torch.Tensor, ctx_attrs: Dict, tensors_to_save_from_forward: Tuple[Any, ...], ) -> Tuple[Any, ...]: @@ -893,8 +658,7 @@ def _linear_setup_ctx( for FSDP2 re-quantization) without having to mutate the structured metadata returned by ``prepare_for_saving``. """ - # ``fwd_outputs`` are the op's user outputs ``(out, new_weight_workspace)``; - # only ``new_weight_workspace`` is needed here, to rebuild the deduped weight. + del out # No-op; kept for symmetry with the compile-time helper signature. inp = fwd_args.inp weight = fwd_args.weight @@ -986,10 +750,6 @@ 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": @@ -1603,68 +1363,6 @@ 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 - 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. @@ -1704,7 +1402,7 @@ def forward( tensors_to_save_from_setup = _linear_setup_ctx( bwd_args, fwd_args, - (out, new_weight_workspace), + out, ctx_attrs, tensors_to_save_from_forward, ) From d5e1f20d4909d5ac883bf1cccd6ef9649aaccc34 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 29 Jul 2026 09:58:38 +0200 Subject: [PATCH 22/28] Keep attributes attached to quantized parameters across _apply Quantized tensors now implement the wrapper-subclass flatten protocol, so nn.Module._apply moves them with torch.utils.swap_tensors instead of the `param.data = ...` path. The swap exchanges the parameter's entire __dict__: that is how the inner buffers reach the surviving object, but it also carries off everything attached to the parameter from the outside. TE relies on several such attributes: _high_precision_init_val and its two accessors (quantized_model_init(preserve_high_precision_init_val=True)), plus main_grad, grad_added_to_main_grad and overwrite_main_grad, which Megatron-Core attaches. They survived before only because `param.data = ...` is a no-op for a wrapper subclass -- the outer tensor is a zero-storage shell and the assignment never touched __dict__, so device moves silently did nothing at all. Snapshot the parameters' __dict__ before delegating to nn.Module._apply and restore the entries the swap dropped, rebinding bound accessors to the surviving parameter. Entries still present afterwards are the tensor's own state, where the post-swap value is the correct one. Covers the two test_sanity grouped-linear high-precision-init tests that broke on B200, and adds a direct test over .cuda() / .cpu() / .half(). Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_sanity.py | 35 +++++++++++++++++++++++ transformer_engine/pytorch/module/base.py | 26 +++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index dff94cef9d..de62e56e53 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -1154,6 +1154,41 @@ def test_quantized_model_init_high_precision_init_val(): ), "clear_high_precision_init_val() not work" +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("move", ["cuda", "cpu", "half"]) +def test_quantized_param_attrs_survive_apply(move): + """Attributes attached to a quantized parameter survive nn.Module._apply. + + Quantized parameters implement the flatten protocol, so ``_apply`` moves them + with ``swap_tensors``, which exchanges the parameter's whole ``__dict__``. + Anything attached from the outside rides out on the discarded tensor unless + the module re-attaches it. + """ + with quantized_model_init(preserve_high_precision_init_val=True): + model = Linear(64, 64) + + weight = model.weight + expected = weight.get_high_precision_init_val() + weight.probe_attr = "attached-from-outside" + + if move == "cuda": + model = model.cuda() + elif move == "cpu": + model = model.cpu() + else: + model = model.half() + + weight = model.weight + assert hasattr(weight, "get_high_precision_init_val"), f"accessor lost by .{move}()" + assert hasattr(weight, "clear_high_precision_init_val"), f"accessor lost by .{move}()" + torch.testing.assert_close(weight.get_high_precision_init_val(), expected, rtol=0, atol=0) + assert weight.probe_attr == "attached-from-outside", f"custom attr lost by .{move}()" + + # The accessor must read the surviving parameter, not the discarded one. + weight.clear_high_precision_init_val() + assert weight.get_high_precision_init_val() is None + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_grouped_linear_single_param_preserves_high_precision_init(monkeypatch): """Grouped MXFP8 and discrete weights produce identical FP32 master initialization.""" diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index a0cdfa1378..5be35f0790 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -927,6 +927,32 @@ def module_setattr(self, name: str, value: Any) -> None: """ super().__setattr__(name, value) + def _apply(self, *args, **kwargs): + """Re-attach attributes that ``swap_tensors`` moves off a quantized parameter. + + ``_apply`` moves wrapper subclasses by exchanging the parameter's whole + ``__dict__``, which carries the inner buffers over but takes externally + attached state (``_high_precision_init_val``, ``main_grad``, ...) with it. + """ + snapshots = { + name: (param, dict(param.__dict__)) + for name, param in self._parameters.items() + if isinstance(param, QuantizedTensorStorage) + } + out = super()._apply(*args, **kwargs) + for name, (old_param, attrs) in snapshots.items(): + new_param = self._parameters.get(name) + if new_param is None: + continue + for key, value in attrs.items(): + # Still present -> tensor state; the post-swap value is the right one. + if key in new_param.__dict__: + continue + if isinstance(value, MethodType) and value.__self__ is old_param: + value = MethodType(value.__func__, new_param) + setattr(new_param, key, value) + return out + @property def output_quantizer_role(self) -> Optional[QuantizerRole]: """Caller-configurable :class:`QuantizerRole` for the forward output quantizer. From 6e61d36e0411747380e543fab27de3729625f868 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 29 Jul 2026 13:53:03 +0200 Subject: [PATCH 23/28] Fail loudly if a buffer vanishes from a parameter during _apply The restore loop keys off "present after the swap": what survived is the tensor's own state, what did not is an externally attached annotation. That holds only as long as every declared buffer really is present afterwards. If one were not, the loop would quietly put the pre-move value back and splice a buffer from the old device (or from before a dtype conversion) into the moved parameter -- silently wrong numerics rather than a crash. Raise instead when a name from _FLATTEN_TENSOR_BUFFERS is about to be restored. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/base.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 5be35f0790..978eb26023 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -944,10 +944,17 @@ def _apply(self, *args, **kwargs): new_param = self._parameters.get(name) if new_param is None: continue + buffers = {attr for attr, _ in type(old_param)._FLATTEN_TENSOR_BUFFERS} for key, value in attrs.items(): # Still present -> tensor state; the post-swap value is the right one. if key in new_param.__dict__: continue + if key in buffers: + raise RuntimeError( + f"Buffer {key} disappeared from {type(self).__name__}.{name} during" + " _apply; restoring the pre-move value would splice stale storage" + " into the parameter" + ) if isinstance(value, MethodType) and value.__self__ is old_param: value = MethodType(value.__func__, new_param) setattr(new_param, key, value) From a24e4e006981323e8cfd458bb153ae9af9cd6c83 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 13:20:17 +0200 Subject: [PATCH 24/28] Raise when a parameter vanishes during _apply instead of skipping it nn.Module._apply only assigns to self._parameters, never removes entries, so a missing parameter after it returns means something unexpected happened. Skipping it silently dropped every attribute attached to that parameter -- the failure this override exists to prevent. Match the buffer check and fail loudly. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/base.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 978eb26023..c346f34236 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -943,7 +943,10 @@ def _apply(self, *args, **kwargs): for name, (old_param, attrs) in snapshots.items(): new_param = self._parameters.get(name) if new_param is None: - continue + raise RuntimeError( + f"{type(self).__name__}.{name} disappeared during _apply; the state" + " attached to it cannot be restored" + ) buffers = {attr for attr, _ in type(old_param)._FLATTEN_TENSOR_BUFFERS} for key, value in attrs.items(): # Still present -> tensor state; the post-swap value is the right one. From a8b88731a5da4e1e40e3feee2113a4fdfb43bbbe Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 13:24:04 +0200 Subject: [PATCH 25/28] Drop the vanished-buffer check from _apply Reverts 6e61d36e. The check guarded a case that cannot arise today: the storages always set every declared buffer attribute, to None when unused, so the key is present whatever the usage flags say. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/base.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index c346f34236..0e922fa072 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -947,17 +947,10 @@ def _apply(self, *args, **kwargs): f"{type(self).__name__}.{name} disappeared during _apply; the state" " attached to it cannot be restored" ) - buffers = {attr for attr, _ in type(old_param)._FLATTEN_TENSOR_BUFFERS} for key, value in attrs.items(): # Still present -> tensor state; the post-swap value is the right one. if key in new_param.__dict__: continue - if key in buffers: - raise RuntimeError( - f"Buffer {key} disappeared from {type(self).__name__}.{name} during" - " _apply; restoring the pre-move value would splice stale storage" - " into the parameter" - ) if isinstance(value, MethodType) and value.__self__ is old_param: value = MethodType(value.__func__, new_param) setattr(new_param, key, value) From 6bb030f507040aa3adb8ae0e9a3282905c8dfb8f Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 13:47:51 +0200 Subject: [PATCH 26/28] Declare flat buffers on the field instead of in a parallel tuple Each storage class listed its tensor buffers twice: once as a field annotation, once as an (attribute, constructor kwarg) pair in _FLATTEN_TENSOR_BUFFERS, in a different order and further down the file. Adding a buffer meant remembering both. Mark the field instead -- _scale_inv: Annotated[torch.Tensor, Buffer("fp8_scale_inv")] -- and collect the declarations in __init_subclass__, which already runs there for the storage registry. _FLATTEN_TENSOR_BUFFERS survives as the derived attribute, so every consumer is untouched, and the collected values are identical to the hand-written tuples for all nine registered classes. Signed-off-by: Pawel Gadzinski --- .../pytorch/quantized_tensor.py | 30 ++++++++++++--- .../float8_blockwise_tensor_storage.py | 21 +++------- .../tensor/storage/float8_tensor_storage.py | 17 +++------ .../tensor/storage/mxfp8_tensor_storage.py | 27 ++++--------- .../tensor/storage/nvfp4_tensor_storage.py | 38 ++++++------------- 5 files changed, 55 insertions(+), 78 deletions(-) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 9c3c83c6d9..6f431bbedf 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -5,7 +5,7 @@ """Pure Python base classes for quantization.""" from __future__ import annotations -from typing import Optional, Tuple, Iterable, Any, Dict, Union +from typing import NamedTuple, Optional, Tuple, Iterable, Any, Dict, Union, get_type_hints import abc import warnings import math @@ -23,7 +23,6 @@ _stride_from_shape, ) - # Custom ops that should pass through __torch_dispatch__ without unwrapping # QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that # handle quantized tensors internally. @@ -34,6 +33,27 @@ _STORAGE_REGISTRY: Dict[str, type] = {} +class Buffer(NamedTuple): + """Marks a storage field as a flat tensor buffer. + + Annotate the field with it -- ``_scale_inv: Annotated[torch.Tensor, + Buffer("fp8_scale_inv")]`` -- and ``__init_subclass__`` collects the + declarations into ``_FLATTEN_TENSOR_BUFFERS``, in field order. + """ + + ctor_kwarg: str + + +def _collect_buffer_fields(cls: type) -> Tuple[Tuple[str, str], ...]: + """Buffers a storage class declares, as ``(attribute, constructor kwarg)``.""" + fields = [] + for attr, hint in get_type_hints(cls, include_extras=True).items(): + for meta in getattr(hint, "__metadata__", ()): + if isinstance(meta, Buffer): + fields.append((attr, meta.ctor_kwarg)) + return tuple(fields) + + class QuantizedTensorStorage: r"""Base class for all TensorStorage classes. @@ -153,9 +173,8 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: # ----- 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. + # Collected from the subclasses' :class:`Buffer` field annotations; 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: @@ -163,6 +182,7 @@ def __init_subclass__(cls, **kwargs) -> None: # 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 + cls._FLATTEN_TENSOR_BUFFERS = _collect_buffer_fields(cls) def _flatten_nontensor_kwargs(self) -> Dict[str, Any]: """Non-tensor constructor kwargs (scalars, dtype, quantizer).""" 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 7cfed561b2..d671122628 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -7,12 +7,12 @@ from __future__ import annotations from collections.abc import Iterable import math -from typing import Optional, Dict, Any, Tuple, Union +from typing import Annotated, Optional, Dict, Any, Tuple, Union import torch import transformer_engine_torch as tex -from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from ...quantized_tensor import Buffer, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType_To_Torch, DType @@ -119,23 +119,14 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): be instantiated directly for performance-critical internal usage. """ - _rowwise_data: Optional[torch.Tensor] - _columnwise_data: Optional[torch.Tensor] + _rowwise_data: Annotated[Optional[torch.Tensor], Buffer("rowwise_data")] + _rowwise_scale_inv: Annotated[Optional[torch.Tensor], Buffer("rowwise_scale_inv")] + _columnwise_data: Annotated[Optional[torch.Tensor], Buffer("columnwise_data")] + _columnwise_scale_inv: Annotated[Optional[torch.Tensor], Buffer("columnwise_scale_inv")] _quantizer: Quantizer _fp8_dtype: DType - _rowwise_scale_inv: Optional[torch.Tensor] - _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 5180562cc2..9b4994a77f 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -5,12 +5,12 @@ """Mixin class holding data specific for Float8Tensor""" from __future__ import annotations -from typing import Any, Dict, Optional, Tuple, Union +from typing import Annotated, Any, Dict, Optional, Tuple, Union import torch import transformer_engine_torch as tex -from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from ...quantized_tensor import Buffer, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch, DType @@ -66,22 +66,15 @@ class Float8TensorStorage(QuantizedTensorStorage): """ - _data: Optional[torch.Tensor] + _data: Annotated[Optional[torch.Tensor], Buffer("data")] _quantizer: Optional[Quantizer] _fp8_dtype: DType - _scale_inv: torch.Tensor # FP8 transpose cache - _transpose: Optional[torch.Tensor] + _transpose: Annotated[Optional[torch.Tensor], Buffer("data_transpose")] _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"), - ) + _scale_inv: Annotated[torch.Tensor, Buffer("fp8_scale_inv")] def __new__( cls, diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 27a44e55fb..936c7f89a4 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -5,14 +5,14 @@ """Mixin class holding data specific for MXFP8Tensor""" from __future__ import annotations -from typing import Optional, Dict, Any, Tuple, Union +from typing import Annotated, Optional, Dict, Any, Tuple, Union from collections.abc import Iterable import math import torch import transformer_engine_torch as tex -from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from ...quantized_tensor import Buffer, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType @@ -66,14 +66,12 @@ class MXFP8TensorStorage(QuantizedTensorStorage): """ - # Row-scaled FP8 data - _rowwise_data: Optional[torch.Tensor] - # Column-scaled FP8 data - _columnwise_data: Optional[torch.Tensor] - # Scaling factors for row-scaled FP8 data - _rowwise_scale_inv: torch.Tensor - # Scaling factors for column-scaled FP8 data - _columnwise_scale_inv: torch.Tensor + # Row-scaled FP8 data and its scaling factors + _rowwise_data: Annotated[Optional[torch.Tensor], Buffer("rowwise_data")] + _rowwise_scale_inv: Annotated[torch.Tensor, Buffer("rowwise_scale_inv")] + # Column-scaled FP8 data and its scaling factors + _columnwise_data: Annotated[Optional[torch.Tensor], Buffer("columnwise_data")] + _columnwise_scale_inv: Annotated[torch.Tensor, Buffer("columnwise_scale_inv")] # Builder class for casting to MXFP8 _quantizer: Optional[Quantizer] @@ -83,15 +81,6 @@ 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 b43c5e93f0..b15565ad1f 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -8,14 +8,14 @@ from collections.abc import Iterable import functools import math -from typing import Any, Dict, Optional, Tuple, Union +from typing import Annotated, Any, Dict, Optional, Tuple, Union import warnings import torch import transformer_engine_torch as tex -from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from ...quantized_tensor import Buffer, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType @@ -80,20 +80,15 @@ class NVFP4TensorStorage(QuantizedTensorStorage): """ - # Row-scaled FP4 data - _rowwise_data: Optional[torch.Tensor] - # Column-scaled FP4 data - _columnwise_data: Optional[torch.Tensor] - # Block scaling factors for row-scaled FP4 data - _rowwise_scale_inv: torch.Tensor - # Block scaling factors for column-scaled FP4 data - _columnwise_scale_inv: torch.Tensor - # Input absolute maximum value (used to compute tensor scale for - # row-scaled FP4 data) - _amax_rowwise: torch.Tensor - # Input absolute maximum value (used to compute tensor scale for - # column-scaled FP4 data) - _amax_columnwise: torch.Tensor + # Row-scaled FP4 data and its block scaling factors + _rowwise_data: Annotated[Optional[torch.Tensor], Buffer("rowwise_data")] + _rowwise_scale_inv: Annotated[torch.Tensor, Buffer("rowwise_scale_inv")] + # Column-scaled FP4 data and its block scaling factors + _columnwise_data: Annotated[Optional[torch.Tensor], Buffer("columnwise_data")] + _columnwise_scale_inv: Annotated[torch.Tensor, Buffer("columnwise_scale_inv")] + # Input absolute maximum values, used to compute the tensor scale + _amax_rowwise: Annotated[torch.Tensor, Buffer("amax_rowwise")] + _amax_columnwise: Annotated[torch.Tensor, Buffer("amax_columnwise")] # Builder class for casting to MXFP8 _quantizer: Optional[Quantizer] @@ -109,17 +104,6 @@ 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], From e0bbb187d6700ed5642dc2eeafaf74e9946160be Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 14:04:00 +0200 Subject: [PATCH 27/28] Name the flat buffers after PyTorch's own term for them "Buffer" collides with nn.Module's buffers, which are a different thing, and _FLATTEN_TENSOR_BUFFERS named a consumer (__tensor_flatten__) rather than the thing itself -- the list has four of them. PyTorch calls exactly this concept "inner tensors", which TensorProto.inner_names() already follows. Also drop the underscore from the two hooks every quantizer has to implement. They were the only members of the extension contract marked private, which is why the tests needed seven protected-access waivers to call them; the members nobody overrides (alloc_tensors, create_metadata) were public already. Buffer -> InnerTensor _FLATTEN_TENSOR_BUFFERS -> _INNER_TENSORS _describe_buffers -> inner_tensor_specs _storage_metadata -> storage_metadata Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 18 ++++----- .../pytorch/dynamo/tensor_proto.py | 12 +++--- .../pytorch/quantized_tensor.py | 38 +++++++++---------- .../pytorch/tensor/float8_blockwise_tensor.py | 4 +- .../pytorch/tensor/float8_tensor.py | 4 +- .../pytorch/tensor/mxfp8_tensor.py | 4 +- .../pytorch/tensor/nvfp4_tensor.py | 6 +-- .../float8_blockwise_tensor_storage.py | 10 ++--- .../tensor/storage/float8_tensor_storage.py | 8 ++-- .../tensor/storage/mxfp8_tensor_storage.py | 10 ++--- .../tensor/storage/nvfp4_tensor_storage.py | 14 +++---- 11 files changed, 62 insertions(+), 66 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index c95f3ae36b..2d25c33f42 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -848,7 +848,7 @@ def _build_from_primitives(quantizer, shape, dtype, device="cpu"): ``__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 + names = tuple(quantizer.inner_tensor_specs(shape)) ctx = quantizer.create_metadata(shape, dtype=dtype) buffers = quantizer.alloc_tensors(shape, device=device) inner = {name: buffers[name] for name in names} @@ -895,7 +895,7 @@ 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 + names = tuple(q.inner_tensor_specs(shape)) def fn(x): t = _build_from_primitives(q, shape, x.dtype, device=x.device) @@ -916,7 +916,7 @@ def fn(x): 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 + bufs = q.inner_tensor_specs(shape) with FakeTensorMode(): alloc = q.alloc_tensors(shape, device="cpu") assert set(alloc) == set(bufs) @@ -937,7 +937,7 @@ def test_storage_flatten_unflatten_roundtrip(factory, shape): _skip_if_dequantize_unsupported(q) tensor = _build_from_primitives(q, shape, torch.bfloat16) - names = tuple(q._describe_buffers(shape)) # pylint: disable=protected-access + names = tuple(q.inner_tensor_specs(shape)) # Fill buffers with deterministic data (empty() may contain NaNs) so the # round-trip can be checked by value via dequantize(). for name in names: @@ -974,7 +974,7 @@ def test_python_alloc_matches_cpp_make_empty(factory, shape, rowwise, columnwise Builds the same quantized tensor twice: via ``Quantizer.make_empty`` (``tex.create_empty_quantized_tensor``, the C++ path) and via the Python - primitives ``_describe_buffers`` + ``create_metadata`` + ``alloc_tensors`` + primitives ``inner_tensor_specs`` + ``create_metadata`` + ``alloc_tensors`` + ``__tensor_unflatten__`` (exactly what ``TensorProto.create_tensor`` does). Checks: @@ -1108,8 +1108,8 @@ def test_tensor_proto_matches_primitives(factory, shape): # 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 + # inner_names + create_inner_tensors match inner_tensor_specs. + bufs = q.inner_tensor_specs(shape) names = tuple(bufs) assert proto.inner_names() == names inner = proto.create_inner_tensors() @@ -1194,9 +1194,7 @@ def test_to_tensor_proto_quantized(factory, shape): 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 + assert proto.inner_names() == tuple(q.inner_tensor_specs(shape)) # 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/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index 5a7c7e59c6..0072bde4e3 100644 --- a/transformer_engine/pytorch/dynamo/tensor_proto.py +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -63,8 +63,8 @@ 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 + ``__tensor_flatten__`` -- i.e. ``_INNER_TENSORS`` order, keeping + only the present buffers. ``inner_tensor_specs`` 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. @@ -72,14 +72,14 @@ def inner_names(self) -> Tuple[str, ...]: 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] + described = list(self.quantizer.inner_tensor_specs(tuple(self.shape)).keys()) + storage_cls = self.quantizer.storage_metadata(self.dtype)["cls"] + flatten_order = [attr for attr, _ in storage_cls._INNER_TENSORS] 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 " + f"_INNER_TENSORS {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) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 6f431bbedf..ef23ff893f 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -33,23 +33,23 @@ _STORAGE_REGISTRY: Dict[str, type] = {} -class Buffer(NamedTuple): +class InnerTensor(NamedTuple): """Marks a storage field as a flat tensor buffer. Annotate the field with it -- ``_scale_inv: Annotated[torch.Tensor, - Buffer("fp8_scale_inv")]`` -- and ``__init_subclass__`` collects the - declarations into ``_FLATTEN_TENSOR_BUFFERS``, in field order. + InnerTensor("fp8_scale_inv")]`` -- and ``__init_subclass__`` collects the + declarations into ``_INNER_TENSORS``, in field order. """ ctor_kwarg: str -def _collect_buffer_fields(cls: type) -> Tuple[Tuple[str, str], ...]: +def _collect_inner_tensor_fields(cls: type) -> Tuple[Tuple[str, str], ...]: """Buffers a storage class declares, as ``(attribute, constructor kwarg)``.""" fields = [] for attr, hint in get_type_hints(cls, include_extras=True).items(): for meta in getattr(hint, "__metadata__", ()): - if isinstance(meta, Buffer): + if isinstance(meta, InnerTensor): fields.append((attr, meta.ctor_kwarg)) return tuple(fields) @@ -173,27 +173,25 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: # ----- PyTorch subclass flatten protocol (torch.compile / TensorProto) ----- - # Collected from the subclasses' :class:`Buffer` field annotations; everything + # Collected from the subclasses' :class:`InnerTensor` field annotations; everything # else returned by :meth:`get_metadata` is treated as non-tensor context. - _FLATTEN_TENSOR_BUFFERS: Tuple[Tuple[str, str], ...] = () + _INNER_TENSORS: 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 - cls._FLATTEN_TENSOR_BUFFERS = _collect_buffer_fields(cls) + cls._INNER_TENSORS = _collect_inner_tensor_fields(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} + tensor_kwargs = {kwarg for _, kwarg in self._INNER_TENSORS} 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 - ] + present = [attr for attr, _ in self._INNER_TENSORS if getattr(self, attr) is not None] ctx = { "cls": type(self).__qualname__, "is_tensor": isinstance(self, QuantizedTensor), @@ -215,7 +213,7 @@ def __tensor_unflatten__( 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: + for attr, kwarg in cls._INNER_TENSORS: kwargs[kwarg] = inner_tensors.get(attr) if not ctx["is_tensor"]: return cls(**kwargs) @@ -450,21 +448,21 @@ def make_empty( # ----- Data-free buffer/metadata primitives backing TensorProto ----- - def _describe_buffers( + def inner_tensor_specs( 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. + ``_INNER_TENSORS`` and respect the quantizer's usage flags. """ raise NotImplementedError( - f"{self.__class__.__name__} does not implement _describe_buffers; " + f"{self.__class__.__name__} does not implement inner_tensor_specs; " "it cannot be used with TensorProto / pure-Python allocation" ) - def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + 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 @@ -474,7 +472,7 @@ def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: ``fp8_dtype``, ``quantizer``, ``fake_dtype``). """ raise NotImplementedError( - f"{self.__class__.__name__} does not implement _storage_metadata; " + f"{self.__class__.__name__} does not implement storage_metadata; " "it cannot be used with TensorProto / pure-Python allocation" ) @@ -492,7 +490,7 @@ def alloc_tensors( 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() + for attr, (buf_shape, buf_dtype) in self.inner_tensor_specs(tuple(shape)).items() } def create_metadata( @@ -505,7 +503,7 @@ def create_metadata( """Build the data-free ``__tensor_unflatten__`` context describing the quantized tensor this quantizer would produce for ``shape`` / ``dtype``. """ - meta = self._storage_metadata(dtype) + meta = self.storage_metadata(dtype) return { "cls": meta["cls"].__qualname__, "is_tensor": not self.internal, diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index a90cd42ac0..2beb487cd9 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -75,7 +75,7 @@ def copy(self) -> Float8BlockQuantizer: # ----- TensorProto / pure-Python allocation ----- - def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + def storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: return { "cls": Float8BlockwiseQTensorStorage if self.internal else Float8BlockwiseQTensor, "nontensor_kwargs": { @@ -86,7 +86,7 @@ def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: }, } - def _describe_buffers( + def inner_tensor_specs( self, shape: Tuple[int, ...] ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: shape = tuple(shape) diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index a4da60d7f3..0bdbd8fefc 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -389,7 +389,7 @@ def supports_only_rowwise_all_gather(self) -> bool: # ----- TensorProto / pure-Python allocation ----- - def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + def storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: return { "cls": Float8TensorStorage if self.internal else Float8Tensor, "nontensor_kwargs": { @@ -399,7 +399,7 @@ def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: }, } - def _describe_buffers( + def inner_tensor_specs( self, shape: Tuple[int, ...] ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: shape = tuple(shape) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index c639795a80..8adcdb5c0c 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -60,7 +60,7 @@ def copy(self) -> MXFP8Quantizer: # ----- TensorProto / pure-Python allocation ----- - def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + def storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: return { "cls": MXFP8TensorStorage if self.internal else MXFP8Tensor, "nontensor_kwargs": { @@ -71,7 +71,7 @@ def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: }, } - def _describe_buffers( + def inner_tensor_specs( self, shape: Tuple[int, ...] ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: shape = tuple(shape) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 866b6aab56..ae40830901 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -347,7 +347,7 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: # ----- TensorProto / pure-Python allocation ----- - def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + def storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: return { "cls": NVFP4TensorStorage if self.internal else NVFP4Tensor, "nontensor_kwargs": { @@ -361,14 +361,14 @@ def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: }, } - def _describe_buffers( + def inner_tensor_specs( 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 + # Order matches NVFP4TensorStorage._INNER_TENSORS (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). 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 d671122628..f2bf573484 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -12,7 +12,7 @@ import transformer_engine_torch as tex -from ...quantized_tensor import Buffer, QuantizedTensorStorage, Quantizer +from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType_To_Torch, DType @@ -119,10 +119,10 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): be instantiated directly for performance-critical internal usage. """ - _rowwise_data: Annotated[Optional[torch.Tensor], Buffer("rowwise_data")] - _rowwise_scale_inv: Annotated[Optional[torch.Tensor], Buffer("rowwise_scale_inv")] - _columnwise_data: Annotated[Optional[torch.Tensor], Buffer("columnwise_data")] - _columnwise_scale_inv: Annotated[Optional[torch.Tensor], Buffer("columnwise_scale_inv")] + _rowwise_data: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_data")] + _rowwise_scale_inv: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_scale_inv")] + _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")] + _columnwise_scale_inv: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_scale_inv")] _quantizer: Quantizer _fp8_dtype: DType _is_2D_scaled: bool diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 9b4994a77f..c82821e88f 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -10,7 +10,7 @@ import transformer_engine_torch as tex -from ...quantized_tensor import Buffer, QuantizedTensorStorage, Quantizer +from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch, DType @@ -66,15 +66,15 @@ class Float8TensorStorage(QuantizedTensorStorage): """ - _data: Annotated[Optional[torch.Tensor], Buffer("data")] + _data: Annotated[Optional[torch.Tensor], InnerTensor("data")] _quantizer: Optional[Quantizer] _fp8_dtype: DType # FP8 transpose cache - _transpose: Annotated[Optional[torch.Tensor], Buffer("data_transpose")] + _transpose: Annotated[Optional[torch.Tensor], InnerTensor("data_transpose")] _transpose_invalid: bool - _scale_inv: Annotated[torch.Tensor, Buffer("fp8_scale_inv")] + _scale_inv: Annotated[torch.Tensor, InnerTensor("fp8_scale_inv")] def __new__( cls, diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 936c7f89a4..e36f7baa37 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -12,7 +12,7 @@ import transformer_engine_torch as tex -from ...quantized_tensor import Buffer, QuantizedTensorStorage, Quantizer +from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType @@ -67,11 +67,11 @@ class MXFP8TensorStorage(QuantizedTensorStorage): """ # Row-scaled FP8 data and its scaling factors - _rowwise_data: Annotated[Optional[torch.Tensor], Buffer("rowwise_data")] - _rowwise_scale_inv: Annotated[torch.Tensor, Buffer("rowwise_scale_inv")] + _rowwise_data: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_data")] + _rowwise_scale_inv: Annotated[torch.Tensor, InnerTensor("rowwise_scale_inv")] # Column-scaled FP8 data and its scaling factors - _columnwise_data: Annotated[Optional[torch.Tensor], Buffer("columnwise_data")] - _columnwise_scale_inv: Annotated[torch.Tensor, Buffer("columnwise_scale_inv")] + _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")] + _columnwise_scale_inv: Annotated[torch.Tensor, InnerTensor("columnwise_scale_inv")] # Builder class for casting to MXFP8 _quantizer: Optional[Quantizer] diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index b15565ad1f..7e0861c967 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -15,7 +15,7 @@ import transformer_engine_torch as tex -from ...quantized_tensor import Buffer, QuantizedTensorStorage, Quantizer +from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType @@ -81,14 +81,14 @@ class NVFP4TensorStorage(QuantizedTensorStorage): """ # Row-scaled FP4 data and its block scaling factors - _rowwise_data: Annotated[Optional[torch.Tensor], Buffer("rowwise_data")] - _rowwise_scale_inv: Annotated[torch.Tensor, Buffer("rowwise_scale_inv")] + _rowwise_data: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_data")] + _rowwise_scale_inv: Annotated[torch.Tensor, InnerTensor("rowwise_scale_inv")] # Column-scaled FP4 data and its block scaling factors - _columnwise_data: Annotated[Optional[torch.Tensor], Buffer("columnwise_data")] - _columnwise_scale_inv: Annotated[torch.Tensor, Buffer("columnwise_scale_inv")] + _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")] + _columnwise_scale_inv: Annotated[torch.Tensor, InnerTensor("columnwise_scale_inv")] # Input absolute maximum values, used to compute the tensor scale - _amax_rowwise: Annotated[torch.Tensor, Buffer("amax_rowwise")] - _amax_columnwise: Annotated[torch.Tensor, Buffer("amax_columnwise")] + _amax_rowwise: Annotated[torch.Tensor, InnerTensor("amax_rowwise")] + _amax_columnwise: Annotated[torch.Tensor, InnerTensor("amax_columnwise")] # Builder class for casting to MXFP8 _quantizer: Optional[Quantizer] From 9a4fa8a377592b4776917c3d244c26a6c95dafdc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 14:30:38 +0200 Subject: [PATCH 28/28] Carry the storage class itself through the flatten context __tensor_flatten__ put the class qualname in the context and __tensor_unflatten__ looked it up in a module-level registry, populated from __init_subclass__. The indirection bought nothing: dynamo bakes the class object into the graph as a constant just as happily, which is what the custom-op branch already relies on. Store type(self) directly and drop _STORAGE_REGISTRY. __init_subclass__ stays for collecting the InnerTensor field annotations. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 8 ++------ transformer_engine/pytorch/dynamo/tensor_proto.py | 6 +----- transformer_engine/pytorch/quantized_tensor.py | 13 +++---------- 3 files changed, 6 insertions(+), 21 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 2d25c33f42..6ab29d5f53 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -29,11 +29,7 @@ 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, - Quantizer, - _STORAGE_REGISTRY, -) +from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, Quantizer from transformer_engine.pytorch.dynamo import TensorProto, to_tensor_proto from transformer_engine.pytorch import ( is_fp8_available, @@ -852,7 +848,7 @@ def _build_from_primitives(quantizer, shape, dtype, device="cpu"): 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"]] + storage_cls = 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() diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index 0072bde4e3..cab7da48e6 100644 --- a/transformer_engine/pytorch/dynamo/tensor_proto.py +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -119,14 +119,10 @@ def assemble(self, inner_tensors: List[torch.Tensor]) -> torch.Tensor: """ if self.quantizer is None: return inner_tensors[0] - 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(), inner_tensors)) - storage_cls = _STORAGE_REGISTRY[ctx["cls"]] + storage_cls = ctx["cls"] return storage_cls.__tensor_unflatten__( inner, ctx, shape, make_contiguous_strides_for(shape) ) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index ef23ff893f..022af2c4c0 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -29,10 +29,6 @@ _quantized_tensor_passthrough_ops: set = set() -#: Maps storage / wrapper class qualname -> class object, for ``__tensor_unflatten__``. -_STORAGE_REGISTRY: Dict[str, type] = {} - - class InnerTensor(NamedTuple): """Marks a storage field as a flat tensor buffer. @@ -179,9 +175,6 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: 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 cls._INNER_TENSORS = _collect_inner_tensor_fields(cls) def _flatten_nontensor_kwargs(self) -> Dict[str, Any]: @@ -193,7 +186,7 @@ def __tensor_flatten__(self) -> Tuple[list, Dict[str, Any]]: """Return ``(inner_tensor_attr_names, context)``; see class comment.""" present = [attr for attr, _ in self._INNER_TENSORS if getattr(self, attr) is not None] ctx = { - "cls": type(self).__qualname__, + "cls": type(self), "is_tensor": isinstance(self, QuantizedTensor), "requires_grad": ( bool(self.requires_grad) if isinstance(self, QuantizedTensor) else False @@ -210,7 +203,7 @@ def __tensor_unflatten__( outer_stride: Optional[Iterable[int]], ) -> QuantizedTensorStorage: """Rebuild a storage / wrapper from flat tensors + context.""" - cls = _STORAGE_REGISTRY[ctx["cls"]] + cls = 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._INNER_TENSORS: @@ -505,7 +498,7 @@ def create_metadata( """ meta = self.storage_metadata(dtype) return { - "cls": meta["cls"].__qualname__, + "cls": meta["cls"], "is_tensor": not self.internal, "requires_grad": requires_grad, "nontensor_kwargs": meta["nontensor_kwargs"],