Check numerical equivalence / closeness between different kernel preferences

Summary:
This PR checks different kernel preferences for Float8Tensor are similar in numerics
(AUTO, TORCH and FBGEMM)

triton implementation and torchao implementation are a bit different right now actually, need to decide if we should fix it or not

1. difference in quantize op
main difference seems to be the triton implementation is using:
```
a_scale = MAX_FP8 / max_abs
then do
a_scale = 1.0 / a_scale
a_fp8 = a * a_scale
```

while torch is doing:
```
a_scale = max_abs / MAX_FP8
a_fp8 = a / a_scale
```

Also the hp_value_lb and hp_value_ub settings are slightly different

triton choose scale and quantize code: https://github.com/pytorch/FBGEMM/blob/a4286c01ef01dad435b2ec8798605127d3032cd8/fbgemm_gpu/experimental/gemm/triton_gemm/fp8_gemm.py#L2382-L2392

torchao choose scale and quantize code:
https://github.com/pytorch/ao/blob/3c466f844684af0fb80014094f2ca8663881eb33/torchao/quantization/quant_primitives.py#L2183
https://github.com/pytorch/ao/blob/3c466f844684af0fb80014094f2ca8663881eb33/torchao/quantization/quant_primitives.py#L2283

2. (potentially) difference in matrix multiplication ops

TORCH and AUTO/FBGEMM are using different quantized mm ops

Added a reverse option to bring sqnr closer:
```
granularity: PerTensor()  sizes: ((128,), 256, 128)  kp: KernelPreference.AUTO tensor(inf, device='cuda:0', dtype=torch.bfloat16)
granularity: PerTensor()  sizes: ((128,), 256, 128)  kp: KernelPreference.FBGEMM tensor(inf, device='cuda:0', dtype=torch.bfloat16)
.granularity: PerTensor()  sizes: ((32, 128), 64, 256)  kp: KernelPreference.AUTO tensor(inf, device='cuda:0', dtype=torch.bfloat16)
granularity: PerTensor()  sizes: ((32, 128), 64, 256)  kp: KernelPreference.FBGEMM tensor(inf, device='cuda:0', dtype=torch.bfloat16)
.granularity: PerRow()  sizes: ((128,), 256, 128)  kp: KernelPreference.AUTO tensor(inf, device='cuda:0', dtype=torch.bfloat16)
granularity: PerRow()  sizes: ((128,), 256, 128)  kp: KernelPreference.FBGEMM tensor(inf, device='cuda:0', dtype=torch.bfloat16)
.granularity: PerRow()  sizes: ((32, 128), 64, 256)  kp: KernelPreference.AUTO tensor(64.5000, device='cuda:0', dtype=torch.bfloat16)
granularity: PerRow()  sizes: ((32, 128), 64, 256)  kp: KernelPreference.FBGEMM tensor(68., device='cuda:0', dtype=torch.bfloat16)
```
Test Plan:
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_kernel_preference_numerical_equivalence

Reviewers:

Subscribers:

Tasks:

Tags:

stack-info: PR: #2651, branch: jerryzh168/stack/15

Author: jerryzh168

test/quantization/quantize_/workflows/float8/test_float8_tensor.py

@@ -268,6 +268,59 @@ def test_slice(self, granularity):
         sqnr = compute_error(res, res_ref)
         self.assertTrue(sqnr > 15, f"sqnr: {sqnr}")
 
+    @common_utils.parametrize("granularity", [PerTensor(), PerRow()])
+    # Inputs are (M,..), K, N
+    @common_utils.parametrize(
+        "sizes",
+        [
+            ((128,), 256, 128),
+            ((32, 128), 64, 256),
+        ],
+    )
+    def test_kernel_preference_numerical_equivalence(self, granularity, sizes):
+        """Test different kernel preferences have the same numerics for float8 dynamic activation
+        and float8 weight config
+        """
+        M, N, K = sizes
+        dtype = torch.bfloat16
+        input_tensor = torch.randn(*M, K, dtype=dtype, device="cuda")
+        # Create a linear layer with bfloat16 dtype
+        model = ToyLinearModel(K, N).eval().to(dtype).to("cuda")
+
+        kernel_preferences = [
+            KernelPreference.TORCH,
+            KernelPreference.AUTO,
+            KernelPreference.FBGEMM,
+        ]
+        quantized_outputs = {}
+        for kp in kernel_preferences:
+            config = Float8DynamicActivationFloat8WeightConfig(
+                granularity=granularity, kernel_preference=kp
+            )
+            quantized_model = copy.deepcopy(model)
+            quantize_(quantized_model, config)
+            quantized_outputs[kp] = quantized_model(input_tensor)
+
+        from torchao.quantization.utils import compute_error
+
+        # comparing numerics between different kernel preferences, using TORCH as the standard
+        kp_and_res = list(quantized_outputs.items())
+        for i in range(1, len(kp_and_res)):
+            kp, res = kp_and_res[i]
+            print(
+                "granularity:",
+                granularity,
+                " sizes:",
+                sizes,
+                " kp:",
+                kp,
+                compute_error(res, kp_and_res[0][1]),
+            )
+            self.assertTrue(
+                compute_error(res, kp_and_res[0][1]) > 28,
+                f"mismatch between {kp=} and {kp_and_res[0]=}, {sizes=}, {granularity=}",
+            )
+
     @common_utils.parametrize("granularity", [PerTensor(), PerRow()])
     def test_slice_preserves_aliasing(self, granularity):
         config = Float8DynamicActivationFloat8WeightConfig(granularity=granularity)

