Reload prepared models from a state dict for all compressors by NehalBhandari · Pull Request #157 · apple/coreai-optimization

NehalBhandari · 2026-10-22T23:42:32Z

Summary

Adds a generic way to reload a prepared model: prepare(..., state_dict=...) now works for every compressor (Quantizer in eager and graph mode, MagnitudePruner, KMeansPalettizer). A model can be prepared (and calibrated or trained) once, saved, and later reloaded to finalize for any backend, without recomputing its compression state.

Changes

_BaseModelCompressor: documents state_dict in the prepare contract and adds a shared _load_prepared_state_dict:
- Strict-loads the state dict; a mismatch re-raises with a hint (same config and model code, and in graph mode the same example inputs, dynamic shapes and torch version).
- Calls a new abstract _uninitialized_compression_modules hook and raises if any compression module is still uninitialized.
Quantizer (eager and graph): Runs prepare as is (since it is pretty quick), then loads the state dict. The costly steps, calibration and QAT, happen after prepare, and their results come back with the state dict. Enable flags are set after loading, so the model ends up as after a regular prepare (observers disabled, fake quantization enabled), even if the state dict was saved mid-calibration.
MagnitudePruner:
- With a state dict, prepare loads the saved masks instead of computing them. On Qwen3-4B this skips about 98 s of a 101 s prepare.
- Each pruning module persists a step_count buffer. On load, the scheduled sparsity is derived by replaying update_sparsity, so a model saved mid-schedule resumes from the saved step.
- step() now raises if called before prepare().
- The normal and reload paths share one prepare flow.
KMeansPalettizer: the existing prepare(state_dict=...) now goes through the shared base path.

Fixes

KMeansPalettizer.prepare(state_dict=...) failed for any config with a quantized LUT (size mismatch on the LUT quantizer's qparams). Skipping k-means also skipped initializing the LUT quantizer. _KMeansFakePalettize._load_from_state_dict now initializes it from the loaded centroids before its saved qparams load.
MagnitudePruner restarted its sparsity schedule from step 0 after reloading a model saved mid-schedule. This was the previously strict-xfail test_save_resume_midschedule, now covered by test_resumes_schedule_mid_way.

crowbat · 2026-10-26T20:28:08Z

+            self._remove_incompatible_fake_palett_modules()
+            self._load_prepared_state_dict(self._model, state_dict)
+
+        self._model.apply(_enable_fake_palett)


I think we should keep self._model.apply(_enable_fake_palett) inside the if case only. The responsibility of setting the buffers when state_dict is present should just automatically fall to the loading of the state dict. It also keeps things symmetrical where the earlier self._model.apply(_disable_fake_palett) only happened if state_dict is None.

I actually think we should set the flags during prepare no matter what the state_dict shows. A user could save a checkpoint mid calibration or training; observer flags are distinct in that they affect further calibration / training steps, so we should consistently return them in their documented state after prepare.

I suppose saving the state dict in the middle of the context managers could lead to unintended compressor states being saved since the context manager cleanup wouldn't have happened. I'm ok then with resetting to post-prepare.

Could we then add a note to palettizer's prepare state_dict docstring saying that it resets to post-prepare states? Similar to what we have for quantization/quantizer.py.

Sure will update the docstring

crowbat · 2026-10-26T21:00:39Z

        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))


This addition of step_count makes PruneImplBase's state dict incompatible with those from before this change. While it's super unlikely that anyone was saving state dicts before, I feel it's worth a mention in changelog saying we added a new buffer field to PruneImplBase.

We could also consider checking if step_count is missing from the state_dict during loading and default the count to 0 in that case, but that may be overkill.

Separately, do you plan to make similar step_count buffer additions to palettizers and quantizers too? Those would both benefit in then being able to easily pick up again if saved in the middle of QAT/PAT training

Yes, because this was an improbable scenario, I did not add the missing key check. I can include a note in the changelog regarding the incompatibility.

