[bc-breaking] Generalize FakeQuantizeConfig beyond intx

 **Summary:** The existing `FakeQuantizeConfig` performs only
intx quantization, but we plan to extend QAT to other dtypes
such as fp8 and nvfp4 in the near future. This is the necessary
refactor before that. Specifically:

```
# New abstract class
FakeQuantizeConfigBase
# Rename
FakeQuantizeConfig -> IntxFakeQuantizeConfig
```

In the future, we will have other types of `FakeQuantizeConfigBase`
for float dtypes that users can pass in instead of the existing
Intx one.

**BC-breaking notes:** For BC, we keep around the old names to
reference the new ones. However, this commit is still BC-breaking
in the sense that a few APIs now accept the abstract
`FakeQuantizeConfigBase` instead. For the most part, this abstract
class will be hidden from the user.

Before:
```
activation_config = FakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False)
weight_config = FakeQuantizeConfig(torch.int4, group_size=32)
```

After:
```
activation_config = IntxFakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False)
weight_config = IntxFakeQuantizeConfig(torch.int4, group_size=32)
```

**Test Plan:**
python test/quantization/test_qat.py

ghstack-source-id: 056170d
Pull Request resolved: #2628

Author: andrewor14

README.md

@@ -180,9 +180,9 @@ Post-training quantization can result in a fast and compact model, but may also
 
 ```python
 from torchao.quantization import quantize_
-from torchao.quantization.qat import FakeQuantizeConfig, IntXQuantizationAwareTrainingConfig
-activation_config = FakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False)
-weight_config = FakeQuantizeConfig(torch.int4, group_size=32)
+from torchao.quantization.qat import IntxFakeQuantizeConfig, IntXQuantizationAwareTrainingConfig
+activation_config = IntxFakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False)
+weight_config = IntxFakeQuantizeConfig(torch.int4, group_size=32)
 qat_config = IntXQuantizationAwareTrainingConfig(activation_config, weight_config),
 quantize_(my_model, qat_config)
 ```

docs/source/api_ref_qat.rst

@@ -24,7 +24,7 @@ Custom QAT APIs
     :toctree: generated/
     :nosignatures:
 
-    FakeQuantizeConfig
+    IntxFakeQuantizeConfig
     FakeQuantizedLinear
     FakeQuantizedEmbedding
     FakeQuantizer

test/prototype/test_parq.py

