diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 1286492a6e..8ebce563ce 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -32,6 +32,9 @@ is_mxfp8_available, is_fp8_block_scaling_available, is_nvfp4_available, + Float8BlockQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, ) from utils import recipe_id @@ -384,3 +387,170 @@ def fn(inp): out = compiled(inp) out.sum().backward() + + +# --------------------------------------------------------------------------- +# Value-opaque quantizers +# --------------------------------------------------------------------------- + + +def _mxfp8(dtype=tex.DType.kFloat8E4M3): + return MXFP8Quantizer(fp8_dtype=dtype) + + +def _blockwise(force_pow_2_scales=True): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=force_pow_2_scales, + ) + + +def _current_scaling(amax_epsilon=0.0): + return Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device=torch.device("cpu"), + amax_epsilon=amax_epsilon, + ) + + +def _nvfp4(with_rht=True): + # Default with_rht=True so the quantize round-trip below exercises the + # derived ``rht_matrix`` tensor (the field most likely to be dropped on + # value-key reconstruction). 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, + ) + + +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, 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.mark.parametrize("factory", _VALUE_QUANTIZERS) +def test_quantizer_value_object(factory): + """Value semantics + ``__fx_repr__`` round-trip via the production FX path.""" + a = factory() + + # ``__fx_repr__`` (used by torch.compile codegen) rebuilds an equal object. + repr_str, globals_ = a.__fx_repr__() + rebuilt = eval(repr_str, dict(globals_)) # pylint: disable=eval-used + assert rebuilt == a and rebuilt is not a + assert hash(rebuilt) == hash(a) + + # The rebuilt quantizer must also *behave* identically, not just compare + # equal: equality only looks at the value key, so a field the kernel needs + # but that is absent from the key (e.g. NVFP4's derived ``rht_matrix``) would + # slip through the checks above and only blow up at quantize time. Run the + # real quantize kernel on both and require bit-exact results. + if torch.cuda.is_available() and _hw_available(a): + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0) + + +def test_value_quantizer_rejects_process_group(): + """A value quantizer holding a live ProcessGroup must refuse to be turned + into a value key / FX constant (raise), not silently drop the group.""" + import torch.distributed as dist # pylint: disable=import-outside-toplevel + + created = not dist.is_initialized() + if created: + dist.init_process_group(backend="gloo", store=dist.HashStore(), rank=0, world_size=1) + try: + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + q.amax_reduction_group = dist.group.WORLD + # Every value-materialization path must reject it (hash, eq, __fx_repr__). + with pytest.raises(TypeError): + hash(q) + with pytest.raises(TypeError): + q.__fx_repr__() + finally: + if created: + dist.destroy_process_group() + + +if _opaque_available: + # A minimal custom op taking a tensor and a value-opaque quantizer that + # quantizes + dequantizes inside it, one per production quantizer class. + # ``test_quantizer_value_object_fullgraph`` drives this under + # ``torch.compile(fullgraph=True)`` so the quantizer is used *inside* the + # graph -- proving the opaque-type registration took effect (a graph break + # would make ``fullgraph=True`` raise). + _qdq_lib = torch.library.Library("test_te_qdq", "DEF") + _QDQ_OPS = {} + for _qcls in ( + MXFP8Quantizer, + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + NVFP4Quantizer, + ): + _op = f"qdq_{_qcls.__name__}" + _qdq_lib.define(f"{_op}(Tensor x, {get_opaque_type_name(_qcls)} q) -> Tensor") + + @torch.library.impl(f"test_te_qdq::{_op}", "CompositeExplicitAutograd", lib=_qdq_lib) + def _qdq_impl(x, q): + return q(x).dequantize() + + @torch.library.register_fake(f"test_te_qdq::{_op}", lib=_qdq_lib) + def _qdq_fake(x, q): + return torch.empty_like(x) + + _QDQ_OPS[_qcls] = getattr(torch.ops.test_te_qdq, _op) + + +@pytest.mark.skipif( + not _opaque_available, + reason="torch.compile opaque-object support requires PyTorch >= 2.11", +) +@pytest.mark.parametrize("factory", _VALUE_QUANTIZERS) +def test_quantizer_value_object_fullgraph(factory): + """Quantizer is usable *inside* a torch.compile(fullgraph=True) graph. + + A custom op quantizes+dequantizes with the (opaque value) quantizer; the + compiled result must match eager. ``fullgraph=True`` raises on any graph + break, so this proves the opaque-type registration actually took effect -- + unlike merely passing the quantizer through. + """ + q = factory() + if not (torch.cuda.is_available() and _hw_available(q)): + pytest.skip("format not supported on this HW") + + op = _QDQ_OPS[type(q)] + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def fn(inp): + return op(inp, q) + + ref = fn(x) + torch._dynamo.reset() + out = torch.compile(fn, fullgraph=True)(x) + torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0) diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py new file mode 100644 index 0000000000..ee860c78e3 --- /dev/null +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""torch.compile glue for Transformer Engine.""" + +from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer + +__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 new file mode 100644 index 0000000000..a342dd1e6c --- /dev/null +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -0,0 +1,134 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Value-opaque quantizers for torch.compile.""" + +from __future__ import annotations +import enum +from typing import Any, Dict, Tuple, get_type_hints + +from ..constants import DType + + +# 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 +# can evaluate ``type(q).__qualname__ in ``, but not set +# membership of a class registered as an opaque type. +_VALUE_OPAQUE_QUALNAMES: set = set() + + +def is_value_opaque_quantizer(quantizer: Any) -> bool: + """Whether *quantizer*'s class is registered as a torch.compile value-opaque + type.""" + return type(quantizer).__qualname__ in _VALUE_OPAQUE_QUALNAMES + + +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. + + Only the value fields (plus derived state via ``_rebuild_derived_state``) + are restored. Non-value attributes such as the deprecated + ``amax_reduction_group`` are deliberately absent on the rebuilt quantizer, + so accessing them fails loudly unless set explicitly. + """ + # Bypass ``__init__`` and restore the value attributes directly: the value + # items already capture every value-defining field (including derived ones), + # and the constructors have heterogeneous signatures / side effects. + obj = cls.__new__(cls) + for name, value in items: + if name == "dtype": + value = DType.cast(value) + object.__setattr__(obj, name, value) + # Restore non-value derived state that ``__init__`` would normally build but + # that cannot live in the value key (e.g. NVFP4's ``rht_matrix`` tensor). + finalize = getattr(obj, "_rebuild_derived_state", None) + if finalize is not None: + finalize() + return obj + + +def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]: + """``__fx_repr__`` for value-opaque quantizers (attached at registration). + + Returns an evaluable expression that rebuilds the quantizer via + :func:`_rebuild_quantizer`, capturing both the helper and the quantizer + class itself in the FX globals so codegen can resolve them with no global + registry and no qualname collisions. + + Raises ``TypeError`` (via :meth:`Quantizer._value_key`) if the quantizer + stores a process group (e.g. a non-``None`` deprecated + ``amax_reduction_group``): live distributed state must never be baked into + the graph as a constant. Pass the reduction group per quantize call instead + of storing it on the quantizer. + """ + cls = type(self) + items = self._value_key()[1] + return ( + f"_rebuild_quantizer({cls.__name__}, {items!r})", + {"_rebuild_quantizer": _rebuild_quantizer, cls.__name__: cls}, + ) + + +def register_value_opaque_quantizer(cls: type) -> None: + """Register a tensorless quantizer class as a torch.compile value opaque type. + + This is the opt-in point for value semantics: it derives the value fields + from the class annotations and stores them on the class (enabling + config-based ``__eq__`` / ``__hash__``, see + :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`), 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 the value semantics still apply, + only the torch.compile specialization is skipped. + + Only plain value types (``int``/``bool``/``float``/``str`` and enums) may + be annotated: anything else (derived tensors, process groups, containers) + cannot be hashed into the value key or rebuilt from its repr, so it must + be left unannotated and rebuilt in ``_rebuild_derived_state`` instead. + This runs once per class at import time, not in any hot path, so resolving + the annotation strings to real types is affordable. + """ + fields = cls._annotated_fields() + resolved = get_type_hints(cls) + for name in fields: + typ = resolved[name] + if typ not in (int, bool, float, str) and not ( + isinstance(typ, type) and issubclass(typ, enum.Enum) + ): + raise TypeError( + f"{cls.__name__} cannot be a torch.compile value quantizer: " + f"annotated field {name!r} ({typ!r}) is not a plain value type " + "(int/bool/float/str/enum). Remove the annotation and rebuild " + "the field in ``_rebuild_derived_state`` instead." + ) + cls._value_field_names = tuple(fields) + # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the + # class, so attach it before registering. + if "__fx_repr__" not in cls.__dict__: + cls.__fx_repr__ = _quantizer_fx_repr + + try: + from torch._library.opaque_object import ( # pylint: disable=import-outside-toplevel + register_opaque_type, + is_opaque_value_type, + ) + except (ImportError, AttributeError): + # Older PyTorch without the opaque-object API: eager value semantics + # still work; torch.compile specialization on the quantizer does not. + return + + try: + if not is_opaque_value_type(cls): + register_opaque_type(cls, typ="value") + except (RuntimeError, TypeError): + # Keep TE importable: neither the opaque-type query nor the registration + # must crash the import, e.g. on PyTorch versions with only partial / + # experimental opaque-object support. + return + + _VALUE_OPAQUE_QUALNAMES.add(cls.__qualname__) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index cfe488aae5..ce820af9c8 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, @@ -408,6 +409,78 @@ def get_usages(self) -> Dict[str, bool]: "columnwise": self.columnwise_usage, } + @classmethod + def _annotated_fields(cls) -> Dict[str, Any]: + """Annotated fields (name -> annotation) across the ``Quantizer`` MRO, + base first. The class annotations are the single source of truth for + what defines a quantizer's value.""" + fields: Dict[str, Any] = {} + for klass in reversed(cls.__mro__): + if issubclass(klass, Quantizer): + fields.update(klass.__dict__.get("__annotations__", {})) + return fields + + def _value_fields(self) -> Optional[Tuple[str, ...]]: + """Value-defining attribute names, or ``None``. + + Computed from the class annotations and stored on the class by + ``register_value_opaque_quantizer``, which also checks that no + annotated field is a derived tensor or a process group. Looked up in + the class's own ``__dict__`` so a subclass of a registered quantizer + does not silently inherit value semantics without registering itself. + ``None`` (any class not registered) keeps identity-based + equality/hashing and graph-breaks under torch.compile when passed to a + custom op, since such a quantizer cannot be baked into the FX graph as + a constant. + """ + return type(self).__dict__.get("_value_field_names") + + def _check_value_has_no_process_group(self) -> None: + # A value quantizer cannot carry live distributed state into the FX + # graph; reject a stored ``amax_reduction_group`` and pass it per + # quantize call instead. + if isinstance(getattr(self, "amax_reduction_group", None), dist_group_type): + raise TypeError( + f"{type(self).__name__} cannot be used as a torch.compile value " + "object: 'amax_reduction_group' holds a torch.distributed.ProcessGroup, " + "which is live distributed state and must not be baked into an FX " + "graph. Pass the amax reduction group per quantize call instead of " + "storing it on the quantizer." + ) + + def _value_key(self) -> Tuple[Any, ...]: + """Hashable, reproducible key identifying this quantizer's value. + + Only valid for value quantizers (``_value_fields()`` is not ``None``). + """ + fields = self._value_fields() # pylint: disable=assignment-from-none + assert fields is not None, f"{type(self).__name__} is not a value quantizer" + self._check_value_has_no_process_group() + items = [] + for name in fields: + value = getattr(self, name) + if name == "dtype": + # ``DType`` is an ``IntEnum``; store the int so the key stays + # plain: hashable and ``repr``-reproducible for FX codegen. + value = int(value) + items.append((name, value)) + return (type(self).__qualname__, tuple(items)) + + def __eq__(self, other: object) -> Any: + # Value quantizers compare by configuration; everything else keeps the + # default identity semantics (returning ``NotImplemented`` makes Python + # fall back to identity). ``_value_key`` rejects a stored ProcessGroup. + if self is other: + return True + if type(self) is not type(other) or self._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 d2d28aecfb..18975c4c6d 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, safe_quantized_repr from ..constants import DType from ..utils import devices_match, round_up_to_nearest_multiple @@ -211,6 +212,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 e90e35c40c..2e0491d49a 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, safe_quantized_repr from ..constants import dist_group_type, DType @@ -207,7 +208,6 @@ class Float8CurrentScalingQuantizer(Quantizer): dtype: DType """amax reduction options""" with_amax_reduction: bool - amax_reduction_group: Optional[dist_group_type] """Options about how to quantize the tensor""" force_pow_2_scales: bool amax_epsilon: float @@ -257,7 +257,8 @@ def copy(self) -> Float8CurrentScalingQuantizer: rowwise=self.rowwise_usage, columnwise=self.columnwise_usage, with_amax_reduction=self.with_amax_reduction, - amax_reduction_group=self.amax_reduction_group, + # Absent on quantizers rebuilt from a value key (deprecated field). + amax_reduction_group=getattr(self, "amax_reduction_group", None), force_pow_2_scales=self.force_pow_2_scales, amax_epsilon=self.amax_epsilon, ) @@ -387,6 +388,9 @@ def supports_only_rowwise_all_gather(self) -> bool: return True +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 33db63d059..3045662216 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, safe_quantized_repr aten = torch.ops.aten @@ -180,6 +181,9 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return MXFP8BlockScaling +register_value_opaque_quantizer(MXFP8Quantizer) + + class MXFP8Tensor(MXFP8TensorStorage, QuantizedTensor): """Experimental tensor class with FP8 data diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index f59af5637b..ccf06ac166 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, safe_quantized_repr aten = torch.ops.aten @@ -119,7 +120,6 @@ class NVFP4Quantizer(Quantizer): with_post_rht_amax: bool """amax reduction options""" with_amax_reduction: bool - amax_reduction_group: Optional[dist_group_type] """2D block scaling, only applicable for weights.""" with_2d_quantization: bool @@ -136,9 +136,8 @@ class NVFP4Quantizer(Quantizer): """NVFP4 4over6 candidate-selection error mode.""" nvfp4_4over6_err_mode: str - """RHT matrix random sign mask""" + """RHT sign mask (0 when sign randomization is disabled)""" rht_matrix_random_sign_mask_t: int - rht_matrix: torch.Tensor def __init__( self, @@ -176,7 +175,7 @@ def __init__( self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( with_random_sign_mask, torch.cuda.current_device() ) - self.rht_matrix = get_rht_matrix(with_random_sign_mask, torch.cuda.current_device()) + self._rebuild_derived_state() def __getstate__(self): """Exclude unpicklable process group from serialized state.""" @@ -184,6 +183,19 @@ def __getstate__(self): state["amax_reduction_group"] = None return state + def _rebuild_derived_state(self) -> None: + """Build the derived ``rht_matrix`` (also used after value-key reconstruction). + + ``rht_matrix`` is a ``torch.Tensor`` derived from the sign mask, so it + cannot be part of the (hashable) value key. ``__init__`` and + ``_rebuild_quantizer`` both call this hook; the ``lru_cache`` on + :func:`get_rht_matrix` makes an already-seen (flag, device) pair a + cheap hit. + """ + self.rht_matrix = get_rht_matrix( + self.rht_matrix_random_sign_mask_t != 0, torch.cuda.current_device() + ) + def update_quantized( self, src: torch.Tensor, @@ -221,7 +233,8 @@ def copy(self) -> NVFP4Quantizer: rowwise=self.rowwise_usage, columnwise=self.columnwise_usage, with_amax_reduction=self.with_amax_reduction, - amax_reduction_group=self.amax_reduction_group, + # Absent on quantizers rebuilt from a value key (deprecated field). + amax_reduction_group=getattr(self, "amax_reduction_group", None), with_rht=self.with_rht, with_post_rht_amax=self.with_post_rht_amax, with_2d_quantization=self.with_2d_quantization, @@ -230,11 +243,10 @@ def copy(self) -> NVFP4Quantizer: nvfp4_use_4over6=self.nvfp4_use_4over6, nvfp4_e4m3_max=self.nvfp4_e4m3_max, nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, + with_random_sign_mask=self.rht_matrix_random_sign_mask_t != 0, ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm - quantizer.rht_matrix = self.rht_matrix - quantizer.rht_matrix_random_sign_mask_t = self.rht_matrix_random_sign_mask_t return quantizer @@ -334,6 +346,9 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return NVFP4BlockScaling +register_value_opaque_quantizer(NVFP4Quantizer) + + class NVFP4Tensor(NVFP4TensorStorage, QuantizedTensor): """Quantized tensor class with FP4 data