diff --git a/changelog.d/157.prepare-state-dict.added b/changelog.d/157.prepare-state-dict.added
new file mode 100644
index 00000000..d399b8d3
--- /dev/null
+++ b/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.
diff --git a/changelog.d/157.prepare-state-dict.changed b/changelog.d/157.prepare-state-dict.changed
new file mode 100644
index 00000000..4934fae8
--- /dev/null
+++ b/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=...)`.
diff --git a/changelog.d/157.prepare-state-dict.fixed b/changelog.d/157.prepare-state-dict.fixed
new file mode 100644
index 00000000..0dc5173b
--- /dev/null
+++ b/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.
diff --git a/src/coreai_opt/base_model_compressor.py b/src/coreai_opt/base_model_compressor.py
index 84362570..614200b8 100644
--- a/src/coreai_opt/base_model_compressor.py
+++ b/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,
diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py
index b074540b..399ec484 100644
--- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py
+++ b/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.
diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py
index 2be31405..0b1d0f02 100644
--- a/src/coreai_opt/palettization/kmeans/palettizer.py
+++ b/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(
diff --git a/src/coreai_opt/pruning/magnitude_pruner.py b/src/coreai_opt/pruning/magnitude_pruner.py
index 035b27c8..22b46314 100644
--- a/src/coreai_opt/pruning/magnitude_pruner.py
+++ b/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."""
diff --git a/src/coreai_opt/pruning/spec/prune.py b/src/coreai_opt/pruning/spec/prune.py
index 7333fdad..4f32b1ee 100644
--- a/src/coreai_opt/pruning/spec/prune.py
+++ b/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:
diff --git a/src/coreai_opt/quantization/_eager/quantizer.py b/src/coreai_opt/quantization/_eager/quantizer.py
index 15e89384..7c81df98 100644
--- a/src/coreai_opt/quantization/_eager/quantizer.py
+++ b/src/coreai_opt/quantization/_eager/quantizer.py
@@ -6,7 +6,7 @@
 import itertools
 import warnings
 from collections import defaultdict
-from collections.abc import Callable, Generator
+from collections.abc import Callable, Generator, Mapping
 from contextlib import contextmanager
 from os import PathLike
 from typing import Any
@@ -172,13 +172,21 @@ def _validate_config(config: QuantizerConfig) -> None:
                 stacklevel=2,
             )
 
-    def prepare(self, example_inputs: tuple[Any, ...]) -> nn.Module:
+    def prepare(
+        self,
+        example_inputs: tuple[Any, ...],
+        *,
+        state_dict: Mapping[str, torch.Tensor] | None = None,
+    ) -> nn.Module:
         """Prepare model for quantization by inserting fake quantizers.
 
         Args:
             example_inputs: Sample inputs used to trace the model and configure
                 quantizers. When activation quantization is in use, these should
                 be representative of the data the model would typically see.
+            state_dict: Optional state dict of a model previously prepared with the
+                same config. When provided, its quantization state (e.g. calibrated
+                qparams) is loaded after preparing.
 
         Returns:
             Model with fake quantizers inserted
@@ -211,6 +219,8 @@ def prepare(self, example_inputs: tuple[Any, ...]) -> nn.Module:
             # pass due to incompatible tensor shapes (e.g., non-divisible block sizes).
             self._remove_disabled_fake_quant_modules()
 
+        if state_dict is not None:
+            self._load_prepared_state_dict(self._model, state_dict)
         self._model.apply(enable_fake_quant)
         self._model.apply(disable_observer)
 
diff --git a/src/coreai_opt/quantization/_graph/quantizer.py b/src/coreai_opt/quantization/_graph/quantizer.py
index 03e9dec0..936ef85c 100644
--- a/src/coreai_opt/quantization/_graph/quantizer.py
+++ b/src/coreai_opt/quantization/_graph/quantizer.py
@@ -962,6 +962,8 @@ def prepare(
         example_inputs: tuple[Any, ...],
         dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = None,
         export_with_no_grad: bool = True,
+        *,
+        state_dict: Mapping[str, torch.Tensor] | None = None,
     ) -> torch.fx.GraphModule:
         """
         Prepare the model for quantization.
@@ -987,6 +989,10 @@ def prepare(
                           at runtime during model export.
             export_with_no_grad: Whether to call torch.export.export within a
                                torch.no_grad() context. Defaults to True.
+            state_dict: Optional state dict of a model previously prepared with the
+                same config, example inputs, dynamic shapes and torch version. When
+                provided, its quantization state (e.g. calibrated qparams) is loaded
+                after preparing.
 
         Returns:
             The prepared GraphModule with quantization annotations and fake quantization
@@ -1093,6 +1099,8 @@ def prepare(
         self._remove_disabled_fake_quant_nodes(prepared_model)
 
         # Set up initial quantization state: fake quant enabled, observers disabled
+        if state_dict is not None:
+            self._load_prepared_state_dict(prepared_model, state_dict)
         prepared_model.apply(enable_fake_quant)
         prepared_model.apply(disable_observer)
 
diff --git a/src/coreai_opt/quantization/base_quantizer.py b/src/coreai_opt/quantization/base_quantizer.py
index f5ccdb12..4edd6cc8 100644
--- a/src/coreai_opt/quantization/base_quantizer.py
+++ b/src/coreai_opt/quantization/base_quantizer.py
@@ -5,8 +5,11 @@
 
 from abc import abstractmethod
 
+import torch
+
 from coreai_opt.base_model_compressor import _BaseModelCompressor
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
+from coreai_opt.quantization.spec.qparams_calculator import StatefulQParamsCalculatorBase
 
 
 class _BaseQuantizer(_BaseModelCompressor):
@@ -23,3 +26,13 @@ def _get_fake_quantize_modules(self) -> dict[str, list[FakeQuantizeImplBase]]:
             instances.
         """
         ...
+
+    def _uninitialized_compression_modules(self, model: torch.nn.Module) -> list[str]:
+        """Return the fake quantizers whose stored qparams are still empty."""
+        return [
+            name
+            for name, module in model.named_modules()
+            if isinstance(module, FakeQuantizeImplBase)
+            and isinstance(module.qparams_calculator, StatefulQParamsCalculatorBase)
+            and module.qparams_calculator.scale.numel() == 0
+        ]
diff --git a/src/coreai_opt/quantization/quantizer.py b/src/coreai_opt/quantization/quantizer.py
index 581edc0a..60ab3da6 100644
--- a/src/coreai_opt/quantization/quantizer.py
+++ b/src/coreai_opt/quantization/quantizer.py
@@ -6,10 +6,12 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Mapping
 from contextlib import contextmanager
 from os import PathLike
 from typing import Any
 
+import torch
 import torch.nn as nn
 from torch import fx
 from torchao.quantization.pt2e import (
@@ -319,6 +321,8 @@ def prepare(
         example_inputs: tuple[Any, ...],
         dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = None,
         export_with_no_grad: bool = True,
+        *,
+        state_dict: Mapping[str, torch.Tensor] | None = None,
     ) -> nn.Module | fx.GraphModule:
         """
         Prepare the model for quantization by inserting fake quantization modules.
@@ -347,6 +351,11 @@ def prepare(
                 Ignored in EAGER mode.
             export_with_no_grad: Whether to export with no_grad (graph mode
                 only). Ignored in EAGER mode.
+            state_dict: Optional state dict of a model previously prepared with the
+                same config (and, in graph mode, the same example inputs, dynamic
+                shapes and torch version), e.g. after calibration or QAT. When
+                provided, its quantization state is loaded after preparing, and fake
+                quantization and observers are left as a regular ``prepare`` leaves them.
 
         Returns:
             The prepared model with fake quantization modules inserted, ready for
@@ -377,12 +386,13 @@ def prepare(
                     UserWarning,
                     stacklevel=2,
                 )
-            prepared_model = self._quantizer.prepare(example_inputs)
+            prepared_model = self._quantizer.prepare(example_inputs, state_dict=state_dict)
         else:
             prepared_model = self._quantizer.prepare(
                 example_inputs,
                 dynamic_shapes=dynamic_shapes,
                 export_with_no_grad=export_with_no_grad,
+                state_dict=state_dict,
             )
 
         self._module_to_fqs = self._get_fake_quantize_modules()
diff --git a/tests/palettization/test_kmeans_palettizer.py b/tests/palettization/test_kmeans_palettizer.py
index 54ef2b74..23d02091 100644
--- a/tests/palettization/test_kmeans_palettizer.py
+++ b/tests/palettization/test_kmeans_palettizer.py
@@ -1686,6 +1686,35 @@ def _raise_assert(self, *args, **kwargs):
         dst = KMeansPalettizer(copy.deepcopy(simple_conv_linear_model), basic_config)
         dst.prepare((simple_model_input,), state_dict=state_dict)
 
+    def test_reproduces_prepared_model_with_quantized_lut(
+        self, simple_conv_linear_model, simple_model_input
+    ):
+        """A model with a quantized LUT reloads, though k-means (which initializes the LUT
+        quantizers) is skipped."""
+        config = KMeansPalettizerConfig(
+            global_config=ModuleKMeansPalettizerConfig(
+                op_state_spec={
+                    "weight": PalettizationSpec(
+                        n_bits=4,
+                        lut_qspec=QuantizationSpec(
+                            dtype=torch.int8, qscheme=QuantizationScheme.SYMMETRIC
+                        ),
+                    )
+                }
+            )
+        )
+        src = KMeansPalettizer(copy.deepcopy(simple_conv_linear_model), config)
+        prepared = src.prepare((simple_model_input,))
+        with torch.no_grad():
+            expected = prepared(simple_model_input)
+
+        dst = KMeansPalettizer(copy.deepcopy(simple_conv_linear_model), config)
+        loaded = dst.prepare((simple_model_input,), state_dict=prepared.state_dict())
+        with torch.no_grad():
+            got = loaded(simple_model_input)
+
+        torch.testing.assert_close(got, expected, rtol=0, atol=0)
+
     def test_missing_centroids_raises(
         self, simple_conv_linear_model, simple_model_input, basic_config
     ):
@@ -1695,7 +1724,7 @@ def test_missing_centroids_raises(
         pruned = {k: v for k, v in state_dict.items() if not k.endswith(".centroids")}
 
         dst = KMeansPalettizer(copy.deepcopy(simple_conv_linear_model), basic_config)
-        with pytest.raises(RuntimeError, match="does not match the palettizer config"):
+        with pytest.raises(RuntimeError, match="State dict does not match"):
             dst.prepare((simple_model_input,), state_dict=pruned)
 
 
diff --git a/tests/pruning/test_magnitude_pruner.py b/tests/pruning/test_magnitude_pruner.py
index b3b60e3b..a7200c68 100644
--- a/tests/pruning/test_magnitude_pruner.py
+++ b/tests/pruning/test_magnitude_pruner.py
@@ -24,6 +24,7 @@
     PolynomialDecaySchedule,
 )
 from coreai_opt.pruning.spec import ChannelStructured, PruneImplBase, Unstructured
+from coreai_opt.pruning.spec.prune import _MagnitudePruneImpl
 
 
 @pytest.fixture
@@ -92,13 +93,16 @@ def test_torch_finalize(self) -> None:
         assert sparsity == 0.5
 
     def test_invalid_state(self) -> None:
-        """Double prepare raises; finalize without prepare raises."""
+        """Double prepare raises; finalize or step without prepare raises."""
         model = nn.Linear(10, 10)
         pruner = MagnitudePruner(model)
 
         with pytest.raises(RuntimeError, match="must be prepared"):
             pruner.finalize()
 
+        with pytest.raises(RuntimeError, match="must be prepared before step"):
+            pruner.step()
+
         pruner.prepare((torch.randn(1, 10),))
 
         with pytest.raises(RuntimeError, match="already been prepared"):
@@ -709,37 +713,66 @@ def test_per_module_schedule_resolution(self) -> None:
         assert self._impl(model[0]).sparsity == 0.5
         assert self._impl(model[1]).sparsity == 0.25
 
-    @pytest.mark.xfail(
-        strict=True,
-        reason="Pruner step count is not persisted in model.state_dict — resumes from 0.",
-    )
-    def test_save_resume_midschedule(self, linear_100x100_unique: nn.Linear) -> None:
-        """Saving the model mid-schedule and resuming should continue from the saved step."""
-        target = 0.5
-        total_iters = 10
 
-        def _build(model: nn.Linear) -> MagnitudePruner:
-            config = MagnitudePrunerConfig(
-                global_config=ModuleMagnitudePrunerConfig(
-                    op_state_spec={"weight": PruningSpec(target_sparsity=target)},
-                    sparsity_schedule=PolynomialDecaySchedule(
-                        begin_step=0, total_iters=total_iters, power=1.0
-                    ),
-                )
+class TestPrepareFromStateDict:
+    """Pruner-specific behavior of ``prepare(state_dict=...)``."""
+
+    def test_skips_mask_computation(self, monkeypatch, linear_100x100_unique: nn.Linear) -> None:
+        """Masks are loaded, not computed, including on the first forward after loading."""
+        example_input = torch.randn(1, 100)
+        source = copy.deepcopy(linear_100x100_unique)
+        MagnitudePruner(source).prepare((example_input,))
+        expected = source(example_input)
+
+        def _raise(*args, **kwargs):
+            raise AssertionError("masks must not be computed when loading from a state dict")
+
+        monkeypatch.setattr(_MagnitudePruneImpl, "compute_mask", staticmethod(_raise))
+        reloaded = MagnitudePruner(copy.deepcopy(linear_100x100_unique)).prepare(
+            (example_input,), state_dict=source.state_dict()
+        )
+
+        torch.testing.assert_close(reloaded(example_input), expected, rtol=0, atol=0)
+
+    def test_resumes_schedule_mid_way(self, monkeypatch, linear_100x100_unique: nn.Linear) -> None:
+        """A model reloaded mid-schedule uses its saved masks and continues exactly like an
+        uninterrupted run."""
+        config = MagnitudePrunerConfig(
+            global_config=ModuleMagnitudePrunerConfig(
+                op_state_spec={"weight": PruningSpec(target_sparsity=0.6)},
+                # Updating every other step makes off-boundary steps reuse the previous sparsity.
+                sparsity_schedule=PolynomialDecaySchedule(
+                    begin_step=0, total_iters=10, power=3.0, update_frequency=2
+                ),
             )
-            p = MagnitudePruner(model, config)
-            p.prepare((torch.randn(1, 100),))
-            return p
-
-        model_orig = copy.deepcopy(linear_100x100_unique)
-        pruner_orig = _build(model_orig)
-        halfway = total_iters // 2
-        for _ in range(halfway):
-            pruner_orig.step()
-
-        state = model_orig.state_dict()
-        model_resumed = copy.deepcopy(linear_100x100_unique)
-        pruner_resumed = _build(model_resumed)
-        model_resumed.load_state_dict(state)
-
-        assert pruner_resumed._step_count == halfway
+        )
+        example_input = torch.randn(1, 100)
+        uninterrupted = copy.deepcopy(linear_100x100_unique)
+        pruner = MagnitudePruner(uninterrupted, config)
+        pruner.prepare((example_input,))
+        # Step 5 is off-boundary: its sparsity is the one from step 4, so restoring it
+        # requires replaying the schedule rather than evaluating step 5 alone.
+        for _ in range(5):
+            pruner.step()
+        saved = copy.deepcopy(uninterrupted.state_dict())
+
+        resumed = copy.deepcopy(linear_100x100_unique)
+        resumed_pruner = MagnitudePruner(resumed, config)
+        resumed_pruner.prepare((example_input,), state_dict=saved)
+        original_impl = uninterrupted.parametrizations["weight"][0]
+        resumed_impl = resumed.parametrizations["weight"][0]
+        assert resumed_pruner._step_count == 5
+        assert resumed_impl.sparsity == original_impl.sparsity
+
+        def _raise(*args, **kwargs):
+            raise AssertionError("the loaded mask must be used, not recomputed")
+
+        with monkeypatch.context() as patch:
+            patch.setattr(_MagnitudePruneImpl, "compute_mask", staticmethod(_raise))
+            resumed(example_input)
+
+        for _ in range(4):
+            pruner.step()
+            resumed_pruner.step()
+            assert resumed_impl.sparsity == original_impl.sparsity
+            assert torch.equal(resumed_impl.mask, original_impl.mask)
diff --git a/tests/test_base_model_compressor.py b/tests/test_base_model_compressor.py
index a38a49f8..f0d23efd 100644
--- a/tests/test_base_model_compressor.py
+++ b/tests/test_base_model_compressor.py
@@ -4,7 +4,9 @@
 # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
 
 import copy
+from collections.abc import Callable
 
+import pytest
 import torch
 import torch.nn as nn
 
@@ -12,6 +14,13 @@
     _COREAI_OPT_PREPARED_ATTR,
     _BaseModelCompressor,
 )
+from coreai_opt.palettization import KMeansPalettizer
+from coreai_opt.pruning import MagnitudePruner
+from coreai_opt.quantization import ModuleQuantizerConfig, Quantizer, QuantizerConfig
+from coreai_opt.quantization.spec import (
+    default_activation_quantization_spec,
+    default_weight_quantization_spec,
+)
 
 
 class SimpleModel(nn.Module):
@@ -79,3 +88,81 @@ def test_exported_graph_module_deepcopy(self):
         graph_module = exported.module()
         _BaseModelCompressor._mark_model_as_prepared(graph_module)
         assert _BaseModelCompressor._is_model_prepared(copy.deepcopy(graph_module))
+
+
+def _quantizer(execution_mode: str) -> Callable[[nn.Module], Quantizer]:
+    config = QuantizerConfig(
+        global_config=ModuleQuantizerConfig(
+            op_state_spec={"weight": default_weight_quantization_spec()},
+            op_input_spec={"*": default_activation_quantization_spec()},
+            op_output_spec=None,
+        ),
+        execution_mode=execution_mode,
+    )
+    return lambda model: Quantizer(model, config)
+
+
+_COMPRESSORS = [
+    pytest.param(KMeansPalettizer, id="palettizer"),
+    pytest.param(_quantizer("eager"), id="quantizer-eager"),
+    pytest.param(_quantizer("graph"), id="quantizer-graph"),
+    pytest.param(MagnitudePruner, id="pruner"),
+]
+
+
+@pytest.mark.parametrize("make_compressor", _COMPRESSORS)
+class TestPrepareFromStateDict:
+    """``prepare(state_dict=...)`` reloads a prepared model's compression state."""
+
+    @staticmethod
+    def _prepare_source(make_compressor, model, example_input):
+        """Prepare ``model`` and, for quantizers, calibrate it on scaled inputs."""
+        compressor = make_compressor(model)
+        prepared = compressor.prepare((example_input,))
+        if isinstance(compressor, Quantizer):
+            with compressor.calibration_mode():
+                prepared(example_input * 3.0)
+        return prepared
+
+    def test_reproduces_prepared_model(self, make_compressor, tmp_path):
+        """A model reloaded from a saved state dict reproduces the source model's output."""
+        torch.manual_seed(0)
+        model, example_input = SimpleModel(), torch.randn(4, 10)
+        source = self._prepare_source(make_compressor, copy.deepcopy(model), example_input)
+        expected = source(example_input)
+        path = tmp_path / "prepared.pt"
+        torch.save(source.state_dict(), path)
+
+        reloaded = make_compressor(copy.deepcopy(model)).prepare(
+            (example_input,), state_dict=torch.load(path)
+        )
+
+        torch.testing.assert_close(reloaded(example_input), expected, rtol=0, atol=0)
+
+    def test_flags_match_regular_prepare(self, make_compressor):
+        """Enable flags are left as a regular ``prepare`` leaves them, whatever was saved."""
+        torch.manual_seed(0)
+        model, example_input = SimpleModel(), torch.randn(4, 10)
+        regular = make_compressor(copy.deepcopy(model)).prepare((example_input,))
+        flipped = {
+            name: 1 - tensor if name.endswith("_enabled") else tensor
+            for name, tensor in regular.state_dict().items()
+        }
+
+        reloaded = make_compressor(copy.deepcopy(model)).prepare(
+            (example_input,), state_dict=flipped
+        )
+
+        for name, tensor in regular.state_dict().items():
+            if name.endswith("_enabled"):
+                assert torch.equal(reloaded.state_dict()[name], tensor), name
+
+    def test_mismatched_state_dict_raises(self, make_compressor):
+        """A state dict that does not match the prepared model raises a clear error."""
+        torch.manual_seed(0)
+        model, example_input = SimpleModel(), torch.randn(4, 10)
+        state_dict = make_compressor(copy.deepcopy(model)).prepare((example_input,)).state_dict()
+        state_dict.pop(next(iter(state_dict)))
+
+        with pytest.raises(RuntimeError, match="State dict does not match"):
+            make_compressor(copy.deepcopy(model)).prepare((example_input,), state_dict=state_dict)
diff --git a/tests/test_joint_compression.py b/tests/test_joint_compression.py
index 6bf73e53..92bba742 100644
--- a/tests/test_joint_compression.py
+++ b/tests/test_joint_compression.py
@@ -16,6 +16,7 @@
 from coreai_opt.palettization.spec.fake_palettize import _FakePalettizeImplBase
 from coreai_opt.quantization import Quantizer
 from tests.fixtures.compression import ParametrizedP4A8CompressionConfigs
+from tests.models.simple import SimpleModel
 from tests.test_utils.general import COREAI_AVAILABLE
 
 batch_size = 128
@@ -85,3 +86,34 @@ def test_p4a8_compression_mnist_accuracy(
         f"Post-calibration accuracy ({post_calib_acc:.2f}%) != "
         f"post-finalize accuracy ({finalized_acc:.2f}%)"
     )
+
+
+@pytest.mark.skipif(not COREAI_AVAILABLE, reason="Requires coreai")
+def test_p4a8_reload_each_stage_from_state_dict(
+    parametrized_p4a8_compression_config: ParametrizedP4A8CompressionConfigs,
+) -> None:
+    """Palettize, finalize to CoreAI, quantize activations and calibrate; then rebuild the
+    same model by reloading each stage from its saved state dict, without re-clustering
+    or re-calibrating."""
+    config = parametrized_p4a8_compression_config
+    torch.manual_seed(0)
+    model = SimpleModel().eval()
+    example_input = torch.randn(4, 1, 28, 28)
+
+    def compress(palett_state=None, quant_state=None):
+        palettizer = KMeansPalettizer(deepcopy(model), config.palett_config)
+        palett_prepared = palettizer.prepare((example_input,), state_dict=palett_state)
+        # Saved before finalize, which replaces the palettization parametrizations.
+        palett_saved = deepcopy(palett_prepared.state_dict())
+        palettized = palettizer.finalize(backend=ExportBackend.CoreAI)
+        quantizer = Quantizer(palettized, config.quant_config)
+        prepared = quantizer.prepare((example_input,), state_dict=quant_state)
+        if quant_state is None:
+            with quantizer.calibration_mode():
+                prepared(example_input * 3.0)
+        return prepared, palett_saved, deepcopy(prepared.state_dict())
+
+    original, palett_saved, quant_saved = compress()
+    reloaded, _, _ = compress(palett_saved, quant_saved)
+
+    torch.testing.assert_close(reloaded(example_input), original(example_input), rtol=0, atol=0)