@@ -30,8 +30,8 @@
 from torchao.prototype.parq.quant.uniform_torchao import _BIT_WIDTH_TO_DTYPE
 from torchao.quantization.granularity import PerGroup
 from torchao.quantization.qat import (
-    FakeQuantizeConfig,
     FromIntXQuantizationAwareTrainingConfig,
+    IntxFakeQuantizeConfig,
     IntXQuantizationAwareTrainingConfig,
 )
 from torchao.quantization.quant_api import (
@@ -393,7 +393,7 @@ def test_int8_dynamic_activation_intx_e2e(
         optimizer.step()
 
         # apply torchao quantized activations on top
-        activation_config = FakeQuantizeConfig(
+        activation_config = IntxFakeQuantizeConfig(
             torch.int8,
             granularity="per_token",
             mapping_type=config.act_mapping_type,

test/quantization/test_qat.py

@@ -21,9 +21,9 @@
 )
 from torchao.quantization.granularity import PerAxis, PerGroup
 from torchao.quantization.qat import (
-    FakeQuantizeConfig,
     FromIntXQuantizationAwareTrainingConfig,
     Int4WeightOnlyEmbeddingQATQuantizer,
+    IntxFakeQuantizeConfig,
     IntXQuantizationAwareTrainingConfig,
 )
 from torchao.quantization.quant_api import (
@@ -282,7 +282,7 @@ def test_identical_to_IntXQuantizationAwareTrainingConfig(
         )
 
         embedding_filter = lambda m, fqn: isinstance(m, torch.nn.Embedding)
-        weight_config = FakeQuantizeConfig(
+        weight_config = IntxFakeQuantizeConfig(
             weight_dtype,
             group_size=group_size,
             is_symmetric=is_symmetric,

torchao/experimental/tests/test_embedding_xbit_quantizer.py

@@ -16,9 +16,9 @@
 from torchao.dtypes import PackedLinearInt8DynamicActivationIntxWeightLayout, QDQLayout
 from torchao.quantization.granularity import PerAxis, PerGroup
 from torchao.quantization.qat import (
-    FakeQuantizeConfig,
     FromIntXQuantizationAwareTrainingConfig,
     Int8DynActInt4WeightQATQuantizer,
+    IntxFakeQuantizeConfig,
     IntXQuantizationAwareTrainingConfig,
 )
 from torchao.quantization.quant_api import (
@@ -538,12 +538,12 @@ def test_identical_to_IntXQuantizationAwareTrainingConfig(
         model = model.to(model_dtype)
         activations = activations.to(model_dtype)
 
-        activation_config = FakeQuantizeConfig(
+        activation_config = IntxFakeQuantizeConfig(
             torch.int8,
             "per_token",
             is_symmetric=is_act_symmetric,
         )
-        weight_config = FakeQuantizeConfig(
+        weight_config = IntxFakeQuantizeConfig(
             weight_dtype,
             group_size=group_size,
             is_symmetric=is_symmetric,

torchao/experimental/tests/test_int8_dynamic_activation_intx_weight.py

@@ -71,7 +71,7 @@ def train_loop(m: torch.nn.Module):
 
 The recommended way to run QAT in torchao is through the `quantize_` API:
 1. **Prepare:** specify how weights and/or activations are to be quantized through
-[`FakeQuantizeConfig`](https://docs.pytorch.org/ao/main/generated/torchao.quantization.qat.FakeQuantizeConfig.html#torchao.quantization.qat.FakeQuantizeConfig) and passing these to [`IntXQuantizationAwareTrainingConfig`](https://docs.pytorch.org/ao/main/generated/torchao.quantization.qat.IntXQuantizationAwareTrainingConfig.html#torchao.quantization.qat.IntXQuantizationAwareTrainingConfig)
+[`IntxFakeQuantizeConfig`](https://docs.pytorch.org/ao/main/generated/torchao.quantization.qat.IntxFakeQuantizeConfig.html#torchao.quantization.qat.IntxFakeQuantizeConfig) and passing these to [`IntXQuantizationAwareTrainingConfig`](https://docs.pytorch.org/ao/main/generated/torchao.quantization.qat.IntXQuantizationAwareTrainingConfig.html#torchao.quantization.qat.IntXQuantizationAwareTrainingConfig)
 2. **Convert:** quantize the model using the standard post-training quantization (PTQ)
 functions such as [`Int8DynamicActivationInt4WeightConfig`](https://docs.pytorch.org/ao/main/generated/torchao.quantization.Int8DynamicActivationInt4WeightConfig.html#torchao.quantization.Int8DynamicActivationInt4WeightConfig)
 
@@ -84,16 +84,16 @@ from torchao.quantization import (
     Int8DynamicActivationInt4WeightConfig,
 )
 from torchao.quantization.qat import (
-    FakeQuantizeConfig,
+    IntxFakeQuantizeConfig,
     FromIntXQuantizationAwareTrainingConfig,
     IntXQuantizationAwareTrainingConfig,
 )
 model = get_model()
 
 # prepare: insert fake quantization ops
 # swaps `torch.nn.Linear` with `FakeQuantizedLinear`
-activation_config = FakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False)
-weight_config = FakeQuantizeConfig(torch.int4, group_size=32)
+activation_config = IntxFakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False)
+weight_config = IntxFakeQuantizeConfig(torch.int4, group_size=32)
 quantize_(
     model,
     IntXQuantizationAwareTrainingConfig(activation_config, weight_config),
@@ -116,8 +116,8 @@ the following with a filter function during the prepare step:
 
 ```
 # first apply linear transformation to the model as above
-activation_config = FakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False)
-weight_config = FakeQuantizeConfig(torch.int4, group_size=32)
+activation_config = IntxFakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False)
+weight_config = IntxFakeQuantizeConfig(torch.int4, group_size=32)
 quantize_(
     model,
     IntXQuantizationAwareTrainingConfig(activation_config, weight_config),

torchao/quantization/qat/README.md

@@ -1,6 +1,5 @@
 from .api import (
     ComposableQATQuantizer,
-    FakeQuantizeConfig,
     FromIntXQuantizationAwareTrainingConfig,
     IntXQuantizationAwareTrainingConfig,
     from_intx_quantization_aware_training,
@@ -11,6 +10,11 @@
     FakeQuantizedEmbedding,
     Int4WeightOnlyEmbeddingQATQuantizer,
 )
+from .fake_quantize_config import (
+    FakeQuantizeConfig,
+    FakeQuantizeConfigBase,
+    IntxFakeQuantizeConfig,
+)
 from .fake_quantizer import FakeQuantizer
 from .linear import (
     FakeQuantizedLinear,
@@ -21,7 +25,7 @@
 
 __all__ = [
     "ComposableQATQuantizer",
-    "FakeQuantizeConfig",
+    "FakeQuantizeConfigBase",
     "FakeQuantizedLinear",
     "FakeQuantizedEmbedding",
     "FakeQuantizer",
@@ -30,8 +34,11 @@
     "Int4WeightOnlyEmbeddingQATQuantizer",
     "Int4WeightOnlyQATQuantizer",
     "Int8DynActInt4WeightQATQuantizer",
+    "IntxFakeQuantizeConfig",
     "IntXQuantizationAwareTrainingConfig",
     "initialize_fake_quantizers",
-    "intx_quantization_aware_training",
+    # for BC
+    "FakeQuantizeConfig",
     "from_intx_quantization_aware_training",
+    "intx_quantization_aware_training",
 ]