diff --git a/docs/src/palettization/overview.md b/docs/src/palettization/overview.md index a232745..617a125 100644 --- a/docs/src/palettization/overview.md +++ b/docs/src/palettization/overview.md @@ -94,6 +94,24 @@ Note that `Sensitive K-means` shares the same `KMeansPalettizer` and `KMeansPale For more details on how to use {class}`~coreai_opt.palettization.config.KMeansPalettizerConfig`, {class}`~coreai_opt.palettization.config.ModuleKMeansPalettizerConfig` to apply different settings to different weights in the model, see [Palettization Config](config.md). +## Loading a Prepared Model + +Preparing a model runs k-means clustering to compute each layer's codebook, which can be expensive for large models. If you have already prepared a model and saved its `state_dict`, reload it into a fresh palettizer with `prepare(..., state_dict=...)` to skip clustering: + +```python +# ---- first run: prepare and save ---- +palettizer = KMeansPalettizer(model, config) +prepared_model = palettizer.prepare(example_inputs) +torch.save(prepared_model.state_dict(), "palettized_state.pt") + +# ---- later run: reload without re-clustering ---- +state_dict = torch.load("palettized_state.pt", weights_only=True) +palettizer = KMeansPalettizer(MyModel().eval(), config) # same config as the saved run +prepared_model = palettizer.prepare(example_inputs, state_dict=state_dict) +``` + +The `config` and model architecture must match the run that produced the state dict; a mismatch raises an error rather than loading a wrong codebook. Sensitivities are not part of the state dict — reload them separately via `prepare(..., sensitivity_path=...)` if a later re-clustering step needs them. + ## Training a Palettized model A `KMeansPalettizer` palettized model can still be fine-tuned in a training pipeline. As palettization is a hard assignment lookup, gradients cannot be propagated for palettized weights, meaning any parameter which is palettized will not update during `optimizer.step()` (the palettization codebook and index assignments will also be fixed). diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index 3b73846..56dea1a 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py @@ -197,6 +197,24 @@ def ensure_initialized(self, tensor: torch.Tensor) -> None: ) self._disabled = True + def check_compatible(self, weight: torch.Tensor) -> bool: + """Return whether ``weight``'s shape is compatible with this module's spec. + + Runs the reshape/block prefix on a meta (zero-storage) copy of ``weight``: + no data is read, no scaling is applied, and nothing is allocated. + + Args: + weight (torch.Tensor): Weight tensor in its original shape. + + Returns: + bool: True if the shape is compatible with the configured spec. + """ + try: + self._reshape_and_block(weight.to("meta")) + except (_IncompatibleClusterDimError, _IncompatibleGranularityError): + return False + return True + def forward_enabled(self, tensor: torch.Tensor) -> torch.Tensor: if self.training: return self._training_strategy.train_forward(self, tensor) @@ -233,17 +251,25 @@ def _maybe_refresh_indices(self, weight: torch.Tensor) -> None: self.indices = self._assign_indices(weight, self.centroids).detach() self._indices_stale = False + def _save_to_state_dict(self, destination, prefix, keep_vars): + """Omit stale indices so they are recomputed from centroids on load.""" + super()._save_to_state_dict(destination, prefix, keep_vars) + if self._indices_stale: + destination.pop(prefix + "indices", None) + 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_key, centroids_key = prefix + "lut", prefix + "centroids" + 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: state_dict[centroids_key] = self._centroids_from_lut(state_dict[lut_key]) + had_indices = index_key in state_dict super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) if self.centroids is not None: self._centroids_initialized = True - self._indices_stale = True + 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) def _centroids_from_lut(self, lut: torch.Tensor) -> torch.Tensor: """Invert ``_reshape_lut_tensor`` to recover ``(num_blocks, num_clusters, @@ -389,6 +415,24 @@ def _scale_reshape_and_block(self, weight: torch.Tensor) -> tuple[list[torch.Ten """ if self.enable_per_channel_scale: weight = self._scale_by_per_channel_scale(weight) + return self._reshape_and_block(weight) + + def _reshape_and_block(self, weight: torch.Tensor) -> tuple[list[torch.Tensor], int]: + """Reshape to 2D and split into per-block tensors (no scaling). + + Shared by the clustering path and the compatibility probe. + + Args: + weight (torch.Tensor): Weight tensor in its original shape. + + Returns: + tuple[list[torch.Tensor], int]: The per-block 2D tensors and the + resolved palettization axis. + + Raises: + _IncompatibleGranularityError: If the shape is incompatible with the granularity. + _IncompatibleClusterDimError: If the shape is incompatible with cluster_dim. + """ axis = self._resolved_axis # Produce a 2d tensor with output channel axis remaining as is, and all other axes flattened @@ -945,6 +989,8 @@ def _unscale_by_per_channel_scale(self, scaled_weight: torch.Tensor) -> torch.Te per channel scales. """ flattened_scaled_weight = scaled_weight.flatten(1) - flattened_unscaled_weight = flattened_scaled_weight * self.per_channel_scale + flattened_unscaled_weight = flattened_scaled_weight * self.per_channel_scale.to( + flattened_scaled_weight.device + ) unscaled_weight = flattened_unscaled_weight.reshape(scaled_weight.shape) return unscaled_weight diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py index 2c75805..8c58ed1 100644 --- a/src/coreai_opt/palettization/kmeans/palettizer.py +++ b/src/coreai_opt/palettization/kmeans/palettizer.py @@ -154,6 +154,7 @@ def prepare( example_inputs: tuple[Any, ...], sensitivity_path: str | None = None, num_workers: int = 1, + state_dict: dict[str, torch.Tensor] | None = None, ) -> torch.nn.Module: """ Prepare the model for palettization. @@ -172,6 +173,9 @@ def prepare( across layers. It is recommended to use more than one worker process to parallelize the clustering, especially when multiple CPUs are available. 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. Returns: The prepared nn.Module with fake palettization @@ -208,18 +212,21 @@ def prepare( sensitivities = torch.load(sensitivity_path, weights_only=True) self._set_sensitivities_in_fake_palettize_modules(sensitivities) - self._model.apply(_disable_fake_palett) + if state_dict is None: + self._model.apply(_disable_fake_palett) - if self._num_workers > 1: - self._calculate_centroids_parallel(num_workers) - else: - self._calculate_centroids_sequential() + if self._num_workers > 1: + self._calculate_centroids_parallel(num_workers) + else: + self._calculate_centroids_sequential() - # 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) + # 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) + self._model.apply(_enable_fake_palett) + else: + self._load_prepared_state_dict(state_dict) # Mark the model as prepared to prevent re-preparation self._mark_model_as_prepared(prepared_model) @@ -229,6 +236,46 @@ 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. + + Raises: + RuntimeError: If a compatible module has no centroids after loading, + meaning the checkpoint does not match the palettizer config. + """ + # 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) + @contextmanager def calibration_mode( self, diff --git a/tests/palettization/test_kmeans_fake_palettize.py b/tests/palettization/test_kmeans_fake_palettize.py index 5fc08ab..a3709de 100644 --- a/tests/palettization/test_kmeans_fake_palettize.py +++ b/tests/palettization/test_kmeans_fake_palettize.py @@ -1860,21 +1860,57 @@ def test_hard_assign_skips_refresh_when_fresh(self): palettizer.hard_assign(weight) assert palettizer.indices is indices_before # untouched - def test_load_from_state_dict_marks_initialized_but_stale(self): - """Loading centroids from a checkpoint marks the module initialized AND - stale (unlike _initialize(), which leaves indices fresh) -- indices are - reconstructed lazily against the loaded centroids on next use. - """ + def test_load_from_state_dict_with_and_without_indices(self): + """Test _centroids_initialized and _indices_stale settings when loading a state dict with + and without indices.""" spec = PalettizationSpec(n_bits=2, granularity=PerTensorGranularity()) source = _KMeansFakePalettize(**spec.__dict__) source._initialize(torch.randn(8, 8)) + state_dict = source.state_dict() + # Loading state dict with indices sets _indices_stale to False target = _KMeansFakePalettize(**spec.__dict__) assert target._centroids_initialized is False - target.load_state_dict(source.state_dict()) + target.load_state_dict(state_dict) + assert target._centroids_initialized is True + assert target._indices_stale is False + + # Loading state dict without indices sets _indices_stale to True + del state_dict["indices"] + target = _KMeansFakePalettize(**spec.__dict__) + target.load_state_dict(state_dict) assert target._centroids_initialized is True assert target._indices_stale is True + def test_load_from_state_dict_missing_scale_is_stale(self): + """With per-channel scale enabled, loading indices but no scale marks stale + so both are recomputed together on next use. + """ + spec = PalettizationSpec( + n_bits=2, granularity=PerTensorGranularity(), enable_per_channel_scale=True + ) + source = _KMeansFakePalettize(**spec.__dict__) + source._initialize(torch.randn(8, 8)) + state_dict = source.state_dict() + del state_dict["per_channel_scale"] + + target = _KMeansFakePalettize(**spec.__dict__) + target.load_state_dict(state_dict) + assert target._centroids_initialized is True + assert target._indices_stale is True + + def test_stale_indices_are_not_saved(self): + """Stale indices are omitted from the state dict and recomputed on load.""" + spec = PalettizationSpec(n_bits=2, granularity=PerTensorGranularity()) + source = _KMeansFakePalettize(**spec.__dict__) + source._initialize(torch.randn(8, 8)) + source._indices_stale = True + assert "indices" not in source.state_dict() + + target = _KMeansFakePalettize(**spec.__dict__) + target.load_state_dict(source.state_dict()) + assert target._indices_stale is True + class TestReinitializeOnEnable: """enable_fake_palett(reinitialize=True) invalidates cached params on a @@ -2219,3 +2255,15 @@ def test_from_cluster_vectors_rejects_element_count_mismatch(): vectors = torch.randn(1, 5, 2) # 5 * 2 = 10 elements, but rows * cols = 8 with pytest.raises(ValueError): palettizer._from_cluster_vectors(vectors, rows=4, cols=2) + + +@pytest.mark.parametrize("axis", [0, 1]) +@pytest.mark.parametrize("group_size", [3, 4]) +def test_check_compatible(axis, group_size): + """Test check_compatible for a variety of axes and group sizes.""" + tensor = torch.randn(8, 8) + spec = PalettizationSpec( + n_bits=2, granularity=PerGroupedChannelGranularity(axis=axis, group_size=group_size) + ) + palettizer = _KMeansFakePalettize(**spec.__dict__) + assert palettizer.check_compatible(torch.randn(8, 8)) is (tensor.shape[axis] % group_size == 0) diff --git a/tests/palettization/test_kmeans_palettizer.py b/tests/palettization/test_kmeans_palettizer.py index 3d6a185..2c218c8 100644 --- a/tests/palettization/test_kmeans_palettizer.py +++ b/tests/palettization/test_kmeans_palettizer.py @@ -1646,3 +1646,52 @@ def test_palettize_multihead_attention( output = prepared_model(simple_mha_model_input) assert output.shape == (1, 10, 64) + + +class TestKMeansPalettizerLoadFromStateDict: + """prepare(state_dict=...) loads a prepared model's buffers instead of clustering.""" + + def test_reproduces_prepared_model( + self, simple_conv_linear_model, simple_model_input, basic_config, temp_dir + ): + """A skip-loaded model reproduces the source model's forward output exactly.""" + src = KMeansPalettizer(copy.deepcopy(simple_conv_linear_model), basic_config) + prepared = src.prepare((simple_model_input,)) + state_dict_path = os.path.join(temp_dir, "palettizer_state_dict.pt") + torch.save(prepared.state_dict(), state_dict_path) + with torch.no_grad(): + expected = prepared(simple_model_input) + + dst = KMeansPalettizer(copy.deepcopy(simple_conv_linear_model), basic_config) + state_dict = torch.load(state_dict_path) + loaded = dst.prepare((simple_model_input,), state_dict=state_dict) + with torch.no_grad(): + got = loaded(simple_model_input) + + assert torch.equal(expected, got) + + def test_skips_clustering( + self, monkeypatch, simple_conv_linear_model, simple_model_input, basic_config + ): + """Loading from a state dict must not run k-means clustering.""" + src = KMeansPalettizer(copy.deepcopy(simple_conv_linear_model), basic_config) + state_dict = src.prepare((simple_model_input,)).state_dict() + + def _raise_assert(self, *args, **kwargs): + raise AssertionError("clustering must not run when loading from a state dict") + + monkeypatch.setattr(_KMeansFakePalettize, "_cluster_to_centroids", _raise_assert) + dst = KMeansPalettizer(copy.deepcopy(simple_conv_linear_model), basic_config) + dst.prepare((simple_model_input,), state_dict=state_dict) + + def test_missing_centroids_raises( + self, simple_conv_linear_model, simple_model_input, basic_config + ): + """Test that a state dict with missing centroids raises an error.""" + src = KMeansPalettizer(copy.deepcopy(simple_conv_linear_model), basic_config) + state_dict = src.prepare((simple_model_input,)).state_dict() + 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"): + dst.prepare((simple_model_input,), state_dict=pruned)