Refactor KMeansPalettizer, add extendable training strategy capability

changelog.d/60.added

@@ -0,0 +1 @@
+Add training-aware palettization to `KMeansPalettizer` via a `training_mode()` context manager and `step()` method, with a `PATSchedule` (`pat_schedule` on the module config) controlling when palettization activates during a training loop. `PalettizationSpec` gains a `training_strategy_config` field for selecting pluggable training-time strategies.

src/coreai_opt/base_model_compressor.py

@@ -11,6 +11,7 @@
 
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
+from enum import Enum
 from os import PathLike
 
 import torch
@@ -21,12 +22,21 @@
 _COREAI_OPT_PREPARED_ATTR = "_coreai_opt_prepared"
 
 
+class _CompressorLifecycle(Enum):
+    """Lifecycle state of a model compressor."""
+
+    IDLE = "idle"
+    TRAINING = "training"
+    CALIBRATING = "calibrating"
+
+
 class _BaseModelCompressor(ABC):
     """
     An abstract base class for implementing model compression techniques.
     """
 
     _supported_modules: tuple[type[torch.nn.Module]]
+    _lifecycle: _CompressorLifecycle
 
     def __init__(self, model: torch.nn.Module, config: CompressionConfig | None = None):
         """
@@ -39,6 +49,7 @@ def __init__(self, model: torch.nn.Module, config: CompressionConfig | None = No
         """
         self._model = model
         self._config = config
+        self._lifecycle = _CompressorLifecycle.IDLE
 
     @staticmethod
     def _is_model_prepared(model: torch.nn.Module) -> bool:

src/coreai_opt/config/compression_config.py

@@ -379,21 +379,11 @@ def _normalize_none_op_configs(self) -> ModuleCompressionConfig:
 
         return self
 
-    def _get_compressor_specific_settings(self) -> dict[str, Any]:
-        """
-        Get compressor-specific settings, excluding base ModuleCompressionConfig fields.
-
-        Returns only fields defined by the concrete subclass (e.g.,
-        enable_fast_kmeans_mode, rounding_precision for palettization), not the base
-        spec and config fields (op_input_spec, op_output_spec, op_state_spec,
-        op_type_config, op_name_config, module_input_spec, module_output_spec,
-        module_state_spec).
-
-        This is useful when constructing arguments for compression operations that need
-        the compression-specific settings but handle the spec fields separately.
+    def _get_fake_module_kwargs(self) -> dict[str, Any]:
+        """Get the settings to forward to the compression simulator module's constructor.
 
         Returns:
-            Dictionary of compressor-specific field names to their values.
+            Dictionary of field names to their values.
         """
         base_field_names = set(ModuleCompressionConfig.model_fields.keys())
         return {k: v for k, v in self.model_dump().items() if k not in base_field_names}

src/coreai_opt/palettization/config/__init__.py

@@ -9,10 +9,12 @@
     KMeansPalettizerConfig,
     ModuleKMeansPalettizerConfig,
     OpKMeansPalettizerConfig,
+    PATSchedule,
 )
 
 __all__ = [
     "KMeansPalettizerConfig",
     "ModuleKMeansPalettizerConfig",
     "OpKMeansPalettizerConfig",
+    "PATSchedule",
 ]

src/coreai_opt/palettization/config/palettization_config.py

@@ -7,7 +7,7 @@
 
 from typing import TYPE_CHECKING, ClassVar, final
 
-from pydantic import PositiveInt, model_validator
+from pydantic import BaseModel, ConfigDict, Field, PositiveInt, model_validator
 
 from coreai_opt.config import (
     CompressionConfig,
@@ -31,6 +31,29 @@
 _PALETTIZATION_SPEC = "palettization_spec"
 
 
+class PATSchedule(BaseModel):
+    """Schedule for enabling palettization-aware training (PAT).
+
+    Defines the step threshold at which a module's fake palettization
+    forward pass becomes active. Used with ``KMeansPalettizer.step()``.
+
+    Attributes:
+        enable_fake_palettize: Step count at which fake palettization is
+            enabled. Must be >= 0.
+
+    Example:
+        >>> schedule = PATSchedule(enable_fake_palettize=500)
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    enable_fake_palettize: int = Field(default=0, ge=0)
+
+    def _compute_state(self, step_count: int) -> bool:
+        """Return whether fake palettization should be active at the given step."""
+        return step_count >= self.enable_fake_palettize
+
+
 class OpKMeansPalettizerConfig(WeightOnlyOpValidationMixin, OpCompressionConfig[PalettizationSpec]):
     """
     Configuration class for palettization at the operation level.
@@ -140,6 +163,11 @@ class ModuleKMeansPalettizerConfig(
             K-means clustering. Higher values preserve more precision but may reduce
             speed benefits. Only used when enable_fast_kmeans_mode is True. Default: 4.
 
+        pat_schedule (PATSchedule | None): Schedule controlling when this
+            module's palettization is active during a training_mode() loop.
+            If None, palettization is active immediately once training_mode()
+            begins. Default: None.
+
     Example:
         >>> config = ModuleKMeansPalettizerConfig()  # Uses defaults
         >>> # Or with custom settings:
@@ -160,6 +188,7 @@ def __init_subclass__(cls, **kwargs):
 
     enable_fast_kmeans_mode: bool = True
     rounding_precision: PositiveInt = 4
+    pat_schedule: PATSchedule | None = None
 
     # Namespace exposing built-in preset constructors.
     presets: ClassVar[_ModuleKMeansPalettizerConfigPresets]
@@ -183,6 +212,12 @@ def validate_fast_kmeans_cluster_dim_constraint(
 
         return self
 
+    def _get_fake_module_kwargs(self) -> dict:
+        """Exclude palettizer-only fields from the fake-palettize constructor args."""
+        base = super()._get_fake_module_kwargs()
+        base.pop("pat_schedule", None)
+        return base
+
 
 @final
 class KMeansPalettizerConfig(CompressionConfig[ModuleKMeansPalettizerConfig]):