Reload prepared models from a state dict for all compressors

changelog.d/157.prepare-state-dict.added

@@ -0,0 +1 @@
+Added a `state_dict` argument to `prepare()` for `Quantizer` (eager and graph mode) and `MagnitudePruner`, matching `KMeansPalettizer`. It reloads a previously prepared model's compression state (calibrated qparams, pruning masks and step count, centroids), validates it against the config, and leaves the model in the same state as a regular `prepare()` call.

changelog.d/157.prepare-state-dict.changed

@@ -0,0 +1 @@
+Pruning modules now store their step count in a `step_count` buffer, so a `MagnitudePruner` state dict saved before this change cannot be reloaded with `prepare(state_dict=...)`.

changelog.d/157.prepare-state-dict.fixed

@@ -0,0 +1 @@
+Fixed `KMeansPalettizer.prepare(state_dict=...)` failing for configs with a quantized LUT, and `MagnitudePruner` restarting its sparsity schedule from step 0 after a model saved mid-schedule was reloaded.

src/coreai_opt/base_model_compressor.py

@@ -10,6 +10,7 @@
 """
 
 from abc import ABC, abstractmethod
+from collections.abc import Mapping
 from contextlib import contextmanager
 from enum import Enum
 from os import PathLike
@@ -105,11 +106,59 @@ def prepare(self, *args, **kwargs) -> torch.nn.Module:
         accuracy may be lower than the baseline model's accuracy; otherwise, it should
         remain the same.
 
+        Every compressor accepts an optional ``state_dict`` keyword argument: the state
+        dict of a model previously prepared with the same config. When provided, the
+        compression state (e.g. calibrated qparams, centroids, pruning masks) is loaded
+        from it via ``_load_prepared_state_dict`` instead of being recomputed. The
+        resulting model is in the same state as after a regular ``prepare`` call (e.g.
+        observers disabled and fake quantization enabled).
+
         Returns:
             The prepared model, ready for evaluation or further compression steps
         """
         pass
 
