[Quantization] support pass MappingType for TorchAoConfig

Author: a120092009

src/diffusers/quantizers/quantization_config.py

@@ -33,6 +33,7 @@
 from packaging import version
 
 from ..utils import is_torch_available, is_torchao_available, logging
+from torchao.quantization.quant_primitives import MappingType
 
 
 if is_torch_available():
@@ -46,6 +47,11 @@ class QuantizationMethod(str, Enum):
     GGUF = "gguf"
     TORCHAO = "torchao"
 
+class CustomJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, MappingType):
+            return obj.name
+        return super().default(obj)
 
 @dataclass
 class QuantizationConfigMixin:
@@ -673,4 +679,4 @@ def __repr__(self):
         ```
         """
         config_dict = self.to_dict()
-        return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n"
+        return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True, cls=CustomJSONEncoder)}\n"

tests/quantization/torchao/test_torchao.py

@@ -77,6 +77,7 @@ def forward(self, input, *args, **kwargs):
     from torchao.dtypes import AffineQuantizedTensor
     from torchao.quantization.linear_activation_quantized_tensor import LinearActivationQuantizedTensor
     from torchao.utils import get_model_size_in_bytes
+    from torchao.quantization.quant_primitives import MappingType
 
 
 @require_torch
@@ -122,6 +123,19 @@ def test_repr(self):
         quantization_repr = repr(quantization_config).replace(" ", "").replace("\n", "")
         self.assertEqual(quantization_repr, expected_repr)
 
+        quantization_config = TorchAoConfig("int4dq", group_size=64, act_mapping_type=MappingType.SYMMETRIC)
+        expected_repr = """TorchAoConfig {
+            "modules_to_not_convert": null,
+            "quant_method": "torchao",
+            "quant_type": "int4dq",
+            "quant_type_kwargs": {
+                "act_mapping_type": "SYMMETRIC",
+                "group_size": 64
+            }
+        }""".replace(" ", "").replace("\n", "")
+        quantization_repr = repr(quantization_config).replace(" ", "").replace("\n", "")
+        self.assertEqual(quantization_repr, expected_repr)
+
 
 # Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners
 @require_torch