From f543cd3c159c4e25137178cd434de588416cf3a0 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 6 Oct 2026 09:10:31 +0000 Subject: [PATCH 1/9] Refactor output dtype handling for Identity and Hybrid quantizers Signed-off-by: Evgeny --- .../fine_grained_quantization.rst | 25 ++++ tests/pytorch/test_hybrid_quantization.py | 120 ++++++++++++++++++ tests/pytorch/test_identity_quantizer.py | 85 ++++++++++++- transformer_engine/pytorch/module/base.py | 7 +- .../pytorch/quantized_tensor.py | 35 ++++- .../pytorch/tensor/hybrid_tensor.py | 67 +++++++--- .../pytorch/tensor/identity_tensor.py | 38 +++--- 7 files changed, 327 insertions(+), 50 deletions(-) diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst index 301b655b3f1..650d30811ec 100644 --- a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -288,6 +288,31 @@ Note that in the second example the columnwise direction is high precision but holds the value reconstructed from MXFP8, not the original input — see `Choosing the columnwise source`_ above. +Output dtype and workspace updates +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Output-dtype handling is opt-in through ``Quantizer.supports_output_dtype``, +which defaults to ``False``. Identity and Hybrid opt in; native quantizers +continue to ignore output-dtype requests, even inside a custom recipe. + +``IdentityQuantizer.quantize(x, dtype=...)`` selects the held output dtype +without changing the quantizer's default or adding an autograd edge to its +stored data. ``HybridQuantizer`` passes a compute-dtype request to its children; +children that support output-dtype binding retain it for later reconstruction. +Native quantizers keep their existing behavior, including ignoring this +argument. Their encodings and default dequantization metadata are unchanged. + +Hybrid's ``rowwise_dequantized`` source is still explicitly decoded into the +source tensor's dtype. This preserves existing native compositions, including +internal storage. If a high-precision child already rounded into BF16, widening +that value to an FP32 source does not recover the discarded precision. + +In-place Identity updates preserve the destination buffer, dtype, shape and +device. An explicit conflicting ``dtype`` with ``out=`` raises an error for +Identity and Hybrid. Module weight caches rebuild these workspaces when an +explicit compute dtype changes. ``HybridQuantizer.make_empty(..., dtype=...)`` +binds the requested dtype on participating custom children just as quantize does. + Example: one format per GEMM ---------------------------- diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 1063c04e9fa..09313bb2618 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -7704,3 +7704,123 @@ def test_prepare_restore_roundtrip_is_identity(self): leftover = saved_obj.restore_from_saved(list(saved_tensors)) assert leftover == [], "restore_from_saved should consume every element" torch.testing.assert_close(saved_obj.dequantize(), expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("internal", [False, True]) +@pytest.mark.parametrize("rowwise", [False, True]) +@pytest.mark.parametrize("requested_dtype", [None, torch.bfloat16, torch.float32]) +def test_hybrid_output_dtype_lifecycle(internal, rowwise, requested_dtype): + """Reconstructing an absent row preserves its resolved child configuration.""" + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(dtype=torch.float16), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + q.internal = internal + q.rowwise_quantizer.internal = internal + q.columnwise_quantizer.internal = internal + q.set_usage(rowwise=rowwise, columnwise=True) + x = torch.randn(32, 64, device="cuda", dtype=torch.float32, requires_grad=True) + out = q.quantize(x, dtype=requested_dtype) + assert q.rowwise_quantizer.dtype == torch.float16 + assert q.columnwise_quantizer.dtype is None + assert not hasattr(q, "_output_dtype") + row_dtype = requested_dtype or torch.float16 + col_dtype = requested_dtype or x.dtype + expected = x.detach().to(row_dtype).to(col_dtype) + torch.testing.assert_close(out._columnwise_storage.dequantize(), expected, rtol=0, atol=0) + if not internal and rowwise: + out.sum().backward() + torch.testing.assert_close(x.grad, torch.ones_like(x), rtol=0, atol=0) + # Exercise both the original quantizer and the result's attached quantizer. + for multiplier, update in [(3, lambda src: q.update_quantized(src, out)), (5, out.quantize_)]: + updated = x.detach() * multiplier + update(updated) + expected = updated.to(row_dtype).to(col_dtype) + torch.testing.assert_close(out._columnwise_storage.dequantize(), expected, rtol=0, atol=0) + if rowwise: + torch.testing.assert_close( + out._rowwise_storage.dequantize(), updated.to(row_dtype), rtol=0, atol=0 + ) + with pytest.raises(ValueError, match="dtype does not match"): + q.quantize(x, out=out, dtype=torch.float16) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("internal", [False, True]) +@pytest.mark.parametrize("source_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("native_format", ["mxfp8", "nvfp4"]) +def test_hybrid_native_dtype_override_preserves_encoding(internal, source_dtype, native_format): + """Native children ignore output requests and keep the old decode-source dtype.""" + x = torch.randn(128, 256, device="cuda", dtype=source_dtype) + if native_format == "nvfp4" and not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + quantizer_type = MXFP8Quantizer if native_format == "mxfp8" else NVFP4Quantizer + options = {"fp8_dtype": tex.DType.kFloat8E4M3} if native_format == "mxfp8" else {} + row = quantizer_type(**options, rowwise=True, columnwise=False) + col = quantizer_type(**options, rowwise=False, columnwise=True) + row.internal = col.internal = internal + q = HybridQuantizer( + rowwise_quantizer=row, + columnwise_quantizer=col, + columnwise_source="rowwise_dequantized", + ) + q.internal = internal + requested_dtype = torch.float32 if source_dtype == torch.bfloat16 else torch.bfloat16 + assert not row.supports_output_dtype + assert row.with_output_dtype(requested_dtype) is row + expected_row = row.quantize(x) + direct_override = row.quantize(x, dtype=requested_dtype) + _assert_storage_data_exact(direct_override, expected_row, context="native dtype override") + assert direct_override._dtype == expected_row._dtype + expected_col = col.quantize(expected_row.dequantize(dtype=source_dtype)) + actual = q.quantize(x, dtype=requested_dtype) + _assert_storage_data_exact(actual._rowwise_storage, expected_row, context="rowwise encoding") + _assert_storage_data_exact( + actual._columnwise_storage, expected_col, context="columnwise encoding" + ) + assert actual._rowwise_storage._dtype == expected_row._dtype + assert actual._columnwise_storage._dtype == expected_col._dtype + + # Native out= and module caches still ignore conflicting output requests. + from transformer_engine.pytorch.module.base import quantize_weight + + assert row.is_compatible_output(expected_row, dtype=requested_dtype) + # Native update_quantized requires a public tensor destination. + public_row = row.copy() + public_row.internal = False + destination = public_row.quantize(x) + expected_destination = public_row.quantize(x) + assert public_row.quantize(x, out=destination, dtype=requested_dtype) is destination + _assert_storage_data_exact( + destination, expected_destination, context="native out dtype override" + ) + assert destination._dtype == direct_override._dtype + reused, new_cache = quantize_weight( + tensor=x, + quantizer=row, + workspace=expected_row, + workspace_dtype=requested_dtype, + cache=True, + update_workspace=False, + ) + assert reused is expected_row and new_cache is None + + +@pytest.mark.parametrize("internal", [False, True]) +@pytest.mark.parametrize("rowwise", [False, True]) +def test_hybrid_empty_dtype_update(internal, rowwise): + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + q.internal = internal + q.set_usage(rowwise=rowwise, columnwise=True) + out = q.make_empty((32, 64), dtype=torch.bfloat16, device="cuda") + x = torch.randn(32, 64, device="cuda", dtype=torch.float32) + q.update_quantized(x, out) + for child in (out._rowwise_storage, out._columnwise_storage): + if child is None: + continue + torch.testing.assert_close(child.dequantize(), x.bfloat16(), rtol=0, atol=0) diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 2c120af0fe3..264e96c879c 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -671,7 +671,7 @@ def test_make_empty_honors_configured_dtype(self): assert tensor._hp_data.dtype == torch.float32 assert tensor.dequantize().dtype == torch.float32 - def test_update_quantized_synchronizes_dtype(self): + def test_update_quantized_preserves_destination_dtype(self): dst = IdentityQuantizer().make_empty( (4, 8), dtype=torch.bfloat16, @@ -681,9 +681,9 @@ def test_update_quantized_synchronizes_dtype(self): quantizer.update_quantized(torch.ones_like(dst._hp_data), dst) - assert dst.dtype == torch.float32 - assert dst._hp_data.dtype == torch.float32 - assert dst.dequantize().dtype == torch.float32 + assert dst.dtype == torch.bfloat16 + assert dst._hp_data.dtype == torch.bfloat16 + assert dst.dequantize().dtype == torch.bfloat16 @pytest.mark.parametrize("noop", [0, 1]) def test_update_quantized_honors_noop(self, noop): @@ -1903,3 +1903,80 @@ def test_load_bf16_state_dict_into_identity_model_exact(self): out_id = model(x) torch.testing.assert_close(out_id, out_ref, rtol=0.0, atol=0.0) + + +@pytest.mark.parametrize("internal", [False, True]) +@pytest.mark.parametrize("configured_dtype", [None, torch.float16]) +def test_identity_output_dtype_lifecycle(internal, configured_dtype): + """Per-call overrides, updates and out= agree without mutating defaults.""" + q = IdentityQuantizer(dtype=configured_dtype) + q.internal = internal + x = torch.randn(32, 64, device="cuda", dtype=torch.float32, requires_grad=True) + out = q.quantize(x, dtype=torch.bfloat16) + assert q.dtype == configured_dtype + assert out._dtype == out._hp_data.dtype == torch.bfloat16 + assert not out._hp_data.requires_grad + if not internal: + out.sum().backward() + torch.testing.assert_close(x.grad, torch.ones_like(x), rtol=0, atol=0) + pointer = out._hp_data.data_ptr() + updated = x.detach() * 3 + assert q.quantize(updated, out=out, dtype=torch.bfloat16) is out + assert out._hp_data.data_ptr() == pointer + torch.testing.assert_close(out.dequantize(), updated.bfloat16(), rtol=0, atol=0) + out.quantize_(updated * 2) + torch.testing.assert_close(out.dequantize(), (updated * 2).bfloat16(), rtol=0, atol=0) + before = out._hp_data.clone() + with pytest.raises(ValueError, match="dtype does not match"): + q.quantize(updated, out=out, dtype=torch.float32) + torch.testing.assert_close(out._hp_data, before, rtol=0, atol=0) + + +def test_identity_weight_workspace_dtype_change(): + """A module rebuilds an Identity workspace when its compute dtype changes.""" + from transformer_engine.pytorch.module.base import quantize_weight + + q = IdentityQuantizer() + x = torch.randn(32, 64, device="cuda", dtype=torch.float32) + out, cache = quantize_weight(tensor=x, quantizer=q, workspace_dtype=torch.bfloat16, cache=True) + pointer = out._hp_data.data_ptr() + updated, new_cache = quantize_weight( + tensor=x * 2, quantizer=q, workspace=cache, workspace_dtype=torch.bfloat16, cache=True + ) + assert updated is out and new_cache is None + assert updated._hp_data.data_ptr() == pointer + torch.testing.assert_close(updated.dequantize(), (x * 2).bfloat16(), rtol=0, atol=0) + changed, new_cache = quantize_weight( + tensor=x, quantizer=q, workspace=cache, workspace_dtype=torch.float32, cache=True + ) + assert new_cache is changed and changed is not out + torch.testing.assert_close(changed.dequantize(), x, rtol=0, atol=0) + + +@pytest.mark.parametrize("internal", [False, True]) +def test_identity_dtype_update_graph_noop(internal): + """Captured FP32-to-BF16 updates retain storage and obey a device noop.""" + q = IdentityQuantizer() + q.internal = internal + x = torch.randn(32, 64, device="cuda", dtype=torch.float32) + out = q.quantize(x, dtype=torch.bfloat16) + pointer = out._hp_data.data_ptr() + noop = torch.zeros(1, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + q.update_quantized(x, out, noop_flag=noop) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + q.update_quantized(x, out, noop_flag=noop) + x.mul_(3) + graph.replay() + torch.testing.assert_close(out.dequantize(), x.bfloat16(), rtol=0, atol=0) + before = out._hp_data.clone() + noop.fill_(1) + x.mul_(2) + graph.replay() + assert out._hp_data.data_ptr() == pointer + torch.testing.assert_close(out.dequantize(), before, rtol=0, atol=0) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 9ead1e73651..955e293a2cc 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -863,7 +863,7 @@ def quantize_weight( fsdp_group: dist_group_type, optional FSDP process group the weights are distributed over. workspace_dtype: torch.dtype, optional - High-precision dtype for debug quantization workspaces. + Requested compute dtype for high-precision quantization workspaces. cache: bool, default = False If ``True`` and a new workspace is created, it will be returned as the second element so the caller can store it. @@ -893,6 +893,11 @@ def quantize_weight( if workspace is not None and quantizer is not None: if not _is_weight_workspace_valid(workspace, quantizer): workspace = None + elif workspace_dtype is not None: + # Custom quantizers may honor a compute dtype. Native quantizers + # retain their existing dtype and cache behavior. + if not quantizer.is_compatible_output(workspace, dtype=workspace_dtype): + workspace = None # FSDP gather on cached workspace if ( diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index f6f2794bdf9..997426cc6a1 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -10,6 +10,7 @@ import enum import warnings import math +from functools import partial import torch from torch.utils._pytree import tree_map @@ -416,6 +417,22 @@ class Quantizer(abc.ABC): """ optimize_for_gemm: bool + # Only high-precision representations opt into output-dtype requests. + # Native quantizers keep ignoring dtype, without casting their inputs. + supports_output_dtype: bool = False + + def with_output_dtype(self, dtype: torch.dtype) -> Quantizer: + """Bind an output default on a copy for participating custom quantizers.""" + if not self.supports_output_dtype: + return self + quantizer = self.copy() + quantizer.dtype = dtype + return quantizer + + def is_compatible_output(self, out: QuantizedTensorStorage, *, dtype=None) -> bool: + """Native outputs accept any request; custom outputs must match its dtype.""" + return not self.supports_output_dtype or dtype is None or out._dtype == dtype + def __init__(self, *, rowwise: bool, columnwise: bool) -> None: self.rowwise_usage = rowwise self.columnwise_usage = columnwise @@ -448,15 +465,25 @@ def quantize( tensor: torch.Tensor, *, out: Optional[QuantizedTensor] = None, - dtype: Optional[torch.dtype] = None, # pylint: disable=unused-argument # used by override + dtype: Optional[torch.dtype] = None, ) -> QuantizedTensor: - """Quantize tensor""" + """Quantize tensor, honoring output dtype only for opted-in quantizers. + + Native quantizers ignore ``dtype``. High precision quantization implementations + accept it in ``quantize_impl``; in-place updates retain destination dtype. + The source tensor is never cast here. + """ if out is not None: + if self.supports_output_dtype and not self.is_compatible_output(out, dtype=dtype): + raise ValueError("Requested output dtype does not match the destination") return self.update_quantized(tensor, out) + impl = self.quantize_impl + if self.supports_output_dtype and dtype is not None: + impl = partial(impl, dtype=dtype) if (not self.internal) and torch.is_grad_enabled(): - result = _QuantizeFunc.apply(tensor, self.quantize_impl) + result = _QuantizeFunc.apply(tensor, impl) else: - result = _QuantizeFunc.forward(None, tensor, self.quantize_impl) + result = _QuantizeFunc.forward(None, tensor, impl) # The amax reduction group must never persist on a tensor's quantizer result_quantizer = getattr(result, "_quantizer", None) if getattr(result_quantizer, "with_amax_reduction", False): diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 26d0798b92b..4d0af47f905 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -82,6 +82,8 @@ class HybridQuantizer(Quantizer): columnwise_quantizer: Quantizer columnwise_source: Literal["original", "rowwise_dequantized"] + supports_output_dtype = True + def __init__( self, *, @@ -157,6 +159,15 @@ def copy(self) -> "HybridQuantizer": quantizer.optimize_for_gemm = self.optimize_for_gemm return quantizer + def with_output_dtype(self, dtype): + """Bind a compute dtype only on children that opt into that contract.""" + quantizer = self.copy() + for name in ("rowwise_quantizer", "columnwise_quantizer"): + child = getattr(quantizer, name) + if child.supports_output_dtype: + setattr(quantizer, name, child.with_output_dtype(dtype)) + return quantizer + @property def with_amax_reduction(self) -> bool: """Whether either sub-quantizer has cross-rank amax reduction enabled.""" @@ -198,37 +209,50 @@ def _columnwise_src_from_rowwise( self, tensor: torch.Tensor, rowwise_result: Optional[Any], + *, + rowwise_quantizer=None, + dtype=None, ) -> torch.Tensor: if rowwise_result is None: - rowwise_result = self.rowwise_quantizer.quantize(tensor) + quantizer = rowwise_quantizer or self.rowwise_quantizer + rowwise_result = quantizer.quantize(tensor, dtype=dtype) + # Decode into the source dtype, preserving any prior Identity/QDQ rounding. return rowwise_result.dequantize(dtype=tensor.dtype) - def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: - # Gate each sub-quantizer call on the parent usage flag. Sub-quantizers - # are pinned to one direction in ``__init__``; the parent flag decides - # whether to invoke them. - rowwise_result = self.rowwise_quantizer.quantize(tensor) if self.rowwise_usage else None + def quantize_impl(self, tensor: torch.Tensor, *, dtype=None) -> QuantizedTensor: + # A result owns the child configuration needed to reconstruct an absent + # rowwise source during updates. No dtype state is kept on the parent. + quantizer = self if dtype is None else self.with_output_dtype(dtype) + rowwise_result = ( + quantizer.rowwise_quantizer.quantize(tensor, dtype=dtype) + if self.rowwise_usage + else None + ) columnwise_src = tensor if self.columnwise_usage and self.columnwise_source == "rowwise_dequantized": - columnwise_src = self._columnwise_src_from_rowwise(tensor, rowwise_result) + columnwise_src = quantizer._columnwise_src_from_rowwise( + tensor, rowwise_result, dtype=dtype + ) columnwise_result = ( - self.columnwise_quantizer.quantize(columnwise_src) if self.columnwise_usage else None + quantizer.columnwise_quantizer.quantize(columnwise_src, dtype=dtype) + if self.columnwise_usage + else None ) if self.internal: return HybridQuantizedTensorStorage( rowwise_storage=rowwise_result, columnwise_storage=columnwise_result, - quantizer=self, - fake_dtype=tensor.dtype, + quantizer=quantizer, + fake_dtype=dtype or tensor.dtype, ) return HybridQuantizedTensor( shape=tensor.shape, - dtype=tensor.dtype, + dtype=dtype or tensor.dtype, rowwise_storage=rowwise_result, columnwise_storage=columnwise_result, - quantizer=self, + quantizer=quantizer, ) def make_empty( @@ -240,19 +264,18 @@ def make_empty( requires_grad: bool = False, pin_memory: bool = False, ) -> Union["HybridQuantizedTensor", HybridQuantizedTensorStorage]: - # Mirror ``quantize_impl``: invoke each sub-quantizer with its own - # ``internal`` setting (no toggle), so the produced sub-storages have - # the same type that ``quantize_impl`` would produce via - # ``sub_quantizer.quantize(tensor)``. + # Bind the same child defaults as quantize(..., dtype=dtype), including + # a rowwise child that is absent but needed as a columnwise source. + quantizer = self.with_output_dtype(dtype) rowwise_empty = ( - self.rowwise_quantizer.make_empty( + quantizer.rowwise_quantizer.make_empty( shape, dtype=dtype, device=device, pin_memory=pin_memory ) if self.rowwise_usage else None ) columnwise_empty = ( - self.columnwise_quantizer.make_empty( + quantizer.columnwise_quantizer.make_empty( shape, dtype=dtype, device=device, pin_memory=pin_memory ) if self.columnwise_usage @@ -263,7 +286,7 @@ def make_empty( return HybridQuantizedTensorStorage( rowwise_storage=rowwise_empty, columnwise_storage=columnwise_empty, - quantizer=self, + quantizer=quantizer, fake_dtype=dtype, ) @@ -274,7 +297,7 @@ def make_empty( device=device, rowwise_storage=rowwise_empty, columnwise_storage=columnwise_empty, - quantizer=self, + quantizer=quantizer, ) def update_quantized( @@ -302,7 +325,9 @@ def update_quantized( columnwise_src = src if self.columnwise_source == "rowwise_dequantized": columnwise_src = self._columnwise_src_from_rowwise( - src, rowwise_result_for_columnwise + src, + rowwise_result_for_columnwise, + rowwise_quantizer=dst._get_quantizer().rowwise_quantizer, ) self.columnwise_quantizer.update_quantized( columnwise_src, dst._columnwise_storage, noop_flag=noop_flag diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index ec171564fe2..359999b8db1 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -45,6 +45,8 @@ class IdentityQuantizer(Quantizer): high-precision buffer serves both directions). """ + supports_output_dtype = True + def __init__( self, *, @@ -66,7 +68,7 @@ def copy(self) -> "IdentityQuantizer": quantizer.optimize_for_gemm = self.optimize_for_gemm return quantizer - def _maybe_cast(self, tensor: torch.Tensor) -> torch.Tensor: + def _maybe_cast(self, tensor: torch.Tensor, *, dtype=None) -> torch.Tensor: # Detach so the held buffer is plain "data" with no autograd graph edge, # mirroring the real quantizers (whose quantize kernels emit fresh, # non-differentiable tensors). Autograd connectivity for the *quantize* @@ -76,12 +78,13 @@ def _maybe_cast(self, tensor: torch.Tensor) -> torch.Tensor: # weight workspace returned across the module Function boundary), which # creates a spurious empty grad edge. out = tensor.detach() - if self.dtype is not None and out.dtype != self.dtype: - return out.to(self.dtype) + dtype = self.dtype if dtype is None else dtype + if dtype is not None and out.dtype != dtype: + return out.to(dtype) return out - def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensorStorage: - data = self._maybe_cast(tensor) + def quantize_impl(self, tensor: torch.Tensor, *, dtype=None) -> QuantizedTensorStorage: + data = self._maybe_cast(tensor, dtype=dtype) if self.internal: return IdentityTensorStorage( hp_data=data, @@ -144,22 +147,17 @@ def update_quantized( raise ValueError( f"IdentityQuantizer can only update IdentityTensorStorage, got {type(dst).__name__}" ) - data = self._maybe_cast(src) - if ( - dst._hp_data is not None - and dst._hp_data.shape == data.shape - and dst._hp_data.dtype == data.dtype - and dst._hp_data.device == data.device - ): - if noop_flag is None: - dst._hp_data.copy_(data) - else: - torch.where(noop_flag == 0, data, dst._hp_data, out=dst._hp_data) + if dst._hp_data is None: + raise ValueError("IdentityQuantizer requires an allocated destination") + if src.shape != dst._hp_data.shape or src.device != dst._hp_data.device: + raise ValueError("Identity source and destination must have matching shape/device") + # The destination owns the representation, including when the source + # or this quantizer's default dtype differs from the original call. + data = self._maybe_cast(src, dtype=dst._hp_data.dtype) + if noop_flag is None: + dst._hp_data.copy_(data) else: - if noop_flag is not None and noop_flag.item() != 0: - return dst - dst._hp_data = data.detach() - dst._dtype = data.dtype + torch.where(noop_flag == 0, data, dst._hp_data, out=dst._hp_data) return dst def calibrate(self, tensor: torch.Tensor) -> None: From 359f6d37473fdcc6666097b4082df6e740a8dcbc Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 6 Oct 2026 12:38:55 +0000 Subject: [PATCH 2/9] Add experimental NVFP4 weight QDQ quantization Signed-off-by: Evgeny --- .../fine_grained_quantization.rst | 32 ++ .../test.sh | 2 + tests/cpp/operator/CMakeLists.txt | 1 + tests/cpp/operator/test_nvfp4_qdq.cu | 74 ++++ tests/pytorch/distributed/run_nvfp4_qdq.py | 85 +++++ tests/pytorch/distributed/test_nvfp4_qdq.py | 34 ++ tests/pytorch/nvfp4/test_nvfp4_qdq.py | 358 ++++++++++++++++++ transformer_engine/common/CMakeLists.txt | 1 + transformer_engine/common/cast/nvfp4/qdq.cu | 103 +++++ .../common/include/transformer_engine/cast.h | 12 + transformer_engine/pytorch/csrc/extensions.h | 2 + .../pytorch/csrc/extensions/cast.cpp | 31 ++ .../pytorch/csrc/extensions/pybind.cpp | 3 + .../pytorch/custom_recipes/qdq.py | 202 ++++++++++ .../custom_recipes/quantizer_factory_zoo.py | 36 ++ 15 files changed, 976 insertions(+) create mode 100644 tests/cpp/operator/test_nvfp4_qdq.cu create mode 100644 tests/pytorch/distributed/run_nvfp4_qdq.py create mode 100644 tests/pytorch/distributed/test_nvfp4_qdq.py create mode 100644 tests/pytorch/nvfp4/test_nvfp4_qdq.py create mode 100644 transformer_engine/common/cast/nvfp4/qdq.cu create mode 100644 transformer_engine/pytorch/custom_recipes/qdq.py diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst index 650d30811ec..497d9e257cf 100644 --- a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -387,3 +387,35 @@ API reference See the :doc:`PyTorch API <../../../api/pytorch>` for ``QuantizerRole``, ``HybridQuantizer``, ``IdentityQuantizer``, and their returned tensor types. See the :doc:`Common API <../../../api/common>` for ``CustomRecipe``. + +Experimental NVFP4 weight QDQ +----------------------------- + +``custom_recipes.qdq.NVFP4QDQQuantizer`` is useful for experiments with +``X[BF16] x W[NVFP4]``: it quantizes weights to NVFP4, then dequantizes them +into a high-precision ``IdentityTensor`` for BF16 forward GEMMs. This preserves +NVFP4 weight rounding error while using BF16 GEMM arithmetic. + +The following factory supports ``Linear`` and ``GroupedLinear`` with discrete +expert parameters. Forward uses high-precision inputs and QDQ weights; +MXFP8 backward uses the original inputs and QDQ weights:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_qdq_weight_fwd_mxfp8_bwd_factory, + ) + + recipe = CustomRecipe(qfactory=nvfp4_qdq_weight_fwd_mxfp8_bwd_factory) + +Output defaults to the source dtype; a module compute-dtype override changes +only the decoded dtype. ``backend="auto"`` fuses plain 1D E2M1/E4M3 QDQ on +SM100 for contiguous, 16-byte-aligned BF16/FP16 tensors with matching output +dtype and flattened matrix dimensions divisible by 16, with ordinary fast math +disabled. Other native-valid modes, including 4over6, 2D scaling, distributed +amax reduction, and mixed dtypes, use native ``quantize().dequantize()``. +``backend="reference"`` forces that composition; ``backend="fused"`` rejects +unsupported cases. Current weight amax is computed when quantization runs. + +Configure ``backend`` or ``nvfp4_options`` on the factory with +``functools.partial``. Existing Hybrid limitations still apply, and RL +convergence has not been validated. diff --git a/qa/L1_pytorch_hybrid_distributed_unittest/test.sh b/qa/L1_pytorch_hybrid_distributed_unittest/test.sh index a2c06266b9b..4c44ee41fa4 100644 --- a/qa/L1_pytorch_hybrid_distributed_unittest/test.sh +++ b/qa/L1_pytorch_hybrid_distributed_unittest/test.sh @@ -25,6 +25,8 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2_hybrid.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py -k "hybrid" || test_fail "hybrid test_torch_fsdp2.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_hybrid_tp_sp.xml $TE_PATH/tests/pytorch/distributed/test_hybrid_tp_sp.py || test_fail "test_hybrid_tp_sp.py" +python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_nvfp4_qdq_distributed.xml $TE_PATH/tests/pytorch/distributed/test_nvfp4_qdq.py || test_fail "test_nvfp4_qdq.py" + if [ "$RET" -ne 0 ]; then echo "Error in the following test cases:$FAILED_CASES" exit 1 diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index c5acd043363..87d65e790e9 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -22,6 +22,7 @@ add_executable(test_operator test_dequantize_mxfp8_grouped.cu test_fused_group_requantize_mxfp8.cu test_dequantize_float8blockwise_grouped.cu + test_nvfp4_qdq.cu test_dequantize_nvfp4.cu test_transpose.cu test_cast_transpose.cu diff --git a/tests/cpp/operator/test_nvfp4_qdq.cu b/tests/cpp/operator/test_nvfp4_qdq.cu new file mode 100644 index 00000000000..98f6551d6f9 --- /dev/null +++ b/tests/cpp/operator/test_nvfp4_qdq.cu @@ -0,0 +1,74 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include + +#include +#include +#include + +#include "../test_common.h" + +using namespace transformer_engine; +using namespace test; + +namespace { +template +void check_qdq(size_t rows, size_t cols) { + cudaDeviceProp prop; + ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); + if (prop.major != 10 || prop.minor != 0) GTEST_SKIP() << "QDQ targets SM100"; + const auto dtype = TypeInfo::dtype; + const std::vector shape{rows, cols}; + Tensor input("input", shape, dtype); + Tensor packed("packed", shape, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING, + DType::kFloat8E4M3); + Tensor reference("reference", shape, dtype); + Tensor output("output", shape, dtype); + Tensor amax("amax", std::vector{1}, DType::kFloat32); + Tensor noop("noop", std::vector{1}, DType::kFloat32); + fillCase(&input, InputsFillCase::uniform); + input.to_cpu(); + float maximum = 0; + for (size_t i = 0; i < rows * cols; ++i) { + maximum = std::max(maximum, std::abs(static_cast(input.rowwise_cpu_dptr()[i]))); + } + packed.set_amax(maximum); + ASSERT_EQ(cudaMemcpy(amax.rowwise_dptr(), &maximum, sizeof(float), cudaMemcpyHostToDevice), + cudaSuccess); + QuantizationConfigWrapper config; + nvte_quantize_v2(input.data(), packed.data(), config, 0); + nvte_dequantize(packed.data(), reference.data(), 0); + nvte_nvfp4_qdq(input.data(), output.data(), amax.data(), nullptr, 0); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + reference.to_cpu(); + output.to_cpu(); + EXPECT_EQ(std::memcmp(reference.rowwise_cpu_dptr(), output.rowwise_cpu_dptr(), + rows * cols * sizeof(T)), 0); + const float one = 1; + ASSERT_EQ(cudaMemcpy(noop.rowwise_dptr(), &one, sizeof(float), cudaMemcpyHostToDevice), + cudaSuccess); + ASSERT_EQ(cudaMemset(input.rowwise_dptr(), 0, rows * cols * sizeof(T)), cudaSuccess); + nvte_nvfp4_qdq(input.data(), output.data(), amax.data(), noop.data(), 0); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + output.to_cpu(); + EXPECT_EQ(std::memcmp(reference.rowwise_cpu_dptr(), output.rowwise_cpu_dptr(), + rows * cols * sizeof(T)), 0); +} +} // namespace + +TEST(NVFP4QDQ, BF16) { + check_qdq(32, 32); + check_qdq(96, 160); + check_qdq(512, 1024); +} + +TEST(NVFP4QDQ, FP16) { + check_qdq(32, 32); + check_qdq(96, 160); +} diff --git a/tests/pytorch/distributed/run_nvfp4_qdq.py b/tests/pytorch/distributed/run_nvfp4_qdq.py new file mode 100644 index 00000000000..4922dffa16a --- /dev/null +++ b/tests/pytorch/distributed/run_nvfp4_qdq.py @@ -0,0 +1,85 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Two-rank NVFP4 QDQ amax reduction and native reference parity.""" + +import os + +import torch +import torch.distributed as dist +from transformer_engine.pytorch.custom_recipes.qdq import NVFP4QDQQuantizer +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer +import transformer_engine_torch as tex + + +def assert_bits(actual, expected): + assert actual.dtype == expected.dtype + assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)) + + +def run_case(options, source_dtype): + rank = dist.get_rank() + torch.manual_seed(123 + rank) + x = torch.randn(128, 256, device="cuda", dtype=source_dtype) * (1 if rank == 0 else 17) + output_dtype = torch.bfloat16 + local_amax = x.float().abs().amax().reshape(1) + global_amax = local_amax.clone() + dist.all_reduce(global_amax, op=dist.ReduceOp.MAX) + assert dist.get_world_size() == 2 + if rank == 0: + assert global_amax.item() > local_amax.item() + native = NVFP4Quantizer(rowwise=True, columnwise=False, **options) + native.internal = True + q = NVFP4QDQQuantizer(native) + q.amax_reduction_group = dist.group.WORLD + q.with_amax_reduction = True + assert q.selected_backend(x, output_dtype) == "reference" + out = q.quantize(x, dtype=output_dtype) + expected = tex.nvfp4_quantize_with_amax(x, native, global_amax, global_amax).dequantize( + dtype=output_dtype + ) + assert_bits(out.dequantize(), expected) + saved = out._get_quantizer() + assert not saved.with_amax_reduction and saved.amax_reduction_group is None + assert q.with_amax_reduction and q.amax_reduction_group is dist.group.WORLD + pointer = out._hp_data.data_ptr() + x.mul_(2) + global_amax.mul_(2) + q.update_quantized(x, out) + expected = tex.nvfp4_quantize_with_amax(x, native, global_amax, global_amax).dequantize( + dtype=output_dtype + ) + assert_bits(out.dequantize(), expected) + assert out._hp_data.data_ptr() == pointer + # Saved workspaces intentionally update locally, without persistent collectives. + out.quantize_(x) + assert_bits(out.dequantize(), native(x).dequantize(dtype=output_dtype)) + + +def main(): + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group("nccl") + try: + count = 0 + for options in ( + {}, + {"with_2d_quantization": True}, + { + "nvfp4_use_4over6": True, + "nvfp4_e4m3_max": 448, + "nvfp4_4over6_err_mode": "MSE", + }, + ): + for source_dtype in (torch.bfloat16, torch.float32): + run_case(options, source_dtype) + count += 1 + dist.barrier() + if dist.get_rank() == 0: + print(f"NVFP4 QDQ distributed: {count} cases passed on 2 ranks", flush=True) + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/pytorch/distributed/test_nvfp4_qdq.py b/tests/pytorch/distributed/test_nvfp4_qdq.py new file mode 100644 index 00000000000..4fffc197290 --- /dev/null +++ b/tests/pytorch/distributed/test_nvfp4_qdq.py @@ -0,0 +1,34 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pytest launcher for distributed NVFP4 QDQ reference parity.""" + +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch +import transformer_engine.pytorch as te + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="QDQ amax reduction requires two GPUs") +def test_nvfp4_qdq_amax_reduction(): + available, reason = te.is_nvfp4_available(return_reason=True) + if not available: + pytest.skip(reason) + script = Path(__file__).with_name("run_nvfp4_qdq.py") + subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nproc_per_node=2", + str(script), + ], + env=os.environ.copy(), + check=True, + ) diff --git a/tests/pytorch/nvfp4/test_nvfp4_qdq.py b/tests/pytorch/nvfp4/test_nvfp4_qdq.py new file mode 100644 index 00000000000..324388ff698 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_qdq.py @@ -0,0 +1,358 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Native parity, lifecycle, and operand-selection tests for weight QDQ.""" + +from functools import partial + +import pytest +import torch +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer +from transformer_engine.pytorch.custom_recipes.qdq import NVFP4QDQQuantizer +from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_qdq_weight_fwd_mxfp8_bwd_factory, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), + reason="QDQ fused implementation targets SM100", +) + + +def native_qdq(x, dtype=None, **options): + q = NVFP4Quantizer(rowwise=True, columnwise=False, **options) + q.internal = True + return q(x).dequantize(dtype=dtype or x.dtype) + + +def assert_bits(actual, expected): + assert actual.dtype == expected.dtype + assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("shape", [(16, 16), (48, 80), (128, 256), (512, 1024)]) +@pytest.mark.parametrize("distribution", ["normal", "zeros", "outliers", "boundaries"]) +def test_plain_exact(dtype, shape, distribution, monkeypatch): + monkeypatch.setenv("NVTE_USE_FAST_MATH", "0") + torch.manual_seed(42) + x = torch.randn(shape, device="cuda", dtype=dtype) + if distribution == "zeros": + x.zero_() + x[:, ::2] = -0.0 + elif distribution == "outliers": + x[0, 0] = 4096 + x[1, :] *= 1e-5 + elif distribution == "boundaries": + values = torch.tensor( + [ + 0, + -0.0, + 0.25, + 0.75, + 1.25, + 1.75, + 2.5, + 3.5, + 5.0, + 6.0, + -6.0, + -5.0, + -2.5, + 1e-7, + 1e-4, + 448.0, + ], + device="cuda", + dtype=dtype, + ) + x.copy_(values.repeat(shape[0], shape[1] // 16)) + q = NVFP4QDQQuantizer(backend="fused") + assert q.selected_backend(x) == "fused" + assert_bits(q(x).dequantize(), native_qdq(x)) + + +@pytest.mark.parametrize("mode", ["MAE", "MSE"]) +@pytest.mark.parametrize("maximum", [256, 448]) +@pytest.mark.parametrize("fast_error", ["0", "1"]) +def test_4over6_fallback(mode, maximum, fast_error, monkeypatch): + monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", fast_error) + x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) + options = dict(nvfp4_use_4over6=True, nvfp4_e4m3_max=maximum, nvfp4_4over6_err_mode=mode) + q = NVFP4QDQQuantizer(NVFP4Quantizer(**options)) + assert q.selected_backend(x) == "reference" + assert_bits(q(x).dequantize(), native_qdq(x, **options)) + q.backend = "fused" + with pytest.raises(ValueError, match="Forced fused"): + q(x) + + +@pytest.mark.parametrize("options", [{"with_2d_quantization": True}, {}]) +def test_dtype_fallback(options): + x = torch.randn(128, 256, device="cuda", dtype=torch.float32) + q = NVFP4QDQQuantizer(NVFP4Quantizer(**options)) + actual = q.quantize(x, dtype=torch.bfloat16) + assert_bits(actual.dequantize(), native_qdq(x, torch.bfloat16, **options)) + assert q.dtype is None + assert actual.dtype == torch.bfloat16 + actual.quantize_(x * 2) + assert_bits(actual.dequantize(), native_qdq(x * 2, torch.bfloat16, **options)) + + +@pytest.mark.parametrize("backend", ["reference", "fused"]) +@pytest.mark.parametrize("internal", [False, True]) +def test_lifecycle_and_graph(backend, internal): + x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) + q = NVFP4QDQQuantizer(backend=backend) + q.internal = internal + y = q(x) + assert y._get_quantizer().backend == backend + address = y._hp_data.data_ptr() + noop = torch.ones(1, device="cuda") + q.update_quantized(x * 3, y, noop_flag=noop) + assert_bits(y.dequantize(), native_qdq(x)) + noop.zero_() + q.update_quantized(x * 2, y, noop_flag=noop) + assert y._hp_data.data_ptr() == address + assert_bits(y.dequantize(), native_qdq(x * 2)) + # Warm up on a side stream, then replay with changed input and noop flag. + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + q.update_quantized(x, y, noop_flag=noop) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + q.update_quantized(x, y, noop_flag=noop) + x.mul_(2) + graph.replay() + assert_bits(y.dequantize(), native_qdq(x)) + previous = y.dequantize().clone() + noop.fill_(1) + x.mul_(2) + graph.replay() + assert_bits(y.dequantize(), previous) + + +def test_ste_and_rng(): + x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16, requires_grad=True) + q = NVFP4QDQQuantizer(backend="reference") + q(x).sum().backward() + assert torch.equal(x.grad, torch.ones_like(x)) + q = NVFP4QDQQuantizer(NVFP4Quantizer(stochastic_rounding=True)) + assert not q.is_requantization_safe() + state = torch.cuda.get_rng_state() + actual = q(x).dequantize() + after = torch.cuda.get_rng_state() + torch.cuda.set_rng_state(state) + expected = native_qdq(x, stochastic_rounding=True) + assert_bits(actual, expected) + assert torch.equal(after, torch.cuda.get_rng_state()) + + +def mx_value(x, *, columnwise=False): + q = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=not columnwise, columnwise=columnwise) + return q(x).dequantize() + + +@pytest.mark.parametrize("backend", ["reference", "fused"]) +@pytest.mark.parametrize("grouped", [False, True]) +@pytest.mark.parametrize("params_dtype", [torch.bfloat16, torch.float32]) +def test_module_operands(backend, grouped, params_dtype, nvfp4_options=None): + if backend == "fused" and params_dtype == torch.float32: + pytest.skip("Mixed FP32 source/BF16 decode intentionally uses native fallback") + torch.manual_seed(12) + factory = nvfp4_qdq_weight_fwd_mxfp8_bwd_factory + recipe = CustomRecipe(qfactory=partial(factory, backend=backend, nvfp4_options=nvfp4_options)) + kwargs = dict(bias=True, params_dtype=params_dtype, device="cuda") + layer = ( + te.GroupedLinear(2, 128, 128, single_grouped_weight=False, **kwargs) + if grouped + else te.Linear(128, 128, **kwargs) + ) + splits = [32, 64] if grouped else [96] + x = torch.randn(sum(splits), 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + dy = torch.randn_like(x) + weights = [getattr(layer, f"weight{i}") for i in range(2)] if grouped else [layer.weight] + biases = [getattr(layer, f"bias{i}") for i in range(2)] if grouped else [layer.bias] + with te.autocast(recipe=recipe), torch.autocast("cuda", dtype=torch.bfloat16): + y = layer(x, splits) if grouped else layer(x) + y.backward(dy) + expected_y, expected_dx = [], [] + for w, b, xs, gs in zip(weights, biases, x.detach().split(splits), dy.split(splits)): + wh = native_qdq(w.detach(), torch.bfloat16, **(nvfp4_options or {})) + expected_y.append(torch.nn.functional.linear(xs, wh, b.detach().to(torch.bfloat16))) + bw = mx_value(wh, columnwise=True) + bx = mx_value(xs, columnwise=True) + grow = mx_value(gs) + gcol = mx_value(gs, columnwise=True) + expected_dx.append(grow @ bw) + dw = gcol.T @ bx + torch.testing.assert_close(w.grad.float(), dw.float(), rtol=0.016, atol=0.002) + torch.testing.assert_close( + b.grad.float(), + gs.float().sum(0).to(b.grad.dtype).float(), + rtol=0.016, + atol=0.002, + ) + torch.testing.assert_close(y, torch.cat(expected_y), rtol=0.016, atol=0.002) + torch.testing.assert_close(x.grad, torch.cat(expected_dx), rtol=0.016, atol=0.002) + + +def test_weight_workspace_dtype_and_reuse(): + from transformer_engine.pytorch.module.base import quantize_weight + from transformer_engine.pytorch.quantization import QuantizerRole + + factory = nvfp4_qdq_weight_fwd_mxfp8_bwd_factory + q = factory(QuantizerRole(module_type="linear", tensor_type="weight")) + w = torch.randn(128, 256, device="cuda", dtype=torch.float32) + out, cache = quantize_weight(tensor=w, quantizer=q, workspace_dtype=torch.bfloat16, cache=True) + expected = native_qdq(w, torch.bfloat16) + assert_bits(out._rowwise_storage.dequantize(), expected) + pointer = out._rowwise_storage._hp_data.data_ptr() + w.mul_(3) + out2, new = quantize_weight( + tensor=w, + quantizer=q, + workspace=cache, + workspace_dtype=torch.bfloat16, + cache=True, + ) + assert new is None and out2 is cache + assert out2._rowwise_storage._hp_data.data_ptr() == pointer + expected = native_qdq(w, torch.bfloat16) + assert_bits(out2._rowwise_storage.dequantize(), expected) + # Hybrid retains native decode-to-source-dtype behavior. Widening the + # BF16 QDQ value preserves its rounding while MXFP8 metadata stays FP32. + expected_col = mx_value(expected.to(w.dtype), columnwise=True) + assert_bits(out2._columnwise_storage.dequantize(), expected_col) + before = out2._rowwise_storage.dequantize().clone() + w.mul_(2) + quantize_weight( + tensor=w, + quantizer=q, + workspace=cache, + update_workspace=False, + workspace_dtype=torch.bfloat16, + ) + assert_bits(out2._rowwise_storage.dequantize(), before) + out3, cache3 = quantize_weight( + tensor=w, + quantizer=q, + workspace=cache, + workspace_dtype=torch.float32, + cache=True, + ) + assert cache3 is not cache and out3._rowwise_storage.dequantize().dtype == torch.float32 + torch.testing.assert_close(out3._rowwise_storage.dequantize(), native_qdq(w), rtol=0, atol=0) + + +def test_noncontiguous_and_fast_math_fallback(monkeypatch): + x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16).T + q = NVFP4QDQQuantizer() + assert q.selected_backend(x) == "reference" + assert_bits(q(x).dequantize(), native_qdq(x)) + monkeypatch.setenv("NVTE_USE_FAST_MATH", "1") + x = x.contiguous() + assert q.selected_backend(x) == "reference" + assert_bits(q(x).dequantize(), native_qdq(x)) + + +@pytest.mark.parametrize("shape", [(0, 128), (128, 0)]) +def test_empty_reference(shape): + x = torch.empty(shape, device="cuda", dtype=torch.bfloat16) + q = NVFP4QDQQuantizer() + assert q.selected_backend(x) == "reference" + assert_bits(q(x).dequantize(), native_qdq(x)) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")]) +def test_exceptional_values(value): + x = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16) + x[0, 0] = value + actual = NVFP4QDQQuantizer(backend="fused")(x).dequantize() + expected = native_qdq(x) + assert torch.equal(torch.isnan(actual), torch.isnan(expected)) + mask = ~torch.isnan(expected) + assert_bits(actual[mask], expected[mask]) + + +def test_current_stream_and_device(): + device = 1 if torch.cuda.device_count() > 1 else 0 + with torch.cuda.device(device), torch.cuda.stream(torch.cuda.Stream(device=device)): + x = torch.randn(48, 80, device=f"cuda:{device}", dtype=torch.bfloat16) + assert_bits(NVFP4QDQQuantizer(backend="fused")(x).dequantize(), native_qdq(x)) + + +def test_copy_preserves_runtime_reduction_state(): + q = NVFP4QDQQuantizer(backend="reference") + group = object() + q.amax_reduction_group = group + q.with_amax_reduction = True + copied = q.copy() + assert copied.amax_reduction_group is group and copied.with_amax_reduction + copied.nvfp4_quantizer.stochastic_rounding = True + copied.with_amax_reduction = False + assert not q.nvfp4_quantizer.stochastic_rounding + assert q.with_amax_reduction + + +@pytest.mark.parametrize("grouped", [False, True]) +def test_4over6_modules(grouped): + test_module_operands( + "reference", + grouped, + torch.bfloat16, + nvfp4_options={ + "nvfp4_use_4over6": True, + "nvfp4_e4m3_max": 448, + "nvfp4_4over6_err_mode": "MSE", + }, + ) + + +@pytest.mark.parametrize("backend", ["reference", "fused"]) +def test_qdq_explicit_dtype_and_destination(backend): + """Output overrides must not inherit the quantizer's conflicting default.""" + q = NVFP4QDQQuantizer(dtype=torch.float16, backend=backend) + x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) + out = q.quantize(x, dtype=torch.bfloat16) + assert q.dtype == torch.float16 + assert_bits(out.dequantize(), native_qdq(x)) + pointer = out._hp_data.data_ptr() + q.quantize(x * 2, out=out, dtype=torch.bfloat16) + assert out._hp_data.data_ptr() == pointer + assert_bits(out.dequantize(), native_qdq(x * 2)) + with pytest.raises(ValueError, match="dtype does not match"): + q.quantize(x, out=out, dtype=torch.float16) + + +@pytest.mark.parametrize("empty", [False, True]) +def test_qdq_columnwise_only_workspace_dtype(empty): + """A transient forward QDQ source must round to the requested compute dtype.""" + from transformer_engine.pytorch.quantization import QuantizerRole + + q = nvfp4_qdq_weight_fwd_mxfp8_bwd_factory( + QuantizerRole(module_type="linear", tensor_type="weight") + ) + q.set_usage(rowwise=False, columnwise=True) + x = torch.randn(128, 256, device="cuda", dtype=torch.float32) + out = ( + q.make_empty(x.shape, dtype=torch.bfloat16, device=x.device) + if empty + else q.quantize(x, dtype=torch.bfloat16) + ) + x.mul_(3) + q.update_quantized(x, out) + assert out._rowwise_storage is None + expected = mx_value(native_qdq(x, torch.bfloat16).float(), columnwise=True) + # Native make_empty honors its allocation dtype; native quantize retains + # the FP32 source dtype. Preserve both behaviors and compare equal decodes. + assert out._columnwise_storage._dtype == (torch.bfloat16 if empty else torch.float32) + assert_bits(out._columnwise_storage.dequantize(dtype=torch.float32), expected) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 529a5548a5a..8b6fd7df7f1 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -358,6 +358,7 @@ list(APPEND transformer_engine_cuda_arch_specific_sources activation/swiglu_grouped.cu activation/swiglu_grouped_bwd.cu activation/swiglu_grouped_dbias.cu + cast/nvfp4/qdq.cu cast/cast.cu cast/cast_dbias.cu cast/cast_grouped.cu diff --git a/transformer_engine/common/cast/nvfp4/qdq.cu b/transformer_engine/common/cast/nvfp4/qdq.cu new file mode 100644 index 00000000000..a7d884cbdc6 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/qdq.cu @@ -0,0 +1,103 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include "../../common.h" +#include "../../util/cuda_runtime.h" +#include "core_nvfp4.cuh" + +namespace transformer_engine { +namespace { +#if FP4_TYPE_SUPPORTED +// One thread owns a complete 1x16 scale block. Both rounded encodings stay in +// registers; only the decoded high-precision values are written to memory. +template +__global__ void qdq_kernel(const T *input, T *output, const float *amax, const float *noop, + size_t blocks) { + if (noop != nullptr && *noop == 1.0f) return; + const size_t block = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (block >= blocks) return; + Vec values; + values.load_from(input + block * 16); + float local_amax = 0.0f; +#pragma unroll + for (int i = 0; i < 16; ++i) { + local_amax = fmaxf(local_amax, fabsf(static_cast(values.data.elt[i]))); + } + using namespace dispatch::nvfp4; + const float global_amax = *amax; + const float encode = core::compute_global_encode_scaling_factor_FP4(global_amax); + const fp8e4m3 scale = core::compute_decoding_scaling_factor(local_amax, encode); + const float inverse = + fminf(1.0f / (static_cast(scale) * (1.0f / encode)), detail::TypeExtrema::max); + // Match native dequantize multiplication order, including zero signs. + const float decode = static_cast(scale) * global_amax * (1.0f / (6.0f * 448.0f)); +#pragma unroll + for (int i = 0; i < 16; i += 4) { + fp4e2m1x4 packed; + ptx::mul_cvt_4x(packed, + make_float2(static_cast(values.data.elt[i]), + static_cast(values.data.elt[i + 1])), + make_float2(static_cast(values.data.elt[i + 2]), + static_cast(values.data.elt[i + 3])), + inverse); + const float4 decoded = static_cast(packed); + values.data.elt[i] = static_cast(decoded.x * decode); + values.data.elt[i + 1] = static_cast(decoded.y * decode); + values.data.elt[i + 2] = static_cast(decoded.z * decode); + values.data.elt[i + 3] = static_cast(decoded.w * decode); + } + values.store_to(output + block * 16); +} +#endif +} // namespace +} // namespace transformer_engine + +void nvte_nvfp4_qdq(const NVTETensor input, NVTETensor output, const NVTETensor amax, + const NVTETensor noop, cudaStream_t stream) { + NVTE_API_CALL(nvte_nvfp4_qdq); + using namespace transformer_engine; +#if FP4_TYPE_SUPPORTED + NVTE_CHECK(cuda::sm_arch() == 100, "NVFP4 QDQ currently supports SM100"); + const auto &in = *convertNVTETensorCheck(input); + auto &out = *convertNVTETensorCheck(output); + const auto &amax_tensor = *convertNVTETensorCheck(amax); + CheckInputTensor(in, "input"); + CheckOutputTensor(out, "output"); + CheckInputTensor(amax_tensor, "amax"); + NVTE_CHECK(in.data.dtype == out.data.dtype && in.data.shape == out.data.shape, + "QDQ input and output must have matching dtype and shape"); + NVTE_CHECK(in.data.dtype == DType::kBFloat16 || in.data.dtype == DType::kFloat16, + "QDQ supports BF16 and FP16"); + NVTE_CHECK(in.data.shape.size() >= 2 && in.numel() > 0 && in.data.shape.back() % 16 == 0 && + (in.numel() / in.data.shape.back()) % 16 == 0, + "QDQ requires nonempty block-aligned rowwise input"); + NVTE_CHECK(amax_tensor.data.dtype == DType::kFloat32 && amax_tensor.numel() == 1, + "QDQ requires a scalar FP32 amax"); + NVTE_CHECK(reinterpret_cast(in.data.dptr) % 16 == 0 && + reinterpret_cast(out.data.dptr) % 16 == 0, + "QDQ requires 16-byte aligned buffers"); + const float *noop_ptr = nullptr; + if (noop != nullptr) { + const auto &flag = *convertNVTETensorCheck(noop); + CheckInputTensor(flag, "noop"); + NVTE_CHECK(flag.data.dtype == DType::kFloat32 && flag.numel() == 1, + "QDQ noop must be a scalar FP32 tensor"); + noop_ptr = static_cast(flag.data.dptr); + } + constexpr int threads = 128; + const size_t blocks = in.numel() / 16; + TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT( + in.data.dtype, T, + qdq_kernel<<<(blocks + threads - 1) / threads, threads, 0, stream>>>( + static_cast(in.data.dptr), static_cast(out.data.dptr), + static_cast(amax_tensor.data.dptr), noop_ptr, blocks);); + NVTE_CHECK_CUDA(cudaGetLastError()); +#else + NVTE_ERROR("NVFP4 QDQ requires a CUDA toolkit with FP4 support"); +#endif +} diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index ca3eb13ed24..b32708c9e9f 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -17,6 +17,18 @@ extern "C" { #endif +/*! \brief Plain rowwise 1x16 NVFP4 QDQ into matching BF16/FP16 output. + * + * Input/output are contiguous, 16-byte aligned, rank >= 2, with both flattened + * dimensions divisible by 16. E2M1 values and E4M3 scales (maximum 448) are + * rounded as native quantization with fast math disabled. No packed output is + * materialized. Caller supplies the current scalar FP32 absolute maximum in + * amax's data buffer. Optional scalar FP32 noop skips writes when equal to 1. + * All buffers are caller-owned and must remain alive until stream completion. + */ +void nvte_nvfp4_qdq(const NVTETensor input, NVTETensor output, const NVTETensor amax, + const NVTETensor noop, cudaStream_t stream); + /* Quantize the tensor * * The type of quantized tensor in the output depends on the scaling mode of the output diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 955763f43b1..86342593a57 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -502,6 +502,8 @@ at::Tensor scaled_aligned_causal_masked_softmax_backward(at::Tensor output_grads * FP8 recipe **************************************************************************************************/ +void nvfp4_qdq(const at::Tensor &input, at::Tensor output, at::Tensor amax, + std::optional noop, bool compute_amax); void compute_amax(const at::Tensor &tensor, at::Tensor &amax); void fused_amax_and_scale_update_after_reduction(const at::Tensor &amax_reduction_buffer, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 06d03cd6e1c..b476fc58cc1 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -52,6 +52,37 @@ void allreduce_nvfp4_amax_tensors(NVFP4Quantizer *nvfp4_quantizer_cpp, } // namespace +void nvfp4_qdq(const at::Tensor &input, at::Tensor output, at::Tensor amax, + std::optional noop, bool compute_amax) { + TORCH_CHECK(input.is_cuda() && output.is_cuda() && amax.is_cuda(), + "NVFP4 QDQ requires CUDA tensors"); + TORCH_CHECK(input.device() == output.device() && input.device() == amax.device(), + "NVFP4 QDQ tensors must share a device"); + at::cuda::CUDAGuard guard(input.device()); + TORCH_CHECK(input.is_contiguous() && output.is_contiguous() && amax.is_contiguous(), + "NVFP4 QDQ requires contiguous tensors"); + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major == 10 && + at::cuda::getCurrentDeviceProperties()->minor == 0, + "NVFP4 QDQ currently supports SM100"); + TORCH_CHECK( + !noop || (noop->is_cuda() && noop->device() == input.device() && noop->is_contiguous()), + "Invalid QDQ noop device/layout"); + auto in = makeTransformerEngineTensor(input); + auto out = makeTransformerEngineTensor(output); + auto amax_cpp = makeTransformerEngineTensor(amax); + std::optional noop_cpp; + if (noop) noop_cpp = makeTransformerEngineTensor(*noop); + const auto stream = at::cuda::getCurrentCUDAStream(); + if (compute_amax) { + TORCH_CHECK(amax.scalar_type() == at::kFloat && amax.numel() == 1, + "QDQ amax must be scalar FP32"); + TensorWrapper reduction_output(nullptr, in.shape(), DType::kFloat32, amax.data_ptr()); + nvte_compute_amax(in.data(), reduction_output.data(), stream); + } + nvte_nvfp4_qdq(in.data(), out.data(), amax_cpp.data(), noop_cpp ? noop_cpp->data() : nullptr, + stream); +} + py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::object &output, std::optional noop_flag) { // Convert quantizer to C++ object diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 2046408eba8..9f52b1ecebb 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -477,6 +477,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()); m.def("get_fused_attn_backend", &transformer_engine::pytorch::get_fused_attn_backend, "Get Fused Attention backend", py::arg("fused_attn_params")); + m.def("nvfp4_qdq", &transformer_engine::pytorch::nvfp4_qdq, "Plain rowwise NVFP4 QDQ", + py::arg("input"), py::arg("output"), py::arg("amax"), py::arg("noop") = py::none(), + py::arg("compute_amax") = true, py::call_guard()); m.def("compute_amax", &transformer_engine::pytorch::compute_amax, "Compute absolute max value in tensor", py::arg("input"), py::arg("amax"), py::call_guard()); diff --git a/transformer_engine/pytorch/custom_recipes/qdq.py b/transformer_engine/pytorch/custom_recipes/qdq.py new file mode 100644 index 00000000000..bc95cac6815 --- /dev/null +++ b/transformer_engine/pytorch/custom_recipes/qdq.py @@ -0,0 +1,202 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Experimental quantize/dequantize quantizers with high-precision storage.""" + +from functools import lru_cache +import os + +import torch +import transformer_engine_torch as tex + +from ..constants import DType +from ..quantized_tensor import Quantizer +from ..tensor.identity_tensor import IdentityTensor +from ..tensor.storage.identity_tensor_storage import IdentityTensorStorage +from ..tensor.nvfp4_tensor import NVFP4Quantizer + + +@lru_cache(maxsize=None) +def _supports_fused_device(device_index): + """Device capabilities are fixed for the process lifetime.""" + return torch.cuda.get_device_capability(device_index) == (10, 0) + + +class NVFP4QDQQuantizer(Quantizer): + """Store native rowwise NVFP4 Q+DQ values in an IdentityTensor. + + ``quantizer`` supplies native NVFP4 options and is copied with rowwise-only + usage. Output defaults to the source dtype; ``dtype`` overrides only the + decode dtype, never the quantization input. The high-precision result can + serve either GEMM orientation. + + ``backend`` is ``auto``, ``reference``, or ``fused``. Auto falls back to + native Q+DQ outside plain contiguous BF16/FP16 1D E4M3 scaling on SM100. + In particular, 4over6 and 2D scaling currently use the reference path. + Forced fused execution raises for unsupported inputs/configurations. + This experimental API is subject to change. + """ + + # Unknown to native fused quantization producers; outputs use native HP GEMM. + custom = True + + supports_output_dtype = True + + def __init__( + self, + quantizer=None, + *, + dtype=None, + backend="auto", + rowwise=True, + columnwise=True, + ): + super().__init__(rowwise=rowwise, columnwise=columnwise) + if backend not in ("auto", "reference", "fused"): + raise ValueError("backend must be auto, reference, or fused") + if quantizer is not None and not isinstance(quantizer, NVFP4Quantizer): + raise TypeError("quantizer must be an NVFP4Quantizer") + self.nvfp4_quantizer = quantizer.copy() if quantizer is not None else NVFP4Quantizer() + self.nvfp4_quantizer.set_usage(rowwise=True, columnwise=False) + self.nvfp4_quantizer.internal = True + self.dtype = dtype + self.backend = backend + + def copy(self): + """Copy options without sharing mutable native quantizer state.""" + result = object.__new__(type(self)) + result.__dict__ = self.__dict__.copy() + # Native NVFP4 state consists of options and an immutable cached RHT + # matrix. A shallow copy isolates option updates without rebuilding + # derived CUDA state for every high-precision result. + # Avoid the native __getstate__ hook: serialization intentionally drops + # process groups, whereas a runtime copy must preserve them. + native = object.__new__(type(self.nvfp4_quantizer)) + native.__dict__ = self.nvfp4_quantizer.__dict__.copy() + result.nvfp4_quantizer = native + return result + + @property + def with_amax_reduction(self): + """Forward distributed amax policy to the native fallback.""" + return self.nvfp4_quantizer.with_amax_reduction + + @with_amax_reduction.setter + def with_amax_reduction(self, value): + self.nvfp4_quantizer.with_amax_reduction = value + + @property + def amax_reduction_group(self): + """Native amax reduction group.""" + return self.nvfp4_quantizer.amax_reduction_group + + @amax_reduction_group.setter + def amax_reduction_group(self, value): + self.nvfp4_quantizer.amax_reduction_group = value + + def selected_backend(self, tensor, dtype=None): + """Return dispatch selection, or reject a forced unsupported fast path.""" + q = self.nvfp4_quantizer + output_dtype = dtype or self.dtype or tensor.dtype + eligible = ( + tensor.is_cuda + and tensor.is_contiguous() + and tensor.ndim >= 2 + and tensor.numel() > 0 + and tensor.shape[-1] % 16 == 0 + and (tensor.numel() // tensor.shape[-1]) % 16 == 0 + and tensor.data_ptr() % 16 == 0 + and tensor.dtype in (torch.bfloat16, torch.float16) + and output_dtype == tensor.dtype + and _supports_fused_device(tensor.device.index) + and q.dtype == DType.kFloat4E2M1 + and q.scale_dtype == DType.kFloat8E4M3 + and q.nvfp4_e4m3_max in (0, 448) + and not any( + ( + q.with_rht, + q.with_post_rht_amax, + q.with_2d_quantization, + q.stochastic_rounding, + q.row_scaled_nvfp4, + q.nvfp4_use_4over6, + q.disable_second_level_scale, + q.with_amax_reduction, + ) + ) + and os.getenv("NVTE_USE_FAST_MATH", "0") == "0" + ) + if self.backend == "fused" and not eligible: + raise ValueError("Forced fused NVFP4 QDQ does not support this input/configuration") + return "fused" if eligible and self.backend != "reference" else "reference" + + def _reference(self, tensor, dtype): + return self.nvfp4_quantizer.quantize(tensor).dequantize(dtype=dtype) + + def _compute(self, tensor, output, noop_flag=None, *, backend=None): + if (backend or self.selected_backend(tensor, output.dtype)) == "fused": + # Framework owns both allocations; amax reduction and QDQ use the + # current stream and may be captured together. + amax = torch.empty(1, dtype=torch.float32, device=tensor.device) + tex.nvfp4_qdq(tensor, output, amax, noop_flag) + else: + data = self._reference(tensor, output.dtype) + if noop_flag is None: + output.copy_(data) + else: + torch.where(noop_flag != 1, data, output, out=output) + + def quantize_impl(self, tensor, *, dtype=None): + dtype = dtype or self.dtype or tensor.dtype + if self.selected_backend(tensor, dtype) == "reference": + data = self._reference(tensor, dtype) + return self._wrap(data) + result = self._wrap(torch.empty(tensor.shape, dtype=dtype, device=tensor.device)) + self._compute(tensor, result._hp_data, backend="fused") + return result + + def make_empty( + self, + shape, + *, + dtype=torch.float32, + device=None, + requires_grad=False, + pin_memory=False, + ): + data = torch.empty( + tuple(shape), + dtype=self.dtype or dtype, + device=device or "cuda", + pin_memory=pin_memory, + ) + return self._wrap(data, requires_grad=requires_grad) + + def _wrap(self, data, requires_grad=False): + if self.internal: + return IdentityTensorStorage(hp_data=data, fake_dtype=data.dtype, quantizer=self) + return IdentityTensor( + data.shape, + data.dtype, + hp_data=data, + quantizer=self, + requires_grad=requires_grad, + device=data.device, + ) + + def update_quantized(self, src, dst, *, noop_flag=None): + if not isinstance(dst, IdentityTensorStorage) or dst._hp_data is None: + raise TypeError("NVFP4 QDQ requires allocated IdentityTensorStorage") + if src.shape != dst._hp_data.shape or src.device != dst._hp_data.device: + raise ValueError("NVFP4 QDQ source and destination must have matching shape/device") + self._compute(src, dst._hp_data, noop_flag) + return dst + + def is_requantization_safe(self): + return self.nvfp4_quantizer.is_requantization_safe() + + def _get_compatible_recipe(self): + from transformer_engine.common.recipe import CustomRecipe + + return CustomRecipe diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py index 0d5fc0c4419..a01f068aec9 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -348,6 +348,42 @@ def nvfp4_row_scaled_fwd_high_precision_bwd_factory( return _plain_nvfp4_quantizer() +def nvfp4_qdq_weight_fwd_mxfp8_bwd_factory( + role: Optional[QuantizerRole], *, backend="auto", nvfp4_options=None +): + """High-precision X and NVFP4-QDQ W forward; MXFP8 backward. + + Applies to Linear and GroupedLinear. The module compute + dtype determines forward precision. Other roles use Identity. Use + functools.partial to select backend or native NVFP4 options, including + 4over6 (reference fallback). + """ + from ..tensor.hybrid_tensor import HybridQuantizer + from ..tensor.identity_tensor import IdentityQuantizer + from ..tensor.nvfp4_tensor import NVFP4Quantizer + from .qdq import NVFP4QDQQuantizer + + if role is None or role.module_type not in ("linear", "grouped_linear"): + return IdentityQuantizer() + if role.tensor_type == "input": + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=mxfp8_factory(role), + columnwise_source="original", + ) + if role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=NVFP4QDQQuantizer( + NVFP4Quantizer(**(nvfp4_options or {})), backend=backend + ), + columnwise_quantizer=mxfp8_factory(role), + columnwise_source="rowwise_dequantized", + ) + if role.tensor_type == "grad_output": + return mxfp8_factory(role) + return IdentityQuantizer() + + # ----------------------------------------------------------------------------- # Linear + Attention Recipes # ----------------------------------------------------------------------------- From 7ee52817ea2317f8deea4933e3dd26c4e58b6259 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Thu, 8 Oct 2026 09:22:23 -0700 Subject: [PATCH 3/9] Add MXFP8 QDQ and shared QDQ quantizer base Signed-off-by: Evgeny --- .../fine_grained_quantization.rst | 82 +++- tests/pytorch/mxfp8/test_mxfp8_qdq.py | 369 ++++++++++++++++++ tests/pytorch/nvfp4/test_nvfp4_qdq.py | 77 +++- transformer_engine/common/CMakeLists.txt | 1 + transformer_engine/common/cast/mxfp8/qdq.cu | 81 ++++ .../common/include/transformer_engine/cast.h | 7 + transformer_engine/pytorch/csrc/extensions.h | 1 + .../pytorch/csrc/extensions/cast.cpp | 17 + .../pytorch/csrc/extensions/pybind.cpp | 3 + .../pytorch/custom_recipes/qdq.py | 245 +++++++++--- .../custom_recipes/quantizer_factory_zoo.py | 48 ++- 11 files changed, 826 insertions(+), 105 deletions(-) create mode 100644 tests/pytorch/mxfp8/test_mxfp8_qdq.py create mode 100644 transformer_engine/common/cast/mxfp8/qdq.cu diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst index 497d9e257cf..d5434cb44d7 100644 --- a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -388,27 +388,54 @@ See the :doc:`PyTorch API <../../../api/pytorch>` for ``QuantizerRole``, ``HybridQuantizer``, ``IdentityQuantizer``, and their returned tensor types. See the :doc:`Common API <../../../api/common>` for ``CustomRecipe``. -Experimental NVFP4 weight QDQ ------------------------------ +Experimental QDQ +---------------- -``custom_recipes.qdq.NVFP4QDQQuantizer`` is useful for experiments with -``X[BF16] x W[NVFP4]``: it quantizes weights to NVFP4, then dequantizes them -into a high-precision ``IdentityTensor`` for BF16 forward GEMMs. This preserves -NVFP4 weight rounding error while using BF16 GEMM arithmetic. +QDQ quantizers model low-precision rounding by quantizing operands and then +dequantizing them into high-precision ``IdentityTensor`` storage. The decoded +operands use high-precision GEMMs; QDQ models quantization error for convergence +experiments, not mixed-format hardware performance. -The following factory supports ``Linear`` and ``GroupedLinear`` with discrete -expert parameters. Forward uses high-precision inputs and QDQ weights; -MXFP8 backward uses the original inputs and QDQ weights:: +Forward and backward examples +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The following factories support ``Linear`` and ``GroupedLinear`` with discrete +expert parameters. With BF16 module compute precision: + +* ``nvfp4_qdq_weight_fwd_bf16_bwd_factory`` models ``X[BF16] x W[NVFP4]`` + forward and uses BF16 backward. Backward consumes the original inputs and + the decoded forward QDQ weights. +* ``mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory`` models ``X[MXFP8] x W[NVFP4]`` + forward and uses native MXFP8 backward. Input QDQ uses block-32 E4M3/E8M0; + weight QDQ uses block-16 E2M1/E4M3 with a tensor scale. Backward quantizes + the original inputs and the decoded forward QDQ weights to MXFP8. + +For example:: from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( - nvfp4_qdq_weight_fwd_mxfp8_bwd_factory, + nvfp4_qdq_weight_fwd_bf16_bwd_factory, + mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory, ) - recipe = CustomRecipe(qfactory=nvfp4_qdq_weight_fwd_mxfp8_bwd_factory) + recipe = CustomRecipe(qfactory=mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory) + +To use BF16 inputs while retaining MXFP8 backward, replace the second factory's +input rowwise ``MXFP8QDQQuantizer`` with ``IdentityQuantizer``. The input +columnwise quantizer and its ``columnwise_source="original"`` stay unchanged. +Both examples use ``columnwise_source="rowwise_dequantized"`` for weights. +Other module roles use ``IdentityQuantizer``. High-precision arithmetic follows +the module compute dtype; select BF16 compute for the BF16 example. + +Configure ``backend`` or ``nvfp4_options`` with ``functools.partial``. Existing +Hybrid limitations still apply, and RL convergence has not been validated. + +Backends and custom QDQ formats +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Output defaults to the source dtype; a module compute-dtype override changes -only the decoded dtype. ``backend="auto"`` fuses plain 1D E2M1/E4M3 QDQ on +only the decoded dtype. For ``NVFP4QDQQuantizer``, ``backend="auto"`` fuses +plain 1D E2M1/E4M3 QDQ on SM100 for contiguous, 16-byte-aligned BF16/FP16 tensors with matching output dtype and flattened matrix dimensions divisible by 16, with ordinary fast math disabled. Other native-valid modes, including 4over6, 2D scaling, distributed @@ -416,6 +443,31 @@ amax reduction, and mixed dtypes, use native ``quantize().dequantize()``. ``backend="reference"`` forces that composition; ``backend="fused"`` rejects unsupported cases. Current weight amax is computed when quantization runs. -Configure ``backend`` or ``nvfp4_options`` on the factory with -``functools.partial``. Existing Hybrid limitations still apply, and RL -convergence has not been validated. +``custom_recipes.qdq.MXFP8QDQQuantizer`` provides the corresponding rowwise +E4M3/block-32 E8M0 QDQ operation. Its ``reference`` backend calls native +``MXFP8Quantizer.quantize().dequantize()``. Its ``fused`` backend supports +contiguous, 16-byte-aligned BF16/FP16 tensors on SM100, with a matching output +dtype, rows divisible by 32, and ordinary fast math disabled. ``auto`` falls +back to the reference when these requirements are not met. +The fused kernel preserves signed zero. Native specialized FP16 quantization +may canonicalize negative zero; the decoded values remain numerically equal. + +Both QDQ classes return high-precision storage with straight-through +gradients. Use them as rowwise children of ``HybridQuantizer`` to select the +backward format independently. ``columnwise_source="rowwise_dequantized"`` +requantizes the forward reconstruction for backward; ``"original"`` uses +the original input. QDQ models quantization error for convergence studies; +the high-precision GEMM does not model mixed-format hardware performance. + +Both formats inherit the experimental ``custom_recipes.qdq.QDQQuantizer`` +base, which supplies allocation, IdentityTensor wrapping, dtype overrides, +update validation, and straight-through gradients. Custom QDQ subclasses +implement ``selected_backend(tensor, dtype)``, ``_reference(tensor, dtype)``, +and ``_compute(tensor, output, noop_flag=None, *, backend=None)``. The reference +must quantize the original input before decoding to the requested dtype. +Compute must write into the supplied output, preserving it when the no-op flag +is one, and honor an explicitly selected backend. Subclasses must also provide +format-appropriate ``copy()`` and ``is_requantization_safe()`` implementations. +Storage hooks should not build autograd graphs; Python implementations can +use ``torch.no_grad()`` since the wrapper supplies the straight-through gradient. +The base and its extension hooks are experimental and subject to change. diff --git a/tests/pytorch/mxfp8/test_mxfp8_qdq.py b/tests/pytorch/mxfp8/test_mxfp8_qdq.py new file mode 100644 index 00000000000..bb012b3da2b --- /dev/null +++ b/tests/pytorch/mxfp8/test_mxfp8_qdq.py @@ -0,0 +1,369 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP8 QDQ native parity, dispatch, storage lifecycle, and API validation.""" + +import pytest +import torch +from transformer_engine.pytorch.custom_recipes.qdq import MXFP8QDQQuantizer +import transformer_engine_torch as tex + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), + reason="MXFP8 QDQ fused implementation targets SM100", +) + + +@pytest.fixture(autouse=True) +def ordinary_math(monkeypatch): + monkeypatch.setenv("NVTE_USE_FAST_MATH", "0") + + +def assert_bits(actual, expected): + assert actual.shape == expected.shape and actual.dtype == expected.dtype + assert torch.equal( + actual.contiguous().view(torch.int16), expected.contiguous().view(torch.int16) + ), ( + (actual != expected).sum().item(), + (actual.float() - expected.float()).abs().max().item(), + ) + + +def operand(backend): + return MXFP8QDQQuantizer(backend=backend) + + +def data(shape, dtype, distribution): + torch.manual_seed(41) + x = torch.randn(shape, device="cuda", dtype=dtype) + if distribution == "zeros": + x.zero_() + x[:, ::2] = -0.0 + elif distribution == "outliers": + x[0, 0] = 4096 + x[1::4] *= 1e-5 + elif distribution == "boundaries": + v = torch.tensor( + [ + 0, + -0.0, + 0.25, + 0.75, + 1.25, + 1.75, + 2.5, + 3.5, + 5, + 6, + -6, + -5, + -2.5, + 1e-7, + 1e-4, + 448, + ], + device="cuda", + dtype=dtype, + ) + x.copy_(v.repeat(shape[0], shape[1] // 16)) + return x + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("shape", [(32, 32), (96, 160), (128, 256), (1536, 1024)]) +@pytest.mark.parametrize("distribution", ["normal", "zeros", "outliers", "boundaries"]) +def test_fused_matches_materialized_q_plus_dq(dtype, shape, distribution): + x = data(shape, dtype, distribution) + fused, reference = operand("fused"), operand("reference") + assert fused.selected_backend(x) == "fused" + expected = reference(x).dequantize() + for _ in range(2): # Exact replay, no stochastic forward state. + actual = fused(x).dequantize() + if dtype == torch.float16: + # Native's specialized FP16 MXFP8 FMA canonicalizes input -0. + # Require exact numerical equality; zero signs are immaterial. + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + else: + assert_bits(actual, expected) + + +def test_inplace_graph_noop_and_ste(): + x = data((128, 256), torch.bfloat16, "normal").requires_grad_() + q = operand("fused") + ref = operand("reference") + y = q(x) + # Differentiable operations use the tensor wrapper. dequantize() exposes + # TE IdentityTensor's detached storage, not its autograd edge. + y.sum().backward() + assert torch.equal(x.grad, torch.ones_like(x)) + ptr = y._hp_data.data_ptr() + noop = torch.ones(1, device="cuda", dtype=torch.float32) + q.update_quantized(x * 2, y, noop_flag=noop) + assert_bits(y.dequantize(), ref(x).dequantize()) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + q.update_quantized(x, y, noop_flag=noop) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + q.update_quantized(x, y, noop_flag=noop) + with torch.no_grad(): + x.mul_(2) + noop.zero_() + graph.replay() + assert y._hp_data.data_ptr() == ptr + assert_bits(y.dequantize(), ref(x).dequantize()) + + +def test_noncontiguous_and_output_dtype_fallback(): + x = data((128, 256), torch.bfloat16, "normal").t() + q = operand("auto") + assert q.selected_backend(x) == "reference" + assert_bits(q(x).dequantize(), operand("reference")(x).dequantize()) + with pytest.raises(ValueError): + operand("fused")(x) + source = x.contiguous().float() + out = q.quantize(source, dtype=torch.bfloat16) + assert out.dequantize().dtype == torch.bfloat16 + q.update_quantized(source * 2, out) + assert_bits(out.dequantize(), q._reference(source * 2, torch.bfloat16)) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_scale_boundaries_and_extremes(dtype): + # Adjacent representable values straddling E8M0 scale transitions. + exponents = range(-120, 119, 13) if dtype == torch.bfloat16 else range(-20, 7, 2) + thresholds = torch.tensor( + [448.0 * 2.0**e for e in exponents], device="cuda", dtype=dtype + ) + maxima = torch.cat( + ( + torch.nextafter(thresholds, torch.zeros_like(thresholds)), + thresholds, + torch.nextafter(thresholds, torch.full_like(thresholds, float("inf"))), + torch.tensor( + [torch.finfo(dtype).max, torch.finfo(dtype).tiny], + device="cuda", + dtype=dtype, + ), + torch.nextafter( + torch.zeros(1, device="cuda", dtype=dtype), + torch.ones(1, device="cuda", dtype=dtype), + ), + ) + ) + # Native materialized MXFP8 requires both logical dimensions aligned. + maxima = torch.nn.functional.pad(maxima, (0, (-maxima.numel()) % 32)) + x = maxima[:, None] * torch.linspace(-1, 1, 32, device="cuda", dtype=dtype) + actual = operand("fused")(x).dequantize() + expected = operand("reference")(x).dequantize() + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_exceptional_values(value, dtype): + x = torch.randn(32, 64, device="cuda", dtype=dtype) + x[0, 0] = value + actual = operand("fused")(x).dequantize() + expected = operand("reference")(x).dequantize() + torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True) + + +@pytest.mark.parametrize("backend", ["reference", "fused"]) +@pytest.mark.parametrize("internal", [False, True]) +def test_copy_workspace_dtype_and_noop(backend, internal): + q = operand(backend) + q.internal = internal + copied = q.copy() + assert copied.mxfp8_quantizer is not q.mxfp8_quantizer + x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + y = q.make_empty(x.shape, dtype=x.dtype, device=x.device) + q.update_quantized(x, y) + before = y.dequantize().clone() + flag = torch.ones(1, device=x.device) + q.update_quantized(2 * x, y, noop_flag=flag) + torch.testing.assert_close(y.dequantize(), before, rtol=0, atol=0) + flag.zero_() + q.update_quantized(2 * x, y, noop_flag=flag) + torch.testing.assert_close( + y.dequantize(), operand("reference")(2 * x).dequantize(), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize( + "bad", + [ + "cpu", + "shape", + "dtype", + "strided", + "alignment", + "columns", + "rank", + "noop_dtype", + "noop_shape", + "noop_cpu", + ], +) +def test_direct_api_rejects_invalid_buffers(bad): + x = torch.ones(32, 64, device="cuda", dtype=torch.bfloat16) + y = torch.empty_like(x) + flag = None + if bad == "cpu": + y = y.cpu() + elif bad == "shape": + y = y[:16] + elif bad == "dtype": + y = y.float() + elif bad == "strided": + y = y.T + elif bad == "alignment": + y = torch.empty(x.numel() + 1, device="cuda", dtype=x.dtype)[1:].view_as(x) + elif bad == "columns": + x = x[:, :16].contiguous() + y = torch.empty_like(x) + elif bad == "rank": + x = x.flatten() + y = torch.empty_like(x) + elif bad == "noop_dtype": + flag = torch.ones(1, device="cuda", dtype=torch.int32) + elif bad == "noop_shape": + flag = torch.ones(2, device="cuda") + elif bad == "noop_cpu": + flag = torch.ones(1) + with pytest.raises((RuntimeError, ValueError)): + tex.mxfp8_qdq(x, y, flag) + + +def test_wrapper_errors_and_fast_math(monkeypatch): + with pytest.raises(ValueError, match="backend"): + MXFP8QDQQuantizer(backend="unknown") + x = torch.ones(32, 64, device="cuda", dtype=torch.bfloat16) + q = operand("auto") + with pytest.raises(TypeError, match="MXFP8QDQQuantizer"): + q.update_quantized(x, torch.empty_like(x)) + with pytest.raises(ValueError, match="MXFP8QDQQuantizer"): + q.update_quantized(x[:16], q(x)) + monkeypatch.setenv("NVTE_USE_FAST_MATH", "1") + assert q.selected_backend(x) == "reference" + with pytest.raises(ValueError, match="Forced fused"): + operand("fused")(x) + + +@pytest.mark.parametrize("shape", [(0, 128), (128, 0)]) +def test_empty_fallback(shape): + x = torch.empty(shape, device="cuda", dtype=torch.bfloat16) + q = operand("auto") + assert q.selected_backend(x) == "reference" + assert q(x).dequantize().shape == x.shape + + +def test_current_stream_and_device(): + device = 1 if torch.cuda.device_count() > 1 else 0 + with torch.cuda.device(device), torch.cuda.stream(torch.cuda.Stream(device=device)): + x = torch.randn(32, 64, device=f"cuda:{device}", dtype=torch.bfloat16) + torch.testing.assert_close( + operand("fused")(x).dequantize(), + operand("reference")(x).dequantize(), + rtol=0, + atol=0, + ) + + +@pytest.mark.parametrize("source", ["original", "rowwise_dequantized"]) +def test_hybrid_backward_source(source): + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=operand("fused"), + columnwise_quantizer=MXFP8Quantizer(tex.DType.kFloat8E4M3), + columnwise_source=source, + ) + y = q(x) + decoded = operand("reference")(x).dequantize() + torch.testing.assert_close(y._rowwise_storage.dequantize(), decoded, rtol=0, atol=0) + native = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=False, columnwise=True) + expected = native(x if source == "original" else decoded) + for name in ("_columnwise_data", "_columnwise_scale_inv"): + assert torch.equal( + getattr(y._columnwise_storage, name), getattr(expected, name) + ) + + +@pytest.mark.parametrize("backend", ["reference", "fused"]) +def test_linear_forward_and_backward(backend): + import transformer_engine.pytorch as te + from transformer_engine.common.recipe import CustomRecipe + + # A standalone custom factory: every operand uses the QDQ wrapper, so + # forward and backward can both be checked against materialized Q+DQ. + def factory(role): + return operand(backend) + + layer = te.Linear(128, 64, bias=False, params_dtype=torch.bfloat16, device="cuda") + x = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + grad = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + with te.autocast(recipe=CustomRecipe(qfactory=factory)): + y = layer(x) + y.backward(grad) + q = operand("reference") + xq, wq, gq = [q(t.detach()).dequantize() for t in (x, layer.weight, grad)] + torch.testing.assert_close(y, xq @ wq.T, rtol=1e-2, atol=2e-3) + torch.testing.assert_close(x.grad, gq @ wq, rtol=1e-2, atol=2e-3) + torch.testing.assert_close(layer.weight.grad, gq.T @ xq, rtol=1e-2, atol=2e-3) + + +def test_custom_qdq_base_lifecycle(): + """A user-defined format inherits storage and autograd without TE kernels.""" + import copy + from transformer_engine.pytorch.custom_recipes.qdq import QDQQuantizer + + class RoundedQDQQuantizer(QDQQuantizer): + def selected_backend(self, tensor, dtype=None): + if self.backend == "fused": + raise ValueError("This example only supports reference execution") + return "reference" + + def _reference(self, tensor, dtype): + return tensor.round().to(dtype) + + @torch.no_grad() + def _compute(self, tensor, output, noop_flag=None, *, backend=None): + data = self._reference(tensor, output.dtype) + if noop_flag is None: + output.copy_(data) + else: + torch.where(noop_flag != 1, data, output, out=output) + + def copy(self): + return copy.copy(self) + + def is_requantization_safe(self): + return True + + with pytest.raises(TypeError): + QDQQuantizer() + x = torch.tensor([[0.25, 1.75]], device="cuda", requires_grad=True) + q = RoundedQDQQuantizer(dtype=torch.bfloat16) + y = q(x) + assert y.dtype == torch.bfloat16 + torch.testing.assert_close(y.dequantize(), x.detach().round().bfloat16()) + y.sum().backward() + torch.testing.assert_close(x.grad, torch.ones_like(x)) + workspace = q.make_empty(x.shape, dtype=x.dtype, device=x.device) + q.update_quantized(x, workspace) + flag = torch.ones(1, device=x.device) + q.update_quantized(x + 2, workspace, noop_flag=flag) + torch.testing.assert_close(workspace.dequantize(), y.dequantize()) + flag.zero_() + q.update_quantized(x + 2, workspace, noop_flag=flag) + torch.testing.assert_close( + workspace.dequantize(), (x.detach() + 2).round().bfloat16() + ) + q.internal = True + assert q(x)._hp_data.dtype == torch.bfloat16 diff --git a/tests/pytorch/nvfp4/test_nvfp4_qdq.py b/tests/pytorch/nvfp4/test_nvfp4_qdq.py index 324388ff698..ba8ed05d71f 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_qdq.py +++ b/tests/pytorch/nvfp4/test_nvfp4_qdq.py @@ -15,7 +15,8 @@ from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.custom_recipes.qdq import NVFP4QDQQuantizer from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( - nvfp4_qdq_weight_fwd_mxfp8_bwd_factory, + mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory, + nvfp4_qdq_weight_fwd_bf16_bwd_factory, ) pytestmark = pytest.mark.skipif( @@ -83,7 +84,9 @@ def test_plain_exact(dtype, shape, distribution, monkeypatch): def test_4over6_fallback(mode, maximum, fast_error, monkeypatch): monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", fast_error) x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) - options = dict(nvfp4_use_4over6=True, nvfp4_e4m3_max=maximum, nvfp4_4over6_err_mode=mode) + options = dict( + nvfp4_use_4over6=True, nvfp4_e4m3_max=maximum, nvfp4_4over6_err_mode=mode + ) q = NVFP4QDQQuantizer(NVFP4Quantizer(**options)) assert q.selected_backend(x) == "reference" assert_bits(q(x).dequantize(), native_qdq(x, **options)) @@ -157,19 +160,30 @@ def test_ste_and_rng(): def mx_value(x, *, columnwise=False): - q = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=not columnwise, columnwise=columnwise) + q = MXFP8Quantizer( + tex.DType.kFloat8E4M3, rowwise=not columnwise, columnwise=columnwise + ) return q(x).dequantize() @pytest.mark.parametrize("backend", ["reference", "fused"]) @pytest.mark.parametrize("grouped", [False, True]) @pytest.mark.parametrize("params_dtype", [torch.bfloat16, torch.float32]) -def test_module_operands(backend, grouped, params_dtype, nvfp4_options=None): +@pytest.mark.parametrize("mixed", [False, True]) +def test_module_operands( + backend, grouped, params_dtype, mixed, nvfp4_options=None +): if backend == "fused" and params_dtype == torch.float32: pytest.skip("Mixed FP32 source/BF16 decode intentionally uses native fallback") torch.manual_seed(12) - factory = nvfp4_qdq_weight_fwd_mxfp8_bwd_factory - recipe = CustomRecipe(qfactory=partial(factory, backend=backend, nvfp4_options=nvfp4_options)) + factory = ( + mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory + if mixed + else nvfp4_qdq_weight_fwd_bf16_bwd_factory + ) + recipe = CustomRecipe( + qfactory=partial(factory, backend=backend, nvfp4_options=nvfp4_options) + ) kwargs = dict(bias=True, params_dtype=params_dtype, device="cuda") layer = ( te.GroupedLinear(2, 128, 128, single_grouped_weight=False, **kwargs) @@ -177,21 +191,31 @@ def test_module_operands(backend, grouped, params_dtype, nvfp4_options=None): else te.Linear(128, 128, **kwargs) ) splits = [32, 64] if grouped else [96] - x = torch.randn(sum(splits), 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + x = torch.randn( + sum(splits), 128, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) dy = torch.randn_like(x) - weights = [getattr(layer, f"weight{i}") for i in range(2)] if grouped else [layer.weight] + weights = ( + [getattr(layer, f"weight{i}") for i in range(2)] if grouped else [layer.weight] + ) biases = [getattr(layer, f"bias{i}") for i in range(2)] if grouped else [layer.bias] with te.autocast(recipe=recipe), torch.autocast("cuda", dtype=torch.bfloat16): y = layer(x, splits) if grouped else layer(x) y.backward(dy) expected_y, expected_dx = [], [] - for w, b, xs, gs in zip(weights, biases, x.detach().split(splits), dy.split(splits)): + for w, b, xs, gs in zip( + weights, biases, x.detach().split(splits), dy.split(splits) + ): wh = native_qdq(w.detach(), torch.bfloat16, **(nvfp4_options or {})) - expected_y.append(torch.nn.functional.linear(xs, wh, b.detach().to(torch.bfloat16))) - bw = mx_value(wh, columnwise=True) - bx = mx_value(xs, columnwise=True) - grow = mx_value(gs) - gcol = mx_value(gs, columnwise=True) + expected_y.append( + torch.nn.functional.linear( + mx_value(xs) if mixed else xs, wh, b.detach().to(torch.bfloat16) + ) + ) + bw = mx_value(wh, columnwise=True) if mixed else wh + bx = mx_value(xs, columnwise=True) if mixed else xs + grow = mx_value(gs) if mixed else gs + gcol = mx_value(gs, columnwise=True) if mixed else gs expected_dx.append(grow @ bw) dw = gcol.T @ bx torch.testing.assert_close(w.grad.float(), dw.float(), rtol=0.016, atol=0.002) @@ -209,10 +233,12 @@ def test_weight_workspace_dtype_and_reuse(): from transformer_engine.pytorch.module.base import quantize_weight from transformer_engine.pytorch.quantization import QuantizerRole - factory = nvfp4_qdq_weight_fwd_mxfp8_bwd_factory + factory = mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory q = factory(QuantizerRole(module_type="linear", tensor_type="weight")) w = torch.randn(128, 256, device="cuda", dtype=torch.float32) - out, cache = quantize_weight(tensor=w, quantizer=q, workspace_dtype=torch.bfloat16, cache=True) + out, cache = quantize_weight( + tensor=w, quantizer=q, workspace_dtype=torch.bfloat16, cache=True + ) expected = native_qdq(w, torch.bfloat16) assert_bits(out._rowwise_storage.dequantize(), expected) pointer = out._rowwise_storage._hp_data.data_ptr() @@ -249,8 +275,13 @@ def test_weight_workspace_dtype_and_reuse(): workspace_dtype=torch.float32, cache=True, ) - assert cache3 is not cache and out3._rowwise_storage.dequantize().dtype == torch.float32 - torch.testing.assert_close(out3._rowwise_storage.dequantize(), native_qdq(w), rtol=0, atol=0) + assert ( + cache3 is not cache + and out3._rowwise_storage.dequantize().dtype == torch.float32 + ) + torch.testing.assert_close( + out3._rowwise_storage.dequantize(), native_qdq(w), rtol=0, atol=0 + ) def test_noncontiguous_and_fast_math_fallback(monkeypatch): @@ -304,11 +335,13 @@ def test_copy_preserves_runtime_reduction_state(): @pytest.mark.parametrize("grouped", [False, True]) -def test_4over6_modules(grouped): +@pytest.mark.parametrize("mixed", [False, True]) +def test_4over6_modules(grouped, mixed): test_module_operands( "reference", grouped, torch.bfloat16, + mixed=mixed, nvfp4_options={ "nvfp4_use_4over6": True, "nvfp4_e4m3_max": 448, @@ -338,7 +371,7 @@ def test_qdq_columnwise_only_workspace_dtype(empty): """A transient forward QDQ source must round to the requested compute dtype.""" from transformer_engine.pytorch.quantization import QuantizerRole - q = nvfp4_qdq_weight_fwd_mxfp8_bwd_factory( + q = mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory( QuantizerRole(module_type="linear", tensor_type="weight") ) q.set_usage(rowwise=False, columnwise=True) @@ -354,5 +387,7 @@ def test_qdq_columnwise_only_workspace_dtype(empty): expected = mx_value(native_qdq(x, torch.bfloat16).float(), columnwise=True) # Native make_empty honors its allocation dtype; native quantize retains # the FP32 source dtype. Preserve both behaviors and compare equal decodes. - assert out._columnwise_storage._dtype == (torch.bfloat16 if empty else torch.float32) + assert out._columnwise_storage._dtype == ( + torch.bfloat16 if empty else torch.float32 + ) assert_bits(out._columnwise_storage.dequantize(dtype=torch.float32), expected) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 33f495af83e..6c489d778f4 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -359,6 +359,7 @@ list(APPEND transformer_engine_cuda_arch_specific_sources activation/swiglu_grouped_bwd.cu activation/swiglu_grouped_dbias.cu cast/nvfp4/qdq.cu + cast/mxfp8/qdq.cu cast/cast.cu cast/cast_dbias.cu cast/cast_grouped.cu diff --git a/transformer_engine/common/cast/mxfp8/qdq.cu b/transformer_engine/common/cast/mxfp8/qdq.cu new file mode 100644 index 00000000000..d67360839c5 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/qdq.cu @@ -0,0 +1,81 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include "../../common.h" +#include "../../util/cuda_runtime.h" +#include "../../util/ptx.cuh" +#include "../../util/ptx_arch_spec.cuh" + +namespace transformer_engine { +namespace { + +// Four adjacent threads own one block of 32. Each thread reads/writes 16B. +template +__global__ void mxfp8_qdq_kernel(const T *input, T *output, const float *noop, size_t blocks) { + if (noop != nullptr && *noop == 1.0f) return; + const size_t thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t block = thread / 4; + if (block >= blocks) return; + const size_t offset = block * 32 + (thread % 4) * 8; + Vec values; + values.load_from(input + offset); + float amax = 0.0f; +#pragma unroll + for (int i = 0; i < 8; ++i) { + amax = fmaxf(amax, fabsf(static_cast(values.data.elt[i]))); + } + const unsigned mask = __activemask(); + amax = fmaxf(amax, __shfl_xor_sync(mask, amax, 1, 4)); + amax = fmaxf(amax, __shfl_xor_sync(mask, amax, 2, 4)); + const e8m0_t exponent = ptx::float_to_e8m0(amax * (1.0f / 448.0f)); + const float encode = ptx::exp2f_rcp(exponent); + const float decode = ptx::exp2f(exponent); +#pragma unroll + for (int i = 0; i < 8; ++i) { + const fp8e4m3 encoded = static_cast(static_cast(values.data.elt[i]) * encode); + values.data.elt[i] = static_cast(static_cast(encoded) * decode); + } + values.store_to(output + offset); +} +} // namespace +} // namespace transformer_engine + +void nvte_mxfp8_qdq(const NVTETensor input, NVTETensor output, const NVTETensor noop, + cudaStream_t stream) { + NVTE_API_CALL(nvte_mxfp8_qdq); + using namespace transformer_engine; + NVTE_CHECK(cuda::sm_arch() == 100, "MXFP8 QDQ currently supports SM100"); + const auto &in = *convertNVTETensorCheck(input); + auto &out = *convertNVTETensorCheck(output); + CheckInputTensor(in, "input"); + CheckOutputTensor(out, "output"); + NVTE_CHECK(in.data.dtype == out.data.dtype && in.data.shape == out.data.shape, + "QDQ input and output must have matching dtype and shape"); + NVTE_CHECK(in.data.dtype == DType::kBFloat16 || in.data.dtype == DType::kFloat16, + "QDQ supports BF16 and FP16"); + NVTE_CHECK(in.data.shape.size() >= 2 && in.numel() > 0 && in.data.shape.back() % 32 == 0, + "QDQ requires nonempty block-32 aligned rows"); + NVTE_CHECK(reinterpret_cast(in.data.dptr) % 16 == 0 && + reinterpret_cast(out.data.dptr) % 16 == 0, + "QDQ requires 16-byte aligned buffers"); + const float *noop_ptr = nullptr; + if (noop != nullptr) { + const auto &flag = *convertNVTETensorCheck(noop); + CheckInputTensor(flag, "noop"); + NVTE_CHECK(flag.data.dtype == DType::kFloat32 && flag.numel() == 1, + "QDQ noop must be a scalar FP32 tensor"); + noop_ptr = static_cast(flag.data.dptr); + } + constexpr int threads = 128; + const size_t blocks = in.numel() / 32; + TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT( + in.data.dtype, T, + mxfp8_qdq_kernel<<<(blocks * 4 + threads - 1) / threads, threads, 0, stream>>>( + static_cast(in.data.dptr), static_cast(out.data.dptr), noop_ptr, + blocks);); + NVTE_CHECK_CUDA(cudaGetLastError()); +} diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index b32708c9e9f..09ad3f24d17 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -29,6 +29,13 @@ extern "C" { void nvte_nvfp4_qdq(const NVTETensor input, NVTETensor output, const NVTETensor amax, const NVTETensor noop, cudaStream_t stream); +/*! \brief Rowwise block-32 E4M3/E8M0 MXFP8 QDQ into matching BF16/FP16 output. + * Input and output are contiguous, with rows divisible by 32. An optional scalar + * FP32 noop skips writes when equal to 1. All buffers are caller-owned. + */ +void nvte_mxfp8_qdq(const NVTETensor input, NVTETensor output, const NVTETensor noop, + cudaStream_t stream); + /* Quantize the tensor * * The type of quantized tensor in the output depends on the scaling mode of the output diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 86342593a57..cee3f572ac1 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -504,6 +504,7 @@ at::Tensor scaled_aligned_causal_masked_softmax_backward(at::Tensor output_grads void nvfp4_qdq(const at::Tensor &input, at::Tensor output, at::Tensor amax, std::optional noop, bool compute_amax); +void mxfp8_qdq(const at::Tensor &input, at::Tensor output, std::optional noop); void compute_amax(const at::Tensor &tensor, at::Tensor &amax); void fused_amax_and_scale_update_after_reduction(const at::Tensor &amax_reduction_buffer, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index b476fc58cc1..69ea9ec09d3 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -83,6 +83,23 @@ void nvfp4_qdq(const at::Tensor &input, at::Tensor output, at::Tensor amax, stream); } +void mxfp8_qdq(const at::Tensor &input, at::Tensor output, std::optional noop) { + TORCH_CHECK(input.is_cuda() && output.is_cuda() && input.device() == output.device(), + "MXFP8 QDQ requires tensors on the same CUDA device"); + at::cuda::CUDAGuard guard(input.device()); + TORCH_CHECK(input.is_contiguous() && output.is_contiguous(), + "MXFP8 QDQ requires contiguous tensors"); + TORCH_CHECK( + !noop || (noop->is_cuda() && noop->device() == input.device() && noop->is_contiguous()), + "Invalid QDQ noop device/layout"); + auto in = makeTransformerEngineTensor(input); + auto out = makeTransformerEngineTensor(output); + std::optional noop_cpp; + if (noop) noop_cpp = makeTransformerEngineTensor(*noop); + nvte_mxfp8_qdq(in.data(), out.data(), noop_cpp ? noop_cpp->data() : nullptr, + at::cuda::getCurrentCUDAStream()); +} + py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::object &output, std::optional noop_flag) { // Convert quantizer to C++ object diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 9f52b1ecebb..cf4c067bf77 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -480,6 +480,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("nvfp4_qdq", &transformer_engine::pytorch::nvfp4_qdq, "Plain rowwise NVFP4 QDQ", py::arg("input"), py::arg("output"), py::arg("amax"), py::arg("noop") = py::none(), py::arg("compute_amax") = true, py::call_guard()); + m.def("mxfp8_qdq", &transformer_engine::pytorch::mxfp8_qdq, "Rowwise block-32 MXFP8 QDQ", + py::arg("input"), py::arg("output"), py::arg("noop") = py::none(), + py::call_guard()); m.def("compute_amax", &transformer_engine::pytorch::compute_amax, "Compute absolute max value in tensor", py::arg("input"), py::arg("amax"), py::call_guard()); diff --git a/transformer_engine/pytorch/custom_recipes/qdq.py b/transformer_engine/pytorch/custom_recipes/qdq.py index bc95cac6815..d5382879128 100644 --- a/transformer_engine/pytorch/custom_recipes/qdq.py +++ b/transformer_engine/pytorch/custom_recipes/qdq.py @@ -4,6 +4,7 @@ """Experimental quantize/dequantize quantizers with high-precision storage.""" +from abc import abstractmethod from functools import lru_cache import os @@ -14,6 +15,7 @@ from ..quantized_tensor import Quantizer from ..tensor.identity_tensor import IdentityTensor from ..tensor.storage.identity_tensor_storage import IdentityTensorStorage +from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer @@ -23,7 +25,176 @@ def _supports_fused_device(device_index): return torch.cuda.get_device_capability(device_index) == (10, 0) -class NVFP4QDQQuantizer(Quantizer): +class QDQQuantizer(Quantizer): + """Experimental base for quantize/dequantize with high-precision storage. + + Subclasses implement ``selected_backend``, ``_reference``, and ``_compute``. + The reference returns decoded values in the requested dtype without casting + the quantization input. Compute writes into caller-owned output, honoring + ``noop_flag == 1`` and an optional preselected backend. Subclasses also + implement ``copy`` and ``is_requantization_safe`` for their format's state. + Storage hooks must not build autograd graphs; Python implementations can + use ``torch.no_grad()`` because the wrapper supplies the gradient. + + The shared lifecycle supplies IdentityTensor storage, dtype overrides, and + straight-through gradients. This extension interface is experimental. + """ + + # Unknown to native fused quantization producers; outputs use native HP GEMM. + custom = True + supports_output_dtype = True + + def __init__(self, *, dtype=None, backend="auto", rowwise=True, columnwise=True): + super().__init__(rowwise=rowwise, columnwise=columnwise) + if backend not in ("auto", "reference", "fused"): + raise ValueError("backend must be auto, reference, or fused") + self.dtype = dtype + self.backend = backend + + @abstractmethod + def selected_backend(self, tensor, dtype=None): + """Return reference or fused, rejecting unsupported forced backends.""" + raise NotImplementedError + + @abstractmethod + def _reference(self, tensor, dtype): + """Return materialized Q+DQ in dtype, preserving the input precision.""" + raise NotImplementedError + + @abstractmethod + def _compute(self, tensor, output, noop_flag=None, *, backend=None): + """Write QDQ into output; a noop flag equal to one preserves output.""" + raise NotImplementedError + + def quantize_impl(self, tensor, *, dtype=None): + dtype = dtype or self.dtype or tensor.dtype + if self.selected_backend(tensor, dtype) == "reference": + data = self._reference(tensor, dtype) + return self._wrap(data) + result = self._wrap( + torch.empty(tensor.shape, dtype=dtype, device=tensor.device) + ) + self._compute(tensor, result._hp_data, backend="fused") + return result + + def make_empty( + self, + shape, + *, + dtype=torch.float32, + device=None, + requires_grad=False, + pin_memory=False, + ): + data = torch.empty( + tuple(shape), + dtype=self.dtype or dtype, + device=device or "cuda", + pin_memory=pin_memory, + ) + return self._wrap(data, requires_grad=requires_grad) + + def _wrap(self, data, requires_grad=False): + if self.internal: + return IdentityTensorStorage( + hp_data=data, fake_dtype=data.dtype, quantizer=self + ) + return IdentityTensor( + data.shape, + data.dtype, + hp_data=data, + quantizer=self, + requires_grad=requires_grad, + device=data.device, + ) + + def update_quantized(self, src, dst, *, noop_flag=None): + if not isinstance(dst, IdentityTensorStorage) or dst._hp_data is None: + raise TypeError( + f"{type(self).__name__} requires allocated IdentityTensorStorage" + ) + if src.shape != dst._hp_data.shape or src.device != dst._hp_data.device: + raise ValueError( + f"{type(self).__name__} source and destination must have matching shape/device" + ) + self._compute(src, dst._hp_data, noop_flag) + return dst + + def _get_compatible_recipe(self): + from transformer_engine.common.recipe import CustomRecipe + + return CustomRecipe + + +class MXFP8QDQQuantizer(QDQQuantizer): + """Rowwise E4M3/block-32 E8M0 QDQ held in high-precision storage. + + ``reference`` explicitly invokes native Q then DQ. ``fused`` keeps both + encodings in registers and requires contiguous matching BF16/FP16 on SM100. + ``auto`` selects the fused implementation where supported. The decoded + buffer can be consumed in either GEMM orientation; quantization is rowwise. + """ + + block_size = 32 + + def __init__(self, *, dtype=None, backend="auto", rowwise=True, columnwise=True): + super().__init__( + dtype=dtype, backend=backend, rowwise=rowwise, columnwise=columnwise + ) + self.mxfp8_quantizer = MXFP8Quantizer( + DType.kFloat8E4M3, rowwise=True, columnwise=False + ) + self.mxfp8_quantizer.internal = True + + def copy(self): + """Copy options and isolate the native reference quantizer state.""" + result = object.__new__(type(self)) + result.__dict__ = self.__dict__.copy() + result.mxfp8_quantizer = self.mxfp8_quantizer.copy() + return result + + def selected_backend(self, tensor, dtype=None): + """Select fused execution, or reject an unsupported forced backend.""" + output_dtype = dtype or self.dtype or tensor.dtype + eligible = ( + tensor.is_cuda + and tensor.is_contiguous() + and tensor.ndim >= 2 + and tensor.numel() > 0 + and tensor.shape[-1] % 32 == 0 + and tensor.data_ptr() % 16 == 0 + and tensor.dtype in (torch.bfloat16, torch.float16) + and output_dtype == tensor.dtype + and _supports_fused_device(tensor.device.index) + and os.getenv("NVTE_USE_FAST_MATH", "0") == "0" + ) + if self.backend == "fused" and not eligible: + raise ValueError( + "Forced fused MXFP8 QDQ does not support this input/configuration" + ) + return "fused" if eligible and self.backend != "reference" else "reference" + + def _reference(self, tensor, dtype): + # Native MXFP8 launch geometry does not support empty tensors. + if tensor.numel() == 0: + return torch.empty_like(tensor, dtype=dtype) + return self.mxfp8_quantizer.quantize(tensor).dequantize(dtype=dtype) + + def _compute(self, tensor, output, noop_flag=None, *, backend=None): + if (backend or self.selected_backend(tensor, output.dtype)) == "fused": + tex.mxfp8_qdq(tensor, output, noop_flag) + else: + data = self._reference(tensor, output.dtype) + if noop_flag is None: + output.copy_(data) + else: + torch.where(noop_flag != 1, data, output, out=output) + + def is_requantization_safe(self): + return True + + +class NVFP4QDQQuantizer(QDQQuantizer): """Store native rowwise NVFP4 Q+DQ values in an IdentityTensor. ``quantizer`` supplies native NVFP4 options and is copied with rowwise-only @@ -38,11 +209,6 @@ class NVFP4QDQQuantizer(Quantizer): This experimental API is subject to change. """ - # Unknown to native fused quantization producers; outputs use native HP GEMM. - custom = True - - supports_output_dtype = True - def __init__( self, quantizer=None, @@ -52,16 +218,16 @@ def __init__( rowwise=True, columnwise=True, ): - super().__init__(rowwise=rowwise, columnwise=columnwise) - if backend not in ("auto", "reference", "fused"): - raise ValueError("backend must be auto, reference, or fused") + super().__init__( + dtype=dtype, backend=backend, rowwise=rowwise, columnwise=columnwise + ) if quantizer is not None and not isinstance(quantizer, NVFP4Quantizer): raise TypeError("quantizer must be an NVFP4Quantizer") - self.nvfp4_quantizer = quantizer.copy() if quantizer is not None else NVFP4Quantizer() + self.nvfp4_quantizer = ( + quantizer.copy() if quantizer is not None else NVFP4Quantizer() + ) self.nvfp4_quantizer.set_usage(rowwise=True, columnwise=False) self.nvfp4_quantizer.internal = True - self.dtype = dtype - self.backend = backend def copy(self): """Copy options without sharing mutable native quantizer state.""" @@ -128,7 +294,9 @@ def selected_backend(self, tensor, dtype=None): and os.getenv("NVTE_USE_FAST_MATH", "0") == "0" ) if self.backend == "fused" and not eligible: - raise ValueError("Forced fused NVFP4 QDQ does not support this input/configuration") + raise ValueError( + "Forced fused NVFP4 QDQ does not support this input/configuration" + ) return "fused" if eligible and self.backend != "reference" else "reference" def _reference(self, tensor, dtype): @@ -147,56 +315,5 @@ def _compute(self, tensor, output, noop_flag=None, *, backend=None): else: torch.where(noop_flag != 1, data, output, out=output) - def quantize_impl(self, tensor, *, dtype=None): - dtype = dtype or self.dtype or tensor.dtype - if self.selected_backend(tensor, dtype) == "reference": - data = self._reference(tensor, dtype) - return self._wrap(data) - result = self._wrap(torch.empty(tensor.shape, dtype=dtype, device=tensor.device)) - self._compute(tensor, result._hp_data, backend="fused") - return result - - def make_empty( - self, - shape, - *, - dtype=torch.float32, - device=None, - requires_grad=False, - pin_memory=False, - ): - data = torch.empty( - tuple(shape), - dtype=self.dtype or dtype, - device=device or "cuda", - pin_memory=pin_memory, - ) - return self._wrap(data, requires_grad=requires_grad) - - def _wrap(self, data, requires_grad=False): - if self.internal: - return IdentityTensorStorage(hp_data=data, fake_dtype=data.dtype, quantizer=self) - return IdentityTensor( - data.shape, - data.dtype, - hp_data=data, - quantizer=self, - requires_grad=requires_grad, - device=data.device, - ) - - def update_quantized(self, src, dst, *, noop_flag=None): - if not isinstance(dst, IdentityTensorStorage) or dst._hp_data is None: - raise TypeError("NVFP4 QDQ requires allocated IdentityTensorStorage") - if src.shape != dst._hp_data.shape or src.device != dst._hp_data.device: - raise ValueError("NVFP4 QDQ source and destination must have matching shape/device") - self._compute(src, dst._hp_data, noop_flag) - return dst - def is_requantization_safe(self): return self.nvfp4_quantizer.is_requantization_safe() - - def _get_compatible_recipe(self): - from transformer_engine.common.recipe import CustomRecipe - - return CustomRecipe diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py index a01f068aec9..f93ed429c7b 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -348,15 +348,16 @@ def nvfp4_row_scaled_fwd_high_precision_bwd_factory( return _plain_nvfp4_quantizer() -def nvfp4_qdq_weight_fwd_mxfp8_bwd_factory( +def nvfp4_qdq_weight_fwd_bf16_bwd_factory( role: Optional[QuantizerRole], *, backend="auto", nvfp4_options=None ): - """High-precision X and NVFP4-QDQ W forward; MXFP8 backward. + """High-precision X and NVFP4-QDQ W forward; BF16 backward. Applies to Linear and GroupedLinear. The module compute - dtype determines forward precision. Other roles use Identity. Use - functools.partial to select backend or native NVFP4 options, including - 4over6 (reference fallback). + dtype determines high-precision arithmetic (BF16 for BF16 compute). + Backward uses original inputs and decoded QDQ weights. Other roles use + Identity. Use functools.partial to select backend or native NVFP4 options, + including 4over6 (reference fallback). """ from ..tensor.hybrid_tensor import HybridQuantizer from ..tensor.identity_tensor import IdentityQuantizer @@ -368,6 +369,43 @@ def nvfp4_qdq_weight_fwd_mxfp8_bwd_factory( if role.tensor_type == "input": return HybridQuantizer( rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="original", + ) + if role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=NVFP4QDQQuantizer( + NVFP4Quantizer(**(nvfp4_options or {})), backend=backend + ), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if role.tensor_type == "grad_output": + return IdentityQuantizer() + return IdentityQuantizer() + + +def mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory( + role: Optional[QuantizerRole], *, backend="auto", nvfp4_options=None +): + """MXFP8-QDQ X and NVFP4-QDQ W forward; native MXFP8 backward. + + Applies to Linear and GroupedLinear. The module compute + dtype determines decoded forward precision. Backward uses original + inputs and decoded QDQ weights. Other roles use Identity. Use + functools.partial to select backend or native NVFP4 options, including + 4over6 (reference fallback). + """ + from ..tensor.hybrid_tensor import HybridQuantizer + from ..tensor.identity_tensor import IdentityQuantizer + from ..tensor.nvfp4_tensor import NVFP4Quantizer + from .qdq import MXFP8QDQQuantizer, NVFP4QDQQuantizer + + if role is None or role.module_type not in ("linear", "grouped_linear"): + return IdentityQuantizer() + if role.tensor_type == "input": + return HybridQuantizer( + rowwise_quantizer=MXFP8QDQQuantizer(backend=backend), columnwise_quantizer=mxfp8_factory(role), columnwise_source="original", ) From 432e4d4d7aca135433ab927a07b80731d93c78d7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:25:23 +0000 Subject: [PATCH 4/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/mxfp8/test_mxfp8_qdq.py | 12 ++--- tests/pytorch/nvfp4/test_nvfp4_qdq.py | 49 +++++-------------- .../pytorch/custom_recipes/qdq.py | 36 ++++---------- 3 files changed, 24 insertions(+), 73 deletions(-) diff --git a/tests/pytorch/mxfp8/test_mxfp8_qdq.py b/tests/pytorch/mxfp8/test_mxfp8_qdq.py index bb012b3da2b..8c86da9a739 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_qdq.py +++ b/tests/pytorch/mxfp8/test_mxfp8_qdq.py @@ -135,9 +135,7 @@ def test_noncontiguous_and_output_dtype_fallback(): def test_scale_boundaries_and_extremes(dtype): # Adjacent representable values straddling E8M0 scale transitions. exponents = range(-120, 119, 13) if dtype == torch.bfloat16 else range(-20, 7, 2) - thresholds = torch.tensor( - [448.0 * 2.0**e for e in exponents], device="cuda", dtype=dtype - ) + thresholds = torch.tensor([448.0 * 2.0**e for e in exponents], device="cuda", dtype=dtype) maxima = torch.cat( ( torch.nextafter(thresholds, torch.zeros_like(thresholds)), @@ -290,9 +288,7 @@ def test_hybrid_backward_source(source): native = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=False, columnwise=True) expected = native(x if source == "original" else decoded) for name in ("_columnwise_data", "_columnwise_scale_inv"): - assert torch.equal( - getattr(y._columnwise_storage, name), getattr(expected, name) - ) + assert torch.equal(getattr(y._columnwise_storage, name), getattr(expected, name)) @pytest.mark.parametrize("backend", ["reference", "fused"]) @@ -362,8 +358,6 @@ def is_requantization_safe(self): torch.testing.assert_close(workspace.dequantize(), y.dequantize()) flag.zero_() q.update_quantized(x + 2, workspace, noop_flag=flag) - torch.testing.assert_close( - workspace.dequantize(), (x.detach() + 2).round().bfloat16() - ) + torch.testing.assert_close(workspace.dequantize(), (x.detach() + 2).round().bfloat16()) q.internal = True assert q(x)._hp_data.dtype == torch.bfloat16 diff --git a/tests/pytorch/nvfp4/test_nvfp4_qdq.py b/tests/pytorch/nvfp4/test_nvfp4_qdq.py index ba8ed05d71f..c792bb859d7 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_qdq.py +++ b/tests/pytorch/nvfp4/test_nvfp4_qdq.py @@ -84,9 +84,7 @@ def test_plain_exact(dtype, shape, distribution, monkeypatch): def test_4over6_fallback(mode, maximum, fast_error, monkeypatch): monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", fast_error) x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) - options = dict( - nvfp4_use_4over6=True, nvfp4_e4m3_max=maximum, nvfp4_4over6_err_mode=mode - ) + options = dict(nvfp4_use_4over6=True, nvfp4_e4m3_max=maximum, nvfp4_4over6_err_mode=mode) q = NVFP4QDQQuantizer(NVFP4Quantizer(**options)) assert q.selected_backend(x) == "reference" assert_bits(q(x).dequantize(), native_qdq(x, **options)) @@ -160,9 +158,7 @@ def test_ste_and_rng(): def mx_value(x, *, columnwise=False): - q = MXFP8Quantizer( - tex.DType.kFloat8E4M3, rowwise=not columnwise, columnwise=columnwise - ) + q = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=not columnwise, columnwise=columnwise) return q(x).dequantize() @@ -170,20 +166,14 @@ def mx_value(x, *, columnwise=False): @pytest.mark.parametrize("grouped", [False, True]) @pytest.mark.parametrize("params_dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("mixed", [False, True]) -def test_module_operands( - backend, grouped, params_dtype, mixed, nvfp4_options=None -): +def test_module_operands(backend, grouped, params_dtype, mixed, nvfp4_options=None): if backend == "fused" and params_dtype == torch.float32: pytest.skip("Mixed FP32 source/BF16 decode intentionally uses native fallback") torch.manual_seed(12) factory = ( - mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory - if mixed - else nvfp4_qdq_weight_fwd_bf16_bwd_factory - ) - recipe = CustomRecipe( - qfactory=partial(factory, backend=backend, nvfp4_options=nvfp4_options) + mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory if mixed else nvfp4_qdq_weight_fwd_bf16_bwd_factory ) + recipe = CustomRecipe(qfactory=partial(factory, backend=backend, nvfp4_options=nvfp4_options)) kwargs = dict(bias=True, params_dtype=params_dtype, device="cuda") layer = ( te.GroupedLinear(2, 128, 128, single_grouped_weight=False, **kwargs) @@ -191,21 +181,15 @@ def test_module_operands( else te.Linear(128, 128, **kwargs) ) splits = [32, 64] if grouped else [96] - x = torch.randn( - sum(splits), 128, device="cuda", dtype=torch.bfloat16, requires_grad=True - ) + x = torch.randn(sum(splits), 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) dy = torch.randn_like(x) - weights = ( - [getattr(layer, f"weight{i}") for i in range(2)] if grouped else [layer.weight] - ) + weights = [getattr(layer, f"weight{i}") for i in range(2)] if grouped else [layer.weight] biases = [getattr(layer, f"bias{i}") for i in range(2)] if grouped else [layer.bias] with te.autocast(recipe=recipe), torch.autocast("cuda", dtype=torch.bfloat16): y = layer(x, splits) if grouped else layer(x) y.backward(dy) expected_y, expected_dx = [], [] - for w, b, xs, gs in zip( - weights, biases, x.detach().split(splits), dy.split(splits) - ): + for w, b, xs, gs in zip(weights, biases, x.detach().split(splits), dy.split(splits)): wh = native_qdq(w.detach(), torch.bfloat16, **(nvfp4_options or {})) expected_y.append( torch.nn.functional.linear( @@ -236,9 +220,7 @@ def test_weight_workspace_dtype_and_reuse(): factory = mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory q = factory(QuantizerRole(module_type="linear", tensor_type="weight")) w = torch.randn(128, 256, device="cuda", dtype=torch.float32) - out, cache = quantize_weight( - tensor=w, quantizer=q, workspace_dtype=torch.bfloat16, cache=True - ) + out, cache = quantize_weight(tensor=w, quantizer=q, workspace_dtype=torch.bfloat16, cache=True) expected = native_qdq(w, torch.bfloat16) assert_bits(out._rowwise_storage.dequantize(), expected) pointer = out._rowwise_storage._hp_data.data_ptr() @@ -275,13 +257,8 @@ def test_weight_workspace_dtype_and_reuse(): workspace_dtype=torch.float32, cache=True, ) - assert ( - cache3 is not cache - and out3._rowwise_storage.dequantize().dtype == torch.float32 - ) - torch.testing.assert_close( - out3._rowwise_storage.dequantize(), native_qdq(w), rtol=0, atol=0 - ) + assert cache3 is not cache and out3._rowwise_storage.dequantize().dtype == torch.float32 + torch.testing.assert_close(out3._rowwise_storage.dequantize(), native_qdq(w), rtol=0, atol=0) def test_noncontiguous_and_fast_math_fallback(monkeypatch): @@ -387,7 +364,5 @@ def test_qdq_columnwise_only_workspace_dtype(empty): expected = mx_value(native_qdq(x, torch.bfloat16).float(), columnwise=True) # Native make_empty honors its allocation dtype; native quantize retains # the FP32 source dtype. Preserve both behaviors and compare equal decodes. - assert out._columnwise_storage._dtype == ( - torch.bfloat16 if empty else torch.float32 - ) + assert out._columnwise_storage._dtype == (torch.bfloat16 if empty else torch.float32) assert_bits(out._columnwise_storage.dequantize(dtype=torch.float32), expected) diff --git a/transformer_engine/pytorch/custom_recipes/qdq.py b/transformer_engine/pytorch/custom_recipes/qdq.py index d5382879128..6c166057d01 100644 --- a/transformer_engine/pytorch/custom_recipes/qdq.py +++ b/transformer_engine/pytorch/custom_recipes/qdq.py @@ -71,9 +71,7 @@ def quantize_impl(self, tensor, *, dtype=None): if self.selected_backend(tensor, dtype) == "reference": data = self._reference(tensor, dtype) return self._wrap(data) - result = self._wrap( - torch.empty(tensor.shape, dtype=dtype, device=tensor.device) - ) + result = self._wrap(torch.empty(tensor.shape, dtype=dtype, device=tensor.device)) self._compute(tensor, result._hp_data, backend="fused") return result @@ -96,9 +94,7 @@ def make_empty( def _wrap(self, data, requires_grad=False): if self.internal: - return IdentityTensorStorage( - hp_data=data, fake_dtype=data.dtype, quantizer=self - ) + return IdentityTensorStorage(hp_data=data, fake_dtype=data.dtype, quantizer=self) return IdentityTensor( data.shape, data.dtype, @@ -110,9 +106,7 @@ def _wrap(self, data, requires_grad=False): def update_quantized(self, src, dst, *, noop_flag=None): if not isinstance(dst, IdentityTensorStorage) or dst._hp_data is None: - raise TypeError( - f"{type(self).__name__} requires allocated IdentityTensorStorage" - ) + raise TypeError(f"{type(self).__name__} requires allocated IdentityTensorStorage") if src.shape != dst._hp_data.shape or src.device != dst._hp_data.device: raise ValueError( f"{type(self).__name__} source and destination must have matching shape/device" @@ -138,12 +132,8 @@ class MXFP8QDQQuantizer(QDQQuantizer): block_size = 32 def __init__(self, *, dtype=None, backend="auto", rowwise=True, columnwise=True): - super().__init__( - dtype=dtype, backend=backend, rowwise=rowwise, columnwise=columnwise - ) - self.mxfp8_quantizer = MXFP8Quantizer( - DType.kFloat8E4M3, rowwise=True, columnwise=False - ) + super().__init__(dtype=dtype, backend=backend, rowwise=rowwise, columnwise=columnwise) + self.mxfp8_quantizer = MXFP8Quantizer(DType.kFloat8E4M3, rowwise=True, columnwise=False) self.mxfp8_quantizer.internal = True def copy(self): @@ -169,9 +159,7 @@ def selected_backend(self, tensor, dtype=None): and os.getenv("NVTE_USE_FAST_MATH", "0") == "0" ) if self.backend == "fused" and not eligible: - raise ValueError( - "Forced fused MXFP8 QDQ does not support this input/configuration" - ) + raise ValueError("Forced fused MXFP8 QDQ does not support this input/configuration") return "fused" if eligible and self.backend != "reference" else "reference" def _reference(self, tensor, dtype): @@ -218,14 +206,10 @@ def __init__( rowwise=True, columnwise=True, ): - super().__init__( - dtype=dtype, backend=backend, rowwise=rowwise, columnwise=columnwise - ) + super().__init__(dtype=dtype, backend=backend, rowwise=rowwise, columnwise=columnwise) if quantizer is not None and not isinstance(quantizer, NVFP4Quantizer): raise TypeError("quantizer must be an NVFP4Quantizer") - self.nvfp4_quantizer = ( - quantizer.copy() if quantizer is not None else NVFP4Quantizer() - ) + self.nvfp4_quantizer = quantizer.copy() if quantizer is not None else NVFP4Quantizer() self.nvfp4_quantizer.set_usage(rowwise=True, columnwise=False) self.nvfp4_quantizer.internal = True @@ -294,9 +278,7 @@ def selected_backend(self, tensor, dtype=None): and os.getenv("NVTE_USE_FAST_MATH", "0") == "0" ) if self.backend == "fused" and not eligible: - raise ValueError( - "Forced fused NVFP4 QDQ does not support this input/configuration" - ) + raise ValueError("Forced fused NVFP4 QDQ does not support this input/configuration") return "fused" if eligible and self.backend != "reference" else "reference" def _reference(self, tensor, dtype): From df8269c92d4aa68a909c4f6583fe618237f7b7cb Mon Sep 17 00:00:00 2001 From: Evgeny Date: Mon, 12 Oct 2026 13:47:55 +0000 Subject: [PATCH 5/9] Clarify QDQ operand labels and factory naming Signed-off-by: Evgeny --- .../fine_grained_quantization.rst | 11 ++++++----- tests/pytorch/nvfp4/test_nvfp4_qdq.py | 10 ++++++---- .../pytorch/custom_recipes/quantizer_factory_zoo.py | 2 +- 3 files changed, 13 insertions(+), 10 deletions(-) diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst index d5434cb44d7..78aee2c4383 100644 --- a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -402,11 +402,12 @@ Forward and backward examples The following factories support ``Linear`` and ``GroupedLinear`` with discrete expert parameters. With BF16 module compute precision: -* ``nvfp4_qdq_weight_fwd_bf16_bwd_factory`` models ``X[BF16] x W[NVFP4]`` +* ``nvfp4_qdq_weight_fwd_bf16_bwd_factory`` models ``X[BF16] x W[NVFP4-QDQ]`` forward and uses BF16 backward. Backward consumes the original inputs and the decoded forward QDQ weights. -* ``mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory`` models ``X[MXFP8] x W[NVFP4]`` - forward and uses native MXFP8 backward. Input QDQ uses block-32 E4M3/E8M0; +* ``mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory`` models + ``X[MXFP8-QDQ] x W[NVFP4-QDQ]`` forward and uses native MXFP8 backward. + Input QDQ uses block-32 E4M3/E8M0; weight QDQ uses block-16 E2M1/E4M3 with a tensor scale. Backward quantizes the original inputs and the decoded forward QDQ weights to MXFP8. @@ -415,10 +416,10 @@ For example:: from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( nvfp4_qdq_weight_fwd_bf16_bwd_factory, - mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory, + mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory, ) - recipe = CustomRecipe(qfactory=mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory) + recipe = CustomRecipe(qfactory=mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory) To use BF16 inputs while retaining MXFP8 backward, replace the second factory's input rowwise ``MXFP8QDQQuantizer`` with ``IdentityQuantizer``. The input diff --git a/tests/pytorch/nvfp4/test_nvfp4_qdq.py b/tests/pytorch/nvfp4/test_nvfp4_qdq.py index c792bb859d7..d67eecf73ad 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_qdq.py +++ b/tests/pytorch/nvfp4/test_nvfp4_qdq.py @@ -15,7 +15,7 @@ from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.custom_recipes.qdq import NVFP4QDQQuantizer from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( - mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory, + mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory, nvfp4_qdq_weight_fwd_bf16_bwd_factory, ) @@ -171,7 +171,9 @@ def test_module_operands(backend, grouped, params_dtype, mixed, nvfp4_options=No pytest.skip("Mixed FP32 source/BF16 decode intentionally uses native fallback") torch.manual_seed(12) factory = ( - mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory if mixed else nvfp4_qdq_weight_fwd_bf16_bwd_factory + mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory + if mixed + else nvfp4_qdq_weight_fwd_bf16_bwd_factory ) recipe = CustomRecipe(qfactory=partial(factory, backend=backend, nvfp4_options=nvfp4_options)) kwargs = dict(bias=True, params_dtype=params_dtype, device="cuda") @@ -217,7 +219,7 @@ def test_weight_workspace_dtype_and_reuse(): from transformer_engine.pytorch.module.base import quantize_weight from transformer_engine.pytorch.quantization import QuantizerRole - factory = mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory + factory = mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory q = factory(QuantizerRole(module_type="linear", tensor_type="weight")) w = torch.randn(128, 256, device="cuda", dtype=torch.float32) out, cache = quantize_weight(tensor=w, quantizer=q, workspace_dtype=torch.bfloat16, cache=True) @@ -348,7 +350,7 @@ def test_qdq_columnwise_only_workspace_dtype(empty): """A transient forward QDQ source must round to the requested compute dtype.""" from transformer_engine.pytorch.quantization import QuantizerRole - q = mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory( + q = mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory( QuantizerRole(module_type="linear", tensor_type="weight") ) q.set_usage(rowwise=False, columnwise=True) diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py index f93ed429c7b..8497a46c2fb 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -385,7 +385,7 @@ def nvfp4_qdq_weight_fwd_bf16_bwd_factory( return IdentityQuantizer() -def mxfp8_nvfp4_qdq_fwd_mxfp8_bwd_factory( +def mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory( role: Optional[QuantizerRole], *, backend="auto", nvfp4_options=None ): """MXFP8-QDQ X and NVFP4-QDQ W forward; native MXFP8 backward. From 1171a0519cfefee0c804b397c654b781e97b15eb Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 13 Oct 2026 09:41:32 +0000 Subject: [PATCH 6/9] Add per-token outer scale to the fused nvfp4 path Signed-off-by: Evgeny --- .../fine_grained_quantization.rst | 10 +- tests/cpp/operator/test_nvfp4_qdq.cu | 27 ++++- tests/pytorch/nvfp4/test_nvfp4_qdq.py | 96 +++++++++++++--- transformer_engine/common/cast/nvfp4/qdq.cu | 107 +++++++++++++----- .../common/include/transformer_engine/cast.h | 10 ++ transformer_engine/pytorch/csrc/extensions.h | 2 + .../pytorch/csrc/extensions/cast.cpp | 18 +++ .../pytorch/csrc/extensions/pybind.cpp | 3 + .../pytorch/custom_recipes/qdq.py | 7 +- 9 files changed, 230 insertions(+), 50 deletions(-) diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst index 78aee2c4383..1331ed8515e 100644 --- a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -436,7 +436,7 @@ Backends and custom QDQ formats Output defaults to the source dtype; a module compute-dtype override changes only the decoded dtype. For ``NVFP4QDQQuantizer``, ``backend="auto"`` fuses -plain 1D E2M1/E4M3 QDQ on +plain 1D E2M1/E4M3 QDQ with tensor or per-row outer scaling on SM100 for contiguous, 16-byte-aligned BF16/FP16 tensors with matching output dtype and flattened matrix dimensions divisible by 16, with ordinary fast math disabled. Other native-valid modes, including 4over6, 2D scaling, distributed @@ -444,6 +444,14 @@ amax reduction, and mixed dtypes, use native ``quantize().dequantize()``. ``backend="reference"`` forces that composition; ``backend="fused"`` rejects unsupported cases. Current weight amax is computed when quantization runs. +For per-token outer scaling, wrap +``NVFP4Quantizer(row_scaled_nvfp4=True)`` in ``NVFP4QDQQuantizer``. Each flattened +row uses its own outer amax, while each 16-value block retains its E4M3 scale. +The fused row-scaled path computes row amax and QDQ in one kernel, without +intermediate amax, packed-value, or scale buffers. The tensor-scaled fused path +retains its separate global amax reduction. Both paths preserve native Q+DQ +rounding and support the same matching BF16/FP16 input and output dtypes. + ``custom_recipes.qdq.MXFP8QDQQuantizer`` provides the corresponding rowwise E4M3/block-32 E8M0 QDQ operation. Its ``reference`` backend calls native ``MXFP8Quantizer.quantize().dequantize()``. Its ``fused`` backend supports diff --git a/tests/cpp/operator/test_nvfp4_qdq.cu b/tests/cpp/operator/test_nvfp4_qdq.cu index 98f6551d6f9..00c4b885b92 100644 --- a/tests/cpp/operator/test_nvfp4_qdq.cu +++ b/tests/cpp/operator/test_nvfp4_qdq.cu @@ -19,7 +19,7 @@ using namespace test; namespace { template -void check_qdq(size_t rows, size_t cols) { +void check_qdq(size_t rows, size_t cols, bool row_scaled = false) { cudaDeviceProp prop; ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); if (prop.major != 10 || prop.minor != 0) GTEST_SKIP() << "QDQ targets SM100"; @@ -39,12 +39,17 @@ void check_qdq(size_t rows, size_t cols) { maximum = std::max(maximum, std::abs(static_cast(input.rowwise_cpu_dptr()[i]))); } packed.set_amax(maximum); + if (row_scaled) packed.set_row_scaled_nvfp4(true); ASSERT_EQ(cudaMemcpy(amax.rowwise_dptr(), &maximum, sizeof(float), cudaMemcpyHostToDevice), cudaSuccess); QuantizationConfigWrapper config; nvte_quantize_v2(input.data(), packed.data(), config, 0); nvte_dequantize(packed.data(), reference.data(), 0); - nvte_nvfp4_qdq(input.data(), output.data(), amax.data(), nullptr, 0); + if (row_scaled) { + nvte_nvfp4_qdq_row_scaled(input.data(), output.data(), nullptr, 0); + } else { + nvte_nvfp4_qdq(input.data(), output.data(), amax.data(), nullptr, 0); + } ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); reference.to_cpu(); output.to_cpu(); @@ -54,7 +59,11 @@ void check_qdq(size_t rows, size_t cols) { ASSERT_EQ(cudaMemcpy(noop.rowwise_dptr(), &one, sizeof(float), cudaMemcpyHostToDevice), cudaSuccess); ASSERT_EQ(cudaMemset(input.rowwise_dptr(), 0, rows * cols * sizeof(T)), cudaSuccess); - nvte_nvfp4_qdq(input.data(), output.data(), amax.data(), noop.data(), 0); + if (row_scaled) { + nvte_nvfp4_qdq_row_scaled(input.data(), output.data(), noop.data(), 0); + } else { + nvte_nvfp4_qdq(input.data(), output.data(), amax.data(), noop.data(), 0); + } ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); output.to_cpu(); EXPECT_EQ(std::memcmp(reference.rowwise_cpu_dptr(), output.rowwise_cpu_dptr(), @@ -72,3 +81,15 @@ TEST(NVFP4QDQ, FP16) { check_qdq(32, 32); check_qdq(96, 160); } + +TEST(NVFP4QDQ, RowScaledBF16) { + check_qdq(32, 32, true); + check_qdq(96, 160, true); + check_qdq(128, 4096, true); +} + +TEST(NVFP4QDQ, RowScaledFP16) { + check_qdq(32, 32, true); + check_qdq(96, 160, true); + check_qdq(32, 16384, true); +} diff --git a/tests/pytorch/nvfp4/test_nvfp4_qdq.py b/tests/pytorch/nvfp4/test_nvfp4_qdq.py index d67eecf73ad..7e3a95f163b 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_qdq.py +++ b/tests/pytorch/nvfp4/test_nvfp4_qdq.py @@ -39,7 +39,8 @@ def assert_bits(actual, expected): @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("shape", [(16, 16), (48, 80), (128, 256), (512, 1024)]) @pytest.mark.parametrize("distribution", ["normal", "zeros", "outliers", "boundaries"]) -def test_plain_exact(dtype, shape, distribution, monkeypatch): +@pytest.mark.parametrize("row_scaled", [False, True]) +def test_plain_exact(dtype, shape, distribution, row_scaled, monkeypatch): monkeypatch.setenv("NVTE_USE_FAST_MATH", "0") torch.manual_seed(42) x = torch.randn(shape, device="cuda", dtype=dtype) @@ -73,18 +74,44 @@ def test_plain_exact(dtype, shape, distribution, monkeypatch): dtype=dtype, ) x.copy_(values.repeat(shape[0], shape[1] // 16)) - q = NVFP4QDQQuantizer(backend="fused") + q = NVFP4QDQQuantizer(NVFP4Quantizer(row_scaled_nvfp4=row_scaled), backend="fused") assert q.selected_backend(x) == "fused" - assert_bits(q(x).dequantize(), native_qdq(x)) + assert_bits(q(x).dequantize(), native_qdq(x, row_scaled_nvfp4=row_scaled)) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("shape", [(16, 4096), (32, 6144), (16, 16384), (2, 8, 80)]) +def test_row_scaled_wide_and_batched(dtype, shape): + torch.manual_seed(27) + x = torch.randn(shape, device="cuda", dtype=dtype) + rows = x.reshape(-1, shape[-1]) + # Separate row scales: a zero row, a tiny row, and a large outlier row. + rows[0].zero_() + rows[0, ::2] = -0.0 + rows[1].mul_(1e-4) + rows[2, 0] = 4096 + q = NVFP4QDQQuantizer(NVFP4Quantizer(row_scaled_nvfp4=True), backend="fused") + actual = q(x).dequantize() + assert_bits(actual, native_qdq(x, row_scaled_nvfp4=True)) + # Changing one row's outer scale must not affect any other row. + previous = actual.reshape_as(rows)[1:].clone() + rows[0].fill_(8192) + assert_bits(q(x).dequantize().reshape_as(rows)[1:], previous) @pytest.mark.parametrize("mode", ["MAE", "MSE"]) @pytest.mark.parametrize("maximum", [256, 448]) @pytest.mark.parametrize("fast_error", ["0", "1"]) -def test_4over6_fallback(mode, maximum, fast_error, monkeypatch): +@pytest.mark.parametrize("row_scaled", [False, True]) +def test_4over6_fallback(mode, maximum, fast_error, row_scaled, monkeypatch): monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", fast_error) x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) - options = dict(nvfp4_use_4over6=True, nvfp4_e4m3_max=maximum, nvfp4_4over6_err_mode=mode) + options = dict( + nvfp4_use_4over6=True, + nvfp4_e4m3_max=maximum, + nvfp4_4over6_err_mode=mode, + row_scaled_nvfp4=row_scaled, + ) q = NVFP4QDQQuantizer(NVFP4Quantizer(**options)) assert q.selected_backend(x) == "reference" assert_bits(q(x).dequantize(), native_qdq(x, **options)) @@ -93,7 +120,9 @@ def test_4over6_fallback(mode, maximum, fast_error, monkeypatch): q(x) -@pytest.mark.parametrize("options", [{"with_2d_quantization": True}, {}]) +@pytest.mark.parametrize( + "options", [{"with_2d_quantization": True}, {"row_scaled_nvfp4": True}, {}] +) def test_dtype_fallback(options): x = torch.randn(128, 256, device="cuda", dtype=torch.float32) q = NVFP4QDQQuantizer(NVFP4Quantizer(**options)) @@ -105,22 +134,49 @@ def test_dtype_fallback(options): assert_bits(actual.dequantize(), native_qdq(x * 2, torch.bfloat16, **options)) +def test_row_scaled_ste_and_dispatch(monkeypatch): + x = torch.randn(16, 80, device="cuda", dtype=torch.bfloat16, requires_grad=True) + q = NVFP4QDQQuantizer(NVFP4Quantizer(row_scaled_nvfp4=True)) + assert q.selected_backend(x) == "fused" + q(x).sum().backward() + assert torch.equal(x.grad, torch.ones_like(x)) + monkeypatch.setenv("NVTE_USE_FAST_MATH", "1") + assert q.selected_backend(x) == "reference" + q.backend = "fused" + with pytest.raises(ValueError, match="Forced fused"): + q(x) + + +def test_row_scaled_low_level_validation(): + x = torch.randn(16, 80, device="cuda", dtype=torch.bfloat16) + output = torch.empty_like(x) + tex.nvfp4_qdq_row_scaled(x, output, None) + assert_bits(output, native_qdq(x, row_scaled_nvfp4=True)) + with pytest.raises(RuntimeError, match="matching dtype and shape"): + tex.nvfp4_qdq_row_scaled(x, output.float(), None) + with pytest.raises(RuntimeError, match="contiguous"): + tex.nvfp4_qdq_row_scaled(x.t(), output.t(), None) + with pytest.raises(RuntimeError, match="scalar FP32"): + tex.nvfp4_qdq_row_scaled(x, output, torch.ones(2, device="cuda")) + + @pytest.mark.parametrize("backend", ["reference", "fused"]) @pytest.mark.parametrize("internal", [False, True]) -def test_lifecycle_and_graph(backend, internal): +@pytest.mark.parametrize("row_scaled", [False, True]) +def test_lifecycle_and_graph(backend, internal, row_scaled): x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) - q = NVFP4QDQQuantizer(backend=backend) + q = NVFP4QDQQuantizer(NVFP4Quantizer(row_scaled_nvfp4=row_scaled), backend=backend) q.internal = internal y = q(x) assert y._get_quantizer().backend == backend address = y._hp_data.data_ptr() noop = torch.ones(1, device="cuda") q.update_quantized(x * 3, y, noop_flag=noop) - assert_bits(y.dequantize(), native_qdq(x)) + assert_bits(y.dequantize(), native_qdq(x, row_scaled_nvfp4=row_scaled)) noop.zero_() q.update_quantized(x * 2, y, noop_flag=noop) assert y._hp_data.data_ptr() == address - assert_bits(y.dequantize(), native_qdq(x * 2)) + assert_bits(y.dequantize(), native_qdq(x * 2, row_scaled_nvfp4=row_scaled)) # Warm up on a side stream, then replay with changed input and noop flag. stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) @@ -133,7 +189,7 @@ def test_lifecycle_and_graph(backend, internal): q.update_quantized(x, y, noop_flag=noop) x.mul_(2) graph.replay() - assert_bits(y.dequantize(), native_qdq(x)) + assert_bits(y.dequantize(), native_qdq(x, row_scaled_nvfp4=row_scaled)) previous = y.dequantize().clone() noop.fill_(1) x.mul_(2) @@ -283,21 +339,27 @@ def test_empty_reference(shape): @pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")]) -def test_exceptional_values(value): - x = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("row_scaled", [False, True]) +def test_exceptional_values(value, dtype, row_scaled): + x = torch.randn(128, 128, device="cuda", dtype=dtype) x[0, 0] = value - actual = NVFP4QDQQuantizer(backend="fused")(x).dequantize() - expected = native_qdq(x) + actual = NVFP4QDQQuantizer(NVFP4Quantizer(row_scaled_nvfp4=row_scaled), backend="fused")( + x + ).dequantize() + expected = native_qdq(x, row_scaled_nvfp4=row_scaled) assert torch.equal(torch.isnan(actual), torch.isnan(expected)) mask = ~torch.isnan(expected) assert_bits(actual[mask], expected[mask]) -def test_current_stream_and_device(): +@pytest.mark.parametrize("row_scaled", [False, True]) +def test_current_stream_and_device(row_scaled): device = 1 if torch.cuda.device_count() > 1 else 0 with torch.cuda.device(device), torch.cuda.stream(torch.cuda.Stream(device=device)): x = torch.randn(48, 80, device=f"cuda:{device}", dtype=torch.bfloat16) - assert_bits(NVFP4QDQQuantizer(backend="fused")(x).dequantize(), native_qdq(x)) + q = NVFP4QDQQuantizer(NVFP4Quantizer(row_scaled_nvfp4=row_scaled), backend="fused") + assert_bits(q(x).dequantize(), native_qdq(x, row_scaled_nvfp4=row_scaled)) def test_copy_preserves_runtime_reduction_state(): diff --git a/transformer_engine/common/cast/nvfp4/qdq.cu b/transformer_engine/common/cast/nvfp4/qdq.cu index a7d884cbdc6..3fbd6a99edd 100644 --- a/transformer_engine/common/cast/nvfp4/qdq.cu +++ b/transformer_engine/common/cast/nvfp4/qdq.cu @@ -8,34 +8,35 @@ #include "../../common.h" #include "../../util/cuda_runtime.h" +#include "../../utils.cuh" #include "core_nvfp4.cuh" namespace transformer_engine { namespace { #if FP4_TYPE_SUPPORTED -// One thread owns a complete 1x16 scale block. Both rounded encodings stay in -// registers; only the decoded high-precision values are written to memory. -template -__global__ void qdq_kernel(const T *input, T *output, const float *amax, const float *noop, - size_t blocks) { - if (noop != nullptr && *noop == 1.0f) return; - const size_t block = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (block >= blocks) return; +// Shared block processing for both outer-scale modes. Both rounded encodings +// stay in registers; only the decoded high-precision values are written. +template +__device__ __forceinline__ void qdq_block(const T *input, T *output, float outer_amax, + float encode) { Vec values; - values.load_from(input + block * 16); + values.load_from(input); float local_amax = 0.0f; #pragma unroll for (int i = 0; i < 16; ++i) { local_amax = fmaxf(local_amax, fabsf(static_cast(values.data.elt[i]))); } using namespace dispatch::nvfp4; - const float global_amax = *amax; - const float encode = core::compute_global_encode_scaling_factor_FP4(global_amax); + if constexpr (!ROW_SCALED_NVFP4) { + // Keep input loads ahead of the encode-scale division, as in the original + // tensor-scaled kernel. Row scaling reuses its CTA's precomputed encode. + encode = core::compute_global_encode_scaling_factor_FP4(outer_amax); + } const fp8e4m3 scale = core::compute_decoding_scaling_factor(local_amax, encode); const float inverse = fminf(1.0f / (static_cast(scale) * (1.0f / encode)), detail::TypeExtrema::max); // Match native dequantize multiplication order, including zero signs. - const float decode = static_cast(scale) * global_amax * (1.0f / (6.0f * 448.0f)); + const float decode = static_cast(scale) * outer_amax * (1.0f / (6.0f * 448.0f)); #pragma unroll for (int i = 0; i < 16; i += 4) { fp4e2m1x4 packed; @@ -51,24 +52,54 @@ __global__ void qdq_kernel(const T *input, T *output, const float *amax, const f values.data.elt[i + 2] = static_cast(decoded.z * decode); values.data.elt[i + 3] = static_cast(decoded.w * decode); } - values.store_to(output + block * 16); + values.store_to(output); +} + +// Tensor scaling distributes blocks across CTAs and reads a precomputed amax. +// Row scaling gives each CTA a complete row and fuses its outer amax reduction. +template +__global__ void qdq_kernel(const T *input, T *output, const float *amax, const float *noop, + size_t blocks) { + if (noop != nullptr && *noop == 1.0f) return; + using namespace dispatch::nvfp4; + if constexpr (ROW_SCALED_NVFP4) { + const size_t row_offset = static_cast(blockIdx.x) * blocks * 16; + float thread_amax = 0.0f; + for (size_t block = threadIdx.x; block < blocks; block += THREADS) { + Vec values; + values.load_from(input + row_offset + block * 16); +#pragma unroll + for (int i = 0; i < 16; ++i) { + thread_amax = fmaxf(thread_amax, fabsf(static_cast(values.data.elt[i]))); + } + } + const float maximum = reduce_max(thread_amax, threadIdx.x / 32); + __shared__ float row_amax; + if (threadIdx.x == 0) row_amax = maximum; + __syncthreads(); + const float encode = core::compute_global_encode_scaling_factor_FP4(row_amax); + for (size_t block = threadIdx.x; block < blocks; block += THREADS) { + qdq_block(input + row_offset + block * 16, + output + row_offset + block * 16, row_amax, encode); + } + } else { + const size_t block = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (block >= blocks) return; + const float outer_amax = *amax; + qdq_block(input + block * 16, output + block * 16, outer_amax, 0.0f); + } } #endif -} // namespace -} // namespace transformer_engine -void nvte_nvfp4_qdq(const NVTETensor input, NVTETensor output, const NVTETensor amax, - const NVTETensor noop, cudaStream_t stream) { - NVTE_API_CALL(nvte_nvfp4_qdq); - using namespace transformer_engine; +template +void launch_qdq(const NVTETensor input, NVTETensor output, const NVTETensor amax, + const NVTETensor noop, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED NVTE_CHECK(cuda::sm_arch() == 100, "NVFP4 QDQ currently supports SM100"); const auto &in = *convertNVTETensorCheck(input); auto &out = *convertNVTETensorCheck(output); - const auto &amax_tensor = *convertNVTETensorCheck(amax); CheckInputTensor(in, "input"); CheckOutputTensor(out, "output"); - CheckInputTensor(amax_tensor, "amax"); NVTE_CHECK(in.data.dtype == out.data.dtype && in.data.shape == out.data.shape, "QDQ input and output must have matching dtype and shape"); NVTE_CHECK(in.data.dtype == DType::kBFloat16 || in.data.dtype == DType::kFloat16, @@ -76,8 +107,14 @@ void nvte_nvfp4_qdq(const NVTETensor input, NVTETensor output, const NVTETensor NVTE_CHECK(in.data.shape.size() >= 2 && in.numel() > 0 && in.data.shape.back() % 16 == 0 && (in.numel() / in.data.shape.back()) % 16 == 0, "QDQ requires nonempty block-aligned rowwise input"); - NVTE_CHECK(amax_tensor.data.dtype == DType::kFloat32 && amax_tensor.numel() == 1, - "QDQ requires a scalar FP32 amax"); + const float *amax_ptr = nullptr; + if constexpr (!ROW_SCALED_NVFP4) { + const auto &amax_tensor = *convertNVTETensorCheck(amax); + CheckInputTensor(amax_tensor, "amax"); + NVTE_CHECK(amax_tensor.data.dtype == DType::kFloat32 && amax_tensor.numel() == 1, + "QDQ requires a scalar FP32 amax"); + amax_ptr = static_cast(amax_tensor.data.dptr); + } NVTE_CHECK(reinterpret_cast(in.data.dptr) % 16 == 0 && reinterpret_cast(out.data.dptr) % 16 == 0, "QDQ requires 16-byte aligned buffers"); @@ -90,14 +127,30 @@ void nvte_nvfp4_qdq(const NVTETensor input, NVTETensor output, const NVTETensor noop_ptr = static_cast(flag.data.dptr); } constexpr int threads = 128; - const size_t blocks = in.numel() / 16; + const size_t blocks = ROW_SCALED_NVFP4 ? in.data.shape.back() / 16 : in.numel() / 16; + const size_t ctas = + ROW_SCALED_NVFP4 ? in.numel() / in.data.shape.back() : (blocks + threads - 1) / threads; TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT( in.data.dtype, T, - qdq_kernel<<<(blocks + threads - 1) / threads, threads, 0, stream>>>( - static_cast(in.data.dptr), static_cast(out.data.dptr), - static_cast(amax_tensor.data.dptr), noop_ptr, blocks);); + qdq_kernel + <<>>(static_cast(in.data.dptr), + static_cast(out.data.dptr), amax_ptr, noop_ptr, blocks);); NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("NVFP4 QDQ requires a CUDA toolkit with FP4 support"); #endif } +} // namespace +} // namespace transformer_engine + +void nvte_nvfp4_qdq(const NVTETensor input, NVTETensor output, const NVTETensor amax, + const NVTETensor noop, cudaStream_t stream) { + NVTE_API_CALL(nvte_nvfp4_qdq); + transformer_engine::launch_qdq(input, output, amax, noop, stream); +} + +void nvte_nvfp4_qdq_row_scaled(const NVTETensor input, NVTETensor output, const NVTETensor noop, + cudaStream_t stream) { + NVTE_API_CALL(nvte_nvfp4_qdq_row_scaled); + transformer_engine::launch_qdq(input, output, nullptr, noop, stream); +} diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 09ad3f24d17..de20c66d202 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -29,6 +29,16 @@ extern "C" { void nvte_nvfp4_qdq(const NVTETensor input, NVTETensor output, const NVTETensor amax, const NVTETensor noop, cudaStream_t stream); +/*! \brief Plain 1x16 NVFP4 QDQ with per-row outer scaling into BF16/FP16. + * Input/output have the same layout and dtype requirements as nvte_nvfp4_qdq. + * Each flattened row (token) computes its own outer amax inside the kernel. + * E2M1 values and E4M3 scales (maximum 448) use native rounding with fast math + * disabled. No intermediate amax, packed values, or scales are materialized. + * Optional scalar FP32 noop skips writes when equal to 1. Buffers are caller-owned. + */ +void nvte_nvfp4_qdq_row_scaled(const NVTETensor input, NVTETensor output, const NVTETensor noop, + cudaStream_t stream); + /*! \brief Rowwise block-32 E4M3/E8M0 MXFP8 QDQ into matching BF16/FP16 output. * Input and output are contiguous, with rows divisible by 32. An optional scalar * FP32 noop skips writes when equal to 1. All buffers are caller-owned. diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index cee3f572ac1..6254d35b853 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -504,6 +504,8 @@ at::Tensor scaled_aligned_causal_masked_softmax_backward(at::Tensor output_grads void nvfp4_qdq(const at::Tensor &input, at::Tensor output, at::Tensor amax, std::optional noop, bool compute_amax); +void nvfp4_qdq_row_scaled(const at::Tensor &input, at::Tensor output, + std::optional noop); void mxfp8_qdq(const at::Tensor &input, at::Tensor output, std::optional noop); void compute_amax(const at::Tensor &tensor, at::Tensor &amax); diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 69ea9ec09d3..783a40b22e0 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -83,6 +83,24 @@ void nvfp4_qdq(const at::Tensor &input, at::Tensor output, at::Tensor amax, stream); } +void nvfp4_qdq_row_scaled(const at::Tensor &input, at::Tensor output, + std::optional noop) { + TORCH_CHECK(input.is_cuda() && output.is_cuda() && input.device() == output.device(), + "Row-scaled NVFP4 QDQ requires CUDA tensors on the same device"); + at::cuda::CUDAGuard guard(input.device()); + TORCH_CHECK(input.is_contiguous() && output.is_contiguous(), + "Row-scaled NVFP4 QDQ requires contiguous tensors"); + TORCH_CHECK( + !noop || (noop->is_cuda() && noop->device() == input.device() && noop->is_contiguous()), + "Invalid QDQ noop device/layout"); + auto in = makeTransformerEngineTensor(input); + auto out = makeTransformerEngineTensor(output); + std::optional noop_cpp; + if (noop) noop_cpp = makeTransformerEngineTensor(*noop); + nvte_nvfp4_qdq_row_scaled(in.data(), out.data(), noop_cpp ? noop_cpp->data() : nullptr, + at::cuda::getCurrentCUDAStream()); +} + void mxfp8_qdq(const at::Tensor &input, at::Tensor output, std::optional noop) { TORCH_CHECK(input.is_cuda() && output.is_cuda() && input.device() == output.device(), "MXFP8 QDQ requires tensors on the same CUDA device"); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index cf4c067bf77..79feb147898 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -480,6 +480,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("nvfp4_qdq", &transformer_engine::pytorch::nvfp4_qdq, "Plain rowwise NVFP4 QDQ", py::arg("input"), py::arg("output"), py::arg("amax"), py::arg("noop") = py::none(), py::arg("compute_amax") = true, py::call_guard()); + m.def("nvfp4_qdq_row_scaled", &transformer_engine::pytorch::nvfp4_qdq_row_scaled, + "Plain NVFP4 QDQ with per-row outer scaling", py::arg("input"), py::arg("output"), + py::arg("noop") = py::none(), py::call_guard()); m.def("mxfp8_qdq", &transformer_engine::pytorch::mxfp8_qdq, "Rowwise block-32 MXFP8 QDQ", py::arg("input"), py::arg("output"), py::arg("noop") = py::none(), py::call_guard()); diff --git a/transformer_engine/pytorch/custom_recipes/qdq.py b/transformer_engine/pytorch/custom_recipes/qdq.py index 6c166057d01..51f58448a2d 100644 --- a/transformer_engine/pytorch/custom_recipes/qdq.py +++ b/transformer_engine/pytorch/custom_recipes/qdq.py @@ -191,7 +191,8 @@ class NVFP4QDQQuantizer(QDQQuantizer): serve either GEMM orientation. ``backend`` is ``auto``, ``reference``, or ``fused``. Auto falls back to - native Q+DQ outside plain contiguous BF16/FP16 1D E4M3 scaling on SM100. + native Q+DQ outside plain contiguous BF16/FP16 1D E4M3 scaling on SM100, + with tensor or per-row outer scaling. Per-row amax is fused into QDQ. In particular, 4over6 and 2D scaling currently use the reference path. Forced fused execution raises for unsupported inputs/configurations. This experimental API is subject to change. @@ -269,7 +270,6 @@ def selected_backend(self, tensor, dtype=None): q.with_post_rht_amax, q.with_2d_quantization, q.stochastic_rounding, - q.row_scaled_nvfp4, q.nvfp4_use_4over6, q.disable_second_level_scale, q.with_amax_reduction, @@ -286,6 +286,9 @@ def _reference(self, tensor, dtype): def _compute(self, tensor, output, noop_flag=None, *, backend=None): if (backend or self.selected_backend(tensor, output.dtype)) == "fused": + if self.nvfp4_quantizer.row_scaled_nvfp4: + tex.nvfp4_qdq_row_scaled(tensor, output, noop_flag) + return # Framework owns both allocations; amax reduction and QDQ use the # current stream and may be captured together. amax = torch.empty(1, dtype=torch.float32, device=tensor.device) From f51212d10e113956f9c919716326683bc2398de4 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 13 Oct 2026 16:05:31 +0000 Subject: [PATCH 7/9] Fix HybridQuantizer dtype validation and column source precision Signed-off-by: Evgeny --- tests/pytorch/nvfp4/test_nvfp4_qdq.py | 101 +++++++++++++ tests/pytorch/test_hybrid_quantization.py | 133 ++++++++++++++++++ .../pytorch/tensor/hybrid_tensor.py | 29 +++- 3 files changed, 261 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_qdq.py b/tests/pytorch/nvfp4/test_nvfp4_qdq.py index 7e3a95f163b..331bdc97a01 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_qdq.py +++ b/tests/pytorch/nvfp4/test_nvfp4_qdq.py @@ -319,6 +319,44 @@ def test_weight_workspace_dtype_and_reuse(): torch.testing.assert_close(out3._rowwise_storage.dequantize(), native_qdq(w), rtol=0, atol=0) +@pytest.mark.parametrize("mixed", [False, True]) +def test_hybrid_qdq_workspace_checks_child_dtype(mixed): + """An FP32 Hybrid must not accept an explicit FP32 request for an FP16 QDQ row.""" + from transformer_engine.pytorch.module.base import quantize_weight + from transformer_engine.pytorch.quantization import QuantizerRole + + factory = ( + mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory + if mixed + else nvfp4_qdq_weight_fwd_bf16_bwd_factory + ) + q = factory(QuantizerRole(module_type="linear", tensor_type="weight"), backend="reference") + q.rowwise_quantizer.dtype = torch.float16 + w = torch.randn(128, 256, device="cuda", dtype=torch.float32) * 0.023 + cache = q.quantize(w) + assert cache._dtype == torch.float32 + assert cache._rowwise_storage._hp_data.dtype == torch.float16 + before = cache._rowwise_storage._hp_data.clone() + assert not q.is_compatible_output(cache, dtype=torch.float32) + with pytest.raises(ValueError, match="dtype does not match"): + q.quantize(w * 3, out=cache, dtype=torch.float32) + assert_bits(cache._rowwise_storage._hp_data, before) + + out, new_cache = quantize_weight( + tensor=w * 3, + quantizer=q, + workspace=cache, + workspace_dtype=torch.float32, + cache=True, + ) + assert new_cache is out and out is not cache + assert q.rowwise_quantizer.dtype == torch.float16 + expected_row = native_qdq(w * 3) + assert_bits(out._rowwise_storage.dequantize(), expected_row) + expected_col = mx_value(expected_row, columnwise=True) if mixed else expected_row + assert_bits(out._columnwise_storage.dequantize(), expected_col) + + def test_noncontiguous_and_fast_math_fallback(monkeypatch): x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16).T q = NVFP4QDQQuantizer() @@ -430,3 +468,66 @@ def test_qdq_columnwise_only_workspace_dtype(empty): # the FP32 source dtype. Preserve both behaviors and compare equal decodes. assert out._columnwise_storage._dtype == (torch.bfloat16 if empty else torch.float32) assert_bits(out._columnwise_storage.dequantize(dtype=torch.float32), expected) + + +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize("rowwise", [False, True]) +@pytest.mark.parametrize( + "source_dtype,output_dtype", + [(torch.bfloat16, torch.float32), (torch.float16, torch.bfloat16)], +) +def test_qdq_columnwise_source_preserves_forward_precision( + mixed, rowwise, source_dtype, output_dtype +): + """Column construction must not narrow a decoded forward value to the source dtype.""" + from transformer_engine.pytorch.quantization import QuantizerRole + + factory = ( + mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory + if mixed + else nvfp4_qdq_weight_fwd_bf16_bwd_factory + ) + q = factory(QuantizerRole(module_type="linear", tensor_type="weight"), backend="reference") + q.set_usage(rowwise=rowwise, columnwise=True) + if source_dtype == torch.float16: + x = torch.full((128, 128), 65504.0, device="cuda", dtype=source_dtype) + else: + generator = torch.Generator(device="cuda").manual_seed(0) + x = (torch.randn(128, 128, device="cuda", generator=generator) * 0.023).to(source_dtype) + out = q.quantize(x, dtype=output_dtype) + expected_row = native_qdq(x, output_dtype) + assert torch.isfinite(expected_row).all() + if rowwise: + assert_bits(out._rowwise_storage.dequantize(), expected_row) + expected_col = mx_value(expected_row.float(), columnwise=True) if mixed else expected_row + assert_bits(out._columnwise_storage.dequantize(dtype=expected_col.dtype), expected_col) + + # The same precision must be retained when reconstructing a dropped row. + if rowwise: + out.update_usage(rowwise_usage=False) + q.set_usage(rowwise=False, columnwise=True) + q.quantize(x, out=out, dtype=output_dtype) + assert_bits(out._columnwise_storage.dequantize(dtype=expected_col.dtype), expected_col) + + +@pytest.mark.parametrize("mixed", [False, True]) +def test_qdq_linear_preserves_bf16_forward_values_from_fp16_weights(mixed): + """Finite BF16 forward weights must stay finite when building the backward operand.""" + factory = ( + mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory + if mixed + else nvfp4_qdq_weight_fwd_bf16_bwd_factory + ) + recipe = CustomRecipe(qfactory=partial(factory, backend="reference")) + layer = te.Linear(128, 128, bias=False, params_dtype=torch.float16, device="cuda") + with torch.no_grad(): + layer.weight.fill_(65504.0) + x = torch.ones(128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with te.autocast(recipe=recipe), torch.autocast("cuda", dtype=torch.bfloat16): + y = layer(x) + y.backward(torch.ones_like(y)) + # QDQ rounds 65504 to finite BF16 65536. Casting that value back to FP16 + # would overflow; both HP and MXFP8 backward weights can retain it exactly. + expected = torch.full_like(x, 128 * 65536.0) + torch.testing.assert_close(y, expected, rtol=0, atol=0) + torch.testing.assert_close(x.grad, expected, rtol=0, atol=0) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 09313bb2618..335400ab1a8 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -7746,6 +7746,128 @@ def test_hybrid_output_dtype_lifecycle(internal, rowwise, requested_dtype): q.quantize(x, out=out, dtype=torch.float16) +@pytest.mark.parametrize("internal", [False, True]) +@pytest.mark.parametrize("direction", ["rowwise", "columnwise"]) +@pytest.mark.parametrize("other_usage", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_hybrid_output_dtype_checks_child_storage(internal, direction, other_usage, dtype): + """A matching parent dtype cannot hide an incompatible allocated HP child.""" + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer( + dtype=torch.float16 if direction == "rowwise" else None + ), + columnwise_quantizer=IdentityQuantizer( + dtype=torch.float16 if direction == "columnwise" else None + ), + ) + q.internal = internal + q.rowwise_quantizer.internal = internal + q.columnwise_quantizer.internal = internal + q.set_usage( + rowwise=direction == "rowwise" or other_usage, + columnwise=direction == "columnwise" or other_usage, + ) + x = torch.randn(32, 64, device="cuda", dtype=dtype) + out = q.quantize(x) + assert out._dtype == dtype + assert getattr(out, f"_{direction}_storage")._hp_data.dtype == torch.float16 + children = [ + child for child in (out._rowwise_storage, out._columnwise_storage) if child is not None + ] + before = [child._hp_data.clone() for child in children] + pointers = [child._hp_data.data_ptr() for child in children] + updated = x.float() * 3 + + assert not q.is_compatible_output(out, dtype=dtype) + with pytest.raises(ValueError, match="dtype does not match"): + q.quantize(updated, out=out, dtype=dtype) + for child, expected, pointer in zip(children, before, pointers): + assert child._hp_data.data_ptr() == pointer + torch.testing.assert_close(child._hp_data, expected, rtol=0, atol=0) + + # Without an override, updates still retain each destination's own dtype. + assert q.is_compatible_output(out) + assert q.quantize(updated, out=out) is out + for child, pointer in zip(children, pointers): + assert child._hp_data.data_ptr() == pointer + torch.testing.assert_close( + child.dequantize(), updated.to(child._hp_data.dtype), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("direction", ["rowwise", "columnwise"]) +def test_hybrid_weight_workspace_checks_child_dtype(direction): + """An incompatible child invalidates the cache even when the parent matches.""" + from transformer_engine.pytorch.module.base import quantize_weight + + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer( + dtype=torch.float16 if direction == "rowwise" else None + ), + columnwise_quantizer=IdentityQuantizer( + dtype=torch.float16 if direction == "columnwise" else None + ), + columnwise_source="rowwise_dequantized", + ) + x = torch.randn(32, 64, device="cuda", dtype=torch.float32) + cache = q.quantize(x) + assert cache._dtype == torch.float32 + assert getattr(cache, f"_{direction}_storage")._hp_data.dtype == torch.float16 + updated = x * 3 + out, new_cache = quantize_weight( + tensor=updated, + quantizer=q, + workspace=cache, + workspace_dtype=torch.float32, + cache=True, + ) + assert new_cache is out and out is not cache + for child in (out._rowwise_storage, out._columnwise_storage): + assert child._hp_data.dtype == torch.float32 + torch.testing.assert_close(child.dequantize(), updated, rtol=0, atol=0) + pointers = [ + out._rowwise_storage._hp_data.data_ptr(), + out._columnwise_storage._hp_data.data_ptr(), + ] + expected = updated * 2 + reused, new_cache = quantize_weight( + tensor=expected, + quantizer=q, + workspace=out, + workspace_dtype=torch.float32, + cache=True, + ) + assert reused is out and new_cache is None + for child, pointer in zip((out._rowwise_storage, out._columnwise_storage), pointers): + assert child._hp_data.data_ptr() == pointer + torch.testing.assert_close(child.dequantize(), expected, rtol=0, atol=0) + assert getattr(q, f"{direction}_quantizer").dtype == torch.float16 + + +@pytest.mark.parametrize("nested", [False, True]) +def test_hybrid_columnwise_identity_inherits_reconstructed_source_dtype(nested): + """An implicit column dtype follows the promoted source, not the original input.""" + row = IdentityQuantizer(dtype=torch.float16) + if nested: + row = HybridQuantizer( + rowwise_quantizer=row, + columnwise_quantizer=IdentityQuantizer(), + ) + q = HybridQuantizer( + rowwise_quantizer=row, + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + q.set_usage(rowwise=False, columnwise=True) + x = torch.tensor([[0.023, -0.031]], dtype=torch.bfloat16) + out = q.quantize(x) + assert out._columnwise_storage._hp_data.dtype == torch.float32 + torch.testing.assert_close( + out._columnwise_storage.dequantize(), x.half().float(), rtol=0, atol=0 + ) + assert q.columnwise_quantizer.dtype is None + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) @pytest.mark.parametrize("internal", [False, True]) @pytest.mark.parametrize("source_dtype", [torch.bfloat16, torch.float32]) @@ -7781,10 +7903,21 @@ def test_hybrid_native_dtype_override_preserves_encoding(internal, source_dtype, ) assert actual._rowwise_storage._dtype == expected_row._dtype assert actual._columnwise_storage._dtype == expected_col._dtype + assert q.is_compatible_output(actual, dtype=requested_dtype) # Native out= and module caches still ignore conflicting output requests. from transformer_engine.pytorch.module.base import quantize_weight + reused, new_cache = quantize_weight( + tensor=x, + quantizer=q, + workspace=actual, + workspace_dtype=requested_dtype, + cache=True, + update_workspace=False, + ) + assert reused is actual and new_cache is None + assert row.is_compatible_output(expected_row, dtype=requested_dtype) # Native update_quantized requires a public tensor destination. public_row = row.copy() diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 4d0af47f905..e964ad98109 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -168,6 +168,20 @@ def with_output_dtype(self, dtype): setattr(quantizer, name, child.with_output_dtype(dtype)) return quantizer + def is_compatible_output(self, out: QuantizedTensorStorage, *, dtype=None) -> bool: + """Validate the parent and each allocated child's output-dtype contract.""" + if not isinstance(out, HybridQuantizedTensorStorage): + return False + if not super().is_compatible_output(out, dtype=dtype): + return False + return all( + storage is None or quantizer.is_compatible_output(storage, dtype=dtype) + for quantizer, storage in ( + (self.rowwise_quantizer, out._rowwise_storage), + (self.columnwise_quantizer, out._columnwise_storage), + ) + ) + @property def with_amax_reduction(self) -> bool: """Whether either sub-quantizer has cross-rank amax reduction enabled.""" @@ -216,8 +230,19 @@ def _columnwise_src_from_rowwise( if rowwise_result is None: quantizer = rowwise_quantizer or self.rowwise_quantizer rowwise_result = quantizer.quantize(tensor, dtype=dtype) - # Decode into the source dtype, preserving any prior Identity/QDQ rounding. - return rowwise_result.dequantize(dtype=tensor.dtype) + # Preserve HP forward values when the source dtype would narrow them. + # Native representations retain their existing decode-to-source behavior. + rowwise_storage = rowwise_result + while isinstance(rowwise_storage, HybridQuantizedTensorStorage): + rowwise_storage = ( + rowwise_storage._rowwise_storage + if rowwise_storage._rowwise_storage is not None + else rowwise_storage._columnwise_storage + ) + columnwise_dtype = tensor.dtype + if rowwise_storage is not None and rowwise_storage._get_quantizer().supports_output_dtype: + columnwise_dtype = torch.promote_types(tensor.dtype, rowwise_storage._dtype) + return rowwise_result.dequantize(dtype=columnwise_dtype) def quantize_impl(self, tensor: torch.Tensor, *, dtype=None) -> QuantizedTensor: # A result owns the child configuration needed to reconstruct an absent From a82d551630c61385a3377f922b5c862222c076c6 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 14 Oct 2026 12:17:43 +0000 Subject: [PATCH 8/9] Enable NVFP4 and MXFP8 QDQ on Ampere+ Signed-off-by: Evgeny --- .../fine_grained_quantization.rst | 41 +++- tests/cpp/operator/test_nvfp4_qdq.cu | 2 +- tests/pytorch/mxfp8/test_mxfp8_qdq.py | 34 ++- tests/pytorch/nvfp4/test_nvfp4_qdq.py | 50 +++- tests/pytorch/nvfp4/test_qdq_emulation.py | 217 ++++++++++++++++++ transformer_engine/common/cast/mxfp8/qdq.cu | 2 +- transformer_engine/common/cast/nvfp4/qdq.cu | 2 +- .../pytorch/csrc/extensions/cast.cpp | 5 +- .../pytorch/custom_recipes/qdq.py | 25 +- 9 files changed, 349 insertions(+), 29 deletions(-) create mode 100644 tests/pytorch/nvfp4/test_qdq_emulation.py diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst index 1331ed8515e..8096db713b3 100644 --- a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -302,10 +302,13 @@ children that support output-dtype binding retain it for later reconstruction. Native quantizers keep their existing behavior, including ignoring this argument. Their encodings and default dequantization metadata are unchanged. -Hybrid's ``rowwise_dequantized`` source is still explicitly decoded into the -source tensor's dtype. This preserves existing native compositions, including -internal storage. If a high-precision child already rounded into BF16, widening -that value to an FP32 source does not recover the discarded precision. +Hybrid's ``rowwise_dequantized`` source uses the source tensor's dtype for +native children, preserving existing native compositions, including internal +storage. For children that support output-dtype handling, it uses +``torch.promote_types`` on the source and stored dtypes to avoid narrowing +decoded forward values. This is the input dtype of the columnwise quantizer, +not necessarily the backward GEMM dtype. Widening a value already rounded into +BF16 does not recover the discarded precision. In-place Identity updates preserve the destination buffer, dtype, shape and device. An explicit conflicting ``dtype`` with ``out=`` raises an error for @@ -437,10 +440,11 @@ Backends and custom QDQ formats Output defaults to the source dtype; a module compute-dtype override changes only the decoded dtype. For ``NVFP4QDQQuantizer``, ``backend="auto"`` fuses plain 1D E2M1/E4M3 QDQ with tensor or per-row outer scaling on -SM100 for contiguous, 16-byte-aligned BF16/FP16 tensors with matching output -dtype and flattened matrix dimensions divisible by 16, with ordinary fast math -disabled. Other native-valid modes, including 4over6, 2D scaling, distributed -amax reduction, and mixed dtypes, use native ``quantize().dequantize()``. +GPUs with compute capability 8.0 or higher for contiguous, 16-byte-aligned +BF16/FP16 tensors with matching output dtype and flattened matrix dimensions +divisible by 16, with ordinary fast math disabled. Other native-valid modes, +including 4over6, 2D scaling, distributed amax reduction, and mixed dtypes, +use native ``quantize().dequantize()`` where its kernels are supported. ``backend="reference"`` forces that composition; ``backend="fused"`` rejects unsupported cases. Current weight amax is computed when quantization runs. @@ -455,12 +459,27 @@ rounding and support the same matching BF16/FP16 input and output dtypes. ``custom_recipes.qdq.MXFP8QDQQuantizer`` provides the corresponding rowwise E4M3/block-32 E8M0 QDQ operation. Its ``reference`` backend calls native ``MXFP8Quantizer.quantize().dequantize()``. Its ``fused`` backend supports -contiguous, 16-byte-aligned BF16/FP16 tensors on SM100, with a matching output -dtype, rows divisible by 32, and ordinary fast math disabled. ``auto`` falls -back to the reference when these requirements are not met. +contiguous, 16-byte-aligned BF16/FP16 tensors on GPUs with compute capability +8.0 or higher, with a matching output dtype, rows divisible by 32, and ordinary +fast math disabled. ``auto`` falls back to the reference when these requirements +are not met. The fused kernel preserves signed zero. Native specialized FP16 quantization may canonicalize negative zero; the decoded values remain numerically equal. +Build Transformer Engine for the target GPU architecture; a build containing +only ``sm_100a`` kernels cannot run those kernels on other architectures. +NVFP4 QDQ requires CUDA Toolkit 12.8 or later for the FP4 conversion types. +On Ampere and Hopper, CUDA emulates the conversions inside the fused kernel. +The existing native packed Q+DQ reference kernels require compute capability +10.0 or higher. On older GPUs, use ``auto`` or ``fused`` with matching BF16/FP16 +and the plain fused configuration described above; a request requiring the +native reference path raises an error. In particular, mixed dtypes, 4over6, +2D scaling, and distributed amax reduction are not supported there. +QDQ's conversion requirements are independent of native low-precision GEMM +support. Recipes using native MXFP8 backward additionally require +``is_mxfp8_available()``. Numerical and performance validation on a particular +GPU does not establish performance on other architectures. + Both QDQ classes return high-precision storage with straight-through gradients. Use them as rowwise children of ``HybridQuantizer`` to select the backward format independently. ``columnwise_source="rowwise_dequantized"`` diff --git a/tests/cpp/operator/test_nvfp4_qdq.cu b/tests/cpp/operator/test_nvfp4_qdq.cu index 00c4b885b92..a8ed5623272 100644 --- a/tests/cpp/operator/test_nvfp4_qdq.cu +++ b/tests/cpp/operator/test_nvfp4_qdq.cu @@ -22,7 +22,7 @@ template void check_qdq(size_t rows, size_t cols, bool row_scaled = false) { cudaDeviceProp prop; ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess); - if (prop.major != 10 || prop.minor != 0) GTEST_SKIP() << "QDQ targets SM100"; + if (prop.major < 10) GTEST_SKIP() << "Native packed Q+DQ reference requires SM10.0+"; const auto dtype = TypeInfo::dtype; const std::vector shape{rows, cols}; Tensor input("input", shape, dtype); diff --git a/tests/pytorch/mxfp8/test_mxfp8_qdq.py b/tests/pytorch/mxfp8/test_mxfp8_qdq.py index 8c86da9a739..2aa1224f369 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_qdq.py +++ b/tests/pytorch/mxfp8/test_mxfp8_qdq.py @@ -9,9 +9,11 @@ from transformer_engine.pytorch.custom_recipes.qdq import MXFP8QDQQuantizer import transformer_engine_torch as tex +# These tests use native packed Q+DQ as their oracle. Portable QDQ coverage is +# in nvfp4/test_qdq_emulation.py and does not require native format availability. pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), - reason="MXFP8 QDQ fused implementation targets SM100", + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0), + reason="MXFP8 QDQ requires compute capability 10.0 or higher", ) @@ -20,6 +22,34 @@ def ordinary_math(monkeypatch): monkeypatch.setenv("NVTE_USE_FAST_MATH", "0") +def test_auto_matches_native_on_available_device(): + x = data((96, 160), torch.bfloat16, "normal") + q = operand("auto") + assert q.selected_backend(x) == "fused" + assert_bits(q(x).dequantize(), operand("reference")(x).dequantize()) + + +@pytest.mark.parametrize( + "capability", [(7, 5), (8, 0), (8, 9), (9, 0), (10, 0), (10, 3), (10, 7), (11, 0), (12, 0)] +) +def test_architecture_dispatch(capability, monkeypatch): + """Exercise dispatch only; numerical execution uses the actual CI device.""" + from transformer_engine.pytorch.custom_recipes.qdq import _supports_fused_device + + x = torch.ones(32, 32, device="cuda", dtype=torch.bfloat16) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: capability) + _supports_fused_device.cache_clear() + try: + expected = "fused" if capability >= (8, 0) else "reference" + assert operand("auto").selected_backend(x) == expected + assert operand("reference").selected_backend(x) == "reference" + if expected == "reference": + with pytest.raises(ValueError, match="Forced fused"): + operand("fused").selected_backend(x) + finally: + _supports_fused_device.cache_clear() + + def assert_bits(actual, expected): assert actual.shape == expected.shape and actual.dtype == expected.dtype assert torch.equal( diff --git a/tests/pytorch/nvfp4/test_nvfp4_qdq.py b/tests/pytorch/nvfp4/test_nvfp4_qdq.py index 331bdc97a01..81d6294c818 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_qdq.py +++ b/tests/pytorch/nvfp4/test_nvfp4_qdq.py @@ -19,10 +19,48 @@ nvfp4_qdq_weight_fwd_bf16_bwd_factory, ) -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), - reason="QDQ fused implementation targets SM100", +# Native packed Q+DQ is this file's oracle. The independently referenced +# portable tests in test_qdq_emulation.py run on Ampere/Hopper as well. +nvfp4_available, nvfp4_reason = te.is_nvfp4_available(return_reason=True) +mxfp8_available, mxfp8_reason = te.is_mxfp8_available(return_reason=True) +pytestmark = pytest.mark.skipif(not nvfp4_available, reason=nvfp4_reason) + + +backward_modes = [ + False, + pytest.param(True, marks=pytest.mark.skipif(not mxfp8_available, reason=mxfp8_reason)), +] + + +@pytest.mark.parametrize("row_scaled", [False, True]) +def test_auto_matches_native_on_available_device(row_scaled, monkeypatch): + monkeypatch.setenv("NVTE_USE_FAST_MATH", "0") + x = torch.randn(48, 80, device="cuda", dtype=torch.bfloat16) + q = NVFP4QDQQuantizer(NVFP4Quantizer(row_scaled_nvfp4=row_scaled)) + assert q.selected_backend(x) == "fused" + assert_bits(q(x).dequantize(), native_qdq(x, row_scaled_nvfp4=row_scaled)) + + +@pytest.mark.parametrize( + "capability", [(7, 5), (8, 0), (8, 9), (9, 0), (10, 0), (10, 3), (10, 7), (11, 0), (12, 0)] ) +def test_architecture_dispatch(capability, monkeypatch): + """Exercise dispatch only; numerical execution uses the actual CI device.""" + from transformer_engine.pytorch.custom_recipes.qdq import _supports_fused_device + + x = torch.ones(16, 16, device="cuda", dtype=torch.bfloat16) + monkeypatch.setenv("NVTE_USE_FAST_MATH", "0") + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: capability) + _supports_fused_device.cache_clear() + try: + expected = "fused" if capability >= (8, 0) else "reference" + assert NVFP4QDQQuantizer().selected_backend(x) == expected + assert NVFP4QDQQuantizer(backend="reference").selected_backend(x) == "reference" + if expected == "reference": + with pytest.raises(ValueError, match="Forced fused"): + NVFP4QDQQuantizer(backend="fused").selected_backend(x) + finally: + _supports_fused_device.cache_clear() def native_qdq(x, dtype=None, **options): @@ -221,7 +259,7 @@ def mx_value(x, *, columnwise=False): @pytest.mark.parametrize("backend", ["reference", "fused"]) @pytest.mark.parametrize("grouped", [False, True]) @pytest.mark.parametrize("params_dtype", [torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize("mixed", backward_modes) def test_module_operands(backend, grouped, params_dtype, mixed, nvfp4_options=None): if backend == "fused" and params_dtype == torch.float32: pytest.skip("Mixed FP32 source/BF16 decode intentionally uses native fallback") @@ -414,7 +452,7 @@ def test_copy_preserves_runtime_reduction_state(): @pytest.mark.parametrize("grouped", [False, True]) -@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize("mixed", backward_modes) def test_4over6_modules(grouped, mixed): test_module_operands( "reference", @@ -510,7 +548,7 @@ def test_qdq_columnwise_source_preserves_forward_precision( assert_bits(out._columnwise_storage.dequantize(dtype=expected_col.dtype), expected_col) -@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize("mixed", backward_modes) def test_qdq_linear_preserves_bf16_forward_values_from_fp16_weights(mixed): """Finite BF16 forward weights must stay finite when building the backward operand.""" factory = ( diff --git a/tests/pytorch/nvfp4/test_qdq_emulation.py b/tests/pytorch/nvfp4/test_qdq_emulation.py new file mode 100644 index 00000000000..8146af1ba5b --- /dev/null +++ b/tests/pytorch/nvfp4/test_qdq_emulation.py @@ -0,0 +1,217 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Portable QDQ numerics and integration, including Ampere/Hopper CI devices. + +The oracle uses PyTorch arithmetic and explicit FP4 rounding, not TE's native +packed quantization kernels, which require Blackwell. The existing NVFP4 L0 +directory invocation discovers this file on every hardware generation. +""" + +from functools import partial + +import pytest +import torch +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.pytorch.custom_recipes.qdq import ( + MXFP8QDQQuantizer, + NVFP4QDQQuantizer, +) +from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_qdq_weight_fwd_bf16_bwd_factory, +) +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import ( + cast_from_fp4x2, + cast_to_fp4x2, +) +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 0), + reason="Portable QDQ requires compute capability 8.0 or higher", +) + + +@pytest.fixture(autouse=True) +def ordinary_math(monkeypatch): + monkeypatch.setenv("NVTE_USE_FAST_MATH", "0") + + +def quantizer(fmt, backend="auto"): + if fmt == "mxfp8": + return MXFP8QDQQuantizer(backend=backend) + return NVFP4QDQQuantizer(NVFP4Quantizer(row_scaled_nvfp4=fmt == "nvfp4_row"), backend=backend) + + +def reference(x, fmt): + """Independent oracle with native rounding, fmax semantics, and decode order.""" + values = x.float().reshape(-1, x.shape[-1]) + width = 32 if fmt == "mxfp8" else 16 + blocks = values.reshape(values.shape[0], -1, width) + magnitudes = torch.fmax(blocks.abs(), torch.zeros_like(blocks)) + block_amax = magnitudes.amax(dim=-1, keepdim=True) + if fmt == "mxfp8": + mantissa, exponent = torch.frexp(block_amax * (1.0 / 448.0)) + exponent = torch.where(mantissa == 0.5, exponent - 1, exponent) + exponent = torch.where(block_amax == 0, -127, exponent).clamp(-127, 127) + exponent = torch.where(torch.isinf(block_amax), 127, exponent) + decode = torch.exp2(exponent.float()) + inverse = torch.exp2(-exponent.float()) + encoded = (blocks * inverse).clamp(-448, 448).to(torch.float8_e4m3fn).float() + else: + outer = ( + magnitudes.amax(dim=(-1, -2), keepdim=True) if fmt == "nvfp4_row" else magnitudes.amax() + ) + # PyTorch FP32 division may use an approximate reciprocal. Round an + # FP64 quotient to FP32 to model CUDA's ordinary div.rn operation. + encode = ((6.0 * 448.0) / outer.double()).float() + encode = encode.clamp(max=torch.finfo(torch.float32).max) + encode = torch.where((outer == 0) | torch.isinf(outer), 1.0, encode) + scale = (block_amax * (encode * (1.0 / 6.0))).clamp(max=448) + scale = scale.to(torch.float8_e4m3fn).float() + reciprocal_encode = (1.0 / encode.double()).float() + inverse = (1.0 / (scale * reciprocal_encode).double()).float() + inverse = inverse.clamp(max=torch.finfo(torch.float32).max) + scaled = blocks * inverse + # Reuse the independent reference's explicit E2M1 tie boundaries. + matrix = scaled.reshape(values.shape) + encoded = cast_from_fp4x2(cast_to_fp4x2(matrix), torch.float32) + # That reference's packing omits the sign of -0; retain it explicitly. + encoded = torch.copysign(encoded.abs(), matrix).reshape(blocks.shape) + encoded = torch.where(torch.isnan(scaled), 6.0, encoded) + decode = scale * outer * (1.0 / (6.0 * 448.0)) + return (encoded * decode).reshape(x.shape).to(x.dtype) + + +def inputs(shape, dtype, distribution): + generator = torch.Generator(device="cuda").manual_seed(123) + x = torch.randn(shape, device="cuda", dtype=dtype, generator=generator) + if distribution == "zeros": + x.zero_() + x[..., ::2] = -0.0 + elif distribution == "outliers": + x.reshape(-1, shape[-1])[0, 0] = 4096 + x.reshape(-1, shape[-1])[1::2] *= 1e-5 + elif distribution == "tiny": + x.fill_(torch.finfo(dtype).tiny / 4) + x[..., ::2].neg_() + x.reshape(-1, shape[-1])[0].zero_() + elif distribution == "boundaries": + values = torch.tensor( + [0, -0.0, 0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5, 6, -6, 448, 1e-7, 1e-4, -5, -0.25], + device="cuda", + dtype=dtype, + ) + x.copy_(values.repeat(x.numel() // 16).reshape(x.shape)) + return x + + +@pytest.mark.parametrize("fmt", ["nvfp4", "nvfp4_row", "mxfp8"]) +@pytest.mark.parametrize("backend", ["auto", "fused"]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("shape", [(16, 32), (48, 160), (2, 8, 4096)]) +@pytest.mark.parametrize("distribution", ["normal", "zeros", "outliers", "boundaries", "tiny"]) +def test_portable_numerics(fmt, backend, dtype, shape, distribution): + x = inputs(shape, dtype, distribution) + q = quantizer(fmt, backend) + assert q.selected_backend(x) == "fused" + actual = q(x).dequantize() + expected = reference(x, fmt) + assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)) + # Validate the independent oracle against the native path where it exists. + native_shape = fmt != "mxfp8" or (x.numel() // x.shape[-1]) % 32 == 0 + if torch.cuda.get_device_capability() >= (10, 0) and native_shape: + native = quantizer(fmt, "reference")(x).dequantize() + torch.testing.assert_close(actual, native, rtol=0, atol=0) + + +@pytest.mark.parametrize("fmt", ["nvfp4", "nvfp4_row", "mxfp8"]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf"), "all_nan"]) +def test_portable_exceptional_values(fmt, dtype, value): + x = inputs((32, 128), dtype, "normal") + if value == "all_nan": + x.fill_(float("nan")) + else: + x[0, 0] = value + actual = quantizer(fmt)(x).dequantize() + expected = reference(x, fmt) + assert torch.equal(torch.isnan(actual), torch.isnan(expected)) + valid = ~torch.isnan(expected) + assert torch.equal(actual[valid].view(torch.int16), expected[valid].view(torch.int16)) + if torch.cuda.get_device_capability() >= (10, 0): + native = quantizer(fmt, "reference")(x).dequantize() + assert torch.equal(torch.isnan(actual), torch.isnan(native)) + torch.testing.assert_close(actual[valid], native[valid], rtol=0, atol=0) + + +@pytest.mark.parametrize("fmt", ["nvfp4", "nvfp4_row", "mxfp8"]) +def test_portable_lifecycle_and_graph(fmt): + x = inputs((32, 128), torch.bfloat16, "normal").requires_grad_() + q = quantizer(fmt) + y = q(x) + y.sum().backward() + assert torch.equal(x.grad, torch.ones_like(x)) + pointer = y._hp_data.data_ptr() + noop = torch.ones(1, device=x.device) + q.update_quantized(x * 2, y, noop_flag=noop) + torch.testing.assert_close(y.dequantize(), reference(x, fmt), rtol=0, atol=0) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + q.update_quantized(x, y, noop_flag=noop) + with torch.no_grad(): + x.mul_(2) + noop.zero_() + graph.replay() + torch.cuda.synchronize() + assert y._hp_data.data_ptr() == pointer + torch.testing.assert_close(y.dequantize(), reference(x, fmt), rtol=0, atol=0) + + +@pytest.mark.parametrize("grouped", [False, True]) +def test_portable_weight_recipe(grouped): + factory = partial(nvfp4_qdq_weight_fwd_bf16_bwd_factory, backend="fused") + recipe = CustomRecipe(qfactory=factory) + options = dict(bias=False, params_dtype=torch.bfloat16, device="cuda") + layer = ( + te.GroupedLinear(2, 128, 128, single_grouped_weight=False, **options) + if grouped + else te.Linear(128, 128, **options) + ) + x = inputs((32, 128), torch.bfloat16, "normal").requires_grad_() + dy = torch.ones_like(x) + weights = [layer.weight0, layer.weight1] if grouped else [layer.weight] + splits = [16, 16] if grouped else [32] + with te.autocast(recipe=recipe), torch.autocast("cuda", dtype=torch.bfloat16): + y = layer(x, splits) if grouped else layer(x) + y.backward(dy) + expected_y, expected_dx = [], [] + for weight, source, grad in zip(weights, x.detach().split(splits), dy.split(splits)): + decoded = reference(weight.detach(), "nvfp4") + expected_y.append(source @ decoded.T) + expected_dx.append(grad @ decoded) + torch.testing.assert_close(weight.grad, grad.T @ source, rtol=0.016, atol=0.002) + torch.testing.assert_close(y, torch.cat(expected_y), rtol=0.016, atol=0.002) + torch.testing.assert_close(x.grad, torch.cat(expected_dx), rtol=0.016, atol=0.002) + + +@pytest.mark.parametrize("fmt", ["nvfp4", "nvfp4_row", "mxfp8"]) +def test_portable_device_and_stream(fmt): + device = 1 if torch.cuda.device_count() > 1 else 0 + with torch.cuda.device(device), torch.cuda.stream(torch.cuda.Stream(device=device)): + x = inputs((32, 128), torch.bfloat16, "normal") + actual = quantizer(fmt)(x).dequantize() + assert torch.equal(actual.view(torch.int16), reference(x, fmt).view(torch.int16)) + + +@pytest.mark.parametrize("fmt", ["nvfp4", "mxfp8"]) +def test_older_device_reference_rejected(fmt, monkeypatch): + x = inputs((32, 128), torch.bfloat16, "normal") + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (9, 0)) + with pytest.raises(RuntimeError, match="Native.*Q\\+DQ requires compute capability 10.0"): + quantizer(fmt, "reference")(x) diff --git a/transformer_engine/common/cast/mxfp8/qdq.cu b/transformer_engine/common/cast/mxfp8/qdq.cu index d67360839c5..c1b954e6cd7 100644 --- a/transformer_engine/common/cast/mxfp8/qdq.cu +++ b/transformer_engine/common/cast/mxfp8/qdq.cu @@ -48,7 +48,7 @@ void nvte_mxfp8_qdq(const NVTETensor input, NVTETensor output, const NVTETensor cudaStream_t stream) { NVTE_API_CALL(nvte_mxfp8_qdq); using namespace transformer_engine; - NVTE_CHECK(cuda::sm_arch() == 100, "MXFP8 QDQ currently supports SM100"); + NVTE_CHECK(cuda::sm_arch() >= 80, "MXFP8 QDQ requires compute capability 8.0 or higher"); const auto &in = *convertNVTETensorCheck(input); auto &out = *convertNVTETensorCheck(output); CheckInputTensor(in, "input"); diff --git a/transformer_engine/common/cast/nvfp4/qdq.cu b/transformer_engine/common/cast/nvfp4/qdq.cu index 3fbd6a99edd..c712aa0474e 100644 --- a/transformer_engine/common/cast/nvfp4/qdq.cu +++ b/transformer_engine/common/cast/nvfp4/qdq.cu @@ -95,7 +95,7 @@ template void launch_qdq(const NVTETensor input, NVTETensor output, const NVTETensor amax, const NVTETensor noop, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED - NVTE_CHECK(cuda::sm_arch() == 100, "NVFP4 QDQ currently supports SM100"); + NVTE_CHECK(cuda::sm_arch() >= 80, "NVFP4 QDQ requires compute capability 8.0 or higher"); const auto &in = *convertNVTETensorCheck(input); auto &out = *convertNVTETensorCheck(output); CheckInputTensor(in, "input"); diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 783a40b22e0..16a5028db10 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -61,9 +61,8 @@ void nvfp4_qdq(const at::Tensor &input, at::Tensor output, at::Tensor amax, at::cuda::CUDAGuard guard(input.device()); TORCH_CHECK(input.is_contiguous() && output.is_contiguous() && amax.is_contiguous(), "NVFP4 QDQ requires contiguous tensors"); - TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major == 10 && - at::cuda::getCurrentDeviceProperties()->minor == 0, - "NVFP4 QDQ currently supports SM100"); + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 8, + "NVFP4 QDQ requires compute capability 8.0 or higher"); TORCH_CHECK( !noop || (noop->is_cuda() && noop->device() == input.device() && noop->is_contiguous()), "Invalid QDQ noop device/layout"); diff --git a/transformer_engine/pytorch/custom_recipes/qdq.py b/transformer_engine/pytorch/custom_recipes/qdq.py index 51f58448a2d..726efc87f5b 100644 --- a/transformer_engine/pytorch/custom_recipes/qdq.py +++ b/transformer_engine/pytorch/custom_recipes/qdq.py @@ -22,7 +22,8 @@ @lru_cache(maxsize=None) def _supports_fused_device(device_index): """Device capabilities are fixed for the process lifetime.""" - return torch.cuda.get_device_capability(device_index) == (10, 0) + # QDQ uses conversion instructions, not architecture-specific GEMM tiling. + return torch.cuda.get_device_capability(device_index) >= (8, 0) class QDQQuantizer(Quantizer): @@ -124,9 +125,11 @@ class MXFP8QDQQuantizer(QDQQuantizer): """Rowwise E4M3/block-32 E8M0 QDQ held in high-precision storage. ``reference`` explicitly invokes native Q then DQ. ``fused`` keeps both - encodings in registers and requires contiguous matching BF16/FP16 on SM100. + encodings in registers and requires contiguous matching BF16/FP16 on GPUs + with compute capability 8.0 or higher. Older GPUs emulate the conversions. ``auto`` selects the fused implementation where supported. The decoded buffer can be consumed in either GEMM orientation; quantization is rowwise. + The native ``reference`` backend requires compute capability 10.0 or higher. """ block_size = 32 @@ -166,6 +169,11 @@ def _reference(self, tensor, dtype): # Native MXFP8 launch geometry does not support empty tensors. if tensor.numel() == 0: return torch.empty_like(tensor, dtype=dtype) + if tensor.is_cuda and torch.cuda.get_device_capability(tensor.device) < (10, 0): + raise RuntimeError( + "Native MXFP8 Q+DQ requires compute capability 10.0 or higher; " + "on Ampere/Hopper use auto or fused with contiguous matching BF16/FP16" + ) return self.mxfp8_quantizer.quantize(tensor).dequantize(dtype=dtype) def _compute(self, tensor, output, noop_flag=None, *, backend=None): @@ -191,9 +199,11 @@ class NVFP4QDQQuantizer(QDQQuantizer): serve either GEMM orientation. ``backend`` is ``auto``, ``reference``, or ``fused``. Auto falls back to - native Q+DQ outside plain contiguous BF16/FP16 1D E4M3 scaling on SM100, - with tensor or per-row outer scaling. Per-row amax is fused into QDQ. + native Q+DQ outside plain contiguous BF16/FP16 1D E4M3 scaling on GPUs with + compute capability 8.0 or higher, with tensor or per-row outer scaling. + Per-row amax is fused into QDQ. In particular, 4over6 and 2D scaling currently use the reference path. + The native reference path requires compute capability 10.0 or higher. Forced fused execution raises for unsupported inputs/configurations. This experimental API is subject to change. """ @@ -282,6 +292,13 @@ def selected_backend(self, tensor, dtype=None): return "fused" if eligible and self.backend != "reference" else "reference" def _reference(self, tensor, dtype): + if tensor.is_cuda and torch.cuda.get_device_capability(tensor.device) < (10, 0): + if tensor.numel() == 0: + return torch.empty_like(tensor, dtype=dtype) + raise RuntimeError( + "Native NVFP4 Q+DQ requires compute capability 10.0 or higher; " + "on Ampere/Hopper use auto or fused with plain contiguous matching BF16/FP16" + ) return self.nvfp4_quantizer.quantize(tensor).dequantize(dtype=dtype) def _compute(self, tensor, output, noop_flag=None, *, backend=None): From 7f5e16fdd6a15e49ded0a3429409f6b1e0f30d25 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 16 Oct 2026 12:56:35 +0000 Subject: [PATCH 9/9] Allow CustomRecipe without FP8 hardware and fix QDQ capability gates Signed-off-by: Evgeny --- .../fine_grained_quantization.rst | 10 ++- tests/pytorch/nvfp4/test_qdq_emulation.py | 77 +++++++++++++++++-- tests/pytorch/test_recipe.py | 41 ++++++++++ transformer_engine/common/cast/nvfp4/qdq.cu | 2 + .../common/include/transformer_engine/cast.h | 7 ++ transformer_engine/common/recipe/__init__.py | 4 + .../pytorch/csrc/extensions/pybind.cpp | 2 + .../pytorch/custom_recipes/qdq.py | 12 ++- transformer_engine/pytorch/quantization.py | 9 ++- 9 files changed, 150 insertions(+), 14 deletions(-) diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst index 8096db713b3..119273e7770 100644 --- a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -430,18 +430,22 @@ columnwise quantizer and its ``columnwise_source="original"`` stay unchanged. Both examples use ``columnwise_source="rowwise_dequantized"`` for weights. Other module roles use ``IdentityQuantizer``. High-precision arithmetic follows the module compute dtype; select BF16 compute for the BF16 example. +``CustomRecipe`` does not require FP8 hardware to enter ``te.autocast``. +Its quantizers and GEMMs must support the target device and TE build; +recipes using only QDQ and high-precision GEMMs can run on Ampere. Configure ``backend`` or ``nvfp4_options`` with ``functools.partial``. Existing Hybrid limitations still apply, and RL convergence has not been validated. Backends and custom QDQ formats -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Output defaults to the source dtype; a module compute-dtype override changes only the decoded dtype. For ``NVFP4QDQQuantizer``, ``backend="auto"`` fuses plain 1D E2M1/E4M3 QDQ with tensor or per-row outer scaling on -GPUs with compute capability 8.0 or higher for contiguous, 16-byte-aligned -BF16/FP16 tensors with matching output dtype and flattened matrix dimensions +GPUs with compute capability 8.0 or higher and a TE build with CUDA 12.8 or +higher for contiguous, 16-byte-aligned BF16/FP16 tensors with matching output +dtype and flattened matrix dimensions divisible by 16, with ordinary fast math disabled. Other native-valid modes, including 4over6, 2D scaling, distributed amax reduction, and mixed dtypes, use native ``quantize().dequantize()`` where its kernels are supported. diff --git a/tests/pytorch/nvfp4/test_qdq_emulation.py b/tests/pytorch/nvfp4/test_qdq_emulation.py index 8146af1ba5b..50083781051 100644 --- a/tests/pytorch/nvfp4/test_qdq_emulation.py +++ b/tests/pytorch/nvfp4/test_qdq_emulation.py @@ -14,6 +14,7 @@ import pytest import torch import transformer_engine.pytorch as te +import transformer_engine_torch as tex from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.custom_recipes.qdq import ( MXFP8QDQQuantizer, @@ -27,13 +28,23 @@ cast_to_fp4x2, ) from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer - +from transformer_engine.pytorch.quantization import FP8GlobalStateManager pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 0), reason="Portable QDQ requires compute capability 8.0 or higher", ) +requires_fp4_build = pytest.mark.skipif( + not tex.nvfp4_qdq_supported(), + reason="NVFP4 QDQ requires a TE build with CUDA 12.8 or higher", +) +portable_formats = [ + pytest.param("nvfp4", marks=requires_fp4_build), + pytest.param("nvfp4_row", marks=requires_fp4_build), + "mxfp8", +] + @pytest.fixture(autouse=True) def ordinary_math(monkeypatch): @@ -109,7 +120,7 @@ def inputs(shape, dtype, distribution): return x -@pytest.mark.parametrize("fmt", ["nvfp4", "nvfp4_row", "mxfp8"]) +@pytest.mark.parametrize("fmt", portable_formats) @pytest.mark.parametrize("backend", ["auto", "fused"]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("shape", [(16, 32), (48, 160), (2, 8, 4096)]) @@ -128,7 +139,7 @@ def test_portable_numerics(fmt, backend, dtype, shape, distribution): torch.testing.assert_close(actual, native, rtol=0, atol=0) -@pytest.mark.parametrize("fmt", ["nvfp4", "nvfp4_row", "mxfp8"]) +@pytest.mark.parametrize("fmt", portable_formats) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf"), "all_nan"]) def test_portable_exceptional_values(fmt, dtype, value): @@ -148,7 +159,7 @@ def test_portable_exceptional_values(fmt, dtype, value): torch.testing.assert_close(actual[valid], native[valid], rtol=0, atol=0) -@pytest.mark.parametrize("fmt", ["nvfp4", "nvfp4_row", "mxfp8"]) +@pytest.mark.parametrize("fmt", portable_formats) def test_portable_lifecycle_and_graph(fmt): x = inputs((32, 128), torch.bfloat16, "normal").requires_grad_() q = quantizer(fmt) @@ -174,6 +185,7 @@ def test_portable_lifecycle_and_graph(fmt): @pytest.mark.parametrize("grouped", [False, True]) +@requires_fp4_build def test_portable_weight_recipe(grouped): factory = partial(nvfp4_qdq_weight_fwd_bf16_bwd_factory, backend="fused") recipe = CustomRecipe(qfactory=factory) @@ -200,7 +212,47 @@ def test_portable_weight_recipe(grouped): torch.testing.assert_close(x.grad, torch.cat(expected_dx), rtol=0.016, atol=0.002) -@pytest.mark.parametrize("fmt", ["nvfp4", "nvfp4_row", "mxfp8"]) +@pytest.mark.parametrize("fmt", portable_formats) +@pytest.mark.parametrize("grouped", [False, True]) +def test_portable_full_qdq_recipe(fmt, grouped, monkeypatch): + """QDQ operands use BF16 GEMMs without requiring native FP8 hardware.""" + monkeypatch.setattr( + FP8GlobalStateManager, + "is_fp8_available", + classmethod(lambda cls: (False, "FP8 hardware unavailable")), + ) + + def factory(role): + if role is not None and role.tensor_type in ("input", "weight"): + return quantizer(fmt) + return te.IdentityQuantizer() + + recipe = CustomRecipe(qfactory=factory) + options = dict(bias=False, params_dtype=torch.bfloat16, device="cuda") + layer = ( + te.GroupedLinear(2, 128, 128, single_grouped_weight=False, **options) + if grouped + else te.Linear(128, 128, **options) + ) + x = inputs((64, 128), torch.bfloat16, "normal").requires_grad_() + dy = torch.ones_like(x) + weights = [layer.weight0, layer.weight1] if grouped else [layer.weight] + splits = [32, 32] if grouped else [64] + with te.autocast(recipe=recipe), torch.autocast("cuda", dtype=torch.bfloat16): + y = layer(x, splits) if grouped else layer(x) + y.backward(dy) + expected_y, expected_dx = [], [] + for weight, source, grad in zip(weights, x.detach().split(splits), dy.split(splits)): + decoded_x = reference(source, fmt) + decoded_w = reference(weight.detach(), fmt) + expected_y.append(decoded_x @ decoded_w.T) + expected_dx.append(grad @ decoded_w) + torch.testing.assert_close(weight.grad, grad.T @ decoded_x, rtol=0.016, atol=0.002) + torch.testing.assert_close(y, torch.cat(expected_y), rtol=0.016, atol=0.002) + torch.testing.assert_close(x.grad, torch.cat(expected_dx), rtol=0.016, atol=0.002) + + +@pytest.mark.parametrize("fmt", portable_formats) def test_portable_device_and_stream(fmt): device = 1 if torch.cuda.device_count() > 1 else 0 with torch.cuda.device(device), torch.cuda.stream(torch.cuda.Stream(device=device)): @@ -215,3 +267,18 @@ def test_older_device_reference_rejected(fmt, monkeypatch): monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (9, 0)) with pytest.raises(RuntimeError, match="Native.*Q\\+DQ requires compute capability 10.0"): quantizer(fmt, "reference")(x) + + +@pytest.mark.parametrize("fp4_supported", [False, True]) +@pytest.mark.parametrize("fmt", ["nvfp4", "nvfp4_row"]) +def test_nvfp4_build_support_dispatch(fp4_supported, fmt, monkeypatch): + from transformer_engine.pytorch.custom_recipes import qdq + + monkeypatch.setattr(qdq, "_NVFP4_QDQ_SUPPORTED", fp4_supported) + x = inputs((32, 128), torch.bfloat16, "normal") + expected = "fused" if fp4_supported else "reference" + assert quantizer(fmt).selected_backend(x) == expected + assert quantizer("mxfp8").selected_backend(x) == "fused" + if not fp4_supported: + with pytest.raises(ValueError, match="Forced fused NVFP4 QDQ"): + quantizer(fmt, "fused").selected_backend(x) diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 446bfda8e59..b25d0612a42 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -26,6 +26,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.quantization import ( + FP8GlobalState, FP8GlobalStateManager, NVFP4BlockScalingRecipeState, QuantizerRole, @@ -820,3 +821,43 @@ def test_fp8_recipe_unpickle_rebuilds_missing_quant_params(recipe): assert restored.fp8_quant_fwd_inp.power_2_scale == recipe.fp8_quant_fwd_inp.power_2_scale assert restored.fp8_quant_fwd_weight.power_2_scale == recipe.fp8_quant_fwd_weight.power_2_scale assert restored.fp8_quant_bwd_grad.power_2_scale == recipe.fp8_quant_bwd_grad.power_2_scale + + +def test_custom_autocast_without_fp8_hardware(monkeypatch): + monkeypatch.setattr(FP8GlobalStateManager, "quantization_state", FP8GlobalState()) + monkeypatch.setattr( + FP8GlobalStateManager, + "is_fp8_available", + classmethod(lambda cls: (False, "FP8 hardware unavailable")), + ) + + def factory(role): + pytest.fail("Autocast must not call the quantizer factory to inspect its requirements") + + recipe = CustomRecipe(qfactory=factory) + with te.autocast(recipe=recipe): + assert FP8GlobalStateManager.is_fp8_enabled() + assert FP8GlobalStateManager.get_fp8_recipe() is recipe + assert not FP8GlobalStateManager.is_fp8_enabled() + assert FP8GlobalStateManager.quantization_state.autocast_depth == 0 + + +@pytest.mark.parametrize( + "recipe", + [ + DelayedScaling(), + Float8CurrentScaling(), + Float8BlockScaling(), + MXFP8BlockScaling(), + NVFP4BlockScaling(), + ], +) +def test_builtin_autocast_still_requires_fp8_hardware(recipe, monkeypatch): + monkeypatch.setattr(FP8GlobalStateManager, "quantization_state", FP8GlobalState()) + monkeypatch.setattr( + FP8GlobalStateManager, + "is_fp8_available", + classmethod(lambda cls: (False, "FP8 hardware unavailable")), + ) + with pytest.raises(AssertionError, match="FP8 hardware unavailable"): + FP8GlobalStateManager.autocast_enter(enabled=True, fp8_recipe=recipe) diff --git a/transformer_engine/common/cast/nvfp4/qdq.cu b/transformer_engine/common/cast/nvfp4/qdq.cu index c712aa0474e..e2f321fba93 100644 --- a/transformer_engine/common/cast/nvfp4/qdq.cu +++ b/transformer_engine/common/cast/nvfp4/qdq.cu @@ -143,6 +143,8 @@ void launch_qdq(const NVTETensor input, NVTETensor output, const NVTETensor amax } // namespace } // namespace transformer_engine +bool nvte_nvfp4_qdq_supported() { return FP4_TYPE_SUPPORTED; } + void nvte_nvfp4_qdq(const NVTETensor input, NVTETensor output, const NVTETensor amax, const NVTETensor noop, cudaStream_t stream) { NVTE_API_CALL(nvte_nvfp4_qdq); diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index de20c66d202..53c344c3a1f 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -17,6 +17,13 @@ extern "C" { #endif +/*! \brief Whether this library was built with NVFP4 QDQ kernels (CUDA 12.8+). + * + * This does not check the device or input requirements of the QDQ operations. + * \return True if NVFP4 QDQ kernels were compiled into the library. + */ +bool nvte_nvfp4_qdq_supported(); + /*! \brief Plain rowwise 1x16 NVFP4 QDQ into matching BF16/FP16 output. * * Input/output are contiguous, 16-byte aligned, rank >= 2, with both flattened diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index f9abf9de87a..58778fda80c 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -647,6 +647,10 @@ class CustomRecipe(Recipe): """ Custom recipe that allows users to provide quantizer factories. + PyTorch autocast does not require FP8 hardware for a custom recipe. The + selected quantizers and GEMMs must support the target device and build; + autocast does not inspect the factory to validate those requirements in advance. + .. warning:: **EXPERIMENTAL**: Custom recipe is experimental, still under active development, and the API is subject to change without notice. Use at your own risk. diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 0535f8563f2..3241dbc1fb6 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -477,6 +477,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()); m.def("get_fused_attn_backend", &transformer_engine::pytorch::get_fused_attn_backend, "Get Fused Attention backend", py::arg("fused_attn_params")); + m.def("nvfp4_qdq_supported", &nvte_nvfp4_qdq_supported, + "Whether the common library was built with NVFP4 QDQ kernels"); m.def("nvfp4_qdq", &transformer_engine::pytorch::nvfp4_qdq, "Plain rowwise NVFP4 QDQ", py::arg("input"), py::arg("output"), py::arg("amax"), py::arg("noop") = py::none(), py::arg("compute_amax") = true, py::call_guard()); diff --git a/transformer_engine/pytorch/custom_recipes/qdq.py b/transformer_engine/pytorch/custom_recipes/qdq.py index 726efc87f5b..409df7e7d5b 100644 --- a/transformer_engine/pytorch/custom_recipes/qdq.py +++ b/transformer_engine/pytorch/custom_recipes/qdq.py @@ -18,6 +18,9 @@ from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer +# Query the common library's build, which may differ from PyTorch's CUDA version. +_NVFP4_QDQ_SUPPORTED = tex.nvfp4_qdq_supported() + @lru_cache(maxsize=None) def _supports_fused_device(device_index): @@ -200,7 +203,8 @@ class NVFP4QDQQuantizer(QDQQuantizer): ``backend`` is ``auto``, ``reference``, or ``fused``. Auto falls back to native Q+DQ outside plain contiguous BF16/FP16 1D E4M3 scaling on GPUs with - compute capability 8.0 or higher, with tensor or per-row outer scaling. + compute capability 8.0 or higher and a TE build with CUDA 12.8 or higher, + with tensor or per-row outer scaling. Per-row amax is fused into QDQ. In particular, 4over6 and 2D scaling currently use the reference path. The native reference path requires compute capability 10.0 or higher. @@ -261,7 +265,8 @@ def selected_backend(self, tensor, dtype=None): q = self.nvfp4_quantizer output_dtype = dtype or self.dtype or tensor.dtype eligible = ( - tensor.is_cuda + _NVFP4_QDQ_SUPPORTED + and tensor.is_cuda and tensor.is_contiguous() and tensor.ndim >= 2 and tensor.numel() > 0 @@ -297,7 +302,8 @@ def _reference(self, tensor, dtype): return torch.empty_like(tensor, dtype=dtype) raise RuntimeError( "Native NVFP4 Q+DQ requires compute capability 10.0 or higher; " - "on Ampere/Hopper use auto or fused with plain contiguous matching BF16/FP16" + "on Ampere/Hopper use a TE build with CUDA 12.8 or higher and " + "auto or fused with plain contiguous matching BF16/FP16" ) return self.nvfp4_quantizer.quantize(tensor).dequantize(dtype=dtype) diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 82e9795d02b..e863c475aeb 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Quantization utilities for TransformerEngine""" + from __future__ import annotations import abc @@ -31,7 +32,6 @@ from .utils import get_device_compute_capability from .jit import jit_fuser - __all__ = [ "autocast", "quantized_model_init", @@ -795,8 +795,11 @@ def autocast_enter( qstate.autocast_depth += 1 if enabled: - fp8_available, reason_for_no_fp8 = cls.is_fp8_available() - assert fp8_available, reason_for_no_fp8 + # A custom factory may use only high-precision or emulated formats. + # Its quantizers and GEMMs determine the required device support. + if not fp8_recipe.custom(): + fp8_available, reason_for_no_fp8 = cls.is_fp8_available() + assert fp8_available, reason_for_no_fp8 if isinstance(fp8_recipe, MXFP8BlockScaling): mxfp8_available, reason_for_no_mxfp8 = cls.is_mxfp8_available() assert mxfp8_available, reason_for_no_mxfp8