Add float8 FakeQuantizeConfig and FakeQuantizer

**Summary:** This commit adds a QAT path for float8, using the
same primitives as `torchao.quantization.Float8Tensor` targeting
PTQ configs like:
- `Float8DynamicActivationFloat8WeightConfig`
- `Float8DynamicActivationInt4WeightConfig`
- `Float8WeightOnlyConfig`

Usage:
```
from torchao.quantization.granularity import PerRow
from torchao.quantization.qat import quantize_, QATConfig

base_config = Float8DynamicActivationFloat8WeightConfig(
    torch.float8_e4m3fn, PerRow(),
)
quantize_(model, QATConfig(base_config, step="prepare"))
quantize_(model, QATConfig(base_config, step="convert"))
```

OR

```
from torchao.quantization.granularity import PerRow
from torchao.quantization.qat import (
    Float8FakeQuantizeConfig,
    QATConfig,
    quantize_,
)

dtype = torch.float8_e4m3fn
granularity = PerRow()
quantize_(model, QATConfig(
    activation_config=Float8FakeQuantizeConfig(dtype, granularity),
    weight_config=Float8FakeQuantizeConfig(dtype, granularity),
    step="prepare",
)

# convert (same as above, not shown)
```

**Test Plan:**
```
python test/quantization/test_qat.py -k test_float8_fake_quantize_config
python test/quantization/test_qat.py -k test_float8_fake_quantize
python test/quantization/test_qat.py -k test_quantize_api_fp8_fp8
python test/quantization/test_qat.py -k test_quantize_api_fp8_int4
```

Author: andrewor14

docs/source/api_ref_qat.rst

@@ -26,10 +26,12 @@ Custom QAT APIs
 
     FakeQuantizeConfigBase
     IntxFakeQuantizeConfig
+    Float8FakeQuantizeConfig
     FakeQuantizedLinear
     FakeQuantizedEmbedding
     FakeQuantizerBase
     IntxFakeQuantizer
+    Float8FakeQuantizer
     linear.enable_linear_fake_quant
     linear.disable_linear_fake_quant
 

test/quantization/test_qat.py

@@ -14,17 +14,21 @@
 
 import torch
 import torch.nn.functional as F
-from parameterized import parameterized
 from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib  # noqa: F401
+from torch.testing._internal.common_utils import (
+    TestCase,
+    instantiate_parametrized_tests,
+    parametrize,
+)
 
 from torchao import quantize_
