From 513e3be1aa86e6032d7d87d359cc74a23929d126 Mon Sep 17 00:00:00 2001 From: Bindu Bhargava R Chintam Date: Sat, 31 Oct 2026 02:04:51 -0700 Subject: [PATCH 1/3] [PyTorch] Add export_components to MXFP8 tensors Frameworks that sync MXFP8 training weights to inference engines read Transformer Engine's internal buffers to get the FP8 values and E8M0 scales, and each one handles the scale padding and the GEMM swizzle on its own. Add MXFP8TensorStorage.export_components(). It returns the row-scaled FP8 values with the tensor's shape and compact E8M0 scales with shape (*shape[:-1], shape[-1] // 32). The values are always a view of the tensor's data. The scales are a view too, unless their padding or GEMM swizzle has to be removed, in which case they are copied. Signed-off-by: Bindu Bhargava R Chintam --- .../low_precision_training/mxfp8/mxfp8.rst | 30 +++ tests/pytorch/test_quantized_tensor.py | 202 ++++++++++++++++++ .../tensor/storage/mxfp8_tensor_storage.py | 52 ++++- 3 files changed, 282 insertions(+), 2 deletions(-) diff --git a/docs/features/low_precision_training/mxfp8/mxfp8.rst b/docs/features/low_precision_training/mxfp8/mxfp8.rst index 1827d42cb61..7c71170288c 100644 --- a/docs/features/low_precision_training/mxfp8/mxfp8.rst +++ b/docs/features/low_precision_training/mxfp8/mxfp8.rst @@ -203,6 +203,36 @@ Quantizer qtensor = quantizer.quantize(x) roundtrip = qtensor.dequantize() +Exporting values and scales +--------------------------- + +Integrations that pass MXFP8 weights to another library, for example when syncing training +weights to an inference engine, can get the FP8 values and E8M0 scales of a PyTorch MXFP8 tensor +with :meth:`~transformer_engine.pytorch.MXFP8TensorStorage.export_components`. It works for +quantized parameters and for tensors returned by a quantizer, and the layout of the result does +not depend on how Transformer Engine stores the tensor: ``data`` has the tensor's shape, and +``scale_inv`` holds one biased exponent per 32 consecutive values along the last dimension, with +shape ``(*shape[:-1], shape[-1] // 32)``. + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + weight = torch.randn(1024, 4096, device="cuda", dtype=torch.bfloat16) + quantizer = te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, columnwise=False) + + data, scale_inv = quantizer(weight).export_components() + # data: torch.float8_e4m3fn with shape (1024, 4096) + # scale_inv: torch.uint8 with shape (1024, 128) + scales = torch.exp2(scale_inv.float() - 127).repeat_interleave(32, dim=-1) + dequantized = data.float() * scales + +The values are a view of the tensor's data, so they change when the tensor is updated in place, +for example by an optimizer step. The scales are a view as well, unless the stored scales are +padded along the last dimension or swizzled for GEMM (see `Swizzling scaling factors`_), in which +case they are copied. Clone the returned tensors to keep a snapshot. + Developer Notes --------------- diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 6cd1bd403b1..6043695b774 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -4,6 +4,7 @@ from collections.abc import Iterable import io +import math from typing import Any, Dict, List, Tuple, Union, Optional import pytest @@ -188,6 +189,28 @@ def make_reference_and_test_tensors( return ref, test +def _swizzle_mxfp8_rowwise_scale_inv(scale_inv: torch.Tensor) -> torch.Tensor: + """Reorder padded MXFP8 row-wise scales into the 128x4 tile layout of cuBLAS""" + rows, cols = scale_inv.shape + row = torch.arange(rows, device=scale_inv.device).view(-1, 1) + col = torch.arange(cols, device=scale_inv.device).view(1, -1) + offsets = ( + ((row // 128) * (cols // 4) + col // 4) * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + col % 4 + ) + out = torch.empty_like(scale_inv).view(-1) + out[offsets.view(-1)] = scale_inv.view(-1) + return out.view(rows, cols) + + +def _decode_mxfp8_components(data: torch.Tensor, scale_inv: torch.Tensor) -> torch.Tensor: + """Dequantize exported MXFP8 values and E8M0 scales with plain PyTorch""" + scales = torch.exp2(scale_inv.to(torch.float32) - 127) + return data.to(torch.float32) * scales.repeat_interleave(32, dim=-1) + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) class TestFloat8Tensor: @@ -777,6 +800,109 @@ def test_mxfp8_fsdp_extract_keeps_partial_columnwise_scale_block(self) -> None: assert extracted_scale_inv.shape == (2, 128) torch.testing.assert_close(extracted_scale_inv, columnwise_scale_inv[:2], rtol=0, atol=0) + @pytest.mark.parametrize("with_gemm_swizzled_scales", [False, True]) + @pytest.mark.parametrize("shape", [(128, 128), (96, 160), (2, 64, 128)]) + def test_mxfp8_export_components_layout( + self, + shape: Tuple[int, ...], + with_gemm_swizzled_scales: bool, + ) -> None: + """MXFP8 export gives FP8 views of the data and compact, unswizzled scales""" + rows = math.prod(shape[:-1]) + scale_cols = shape[-1] // 32 + padded_shape = (math.ceil(rows / 128) * 128, math.ceil(scale_cols / 4) * 4) + compact_scale_inv = torch.randint(0, 255, padded_shape, dtype=torch.uint8) + stored_scale_inv = compact_scale_inv + if with_gemm_swizzled_scales: + stored_scale_inv = _swizzle_mxfp8_rowwise_scale_inv(compact_scale_inv) + data = torch.randint(0, 256, shape, dtype=torch.uint8) + storage = MXFP8TensorStorage( + rowwise_data=data, + rowwise_scale_inv=stored_scale_inv, + columnwise_data=None, + columnwise_scale_inv=None, + fp8_dtype=te.DType.kFloat8E4M3, + quantizer=None, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + fake_dtype=torch.bfloat16, + ) + + exported = storage.export_components() + + assert exported.data.dtype == torch.float8_e4m3fn + assert exported.data.shape == shape + assert exported.data.data_ptr() == data.data_ptr() + assert exported.scale_inv.dtype == torch.uint8 + assert exported.scale_inv.shape == (*shape[:-1], scale_cols) + assert exported.scale_inv.is_contiguous() + torch.testing.assert_close( + exported.scale_inv.view(rows, scale_cols), + compact_scale_inv[:rows, :scale_cols], + rtol=0, + atol=0, + ) + + # Scales are copied only when their layout has to change + is_view = exported.scale_inv.data_ptr() == stored_scale_inv.data_ptr() + assert is_view == (not with_gemm_swizzled_scales and scale_cols % 4 == 0) + + def test_mxfp8_export_components_views(self) -> None: + """Exported MXFP8 views follow in-place updates of the tensor""" + data = torch.randint(0, 256, (128, 128), dtype=torch.uint8) + scale_inv = torch.randint(0, 255, (128, 4), dtype=torch.uint8) + storage = MXFP8TensorStorage( + rowwise_data=data, + rowwise_scale_inv=scale_inv, + columnwise_data=None, + columnwise_scale_inv=None, + fp8_dtype=te.DType.kFloat8E5M2, + quantizer=None, + with_gemm_swizzled_scales=False, + fake_dtype=torch.bfloat16, + ) + + exported = storage.export_components() + data.add_(1) + scale_inv.add_(1) + + assert exported.data.dtype == torch.float8_e5m2 + torch.testing.assert_close(exported.data.view(torch.uint8), data, rtol=0, atol=0) + torch.testing.assert_close(exported.scale_inv, scale_inv, rtol=0, atol=0) + + def test_mxfp8_export_components_empty(self) -> None: + """MXFP8 export of an empty tensor gives empty values and scales""" + storage = MXFP8TensorStorage( + rowwise_data=torch.empty((0, 64), dtype=torch.uint8), + rowwise_scale_inv=torch.empty((0, 4), dtype=torch.uint8), + columnwise_data=None, + columnwise_scale_inv=None, + fp8_dtype=te.DType.kFloat8E4M3, + quantizer=None, + with_gemm_swizzled_scales=False, + fake_dtype=torch.bfloat16, + ) + + exported = storage.export_components() + + assert exported.data.shape == (0, 64) + assert exported.scale_inv.shape == (0, 2) + + def test_mxfp8_export_components_requires_rowwise_data(self) -> None: + """MXFP8 export needs the row-scaled data""" + storage = MXFP8TensorStorage( + rowwise_data=None, + rowwise_scale_inv=None, + columnwise_data=torch.zeros((64, 64), dtype=torch.uint8), + columnwise_scale_inv=torch.zeros((4, 128), dtype=torch.uint8), + fp8_dtype=te.DType.kFloat8E4M3, + quantizer=None, + with_gemm_swizzled_scales=False, + fake_dtype=torch.bfloat16, + ) + + with pytest.raises(ValueError, match="row-scaled data"): + storage.export_components() + @pytest.mark.parametrize("quantization", _quantization_list) @pytest.mark.parametrize("dim", [0, 1]) def test_chunk( @@ -1096,6 +1222,82 @@ def test_mxfp8_dequantize_columnwise_only_quantized_separately( with pytest.raises(AssertionError): torch.testing.assert_close(x_deq, -x_ref, **_tols[fp8_dtype]) + @pytest.mark.parametrize( + "fp8_dtype, with_2d_quantization", + [ + (te.DType.kFloat8E4M3, False), + (te.DType.kFloat8E5M2, False), + (te.DType.kFloat8E4M3, True), + ], + ) + @pytest.mark.parametrize("dims", [[128, 256], [96, 160], [4, 64, 128]]) + def test_mxfp8_export_components_matches_dequantize( + self, + fp8_dtype: te.DType, + with_2d_quantization: bool, + dims: DimsType, + ) -> None: + """Exported MXFP8 values and scales decode to the dequantized tensor""" + x = torch.randn(_to_list(dims), dtype=torch.bfloat16, device="cuda") + quantizer = MXFP8Quantizer( + fp8_dtype=fp8_dtype, + rowwise=True, + columnwise=False, + with_2d_quantization=with_2d_quantization, + ) + x_mxfp8 = quantizer(x) + + data, scale_inv = x_mxfp8.export_components() + + torch.testing.assert_close( + _decode_mxfp8_components(data, scale_inv), + x_mxfp8.dequantize(dtype=torch.float32), + rtol=0, + atol=0, + ) + + @pytest.mark.parametrize("dims", [[128, 256], [96, 160], [4, 64, 128]]) + def test_mxfp8_export_components_unswizzles_gemm_scales(self, dims: DimsType) -> None: + """MXFP8 export gives the same scales whether or not they are swizzled for GEMM""" + x = torch.randn(_to_list(dims), dtype=torch.bfloat16, device="cuda") + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=False) + gemm_quantizer = quantizer.copy() + gemm_quantizer.optimize_for_gemm = True + x_mxfp8 = quantizer(x) + x_mxfp8_gemm = gemm_quantizer(x) + assert not x_mxfp8._with_gemm_swizzled_scales + assert x_mxfp8_gemm._with_gemm_swizzled_scales + + expected = x_mxfp8.export_components() + exported = x_mxfp8_gemm.export_components() + + torch.testing.assert_close( + exported.data.view(torch.uint8), expected.data.view(torch.uint8), rtol=0, atol=0 + ) + torch.testing.assert_close(exported.scale_inv, expected.scale_inv, rtol=0, atol=0) + + def test_mxfp8_export_components_of_quantized_weight(self) -> None: + """MXFP8 export of a quantized parameter after a training step""" + recipe = transformer_engine.common.recipe.MXFP8BlockScaling() + with te.quantized_model_init(enabled=True, recipe=recipe): + linear = te.Linear(128, 256, params_dtype=torch.bfloat16, device="cuda") + x = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda", requires_grad=True) + with te.autocast(enabled=True, recipe=recipe): + y = linear(x) + y.sum().backward() + weight = linear.weight + assert isinstance(weight, MXFP8Tensor) + + data, scale_inv = weight.export_components() + + assert data.data_ptr() == weight._rowwise_data.data_ptr() + torch.testing.assert_close( + _decode_mxfp8_components(data, scale_inv), + weight.dequantize(dtype=torch.float32), + rtol=0, + atol=0, + ) + @pytest.mark.parametrize("quantization", _quantization_list) @pytest.mark.parametrize("usage", ["rowwise", "columnwise", "both"]) diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 4d18560dd26..17a012ee017 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -5,7 +5,7 @@ """Mixin class holding data specific for MXFP8Tensor""" from __future__ import annotations -from typing import Annotated, Optional, Dict, Any, Tuple, Union +from typing import Annotated, NamedTuple, Optional, Dict, Any, Tuple, Union from collections.abc import Iterable import math import torch @@ -18,11 +18,12 @@ from ...constants import ( TE_DType as torch_to_transformer_engine_dtype, + TE_DType_To_Torch, MXFP8_BLOCK_SCALING_SIZE, DType, ) -from ...utils import _empty_tensor +from ...utils import _empty_tensor, round_up_to_nearest_multiple class _FromMXFP8Func(torch.autograd.Function): @@ -61,6 +62,23 @@ def backward( return grad, None +class MXFP8Components(NamedTuple): + """FP8 values and E8M0 scales of an MXFP8 tensor + + Each block of 32 consecutive values along the last dimension + shares one scale, so the dequantized value at ``[..., j]`` is + ``data[..., j] * 2 ** (scale_inv[..., j // 32] - 127)``. + + """ + + # FP8 values (torch.float8_e4m3fn or torch.float8_e5m2) with the + # tensor's shape + data: torch.Tensor + # E8M0 scales as biased exponents (torch.uint8) with shape + # (*shape[:-1], shape[-1] // 32) + scale_inv: torch.Tensor + + class MXFP8TensorStorage(QuantizedTensorStorage): """Mixin class that holds data attributes of MXFP8Tensor. @@ -198,6 +216,36 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: return torch.empty(self.size(), dtype=dtype, device=self.device) return _FromMXFP8Func.forward(None, self, dtype) + def export_components(self) -> MXFP8Components: + """Get the FP8 values and E8M0 scales for use outside Transformer Engine + + The row-scaled data is exported, with compact scales: no + padding and no swizzling for GEMMs. The values are a view of + this tensor's data, and so are the scales unless removing + their padding or GEMM swizzling needs a copy. Views follow + in-place updates of this tensor, e.g. optimizer steps, so + clone them to keep a snapshot. + + """ + if self._rowwise_data is None: + raise ValueError( + "MXFP8 export requires row-scaled data, which this tensor does not have" + ) + shape = self._rowwise_data.shape + rows = math.prod(shape[:-1]) + scale_cols = shape[-1] // MXFP8_BLOCK_SCALING_SIZE + + data = self._rowwise_data.view(TE_DType_To_Torch[self._fp8_dtype]) + scale_inv = self._rowwise_scale_inv + if self._with_gemm_swizzled_scales: + # Undo the 128x4 tile layout that cuBLAS expects for block scales + padded_rows = round_up_to_nearest_multiple(rows, 128) + padded_cols = round_up_to_nearest_multiple(scale_cols, 4) + scale_inv = scale_inv.reshape(padded_rows // 128, padded_cols // 4, 32, 4, 4) + scale_inv = scale_inv.permute(0, 3, 2, 1, 4).reshape(padded_rows, padded_cols) + scale_inv = scale_inv[:rows, :scale_cols].contiguous() + return MXFP8Components(data=data, scale_inv=scale_inv.view(*shape[:-1], scale_cols)) + def size(self, *args, **kwargs): # pylint: disable=missing-function-docstring if self._rowwise_data is not None: From b6de4aa8e1a7953e021919024e4a95dab249751c Mon Sep 17 00:00:00 2001 From: Bindu Bhargava R Chintam Date: Sat, 31 Oct 2026 16:00:00 -0700 Subject: [PATCH 2/3] [PyTorch] Check that exported MXFP8 views follow optimizer updates test_mxfp8_export_components_of_quantized_weight now exports the weight before a FusedAdam step with FP32 master weights. After the step, the exported values and scales must equal a fresh rowwise quantization of the new master weights, and must differ from what they decoded to before the step. Signed-off-by: Bindu Bhargava R Chintam --- tests/pytorch/test_quantized_tensor.py | 31 +++++++++++++++++++++----- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 6043695b774..8c5738260d8 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -1277,7 +1277,7 @@ def test_mxfp8_export_components_unswizzles_gemm_scales(self, dims: DimsType) -> torch.testing.assert_close(exported.scale_inv, expected.scale_inv, rtol=0, atol=0) def test_mxfp8_export_components_of_quantized_weight(self) -> None: - """MXFP8 export of a quantized parameter after a training step""" + """MXFP8 export of a quantized parameter follows its optimizer updates""" recipe = transformer_engine.common.recipe.MXFP8BlockScaling() with te.quantized_model_init(enabled=True, recipe=recipe): linear = te.Linear(128, 256, params_dtype=torch.bfloat16, device="cuda") @@ -1291,12 +1291,33 @@ def test_mxfp8_export_components_of_quantized_weight(self) -> None: data, scale_inv = weight.export_components() assert data.data_ptr() == weight._rowwise_data.data_ptr() + decoded_before_step = _decode_mxfp8_components(data, scale_inv) torch.testing.assert_close( - _decode_mxfp8_components(data, scale_inv), - weight.dequantize(dtype=torch.float32), - rtol=0, - atol=0, + decoded_before_step, weight.dequantize(dtype=torch.float32), rtol=0, atol=0 + ) + + # FusedAdam writes the updated master weights back into the parameter's + # buffers, so the exported views must now hold the new values and scales + optimizer = te.optimizers.FusedAdam( + [weight], + lr=0.1, + master_weights=True, + master_weight_dtype=torch.float32, + use_decoupled_grad=True, + ) + weight.decoupled_grad = torch.randn(weight.shape, dtype=torch.float32, device="cuda") + optimizer.step() + quantizer = weight._quantizer.copy() + quantizer.set_usage(rowwise=True, columnwise=False) + expected = quantizer( + optimizer.get_unscaled_state(weight, "master_param") + ).export_components() + + torch.testing.assert_close( + data.view(torch.uint8), expected.data.view(torch.uint8), rtol=0, atol=0 ) + torch.testing.assert_close(scale_inv, expected.scale_inv, rtol=0, atol=0) + assert not torch.equal(_decode_mxfp8_components(data, scale_inv), decoded_before_step) @pytest.mark.parametrize("quantization", _quantization_list) From ed0362adf23a876f7c0f4e786e022db720ed3d42 Mon Sep 17 00:00:00 2001 From: Bindu Bhargava R Chintam Date: Sat, 7 Nov 2026 17:25:08 -0800 Subject: [PATCH 3/3] [PyTorch] Export a scale for a partial last MXFP8 block Since #2894, MXFP8 accepts a last dimension that is a multiple of 16, and a partial last block keeps its own scale, so rowwise tensors carry ceil(K / 32) scale columns. export_components now returns all of them instead of K // 32. Tests cover K = 48 and K = 16, including rows that are not a multiple of 32 and GEMM-swizzled scales. Signed-off-by: Bindu Bhargava R Chintam --- .../low_precision_training/mxfp8/mxfp8.rst | 4 +-- tests/pytorch/test_quantized_tensor.py | 31 ++++++++++++++++--- .../tensor/storage/mxfp8_tensor_storage.py | 9 +++--- 3 files changed, 34 insertions(+), 10 deletions(-) diff --git a/docs/features/low_precision_training/mxfp8/mxfp8.rst b/docs/features/low_precision_training/mxfp8/mxfp8.rst index 07a48e57277..11c89e5ee0b 100644 --- a/docs/features/low_precision_training/mxfp8/mxfp8.rst +++ b/docs/features/low_precision_training/mxfp8/mxfp8.rst @@ -221,7 +221,7 @@ with :meth:`~transformer_engine.pytorch.MXFP8TensorStorage.export_components`. I quantized parameters and for tensors returned by a quantizer, and the layout of the result does not depend on how Transformer Engine stores the tensor: ``data`` has the tensor's shape, and ``scale_inv`` holds one biased exponent per 32 consecutive values along the last dimension, with -shape ``(*shape[:-1], shape[-1] // 32)``. +shape ``(*shape[:-1], ceil(shape[-1] / 32))``; the last block may be shorter than 32. .. code-block:: python @@ -234,7 +234,7 @@ shape ``(*shape[:-1], shape[-1] // 32)``. data, scale_inv = quantizer(weight).export_components() # data: torch.float8_e4m3fn with shape (1024, 4096) # scale_inv: torch.uint8 with shape (1024, 128) - scales = torch.exp2(scale_inv.float() - 127).repeat_interleave(32, dim=-1) + scales = torch.exp2(scale_inv.float() - 127).repeat_interleave(32, dim=-1)[..., : data.shape[-1]] dequantized = data.float() * scales The values are a view of the tensor's data, so they change when the tensor is updated in place, diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 65067982cf9..2151535f4f7 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -208,7 +208,8 @@ def _swizzle_mxfp8_rowwise_scale_inv(scale_inv: torch.Tensor) -> torch.Tensor: def _decode_mxfp8_components(data: torch.Tensor, scale_inv: torch.Tensor) -> torch.Tensor: """Dequantize exported MXFP8 values and E8M0 scales with plain PyTorch""" scales = torch.exp2(scale_inv.to(torch.float32) - 127) - return data.to(torch.float32) * scales.repeat_interleave(32, dim=-1) + scales = scales.repeat_interleave(32, dim=-1)[..., : data.shape[-1]] + return data.to(torch.float32) * scales @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @@ -840,7 +841,7 @@ def test_mxfp8_fsdp_extract_keeps_partial_columnwise_scale_block(self) -> None: torch.testing.assert_close(extracted_scale_inv, columnwise_scale_inv[:2], rtol=0, atol=0) @pytest.mark.parametrize("with_gemm_swizzled_scales", [False, True]) - @pytest.mark.parametrize("shape", [(128, 128), (96, 160), (2, 64, 128)]) + @pytest.mark.parametrize("shape", [(128, 128), (96, 160), (2, 64, 128), (100, 48)]) def test_mxfp8_export_components_layout( self, shape: Tuple[int, ...], @@ -848,7 +849,7 @@ def test_mxfp8_export_components_layout( ) -> None: """MXFP8 export gives FP8 views of the data and compact, unswizzled scales""" rows = math.prod(shape[:-1]) - scale_cols = shape[-1] // 32 + scale_cols = math.ceil(shape[-1] / 32) padded_shape = (math.ceil(rows / 128) * 128, math.ceil(scale_cols / 4) * 4) compact_scale_inv = torch.randint(0, 255, padded_shape, dtype=torch.uint8) stored_scale_inv = compact_scale_inv @@ -1330,7 +1331,29 @@ def test_mxfp8_export_components_matches_dequantize( atol=0, ) - @pytest.mark.parametrize("dims", [[128, 256], [96, 160], [4, 64, 128]]) + @pytest.mark.parametrize("fp8_dtype", _fp8_dtypes) + @pytest.mark.parametrize("dims", [[100, 48], [32, 16]]) + def test_mxfp8_export_components_partial_block( + self, + fp8_dtype: te.DType, + dims: DimsType, + ) -> None: + """A last block shorter than 32 values keeps its own exported scale""" + x = torch.randn(_to_list(dims), dtype=torch.bfloat16, device="cuda") + quantizer = MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=False) + x_mxfp8 = quantizer(x) + + data, scale_inv = x_mxfp8.export_components() + + assert scale_inv.shape == (*dims[:-1], math.ceil(dims[-1] / 32)) + torch.testing.assert_close( + _decode_mxfp8_components(data, scale_inv), + x_mxfp8.dequantize(dtype=torch.float32), + rtol=0, + atol=0, + ) + + @pytest.mark.parametrize("dims", [[128, 256], [96, 160], [4, 64, 128], [100, 48]]) def test_mxfp8_export_components_unswizzles_gemm_scales(self, dims: DimsType) -> None: """MXFP8 export gives the same scales whether or not they are swizzled for GEMM""" x = torch.randn(_to_list(dims), dtype=torch.bfloat16, device="cuda") diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 17a012ee017..239467e0e2e 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -23,7 +23,7 @@ DType, ) -from ...utils import _empty_tensor, round_up_to_nearest_multiple +from ...utils import _empty_tensor, ceil_div, round_up_to_nearest_multiple class _FromMXFP8Func(torch.autograd.Function): @@ -66,7 +66,8 @@ class MXFP8Components(NamedTuple): """FP8 values and E8M0 scales of an MXFP8 tensor Each block of 32 consecutive values along the last dimension - shares one scale, so the dequantized value at ``[..., j]`` is + shares one scale (the last block may be shorter), so the + dequantized value at ``[..., j]`` is ``data[..., j] * 2 ** (scale_inv[..., j // 32] - 127)``. """ @@ -75,7 +76,7 @@ class MXFP8Components(NamedTuple): # tensor's shape data: torch.Tensor # E8M0 scales as biased exponents (torch.uint8) with shape - # (*shape[:-1], shape[-1] // 32) + # (*shape[:-1], ceil(shape[-1] / 32)) scale_inv: torch.Tensor @@ -233,7 +234,7 @@ def export_components(self) -> MXFP8Components: ) shape = self._rowwise_data.shape rows = math.prod(shape[:-1]) - scale_cols = shape[-1] // MXFP8_BLOCK_SCALING_SIZE + scale_cols = ceil_div(shape[-1], MXFP8_BLOCK_SCALING_SIZE) data = self._rowwise_data.view(TE_DType_To_Torch[self._fp8_dtype]) scale_inv = self._rowwise_scale_inv