Refactor KMeansPalettizer, add extendable training strategy capability by crowbat · Pull Request #60 · apple/coreai-optimization

crowbat · 2026-07-29T22:50:18Z

No description provided.

crowbat · 2026-08-04T22:04:10Z

+        raise NotImplementedError
+
+
+@TrainingStrategy.register("default")


Now that we have TrainingStrategyConfig being the main object users use when configuring PalettizationSpec, associating TrainingStrategy with a registry is not needed anymore (this was based on a previous design). I will remove the registry association.

u-simha

Mostly looks good, have left some comments.

One thing to note, I think this might be backward-incompatible for an older palettization checkpoint, since we have added / modified buffers?

u-simha · 2026-08-04T22:55:29Z

+    def _refresh_indices(self, weight: torch.Tensor) -> None:
+        """Recompute indices from the current centroids, without re-clustering."""
+        self.indices = self._assign_indices(weight, self.centroids).detach()
+        self._indices_stale = False


Should we re-assign only if the indices are stale? And have a flag to force re-assignment

self._indices_stale is meant to be the flag for forcing re-assignment which is checked in the forward pass during hard_assign. _refresh_indices is only ever called there if self._indices_stale is True, so I think this covers your concern?

u-simha · 2026-08-04T23:13:01Z

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


Nit: lut -> fq_lut

u-simha · 2026-08-04T23:14:30Z

-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        weight = original_weights.cpu()
+    @property
+    def quantized_lut(self) -> torch.Tensor | None:


Minor suggestion: this can be reused in the lut function with a fq flag

This part, along with lut_quantization_scale and lut_quantization_zero_point were made into properties to try to preserve some backwards comaptibleness with the previous implementation which stored all 3 as separate buffers. Though there is still incompatibility when it comes to being able to save and load actual buffer values so perhaps we can consider deprecating them instead

u-simha · 2026-08-04T23:15:10Z

-        if self.enable_per_channel_scale:
-            weight = self._scale_by_per_channel_scale(weight)
+    @property
+    def lut_quantization_scale(self) -> torch.Tensor | None:


Can we have lut_quantization_qparams and return both the scale & zero point together? Any reason this is separated?

u-simha · 2026-08-04T23:19:40Z

+        with palettizer.training_mode():
+            prepared(torch.randn(2, 16)).sum().backward()
+
+        assert prepared.palettized.parametrizations.weight.original.grad is None
+        assert prepared.head.weight.grad is not None


Can we have a bit more extensive test that tests that we don't break gradients within and outside of the training mode context:

coreai-optimization/tests/quantization/test_fake_quantize.py

Line 228 in 7ad2df3

def test_qat_gradient_flow(dtype, granularity, qformulation, compression_target_tensor):

I saw a couple places we call .detach() which sometimes breaks the gradients if it is part of the backward pass (or can short circuit and have the gradients flowing through unintended variables - this happened in quantization at some point)

Sounds good, I'll add some coverage for this too

u-simha · 2026-08-06T01:48:32Z

+                        self._fp_to_schedule[param] = schedule
+                    break
+
+    def _resolve_schedule(self, module_name: str) -> PATSchedule | None:


Would this follow the module priority while applying the PAT schedule? I think I should use this for the QAT schedule too

This part of the code is pretty much a mirror of what we have in quantizer: https://github.com/apple/coreai-optimization/blob/main/src/coreai_opt/quantization/quantizer.py#L184

so it may inherently have the same issue as QAT

u-simha · 2026-08-06T01:49:40Z

+
+    def _apply_schedule(self) -> None:
+        for fp_module, schedule in self._fp_to_schedule.items():
+            fp_module.enable_fake_palett(schedule._compute_state(self._step_count))


Would it be better to do fp_module.apply(enable_fake_palett) that way even children modules have the schedule applied?

Given that fp_to_schedule is supposed to carry only leaf level modules mapping to schedules, it may be better to keep it as is to ensure we are setting the schedule for exactly the module we intend and no more accidentally

u-simha · 2026-08-06T01:52:20Z


+    @abstractmethod
    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Apply fake palettization to input tensor"""


I don't fully follow the reason for this logic to be moved to the downstream class. Is it because of the training strategy? I would think that would be generic across all fake palletize, and not specific to KMeans

At first, due to the removal of observer_enabled, most of this function simply got deleted and what was left seemed to get overridden anyways by much of the newly added things in kmeans_fake_palettizer, so it felt like this function wasn't lifting any weight.

But I think we can reframe some of what is in kmeans_fake_palettizer to be generic. I'll try some alternatives to see

u-simha · 2026-08-06T01:55:07Z

-                self._model.apply(_enable_observer)
+    @contextmanager
+    def training_mode(self):
+        """Context manager wrapping a training loop. Mutually exclusive with


We should check that the model is prepared, similar to what we do for quantization

u-simha · 2026-08-06T01:58:12Z

+        self._mode = "training"
+        try:
+            self._model.train()
+            self._build_fp_to_schedule()


Can we cache the fp_to_schedule since it shouldn't change across entering the training mode context?

I would say a common way to call this code would be to enter it on every batch step / epoch step and exit it while doing eval

Internally _build_fp_to_schedule() does check whether self._fp_to_schedule exists and returns that if it exists already. I can rename the function to _get_fp_to_schedule() to be less misleading

crowbat

Thanks for your comments @u-simha . I think on the backwards compatibility let's talk more about the best way to handle this since a good amount has changed.

crowbat · 2026-08-06T22:30:04Z

+    def _refresh_indices(self, weight: torch.Tensor) -> None:
+        """Recompute indices from the current centroids, without re-clustering."""
+        self.indices = self._assign_indices(weight, self.centroids).detach()
+        self._indices_stale = False


self._indices_stale is meant to be the flag for forcing re-assignment which is checked in the forward pass during hard_assign. _refresh_indices is only ever called there if self._indices_stale is True, so I think this covers your concern?

crowbat · 2026-08-06T22:32:17Z

+        self._mode = "training"
+        try:
+            self._model.train()
+            self._build_fp_to_schedule()


Internally _build_fp_to_schedule() does check whether self._fp_to_schedule exists and returns that if it exists already. I can rename the function to _get_fp_to_schedule() to be less misleading

crowbat · 2026-08-06T23:53:08Z


+    @abstractmethod
    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Apply fake palettization to input tensor"""


At first, due to the removal of observer_enabled, most of this function simply got deleted and what was left seemed to get overridden anyways by much of the newly added things in kmeans_fake_palettizer, so it felt like this function wasn't lifting any weight.

But I think we can reframe some of what is in kmeans_fake_palettizer to be generic. I'll try some alternatives to see

crowbat · 2026-08-07T00:23:23Z

-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        weight = original_weights.cpu()
+    @property
+    def quantized_lut(self) -> torch.Tensor | None:


This part, along with lut_quantization_scale and lut_quantization_zero_point were made into properties to try to preserve some backwards comaptibleness with the previous implementation which stored all 3 as separate buffers. Though there is still incompatibility when it comes to being able to save and load actual buffer values so perhaps we can consider deprecating them instead

crowbat · 2026-08-07T00:51:07Z


        self._num_workers = 1

+        self._mode: str = "idle"  # "idle" | "training" | "calibrating"


I'll define it in BaseModelCompressor but not add it to pruning or quantizers yet (that can come separately)

crowbat · 2026-08-07T01:07:36Z

+                        self._fp_to_schedule[param] = schedule
+                    break
+
+    def _resolve_schedule(self, module_name: str) -> PATSchedule | None:


This part of the code is pretty much a mirror of what we have in quantizer: https://github.com/apple/coreai-optimization/blob/main/src/coreai_opt/quantization/quantizer.py#L184

so it may inherently have the same issue as QAT

crowbat · 2026-08-07T01:14:56Z

+
+    def _apply_schedule(self) -> None:
+        for fp_module, schedule in self._fp_to_schedule.items():
+            fp_module.enable_fake_palett(schedule._compute_state(self._step_count))


Given that fp_to_schedule is supposed to carry only leaf level modules mapping to schedules, it may be better to keep it as is to ensure we are setting the schedule for exactly the module we intend and no more accidentally

crowbat · 2026-08-07T01:21:46Z

+        with palettizer.training_mode():
+            prepared(torch.randn(2, 16)).sum().backward()
+
+        assert prepared.palettized.parametrizations.weight.original.grad is None
+        assert prepared.head.weight.grad is not None


Sounds good, I'll add some coverage for this too

crowbat requested review from anotheranshu, dengqiaoyu, pkmandke and u-simha July 29, 2026 22:50

crowbat assigned u-simha Jul 29, 2026

crowbat force-pushed the u/k_hsieh/update_kmeans_palettization branch from 67ade25 to 2a9325e Compare July 29, 2026 22:56

crowbat commented Aug 4, 2026

View reviewed changes

crowbat force-pushed the u/k_hsieh/update_kmeans_palettization branch from 2a9325e to 44fbc11 Compare August 4, 2026 22:15

u-simha reviewed Aug 6, 2026

View reviewed changes

crowbat commented Aug 7, 2026

View reviewed changes

crowbat force-pushed the u/k_hsieh/update_kmeans_palettization branch from b44e145 to 734573b Compare August 7, 2026 01:55

crowbat added 3 commits August 7, 2026 15:44

Refactor KMeansPalettizer, add extendable training strategy capability

025c0e1

Address review comments

1d00647

Restore compatibility with legacy palettizer checkpoints

a5d23ad

crowbat force-pushed the u/k_hsieh/update_kmeans_palettization branch from 734573b to a5d23ad Compare August 7, 2026 23:47

		raise NotImplementedError


		@TrainingStrategy.register("default")


		self._num_workers = 1

		self._mode: str = "idle" # "idle" \| "training" \| "calibrating"

Uh oh!

Conversation

crowbat commented Jul 29, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

u-simha left a comment

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!

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!

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!

crowbat left a comment

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants