Config serde

stack-info: PR: #1806, branch: drisspg/stack/41

Author: drisspg

test/quantization/test_config_serilization.py

@@ -0,0 +1,163 @@
+import json
+
+import pytest
+import torch
+
+from torchao.quantization.quant_api import (
+    Float8WeightOnlyConfig,
+    UIntXWeightOnlyConfig,
+    Int4DynamicActivationInt4WeightConfig,
+    Int4WeightOnlyConfig,
+    Int8DynamicActivationInt4WeightConfig,
+    Int8DynamicActivationInt8WeightConfig,
+    Int8WeightOnlyConfig
+)
+
+# Define test configurations as fixtures
+configs = [
+    # Float8DynamicActivationFloat8WeightConfig(),
+    Float8WeightOnlyConfig(
+        weight_dtype=torch.float8_e4m3fn,
+    ),
+    UIntXWeightOnlyConfig(dtype=torch.uint1),
+    # Int4DynamicActivationInt4WeightConfig(),
+    Int4WeightOnlyConfig(
+        group_size=32,
+    ),
+    # Int8DynamicActivationInt4WeightConfig(
+    #     group_size=64,
+    # ),
+    # Int8DynamicActivationInt8WeightConfig(),
+    # Int8WeightOnlyConfig(
+    #     group_size=128,
+    # ),
+    # UIntXWeightOnlyConfig(
+    #     bit_width=4,
+    #     group_size=32,
+    # ),
+    # Float8StaticActivationFloat8WeightConfig(
+    #     activation_dtype=torch.float8_e4m3fn,
+    #     weight_dtype=torch.float8_e4m3fn,
+    # ),
+]
+
+
+# Create ids for better test naming
+def get_config_ids(configs):
+    return [config.__class__.__name__ for config in configs]
+
+
+# Parametrized tests
+@pytest.mark.parametrize("config", configs, ids=get_config_ids)
+def test_to_dict_serialization(config):
+    """Test that all configs can be serialized to a dictionary."""
+    # Test to_dict method exists and returns a dict
+    assert hasattr(
+        config, "to_dict"
+    ), f"{config.__class__.__name__} missing to_dict method"
+    result = config.to_dict()
+    assert isinstance(result, dict)
+
+    # Check that all essential attributes are present in the dict
+    for attr_name in config.__dict__:
+        if not attr_name.startswith("_"):  # Skip private attributes
+            assert attr_name in result, f"{attr_name} missing in serialized dict"
+
+
+@pytest.mark.parametrize("config", configs, ids=get_config_ids)
+def test_to_json_serialization(config):
+    """Test that all configs can be serialized to JSON."""
+    # Test to_json method exists and returns a string
+    assert hasattr(
+        config, "to_json"
+    ), f"{config.__class__.__name__} missing to_json method"
+    json_str = config.to_json()
+    assert isinstance(json_str, str)
+
+    # Verify it's valid JSON
+    try:
+        parsed = json.loads(json_str)
+        assert isinstance(parsed, dict)
+    except json.JSONDecodeError as e:
+        pytest.fail(f"Invalid JSON for {config.__class__.__name__}: {e}")
+
+
+@pytest.mark.parametrize("config", configs, ids=get_config_ids)
+def test_from_dict_deserialization(config):
+    """Test that all configs can be deserialized from a dictionary."""
+    # Get the class of the instance
+    cls = config.__class__
+
+    # Serialize to dict
+    data = config.to_dict()
+
+    # Test from_dict class method exists
+    assert hasattr(cls, "from_dict"), f"{cls.__name__} missing from_dict class method"
+
+    # Deserialize back to instance
+    deserialized = cls.from_dict(data)
+
+    # Check it's the right class
+    assert isinstance(deserialized, cls)
+
+    # Compare key attributes
+    for attr_name in config.__dict__:
+        if not attr_name.startswith("_"):  # Skip private attributes
+            original_value = getattr(config, attr_name)
+            deserialized_value = getattr(deserialized, attr_name)
+
+            # Special handling for torch dtypes
+            if (
+                hasattr(original_value, "__module__")
+                and original_value.__module__ == "torch"
+            ):
+                assert str(original_value) == str(
+                    deserialized_value
+                ), f"Attribute {attr_name} mismatch for {cls.__name__}"
+            else:
+                assert (
+                    original_value == deserialized_value
+                ), f"Attribute {attr_name} mismatch for {cls.__name__}"
+
+
+@pytest.mark.parametrize("config", configs, ids=get_config_ids)
+def test_from_json_deserialization(config):
+    """Test that all configs can be deserialized from JSON."""
+    # Get the class of the instance
+    cls = config.__class__
+
+    # Serialize to JSON
+    json_str = config.to_json()
+
+    # Test from_json class method exists
+    assert hasattr(cls, "from_json"), f"{cls.__name__} missing from_json class method"
+
+    # Deserialize back to instance
+    deserialized = cls.from_json(json_str)
+
+    # Check it's the right class
+    assert isinstance(deserialized, cls)
+
+    # Verify the instance is equivalent to the original
+    # This assumes __eq__ is properly implemented
+    assert (
+        config == deserialized
+    ), f"Deserialized instance doesn't match original for {cls.__name__}"
+
+
+@pytest.mark.parametrize("config", configs, ids=get_config_ids)
+def test_round_trip_equivalence(config):
+    """Test complete serialization and deserialization round trip."""
+    # JSON round trip
+    json_str = config.to_json()
+    deserialized_from_json = config.__class__.from_json(json_str)
+    assert (
+        config == deserialized_from_json
+    ), f"JSON round trip failed for {config.__class__.__name__}"
+
+    # Dict round trip
+    data_dict = config.to_dict()
+    deserialized_from_dict = config.__class__.from_dict(data_dict)
+    assert (
+        config == deserialized_from_dict
+    ), f"Dict round trip failed for {config.__class__.__name__}"

torchao/core/config.py

@@ -1,29 +1,56 @@
-import abc
+from typing import Any, Dict
 
+import torch
+from pydantic import BaseModel, field_serializer, model_validator
 
-class AOBaseConfig(abc.ABC):
-    """
-    If a workflow config inherits from this then `quantize_` knows
-    how to a apply it to a model. For example::
-
-        # user facing code
-        class WorkflowFooConfig(AOBaseConfig): ...
-            # configuration for workflow `Foo` is defined here
-            bar = 'baz'
-
-        # non user facing code
-        @register_quantize_module_handler(WorkflowFooConfig)
-        def _transform(
-            mod: torch.nn.Module,
-            config: WorkflowFooConfig,
-        ) -> torch.nn.Module:
-            # the transform is implemented here, usually a tensor sublass
-            # weight swap or a module swap
-            ...
-
-        # then, the user calls `quantize_` with a config, and `_transform` is called
-        # under the hood by `quantize_.
 
+class AOBaseConfig(BaseModel):
+    """
+    Base configuration class with native Pydantic handling for torch.dtype.
     """
 
-    pass
+    model_config = {
+        "arbitrary_types_allowed": True,
+        "validate_assignment": True,
+        "extra": "forbid",
+        "validate_default": True,
+        "populate_by_name": True,
+    }
+
+    @field_serializer("*")
+    def serialize_torch_dtype(self, v, _info):
+        if isinstance(v, torch.dtype):
+            return str(v)
+        return v
+
+    @model_validator(mode="before")
+    @classmethod
+    def convert_dtypes(cls, data: Any) -> Any:
+        """Simple converter for torch dtype strings"""
+        if isinstance(data, str) and data.startswith("torch."):
+            dtype_name = data.split("torch.")[1]
+            if hasattr(torch, dtype_name):
+                return getattr(torch, dtype_name)
+        elif isinstance(data, dict):
+            return {k: cls.convert_dtypes(v) for k, v in data.items()}
+        elif isinstance(data, list):
+            return [cls.convert_dtypes(item) for item in data]
+        return data
+
+    def to_dict(self) -> dict:
+        """Convert the configuration to a dictionary"""
+        return self.model_dump()
+
+    def to_json(self) -> str:
+        """Convert the configuration to a JSON string."""
+        return self.model_dump_json()
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "AOBaseConfig":
+        """Create a configuration from a dictionary."""
+        return cls.model_validate(data)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "AOBaseConfig":
+        """Create a configuration from a JSON string."""
+        return cls.model_validate_json(json_str)

torchao/quantization/quant_api.py

@@ -600,7 +600,6 @@ def _int8_symm_per_token_quant(x: torch.Tensor) -> torch.Tensor:
     )
 
 
-@dataclass
 class Int8DynamicActivationInt4WeightConfig(AOBaseConfig):
     """Configuration for applying int8 dynamic per token asymmetric activation quantization and int4 per group weight symmetric quantization to linear
     This is used to produce a model for executorch backend, but currently executorch did not
@@ -681,7 +680,6 @@ def _int8_dynamic_activation_int4_weight_transform(
     return module
 
 
-@dataclass
 class Int4DynamicActivationInt4WeightConfig(AOBaseConfig):
     """Applies int4 dynamic per token symmetric activation quantization and int4 per row weight symmetric quantization to linear
 
@@ -787,7 +785,6 @@ def _gemlite_uintx_weight_only_transform(
     return module
 
 
-@dataclass
 class Int4WeightOnlyConfig(AOBaseConfig):
     """
     Configuration for applying uint4 weight-only asymmetric per-group quantization to linear layers, using
@@ -893,7 +890,6 @@ def _int4_weight_only_transform(
     return module
 
 
-@dataclass
 class Int8WeightOnlyConfig(AOBaseConfig):
     """
     Configuration for applying int8 weight-only symmetric per-channel quantization to linear layers.
@@ -1007,7 +1003,6 @@ def _int4_symm_per_token_quant_cutlass(x: torch.Tensor) -> torch.Tensor:
     )
 
 
-@dataclass
 class Int8DynamicActivationInt8WeightConfig(AOBaseConfig):
     """
     Configuration for applying int8 dynamic symmetric per-token activation and int8 per-channel weight
@@ -1092,7 +1087,6 @@ def int8_dynamic_activation_int8_semi_sparse_weight():
     return int8_dynamic_activation_int8_weight(layout=SemiSparseLayout())
 
 
-@dataclass
 class Float8WeightOnlyConfig(AOBaseConfig):
     """
     Configuration for applying float8 weight-only symmetric per-channel quantization to linear layers.
@@ -1408,7 +1402,6 @@ def _float8_static_activation_float8_weight_transform(
     return module
 
 
-@dataclass
 class UIntXWeightOnlyConfig(AOBaseConfig):
     """
     Configuration for applying uintx weight-only asymmetric per-group quantization to linear layers, using uintx quantization where