Config serde

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

Author: drisspg

test/quantization/test_config_serde.py

@@ -0,0 +1,170 @@
+import json
+
+import pytest
+import torch
+
+from torchao.quantization.quant_api import (
+    Float8WeightOnlyConfig,
+    FPXWeightOnlyConfig,
+    GemliteUIntXWeightOnlyConfig,
+    Int4WeightOnlyConfig,
+    Int8WeightOnlyConfig,
+    UIntXWeightOnlyConfig,
+)
+
+# 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(
+        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),
+    # 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,84 @@
-import abc
+from typing import Any, Dict
 
+import torch
+from pydantic import BaseModel, 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 = {
+        "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

@@ -18,7 +18,6 @@
 import logging
 import types
 import warnings
-from dataclasses import dataclass
 from typing import Any, Callable, Optional, Tuple, Union
 
 import torch
@@ -600,7 +599,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 +679,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
 
@@ -738,7 +735,6 @@ def _int4_dynamic_activation_int4_weight_transform(
     return module
 
 
-@dataclass
 class GemliteUIntXWeightOnlyConfig(AOBaseConfig):
     """
     applies weight only 4 or 8 bit integer quantization and utilizes the gemlite triton kernel and its associated weight packing format.
@@ -787,7 +783,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
@@ -811,7 +806,9 @@ class Int4WeightOnlyConfig(AOBaseConfig):
     """
 
     group_size: int = 128
-    layout: Optional[TensorCoreTiledLayout] = TensorCoreTiledLayout(inner_k_tiles=8)
+    layout: Optional[Union[TensorCoreTiledLayout, Int4CPULayout]] = (
+        TensorCoreTiledLayout(inner_k_tiles=8)
+    )
     use_hqq: bool = False
     zero_point_domain: Optional[ZeroPointDomain] = ZeroPointDomain.NONE
 
@@ -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.
@@ -1245,7 +1239,6 @@ def _fp8_mm_compat(weight: torch.Tensor) -> bool:
     return is_compatible
 
 
-@dataclass
 class Float8DynamicActivationFloat8WeightConfig(AOBaseConfig):
     """
     Configuration for applying float8 dynamic symmetric quantization to both activations and weights of linear layers.
@@ -1327,7 +1320,6 @@ def _float8_dynamic_activation_float8_weight_transform(
     return module
 
 
-@dataclass
 class Float8StaticActivationFloat8WeightConfig(AOBaseConfig):
     """
     Configuration for applying float8 static symmetric quantization to
@@ -1408,7 +1400,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
@@ -1499,7 +1490,6 @@ def _uintx_weight_only_transform(
     return module
 
 
-@dataclass
 class FPXWeightOnlyConfig(AOBaseConfig):
     """Sub-byte floating point dtypes defined by `ebits`: exponent bits and `mbits`: mantissa bits
     e.g. fp6_e3_m2, fp6_e2_m3, ...