dengqiaoyu · 2026-10-27T00:54:56Z

            if self.mask.device != weight.device or self.mask.dtype != weight.dtype:
                self.mask = self.mask.to(device=weight.device, dtype=weight.dtype)


Since we are supporting _load_from_state_dict, should we always ensure that the mask is on the same device as the weight, rather than only doing so when self._dirty is set?

dengqiaoyu · 2026-10-27T01:05:25Z

+        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)


We can also consider adding a version checker to ensure compatibility:

_version = 2 step_key = prefix + "step_count" if local_metadata.get("version", 1) < 2 and step_key not in state_dict: # Checkpoints saved before step_count existed keep the current step. state_dict[step_key] = self.step_count

dengqiaoyu · 2026-10-27T01:16:50Z

        Returns:
            Model with fake quantizers inserted
        """
        if self._is_model_prepared(self._model):


Since we are overloading prepare here to load the state_dict, I think we should also prevent the prepared model from being prepare loaded with correct error message, which should not be a no-op.

Suggested change

if self._is_model_prepared(self._model):

if state_dict is not None:

raise RuntimeError(

"Model has already been prepared, so prepare() cannot load state_dict. "

"Pass the unprepared model, or call load_state_dict() on the prepared one."

)

dengqiaoyu · 2026-10-27T01:23:31Z

+
+        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:


Can we also add a test to verify that calling load_state_dict after prepare((example_input,)) still works for the same check we do below?

dengqiaoyu · 2026-10-27T01:33:46Z

+                qparams) is loaded after preparing.

        Returns:
            Model with fake quantizers inserted


I think we should add a new "RuntimeError" section for all three specific implementations:

RuntimeError: If ``state_dict`` does not match the prepared model

NehalBhandari added 2 commits October 22, 2026 16:40

Reload prepared models from a state dict for all compressors

a7d0253

Add changelog

f2d4c74

NehalBhandari requested review from crowbat, dengqiaoyu and u-simha and removed request for crowbat, dengqiaoyu and u-simha October 23, 2026 17:53

NehalBhandari assigned crowbat and u-simha Oct 23, 2026

crowbat reviewed Oct 26, 2026

View reviewed changes

Review comments

263bcfa

dengqiaoyu reviewed Oct 27, 2026

View reviewed changes

Navigation Menu

Repository navigation

Reload prepared models from a state dict for all compressors - #157

Reload prepared models from a state dict for all compressors#157
NehalBhandari wants to merge 3 commits into
apple:mainfrom
NehalBhandari:u/nehal_bhandari/prepared-state-dict-reload

NehalBhandari commented Oct 22, 2026 •

edited

Loading

Uh oh!

crowbat Oct 26, 2026

Uh oh!

NehalBhandari Oct 26, 2026 •

edited

Loading

Uh oh!

crowbat Oct 26, 2026

Uh oh!

NehalBhandari Oct 26, 2026

Uh oh!

Uh oh!

Uh oh!

crowbat Oct 26, 2026

Uh oh!

NehalBhandari Oct 26, 2026

Uh oh!

dengqiaoyu Oct 27, 2026

Uh oh!

dengqiaoyu Oct 27, 2026

Uh oh!

dengqiaoyu Oct 27, 2026

Uh oh!

dengqiaoyu Oct 27, 2026

Uh oh!

dengqiaoyu Oct 27, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		if self.mask.device != weight.device or self.mask.dtype != weight.dtype:
		self.mask = self.mask.to(device=weight.device, dtype=weight.dtype)

+        if self._is_model_prepared(self._model):
+            if state_dict is not None:
+                raise RuntimeError(
+                    "Model has already been prepared, so prepare() cannot load state_dict. "
+                    "Pass the unprepared model, or call load_state_dict() on the prepared one."
+                )


		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:

Conversation

NehalBhandari commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Fixes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NehalBhandari Oct 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

NehalBhandari commented Oct 22, 2026 •

edited

Loading

NehalBhandari Oct 26, 2026 •

edited

Loading