Skip to content

Re-cluster palettization centroids from the warmed-up weights when a PATSchedule enables fake palettization mid-training - #86

Merged
crowbat merged 3 commits into
apple:mainfrom
crowbat:u/k_hsieh/recompute_centroids_when_enabling_fake_palettize_for_pat_schedule
Sep 4, 2026
Merged

Re-cluster palettization centroids from the warmed-up weights when a PATSchedule enables fake palettization mid-training#86
crowbat merged 3 commits into
apple:mainfrom
crowbat:u/k_hsieh/recompute_centroids_when_enabling_fake_palettize_for_pat_schedule

Conversation

@crowbat

@crowbat crowbat commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

Currently if enable_fake_palettize > 0 in PATSchedule, centroids are initialized at the start and reused when the module is later enabled via PATSchedule and palettizer.step(). Instead, recompute centroids based on trained updated weights that changed during step < enable_fake_palettize.

new_state = schedule._compute_state(self._step_count)
# On the warm-up -> enabled transition, re-cluster from the current
# (warmed-up) weights instead of the frozen prepare-time centroids.
if new_state and not fp_module.fake_palett_enabled[0]:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Have you considered pushing the call to mark the centroids stale down into enable_fake_palett? That way it's not tied exclusively to the scheduler; if the user directly enables the fake palettizer we can trigger a centroid update.

@crowbat
crowbat force-pushed the u/k_hsieh/recompute_centroids_when_enabling_fake_palettize_for_pat_schedule branch from 83b3b6b to 5557765 Compare September 3, 2026 21:05
@crowbat
crowbat merged commit f6ba0ee into apple:main Sep 4, 2026
13 checks passed
@crowbat
crowbat deleted the u/k_hsieh/recompute_centroids_when_enabling_fake_palettize_for_pat_schedule branch September 4, 2026 00:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants