Rename AOPerModuleConfig to ModuleFqnToConfig (#2243)

* Rename AOPerModuleConfig to ModuleFqnToConfig

Summary:
to be more explicit on what this config means

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:

* renaming rest

Author: jerryzh168

test/quantization/test_config_serialization.py

@@ -20,7 +20,6 @@
     config_to_dict,
 )
 from torchao.quantization.quant_api import (
-    AOPerModuleConfig,
     Float8DynamicActivationFloat8WeightConfig,
     Float8WeightOnlyConfig,
     FPXWeightOnlyConfig,
@@ -30,6 +29,7 @@
     Int8DynamicActivationInt4WeightConfig,
     Int8DynamicActivationInt8WeightConfig,
     Int8WeightOnlyConfig,
+    ModuleFqnToConfig,
     PerRow,
     UIntXWeightOnlyConfig,
 )
@@ -68,9 +68,9 @@
     # Sparsity configs
     SemiSparseWeightConfig(),
     BlockSparseWeightConfig(blocksize=128),
-    AOPerModuleConfig({}),
-    AOPerModuleConfig({"_default": Int4WeightOnlyConfig(), "linear1": None}),
-    AOPerModuleConfig(
+    ModuleFqnToConfig({}),
+    ModuleFqnToConfig({"_default": Int4WeightOnlyConfig(), "linear1": None}),
+    ModuleFqnToConfig(
         {
             "linear1": Int4WeightOnlyConfig(),
             "linear2": Int8DynamicActivationInt4WeightConfig(),

test/quantization/test_quant_api.py

@@ -38,11 +38,11 @@
     PerGroup,
 )
 from torchao.quantization.quant_api import (
-    AOPerModuleConfig,
     Int4WeightOnlyConfig,
     Int8DynamicActivationInt4WeightConfig,
     Int8WeightOnlyConfig,
     IntxWeightOnlyConfig,
+    ModuleFqnToConfig,
     Quantizer,
     TwoStepQuantizer,
     _replace_with_custom_fn_if_matches_filter,
@@ -946,10 +946,10 @@ def test_workflow_e2e_numerics(self, config):
         assert sqnr >= 16.5, f"SQNR {sqnr} is too low"
 
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
-    def test_ao_per_module_config_default(self):
+    def test_module_fqn_to_config_default(self):
         config1 = Int4WeightOnlyConfig(group_size=32)
         config2 = Int8WeightOnlyConfig()
-        config = AOPerModuleConfig({"_default": config1, "linear2": config2})
+        config = ModuleFqnToConfig({"_default": config1, "linear2": config2})
         model = ToyLinearModel().cuda().to(dtype=torch.bfloat16)
         example_inputs = model.example_inputs(device="cuda", dtype=torch.bfloat16)
         quantize_(model, config)
@@ -960,10 +960,10 @@ def test_ao_per_module_config_default(self):
         assert isinstance(model.linear2.weight._layout, PlainLayout)
 
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
-    def test_ao_per_module_config_module_name(self):
+    def test_module_fqn_to_config_module_name(self):
         config1 = Int4WeightOnlyConfig(group_size=32)
         config2 = Int8WeightOnlyConfig()
-        config = AOPerModuleConfig({"linear1": config1, "linear2": config2})
+        config = ModuleFqnToConfig({"linear1": config1, "linear2": config2})
         model = ToyLinearModel().cuda().to(dtype=torch.bfloat16)
         example_inputs = model.example_inputs(device="cuda", dtype=torch.bfloat16)
         quantize_(model, config)
@@ -974,7 +974,7 @@ def test_ao_per_module_config_module_name(self):
         assert isinstance(model.linear2.weight._layout, PlainLayout)
 
     @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_6, "Need torch 2.6+")
-    def test_ao_per_module_config_embedding_linear(self):
+    def test_module_fqn_to_config_embedding_linear(self):
         weight_dtype = torch.int8
         granularity = PerGroup(8)
         mapping_type = MappingType.SYMMETRIC
@@ -987,7 +987,7 @@ def test_ao_per_module_config_embedding_linear(self):
         # example model linear is Linear(16, 8)
         linear_config = Int8DynamicActivationInt4WeightConfig(group_size=16)
 
-        config = AOPerModuleConfig({"emb": embedding_config, "linear": linear_config})
+        config = ModuleFqnToConfig({"emb": embedding_config, "linear": linear_config})
         indices = torch.randint(0, 10, (32,))
         indices = indices.unsqueeze(0)
         example_inputs = (indices,)
@@ -1006,9 +1006,9 @@ def test_ao_per_module_config_embedding_linear(self):
         assert isinstance(model.linear.weight, LinearActivationQuantizedTensor)
 
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
-    def test_ao_per_module_config_skip(self):
+    def test_module_fqn_to_config_skip(self):
         config1 = Int4WeightOnlyConfig(group_size=32)
-        config = AOPerModuleConfig({"_default": config1, "linear2": None})
+        config = ModuleFqnToConfig({"_default": config1, "linear2": None})
         model = ToyLinearModel().cuda().to(dtype=torch.bfloat16)
         example_inputs = model.example_inputs(device="cuda", dtype=torch.bfloat16)
         quantize_(model, config)

torchao/quantization/__init__.py

@@ -39,7 +39,6 @@
     AffineQuantizedObserverBase,
 )
 from .quant_api import (
-    AOPerModuleConfig,
     CutlassInt4PackedLayout,
     Float8DynamicActivationFloat8SemiSparseWeightConfig,
     Float8DynamicActivationFloat8WeightConfig,
@@ -55,6 +54,7 @@
     Int8DynamicActivationIntxWeightConfig,
     Int8WeightOnlyConfig,
     IntxWeightOnlyConfig,
+    ModuleFqnToConfig,
     PlainLayout,
     TensorCoreTiledLayout,
     UIntXWeightOnlyConfig,
@@ -147,7 +147,7 @@
     "IntxWeightOnlyConfig",
     "FPXWeightOnlyConfig",
     "GemliteUIntXWeightOnlyConfig",
-    "AOPerModuleConfig",
+    "ModuleFqnToConfig",
     # smooth quant - subject to change
     "get_scale",
     "SmoothFakeDynQuantMixin",

torchao/quantization/quant_api.py

@@ -136,6 +136,7 @@
     "Int8DynActInt4WeightQuantizer",
     "Int8DynActInt4WeightGPTQQuantizer",
     "Float8DynamicActivationFloat8SemiSparseWeightConfig",
+    "ModuleFqnToConfig",
 ]
 
 LAYOUT_TO_ZERO_POINT_DOMAIN = {
@@ -596,10 +597,10 @@ def quantize_(
     """
     filter_fn = _is_linear if filter_fn is None else filter_fn
 
-    if isinstance(config, AOPerModuleConfig):
+    if isinstance(config, ModuleFqnToConfig):
         _replace_with_custom_fn_if_matches_filter_with_name(
             model,
-            _ao_per_module_config_handler,
+            _module_fqn_to_config_handler,
             filter_fn,
             device=device,
             extra_args=(config,),
@@ -2002,7 +2003,7 @@ def _fpx_weight_only_transform(
 
 
 @dataclass
-class AOPerModuleConfig(AOBaseConfig):
+class ModuleFqnToConfig(AOBaseConfig):
     """Per module configurations for torchao quantize_ API
 
     Args:
@@ -2018,8 +2019,8 @@ class AOPerModuleConfig(AOBaseConfig):
     )
 
 
-def _ao_per_module_config_handler(
-    module: torch.nn.Module, module_fqn: str, config: AOPerModuleConfig
+def _module_fqn_to_config_handler(
+    module: torch.nn.Module, module_fqn: str, config: ModuleFqnToConfig
 ):
     c = None
     if module_fqn in config.module_fqn_to_config: