Config serde

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

Author: drisspg

setup.py

@@ -327,6 +327,9 @@ def get_extensions():
     version=version + version_suffix,
     packages=find_packages(),
     include_package_data=True,
+    install_requires=[
+        "pydantic>=2",
+    ],
     package_data={
         "torchao.kernel.configs": ["*.pkl"],
     },

test/dtypes/test_affine_quantized.py

@@ -60,7 +60,7 @@ def get_quantization_functions(
             if device == "cuda":
                 base_functions.append(
                     int8_dynamic_activation_int4_weight(
-                        group_size=None,
+                        group_size=32,
                         mapping_type=MappingType.SYMMETRIC,
                         act_mapping_type=MappingType.SYMMETRIC,
                         layout=CutlassInt4PackedLayout(),

test/quantization/test_config_serialization.py

@@ -0,0 +1,177 @@
+import json
+
+import pytest
+import torch
+
+from torchao.quantization.quant_api import (
+    Float8DynamicActivationFloat8WeightConfig,
+    Float8WeightOnlyConfig,
+    FPXWeightOnlyConfig,
+    GemliteUIntXWeightOnlyConfig,
+    Int4DynamicActivationInt4WeightConfig,
+    Int4WeightOnlyConfig,
+    Int8DynamicActivationInt4WeightConfig,
+    Int8DynamicActivationInt8WeightConfig,
+    Int8WeightOnlyConfig,
+    PerRow,
+    UIntXWeightOnlyConfig,
+)
+
+# Define test configurations as fixtures
+configs = [
+    Float8DynamicActivationFloat8WeightConfig(),
+    Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()),
+    Float8WeightOnlyConfig(
+        weight_dtype=torch.float8_e4m3fn,
+    ),
+    UIntXWeightOnlyConfig(dtype=torch.uint1),
+    Int4DynamicActivationInt4WeightConfig(),
+    Int4WeightOnlyConfig(
+        group_size=32,
+    ),
+    Int8DynamicActivationInt4WeightConfig(
+        group_size=64,
+    ),
+    Int8DynamicActivationInt8WeightConfig(),
+    # Int8DynamicActivationInt8WeightConfig(layout=SemiSparseLayout()),
+    Int8WeightOnlyConfig(
+        group_size=128,
+    ),
+    UIntXWeightOnlyConfig(
+        dtype=torch.uint3,
+        group_size=32,
+        use_hqq=True,
+    ),
+    GemliteUIntXWeightOnlyConfig(
+        group_size=128,  # Optional, has default of 64
+        bit_width=8,  # Optional, has default of 4
+        packing_bitwidth=8,  # Optional, has default of 32
+        contiguous=True,  # Optional, has default of None
+    ),
+    FPXWeightOnlyConfig(ebits=4, mbits=8),
+]
+
+
+# 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__}"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

torchao/core/config.py

@@ -1,29 +1,84 @@
-import abc
+from typing import Any, Dict
 
+import torch
+from pydantic import BaseModel, ConfigDict, field_serializer, model_validator
 
-class AOBaseConfig(abc.ABC):
+
+class AOBaseConfig(BaseModel):
     """
-    If a workflow config inherits from this then `quantize_` knows
-    how to a apply it to a model. For example::
+    Base configuration class for TorchAO quantization workflows with native Pydantic handling for torch.dtype.
+
+    When a workflow configuration inherits from AOBaseConfig, the `quantize_` function can automatically
+    apply the appropriate transformation to a model based on the configuration type.
 
-        # user facing code
-        class WorkflowFooConfig(AOBaseConfig): ...
-            # configuration for workflow `Foo` is defined here
-            bar = 'baz'
+    Usage example:
+        # 1. Define a configuration class for your workflow
+        class WorkflowFooConfig(AOBaseConfig):
+            # Configuration parameters for workflow 'Foo'
+            bar: str = 'baz'
 
-        # non user facing code
+        # 2. Register a handler for this configuration (internal implementation)
         @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
+            # Implementation of the transformation logic
+            # Typically performs tensor subclass weight swapping or module replacement
             ...
 
-        # then, the user calls `quantize_` with a config, and `_transform` is called
-        # under the hood by `quantize_.
+        # 3. Apply the configuration to a model
+        # The user simply calls `quantize_` with a model and config instance
+        # The appropriate handler is automatically selected based on the config type
+        model = ...
+        quantized_model = quantize_(model, WorkflowFooConfig(bar='custom_value'))
 
+    Note on serialization, if you add a new AOBaseConfig and want to support serialization,
+    please add a test in test/quantization/test_config_serialization.py
     """
 
-    pass
+    model_config = ConfigDict(
+        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)