Improve QAT nvfp4 numerics

**Summary:** Similar to #2986,
this commit improves the prepare vs convert SQNR of NVFP4 QAT
from 12 to 36 with `use_per_tensor_scale`, and 12 to inf without.
This is achieved by mimicking the PTQ flow more closely,
in particular, in descending order of significance:

1. Simulate `f4_unpacked_to_f32` and `f32_to_f4_unpacked`,
   but in `torch.int32` instead of `torch.uint8`
2. Do not cast intermediate fake quantized values to original
   dtype, e.g. bf16 which loses some fidelity from fp32
3. Fake round blockwise scales to float8

**Test Plan:**
```
python test/quantization/test_qat.py -k test_qat_nvfp4
python test/quantization/test_qat.py -k test_quantize_api_nvfp4
```

End-to-end tests TBD.

ghstack-source-id: 04f6bce
Pull Request resolved: #3050

Author: andrewor14

test/quantization/test_qat.py

@@ -1910,7 +1910,6 @@ def _test_quantize_api_against_ptq(
         quantize_(m, QATConfig(base_config, step="prepare"), filter_fn)
         out_prepared = m(*example_inputs)
         prepare_sqnr = compute_error(out_prepared, out_baseline)
-
         self.assertGreaterEqual(prepare_sqnr, target_prepare_sqnr)
 
         # compare convert
@@ -2086,9 +2085,14 @@ def test_quantize_api_nvfp4(self, use_per_tensor_scale: bool):
         """
         from torchao.prototype.mx_formats import NVFP4InferenceConfig
 
+        if use_per_tensor_scale:
+            target_prepare_sqnr = 36
+        else:
+            target_prepare_sqnr = float("inf")
+
         self._test_quantize_api_against_ptq(
             NVFP4InferenceConfig(use_dynamic_per_tensor_scale=use_per_tensor_scale),
-            target_prepare_sqnr=12,
+            target_prepare_sqnr=target_prepare_sqnr,
             target_convert_sqnr=float("inf"),
         )
 
@@ -2116,7 +2120,7 @@ def test_qat_nvfp4(self, use_per_tensor_scale: bool):
         out = m(*x)
         baseline_out = baseline_model(*x)
         sqnr = compute_error(out, baseline_out).item()
-        self.assertGreater(sqnr, 24)
+        self.assertGreater(sqnr, 10)
 
     @unittest.skipIf(not _CUDA_IS_AVAILABLE, "skipping when cuda is not available")
     @unittest.skipIf(

torchao/prototype/custom_fp_utils.py

@@ -24,7 +24,9 @@ def _n_ones(n: int) -> int:
 F32_EXP_BIAS = _n_ones(EBITS_F32 - 1)
 
 
-def _f32_to_floatx_unpacked(x: Tensor, ebits: int, mbits: int) -> Tensor:
+def _f32_to_floatx_unpacked(
+    x: Tensor, ebits: int, mbits: int, fake_quantize: bool = False
+) -> Tensor:
     """Convert FP32 numbers to sub-byte floating point numbers with the given
     number of exponent and mantissa bits.
 
@@ -105,7 +107,8 @@ def _f32_to_floatx_unpacked(x: Tensor, ebits: int, mbits: int) -> Tensor:
     denormal_x = x + denorm_mask_float
     denormal_x = denormal_x.view(torch.int32)
     denormal_x -= denorm_mask_int
-    denormal_x = denormal_x.to(torch.uint8)
+    if not fake_quantize:
+        denormal_x = denormal_x.to(torch.uint8)
 
     #
     # branch 3: stay in normal range, adjust the exponent and round
@@ -120,31 +123,41 @@ def _f32_to_floatx_unpacked(x: Tensor, ebits: int, mbits: int) -> Tensor:
     normal_x += mant_odd
     # take the bits!
     normal_x = normal_x >> (MBITS_F32 - mbits)
-    normal_x = normal_x.to(torch.uint8)
+    if not fake_quantize:
+        normal_x = normal_x.to(torch.uint8)
 
     #
     # combine the branches
     #
-    x = torch.full_like(x, max_int, dtype=torch.uint8)
+    if fake_quantize:
+        x = torch.full_like(x, max_int, dtype=torch.int32)
+    else:
+        x = torch.full_like(x, max_int, dtype=torch.uint8)
     x = torch.where(denormal_mask, denormal_x, x)
     x = torch.where(normal_mask, normal_x, x)
 
     # add sign back
     sign_lp = sign >> (MBITS_F32 + EBITS_F32 - mbits - ebits)
-    sign_lp = sign_lp.to(torch.uint8)
+    if not fake_quantize:
+        sign_lp = sign_lp.to(torch.uint8)
     # Right shift of a negative signed integer can fill the least significant
     # bits with either 1s or 0s, depending on the implementation. Since PyTorch
     # doesn't have an uint32 dtype, we mask out these bits to get just the
     # f4 sign bit
     sign_lp = sign_lp & sign_mask
     x = x | sign_lp
 
-    return x.to(torch.uint8)
+    if fake_quantize:
+        return x
+    else:
+        return x.to(torch.uint8)
 
 
 # TODO(future): check if LUT for everything is faster than bit shifting,
 # especially for fp4 (only 2^4=16 unique values).
-def _floatx_unpacked_to_f32(x: Tensor, ebits: int, mbits: int) -> Tensor:
+def _floatx_unpacked_to_f32(
+    x: Tensor, ebits: int, mbits: int, fake_quantize: bool = False
+) -> Tensor:
     """Convert sub-byte floating point numbers with the given number of exponent
     and mantissa bits to FP32.
 
@@ -154,7 +167,8 @@ def _floatx_unpacked_to_f32(x: Tensor, ebits: int, mbits: int) -> Tensor:
       fp6: bits 0-1 empty and bits 2-7 in fp6_e2m3 or fp6_e3m2 encoding
     Output: torch.Tensor of dtype fp32 with the dequantized value
     """
-    assert x.dtype == torch.uint8
+    if not fake_quantize:
+        assert x.dtype == torch.uint8
     assert 1 + ebits + mbits <= 8
 
     sign_mask = 1 << (ebits + mbits)

torchao/prototype/mx_formats/kernels.py

@@ -65,13 +65,13 @@ def get_bits(x: torch.Tensor) -> str:
 ZERO_POINT_FIVE_BITS_F32 = 0x3F000000
 
 
-def f32_to_f4_unpacked(x):
+def f32_to_f4_unpacked(x, fake_quantize: bool = False):
     """
     Input: torch.Tensor of dtype torch.float
     Output: torch.Tensor of dtype torch.uint8, with bits 0-3 empty and
       bits 4-7 in fp4_e2m1
     """
-    return _f32_to_floatx_unpacked(x, EBITS_F4_E2M1, MBITS_F4_E2M1)
+    return _f32_to_floatx_unpacked(x, EBITS_F4_E2M1, MBITS_F4_E2M1, fake_quantize)
 
 
 def f32_to_f6_e2m3_unpacked(x):
@@ -92,13 +92,13 @@ def f32_to_f6_e3m2_unpacked(x):
     return _f32_to_floatx_unpacked(x, EBITS_F6_E3M2, MBITS_F6_E3M2)
 
 
-def f4_unpacked_to_f32(x: torch.Tensor):
+def f4_unpacked_to_f32(x: torch.Tensor, fake_quantize: bool = False):
     """
     Input: torch.Tensor of dtype uint8, with bits 0-3 empty and bits 4-7
       containing an fp4_e2m1 encoding
     Output: torch.Tensor of dtype fp32 with the dequantized value
     """
-    return _floatx_unpacked_to_f32(x, EBITS_F4_E2M1, MBITS_F4_E2M1)
+    return _floatx_unpacked_to_f32(x, EBITS_F4_E2M1, MBITS_F4_E2M1, fake_quantize)
 
 
 def f6_e2m3_unpacked_to_f32(x: torch.Tensor):

torchao/prototype/mx_formats/nvfp4_tensor.py

@@ -798,7 +798,6 @@ def _nvfp4_quantize(
     assert data_hp.is_contiguous(), "Only support contiguous data for now"
     assert block_size == 16, "NVFP4 requires block_size=16"
 
-    orig_dtype = data_hp.dtype
     orig_shape = data_hp.shape
     # Convert to float32 early for consistent precision with Triton implementation
     data_hp = data_hp.float().reshape(orig_shape[0], -1, block_size)
@@ -834,7 +833,7 @@ def _nvfp4_quantize(
     data_scaled = torch.clamp(data_scaled, -F4_E2M1_MAX, F4_E2M1_MAX)
     data_scaled = data_scaled.view(orig_shape)
     if skip_dtype_cast_and_packing:
-        return out_scales.to(torch.float32), data_scaled.to(orig_dtype)
+        return _Float8Round.apply(out_scales), data_scaled
     else:
         data_lp = f32_to_f4_unpacked(data_scaled)
         # TODO: NotImplementedError: "copy_kernel" not implemented for 'Float4_e2m1fn_x2'

torchao/prototype/qat/nvfp4.py

@@ -2,6 +2,10 @@
 
 import torch
 
+from torchao.prototype.mx_formats.kernels import (
+    f4_unpacked_to_f32,
+    f32_to_f4_unpacked,
+)
 from torchao.prototype.mx_formats.nvfp4_tensor import (
     _nvfp4_quantize,
     per_tensor_amax_to_scale,
@@ -56,13 +60,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             per_tensor_scale=per_tensor_scale,
             skip_dtype_cast_and_packing=True,
         )
+        q = f32_to_f4_unpacked(q, fake_quantize=True)
         if self.config.use_per_tensor_scale:
             scale = scale * per_tensor_scale
-        assert q.dtype == x.dtype
         assert scale.dtype == torch.float32
 
         # dequantize
         M, K = q.shape[0], q.shape[1]
+        q = f4_unpacked_to_f32(q, fake_quantize=True)
         q = q.view(M, K // block_size, block_size)
         scale = scale.view(M, K // block_size, 1)
         dq = q * scale