+    def _load_prepared_state_dict(
+        self,
+        model: torch.nn.Module,
+        state_dict: Mapping[str, torch.Tensor],
+    ) -> None:
+        """Strict-load a previously prepared model's state into ``model`` and validate it.
+
+        The model's structure comes from the config, never from ``state_dict``. After
+        loading, every compression module must be initialized. Flags are buffers, so they
+        are loaded too; ``prepare`` sets them after calling this, as it does without a
+        state dict.
+
+        Args:
+            model (torch.nn.Module): The prepared model to load into.
+            state_dict (Mapping[str, torch.Tensor]): State dict of a model previously
+                prepared with the same config.
+
+        Raises:
+            RuntimeError: If ``state_dict`` does not match the prepared model's keys or
+                shapes, or leaves a compression module uninitialized.
+        """
+        try:
+            model.load_state_dict(state_dict)
+        except RuntimeError as e:
+            raise RuntimeError(
+                "State dict does not match the prepared model. It must come from a model "
+                "prepared with the same config and model code (and, in graph mode, the same "
+                f"example inputs, dynamic shapes and torch version): {e}"
+            ) from e
+
+        uninitialized = self._uninitialized_compression_modules(model)
+        if uninitialized:
+            raise RuntimeError(
+                "State dict does not match the compression config. These compression "
+                f"modules are uninitialized after loading: {uninitialized}"
+            )
+
+    @abstractmethod
+    def _uninitialized_compression_modules(self, model: torch.nn.Module) -> list[str]:
+        """Return names of compression modules in ``model`` whose state is uninitialized."""
+
     @abstractmethod
     def finalize(
         self,

src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py

@@ -263,7 +263,7 @@ def _save_to_state_dict(self, destination, prefix, keep_vars):
 
     def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
         """Load centroids from a checkpoint, reconstructing them from a legacy
-        ``lut`` buffer when present.
+        ``lut`` buffer when present, and initialize an uninitialized LUT quantizer from them.
         """
         lut_key, centroids_key, index_key = prefix + "lut", prefix + "centroids", prefix + "indices"
         if centroids_key not in state_dict and lut_key in state_dict:
@@ -275,6 +275,15 @@ def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
             scale_ok = not self.enable_per_channel_scale or self.per_channel_scale is not None
             self._indices_stale = not (had_indices and scale_ok)
 
+            lut_scale = None
+            if self._lut_fake_quantizer is not None:
+                lut_scale = getattr(self._lut_fake_quantizer.qparams_calculator, "scale", None)
+            if lut_scale is not None and lut_scale.numel() == 0:
+                # Initialize the LUT quantizer as k-means would, so its qparams buffers take their
+                # shapes before its own saved qparams load: load_state_dict loads child
+                # modules after this hook.
+                self.quantize_lut(self._raw_lut(self.centroids))
+
     def _centroids_from_lut(self, lut: torch.Tensor) -> torch.Tensor:
         """Invert ``_reshape_lut_tensor`` to recover ``(num_blocks, num_clusters,
         cluster_dim)`` centroids from a stored 4D LUT tensor.

src/coreai_opt/palettization/kmeans/palettizer.py

@@ -6,7 +6,7 @@
 import logging
 import os
 import tempfile
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from contextlib import contextmanager
 from dataclasses import dataclass
 from os import PathLike
@@ -160,7 +160,7 @@ def prepare(
         example_inputs: tuple[Any, ...],
         sensitivity_path: str | None = None,
         num_workers: int = 1,
-        state_dict: dict[str, torch.Tensor] | None = None,
+        state_dict: Mapping[str, torch.Tensor] | None = None,
     ) -> torch.nn.Module:
         """
         Prepare the model for palettization.
@@ -182,7 +182,8 @@ def prepare(
                 main process instead, concurrently with the workers. Defaults to ``1``.
             state_dict: Optional state dict from a model previously prepared
                 with the same config. When provided, buffers are loaded from it
-                instead of running k-means clustering.
+                instead of running k-means clustering. Fake palettization is left
+                enabled, as a regular ``prepare`` leaves it.
 
         Returns:
             The prepared nn.Module with fake palettization
@@ -227,10 +228,11 @@ def prepare(
             # Remove FakePalettize modules that were disabled during the forward
             # pass due to incompatible granularity or cluster dimensions.
             self._remove_disabled_fake_palett_modules(self._model)
-
-            self._model.apply(_enable_fake_palett)
         else:
-            self._load_prepared_state_dict(state_dict)
+            self._remove_incompatible_fake_palett_modules()
+            self._load_prepared_state_dict(self._model, state_dict)
+
+        self._model.apply(_enable_fake_palett)
 
         # Mark the model as prepared to prevent re-preparation
         self._mark_model_as_prepared(prepared_model)
@@ -240,45 +242,24 @@ def prepare(
 
         return self._model
 
-    def _load_prepared_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None:
-        """Load a previously-prepared model's buffers instead of computing centroids.
-
-        Prunes the same incompatible modules a normal ``prepare()`` would (decided
-        from config and weight shapes, not from ``state_dict``), strict-loads, then
-        verifies every surviving module received its centroids.
-
-        Args:
-            state_dict (dict[str, torch.Tensor]): State dict from a model
-                previously prepared with the same config.
+    def _remove_incompatible_fake_palett_modules(self) -> None:
+        """Remove the modules a normal ``prepare()`` would remove, without clustering.
 
-        Raises:
-            RuntimeError: If a compatible module has no centroids after loading,
-                meaning the checkpoint does not match the palettizer config.
+        Uses a shape-only (meta) compatibility probe, so the structure comes from the
+        config and weight shapes, never from a state dict.
         """
-        # Disable exactly the modules a normal prepare() would remove, using a
-        # shape-only (meta) compatibility probe. Structure comes from the config,
-        # never from the state dict.
         for info in self._collect_fake_palett_info(to_cpu=False):
             compatible = info.fp_module.check_compatible(info.weight)
             info.fp_module._disabled = not compatible
         self._remove_disabled_fake_palett_modules(self._model)
 
-        self._model.load_state_dict(state_dict)
-
-        # Verify all surviving modules have centroids
-        missing_centroids = []
-        for info in self._collect_fake_palett_info(to_cpu=False):
-            fp = info.fp_module
-            if fp.centroids is None:
-                missing_centroids.append(fp)
-
-        if missing_centroids:
-            error_msg = (
-                "Error when loading state dict. The following palettized weights are missing "
-                "centroids/lut. State dict does not match the palettizer config: "
-                f"{[fp.tensor_fqn for fp in missing_centroids]}"
-            )
-            raise RuntimeError(error_msg)
+    def _uninitialized_compression_modules(self, model: torch.nn.Module) -> list[str]:
+        """Return the palettized weights that have no centroids."""
+        return [
+            info.fp_module.tensor_fqn
+            for info in self._collect_fake_palett_info(to_cpu=False)
+            if info.fp_module.centroids is None
+        ]
 
     @contextmanager
     def calibration_mode(

src/coreai_opt/pruning/magnitude_pruner.py

@@ -6,6 +6,7 @@
 """Magnitude pruner implementation."""
 
 import logging
+from collections.abc import Mapping
 from typing import Any, NamedTuple
 
 import torch
@@ -73,8 +74,8 @@ class MagnitudePruner(_BasePruner, _EagerCompressionComponentBuilderMixin):
         >>> pruner.finalize()
     """
 
-    _step_count: int
     _scheduled_modules: list[_ScheduledModule]
+    _prune_impls: list[PruneImplBase]
 
     def __init__(self, model: torch.nn.Module, config: MagnitudePrunerConfig | None = None):
         if config is None:
@@ -91,15 +92,32 @@ def __init__(self, model: torch.nn.Module, config: MagnitudePrunerConfig | None
             supported_ops_registry=_PrunerSupportedOpsRegistry,
             optimization_type_name="prune",
         )
-        self._step_count = 0
         self._scheduled_modules = []
+        self._prune_impls = []
 
-    def prepare(self, example_inputs: tuple[Any, ...]) -> torch.nn.Module:
+    @property
+    def _step_count(self) -> int:
+        """Number of ``step()`` calls, persisted in each pruning module's ``step_count`` buffer.
+
+        Every pruning module holds the same count; it is 0 when nothing is pruned.
+        """
+        return int(self._prune_impls[0].step_count) if self._prune_impls else 0
+
+    def prepare(
+        self,
+        example_inputs: tuple[Any, ...],
+        *,
+        state_dict: Mapping[str, torch.Tensor] | None = None,
+    ) -> torch.nn.Module:
         """Prepare the model for pruning.
 
         Args:
             example_inputs (tuple[Any, ...]): Sample inputs to trace the
                 model and configure pruning parametrizations.
+            state_dict (Mapping[str, torch.Tensor] | None): Optional state dict of a
+                model previously prepared with the same config. When provided, its
+                pruning masks and step count are loaded instead of computing the masks,
+                so a model saved mid-schedule resumes from the saved step.
 
         Returns:
             torch.nn.Module: The prepared model with pruning parametrizations.
@@ -117,11 +135,17 @@ def prepare(self, example_inputs: tuple[Any, ...]) -> torch.nn.Module:
 
         self._mark_model_as_prepared(prepared_model)
         self._model = prepared_model
-        # Apply schedule before we run a forward pass to initialize the parameterizations
+        self._prune_impls = [m for m in prepared_model.modules() if isinstance(m, PruneImplBase)]
         self._build_scheduled_modules()
+        for entry in self._scheduled_modules:
+            entry.impl.update_sparsity(self._step_count)
 
-        with _move_model_to_eval(prepared_model), torch.no_grad():
-            prepared_model(*example_inputs)
+        if state_dict is None:
+            # Run a forward pass to compute the masks of every parametrization.
+            with _move_model_to_eval(prepared_model), torch.no_grad():
+                prepared_model(*example_inputs)
+        else:
+            self._load_prepared_state_dict(prepared_model, state_dict)
 
         return self._model
 
@@ -131,8 +155,14 @@ def step(self) -> None:
         Increments the step count, then recomputes and applies the mask for
         every parametrization with a configured ``sparsity_schedule``. Safe to
         call when no schedule is configured (no-op).
+
+        Raises:
+            RuntimeError: If the model has not been prepared.
         """
-        self._step_count += 1
+        if not self._is_model_prepared(self._model):
+            raise RuntimeError("Model must be prepared before step(). Call prepare() first.")
+        for impl in self._prune_impls:
+            impl.step_count += 1
         self._apply_schedule()
 
     def finalize(
@@ -178,7 +208,7 @@ def finalize(
         return model
 
     def _build_scheduled_modules(self) -> None:
-        """Attach schedules to each scheduled PruneImplBase and apply the step-0 state."""
+        """Attach schedules to each scheduled PruneImplBase."""
         config_dict = self._config.build_module_config_dict(self._model)
         for name, module in self._model.named_modules():
             if not hasattr(module, "parametrizations"):
@@ -198,7 +228,14 @@ def _build_scheduled_modules(self) -> None:
                 self._scheduled_modules.append(
                     _ScheduledModule(impl=impl, module=module, param_name=param_name)
                 )
-        self._apply_schedule()
+
+    def _uninitialized_compression_modules(self, model: torch.nn.Module) -> list[str]:
+        """Return the pruning modules whose mask is still empty."""
+        return [
+            name
+            for name, module in model.named_modules()
+            if isinstance(module, PruneImplBase) and module.mask.numel() == 0
+        ]
 
     def _apply_schedule(self) -> None:
         """For each scheduled impl, advance its sparsity and materialize the mask."""

src/coreai_opt/pruning/spec/prune.py

@@ -37,6 +37,7 @@ class PruneImplBase(CompressionSimulatorBase):
     """
 
     schedule: SparsityScheduleBase | None = None
+    step_count: torch.Tensor
     _sparsity: float
     _target_sparsity: float
     _pruning_scheme: PruningScheme
@@ -55,6 +56,38 @@ def __init__(
         self.schedule = None
         self._dirty = True
         self.register_buffer("mask", torch.empty(0))
+        # Persisted so a reloaded model resumes mid-schedule; the sparsity is derived from it.
+        self.register_buffer("step_count", torch.tensor(0, dtype=torch.int64))
+
+    def _load_from_state_dict(
+        self,
+        state_dict: dict[str, Any],
+        prefix: str,
+        local_metadata: dict[str, Any],
+        strict: bool,
+        missing_keys: list[str],
+        unexpected_keys: list[str],
+        error_msgs: list[str],
+    ) -> None:
+        """Load a saved mask and step count, and derive the scheduled sparsity from the step.
+
+        A mask that has not been computed yet takes the saved shape.
+        """
+        mask_key = prefix + "mask"
+        # The mask is empty until first computed; take the saved shape so it loads correctly
+        if mask_key in state_dict and self.mask.numel() == 0:
+            self.mask = torch.empty_like(state_dict[mask_key], device=self.mask.device)
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+        )
+        if self.schedule is not None:
+            # Replayed from step 0, since off-boundary steps of some schedules reuse the
+            # previous sparsity. Fine for a one-off restore.
+            for step in range(int(self.step_count) + 1):
+                self.update_sparsity(step)
+        if mask_key in state_dict and self.mask.numel() > 0:
+            # A loaded mask is current, this prevents it from being recomputed during forward pass
+            self._dirty = False
 
     @property
     def sparsity(self) -> float: