diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 5ae657c5d2..66857d8125 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -36,6 +36,7 @@ ) from transformer_engine.pytorch import ( QuantizedTensor, + Float8BlockQuantizer, Float8CurrentScalingQuantizer, Float8Quantizer, MXFP8Quantizer, @@ -58,6 +59,9 @@ fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) # Supported data types _dtypes: list[torch.dtype] = [torch.float32, torch.float16] @@ -76,6 +80,8 @@ if nvfp4_available: _quantization_list.append("nvfp4") _quantization_list.append("nvfp4_4over6") +if fp8_block_scaling_available: + _quantization_list.append("fp8_block_scaling") @pytest.fixture(autouse=True, scope="function") @@ -110,6 +116,8 @@ def maybe_skip_quantization( and not nvfp4_available ): pytest.skip(reason_for_no_nvfp4) + if quantization == "fp8_block_scaling" and not fp8_block_scaling_available: + pytest.skip(reason_for_no_fp8_block_scaling) # Check dims if dims is not None: @@ -121,6 +129,9 @@ def maybe_skip_quantization( elif quantization == "mxfp8": if math.prod(dims[:-1]) % 32 != 0 or dims[-1] % 32 != 0: pytest.skip("MXFP8 GEMMs require dims that are divisible by 32") + elif quantization == "fp8_block_scaling": + if math.prod(dims[:-1]) % 128 != 0 or dims[-1] % 128 != 0: + pytest.skip("FP8 block scaling requires dims that are divisible by 128") elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") @@ -186,6 +197,17 @@ def make_reference_and_test_tensors( test = quantizer(test) elif quantization == "mxfp8": test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test) + elif quantization == "fp8_block_scaling": + tensor_type = "input" + if quantizer_role is not None: + tensor_type = quantizer_role.tensor_type + # Weights use 2D (128x128) blocks; everything else 1D (1x128). + test = Float8BlockQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2 if tensor_type == "weight" else 1, + )(test) elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht"): tensor_type = "input" if quantizer_role is not None: @@ -1017,8 +1039,9 @@ def _test_basic_linear( ) torch.testing.assert_close(dw_test, w_ref.grad, **tols) - @pytest.mark.parametrize("weight_shape", ((64, 32), (3, 5))) - @pytest.mark.parametrize("in_shape", ((-1,), (5, 1, -1), (4, 2, 4, -1))) + # (128, 128) + a 128-token in_shape keep FP8 block scaling (128-divisible dims) unskipped. + @pytest.mark.parametrize("weight_shape", ((64, 32), (3, 5), (128, 128))) + @pytest.mark.parametrize("in_shape", ((-1,), (5, 1, -1), (4, 2, 4, -1), (128, -1))) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) @@ -3360,7 +3383,7 @@ def test_layernorm_mlp( quantization: Optional[str], device: torch.device = "cuda", hidden_size: int = 256, - sequence_length: int = 48, + sequence_length: int = 64, batch_size: int = 4, ffn_hidden_size: int = 384, layernorm_epsilon: float = 1e-5, diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index f068f3581e..21601d8cdd 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -114,6 +114,7 @@ def quantization_tols(name: str) -> dict[str, float]: "fp8_delayed_scaling", "fp8_current_scaling", "fp8_blockwise", + "fp8_block_scaling", "mxfp8", "mxfp8_block_scaling", ): diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index af3f8b5930..d3e2dc0a13 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -27,6 +27,7 @@ from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe from ...quantized_tensor import QuantizedTensorStorage from ...tensor import ( + Float8BlockQuantizer, Float8CurrentScalingQuantizer, MXFP8Quantizer, MXFP8Tensor, @@ -791,9 +792,10 @@ def _is_graph_safe_path_supported( * FP8 per-tensor current scaling is backed by grouped current-scaling quantization (``tex.group_quantize``) and cuBLASLt grouped GEMM with per-batch scalar FP8 scaling, which are supported on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0). - Every other quantization recipe (fp8 delayed scaling, fp8 block scaling, ...) - falls back to the legacy flow because the corresponding grouped quantization kernels are - missing. + * FP8 block scaling uses the grouped-tensor path on Hopper (CC 9.0) with cuBLAS 13.4+; + it is Hopper-only (no MXFP8-broadcast emulation), so elsewhere it falls back. + * Every other quantization recipe (fp8 delayed scaling, ...) falls back to the legacy flow + because the corresponding grouped quantization kernels are missing. * Unquantized compute supports BF16/FP16 on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0) -- FP32 is excluded because the cuBLASLt grouped GEMM doesn't support it. * Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it @@ -812,6 +814,12 @@ def _is_graph_safe_path_supported( ): return False return True + if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): + # Grouped FP8 block scaling is Hopper-only and needs cuBLAS 13.4+; elsewhere + # fall back to the split-quantize (MXFP8-emulated) flow. + if get_device_compute_capability() >= (10, 0): + return False + return tex.get_cublasLt_version() >= 130400 # MXFP8 and NVFP4 grouped quantization kernels require Blackwell. if not (10, 0) <= get_device_compute_capability() <= (11, 0): return False @@ -1666,11 +1674,12 @@ def _fuser_backward_grouped_tensor( ) grad_output_quantizer.optimize_for_gemm = True - if ( - has_bias - and not self._scale_bias - and isinstance(grad_output_quantizer, MXFP8Quantizer) - ): + # FP8 block scaling computes dbias in the rowwise (dgrad) pass, so only fuse + # when dgrad is required. + fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or ( + isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.input_requires_grad + ) + if has_bias and not self._scale_bias and fuse_bgrad: grouped_dy, dbias_packed = tex.bgrad_group_quantize( dy_2d, grad_output_quantizer, num_groups, split_sizes ) diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 849c900f95..5106ec9e0a 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -275,11 +275,6 @@ def reset_recipe_state( if num_quantizers == 0: continue - if recipe.float8_block_scaling(): - raise NotImplementedError( - "Fusible operations do not support FP8 block scaling recipe" - ) - # Construct quantization recipe state roles = self.get_quantizer_roles(mode) # pylint: disable=assignment-from-none if roles is not None: diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index b24a4e9144..464a7c3b23 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -5,6 +5,7 @@ """Mixin class holding data specific for Float8BlockwiseQTensor""" from __future__ import annotations +from collections.abc import Iterable import math from typing import Optional, Dict, Any, Tuple, Union import torch @@ -323,6 +324,61 @@ def size(self, *args, **kwargs): reordered.append(dims[0]) return torch.Size(reordered) + def view(self, shape): + """Reshape the leading (token) dims without dequantizing. + + Mirrors ``MXFP8TensorStorage.view``. Only leading dims may change: block tiling fixes + the inner dim (1D) or inner two (2D), so scale-inv stays valid. Columnwise data is + stored transposed, viewed as ``[inner, *leading]``. + """ + cur_shape = self.size() + if shape is None or shape == cur_shape: + return self + # Canonicalize shape + if not isinstance(shape, Iterable): + shape = [shape] + elif len(shape) == 1 and isinstance(shape[0], Iterable): + shape = shape[0] + shape = list(shape) + if -1 in shape: + d_inferred = -math.prod(cur_shape) // math.prod(shape) + for i, d in enumerate(shape): + if d == -1: + shape[i] = d_inferred + break + if shape == list(cur_shape): + return self + + if self._is_2D_scaled: + if shape[-2:] != list(cur_shape)[-2:]: + raise RuntimeError( + "Float8BlockwiseQTensorStorage (2D block scaling) cannot reshape the inner " + f"two dimensions (attempted {tuple(cur_shape)} -> {tuple(shape)})" + ) + elif shape[-1] != cur_shape[-1]: + raise RuntimeError( + "Float8BlockwiseQTensorStorage (1D block scaling) cannot reshape the inner " + f"dimension (attempted {tuple(cur_shape)} -> {tuple(shape)})" + ) + + new_rowwise_data = None + if self._rowwise_data is not None: + new_rowwise_data = self._rowwise_data.view(*shape) + new_columnwise_data = None + if self._columnwise_data is not None: + new_columnwise_data = self._columnwise_data.view([shape[-1], *shape[:-1]]) + + return Float8BlockwiseQTensorStorage( + new_rowwise_data, + self._rowwise_scale_inv, + new_columnwise_data, + self._columnwise_scale_inv, + self._fp8_dtype, + self._quantizer, + self._is_2D_scaled, + fake_dtype=self._dtype, + ) + @property def device(self): """Return the device of the tensor. Define this to avoid expensive PyObject lookups."""