torchao/quantization/quant_api.py

@@ -1640,7 +1640,7 @@ class Float8DynamicActivationFloat8WeightConfig(AOBaseConfig):
     weight_dtype: torch.dtype = e4m3_dtype
     granularity: Optional[Union[FP8Granularity, List[FP8Granularity]]] = None
     mm_config: Optional[Float8MMConfig] = None
-    activation_value_lb: Optional[float] = None
+    activation_value_lb: Optional[float] = 1e-12
     activation_value_ub: Optional[float] = None
     kernel_preference: KernelPreference = KernelPreference.AUTO
     set_inductor_config: bool = True

torchao/quantization/quant_primitives.py

@@ -2185,8 +2185,9 @@ def _choose_scale_float8(
     block_size: List[int],
     float8_dtype: torch.dtype = torch.float8_e4m3fn,
     scale_dtype: torch.dtype = torch.float32,
-    hp_value_lb: Optional[float] = None,
+    hp_value_lb: Optional[float] = 1e-12,
     hp_value_ub: Optional[float] = None,
+    reverse: bool = False,
 ) -> torch.Tensor:
     """
     Calculates float8 scaling factor for the given high precision tensor, using tensorwise granularity.
@@ -2214,7 +2215,11 @@ def _choose_scale_float8(
         max_abs = tensor_reshaped.abs().amax(dim=reduction_dims, keepdim=True)
         if hp_value_lb is not None or hp_value_ub is not None:
             max_abs = torch.clamp(max_abs, min=hp_value_lb, max=hp_value_ub)
-        scale = max_abs / quant_max
+        if reverse:
+            scale = quant_max / max_abs.to(torch.float32)
+            scale[scale == float("inf")] = 1.0
+        else:
+            scale = max_abs / quant_max
         # Reshape scale back to match the expected output shape
         # The scale tensor should have the same shape as the input divided by block_size
         output_shape = [
@@ -2284,6 +2289,7 @@ def _quantize_affine_float8(
     tensor: torch.Tensor,
     scale: torch.Tensor,
     float8_dtype: torch.dtype = torch.float8_e4m3fn,
+    reverse: bool = False,
 ) -> torch.Tensor:
     """
     Quantizes the high precision floating point tensor to a float8 tensor, using the given scaling factor.
@@ -2293,7 +2299,10 @@ def _quantize_affine_float8(
     # Expand scale to match tensor dimensions for block-wise quantization
     scale_expanded = _expand_scale_to_tensor_shape(scale, tensor.shape)
 
-    tensor_scaled = tensor_fp32 / scale_expanded
+    if reverse:
+        tensor_scaled = tensor_fp32 * scale_expanded
+    else:
+        tensor_scaled = tensor_fp32 / scale_expanded
     max_value = torch.finfo(float8_dtype).max
     tensor_clamped = tensor_scaled.clamp(min=-max_value, max=max_value)
     fp8_tensor = tensor_clamped.to(float8_dtype)

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

@@ -169,7 +169,7 @@ def to_float8(
         float8_dtype: torch.dtype = torch.float8_e4m3fn,
         granularity: FP8Granularity = PerRow(),
         mm_config: Optional[Float8MMConfig] = None,
-        hp_value_lb: Optional[float] = None,
+        hp_value_lb: Optional[float] = 1e-12,
         hp_value_ub: Optional[float] = None,
         kernel_preference: KernelPreference = KernelPreference.AUTO,
         act_quant_kwargs: Optional[QuantizeTensorToFloat8Kwargs] = None,
@@ -209,8 +209,10 @@ def to_float8(
                 block_size=block_size,
                 hp_value_lb=hp_value_lb,
                 hp_value_ub=hp_value_ub,
+                reverse=True,
             )
-            data = _quantize_affine_float8(hp_tensor, scale, float8_dtype)
+            data = _quantize_affine_float8(hp_tensor, scale, float8_dtype, reverse=True)
+            scale = 1.0 / scale
 
         hp_dtype = hp_tensor.dtype
         return Float8Tensor(