[bc-breaking] Generalize QAT configs beyond intx quantization

README.md

@@ -180,10 +180,10 @@ 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)
-qat_config = IntXQuantizationAwareTrainingConfig(activation_config, weight_config),
+from torchao.quantization.qat import IntxFakeQuantizeConfig, QuantizationAwareTrainingConfig
+activation_config = IntxFakeQuantizeConfig(torch.int8, "per_token", is_symmetric=False)
+weight_config = IntxFakeQuantizeConfig(torch.int4, group_size=32)
+qat_config = QuantizationAwareTrainingConfig(activation_config, weight_config),
 quantize_(my_model, qat_config)
 ```
 

docs/source/api_ref_qat.rst

@@ -15,16 +15,16 @@ please refer to the `QAT README <https://github.com/pytorch/ao/blob/main/torchao
     :toctree: generated/
     :nosignatures:
 
-    IntXQuantizationAwareTrainingConfig
-    FromIntXQuantizationAwareTrainingConfig
+    QuantizationAwareTrainingConfig
+    FromQuantizationAwareTrainingConfig
 
 Custom QAT APIs
 ---------------
 .. autosummary::
     :toctree: generated/
     :nosignatures:
 
-    FakeQuantizeConfig
+    IntxFakeQuantizeConfig
     FakeQuantizedLinear
     FakeQuantizedEmbedding
     FakeQuantizer

test/prototype/test_parq.py

@@ -30,9 +30,9 @@
 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,
-    IntXQuantizationAwareTrainingConfig,
+    FromQuantizationAwareTrainingConfig,
+    IntxFakeQuantizeConfig,
+    QuantizationAwareTrainingConfig,
 )
 from torchao.quantization.quant_api import (
     Int8DynamicActivationIntxWeightConfig,
@@ -393,15 +393,15 @@ 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,
         )
         filter_fn = optimizer.get_filter_fn(model)
         quantize_(
             model,
-            IntXQuantizationAwareTrainingConfig(activation_config=activation_config),
+            QuantizationAwareTrainingConfig(activation_config=activation_config),
             filter_fn=filter_fn,
         )
         out = model(x)
@@ -410,7 +410,7 @@ def test_int8_dynamic_activation_intx_e2e(
         # equivalent to torchao's convert step
         model.eval()
         optimizer.restore_latent_params()
-        quantize_(model, FromIntXQuantizationAwareTrainingConfig(), filter_fn=filter_fn)
+        quantize_(model, FromQuantizationAwareTrainingConfig(), filter_fn=filter_fn)
         quantize_(model, config, filter_fn=filter_fn)
         converted_out = model(x)
         torch.testing.assert_close(converted_out, ref_out, atol=0, rtol=0)