[PyTorch] Add export_components to MXFP8 tensors

docs/features/low_precision_training/mxfp8/mxfp8.rst

@@ -212,6 +212,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], ceil(shape[-1] / 32))``; the last block may be shorter than 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)[..., : 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,
+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
 ---------------
 

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,29 @@ 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)
+    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)
 class TestFloat8Tensor:
 
@@ -816,6 +840,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), (100, 48)])
+    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 = 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
+        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(
@@ -1170,6 +1297,125 @@ 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("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")
+        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 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")
+        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()
+        decoded_before_step = _decode_mxfp8_components(data, scale_inv)
+        torch.testing.assert_close(
+            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)
 @pytest.mark.parametrize("usage", ["rowwise", "columnwise", "both"])

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, ceil_div, round_up_to_nearest_multiple
 
 
 class _FromMXFP8Func(torch.autograd.Function):
@@ -61,6 +62,24 @@ 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 (the last block may be shorter), 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], ceil(shape[-1] / 32))
+    scale_inv: torch.Tensor
+
+
 class MXFP8TensorStorage(QuantizedTensorStorage):
     """Mixin class that holds data attributes of MXFP8Tensor.
 
@@ -198,6 +217,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 = 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
+        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: