Address review comments

Author: crowbat

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/palettization/kmeans/kmeans_fake_palettize.py

@@ -155,23 +155,19 @@ def sensitivities(self, value: torch.Tensor | None) -> None:
         self._centroids_initialized = False
         self._indices_stale = True
 
-    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        if self._disabled:
-            return tensor
-
-        if not self._centroids_initialized:
-            try:
-                self._initialize(tensor.detach())
-            except (_IncompatibleClusterDimError, _IncompatibleGranularityError) as e:
-                logger.warning(
-                    f"Tensor incompatible with configured spec: {e}. Skipping palettization."
-                )
-                self._disabled = True
-                return tensor
-
-        if self.fake_palett_enabled[0] == 0:
-            return tensor
+    def ensure_initialized(self, tensor: torch.Tensor) -> None:
+        """Cluster centroids on first use; disable on an incompatible tensor."""
+        if self._centroids_initialized:
+            return
+        try:
+            self._initialize(tensor.detach())
+        except (_IncompatibleClusterDimError, _IncompatibleGranularityError) as e:
+            logger.warning(
+                f"Tensor incompatible with configured spec: {e}. Skipping palettization."
+            )
+            self._disabled = True
 
+    def forward_enabled(self, tensor: torch.Tensor) -> torch.Tensor:
         if self.training:
             return self._training_strategy.train_forward(self, tensor)
         return self.hard_assign(tensor)
@@ -223,14 +219,14 @@ def lut(self) -> torch.Tensor | None:
 
         raw_lut = self._raw_lut(self.centroids)
         if self._lut_fake_quantizer is None:
-            lut = raw_lut
-        else:
-            orig_dtype = raw_lut.dtype
-            scale, zero_point, minval = self._lut_fake_quantizer.qparams_calculator.get_qparams()
-            lut = self._lut_fake_quantizer._fused_fake_quant_dequant(
-                raw_lut.to(torch.float32), scale, zero_point, minval
-            ).to(orig_dtype)
-        return self._reshape_lut_tensor(lut)
+            return self._reshape_lut_tensor(raw_lut)
+
+        orig_dtype = raw_lut.dtype
+        scale, zero_point, minval = self._lut_fake_quantizer.qparams_calculator.get_qparams()
+        fq_lut = self._lut_fake_quantizer._fused_fake_quant_dequant(
+            raw_lut.to(torch.float32), scale, zero_point, minval
+        ).to(orig_dtype)
+        return self._reshape_lut_tensor(fq_lut)
 
     @property
     def quantized_lut(self) -> torch.Tensor | None:

src/coreai_opt/palettization/kmeans/palettizer.py

@@ -29,8 +29,10 @@
 from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor
 from coreai_opt._utils.torch_utils import (
     move_model_to_eval as _move_model_to_eval,
+    move_model_to_train as _move_model_to_train,
     remove_compression_parametrizations as _remove_compression_parametrizations,
 )
+from coreai_opt.base_model_compressor import _CompressorLifecycle
 from coreai_opt.common import ExportBackend
 from coreai_opt.config.compression_config import ModuleCompressionConfig, ModuleConfigDict
 from coreai_opt.config.spec import CompressionTargetTensor
@@ -146,7 +148,6 @@ def __init__(self, model: torch.nn.Module, config: KMeansPalettizerConfig | None
 
         self._num_workers = 1
 
-        self._mode: str = "idle"  # "idle" | "training" | "calibrating"
         self._step_count: int = 0
         self._fp_to_schedule: dict[_KMeansFakePalettize, PATSchedule] = {}
         self._module_config_dict: ModuleConfigDict[ModuleCompressionConfig] = {}
@@ -288,9 +289,11 @@ def calibration_mode(
                 "Model must be prepared before entering calibration mode. Call prepare() first."
             )
 
-        if self._mode != "idle":
-            raise RuntimeError(f"Cannot enter calibration_mode() while palettizer is {self._mode}")
-        self._mode = "calibrating"
+        if self._lifecycle is not _CompressorLifecycle.IDLE:
+            raise RuntimeError(
+                f"Cannot enter calibration_mode() while palettizer is {self._lifecycle.value}"
+            )
+        self._lifecycle = _CompressorLifecycle.CALIBRATING
         try:
             # Save model checkpoint before modifying gradients
             checkpoint_path = self._save_model_checkpoint(self._model)
@@ -346,30 +349,34 @@ def step(self, output: torch.Tensor, target: torch.Tensor):
                     # Restore normal operation
                     self._model.apply(_enable_fake_palett)
         finally:
-            self._mode = "idle"
+            self._lifecycle = _CompressorLifecycle.IDLE
 
     @contextmanager
     def training_mode(self):
         """Context manager wrapping a training loop. Mutually exclusive with
         calibration_mode().
         """
-        if self._mode != "idle":
-            raise RuntimeError(f"Cannot enter training_mode() while palettizer is {self._mode}")
-        self._mode = "training"
-        try:
-            self._model.train()
-            self._build_fp_to_schedule()
-            self._apply_schedule()
+        if not self._is_model_prepared(self._model):
+            raise RuntimeError(
+                "Model must be prepared before entering training mode. Call prepare() first."
+            )
+
+        if self._lifecycle is not _CompressorLifecycle.IDLE:
+            raise RuntimeError(
+                f"Cannot enter training_mode() while palettizer is {self._lifecycle.value}"
+            )
+        self._lifecycle = _CompressorLifecycle.TRAINING
+        with _move_model_to_train(self._model):
             try:
+                self._build_fp_to_schedule()
+                self._apply_schedule()
                 yield self
             finally:
-                self._model.eval()
-        finally:
-            self._mode = "idle"
+                self._lifecycle = _CompressorLifecycle.IDLE
 
     def step(self) -> None:
         """Advance the schedule by one step. Must be called inside training_mode()."""
-        if self._mode != "training":
+        if self._lifecycle is not _CompressorLifecycle.TRAINING:
             raise RuntimeError("step() must be called inside a training_mode() context.")
         self._step_count += 1
         self._apply_schedule()

src/coreai_opt/palettization/spec/fake_palettize.py

@@ -56,9 +56,36 @@ def is_disabled(self) -> bool:
         """Return True if fake palettization has been disabled."""
         return self._disabled
 
-    @abstractmethod
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Apply fake palettization to input tensor."""
+        """Fake-palettize ``tensor`` through the enable/disable lifecycle.
+
+        Delegates representation-specific work to ``ensure_initialized`` and
+        ``forward_enabled``.
+        """
+        if self._disabled:
+            return tensor
+
+        self.ensure_initialized(tensor)
+
+        # Check for self._disabled again in case ensure_initialized disabled the palettizer.
+        if self._disabled:
+            return tensor
+        if self.fake_palett_enabled[0] == 0:
+            return tensor
+        return self.forward_enabled(tensor)
+
+    @abstractmethod
+    def ensure_initialized(self, tensor: torch.Tensor) -> None:
+        """Initialize compression parameters from ``tensor`` on first use.
+
+        Set ``self._disabled = True`` if ``tensor`` is incompatible with the
+        configured spec.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def forward_enabled(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Return the palettized output for ``tensor`` when enabled."""
         raise NotImplementedError()
 
     @abstractmethod

tests/palettization/test_pat_schedule.py

@@ -19,6 +19,7 @@
 import torch.nn.utils.parametrize as P
 from pydantic import ValidationError
 
+from coreai_opt.base_model_compressor import _CompressorLifecycle
 from coreai_opt.palettization import (
     KMeansPalettizer,
     KMeansPalettizerConfig,
@@ -114,13 +115,31 @@ def test_frozen(self):
 class TestTrainingMode:
     """Runtime behavior of KMeansPalettizer.training_mode()."""
 
+    def test_training_mode_requires_prepared_model(self):
+        config = KMeansPalettizerConfig(
+            global_config=ModuleKMeansPalettizerConfig(
+                op_state_spec={"weight": default_weight_palettization_spec()},
+            )
+        )
+        palettizer = KMeansPalettizer(ToyModel(), config)
+        with pytest.raises(RuntimeError, match="Model must be prepared"):
+            with palettizer.training_mode():
+                pass
+
     def test_entry_trains_exit_evals(self):
         palettizer, prepared = _prepared_palettizer()
         prepared.eval()
         with palettizer.training_mode():
             assert prepared.training is True
         assert prepared.training is False
 
+    def test_entry_from_train_stays_train_on_exit(self):
+        palettizer, prepared = _prepared_palettizer()
+        prepared.train()
+        with palettizer.training_mode():
+            assert prepared.training is True
+        assert prepared.training is True
+
     def test_default_no_schedule_stays_enabled(self):
         palettizer, prepared = _prepared_palettizer()
         with palettizer.training_mode():
@@ -161,9 +180,9 @@ def test_mode_restored_to_idle_on_exception(self):
         palettizer, _ = _prepared_palettizer()
         with pytest.raises(ValueError, match="boom"):
             with palettizer.training_mode():
-                assert palettizer._mode == "training"
+                assert palettizer._lifecycle is _CompressorLifecycle.TRAINING
                 raise ValueError("boom")
-        assert palettizer._mode == "idle"
+        assert palettizer._lifecycle is _CompressorLifecycle.IDLE
 
 
 class TestStep:
@@ -298,10 +317,12 @@ def forward(self, x):
 class TestDefaultStrategyTraining:
     """The default training strategy's behavior inside a training_mode() loop."""
 
-    def test_freezes_palettized_weight_but_trains_the_rest(self):
-        """The default strategy reconstructs from frozen centroids, so a
-        palettized weight receives no gradient during training, while
-        non-palettized parameters still train against the palettized values.
+    @pytest.mark.parametrize("use_training_mode_ctx", [True, False])
+    def test_gradient_flow_with_and_without_training_mode_context(self, use_training_mode_ctx):
+        """Training-time gradients must flow through the palettized layer whether
+        or not the training_mode() context is active: the input and downstream
+        params receive gradients while the frozen palettized weight receives none
+        (no unintended gradient path).
         """
         config = KMeansPalettizerConfig(
             module_name_configs={
@@ -314,8 +335,16 @@ def test_freezes_palettized_weight_but_trains_the_rest(self):
         palettizer = KMeansPalettizer(MixedModel(), config)
         prepared = palettizer.prepare((torch.randn(2, 16),))
 
-        with palettizer.training_mode():
-            prepared(torch.randn(2, 16)).sum().backward()
+        x = torch.randn(2, 16, requires_grad=True)
+        if use_training_mode_ctx:
+            with palettizer.training_mode():
+                prepared(x).sum().backward()
+        else:
+            prepared.train()
+            prepared(x).sum().backward()
 
+        # frozen palettized weight gets no gradient (no unintended path)
         assert prepared.palettized.parametrizations.weight.original.grad is None
+        # graph intact: gradients reach the input (through the palettized layer) and downstream
+        assert x.grad is not None
         assert prepared.head.weight.grad is not None