From f3401df703909f211dc8617e8832ba2f2d3592a0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Sat, 6 Jun 2026 14:11:29 +0200 Subject: [PATCH 001/108] [PyTorch] Make tensorless quantizers opaque value objects for torch.compile Give tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling, NVFP4) value-object semantics so torch.compile can treat them as baked-in constants: - Add opt-in value identity to the base Quantizer (_value_fields / _value_key / __eq__ / __hash__). Quantizers holding live tensors (delayed-scaling Float8Quantizer) and custom quantizers keep identity semantics. - New transformer_engine/pytorch/dynamo.py houses the torch.compile glue: __fx_repr__, value-key reconstruction and register_value_opaque_quantizer (gracefully a no-op on PyTorch builds without the opaque-object API). - Register the four tensorless quantizers as value opaque types. Also fix CustomRecipe state caching in TransformerEngineBaseModule: set_meta_tensor now rebuilds quantizers when the CustomRecipe instance changes (e.g. nested te.autocast regions) instead of reusing the first recipe's state, since every CustomRecipe shares the CustomRecipeState type but carries its own qfactory. Move the quantizer value-object tests into tests/pytorch/test_torch_compile.py and add that file to the L0 pytorch unittest QA suite. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 154 ++++++++++++++++++ transformer_engine/pytorch/dynamo.py | 120 ++++++++++++++ .../pytorch/quantized_tensor.py | 67 ++++++++ .../pytorch/tensor/float8_blockwise_tensor.py | 7 + .../pytorch/tensor/float8_tensor.py | 9 + .../pytorch/tensor/mxfp8_tensor.py | 7 + .../pytorch/tensor/nvfp4_tensor.py | 25 +++ 7 files changed, 389 insertions(+) create mode 100644 transformer_engine/pytorch/dynamo.py diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 1286492a6e..8adcc88e61 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -24,6 +24,7 @@ from transformer_engine.common import recipe from transformer_engine.pytorch.constants import FP8FwdTensorIdx, FP8BwdTensorIdx from transformer_engine.pytorch.module.base import TransformerEngineBaseModule +from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.quantization import QuantizerRole @@ -32,6 +33,14 @@ is_mxfp8_available, is_fp8_block_scaling_available, is_nvfp4_available, + Float8Quantizer, + Float8BlockQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) +from transformer_engine.pytorch.dynamo import ( + register_value_opaque_quantizer, + _quantizer_from_value_key, ) from utils import recipe_id @@ -384,3 +393,148 @@ def fn(inp): out = compiled(inp) out.sum().backward() + + +# --------------------------------------------------------------------------- +# Value-opaque quantizers: eager value semantics + FX reconstruction +# +# The tensorless quantizers (current-scaling FP8, FP8 blockwise, MXFP8, NVFP4) +# are torch.compile *value* opaque types: they provide value-based +# ``__eq__`` / ``__hash__`` and an evaluable ``__fx_repr__`` (see +# ``torch._library.opaque_object`` Note [Opaque Objects]). These tests exercise +# the eager value semantics and the FX reconstruction round-trip. They are +# CPU-friendly except for NVFP4 (whose constructor touches the current CUDA +# device). +# --------------------------------------------------------------------------- + + +def _mxfp8(dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True): + return MXFP8Quantizer(fp8_dtype=dtype, rowwise=rowwise, columnwise=columnwise) + + +def _blockwise(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=True, block_scaling_dim=2): + return Float8BlockQuantizer( + fp8_dtype=dtype, + rowwise=True, + columnwise=True, + force_pow_2_scales=force_pow_2_scales, + block_scaling_dim=block_scaling_dim, + ) + + +def _current_scaling(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=False, amax_epsilon=0.0): + return Float8CurrentScalingQuantizer( + fp8_dtype=dtype, + device=torch.device("cpu"), + force_pow_2_scales=force_pow_2_scales, + amax_epsilon=amax_epsilon, + ) + + +# (factory, kwargs_for_a_different_but_valid_config) +_CPU_VALUE_QUANTIZERS = [ + pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), + pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"), + pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), +] + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_value_equality(factory, other_kwargs): + """Same config -> equal & same hash; different config -> not equal.""" + a = factory() + b = factory() + assert a is not b + assert a == b + assert hash(a) == hash(b) + + c = factory(**other_kwargs) + assert a != c + # Usage flags participate in the value. + d = factory() + d.set_usage(rowwise=False, columnwise=False) + assert a != d + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_usable_in_set_and_dict(factory, other_kwargs): + a = factory() + b = factory() + c = factory(**other_kwargs) + assert len({a, b, c}) == 2 + mapping = {a: "x"} + assert mapping[b] == "x" + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_cross_type_inequality(factory, other_kwargs): + a = factory() + other = _current_scaling() if not isinstance(a, Float8CurrentScalingQuantizer) else _mxfp8() + assert a != other + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_fx_repr_roundtrip(factory, other_kwargs): + """``__fx_repr__`` returns an evaluable expression rebuilding an equal object.""" + a = factory() + repr_str, globals_ = a.__fx_repr__() + assert isinstance(repr_str, str) + assert isinstance(globals_, dict) + rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used + assert rebuilt == a + assert rebuilt is not a + assert hash(rebuilt) == hash(a) + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_value_key_reconstruction(factory, other_kwargs): + a = factory() + rebuilt = _quantizer_from_value_key(a._value_key()) + assert type(rebuilt) is type(a) + assert rebuilt == a + # The deprecated amax-reduction process group is never carried in the value. + assert getattr(rebuilt, "amax_reduction_group", None) is None + + +def test_quantizer_delayed_scaling_keeps_identity_semantics(): + """Float8Quantizer holds live tensors -> identity (not value) semantics.""" + scale = torch.ones(1) + amax = torch.zeros(1) + a = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3) + b = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3) + assert a._value_fields() is None + assert a == a + assert a != b # distinct instances are not equal despite identical config + assert hash(a) == object.__hash__(a) + + +@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) +def test_quantizer_registration_is_idempotent_and_tolerant(factory, other_kwargs): + """Re-registering must not raise, regardless of PyTorch opaque-object support.""" + cls = type(factory()) + register_value_opaque_quantizer(cls) + register_value_opaque_quantizer(cls) + + if not _opaque_available: + pytest.skip("PyTorch build without opaque-object API") + from torch._library.opaque_object import is_opaque_value_type + + assert is_opaque_value_type(cls) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4Quantizer requires CUDA") +def test_quantizer_nvfp4_value_semantics(): + a = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + b = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + assert a == b + assert hash(a) == hash(b) + + c = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1, with_rht=not a.with_rht) + assert a != c + + rebuilt = _quantizer_from_value_key(a._value_key()) + assert rebuilt == a + assert rebuilt.amax_reduction_group is None + + repr_str, globals_ = a.__fx_repr__() + assert eval(repr_str, dict(globals_)) == a # pylint: disable=eval-used diff --git a/transformer_engine/pytorch/dynamo.py b/transformer_engine/pytorch/dynamo.py new file mode 100644 index 0000000000..26aff6f59c --- /dev/null +++ b/transformer_engine/pytorch/dynamo.py @@ -0,0 +1,120 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""torch.compile glue for Transformer Engine quantizers. + +This module isolates the torch.compile-specific plumbing that turns a +*tensorless* quantizer into a torch.compile **value** opaque type: + + * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used + by FX codegen and registers the quantizer class with + ``torch._library.opaque_object``. It is a no-op (other than populating the + local registry) on PyTorch builds without the opaque-object API, so + importing Transformer Engine never fails on older PyTorch -- only + torch.compile specialization on the quantizer is unavailable there. + * :func:`_quantizer_from_value_key` -- rebuilds a quantizer constant from its + value key inside the generated FX graph. + +The eager value semantics (``__eq__`` / ``__hash__`` / ``_value_key`` / +``_value_fields``) live on the quantizer itself; see +:class:`transformer_engine.pytorch.quantized_tensor.Quantizer`. + +See ``torch._library.opaque_object`` Note [Opaque Objects] for the contract a +value-typed opaque object must satisfy (``__eq__`` / ``__hash__`` / +``__fx_repr__``). +""" + +from __future__ import annotations +from typing import Any, Dict, Tuple + +from .constants import DType + + +# Maps a quantizer class qualname to the class object. A value key stores only +# the qualname, so reconstruction looks the class up here. Populated by +# ``register_value_opaque_quantizer`` at import time of each tensor module; this +# avoids importing the tensor modules into this module (which would create an +# import cycle). +_QUANTIZER_VALUE_REGISTRY: Dict[str, type] = {} + + +def _quantizer_from_value_key(key: Tuple[Any, ...]) -> Any: + """Rebuild a tensorless quantizer from its value key. + + Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the + generated FX code calls this to materialize the quantizer constant. The + deprecated amax-reduction process group is never part of the value, so a + reconstructed quantizer always starts with no stored group. + """ + qualname, items = key[0], key[1] + cls = _QUANTIZER_VALUE_REGISTRY[qualname] + # Bypass ``__init__`` and restore the value attributes directly: the value + # key already captures every value-defining field (including derived ones), + # and the constructors have heterogeneous signatures / side effects. + obj = cls.__new__(cls) + field_names = set() + for name, value in items: + if name == "dtype": + value = DType.cast(value) + object.__setattr__(obj, name, value) + field_names.add(name) + # The deprecated amax-reduction process group is excluded from the value; + # restore it as ``None`` for quantizers that still carry the fallback so + # attribute access keeps working. + if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"): + object.__setattr__(obj, "amax_reduction_group", None) + return obj + + +def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]: + """``__fx_repr__`` for value-opaque quantizers (attached at registration). + + Returns an evaluable expression that rebuilds the quantizer via + :func:`_quantizer_from_value_key`, together with the globals needed to + evaluate it. + """ + return ( + f"_quantizer_from_value_key({self._value_key()!r})", + {"_quantizer_from_value_key": _quantizer_from_value_key}, + ) + + +def register_value_opaque_quantizer(cls: type) -> None: + """Register a tensorless quantizer class as a torch.compile value opaque type. + + Attaches ``__fx_repr__`` and registers the class with + ``torch._library.opaque_object``. Safe to call on any PyTorch build: on + versions without the opaque-object API it only records the class in the + local registry and attaches ``__fx_repr__`` (both harmless), so Transformer + Engine keeps importing and running in eager mode. + + The quantizer class must already provide value ``__eq__`` / ``__hash__`` and + a non-``None`` ``_value_fields`` (see + :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`). + """ + _QUANTIZER_VALUE_REGISTRY[cls.__qualname__] = cls + + # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the + # class, so attach it before registering. + if "__fx_repr__" not in cls.__dict__: + cls.__fx_repr__ = _quantizer_fx_repr + + try: + from torch._library.opaque_object import ( # pylint: disable=import-outside-toplevel + register_opaque_type, + is_opaque_value_type, + ) + except (ImportError, AttributeError): + # Older PyTorch without the opaque-object API: eager value semantics + # still work; torch.compile specialization on the quantizer does not. + return + + if is_opaque_value_type(cls): + return + + try: + register_opaque_type(cls, typ="value") + except (ImportError, AttributeError, RuntimeError, TypeError): + # Tolerate partial / experimental opaque-object support. + pass diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index cfe488aae5..f37b3cc63d 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -408,6 +408,73 @@ def get_usages(self) -> Dict[str, bool]: "columnwise": self.columnwise_usage, } + # ----- Value-object identity (torch.compile opaque value support) ----- + # A *tensorless* quantizer (one whose entire state is a handful of plain, + # reproducible scalars -- no live tensors, no process groups) behaves like a + # value: two instances with the same configuration are interchangeable. Such + # quantizers opt into value-based ``__eq__`` / ``__hash__`` by overriding + # ``_value_fields``. Quantizers that keep the default (e.g. delayed-scaling + # ``Float8Quantizer``, which holds live scale/amax tensors, and any custom + # quantizer) retain the default identity semantics. + # + # This is the eager-side half of registering the quantizer as a torch.compile + # *value* opaque type; the torch.compile glue (``__fx_repr__``, FX + # reconstruction and ``register_opaque_type``) lives in + # ``transformer_engine.pytorch.dynamo``. + + #: Attributes shared by every quantizer that take part in value identity. + _BASE_VALUE_FIELDS: Tuple[str, ...] = ( + "rowwise_usage", + "columnwise_usage", + "internal", + "optimize_for_gemm", + ) + + def _value_fields(self) -> Optional[Tuple[str, ...]]: + """Subclass-specific value-defining attribute names, or ``None``. + + Returning ``None`` (the default) means the quantizer is *not* a value + object and keeps identity-based equality/hashing. Tensorless quantizers + override this to return the tuple of attribute names that, together with + :attr:`_BASE_VALUE_FIELDS`, fully determine their value (excluding + non-value state such as a deprecated amax-reduction process group). + """ + return None + + def _value_key(self) -> Tuple[Any, ...]: + """Hashable, reproducible key identifying this quantizer's value. + + Only valid for value quantizers (``_value_fields()`` is not ``None``). + """ + fields = self._value_fields() # pylint: disable=assignment-from-none + assert fields is not None, f"{type(self).__name__} is not a value quantizer" + items = [] + for name in self._BASE_VALUE_FIELDS + tuple(fields): + value = getattr(self, name) + if name == "dtype": + # ``DType`` is an ``IntEnum``; store the int so the key stays + # plain: hashable and ``repr``-reproducible for FX codegen. + value = int(value) + items.append((name, value)) + return (type(self).__qualname__, tuple(items)) + + def __eq__(self, other: object) -> Any: + # Value quantizers compare by configuration; everything else keeps the + # default identity semantics (returning ``NotImplemented`` makes Python + # fall back to identity). + if self is other: + return True + if self._value_fields() is None or type(self) is not type(other): + return NotImplemented + if other._value_fields() is None: + return NotImplemented + return self._value_key() == other._value_key() + + def __hash__(self) -> int: + if self._value_fields() is None: + return object.__hash__(self) + return hash(self._value_key()) + class QuantizedTensor(torch.Tensor): """Abstract base class for tensor with quantized data diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index ba46508d74..92c22acd0b 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -14,6 +14,7 @@ from transformer_engine.common.recipe import Float8BlockScaling, Recipe from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc from ..constants import DType from ..utils import devices_match, round_up_to_nearest_multiple @@ -69,6 +70,9 @@ def copy(self) -> Float8BlockQuantizer: return quantizer + def _value_fields(self) -> Tuple[str, ...]: + return ("dtype", "block_len", "amax_epsilon", "force_pow_2_scales", "block_scaling_dim") + def update_quantized( self, src: torch.Tensor, @@ -211,6 +215,9 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return Float8BlockScaling +register_value_opaque_quantizer(Float8BlockQuantizer) + + class Float8BlockwiseQTensor(Float8BlockwiseQTensorStorage, QuantizedTensor): """Tensor class with FP8 data quantized via NxN blocks or 1xN blocks. diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index e26abf7df0..56b1ecfb09 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -18,6 +18,7 @@ from ..utils import canonicalize_process_group, devices_match from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc from ..constants import dist_group_type, DType @@ -386,6 +387,14 @@ def supports_only_rowwise_all_gather(self) -> bool: """ return True + def _value_fields(self) -> Tuple[str, ...]: + # ``amax_reduction_group`` is intentionally excluded: it is a deprecated + # process group (not a value) and is restored as ``None`` on rebuild. + return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction") + + +register_value_opaque_quantizer(Float8CurrentScalingQuantizer) + class Float8Tensor(Float8TensorStorage, QuantizedTensor): """Experimental tensor class with FP8 data diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index d759aaf5c4..a3746b3088 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -18,6 +18,7 @@ from ..utils import devices_match, round_up_to_nearest_multiple from .storage.mxfp8_tensor_storage import MXFP8TensorStorage, _FromMXFP8Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc aten = torch.ops.aten @@ -57,6 +58,9 @@ def copy(self) -> MXFP8Quantizer: return quantizer + def _value_fields(self) -> Tuple[str, ...]: + return ("dtype",) + def update_quantized( self, src: torch.Tensor, @@ -1058,3 +1062,6 @@ def backward( ) return dgrad, None return grad.view(ctx.shape), None + + +register_value_opaque_quantizer(MXFP8Quantizer) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index aa92be004f..573c78907c 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -23,6 +23,7 @@ from .storage.nvfp4_tensor_storage import NVFP4TensorStorage, _FromNVFP4Func from ..quantized_tensor import QuantizedTensor, Quantizer +from ..dynamo import register_value_opaque_quantizer from ._quantization_helpers import _IdentityFunc aten = torch.ops.aten @@ -333,6 +334,30 @@ def _canonicalized_amax_reduction_group(self) -> dist_group_type: def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return NVFP4BlockScaling + def _value_fields(self) -> Tuple[str, ...]: + # ``amax_reduction_group`` is intentionally excluded: it is a deprecated + # process group (not a value) and is restored as ``None`` on rebuild. + # ``rht_matrix_random_sign_mask_t`` is derived (from + # ``_with_random_sign_mask`` and the device) but is stored verbatim so + # reconstruction does not need to touch the device. + return ( + "dtype", + "with_rht", + "with_post_rht_amax", + "with_2d_quantization", + "stochastic_rounding", + "row_scaled_nvfp4", + "nvfp4_use_4over6", + "nvfp4_e4m3_max", + "nvfp4_4over6_err_mode", + "_with_random_sign_mask", + "rht_matrix_random_sign_mask_t", + "with_amax_reduction", + ) + + +register_value_opaque_quantizer(NVFP4Quantizer) + class NVFP4Tensor(NVFP4TensorStorage, QuantizedTensor): """Quantized tensor class with FP4 data From c4ad54c1b1f28d9d3402cff92ce244fb9f1a2bc4 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Sat, 6 Jun 2026 14:51:06 +0200 Subject: [PATCH 002/108] [PyTorch] Drop quantizer value registry; reconstruct via __fx_repr__ globals Follow-up to the value-opaque quantizer support: - Remove the module-level _QUANTIZER_VALUE_REGISTRY (qualname -> class) and _quantizer_from_value_key. __fx_repr__ now captures the quantizer class directly in the FX globals and reconstructs via _rebuild_quantizer(cls, items), matching how PyTorch's own value opaque types (e.g. DTensor placements) reconstruct themselves. This removes global mutable state and the qualname collision risk. - Consolidate the quantizer value-object tests in test_torch_compile.py down to two functions and exercise reconstruction through the public __fx_repr__ path instead of internal helpers. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 138 +++--------------- transformer_engine/pytorch/dynamo.py | 53 +++---- .../pytorch/quantized_tensor.py | 22 +-- 3 files changed, 51 insertions(+), 162 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 8adcc88e61..4491e23881 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -36,11 +36,6 @@ Float8Quantizer, Float8BlockQuantizer, MXFP8Quantizer, - NVFP4Quantizer, -) -from transformer_engine.pytorch.dynamo import ( - register_value_opaque_quantizer, - _quantizer_from_value_key, ) from utils import recipe_id @@ -396,145 +391,60 @@ def fn(inp): # --------------------------------------------------------------------------- -# Value-opaque quantizers: eager value semantics + FX reconstruction +# Value-opaque quantizers # -# The tensorless quantizers (current-scaling FP8, FP8 blockwise, MXFP8, NVFP4) -# are torch.compile *value* opaque types: they provide value-based -# ``__eq__`` / ``__hash__`` and an evaluable ``__fx_repr__`` (see -# ``torch._library.opaque_object`` Note [Opaque Objects]). These tests exercise -# the eager value semantics and the FX reconstruction round-trip. They are -# CPU-friendly except for NVFP4 (whose constructor touches the current CUDA -# device). +# Tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling) are +# torch.compile *value* opaque types: value-based ``__eq__`` / ``__hash__`` plus +# an evaluable ``__fx_repr__`` that rebuilds an equal object (see +# ``torch._library.opaque_object`` Note [Opaque Objects]). # --------------------------------------------------------------------------- -def _mxfp8(dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True): - return MXFP8Quantizer(fp8_dtype=dtype, rowwise=rowwise, columnwise=columnwise) +def _mxfp8(dtype=tex.DType.kFloat8E4M3): + return MXFP8Quantizer(fp8_dtype=dtype) -def _blockwise(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=True, block_scaling_dim=2): +def _blockwise(force_pow_2_scales=True): return Float8BlockQuantizer( - fp8_dtype=dtype, + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, force_pow_2_scales=force_pow_2_scales, - block_scaling_dim=block_scaling_dim, ) -def _current_scaling(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=False, amax_epsilon=0.0): +def _current_scaling(amax_epsilon=0.0): return Float8CurrentScalingQuantizer( - fp8_dtype=dtype, + fp8_dtype=tex.DType.kFloat8E4M3, device=torch.device("cpu"), - force_pow_2_scales=force_pow_2_scales, amax_epsilon=amax_epsilon, ) -# (factory, kwargs_for_a_different_but_valid_config) -_CPU_VALUE_QUANTIZERS = [ +# (factory, kwargs producing a different-but-valid config) +_VALUE_QUANTIZERS = [ pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"), pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), ] -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_value_equality(factory, other_kwargs): - """Same config -> equal & same hash; different config -> not equal.""" - a = factory() - b = factory() +@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) +def test_quantizer_value_object(factory, other_kwargs): + """Value semantics + ``__fx_repr__`` round-trip via the production FX path.""" + a, b = factory(), factory() + # Same config -> equal, same hash, interchangeable as a dict/set key. assert a is not b assert a == b assert hash(a) == hash(b) + assert {a: "x"}[b] == "x" + # Different config -> not equal. + assert a != factory(**other_kwargs) - c = factory(**other_kwargs) - assert a != c - # Usage flags participate in the value. - d = factory() - d.set_usage(rowwise=False, columnwise=False) - assert a != d - - -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_usable_in_set_and_dict(factory, other_kwargs): - a = factory() - b = factory() - c = factory(**other_kwargs) - assert len({a, b, c}) == 2 - mapping = {a: "x"} - assert mapping[b] == "x" - - -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_cross_type_inequality(factory, other_kwargs): - a = factory() - other = _current_scaling() if not isinstance(a, Float8CurrentScalingQuantizer) else _mxfp8() - assert a != other - - -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_fx_repr_roundtrip(factory, other_kwargs): - """``__fx_repr__`` returns an evaluable expression rebuilding an equal object.""" - a = factory() + # ``__fx_repr__`` (used by torch.compile codegen) rebuilds an equal object. repr_str, globals_ = a.__fx_repr__() - assert isinstance(repr_str, str) - assert isinstance(globals_, dict) rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used - assert rebuilt == a - assert rebuilt is not a + assert rebuilt == a and rebuilt is not a assert hash(rebuilt) == hash(a) - - -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_value_key_reconstruction(factory, other_kwargs): - a = factory() - rebuilt = _quantizer_from_value_key(a._value_key()) - assert type(rebuilt) is type(a) - assert rebuilt == a - # The deprecated amax-reduction process group is never carried in the value. + # The deprecated amax-reduction group is never part of the value. assert getattr(rebuilt, "amax_reduction_group", None) is None - - -def test_quantizer_delayed_scaling_keeps_identity_semantics(): - """Float8Quantizer holds live tensors -> identity (not value) semantics.""" - scale = torch.ones(1) - amax = torch.zeros(1) - a = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3) - b = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3) - assert a._value_fields() is None - assert a == a - assert a != b # distinct instances are not equal despite identical config - assert hash(a) == object.__hash__(a) - - -@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS) -def test_quantizer_registration_is_idempotent_and_tolerant(factory, other_kwargs): - """Re-registering must not raise, regardless of PyTorch opaque-object support.""" - cls = type(factory()) - register_value_opaque_quantizer(cls) - register_value_opaque_quantizer(cls) - - if not _opaque_available: - pytest.skip("PyTorch build without opaque-object API") - from torch._library.opaque_object import is_opaque_value_type - - assert is_opaque_value_type(cls) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4Quantizer requires CUDA") -def test_quantizer_nvfp4_value_semantics(): - a = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) - b = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) - assert a == b - assert hash(a) == hash(b) - - c = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1, with_rht=not a.with_rht) - assert a != c - - rebuilt = _quantizer_from_value_key(a._value_key()) - assert rebuilt == a - assert rebuilt.amax_reduction_group is None - - repr_str, globals_ = a.__fx_repr__() - assert eval(repr_str, dict(globals_)) == a # pylint: disable=eval-used diff --git a/transformer_engine/pytorch/dynamo.py b/transformer_engine/pytorch/dynamo.py index 26aff6f59c..16f73920ae 100644 --- a/transformer_engine/pytorch/dynamo.py +++ b/transformer_engine/pytorch/dynamo.py @@ -9,12 +9,14 @@ * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used by FX codegen and registers the quantizer class with - ``torch._library.opaque_object``. It is a no-op (other than populating the - local registry) on PyTorch builds without the opaque-object API, so - importing Transformer Engine never fails on older PyTorch -- only - torch.compile specialization on the quantizer is unavailable there. - * :func:`_quantizer_from_value_key` -- rebuilds a quantizer constant from its - value key inside the generated FX graph. + ``torch._library.opaque_object``. It is a no-op on PyTorch builds without + the opaque-object API, so importing Transformer Engine never fails on older + PyTorch -- only torch.compile specialization on the quantizer is + unavailable there. + * :func:`_rebuild_quantizer` -- rebuilds a quantizer constant from its value + items inside the generated FX graph. The quantizer class is captured + directly in the FX globals (see :func:`_quantizer_fx_repr`), so no global + class registry is needed. The eager value semantics (``__eq__`` / ``__hash__`` / ``_value_key`` / ``_value_fields``) live on the quantizer itself; see @@ -22,7 +24,10 @@ See ``torch._library.opaque_object`` Note [Opaque Objects] for the contract a value-typed opaque object must satisfy (``__eq__`` / ``__hash__`` / -``__fx_repr__``). +``__fx_repr__``). The ``__fx_repr__`` contract -- ``(repr_string, {name: type})`` +where ``repr_string`` references the names in the dict -- is exactly how +PyTorch's own value opaque types (e.g. DTensor placements) reconstruct +themselves, including across the on-disk compile cache. """ from __future__ import annotations @@ -31,26 +36,16 @@ from .constants import DType -# Maps a quantizer class qualname to the class object. A value key stores only -# the qualname, so reconstruction looks the class up here. Populated by -# ``register_value_opaque_quantizer`` at import time of each tensor module; this -# avoids importing the tensor modules into this module (which would create an -# import cycle). -_QUANTIZER_VALUE_REGISTRY: Dict[str, type] = {} - - -def _quantizer_from_value_key(key: Tuple[Any, ...]) -> Any: - """Rebuild a tensorless quantizer from its value key. +def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: + """Rebuild a tensorless quantizer of type *cls* from its value items. Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the generated FX code calls this to materialize the quantizer constant. The deprecated amax-reduction process group is never part of the value, so a reconstructed quantizer always starts with no stored group. """ - qualname, items = key[0], key[1] - cls = _QUANTIZER_VALUE_REGISTRY[qualname] # Bypass ``__init__`` and restore the value attributes directly: the value - # key already captures every value-defining field (including derived ones), + # items already capture every value-defining field (including derived ones), # and the constructors have heterogeneous signatures / side effects. obj = cls.__new__(cls) field_names = set() @@ -71,12 +66,15 @@ def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]: """``__fx_repr__`` for value-opaque quantizers (attached at registration). Returns an evaluable expression that rebuilds the quantizer via - :func:`_quantizer_from_value_key`, together with the globals needed to - evaluate it. + :func:`_rebuild_quantizer`, capturing both the helper and the quantizer + class itself in the FX globals so codegen can resolve them with no global + registry and no qualname collisions. """ + cls = type(self) + items = self._value_key()[1] return ( - f"_quantizer_from_value_key({self._value_key()!r})", - {"_quantizer_from_value_key": _quantizer_from_value_key}, + f"_rebuild_quantizer({cls.__name__}, {items!r})", + {"_rebuild_quantizer": _rebuild_quantizer, cls.__name__: cls}, ) @@ -85,16 +83,13 @@ def register_value_opaque_quantizer(cls: type) -> None: Attaches ``__fx_repr__`` and registers the class with ``torch._library.opaque_object``. Safe to call on any PyTorch build: on - versions without the opaque-object API it only records the class in the - local registry and attaches ``__fx_repr__`` (both harmless), so Transformer - Engine keeps importing and running in eager mode. + versions without the opaque-object API it only attaches ``__fx_repr__`` + (harmless), so Transformer Engine keeps importing and running in eager mode. The quantizer class must already provide value ``__eq__`` / ``__hash__`` and a non-``None`` ``_value_fields`` (see :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`). """ - _QUANTIZER_VALUE_REGISTRY[cls.__qualname__] = cls - # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the # class, so attach it before registering. if "__fx_repr__" not in cls.__dict__: diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index f37b3cc63d..6809893a40 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -408,20 +408,6 @@ def get_usages(self) -> Dict[str, bool]: "columnwise": self.columnwise_usage, } - # ----- Value-object identity (torch.compile opaque value support) ----- - # A *tensorless* quantizer (one whose entire state is a handful of plain, - # reproducible scalars -- no live tensors, no process groups) behaves like a - # value: two instances with the same configuration are interchangeable. Such - # quantizers opt into value-based ``__eq__`` / ``__hash__`` by overriding - # ``_value_fields``. Quantizers that keep the default (e.g. delayed-scaling - # ``Float8Quantizer``, which holds live scale/amax tensors, and any custom - # quantizer) retain the default identity semantics. - # - # This is the eager-side half of registering the quantizer as a torch.compile - # *value* opaque type; the torch.compile glue (``__fx_repr__``, FX - # reconstruction and ``register_opaque_type``) lives in - # ``transformer_engine.pytorch.dynamo``. - #: Attributes shared by every quantizer that take part in value identity. _BASE_VALUE_FIELDS: Tuple[str, ...] = ( "rowwise_usage", @@ -433,11 +419,9 @@ def get_usages(self) -> Dict[str, bool]: def _value_fields(self) -> Optional[Tuple[str, ...]]: """Subclass-specific value-defining attribute names, or ``None``. - Returning ``None`` (the default) means the quantizer is *not* a value - object and keeps identity-based equality/hashing. Tensorless quantizers - override this to return the tuple of attribute names that, together with - :attr:`_BASE_VALUE_FIELDS`, fully determine their value (excluding - non-value state such as a deprecated amax-reduction process group). + Returning ``None`` (the default) means the quantizer cannot be represented as + a value opaque object and keeps identity-based equality/hashing. + This also means, that torch.compile will not be able to optimize the quantizer. """ return None From a06324bd4f10354d76b5fbec133f07709bd6a1da Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Sun, 7 Jun 2026 15:42:31 +0200 Subject: [PATCH 003/108] [PyTorch] Split dynamo.py into a dynamo/ package Replace the single dynamo.py module with a dynamo/ package so the torch.compile glue can grow with a clear responsibility split across the stacked branches. This branch owns the value-opaque quantizer layer. * dynamo/quantizer_opaque.py -- register_value_opaque_quantizer and helpers * dynamo/__init__.py -- re-exports the public API so callers keep importing from transformer_engine.pytorch.dynamo unchanged Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/__init__.py | 18 ++++++++++++++++++ .../{dynamo.py => dynamo/quantizer_opaque.py} | 7 +++---- 2 files changed, 21 insertions(+), 4 deletions(-) create mode 100644 transformer_engine/pytorch/dynamo/__init__.py rename transformer_engine/pytorch/{dynamo.py => dynamo/quantizer_opaque.py} (95%) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py new file mode 100644 index 0000000000..44ca61d470 --- /dev/null +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""torch.compile glue for Transformer Engine. + +Public API is re-exported here so callers keep importing from +``transformer_engine.pytorch.dynamo`` regardless of the internal module layout: + + * :mod:`.quantizer_opaque` -- make a tensorless quantizer a torch.compile + *value* opaque type (:func:`register_value_opaque_quantizer`). +""" + +from .quantizer_opaque import register_value_opaque_quantizer + +__all__ = [ + "register_value_opaque_quantizer", +] diff --git a/transformer_engine/pytorch/dynamo.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py similarity index 95% rename from transformer_engine/pytorch/dynamo.py rename to transformer_engine/pytorch/dynamo/quantizer_opaque.py index 16f73920ae..c7455846ee 100644 --- a/transformer_engine/pytorch/dynamo.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -2,10 +2,9 @@ # # See LICENSE for license information. -"""torch.compile glue for Transformer Engine quantizers. +"""Value-opaque quantizers for torch.compile. -This module isolates the torch.compile-specific plumbing that turns a -*tensorless* quantizer into a torch.compile **value** opaque type: +Turns a *tensorless* quantizer into a torch.compile **value** opaque type: * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used by FX codegen and registers the quantizer class with @@ -33,7 +32,7 @@ class registry is needed. from __future__ import annotations from typing import Any, Dict, Tuple -from .constants import DType +from ..constants import DType def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: From ea5b396b9e1e3ee67905c58b50e30989f5ab095d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 8 Jun 2026 12:39:40 +0200 Subject: [PATCH 004/108] [PyTorch] Raise in quantizer __fx_repr__ when a process group is stored A value-opaque quantizer must not carry live distributed state. Scan the quantizer attributes in __fx_repr__ and raise TypeError if any holds a torch.distributed.ProcessGroup (e.g. a non-None deprecated amax_reduction_group), so it cannot be silently baked into a torch.compile FX graph. Clarify the related comments accordingly. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 7 -- transformer_engine/pytorch/dynamo/__init__.py | 9 +-- .../pytorch/dynamo/quantizer_opaque.py | 70 ++++++++++--------- .../pytorch/quantized_tensor.py | 4 +- .../pytorch/tensor/float8_tensor.py | 3 +- .../pytorch/tensor/nvfp4_tensor.py | 3 +- 6 files changed, 45 insertions(+), 51 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 4491e23881..9a7a4a356a 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -392,11 +392,6 @@ def fn(inp): # --------------------------------------------------------------------------- # Value-opaque quantizers -# -# Tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling) are -# torch.compile *value* opaque types: value-based ``__eq__`` / ``__hash__`` plus -# an evaluable ``__fx_repr__`` that rebuilds an equal object (see -# ``torch._library.opaque_object`` Note [Opaque Objects]). # --------------------------------------------------------------------------- @@ -446,5 +441,3 @@ def test_quantizer_value_object(factory, other_kwargs): 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 diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index 44ca61d470..aae8b9cff6 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -2,14 +2,7 @@ # # See LICENSE for license information. -"""torch.compile glue for Transformer Engine. - -Public API is re-exported here so callers keep importing from -``transformer_engine.pytorch.dynamo`` regardless of the internal module layout: - - * :mod:`.quantizer_opaque` -- make a tensorless quantizer a torch.compile - *value* opaque type (:func:`register_value_opaque_quantizer`). -""" +"""torch.compile glue for Transformer Engine.""" from .quantizer_opaque import register_value_opaque_quantizer diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index c7455846ee..55bc8326e6 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -2,46 +2,35 @@ # # See LICENSE for license information. -"""Value-opaque quantizers for torch.compile. - -Turns a *tensorless* quantizer into a torch.compile **value** opaque type: - - * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used - by FX codegen and registers the quantizer class with - ``torch._library.opaque_object``. It is a no-op on PyTorch builds without - the opaque-object API, so importing Transformer Engine never fails on older - PyTorch -- only torch.compile specialization on the quantizer is - unavailable there. - * :func:`_rebuild_quantizer` -- rebuilds a quantizer constant from its value - items inside the generated FX graph. The quantizer class is captured - directly in the FX globals (see :func:`_quantizer_fx_repr`), so no global - class registry is needed. - -The eager value semantics (``__eq__`` / ``__hash__`` / ``_value_key`` / -``_value_fields``) live on the quantizer itself; see -:class:`transformer_engine.pytorch.quantized_tensor.Quantizer`. - -See ``torch._library.opaque_object`` Note [Opaque Objects] for the contract a -value-typed opaque object must satisfy (``__eq__`` / ``__hash__`` / -``__fx_repr__``). The ``__fx_repr__`` contract -- ``(repr_string, {name: type})`` -where ``repr_string`` references the names in the dict -- is exactly how -PyTorch's own value opaque types (e.g. DTensor placements) reconstruct -themselves, including across the on-disk compile cache. -""" +"""Value-opaque quantizers for torch.compile.""" from __future__ import annotations from typing import Any, Dict, Tuple -from ..constants import DType +from ..constants import DType, dist_group_type + + +def _contains_process_group(value: Any) -> bool: + """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``. + + Checks the value directly and one level of ``tuple``/``list`` nesting, which + covers the shapes a quantizer value field could plausibly take. + """ + if isinstance(value, dist_group_type): + return True + if isinstance(value, (tuple, list)): + return any(_contains_process_group(item) for item in value) + return False def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: """Rebuild a tensorless quantizer of type *cls* from its value items. Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the - generated FX code calls this to materialize the quantizer constant. The - deprecated amax-reduction process group is never part of the value, so a - reconstructed quantizer always starts with no stored group. + generated FX code calls this to materialize the quantizer constant. A + quantizer that actually stores a process group never reaches this path: + ``__fx_repr__`` raises for it. The deprecated amax-reduction group is not a + value field, so the rebuilt quantizer simply has no group attribute. """ # Bypass ``__init__`` and restore the value attributes directly: the value # items already capture every value-defining field (including derived ones), @@ -53,9 +42,9 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: value = DType.cast(value) object.__setattr__(obj, name, value) field_names.add(name) - # The deprecated amax-reduction process group is excluded from the value; - # restore it as ``None`` for quantizers that still carry the fallback so - # attribute access keeps working. + # The deprecated amax-reduction group is not a value field. Quantizers that + # actually hold a group error out in ``__fx_repr__`` before reaching here, so + # this only initializes the (groupless) attribute to keep access working. if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"): object.__setattr__(obj, "amax_reduction_group", None) return obj @@ -68,8 +57,23 @@ def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]: :func:`_rebuild_quantizer`, capturing both the helper and the quantizer class itself in the FX globals so codegen can resolve them with no global registry and no qualname collisions. + + Raises ``TypeError`` if the quantizer stores a process group (e.g. a + non-``None`` deprecated ``amax_reduction_group``): live distributed state + must never be baked into the graph as a constant, so such a quantizer cannot + be used with ``torch.compile``. Pass the reduction group per quantize call + instead of storing it on the quantizer. """ cls = type(self) + for name, value in vars(self).items(): + if _contains_process_group(value): + raise TypeError( + f"{cls.__name__} cannot be used with torch.compile: attribute " + f"{name!r} holds a torch.distributed.ProcessGroup, which is live " + "distributed state and must not be baked into an FX graph as a " + "constant. Pass the amax reduction group per quantize call instead " + "of storing it on the quantizer." + ) items = self._value_key()[1] return ( f"_rebuild_quantizer({cls.__name__}, {items!r})", diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 6809893a40..3612c1080d 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -421,7 +421,9 @@ def _value_fields(self) -> Optional[Tuple[str, ...]]: Returning ``None`` (the default) means the quantizer cannot be represented as a value opaque object and keeps identity-based equality/hashing. - This also means, that torch.compile will not be able to optimize the quantizer. + This also means that passing such a quantizer as an argument to a custom op + causes a graph break under torch.compile, since it cannot be baked into the + FX graph as a constant. """ return None diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 56b1ecfb09..0310c3855c 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -389,7 +389,8 @@ def supports_only_rowwise_all_gather(self) -> bool: def _value_fields(self) -> Tuple[str, ...]: # ``amax_reduction_group`` is intentionally excluded: it is a deprecated - # process group (not a value) and is restored as ``None`` on rebuild. + # process group (not a value). If one is actually stored, ``__fx_repr__`` + # raises so it can never be baked into a torch.compile graph. return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction") diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 573c78907c..4bca783922 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -336,7 +336,8 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: def _value_fields(self) -> Tuple[str, ...]: # ``amax_reduction_group`` is intentionally excluded: it is a deprecated - # process group (not a value) and is restored as ``None`` on rebuild. + # process group (not a value). If one is actually stored, ``__fx_repr__`` + # raises so it can never be baked into a torch.compile graph. # ``rht_matrix_random_sign_mask_t`` is derived (from # ``_with_random_sign_mask`` and the device) but is stored verbatim so # reconstruction does not need to touch the device. From aa65e34e2ed8ba784543caa703ef7962062d3546 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 8 Jun 2026 16:17:36 +0200 Subject: [PATCH 005/108] [PyTorch] Cover NVFP4 in quantizer value-object test NVFP4Quantizer is registered as a value-opaque quantizer but was missing from the value-semantics / __fx_repr__ round-trip test. Add it to _VALUE_QUANTIZERS (skipped without CUDA, which it needs to construct). Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 9a7a4a356a..3405001e04 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -27,7 +27,7 @@ from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer -from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, @@ -416,11 +416,29 @@ def _current_scaling(amax_epsilon=0.0): ) +def _nvfp4(with_rht=False): + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_rht=with_rht, + ) + + # (factory, kwargs producing a different-but-valid config) _VALUE_QUANTIZERS = [ pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"), pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), + pytest.param( + _nvfp4, + {"with_rht": True}, + id="nvfp4", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), + reason="NVFP4Quantizer requires CUDA to construct", + ), + ), ] From e1b1db6b1a9de58a6868f1a706659733792476cc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 16 Jun 2026 17:13:15 +0200 Subject: [PATCH 006/108] Reject a value quantizer that carries an amax reduction group in __eq__/__hash__ The amax reduction group is excluded from the value key, so a value quantizer that stored one would compare/hash equal to a groupless one and let torch.compile reuse a graph that skips the reduction. __eq__/__hash__ now raise (mirroring __fx_repr__, which already rejects any process-group-bearing quantizer). The group should be passed per quantize call, not stored on the quantizer. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/quantized_tensor.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 3612c1080d..b7357b0e7b 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -444,6 +444,18 @@ def _value_key(self) -> Tuple[Any, ...]: items.append((name, value)) return (type(self).__qualname__, tuple(items)) + def _check_value_has_no_amax_reduction_group(self) -> None: + # The amax reduction group is not part of the value key, so a value + # quantizer that stores one would compare/hash equal to a groupless one + # and let torch.compile reuse a graph that skips the reduction. Reject it + # (mirrors ``__fx_repr__``); pass the group per quantize call instead. + if getattr(self, "amax_reduction_group", None) is not None: + raise TypeError( + f"{type(self).__name__} with a non-None amax_reduction_group cannot be " + "used as a value object; pass the amax reduction group per quantize call " + "instead of storing it on the quantizer." + ) + def __eq__(self, other: object) -> Any: # Value quantizers compare by configuration; everything else keeps the # default identity semantics (returning ``NotImplemented`` makes Python @@ -454,11 +466,14 @@ def __eq__(self, other: object) -> Any: return NotImplemented if other._value_fields() is None: return NotImplemented + self._check_value_has_no_amax_reduction_group() + other._check_value_has_no_amax_reduction_group() return self._value_key() == other._value_key() def __hash__(self) -> int: if self._value_fields() is None: return object.__hash__(self) + self._check_value_has_no_amax_reduction_group() return hash(self._value_key()) From 8c33d0ec4dd8ff255a75827fd643d3619e6ae9af Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 16 Jun 2026 18:03:29 +0200 Subject: [PATCH 007/108] Recognize value-opaque quantizers via a class flag Add is_value_opaque_quantizer() + the _te_compile_value_opaque flag stamped at registration, so dynamo-traced code can detect registered quantizers (and fall back to eager for unregistered ones). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/__init__.py | 3 ++- .../pytorch/dynamo/quantizer_opaque.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index aae8b9cff6..ee860c78e3 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -4,8 +4,9 @@ """torch.compile glue for Transformer Engine.""" -from .quantizer_opaque import register_value_opaque_quantizer +from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer __all__ = [ "register_value_opaque_quantizer", + "is_value_opaque_quantizer", ] diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 55bc8326e6..6cb9552b71 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -10,6 +10,17 @@ from ..constants import DType, dist_group_type +# Class attribute stamped on quantizers registered as torch.compile value-opaque +# types. +_VALUE_OPAQUE_FLAG = "_te_compile_value_opaque" + + +def is_value_opaque_quantizer(quantizer: Any) -> bool: + """Whether *quantizer*'s class is registered as a torch.compile value-opaque + type.""" + return getattr(quantizer, _VALUE_OPAQUE_FLAG, False) + + def _contains_process_group(value: Any) -> bool: """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``. @@ -93,6 +104,10 @@ def register_value_opaque_quantizer(cls: type) -> None: a non-``None`` ``_value_fields`` (see :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`). """ + # Stamp the class so it can be recognized as value-opaque in dynamo-traced + # code (used to fall back to eager for unregistered quantizers). + setattr(cls, _VALUE_OPAQUE_FLAG, True) + # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the # class, so attach it before registering. if "__fx_repr__" not in cls.__dict__: From 945f62dadd18d5d0bb7b68c54f4bfb9767699521 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 11:34:37 +0200 Subject: [PATCH 008/108] Address review: narrow opaque-type except, add fullgraph test, fix nvfp4 value key - Narrow register_opaque_type except to (RuntimeError, TypeError): the API is already imported above, so ImportError/AttributeError there only mask real errors. - Add test_quantizer_value_object_fullgraph exercising torch.compile(fullgraph=True) end-to-end to verify opaque-type registration took effect. - Restore missing NVFP4Quantizer._with_random_sign_mask assignment required by _value_fields()/_value_key(). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 18 ++++++++++++++++++ .../pytorch/dynamo/quantizer_opaque.py | 5 +++-- .../pytorch/tensor/nvfp4_tensor.py | 1 + 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 3405001e04..5e1f753f08 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -459,3 +459,21 @@ def test_quantizer_value_object(factory, other_kwargs): rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used assert rebuilt == a and rebuilt is not a assert hash(rebuilt) == hash(a) + + +@pytest.mark.skipif( + not _opaque_available, + reason="torch.compile opaque-object support requires PyTorch >= 2.11", +) +@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) +def test_quantizer_value_object_fullgraph(factory, other_kwargs): + """Quantizer survives torch.compile(fullgraph=True) - verifies registration took effect.""" + + def fn(quantizer): + return quantizer + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + quantizer = factory() + assert compiled(quantizer) is quantizer diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 6cb9552b71..6d630d665c 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -128,6 +128,7 @@ def register_value_opaque_quantizer(cls: type) -> None: try: register_opaque_type(cls, typ="value") - except (ImportError, AttributeError, RuntimeError, TypeError): - # Tolerate partial / experimental opaque-object support. + except (RuntimeError, TypeError): + # Keep TE importable: registration must never crash the import, e.g. on + # PyTorch versions with only partial / experimental opaque-object support. pass diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 4bca783922..ffc5f97eca 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -174,6 +174,7 @@ def __init__( self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") + self._with_random_sign_mask = with_random_sign_mask self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( with_random_sign_mask, torch.cuda.current_device() ) From e3c8f430883762ac7be289616bf79f40eb522a0d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 12:05:01 +0200 Subject: [PATCH 009/108] Restore NVFP4 rht_matrix on value-key rebuild; assert quantize round-trip _rebuild_quantizer only restores value-key fields, so a reconstructed NVFP4Quantizer was missing the derived rht_matrix tensor (not hashable, so not in the value key) and failed at copy()/quantize time. Add a _rebuild_derived_state hook (called by _rebuild_quantizer) that NVFP4Quantizer uses to rebuild rht_matrix from _with_random_sign_mask (lru_cache -> cheap). Extend test_quantizer_value_object to also quantize with the original and the rebuilt quantizer and require bit-exact results (gated on HW support), so a field the kernel needs but the value key omits can no longer slip through. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 27 +++++++++++++++++-- .../pytorch/dynamo/quantizer_opaque.py | 5 ++++ .../pytorch/tensor/nvfp4_tensor.py | 10 +++++++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 5e1f753f08..da4f96d2f3 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -416,7 +416,10 @@ def _current_scaling(amax_epsilon=0.0): ) -def _nvfp4(with_rht=False): +def _nvfp4(with_rht=True): + # Default with_rht=True so the quantize round-trip below exercises the + # derived ``rht_matrix`` tensor (the field most likely to be dropped on + # value-key reconstruction). return NVFP4Quantizer( fp4_dtype=tex.DType.kFloat4E2M1, rowwise=True, @@ -425,6 +428,17 @@ def _nvfp4(with_rht=False): ) +def _hw_available(quantizer): + """Whether this HW can actually run the quantize kernel for *quantizer*.""" + if isinstance(quantizer, MXFP8Quantizer): + return mxfp8_available + if isinstance(quantizer, NVFP4Quantizer): + return nvfp4_available + if isinstance(quantizer, Float8BlockQuantizer): + return fp8_block_scaling_available + return fp8_available # Float8CurrentScalingQuantizer + + # (factory, kwargs producing a different-but-valid config) _VALUE_QUANTIZERS = [ pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"), @@ -432,7 +446,7 @@ def _nvfp4(with_rht=False): pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"), pytest.param( _nvfp4, - {"with_rht": True}, + {"with_rht": False}, id="nvfp4", marks=pytest.mark.skipif( not torch.cuda.is_available(), @@ -460,6 +474,15 @@ def test_quantizer_value_object(factory, other_kwargs): assert rebuilt == a and rebuilt is not a assert hash(rebuilt) == hash(a) + # The rebuilt quantizer must also *behave* identically, not just compare + # equal: equality only looks at the value key, so a field the kernel needs + # but that is absent from the key (e.g. NVFP4's derived ``rht_matrix``) would + # slip through the checks above and only blow up at quantize time. Run the + # real quantize kernel on both and require bit-exact results. + if torch.cuda.is_available() and _hw_available(a): + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0) + @pytest.mark.skipif( not _opaque_available, diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 6d630d665c..97532b2790 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -58,6 +58,11 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: # this only initializes the (groupless) attribute to keep access working. if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"): object.__setattr__(obj, "amax_reduction_group", None) + # Restore non-value derived state that ``__init__`` would normally build but + # that cannot live in the value key (e.g. NVFP4's ``rht_matrix`` tensor). + finalize = getattr(obj, "_rebuild_derived_state", None) + if finalize is not None: + finalize() return obj diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index ffc5f97eca..f2f30cdcc5 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -186,6 +186,16 @@ def __getstate__(self): state["amax_reduction_group"] = None return state + def _rebuild_derived_state(self) -> None: + """Restore the derived ``rht_matrix`` after value-key reconstruction. + + ``rht_matrix`` is a ``torch.Tensor`` built from ``_with_random_sign_mask`` + and the device, so it cannot be part of the (hashable) value key. + ``_rebuild_quantizer`` calls this hook to rebuild it; the ``lru_cache`` on + :func:`get_rht_matrix` makes an already-seen (flag, device) a cheap hit. + """ + self.rht_matrix = get_rht_matrix(self._with_random_sign_mask, torch.cuda.current_device()) + def update_quantized( self, src: torch.Tensor, From 3f6862137b30aa30c2980cc1cb9cc750599b02fa Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 14:46:49 +0200 Subject: [PATCH 010/108] Enforce process-group rejection in _value_key, not __fx_repr__; add test Move the ProcessGroup guard out of the (overridable) __fx_repr__ into Quantizer._value_key -- the single point every value-materialization path (__eq__/__hash__/__fx_repr__) goes through -- so a custom __fx_repr__ can no longer bypass it. Generalizes the old amax-only check to any field holding a ProcessGroup. Add a test that a value quantizer carrying a live group raises. Addresses review on NVIDIA#3152. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 21 ++++++++ .../pytorch/dynamo/quantizer_opaque.py | 34 +++---------- .../pytorch/quantized_tensor.py | 50 +++++++++++++------ 3 files changed, 61 insertions(+), 44 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index da4f96d2f3..dc98e13708 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -484,6 +484,27 @@ def test_quantizer_value_object(factory, other_kwargs): torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0) +def test_value_quantizer_rejects_process_group(): + """A value quantizer holding a live ProcessGroup must refuse to be turned + into a value key / FX constant (raise), not silently drop the group.""" + import torch.distributed as dist # pylint: disable=import-outside-toplevel + + created = not dist.is_initialized() + if created: + dist.init_process_group(backend="gloo", store=dist.HashStore(), rank=0, world_size=1) + try: + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + q.amax_reduction_group = dist.group.WORLD + # Every value-materialization path must reject it (hash, eq, __fx_repr__). + with pytest.raises(TypeError): + hash(q) + with pytest.raises(TypeError): + q.__fx_repr__() + finally: + if created: + dist.destroy_process_group() + + @pytest.mark.skipif( not _opaque_available, reason="torch.compile opaque-object support requires PyTorch >= 2.11", diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 97532b2790..37e718689d 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -7,7 +7,7 @@ from __future__ import annotations from typing import Any, Dict, Tuple -from ..constants import DType, dist_group_type +from ..constants import DType # Class attribute stamped on quantizers registered as torch.compile value-opaque @@ -21,19 +21,6 @@ def is_value_opaque_quantizer(quantizer: Any) -> bool: return getattr(quantizer, _VALUE_OPAQUE_FLAG, False) -def _contains_process_group(value: Any) -> bool: - """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``. - - Checks the value directly and one level of ``tuple``/``list`` nesting, which - covers the shapes a quantizer value field could plausibly take. - """ - if isinstance(value, dist_group_type): - return True - if isinstance(value, (tuple, list)): - return any(_contains_process_group(item) for item in value) - return False - - def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: """Rebuild a tensorless quantizer of type *cls* from its value items. @@ -74,22 +61,13 @@ def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]: class itself in the FX globals so codegen can resolve them with no global registry and no qualname collisions. - Raises ``TypeError`` if the quantizer stores a process group (e.g. a - non-``None`` deprecated ``amax_reduction_group``): live distributed state - must never be baked into the graph as a constant, so such a quantizer cannot - be used with ``torch.compile``. Pass the reduction group per quantize call - instead of storing it on the quantizer. + Raises ``TypeError`` (via :meth:`Quantizer._value_key`) if the quantizer + stores a process group (e.g. a non-``None`` deprecated + ``amax_reduction_group``): live distributed state must never be baked into + the graph as a constant. Pass the reduction group per quantize call instead + of storing it on the quantizer. """ cls = type(self) - for name, value in vars(self).items(): - if _contains_process_group(value): - raise TypeError( - f"{cls.__name__} cannot be used with torch.compile: attribute " - f"{name!r} holds a torch.distributed.ProcessGroup, which is live " - "distributed state and must not be baked into an FX graph as a " - "constant. Pass the amax reduction group per quantize call instead " - "of storing it on the quantizer." - ) items = self._value_key()[1] return ( f"_rebuild_quantizer({cls.__name__}, {items!r})", diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index b7357b0e7b..033a35f8e1 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -16,6 +16,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe +from transformer_engine.pytorch.constants import dist_group_type from transformer_engine.pytorch.tensor._quantization_helpers import ( _QuantizeFunc, _IdentityFunc, @@ -23,6 +24,19 @@ ) +def _contains_process_group(value: Any) -> bool: + """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``. + + Checks the value directly and one level of ``tuple``/``list`` nesting, which + covers the shapes a quantizer value field could plausibly take. + """ + if isinstance(value, dist_group_type): + return True + if isinstance(value, (tuple, list)): + return any(_contains_process_group(item) for item in value) + return False + + # Custom ops that should pass through __torch_dispatch__ without unwrapping # QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that # handle quantized tensors internally. @@ -427,6 +441,24 @@ def _value_fields(self) -> Optional[Tuple[str, ...]]: """ return None + def _check_value_has_no_process_group(self) -> None: + # A value quantizer is baked into the FX graph as a constant via its + # value key, which cannot carry live distributed state. Enforced here -- + # the single point every value-materialization path (``__eq__`` / + # ``__hash__`` / ``__fx_repr__``) goes through -- so a custom + # ``__fx_repr__`` cannot bypass it. Reject any field holding a + # ProcessGroup (e.g. the deprecated ``amax_reduction_group``) rather than + # silently dropping it; pass the reduction group per quantize call. + for name, value in vars(self).items(): + if _contains_process_group(value): + raise TypeError( + f"{type(self).__name__} cannot be used as a torch.compile value " + f"object: attribute {name!r} holds a torch.distributed.ProcessGroup, " + "which is live distributed state and must not be baked into an FX " + "graph. Pass the amax reduction group per quantize call instead of " + "storing it on the quantizer." + ) + def _value_key(self) -> Tuple[Any, ...]: """Hashable, reproducible key identifying this quantizer's value. @@ -434,6 +466,7 @@ def _value_key(self) -> Tuple[Any, ...]: """ fields = self._value_fields() # pylint: disable=assignment-from-none assert fields is not None, f"{type(self).__name__} is not a value quantizer" + self._check_value_has_no_process_group() items = [] for name in self._BASE_VALUE_FIELDS + tuple(fields): value = getattr(self, name) @@ -444,36 +477,21 @@ def _value_key(self) -> Tuple[Any, ...]: items.append((name, value)) return (type(self).__qualname__, tuple(items)) - def _check_value_has_no_amax_reduction_group(self) -> None: - # The amax reduction group is not part of the value key, so a value - # quantizer that stores one would compare/hash equal to a groupless one - # and let torch.compile reuse a graph that skips the reduction. Reject it - # (mirrors ``__fx_repr__``); pass the group per quantize call instead. - if getattr(self, "amax_reduction_group", None) is not None: - raise TypeError( - f"{type(self).__name__} with a non-None amax_reduction_group cannot be " - "used as a value object; pass the amax reduction group per quantize call " - "instead of storing it on the quantizer." - ) - def __eq__(self, other: object) -> Any: # Value quantizers compare by configuration; everything else keeps the # default identity semantics (returning ``NotImplemented`` makes Python - # fall back to identity). + # fall back to identity). ``_value_key`` rejects a stored ProcessGroup. if self is other: return True if self._value_fields() is None or type(self) is not type(other): return NotImplemented if other._value_fields() is None: return NotImplemented - self._check_value_has_no_amax_reduction_group() - other._check_value_has_no_amax_reduction_group() return self._value_key() == other._value_key() def __hash__(self) -> int: if self._value_fields() is None: return object.__hash__(self) - self._check_value_has_no_amax_reduction_group() return hash(self._value_key()) From 32d17683f036390a6825902ebacba0e7bf592050 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 15:08:39 +0200 Subject: [PATCH 011/108] Strengthen fullgraph test: quantize/dequantize via a custom op, not passthrough Replace the trivial pass-through fullgraph test with one that drives each production quantizer through a minimal custom op (quantize + dequantize) under torch.compile(fullgraph=True) and compares to eager -- so the opaque-type registration is actually exercised inside the graph (a graph break would make fullgraph=True raise). Op registration sits right before the test. Also drop stale comments referencing the old __fx_repr__-side process-group guard. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 54 ++++++++++++++++--- .../pytorch/dynamo/quantizer_opaque.py | 10 ++-- .../pytorch/tensor/nvfp4_tensor.py | 3 +- 3 files changed, 51 insertions(+), 16 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index dc98e13708..63cb82eca8 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -505,19 +505,59 @@ def test_value_quantizer_rejects_process_group(): dist.destroy_process_group() +if _opaque_available: + # A minimal custom op taking a tensor and a value-opaque quantizer that + # quantizes + dequantizes inside it, one per production quantizer class. + # ``test_quantizer_value_object_fullgraph`` drives this under + # ``torch.compile(fullgraph=True)`` so the quantizer is used *inside* the + # graph -- proving the opaque-type registration took effect (a graph break + # would make ``fullgraph=True`` raise). + _qdq_lib = torch.library.Library("test_te_qdq", "DEF") + _QDQ_OPS = {} + for _qcls in ( + MXFP8Quantizer, + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + NVFP4Quantizer, + ): + _op = f"qdq_{_qcls.__name__}" + _qdq_lib.define(f"{_op}(Tensor x, {get_opaque_type_name(_qcls)} q) -> Tensor") + + @torch.library.impl(f"test_te_qdq::{_op}", "CompositeExplicitAutograd", lib=_qdq_lib) + def _qdq_impl(x, q): + return q(x).dequantize() + + @torch.library.register_fake(f"test_te_qdq::{_op}", lib=_qdq_lib) + def _qdq_fake(x, q): + return torch.empty_like(x) + + _QDQ_OPS[_qcls] = getattr(torch.ops.test_te_qdq, _op) + + @pytest.mark.skipif( not _opaque_available, reason="torch.compile opaque-object support requires PyTorch >= 2.11", ) @pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS) def test_quantizer_value_object_fullgraph(factory, other_kwargs): - """Quantizer survives torch.compile(fullgraph=True) - verifies registration took effect.""" + """Quantizer is usable *inside* a torch.compile(fullgraph=True) graph. - def fn(quantizer): - return quantizer + A custom op quantizes+dequantizes with the (opaque value) quantizer; the + compiled result must match eager. ``fullgraph=True`` raises on any graph + break, so this proves the opaque-type registration actually took effect -- + unlike merely passing the quantizer through. + """ + q = factory() + if not (torch.cuda.is_available() and _hw_available(q)): + pytest.skip("format not supported on this HW") - torch._dynamo.reset() - compiled = torch.compile(fn, fullgraph=True) + op = _QDQ_OPS[type(q)] + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def fn(inp): + return op(inp, q) - quantizer = factory() - assert compiled(quantizer) is quantizer + ref = fn(x) + torch._dynamo.reset() + out = torch.compile(fn, fullgraph=True)(x) + torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 37e718689d..409cb979d1 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -25,10 +25,7 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: """Rebuild a tensorless quantizer of type *cls* from its value items. Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the - generated FX code calls this to materialize the quantizer constant. A - quantizer that actually stores a process group never reaches this path: - ``__fx_repr__`` raises for it. The deprecated amax-reduction group is not a - value field, so the rebuilt quantizer simply has no group attribute. + generated FX code calls this to materialize the quantizer constant. """ # Bypass ``__init__`` and restore the value attributes directly: the value # items already capture every value-defining field (including derived ones), @@ -40,9 +37,8 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any: value = DType.cast(value) object.__setattr__(obj, name, value) field_names.add(name) - # The deprecated amax-reduction group is not a value field. Quantizers that - # actually hold a group error out in ``__fx_repr__`` before reaching here, so - # this only initializes the (groupless) attribute to keep access working. + # The deprecated amax-reduction group is not a value field; initialize it to + # None so attribute access keeps working on the rebuilt quantizer. if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"): object.__setattr__(obj, "amax_reduction_group", None) # Restore non-value derived state that ``__init__`` would normally build but diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index f2f30cdcc5..d9cd534606 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -347,8 +347,7 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: def _value_fields(self) -> Tuple[str, ...]: # ``amax_reduction_group`` is intentionally excluded: it is a deprecated - # process group (not a value). If one is actually stored, ``__fx_repr__`` - # raises so it can never be baked into a torch.compile graph. + # process group, not a value (``_value_key`` rejects a stored group). # ``rht_matrix_random_sign_mask_t`` is derived (from # ``_with_random_sign_mask`` and the device) but is stored verbatim so # reconstruction does not need to touch the device. From 28bde9e7040a8b4c87b992cd5717b7518d19e000 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 15:29:11 +0200 Subject: [PATCH 012/108] Clarify comments: rht_matrix_random_sign_mask_t derivation; why the opaque flag - rht_matrix_random_sign_mask_t is a device-independent int derived from _with_random_sign_mask (the device only places a throwaway tensor); fix the misleading comment. - Explain why registration uses a class attribute, not a registry set: is_value_opaque_quantizer is traced inside the compile graph and dynamo can bake a getattr constant but cannot do 'type(q) in set' on the opaque class. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/quantizer_opaque.py | 8 ++++++-- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 6 +++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 409cb979d1..33e831ef38 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -10,8 +10,12 @@ from ..constants import DType -# Class attribute stamped on quantizers registered as torch.compile value-opaque -# types. +# Registration marks the class with this attribute instead of recording it in a +# module-level set. ``is_value_opaque_quantizer`` runs *inside* the torch.compile +# graph (``Linear.forward`` consults it): Dynamo can trace a ``getattr`` on the +# opaque quantizer and bake the result as a constant, but cannot evaluate +# ``type(q) in some_set`` -- it has no equality/hash rules for the opaque class +# object, so a set/dict lookup graph-breaks under ``fullgraph=True``. _VALUE_OPAQUE_FLAG = "_te_compile_value_opaque" diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index d9cd534606..c8c0a7b854 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -348,9 +348,9 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: def _value_fields(self) -> Tuple[str, ...]: # ``amax_reduction_group`` is intentionally excluded: it is a deprecated # process group, not a value (``_value_key`` rejects a stored group). - # ``rht_matrix_random_sign_mask_t`` is derived (from - # ``_with_random_sign_mask`` and the device) but is stored verbatim so - # reconstruction does not need to touch the device. + # ``rht_matrix_random_sign_mask_t`` is a device-independent int derived + # from ``_with_random_sign_mask``; kept in the key so the rebuilt + # quantizer carries it without recomputation. return ( "dtype", "with_rht", From 2c3c5df0fb3488b8e95670066ac53f958640a6de Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 15:30:44 +0200 Subject: [PATCH 013/108] Reword opaque-flag comment: self-contained, no Linear reference Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/quantizer_opaque.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 33e831ef38..98349e12ba 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -10,12 +10,12 @@ from ..constants import DType -# Registration marks the class with this attribute instead of recording it in a -# module-level set. ``is_value_opaque_quantizer`` runs *inside* the torch.compile -# graph (``Linear.forward`` consults it): Dynamo can trace a ``getattr`` on the -# opaque quantizer and bake the result as a constant, but cannot evaluate -# ``type(q) in some_set`` -- it has no equality/hash rules for the opaque class -# object, so a set/dict lookup graph-breaks under ``fullgraph=True``. +# Registration marks the class with this attribute rather than recording it in a +# module-level set. It looks odd but is a deliberate workaround: the check must +# stay traceable when it runs inside a torch.compile graph -- Dynamo can bake a +# ``getattr`` on the opaque quantizer into a constant, but cannot evaluate +# ``type(q) in some_set`` (no equality/hash rules for the opaque class object), +# which would graph-break under ``fullgraph=True``. _VALUE_OPAQUE_FLAG = "_te_compile_value_opaque" From 826f271ebb28466cdd94b9ee5cc6875eb61d4a97 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 15:38:53 +0200 Subject: [PATCH 014/108] Cover is_opaque_value_type with the import-safety guard too is_opaque_value_type(cls) sat between the import guard and the register_opaque_type guard, so on a partial/experimental opaque-object build it could raise RuntimeError/TypeError and crash TE import. Move it inside the same except so the 'registration never crashes import' promise holds for both calls. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/quantizer_opaque.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index 98349e12ba..8b8b3caa69 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -106,12 +106,11 @@ def register_value_opaque_quantizer(cls: type) -> None: # still work; torch.compile specialization on the quantizer does not. return - if is_opaque_value_type(cls): - return - try: - register_opaque_type(cls, typ="value") + if not is_opaque_value_type(cls): + register_opaque_type(cls, typ="value") except (RuntimeError, TypeError): - # Keep TE importable: registration must never crash the import, e.g. on - # PyTorch versions with only partial / experimental opaque-object support. + # Keep TE importable: neither the opaque-type query nor the registration + # must crash the import, e.g. on PyTorch versions with only partial / + # experimental opaque-object support. pass From ad1ccce0e53e01833c4342d3ed8eb342a0064929 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 16 Jun 2026 17:29:49 +0200 Subject: [PATCH 015/108] 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 | 251 ++++++++++++++++ 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, 1007 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 63cb82eca8..4a3f22a291 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 @@ -28,6 +29,9 @@ from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer +from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, _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, @@ -473,6 +477,8 @@ def test_quantizer_value_object(factory, other_kwargs): 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 @@ -561,3 +567,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 fed367bce5..d479d2f67b 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -68,6 +68,7 @@ prepare_for_saving, restore_from_func_ctx, ) +from ..dynamo import TensorProto from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -92,7 +93,7 @@ class LinearFwdArgs: # --- Differentiable tensors (also passed positionally to autograd) --- weight: TensorOrQuantized - inp: torch.Tensor + inp: TensorOrQuantized bias: Optional[torch.Tensor] # --- Non-differentiable cached tensors --- @@ -301,7 +302,15 @@ def _linear_forward_impl( # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) - backward_needs_input = is_grad_enabled and weight.requires_grad + # Use the requires-grad flags captured into ``args`` at op-call time rather + # than the live tensors': the fake impl (``_linear_forward_impl_fake``) keys + # the number of FP8 inner buffers it emits off ``args.*_requires_grad``, so + # the real impl must agree to keep the custom-op output arity stable. Under + # ``torch.compile`` with CUDA-graph trees (``mode="reduce-overhead"``) the + # static graph inputs are detached during capture, so live + # ``weight.requires_grad`` / ``inp.requires_grad`` flip to False mid-capture + # and would otherwise diverge from the fake (schema/arity mismatch). + backward_needs_input = is_grad_enabled and args.weight_requires_grad with_input_all_gather_nccl = ( parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop ) @@ -418,7 +427,7 @@ def _linear_forward_impl( # No need to set the quantizer states if weight is already quantized # for debug mode we create quantizer every iteration, thus we need to set the quantizer states if weight_quantizer is not None and (not isinstance(weight, QuantizedTensor) or debug): - columnwise_usage = is_grad_enabled and inp.requires_grad and not is_fsdp2 + columnwise_usage = is_grad_enabled and args.input_requires_grad and not is_fsdp2 if backward_override is not None: columnwise_usage = False if not columnwise_usage: @@ -622,6 +631,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, @@ -1311,6 +1518,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 033a35f8e1..ce358aaaf4 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -43,6 +43,10 @@ def _contains_process_group(value: Any) -> bool: _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. @@ -146,6 +150,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], @@ -363,6 +425,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 92c22acd0b..c816b4fb04 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: def _value_fields(self) -> Tuple[str, ...]: return ("dtype", "block_len", "amax_epsilon", "force_pow_2_scales", "block_scaling_dim") + # ----- TensorProto / pure-Python allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": Float8BlockwiseQTensorStorage if self.internal else Float8BlockwiseQTensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "is_2D_scaled": self.block_scaling_dim == 2, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # Blockwise FP8 scales are FP32; columnwise data is stored transposed. + if self.rowwise_usage: + buffers["_rowwise_data"] = (shape, torch.uint8) + buffers["_rowwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=False)), + torch.float32, + ) + if self.columnwise_usage: + buffers["_columnwise_data"] = (tuple(self.get_columnwise_shape(shape)), torch.uint8) + buffers["_columnwise_scale_inv"] = ( + tuple(self.get_scale_shape(shape, columnwise=True)), + torch.float32, + ) + return buffers + def update_quantized( self, src: torch.Tensor, diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 0310c3855c..6d3a53b3d7 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -4,7 +4,7 @@ """Tensor class with FP8 data""" from __future__ import annotations -from typing import Any, Optional, Tuple, Iterable, Union +from typing import Any, Dict, Optional, Tuple, Iterable, Union import warnings import torch from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState @@ -15,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 @@ -393,6 +393,36 @@ def _value_fields(self) -> Tuple[str, ...]: # raises so it can never be baked into a torch.compile graph. return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction") + # ----- TensorProto / pure-Python allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": Float8TensorStorage if self.internal else Float8Tensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # Mirror the C++ quantizer allocation (csrc/quantizer.cpp): on non-TN-capable + # archs (Blackwell+) a single ``_data`` buffer backs both row- and column-wise + # usage and no separate transpose is materialized. This must match what the + # real kernel produces so the torch.compile fake layout lines up slot-for-slot. + non_tn = is_non_tn_fp8_gemm_supported() + if self.rowwise_usage or non_tn: + buffers["_data"] = (shape, torch.uint8) + if self.columnwise_usage and not non_tn: + buffers["_transpose"] = ((shape[-1], *shape[:-1]), torch.uint8) + # Per-tensor scale-inv is always present for current scaling. + buffers["_scale_inv"] = ((1,), torch.float32) + return buffers + register_value_opaque_quantizer(Float8CurrentScalingQuantizer) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index a3746b3088..3e41bf7bb1 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 @@ -61,6 +61,40 @@ def copy(self) -> MXFP8Quantizer: def _value_fields(self) -> Tuple[str, ...]: return ("dtype",) + # ----- TensorProto / pure-Python allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": MXFP8TensorStorage if self.internal else MXFP8Tensor, + "nontensor_kwargs": { + "fp8_dtype": self.dtype, + "quantizer": self, + "with_gemm_swizzled_scales": self.optimize_for_gemm, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # 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 c8c0a7b854..8827b61467 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 @@ -366,6 +366,49 @@ def _value_fields(self) -> Tuple[str, ...]: "with_amax_reduction", ) + # ----- TensorProto / pure-Python allocation ----- + + def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: + return { + "cls": NVFP4TensorStorage if self.internal else NVFP4Tensor, + "nontensor_kwargs": { + "fp4_dtype": self.dtype, + "quantizer": self, + "with_gemm_swizzled_scales": self.optimize_for_gemm, + "row_scaled_nvfp4": self.row_scaled_nvfp4, + "nvfp4_use_4over6": self.nvfp4_use_4over6, + "nvfp4_e4m3_max": self.nvfp4_e4m3_max, + "fake_dtype": fake_dtype, + }, + } + + def _describe_buffers( + self, shape: Tuple[int, ...] + ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]: + shape = tuple(shape) + buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {} + # FP4 data packs 2 values per byte (uint8); block scales are E4M3 stored + # as uint8; amax buffers are FP32 (per-row when row-scaled, else scalar). + 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 f7a3dae70b..ec2c40ef9a 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -35,6 +35,15 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): _columnwise_scale_inv: Optional[torch.Tensor] _is_2D_scaled: bool + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_rowwise_data", "rowwise_data"), + ("_rowwise_scale_inv", "rowwise_scale_inv"), + ("_columnwise_data", "columnwise_data"), + ("_columnwise_scale_inv", "columnwise_scale_inv"), + ) + def __new__( cls, rowwise_data: Optional[torch.Tensor], diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index a97162f91c..1419a02559 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -75,6 +75,14 @@ class Float8TensorStorage(QuantizedTensorStorage): _transpose: Optional[torch.Tensor] _transpose_invalid: bool + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_data", "data"), + ("_transpose", "data_transpose"), + ("_scale_inv", "fp8_scale_inv"), + ) + def __new__( cls, *args, diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index ea592cd989..a2a3bf2f4c 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -82,6 +82,15 @@ class MXFP8TensorStorage(QuantizedTensorStorage): # GEMM _with_gemm_swizzled_scales: bool + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_rowwise_data", "rowwise_data"), + ("_rowwise_scale_inv", "rowwise_scale_inv"), + ("_columnwise_data", "columnwise_data"), + ("_columnwise_scale_inv", "columnwise_scale_inv"), + ) + def __new__( cls, rowwise_data: Optional[torch.Tensor], diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 53bb5e7c11..5ed6d1d641 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -108,6 +108,17 @@ class NVFP4TensorStorage(QuantizedTensorStorage): # Global E4M3 scale bound used by this NVFP4 tensor _nvfp4_e4m3_max: int + # (attribute_name, constructor_kwarg) for each tensor buffer; drives + # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage). + _FLATTEN_TENSOR_BUFFERS = ( + ("_rowwise_data", "rowwise_data"), + ("_rowwise_scale_inv", "rowwise_scale_inv"), + ("_columnwise_data", "columnwise_data"), + ("_columnwise_scale_inv", "columnwise_scale_inv"), + ("_amax_rowwise", "amax_rowwise"), + ("_amax_columnwise", "amax_columnwise"), + ) + def __new__( cls, rowwise_data: Optional[torch.Tensor], From ea3df7a22dfb9c803dfe315d40fb22c99af76b89 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 22 Jun 2026 12:45:34 +0200 Subject: [PATCH 016/108] [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 d479d2f67b..2b2ee05a3e 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -607,12 +607,24 @@ def _linear_forward_impl( if is_fsdp2 and weightmat is not weight: wt_save = None - # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` - # rebuilds the refs from ``inp`` / ``weight`` / ``bias``. - # Needed for torch.compile to work correctly. + # Dedup save slots that alias forward inputs or other op returns; + # ``_linear_setup_ctx`` rebuilds the refs. Needed because a custom op may + # not return a tensor that aliases an input or another return: the cached + # FP8 weight is the same tensor as ``new_weight_workspace`` (a return, on a + # cache miss) or ``weight_workspace`` (an input, on a cache hit). + if wt_save is None: + wt_alias = None + elif wt_save is weight: + wt_alias = "weight" + elif new_weight_workspace is not None and wt_save is new_weight_workspace: + wt_alias = "new_workspace" + elif args.weight_workspace is not None and wt_save is args.weight_workspace: + wt_alias = "weight_workspace" + else: + wt_alias = None saved_tensor_aliases = ( "inp" if saved_inputmat is inp else None, - "weight" if wt_save is weight else None, + wt_alias, "weight", # ``saved_weight`` slot is always the weight parameter "bias" if bias is not None else None, ) @@ -798,13 +810,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: @@ -832,7 +851,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, ...]: @@ -845,7 +864,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 @@ -935,6 +955,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": @@ -1619,7 +1643,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 4997929645510e3157e9cbe7fe5b260d438273a5 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 22 Jun 2026 13:17:19 +0200 Subject: [PATCH 017/108] [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 8827b61467..0cc436c399 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -389,14 +389,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)), @@ -406,6 +406,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 50c11cd415466267893f351e276ff8a432a63320 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 08:00:58 +0200 Subject: [PATCH 018/108] 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 4a3f22a291..0c8f34e5e2 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, @@ -585,8 +584,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", ), ), ] @@ -603,7 +602,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): @@ -807,7 +809,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 2b2ee05a3e..473483faff 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -766,8 +766,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 ce358aaaf4..86c4c94f7c 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -176,7 +176,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 3e41bf7bb1..f804a96f24 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -79,8 +79,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 ff48e52a0d794acc17717575329f142e4cb09272 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 29 Jun 2026 16:25:44 +0200 Subject: [PATCH 019/108] [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 0cc436c399..8d51480f51 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -391,15 +391,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 e1e271cc5f51835a6e28b10c4cbb96efefd74da9 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 15 Jun 2026 13:28:53 +0200 Subject: [PATCH 020/108] [PyTorch] torch.compile: wrap pybind11 UB methods as compile-time constants; fix SP memory leak; test suite hook-up MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wrap CommOverlapCore pybind11 methods that return compile-time constants so torch.compile(fullgraph=True) can trace through them without graph breaks: - `is_fp8_ubuf()` → `ub_is_fp8()` / `get_ub_is_fp8()` in base.py; `_ub_is_fp8()` in gemm.py - `with_cublasmp()` → `ub_is_cublasmp()` in base.py All callers in linear.py, layernorm_linear.py, layernorm_mlp.py, base.py, gemm.py, userbuffers_backward_linear.py and userbuffers_forward_linear.py updated. Fix quantized grad_output not being freed early for column-parallel SP backward. Row-parallel SP already called clear_tensor_data(grad_output) to release the gathered tensor; column-parallel SP quantizes grad_output to Float8TensorStorage but never freed it before returning. Under torch.compile reduce-overhead this leaves 3 live pool tensors at recording end and triggers "Detected 3 tensor(s) in the cudagraph pool not tracked as outputs". Extend the existing clear_tensor_data guard to cover both parallel modes. Fix custom-recipe quantizer state being re-initialised on every forward call even when the recipe object has not changed. The existing early-exit for CustomRecipeState was missing an identity check on the recipe object, so any repeated call with the same recipe would bypass the early-return and rebuild quantizers unnecessarily. Add `if recipe_state.recipe is recipe: return` to restore the intended caching behaviour. Add test_torch_compile.py to L0_pytorch_unittest so the autocast and existing compile tests run in CI. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski (cherry picked from commit bfce3a7d02a0c87e0c3472bd40f2fadf68a0e6a4) --- transformer_engine/pytorch/module/layernorm_linear.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 43799b003c..1214edb937 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -424,7 +424,11 @@ def forward( if ub_overlap_rs_fprop: # cuBLASMp writes the reduce-scattered output directly into the # GEMM output tensor; Userbuffers writes it into the extra-output buffer. - out = gemm_out if ub_obj is not None and ub_obj.with_cublasmp() else reduce_scatter_out + out = ( + gemm_out + if ub_obj is not None and ub_obj.with_cublasmp() + else reduce_scatter_out + ) elif parallel_mode == "row" and tp_size > 1: nvtx_range_push(f"{nvtx_label}.row_parallel_comm") out = gemm_out From 598a07cdc546fd6c9a90c43d640bfc696eef726f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 15 Jun 2026 14:42:37 +0000 Subject: [PATCH 021/108] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci (cherry picked from commit afe364bb6ef8004a03129b9e39047df6871317a6) --- transformer_engine/pytorch/module/layernorm_linear.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 1214edb937..43799b003c 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -424,11 +424,7 @@ def forward( if ub_overlap_rs_fprop: # cuBLASMp writes the reduce-scattered output directly into the # GEMM output tensor; Userbuffers writes it into the extra-output buffer. - out = ( - gemm_out - if ub_obj is not None and ub_obj.with_cublasmp() - else reduce_scatter_out - ) + out = gemm_out if ub_obj is not None and ub_obj.with_cublasmp() else reduce_scatter_out elif parallel_mode == "row" and tp_size > 1: nvtx_range_push(f"{nvtx_label}.row_parallel_comm") out = gemm_out From 05af2a0a15e0fc694f30b84f0d637771fe3b1a25 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 16 Jun 2026 16:32:11 +0200 Subject: [PATCH 022/108] Provide explicit QuantizerRoles in torch.compile custom-recipe test ToyLinear now overrides get_quantizer_roles so CustomRecipeState doesn't hit the no-roles warning, which graph-breaks under fullgraph=True. qfactory dispatches on role.tensor_type instead of a pre-baked string key. Signed-off-by: Pawel Gadzinski (cherry picked from commit 22f80e40846365bd46bb97a5febb10ca02889136) --- tests/pytorch/test_torch_compile.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 0c8f34e5e2..f88bf9ae2c 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -31,6 +31,7 @@ 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.quantization import QuantizerRole from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, From 9bd16fd2c69b3ae0c0fa2f5f2adcc056a120d82a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 16 Jun 2026 17:31:37 +0200 Subject: [PATCH 023/108] Add torch.compile custom-op path for Linear Squashed PR #9 (linear_compile) onto the rebased base. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski (cherry picked from commit 84dbc6b2b704a1ca82dd0da37361e62228d9935a) --- .../distributed/run_layer_with_overlap.py | 28 + tests/pytorch/distributed/run_numerics.py | 51 +- .../distributed/test_comm_gemm_overlap.py | 40 + tests/pytorch/test_torch_compile.py | 343 ++++- transformer_engine/pytorch/dynamo/__init__.py | 2 + .../pytorch/dynamo/custom_op.py | 1322 +++++++++++++++++ transformer_engine/pytorch/module/linear.py | 203 ++- transformer_engine/pytorch/utils.py | 14 + 8 files changed, 1968 insertions(+), 35 deletions(-) create mode 100644 transformer_engine/pytorch/dynamo/custom_op.py diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index 46795415e5..e65824ce85 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -200,6 +200,19 @@ def _parse_args(argv=None, namespace=None): parser.add_argument( "--use-cuda-graphs", action="store_true", default=False, help="Use CUDA Graphs." ) + parser.add_argument( + "--compile", + action="store_true", + default=False, + help="Wrap each layer in torch.compile (tests Userbuffers on the compiled path).", + ) + parser.add_argument( + "--compile-mode", + type=str, + default="default", + choices=["default", "reduce-overhead"], + help="torch.compile mode used when --compile is set.", + ) parser.add_argument( "--ub-cfg", type=str, default=None, help="Optional TP config yaml file input." ) @@ -485,6 +498,9 @@ def dist_print(msg, src=None, end="\n", debug=False, error=False): torch.testing.assert_close(test_param, ref_param, rtol=0.0, atol=0.0) dist_print("Copied parameters from test model to reference model...", debug=True) + if opts.compile and opts.use_cuda_graphs: + raise ValueError("--compile and --use-cuda-graphs are mutually exclusive.") + # Fp8 recipe setup fp8_format = Format.HYBRID fp8_recipe = None @@ -535,6 +551,18 @@ def run_fwd_bwd(model, x): loss.backward() return out + if opts.compile: + for i, layer in enumerate(test_model.layers): + # dynamic=False for now: symbolic shapes would land in an OpaqueValueBundle + # op arg whose hash chokes on non-nested SymInt (see run_numerics). + test_model.layers[i] = torch.compile( + layer, fullgraph=True, mode=opts.compile_mode, dynamic=False + ) + dist_print( + f"Compiled test model layers with torch.compile (mode={opts.compile_mode})...", + debug=True, + ) + torch_rng_state = torch.get_rng_state() cuda_rng_state = torch.cuda.get_rng_state(torch.device(f"cuda:{LOCAL_RANK}")) if opts.use_cuda_graphs: diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index fe02f990b4..1590979ba3 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -310,22 +310,45 @@ def _copy_params(model_distributed, model_single): def _apply_models( - model_single_node, model_distributed, input_single_node, input_distributed, **kwargs + model_single_node, + model_distributed, + input_single_node, + input_distributed, + use_compile=False, + compile_mode="default", + **kwargs, ): _alloc_main_grad(model_single_node, model_distributed) # for fuse_wgrad_accumulation=True input_single_node.requires_grad_() input_distributed.requires_grad_() + forward_single_node = model_single_node + forward_distributed = model_distributed + if use_compile: + # Each parametrized case compiles the same module.forward code object with + # a different shape/recipe; with dynamic=False those guards accumulate and + # eventually trip Dynamo's recompile_limit. Reset so every case starts from + # a clean compile cache (mirrors the single-GPU torch.compile tests). + torch._dynamo.reset() + # dynamic=False for now: a symbolic shape would land in an OpaqueValueBundle + # (value-opaque op arg) whose hash chokes on non-nested SymInt. Force static + # shapes (recompile per shape) until the bundle handles symbolic shapes. + forward_single_node = torch.compile( + model_single_node, fullgraph=True, mode=compile_mode, dynamic=False + ) + forward_distributed = torch.compile( + model_distributed, fullgraph=True, mode=compile_mode, dynamic=False + ) with te.autocast( enabled=QUANTIZATION is not None, recipe=quantization_recipe(), ): - output_single_node = model_single_node(input_single_node, **kwargs) + output_single_node = forward_single_node(input_single_node, **kwargs) with te.autocast( enabled=QUANTIZATION is not None, recipe=quantization_recipe(), amax_reduction_group=NCCL_WORLD, ): - output_distributed = model_distributed(input_distributed, **kwargs) + output_distributed = forward_distributed(input_distributed, **kwargs) return output_single_node, output_distributed @@ -641,12 +664,20 @@ def test_quantized_all_gather(): # Linear # ############################################ @run_distributed_test() -def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs): +def _test_linear( + parallel_mode=None, + sequence_parallel=False, + use_compile=False, + compile_mode="default", + **kwargs, +): """Test the linear layer with specified parallel mode and sequence parallelization. Args: parallel_mode (str): 'row' or 'column' parallelism. sequence_parallel (bool): Enable sequence parallelism if True. + use_compile (bool): Wrap the modules in ``torch.compile`` before running. + compile_mode (str): ``torch.compile`` mode ("default" or "reduce-overhead"). kwargs (dict): Additional arguments for the linear layer. """ # Set parameter data type @@ -696,7 +727,12 @@ def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs): # Apply models output_single_node, output_distributed = _apply_models( - model_single_node, model_distributed, input_single_node, input_distributed + model_single_node, + model_distributed, + input_single_node, + input_distributed, + use_compile=use_compile, + compile_mode=compile_mode, ) if "return_bias" in kwargs: @@ -740,6 +776,8 @@ def test_linear(): {"params_dtype": torch.float16 if QUANTIZATION != "nvfp4" else torch.bfloat16}, {"delay_wgrad_compute": True}, {"save_original_input": True}, + {"use_compile": True}, + {"use_compile": True, "compile_mode": "reduce-overhead"}, ] for kwargs in kwargs_list: @@ -747,6 +785,9 @@ def test_linear(): continue if kwargs.get("delay_wgrad_compute", False) and NVTE_TEST_NVINSPECT_ENABLED: continue + # debug instrumentation forces the eager fallback, so compile is a no-op there. + if kwargs.get("use_compile", False) and NVTE_TEST_NVINSPECT_ENABLED: + continue for parallel_mode in ["column", "row"]: for sequence_parallel in [False, True]: _test_linear(parallel_mode, sequence_parallel, **kwargs) diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 6b1ad870e9..6521ed7dbc 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -111,6 +111,8 @@ def _run_layer_with_overlap( quantization, num_layers=1, use_cublasmp=False, + compile=False, + compile_mode="default", ): test_path = TEST_ROOT / "run_layer_with_overlap.py" test_cmd = LAUNCH_CMD + [ @@ -129,6 +131,10 @@ def _run_layer_with_overlap( if overlap_rs_dgrad: test_cmd.append("--overlap-rs-dgrad") + if compile: + test_cmd.append("--compile") + test_cmd.append(f"--compile-mode={compile_mode}") + if fp8: if quantization in ("fp8_delayed_scaling", "fp8_current_scaling") and not fp8_available: pytest.skip(reason_for_no_fp8) @@ -281,6 +287,40 @@ def test_layers_with_overlap_bf16( ) +@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"]) +@pytest.mark.parametrize( + "linear_parallel_mode,overlap_rs_dgrad", + [ + ("row", False), + ("column", False), + ("column", True), + ], + ids=[ + "ROW-PARALLEL", + "COL-PARALLEL - BULK DGRAD/WGRAD", + "COL-PARALLEL - DGRAD+RS", + ], +) +def test_linear_with_overlap_compile(linear_parallel_mode, overlap_rs_dgrad, compile_mode): + """te.Linear comm+GEMM overlap (Userbuffers) under torch.compile (BF16). + + Userbuffers is expected to stay on Linear's compiled custom-op path (the + collective lives inside the opaque op), so this checks that torch.compile + + Userbuffers stays numerically correct against the eager, non-overlap reference. + ``compile_mode="reduce-overhead"`` additionally exercises CUDA-graph trees on + top of the Userbuffers collectives. + """ + _run_layer_with_overlap( + te.Linear.__name__, + linear_parallel_mode, + overlap_rs_dgrad, + False, + None, + compile=True, + compile_mode=compile_mode, + ) + + @pytest.mark.parametrize("use_cublasmp", (False, True)) @pytest.mark.parametrize( "quantization", diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index f88bf9ae2c..a0501cca8e 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -3,9 +3,11 @@ # See LICENSE for license information. import abc +import contextlib import pytest import torch +from torch._dynamo.utils import counters from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode try: @@ -31,7 +33,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.quantization import QuantizerRole from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, @@ -88,6 +89,75 @@ def nvfp4_4over6(): _all_recipes.append(nvfp4_row_scaled()) +# torch.compile modes exercised by the te.Linear tests: the default backend and +# "reduce-overhead" (CUDA-graph trees), to ensure the custom-op path is +# CUDA-graph capturable. +_compile_modes = ["default", "reduce-overhead"] + + +def _cudagraph_warmup(fn, inp, *, backward: bool) -> None: + """ + Force TE's lazily-created global scratch to be allocated before capture. + """ + out = fn(inp) + if backward: + out.sum().backward() + + +@contextlib.contextmanager +def _assert_no_cudagraph_skips(enabled: bool): + """Assert ``torch.compile(mode="reduce-overhead")`` actually captured CUDA + graphs for every graph instead of silently running it eagerly. + + Inductor bumps ``counters["inductor"]["cudagraph_skips"]`` whenever it + declines to capture a cudagraph (input mutation, CPU scalars, cudagraph-unsafe + ops, ...) and falls back to eager for that graph. ``fullgraph=True`` only rules + out *dynamo* graph breaks, not these *inductor*-level skips, so this guards that + the reduce-overhead path didn't degrade to eager. No-op when ``enabled`` is + False (e.g. the default backend, where cudagraphs don't apply). + """ + before = counters["inductor"]["cudagraph_skips"] + yield + if enabled: + skipped = counters["inductor"]["cudagraph_skips"] - before + assert skipped == 0, ( + f"reduce-overhead fell back to eager: {skipped} cudagraph skip(s); " + "see the 'skipping cudagraphs due to ...' log for the reason" + ) + + +# bf16 output tolerance: eager and compiled run the same kernels, so they should +# agree closely; the slack only absorbs reduction-order / cuda-graph differences. +_EAGER_ATOL, _EAGER_RTOL = 1e-2, 1.6e-2 + + +def _assert_close_eager_compiled(fn, compiled, model, base): + """Run ``fn`` eagerly and ``compiled`` on identical inputs; assert the + forward output and the input / weight gradients match. + + Guards the compiled custom-op path against silently diverging from eager + execution -- a wrong-but-same-shape result would slip past shape / grad + presence checks alone. + """ + inp_eager = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_eager = fn(inp_eager) + out_eager.sum().backward() + ref_out = out_eager.detach().clone() + ref_wgrad = model.weight.grad.detach().clone() + ref_igrad = inp_eager.grad.detach().clone() + + inp_compiled = base.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + # Clone before a later cuda-graph replay overwrites the static output buffer. + out_compiled = compiled(inp_compiled).clone() + out_compiled.sum().backward() + + torch.testing.assert_close(out_compiled, ref_out, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close(inp_compiled.grad, ref_igrad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + torch.testing.assert_close(model.weight.grad, ref_wgrad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL) + + # --------------------------------------------------------------------------- # ToyQuantizer – opaque value-type quantizer for torch.compile # (requires torch opaque object support, not available in older PyTorch) @@ -724,9 +794,13 @@ 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. + # inner_names follows the storage's canonical __tensor_flatten__ order (the + # order the real op flattens its outputs to), while create_inner_tensors + # matches the _describe_buffers geometry (a name->shape/dtype mapping). bufs = q._describe_buffers(shape) # pylint: disable=protected-access - names = tuple(bufs) + direct = _build_from_primitives(q, shape, torch.bfloat16) + names = tuple(direct.__tensor_flatten__()[0]) + assert set(names) == set(bufs) assert proto.inner_names() == names inner = proto.create_inner_tensors() assert len(inner) == len(names) @@ -736,7 +810,6 @@ def test_tensor_proto_matches_primitives(factory, 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) @@ -817,3 +890,265 @@ def test_to_tensor_proto_quantized(factory, shape): assert _signature(proto.create_tensor(), proto.inner_names()) == _signature( tensor, proto.inner_names() ) + + +# --------------------------------------------------------------------------- +# te.Linear +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.parametrize("compile_mode", _compile_modes) +@pytest.mark.parametrize( + "fp8_recipe", + [None, *_all_recipes], + ids=lambda r: "bf16" if r is None else type(r).__name__, +) +def test_te_linear_compiles(fp8_recipe, compile_mode): + """ + torch.compile(fullgraph=True) of ``te.Linear`` under every built-in + recipe (plus the bf16-only baseline with no autocast), for both the default + backend and ``mode="reduce-overhead"`` (CUDA-graph trees). + """ + if fp8_recipe is not None and not fp8_available: + pytest.skip(reason_for_no_fp8) + + dtype = torch.bfloat16 + device = "cuda" + + # FP8 GEMMs require leading dimensions divisible by 16. + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + def fn(inp): + if fp8_recipe is None: + return model(inp) + with te.autocast(recipe=fp8_recipe): + return model(inp) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup( + fn, + torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), + backward=True, + ) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + # ``reduce-overhead`` warms up on the first call(s) and replays a captured + # CUDA graph afterwards, so iterate a few times to actually exercise replay. + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + base = torch.randn(32, 64, dtype=dtype, device=device) + _assert_close_eager_compiled(fn, compiled, model, base) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_linear_compile_with_quantized_fp8_weight(compile_mode): + """torch.compile should handle Linear weights initialized as FP8 tensors, + for both the default backend and ``mode="reduce-overhead"``. + + Exercises the two-tier op + ``register_torch_dispatch`` flattening of a + ``Float8Tensor`` weight *input* in + :mod:`transformer_engine.pytorch.dynamo`. + """ + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + + with te.quantized_model_init(enabled=True, recipe=fp8_recipe): + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + assert isinstance(model.weight, te.Float8Tensor) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup( + fn, + torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), + backward=True, + ) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + base = torch.randn(32, 64, dtype=dtype, device=device) + _assert_close_eager_compiled(fn, compiled, model, base) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_linear_compile_with_fp8_output(compile_mode): + """torch.compile of ``te.Linear(..., fp8_output=True)`` without gradient: + forward returns a :class:`Float8Tensor`. Covers the default backend and + ``mode="reduce-overhead"``. + + Exercises the output-rewrap path in + :mod:`transformer_engine.pytorch.dynamo`: when an output quantizer is + active, the op returns the flat inner data tensors and the framework + rewraps them into a ``Float8Tensor`` via ``__tensor_unflatten__``. A + differentiable FP8 output is unsupported under compile (``Linear.forward`` + falls back to eager), so this test covers the supported case: an FP8 output + that does not require grad (inference / ``torch.no_grad``). + """ + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, fp8_output=True) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + with torch.no_grad(): + _cudagraph_warmup(fn, torch.randn(32, 64, dtype=dtype, device=device), backward=False) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + n_iters = 3 if compile_mode == "reduce-overhead" else 1 + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for _ in range(n_iters): + inp = torch.randn(32, 64, dtype=dtype, device=device) + with torch.no_grad(): + out_eager = fn(inp) + out = compiled(inp) + assert isinstance( + out, te.Float8Tensor + ), f"expected Float8Tensor output, got {type(out).__name__}" + assert out.shape == (32, 32) + assert ( + out._quantizer is not None + ), "FP8 output lost its quantizer on the torch.compile path" + # The rewrap rebuilt a fully-functional Float8Tensor: dequantizing it + # outside the compiled region exercises scale + data + dtype wiring. + deq = out.dequantize() + assert deq.shape == (32, 32) + assert deq.dtype == dtype + # Compiled FP8 output must match the eager FP8 output value-wise. + torch.testing.assert_close( + deq, out_eager.dequantize(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL + ) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("compile_mode", _compile_modes) +def test_te_linear_compile_is_first_microbatch(compile_mode): + """torch.compile of ``te.Linear`` across a multi-step microbatch schedule that + drives FP8 weight caching via ``is_first_microbatch``, for the default backend + and ``mode="reduce-overhead"`` (CUDA-graph trees). + + ``is_first_microbatch=True`` quantizes and caches the FP8 weight; subsequent + ``False`` steps must reuse the cached FP8 weight instead of re-quantizing. This + exercises that cache path under compile and checks it stays numerically aligned + with eager. ``is_first_microbatch`` is a Python bool, so each distinct value is + its own dynamo guard/graph. + """ + dtype = torch.bfloat16 + device = "cuda" + fp8_recipe = recipe.Float8CurrentScaling() + model = te.Linear(64, 32, params_dtype=dtype, device=device) + + # First microbatch caches the FP8 weight, the rest reuse the cache. + schedule = [True, False, False] + is_first = schedule[0] # rebound each step; closed over by ``fn``. + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp, is_first_microbatch=is_first) + + torch._dynamo.reset() + if compile_mode == "reduce-overhead": + _cudagraph_warmup( + fn, + torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True), + backward=True, + ) + model.zero_grad(set_to_none=True) + compiled = torch.compile(fn, fullgraph=True, mode=compile_mode) + + with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"): + for is_first in schedule: + base = torch.randn(32, 64, dtype=dtype, device=device) + _assert_close_eager_compiled(fn, compiled, model, base) + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +def test_te_linear_dynamic_shapes(): + """torch.compile(dynamic=True) of ``te.Linear`` with varying batch sizes. + + Verifies that the compiled graph handles symbolic (dynamic) leading + dimensions without graph breaks or recompilations after the initial trace. + Key correctness property: a graph compiled for batch=16 must produce + numerically correct results for batch=32 without triggering a recompile. + + This exercises two fixes for dynamic shapes: + 1. ``_linear_setup_ctx`` no longer stores ``inp_shape`` in the value bundle + (torch.Size with SymInt dims is not hashable in OpaqueValueBundle). + 2. ``_linear_backward_impl_fake`` derives dgrad shape from grad_output + + weight + SP config instead of relying on the stored ``inp_shape``. + 3. ``_linear_backward`` reconstructs ``inp_shape`` on-the-fly from the same + tensor sources when it is None (compiled mode). + + FP8 + dynamic=True is tracked separately (requires resolving + ``UnsafeScriptObjectError`` for TorchScript quantizer objects with Dynamo). + """ + dtype = torch.bfloat16 + device = "cuda" + in_features, out_features = 64, 32 + model = te.Linear(in_features, out_features, params_dtype=dtype, device=device) + + def fn(inp): + return model(inp) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + batch_sizes = [16, 32, 48] + + for i, batch in enumerate(batch_sizes): + inp = torch.randn(batch, in_features, dtype=dtype, device=device, requires_grad=True) + # Mark batch dim as dynamic so Dynamo traces once and reuses across batch sizes. + torch._dynamo.mark_dynamic(inp, 0) + out = compiled(inp) + assert out.shape == (batch, out_features), f"wrong output shape for batch={batch}" + out.sum().backward() + assert inp.grad is not None, f"no input gradient for batch={batch}" + assert inp.grad.shape == inp.shape, f"wrong grad shape for batch={batch}" + + # Verify numerics against eager on each distinct batch size. + inp_eager = inp.detach().clone().requires_grad_(True) + model.zero_grad(set_to_none=True) + out_eager = model(inp_eager) + out_eager.sum().backward() + torch.testing.assert_close( + out.detach(), out_eager.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL, + msg=f"forward mismatch at batch={batch}", + ) + torch.testing.assert_close( + inp.grad, inp_eager.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL, + msg=f"dgrad mismatch at batch={batch}", + ) + + if i == 0: + # After the first (tracing) call, record the recompile counter + # baseline -- subsequent batch sizes must not trigger recompiles. + recompile_count_baseline = counters["stats"].get("recompile_reasons", 0) + + recompile_count_after = counters["stats"].get("recompile_reasons", 0) + assert recompile_count_after == recompile_count_baseline, ( + f"Unexpected recompilation(s) across different batch sizes: " + f"{recompile_count_after - recompile_count_baseline} recompile(s) detected" + ) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index f932a7d9c3..66e0525a9e 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,10 +6,12 @@ from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer from .tensor_proto import TensorProto, to_tensor_proto +from .custom_op import register_custom_op __all__ = [ "register_value_opaque_quantizer", "is_value_opaque_quantizer", "TensorProto", "to_tensor_proto", + "register_custom_op", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py new file mode 100644 index 0000000000..ef943d4e8e --- /dev/null +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -0,0 +1,1322 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""torch.compile custom-op framework for Transformer Engine.""" + +from __future__ import annotations +import dataclasses +import warnings +from enum import Enum +from typing import ( + Any, + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + Union, + get_args, + get_origin, + get_type_hints, +) + +import torch + +from .tensor_proto import TensorProto, to_tensor_proto, _contiguous_stride +from ..quantized_tensor import ( + QuantizedTensor, + QuantizedTensorStorage, + Quantizer, + _STORAGE_REGISTRY, + _quantized_tensor_passthrough_ops, + prepare_for_saving, +) + +_TE_OP_NAMESPACE = "transformer_engine_compile" +_TE_LIB = torch.library.Library(_TE_OP_NAMESPACE, "FRAGMENT") # noqa: TOR901 + + +# ``None`` entries in an op's flat ``Tensor[]`` return are smuggled through a +# 0-element uint8 tensor: a non-nullable ``Tensor[]`` schema is required for +# ``register_autograd`` to attach a ``grad_fn`` to the outputs. +_NONE_SENTINEL_DTYPE = torch.uint8 + + +def _encode_none(t: Optional[torch.Tensor]) -> torch.Tensor: + """Replace ``None`` with a 0-element uint8 sentinel tensor.""" + if t is None: + return torch.empty(0, dtype=_NONE_SENTINEL_DTYPE) + return t + + +def _decode_none(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Inverse of :func:`_encode_none`.""" + if t is None: + return None + if t.numel() == 0 and t.dtype == _NONE_SENTINEL_DTYPE: + return None + return t + + +# --------------------------------------------------------------------------- # +# OpaqueValueBundle: bundle of simple / value-opaque Python values +# --------------------------------------------------------------------------- # + + +class OpaqueValueBundle: + """Opaque value-type bundle of simple Python values. + + Wraps a ``{name: value}`` dict so many small non-Tensor args pass through a + single custom-op input; registered as a torch.compile *value* opaque type + (Dynamo specializes the graph on its contents). Allowed values: primitives + in :attr:`PRIMITIVE_TYPES` (incl. ``torch.Size``), ``enum.Enum``, any + registered value-opaque type (e.g. TE quantizers), plus nested tuples / + lists / dicts thereof (so a bundle can carry a ``__tensor_flatten__`` + context verbatim). + """ + + PRIMITIVE_TYPES: Tuple[type, ...] = ( + type(None), + bool, + int, + float, + str, + torch.dtype, + torch.device, + torch.Size, + ) + + @classmethod + def is_simple_value(cls, value: Any) -> bool: + """Whether ``value`` may be stored inside an instance (recursive).""" + if isinstance(value, cls.PRIMITIVE_TYPES): + return True + if isinstance(value, Enum): + return True + if _is_opaque_value_type(type(value)): + return True + if isinstance(value, dict): + return all( + isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items() + ) + if isinstance(value, (list, tuple)): + return all(cls.is_simple_value(v) for v in value) + return False + + @classmethod + def _to_hashable(cls, value: Any) -> Any: + if isinstance(value, dict): + return tuple(sorted((k, cls._to_hashable(v)) for k, v in value.items())) + if isinstance(value, (list, tuple, torch.Size)): + return tuple(cls._to_hashable(v) for v in value) + return value + + @classmethod + def _fmt_simple(cls, value: Any) -> str: + """Repr for a value, evaluable in a context with ``torch`` globals.""" + if isinstance(value, torch.dtype): + return f"__import__('torch').{str(value).split('.')[-1]}" + if isinstance(value, torch.device): + return f"__import__('torch').device({str(value)!r})" + if isinstance(value, torch.Size): + return f"__import__('torch').Size({list(value)!r})" + # Enum before primitives: IntEnum is also ``int`` but must render as + # ``EnumName.MEMBER`` (the Enum class is added to globals by ``_collect``). + if isinstance(value, Enum): + return f"{type(value).__name__}.{value.name}" + if isinstance(value, dict): + body = ", ".join(f"{k!r}: {cls._fmt_simple(v)}" for k, v in value.items()) + return f"{{{body}}}" + if isinstance(value, list): + return "[" + ", ".join(cls._fmt_simple(v) for v in value) + "]" + if isinstance(value, tuple): + body = ", ".join(cls._fmt_simple(v) for v in value) + return f"({body},)" if len(value) == 1 else f"({body})" + if _is_opaque_value_type(type(value)): + return value.__fx_repr__()[0] + return repr(value) + + def __init__(self, data: Optional[Dict[str, Any]] = None) -> None: + data = dict(data) if data else {} + for k, v in data.items(): + if not OpaqueValueBundle.is_simple_value(v): + raise TypeError( + f"OpaqueValueBundle field '{k}' has unsupported type " + f"{type(v).__name__}; only simple primitives, Enum, " + "torch.Size, registered value-opaque types and nested " + "tuples / lists / dicts thereof are allowed." + ) + self._data: Dict[str, Any] = data + self._frozen: Tuple[Tuple[str, Any], ...] = tuple( + (k, OpaqueValueBundle._to_hashable(v)) for k, v in sorted(data.items()) + ) + + def __getitem__(self, key: str) -> Any: + return self._data[key] + + def __getattr__(self, name: str) -> Any: + try: + return self._data[name] + except KeyError as e: + raise AttributeError(name) from e + + def get(self, key: str, default: Any = None) -> Any: + """Return ``self._data.get(key, default)``.""" + return self._data.get(key, default) + + def as_dict(self) -> Dict[str, Any]: + """Return a shallow copy of the stored mapping.""" + return dict(self._data) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, OpaqueValueBundle): + return NotImplemented + return self._frozen == other._frozen + + def __hash__(self) -> int: + return hash(self._frozen) + + def __fx_repr__(self) -> Tuple[str, Dict[str, Any]]: + items = ", ".join( + f"{k!r}: {OpaqueValueBundle._fmt_simple(v)}" for k, v in self._data.items() + ) + globals_: Dict[str, Any] = {"OpaqueValueBundle": OpaqueValueBundle} + + def _collect(value: Any) -> None: + if isinstance(value, dict): + for v in value.values(): + _collect(v) + return + if isinstance(value, (list, tuple)): + for v in value: + _collect(v) + return + if isinstance(value, Enum): + globals_[type(value).__name__] = type(value) + return + if isinstance(value, OpaqueValueBundle.PRIMITIVE_TYPES): + return + if _is_opaque_value_type(type(value)): + _, extra = value.__fx_repr__() + globals_.update(extra) + + for v in self._data.values(): + _collect(v) + return (f"OpaqueValueBundle({{{items}}})", globals_) + + +try: + from torch._library.opaque_object import ( # pylint: disable=import-outside-toplevel + get_opaque_type_name, + is_opaque_value_type as _is_opaque_value_type, + is_opaque_reference_type as _is_opaque_reference_type, + register_opaque_type, + ) + + register_opaque_type(OpaqueValueBundle, typ="value") + _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) +except Exception: # pylint: disable=broad-exception-caught # pragma: no cover - older torch without opaque_object + _is_opaque_value_type = None + _is_opaque_reference_type = None + _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None + + +def _pg_pickle_stub(*args: Any) -> None: # pragma: no cover + raise RuntimeError("ProcessGroup cannot be unpickled — cache-key use only") + + +def _ensure_distributed_opaque_types() -> None: + """Register ``torch.distributed.ProcessGroup`` as a *reference* opaque type. + + A process group is live distributed state: unlike a value-opaque quantizer + (which Dynamo bakes into the graph as a constant), it must be carried through + the custom op as a graph *input*. PyTorch supports this via + ``register_opaque_type(ProcessGroup, typ="reference")`` but only auto-runs it + when ``torch.distributed.tensor`` (DTensor) is imported; TE may not import + that, so trigger the same idempotent registration here. Best-effort: on + builds without the opaque-object / distributed APIs this is a no-op and the + process-group field simply falls back to eager under torch.compile. + + Also registers a ``copyreg`` reducer that lets ``FxGraphCachePickler`` hash + graphs containing a ``ProcessGroup`` input without crashing. Without this, + inductor logs "Failed to pickle cache key" warnings and bypasses the FX + graph disk cache for every distributed compiled call. The reducer encodes + the group as (world_size, rank, backend) — enough to distinguish configs — + and raises on reconstruct since deserialization is never needed for hashing. + """ + if _is_opaque_reference_type is None: + return + try: # pylint: disable=import-outside-toplevel + from torch.distributed.device_mesh import _register_distributed_opaque_types + + _register_distributed_opaque_types() + except Exception: # pylint: disable=broad-exception-caught + pass + + # Workaround for PyTorch issue: FxGraphCachePickler handles FakeScriptObject + # but not the real ProcessGroup that appears in example_inputs at inductor + # compile time. Register a copyreg reducer so the pickler can hash the key. + try: # pylint: disable=import-outside-toplevel + import copyreg + import torch.distributed as dist + from torch._C._distributed_c10d import ProcessGroup + + if ProcessGroup not in copyreg.dispatch_table: + + def _pg_reduce(pg: ProcessGroup) -> tuple: # type: ignore[valid-type] + try: + return _pg_pickle_stub, ( + dist.get_world_size(pg), + dist.get_rank(pg), + dist.get_backend(pg), + ) + except Exception: # pylint: disable=broad-exception-caught + return _pg_pickle_stub, (id(pg),) + + copyreg.pickle(ProcessGroup, _pg_reduce) + except Exception: # pylint: disable=broad-exception-caught + pass + + +_ensure_distributed_opaque_types() + + +# --------------------------------------------------------------------------- # +# Storage flatten / unflatten (value-opaque quantizer; no ProcessGroup) +# --------------------------------------------------------------------------- # + + +def _storage_flatten(value: Any) -> Tuple["OpaqueValueBundle", List[torch.Tensor]]: + """Split a ``QuantizedTensor`` / bare storage into ``(meta, Tensor[])``. + + The flatten context (embedding the value-opaque quantizer) plus inner names + and -- for a wrapper subclass -- the outer geometry are stashed in the bundle + so :func:`_storage_unflatten` can rebuild without PyTorch's ``outer_size``. + """ + inner_names, ctx = value.__tensor_flatten__() + meta = dict(ctx) + meta["_inner_names"] = list(inner_names) + if isinstance(value, torch.Tensor): + meta["_outer_shape"] = torch.Size(value.shape) + tensors = [getattr(value, name) for name in inner_names] + return OpaqueValueBundle(meta), tensors + + +def _storage_unflatten(meta: Any, tensors: List[torch.Tensor]) -> Any: + """Inverse of :func:`_storage_flatten`.""" + meta_dict = meta.as_dict() if isinstance(meta, OpaqueValueBundle) else dict(meta) + inner_names = meta_dict["_inner_names"] + inner = dict(zip(inner_names, tensors)) + outer_shape = meta_dict.get("_outer_shape") + stride = _contiguous_stride(tuple(outer_shape)) if outer_shape is not None else None + return QuantizedTensorStorage.__tensor_unflatten__(inner, meta_dict, outer_shape, stride) + + +# --------------------------------------------------------------------------- # +# Field buckets: dataclass field <-> flat torch.library slot(s) +# --------------------------------------------------------------------------- # + + +def _strip_optional(annot: Any) -> Tuple[Any, bool]: + """If ``annot`` is ``Optional[X]`` return ``(X, True)``; else ``(annot, False)``.""" + if get_origin(annot) is Union: + args = get_args(annot) + if type(None) in args: + non_none = [a for a in args if a is not type(None)] + if len(non_none) == 1: + return non_none[0], True + return annot, False + + +class _Bucket: + """Maps one (or, for the aggregating bucket, several) dataclass field(s) + to/from a contiguous run of custom-op schema *slots*. + + A custom op only takes flat, simply-typed arguments, but a TE op takes a + single ``@dataclass`` of mixed fields. Each bucket knows how to translate + its kind of field both ways. ``try_build`` and ``schema_slots`` run once at + registration (to build the op's schema); ``pack`` and ``unpack`` run on each + call and must agree on the slot layout that ``schema_slots`` declares. + """ + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_Bucket"]: + """Decide whether this bucket type handles the field ``name`` given its + type annotation ``annot``; return a configured bucket if so, else + ``None`` so the next candidate is tried. + + Called once per field at registration, in :data:`_FIELD_BUCKETS` + priority order. + """ + raise NotImplementedError + + def schema_slots(self) -> List[Tuple[str, str]]: + """Declare the schema slots this field occupies, each as a + ``(slot_name, schema_type)`` pair (e.g. ``("bias", "Tensor?")``). + + Concatenated across all buckets to form the op's schema string. + """ + raise NotImplementedError + + def pack(self, owner: Any) -> List[Tuple[str, Any]]: + """Read this field from the dataclass ``owner`` and produce the concrete + value for each of its schema slots, as ``(slot_name, value)`` pairs. + + Composite values are flattened to fit the (tensor-only) slots: e.g. a + quantized tensor is split into its plain inner buffers plus a metadata + bundle. Inverse of :meth:`unpack`. + """ + raise NotImplementedError + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + """Read this field's slots back from the op arguments ``args`` and write + the reconstructed field value into ``kwargs`` (rebuilding any flattened + composite). The filled ``kwargs`` are then used to rebuild the original + dataclass for the eager implementation. Inverse of :meth:`pack`. + """ + raise NotImplementedError + + def grad_slot(self) -> Optional[int]: + """Index (within this bucket's :meth:`schema_slots`) of the slot that + carries a gradient, or ``None`` if the field is not differentiable. + + Used to map ``input_tensors_for_grad`` names onto backward grad-output + positions. Non-tensor buckets (quantizers, metadata) return ``None``. + """ + return None + + +class _UniversalKind(Enum): + """What a universal-tensor slot group carries, tagged in its ``__meta``.""" + + NONE = "none" + TENSOR = "tensor" + STORAGE = "storage" + + +class _UniversalTensorBucket(_Bucket): + """``Tensor | QuantizedTensorStorage`` (also subclass tensor) field. + + Three slots regardless of value: ```` (``Tensor?`` -- plain / subclass + tensor passes through, ``None`` for bare storage), ``__tensors`` + (``Tensor[]`` flat inner tensors when flattened), ``__meta`` + (``OpaqueValueBundle`` flatten metadata + a ``__kind__`` marker). + """ + + KIND_KEY = "__kind__" + + def __init__(self, name: str) -> None: + self.name = name + + def slot_name(self) -> str: + """Primary slot name for a plain / subclass tensor.""" + return self.name + + def slot_tensors(self) -> str: + """Flat inner-tensor slot name.""" + return self.name + "__tensors" + + def slot_meta(self) -> str: + """Flatten-metadata slot name.""" + return self.name + "__meta" + + def schema_slots(self) -> List[Tuple[str, str]]: + return [ + (self.slot_name(), "Tensor?"), + (self.slot_tensors(), "Tensor[]"), + (self.slot_meta(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME), + ] + + @staticmethod + def _is_tensor_storage_union(annot: Any) -> bool: + if get_origin(annot) is not Union: + return False + members = [a for a in get_args(annot) if a is not type(None)] + if torch.Tensor not in members: + return False + return any( + isinstance(m, type) and issubclass(m, QuantizedTensorStorage) for m in members + ) + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_UniversalTensorBucket"]: + if cls._is_tensor_storage_union(annot): + return cls(name) + return None + + def pack(self, owner: Any) -> List[Tuple[str, Any]]: + value = getattr(owner, self.name) + if value is None: + return [ + (self.slot_name(), None), + (self.slot_tensors(), []), + (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _UniversalKind.NONE})), + ] + if isinstance(value, torch.Tensor): + # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the + # ``Tensor?`` slot; subclass flattening (if any) is done by the + # outer op's ``register_torch_dispatch`` rule. + return [ + (self.slot_name(), value), + (self.slot_tensors(), []), + (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _UniversalKind.TENSOR})), + ] + if isinstance(value, QuantizedTensorStorage): + meta, tensors = _storage_flatten(value) + meta._data[self.KIND_KEY] = _UniversalKind.STORAGE + return [ + (self.slot_name(), None), + (self.slot_tensors(), list(tensors)), + (self.slot_meta(), meta), + ] + raise TypeError( + f"field {self.name!r} expected None, torch.Tensor, or " + f"QuantizedTensorStorage, got {type(value).__name__}" + ) + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + meta = args[self.slot_meta()] + kind = meta.get(self.KIND_KEY) + if kind == _UniversalKind.NONE: + kwargs[self.name] = None + elif kind == _UniversalKind.TENSOR: + kwargs[self.name] = args[self.slot_name()] + else: + kwargs[self.name] = _storage_unflatten(meta, args[self.slot_tensors()]) + + def grad_slot(self) -> Optional[int]: + # Gradient flows to the plain / subclass tensor slot (``slot_name``, + # the first of the three). + return 0 + + +class _TensorBucket(_Bucket): + """``Tensor`` / ``Optional[Tensor]`` -> single ``Tensor`` / ``Tensor?`` slot.""" + + def __init__(self, name: str, is_optional: bool) -> None: + self.name = name + self.type_str = "Tensor?" if is_optional else "Tensor" + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_TensorBucket"]: + stripped, is_optional = _strip_optional(annot) + if stripped is torch.Tensor: + return cls(name, is_optional) + return None + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.name, self.type_str)] + + def pack(self, owner: Any) -> List[Tuple[str, Any]]: + return [(self.name, getattr(owner, self.name))] + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = args[self.name] + + def grad_slot(self) -> Optional[int]: + return 0 + + +class _QuantizerBucket(_Bucket): + """``Quantizer`` / ``Optional[Quantizer]`` -> one own ``OpaqueValueBundle`` slot. + + Each quantizer gets its own dedicated slot. The field is annotated with the + base ``Quantizer`` (not itself a registered opaque type), so the simple + bundle would not claim it. + """ + + KEY = "q" + + def __init__(self, name: str) -> None: + self.name = name + + def slot(self) -> str: + """Opaque quantizer metadata slot name.""" + return self.name + "__q" + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_QuantizerBucket"]: + stripped, _ = _strip_optional(annot) + if not isinstance(stripped, type): + return None + if issubclass(stripped, Quantizer): + return cls(name) + return None + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] + + def pack(self, owner: Any) -> List[Tuple[str, Any]]: + return [(self.slot(), OpaqueValueBundle({self.KEY: getattr(owner, self.name)}))] + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = args[self.slot()][self.KEY] + + +class _ReferenceOpaqueBucket(_Bucket): + """``ProcessGroup`` (or any reference-opaque type) -> one own opaque slot. + + A reference-opaque object is live, stateful black-box data (e.g. a + ``torch.distributed.ProcessGroup``): it cannot be specialized on or baked + into the graph as a constant the way a value-opaque quantizer is. torch.compile + instead carries it through as a graph *input*, so it passes straight through + its own schema slot (no ``OpaqueValueBundle`` wrapper). The field is annotated + with a concrete type registered via ``register_opaque_type(..., typ="reference")``. + + On the fake / setup-context path the slot holds a ``FakeScriptObject`` (or + ``None``); it is assigned to the field verbatim, so the fake impl must never + read the object's contents. + """ + + def __init__(self, name: str, type_name: str, is_optional: bool) -> None: + self.name = name + self.type_str = f"{type_name}?" if is_optional else type_name + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_ReferenceOpaqueBucket"]: + if _is_opaque_reference_type is None: + return None + stripped, is_optional = _strip_optional(annot) + if not isinstance(stripped, type): + return None + if _is_opaque_reference_type(stripped): + return cls(name, get_opaque_type_name(stripped), is_optional) + return None + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.name, self.type_str)] + + def pack(self, owner: Any) -> List[Tuple[str, Any]]: + return [(self.name, getattr(owner, self.name))] + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = args[self.name] + + +class _SimpleBundleBucket(_Bucket): + """Aggregates every simple-typed field into a single OpaqueValueBundle.""" + + SLOT = "_simple_meta" + + def __init__(self, names: List[str]) -> None: + self.names = list(names) + + @classmethod + def matches_field(cls, annot: Any) -> bool: + """Whether ``annot`` (Optional-aware, recursive) is bundle-simple.""" + annot, _ = _strip_optional(annot) + if annot in OpaqueValueBundle.PRIMITIVE_TYPES: + return True + if isinstance(annot, type) and issubclass(annot, Enum): + return True + if ( + isinstance(annot, type) + and _is_opaque_value_type is not None + and _is_opaque_value_type(annot) + ): + return True + if get_origin(annot) in (tuple, list): + inner = [a for a in get_args(annot) if a is not Ellipsis] + return bool(inner) and all(cls.matches_field(a) for a in inner) + return False + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] + + def pack(self, owner: Any) -> List[Tuple[str, Any]]: + return [(self.SLOT, OpaqueValueBundle({n: getattr(owner, n) for n in self.names}))] + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + if self.SLOT not in args: + return + meta = args[self.SLOT] + for n in self.names: + kwargs[n] = meta[n] + + +class _UnknownBucket(_Bucket): + """Fallback for fields no other bucket claims. + + Emits no slot; pack rejects non-trivial values (anything other than + ``None`` / all-``None`` sequence); unpack restores the field as ``None``. + """ + + def __init__(self, name: str, owner_cls_name: str) -> None: + self.name = name + self.owner_cls_name = owner_cls_name + + @staticmethod + def _is_trivial(value: Any) -> bool: + if value is None: + return True + if isinstance(value, (list, tuple)): + return all(v is None for v in value) + return False + + def schema_slots(self) -> List[Tuple[str, str]]: + return [] + + def pack(self, owner: Any) -> List[Tuple[str, Any]]: + value = getattr(owner, self.name, None) + if not self._is_trivial(value): + raise TypeError( + f"{self.owner_cls_name} field {self.name!r} has a type not " + "supported by torch.compile (not Tensor, simple, or Quantizer) " + "and carries a non-trivial value; add a matching bucket in " + "dynamo.py to handle it." + ) + return [] + + def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + kwargs[self.name] = None + + +# Buckets, in priority order, owning ``try_build`` for a single field. +_FIELD_BUCKETS: Tuple[type, ...] = ( + _UniversalTensorBucket, + _TensorBucket, + _ReferenceOpaqueBucket, + _QuantizerBucket, +) + + +def _resolved_field_annotations(cls: type) -> List[Tuple[str, Any]]: + """Return ``[(field_name, resolved_type), ...]`` for a dataclass.""" + if not dataclasses.is_dataclass(cls): + raise TypeError(f"{cls.__name__} must be a @dataclass to be a TE op arg container.") + try: + hints = get_type_hints(cls) + except Exception: # pylint: disable=broad-exception-caught + hints = {} + return [(f.name, hints.get(f.name, f.type)) for f in dataclasses.fields(cls)] + + +def _get_buckets(cls: type) -> List[_Bucket]: + """Build the bucket list for a dataclass from its field annotations.""" + if _OPAQUE_VALUE_BUNDLE_TYPE_NAME is None: + raise RuntimeError( + f"{cls.__name__} cannot be turned into a TE custom op: OpaqueValueBundle " + "is not registered as a torch._library value-opaque type (PyTorch build " + "without opaque-object support)." + ) + buckets: List[_Bucket] = [] + simple_names: List[str] = [] + for name, annot in _resolved_field_annotations(cls): + built: Optional[_Bucket] = None + for bucket_cls in _FIELD_BUCKETS: + built = bucket_cls.try_build(name, annot) + if built is not None: + break + if built is not None: + buckets.append(built) + elif _SimpleBundleBucket.matches_field(annot): + simple_names.append(name) + else: + buckets.append(_UnknownBucket(name, cls.__name__)) + if simple_names: + buckets.append(_SimpleBundleBucket(simple_names)) + return buckets + + +def _tensor_field_names(buckets: List[_Bucket]) -> List[str]: + """Names of fields carrying tensors (for building the proto view).""" + return [b.name for b in buckets if isinstance(b, (_TensorBucket, _UniversalTensorBucket))] + + +def _build_schema(buckets: List[_Bucket]) -> Tuple[str, List[str]]: + """Return ``(schema_arg_str, slot_names)`` for a bucket list.""" + spec = [slot for b in buckets for slot in b.schema_slots()] + names = [name for name, _ in spec] + schema_str = "(" + ", ".join(f"{type_str} {name}" for name, type_str in spec) + ")" + return schema_str, names + + +def _pack(obj: Any, buckets: List[_Bucket]) -> Dict[str, Any]: + """Build the op's flat ``{slot_name: value}`` argument dict from an args + dataclass ``obj`` (e.g. ``LinearFwdArgs``), by collecting every bucket's + packed slot(s). Inverse of :func:`_unpack`. + """ + out: Dict[str, Any] = {} + for bucket in buckets: + for name, value in bucket.pack(obj): + out[name] = value + return out + + +def _unpack(cls: type, args: Dict[str, Any], buckets: List[_Bucket]) -> Any: + """Rebuild a fresh args dataclass ``cls`` (e.g. ``LinearFwdArgs``) from the + op's flat slot ``args`` dict, by letting every bucket restore its field(s). + Inverse of :func:`_pack`. + """ + kwargs: Dict[str, Any] = {} + for bucket in buckets: + bucket.unpack(args, kwargs) + obj = cls.__new__(cls) + for k, v in kwargs.items(): + object.__setattr__(obj, k, v) + return obj + + +def _proto_view(obj: Any, tensor_field_names: Sequence[str]) -> Any: + """Copy of dataclass ``obj`` with each tensor field replaced by a :class:`TensorProto`. + + Only tensor fields have a ``TensorProto`` equivalent, so quantizer / scalar + fields are simply carried over unchanged; the fake impl works purely on + geometry. Built with :func:`dataclasses.replace` (the only such construction + Dynamo can trace). + """ + overrides: Dict[str, Any] = {} + for name in tensor_field_names: + value = getattr(obj, name, None) + if value is not None and not isinstance(value, TensorProto): + overrides[name] = to_tensor_proto(value) + if not overrides: + return obj + return dataclasses.replace(obj, **overrides) + + +# --------------------------------------------------------------------------- # +# Op outputs <-> flat ``Tensor[]`` payload: this is how an op returns / saves +# quantized tensors (and wrapper subclasses). Outputs are flattened to their +# inner buffers on the way out and rebuilt via ``__tensor_unflatten__`` on the +# way back; on the fake side a TensorProto supplies the geometry. +# --------------------------------------------------------------------------- # + + +def _proto_slot_count(proto: Optional[TensorProto]) -> int: + """Flat ``Tensor[]`` slots the value for ``proto`` occupies.""" + if proto is None: + return 1 + return len(proto.inner_names()) + + +def _proto_reassemble( + proto: Optional[TensorProto], + chunk: List[Optional[torch.Tensor]], +) -> Optional[Union[torch.Tensor, QuantizedTensorStorage]]: + """Rebuild the value described by ``proto`` from its flat tensors ``chunk``. + + ``proto`` describes one output: ``None`` (-> ``None``), a plain tensor + (``chunk`` is the single tensor, returned as-is), or a quantized tensor + (``chunk`` are its inner tensors, reassembled into the wrapper subclass via + ``__tensor_unflatten__``). + """ + if proto is None: + return None + if proto.quantizer is None: + return chunk[0] + inner_names = proto.inner_names() + meta = proto.create_metadata() + shape = tuple(proto.shape) + stride = _contiguous_stride(shape) + storage_cls = _STORAGE_REGISTRY[meta["cls"]] + inner_dict = dict(zip(inner_names, chunk)) + return storage_cls.__tensor_unflatten__(inner_dict, meta, shape, stride) + + +def _value_to_flat_tensors( + value: Optional[Union[torch.Tensor, QuantizedTensorStorage, TensorProto]], +) -> List[torch.Tensor]: + """Return the flat ``Tensor[]`` slots that represent one op output ``value``. + + Inverse of :func:`_proto_reassemble`; the slot count matches + :func:`_proto_slot_count`. + """ + if value is None: + return [_encode_none(None)] + if isinstance(value, TensorProto): + return [_encode_none(t) for t in value.create_inner_tensors()] + if isinstance(value, torch.Tensor): + if type(value) is not torch.Tensor and hasattr( # pylint: disable=unidiomatic-typecheck + value, "__tensor_flatten__" + ): + inner_names, _ = value.__tensor_flatten__() + return [_encode_none(getattr(value, n)) for n in inner_names] + return [_encode_none(value)] + if hasattr(value, "__tensor_flatten__"): + inner_names, _ = value.__tensor_flatten__() + return [_encode_none(getattr(value, n)) for n in inner_names] + raise TypeError( + f"unsupported value type {type(value).__name__}; expected None / " + "torch.Tensor / tensor subclass / bare storage / TensorProto." + ) + + +# Trailing slots in every fwd-impl return: ``tensors_to_save, tensor_objects, +# ctx_attrs``. User-output count is ``len(result) - this``. +_FWD_TRAILING_SLOTS = 3 + + +def _format_fwd_result(result: Any) -> List[torch.Tensor]: + """Pack a fwd-impl return tuple into the op's ``Tensor[]`` payload. + + User outputs first, then saved-for-backward tensors in declaration order. + """ + num_outputs = len(result) - _FWD_TRAILING_SLOTS + flat: List[torch.Tensor] = [] + for value in result[:num_outputs]: + flat.extend(_value_to_flat_tensors(value)) + saved = result[num_outputs] or () + for value in saved: + flat.extend(_value_to_flat_tensors(value)) + return flat + + +def _format_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List[torch.Tensor]: + """Pack a backward-impl return tuple into the op's ``Tensor[]`` payload. + + Each grad occupies exactly one slot (validated against ``num_grad_inputs``); + a :class:`TensorProto` grad is materialized into a single tensor. + """ + grads = list(grads) + if len(grads) != num_grad_inputs: + raise RuntimeError( + f"{op_qualname} expected backward_impl to return {num_grad_inputs} grads " + f"(one per input_tensors_for_grad entry), got {len(grads)}" + ) + out: List[torch.Tensor] = [] + for g in grads: + if isinstance(g, TensorProto): + out.append(_encode_none(g.create_tensor())) + else: + out.append(_encode_none(g)) + return out + + +def _split_fwd_fake_result( + result: Tuple[Any, ...], +) -> Tuple[List[Any], List[Any], Dict[str, Any]]: + """Slice a fwd fake-impl return into ``(user_fakes, saved_fakes, ctx_attrs)``.""" + num_outputs = len(result) - _FWD_TRAILING_SLOTS + saved = result[num_outputs] + ctx_attrs = result[num_outputs + 2] + user_fakes = list(result[:num_outputs]) + saved_fakes = list(saved) if saved is not None else [] + ctx_attrs = dict(ctx_attrs) if ctx_attrs else {} + return user_fakes, saved_fakes, ctx_attrs + + +# --------------------------------------------------------------------------- # +# Op registration +# --------------------------------------------------------------------------- # + + +def _resolve_grad_targets( + fwd_buckets: List[_Bucket], + fwd_arg_type: type, + input_tensors_for_grad: List[str], +) -> Tuple[List[Any], List[int]]: + """Validate ``input_tensors_for_grad`` and resolve the grad-output layout. + + Returns ``(fwd_slot_defaults, grad_targets)``: the per-slot no-grad template + (``[]`` for ``Tensor[]`` slots, ``None`` otherwise) and, for each requested + input name, the schema-slot index its gradient maps to. + """ + fwd_slot_defaults: List[Any] = [] + name_to_slot: Dict[str, int] = {} + slot_offset = 0 + for bucket in fwd_buckets: + slots = bucket.schema_slots() + for _, type_str in slots: + fwd_slot_defaults.append([] if type_str.endswith("[]") else None) + grad_slot = bucket.grad_slot() + if grad_slot is not None: + name_to_slot[bucket.name] = slot_offset + grad_slot + slot_offset += len(slots) + + unknown = [n for n in input_tensors_for_grad if n not in name_to_slot] + if unknown: + raise ValueError( + f"input_tensors_for_grad contains names not in {fwd_arg_type.__name__} " + f"schema: {unknown}" + ) + grad_targets = [name_to_slot[n] for n in input_tensors_for_grad] + return fwd_slot_defaults, grad_targets + + +def _register_kernel( + *, + op_name: str, + op_qualname: str, + arg_type: type, + arg_names: List[str], + buckets: List[_Bucket], + tensor_field_names: List[str], + impl: Callable[[Any], Any], + fake_impl: Callable[[Any], Any], + format_result: Callable[[Any], List[torch.Tensor]], +) -> None: + """Wire the real ``impl`` + the ``fake_impl`` (proto) into the library. + + The real kernel rebuilds the dataclass and runs ``impl``; the fake kernel + runs the proto fake impl on the :func:`_proto_view`. Both go through + ``format_result``. + """ + + def _impl(*flat: Any) -> List[torch.Tensor]: + kwargs = dict(zip(arg_names, flat)) + obj = _unpack(arg_type, kwargs, buckets) + return format_result(impl(obj)) + + def _fake(*flat: Any) -> List[torch.Tensor]: + kwargs = dict(zip(arg_names, flat)) + obj = _unpack(arg_type, kwargs, buckets) + proto_obj = _proto_view(obj, tensor_field_names) + return format_result(fake_impl(proto_obj)) + + _TE_LIB.impl(op_name, _impl, "CompositeExplicitAutograd") + torch.library.register_fake(op_qualname, _fake, lib=_TE_LIB) + + +def _register_autograd_for_op( + *, + fwd_op_name: str, + bwd_op_name: str, + fwd_arg_type: type, + fwd_arg_names: List[str], + fwd_buckets: List[_Bucket], + fwd_tensor_field_names: List[str], + bwd_arg_names: List[str], + bwd_buckets: List[_Bucket], + fwd_slot_defaults: List[Any], + grad_targets: List[int], + setup_context_user: Callable[..., Any], + backward_obj_type: type, + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> None: + """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op_name``. + + ``setup_context`` re-runs the proto fwd fake impl to recover output / saved + templates, reassembles each flat output chunk, and hands the saved tuple + + ``ctx_attrs`` to the module's ``setup_context``. + """ + fwd_qualname = f"{_TE_OP_NAMESPACE}::{fwd_op_name}" + + def _setup_context(ctx, inputs, output): + ctx._te_fwd_tensor_list_lengths = { + i: len(value) for i, value in enumerate(inputs) if isinstance(value, list) + } + kwargs = dict(zip(fwd_arg_names, inputs)) + fwd_obj = _unpack(fwd_arg_type, kwargs, fwd_buckets) + proto_obj = _proto_view(fwd_obj, fwd_tensor_field_names) + + user_fakes, saved_fakes, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(proto_obj)) + + cursor = 0 + user_outputs: List[Any] = [] + for proto in user_fakes: + n = _proto_slot_count(proto) + chunk = [_decode_none(t) for t in output[cursor : cursor + n]] + cursor += n + user_outputs.append(_proto_reassemble(proto, chunk)) + + saved_list: List[Any] = [] + for proto in saved_fakes: + n = _proto_slot_count(proto) + chunk = [_decode_none(t) for t in output[cursor : cursor + n]] + cursor += n + saved_list.append(_proto_reassemble(proto, chunk)) + + bwd_obj = backward_obj_type() + tensors_to_save_from_setup = setup_context_user( + bwd_obj, + fwd_obj, + user_outputs[0] if len(user_fakes) == 1 else tuple(user_outputs), + ctx_attrs, + tuple(saved_list), + ) + tensors_to_save, tensor_objects = prepare_for_saving( + *(tensors_to_save_from_setup or ()) + ) + ctx.tensor_objects = tensor_objects + ctx.save_for_backward(*tensors_to_save) + ctx.bwd_obj = bwd_obj + + def _autograd_backward(ctx, *grad_outputs): + bwd_obj = ctx.bwd_obj + if hasattr(bwd_obj, "setup_saved_tensors"): + bwd_obj.setup_saved_tensors(ctx) + ctx.tensor_objects = None + per_output_grads = grad_outputs[0] + bwd_obj.grad_output = _decode_none(per_output_grads[0]) + kwargs = _pack(bwd_obj, bwd_buckets) + bwd_args_flat = [kwargs[name] for name in bwd_arg_names] + bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), bwd_op_name) + grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] + out: List[Any] = list(fwd_slot_defaults) + tensor_list_lengths = getattr(ctx, "_te_fwd_tensor_list_lengths", {}) + for pos, length in tensor_list_lengths.items(): + if isinstance(out[pos], list): + out[pos] = [None] * length + for pos, g in zip(grad_targets, grads): + out[pos] = g + return tuple(out) + + torch.library.register_autograd( + fwd_qualname, + _autograd_backward, + setup_context=_setup_context, + lib=_TE_LIB, + ) + + +def _collect_universal_slot_offsets(buckets: List[_Bucket]) -> List[int]: + """Start index of each ``_UniversalTensorBucket`` group in the flat args.""" + offsets: List[int] = [] + pos = 0 + for bucket in buckets: + if isinstance(bucket, _UniversalTensorBucket): + offsets.append(pos) + pos += len(bucket.schema_slots()) + return offsets + + +def _flatten_subclass_into_slots( + new_args: List[Any], slot_offsets: List[int], subclass: type +) -> None: + """Rewrite each universal-bucket group whose ``Tensor?`` slot holds an + instance of ``subclass`` into the storage layout (3 slots: name / tensors / meta). + """ + for offset in slot_offsets: + val = new_args[offset] + if val is None or not isinstance(val, subclass): + continue + meta, tensors = _storage_flatten(val) + meta._data[_UniversalTensorBucket.KIND_KEY] = _UniversalKind.STORAGE + new_args[offset] = None + new_args[offset + 1] = list(tensors) + new_args[offset + 2] = meta + + +def _register_outer_forwarder( + *, + outer_op_name: str, + inner_op_name: str, + buckets: Optional[List[_Bucket]] = None, + subclass_list: Optional[List[type]] = None, +) -> None: + """Register the outer op's default kernel + fake: forward to the inner op, + optionally flattening registered subclass inputs in place first. + """ + inner_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_op_name) + input_flatten_enabled = bool(subclass_list) and buckets is not None + slot_offsets = _collect_universal_slot_offsets(buckets) if input_flatten_enabled else [] + + def _forward(*flat: Any) -> List[torch.Tensor]: + if not input_flatten_enabled: + return inner_op(*flat) + new_args = list(flat) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, slot_offsets, sub) + return inner_op(*new_args) + + _TE_LIB.impl(outer_op_name, _forward, "CompositeExplicitAutograd") + torch.library.register_fake(f"{_TE_OP_NAMESPACE}::{outer_op_name}", _forward, lib=_TE_LIB) + + +def _all_quantized_tensor_subclasses() -> List[type]: + """Return every imported ``QuantizedTensor`` wrapper subclass.""" + import transformer_engine.pytorch.tensor # noqa: F401 pylint: disable=import-outside-toplevel,unused-import + return [cls for cls in _STORAGE_REGISTRY.values() if issubclass(cls, QuantizedTensor)] + + +def register_custom_op( + *, + op_name: str, + input_tensors_for_grad: List[str], + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + setup_context: Callable[..., Any], + backward_arg_type: type, + backward_obj: type, + backward_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> Optional[Callable[..., Any]]: + """Register a TE module's forward + backward as torch custom ops. + + Always two-tier: an inner ``_base`` op carries the real schema / + autograd, and an outer ```` op forwards to it, flattening any + quantized-tensor wrapper inputs first via ``register_torch_dispatch`` (an + empty subclass list simply makes the outer op a pass-through, so a pure + plain-tensor / bf16 call goes straight through). + + Returns ``forward_fn(fwd_arg_type_instance)`` -- a drop-in for + ``Function.apply`` under ``torch.compiler.is_compiling()`` that dispatches + through the outer op and returns the user-facing outputs. + + Registration touches experimental ``torch.library`` / opaque-object APIs + that may be missing on older PyTorch. If it fails, this warns once and + returns ``None`` instead of raising, so callers can fall back to eager under + ``torch.compile`` (a graph break) rather than breaking import. + """ + try: + return _register_custom_op_impl( + op_name=op_name, + input_tensors_for_grad=input_tensors_for_grad, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + setup_context=setup_context, + backward_arg_type=backward_arg_type, + backward_obj=backward_obj, + backward_impl=backward_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_fake_impl=bwd_fake_impl, + ) + except (ImportError, AttributeError, RuntimeError, TypeError) as e: + warnings.warn( + f"Could not register the torch.compile custom op '{op_name}' " + f"({type(e).__name__}: {e}); modules using it will fall back to eager " + "execution under torch.compile (a graph break, incompatible with " + "fullgraph=True)." + ) + return None + + +def _register_custom_op_impl( + *, + op_name: str, + input_tensors_for_grad: List[str], + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + setup_context: Callable[..., Any], + backward_arg_type: type, + backward_obj: type, + backward_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], +) -> Callable[..., Any]: + """Body of :func:`register_custom_op`; see it for semantics.""" + outer_fwd_name = op_name + outer_bwd_name = f"{op_name}_backward" + inner_fwd_name = f"{op_name}_base" + inner_bwd_name = f"{outer_bwd_name}_base" + subclass_list = _all_quantized_tensor_subclasses() + + fwd_buckets = _get_buckets(fwd_arg_type) + bwd_buckets = _get_buckets(backward_arg_type) + fwd_tensor_field_names = _tensor_field_names(fwd_buckets) + bwd_tensor_field_names = _tensor_field_names(bwd_buckets) + + fwd_schema_args, fwd_arg_names = _build_schema(fwd_buckets) + bwd_schema_args, bwd_arg_names = _build_schema(bwd_buckets) + + num_grad_inputs = len(input_tensors_for_grad) + fwd_slot_defaults, grad_targets = _resolve_grad_targets( + fwd_buckets, fwd_arg_type, input_tensors_for_grad + ) + + _TE_LIB.define(f"{inner_fwd_name}{fwd_schema_args} -> Tensor[]") + _TE_LIB.define(f"{inner_bwd_name}{bwd_schema_args} -> Tensor[]") + _TE_LIB.define(f"{outer_fwd_name}{fwd_schema_args} -> Tensor[]") + _TE_LIB.define(f"{outer_bwd_name}{bwd_schema_args} -> Tensor[]") + + inner_fwd_qualname = f"{_TE_OP_NAMESPACE}::{inner_fwd_name}" + inner_bwd_qualname = f"{_TE_OP_NAMESPACE}::{inner_bwd_name}" + outer_fwd_qualname = f"{_TE_OP_NAMESPACE}::{outer_fwd_name}" + outer_bwd_qualname = f"{_TE_OP_NAMESPACE}::{outer_bwd_name}" + + _register_kernel( + op_name=inner_fwd_name, + op_qualname=inner_fwd_qualname, + arg_type=fwd_arg_type, + arg_names=fwd_arg_names, + buckets=fwd_buckets, + tensor_field_names=fwd_tensor_field_names, + impl=fwd_impl, + fake_impl=fwd_fake_impl, + format_result=_format_fwd_result, + ) + _register_kernel( + op_name=inner_bwd_name, + op_qualname=inner_bwd_qualname, + arg_type=backward_arg_type, + arg_names=bwd_arg_names, + buckets=bwd_buckets, + tensor_field_names=bwd_tensor_field_names, + impl=backward_impl, + fake_impl=bwd_fake_impl, + format_result=lambda g: _format_bwd_result(g, num_grad_inputs, inner_bwd_qualname), + ) + + autograd_common = { + "fwd_arg_type": fwd_arg_type, + "fwd_arg_names": fwd_arg_names, + "fwd_buckets": fwd_buckets, + "fwd_tensor_field_names": fwd_tensor_field_names, + "bwd_arg_names": bwd_arg_names, + "bwd_buckets": bwd_buckets, + "fwd_slot_defaults": fwd_slot_defaults, + "grad_targets": grad_targets, + "setup_context_user": setup_context, + "backward_obj_type": backward_obj, + "fwd_fake_impl": fwd_fake_impl, + } + _register_autograd_for_op( + fwd_op_name=inner_fwd_name, bwd_op_name=inner_bwd_name, **autograd_common + ) + _register_autograd_for_op( + fwd_op_name=outer_fwd_name, bwd_op_name=outer_bwd_name, **autograd_common + ) + + _register_outer_forwarder( + outer_op_name=outer_fwd_name, + inner_op_name=inner_fwd_name, + buckets=fwd_buckets, + subclass_list=list(subclass_list), + ) + _register_outer_forwarder(outer_op_name=outer_bwd_name, inner_op_name=inner_bwd_name) + + inner_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_fwd_name) + inner_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_bwd_name) + outer_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_fwd_name) + outer_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_bwd_name) + + fwd_slot_offsets = _collect_universal_slot_offsets(fwd_buckets) + bwd_slot_offsets = _collect_universal_slot_offsets(bwd_buckets) + + def _fwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) + return inner_fwd_op(*new_args) + + def _bwd_rule(mode, func, types, args, kwargs): + del mode, func, types, kwargs + new_args = list(args) + for sub in subclass_list: + _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) + return inner_bwd_op(*new_args) + + for sub in subclass_list: + torch.library.register_torch_dispatch(outer_fwd_qualname, sub, _fwd_rule, lib=_TE_LIB) + torch.library.register_torch_dispatch(outer_bwd_qualname, sub, _bwd_rule, lib=_TE_LIB) + + _quantized_tensor_passthrough_ops.add(outer_fwd_op.default) + _quantized_tensor_passthrough_ops.add(outer_bwd_op.default) + _quantized_tensor_passthrough_ops.add(inner_fwd_op.default) + _quantized_tensor_passthrough_ops.add(inner_bwd_op.default) + + def forward_fn(fwd_args): + proto_obj = _proto_view(fwd_args, fwd_tensor_field_names) + user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(proto_obj)) + kwargs = _pack(fwd_args, fwd_buckets) + flat_in = [kwargs[name] for name in fwd_arg_names] + result = outer_fwd_op(*flat_in) + + cursor = 0 + outputs: List[Any] = [] + for proto in user_fakes: + n = _proto_slot_count(proto) + chunk = [_decode_none(t) for t in result[cursor : cursor + n]] + cursor += n + outputs.append(_proto_reassemble(proto, chunk)) + + if len(outputs) == 1: + return outputs[0] + return tuple(outputs) + + return forward_fn diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 473483faff..d1ea805077 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -3,10 +3,12 @@ # See LICENSE for license information. """Linear API""" + from dataclasses import dataclass from typing import Any, Callable, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op +import math import warnings import weakref @@ -38,10 +40,10 @@ divide, init_method_constant, needs_quantized_gemm, - assert_dim_for_fp8_exec, nvtx_range_pop, nvtx_range_push, get_nvtx_range_context, + warn_compile_eager_fallback, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -58,9 +60,10 @@ from ..cpp_extensions import ( general_gemm, ) +from ..cpp_extensions.gemm import get_cublas_workspace from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type -from ..jit import no_torch_dynamo from ..graph import is_graph_capturing +from ..jit import no_torch_dynamo from ..quantized_tensor import ( QuantizedTensor, QuantizedTensorStorage, @@ -68,7 +71,7 @@ prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorProto +from ..dynamo import TensorProto, register_custom_op, is_value_opaque_quantizer from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -97,7 +100,24 @@ class LinearFwdArgs: bias: Optional[torch.Tensor] # --- Non-differentiable cached tensors --- - weight_workspace: Optional[torch.Tensor] + # Same union as ``weight`` so a cached quantized workspace is flattened to its + # inner tensors on the way into the op (symmetric with ``new_weight_workspace`` + # on the way out); a plain ``Tensor?`` slot can't carry a quantized subclass + # across the torch.compile custom-op boundary. + weight_workspace: Optional[TensorOrQuantized] + + # --- CUDA-graph workspace pinning (torch.compile / reduce-overhead) --- + # Fetched in the *traced* module forward and threaded in as op inputs purely so + # the process-global, lru_cached cuBLAS / NVFP4-RHT workspaces become graph + # inputs: they are then allocated at trace time in the normal allocator (external + # to the cudagraph private pool) instead of being created inside the op during + # capture, where a persistent allocation trips check_memory_pool ("tensor not + # tracked as outputs"). The op body never reads these; general_gemm / the + # quantizer fetch the same lru_cached globals by address. Pinning the workspace + # in the forward also covers the backward GEMM, which reuses the same cached + # global. None on eager / non-compiled paths. + cublas_workspace: Optional[torch.Tensor] + rht_matrix: Optional[torch.Tensor] # --- requires_grad flags (cached so backward does not re-query) --- input_requires_grad: bool @@ -131,7 +151,9 @@ class LinearFwdArgs: # --- Tensor / sequence parallelism --- parallel_mode: Optional[str] - tp_group: Optional[Any] + # ProcessGroup is a *reference*-opaque type: carried through the torch.compile + # custom op as a graph input (never baked into the graph as a constant). + tp_group: Optional[dist_group_type] tp_size: int tensor_parallel: bool sequence_parallel: bool @@ -159,6 +181,39 @@ class LinearFwdArgs: cpu_offloading: bool is_grad_enabled: bool + def compile_unsupported_reason(self) -> Optional[str]: + """Reason this config can't use the torch.compile custom-op path (else None).""" + if self.debug: + return "debug instrumentation (nvidia-dlfw-inspect)" + if self.fsdp_group is not None and self.is_grad_enabled: + return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP" + if ( + self.fp8_output + and self.is_grad_enabled + and (self.input_requires_grad or self.weight_requires_grad) + ): + return "differentiable fp8_output=True" + if self.cpu_offloading: + return "CPU activation offloading" + if self.wgrad_store is not None: + # Non-None only when delayed wgrad compute is on (see Linear.forward). + return "delayed wgrad compute (wgrad_store)" + if self.fuse_wgrad_accumulation: + return "fuse_wgrad_accumulation (main_grad)" + for quantizer in ( + self.input_quantizer, + self.weight_quantizer, + self.output_quantizer, + self.grad_input_quantizer, + self.grad_weight_quantizer, + self.grad_output_quantizer, + ): + # e.g. delayed-scaling Float8Quantizer and unregistered custom-recipe + # quantizers are not value-opaque and can't cross the custom-op boundary. + if quantizer is not None and not is_value_opaque_quantizer(quantizer): + return "a quantizer not registered as a torch.compile value-opaque type" + return None + @dataclass(slots=True) class LinearBwdArgs: @@ -196,7 +251,8 @@ class LinearBwdArgs: # --- Tensor / sequence parallelism --- parallel_mode: Optional[str] = None - tp_group: Optional[Any] = None + # Reference-opaque ProcessGroup (graph input), see LinearFwdArgs.tp_group. + tp_group: Optional[dist_group_type] = None tp_size: int = 1 tensor_parallel: bool = False sequence_parallel: bool = False @@ -296,9 +352,7 @@ def _linear_forward_impl( if ub_name is not None: nvtx_label = f"{nvtx_label}.{ub_name}" - # Make sure input dimensions are compatible - out_features, in_features = weight.shape - assert inp.shape[-1] == in_features, "GEMM not possible" + out_features = weight.shape[0] # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) @@ -340,7 +394,6 @@ def _linear_forward_impl( inputmat_total = None # Input tensor to pass to GEMM (gathered) own_quantized_input = False if fp8: - assert_dim_for_fp8_exec(inputmat, weight) if save_original_input: assert not isinstance( input_quantizer, Float8Quantizer @@ -887,7 +940,10 @@ def _linear_setup_ctx( bwd_args.use_bias = bias is not None bwd_args.requires_dgrad = fwd_args.input_requires_grad bwd_args.requires_wgrad = fwd_args.weight_requires_grad - bwd_args.inp_shape = inp.shape + # Don't store inp_shape in the value bundle: under torch.compile(dynamic=True) + # inp.shape contains SymInt dims which are not hashable in OpaqueValueBundle. + # The backward reconstructs inp_shape from grad_output + weight + SP config. + bwd_args.inp_shape = None # Numerical / dtype config bwd_args.activation_dtype = fwd_args.activation_dtype @@ -1035,6 +1091,19 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. ) nvtx_range_pop(f"{nvtx_label}.fsdp_gather") + # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes). + if bwd_args.inp_shape is None: + _w = saved_weight if saved_weight is not None else weight_fp8 + in_features = _w.shape[-1] + go_leading = grad_output.shape[0] + if bwd_args.parallel_mode == "column" and bwd_args.sequence_parallel: + inp_leading = go_leading // bwd_args.tp_size + elif bwd_args.parallel_mode == "row" and bwd_args.sequence_parallel: + inp_leading = go_leading * bwd_args.tp_size + else: + inp_leading = go_leading + bwd_args.inp_shape = torch.Size([inp_leading, *grad_output.shape[1:-1], in_features]) + # Configure Userbuffers communication (comm+GEMM overlap) bwd_args.ub_obj_gradout = None ub_obj_dgrad = None @@ -1574,8 +1643,19 @@ def _linear_backward_impl_fake( dgrad = None if args.requires_dgrad: # dgrad has the logical input shape and may be quantized for the next op. + # Derive shape from grad_output + weight + SP config instead of args.inp_shape: + # inp_shape is not stored in the value bundle under dynamic shapes (SymInt is + # not hashable in OpaqueValueBundle), so we reconstruct it here. + _in_features = weight.shape[-1] + _go_leading = args.grad_output.shape[0] + if args.parallel_mode == "column" and args.sequence_parallel: + _dgrad_leading = _go_leading // args.tp_size + elif args.parallel_mode == "row" and args.sequence_parallel: + _dgrad_leading = _go_leading * args.tp_size + else: + _dgrad_leading = _go_leading dgrad = TensorProto( - shape=tuple(args.inp_shape), + shape=(_dgrad_leading, *args.grad_output.shape[1:-1], _in_features), dtype=out_dtype, quantizer=args.grad_input_quantizer, device=args.grad_output.device, @@ -1603,6 +1683,21 @@ def _linear_backward_impl_fake( return wgrad, dgrad, grad_bias +# Custom op used under ``torch.compile``. +_linear_op = register_custom_op( + op_name="linear", + input_tensors_for_grad=["weight", "inp", "bias"], + fwd_arg_type=LinearFwdArgs, + fwd_impl=_linear_forward_impl, + fwd_fake_impl=_linear_forward_impl_fake, + setup_context=_linear_setup_ctx, + backward_arg_type=LinearBwdArgs, + backward_obj=LinearBwdArgs, + backward_impl=_linear_backward, + bwd_fake_impl=_linear_backward_impl_fake, +) + + class _Linear(torch.autograd.Function): """Linear semi-top level module Calls custom cuda extensions. @@ -1687,6 +1782,20 @@ def backward( return result +@no_torch_dynamo() +def _linear_eager( + weight_tensor: torch.Tensor, + inp: torch.Tensor, + bias: Optional[torch.Tensor], + fwd_args: LinearFwdArgs, + is_grad_enabled: bool, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Run ``_Linear`` eagerly, bypassing Dynamo.""" + if is_grad_enabled: + return _Linear.apply(weight_tensor, inp, bias, fwd_args) + return _Linear.forward(None, weight_tensor, inp, bias, fwd_args) + + class Linear(TransformerEngineBaseModule): """Applies a linear transformation to the incoming data :math:`y = xA^T + b` @@ -2084,7 +2193,6 @@ def reset_parameters(self, defer_init=False): elif self.parallel_mode == "column": set_tensor_model_parallel_attributes(getattr(self, bias), True, 0, 1) - @no_torch_dynamo() def forward( self, inp: torch.Tensor, @@ -2159,12 +2267,7 @@ def forward( grad_output_quantizer, ) = quantizers - if is_grad_enabled: - linear_fn = _Linear.apply - autograd_ctx = [] - else: - linear_fn = _Linear.forward - autograd_ctx = [None] + use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight" weight_workspace = ( @@ -2204,16 +2307,58 @@ def forward( ub_bulk_dgrad = self.ub_bulk_dgrad ub_bulk_wgrad = self.ub_bulk_wgrad + torch._check( + inp.shape[-1] == weight_tensor.shape[-1], + lambda: "GEMM not possible: input last dim must equal in_features", + ) + if self.fp8: + torch._check( + math.prod(inp.shape[:-1]) % 8 == 0, + lambda: ( + "FP8 execution requires the product of all input dimensions except" + " the last to be divisible by 8" + ), + ) + torch._check( + inp.shape[-1] % 16 == 0, + lambda: "FP8 execution requires the input last dimension to be divisible by 16", + ) + torch._check( + weight_tensor.shape[0] % 16 == 0, + lambda: "FP8 execution requires out_features to be divisible by 16", + ) + torch._check( + weight_tensor.shape[1] % 16 == 0, + lambda: "FP8 execution requires in_features to be divisible by 16", + ) + linear_bias_tensor = ( bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None ) wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None + + # Pin the lazily-cached cuBLAS (and NVFP4-RHT) workspaces as op inputs so + # they are materialized at trace time (external to the cudagraph pool) + # rather than inside the op during capture. See LinearFwdArgs for details. + cublas_workspace = None + rht_matrix = None + if use_compiled_op: + cublas_workspace = get_cublas_workspace(inp.device.index, False, False) + from ..tensor.nvfp4_tensor import NVFP4Quantizer, get_rht_matrix + + if isinstance(input_quantizer, NVFP4Quantizer): + rht_matrix = get_rht_matrix( + input_quantizer._with_random_sign_mask, inp.device.index + ) + fwd_args = LinearFwdArgs( # tensors weight=weight_tensor, inp=inp, bias=linear_bias_tensor, weight_workspace=weight_workspace, + cublas_workspace=cublas_workspace, + rht_matrix=rht_matrix, # requires_grad flags input_requires_grad=inp.requires_grad, weight_requires_grad=weight_tensor.requires_grad, @@ -2268,13 +2413,19 @@ def forward( cpu_offloading=is_cpu_offload_enabled(), is_grad_enabled=is_grad_enabled, ) - out, new_weight_workspace = linear_fn( - *autograd_ctx, - weight_tensor, - inp, - linear_bias_tensor, - fwd_args, - ) + + if use_compiled_op: + fallback_reason = fwd_args.compile_unsupported_reason() + if fallback_reason is not None: + warn_compile_eager_fallback(fallback_reason) + use_compiled_op = False + + if use_compiled_op: + out, new_weight_workspace = _linear_op(fwd_args) + else: + out, new_weight_workspace = _linear_eager( + weight_tensor, inp, linear_bias_tensor, fwd_args, is_grad_enabled + ) if new_weight_workspace is not None and cache_name is not None: if isinstance(new_weight_workspace, torch.Tensor): diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index ffbfbc1fdd..8eeaf2417d 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -26,6 +26,20 @@ ] +def warn_compile_eager_fallback(reason: str) -> None: + """Warn that a TE module is running eagerly under ``torch.compile``. + + Emitted when ``reason`` is unsupported on the module's compiled custom-op + path. Python's default warning filter dedups identical messages, so each + distinct ``reason`` is surfaced once. + """ + warnings.warn( + f"Falling back to eager execution under torch.compile: {reason} is " + "unsupported on the compiled path (graph-breaks under fullgraph=True).", + stacklevel=2, + ) + + @functools.lru_cache(maxsize=None) def get_cached_ones_tensor( num_elements: int, From fdac6597edbca38fdcd8375af2d308596cafe809 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 22 Jun 2026 12:48:26 +0200 Subject: [PATCH 024/108] [PyTorch] torch.compile: register TE custom ops via torch.library.custom_op Replace the low-level torch.library.Library (_TE_LIB.define/.impl + functional register_fake/register_autograd/register_torch_dispatch with lib=) with the standard torch.library.custom_op API, passing the dynamically built schema explicitly via schema=. register_fake/register_autograd/register_torch_dispatch are now methods on the returned CustomOpDef. Drops the TOR901 Library usage and is robust to re-registration (get_library_allowing_overwrite). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski (cherry picked from commit d590560c9a2ed26072ea48b0c8e3934ad62b9a29) --- .../pytorch/dynamo/custom_op.py | 82 +++++++++---------- 1 file changed, 41 insertions(+), 41 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index ef943d4e8e..2b1b34643c 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -35,7 +35,6 @@ ) _TE_OP_NAMESPACE = "transformer_engine_compile" -_TE_LIB = torch.library.Library(_TE_OP_NAMESPACE, "FRAGMENT") # noqa: TOR901 # ``None`` entries in an op's flat ``Tensor[]`` return are smuggled through a @@ -939,7 +938,7 @@ def _resolve_grad_targets( def _register_kernel( *, op_name: str, - op_qualname: str, + schema_str: str, arg_type: type, arg_names: List[str], buckets: List[_Bucket], @@ -947,8 +946,9 @@ def _register_kernel( impl: Callable[[Any], Any], fake_impl: Callable[[Any], Any], format_result: Callable[[Any], List[torch.Tensor]], -) -> None: - """Wire the real ``impl`` + the ``fake_impl`` (proto) into the library. +) -> Any: + """Define the op via ``torch.library.custom_op`` with the real ``impl`` + the + ``fake_impl`` (proto), returning the ``CustomOpDef``. The real kernel rebuilds the dataclass and runs ``impl``; the fake kernel runs the proto fake impl on the :func:`_proto_view`. Both go through @@ -966,13 +966,16 @@ def _fake(*flat: Any) -> List[torch.Tensor]: proto_obj = _proto_view(obj, tensor_field_names) return format_result(fake_impl(proto_obj)) - _TE_LIB.impl(op_name, _impl, "CompositeExplicitAutograd") - torch.library.register_fake(op_qualname, _fake, lib=_TE_LIB) + op = torch.library.custom_op( + f"{_TE_OP_NAMESPACE}::{op_name}", _impl, mutates_args=(), schema=schema_str + ) + op.register_fake(_fake) + return op def _register_autograd_for_op( *, - fwd_op_name: str, + fwd_op: Any, bwd_op_name: str, fwd_arg_type: type, fwd_arg_names: List[str], @@ -992,7 +995,6 @@ def _register_autograd_for_op( templates, reassembles each flat output chunk, and hands the saved tuple + ``ctx_attrs`` to the module's ``setup_context``. """ - fwd_qualname = f"{_TE_OP_NAMESPACE}::{fwd_op_name}" def _setup_context(ctx, inputs, output): ctx._te_fwd_tensor_list_lengths = { @@ -1054,12 +1056,7 @@ def _autograd_backward(ctx, *grad_outputs): out[pos] = g return tuple(out) - torch.library.register_autograd( - fwd_qualname, - _autograd_backward, - setup_context=_setup_context, - lib=_TE_LIB, - ) + fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context) def _collect_universal_slot_offsets(buckets: List[_Bucket]) -> List[int]: @@ -1093,12 +1090,14 @@ def _flatten_subclass_into_slots( def _register_outer_forwarder( *, outer_op_name: str, + schema_str: str, inner_op_name: str, buckets: Optional[List[_Bucket]] = None, subclass_list: Optional[List[type]] = None, -) -> None: - """Register the outer op's default kernel + fake: forward to the inner op, - optionally flattening registered subclass inputs in place first. +) -> Any: + """Define the outer op via ``torch.library.custom_op``: forward to the inner + op, optionally flattening registered subclass inputs in place first. Returns + the ``CustomOpDef``. """ inner_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_op_name) input_flatten_enabled = bool(subclass_list) and buckets is not None @@ -1112,8 +1111,11 @@ def _forward(*flat: Any) -> List[torch.Tensor]: _flatten_subclass_into_slots(new_args, slot_offsets, sub) return inner_op(*new_args) - _TE_LIB.impl(outer_op_name, _forward, "CompositeExplicitAutograd") - torch.library.register_fake(f"{_TE_OP_NAMESPACE}::{outer_op_name}", _forward, lib=_TE_LIB) + op = torch.library.custom_op( + f"{_TE_OP_NAMESPACE}::{outer_op_name}", _forward, mutates_args=(), schema=schema_str + ) + op.register_fake(_forward) + return op def _all_quantized_tensor_subclasses() -> List[type]: @@ -1208,19 +1210,14 @@ def _register_custom_op_impl( fwd_buckets, fwd_arg_type, input_tensors_for_grad ) - _TE_LIB.define(f"{inner_fwd_name}{fwd_schema_args} -> Tensor[]") - _TE_LIB.define(f"{inner_bwd_name}{bwd_schema_args} -> Tensor[]") - _TE_LIB.define(f"{outer_fwd_name}{fwd_schema_args} -> Tensor[]") - _TE_LIB.define(f"{outer_bwd_name}{bwd_schema_args} -> Tensor[]") + fwd_schema = f"{fwd_schema_args} -> Tensor[]" + bwd_schema = f"{bwd_schema_args} -> Tensor[]" - inner_fwd_qualname = f"{_TE_OP_NAMESPACE}::{inner_fwd_name}" inner_bwd_qualname = f"{_TE_OP_NAMESPACE}::{inner_bwd_name}" - outer_fwd_qualname = f"{_TE_OP_NAMESPACE}::{outer_fwd_name}" - outer_bwd_qualname = f"{_TE_OP_NAMESPACE}::{outer_bwd_name}" - _register_kernel( + inner_fwd_def = _register_kernel( op_name=inner_fwd_name, - op_qualname=inner_fwd_qualname, + schema_str=fwd_schema, arg_type=fwd_arg_type, arg_names=fwd_arg_names, buckets=fwd_buckets, @@ -1231,7 +1228,7 @@ def _register_custom_op_impl( ) _register_kernel( op_name=inner_bwd_name, - op_qualname=inner_bwd_qualname, + schema_str=bwd_schema, arg_type=backward_arg_type, arg_names=bwd_arg_names, buckets=bwd_buckets, @@ -1241,6 +1238,17 @@ def _register_custom_op_impl( format_result=lambda g: _format_bwd_result(g, num_grad_inputs, inner_bwd_qualname), ) + outer_fwd_def = _register_outer_forwarder( + outer_op_name=outer_fwd_name, + schema_str=fwd_schema, + inner_op_name=inner_fwd_name, + buckets=fwd_buckets, + subclass_list=list(subclass_list), + ) + outer_bwd_def = _register_outer_forwarder( + outer_op_name=outer_bwd_name, schema_str=bwd_schema, inner_op_name=inner_bwd_name + ) + autograd_common = { "fwd_arg_type": fwd_arg_type, "fwd_arg_names": fwd_arg_names, @@ -1255,19 +1263,11 @@ def _register_custom_op_impl( "fwd_fake_impl": fwd_fake_impl, } _register_autograd_for_op( - fwd_op_name=inner_fwd_name, bwd_op_name=inner_bwd_name, **autograd_common + fwd_op=inner_fwd_def, bwd_op_name=inner_bwd_name, **autograd_common ) _register_autograd_for_op( - fwd_op_name=outer_fwd_name, bwd_op_name=outer_bwd_name, **autograd_common - ) - - _register_outer_forwarder( - outer_op_name=outer_fwd_name, - inner_op_name=inner_fwd_name, - buckets=fwd_buckets, - subclass_list=list(subclass_list), + fwd_op=outer_fwd_def, bwd_op_name=outer_bwd_name, **autograd_common ) - _register_outer_forwarder(outer_op_name=outer_bwd_name, inner_op_name=inner_bwd_name) inner_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_fwd_name) inner_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_bwd_name) @@ -1292,8 +1292,8 @@ def _bwd_rule(mode, func, types, args, kwargs): return inner_bwd_op(*new_args) for sub in subclass_list: - torch.library.register_torch_dispatch(outer_fwd_qualname, sub, _fwd_rule, lib=_TE_LIB) - torch.library.register_torch_dispatch(outer_bwd_qualname, sub, _bwd_rule, lib=_TE_LIB) + outer_fwd_def.register_torch_dispatch(sub, _fwd_rule) + outer_bwd_def.register_torch_dispatch(sub, _bwd_rule) _quantized_tensor_passthrough_ops.add(outer_fwd_op.default) _quantized_tensor_passthrough_ops.add(outer_bwd_op.default) From 554f5a8fcd0f8316e5243fde4449307f87e70ef6 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 30 Jun 2026 14:06:46 +0200 Subject: [PATCH 025/108] [PyTorch] custom_op: note pytorch/pytorch#187434 enables dropping None sentinel Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 2b1b34643c..68b2aae59e 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -40,6 +40,10 @@ # ``None`` entries in an op's flat ``Tensor[]`` return are smuggled through a # 0-element uint8 tensor: a non-nullable ``Tensor[]`` schema is required for # ``register_autograd`` to attach a ``grad_fn`` to the outputs. +# +# TODO: once https://github.com/pytorch/pytorch/pull/187434 lands, a nullable +# ``Tensor?[]`` return schema lets ``None`` pass through directly and this +# sentinel encoding (``_encode_none`` / ``_decode_none``) can be removed. _NONE_SENTINEL_DTYPE = torch.uint8 From 6d3eecea21704f32d98cf80a9c57572e2b0e85c0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 16 Jun 2026 17:29:49 +0200 Subject: [PATCH 026/108] 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 027/108] [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 028/108] [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 029/108] 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 030/108] [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 031/108] [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 032/108] [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 033/108] [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 034/108] [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 035/108] [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 036/108] [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 037/108] [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 038/108] [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 039/108] [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 040/108] [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 ce230b72f60414904270866d3543eee197270c07 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 12:23:21 +0200 Subject: [PATCH 041/108] [PyTorch] Rename _Bucket.pack/unpack to to_slots/from_slots Address review: 'pack' wrongly connotes aggregating into one object, whereas the method fans a dataclass field out across schema slots. to_slots/from_slots is directional and matches the existing schema_slots / grad_slot vocabulary. Free functions _pack/_unpack keep their names (they pack the whole dataclass). Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 68b2aae59e..ce01d4ca9e 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -340,8 +340,8 @@ class _Bucket: A custom op only takes flat, simply-typed arguments, but a TE op takes a single ``@dataclass`` of mixed fields. Each bucket knows how to translate its kind of field both ways. ``try_build`` and ``schema_slots`` run once at - registration (to build the op's schema); ``pack`` and ``unpack`` run on each - call and must agree on the slot layout that ``schema_slots`` declares. + registration (to build the op's schema); ``to_slots`` and ``from_slots`` run + on each call and must agree on the slot layout that ``schema_slots`` declares. """ @classmethod @@ -363,21 +363,21 @@ def schema_slots(self) -> List[Tuple[str, str]]: """ raise NotImplementedError - def pack(self, owner: Any) -> List[Tuple[str, Any]]: + def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: """Read this field from the dataclass ``owner`` and produce the concrete value for each of its schema slots, as ``(slot_name, value)`` pairs. Composite values are flattened to fit the (tensor-only) slots: e.g. a quantized tensor is split into its plain inner buffers plus a metadata - bundle. Inverse of :meth:`unpack`. + bundle. Inverse of :meth:`from_slots`. """ raise NotImplementedError - def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: """Read this field's slots back from the op arguments ``args`` and write the reconstructed field value into ``kwargs`` (rebuilding any flattened composite). The filled ``kwargs`` are then used to rebuild the original - dataclass for the eager implementation. Inverse of :meth:`pack`. + dataclass for the eager implementation. Inverse of :meth:`to_slots`. """ raise NotImplementedError @@ -449,7 +449,7 @@ def try_build(cls, name: str, annot: Any) -> Optional["_UniversalTensorBucket"]: return cls(name) return None - def pack(self, owner: Any) -> List[Tuple[str, Any]]: + def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: value = getattr(owner, self.name) if value is None: return [ @@ -479,7 +479,7 @@ def pack(self, owner: Any) -> List[Tuple[str, Any]]: f"QuantizedTensorStorage, got {type(value).__name__}" ) - def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: meta = args[self.slot_meta()] kind = meta.get(self.KIND_KEY) if kind == _UniversalKind.NONE: @@ -512,10 +512,10 @@ def try_build(cls, name: str, annot: Any) -> Optional["_TensorBucket"]: def schema_slots(self) -> List[Tuple[str, str]]: return [(self.name, self.type_str)] - def pack(self, owner: Any) -> List[Tuple[str, Any]]: + def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: return [(self.name, getattr(owner, self.name))] - def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = args[self.name] def grad_slot(self) -> Optional[int]: @@ -551,10 +551,10 @@ def try_build(cls, name: str, annot: Any) -> Optional["_QuantizerBucket"]: def schema_slots(self) -> List[Tuple[str, str]]: return [(self.slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] - def pack(self, owner: Any) -> List[Tuple[str, Any]]: + def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: return [(self.slot(), OpaqueValueBundle({self.KEY: getattr(owner, self.name)}))] - def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = args[self.slot()][self.KEY] @@ -591,10 +591,10 @@ def try_build(cls, name: str, annot: Any) -> Optional["_ReferenceOpaqueBucket"]: def schema_slots(self) -> List[Tuple[str, str]]: return [(self.name, self.type_str)] - def pack(self, owner: Any) -> List[Tuple[str, Any]]: + def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: return [(self.name, getattr(owner, self.name))] - def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = args[self.name] @@ -628,10 +628,10 @@ def matches_field(cls, annot: Any) -> bool: def schema_slots(self) -> List[Tuple[str, str]]: return [(self.SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] - def pack(self, owner: Any) -> List[Tuple[str, Any]]: + def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: return [(self.SLOT, OpaqueValueBundle({n: getattr(owner, n) for n in self.names}))] - def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: if self.SLOT not in args: return meta = args[self.SLOT] @@ -642,8 +642,8 @@ def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: class _UnknownBucket(_Bucket): """Fallback for fields no other bucket claims. - Emits no slot; pack rejects non-trivial values (anything other than - ``None`` / all-``None`` sequence); unpack restores the field as ``None``. + Emits no slot; ``to_slots`` rejects non-trivial values (anything other + than ``None`` / all-``None`` sequence); ``from_slots`` restores the field as ``None``. """ def __init__(self, name: str, owner_cls_name: str) -> None: @@ -661,7 +661,7 @@ def _is_trivial(value: Any) -> bool: def schema_slots(self) -> List[Tuple[str, str]]: return [] - def pack(self, owner: Any) -> List[Tuple[str, Any]]: + def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: value = getattr(owner, self.name, None) if not self._is_trivial(value): raise TypeError( @@ -672,7 +672,7 @@ def pack(self, owner: Any) -> List[Tuple[str, Any]]: ) return [] - def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = None @@ -743,7 +743,7 @@ def _pack(obj: Any, buckets: List[_Bucket]) -> Dict[str, Any]: """ out: Dict[str, Any] = {} for bucket in buckets: - for name, value in bucket.pack(obj): + for name, value in bucket.to_slots(obj): out[name] = value return out @@ -755,7 +755,7 @@ def _unpack(cls: type, args: Dict[str, Any], buckets: List[_Bucket]) -> Any: """ kwargs: Dict[str, Any] = {} for bucket in buckets: - bucket.unpack(args, kwargs) + bucket.from_slots(args, kwargs) obj = cls.__new__(cls) for k, v in kwargs.items(): object.__setattr__(obj, k, v) From 4600cc4b2a013be93926928be0027878740383e7 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 12:34:32 +0200 Subject: [PATCH 042/108] [PyTorch] Drop dead weight_fp8 fallbacks for saved_weight in backward Address review: the saved_weight slot is unconditionally aliased to the weight parameter in forward (setup_ctx sets saved_weight = weight), so it is never None in backward. Simplify both the real backward inp_shape reconstruction and the fake backward. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index d1ea805077..9cb3def582 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1093,8 +1093,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes). if bwd_args.inp_shape is None: - _w = saved_weight if saved_weight is not None else weight_fp8 - in_features = _w.shape[-1] + in_features = saved_weight.shape[-1] go_leading = grad_output.shape[0] if bwd_args.parallel_mode == "column" and bwd_args.sequence_parallel: inp_leading = go_leading // bwd_args.tp_size @@ -1631,7 +1630,7 @@ 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 + weight = args.saved_weight out_dtype = args.activation_dtype out_features, in_features = weight.shape From a5dccf54b76538994ebcfe893bf4a6344ded8bdc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 12:37:25 +0200 Subject: [PATCH 043/108] [PyTorch] Drop underscore prefixes on fake-backward dgrad locals Address review: these are plain locals, not throwaways; the prefix was inconsistent with siblings. Also reuse the existing in_features instead of recomputing it as _in_features. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 9cb3def582..3ab2b4d217 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1645,16 +1645,15 @@ def _linear_backward_impl_fake( # Derive shape from grad_output + weight + SP config instead of args.inp_shape: # inp_shape is not stored in the value bundle under dynamic shapes (SymInt is # not hashable in OpaqueValueBundle), so we reconstruct it here. - _in_features = weight.shape[-1] - _go_leading = args.grad_output.shape[0] + go_leading = args.grad_output.shape[0] if args.parallel_mode == "column" and args.sequence_parallel: - _dgrad_leading = _go_leading // args.tp_size + dgrad_leading = go_leading // args.tp_size elif args.parallel_mode == "row" and args.sequence_parallel: - _dgrad_leading = _go_leading * args.tp_size + dgrad_leading = go_leading * args.tp_size else: - _dgrad_leading = _go_leading + dgrad_leading = go_leading dgrad = TensorProto( - shape=(_dgrad_leading, *args.grad_output.shape[1:-1], _in_features), + shape=(dgrad_leading, *args.grad_output.shape[1:-1], in_features), dtype=out_dtype, quantizer=args.grad_input_quantizer, device=args.grad_output.device, From acdaa91bdadd181ced758b53e6ebbf0caf63a2ba Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 12:39:22 +0200 Subject: [PATCH 044/108] [PyTorch] Trim verbose cublas_workspace field comment Address review: drop the cudagraph-pool mechanics from the field comment; keep just what the field is and why it's threaded through the op. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 3ab2b4d217..3e314e23ee 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -106,16 +106,11 @@ class LinearFwdArgs: # across the torch.compile custom-op boundary. weight_workspace: Optional[TensorOrQuantized] - # --- CUDA-graph workspace pinning (torch.compile / reduce-overhead) --- - # Fetched in the *traced* module forward and threaded in as op inputs purely so - # the process-global, lru_cached cuBLAS / NVFP4-RHT workspaces become graph - # inputs: they are then allocated at trace time in the normal allocator (external - # to the cudagraph private pool) instead of being created inside the op during - # capture, where a persistent allocation trips check_memory_pool ("tensor not - # tracked as outputs"). The op body never reads these; general_gemm / the - # quantizer fetch the same lru_cached globals by address. Pinning the workspace - # in the forward also covers the backward GEMM, which reuses the same cached - # global. None on eager / non-compiled paths. + # Workspace pinning (torch.compile). The process-global, lru_cached cuBLAS / + # NVFP4-RHT workspaces are fetched in the traced forward and threaded in as op + # inputs so they are allocated at trace time rather than lazily inside the op. + # The op body never reads them (general_gemm / the quantizer fetch the same + # globals by address). None on eager / non-compiled paths. cublas_workspace: Optional[torch.Tensor] rht_matrix: Optional[torch.Tensor] From d405c3bb76ee0e21465e641a5590b2f04b2e6128 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 12:47:10 +0200 Subject: [PATCH 045/108] [PyTorch] Handle PEP 604 X | Y unions in annotation introspection Address review: get_origin returns types.UnionType for X | None / X | Y but typing.Union for Optional[X] / Union[...], so the annotation-driven bucket dispatch silently failed to recognize the PEP 604 form. Add an _is_union helper that accepts both and use it in _strip_optional and the universal tensor-storage-union check. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index ce01d4ca9e..5d0deccfbf 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -6,6 +6,7 @@ from __future__ import annotations import dataclasses +import types import warnings from enum import Enum from typing import ( @@ -322,9 +323,19 @@ def _storage_unflatten(meta: Any, tensors: List[torch.Tensor]) -> Any: # --------------------------------------------------------------------------- # +def _is_union(annot: Any) -> bool: + """True for both ``typing.Union[...]`` / ``Optional[...]`` and PEP 604 ``X | Y``. + + ``get_origin`` returns ``typing.Union`` for the former but ``types.UnionType`` + for the latter, so the two syntaxes must be checked separately. + """ + origin = get_origin(annot) + return origin is Union or origin is getattr(types, "UnionType", ()) + + def _strip_optional(annot: Any) -> Tuple[Any, bool]: """If ``annot`` is ``Optional[X]`` return ``(X, True)``; else ``(annot, False)``.""" - if get_origin(annot) is Union: + if _is_union(annot): args = get_args(annot) if type(None) in args: non_none = [a for a in args if a is not type(None)] @@ -434,7 +445,7 @@ def schema_slots(self) -> List[Tuple[str, str]]: @staticmethod def _is_tensor_storage_union(annot: Any) -> bool: - if get_origin(annot) is not Union: + if not _is_union(annot): return False members = [a for a in get_args(annot) if a is not type(None)] if torch.Tensor not in members: From b1b24bf532084ca33854af95c1a91dbfbafa5510 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 12:52:03 +0200 Subject: [PATCH 046/108] [PyTorch] Note None in _UniversalTensorBucket docstring Address review: the field is Optional; record that in the type line and note the None case is tagged _UniversalKind.NONE. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 5d0deccfbf..f688bcab4f 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -411,12 +411,13 @@ class _UniversalKind(Enum): class _UniversalTensorBucket(_Bucket): - """``Tensor | QuantizedTensorStorage`` (also subclass tensor) field. + """``Tensor | QuantizedTensorStorage | None`` (also subclass tensor) field. Three slots regardless of value: ```` (``Tensor?`` -- plain / subclass tensor passes through, ``None`` for bare storage), ``__tensors`` (``Tensor[]`` flat inner tensors when flattened), ``__meta`` - (``OpaqueValueBundle`` flatten metadata + a ``__kind__`` marker). + (``OpaqueValueBundle`` flatten metadata + a ``__kind__`` marker). A ``None`` + field is tagged ``_UniversalKind.NONE`` with the other two slots empty. """ KIND_KEY = "__kind__" From 94fcf8204f0ea7c43b09d13ceaf8917132ef5d17 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:00:58 +0200 Subject: [PATCH 047/108] [PyTorch] Match universal-tensor field by exact TensorOrQuantized member set Address review: the previous heuristic (union containing torch.Tensor + some storage subclass) silently accepted junk unions like Union[Tensor, Float8Tensor, list]. Match the exact {torch.Tensor, QuantizedTensorStorage} member set instead; a bare quantized annotation or extra members are now rejected. Subclass tensor *values* still pass through the Tensor? slot at runtime. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index f688bcab4f..99283b3ef7 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -444,16 +444,18 @@ def schema_slots(self) -> List[Tuple[str, str]]: (self.slot_meta(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME), ] - @staticmethod - def _is_tensor_storage_union(annot: Any) -> bool: + # Canonical "plain tensor or quantized tensor" field annotation (the + # ``TensorOrQuantized`` alias in module code). Matched by exact member set, + # so a bare quantized annotation or an accidental extra union member is + # rejected rather than silently taken as a universal-tensor field. + _MEMBERS = frozenset({torch.Tensor, QuantizedTensorStorage}) + + @classmethod + def _is_tensor_storage_union(cls, annot: Any) -> bool: if not _is_union(annot): return False - members = [a for a in get_args(annot) if a is not type(None)] - if torch.Tensor not in members: - return False - return any( - isinstance(m, type) and issubclass(m, QuantizedTensorStorage) for m in members - ) + members = frozenset(a for a in get_args(annot) if a is not type(None)) + return members == cls._MEMBERS @classmethod def try_build(cls, name: str, annot: Any) -> Optional["_UniversalTensorBucket"]: From 75f41a9ce518763cc2d7d784962f0083bb55905b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:11:30 +0200 Subject: [PATCH 048/108] [PyTorch] Keep OpaqueValueBundle _frozen consistent when tagging __kind__ Address review: to_slots / _flatten_subclass_into_slots mutated meta._data[__kind__] after the bundle was built, leaving _frozen (the hash/eq key) without __kind__. Merge the kind into the meta dict via a new _storage_flatten(extra_meta=...) arg so it is present at construction and _frozen stays consistent. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 99283b3ef7..f68f6040ef 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -292,18 +292,24 @@ def _pg_reduce(pg: ProcessGroup) -> tuple: # type: ignore[valid-type] # --------------------------------------------------------------------------- # -def _storage_flatten(value: Any) -> Tuple["OpaqueValueBundle", List[torch.Tensor]]: +def _storage_flatten( + value: Any, extra_meta: Optional[Dict[str, Any]] = None +) -> Tuple["OpaqueValueBundle", List[torch.Tensor]]: """Split a ``QuantizedTensor`` / bare storage into ``(meta, Tensor[])``. The flatten context (embedding the value-opaque quantizer) plus inner names and -- for a wrapper subclass -- the outer geometry are stashed in the bundle so :func:`_storage_unflatten` can rebuild without PyTorch's ``outer_size``. + ``extra_meta`` is merged in before the bundle is built (so its ``_frozen`` + hash key stays consistent) -- used to tag the universal-slot ``__kind__``. """ inner_names, ctx = value.__tensor_flatten__() meta = dict(ctx) meta["_inner_names"] = list(inner_names) if isinstance(value, torch.Tensor): meta["_outer_shape"] = torch.Size(value.shape) + if extra_meta: + meta.update(extra_meta) tensors = [getattr(value, name) for name in inner_names] return OpaqueValueBundle(meta), tensors @@ -481,8 +487,7 @@ def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _UniversalKind.TENSOR})), ] if isinstance(value, QuantizedTensorStorage): - meta, tensors = _storage_flatten(value) - meta._data[self.KIND_KEY] = _UniversalKind.STORAGE + meta, tensors = _storage_flatten(value, {self.KIND_KEY: _UniversalKind.STORAGE}) return [ (self.slot_name(), None), (self.slot_tensors(), list(tensors)), @@ -1098,8 +1103,9 @@ def _flatten_subclass_into_slots( val = new_args[offset] if val is None or not isinstance(val, subclass): continue - meta, tensors = _storage_flatten(val) - meta._data[_UniversalTensorBucket.KIND_KEY] = _UniversalKind.STORAGE + meta, tensors = _storage_flatten( + val, {_UniversalTensorBucket.KIND_KEY: _UniversalKind.STORAGE} + ) new_args[offset] = None new_args[offset + 1] = list(tensors) new_args[offset + 2] = meta From c6cd8e1656893e6016026e736b2ef5cc98485062 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:12:43 +0200 Subject: [PATCH 049/108] [PyTorch] Collapse _QuantizerBucket.try_build guard into one condition Address review nit: merge the isinstance/issubclass checks, matching the pattern used elsewhere (e.g. _is_tensor_storage_union). Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index f68f6040ef..e11a2e794d 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -561,9 +561,7 @@ def slot(self) -> str: @classmethod def try_build(cls, name: str, annot: Any) -> Optional["_QuantizerBucket"]: stripped, _ = _strip_optional(annot) - if not isinstance(stripped, type): - return None - if issubclass(stripped, Quantizer): + if isinstance(stripped, type) and issubclass(stripped, Quantizer): return cls(name) return None From 754f0a470d251521075dc1631a9f4634e4b24799 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:17:29 +0200 Subject: [PATCH 050/108] [PyTorch] Rename _UniversalTensorBucket -> _TensorOrQuantizedBucket Address review: 'universal' was abstract; name the bucket after the TensorOrQuantized union it matches (also _UniversalKind -> _TensorOrQuantizedKind, _collect_universal_slot_offsets -> ...). Add a note that the field buckets are mutually exclusive on annotations, so _FIELD_BUCKETS order is iteration, not priority. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 48 ++++++++++--------- 1 file changed, 26 insertions(+), 22 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index e11a2e794d..76e7091e41 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -301,7 +301,7 @@ def _storage_flatten( and -- for a wrapper subclass -- the outer geometry are stashed in the bundle so :func:`_storage_unflatten` can rebuild without PyTorch's ``outer_size``. ``extra_meta`` is merged in before the bundle is built (so its ``_frozen`` - hash key stays consistent) -- used to tag the universal-slot ``__kind__``. + hash key stays consistent) -- used to tag the tensor-or-quantized slot ``__kind__``. """ inner_names, ctx = value.__tensor_flatten__() meta = dict(ctx) @@ -408,22 +408,22 @@ def grad_slot(self) -> Optional[int]: return None -class _UniversalKind(Enum): - """What a universal-tensor slot group carries, tagged in its ``__meta``.""" +class _TensorOrQuantizedKind(Enum): + """What a tensor-or-quantized slot group carries, tagged in its ``__meta``.""" NONE = "none" TENSOR = "tensor" STORAGE = "storage" -class _UniversalTensorBucket(_Bucket): +class _TensorOrQuantizedBucket(_Bucket): """``Tensor | QuantizedTensorStorage | None`` (also subclass tensor) field. Three slots regardless of value: ```` (``Tensor?`` -- plain / subclass tensor passes through, ``None`` for bare storage), ``__tensors`` (``Tensor[]`` flat inner tensors when flattened), ``__meta`` (``OpaqueValueBundle`` flatten metadata + a ``__kind__`` marker). A ``None`` - field is tagged ``_UniversalKind.NONE`` with the other two slots empty. + field is tagged ``_TensorOrQuantizedKind.NONE`` with the other two slots empty. """ KIND_KEY = "__kind__" @@ -453,7 +453,7 @@ def schema_slots(self) -> List[Tuple[str, str]]: # Canonical "plain tensor or quantized tensor" field annotation (the # ``TensorOrQuantized`` alias in module code). Matched by exact member set, # so a bare quantized annotation or an accidental extra union member is - # rejected rather than silently taken as a universal-tensor field. + # rejected rather than silently taken as a tensor-or-quantized field. _MEMBERS = frozenset({torch.Tensor, QuantizedTensorStorage}) @classmethod @@ -464,7 +464,7 @@ def _is_tensor_storage_union(cls, annot: Any) -> bool: return members == cls._MEMBERS @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_UniversalTensorBucket"]: + def try_build(cls, name: str, annot: Any) -> Optional["_TensorOrQuantizedBucket"]: if cls._is_tensor_storage_union(annot): return cls(name) return None @@ -475,7 +475,7 @@ def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: return [ (self.slot_name(), None), (self.slot_tensors(), []), - (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _UniversalKind.NONE})), + (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.NONE})), ] if isinstance(value, torch.Tensor): # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the @@ -484,10 +484,10 @@ def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: return [ (self.slot_name(), value), (self.slot_tensors(), []), - (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _UniversalKind.TENSOR})), + (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.TENSOR})), ] if isinstance(value, QuantizedTensorStorage): - meta, tensors = _storage_flatten(value, {self.KIND_KEY: _UniversalKind.STORAGE}) + meta, tensors = _storage_flatten(value, {self.KIND_KEY: _TensorOrQuantizedKind.STORAGE}) return [ (self.slot_name(), None), (self.slot_tensors(), list(tensors)), @@ -501,9 +501,9 @@ def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: meta = args[self.slot_meta()] kind = meta.get(self.KIND_KEY) - if kind == _UniversalKind.NONE: + if kind == _TensorOrQuantizedKind.NONE: kwargs[self.name] = None - elif kind == _UniversalKind.TENSOR: + elif kind == _TensorOrQuantizedKind.TENSOR: kwargs[self.name] = args[self.slot_name()] else: kwargs[self.name] = _storage_unflatten(meta, args[self.slot_tensors()]) @@ -694,8 +694,12 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: # Buckets, in priority order, owning ``try_build`` for a single field. +# These buckets are mutually exclusive on annotations (a plain ``torch.Tensor`` +# matches only ``_TensorBucket``; the ``TensorOrQuantized`` union only +# ``_TensorOrQuantizedBucket``; etc.), so the order is just iteration, not a +# priority ranking -- no annotation can be claimed by more than one. _FIELD_BUCKETS: Tuple[type, ...] = ( - _UniversalTensorBucket, + _TensorOrQuantizedBucket, _TensorBucket, _ReferenceOpaqueBucket, _QuantizerBucket, @@ -742,7 +746,7 @@ def _get_buckets(cls: type) -> List[_Bucket]: def _tensor_field_names(buckets: List[_Bucket]) -> List[str]: """Names of fields carrying tensors (for building the proto view).""" - return [b.name for b in buckets if isinstance(b, (_TensorBucket, _UniversalTensorBucket))] + return [b.name for b in buckets if isinstance(b, (_TensorBucket, _TensorOrQuantizedBucket))] def _build_schema(buckets: List[_Bucket]) -> Tuple[str, List[str]]: @@ -1080,12 +1084,12 @@ def _autograd_backward(ctx, *grad_outputs): fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context) -def _collect_universal_slot_offsets(buckets: List[_Bucket]) -> List[int]: - """Start index of each ``_UniversalTensorBucket`` group in the flat args.""" +def _collect_tensor_or_quantized_slot_offsets(buckets: List[_Bucket]) -> List[int]: + """Start index of each ``_TensorOrQuantizedBucket`` group in the flat args.""" offsets: List[int] = [] pos = 0 for bucket in buckets: - if isinstance(bucket, _UniversalTensorBucket): + if isinstance(bucket, _TensorOrQuantizedBucket): offsets.append(pos) pos += len(bucket.schema_slots()) return offsets @@ -1094,7 +1098,7 @@ def _collect_universal_slot_offsets(buckets: List[_Bucket]) -> List[int]: def _flatten_subclass_into_slots( new_args: List[Any], slot_offsets: List[int], subclass: type ) -> None: - """Rewrite each universal-bucket group whose ``Tensor?`` slot holds an + """Rewrite each tensor-or-quantized-bucket group whose ``Tensor?`` slot holds an instance of ``subclass`` into the storage layout (3 slots: name / tensors / meta). """ for offset in slot_offsets: @@ -1102,7 +1106,7 @@ def _flatten_subclass_into_slots( if val is None or not isinstance(val, subclass): continue meta, tensors = _storage_flatten( - val, {_UniversalTensorBucket.KIND_KEY: _UniversalKind.STORAGE} + val, {_TensorOrQuantizedBucket.KIND_KEY: _TensorOrQuantizedKind.STORAGE} ) new_args[offset] = None new_args[offset + 1] = list(tensors) @@ -1123,7 +1127,7 @@ def _register_outer_forwarder( """ inner_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_op_name) input_flatten_enabled = bool(subclass_list) and buckets is not None - slot_offsets = _collect_universal_slot_offsets(buckets) if input_flatten_enabled else [] + slot_offsets = _collect_tensor_or_quantized_slot_offsets(buckets) if input_flatten_enabled else [] def _forward(*flat: Any) -> List[torch.Tensor]: if not input_flatten_enabled: @@ -1296,8 +1300,8 @@ def _register_custom_op_impl( outer_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_fwd_name) outer_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_bwd_name) - fwd_slot_offsets = _collect_universal_slot_offsets(fwd_buckets) - bwd_slot_offsets = _collect_universal_slot_offsets(bwd_buckets) + fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_buckets) + bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_buckets) def _fwd_rule(mode, func, types, args, kwargs): del mode, func, types, kwargs From 6536692feac133bff380700abe033970461d0439 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:22:09 +0200 Subject: [PATCH 051/108] [PyTorch] Document that _SimpleBundleBucket is one-per-op Address review: unlike the per-field buckets, exactly one aggregating bucket exists per op (owns the shared _simple_meta slot). Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 76e7091e41..00d4b5064b 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -616,7 +616,12 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: class _SimpleBundleBucket(_Bucket): - """Aggregates every simple-typed field into a single OpaqueValueBundle.""" + """Aggregates every simple-typed field into a single OpaqueValueBundle. + + Unlike the per-field buckets, exactly one of these exists per op: it owns the + single shared ``_simple_meta`` slot, and ``_get_buckets`` builds it once from + all simple-typed field names collected across the dataclass. + """ SLOT = "_simple_meta" From 5b752be533a4eee523c501e903b546f2adcc6b47 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:26:17 +0200 Subject: [PATCH 052/108] [PyTorch] Rename _UnknownBucket -> _UnsupportedBucket, document behavior Address review: 'Unknown' misread as 'accepts anything'; the bucket is for fields of an unsupported type, tolerated only when the runtime value is trivial (None / all-None) and rejected otherwise. Clarify why the check must run at call time rather than in _get_buckets. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 00d4b5064b..ae05e41190 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -661,11 +661,19 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[n] = meta[n] -class _UnknownBucket(_Bucket): - """Fallback for fields no other bucket claims. - - Emits no slot; ``to_slots`` rejects non-trivial values (anything other - than ``None`` / all-``None`` sequence); ``from_slots`` restores the field as ``None``. +class _UnsupportedBucket(_Bucket): + """Fallback for fields whose type no other bucket can encode. + + Such a field cannot cross the op boundary, so it emits no slot and is + tolerated only when its runtime value carries nothing: ``to_slots`` accepts + ``None`` / an all-``None`` sequence (e.g. an unset ``Optional[Any]`` field + like ``fsdp_group`` or an empty ``fsdp_shapes`` on the compiled path) and + ``from_slots`` restores it as ``None``. A non-trivial value means the config + is genuinely unsupported under torch.compile, and ``to_slots`` raises. + + The check must run at call time (not in ``_get_buckets``): the annotation + alone -- e.g. ``Optional[Any]`` -- is valid when the value is ``None``, so + only the runtime value can decide. """ def __init__(self, name: str, owner_cls_name: str) -> None: @@ -743,7 +751,7 @@ def _get_buckets(cls: type) -> List[_Bucket]: elif _SimpleBundleBucket.matches_field(annot): simple_names.append(name) else: - buckets.append(_UnknownBucket(name, cls.__name__)) + buckets.append(_UnsupportedBucket(name, cls.__name__)) if simple_names: buckets.append(_SimpleBundleBucket(simple_names)) return buckets From 748265875b5d26afd0761ad9e046b8ba42cb69ae Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:28:40 +0200 Subject: [PATCH 053/108] [PyTorch] Keep _UnsupportedBucket docstring framework-generic Drop the linear-specific field names (fsdp_group / fsdp_shapes) from the generic custom-op framework docstring. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index ae05e41190..bdfad0261b 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -666,10 +666,10 @@ class _UnsupportedBucket(_Bucket): Such a field cannot cross the op boundary, so it emits no slot and is tolerated only when its runtime value carries nothing: ``to_slots`` accepts - ``None`` / an all-``None`` sequence (e.g. an unset ``Optional[Any]`` field - like ``fsdp_group`` or an empty ``fsdp_shapes`` on the compiled path) and - ``from_slots`` restores it as ``None``. A non-trivial value means the config - is genuinely unsupported under torch.compile, and ``to_slots`` raises. + ``None`` / an all-``None`` sequence (e.g. an unset ``Optional[Any]`` field, + or an empty list, on the compiled path) and ``from_slots`` restores it as + ``None``. A non-trivial value means the config is genuinely unsupported + under torch.compile, and ``to_slots`` raises. The check must run at call time (not in ``_get_buckets``): the annotation alone -- e.g. ``Optional[Any]`` -- is valid when the value is ``None``, so From 011c268a8fc83b6f27ee2610587ee2495783cc98 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:33:33 +0200 Subject: [PATCH 054/108] [PyTorch] Make _Bucket.to_slots return a dict Address review: to_slots returned a list of (slot, value) pairs while from_slots reads a dict; return a {slot: value} dict for symmetry, and simplify _pack to out.update(bucket.to_slots(obj)). Slot names are unique per bucket and final arg order comes from arg_names, so no behavior change. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 59 +++++++++---------- 1 file changed, 29 insertions(+), 30 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index bdfad0261b..7c4a7d91a9 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -380,9 +380,9 @@ def schema_slots(self) -> List[Tuple[str, str]]: """ raise NotImplementedError - def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: + def to_slots(self, owner: Any) -> Dict[str, Any]: """Read this field from the dataclass ``owner`` and produce the concrete - value for each of its schema slots, as ``(slot_name, value)`` pairs. + value for each of its schema slots, as a ``{slot_name: value}`` dict. Composite values are flattened to fit the (tensor-only) slots: e.g. a quantized tensor is split into its plain inner buffers plus a metadata @@ -469,30 +469,30 @@ def try_build(cls, name: str, annot: Any) -> Optional["_TensorOrQuantizedBucket" return cls(name) return None - def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: + def to_slots(self, owner: Any) -> Dict[str, Any]: value = getattr(owner, self.name) if value is None: - return [ - (self.slot_name(), None), - (self.slot_tensors(), []), - (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.NONE})), - ] + return { + self.slot_name(): None, + self.slot_tensors(): [], + self.slot_meta(): OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.NONE}), + } if isinstance(value, torch.Tensor): # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the # ``Tensor?`` slot; subclass flattening (if any) is done by the # outer op's ``register_torch_dispatch`` rule. - return [ - (self.slot_name(), value), - (self.slot_tensors(), []), - (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.TENSOR})), - ] + return { + self.slot_name(): value, + self.slot_tensors(): [], + self.slot_meta(): OpaqueValueBundle({self.KIND_KEY: _TensorOrQuantizedKind.TENSOR}), + } if isinstance(value, QuantizedTensorStorage): meta, tensors = _storage_flatten(value, {self.KIND_KEY: _TensorOrQuantizedKind.STORAGE}) - return [ - (self.slot_name(), None), - (self.slot_tensors(), list(tensors)), - (self.slot_meta(), meta), - ] + return { + self.slot_name(): None, + self.slot_tensors(): list(tensors), + self.slot_meta(): meta, + } raise TypeError( f"field {self.name!r} expected None, torch.Tensor, or " f"QuantizedTensorStorage, got {type(value).__name__}" @@ -531,8 +531,8 @@ def try_build(cls, name: str, annot: Any) -> Optional["_TensorBucket"]: def schema_slots(self) -> List[Tuple[str, str]]: return [(self.name, self.type_str)] - def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: - return [(self.name, getattr(owner, self.name))] + def to_slots(self, owner: Any) -> Dict[str, Any]: + return {self.name: getattr(owner, self.name)} def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = args[self.name] @@ -568,8 +568,8 @@ def try_build(cls, name: str, annot: Any) -> Optional["_QuantizerBucket"]: def schema_slots(self) -> List[Tuple[str, str]]: return [(self.slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] - def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: - return [(self.slot(), OpaqueValueBundle({self.KEY: getattr(owner, self.name)}))] + def to_slots(self, owner: Any) -> Dict[str, Any]: + return {self.slot(): OpaqueValueBundle({self.KEY: getattr(owner, self.name)})} def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = args[self.slot()][self.KEY] @@ -608,8 +608,8 @@ def try_build(cls, name: str, annot: Any) -> Optional["_ReferenceOpaqueBucket"]: def schema_slots(self) -> List[Tuple[str, str]]: return [(self.name, self.type_str)] - def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: - return [(self.name, getattr(owner, self.name))] + def to_slots(self, owner: Any) -> Dict[str, Any]: + return {self.name: getattr(owner, self.name)} def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = args[self.name] @@ -650,8 +650,8 @@ def matches_field(cls, annot: Any) -> bool: def schema_slots(self) -> List[Tuple[str, str]]: return [(self.SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)] - def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: - return [(self.SLOT, OpaqueValueBundle({n: getattr(owner, n) for n in self.names}))] + def to_slots(self, owner: Any) -> Dict[str, Any]: + return {self.SLOT: OpaqueValueBundle({n: getattr(owner, n) for n in self.names})} def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: if self.SLOT not in args: @@ -691,7 +691,7 @@ def _is_trivial(value: Any) -> bool: def schema_slots(self) -> List[Tuple[str, str]]: return [] - def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: + def to_slots(self, owner: Any) -> Dict[str, Any]: value = getattr(owner, self.name, None) if not self._is_trivial(value): raise TypeError( @@ -700,7 +700,7 @@ def to_slots(self, owner: Any) -> List[Tuple[str, Any]]: "and carries a non-trivial value; add a matching bucket in " "dynamo.py to handle it." ) - return [] + return {} def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = None @@ -777,8 +777,7 @@ def _pack(obj: Any, buckets: List[_Bucket]) -> Dict[str, Any]: """ out: Dict[str, Any] = {} for bucket in buckets: - for name, value in bucket.to_slots(obj): - out[name] = value + out.update(bucket.to_slots(obj)) return out From e5d49853d4e0adf2f12c198b3e8dadf7cd949533 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:40:24 +0200 Subject: [PATCH 055/108] [PyTorch] Move proto reassembly into TensorProto.assemble Address review: the general 'rebuild a tensor from inner buffers' logic was duplicated between custom_op._proto_reassemble and TensorProto.create_tensor. Extract TensorProto.assemble(inner_tensors); create_tensor becomes assemble(create_inner_tensors()) and _proto_reassemble becomes a thin op-boundary wrapper (None -> None else proto.assemble(chunk)). The general primitive now lives in tensor_proto.py. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 17 +++------- .../pytorch/dynamo/tensor_proto.py | 34 +++++++++++++------ 2 files changed, 27 insertions(+), 24 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 7c4a7d91a9..57309f085b 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -834,22 +834,13 @@ def _proto_reassemble( ) -> Optional[Union[torch.Tensor, QuantizedTensorStorage]]: """Rebuild the value described by ``proto`` from its flat tensors ``chunk``. - ``proto`` describes one output: ``None`` (-> ``None``), a plain tensor - (``chunk`` is the single tensor, returned as-is), or a quantized tensor - (``chunk`` are its inner tensors, reassembled into the wrapper subclass via - ``__tensor_unflatten__``). + ``proto is None`` -> ``None`` (op-boundary sentinel for an absent output); + otherwise delegates to :meth:`TensorProto.assemble`, which returns a plain + tensor as-is or reassembles a quantized tensor from its inner buffers. """ if proto is None: return None - if proto.quantizer is None: - return chunk[0] - inner_names = proto.inner_names() - meta = proto.create_metadata() - shape = tuple(proto.shape) - stride = _contiguous_stride(shape) - storage_cls = _STORAGE_REGISTRY[meta["cls"]] - inner_dict = dict(zip(inner_names, chunk)) - return storage_cls.__tensor_unflatten__(inner_dict, meta, shape, stride) + return proto.assemble(chunk) def _value_to_flat_tensors( diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index 911b4151bf..1393aeee13 100644 --- a/transformer_engine/pytorch/dynamo/tensor_proto.py +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -117,11 +117,31 @@ 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 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 the op's flat ``Tensor[]`` payload). + + 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: + 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"]] + return storage_cls.__tensor_unflatten__(inner, ctx, shape, _contiguous_stride(shape)) + 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__``. + 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") @@ -131,15 +151,7 @@ def create_tensor(self) -> torch.Tensor: 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)) + return self.assemble(self.create_inner_tensors()) def to_tensor_proto(tensor: Any) -> TensorProto: From a01c77433a4271793084f8af517595342a63ba6a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:44:09 +0200 Subject: [PATCH 056/108] [PyTorch] Dedup flatten branches in _value_to_flat_tensors Address review suggestion: check __tensor_flatten__ first (covers both subclass tensors and bare storage), then plain torch.Tensor. Removes the duplicated flatten block and the redundant type-is-not-Tensor guard (a plain torch.Tensor has no __tensor_flatten__). Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 57309f085b..6edf415b10 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -855,16 +855,11 @@ def _value_to_flat_tensors( return [_encode_none(None)] if isinstance(value, TensorProto): return [_encode_none(t) for t in value.create_inner_tensors()] - if isinstance(value, torch.Tensor): - if type(value) is not torch.Tensor and hasattr( # pylint: disable=unidiomatic-typecheck - value, "__tensor_flatten__" - ): - inner_names, _ = value.__tensor_flatten__() - return [_encode_none(getattr(value, n)) for n in inner_names] - return [_encode_none(value)] if hasattr(value, "__tensor_flatten__"): inner_names, _ = value.__tensor_flatten__() return [_encode_none(getattr(value, n)) for n in inner_names] + if isinstance(value, torch.Tensor): + return [_encode_none(value)] raise TypeError( f"unsupported value type {type(value).__name__}; expected None / " "torch.Tensor / tensor subclass / bare storage / TensorProto." From 858603c4dc890888d4437808cab207a761cb8a43 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:45:18 +0200 Subject: [PATCH 057/108] [PyTorch] Use explicit None check for saved-tensors slot (review nit) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 6edf415b10..af1ab4c48c 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -880,9 +880,10 @@ def _format_fwd_result(result: Any) -> List[torch.Tensor]: flat: List[torch.Tensor] = [] for value in result[:num_outputs]: flat.extend(_value_to_flat_tensors(value)) - saved = result[num_outputs] or () - for value in saved: - flat.extend(_value_to_flat_tensors(value)) + saved = result[num_outputs] + if saved is not None: + for value in saved: + flat.extend(_value_to_flat_tensors(value)) return flat From 55de0e2a7e1df4e3bb8a23d66875ff45bc55b302 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 13:59:29 +0200 Subject: [PATCH 058/108] [PyTorch] Drop dead tensor_objects slot from fwd-impl return contract The middle trailing slot was always None in Linear (real + fake) and never read by the framework. Reduce _FWD_TRAILING_SLOTS to 2, fix the ctx_attrs index in _split_fwd_fake_result, and drop the None slot from Linear's returns/signatures. Contract is now (*user_outputs, tensors_to_save, ctx_attrs). Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 8 ++++---- transformer_engine/pytorch/module/linear.py | 12 ++++++------ 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index af1ab4c48c..1fe6a054b1 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -866,9 +866,9 @@ def _value_to_flat_tensors( ) -# Trailing slots in every fwd-impl return: ``tensors_to_save, tensor_objects, -# ctx_attrs``. User-output count is ``len(result) - this``. -_FWD_TRAILING_SLOTS = 3 +# Trailing slots in every fwd-impl return: ``tensors_to_save, ctx_attrs``. +# User-output count is ``len(result) - this``. +_FWD_TRAILING_SLOTS = 2 def _format_fwd_result(result: Any) -> List[torch.Tensor]: @@ -914,7 +914,7 @@ def _split_fwd_fake_result( """Slice a fwd fake-impl return into ``(user_fakes, saved_fakes, ctx_attrs)``.""" num_outputs = len(result) - _FWD_TRAILING_SLOTS saved = result[num_outputs] - ctx_attrs = result[num_outputs + 2] + ctx_attrs = result[num_outputs + 1] user_fakes = list(result[:num_outputs]) saved_fakes = list(saved) if saved is not None else [] ctx_attrs = dict(ctx_attrs) if ctx_attrs else {} diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 3e314e23ee..79404fbe79 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -307,13 +307,13 @@ def _check_fp8_reduce_and_update(): def _linear_forward_impl( args: LinearFwdArgs, -) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], None, Optional[Dict]]: +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], Optional[Dict]]: """Forward implementation for the linear layer. - Returns ``(out, new_weight_workspace, tensors_to_save_from_forward, None, + Returns ``(out, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs)``. ``new_weight_workspace`` is the freshly produced FP8 weight workspace (returned alongside ``out`` so the caller can refresh its - cache). The last three are ``None`` when gradients are disabled. + cache). The last two are ``None`` when gradients are disabled. """ weight = args.weight @@ -688,12 +688,12 @@ def _linear_forward_impl( "saved_tensor_aliases": saved_tensor_aliases, } - return out, new_weight_workspace, tensors_to_save_from_forward, None, ctx_attrs + return out, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs def _linear_forward_impl_fake( args: LinearFwdArgs, -) -> Tuple[TensorProto, Optional[TensorProto], Optional[Tuple[Any, ...]], None, Optional[Dict]]: +) -> Tuple[TensorProto, Optional[TensorProto], Optional[Tuple[Any, ...]], 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.""" @@ -892,7 +892,7 @@ def _linear_forward_impl_fake( "saved_tensor_aliases": saved_tensor_aliases, } - return out, new_weight_workspace, tensors_to_save_from_forward, None, ctx_attrs + return out, new_weight_workspace, tensors_to_save_from_forward, ctx_attrs def _linear_setup_ctx( From 59470356122306b24e934dc57fc78ea049d12cdb Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:03:17 +0200 Subject: [PATCH 059/108] [PyTorch] Validate fwd-impl return contract at the boundary Add _check_fwd_result: enforce (*user_outputs, tensors_to_save, ctx_attrs) shape and trailing-slot types with a clear error for op authors, called from both _format_fwd_result and _split_fwd_fake_result. Traceable (isinstance on non-tensors), so no graph break under compile. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 1fe6a054b1..56bad807dd 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -871,11 +871,31 @@ def _value_to_flat_tensors( _FWD_TRAILING_SLOTS = 2 +def _check_fwd_result(result: Any) -> None: + """Validate a fwd-impl return against the + ``(*user_outputs, tensors_to_save, ctx_attrs)`` contract, with a clear + message for op authors (user-output *types* are checked later, by + :func:`_value_to_flat_tensors`). + """ + if not isinstance(result, tuple) or len(result) < _FWD_TRAILING_SLOTS: + raise TypeError( + f"fwd impl must return a tuple of >= {_FWD_TRAILING_SLOTS} elements " + "(*user_outputs, tensors_to_save, ctx_attrs); " + f"got {type(result).__name__}" + ) + tensors_to_save, ctx_attrs = result[-2], result[-1] + if tensors_to_save is not None and not isinstance(tensors_to_save, (list, tuple)): + raise TypeError("fwd impl 'tensors_to_save' slot must be a list/tuple or None") + if ctx_attrs is not None and not isinstance(ctx_attrs, dict): + raise TypeError("fwd impl 'ctx_attrs' slot must be a dict or None") + + def _format_fwd_result(result: Any) -> List[torch.Tensor]: """Pack a fwd-impl return tuple into the op's ``Tensor[]`` payload. User outputs first, then saved-for-backward tensors in declaration order. """ + _check_fwd_result(result) num_outputs = len(result) - _FWD_TRAILING_SLOTS flat: List[torch.Tensor] = [] for value in result[:num_outputs]: @@ -912,6 +932,7 @@ def _split_fwd_fake_result( result: Tuple[Any, ...], ) -> Tuple[List[Any], List[Any], Dict[str, Any]]: """Slice a fwd fake-impl return into ``(user_fakes, saved_fakes, ctx_attrs)``.""" + _check_fwd_result(result) num_outputs = len(result) - _FWD_TRAILING_SLOTS saved = result[num_outputs] ctx_attrs = result[num_outputs + 1] From d1c08275963514600322714a13023d4565265cb7 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:05:44 +0200 Subject: [PATCH 060/108] [PyTorch] Run fwd-result check only on the fake path (compile-time) _format_fwd_result is the real kernel formatter and runs per compiled call; drop the check there to avoid per-call overhead. Keep it only in _split_fwd_fake_result, which runs at trace time -- the real impl must return the same shape as the fake, so validating the fake covers both at zero runtime cost. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 56bad807dd..9a5c112237 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -876,6 +876,11 @@ def _check_fwd_result(result: Any) -> None: ``(*user_outputs, tensors_to_save, ctx_attrs)`` contract, with a clear message for op authors (user-output *types* are checked later, by :func:`_value_to_flat_tensors`). + + Only called on the fake path (:func:`_split_fwd_fake_result`), which runs at + trace/compile time -- so this is a compile-time check with no per-call cost. + The real impl must return the same shape as the fake, so validating the fake + covers both. """ if not isinstance(result, tuple) or len(result) < _FWD_TRAILING_SLOTS: raise TypeError( @@ -895,7 +900,6 @@ def _format_fwd_result(result: Any) -> List[torch.Tensor]: User outputs first, then saved-for-backward tensors in declaration order. """ - _check_fwd_result(result) num_outputs = len(result) - _FWD_TRAILING_SLOTS flat: List[torch.Tensor] = [] for value in result[:num_outputs]: From b70bf8c694c52fc1f9f99bbf3555ca640b5fbd1b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:14:55 +0200 Subject: [PATCH 061/108] [PyTorch] Document the custom-op callable + arg-container contract Address review: register_custom_op now documents the fwd/backward arg dataclasses (field annotations -> schema), how the backward container is populated (setup_context + the optional setup_saved_tensors hook + grad_output), and each callable's contract including the fwd-impl return shape (*user_outputs, tensors_to_save, ctx_attrs). Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 9a5c112237..8869698466 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1199,6 +1199,48 @@ def register_custom_op( ``Function.apply`` under ``torch.compiler.is_compiling()`` that dispatches through the outer op and returns the user-facing outputs. + Arg containers. ``fwd_arg_type`` and ``backward_arg_type`` are ``@dataclass``es + whose *field annotations* define the op schema: each field maps to one or more + flat schema slots (tensor fields cross the boundary as tensors, quantizers ride + as value-opaque objects, simple values are bundled -- see the ``_Bucket`` + classes). The caller builds a ``fwd_arg_type`` instance and passes it to the + returned ``forward_fn``. + + How the backward container is populated. ``setup_context`` fills the + ``backward_arg_type`` instance's non-tensor fields (quantizers, config) from + forward state and returns the tensors to persist; the framework saves them + (``ctx.save_for_backward``). Before ``backward_impl`` runs, the framework + restores those tensors into the container's *tensor* fields by calling its + optional ``setup_saved_tensors(self, ctx)`` hook (invoked only if defined), + and sets ``grad_output`` directly. So ``backward_impl`` receives a + fully-populated ``backward_arg_type``. + + Callable contracts: + + * ``fwd_impl(fwd_args) -> (*user_outputs, tensors_to_save, ctx_attrs)`` -- the + real forward. ``user_outputs``: op outputs (tensor / quantized / ``None``); + ``tensors_to_save``: list/tuple (or ``None``) of tensors for backward; + ``ctx_attrs``: dict (or ``None``) of plain metadata for ``setup_context``. + The trailing two slots are fixed (``_FWD_TRAILING_SLOTS``); everything + before them is a user output. + * ``fwd_fake_impl(fwd_args)`` -- data-free traceable twin of ``fwd_impl``: + same return shape, but tensor outputs are :class:`TensorProto`. Must match + ``fwd_impl``'s shape (checked at compile time by ``_check_fwd_result``). + * ``setup_context(bwd_obj, fwd_args, user_outputs, ctx_attrs, saved) + -> tensors_to_save`` -- populate ``bwd_obj`` from forward state; return the + tensors to persist across the boundary. + * ``backward_impl(bwd_args) -> grads`` -- exactly one grad per + ``input_tensors_for_grad`` entry, in that order (``None`` for a + non-differentiable input). + * ``bwd_fake_impl(bwd_args)`` -- data-free twin of ``backward_impl`` returning + :class:`TensorProto` grads. + * ``backward_obj.setup_saved_tensors(ctx)`` -- optional hook on the backward + container (see above); skipped if absent. + + ``input_tensors_for_grad`` lists the ``fwd_arg_type`` fields that receive + gradients (this fixes the backward grad order). ``backward_obj`` is the type + instantiated to hold the backward args (usually == ``backward_arg_type``). + Registration touches experimental ``torch.library`` / opaque-object APIs that may be missing on older PyTorch. If it fails, this warns once and returns ``None`` instead of raising, so callers can fall back to eager under From db5d90260a79300129fa525be0240de3ac0bda48 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:25:59 +0200 Subject: [PATCH 062/108] [PyTorch] Fix eager _Linear.forward unpack after tensor_objects slot removal Commit 55de0e2a7 dropped the middle slot from _linear_forward_impl's return but missed the eager path, which still unpacked a 5-tuple -> ValueError under _Linear.apply (non-compiled). The compiled path was unaffected (indexes by _FWD_TRAILING_SLOTS), so the earlier smoke test missed it; a compiled-vs-eager grad-parity check now covers both. 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 79404fbe79..7207bc586f 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1722,7 +1722,6 @@ def forward( out, new_weight_workspace, tensors_to_save_from_forward, - _, ctx_attrs, ) = _linear_forward_impl(fwd_args) if ctx is not None: From 246a74fa3c12f38b089c25c66a3887536a57bb4c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:25:59 +0200 Subject: [PATCH 063/108] [PyTorch] Drop redundant fwd_slot_defaults in backward grad layout Address review: the per-slot [] / None template was immediately overwritten -- every Tensor[] slot is always recorded in _te_fwd_tensor_list_lengths, so its [] default never survives. Init the grad-output list as [None]*slot_count and let that loop shape the Tensor[] slots; _resolve_grad_targets now returns (slot_count, grad_targets) and the isinstance guard is gone. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 27 +++++++++---------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 8869698466..4110643e45 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -955,20 +955,17 @@ def _resolve_grad_targets( fwd_buckets: List[_Bucket], fwd_arg_type: type, input_tensors_for_grad: List[str], -) -> Tuple[List[Any], List[int]]: +) -> Tuple[int, List[int]]: """Validate ``input_tensors_for_grad`` and resolve the grad-output layout. - Returns ``(fwd_slot_defaults, grad_targets)``: the per-slot no-grad template - (``[]`` for ``Tensor[]`` slots, ``None`` otherwise) and, for each requested - input name, the schema-slot index its gradient maps to. + Returns ``(slot_count, grad_targets)``: the total number of input schema + slots and, for each requested input name, the schema-slot index its gradient + maps to. """ - fwd_slot_defaults: List[Any] = [] name_to_slot: Dict[str, int] = {} slot_offset = 0 for bucket in fwd_buckets: slots = bucket.schema_slots() - for _, type_str in slots: - fwd_slot_defaults.append([] if type_str.endswith("[]") else None) grad_slot = bucket.grad_slot() if grad_slot is not None: name_to_slot[bucket.name] = slot_offset + grad_slot @@ -981,7 +978,7 @@ def _resolve_grad_targets( f"schema: {unknown}" ) grad_targets = [name_to_slot[n] for n in input_tensors_for_grad] - return fwd_slot_defaults, grad_targets + return slot_offset, grad_targets def _register_kernel( @@ -1032,7 +1029,7 @@ def _register_autograd_for_op( fwd_tensor_field_names: List[str], bwd_arg_names: List[str], bwd_buckets: List[_Bucket], - fwd_slot_defaults: List[Any], + slot_count: int, grad_targets: List[int], setup_context_user: Callable[..., Any], backward_obj_type: type, @@ -1096,11 +1093,13 @@ def _autograd_backward(ctx, *grad_outputs): bwd_args_flat = [kwargs[name] for name in bwd_arg_names] bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), bwd_op_name) grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] - out: List[Any] = list(fwd_slot_defaults) + # One grad per input schema slot: default None, but a ``Tensor[]`` slot + # (always recorded in ``_te_fwd_tensor_list_lengths``) needs a + # list-shaped no-grad of matching length. + out: List[Any] = [None] * slot_count tensor_list_lengths = getattr(ctx, "_te_fwd_tensor_list_lengths", {}) for pos, length in tensor_list_lengths.items(): - if isinstance(out[pos], list): - out[pos] = [None] * length + out[pos] = [None] * length for pos, g in zip(grad_targets, grads): out[pos] = g return tuple(out) @@ -1298,7 +1297,7 @@ def _register_custom_op_impl( bwd_schema_args, bwd_arg_names = _build_schema(bwd_buckets) num_grad_inputs = len(input_tensors_for_grad) - fwd_slot_defaults, grad_targets = _resolve_grad_targets( + slot_count, grad_targets = _resolve_grad_targets( fwd_buckets, fwd_arg_type, input_tensors_for_grad ) @@ -1348,7 +1347,7 @@ def _register_custom_op_impl( "fwd_tensor_field_names": fwd_tensor_field_names, "bwd_arg_names": bwd_arg_names, "bwd_buckets": bwd_buckets, - "fwd_slot_defaults": fwd_slot_defaults, + "slot_count": slot_count, "grad_targets": grad_targets, "setup_context_user": setup_context, "backward_obj_type": backward_obj, From d55d38c2fe834ed9a353559020d17d7bbfd53ffa Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:35:32 +0200 Subject: [PATCH 064/108] [PyTorch] Drop redundant fwd_arg_type from _resolve_grad_targets Address review: the buckets are built from the arg dataclass and already carry its field names, so the type was passed only for its name in the error. Remove it and make the error diagnostic instead -- list the actual differentiable field names (from the buckets), which also tells an author whether a missing name is their typo or a grad_slot gap. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 4110643e45..df0e82c718 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -953,11 +953,13 @@ def _split_fwd_fake_result( def _resolve_grad_targets( fwd_buckets: List[_Bucket], - fwd_arg_type: type, input_tensors_for_grad: List[str], ) -> Tuple[int, List[int]]: """Validate ``input_tensors_for_grad`` and resolve the grad-output layout. + ``fwd_buckets`` already encode the arg dataclass's fields (they are built + from it), so the type itself is not needed here. + Returns ``(slot_count, grad_targets)``: the total number of input schema slots and, for each requested input name, the schema-slot index its gradient maps to. @@ -974,8 +976,8 @@ def _resolve_grad_targets( unknown = [n for n in input_tensors_for_grad if n not in name_to_slot] if unknown: raise ValueError( - f"input_tensors_for_grad contains names not in {fwd_arg_type.__name__} " - f"schema: {unknown}" + f"input_tensors_for_grad {unknown} are not differentiable input slots; " + f"differentiable fields are {sorted(name_to_slot)}" ) grad_targets = [name_to_slot[n] for n in input_tensors_for_grad] return slot_offset, grad_targets @@ -1297,9 +1299,7 @@ def _register_custom_op_impl( bwd_schema_args, bwd_arg_names = _build_schema(bwd_buckets) num_grad_inputs = len(input_tensors_for_grad) - slot_count, grad_targets = _resolve_grad_targets( - fwd_buckets, fwd_arg_type, input_tensors_for_grad - ) + slot_count, grad_targets = _resolve_grad_targets(fwd_buckets, input_tensors_for_grad) fwd_schema = f"{fwd_schema_args} -> Tensor[]" bwd_schema = f"{bwd_schema_args} -> Tensor[]" From c4c831389ffe701489d752933a6d64b081b7b2a6 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:44:25 +0200 Subject: [PATCH 065/108] [PyTorch] Split input_tensors_for_grad validation into existence + differentiability Address review: check field existence early at the API boundary (_register_custom_op_impl, using dataclasses.fields) with a clear 'names not in ' error, and keep the differentiability test in _resolve_grad_targets rephrased per the suggestion ('contains non-differentiable fields'). The two conditions now report distinctly. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index df0e82c718..fa21a5ece8 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -973,11 +973,10 @@ def _resolve_grad_targets( name_to_slot[bucket.name] = slot_offset + grad_slot slot_offset += len(slots) - unknown = [n for n in input_tensors_for_grad if n not in name_to_slot] - if unknown: + non_differentiable = [n for n in input_tensors_for_grad if n not in name_to_slot] + if non_differentiable: raise ValueError( - f"input_tensors_for_grad {unknown} are not differentiable input slots; " - f"differentiable fields are {sorted(name_to_slot)}" + f"input_tensors_for_grad contains non-differentiable fields: {non_differentiable}" ) grad_targets = [name_to_slot[n] for n in input_tensors_for_grad] return slot_offset, grad_targets @@ -1284,6 +1283,17 @@ def _register_custom_op_impl( bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], ) -> Callable[..., Any]: """Body of :func:`register_custom_op`; see it for semantics.""" + # Existence check at the API boundary: every ``input_tensors_for_grad`` name + # must be an actual field of ``fwd_arg_type`` (differentiability -- whether + # that field can carry a gradient -- is checked later in + # :func:`_resolve_grad_targets`). + fwd_field_names = {f.name for f in dataclasses.fields(fwd_arg_type)} + missing = [n for n in input_tensors_for_grad if n not in fwd_field_names] + if missing: + raise ValueError( + f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}" + ) + outer_fwd_name = op_name outer_bwd_name = f"{op_name}_backward" inner_fwd_name = f"{op_name}_base" From 197325d9ea39e33636ba9d7b62711e3a9f6ab979 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:47:36 +0200 Subject: [PATCH 066/108] [PyTorch] Rename _pack/_unpack -> _args_to_slots/_args_from_slots Address review: _unpack constructs an args dataclass, so 'unpack' read awkwardly. Name the pair after the aggregate of the bucket to_slots / from_slots methods they drive. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index fa21a5ece8..d286aa4ea1 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -770,10 +770,10 @@ def _build_schema(buckets: List[_Bucket]) -> Tuple[str, List[str]]: return schema_str, names -def _pack(obj: Any, buckets: List[_Bucket]) -> Dict[str, Any]: +def _args_to_slots(obj: Any, buckets: List[_Bucket]) -> Dict[str, Any]: """Build the op's flat ``{slot_name: value}`` argument dict from an args dataclass ``obj`` (e.g. ``LinearFwdArgs``), by collecting every bucket's - packed slot(s). Inverse of :func:`_unpack`. + packed slot(s). Inverse of :func:`_args_from_slots`. """ out: Dict[str, Any] = {} for bucket in buckets: @@ -781,10 +781,10 @@ def _pack(obj: Any, buckets: List[_Bucket]) -> Dict[str, Any]: return out -def _unpack(cls: type, args: Dict[str, Any], buckets: List[_Bucket]) -> Any: +def _args_from_slots(cls: type, args: Dict[str, Any], buckets: List[_Bucket]) -> Any: """Rebuild a fresh args dataclass ``cls`` (e.g. ``LinearFwdArgs``) from the op's flat slot ``args`` dict, by letting every bucket restore its field(s). - Inverse of :func:`_pack`. + Inverse of :func:`_args_to_slots`. """ kwargs: Dict[str, Any] = {} for bucket in buckets: @@ -1004,12 +1004,12 @@ def _register_kernel( def _impl(*flat: Any) -> List[torch.Tensor]: kwargs = dict(zip(arg_names, flat)) - obj = _unpack(arg_type, kwargs, buckets) + obj = _args_from_slots(arg_type, kwargs, buckets) return format_result(impl(obj)) def _fake(*flat: Any) -> List[torch.Tensor]: kwargs = dict(zip(arg_names, flat)) - obj = _unpack(arg_type, kwargs, buckets) + obj = _args_from_slots(arg_type, kwargs, buckets) proto_obj = _proto_view(obj, tensor_field_names) return format_result(fake_impl(proto_obj)) @@ -1048,7 +1048,7 @@ def _setup_context(ctx, inputs, output): i: len(value) for i, value in enumerate(inputs) if isinstance(value, list) } kwargs = dict(zip(fwd_arg_names, inputs)) - fwd_obj = _unpack(fwd_arg_type, kwargs, fwd_buckets) + fwd_obj = _args_from_slots(fwd_arg_type, kwargs, fwd_buckets) proto_obj = _proto_view(fwd_obj, fwd_tensor_field_names) user_fakes, saved_fakes, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(proto_obj)) @@ -1090,7 +1090,7 @@ def _autograd_backward(ctx, *grad_outputs): ctx.tensor_objects = None per_output_grads = grad_outputs[0] bwd_obj.grad_output = _decode_none(per_output_grads[0]) - kwargs = _pack(bwd_obj, bwd_buckets) + kwargs = _args_to_slots(bwd_obj, bwd_buckets) bwd_args_flat = [kwargs[name] for name in bwd_arg_names] bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), bwd_op_name) grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] @@ -1404,7 +1404,7 @@ def _bwd_rule(mode, func, types, args, kwargs): def forward_fn(fwd_args): proto_obj = _proto_view(fwd_args, fwd_tensor_field_names) user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(proto_obj)) - kwargs = _pack(fwd_args, fwd_buckets) + kwargs = _args_to_slots(fwd_args, fwd_buckets) flat_in = [kwargs[name] for name in fwd_arg_names] result = outer_fwd_op(*flat_in) From f4e6dbe9298619fcc832ca61027b7c962db136ca Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:51:22 +0200 Subject: [PATCH 067/108] [PyTorch] Merge backward_obj into backward_arg_type Address review: the two were always the same class (and had to be structurally compatible, since bwd_obj is packed with buckets derived from backward_arg_type). Drop the redundant backward_obj param and instantiate backward_arg_type() directly. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 12 +++++------- transformer_engine/pytorch/module/linear.py | 1 - 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index d286aa4ea1..cf2c068e03 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1182,7 +1182,6 @@ def register_custom_op( fwd_impl: Callable[[Any], Any], setup_context: Callable[..., Any], backward_arg_type: type, - backward_obj: type, backward_impl: Callable[[Any], Any], fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], @@ -1234,12 +1233,13 @@ def register_custom_op( non-differentiable input). * ``bwd_fake_impl(bwd_args)`` -- data-free twin of ``backward_impl`` returning :class:`TensorProto` grads. - * ``backward_obj.setup_saved_tensors(ctx)`` -- optional hook on the backward + * ``backward_arg_type.setup_saved_tensors(ctx)`` -- optional hook on the backward container (see above); skipped if absent. ``input_tensors_for_grad`` lists the ``fwd_arg_type`` fields that receive - gradients (this fixes the backward grad order). ``backward_obj`` is the type - instantiated to hold the backward args (usually == ``backward_arg_type``). + gradients (this fixes the backward grad order). ``backward_arg_type`` is both + the schema source and the type instantiated (``backward_arg_type()``) to hold + the backward args, so it must be constructible with no arguments. Registration touches experimental ``torch.library`` / opaque-object APIs that may be missing on older PyTorch. If it fails, this warns once and @@ -1254,7 +1254,6 @@ def register_custom_op( fwd_impl=fwd_impl, setup_context=setup_context, backward_arg_type=backward_arg_type, - backward_obj=backward_obj, backward_impl=backward_impl, fwd_fake_impl=fwd_fake_impl, bwd_fake_impl=bwd_fake_impl, @@ -1277,7 +1276,6 @@ def _register_custom_op_impl( fwd_impl: Callable[[Any], Any], setup_context: Callable[..., Any], backward_arg_type: type, - backward_obj: type, backward_impl: Callable[[Any], Any], fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], @@ -1360,7 +1358,7 @@ def _register_custom_op_impl( "slot_count": slot_count, "grad_targets": grad_targets, "setup_context_user": setup_context, - "backward_obj_type": backward_obj, + "backward_obj_type": backward_arg_type, "fwd_fake_impl": fwd_fake_impl, } _register_autograd_for_op( diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 7207bc586f..3b5261c980 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1685,7 +1685,6 @@ def _linear_backward_impl_fake( fwd_fake_impl=_linear_forward_impl_fake, setup_context=_linear_setup_ctx, backward_arg_type=LinearBwdArgs, - backward_obj=LinearBwdArgs, backward_impl=_linear_backward, bwd_fake_impl=_linear_backward_impl_fake, ) From 73348113dcb5cb1a5cda169039658cb5b799a6c9 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:53:16 +0200 Subject: [PATCH 068/108] [PyTorch] Check len(user_outputs) not len(user_fakes) (review nit) Same length (user_outputs is built one-per user_fakes), but keying the condition off the value being indexed reads clearer. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index cf2c068e03..d3d70d9e72 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1072,7 +1072,7 @@ def _setup_context(ctx, inputs, output): tensors_to_save_from_setup = setup_context_user( bwd_obj, fwd_obj, - user_outputs[0] if len(user_fakes) == 1 else tuple(user_outputs), + user_outputs[0] if len(user_outputs) == 1 else tuple(user_outputs), ctx_attrs, tuple(saved_list), ) From c18b8ac0d437a694260ef065797436dc8e93a076 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:58:27 +0200 Subject: [PATCH 069/108] [PyTorch] Pass resolved bwd_op object to _register_autograd_for_op Address review: the backward op was passed by name and re-resolved via getattr on every backward, while the same op was also resolved separately for the passthrough set. Resolve the inner/outer bwd ops once (before wiring autograd) and pass the object; the closure uses it directly. Removes the duplicate resolution. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index d3d70d9e72..c4886c340a 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1023,7 +1023,7 @@ def _fake(*flat: Any) -> List[torch.Tensor]: def _register_autograd_for_op( *, fwd_op: Any, - bwd_op_name: str, + bwd_op: Any, fwd_arg_type: type, fwd_arg_names: List[str], fwd_buckets: List[_Bucket], @@ -1036,7 +1036,7 @@ def _register_autograd_for_op( backward_obj_type: type, fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], ) -> None: - """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op_name``. + """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op``. ``setup_context`` re-runs the proto fwd fake impl to recover output / saved templates, reassembles each flat output chunk, and hands the saved tuple + @@ -1092,7 +1092,6 @@ def _autograd_backward(ctx, *grad_outputs): bwd_obj.grad_output = _decode_none(per_output_grads[0]) kwargs = _args_to_slots(bwd_obj, bwd_buckets) bwd_args_flat = [kwargs[name] for name in bwd_arg_names] - bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), bwd_op_name) grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] # One grad per input schema slot: default None, but a ``Tensor[]`` slot # (always recorded in ``_te_fwd_tensor_list_lengths``) needs a @@ -1361,18 +1360,14 @@ def _register_custom_op_impl( "backward_obj_type": backward_arg_type, "fwd_fake_impl": fwd_fake_impl, } - _register_autograd_for_op( - fwd_op=inner_fwd_def, bwd_op_name=inner_bwd_name, **autograd_common - ) - _register_autograd_for_op( - fwd_op=outer_fwd_def, bwd_op_name=outer_bwd_name, **autograd_common - ) - inner_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_fwd_name) inner_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_bwd_name) outer_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_fwd_name) outer_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_bwd_name) + _register_autograd_for_op(fwd_op=inner_fwd_def, bwd_op=inner_bwd_op, **autograd_common) + _register_autograd_for_op(fwd_op=outer_fwd_def, bwd_op=outer_bwd_op, **autograd_common) + fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_buckets) bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_buckets) From a550a6b93055aaf4f8bcfdcd77a24a3f2ec54597 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 14:59:45 +0200 Subject: [PATCH 070/108] [PyTorch] Drop redundant 'val is None or' in subclass flatten check (nit) isinstance(None, subclass) is already False. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index c4886c340a..e5a9f61ad0 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1126,7 +1126,7 @@ def _flatten_subclass_into_slots( """ for offset in slot_offsets: val = new_args[offset] - if val is None or not isinstance(val, subclass): + if not isinstance(val, subclass): continue meta, tensors = _storage_flatten( val, {_TensorOrQuantizedBucket.KIND_KEY: _TensorOrQuantizedKind.STORAGE} From 844832bea7acc0c1e98fe19dfddd4b842943c605 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 15:03:36 +0200 Subject: [PATCH 071/108] [PyTorch] Drop redundant list() copy of storage inner tensors (nit) _storage_flatten already returns a fresh, unaliased list, so the extra list() copy at both call sites was unnecessary. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index e5a9f61ad0..50121885de 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -490,7 +490,7 @@ def to_slots(self, owner: Any) -> Dict[str, Any]: meta, tensors = _storage_flatten(value, {self.KIND_KEY: _TensorOrQuantizedKind.STORAGE}) return { self.slot_name(): None, - self.slot_tensors(): list(tensors), + self.slot_tensors(): tensors, self.slot_meta(): meta, } raise TypeError( @@ -1132,7 +1132,7 @@ def _flatten_subclass_into_slots( val, {_TensorOrQuantizedBucket.KIND_KEY: _TensorOrQuantizedKind.STORAGE} ) new_args[offset] = None - new_args[offset + 1] = list(tensors) + new_args[offset + 1] = tensors new_args[offset + 2] = meta From c752bbbb4fd7501328ebd2dd673cbc35dc480255 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 15:09:51 +0200 Subject: [PATCH 072/108] [PyTorch] Rename two-tier op inner/outer -> base/wrapper Address review: 'outer'/'inner' didn't match the actual op names -- the inner op is registered as _base. Rename to base/wrapper so the identifiers line up with the op names (base op = _base with the real schema/autograd; wrapper op = that forwards to it, flattening subclass inputs). Also disambiguates 'inner', now reserved for the tensor-flatten sense (inner_names / inner_tensors); the flatten-geometry names (inner_names, outer_shape, outer_size) are left unchanged. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 84 +++++++++---------- 1 file changed, 42 insertions(+), 42 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 50121885de..82c197d6e4 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -480,7 +480,7 @@ def to_slots(self, owner: Any) -> Dict[str, Any]: if isinstance(value, torch.Tensor): # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the # ``Tensor?`` slot; subclass flattening (if any) is done by the - # outer op's ``register_torch_dispatch`` rule. + # wrapper op's ``register_torch_dispatch`` rule. return { self.slot_name(): value, self.slot_tensors(): [], @@ -1136,32 +1136,32 @@ def _flatten_subclass_into_slots( new_args[offset + 2] = meta -def _register_outer_forwarder( +def _register_wrapper_op( *, - outer_op_name: str, + wrapper_op_name: str, schema_str: str, - inner_op_name: str, + base_op_name: str, buckets: Optional[List[_Bucket]] = None, subclass_list: Optional[List[type]] = None, ) -> Any: - """Define the outer op via ``torch.library.custom_op``: forward to the inner + """Define the wrapper op via ``torch.library.custom_op``: forward to the base op, optionally flattening registered subclass inputs in place first. Returns the ``CustomOpDef``. """ - inner_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_op_name) + base_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_op_name) input_flatten_enabled = bool(subclass_list) and buckets is not None slot_offsets = _collect_tensor_or_quantized_slot_offsets(buckets) if input_flatten_enabled else [] def _forward(*flat: Any) -> List[torch.Tensor]: if not input_flatten_enabled: - return inner_op(*flat) + return base_op(*flat) new_args = list(flat) for sub in subclass_list: _flatten_subclass_into_slots(new_args, slot_offsets, sub) - return inner_op(*new_args) + return base_op(*new_args) op = torch.library.custom_op( - f"{_TE_OP_NAMESPACE}::{outer_op_name}", _forward, mutates_args=(), schema=schema_str + f"{_TE_OP_NAMESPACE}::{wrapper_op_name}", _forward, mutates_args=(), schema=schema_str ) op.register_fake(_forward) return op @@ -1187,15 +1187,15 @@ def register_custom_op( ) -> Optional[Callable[..., Any]]: """Register a TE module's forward + backward as torch custom ops. - Always two-tier: an inner ``_base`` op carries the real schema / - autograd, and an outer ```` op forwards to it, flattening any + Always two-tier: an base ``_base`` op carries the real schema / + autograd, and an wrapper ```` op forwards to it, flattening any quantized-tensor wrapper inputs first via ``register_torch_dispatch`` (an - empty subclass list simply makes the outer op a pass-through, so a pure + empty subclass list simply makes the wrapper op a pass-through, so a pure plain-tensor / bf16 call goes straight through). Returns ``forward_fn(fwd_arg_type_instance)`` -- a drop-in for ``Function.apply`` under ``torch.compiler.is_compiling()`` that dispatches - through the outer op and returns the user-facing outputs. + through the wrapper op and returns the user-facing outputs. Arg containers. ``fwd_arg_type`` and ``backward_arg_type`` are ``@dataclass``es whose *field annotations* define the op schema: each field maps to one or more @@ -1291,10 +1291,10 @@ def _register_custom_op_impl( f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}" ) - outer_fwd_name = op_name - outer_bwd_name = f"{op_name}_backward" - inner_fwd_name = f"{op_name}_base" - inner_bwd_name = f"{outer_bwd_name}_base" + wrapper_fwd_name = op_name + wrapper_bwd_name = f"{op_name}_backward" + base_fwd_name = f"{op_name}_base" + base_bwd_name = f"{wrapper_bwd_name}_base" subclass_list = _all_quantized_tensor_subclasses() fwd_buckets = _get_buckets(fwd_arg_type) @@ -1311,10 +1311,10 @@ def _register_custom_op_impl( fwd_schema = f"{fwd_schema_args} -> Tensor[]" bwd_schema = f"{bwd_schema_args} -> Tensor[]" - inner_bwd_qualname = f"{_TE_OP_NAMESPACE}::{inner_bwd_name}" + base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" - inner_fwd_def = _register_kernel( - op_name=inner_fwd_name, + base_fwd_def = _register_kernel( + op_name=base_fwd_name, schema_str=fwd_schema, arg_type=fwd_arg_type, arg_names=fwd_arg_names, @@ -1325,7 +1325,7 @@ def _register_custom_op_impl( format_result=_format_fwd_result, ) _register_kernel( - op_name=inner_bwd_name, + op_name=base_bwd_name, schema_str=bwd_schema, arg_type=backward_arg_type, arg_names=bwd_arg_names, @@ -1333,18 +1333,18 @@ def _register_custom_op_impl( tensor_field_names=bwd_tensor_field_names, impl=backward_impl, fake_impl=bwd_fake_impl, - format_result=lambda g: _format_bwd_result(g, num_grad_inputs, inner_bwd_qualname), + format_result=lambda g: _format_bwd_result(g, num_grad_inputs, base_bwd_qualname), ) - outer_fwd_def = _register_outer_forwarder( - outer_op_name=outer_fwd_name, + wrapper_fwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_fwd_name, schema_str=fwd_schema, - inner_op_name=inner_fwd_name, + base_op_name=base_fwd_name, buckets=fwd_buckets, subclass_list=list(subclass_list), ) - outer_bwd_def = _register_outer_forwarder( - outer_op_name=outer_bwd_name, schema_str=bwd_schema, inner_op_name=inner_bwd_name + wrapper_bwd_def = _register_wrapper_op( + wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op_name=base_bwd_name ) autograd_common = { @@ -1360,13 +1360,13 @@ def _register_custom_op_impl( "backward_obj_type": backward_arg_type, "fwd_fake_impl": fwd_fake_impl, } - inner_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_fwd_name) - inner_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_bwd_name) - outer_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_fwd_name) - outer_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_bwd_name) + base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) + base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) + wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) + wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) - _register_autograd_for_op(fwd_op=inner_fwd_def, bwd_op=inner_bwd_op, **autograd_common) - _register_autograd_for_op(fwd_op=outer_fwd_def, bwd_op=outer_bwd_op, **autograd_common) + _register_autograd_for_op(fwd_op=base_fwd_def, bwd_op=base_bwd_op, **autograd_common) + _register_autograd_for_op(fwd_op=wrapper_fwd_def, bwd_op=wrapper_bwd_op, **autograd_common) fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_buckets) bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_buckets) @@ -1376,30 +1376,30 @@ def _fwd_rule(mode, func, types, args, kwargs): new_args = list(args) for sub in subclass_list: _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub) - return inner_fwd_op(*new_args) + return base_fwd_op(*new_args) def _bwd_rule(mode, func, types, args, kwargs): del mode, func, types, kwargs new_args = list(args) for sub in subclass_list: _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub) - return inner_bwd_op(*new_args) + return base_bwd_op(*new_args) for sub in subclass_list: - outer_fwd_def.register_torch_dispatch(sub, _fwd_rule) - outer_bwd_def.register_torch_dispatch(sub, _bwd_rule) + wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) + wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule) - _quantized_tensor_passthrough_ops.add(outer_fwd_op.default) - _quantized_tensor_passthrough_ops.add(outer_bwd_op.default) - _quantized_tensor_passthrough_ops.add(inner_fwd_op.default) - _quantized_tensor_passthrough_ops.add(inner_bwd_op.default) + _quantized_tensor_passthrough_ops.add(wrapper_fwd_op.default) + _quantized_tensor_passthrough_ops.add(wrapper_bwd_op.default) + _quantized_tensor_passthrough_ops.add(base_fwd_op.default) + _quantized_tensor_passthrough_ops.add(base_bwd_op.default) def forward_fn(fwd_args): proto_obj = _proto_view(fwd_args, fwd_tensor_field_names) user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(proto_obj)) kwargs = _args_to_slots(fwd_args, fwd_buckets) flat_in = [kwargs[name] for name in fwd_arg_names] - result = outer_fwd_op(*flat_in) + result = wrapper_fwd_op(*flat_in) cursor = 0 outputs: List[Any] = [] From 1fe18b25dc9e6807ef91d40882db1b40e84800b4 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 15:12:29 +0200 Subject: [PATCH 073/108] [PyTorch] Pass resolved base_op object to _register_wrapper_op Same as the bwd_op change: the base op was passed by name and resolved via getattr inside _register_wrapper_op, while it was also resolved again below for autograd/passthrough. Resolve base_fwd_op/base_bwd_op once and pass the object; drop the duplicate resolution. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 82c197d6e4..d155fc3d12 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1140,7 +1140,7 @@ def _register_wrapper_op( *, wrapper_op_name: str, schema_str: str, - base_op_name: str, + base_op: Any, buckets: Optional[List[_Bucket]] = None, subclass_list: Optional[List[type]] = None, ) -> Any: @@ -1148,7 +1148,6 @@ def _register_wrapper_op( op, optionally flattening registered subclass inputs in place first. Returns the ``CustomOpDef``. """ - base_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_op_name) input_flatten_enabled = bool(subclass_list) and buckets is not None slot_offsets = _collect_tensor_or_quantized_slot_offsets(buckets) if input_flatten_enabled else [] @@ -1336,15 +1335,18 @@ def _register_custom_op_impl( format_result=lambda g: _format_bwd_result(g, num_grad_inputs, base_bwd_qualname), ) + base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) + base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) + wrapper_fwd_def = _register_wrapper_op( wrapper_op_name=wrapper_fwd_name, schema_str=fwd_schema, - base_op_name=base_fwd_name, + base_op=base_fwd_op, buckets=fwd_buckets, subclass_list=list(subclass_list), ) wrapper_bwd_def = _register_wrapper_op( - wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op_name=base_bwd_name + wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op ) autograd_common = { @@ -1360,8 +1362,6 @@ def _register_custom_op_impl( "backward_obj_type": backward_arg_type, "fwd_fake_impl": fwd_fake_impl, } - base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) - base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) From dd99206d797fc310b9e886d0ce0fb023a4964324 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 15 Jul 2026 16:49:42 +0200 Subject: [PATCH 074/108] [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 33f21109280c4de0edcf510382069b7f1409c482 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 16 Jul 2026 12:44:27 +0200 Subject: [PATCH 075/108] [PyTorch] Rename field _Bucket -> _Adapter in custom-op framework Address review: 'bucket' evoked a collector, but the class is a per-field translator between a dataclass field and the op's schema slots (declare slots + pack/unpack). Rename _Bucket and subclasses to _Adapter (+ _FIELD_ADAPTERS, _get_adapters). Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 196 +++++++++--------- 1 file changed, 98 insertions(+), 98 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 1abeeda9ea..26060069dc 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -327,7 +327,7 @@ def _storage_unflatten(meta: Any, tensors: List[torch.Tensor]) -> Any: # --------------------------------------------------------------------------- # -# Field buckets: dataclass field <-> flat torch.library slot(s) +# Field adapters: dataclass field <-> flat torch.library slot(s) # --------------------------------------------------------------------------- # @@ -352,24 +352,24 @@ def _strip_optional(annot: Any) -> Tuple[Any, bool]: return annot, False -class _Bucket: - """Maps one (or, for the aggregating bucket, several) dataclass field(s) +class _Adapter: + """Maps one (or, for the aggregating adapter, several) dataclass field(s) to/from a contiguous run of custom-op schema *slots*. A custom op only takes flat, simply-typed arguments, but a TE op takes a - single ``@dataclass`` of mixed fields. Each bucket knows how to translate + single ``@dataclass`` of mixed fields. Each adapter knows how to translate its kind of field both ways. ``try_build`` and ``schema_slots`` run once at registration (to build the op's schema); ``to_slots`` and ``from_slots`` run on each call and must agree on the slot layout that ``schema_slots`` declares. """ @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_Bucket"]: - """Decide whether this bucket type handles the field ``name`` given its - type annotation ``annot``; return a configured bucket if so, else + def try_build(cls, name: str, annot: Any) -> Optional["_Adapter"]: + """Decide whether this adapter type handles the field ``name`` given its + type annotation ``annot``; return a configured adapter if so, else ``None`` so the next candidate is tried. - Called once per field at registration, in :data:`_FIELD_BUCKETS` + Called once per field at registration, in :data:`_FIELD_ADAPTERS` priority order. """ raise NotImplementedError @@ -378,7 +378,7 @@ def schema_slots(self) -> List[Tuple[str, str]]: """Declare the schema slots this field occupies, each as a ``(slot_name, schema_type)`` pair (e.g. ``("bias", "Tensor?")``). - Concatenated across all buckets to form the op's schema string. + Concatenated across all adapters to form the op's schema string. """ raise NotImplementedError @@ -401,11 +401,11 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: raise NotImplementedError def grad_slot(self) -> Optional[int]: - """Index (within this bucket's :meth:`schema_slots`) of the slot that + """Index (within this adapter's :meth:`schema_slots`) of the slot that carries a gradient, or ``None`` if the field is not differentiable. Used to map ``input_tensors_for_grad`` names onto backward grad-output - positions. Non-tensor buckets (quantizers, metadata) return ``None``. + positions. Non-tensor adapters (quantizers, metadata) return ``None``. """ return None @@ -418,7 +418,7 @@ class _TensorOrQuantizedKind(Enum): STORAGE = "storage" -class _TensorOrQuantizedBucket(_Bucket): +class _TensorOrQuantizedAdapter(_Adapter): """``Tensor | QuantizedTensorStorage | None`` (also subclass tensor) field. Three slots regardless of value: ```` (``Tensor?`` -- plain / subclass @@ -466,7 +466,7 @@ def _is_tensor_storage_union(cls, annot: Any) -> bool: return members == cls._MEMBERS @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_TensorOrQuantizedBucket"]: + def try_build(cls, name: str, annot: Any) -> Optional["_TensorOrQuantizedAdapter"]: if cls._is_tensor_storage_union(annot): return cls(name) return None @@ -516,7 +516,7 @@ def grad_slot(self) -> Optional[int]: return 0 -class _TensorBucket(_Bucket): +class _TensorAdapter(_Adapter): """``Tensor`` / ``Optional[Tensor]`` -> single ``Tensor`` / ``Tensor?`` slot.""" def __init__(self, name: str, is_optional: bool) -> None: @@ -524,7 +524,7 @@ def __init__(self, name: str, is_optional: bool) -> None: self.type_str = "Tensor?" if is_optional else "Tensor" @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_TensorBucket"]: + def try_build(cls, name: str, annot: Any) -> Optional["_TensorAdapter"]: stripped, is_optional = _strip_optional(annot) if stripped is torch.Tensor: return cls(name, is_optional) @@ -543,7 +543,7 @@ def grad_slot(self) -> Optional[int]: return 0 -class _QuantizerBucket(_Bucket): +class _QuantizerAdapter(_Adapter): """``Quantizer`` / ``Optional[Quantizer]`` -> one own ``OpaqueValueBundle`` slot. Each quantizer gets its own dedicated slot. The field is annotated with the @@ -561,7 +561,7 @@ def slot(self) -> str: return self.name + "__q" @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_QuantizerBucket"]: + def try_build(cls, name: str, annot: Any) -> Optional["_QuantizerAdapter"]: stripped, _ = _strip_optional(annot) if isinstance(stripped, type) and issubclass(stripped, Quantizer): return cls(name) @@ -577,7 +577,7 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = args[self.slot()][self.KEY] -class _ReferenceOpaqueBucket(_Bucket): +class _ReferenceOpaqueAdapter(_Adapter): """``ProcessGroup`` (or any reference-opaque type) -> one own opaque slot. A reference-opaque object is live, stateful black-box data (e.g. a @@ -597,7 +597,7 @@ def __init__(self, name: str, type_name: str, is_optional: bool) -> None: self.type_str = f"{type_name}?" if is_optional else type_name @classmethod - def try_build(cls, name: str, annot: Any) -> Optional["_ReferenceOpaqueBucket"]: + def try_build(cls, name: str, annot: Any) -> Optional["_ReferenceOpaqueAdapter"]: if _is_opaque_reference_type is None: return None stripped, is_optional = _strip_optional(annot) @@ -617,11 +617,11 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = args[self.name] -class _SimpleBundleBucket(_Bucket): +class _SimpleBundleAdapter(_Adapter): """Aggregates every simple-typed field into a single OpaqueValueBundle. - Unlike the per-field buckets, exactly one of these exists per op: it owns the - single shared ``_simple_meta`` slot, and ``_get_buckets`` builds it once from + Unlike the per-field adapters, exactly one of these exists per op: it owns the + single shared ``_simple_meta`` slot, and ``_get_adapters`` builds it once from all simple-typed field names collected across the dataclass. """ @@ -663,8 +663,8 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[n] = meta[n] -class _UnsupportedBucket(_Bucket): - """Fallback for fields whose type no other bucket can encode. +class _UnsupportedAdapter(_Adapter): + """Fallback for fields whose type no other adapter can encode. Such a field cannot cross the op boundary, so it emits no slot and is tolerated only when its runtime value carries nothing: ``to_slots`` accepts @@ -673,7 +673,7 @@ class _UnsupportedBucket(_Bucket): ``None``. A non-trivial value means the config is genuinely unsupported under torch.compile, and ``to_slots`` raises. - The check must run at call time (not in ``_get_buckets``): the annotation + The check must run at call time (not in ``_get_adapters``): the annotation alone -- e.g. ``Optional[Any]`` -- is valid when the value is ``None``, so only the runtime value can decide. """ @@ -699,7 +699,7 @@ def to_slots(self, owner: Any) -> Dict[str, Any]: raise TypeError( f"{self.owner_cls_name} field {self.name!r} has a type not " "supported by torch.compile (not Tensor, simple, or Quantizer) " - "and carries a non-trivial value; add a matching bucket in " + "and carries a non-trivial value; add a matching adapter in " "dynamo.py to handle it." ) return {} @@ -708,16 +708,16 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: kwargs[self.name] = None -# Buckets, in priority order, owning ``try_build`` for a single field. -# These buckets are mutually exclusive on annotations (a plain ``torch.Tensor`` -# matches only ``_TensorBucket``; the ``TensorOrQuantized`` union only -# ``_TensorOrQuantizedBucket``; etc.), so the order is just iteration, not a +# Adapters, in priority order, owning ``try_build`` for a single field. +# These adapters are mutually exclusive on annotations (a plain ``torch.Tensor`` +# matches only ``_TensorAdapter``; the ``TensorOrQuantized`` union only +# ``_TensorOrQuantizedAdapter``; etc.), so the order is just iteration, not a # priority ranking -- no annotation can be claimed by more than one. -_FIELD_BUCKETS: Tuple[type, ...] = ( - _TensorOrQuantizedBucket, - _TensorBucket, - _ReferenceOpaqueBucket, - _QuantizerBucket, +_FIELD_ADAPTERS: Tuple[type, ...] = ( + _TensorOrQuantizedAdapter, + _TensorAdapter, + _ReferenceOpaqueAdapter, + _QuantizerAdapter, ) @@ -732,65 +732,65 @@ def _resolved_field_annotations(cls: type) -> List[Tuple[str, Any]]: return [(f.name, hints.get(f.name, f.type)) for f in dataclasses.fields(cls)] -def _get_buckets(cls: type) -> List[_Bucket]: - """Build the bucket list for a dataclass from its field annotations.""" +def _get_adapters(cls: type) -> List[_Adapter]: + """Build the adapter list for a dataclass from its field annotations.""" if _OPAQUE_VALUE_BUNDLE_TYPE_NAME is None: raise RuntimeError( f"{cls.__name__} cannot be turned into a TE custom op: OpaqueValueBundle " "is not registered as a torch._library value-opaque type (PyTorch build " "without opaque-object support)." ) - buckets: List[_Bucket] = [] + adapters: List[_Adapter] = [] simple_names: List[str] = [] for name, annot in _resolved_field_annotations(cls): - built: Optional[_Bucket] = None - for bucket_cls in _FIELD_BUCKETS: - built = bucket_cls.try_build(name, annot) + built: Optional[_Adapter] = None + for adapter_cls in _FIELD_ADAPTERS: + built = adapter_cls.try_build(name, annot) if built is not None: break if built is not None: - buckets.append(built) - elif _SimpleBundleBucket.matches_field(annot): + adapters.append(built) + elif _SimpleBundleAdapter.matches_field(annot): simple_names.append(name) else: - buckets.append(_UnsupportedBucket(name, cls.__name__)) + adapters.append(_UnsupportedAdapter(name, cls.__name__)) if simple_names: - buckets.append(_SimpleBundleBucket(simple_names)) - return buckets + adapters.append(_SimpleBundleAdapter(simple_names)) + return adapters -def _tensor_field_names(buckets: List[_Bucket]) -> List[str]: +def _tensor_field_names(adapters: List[_Adapter]) -> List[str]: """Names of fields carrying tensors (for building the proto view).""" - return [b.name for b in buckets if isinstance(b, (_TensorBucket, _TensorOrQuantizedBucket))] + return [b.name for b in adapters if isinstance(b, (_TensorAdapter, _TensorOrQuantizedAdapter))] -def _build_schema(buckets: List[_Bucket]) -> Tuple[str, List[str]]: - """Return ``(schema_arg_str, slot_names)`` for a bucket list.""" - spec = [slot for b in buckets for slot in b.schema_slots()] +def _build_schema(adapters: List[_Adapter]) -> Tuple[str, List[str]]: + """Return ``(schema_arg_str, slot_names)`` for a adapter list.""" + spec = [slot for b in adapters for slot in b.schema_slots()] names = [name for name, _ in spec] schema_str = "(" + ", ".join(f"{type_str} {name}" for name, type_str in spec) + ")" return schema_str, names -def _args_to_slots(obj: Any, buckets: List[_Bucket]) -> Dict[str, Any]: +def _args_to_slots(obj: Any, adapters: List[_Adapter]) -> Dict[str, Any]: """Build the op's flat ``{slot_name: value}`` argument dict from an args - dataclass ``obj`` (e.g. ``LinearFwdArgs``), by collecting every bucket's + dataclass ``obj`` (e.g. ``LinearFwdArgs``), by collecting every adapter's packed slot(s). Inverse of :func:`_args_from_slots`. """ out: Dict[str, Any] = {} - for bucket in buckets: - out.update(bucket.to_slots(obj)) + for adapter in adapters: + out.update(adapter.to_slots(obj)) return out -def _args_from_slots(cls: type, args: Dict[str, Any], buckets: List[_Bucket]) -> Any: +def _args_from_slots(cls: type, args: Dict[str, Any], adapters: List[_Adapter]) -> Any: """Rebuild a fresh args dataclass ``cls`` (e.g. ``LinearFwdArgs``) from the - op's flat slot ``args`` dict, by letting every bucket restore its field(s). + op's flat slot ``args`` dict, by letting every adapter restore its field(s). Inverse of :func:`_args_to_slots`. """ kwargs: Dict[str, Any] = {} - for bucket in buckets: - bucket.from_slots(args, kwargs) + for adapter in adapters: + adapter.from_slots(args, kwargs) obj = cls.__new__(cls) for k, v in kwargs.items(): object.__setattr__(obj, k, v) @@ -954,12 +954,12 @@ def _split_fwd_fake_result( def _resolve_grad_targets( - fwd_buckets: List[_Bucket], + fwd_adapters: List[_Adapter], input_tensors_for_grad: List[str], ) -> Tuple[int, List[int]]: """Validate ``input_tensors_for_grad`` and resolve the grad-output layout. - ``fwd_buckets`` already encode the arg dataclass's fields (they are built + ``fwd_adapters`` already encode the arg dataclass's fields (they are built from it), so the type itself is not needed here. Returns ``(slot_count, grad_targets)``: the total number of input schema @@ -968,11 +968,11 @@ def _resolve_grad_targets( """ name_to_slot: Dict[str, int] = {} slot_offset = 0 - for bucket in fwd_buckets: - slots = bucket.schema_slots() - grad_slot = bucket.grad_slot() + for adapter in fwd_adapters: + slots = adapter.schema_slots() + grad_slot = adapter.grad_slot() if grad_slot is not None: - name_to_slot[bucket.name] = slot_offset + grad_slot + name_to_slot[adapter.name] = slot_offset + grad_slot slot_offset += len(slots) non_differentiable = [n for n in input_tensors_for_grad if n not in name_to_slot] @@ -990,7 +990,7 @@ def _register_kernel( schema_str: str, arg_type: type, arg_names: List[str], - buckets: List[_Bucket], + adapters: List[_Adapter], tensor_field_names: List[str], impl: Callable[[Any], Any], fake_impl: Callable[[Any], Any], @@ -1006,12 +1006,12 @@ def _register_kernel( def _impl(*flat: Any) -> List[torch.Tensor]: kwargs = dict(zip(arg_names, flat)) - obj = _args_from_slots(arg_type, kwargs, buckets) + obj = _args_from_slots(arg_type, kwargs, adapters) return format_result(impl(obj)) def _fake(*flat: Any) -> List[torch.Tensor]: kwargs = dict(zip(arg_names, flat)) - obj = _args_from_slots(arg_type, kwargs, buckets) + obj = _args_from_slots(arg_type, kwargs, adapters) proto_obj = _proto_view(obj, tensor_field_names) return format_result(fake_impl(proto_obj)) @@ -1028,10 +1028,10 @@ def _register_autograd_for_op( bwd_op: Any, fwd_arg_type: type, fwd_arg_names: List[str], - fwd_buckets: List[_Bucket], + fwd_adapters: List[_Adapter], fwd_tensor_field_names: List[str], bwd_arg_names: List[str], - bwd_buckets: List[_Bucket], + bwd_adapters: List[_Adapter], slot_count: int, grad_targets: List[int], setup_context_user: Callable[..., Any], @@ -1050,7 +1050,7 @@ def _setup_context(ctx, inputs, output): i: len(value) for i, value in enumerate(inputs) if isinstance(value, list) } kwargs = dict(zip(fwd_arg_names, inputs)) - fwd_obj = _args_from_slots(fwd_arg_type, kwargs, fwd_buckets) + fwd_obj = _args_from_slots(fwd_arg_type, kwargs, fwd_adapters) proto_obj = _proto_view(fwd_obj, fwd_tensor_field_names) user_fakes, saved_fakes, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(proto_obj)) @@ -1092,7 +1092,7 @@ def _autograd_backward(ctx, *grad_outputs): ctx.tensor_objects = None per_output_grads = grad_outputs[0] bwd_obj.grad_output = _decode_none(per_output_grads[0]) - kwargs = _args_to_slots(bwd_obj, bwd_buckets) + kwargs = _args_to_slots(bwd_obj, bwd_adapters) bwd_args_flat = [kwargs[name] for name in bwd_arg_names] grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)] # One grad per input schema slot: default None, but a ``Tensor[]`` slot @@ -1109,21 +1109,21 @@ def _autograd_backward(ctx, *grad_outputs): fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context) -def _collect_tensor_or_quantized_slot_offsets(buckets: List[_Bucket]) -> List[int]: - """Start index of each ``_TensorOrQuantizedBucket`` group in the flat args.""" +def _collect_tensor_or_quantized_slot_offsets(adapters: List[_Adapter]) -> List[int]: + """Start index of each ``_TensorOrQuantizedAdapter`` group in the flat args.""" offsets: List[int] = [] pos = 0 - for bucket in buckets: - if isinstance(bucket, _TensorOrQuantizedBucket): + for adapter in adapters: + if isinstance(adapter, _TensorOrQuantizedAdapter): offsets.append(pos) - pos += len(bucket.schema_slots()) + pos += len(adapter.schema_slots()) return offsets def _flatten_subclass_into_slots( new_args: List[Any], slot_offsets: List[int], subclass: type ) -> None: - """Rewrite each tensor-or-quantized-bucket group whose ``Tensor?`` slot holds an + """Rewrite each tensor-or-quantized-adapter group whose ``Tensor?`` slot holds an instance of ``subclass`` into the storage layout (3 slots: name / tensors / meta). """ for offset in slot_offsets: @@ -1131,7 +1131,7 @@ def _flatten_subclass_into_slots( if not isinstance(val, subclass): continue meta, tensors = _storage_flatten( - val, {_TensorOrQuantizedBucket.KIND_KEY: _TensorOrQuantizedKind.STORAGE} + val, {_TensorOrQuantizedAdapter.KIND_KEY: _TensorOrQuantizedKind.STORAGE} ) new_args[offset] = None new_args[offset + 1] = tensors @@ -1143,15 +1143,15 @@ def _register_wrapper_op( wrapper_op_name: str, schema_str: str, base_op: Any, - buckets: Optional[List[_Bucket]] = None, + adapters: Optional[List[_Adapter]] = None, subclass_list: Optional[List[type]] = None, ) -> Any: """Define the wrapper op via ``torch.library.custom_op``: forward to the base op, optionally flattening registered subclass inputs in place first. Returns the ``CustomOpDef``. """ - input_flatten_enabled = bool(subclass_list) and buckets is not None - slot_offsets = _collect_tensor_or_quantized_slot_offsets(buckets) if input_flatten_enabled else [] + input_flatten_enabled = bool(subclass_list) and adapters is not None + slot_offsets = _collect_tensor_or_quantized_slot_offsets(adapters) if input_flatten_enabled else [] def _forward(*flat: Any) -> List[torch.Tensor]: if not input_flatten_enabled: @@ -1201,7 +1201,7 @@ def register_custom_op( Arg containers. ``fwd_arg_type`` and ``backward_arg_type`` are ``@dataclass``es whose *field annotations* define the op schema: each field maps to one or more flat schema slots (tensor fields cross the boundary as tensors, quantizers ride - as value-opaque objects, simple values are bundled -- see the ``_Bucket`` + as value-opaque objects, simple values are bundled -- see the ``_Adapter`` classes). The caller builds a ``fwd_arg_type`` instance and passes it to the returned ``forward_fn``. @@ -1298,16 +1298,16 @@ def _register_custom_op_impl( base_bwd_name = f"{wrapper_bwd_name}_base" subclass_list = _all_quantized_tensor_subclasses() - fwd_buckets = _get_buckets(fwd_arg_type) - bwd_buckets = _get_buckets(backward_arg_type) - fwd_tensor_field_names = _tensor_field_names(fwd_buckets) - bwd_tensor_field_names = _tensor_field_names(bwd_buckets) + fwd_adapters = _get_adapters(fwd_arg_type) + bwd_adapters = _get_adapters(backward_arg_type) + fwd_tensor_field_names = _tensor_field_names(fwd_adapters) + bwd_tensor_field_names = _tensor_field_names(bwd_adapters) - fwd_schema_args, fwd_arg_names = _build_schema(fwd_buckets) - bwd_schema_args, bwd_arg_names = _build_schema(bwd_buckets) + fwd_schema_args, fwd_arg_names = _build_schema(fwd_adapters) + bwd_schema_args, bwd_arg_names = _build_schema(bwd_adapters) num_grad_inputs = len(input_tensors_for_grad) - slot_count, grad_targets = _resolve_grad_targets(fwd_buckets, input_tensors_for_grad) + slot_count, grad_targets = _resolve_grad_targets(fwd_adapters, input_tensors_for_grad) fwd_schema = f"{fwd_schema_args} -> Tensor[]" bwd_schema = f"{bwd_schema_args} -> Tensor[]" @@ -1319,7 +1319,7 @@ def _register_custom_op_impl( schema_str=fwd_schema, arg_type=fwd_arg_type, arg_names=fwd_arg_names, - buckets=fwd_buckets, + adapters=fwd_adapters, tensor_field_names=fwd_tensor_field_names, impl=fwd_impl, fake_impl=fwd_fake_impl, @@ -1330,7 +1330,7 @@ def _register_custom_op_impl( schema_str=bwd_schema, arg_type=backward_arg_type, arg_names=bwd_arg_names, - buckets=bwd_buckets, + adapters=bwd_adapters, tensor_field_names=bwd_tensor_field_names, impl=backward_impl, fake_impl=bwd_fake_impl, @@ -1344,7 +1344,7 @@ def _register_custom_op_impl( wrapper_op_name=wrapper_fwd_name, schema_str=fwd_schema, base_op=base_fwd_op, - buckets=fwd_buckets, + adapters=fwd_adapters, subclass_list=list(subclass_list), ) wrapper_bwd_def = _register_wrapper_op( @@ -1354,10 +1354,10 @@ def _register_custom_op_impl( autograd_common = { "fwd_arg_type": fwd_arg_type, "fwd_arg_names": fwd_arg_names, - "fwd_buckets": fwd_buckets, + "fwd_adapters": fwd_adapters, "fwd_tensor_field_names": fwd_tensor_field_names, "bwd_arg_names": bwd_arg_names, - "bwd_buckets": bwd_buckets, + "bwd_adapters": bwd_adapters, "slot_count": slot_count, "grad_targets": grad_targets, "setup_context_user": setup_context, @@ -1370,8 +1370,8 @@ def _register_custom_op_impl( _register_autograd_for_op(fwd_op=base_fwd_def, bwd_op=base_bwd_op, **autograd_common) _register_autograd_for_op(fwd_op=wrapper_fwd_def, bwd_op=wrapper_bwd_op, **autograd_common) - fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_buckets) - bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_buckets) + fwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(fwd_adapters) + bwd_slot_offsets = _collect_tensor_or_quantized_slot_offsets(bwd_adapters) def _fwd_rule(mode, func, types, args, kwargs): del mode, func, types, kwargs @@ -1399,7 +1399,7 @@ def _bwd_rule(mode, func, types, args, kwargs): def forward_fn(fwd_args): proto_obj = _proto_view(fwd_args, fwd_tensor_field_names) user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(proto_obj)) - kwargs = _args_to_slots(fwd_args, fwd_buckets) + kwargs = _args_to_slots(fwd_args, fwd_adapters) flat_in = [kwargs[name] for name in fwd_arg_names] result = wrapper_fwd_op(*flat_in) From 8ecd10ceb433e9bde8ad8a58b587e688d5f8bce6 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 16 Jul 2026 13:09:17 +0200 Subject: [PATCH 076/108] [PyTorch] Add architecture overview to custom_op.py module docstring Address review: document the abstractions and how they fit together -- the args dataclass vs a custom op's flat slots, the per-field _Adapter mapping, the fake/TensorProto output-assembly flow (forward_fn + setup_context), and the two-tier base/wrapper op that lets a QuantizedTensor subclass be an op input. Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 80 ++++++++++++++++++- 1 file changed, 79 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 26060069dc..6bcd00c9bc 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -2,7 +2,85 @@ # # See LICENSE for license information. -"""torch.compile custom-op framework for Transformer Engine.""" +"""torch.compile custom-op framework for Transformer Engine. + +Turns a TE module's eager forward/backward into ``torch.library`` custom ops so +``torch.compile(fullgraph=True)`` traces them as single graph nodes -- no graph +break into the eager ``autograd.Function``. ``register_custom_op`` is the entry +point (its docstring documents the per-callable contract); ``module/linear.py`` +is the first user. Internal framework API -- exported from +``transformer_engine.pytorch.dynamo``, not re-exported at the top level. + +A TE op's forward/backward is written as a plain impl over a single *args +dataclass* (``fwd_arg_type`` / ``backward_arg_type``, e.g. ``LinearFwdArgs``): its +fields are a mix of tensors, quantized tensors, quantizers, process groups, +scalars and other Python values. The forward impl returns a tuple (user outputs + +saved-for-backward tensors + ctx metadata); the backward impl returns one gradient +per differentiable input. + +A ``torch.library`` custom op is narrower: it takes a flat list of schema slots +-- tensors / ``Tensor[]`` plus, via torch's opaque-object support, value-opaque +and reference-opaque objects -- and returns a flat ``Tensor[]``. + +Bridging the two takes three parts (below): per-field *adapters* map the args +dataclass onto the op's input slots; *fake impls* on data-free protos give the +output geometry and reassemble the op's flat return; and a *two-tier op* lets a +quantized-tensor subclass be an op input. + +Field <-> slot mapping. This mapping turns each field of the args dataclass into +the op's flat input slots, in a way that suits the field's type. A field's type +annotation selects the one ``_Adapter`` that handles it; that adapter declares the +slot(s) the field needs, packs the field's value into them on the way into the op, +and unpacks it back on the way out. The kinds -- and how each represents its field +as op inputs: + + * ``_TensorAdapter`` -- a plain ``Tensor`` / ``Optional[Tensor]``: one tensor + slot. + * ``_TensorOrQuantizedAdapter`` -- a field that may be a plain tensor, a bare + quantized storage, or ``None``: three slots (the tensor, its flat inner + buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers. + * ``_QuantizerAdapter`` -- a quantizer, baked into the graph as a value-opaque + constant. + * ``_ReferenceOpaqueAdapter`` -- a ProcessGroup, carried as a live opaque graph + input. + * ``_SimpleBundleAdapter`` -- every remaining simple value (scalars, enums, + sizes, nested collections of them), gathered into one ``OpaqueValueBundle`` + slot. + * ``_UnsupportedAdapter`` -- fallback for a field no adapter can encode; allowed + only when its value is trivial (``None`` / all-``None``) at call time. + +What runs where. Each op registers a data-free fake (``register_fake``) so it +traces under ``torch.compile`` without allocating. ``register_custom_op`` returns +``forward_fn`` -- the drop-in for the eager ``autograd.Function.apply``. A forward +call through it: + + * runs the fake ``fwd_fake_impl`` on ``TensorProto`` descriptors (data-free; see + ``tensor_proto.py``) to get the outputs' geometry in pure Python; + * calls the *forward op* -- which runs the real ``fwd_impl`` -- for a flat + ``Tensor[]`` payload; + * rebuilds the structured user outputs from that payload, sliced and reassembled + per the fake's output descriptors (``_proto_reassemble``; + ``_value_to_flat_tensors`` is the pack-side inverse). + +Autograd, registered on the op, drives backward: + + * ``setup_context`` (run when the forward is taped) re-runs ``fwd_fake_impl`` for + the saved-tensor descriptors and a ``ctx_attrs`` dict, reassembles the saved + tensors from the op's flat output, then calls the user ``setup_context`` to + fill the backward args from forward state + ``ctx_attrs`` (e.g. saved-tensor + aliases) and return the tensors to persist; + * on ``backward()`` the backward args container's optional ``setup_saved_tensors`` + hook restores those saved tensors, then the *backward op* runs the real + ``backward_impl`` and returns the flat grads (``bwd_fake_impl`` is its + data-free fake). + +Two-tier op (``base`` + ``wrapper``), so a ``QuantizedTensor`` subclass can be an +op *input*. The ``_base`` op carries the real schema + autograd; a custom op +can't take a tensor-subclass input directly, so the ```` wrapper intercepts +those via ``register_torch_dispatch`` and flattens each into the base op's slots +(``_flatten_subclass_into_slots``) before forwarding. An empty subclass list makes +the wrapper a pass-through (plain / bf16 calls go straight through). +""" from __future__ import annotations import dataclasses From 7aefcdb20c1e5703624a0ca929edc8e9ee80d12a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 16 Jul 2026 13:13:38 +0200 Subject: [PATCH 077/108] [PyTorch] Decouple TensorOrQuantized comment from the custom-op detail Address review: the alias comment now just explains the union as an honest type (fields may hold a bare QuantizedTensorStorage); the annotation-driven op-schema machinery is documented in the custom-op framework, not leaked into module code. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 962ff46b5b..f650d64397 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -87,11 +87,9 @@ __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. +# Fields with this union may hold a *bare* ``QuantizedTensorStorage`` (the +# internal-quantizer optimization; see its docstring), not just a plain / +# subclass tensor -- hence the union rather than a plain ``torch.Tensor``. TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] From 2b0ab1ee183be416dbc1232fc4b4a81ea3003d0f Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 16 Jul 2026 13:25:04 +0200 Subject: [PATCH 078/108] [PyTorch] Drop dead getattr guard for types.UnionType in _is_union Address review: TE requires Python >= 3.10 (build_tools/utils.py), where types.UnionType always exists, so use it directly. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 6bcd00c9bc..2e26a51ee0 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -416,7 +416,7 @@ def _is_union(annot: Any) -> bool: for the latter, so the two syntaxes must be checked separately. """ origin = get_origin(annot) - return origin is Union or origin is getattr(_types, "UnionType", ()) + return origin is Union or origin is _types.UnionType def _strip_optional(annot: Any) -> Tuple[Any, bool]: From f26b05935de5474de7a85d0001a21d3dc228424d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 16 Jul 2026 13:26:01 +0200 Subject: [PATCH 079/108] [PyTorch] Fix _SimpleBundleAdapter doc: at most one per op, not exactly one Address review: _get_adapters only appends it when there are simple-typed fields, so an op with none has zero. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 2e26a51ee0..1a10491f60 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -698,9 +698,10 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: class _SimpleBundleAdapter(_Adapter): """Aggregates every simple-typed field into a single OpaqueValueBundle. - Unlike the per-field adapters, exactly one of these exists per op: it owns the - single shared ``_simple_meta`` slot, and ``_get_adapters`` builds it once from - all simple-typed field names collected across the dataclass. + Unlike the per-field adapters, at most one of these exists per op (none if the + dataclass has no simple-typed fields): it owns the single shared + ``_simple_meta`` slot, and ``_get_adapters`` builds it once from all + simple-typed field names collected across the dataclass. """ SLOT = "_simple_meta" From 2d9de11118459ea7d8a97b4793500f23cc42a6dc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 16 Jul 2026 13:26:58 +0200 Subject: [PATCH 080/108] [PyTorch] Mention reference-opaque types in unsupported-field error Address review nit: ProcessGroup (reference-opaque) is supported too; the error listing 'Tensor, simple, or Quantizer' omitted it. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 1a10491f60..6243eb40c7 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -777,9 +777,9 @@ def to_slots(self, owner: Any) -> Dict[str, Any]: if not self._is_trivial(value): raise TypeError( f"{self.owner_cls_name} field {self.name!r} has a type not " - "supported by torch.compile (not Tensor, simple, or Quantizer) " - "and carries a non-trivial value; add a matching adapter in " - "dynamo.py to handle it." + "supported by torch.compile (not Tensor, simple, Quantizer, or a " + "reference-opaque type such as ProcessGroup) and carries a " + "non-trivial value; add a matching adapter in dynamo.py to handle it." ) return {} From a6cca41e091518605946811e0708af18c2030d1c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 16 Jul 2026 14:22:26 +0200 Subject: [PATCH 081/108] [PyTorch] Guard Float8 __repr__ against materializing fake tensors under tracing Float8Tensor / Float8TensorStorage __repr__ called self._scale_inv.item() inside the f-string. On a fake/meta/functional tensor .item() does not raise; it silently allocates an unbacked SymFloat into the active ShapeEnv, which later crashes inductor with PendingUnbackedSymbolNotFound. AOTAutograd repr's the fake quantized tensor while logging graph metadata, so the scalar- materializing repr leaked an unbacked symbol during compile. Add tensor_can_be_materialized() to detect fake/meta/functional tensors and have both reprs fall back to the metadata-only safe_quantized_repr under tracing, before touching .item(). The previous try/except fallback fired too late -- the symbol was already leaked. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski --- .../pytorch/tensor/_quantization_helpers.py | 37 +++++++++++++++++++ .../pytorch/tensor/float8_tensor.py | 11 +++++- .../tensor/storage/float8_tensor_storage.py | 7 +++- 3 files changed, 53 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py index 1b08039dda..0d29ede2b0 100644 --- a/transformer_engine/pytorch/tensor/_quantization_helpers.py +++ b/transformer_engine/pytorch/tensor/_quantization_helpers.py @@ -85,6 +85,43 @@ def _stride_from_shape(shape: list[int]): return list(reversed(rstride)) +def tensor_can_be_materialized(t) -> bool: + """Whether ``t`` holds concrete data that ``.item()`` / ``.tolist()`` can read + without side effects. + + A ``__repr__`` must never mutate tracing state. On a fake / meta / functional + tensor (torch.compile / export tracing) ``.item()`` does *not* raise -- it + silently allocates an *unbacked* SymInt/SymFloat into the active ShapeEnv, + which later crashes inductor with ``PendingUnbackedSymbolNotFound``. (torch's + AOTAutograd repr's the fake quantized tensor while logging graph metadata, so + a scalar-materializing ``__repr__`` leaks an unbacked symbol during compile.) + So detect those tensors and fall back to a metadata-only repr instead. + """ + if not isinstance(t, torch.Tensor): + return False + if getattr(t, "is_meta", False): + return False + try: + from torch._subclasses.fake_tensor import ( # pylint: disable=import-outside-toplevel + FakeTensor, + ) + + if isinstance(t, FakeTensor): + return False + except Exception: # pylint: disable=broad-except + pass + try: + from torch._subclasses.functional_tensor import ( # pylint: disable=import-outside-toplevel + FunctionalTensor, + ) + + if isinstance(t, FunctionalTensor): + return False + except Exception: # pylint: disable=broad-except + pass + return True + + def safe_quantized_repr(obj, cls_name, extras=None, error=None): """Metadata-only repr fallback for quantized tensors whose data cannot be materialized for any reason. diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index a4da60d7f3..25056f2427 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -19,7 +19,11 @@ from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func from ..quantized_tensor import QuantizedTensor, Quantizer from ..dynamo import register_value_opaque_quantizer -from ._quantization_helpers import _IdentityFunc, safe_quantized_repr +from ._quantization_helpers import ( + _IdentityFunc, + safe_quantized_repr, + tensor_can_be_materialized, +) from ..constants import dist_group_type, DType aten = torch.ops.aten @@ -457,6 +461,11 @@ class Float8Tensor(Float8TensorStorage, QuantizedTensor): amax_reduction_group: Optional[dist_group_type] = None def __repr__(self, *, tensor_contents=None): + # A fake/meta/functional scale_inv cannot be materialized without leaking + # an unbacked symbol into the ShapeEnv (see tensor_can_be_materialized); + # fall back to a metadata-only repr under tracing. + if not tensor_can_be_materialized(self._scale_inv): + return safe_quantized_repr(self, "Float8Tensor") try: return ( "Float8Tensor(" diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index e1078674e7..1899f7046a 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -12,7 +12,7 @@ import transformer_engine_torch as tex from ...quantized_tensor import QuantizedTensorStorage, Quantizer -from .._quantization_helpers import safe_quantized_repr +from .._quantization_helpers import safe_quantized_repr, tensor_can_be_materialized from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch, DType @@ -218,6 +218,11 @@ def view(self, shape: torch.Size): ) def __repr__(self): + # A fake/meta/functional scale_inv cannot be materialized without leaking + # an unbacked symbol into the ShapeEnv (see tensor_can_be_materialized); + # fall back to a metadata-only repr under tracing. + if not tensor_can_be_materialized(self._scale_inv): + return safe_quantized_repr(self, "Float8TensorStorage") try: return ( "Float8TensorStorage(" From d1b6b6eae4ccffaea8248aaf2fe0f9f1e59761e4 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 16 Jul 2026 14:22:26 +0200 Subject: [PATCH 082/108] [PyTorch] Compute quantized subclasses inside _register_wrapper_op _register_wrapper_op took a subclass_list parameter that the sole caller filled with _all_quantized_tensor_subclasses(). The helper is idempotent and cheap (registry filter over an already-imported module), so drop the parameter and call the helper directly inside the op. This keeps "which subclasses to flatten" encapsulated in the helper instead of surfacing it as a caller- provided argument. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 6243eb40c7..b24ab2f11d 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -1223,12 +1223,20 @@ def _register_wrapper_op( schema_str: str, base_op: Any, adapters: Optional[List[_Adapter]] = None, - subclass_list: Optional[List[type]] = None, ) -> Any: """Define the wrapper op via ``torch.library.custom_op``: forward to the base - op, optionally flattening registered subclass inputs in place first. Returns - the ``CustomOpDef``. + op, first flattening any ``QuantizedTensor`` subclass input into the base op's + slots. + + A ``torch.library`` op cannot take a tensor subclass directly, so each such + input is unpacked into its tensor-or-quantized slots + (``_flatten_subclass_into_slots``) before forwarding to the base op -- see the + two-tier op note in the module docstring. The subclasses to flatten are the + live ``QuantizedTensor`` wrappers (e.g. ``Float8Tensor``), obtained from the + registry via ``_all_quantized_tensor_subclasses()``. With no adapters the + wrapper is a plain pass-through. Returns the ``CustomOpDef``. """ + subclass_list = _all_quantized_tensor_subclasses() input_flatten_enabled = bool(subclass_list) and adapters is not None slot_offsets = _collect_tensor_or_quantized_slot_offsets(adapters) if input_flatten_enabled else [] @@ -1424,7 +1432,6 @@ def _register_custom_op_impl( schema_str=fwd_schema, base_op=base_fwd_op, adapters=fwd_adapters, - subclass_list=list(subclass_list), ) wrapper_bwd_def = _register_wrapper_op( wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op From a9858cda69cf47183e61e2e8cdabf4da3eb8bede Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 16 Jul 2026 15:38:04 +0200 Subject: [PATCH 083/108] [PyTorch] Reword nullable-return note to not trip pylint fixme (W0511) The "TODO:" prefix on the sentinel-encoding note triggered pylint's W0511 (fixme) and dropped the lint score. Reword as a plain forward-looking note -- same information, no fixme keyword. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/custom_op.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index b24ab2f11d..85ca71b785 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -122,8 +122,8 @@ # 0-element uint8 tensor: a non-nullable ``Tensor[]`` schema is required for # ``register_autograd`` to attach a ``grad_fn`` to the outputs. # -# TODO: once https://github.com/pytorch/pytorch/pull/187434 lands, a nullable -# ``Tensor?[]`` return schema lets ``None`` pass through directly and this +# Once https://github.com/pytorch/pytorch/pull/187434 lands, a nullable +# ``Tensor?[]`` return schema will let ``None`` pass through directly and this # sentinel encoding (``_encode_none`` / ``_decode_none``) can be removed. _NONE_SENTINEL_DTYPE = torch.uint8 From 9bce3ad2982d962d8bce61068beac0e605fe0b3d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 16 Jul 2026 17:11:50 +0200 Subject: [PATCH 084/108] [PyTorch] Drop dead rht_matrix field from LinearFwdArgs The NVFP4 RHT matrix was pinned as an op input but never read by the op body: the reconstructed quantizer that crosses the op boundary already carries its own compile-time-baked RHT matrix (built once in NVFP4Quantizer.__init__). Worse, calling the @lru_cache'd get_rht_matrix from the traced forward made Dynamo ignore the cache wrapper and trace the 16x16 RHT construction (eye + mul + mm) into the graph, which Inductor does not constant-fold, so it was rebuilt on every compiled invocation. Remove the traced RHT fetch and the now-unused LinearFwdArgs.rht_matrix field. Numerically identical and CUDA-graph-safe. The cuBLAS workspace pin is kept (its only consumer, general_gemm, allocates it lazily inside the op with no eager pre-warm). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 25 +++++++-------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index f650d64397..8a5141b58b 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -109,13 +109,12 @@ class LinearFwdArgs: # across the torch.compile custom-op boundary. weight_workspace: Optional[TensorOrQuantized] - # Workspace pinning (torch.compile). The process-global, lru_cached cuBLAS / - # NVFP4-RHT workspaces are fetched in the traced forward and threaded in as op - # inputs so they are allocated at trace time rather than lazily inside the op. - # The op body never reads them (general_gemm / the quantizer fetch the same - # globals by address). None on eager / non-compiled paths. + # Workspace pinning (torch.compile). The process-global, lru_cached cuBLAS + # workspace is fetched in the traced forward and threaded in as an op input so + # it is allocated at trace time rather than lazily inside the op. The op body + # never reads it (general_gemm fetches the same global by address). None on + # eager / non-compiled paths. cublas_workspace: Optional[torch.Tensor] - rht_matrix: Optional[torch.Tensor] # --- requires_grad flags (cached so backward does not re-query) --- input_requires_grad: bool @@ -2329,19 +2328,12 @@ def forward( ) wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None - # Pin the lazily-cached cuBLAS (and NVFP4-RHT) workspaces as op inputs so - # they are materialized at trace time (external to the cudagraph pool) - # rather than inside the op during capture. See LinearFwdArgs for details. + # Pin the lazily-cached cuBLAS workspace as an op input so it is + # materialized at trace time (external to the cudagraph pool) rather + # than inside the op during capture. See LinearFwdArgs for details. cublas_workspace = None - rht_matrix = None if use_compiled_op: cublas_workspace = get_cublas_workspace(inp.device.index, False, False) - from ..tensor.nvfp4_tensor import NVFP4Quantizer, get_rht_matrix - - if isinstance(input_quantizer, NVFP4Quantizer): - rht_matrix = get_rht_matrix( - input_quantizer._with_random_sign_mask, inp.device.index - ) fwd_args = LinearFwdArgs( # tensors @@ -2350,7 +2342,6 @@ def forward( bias=linear_bias_tensor, weight_workspace=weight_workspace, cublas_workspace=cublas_workspace, - rht_matrix=rht_matrix, # requires_grad flags input_requires_grad=inp.requires_grad, weight_requires_grad=weight_tensor.requires_grad, From 7e7201393b9412cf3275d9fb3f719e8de90d9545 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 28 Jul 2026 14:28:01 +0200 Subject: [PATCH 085/108] 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 086/108] 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 087/108] 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 088/108] 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 089/108] 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 cb79da2024eece081d0de8d714e8d585b3c92a70 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 28 Jul 2026 15:44:39 +0200 Subject: [PATCH 090/108] Fix three fake/eager divergences in the Linear fake forward Carried over from the TensorProto PR (#3153), where these impls used to live without a caller; they belong here, with the custom-op registration that consumes them. - Weight workspace: quantize_weight returns a fresh workspace on every cache miss with cache=True, not only when update_workspace is set; it discards a cached workspace that fails _is_weight_workspace_valid; and it quantizes persistent workspaces with internal=False so the cache holds wrapper tensors. The fake did none of the three. - backward_override="dequantized" forces save_original_input=False in the eager forward; the fake only handled "high_precision", so it aliased the original input where eager saves a rowwise-only quantized one. - The output's requires_grad ignored the bias, describing the output of a bias-only-trainable Linear as non-differentiable. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 31 +++++++++++++-------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 8a5141b58b..975ad292b2 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -27,6 +27,7 @@ is_ub_initialized, using_cublasmp_backend, quantize_weight, + _is_weight_workspace_valid, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -71,7 +72,7 @@ prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorProto, register_custom_op, is_value_opaque_quantizer +from ..dynamo import TensorProto, to_tensor_proto, register_custom_op, is_value_opaque_quantizer from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -717,6 +718,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 @@ -778,21 +781,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( @@ -814,7 +822,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 d5e1f20d4909d5ac883bf1cccd6ef9649aaccc34 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 29 Jul 2026 09:58:38 +0200 Subject: [PATCH 091/108] 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 cfe3bf2960f14194682df86fa44ec1d0e28de5f7 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 29 Jul 2026 12:23:28 +0200 Subject: [PATCH 092/108] [PyTorch] Drop _STORAGE_REGISTRY; carry the class itself in the flatten context The __tensor_flatten__ context stored the storage / wrapper class as a qualname string, which __tensor_unflatten__ then resolved through a global registry populated by __init_subclass__. The indirection only existed because OpaqueValueBundle could not carry a class object across the custom-op boundary. Teach the bundle to handle classes exactly as it already handles Enum (render by name, add the class to the FX graph globals) and store type(self) in the context directly. Quantizer._storage_metadata already returned the class, so create_metadata no longer has to stringify it. _all_quantized_tensor_subclasses now walks QuantizedTensor.__subclasses__() instead of filtering the registry. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 4 +-- .../pytorch/dynamo/custom_op.py | 26 +++++++++++++++---- .../pytorch/dynamo/tensor_proto.py | 6 +---- .../pytorch/quantized_tensor.py | 16 +++--------- 4 files changed, 27 insertions(+), 25 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index a0501cca8e..7c684e1678 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -31,7 +31,7 @@ from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer -from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, _STORAGE_REGISTRY +from transformer_engine.pytorch.quantized_tensor import QuantizedTensor from transformer_engine.pytorch.dynamo import TensorProto, to_tensor_proto from transformer_engine.pytorch import ( is_fp8_available, @@ -672,7 +672,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/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 85ca71b785..314202311d 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -110,7 +110,6 @@ QuantizedTensor, QuantizedTensorStorage, Quantizer, - _STORAGE_REGISTRY, _quantized_tensor_passthrough_ops, prepare_for_saving, ) @@ -155,10 +154,10 @@ class OpaqueValueBundle: Wraps a ``{name: value}`` dict so many small non-Tensor args pass through a single custom-op input; registered as a torch.compile *value* opaque type (Dynamo specializes the graph on its contents). Allowed values: primitives - in :attr:`PRIMITIVE_TYPES` (incl. ``torch.Size``), ``enum.Enum``, any - registered value-opaque type (e.g. TE quantizers), plus nested tuples / + in :attr:`PRIMITIVE_TYPES` (incl. ``torch.Size``), ``enum.Enum``, classes, + any registered value-opaque type (e.g. TE quantizers), plus nested tuples / lists / dicts thereof (so a bundle can carry a ``__tensor_flatten__`` - context verbatim). + context verbatim -- including its ``cls`` entry). """ PRIMITIVE_TYPES: Tuple[type, ...] = ( @@ -179,6 +178,8 @@ def is_simple_value(cls, value: Any) -> bool: return True if isinstance(value, Enum): return True + if isinstance(value, type): + return True if _is_opaque_value_type(type(value)): return True if isinstance(value, dict): @@ -210,6 +211,10 @@ def _fmt_simple(cls, value: Any) -> str: # ``EnumName.MEMBER`` (the Enum class is added to globals by ``_collect``). if isinstance(value, Enum): return f"{type(value).__name__}.{value.name}" + # Class objects (e.g. the flatten context's ``cls``) render by name; the + # class itself is added to globals by ``_collect``. + if isinstance(value, type): + return value.__name__ if isinstance(value, dict): body = ", ".join(f"{k!r}: {cls._fmt_simple(v)}" for k, v in value.items()) return f"{{{body}}}" @@ -280,6 +285,9 @@ def _collect(value: Any) -> None: if isinstance(value, Enum): globals_[type(value).__name__] = type(value) return + if isinstance(value, type): + globals_[value.__name__] = value + return if isinstance(value, OpaqueValueBundle.PRIMITIVE_TYPES): return if _is_opaque_value_type(type(value)): @@ -1258,7 +1266,15 @@ def _forward(*flat: Any) -> List[torch.Tensor]: def _all_quantized_tensor_subclasses() -> List[type]: """Return every imported ``QuantizedTensor`` wrapper subclass.""" import transformer_engine.pytorch.tensor # noqa: F401 pylint: disable=import-outside-toplevel,unused-import - return [cls for cls in _STORAGE_REGISTRY.values() if issubclass(cls, QuantizedTensor)] + + found: List[type] = [] + stack = list(QuantizedTensor.__subclasses__()) + while stack: + cls = stack.pop() + if cls not in found: + found.append(cls) + stack.extend(cls.__subclasses__()) + return found def register_custom_op( diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index 5a7c7e59c6..506abc9bf6 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 d9740c820e..ef0d98659b 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -30,10 +30,6 @@ _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. @@ -158,12 +154,6 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: # :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} @@ -175,7 +165,7 @@ def __tensor_flatten__(self) -> Tuple[list, Dict[str, Any]]: attr for attr, _ in self._FLATTEN_TENSOR_BUFFERS 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 @@ -192,7 +182,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._FLATTEN_TENSOR_BUFFERS: @@ -487,7 +477,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"], From 63433178debfee9323e910ca620eb4226ebf39fb Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 29 Jul 2026 12:23:42 +0200 Subject: [PATCH 093/108] [PyTorch] Make to_tensor_proto idempotent so the weight workspace keeps its quantizer to_tensor_proto reads the quantizer from _quantizer, an attribute of the real TE tensors; a TensorProto exposes it as quantizer. Re-describing a proto therefore produced a proto with quantizer=None. _proto_view converts weight_workspace to a TensorProto before the fake impl runs, and the fake Linear forward re-describes it on the FP8 weight-cache hit, so update_usage(rowwise_usage=True) hit a proto it believed to be unquantized: ValueError: update_usage called on a non-quantized TensorProto which dynamo surfaced as an Unsupported observed exception. This only fired with is_first_microbatch, the sole path feeding a cached workspace into the op, and made test_te_linear_compile_is_first_microbatch fail for both compile modes. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/tensor_proto.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index 506abc9bf6..c025fd8df8 100644 --- a/transformer_engine/pytorch/dynamo/tensor_proto.py +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -150,7 +150,19 @@ def to_tensor_proto(tensor: Any) -> TensorProto: Works for plain ``torch.Tensor`` and for ``QuantizedTensorStorage`` / ``QuantizedTensor``. A *bare* storage exposes its (fake) dtype via ``_dtype`` rather than ``.dtype``. + + Idempotent: a proto-view field (``_proto_view``) may be re-described by a + fake impl, and a ``TensorProto`` holds its quantizer as ``quantizer``, not + ``_quantizer`` -- rebuilding it blindly would silently drop it. """ + if isinstance(tensor, TensorProto): + return TensorProto( + shape=tuple(tensor.shape), + dtype=tensor.dtype, + quantizer=tensor.quantizer, + requires_grad=tensor.requires_grad, + device=tensor.device, + ) requires_grad = bool(getattr(tensor, "requires_grad", False)) dtype = getattr(tensor, "dtype", None) if dtype is None: From c0ab88d5d6a6ca97ba3a362bebdbee0c5c6c1eeb Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 29 Jul 2026 13:32:51 +0200 Subject: [PATCH 094/108] [PyTorch] Stop re-describing the weight workspace proto in the Linear fake forward 63433178d made to_tensor_proto idempotent to stop the FP8 weight-cache hit from losing its quantizer. That papered over the actual defect: _proto_view has already turned weight_workspace into a TensorProto by the time the fake forward runs, so calling to_tensor_proto on it was a no-op that only ever had the chance to be wrong. Copy the proto with dataclasses.replace instead -- keeping the fresh-quantizer-copy semantics update_usage relies on -- and drop the idempotency branch, which now has no caller (the other call site is guarded by an isinstance check). _is_weight_workspace_valid goes with it: it dispatches on storage types (Float8TensorStorage, ...), so on a proto every branch fell through and it returned True unconditionally. The fake path therefore never detected a stale workspace; that gap is now stated in a comment rather than hidden behind a call that looked like it worked. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/dynamo/tensor_proto.py | 15 ++++----------- transformer_engine/pytorch/module/linear.py | 14 +++++++------- 2 files changed, 11 insertions(+), 18 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py index c025fd8df8..91a1b2ba94 100644 --- a/transformer_engine/pytorch/dynamo/tensor_proto.py +++ b/transformer_engine/pytorch/dynamo/tensor_proto.py @@ -151,18 +151,11 @@ def to_tensor_proto(tensor: Any) -> TensorProto: ``QuantizedTensor``. A *bare* storage exposes its (fake) dtype via ``_dtype`` rather than ``.dtype``. - Idempotent: a proto-view field (``_proto_view``) may be re-described by a - fake impl, and a ``TensorProto`` holds its quantizer as ``quantizer``, not - ``_quantizer`` -- rebuilding it blindly would silently drop it. + Not for re-describing a ``TensorProto``: a proto holds its quantizer as + ``quantizer``, not ``_quantizer``, so it would come back unquantized. Fake + impls already receive protos from ``_proto_view`` -- copy those with + ``dataclasses.replace``. """ - if isinstance(tensor, TensorProto): - return TensorProto( - shape=tuple(tensor.shape), - dtype=tensor.dtype, - quantizer=tensor.quantizer, - requires_grad=tensor.requires_grad, - device=tensor.device, - ) requires_grad = bool(getattr(tensor, "requires_grad", False)) dtype = getattr(tensor, "dtype", None) if dtype is None: diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 975ad292b2..ddf2b19bfa 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -4,7 +4,7 @@ """Linear API""" -from dataclasses import dataclass +from dataclasses import dataclass, replace as dataclass_replace from typing import Any, Callable, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op @@ -27,7 +27,6 @@ is_ub_initialized, using_cublasmp_backend, quantize_weight, - _is_weight_workspace_valid, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -72,7 +71,7 @@ prepare_for_saving, restore_from_func_ctx, ) -from ..dynamo import TensorProto, to_tensor_proto, register_custom_op, is_value_opaque_quantizer +from ..dynamo import TensorProto, register_custom_op, is_value_opaque_quantizer from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import clear_columnwise_cache, is_custom @@ -782,12 +781,13 @@ def _linear_forward_impl_fake( weightmat_aliases_weight = True else: weightmat_is_storage = True + # ``_proto_view`` already turned this field into a ``TensorProto``, so + # the real ``_is_weight_workspace_valid`` cannot run here: it dispatches + # on storage types and would accept any proto unconditionally. 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) + # Copy, so the ``update_usage`` below stays off the input proto. + weightmat = dataclass_replace(workspace) else: weightmat = TensorProto( shape=tuple(weight.shape), From 6e61d36e0411747380e543fab27de3729625f868 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 29 Jul 2026 13:53:03 +0200 Subject: [PATCH 095/108] 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 693f5925865044623ceca2ba80fde2b29f44ab2c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 12:12:01 +0200 Subject: [PATCH 096/108] [PyTorch] Warn once when the torch.compile registrations fall back to eager Registration of the opaque quantizer types, of OpaqueValueBundle and of the custom ops is best-effort so that TE stays importable on PyTorch builds without the opaque-object APIs. Failing silently makes the fallback resurface much later as an obscure error inside the compiled region, so warn once per process instead and point at the PyTorch version. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Pawel Gadzinski --- .../pytorch/dynamo/custom_op.py | 12 ++++---- .../pytorch/dynamo/quantizer_opaque.py | 29 +++++++++++++++++-- 2 files changed, 32 insertions(+), 9 deletions(-) diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 314202311d..9064963062 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -85,7 +85,6 @@ from __future__ import annotations import dataclasses import types as _types -import warnings from enum import Enum from typing import ( Any, @@ -105,6 +104,7 @@ from torch._prims_common import make_contiguous_strides_for +from .quantizer_opaque import warn_compile_unsupported from .tensor_proto import TensorProto, to_tensor_proto from ..quantized_tensor import ( QuantizedTensor, @@ -309,7 +309,8 @@ def _collect(value: Any) -> None: register_opaque_type(OpaqueValueBundle, typ="value") _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) -except Exception: # pylint: disable=broad-exception-caught # pragma: no cover - older torch without opaque_object +except Exception as e: # pylint: disable=broad-exception-caught # pragma: no cover - older torch without opaque_object + warn_compile_unsupported(f"could not register OpaqueValueBundle as an opaque type ({e})") _is_opaque_value_type = None _is_opaque_reference_type = None _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None @@ -1362,11 +1363,8 @@ def register_custom_op( bwd_fake_impl=bwd_fake_impl, ) except (ImportError, AttributeError, RuntimeError, TypeError) as e: - warnings.warn( - f"Could not register the torch.compile custom op '{op_name}' " - f"({type(e).__name__}: {e}); modules using it will fall back to eager " - "execution under torch.compile (a graph break, incompatible with " - "fullgraph=True)." + warn_compile_unsupported( + f"could not register the custom op '{op_name}' ({type(e).__name__}: {e})" ) return None diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index a342dd1e6c..96770e5e4e 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -6,11 +6,34 @@ from __future__ import annotations import enum +import warnings from typing import Any, Dict, Tuple, get_type_hints from ..constants import DType +_warned_compile_unsupported = False + + +def warn_compile_unsupported(reason: str) -> None: + """Warn once per process that TE's torch.compile path is off. + + The registrations below all work or all fail together, so one message is + enough. + """ + global _warned_compile_unsupported # pylint: disable=global-statement + if _warned_compile_unsupported: + return + _warned_compile_unsupported = True + warnings.warn( + "Transformer Engine torch.compile support is disabled: " + f"{reason}. Modules will fall back to eager execution under " + "torch.compile, i.e. a graph break, which is incompatible with " + "fullgraph=True. Use a newer PyTorch build.", + stacklevel=3, + ) + + # Qualnames of the registered quantizer classes. The set holds strings rather # than the classes themselves so that ``is_value_opaque_quantizer`` can be # called inside a ``torch.compile``'d function without a graph break: Dynamo @@ -117,18 +140,20 @@ def register_value_opaque_quantizer(cls: type) -> None: register_opaque_type, is_opaque_value_type, ) - except (ImportError, AttributeError): + except (ImportError, AttributeError) as e: # Older PyTorch without the opaque-object API: eager value semantics # still work; torch.compile specialization on the quantizer does not. + warn_compile_unsupported(f"this PyTorch build has no opaque-object API ({e})") return try: if not is_opaque_value_type(cls): register_opaque_type(cls, typ="value") - except (RuntimeError, TypeError): + except (RuntimeError, TypeError) as e: # Keep TE importable: neither the opaque-type query nor the registration # must crash the import, e.g. on PyTorch versions with only partial / # experimental opaque-object support. + warn_compile_unsupported(f"could not register {cls.__name__} as an opaque type ({e})") return _VALUE_OPAQUE_QUALNAMES.add(cls.__qualname__) From 1523d15c5c526a1d3c0265590a65e52fe18d5aad Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 12:14:07 +0200 Subject: [PATCH 097/108] [PyTorch] Trim the weight-workspace validity comment in the fake forward Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index ddf2b19bfa..fd4d043582 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -781,9 +781,8 @@ def _linear_forward_impl_fake( weightmat_aliases_weight = True else: weightmat_is_storage = True - # ``_proto_view`` already turned this field into a ``TensorProto``, so - # the real ``_is_weight_workspace_valid`` cannot run here: it dispatches - # on storage types and would accept any proto unconditionally. + # No ``_is_weight_workspace_valid`` here: it dispatches on storage + # types and would accept any proto. workspace = args.weight_workspace if workspace is not None: # Copy, so the ``update_usage`` below stays off the input proto. From 4c4519c2babee66e2e10c9d78ab2166061de253d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 12:16:21 +0200 Subject: [PATCH 098/108] [PyTorch] Drop the weight-workspace comment in the fake forward Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 1471e83695..51ea87caf9 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -783,8 +783,6 @@ def _linear_forward_impl_fake( weightmat_aliases_weight = True else: weightmat_is_storage = True - # No ``_is_weight_workspace_valid`` here: it dispatches on storage - # types and would accept any proto. workspace = args.weight_workspace if workspace is not None: # Copy, so the ``update_usage`` below stays off the input proto. From a24e4e006981323e8cfd458bb153ae9af9cd6c83 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 13:20:17 +0200 Subject: [PATCH 099/108] 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 100/108] 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 101/108] 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 102/108] 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 103/108] 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"], From 7514cf968d16c03435439e7f4316f7d1bc8dc8f6 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 16:32:21 +0200 Subject: [PATCH 104/108] Enable post-RHT amax in the NVFP4 value-object test factory The quantize kernel rejects with_rht=True without with_post_rht_amax=True (pre-RHT amax unsupported); mirror the recipe, which always sets both together. Lost in a 3-way merge that took the HEAD side. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 19a59b99aa..c3f4a83358 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -759,12 +759,14 @@ def _current_scaling(amax_epsilon=0.0): def _nvfp4(with_rht=True): # Default with_rht=True so the quantize round-trip below exercises the # derived ``rht_matrix`` tensor (the field most likely to be dropped on - # value-key reconstruction). + # value-key reconstruction). Post-RHT amax is required by the kernel + # whenever RHT is on (pre-RHT amax is unsupported). return NVFP4Quantizer( fp4_dtype=tex.DType.kFloat4E2M1, rowwise=True, columnwise=True, with_rht=with_rht, + with_post_rht_amax=with_rht, ) From 58f37753c64d6cf6a9817637c4b97d77c3a71ad0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 16:32:27 +0200 Subject: [PATCH 105/108] Reformat test_torch_compile.py with black Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_torch_compile.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index c3f4a83358..6c587511c1 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1528,11 +1528,17 @@ def fn(inp): out_eager = model(inp_eager) out_eager.sum().backward() torch.testing.assert_close( - out.detach(), out_eager.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL, + out.detach(), + out_eager.detach(), + atol=_EAGER_ATOL, + rtol=_EAGER_RTOL, msg=f"forward mismatch at batch={batch}", ) torch.testing.assert_close( - inp.grad, inp_eager.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL, + inp.grad, + inp_eager.grad, + atol=_EAGER_ATOL, + rtol=_EAGER_RTOL, msg=f"dgrad mismatch at batch={batch}", ) @@ -1543,6 +1549,6 @@ def fn(inp): recompile_count_after = counters["stats"].get("recompile_reasons", 0) assert recompile_count_after == recompile_count_baseline, ( - f"Unexpected recompilation(s) across different batch sizes: " + "Unexpected recompilation(s) across different batch sizes: " f"{recompile_count_after - recompile_count_baseline} recompile(s) detected" ) From ef427f3790e5591a37e86287632a24d2aae67e76 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 16:36:59 +0200 Subject: [PATCH 106/108] [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. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/module/linear.py | 31 ++++++++++++++++----- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 51ea87caf9..b2b4ec912c 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -669,12 +669,24 @@ def _linear_forward_impl( if is_dist_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. A custom op may not return a + # tensor aliasing an input or another return, and the cached FP8 weight + # is the same object as ``new_weight_workspace`` (cache miss) or + # ``weight_workspace`` (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, ) @@ -906,7 +918,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, ...]: @@ -919,7 +931,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`` is ``(out, new_weight_workspace)``; only the latter is used, + # to rebuild the deduped weight save slot. inp = fwd_args.inp weight = fwd_args.weight @@ -1014,6 +1027,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": @@ -1763,7 +1780,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 2dc0bcfc4480422134aad2bc1f1076a35e7f577b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 31 Jul 2026 17:06:04 +0200 Subject: [PATCH 107/108] Take requires_grad from the args bundle, not from the input tensors _linear_forward_impl decided columnwise usage from inp.requires_grad, while its fake twin used args.input_requires_grad -- a value read outside the graph and baked into it as a constant. Under mode="reduce-overhead" cudagraph_trees swaps the inputs for static placeholders while recording, and those carry requires_grad=False, so the real impl quantized the weight without columnwise data and returned one buffer fewer than the fake had promised. Every later slot shifted by one and inductor's assert_size_stride landed on the 1-D scale_inv where it expected the 2-D transpose: AssertionError: wrong number of dimensions1 for op: torch.ops.transformer_engine_compile.linear.default Read both flags from LinearFwdArgs so the two impls decide from one source. backward_needs_input is the same trap and changes with it, though it happens not to fire today: the weight survives the swap as a static parameter. Fixes test_te_linear_compiles[*-reduce-overhead] for every recipe whose inner-tensor count depends on columnwise usage -- Float8CurrentScaling on TN-capable archs, MXFP8/NVFP4/Float8BlockScaling on Blackwell. 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 b2b4ec912c..60be12708b 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -362,7 +362,7 @@ 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 + 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 ) @@ -486,7 +486,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: From 47229f0fca257010d73b0d8247184d92b43c73b5 Mon Sep 17 00:00:00 2001 From: Masato Shinokawa Date: Tue, 4 Aug 2026 02:22:34 -0700 Subject: [PATCH 108/108] Carry symbolic shapes through compile custom ops --- tests/pytorch/test_torch_compile.py | 12 +++---- .../pytorch/dynamo/custom_op.py | 32 +++++++++++++++++++ transformer_engine/pytorch/module/linear.py | 29 ++--------------- 3 files changed, 38 insertions(+), 35 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 6c587511c1..0f79ed58ef 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1488,13 +1488,9 @@ def test_te_linear_dynamic_shapes(): Key correctness property: a graph compiled for batch=16 must produce numerically correct results for batch=32 without triggering a recompile. - This exercises two fixes for dynamic shapes: - 1. ``_linear_setup_ctx`` no longer stores ``inp_shape`` in the value bundle - (torch.Size with SymInt dims is not hashable in OpaqueValueBundle). - 2. ``_linear_backward_impl_fake`` derives dgrad shape from grad_output + - weight + SP config instead of relying on the stored ``inp_shape``. - 3. ``_linear_backward`` reconstructs ``inp_shape`` on-the-fly from the same - tensor sources when it is None (compiled mode). + ``LinearBwdArgs.inp_shape`` crosses the custom-op boundary through a + ``SymInt[]`` schema slot rather than the value-opaque metadata bundle. The + backward fake and real implementation both consume that symbolic shape. FP8 + dynamic=True is tracked separately (requires resolving ``UnsafeScriptObjectError`` for TorchScript quantizer objects with Dynamo). @@ -1508,7 +1504,7 @@ def fn(inp): return model(inp) torch._dynamo.reset() - compiled = torch.compile(fn, fullgraph=True) + compiled = torch.compile(fn, fullgraph=True, dynamic=True) batch_sizes = [16, 32, 48] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 9064963062..19f51da3e7 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -630,6 +630,37 @@ def grad_slot(self) -> Optional[int]: return 0 +class _SymIntSequenceAdapter(_Adapter): + """``torch.Size`` / ``Optional[torch.Size]`` -> ``SymInt[]`` schema slot. + + Shape values must remain graph data when dimensions are symbolic. Putting a + ``torch.Size`` containing ``SymInt`` values in the value-opaque simple bundle + would instead try to specialize and hash it as a Python constant. + """ + + def __init__(self, name: str, is_optional: bool) -> None: + self.name = name + self.type_str = "SymInt[]?" if is_optional else "SymInt[]" + + @classmethod + def try_build(cls, name: str, annot: Any) -> Optional["_SymIntSequenceAdapter"]: + stripped, is_optional = _strip_optional(annot) + if stripped is torch.Size: + return cls(name, is_optional) + return None + + def schema_slots(self) -> List[Tuple[str, str]]: + return [(self.name, self.type_str)] + + def to_slots(self, owner: Any) -> Dict[str, Any]: + value = getattr(owner, self.name) + return {self.name: None if value is None else list(value)} + + def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: + value = args[self.name] + kwargs[self.name] = None if value is None else torch.Size(value) + + class _QuantizerAdapter(_Adapter): """``Quantizer`` / ``Optional[Quantizer]`` -> one own ``OpaqueValueBundle`` slot. @@ -804,6 +835,7 @@ def from_slots(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None: _FIELD_ADAPTERS: Tuple[type, ...] = ( _TensorOrQuantizedAdapter, _TensorAdapter, + _SymIntSequenceAdapter, _ReferenceOpaqueAdapter, _QuantizerAdapter, ) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 60be12708b..96895cbc95 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -955,10 +955,7 @@ def _linear_setup_ctx( bwd_args.use_bias = bias is not None bwd_args.requires_dgrad = fwd_args.input_requires_grad bwd_args.requires_wgrad = fwd_args.weight_requires_grad - # Don't store inp_shape in the value bundle: under torch.compile(dynamic=True) - # inp.shape contains SymInt dims which are not hashable in OpaqueValueBundle. - # The backward reconstructs inp_shape from grad_output + weight + SP config. - bwd_args.inp_shape = None + bwd_args.inp_shape = inp.shape # Numerical / dtype config bwd_args.activation_dtype = fwd_args.activation_dtype @@ -1110,18 +1107,6 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. ) nvtx_range_pop(f"{nvtx_label}.fsdp_gather") - # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes). - if bwd_args.inp_shape is None: - in_features = saved_weight.shape[-1] - go_leading = grad_output.shape[0] - if bwd_args.parallel_mode == "column" and bwd_args.sequence_parallel: - inp_leading = go_leading // bwd_args.tp_size - elif bwd_args.parallel_mode == "row" and bwd_args.sequence_parallel: - inp_leading = go_leading * bwd_args.tp_size - else: - inp_leading = go_leading - bwd_args.inp_shape = torch.Size([inp_leading, *grad_output.shape[1:-1], in_features]) - # Configure Userbuffers communication (comm+GEMM overlap) bwd_args.ub_obj_gradout = None ub_obj_dgrad = None @@ -1689,18 +1674,8 @@ def _linear_backward_impl_fake( dgrad = None if args.requires_dgrad: # dgrad has the logical input shape and may be quantized for the next op. - # Derive shape from grad_output + weight + SP config instead of args.inp_shape: - # inp_shape is not stored in the value bundle under dynamic shapes (SymInt is - # not hashable in OpaqueValueBundle), so we reconstruct it here. - go_leading = args.grad_output.shape[0] - if args.parallel_mode == "column" and args.sequence_parallel: - dgrad_leading = go_leading // args.tp_size - elif args.parallel_mode == "row" and args.sequence_parallel: - dgrad_leading = go_leading * args.tp_size - else: - dgrad_leading = go_leading dgrad = TensorProto( - shape=(dgrad_leading, *args.grad_output.shape[1:-1], in_features), + shape=tuple(args.inp_shape), dtype=out_dtype, quantizer=args.grad_input_quantizer, device=args.grad_output.device,