-from torchao.float8.config import ScalingGranularity
-from torchao.float8.float8_scaling_utils import hp_tensor_to_float8_dynamic
-from torchao.float8.float8_training_tensor import LinearMMConfig
+from torchao.quantization import Float8Tensor
 from torchao.quantization.granularity import (
+    Granularity,
     PerAxis,
     PerGroup,
     PerRow,
+    PerTensor,
     PerToken,
 )
 from torchao.quantization.linear_quant_modules import (
@@ -43,11 +47,12 @@
     FakeQuantizedEmbedding,
 )
 from torchao.quantization.qat.fake_quantize_config import (
+    Float8FakeQuantizeConfig,
     IntxFakeQuantizeConfig,
 )
 from torchao.quantization.qat.fake_quantizer import (
+    Float8FakeQuantizer,
     IntxFakeQuantizer,
-    _Float8RowwiseActivationFakeQuantizer,
 )
 from torchao.quantization.qat.linear import (
     FakeQuantizedLinear,
@@ -58,10 +63,11 @@
 from torchao.quantization.qat.utils import (
     _fake_quantize_per_channel_group,
     _fake_quantize_per_token,
-    _Float8RowwiseFakeQuantize,
     _get_qmin_qmax,
 )
 from torchao.quantization.quant_api import (
+    Float8DynamicActivationFloat8WeightConfig,
+    Float8DynamicActivationInt4WeightConfig,
     Int8DynamicActivationInt4WeightConfig,
 )
 from torchao.quantization.quant_primitives import (
@@ -83,6 +89,10 @@
     get_groupwise_affine_qparams,
     groupwise_affine_quantize_tensor,
 )
+from torchao.utils import (
+    _is_fbgemm_genai_gpu_available,
+    is_sm_at_least_89,
+)
 
 # TODO: put this in a common test utils file
 _CUDA_IS_AVAILABLE = torch.cuda.is_available()
@@ -193,7 +203,7 @@ def forward(self, x):
         return x
 
 
-class TestQAT(unittest.TestCase):
+class TestQAT(TestCase):
     SEED = 123
 
     def test_fake_quantize_per_channel_group(self):
@@ -1420,7 +1430,7 @@ def test_qat_linear_bias(self):
         example_inputs = m.example_inputs()
         m(*example_inputs)
 
-    @parameterized.expand([(torch.float32,), (torch.bfloat16,), (torch.float16,)])
+    @parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
     def test_fake_quantize_per_token_vs_convert(self, dtype: torch.dtype):
         """
         Test that the following produce the exact same numerics:
@@ -1437,7 +1447,7 @@ def test_fake_quantize_per_token_vs_convert(self, dtype: torch.dtype):
         baseline_out = per_token_dynamic_quant(x)
         torch.testing.assert_close(fake_quantizer_out, baseline_out, atol=0, rtol=0)
 
-    @parameterized.expand([(torch.float32,), (torch.bfloat16,), (torch.float16,)])
+    @parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
     def test_qat_8da4w_prepare_vs_convert(self, dtype: torch.dtype):
         """
         Test that the prepare and convert steps of Int8DynActInt4QATQuantizer produces
@@ -1548,7 +1558,7 @@ def test_qat_8da4w_eps(self):
         actual_out = converted_model.linear1(x)
         torch.testing.assert_close(expected_out, actual_out, atol=0, rtol=0)
 
-    @parameterized.expand([(True,), (False,)])
+    @parametrize("is_symmetric", [True, False])
     def test_fake_quantizer_range_learning(self, is_symmetric):
         """
         Test that range learning requires `IntxFakeQuantizer`s to be initialized correctly.
@@ -1589,7 +1599,7 @@ def test_fake_quantizer_range_learning(self, is_symmetric):
             self.assertTrue(fake_quantizer.zero_point.requires_grad)
         fake_quantizer(*example_inputs)
 
-    @parameterized.expand([(True,), (False,)])
+    @parametrize("is_symmetric", [True, False])
     def test_qat_range_learning(self, is_symmetric):
         """
         Test end-to-end QAT flow with range learning.
@@ -1664,24 +1674,6 @@ def test_qat_range_learning(self, is_symmetric):
             self.assertNotEqual(torch.count_nonzero(new_weight.grad), 0)
             self.assertFalse(torch.equal(new_weight, prev_weight))
 
-    def test_float8_rowwise_fake_quantize(self):
-        """
-        Test that `_Float8RowwiseFakeQuantize` is numerically close to `Float8TrainingTensor`.
-        """
-        torch.manual_seed(self.SEED)
-        dtype = torch.float8_e4m3fn
-        x = torch.randn(32, 64)
-        axiswise_dim = 0
-        out = _Float8RowwiseFakeQuantize.apply(x, dtype, axiswise_dim)
-        out_expected = hp_tensor_to_float8_dynamic(
-            x,
-            dtype,
-            LinearMMConfig(),
-            scaling_granularity=ScalingGranularity.AXISWISE,
-            axiswise_dim=axiswise_dim,
-        ).to_original_precision()
-        torch.testing.assert_close(out, out_expected, atol=0, rtol=0)
-
     def test_qat_fp8a4w_quantizer(self):
         """
         Test basic model training with `Float8ActInt4WeightQATQuantizer`.
@@ -1693,7 +1685,8 @@ def test_qat_fp8a4w_quantizer(self):
         for linear in [m.linear1, m.sub.linear, m.linear2]:
             self.assertIsInstance(linear, FakeQuantizedLinear)
             self.assertIsInstance(
-                linear.activation_fake_quantizer, _Float8RowwiseActivationFakeQuantizer
+                linear.activation_fake_quantizer,
+                Float8FakeQuantizer,
             )
             self.assertIsInstance(linear.weight_fake_quantizer, IntxFakeQuantizer)
         prev_weight = copy.deepcopy(m.linear1.weight)
@@ -1833,6 +1826,82 @@ def test_qat_api_convert_no_quantization(self):
         baseline_out = baseline_model(*x2)
         torch.testing.assert_close(out, baseline_out, atol=0, rtol=0)
 
+    def test_float8_fake_quantize_config(self):
+        """
+        Test that the correct errors are thrown if `Float8FakeQuantizeConfig` is not instantiated properly.
+        """
+        # OK
+        Float8FakeQuantizeConfig(torch.float8_e4m3fn)
+        Float8FakeQuantizeConfig(torch.float8_e4m3fn, PerRow())
+        Float8FakeQuantizeConfig(torch.float8_e4m3fn, PerTensor())
+
+        with self.assertRaisesRegex(ValueError, "not a float8 dtype"):
+            Float8FakeQuantizeConfig(torch.int8)
+        with self.assertRaisesRegex(
+            ValueError, "Please specify the granularity object instead of the class"
+        ):
+            Float8FakeQuantizeConfig(granularity=PerRow)
+        with self.assertRaisesRegex(
+            ValueError, "Expected PerRow or PerTensor granularity"
+        ):
+            Float8FakeQuantizeConfig(granularity=PerToken())
+
+    @parametrize("granularity", [PerTensor(), PerRow()])
+    def test_float8_fake_quantize(self, granularity: Granularity):
+        """
+        Test that `Float8FakeQuantizer` is numerically close to `Float8Tensor`.
+        """
+        dtype = torch.float8_e4m3fn
+        fq_config = Float8FakeQuantizeConfig(dtype, granularity)
+        fake_quantizer = Float8FakeQuantizer(fq_config)
+        torch.manual_seed(self.SEED)
+        x = torch.randn(32, 64)
+        out = fake_quantizer(x)
+        out_expected = Float8Tensor.to_float8(x, dtype, granularity).dequantize()
+        sqnr = compute_error(out, out_expected)
+        self.assertGreater(sqnr, 16)
+
+    @parametrize("granularity", [PerTensor(), PerRow()])
+    @unittest.skipIf(not _CUDA_IS_AVAILABLE, "skipping when cuda is not available")
+    @unittest.skipIf(not is_sm_at_least_89(), "Need sm89+")
+    def test_quantize_api_fp8_fp8(self, granularity: Granularity):
+        """
+        Test the following:
+            quantize_(model, QATConfig(Float8DynamicActivationFloat8Weight(), step="prepare"))
+            quantize_(model, QATConfig(Float8DynamicActivationFloat8Weight(), step="convert"))
+        """
+        torch.manual_seed(self.SEED)
+        m = M().to(torch.bfloat16).cuda()
+        example_inputs = (m.example_inputs()[0].to(torch.bfloat16).cuda(),)
+        base_config = Float8DynamicActivationFloat8WeightConfig(granularity=granularity)
+        quantize_(m, QATConfig(base_config, step="prepare"))
+        m(*example_inputs)
+        quantize_(m, QATConfig(base_config, step="convert"))
+        m(*example_inputs)
+
+    @unittest.skipIf(not _CUDA_IS_AVAILABLE, "skipping when cuda is not available")
+    @unittest.skipIf(not is_sm_at_least_89(), "Need sm89+")
+    @unittest.skipIf(
+        not _is_fbgemm_genai_gpu_available(), "Requires fbgemm-gpu-genai >= 1.2.0"
+    )
+    def test_quantize_api_fp8_int4(self):
+        """
+        Test the following:
+            quantize_(model, QATConfig(Float8DynamicActivationInt4WeightConfig(), step="prepare"))
+            quantize_(model, QATConfig(Float8DynamicActivationInt4WeightConfig(), step="convert"))
+        """
+        torch.manual_seed(self.SEED)
+        m = M().to(torch.bfloat16).cuda()
+        example_inputs = (m.example_inputs()[0].to(torch.bfloat16).cuda(),)
+        base_config = Float8DynamicActivationInt4WeightConfig(group_size=128)
+        quantize_(m, QATConfig(base_config, step="prepare"))
+        m(*example_inputs)
+        quantize_(m, QATConfig(base_config, step="convert"))
+        m(*example_inputs)
+
+
+instantiate_parametrized_tests(TestQAT)
+
 
 if __name__ == "__main__":
     unittest.main()

torchao/quantization/qat/__init__.py

@@ -15,11 +15,13 @@
 from .fake_quantize_config import (
     FakeQuantizeConfig,
     FakeQuantizeConfigBase,
+    Float8FakeQuantizeConfig,
     IntxFakeQuantizeConfig,
 )
 from .fake_quantizer import (
     FakeQuantizer,
     FakeQuantizerBase,
+    Float8FakeQuantizer,
     IntxFakeQuantizer,
 )
 from .linear import (
@@ -34,6 +36,8 @@
     "QATStep",
     "FakeQuantizeConfigBase",
     "FakeQuantizerBase",
+    "Float8FakeQuantizeConfig",
+    "Float8FakeQuantizer",
     "IntxFakeQuantizeConfig",
     "IntxFakeQuantizer",
     "FakeQuantizedLinear",

torchao/quantization/qat/fake_quantize_config.py

@@ -11,10 +11,17 @@
 import torch
 
 from torchao.core.config import AOBaseConfig
+from torchao.float8.config import e4m3_dtype
+from torchao.float8.inference import (
+    FP8Granularity,
+    _normalize_granularity,
+)
 from torchao.quantization.granularity import (
     Granularity,
     PerAxis,
     PerGroup,
+    PerRow,
+    PerTensor,
     PerToken,
 )
 from torchao.quantization.quant_primitives import (
@@ -24,6 +31,7 @@
     TorchAODType,
     ZeroPointDomain,
 )
+from torchao.utils import _is_float8_type
 
 from .utils import _log_deprecation_warning
 
@@ -36,6 +44,39 @@ class FakeQuantizeConfigBase(abc.ABC):
     pass
 
 
+@dataclass
+class Float8FakeQuantizeConfig(FakeQuantizeConfigBase):
+    """
+    Config for float8 fake quantization, targeting :class:`~torchao.quantization.Float8Tensor`.
+
+    Args:
+       dtype (torch.dtype): the dtype for float8 Tensor
+       granularity (FP8Granularity): the granularity for the Tensor, currently either PerRow() or PerTensor()
+       hp_value_lb (Optional[float]): the lower bound for high precision floating point value for calculating scale
+       hp_value_ub (Optional[float]): the upper bound for high precision floating point value for calculating scale
+    """
+
+    dtype: torch.dtype = e4m3_dtype
+    granularity: FP8Granularity = PerRow()
+    hp_value_lb: Optional[float] = None
+    hp_value_ub: Optional[float] = None
+
+    def __post_init__(self):
+        """
+        Verify dtype and granularity are the ones we support.
+        """
+        if not _is_float8_type(self.dtype):
+            raise ValueError(f"{self.dtype} is not a float8 dtype")
+        if isinstance(self.granularity, type):
+            raise ValueError(
+                "Please specify the granularity object instead of the class, e.g. PerRow() instead of PerRow"
+            )
+        if type(self.granularity) not in [PerRow, PerTensor]:
+            raise ValueError(
+                f"Expected PerRow or PerTensor granularity, got {self.granularity}"
+            )
+
+
 @dataclass
 class IntxFakeQuantizeConfig(FakeQuantizeConfigBase):
     """
@@ -279,6 +320,7 @@ def __post_init__(self):
         _log_deprecation_warning(self)
 
 
+# TODO: rewrite using registration API?
 def _infer_fake_quantize_configs(
     base_config: AOBaseConfig,
 ) -> Tuple[Optional[FakeQuantizeConfigBase], Optional[FakeQuantizeConfigBase]]:
@@ -291,6 +333,8 @@ def _infer_fake_quantize_configs(
     """
     # avoid circular imports
     from torchao.quantization import (
+        Float8DynamicActivationFloat8WeightConfig,
+        Float8DynamicActivationInt4WeightConfig,
         Int4WeightOnlyConfig,
         Int8DynamicActivationInt4WeightConfig,
     )
@@ -302,18 +346,45 @@ def _infer_fake_quantize_configs(
             is_symmetric=base_config.act_mapping_type == MappingType.SYMMETRIC,
         )
         weight_config = IntxFakeQuantizeConfig(
-            dtype=TorchAODType.INT4,
+            dtype=torch.int4,
             group_size=base_config.group_size,
             is_symmetric=base_config.mapping_type == MappingType.SYMMETRIC,
         )
-        return (act_config, weight_config)
     elif isinstance(base_config, Int4WeightOnlyConfig):
+        if base_config.version != 2:
+            raise ValueError(f"Only version 2 of {type(base_config)} is supported")
+        act_config = None
+        weight_config = IntxFakeQuantizeConfig(
+            dtype=torch.int4,
+            group_size=base_config.group_size,
+            is_symmetric=True,
+        )
+    elif isinstance(base_config, Float8DynamicActivationFloat8WeightConfig):
+        if base_config.version != 2:
+            raise ValueError(f"Only version 2 of {type(base_config)} is supported")
+        (act_granularity, weight_granularity) = _normalize_granularity(
+            base_config.granularity
+        )
+        act_config = Float8FakeQuantizeConfig(
+            dtype=base_config.activation_dtype,
+            granularity=act_granularity,
+            hp_value_lb=base_config.activation_value_lb,
+            hp_value_ub=base_config.activation_value_ub,
+        )
+        weight_config = Float8FakeQuantizeConfig(
+            dtype=base_config.weight_dtype,
+            granularity=weight_granularity,
+        )
+    elif isinstance(base_config, Float8DynamicActivationInt4WeightConfig):
+        act_config = Float8FakeQuantizeConfig(
+            dtype=torch.float8_e4m3fn,
+            granularity=PerRow(),
+        )
         weight_config = IntxFakeQuantizeConfig(
-            dtype=torch.uint4,
+            dtype=torch.int4,
             group_size=base_config.group_size,
-            is_symmetric=False,
-            zero_point_domain=base_config.zero_point_domain,
+            is_symmetric=True,
         )
-        return (None, weight_config)
     else:
         raise ValueError("Unexpected base config: %s" % base_config)
+    return (act_config, weight_config)