From 025c0e13887b04feae1154908b68732aa66d8373 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:44:57 -0700 Subject: [PATCH 1/4] Refactor KMeansPalettizer, add extendable training strategy capability --- changelog.d/60.added | 1 + src/coreai_opt/base_model_compressor.py | 11 + src/coreai_opt/config/compression_config.py | 16 +- .../palettization/config/__init__.py | 2 + .../config/palettization_config.py | 37 +- .../kmeans/kmeans_fake_palettize.py | 345 +++++++---- .../palettization/kmeans/palettizer.py | 198 ++++--- src/coreai_opt/palettization/spec/__init__.py | 4 + .../palettization/spec/fake_palettize.py | 125 +--- src/coreai_opt/palettization/spec/spec.py | 13 + .../palettization/spec/training_strategy.py | 101 ++++ .../test_kmeans_fake_palettize.py | 547 +++++++++++++++--- tests/palettization/test_kmeans_palettizer.py | 85 +-- tests/palettization/test_kmeans_parallel.py | 8 +- tests/palettization/test_pat_schedule.py | 350 +++++++++++ tests/test_compression_config.py | 6 +- 16 files changed, 1393 insertions(+), 456 deletions(-) create mode 100644 changelog.d/60.added create mode 100644 src/coreai_opt/palettization/spec/training_strategy.py create mode 100644 tests/palettization/test_pat_schedule.py diff --git a/changelog.d/60.added b/changelog.d/60.added new file mode 100644 index 0000000..4238628 --- /dev/null +++ b/changelog.d/60.added @@ -0,0 +1 @@ +Add training-aware palettization to `KMeansPalettizer` via a `training_mode()` context manager and `step()` method, with a `PATSchedule` (`pat_schedule` on the module config) controlling when palettization activates during a training loop. `PalettizationSpec` gains a `training_strategy_config` field for selecting pluggable training-time strategies. diff --git a/src/coreai_opt/base_model_compressor.py b/src/coreai_opt/base_model_compressor.py index 0c084e1..8436257 100644 --- a/src/coreai_opt/base_model_compressor.py +++ b/src/coreai_opt/base_model_compressor.py @@ -11,6 +11,7 @@ from abc import ABC, abstractmethod from contextlib import contextmanager +from enum import Enum from os import PathLike import torch @@ -21,12 +22,21 @@ _COREAI_OPT_PREPARED_ATTR = "_coreai_opt_prepared" +class _CompressorLifecycle(Enum): + """Lifecycle state of a model compressor.""" + + IDLE = "idle" + TRAINING = "training" + CALIBRATING = "calibrating" + + class _BaseModelCompressor(ABC): """ An abstract base class for implementing model compression techniques. """ _supported_modules: tuple[type[torch.nn.Module]] + _lifecycle: _CompressorLifecycle def __init__(self, model: torch.nn.Module, config: CompressionConfig | None = None): """ @@ -39,6 +49,7 @@ def __init__(self, model: torch.nn.Module, config: CompressionConfig | None = No """ self._model = model self._config = config + self._lifecycle = _CompressorLifecycle.IDLE @staticmethod def _is_model_prepared(model: torch.nn.Module) -> bool: diff --git a/src/coreai_opt/config/compression_config.py b/src/coreai_opt/config/compression_config.py index e8629b3..14a2ab4 100644 --- a/src/coreai_opt/config/compression_config.py +++ b/src/coreai_opt/config/compression_config.py @@ -379,21 +379,11 @@ def _normalize_none_op_configs(self) -> ModuleCompressionConfig: return self - def _get_compressor_specific_settings(self) -> dict[str, Any]: - """ - Get compressor-specific settings, excluding base ModuleCompressionConfig fields. - - Returns only fields defined by the concrete subclass (e.g., - enable_fast_kmeans_mode, rounding_precision for palettization), not the base - spec and config fields (op_input_spec, op_output_spec, op_state_spec, - op_type_config, op_name_config, module_input_spec, module_output_spec, - module_state_spec). - - This is useful when constructing arguments for compression operations that need - the compression-specific settings but handle the spec fields separately. + def _get_fake_module_kwargs(self) -> dict[str, Any]: + """Get the settings to forward to the compression simulator module's constructor. Returns: - Dictionary of compressor-specific field names to their values. + Dictionary of field names to their values. """ base_field_names = set(ModuleCompressionConfig.model_fields.keys()) return {k: v for k, v in self.model_dump().items() if k not in base_field_names} diff --git a/src/coreai_opt/palettization/config/__init__.py b/src/coreai_opt/palettization/config/__init__.py index 1161cc3..73b845d 100644 --- a/src/coreai_opt/palettization/config/__init__.py +++ b/src/coreai_opt/palettization/config/__init__.py @@ -9,10 +9,12 @@ KMeansPalettizerConfig, ModuleKMeansPalettizerConfig, OpKMeansPalettizerConfig, + PATSchedule, ) __all__ = [ "KMeansPalettizerConfig", "ModuleKMeansPalettizerConfig", "OpKMeansPalettizerConfig", + "PATSchedule", ] diff --git a/src/coreai_opt/palettization/config/palettization_config.py b/src/coreai_opt/palettization/config/palettization_config.py index 083a6e9..f5c1c56 100644 --- a/src/coreai_opt/palettization/config/palettization_config.py +++ b/src/coreai_opt/palettization/config/palettization_config.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, ClassVar, final -from pydantic import PositiveInt, model_validator +from pydantic import BaseModel, ConfigDict, Field, PositiveInt, model_validator from coreai_opt.config import ( CompressionConfig, @@ -31,6 +31,29 @@ _PALETTIZATION_SPEC = "palettization_spec" +class PATSchedule(BaseModel): + """Schedule for enabling palettization-aware training (PAT). + + Defines the step threshold at which a module's fake palettization + forward pass becomes active. Used with ``KMeansPalettizer.step()``. + + Attributes: + enable_fake_palettize: Step count at which fake palettization is + enabled. Must be >= 0. + + Example: + >>> schedule = PATSchedule(enable_fake_palettize=500) + """ + + model_config = ConfigDict(frozen=True) + + enable_fake_palettize: int = Field(default=0, ge=0) + + def _compute_state(self, step_count: int) -> bool: + """Return whether fake palettization should be active at the given step.""" + return step_count >= self.enable_fake_palettize + + class OpKMeansPalettizerConfig(WeightOnlyOpValidationMixin, OpCompressionConfig[PalettizationSpec]): """ Configuration class for palettization at the operation level. @@ -140,6 +163,11 @@ class ModuleKMeansPalettizerConfig( K-means clustering. Higher values preserve more precision but may reduce speed benefits. Only used when enable_fast_kmeans_mode is True. Default: 4. + pat_schedule (PATSchedule | None): Schedule controlling when this + module's palettization is active during a training_mode() loop. + If None, palettization is active immediately once training_mode() + begins. Default: None. + Example: >>> config = ModuleKMeansPalettizerConfig() # Uses defaults >>> # Or with custom settings: @@ -160,6 +188,7 @@ def __init_subclass__(cls, **kwargs): enable_fast_kmeans_mode: bool = True rounding_precision: PositiveInt = 4 + pat_schedule: PATSchedule | None = None # Namespace exposing built-in preset constructors. presets: ClassVar[_ModuleKMeansPalettizerConfigPresets] @@ -183,6 +212,12 @@ def validate_fast_kmeans_cluster_dim_constraint( return self + def _get_fake_module_kwargs(self) -> dict: + """Exclude palettizer-only fields from the fake-palettize constructor args.""" + base = super()._get_fake_module_kwargs() + base.pop("pat_schedule", None) + return base + @final class KMeansPalettizerConfig(CompressionConfig[ModuleKMeansPalettizerConfig]): diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index 478f7aa..c02bc1f 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py @@ -16,8 +16,15 @@ PerGroupedChannelGranularity, PerTensorGranularity, ) -from coreai_opt.palettization.spec.errors import _IncompatibleClusterDimError +from coreai_opt.palettization.spec.errors import ( + _IncompatibleClusterDimError, + _IncompatibleGranularityError, +) from coreai_opt.palettization.spec.fake_palettize import _FakePalettizeImplBase +from coreai_opt.palettization.spec.training_strategy import ( + DefaultTrainingConfig, + TrainingStrategyConfig, +) from coreai_opt.quantization.spec import ( PerChannelGranularity as _QuantPerChannelGranularity, QuantizationComponentFactory, @@ -45,7 +52,8 @@ class _KMeansFakePalettize(_FakePalettizeImplBase): The workflow proceeds in two steps: - 1. ``_calculate_centroids()``: Clusters weights using k-means, returns LUT and indices + 1. ``_initialize()``: Clusters weights using k-means and computes the + resulting LUT and indices 2. ``_palettize()``: Reconstructs palettized weights from LUT and indices Example: @@ -58,8 +66,8 @@ class _KMeansFakePalettize(_FakePalettizeImplBase): ... ) >>> palettizer = _KMeansFakePalettize(**spec.__dict__) >>> weight = torch.randn(4, 4) - >>> lut, indices = palettizer._calculate_centroids(weight) - >>> palettized_weight = palettizer._palettize(lut, indices, weight) + >>> palettizer._initialize(weight) + >>> palettized_weight = palettizer._palettize(palettizer.lut, palettizer.indices, weight) """ def __init__( @@ -73,7 +81,7 @@ def __init__( enable_fast_kmeans_mode: bool = True, rounding_precision: int = 4, op_to_optimize: Callable | None = None, - **kwargs, + training_strategy_config: TrainingStrategyConfig | None = None, ): super().__init__( n_bits=n_bits, @@ -81,13 +89,11 @@ def __init__( granularity=granularity, cluster_dim=cluster_dim, enable_per_channel_scale=enable_per_channel_scale, - **kwargs, ) self.enable_fast_kmeans_mode = enable_fast_kmeans_mode self.rounding_precision = rounding_precision self._sensitivities = sensitivities - self._centroids_stale = False # Create LUT fake quantizer if LUT quantization is enabled. # Use PerChannelGranularity(axis=0) so the stacked LUT tensor @@ -126,6 +132,15 @@ def __init__( update={"axis": self.reshape_strategy.default_axis} ) + self.register_buffer("centroids", None) + self._centroids_initialized: bool = False + self._indices_stale: bool = True + + # Resolve and construct the training strategy from its paired config. + if training_strategy_config is None: + training_strategy_config = DefaultTrainingConfig() + self._training_strategy = training_strategy_config.build_strategy() + @property def sensitivities(self) -> torch.Tensor | None: """Get the sensitivity values used for weighted k-means clustering.""" @@ -133,71 +148,173 @@ def sensitivities(self) -> torch.Tensor | None: @sensitivities.setter def sensitivities(self, value: torch.Tensor | None) -> None: - """Set sensitivity values and mark centroids as stale. - - When sensitivities are updated, the LUT and indices become stale and must - be recomputed via _calculate_centroids before use. This is typically done - by enabling the observer and running a forward pass through the model. + """Set sensitivity values. Centroids and indices are recomputed + on the next forward call. """ self._sensitivities = value - self._centroids_stale = True - - def forward(self, tensor: torch.Tensor) -> torch.Tensor: + self._centroids_initialized = False + self._indices_stale = True + + def ensure_initialized(self, tensor: torch.Tensor) -> None: + """Cluster centroids on first use; disable on an incompatible tensor.""" + if self._centroids_initialized: + return + try: + self._initialize(tensor.detach()) + except (_IncompatibleClusterDimError, _IncompatibleGranularityError) as e: + logger.warning( + f"Tensor incompatible with configured spec: {e}. Skipping palettization." + ) + self._disabled = True + + def forward_enabled(self, tensor: torch.Tensor) -> torch.Tensor: + if self.training: + return self._training_strategy.train_forward(self, tensor) + return self.hard_assign(tensor) + + def _initialize(self, weight: torch.Tensor) -> None: + """(Re-)initialize centroids and indices from k-means. + + Also seeds the LUT quantizer's observed qparams from these initial + centroids (see ``quantize_lut()``) — without this, a strategy that + never otherwise calls ``quantize_lut()`` (e.g. the default one-shot + strategy) would leave it permanently unobserved, and ``lut`` (which + reads frozen qparams via ``get_qparams()``, never observes) would + have nothing valid to read. + + Use ``_refresh_indices()`` instead to recompute indices from + already-updated centroids, without re-clustering. ``lut`` is never + stored — it's derived fresh from ``centroids`` on every access. """ - Apply fake palettization to input tensor. + self.centroids, indices = self._cluster_to_centroids(weight, self._sensitivities) + self.indices = indices.detach() + + # Run self.quantize_lut to seed lut quantizer qparams + self.quantize_lut(self._raw_lut(self.centroids)) - Overrides base class to add stale centroids warning when sensitivities - have been updated but centroids have not been recomputed. + self._centroids_initialized = True + self._indices_stale = False + + 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 + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Mark centroids as initialized after loading them from a checkpoint.""" + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + if self.centroids is not None: + self._centroids_initialized = True + self._indices_stale = True + + @property + def lut(self) -> torch.Tensor | None: + """Lookup table dequantized from ``centroids`` using the LUT + quantizer's frozen qparams (read via ``get_qparams()``, never + re-observed here — see ``quantize_lut()`` for the training-time + observing path). Cheap (independent of weight size), so never cached. """ - # Check for stale centroids when observer is disabled and we have - # initialized LUT/indices that would be used - if ( - self._centroids_stale - and self._initialized - and self.observer_enabled[0] == 0 - and self.fake_palett_enabled[0] == 1 - ): - logger.warning( - "Sensitivities were updated but centroids have not been recomputed. " - "The current LUT and indices do not reflect the new sensitivity values." - "Enable observer and run a forward pass to recompute centroids, or use " - "calibration_mode() which handles this automatically." - ) + if self.centroids is None: + return None - return super().forward(tensor) + raw_lut = self._raw_lut(self.centroids) + if self._lut_fake_quantizer is None: + return self._reshape_lut_tensor(raw_lut) - @torch.no_grad() - def _calculate_centroids( - self, original_weights: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - weight = original_weights.cpu() + orig_dtype = raw_lut.dtype + scale, zero_point, minval = self._lut_fake_quantizer.qparams_calculator.get_qparams() + fq_lut = self._lut_fake_quantizer._fused_fake_quant_dequant( + raw_lut.to(torch.float32), scale, zero_point, minval + ).to(orig_dtype) + return self._reshape_lut_tensor(fq_lut) - if self.enable_per_channel_scale: - weight = self._scale_by_per_channel_scale(weight) + @property + def quantized_lut(self) -> torch.Tensor | None: + """LUT quantized against ``lut_qspec``, derived fresh from + ``centroids``. ``None`` if no LUT quantizer is configured. + + Reads qparams via ``get_qparams()`` (a pure buffer read) rather than + calling the qparams calculator directly — some calculators (e.g. + moving-average) mutate running statistics on every call, so calling + them again here would drift the observer relative to ``lut``'s own + computation and make repeated reads mutually inconsistent. This + reflects the calculator's state as of the last access to ``lut`` + (which does call it, to legitimately observe the current centroids). + """ + if self.centroids is None or self._lut_fake_quantizer is None: + return None + raw_lut = self._raw_lut(self.centroids) + scale, zero_point, minval = self._lut_fake_quantizer.qparams_calculator.get_qparams() + quantized = self._lut_fake_quantizer.quantize(raw_lut, scale, zero_point, minval) + return self._reshape_lut_tensor(quantized.detach()) - # Reshape weight into 2D matrix for clustering - axis = self.granularity.axis if self.granularity.axis else 0 - weight = self.reshape_strategy.reshape_for_kmeans(weight, axis) + @property + def lut_quantization_scale(self) -> torch.Tensor | None: + """Quantization scale for ``quantized_lut``. ``None`` if no LUT + quantizer is configured. See ``quantized_lut`` for why this reads + ``get_qparams()`` instead of invoking the calculator directly. + """ + if self.centroids is None or self._lut_fake_quantizer is None: + return None + scale, _, _ = self._lut_fake_quantizer.qparams_calculator.get_qparams() + return self._reshape_lut_tensor(scale.detach()) + + @property + def lut_quantization_zero_point(self) -> torch.Tensor | None: + """Quantization zero point for ``quantized_lut``. ``None`` if no LUT + quantizer is configured or the quantization scheme has no zero + point. See ``quantized_lut`` for why this reads ``get_qparams()`` + instead of invoking the calculator directly. + """ + if self.centroids is None or self._lut_fake_quantizer is None: + return None + _, zero_point, _ = self._lut_fake_quantizer.qparams_calculator.get_qparams() + return self._reshape_lut_tensor(zero_point.detach()) if zero_point is not None else None + + def hard_assign(self, weight: torch.Tensor) -> torch.Tensor: + """Nearest-centroid reconstruction against the current centroids, + refreshing indices first if stale. + """ + if self._indices_stale: + self._refresh_indices(weight) + return self._palettize(self.lut, self.indices, weight) - # Validate cluster_dim divisibility along output channel axis (axis 0). + def _blocks_to_cluster(self, weight_2d: torch.Tensor, axis: int) -> list[torch.Tensor]: + """Validate cluster_dim divisibility and split a 2D weight/sensitivity + tensor into per-partition blocks. + """ if self.cluster_dim > 1: - # For per-grouped-channel axis=0, each block has group_size rows, - # so we must check group_size divisibility, not the full weight dim. if isinstance(self.granularity, PerGroupedChannelGranularity) and axis == 0: weight_dim = self.granularity.group_size else: - weight_dim = weight.shape[0] + weight_dim = weight_2d.shape[0] if weight_dim % self.cluster_dim != 0: raise _IncompatibleClusterDimError( f"Tensor dimension {weight_dim} along output channel axis " f"is not divisible by cluster_dim {self.cluster_dim}." ) + return self.granularity.get_blocks_to_cluster(weight_2d) - block_weights_to_cluster = self.granularity.get_blocks_to_cluster(weight) + @torch.no_grad() + def _cluster_to_centroids( + self, original_weights: torch.Tensor, sensitivities: torch.Tensor | None = None + ) -> tuple[torch.Tensor, torch.Tensor]: + """Cluster weight (+ optional sensitivities) into centroids via k-means. + + Returns ``(centroids, indices)``: ``centroids`` is a ``(num_blocks, + num_clusters, cluster_dim)`` tensor, and ``indices`` is the per-element + cluster assignment produced directly by the clustering algorithm. + """ + weight = original_weights.cpu() + if self.enable_per_channel_scale: + weight = self._scale_by_per_channel_scale(weight) + + axis = self.granularity.axis if self.granularity.axis else 0 + weight = self.reshape_strategy.reshape_for_kmeans(weight, axis) + block_weights_to_cluster = self._blocks_to_cluster(weight, axis) - # Reshape sensitivities if available - if self.sensitivities is not None: - sensitivities = self.sensitivities.cpu() + if sensitivities is not None: + sensitivities = sensitivities.cpu() # numpy has no bfloat16 dtype, so cluster bf16 sensitivities as # float32, matching the block-weight handling in _cluster_weights_1d. if sensitivities.dtype == torch.bfloat16: @@ -207,10 +324,9 @@ def _calculate_centroids( else: block_sensitivities = [None] * len(block_weights_to_cluster) - lut = [] - indices = [] num_clusters = 2**self.n_bits - + centroids_per_block = [] + block_indices = [] for block_weight, block_sensitivity in zip( block_weights_to_cluster, block_sensitivities, strict=True ): @@ -218,35 +334,66 @@ def _calculate_centroids( centroids, clusters = self._cluster_weights_1d(block_weight, block_sensitivity) else: centroids, clusters = self._cluster_weights_2d(block_weight, block_sensitivity) - centroids = self._pad_lut_to_num_clusters(centroids, num_clusters) + centroids_per_block.append(centroids.to(weight.dtype)) + block_indices.append(self._build_block_indices(clusters, block_weight).to(torch.uint8)) - lut.append(centroids.to(weight.dtype)) - block_indices = self._build_block_indices(clusters, block_weight) - block_indices = block_indices.to(torch.uint8) - indices.append(block_indices) + stacked = torch.stack(centroids_per_block) + # Keep a trailing vector dimension so shape is (num_blocks, + # num_clusters, cluster_dim) for both scalar and vector palettization. + centroids_pnd = stacked if self.cluster_dim > 1 else stacked.unsqueeze(-1) + centroids_pnd = centroids_pnd.detach().clone().to(original_weights.device) - # Handle concatenation based on granularity axis and convert to the shape - # of original weight tensor (axis 0 reduced by cluster_dim for vector case) - indices = torch.cat(indices, dim=axis) - indices_shape = list(original_weights.shape) - indices_shape[0] = indices_shape[0] // self.cluster_dim - indices = self.reshape_strategy.reshape_to_original( - indices, axis, torch.Size(indices_shape) - ) + indices = self._combine_block_indices(block_indices, axis, original_weights) + return centroids_pnd, indices + + @torch.no_grad() + def _assign_indices( + self, original_weights: torch.Tensor, centroids: torch.Tensor + ) -> torch.Tensor: + """Nearest-centroid hard assignment of ``original_weights`` against a + given ``centroids`` (P, K, D) tensor. + """ + weight = original_weights.detach().cpu() + if self.enable_per_channel_scale: + weight = self._scale_by_per_channel_scale(weight) - # Combine LUTs for all blocks into single tensor - lut = torch.stack(lut) + axis = self.granularity.axis if self.granularity.axis else 0 + weight_2d = self.reshape_strategy.reshape_for_kmeans(weight, axis) + blocks = self._blocks_to_cluster(weight_2d, axis) + centroids_cpu = centroids.detach().cpu().float() - # Quantize the entire stacked LUT in one shot - lut = self._quantize_lut(lut) + block_indices = [] + for block_idx, block_weight in enumerate(blocks): + vec = self._vectorize(block_weight) + dist = torch.cdist(vec.float(), centroids_cpu[block_idx]) + clusters = dist.argmin(dim=-1) + block_indices.append(self._build_block_indices(clusters, block_weight).to(torch.uint8)) - lut = self._reshape_lut_tensor(lut) + # self.indices is always CPU-resident, matching self.lut. + return self._combine_block_indices(block_indices, axis, original_weights) - # Clear stale flag since centroids are now up-to-date with sensitivities - self._centroids_stale = False + def _combine_block_indices( + self, + block_indices: list[torch.Tensor], + axis: int, + original_weights: torch.Tensor, + ) -> torch.Tensor: + """Concatenate per-block indices and reshape to the original weight + shape (axis 0 reduced by ``cluster_dim`` for vector palettization). + """ + indices = torch.cat(block_indices, dim=axis) + indices_shape = list(original_weights.shape) + indices_shape[0] = indices_shape[0] // self.cluster_dim + return self.reshape_strategy.reshape_to_original(indices, axis, torch.Size(indices_shape)) - return lut, indices + @torch.no_grad() + def _raw_lut(self, centroids: torch.Tensor) -> torch.Tensor: + """Reshape ``centroids`` (P, K, D) to the pre-quantization LUT shape + ``(P, K[, D])``, detached. + """ + centroids = centroids.detach() + return centroids if self.cluster_dim > 1 else centroids.squeeze(-1) def _palettize( self, lut: torch.Tensor, indices: torch.Tensor, original_weights: torch.Tensor @@ -269,8 +416,10 @@ def _palettize( clustered_weight = None axis = self.granularity.axis if self.granularity.axis else 0 + lut = lut.to(indices.device) + # Reshape indices back to 2D for block processing (reverse of - # reshape_to_original in _calculate_centroids) + # reshape_to_original in _assign_indices) indices = self.reshape_strategy.reshape_for_kmeans(indices, axis) # Cast to int for indexing since PyTorch treats uint8 as a boolean mask indices = indices.int() @@ -283,7 +432,7 @@ def _palettize( flat_lut = lut.squeeze() clustered_weight = flat_lut[indices] if self.cluster_dim > 1: - clustered_weight = self._devectorize(clustered_weight) + clustered_weight = self._lookup_result_to_block(clustered_weight) elif isinstance(self.granularity, PerGroupedChannelGranularity): # Per-grouped-channel granularity: multiple LUTs for different blocks depalett_block_weights = [] @@ -318,7 +467,7 @@ def _palettize( depalett_block_weight = block_lut[block_indices] if self.cluster_dim > 1: - depalett_block_weight = self._devectorize(depalett_block_weight) + depalett_block_weight = self._lookup_result_to_block(depalett_block_weight) depalett_block_weights.append(depalett_block_weight) @@ -327,8 +476,6 @@ def _palettize( # Unknown granularity raise ValueError(f"Unsupported granularity: {self.granularity}") - clustered_weight.to(original_weights.dtype) - # Reshape to original weight shape clustered_weight = self.reshape_strategy.reshape_to_original( clustered_weight, axis, original_weights.shape @@ -337,7 +484,7 @@ def _palettize( if self.enable_per_channel_scale: clustered_weight = self._unscale_by_per_channel_scale(clustered_weight) - return clustered_weight.to(original_weights.device) + return clustered_weight.to(original_weights.device, original_weights.dtype) def _pad_lut_to_num_clusters( self, @@ -373,17 +520,15 @@ def _pad_lut_to_num_clusters( padded_lut[: len(centroids)] = centroids return padded_lut - def _quantize_lut( + def quantize_lut( self, lut: torch.Tensor, ) -> torch.Tensor: - """Quantize the stacked LUT tensor and populate export buffers. + """Quantize the stacked LUT tensor and dequantize it back via STE. Computes per-block quantization parameters on the stacked LUT of shape ``(num_blocks, num_clusters[, cluster_dim])``, quantizes it, then - dequantizes back to the original dtype for STE-style training. Stores - ``quantized_lut``, ``lut_quantization_scale``, and - ``lut_quantization_zero_point`` as reshaped/detached buffers for export. + dequantizes back to the original dtype via a fused STE op. If no LUT fake quantizer is configured, returns the input unchanged. @@ -397,16 +542,11 @@ def _quantize_lut( if self._lut_fake_quantizer is None: return lut + orig_dtype = lut.dtype scale, zero_point, minval = self._lut_fake_quantizer.qparams_calculator(lut) - quantized_lut = self._lut_fake_quantizer.quantize(lut, scale, zero_point, minval) - lut = self._lut_fake_quantizer.dequantize( - quantized_lut, scale, zero_point, minval, output_dtype=lut.dtype - ) - self.quantized_lut = self._reshape_lut_tensor(quantized_lut.detach()) - self.lut_quantization_scale = self._reshape_lut_tensor(scale.detach()) - self.lut_quantization_zero_point = ( - self._reshape_lut_tensor(zero_point.detach()) if zero_point is not None else None - ) + lut = self._lut_fake_quantizer._fused_fake_quant_dequant( + lut.to(torch.float32), scale, zero_point, minval + ).to(orig_dtype) return lut def _cluster_weights_1d( @@ -511,19 +651,18 @@ def _cluster_weights_2d( return centroids, labels def _vectorize(self, tensor: torch.Tensor) -> torch.Tensor: - """Reshape a 2D tensor into (N, cluster_dim) vectors for vector k-means. + """Reshape a 2D tensor into (N, cluster_dim) vectors for k-means. Vectors are always formed along axis 0 (output channel axis). This transposes - the tensor so consecutive elements along axis 0 are grouped into vectors. + the tensor so consecutive elements along axis 0 are grouped into vectors. For + ``cluster_dim == 1`` (scalar palettization), this is just a flatten. """ + if self.cluster_dim == 1: + return tensor.reshape(-1, 1) return tensor.transpose(0, 1).reshape(-1, self.cluster_dim) - def _devectorize(self, looked_up: torch.Tensor) -> torch.Tensor: - """Reshape vector LUT lookup result back to 2D weight shape. - - Reverses the vectorization along axis 0 (output channel axis). - Since vectorization is always along axis 0, looked_up always has shape - (rows // cluster_dim, cols, cluster_dim). + def _lookup_result_to_block(self, looked_up: torch.Tensor) -> torch.Tensor: + """Reshape a vector LUT lookup result back to 2D weight shape. Args: looked_up: Result of LUT[indices] of shape diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py index a892d7e..b85d7c6 100644 --- a/src/coreai_opt/palettization/kmeans/palettizer.py +++ b/src/coreai_opt/palettization/kmeans/palettizer.py @@ -16,6 +16,7 @@ import torch.nn.utils.parametrize as P from tqdm import tqdm +from coreai_opt._utils.config_utils import ConfigLevel as _ConfigLevel from coreai_opt._utils.eager_utils import ( EagerCompressionComponentBuilderMixin as _EagerCompressionComponentBuilderMixin, ) @@ -27,21 +28,22 @@ ) from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor from coreai_opt._utils.torch_utils import ( + move_model_to_train as _move_model_to_train, remove_compression_parametrizations as _remove_compression_parametrizations, ) +from coreai_opt.base_model_compressor import _CompressorLifecycle from coreai_opt.common import ExportBackend -from coreai_opt.config.compression_config import ModuleCompressionConfig +from coreai_opt.config.compression_config import ModuleCompressionConfig, ModuleConfigDict from coreai_opt.config.spec import CompressionTargetTensor from coreai_opt.config.spec.base import CompressionSpec from coreai_opt.palettization.base_palettizer import _BasePalettizer from coreai_opt.palettization.config.palettization_config import ( KMeansPalettizerConfig, + PATSchedule, ) from coreai_opt.palettization.spec.fake_palettize import ( _disable_fake_palett, - _disable_observer, _enable_fake_palett, - _enable_observer, ) from ._prepare_for_export import ( @@ -95,7 +97,8 @@ def _calculate_centroids_for_module( fp_module, weight, layer_name = args try: - fp_module(weight) + with torch.no_grad(): + fp_module(weight) except Exception as e: raise RuntimeError(f"Centroid calculation failed for layer {layer_name!r}") from e @@ -139,6 +142,10 @@ def __init__(self, model: torch.nn.Module, config: KMeansPalettizerConfig | None self._num_workers = 1 + self._step_count: int = 0 + self._fp_to_schedule: dict[_KMeansFakePalettize, PATSchedule] = {} + self._module_config_dict: ModuleConfigDict[ModuleCompressionConfig] = {} + @classmethod def get_op_type_resolver(cls) -> Callable[[Callable], str | None]: """Return a function that maps a torch function to its palettizable op type.""" @@ -187,6 +194,9 @@ def prepare( # Save so calibration_mode's recompute can use the same parallelism. self._num_workers = num_workers + # Cache config dict before prepare() modifies module types. + self._module_config_dict = self._config.build_module_config_dict(self._model) + # Prepare the model logger.info("Preparing model for palettization") prepared_model = self._handler.prepare(self._model, example_inputs=example_inputs) @@ -200,7 +210,6 @@ def prepare( sensitivities = torch.load(sensitivity_path, weights_only=True) self._set_sensitivities_in_fake_palettize_modules(sensitivities) - self._model.apply(_enable_observer) self._model.apply(_disable_fake_palett) if self._num_workers > 1: @@ -213,7 +222,6 @@ def prepare( self._remove_disabled_fake_palett_modules(self._model) self._model.apply(_enable_fake_palett) - self._model.apply(_disable_observer) # Mark the model as prepared to prevent re-preparation self._mark_model_as_prepared(prepared_model) @@ -272,64 +280,124 @@ def calibration_mode( "Model must be prepared before entering calibration mode. Call prepare() first." ) - # Save model checkpoint before modifying gradients - checkpoint_path = self._save_model_checkpoint(self._model) - self._model.zero_grad() - - # Helper class for loss computation - class CalibrationHelper: - def __init__(self, loss_fn): - self.loss_fn = loss_fn - self.step_called = False - - def step(self, output: torch.Tensor, target: torch.Tensor): - """Compute loss and backward pass.""" - loss = self.loss_fn(output, target) - loss.backward() - self.step_called = True + if self._lifecycle is not _CompressorLifecycle.IDLE: + raise RuntimeError( + f"Cannot enter calibration_mode() while palettizer is {self._lifecycle.value}" + ) + self._lifecycle = _CompressorLifecycle.CALIBRATING + try: + # Save model checkpoint before modifying gradients + checkpoint_path = self._save_model_checkpoint(self._model) + self._model.zero_grad() + + # Helper class for loss computation + class CalibrationHelper: + def __init__(self, loss_fn): + self.loss_fn = loss_fn + self.step_called = False + + def step(self, output: torch.Tensor, target: torch.Tensor): + """Compute loss and backward pass.""" + loss = self.loss_fn(output, target) + loss.backward() + self.step_called = True + + # Disable fake palettization for sensitivity computation + self._model.apply(_disable_fake_palett) + + calibration_helper = CalibrationHelper(loss_fn) + + with self._register_grad_square_hooks(self._model): + try: + yield calibration_helper + finally: + # Ensure step() was called at least once + if not calibration_helper.step_called: + raise RuntimeError( + "calibration_mode requires at least one call to step(). " + "No calibration data was processed." + ) - # Disable observers and fake palettization for sensitivity computation - self._model.apply(_disable_observer) - self._model.apply(_disable_fake_palett) + # Construct sensitivities + sensitivities = self._construct_sensitivities(sensitivity_path) - calibration_helper = CalibrationHelper(loss_fn) + # Restore model from checkpoint + self._load_model_checkpoint(self._model, checkpoint_path) - with self._register_grad_square_hooks(self._model): - try: - yield calibration_helper - finally: - # Ensure step() was called at least once - if not calibration_helper.step_called: - raise RuntimeError( - "calibration_mode requires at least one call to step(). " - "No calibration data was processed." - ) + # Set sensitivities in fake palettize modules + self._set_sensitivities_in_fake_palettize_modules(sensitivities) - # Construct sensitivities - sensitivities = self._construct_sensitivities(sensitivity_path) + # Zero out gradients to clean up squared gradient values from hooks + self._model.zero_grad() - # Restore model from checkpoint - self._load_model_checkpoint(self._model, checkpoint_path) + # Recompute centroids with sensitivities, matching the + # parallelism the user opted into at prepare() time. + if self._num_workers > 1: + self._calculate_centroids_parallel(self._num_workers) + else: + self._calculate_centroids_sequential() - # Set sensitivities in fake palettize modules - self._set_sensitivities_in_fake_palettize_modules(sensitivities) + # Restore normal operation + self._model.apply(_enable_fake_palett) + finally: + self._lifecycle = _CompressorLifecycle.IDLE - # Zero out gradients to clean up squared gradient values from hooks - self._model.zero_grad() + @contextmanager + def training_mode(self): + """Context manager wrapping a training loop. Mutually exclusive with + calibration_mode(). + """ + if not self._is_model_prepared(self._model): + raise RuntimeError( + "Model must be prepared before entering training mode. Call prepare() first." + ) - # Enable observers to recompute LUTs with sensitivities - self._model.apply(_enable_observer) + if self._lifecycle is not _CompressorLifecycle.IDLE: + raise RuntimeError( + f"Cannot enter training_mode() while palettizer is {self._lifecycle.value}" + ) + self._lifecycle = _CompressorLifecycle.TRAINING + with _move_model_to_train(self._model): + try: + self._build_fp_to_schedule() + self._apply_schedule() + yield self + finally: + self._lifecycle = _CompressorLifecycle.IDLE - # Recompute centroids with sensitivities, matching the - # parallelism the user opted into at prepare() time. - if self._num_workers > 1: - self._calculate_centroids_parallel(self._num_workers) - else: - self._calculate_centroids_sequential() + def step(self) -> None: + """Advance the schedule by one step. Must be called inside training_mode().""" + if self._lifecycle is not _CompressorLifecycle.TRAINING: + raise RuntimeError("step() must be called inside a training_mode() context.") + self._step_count += 1 + self._apply_schedule() - # Restore normal operation - self._model.apply(_enable_fake_palett) - self._model.apply(_disable_observer) + def _build_fp_to_schedule(self) -> None: + if self._fp_to_schedule: + return + for module_name, module in self._model.named_modules(remove_duplicate=True): + if not P.is_parametrized(module): + continue + for parametrizations in module.parametrizations.values(): + for param in parametrizations: + if not isinstance(param, _KMeansFakePalettize): + continue + schedule = self._resolve_schedule(module_name) + if schedule is not None: + self._fp_to_schedule[param] = schedule + break + + def _resolve_schedule(self, module_name: str) -> PATSchedule | None: + """Look up the PAT schedule for a module via the config hierarchy.""" + for level in _ConfigLevel.priority_order(): + config = self._module_config_dict[level].get(module_name) + if config is not None: + return config.pat_schedule + return None + + 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)) def _validate_mmap_dir_constraints( self, @@ -425,7 +493,7 @@ def _spec_to_partial( # Serialize the spec, then layer in the owning module's compressor-specific # settings (e.g. enable_fast_kmeans_mode, rounding_precision). args = spec.model_dump_preserve_objects() - args.update(module_config._get_compressor_specific_settings()) + args.update(module_config._get_fake_module_kwargs()) return _KMeansFakePalettize.with_args(**args) def _collect_fake_palett_info(self, *, to_cpu: bool) -> list[_FakePalettInfo]: @@ -720,17 +788,15 @@ def _save_model_checkpoint(model: torch.nn.Module) -> str: @staticmethod def _load_model_checkpoint(model: torch.nn.Module, checkpoint_path: str): """Restore model checkpoint from specified checkpoint path.""" - if checkpoint_path is not None and os.path.exists(checkpoint_path): - logger.debug( - f"Restoring model from checkpoint {checkpoint_path} " - "before setting sensitivities and recomputing centroids" - ) - model.load_state_dict(torch.load(checkpoint_path, weights_only=True)) - # Clean up temporary checkpoint file - logger.debug(f"Removing temporary checkpoint {checkpoint_path}") - os.unlink(checkpoint_path) - else: - logger.error(f"Failed to load model checkpoint from path: {checkpoint_path}") + if checkpoint_path is None or not os.path.exists(checkpoint_path): + raise RuntimeError(f"Failed to load model checkpoint from path: {checkpoint_path}") + logger.debug( + f"Restoring model from checkpoint {checkpoint_path} " + "before setting sensitivities and recomputing centroids" + ) + model.load_state_dict(torch.load(checkpoint_path, weights_only=True)) + logger.debug(f"Removing temporary checkpoint {checkpoint_path}") + os.unlink(checkpoint_path) @staticmethod def _remove_disabled_fake_palett_modules(model: torch.nn.Module) -> None: diff --git a/src/coreai_opt/palettization/spec/__init__.py b/src/coreai_opt/palettization/spec/__init__.py index 82e6818..fc7f57a 100644 --- a/src/coreai_opt/palettization/spec/__init__.py +++ b/src/coreai_opt/palettization/spec/__init__.py @@ -11,11 +11,15 @@ PerTensorGranularity, ) from .spec import PalettizationSpec, default_weight_palettization_spec +from .training_strategy import DefaultTrainingConfig, TrainingStrategy, TrainingStrategyConfig __all__ = [ + "DefaultTrainingConfig", "PalettizationGranularity", "PalettizationSpec", "PerGroupedChannelGranularity", "PerTensorGranularity", + "TrainingStrategy", + "TrainingStrategyConfig", "default_weight_palettization_spec", ] diff --git a/src/coreai_opt/palettization/spec/fake_palettize.py b/src/coreai_opt/palettization/spec/fake_palettize.py index ee4f7e3..52b3180 100644 --- a/src/coreai_opt/palettization/spec/fake_palettize.py +++ b/src/coreai_opt/palettization/spec/fake_palettize.py @@ -5,7 +5,6 @@ from __future__ import annotations -import logging from abc import abstractmethod import torch @@ -19,28 +18,17 @@ from coreai_opt.palettization.spec import ( PalettizationGranularity, ) -from coreai_opt.palettization.spec.errors import ( - _IncompatibleClusterDimError, - _IncompatibleGranularityError, -) from coreai_opt.quantization.spec import QuantizationSpec -logger = logging.getLogger(__name__) - class _FakePalettizeImplBase(CompressionSimulatorBase, nn.Module): """Base class for fake palettization implementations with clustering and reconstruction methods. """ - lut: torch.Tensor indices: torch.Tensor per_channel_scale: torch.Tensor | None - quantized_lut: torch.Tensor | None - lut_quantization_scale: torch.Tensor | None - lut_quantization_zero_point: torch.Tensor | None fake_palett_enabled: torch.Tensor - observer_enabled: torch.Tensor def __init__( self, @@ -59,84 +47,54 @@ def __init__( self.enable_per_channel_scale = enable_per_channel_scale self.register_buffer("fake_palett_enabled", torch.tensor([1], dtype=torch.uint8)) - self.register_buffer("observer_enabled", torch.tensor([1], dtype=torch.uint8)) self._disabled = False - self.register_buffer("lut", None) self.register_buffer("indices", None) self.register_buffer("per_channel_scale", None) - self.register_buffer("quantized_lut", None) - self.register_buffer("lut_quantization_scale", None) - self.register_buffer("lut_quantization_zero_point", None) - - @property - def _initialized(self) -> bool: - """Return True if lut and indices have been initialized (not None).""" - return self.lut is not None and self.indices is not None def is_disabled(self) -> bool: """Return True if fake palettization has been disabled.""" return self._disabled def forward(self, tensor: torch.Tensor) -> torch.Tensor: - """Apply fake palettization to input tensor""" - # If permanently disabled due to incompatibility, return original tensor + """Fake-palettize ``tensor`` through the enable/disable lifecycle. + + Delegates representation-specific work to ``ensure_initialized`` and + ``forward_enabled``. + """ if self._disabled: return tensor - if self.observer_enabled[0] == 1: - # Cluster weights - try: - lut, indices = self._calculate_centroids(tensor) - except _IncompatibleGranularityError as e: - logger.warning( - f"Tensor incompatible with granularity: {e}. Skipping palettization." - ) - self._disabled = True - return tensor - except _IncompatibleClusterDimError as e: - logger.warning( - f"Tensor incompatible with cluster_dim: {e}. Skipping palettization." - ) - self._disabled = True - return tensor - - self.lut = lut.detach() - self.indices = indices.detach() - else: - # Check that recomputed statistics exist - if not self._initialized: - # Not initialized yet, return original tensor - return tensor - - if self.fake_palett_enabled[0] == 1: - return self._palettize(lut=self.lut, indices=self.indices, original_weights=tensor) - - return tensor + self.ensure_initialized(tensor) - @abstractmethod - def _palettize( - self, lut: torch.Tensor, indices: torch.Tensor, original_weights: torch.Tensor - ) -> torch.Tensor: - """Reconstruct palettized weights from lookup table and indices.""" - raise NotImplementedError() + # Check for self._disabled again in case ensure_initialized disabled the palettizer. + if self._disabled: + return tensor + if self.fake_palett_enabled[0] == 0: + return tensor + return self.forward_enabled(tensor) @abstractmethod - def _calculate_centroids(self, weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Cluster weights and return lookup table (LUT) and corresponding indices. + def ensure_initialized(self, tensor: torch.Tensor) -> None: + """Initialize compression parameters from ``tensor`` on first use. - If tensor is incompatible with the specified granularity, this method - should set self._disabled = True and return dummy values. + Set ``self._disabled = True`` if ``tensor`` is incompatible with the + configured spec. + """ + raise NotImplementedError() - Args: - weight: The weight tensor to cluster + @abstractmethod + def forward_enabled(self, tensor: torch.Tensor) -> torch.Tensor: + """Return the palettized output for ``tensor`` when enabled.""" + raise NotImplementedError() - Returns: - Tuple of (lut, indices) where: - - lut: Lookup table of cluster centroids - - indices: Index tensor mapping each weight to its cluster + @abstractmethod + def _palettize( + self, lut: torch.Tensor, indices: torch.Tensor, original_weights: torch.Tensor + ) -> torch.Tensor: + """Reconstruct palettized weights from lookup table and indices. - LUT shape must be of the following form: + ``lut`` shape must be of the following form: [NUM_LUT_AXIS_0, NUM_LUT_AXIS_1, NUM_PALETTES, VECTOR_SIZE] where, NUM_LUT_* is the number of LUTs for the corresponding axis. The computation @@ -148,7 +106,7 @@ def _calculate_centroids(self, weight: torch.Tensor) -> tuple[torch.Tensor, torc VECTOR_SIZE is lut.shape[-1] and is added to support vector palettization. When VECTOR_SIZE is 1, it is scalar palettization. - Indices shape much match the shape of the palettized weight. + ``indices`` shape must match the shape of the palettized weight. """ raise NotImplementedError() @@ -166,7 +124,7 @@ def _load_from_state_dict( ): """Custom state dict loading for palettization-specific buffers. - This method handles the loading of palettization-specific buffers (lut, indices, + This method handles the loading of palettization-specific buffers (indices, per_channel_scale) that may be dynamically created during forward passes. By registering them here, we ensure they are properly loaded from saved checkpoints and don't generate unexpected key warnings. @@ -175,12 +133,9 @@ def _load_from_state_dict( load_state_dict, etc.) and should not be called directly. """ buffer_names = { - "lut", "indices", "per_channel_scale", - "quantized_lut", - "lut_quantization_scale", - "lut_quantization_zero_point", + "centroids", } for buffer_name in buffer_names: @@ -202,12 +157,6 @@ def enable_fake_palett(self, enabled: bool = True) -> None: def disable_fake_palett(self): self.enable_fake_palett(False) - def enable_observer(self, enabled: bool = True) -> None: - self.observer_enabled[0] = 1 if enabled else 0 - - def disable_observer(self): - self.enable_observer(False) - def _enable_fake_palett(mod): """Enable fake palettization for the module.""" @@ -219,15 +168,3 @@ def _disable_fake_palett(mod): """Disable fake palettization for the module.""" if isinstance(mod, _FakePalettizeImplBase): mod.disable_fake_palett() - - -def _enable_observer(mod): - """Enable observation for this module.""" - if isinstance(mod, _FakePalettizeImplBase): - mod.enable_observer() - - -def _disable_observer(mod): - """Disable observation for this module.""" - if isinstance(mod, _FakePalettizeImplBase): - mod.disable_observer() diff --git a/src/coreai_opt/palettization/spec/spec.py b/src/coreai_opt/palettization/spec/spec.py index d3b4734..fa9307c 100644 --- a/src/coreai_opt/palettization/spec/spec.py +++ b/src/coreai_opt/palettization/spec/spec.py @@ -23,6 +23,7 @@ ) from .granularity import PalettizationGranularity, PerTensorGranularity +from .training_strategy import DefaultTrainingConfig, TrainingStrategyConfig _SUPPORTED_LUT_DTYPES = {torch.int8, torch.uint8, torch.float8_e4m3fn, torch.float8_e5m2} @@ -60,6 +61,14 @@ class PalettizationSpec(CompressionSpec): enable_per_channel_scale: When set to True, weights are normalized along the output channels using per-channel scales before being palettized. Default: False. + training_strategy_config: Which training-time behavior this weight's + fake-palettize module uses, and that strategy's settings. + ``DefaultTrainingConfig()`` is post-training, one-shot k-means + (today's KMeansPalettizer behavior). Additional strategies are added + by subclassing ``TrainingStrategyConfig`` (registered via + ``TrainingStrategyConfig.register()``) and pointing its + ``_strategy_cls`` at a ``TrainingStrategy`` subclass. Default: + DefaultTrainingConfig(). Example: >>> # Basic 4-bit palettization @@ -91,6 +100,10 @@ class PalettizationSpec(CompressionSpec): ] = PerTensorGranularity() cluster_dim: PositiveInt = 1 enable_per_channel_scale: bool = False + training_strategy_config: Annotated[ + TrainingStrategyConfig, + BeforeValidator(TrainingStrategyConfig.maybe_build_from_dict), + ] = DefaultTrainingConfig() # Private attribute for compression type _compression_type: CompressionType = PrivateAttr(default=CompressionType.PALETTIZATION) diff --git a/src/coreai_opt/palettization/spec/training_strategy.py b/src/coreai_opt/palettization/spec/training_strategy.py new file mode 100644 index 0000000..2c58b43 --- /dev/null +++ b/src/coreai_opt/palettization/spec/training_strategy.py @@ -0,0 +1,101 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Pluggable training-time behavior for fake-palettize modules.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, ClassVar + +import torch +from pydantic import BaseModel, ConfigDict, model_serializer + +from coreai_opt._utils.registry_utils import ConfigRegistryMixin as _ConfigRegistryMixin + +if TYPE_CHECKING: + from coreai_opt.palettization.kmeans.kmeans_fake_palettize import _KMeansFakePalettize + + +class TrainingStrategy(ABC): + """Contract for a fake-palettize module's training-time forward pass.""" + + @abstractmethod + def train_forward(self, module: _KMeansFakePalettize, weight: torch.Tensor) -> torch.Tensor: + """Compute this training step's output for ``weight``. + + May mutate ``module.centroids`` in place (must set + ``module._indices_stale = True`` if it does). Use ``module.quantize_lut()`` + to fake-quantize intermediate centroids against the module's + configured LUT quantizer. Must leave ``module.centroids`` such that + ``module.hard_assign(weight)`` produces a sensible result at eval time. + """ + raise NotImplementedError + + +class _DefaultTrainingStrategy(TrainingStrategy): + """Post-training, one-shot k-means — today's KMeansPalettizer behavior. + + This strategy does not train the palettized weights. Inside a + ``training_mode()`` loop it reconstructs each forward from the frozen + centroids/indices computed at ``prepare()`` time, so no gradient flows back + to a palettized weight — palettized weights stay fixed at their one-shot + k-means values. Once a module's ``pat_schedule`` has enabled fake + palettization, the rest of the model still trains normally and its forward + pass sees the palettized weights, so non-palettized parameters adapt around + them (palettization-aware fine-tuning of the rest of the network). To learn + the palettized weights/centroids themselves, register a custom + ``TrainingStrategy``. + """ + + def train_forward(self, module: _KMeansFakePalettize, weight: torch.Tensor) -> torch.Tensor: + return module.hard_assign(weight) + + +class TrainingStrategyConfig(BaseModel, _ConfigRegistryMixin): + """Base class for a fake-palettize module's training-strategy settings. + + Each subclass points ``_strategy_cls`` at its paired ``TrainingStrategy`` + behavior class; ``build_strategy()`` constructs that strategy from this + config's own fields. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + # Each subclass points this at its paired TrainingStrategy behavior class. + _strategy_cls: ClassVar[type[TrainingStrategy]] + + @model_serializer + def _serialize_model(self) -> dict[str, Any]: + """Custom serializer that includes the registry type.""" + data = {} + + for field_name in type(self).model_fields: + data[field_name] = getattr(self, field_name) + + # Find the registry key for this class type + registry_key = None + # Use the base class registry instead of instance registry + for key, registered_class in TrainingStrategyConfig.REGISTRY.items(): + if registered_class is type(self): + registry_key = key + break + + if registry_key is not None: + data["type"] = registry_key + + return data + + def build_strategy(self) -> TrainingStrategy: + """Construct this config's paired ``TrainingStrategy`` behavior instance.""" + kwargs = {k: v for k, v in self.model_dump().items() if k != "type"} + return self._strategy_cls(**kwargs) + + +@TrainingStrategyConfig.register("default") +class DefaultTrainingConfig(TrainingStrategyConfig): + """Settings for the default, post-training one-shot k-means strategy. No fields.""" + + _strategy_cls = _DefaultTrainingStrategy diff --git a/tests/palettization/test_kmeans_fake_palettize.py b/tests/palettization/test_kmeans_fake_palettize.py index 0e986a1..a08357b 100644 --- a/tests/palettization/test_kmeans_fake_palettize.py +++ b/tests/palettization/test_kmeans_fake_palettize.py @@ -17,15 +17,19 @@ _KMeansPalettizerSupportedOpsRegistry, ) from coreai_opt.palettization.spec import ( + DefaultTrainingConfig, PalettizationSpec, PerGroupedChannelGranularity, PerTensorGranularity, + TrainingStrategy, + TrainingStrategyConfig, ) from coreai_opt.palettization.spec.errors import ( _IncompatibleClusterDimError, _IncompatibleGranularityError, ) from coreai_opt.palettization.spec.spec import _SUPPORTED_LUT_DTYPES +from coreai_opt.palettization.spec.training_strategy import _DefaultTrainingStrategy from coreai_opt.quantization.spec import QuantizationScheme, QuantizationSpec @@ -41,6 +45,17 @@ def _make_lut_qspec( return QuantizationSpec(dtype=lut_dtype, qscheme=lut_qscheme) +def _initialize_and_get_lut_indices( + palettizer: _KMeansFakePalettize, weight: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Initialize the palettizer from ``weight`` and return its ``(lut, indices)``, + for tests that want the clustering result without going through the + lazy-init ``forward()`` path. + """ + palettizer._initialize(weight) + return palettizer.lut, palettizer.indices + + def _valid_lut_dtype_qscheme_combinations(): """Return valid (dtype, qscheme) pairs for LUT quantization. @@ -92,7 +107,7 @@ def test__calculate_centroids_simple_per_tensor(self, lut_dtype, weight_dtype): ) # Verify centroid properties - lut, indices = palettizer._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) assert lut.shape == (1, 1, 2**spec.n_bits, 1) assert lut.dtype == weight.dtype @@ -130,7 +145,7 @@ def test__calculate_centroids_bfloat16_sensitivities(self, cluster_dim): # Positive sensitivities prevent a zero-total-weight cluster from yielding a NaN centroid. palettizer.sensitivities = torch.rand_like(weight) + 1.0 - lut, indices = palettizer._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) assert lut.dtype == weight.dtype assert lut.shape == (1, 1, 2**spec.n_bits, cluster_dim) @@ -155,7 +170,7 @@ def test_cluster_weights_1d_with_few_unique_values(self): enable_per_channel_scale=spec.enable_per_channel_scale, ) - centroids, indices = palettizer._calculate_centroids(weight) + centroids, indices = _initialize_and_get_lut_indices(palettizer, weight) assert centroids.shape == (1, 1, 2**spec.n_bits, 1) @@ -199,7 +214,7 @@ def test__calculate_centroids_per_grouped_channel_axis_0(self, lut_dtype): enable_per_channel_scale=spec.enable_per_channel_scale, ) - lut, indices = palettizer._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) # Test weight reconstruction for per-grouped-channel granularity palettized_weight = palettizer._palettize(lut, indices, weight) @@ -241,7 +256,7 @@ def test__calculate_centroids_per_grouped_channel_axis_1(self, lut_dtype): enable_per_channel_scale=spec.enable_per_channel_scale, ) - lut, indices = palettizer._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) # Test weight reconstruction for per-grouped-channel granularity palettized_weight = palettizer._palettize(lut, indices, weight) @@ -291,8 +306,10 @@ def test_cluster_weights_1d_fast_mode_vs_regular(self, dtype): enable_fast_kmeans_mode=False, ) - centroids_fast, clusters_fast = palettizer_fast._calculate_centroids(weight) - centroids_regular, clusters_regular = palettizer_regular._calculate_centroids(weight) + centroids_fast, clusters_fast = _initialize_and_get_lut_indices(palettizer_fast, weight) + centroids_regular, clusters_regular = _initialize_and_get_lut_indices( + palettizer_regular, weight + ) # Both should produce valid results assert len(centroids_fast) <= 4 @@ -360,7 +377,7 @@ def test_group_size_divisibility_validation_divisible(self, weight_shape, axis, enable_per_channel_scale=False, ) - lut, indices = fake_palettize._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(fake_palettize, weight) assert lut is not None assert indices is not None num_blocks = fake_palettize.granularity.num_blocks_to_cluster(weight) @@ -389,7 +406,7 @@ def test_group_size_divisibility_validation_indivisible( ) with pytest.raises(_IncompatibleGranularityError): - fake_palettize._calculate_centroids(weight) + _initialize_and_get_lut_indices(fake_palettize, weight) def test_group_size_divisibility_validation_insufficient_dimensions(self): """Test that validation warns for insufficient parameter dimensions.""" @@ -408,7 +425,7 @@ def test_group_size_divisibility_validation_insufficient_dimensions(self): # Should warn and skip because parameter is 1D but axis 1 was specified with pytest.raises(_IncompatibleGranularityError): - fake_palettize._calculate_centroids(weight) + _initialize_and_get_lut_indices(fake_palettize, weight) def test_disabled_flag_behavior(self): """Test that _disabled flag works correctly and permanently @@ -429,9 +446,9 @@ def test_disabled_flag_behavior(self): # Initially not disabled assert fake_palettize._disabled is False - # Spy on _calculate_centroids to verify it's called - original__calculate_centroids = fake_palettize._calculate_centroids - fake_palettize._calculate_centroids = Mock(side_effect=original__calculate_centroids) + # Spy on _cluster_to_centroids to verify it's called + original__cluster_to_centroids = fake_palettize._cluster_to_centroids + fake_palettize._cluster_to_centroids = Mock(side_effect=original__cluster_to_centroids) # Call forward - this should trigger the _IncompatibleGranularityError # and set _disabled = True @@ -441,8 +458,8 @@ def test_disabled_flag_behavior(self): assert fake_palettize._disabled is True assert torch.equal(result, weight) # Should return original tensor - # _calculate_centroids should have been called once - assert fake_palettize._calculate_centroids.call_count == 1 + # _cluster_to_centroids should have been called once + assert fake_palettize._cluster_to_centroids.call_count == 1 # Call forward again - should immediately return original tensor result2 = fake_palettize.forward(weight) @@ -451,15 +468,14 @@ def test_disabled_flag_behavior(self): assert fake_palettize._disabled is True assert torch.equal(result2, weight) - # _calculate_centroids should still have been called only once + # _cluster_to_centroids should still have been called only once # (not called the second time) - assert fake_palettize._calculate_centroids.call_count == 1 + assert fake_palettize._cluster_to_centroids.call_count == 1 @pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available") def test_mps_device_handling(self): - """ - Test that palettization preserves MPS device for weights while - using CPU for computation. + """Palettization keeps the LUT on the weight's device while the weight-sized + indices stay CPU-resident, and reconstruction returns on the original device. """ # Create weights on MPS device weight = torch.randn(4, 8, dtype=torch.float32, device="mps") @@ -477,14 +493,15 @@ def test_mps_device_handling(self): enable_per_channel_scale=spec.enable_per_channel_scale, ) - # Calculate centroids - this should move computation to CPU internally - lut, indices = palettizer._calculate_centroids(weight) + # Clustering runs on CPU internally, but centroids are placed on the weight's device. + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) - # Verify LUT and indices are on CPU (expected behavior) - assert lut.device.type == "cpu" + # The LUT follows the weight's device (kept on-device so training-time LUT + # quantization stays on-device); the weight-sized indices stay CPU-resident. + assert lut.device.type == "mps" assert indices.device.type == "cpu" - # Palettize the weights - this should return result on original device (MPS) + # Reconstruction gathers against the CPU indices and returns on the original device (MPS). palettized_weight = palettizer._palettize(lut, indices, weight) # Verify the palettized weights are back on MPS @@ -605,8 +622,8 @@ def test_quantized_lut_centroid(self, lut_dtype, lut_qscheme): palettizer_base = _KMeansFakePalettize(**spec_base.__dict__) palettizer_quant = _KMeansFakePalettize(**spec_quant.__dict__) - lut_base, indices_base = palettizer_base._calculate_centroids(weight) - lut_quant, indices_quant = palettizer_quant._calculate_centroids(weight) + lut_base, indices_base = _initialize_and_get_lut_indices(palettizer_base, weight) + lut_quant, indices_quant = _initialize_and_get_lut_indices(palettizer_quant, weight) # LUT shapes should be the same assert lut_base.shape == lut_quant.shape @@ -694,7 +711,7 @@ def test_quantized_lut_exact_values(self, lut_dtype, lut_qscheme, expected_lut): ) palettizer = _KMeansFakePalettize(**spec.__dict__) - lut, indices = palettizer._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) # LUT shape: [1, 1, 4, 1] for per-tensor with 2-bit, cluster_dim=1 assert lut.shape == (1, 1, 4, 1) @@ -751,7 +768,7 @@ def test_quantized_lut_with_sensitivities(self, lut_dtype, lut_qscheme): enable_fast_kmeans_mode=True, ) - lut, indices = palettizer._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) palettized = palettizer._palettize(lut, indices, weight) # Shape should be preserved @@ -884,7 +901,7 @@ def test_scaling_affects_clustering(self): ) palettizer_with_scale = _KMeansFakePalettize(**spec_with_scale.__dict__) - lut, indices = palettizer_with_scale._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer_with_scale, weight) palettized_with_scale = palettizer_with_scale._palettize(lut, indices, weight) mse_with_scale = torch.mean((weight - palettized_with_scale) ** 2) @@ -909,30 +926,27 @@ def test_state_dict_after_initialization(self, enable_per_channel_scale): weight = torch.randn(4, 8) # Initially, the module should not be initialized - assert not palettizer._initialized + assert palettizer.indices is None # Before initialization, buffers should be None # Note: state_dict() only includes buffers that are not None initial_state_dict = palettizer.state_dict() - assert "lut" not in initial_state_dict assert "indices" not in initial_state_dict assert "per_channel_scale" not in initial_state_dict # Initialize by running forward pass palettizer.forward(weight) - # After initialization, check the module is initialized - assert palettizer._initialized + # After initialization, check indices exists + assert palettizer.indices is not None + assert isinstance(palettizer.lut, torch.Tensor) + assert palettizer.lut.shape == (1, 1, 4, 1) # 2^2 = 4 clusters + assert palettizer.lut.dtype == weight.dtype + # Check state_dict contains proper values state_dict = palettizer.state_dict() - # LUT should be a tensor with shape (1, 1, num_clusters, 1) - assert state_dict["lut"] is not None - assert isinstance(state_dict["lut"], torch.Tensor) - assert state_dict["lut"].shape == (1, 1, 4, 1) # 2^2 = 4 clusters - assert state_dict["lut"].dtype == weight.dtype - # Indices should be a tensor with same shape as weight assert state_dict["indices"] is not None assert isinstance(state_dict["indices"], torch.Tensor) @@ -950,55 +964,33 @@ def test_state_dict_after_initialization(self, enable_per_channel_scale): assert palettizer.per_channel_scale is None # Verify we can access params without error - assert torch.equal(palettizer.lut, state_dict["lut"]) assert torch.equal(palettizer.indices, state_dict["indices"]) - def test_observer_modes(self): - """Test that initialization behavior respects observer modes.""" + def test_lazy_initialization_regardless_of_fake_palett_enabled(self): + """Test that centroid/lut initialization happens on first forward + regardless of ``fake_palett_enabled``. + """ spec = PalettizationSpec( n_bits=2, granularity=PerTensorGranularity(), enable_per_channel_scale=False ) palettizer = _KMeansFakePalettize(**spec.__dict__) weight = torch.randn(3, 4) - # Test with observer enabled, fake_palett disabled - palettizer.enable_observer(True) palettizer.enable_fake_palett(False) - # Should initialize but return original tensor output = palettizer.forward(weight) - assert palettizer._initialized - # Should return original since fake_palett disabled + assert palettizer.indices is not None assert torch.equal(output, weight) # State dict should have initialized values state_dict = palettizer.state_dict() - assert state_dict["lut"] is not None assert state_dict["indices"] is not None - # Test with observer disabled palettizer2 = _KMeansFakePalettize(**spec.__dict__) - palettizer2.enable_observer(False) - # Should not initialize output2 = palettizer2.forward(weight) - assert not palettizer2._initialized - # Should return original since not initialized - assert torch.equal(output2, weight) - - # Test with both disabled - palettizer2 = _KMeansFakePalettize(**spec.__dict__) - palettizer2.enable_observer(False) - - # Should not enabled - palettizer3 = _KMeansFakePalettize(**spec.__dict__) - palettizer3.enable_observer(True) - palettizer3.enable_fake_palett(True) - - output3 = palettizer3.forward(weight) - assert palettizer3._initialized - # Should return different output - assert not torch.equal(output3, weight) + assert palettizer2.indices is not None + assert not torch.equal(output2, weight) def test_save_load_state_dict_preserves_palettization(self): """Test that saving and loading state_dict preserves palettization behavior.""" @@ -1022,10 +1014,7 @@ def test_save_load_state_dict_preserves_palettization(self): original_output = original_palettizer.forward(weight) # Verify it's initialized - assert original_palettizer._initialized - - # Disable observer - original_palettizer.enable_observer(False) + assert original_palettizer.indices is not None # Save state_dict to temporary file with tempfile.NamedTemporaryFile(delete=False, suffix=".pth") as tmp_file: @@ -1037,17 +1026,16 @@ def test_save_load_state_dict_preserves_palettization(self): loaded_palettizer = _KMeansFakePalettize(**spec.__dict__) # Verify it starts uninitialized - assert not loaded_palettizer._initialized + assert loaded_palettizer.indices is None # Load state_dict state_dict = torch.load(temp_path, weights_only=True) loaded_palettizer.load_state_dict(state_dict) # Verify it's now initialized after loading - assert loaded_palettizer._initialized + assert loaded_palettizer.indices is not None assert loaded_palettizer.fake_palett_enabled.item() == 1 - assert loaded_palettizer.observer_enabled.item() == 0 loaded_output = loaded_palettizer.forward(weight) @@ -1096,14 +1084,16 @@ def test_sensitivities_reduce_error_for_sensitive_values(self): ) palettizer_with_sens = _KMeansFakePalettize(**spec.__dict__, sensitivities=sensitivities) - lut_with_sens, indices_with_sens = palettizer_with_sens._calculate_centroids(weight) + lut_with_sens, indices_with_sens = _initialize_and_get_lut_indices( + palettizer_with_sens, weight + ) palettized_with_sens = palettizer_with_sens._palettize( lut_with_sens, indices_with_sens, weight ) # Palettize without sensitivities palettizer_no_sens = _KMeansFakePalettize(**spec.__dict__, sensitivities=None) - lut_no_sens, indices_no_sens = palettizer_no_sens._calculate_centroids(weight) + lut_no_sens, indices_no_sens = _initialize_and_get_lut_indices(palettizer_no_sens, weight) palettized_no_sens = palettizer_no_sens._palettize(lut_no_sens, indices_no_sens, weight) # Calculate reconstruction errors for each region @@ -1152,7 +1142,7 @@ def test_sensitivities_affect_centroid_placement(self): ) palettizer_with_sens = _KMeansFakePalettize(**spec.__dict__, sensitivities=sensitivities) - lut_with_sens, _ = palettizer_with_sens._calculate_centroids(weight) + lut_with_sens, _ = _initialize_and_get_lut_indices(palettizer_with_sens, weight) # The centroids should be more concentrated around the low values (1-4) # because they have high sensitivity @@ -1191,7 +1181,7 @@ def test_fast_vs_regular_kmeans_with_sensitivities_consistency(self): sensitivities=sensitivities.clone(), enable_fast_kmeans_mode=True, ) - lut_fast, indices_fast = palettizer_fast._calculate_centroids(weight) + lut_fast, indices_fast = _initialize_and_get_lut_indices(palettizer_fast, weight) palettized_fast = palettizer_fast._palettize(lut_fast, indices_fast, weight) # Regular mode @@ -1200,7 +1190,7 @@ def test_fast_vs_regular_kmeans_with_sensitivities_consistency(self): sensitivities=sensitivities.clone(), enable_fast_kmeans_mode=False, ) - lut_regular, indices_regular = palettizer_regular._calculate_centroids(weight) + lut_regular, indices_regular = _initialize_and_get_lut_indices(palettizer_regular, weight) palettized_regular = palettizer_regular._palettize(lut_regular, indices_regular, weight) # Calculate reconstruction errors @@ -1241,11 +1231,11 @@ def test_uniform_sensitivities_equivalent_to_no_sensitivities(self): palettizer_uniform = _KMeansFakePalettize( **spec.__dict__, sensitivities=sensitivities_uniform ) - lut_uniform, indices_uniform = palettizer_uniform._calculate_centroids(weight) + lut_uniform, indices_uniform = _initialize_and_get_lut_indices(palettizer_uniform, weight) # Without sensitivities palettizer_none = _KMeansFakePalettize(**spec.__dict__, sensitivities=None) - lut_none, indices_none = palettizer_none._calculate_centroids(weight) + lut_none, indices_none = _initialize_and_get_lut_indices(palettizer_none, weight) # Errors should be very close assert torch.equal(lut_none, lut_uniform), ( @@ -1289,7 +1279,7 @@ def test_sensitivities_with_per_channel_scale(self, enable_fast_kmeans): enable_fast_kmeans_mode=enable_fast_kmeans, ) - lut, indices = palettizer._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) palettized = palettizer._palettize(lut, indices, weight) # Verify shape is preserved @@ -1320,7 +1310,7 @@ def test_vector_palettization_per_tensor(self, cluster_dim): palettizer = _KMeansFakePalettize(**spec.__dict__) - lut, indices = palettizer._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) # LUT shape: [1, 1, 2^n_bits, cluster_dim] num_clusters = 2**spec.n_bits @@ -1353,7 +1343,7 @@ def test_vector_palettization_per_grouped_channel(self, cluster_dim, axis): palettizer = _KMeansFakePalettize(**spec.__dict__) - lut, indices = palettizer._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) # LUT shape depends on axis num_clusters = 2**spec.n_bits @@ -1417,7 +1407,7 @@ def test_vector_palettization_incompatible_dim_raises(self, weight_shape, granul palettizer = _KMeansFakePalettize(**spec.__dict__) with pytest.raises(_IncompatibleClusterDimError): - palettizer._calculate_centroids(weight) + _initialize_and_get_lut_indices(palettizer, weight) @pytest.mark.parametrize("lut_dtype", _SUPPORTED_LUT_DTYPES) def test_vector_palettization_with_quantized_lut(self, lut_dtype): @@ -1433,7 +1423,7 @@ def test_vector_palettization_with_quantized_lut(self, lut_dtype): palettizer = _KMeansFakePalettize(**spec.__dict__) - lut, indices = palettizer._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) # LUT shape: [1, 1, num_clusters, cluster_dim] assert lut.shape == (1, 1, 2**spec.n_bits, 2) @@ -1467,7 +1457,7 @@ def test_vector_palettization_with_per_channel_scale(self, granularity): palettizer = _KMeansFakePalettize(**spec.__dict__) - lut, indices = palettizer._calculate_centroids(weight) + lut, indices = _initialize_and_get_lut_indices(palettizer, weight) # Per-channel scale should be populated assert palettizer.per_channel_scale is not None @@ -1495,7 +1485,7 @@ def test_vector_palettization_reconstruction_quality(self): cluster_dim=1, ) palettizer_scalar = _KMeansFakePalettize(**spec_scalar.__dict__) - lut_s, idx_s = palettizer_scalar._calculate_centroids(weight) + lut_s, idx_s = _initialize_and_get_lut_indices(palettizer_scalar, weight) palettized_scalar = palettizer_scalar._palettize(lut_s, idx_s, weight) mse_scalar = torch.mean((weight - palettized_scalar) ** 2) @@ -1506,7 +1496,7 @@ def test_vector_palettization_reconstruction_quality(self): cluster_dim=2, ) palettizer_vector = _KMeansFakePalettize(**spec_vector.__dict__) - lut_v, idx_v = palettizer_vector._calculate_centroids(weight) + lut_v, idx_v = _initialize_and_get_lut_indices(palettizer_vector, weight) palettized_vector = palettizer_vector._palettize(lut_v, idx_v, weight) mse_vector = torch.mean((weight - palettized_vector) ** 2) @@ -1516,3 +1506,378 @@ def test_vector_palettization_reconstruction_quality(self): assert mse_vector < mse_scalar, ( f"Expected Vector MSE: {mse_vector:.4f} to be lower than Scale MSE: {mse_scalar:.4f}" ) + + +class TestDerivedLutProperties: + """`lut`/`quantized_lut`/`lut_quantization_scale`/`lut_quantization_zero_point` + are computed from `centroids` on each access. They read frozen qparams + (never re-observing), so repeated reads are stable and mutually consistent. + """ + + @staticmethod + def _make_palettizer( + *, lut_qspec: QuantizationSpec | None, n_bits: int = 2 + ) -> _KMeansFakePalettize: + spec = PalettizationSpec( + n_bits=n_bits, + granularity=PerTensorGranularity(), + lut_qspec=lut_qspec, + ) + return _KMeansFakePalettize(**spec.__dict__) + + def test_none_before_initialization(self): + """With no centroids yet, every derived property is None.""" + palettizer = self._make_palettizer(lut_qspec=_make_lut_qspec(torch.int8)) + assert palettizer.centroids is None + assert palettizer.lut is None + assert palettizer.quantized_lut is None + assert palettizer.lut_quantization_scale is None + assert palettizer.lut_quantization_zero_point is None + + def test_valid_after_initialize_only(self): + """A single _initialize() (no further quantize_lut() calls) seeds the LUT + observer, so the derived properties are already valid. + """ + palettizer = self._make_palettizer(lut_qspec=_make_lut_qspec(torch.int8)) + palettizer._initialize(torch.randn(8, 8)) + + assert palettizer.lut is not None + assert torch.isfinite(palettizer.lut).all() + assert palettizer.quantized_lut is not None + scale = palettizer.lut_quantization_scale + assert scale is not None and (scale > 0).all() + + def test_passthrough_without_lut_quantizer(self): + """With no lut_qspec, lut is the raw (reshaped) centroids and the + quantization properties are None. + """ + palettizer = self._make_palettizer(lut_qspec=None) + palettizer._initialize(torch.randn(8, 8)) + + assert palettizer.lut is not None + assert palettizer.quantized_lut is None + assert palettizer.lut_quantization_scale is None + assert palettizer.lut_quantization_zero_point is None + + def test_repeated_reads_idempotent_with_moving_average(self): + """Reading the derived properties must not perturb a moving-average LUT + observer -- repeated reads over frozen centroids are byte-identical. + """ + lut_qspec = QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.SYMMETRIC, + qparam_calculator_cls="moving_average", + averaging_constant=0.5, + ) + palettizer = self._make_palettizer(lut_qspec=lut_qspec) + palettizer._initialize(torch.randn(8, 8)) + + # Advance the observer with several distinct centroid snapshots, as a + # training loop's quantize_lut() calls would. + for _ in range(5): + palettizer.centroids = palettizer.centroids + torch.randn_like(palettizer.centroids) + palettizer.quantize_lut(palettizer._raw_lut(palettizer.centroids)) + + # Centroids now frozen: two consecutive reads must match exactly. + first = ( + palettizer.lut, + palettizer.quantized_lut, + palettizer.lut_quantization_scale, + ) + second = ( + palettizer.lut, + palettizer.quantized_lut, + palettizer.lut_quantization_scale, + ) + for a, b in zip(first, second, strict=True): + assert torch.equal(a, b) + + def test_lut_consistent_with_quantized_lut_and_scale(self): + """The dequantized `lut` matches `scale * quantized_lut` (int8 symmetric), + confirming all four properties derive from the same frozen qparams. + """ + palettizer = self._make_palettizer(lut_qspec=_make_lut_qspec(torch.int8)) + palettizer._initialize(torch.randn(8, 8)) + + lut = palettizer.lut + quantized_lut = palettizer.quantized_lut + scale = palettizer.lut_quantization_scale + torch.testing.assert_close( + lut.flatten(), + scale.flatten() * quantized_lut.flatten().float(), + atol=1e-4, + rtol=1e-4, + ) + + +class TestTrainingStrategy: + """The training-strategy config registry and its paired behavior classes. + + Scope is the generic OSS machinery (default strategy only); concrete + strategies defined elsewhere are tested with those strategies. + """ + + def test_default_config_points_at_default_strategy(self): + assert DefaultTrainingConfig._strategy_cls is _DefaultTrainingStrategy + assert issubclass(DefaultTrainingConfig._strategy_cls, TrainingStrategy) + + def test_build_from_dict_unregistered_raises(self): + with pytest.raises(KeyError): + TrainingStrategyConfig.maybe_build_from_dict({"type": "nonexistent"}) + + def test_default_config_builds_default_strategy(self): + assert isinstance(DefaultTrainingConfig().build_strategy(), _DefaultTrainingStrategy) + + def test_maybe_build_from_dict(self): + config = TrainingStrategyConfig.maybe_build_from_dict({"type": "default"}) + assert isinstance(config, DefaultTrainingConfig) + + def test_serialize_injects_type(self): + assert DefaultTrainingConfig().model_dump() == {"type": "default"} + + def test_spec_default_is_default_config(self): + assert isinstance(PalettizationSpec().training_strategy_config, DefaultTrainingConfig) + + def test_spec_parses_dict_via_discriminated_field(self): + spec = PalettizationSpec(training_strategy_config={"type": "default"}) + assert isinstance(spec.training_strategy_config, DefaultTrainingConfig) + + def test_module_builds_strategy_from_config(self): + spec = PalettizationSpec( + n_bits=2, + granularity=PerTensorGranularity(), + training_strategy_config=DefaultTrainingConfig(), + ) + palettizer = _KMeansFakePalettize(**spec.__dict__) + assert isinstance(palettizer._training_strategy, _DefaultTrainingStrategy) + + def test_module_defaults_to_default_strategy(self): + palettizer = _KMeansFakePalettize( + n_bits=2, + lut_qspec=None, + granularity=PerTensorGranularity(), + cluster_dim=1, + enable_per_channel_scale=False, + ) + assert isinstance(palettizer._training_strategy, _DefaultTrainingStrategy) + + def test_default_strategy_train_matches_eval_output(self): + """The default strategy's train_forward() delegates to hard_assign(), so + the forward output is identical in train and eval mode. + """ + spec = PalettizationSpec(n_bits=2, granularity=PerTensorGranularity()) + palettizer = _KMeansFakePalettize(**spec.__dict__) + weight = torch.randn(8, 8) + palettizer._initialize(weight) + + palettizer.train() + out_train = palettizer(weight) + palettizer.eval() + out_eval = palettizer(weight) + assert torch.equal(out_train, out_eval) + + +class TestLazyInitAndStaleness: + """Lazy centroid init and index staleness tracking on _KMeansFakePalettize.""" + + @staticmethod + def _make() -> _KMeansFakePalettize: + spec = PalettizationSpec(n_bits=2, granularity=PerTensorGranularity()) + return _KMeansFakePalettize(**spec.__dict__) + + def test_initial_flags(self): + palettizer = self._make() + assert palettizer._centroids_initialized is False + assert palettizer._indices_stale is True + + def test_first_forward_initializes(self): + palettizer = self._make() + palettizer(torch.randn(8, 8)) + assert palettizer._centroids_initialized is True + assert palettizer._indices_stale is False + + def test_setting_sensitivities_resets_flags(self): + palettizer = self._make() + weight = torch.randn(8, 8) + palettizer(weight) + assert palettizer._centroids_initialized is True + + palettizer.sensitivities = torch.rand_like(weight) + 1.0 + assert palettizer._centroids_initialized is False + assert palettizer._indices_stale is True + + # Next forward re-initializes. + palettizer(weight) + assert palettizer._centroids_initialized is True + assert palettizer._indices_stale is False + + def test_hard_assign_refreshes_when_stale(self): + palettizer = self._make() + weight = torch.randn(8, 8) + palettizer._initialize(weight) # leaves indices fresh + + palettizer._indices_stale = True + indices_before = palettizer.indices + palettizer.hard_assign(weight) + assert palettizer._indices_stale is False + assert palettizer.indices is not indices_before # recomputed + + def test_hard_assign_skips_refresh_when_fresh(self): + palettizer = self._make() + weight = torch.randn(8, 8) + palettizer._initialize(weight) # _indices_stale is False + + indices_before = palettizer.indices + 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. + """ + spec = PalettizationSpec(n_bits=2, granularity=PerTensorGranularity()) + source = _KMeansFakePalettize(**spec.__dict__) + source._initialize(torch.randn(8, 8)) + + target = _KMeansFakePalettize(**spec.__dict__) + assert target._centroids_initialized is False + target.load_state_dict(source.state_dict()) + assert target._centroids_initialized is True + assert target._indices_stale is True + + +class TestQuantizeLutSTE: + """quantize_lut() fake-quantizes the LUT through a straight-through estimator.""" + + @staticmethod + def _make(lut_qspec: QuantizationSpec | None) -> _KMeansFakePalettize: + spec = PalettizationSpec( + n_bits=2, + granularity=PerTensorGranularity(), + lut_qspec=lut_qspec, + ) + return _KMeansFakePalettize(**spec.__dict__) + + def test_gradient_flows_through(self): + """quantize_lut() passes gradient through to the input LUT.""" + palettizer = self._make(_make_lut_qspec(torch.int8)) + lut = torch.randn(1, 4, requires_grad=True) + + palettizer.quantize_lut(lut).sum().backward() + + assert lut.grad is not None + assert torch.count_nonzero(lut.grad) > 0 + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) + def test_output_dtype_matches_input(self, dtype): + palettizer = self._make(_make_lut_qspec(torch.int8)) + lut = torch.randn(1, 4, dtype=dtype) + assert palettizer.quantize_lut(lut).dtype == dtype + + def test_noop_without_lut_quantizer(self): + palettizer = self._make(lut_qspec=None) + lut = torch.randn(1, 4) + assert palettizer.quantize_lut(lut) is lut + + +class TestClustering: + """The k-means clustering utilities: _cluster_to_centroids, _assign_indices, + and _blocks_to_cluster. + """ + + def test_cluster_to_centroids_shapes_scalar(self): + weight = torch.randn(4, 16) + spec = PalettizationSpec(n_bits=2, granularity=PerTensorGranularity(), cluster_dim=1) + palettizer = _KMeansFakePalettize(**spec.__dict__) + + centroids, indices = palettizer._cluster_to_centroids(weight) + # (num_blocks, num_clusters, cluster_dim) + assert centroids.shape == (1, 4, 1) + assert indices.shape == weight.shape + + def test_cluster_to_centroids_shapes_vector(self): + weight = torch.randn(4, 16) + spec = PalettizationSpec(n_bits=2, granularity=PerTensorGranularity(), cluster_dim=2) + palettizer = _KMeansFakePalettize(**spec.__dict__) + + centroids, indices = palettizer._cluster_to_centroids(weight) + assert centroids.shape == (1, 4, 2) + # Vector palettization reduces axis 0 by cluster_dim. + assert indices.shape == (weight.shape[0] // 2, weight.shape[1]) + + def test_indices_consistent_with_centroids(self): + """The indices from _cluster_to_centroids agree with a nearest-centroid + reassignment against the returned centroids. Well-separated values keep + this unambiguous under fast-mode fp16 rounding. + """ + row = torch.tensor([-10.0, 0.0, 10.0, 20.0]) + weight = row.repeat(4, 4) # (4, 16), four well-separated values + spec = PalettizationSpec(n_bits=2, granularity=PerTensorGranularity()) + palettizer = _KMeansFakePalettize(**spec.__dict__) + + centroids, indices = palettizer._cluster_to_centroids(weight) + reassigned = palettizer._assign_indices(weight, centroids) + assert torch.equal(indices, reassigned) + + def test_blocks_to_cluster_raises_on_indivisible_cluster_dim(self): + spec = PalettizationSpec(n_bits=2, granularity=PerTensorGranularity(), cluster_dim=3) + palettizer = _KMeansFakePalettize(**spec.__dict__) + weight_2d = torch.randn(4, 16) # axis-0 size 4 not divisible by cluster_dim 3 + with pytest.raises(_IncompatibleClusterDimError): + palettizer._blocks_to_cluster(weight_2d, axis=0) + + +def _accelerator_device() -> str | None: + """Return an available accelerator device type ("cuda" or "mps"), else None.""" + if torch.cuda.is_available(): + return "cuda" + if torch.backends.mps.is_available(): + return "mps" + return None + + +@pytest.mark.skipif(_accelerator_device() is None, reason="requires a CUDA or MPS accelerator") +def test_device_placement_on_accelerator(): + """LUT quantization and reconstruction device behavior on an accelerator. + + Checks in one pass that: + - centroids follow the model device; ``indices`` stays on CPU (memory), + - ``_raw_lut`` and the ``lut`` property follow the device (not forced to CPU), + - ``quantize_lut`` (the training-side observe path) stays on the input device + with no forced CPU round-trip, + - ``hard_assign`` (eval) gathers against the CPU ``indices`` and returns on the + weight's device -- no accelerator/cpu device mismatch. + """ + device = _accelerator_device() + + spec = PalettizationSpec( + n_bits=2, + granularity=PerTensorGranularity(), + cluster_dim=1, + lut_qspec=_make_lut_qspec(torch.int8), + ) + palettizer = _KMeansFakePalettize(**spec.__dict__) + weight = torch.randn(8, 8, device=device) + + # Clusters on CPU, places centroids on the weight's device, and seeds the LUT + # quantizer on-device via quantize_lut(_raw_lut(centroids)). + palettizer._initialize(weight) + + # centroids follow the model device; the weight-sized indices stay CPU-resident. + assert palettizer.centroids.device.type == device + assert palettizer.indices.device.type == "cpu" + + # _raw_lut and the lut property follow the centroids' device (not forced CPU). + raw_lut = palettizer._raw_lut(palettizer.centroids) + assert raw_lut.device.type == device + assert palettizer.lut.device.type == device + + # Quantizing an on-device LUT stays on-device. + assert palettizer.quantize_lut(raw_lut).device.type == device + + # Eval path: reconstruction gathers against the CPU indices and returns on the + # weight's device. + out = palettizer.hard_assign(weight) + assert out.device.type == device + assert out.shape == weight.shape diff --git a/tests/palettization/test_kmeans_palettizer.py b/tests/palettization/test_kmeans_palettizer.py index 58ff74c..5379ee8 100644 --- a/tests/palettization/test_kmeans_palettizer.py +++ b/tests/palettization/test_kmeans_palettizer.py @@ -27,11 +27,7 @@ default_weight_palettization_spec, ) from coreai_opt.palettization.spec.errors import _IncompatibleGranularityError -from coreai_opt.palettization.spec.fake_palettize import ( - _disable_observer, - _enable_observer, - _FakePalettizeImplBase, -) +from coreai_opt.palettization.spec.fake_palettize import _FakePalettizeImplBase from coreai_opt.quantization.spec import QuantizationScheme, QuantizationSpec @@ -793,30 +789,26 @@ def test_calibration_mode_recomputes_centroids( assert luts_changed, "LUTs were not recomputed after calibration" - def test_calibration_mode_observer_states( + def test_calibration_mode_fake_palett_state( self, simple_conv_linear_model, basic_config, simple_model_input ): - """Test that observer and fake palettize states are managed correctly.""" + """Test that fake palettize state is managed correctly around calibration.""" palettizer = KMeansPalettizer(simple_conv_linear_model, basic_config) prepared_model = palettizer.prepare((simple_model_input,)) - # After prepare, observers should be disabled and fake palettize enabled + # After prepare, fake palettize should be enabled for module in prepared_model.modules(): if isinstance(module, _KMeansFakePalettize): - assert not module.observer_enabled, "Observer should be disabled after prepare" assert module.fake_palett_enabled, "Fake palettize should be enabled after prepare" # Create dummy batch dummy_input = simple_model_input dummy_target = torch.randint(0, 10, (1,)) - # Inside calibration context, observers and fake palettize are disabled + # Inside calibration context, fake palettize is disabled with palettizer.calibration_mode(loss_fn=nn.functional.cross_entropy) as skm: for module in prepared_model.modules(): if isinstance(module, _KMeansFakePalettize): - assert not module.observer_enabled, ( - "Observer should be disabled during calibration" - ) assert not module.fake_palett_enabled, ( "Fake palettize should be disabled during calibration" ) @@ -824,10 +816,9 @@ def test_calibration_mode_observer_states( output = prepared_model(dummy_input) skm.step(output, dummy_target) - # After calibration, observers should be disabled and fake palettize enabled + # After calibration, fake palettize should be enabled again for module in prepared_model.modules(): if isinstance(module, _KMeansFakePalettize): - assert not module.observer_enabled, "Observer should be disabled after calibration" assert module.fake_palett_enabled, ( "Fake palettize should be enabled after calibration" ) @@ -1024,70 +1015,6 @@ def test_calibration_mode_saves_sensitivities_to_path( f"Sensitivity shape for {name} should match parameter shape" ) - def test_stale_centroids_warning_when_sensitivities_set( - self, simple_conv_linear_model, basic_config, simple_model_input, caplog - ): - """ - Test that a warning is logged when sensitivities are set on a - _KMeansFakePalettize module but centroids are not recomputed before use. - """ - palettizer = KMeansPalettizer(simple_conv_linear_model, basic_config) - prepared_model = palettizer.prepare((simple_model_input,)) - - # Find a _KMeansFakePalettize module - fake_palettize_module = None - for module in prepared_model.modules(): - if isinstance(module, _KMeansFakePalettize): - fake_palettize_module = module - break - - assert fake_palettize_module is not None, "Should have a _KMeansFakePalettize module" - - # Verify centroids are not stale initially - assert not fake_palettize_module._centroids_stale - - # Manually set sensitivities without going through calibration_mode - # This simulates a user incorrectly setting sensitivities directly - fake_sensitivities = torch.ones_like(prepared_model.conv.parametrizations.weight.original) - fake_palettize_module.sensitivities = fake_sensitivities - - # Verify centroids are now stale - assert fake_palettize_module._centroids_stale - - # Run a forward pass with observer disabled (default after prepare) - # This should log a warning about stale centroids - with caplog.at_level(logging.WARNING): - _ = prepared_model(simple_model_input) - - # Check that the warning was logged - assert any( - "Sensitivities were updated but centroids have not been recomputed" in record.message - for record in caplog.records - ), "Expected warning about stale centroids was not logged" - - # Now recompute centroids by enabling observer and running a forward pass - prepared_model.apply(_enable_observer) - - with torch.no_grad(): - _ = prepared_model(simple_model_input) - - # Disable observer for normal operation - prepared_model.apply(_disable_observer) - - # Verify centroids are no longer stale - assert not fake_palettize_module._centroids_stale - - # Clear the log and run another forward pass - caplog.clear() - with caplog.at_level(logging.WARNING): - _ = prepared_model(simple_model_input) - - # Verify the warning is NOT logged after recomputing centroids - assert not any( - "Sensitivities were updated but centroids have not been recomputed" in record.message - for record in caplog.records - ), "Warning should not be logged after centroids are recomputed" - def test_save_sensitivities_before_prepare_raises_error( self, simple_conv_linear_model, basic_config, tmp_path ): diff --git a/tests/palettization/test_kmeans_parallel.py b/tests/palettization/test_kmeans_parallel.py index b305ef5..d84abf4 100644 --- a/tests/palettization/test_kmeans_parallel.py +++ b/tests/palettization/test_kmeans_parallel.py @@ -227,7 +227,7 @@ def test_calibration_mode_respects_prepare_num_workers( Runs calibration on two copies of the same model -- one prepared with ``num_workers=1``, the other with ``num_workers=2`` -- and asserts the - post-calibration LUTs, indices, observer/fake_palett state, and model + post-calibration LUTs, indices, fake_palett state, and model outputs match. Same calibration inputs are used for both so any divergence would point to the recompute path itself rather than RNG. """ @@ -258,13 +258,9 @@ def test_calibration_mode_respects_prepare_num_workers( torch.testing.assert_close(seq_fp.indices, par_fp.indices) # After calibration_mode exits, both paths should have restored - # fake_palett=on / observer=off. The parallel path swaps fp_modules - # into parametrization slots, so the subsequent apply(_enable_fake_palett) - # and apply(_disable_observer) must reach the new modules. + # fake_palett=on. assert seq_fp.fake_palett_enabled.item() == 1 assert par_fp.fake_palett_enabled.item() == 1 - assert seq_fp.observer_enabled.item() == 0 - assert par_fp.observer_enabled.item() == 0 with torch.no_grad(): torch.testing.assert_close( diff --git a/tests/palettization/test_pat_schedule.py b/tests/palettization/test_pat_schedule.py new file mode 100644 index 0000000..8ad92a5 --- /dev/null +++ b/tests/palettization/test_pat_schedule.py @@ -0,0 +1,350 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Tests for PAT (palettization-aware training) schedule runtime behavior. + +Covers the PATSchedule config object plus the step-based training control on +KMeansPalettizer: training_mode(), step(), and the _mode state machine shared +with calibration_mode(). +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.nn.utils.parametrize as P +from pydantic import ValidationError + +from coreai_opt.base_model_compressor import _CompressorLifecycle +from coreai_opt.palettization import ( + KMeansPalettizer, + KMeansPalettizerConfig, + ModuleKMeansPalettizerConfig, +) +from coreai_opt.palettization.config import PATSchedule +from coreai_opt.palettization.kmeans.kmeans_fake_palettize import _KMeansFakePalettize +from coreai_opt.palettization.spec import default_weight_palettization_spec + + +class ToyModel(nn.Module): + """Single palettizable Linear layer.""" + + def __init__(self): + super().__init__() + self.linear = nn.Linear(16, 8) + + def forward(self, x): + return self.linear(x) + + +def _example_input() -> tuple[torch.Tensor]: + return (torch.randn(2, 16),) + + +def _prepared_palettizer( + schedule: PATSchedule | None = None, +) -> tuple[KMeansPalettizer, nn.Module]: + """Build a palettizer over ToyModel (optionally with a PAT schedule) and + prepare it. Returns ``(palettizer, prepared_model)``. + """ + config = KMeansPalettizerConfig( + global_config=ModuleKMeansPalettizerConfig( + op_state_spec={"weight": default_weight_palettization_spec()}, + pat_schedule=schedule, + ) + ) + palettizer = KMeansPalettizer(ToyModel(), config) + prepared = palettizer.prepare(_example_input()) + return palettizer, prepared + + +def _fake_palett_modules(model: nn.Module) -> list[_KMeansFakePalettize]: + """Return every _KMeansFakePalettize parametrization in the model.""" + modules = [] + for _, module in model.named_modules(): + if not P.is_parametrized(module): + continue + for parametrizations in module.parametrizations.values(): + for param in parametrizations: + if isinstance(param, _KMeansFakePalettize): + modules.append(param) + return modules + + +def _all_enabled(model: nn.Module) -> bool: + return all(m.fake_palett_enabled[0].item() == 1 for m in _fake_palett_modules(model)) + + +def _all_disabled(model: nn.Module) -> bool: + return all(m.fake_palett_enabled[0].item() == 0 for m in _fake_palett_modules(model)) + + +class TestPATSchedule: + """Pure-logic tests for the PATSchedule config object.""" + + def test_default_enable_fake_palettize_is_zero(self): + assert PATSchedule().enable_fake_palettize == 0 + + @pytest.mark.parametrize( + "threshold, step_count, expected", + [ + (0, 0, True), # active immediately when threshold is 0 + (5, 4, False), # before threshold + (5, 5, True), # exactly at threshold + (5, 6, True), # after threshold + ], + ) + def test_compute_state(self, threshold, step_count, expected): + schedule = PATSchedule(enable_fake_palettize=threshold) + assert schedule._compute_state(step_count) is expected + + def test_negative_threshold_rejected(self): + with pytest.raises(ValidationError): + PATSchedule(enable_fake_palettize=-1) + + def test_frozen(self): + schedule = PATSchedule(enable_fake_palettize=3) + with pytest.raises(ValidationError): + schedule.enable_fake_palettize = 5 + + +class TestTrainingMode: + """Runtime behavior of KMeansPalettizer.training_mode().""" + + def test_training_mode_requires_prepared_model(self): + config = KMeansPalettizerConfig( + global_config=ModuleKMeansPalettizerConfig( + op_state_spec={"weight": default_weight_palettization_spec()}, + ) + ) + palettizer = KMeansPalettizer(ToyModel(), config) + with pytest.raises(RuntimeError, match="Model must be prepared"): + with palettizer.training_mode(): + pass + + def test_entry_trains_exit_evals(self): + palettizer, prepared = _prepared_palettizer() + prepared.eval() + with palettizer.training_mode(): + assert prepared.training is True + assert prepared.training is False + + def test_entry_from_train_stays_train_on_exit(self): + palettizer, prepared = _prepared_palettizer() + prepared.train() + with palettizer.training_mode(): + assert prepared.training is True + assert prepared.training is True + + def test_default_no_schedule_stays_enabled(self): + palettizer, prepared = _prepared_palettizer() + with palettizer.training_mode(): + assert _all_enabled(prepared) + + def test_schedule_below_threshold_gates_off_on_entry(self): + palettizer, prepared = _prepared_palettizer(PATSchedule(enable_fake_palettize=5)) + # prepare() leaves fake palettization enabled... + assert _all_enabled(prepared) + # ...but entering training_mode applies the schedule at step 0 (< 5). + with palettizer.training_mode(): + assert _all_disabled(prepared) + + def test_nested_training_mode_raises(self): + palettizer, _ = _prepared_palettizer() + with palettizer.training_mode(): + with pytest.raises(RuntimeError, match="Cannot enter training_mode"): + with palettizer.training_mode(): + pass + + def test_calibration_inside_training_raises(self): + palettizer, _ = _prepared_palettizer() + with pytest.raises(RuntimeError, match="Cannot enter calibration_mode"): + with palettizer.training_mode(): + with palettizer.calibration_mode(loss_fn=F.mse_loss): + pass + + def test_training_inside_calibration_raises(self): + palettizer, prepared = _prepared_palettizer() + (example,) = _example_input() + with pytest.raises(RuntimeError, match="Cannot enter training_mode"): + with palettizer.calibration_mode(loss_fn=F.mse_loss) as skm: + skm.step(prepared(example), torch.randn(2, 8)) + with palettizer.training_mode(): + pass + + def test_mode_restored_to_idle_on_exception(self): + palettizer, _ = _prepared_palettizer() + with pytest.raises(ValueError, match="boom"): + with palettizer.training_mode(): + assert palettizer._lifecycle is _CompressorLifecycle.TRAINING + raise ValueError("boom") + assert palettizer._lifecycle is _CompressorLifecycle.IDLE + + +class TestStep: + """Runtime behavior of KMeansPalettizer.step().""" + + def test_increments_counter(self): + palettizer, _ = _prepared_palettizer() + with palettizer.training_mode(): + assert palettizer._step_count == 0 + palettizer.step() + assert palettizer._step_count == 1 + palettizer.step() + assert palettizer._step_count == 2 + + def test_outside_training_mode_raises(self): + palettizer, _ = _prepared_palettizer() + with pytest.raises(RuntimeError, match="must be called inside a training_mode"): + palettizer.step() + + def test_counter_monotonic_across_loops(self): + palettizer, _ = _prepared_palettizer() + with palettizer.training_mode(): + palettizer.step() + palettizer.step() + assert palettizer._step_count == 2 + # A second training_mode() loop continues counting, never resets. + with palettizer.training_mode(): + palettizer.step() + assert palettizer._step_count == 3 + + def test_crossing_threshold_enables_at_exact_step(self): + palettizer, prepared = _prepared_palettizer(PATSchedule(enable_fake_palettize=2)) + with palettizer.training_mode(): + assert _all_disabled(prepared) # step 0 < 2 + palettizer.step() + assert _all_disabled(prepared) # step 1 < 2 + palettizer.step() + assert _all_enabled(prepared) # step 2 == threshold + + def test_noop_without_schedule(self): + palettizer, prepared = _prepared_palettizer() + with palettizer.training_mode(): + palettizer.step() # no schedule configured -> does not raise + assert _all_enabled(prepared) + assert palettizer._step_count == 1 + + +class MultiLayerModel(nn.Module): + """Conv2d + two Linear layers, for exercising per-module schedule resolution.""" + + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(1, 2, 3, padding=1) + self.linear1 = nn.Linear(2 * 8 * 8, 8) + self.linear2 = nn.Linear(8, 4) + + def forward(self, x): + x = self.conv(x).flatten(1) + return self.linear2(self.linear1(x)) + + +def _fp_for(model: nn.Module, name: str) -> _KMeansFakePalettize: + """Return the fake-palettize module parametrizing ``.weight``.""" + module = model.get_submodule(name) + for parametrizations in module.parametrizations.values(): + for param in parametrizations: + if isinstance(param, _KMeansFakePalettize): + return param + raise AssertionError(f"no fake-palettize module for {name}") + + +class TestPerModuleSchedule: + """Schedule resolution across the config hierarchy.""" + + def test_name_over_type_over_global(self): + """One test exercising all three levels: module_name_configs beats + module_type_configs beats global_config. + """ + + def _module_config(threshold: int) -> ModuleKMeansPalettizerConfig: + return ModuleKMeansPalettizerConfig( + op_state_spec={"weight": default_weight_palettization_spec()}, + pat_schedule=PATSchedule(enable_fake_palettize=threshold), + ) + + config = KMeansPalettizerConfig( + global_config=_module_config(3), # conv falls through to here + module_type_configs={nn.Linear: _module_config(2)}, # linear2 + module_name_configs={"linear1": _module_config(0)}, # linear1 + ) + palettizer = KMeansPalettizer(MultiLayerModel(), config) + prepared = palettizer.prepare((torch.randn(2, 1, 8, 8),)) + + fp_conv = _fp_for(prepared, "conv") + fp_l1 = _fp_for(prepared, "linear1") + fp_l2 = _fp_for(prepared, "linear2") + + with palettizer.training_mode(): + palettizer.step() # step_count == 1 + # linear1 (name, threshold 0) active; linear2 (type, 2) and conv + # (global, 3) not yet. + assert fp_l1.fake_palett_enabled[0].item() == 1 + assert fp_l2.fake_palett_enabled[0].item() == 0 + assert fp_conv.fake_palett_enabled[0].item() == 0 + + palettizer.step() # step_count == 2 + # linear2 crosses its type-level threshold; conv's global threshold + # (3) still far off -> confirms each layer used its own level. + assert fp_l1.fake_palett_enabled[0].item() == 1 + assert fp_l2.fake_palett_enabled[0].item() == 1 + assert fp_conv.fake_palett_enabled[0].item() == 0 + + palettizer.step() # step_count == 3 + # conv's global threshold is crossed. + assert fp_l1.fake_palett_enabled[0].item() == 1 + assert fp_l2.fake_palett_enabled[0].item() == 1 + assert fp_conv.fake_palett_enabled[0].item() == 1 + + +class MixedModel(nn.Module): + """One palettized Linear feeding a non-palettized Linear head.""" + + def __init__(self): + super().__init__() + self.palettized = nn.Linear(16, 8, bias=False) + self.head = nn.Linear(8, 4) + + def forward(self, x): + return self.head(self.palettized(x)) + + +class TestDefaultStrategyTraining: + """The default training strategy's behavior inside a training_mode() loop.""" + + @pytest.mark.parametrize("use_training_mode_ctx", [True, False]) + def test_gradient_flow_with_and_without_training_mode_context(self, use_training_mode_ctx): + """Training-time gradients must flow through the palettized layer whether + or not the training_mode() context is active: the input and downstream + params receive gradients while the frozen palettized weight receives none + (no unintended gradient path). + """ + config = KMeansPalettizerConfig( + module_name_configs={ + "palettized": ModuleKMeansPalettizerConfig( + op_state_spec={"weight": default_weight_palettization_spec()} + ) + }, + global_config=None, # only the named layer is palettized + ) + palettizer = KMeansPalettizer(MixedModel(), config) + prepared = palettizer.prepare((torch.randn(2, 16),)) + + x = torch.randn(2, 16, requires_grad=True) + if use_training_mode_ctx: + with palettizer.training_mode(): + prepared(x).sum().backward() + else: + prepared.train() + prepared(x).sum().backward() + + # frozen palettized weight gets no gradient (no unintended path) + assert prepared.palettized.parametrizations.weight.original.grad is None + # graph intact: gradients reach the input (through the palettized layer) and downstream + assert x.grad is not None + assert prepared.head.weight.grad is not None diff --git a/tests/test_compression_config.py b/tests/test_compression_config.py index 9655648..c39518e 100644 --- a/tests/test_compression_config.py +++ b/tests/test_compression_config.py @@ -449,9 +449,9 @@ def test_generic_type_extraction(): assert base_cls is None -def test_get_compressor_specific_settings(): +def test_get_fake_module_kwargs(): """ - Test that _get_compressor_specific_settings returns only subclass-defined fields. + Test that _get_fake_module_kwargs returns only subclass-defined fields. """ # Create a config with both base class fields and subclass fields config = ModuleFooCompressionConfig( @@ -469,7 +469,7 @@ def test_get_compressor_specific_settings(): ) # Get compressor-specific settings - settings = config._get_compressor_specific_settings() + settings = config._get_fake_module_kwargs() # Should only contain subclass-defined fields assert "compression_ratio" in settings From 1d00647d4cc56e0bf2482591803fe5497f17fac3 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:47:16 -0700 Subject: [PATCH 2/4] Address review comments --- .../kmeans/kmeans_fake_palettize.py | 14 ++++++++------ .../palettization/spec/fake_palettize.py | 5 +++++ .../palettization/spec/training_strategy.py | 17 +++++++---------- 3 files changed, 20 insertions(+), 16 deletions(-) diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index c02bc1f..402713b 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py @@ -195,10 +195,13 @@ def _initialize(self, weight: torch.Tensor) -> None: self._centroids_initialized = True self._indices_stale = False - 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 + def _maybe_refresh_indices(self, weight: torch.Tensor) -> None: + """Recompute indices from the current centroids if self._indices_stale is true, + without re-clustering. + """ + if self._indices_stale: + self.indices = self._assign_indices(weight, self.centroids).detach() + self._indices_stale = False def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): """Mark centroids as initialized after loading them from a checkpoint.""" @@ -275,8 +278,7 @@ def hard_assign(self, weight: torch.Tensor) -> torch.Tensor: """Nearest-centroid reconstruction against the current centroids, refreshing indices first if stale. """ - if self._indices_stale: - self._refresh_indices(weight) + self._maybe_refresh_indices(weight) return self._palettize(self.lut, self.indices, weight) def _blocks_to_cluster(self, weight_2d: torch.Tensor, axis: int) -> list[torch.Tensor]: diff --git a/src/coreai_opt/palettization/spec/fake_palettize.py b/src/coreai_opt/palettization/spec/fake_palettize.py index 52b3180..52bda3d 100644 --- a/src/coreai_opt/palettization/spec/fake_palettize.py +++ b/src/coreai_opt/palettization/spec/fake_palettize.py @@ -88,6 +88,11 @@ def forward_enabled(self, tensor: torch.Tensor) -> torch.Tensor: """Return the palettized output for ``tensor`` when enabled.""" raise NotImplementedError() + @abstractmethod + def hard_assign(self, weight: torch.Tensor) -> torch.Tensor: + """Return the hard-assigned (deployable) palettized reconstruction of ``weight``.""" + raise NotImplementedError() + @abstractmethod def _palettize( self, lut: torch.Tensor, indices: torch.Tensor, original_weights: torch.Tensor diff --git a/src/coreai_opt/palettization/spec/training_strategy.py b/src/coreai_opt/palettization/spec/training_strategy.py index 2c58b43..e201664 100644 --- a/src/coreai_opt/palettization/spec/training_strategy.py +++ b/src/coreai_opt/palettization/spec/training_strategy.py @@ -16,21 +16,18 @@ from coreai_opt._utils.registry_utils import ConfigRegistryMixin as _ConfigRegistryMixin if TYPE_CHECKING: - from coreai_opt.palettization.kmeans.kmeans_fake_palettize import _KMeansFakePalettize + from coreai_opt.palettization.spec.fake_palettize import _FakePalettizeImplBase class TrainingStrategy(ABC): """Contract for a fake-palettize module's training-time forward pass.""" @abstractmethod - def train_forward(self, module: _KMeansFakePalettize, weight: torch.Tensor) -> torch.Tensor: - """Compute this training step's output for ``weight``. - - May mutate ``module.centroids`` in place (must set - ``module._indices_stale = True`` if it does). Use ``module.quantize_lut()`` - to fake-quantize intermediate centroids against the module's - configured LUT quantizer. Must leave ``module.centroids`` such that - ``module.hard_assign(weight)`` produces a sensible result at eval time. + def train_forward(self, module: _FakePalettizeImplBase, weight: torch.Tensor) -> torch.Tensor: + """Return the training-time output for ``weight``. + + Defines how the palettized weight behaves during training (e.g. frozen, + straight-through, soft assignment). """ raise NotImplementedError @@ -50,7 +47,7 @@ class _DefaultTrainingStrategy(TrainingStrategy): ``TrainingStrategy``. """ - def train_forward(self, module: _KMeansFakePalettize, weight: torch.Tensor) -> torch.Tensor: + def train_forward(self, module: _FakePalettizeImplBase, weight: torch.Tensor) -> torch.Tensor: return module.hard_assign(weight) From a5d23ad39907fe90eee88c92a2f3eeb6435c0708 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Fri, 7 Aug 2026 16:37:13 -0700 Subject: [PATCH 3/4] Restore compatibility with legacy palettizer checkpoints --- .../kmeans/kmeans_fake_palettize.py | 23 ++++- .../palettization/spec/fake_palettize.py | 25 +++++ .../test_kmeans_fake_palettize.py | 91 ++++++++++++++++++- 3 files changed, 135 insertions(+), 4 deletions(-) diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index 402713b..b48db2a 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py @@ -156,8 +156,11 @@ def sensitivities(self, value: torch.Tensor | None) -> None: self._indices_stale = True def ensure_initialized(self, tensor: torch.Tensor) -> None: - """Cluster centroids on first use; disable on an incompatible tensor.""" - if self._centroids_initialized: + """Cluster centroids on first use; disable on an incompatible tensor. + + Re-clusters on every call while ``observer_enabled`` is set. + """ + if self._centroids_initialized and self.observer_enabled[0] == 0: return try: self._initialize(tensor.detach()) @@ -204,12 +207,26 @@ def _maybe_refresh_indices(self, weight: torch.Tensor) -> None: self._indices_stale = False def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): - """Mark centroids as initialized after loading them from a checkpoint.""" + """Load centroids from a checkpoint, reconstructing them from a legacy + ``lut`` buffer when present. + """ + lut_key, centroids_key = prefix + "lut", prefix + "centroids" + 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]) super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) if self.centroids is not None: self._centroids_initialized = True self._indices_stale = True + 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. + """ + ungrouped_dim = 0 if self.granularity.axis == 1 else 1 + centroids = lut.squeeze(-1) if self.cluster_dim == 1 else lut + centroids = centroids.squeeze(ungrouped_dim) + return centroids.unsqueeze(-1) if self.cluster_dim == 1 else centroids + @property def lut(self) -> torch.Tensor | None: """Lookup table dequantized from ``centroids`` using the LUT diff --git a/src/coreai_opt/palettization/spec/fake_palettize.py b/src/coreai_opt/palettization/spec/fake_palettize.py index 52bda3d..b2e09ca 100644 --- a/src/coreai_opt/palettization/spec/fake_palettize.py +++ b/src/coreai_opt/palettization/spec/fake_palettize.py @@ -29,6 +29,7 @@ class _FakePalettizeImplBase(CompressionSimulatorBase, nn.Module): indices: torch.Tensor per_channel_scale: torch.Tensor | None fake_palett_enabled: torch.Tensor + observer_enabled: torch.Tensor def __init__( self, @@ -47,6 +48,11 @@ def __init__( self.enable_per_channel_scale = enable_per_channel_scale self.register_buffer("fake_palett_enabled", torch.tensor([1], dtype=torch.uint8)) + # Non-persistent (kept out of new checkpoints); when set to 1 (at runtime or + # via a legacy checkpoint) the forward pass re-clusters centroids every call. + self.register_buffer( + "observer_enabled", torch.tensor([0], dtype=torch.uint8), persistent=False + ) self._disabled = False self.register_buffer("indices", None) @@ -156,6 +162,25 @@ def _load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ) + obs_key = prefix + "observer_enabled" + if obs_key in state_dict: + self.observer_enabled.copy_(state_dict[obs_key]) + if obs_key in unexpected_keys: + unexpected_keys.remove(obs_key) + + # Accept (and ignore) buffers from legacy checkpoints that are now + # derived properties, so old state dicts load without unexpected-key errors. + # ``lut`` is consumed by the subclass to reconstruct ``centroids``. + for legacy_name in ( + "lut", + "quantized_lut", + "lut_quantization_scale", + "lut_quantization_zero_point", + ): + prefixed_key = prefix + legacy_name + if prefixed_key in unexpected_keys: + unexpected_keys.remove(prefixed_key) + def enable_fake_palett(self, enabled: bool = True) -> None: self.fake_palett_enabled[0] = 1 if enabled else 0 diff --git a/tests/palettization/test_kmeans_fake_palettize.py b/tests/palettization/test_kmeans_fake_palettize.py index a08357b..d068629 100644 --- a/tests/palettization/test_kmeans_fake_palettize.py +++ b/tests/palettization/test_kmeans_fake_palettize.py @@ -943,7 +943,6 @@ def test_state_dict_after_initialization(self, enable_per_channel_scale): assert palettizer.lut.shape == (1, 1, 4, 1) # 2^2 = 4 clusters assert palettizer.lut.dtype == weight.dtype - # Check state_dict contains proper values state_dict = palettizer.state_dict() @@ -1047,6 +1046,96 @@ def test_save_load_state_dict_preserves_palettization(self): os.unlink(temp_path) +class TestLegacyCheckpointCompat: + """Backward compatibility with pre-refactor checkpoints, which stored + ``observer_enabled``/``lut``/``quantized_lut``/``lut_quantization_scale``/ + ``lut_quantization_zero_point`` as buffers instead of deriving them. + """ + + @staticmethod + def _legacy_state_dict(ref: _KMeansFakePalettize) -> dict: + """Rewrite an initialized module's state dict into the pre-refactor + layout: drop ``centroids`` and add the old ``lut`` (+ quant) buffers. + """ + sd = ref.state_dict() + sd.pop("centroids") + sd["lut"] = ref.lut + if ref.quantized_lut is not None: + sd["quantized_lut"] = ref.quantized_lut + sd["lut_quantization_scale"] = ref.lut_quantization_scale + if ref.lut_quantization_zero_point is not None: + sd["lut_quantization_zero_point"] = ref.lut_quantization_zero_point + return sd + + @pytest.mark.parametrize("lut_dtype", [None, torch.int8]) + def test_load_reconstructs_centroids_from_legacy_lut(self, lut_dtype): + spec = PalettizationSpec( + n_bits=2, + granularity=PerTensorGranularity(), + lut_qspec=_make_lut_qspec(lut_dtype), + enable_per_channel_scale=False, + ) + weight = torch.randn(4, 8) + ref = _KMeansFakePalettize(**spec.__dict__) + ref.forward(weight) + ref_out = ref.hard_assign(weight) + + legacy_sd = self._legacy_state_dict(ref) + assert "centroids" not in legacy_sd and "lut" in legacy_sd + + loaded = _KMeansFakePalettize(**spec.__dict__) + if lut_dtype is not None: + # Loading a LUT-quantized checkpoint requires the LUT quantizer's + # qparams buffers to be sized first (pre-existing quantization + # requirement, independent of this reconstruction path). + loaded.forward(weight) + loaded.load_state_dict(legacy_sd) # strict=True: must not raise on legacy keys + + # centroids reconstructed and the derived LUT reproduces the stored one + assert loaded.centroids is not None + assert torch.equal(loaded.lut, ref.lut) + if lut_dtype is None: + # plain palettization is bit-exact end to end + assert torch.equal(loaded.hard_assign(weight), ref_out) + + def test_observer_enabled_defaults_off_and_loads(self): + spec = PalettizationSpec( + n_bits=2, granularity=PerTensorGranularity(), enable_per_channel_scale=False + ) + p = _KMeansFakePalettize(**spec.__dict__) + assert p.observer_enabled.item() == 0 + # Non-persistent: not serialized into new checkpoints. + assert "observer_enabled" not in p.state_dict() + + p.forward(torch.randn(4, 8)) + sd = p.state_dict() + assert "observer_enabled" not in sd + # A legacy checkpoint that carries observer_enabled=1 is honored on load. + sd["observer_enabled"] = torch.tensor([1], dtype=torch.uint8) + loaded = _KMeansFakePalettize(**spec.__dict__) + loaded.load_state_dict(sd) + assert loaded.observer_enabled.item() == 1 + + def test_observer_enabled_recomputes_every_forward(self): + spec = PalettizationSpec( + n_bits=2, granularity=PerTensorGranularity(), enable_per_channel_scale=False + ) + p = _KMeansFakePalettize(**spec.__dict__) + low = torch.linspace(0.0, 1.0, 32).reshape(4, 8) + high = torch.linspace(10.0, 11.0, 32).reshape(4, 8) + + # observer off (default): centroids are computed once and not recomputed. + p.forward(low) + first = p.centroids.clone() + p.forward(high) + assert torch.equal(p.centroids, first) + + # observer on: re-clusters against the new tensor on every forward. + p.observer_enabled[0] = 1 + p.forward(high) + assert not torch.equal(p.centroids, first) + + class TestSensitivityBasedPalettization: """Test cases for sensitivity-based (weighted) k-means palettization.""" From 4fab50ef404eb9e96ee9e9cabe599937edfb0f7b Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Tue, 11 Aug 2026 14:33:49 -0700 Subject: [PATCH 4/4] Address review comments --- .../kmeans/kmeans_fake_palettize.py | 5 +++ tests/conftest.py | 10 ++++++ .../test_kmeans_fake_palettize.py | 34 +++++++------------ 3 files changed, 27 insertions(+), 22 deletions(-) diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index b48db2a..218c233 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py @@ -222,6 +222,11 @@ 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. """ + if lut.ndim != 4: + raise ValueError( + "Legacy 'lut' buffer must be 4D (num_blocks_axis0, num_blocks_axis1, " + f"num_clusters, cluster_dim); got shape {tuple(lut.shape)}." + ) ungrouped_dim = 0 if self.granularity.axis == 1 else 1 centroids = lut.squeeze(-1) if self.cluster_dim == 1 else lut centroids = centroids.squeeze(ungrouped_dim) diff --git a/tests/conftest.py b/tests/conftest.py index bb0ed47..71e4afa 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -133,6 +133,16 @@ def temp_dir(): yield tmpdir +@pytest.fixture +def accelerator_device() -> str: + """The available accelerator device type ("cuda" or "mps"); skip if neither.""" + if torch.cuda.is_available(): + return "cuda" + if torch.backends.mps.is_available(): + return "mps" + pytest.skip("requires a CUDA or MPS accelerator") + + @pytest.fixture(scope="function") def mnist_pretrained_model(custom_test_mnist_model): """Load the committed 1-epoch MNIST checkpoint into a fresh model.""" diff --git a/tests/palettization/test_kmeans_fake_palettize.py b/tests/palettization/test_kmeans_fake_palettize.py index d068629..48b292e 100644 --- a/tests/palettization/test_kmeans_fake_palettize.py +++ b/tests/palettization/test_kmeans_fake_palettize.py @@ -1633,6 +1633,7 @@ def test_valid_after_initialize_only(self): assert palettizer.lut is not None assert torch.isfinite(palettizer.lut).all() assert palettizer.quantized_lut is not None + assert not palettizer.quantized_lut.dtype.is_floating_point scale = palettizer.lut_quantization_scale assert scale is not None and (scale > 0).all() @@ -1681,21 +1682,20 @@ def test_repeated_reads_idempotent_with_moving_average(self): for a, b in zip(first, second, strict=True): assert torch.equal(a, b) - def test_lut_consistent_with_quantized_lut_and_scale(self): - """The dequantized `lut` matches `scale * quantized_lut` (int8 symmetric), - confirming all four properties derive from the same frozen qparams. + def test_lut_small_quantization_error(self): + """The dequantized `lut` reconstructs the raw centroids within one + quantization step. """ palettizer = self._make_palettizer(lut_qspec=_make_lut_qspec(torch.int8)) palettizer._initialize(torch.randn(8, 8)) - lut = palettizer.lut - quantized_lut = palettizer.quantized_lut - scale = palettizer.lut_quantization_scale + raw = palettizer._raw_lut(palettizer.centroids) + scale = palettizer.lut_quantization_scale.max().item() torch.testing.assert_close( - lut.flatten(), - scale.flatten() * quantized_lut.flatten().float(), - atol=1e-4, - rtol=1e-4, + palettizer.lut.squeeze(), + raw.squeeze(), + atol=scale, + rtol=0, ) @@ -1917,17 +1917,7 @@ def test_blocks_to_cluster_raises_on_indivisible_cluster_dim(self): palettizer._blocks_to_cluster(weight_2d, axis=0) -def _accelerator_device() -> str | None: - """Return an available accelerator device type ("cuda" or "mps"), else None.""" - if torch.cuda.is_available(): - return "cuda" - if torch.backends.mps.is_available(): - return "mps" - return None - - -@pytest.mark.skipif(_accelerator_device() is None, reason="requires a CUDA or MPS accelerator") -def test_device_placement_on_accelerator(): +def test_device_placement_on_accelerator(accelerator_device): """LUT quantization and reconstruction device behavior on an accelerator. Checks in one pass that: @@ -1938,7 +1928,7 @@ def test_device_placement_on_accelerator(): - ``hard_assign`` (eval) gathers against the CPU ``indices`` and returns on the weight's device -- no accelerator/cpu device mismatch. """ - device = _accelerator_device() + device = accelerator_device spec = PalettizationSpec( n_bits=2,