From 8a180a08489d15948eb0793e1cdea51bd94c94f5 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 11 Aug 2026 14:45:21 -0700 Subject: [PATCH 1/2] up --- backends/apple/coreai/palettizer/__init__.py | 11 ++ .../apple/coreai/palettizer/palettizer.py | 107 ++++++++++++ backends/apple/coreai/test/test_palettizer.py | 161 ++++++++++++++++++ 3 files changed, 279 insertions(+) create mode 100644 backends/apple/coreai/palettizer/__init__.py create mode 100644 backends/apple/coreai/palettizer/palettizer.py create mode 100644 backends/apple/coreai/test/test_palettizer.py diff --git a/backends/apple/coreai/palettizer/__init__.py b/backends/apple/coreai/palettizer/__init__.py new file mode 100644 index 00000000000..4c1134b3336 --- /dev/null +++ b/backends/apple/coreai/palettizer/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.backends.apple.coreai.palettizer.palettizer import CoreAIPalettizer + +__all__ = [ + "CoreAIPalettizer", +] diff --git a/backends/apple/coreai/palettizer/palettizer.py b/backends/apple/coreai/palettizer/palettizer.py new file mode 100644 index 00000000000..26bbb2b5cfd --- /dev/null +++ b/backends/apple/coreai/palettizer/palettizer.py @@ -0,0 +1,107 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Core AI palettization front-end for ExecuTorch. + +Palettization replaces a weight with a small learned lookup table plus an index +per element, rather than a scale and zero-point. It sits alongside +quantization:: + + prepare -> (optional sensitivity calibration) -> finalize -> export + +Unlike :class:`~executorch.backends.apple.coreai.quantizer.quantizer.CoreAIQuantizer`, +:meth:`finalize` returns an eager ``nn.Module`` whose weights carry a +parametrization, not an fx graph. The compression becomes a ``lut_to_dense`` op +only once the model is exported, so nothing here needs to run before +``torch.export`` for the backend's sake; the ordinary +``to_edge_transform_and_lower`` path lowers it. + +Palettization is weight-only by construction, so there is no activation +counterpart to configure. + +``coreai_opt`` (the ``coreai-optimization`` package) is imported lazily so this +module stays importable in environments without it installed. +""" + +from __future__ import annotations + +from typing import Any, Callable, Optional, Sequence + +from torch import nn + + +class CoreAIPalettizer: + """K-means palettization front-end for Core AI, wrapping ``coreai_opt``. + + Args: + model: the ``nn.Module`` to palettize. + config: a ``coreai_opt.palettization.KMeansPalettizerConfig``. Defaults + to ``coreai_opt``'s own default when ``None``. Ready-made + configurations are available as + ``KMeansPalettizerConfig.presets.w4()``, ``.w6()`` and ``.w8()``. + """ + + def __init__(self, model: nn.Module, config: Optional[Any] = None) -> None: + from coreai_opt.palettization import KMeansPalettizer + + self._palettizer = KMeansPalettizer(model, config) + self._prepared: Optional[nn.Module] = None + + def prepare( + self, + example_inputs: Sequence[Any], + sensitivity_path: Optional[str] = None, + num_workers: int = 1, + ) -> nn.Module: + """Cluster the weights and attach the palette parametrizations. + + Args: + example_inputs: Sample inputs used to trace the model. + sensitivity_path: Sensitivities saved by a previous + :meth:`calibration_mode` run, to weight the clustering. + num_workers: Parallelism for clustering, which dominates the cost + on large models. + """ + self._prepared = self._palettizer.prepare( + tuple(example_inputs), + sensitivity_path=sensitivity_path, + num_workers=num_workers, + ) + return self._prepared + + def calibration_mode(self, *, loss_fn: Callable, sensitivity_path=None): + """Context manager for sensitivity-weighted clustering (SqueezeLLM). + + Optional, and unlike the quantizer's calibration this is not needed for + a plain run: k-means clusters the weights, which requires no data. + Supplying a loss function collects squared gradients as per-element + sensitivities so clustering favors the weights that matter most. + """ + return self._palettizer.calibration_mode( + loss_fn=loss_fn, sensitivity_path=sensitivity_path + ) + + def training_mode(self): + """Context manager for palettization-aware training.""" + return self._palettizer.training_mode() + + def finalize(self, mmap_dir=None) -> nn.Module: + """Produce the export-ready palettized model. + + Returns an eager ``nn.Module``, not a graph: the palette surfaces as a + ``lut_to_dense`` op when the result is exported. + + Args: + mmap_dir: Directory to memory-map large palettes through, rather + than holding them in memory. + """ + from coreai_opt.common import ExportBackend + + if self._prepared is None: + raise RuntimeError("Call prepare() before finalize().") + return self._palettizer.finalize( + backend=ExportBackend.CoreAI, mmap_dir=mmap_dir + ) diff --git a/backends/apple/coreai/test/test_palettizer.py b/backends/apple/coreai/test/test_palettizer.py new file mode 100644 index 00000000000..52c532357a0 --- /dev/null +++ b/backends/apple/coreai/test/test_palettizer.py @@ -0,0 +1,161 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""K-means palettization through the full ExecuTorch lowering path. + +Palettization compresses a weight into a lookup table plus an index per +element, so the observables differ from quantization: a ``lut_to_dense`` op and +``lut``/``indices`` state, rather than a scale, zero-point and quantized +weight. +""" + +import unittest +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from executorch.backends.apple.coreai import ( + CoreAIPartitioner, + get_default_compile_config, + get_default_passes, +) +from executorch.backends.apple.coreai.palettizer import CoreAIPalettizer +from executorch.exir import to_edge_transform_and_lower + + +def _model(): + return nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 64)).eval() + + +def _op_count(exported_program, prefix): + """How many call_function targets start with ``prefix``. + + Matched by prefix because targets carry an overload suffix, e.g. + ``lut_to_dense.default``. + """ + return sum( + 1 + for node in exported_program.graph.nodes + if node.op == "call_function" + and getattr(node.target, "__name__", str(node.target)).startswith(prefix) + ) + + +@dataclass(frozen=True) +class PalettizeCase: + """One palettization preset and the palette it should produce.""" + + name: str + preset: str + n_bits: int + + @property + def palette_entries(self) -> int: + return 2**self.n_bits + + def config(self): + from coreai_opt.palettization import KMeansPalettizerConfig + + return getattr(KMeansPalettizerConfig.presets, self.preset)() + + +class PalettizationTest(unittest.TestCase): + """The presets Apple ships: w4, w6 and w8.""" + + CASES = ( + PalettizeCase(name="w4", preset="w4", n_bits=4), + PalettizeCase(name="w6", preset="w6", n_bits=6), + PalettizeCase(name="w8", preset="w8", n_bits=8), + ) + + def setUp(self): + self.example_inputs = (torch.randn(2, 64),) + + def _finalize(self, case): + palettizer = CoreAIPalettizer(_model(), case.config()) + palettizer.prepare(self.example_inputs) + return palettizer.finalize() + + def _export(self, case): + return torch.export.export(self._finalize(case), self.example_inputs) + + def _pte_size(self, case): + lowered = to_edge_transform_and_lower( + self._export(case), + transform_passes=get_default_passes(), + partitioner=[CoreAIPartitioner()], + compile_config=get_default_compile_config(), + ) + return len(bytes(lowered.to_executorch().buffer)) + + def test_finalize_returns_an_eager_module(self): + """Unlike the quantizer's convert(), which returns an fx graph.""" + finalized = self._finalize(self.CASES[0]) + self.assertIsInstance(finalized, nn.Module) + self.assertFalse(hasattr(finalized, "graph")) + + def test_weights_become_a_lookup_table(self): + for case in self.CASES: + with self.subTest(case.name): + state = self._export(case).state_dict + luts = [k for k in state if k.endswith(".lut")] + indices = [k for k in state if k.endswith(".indices")] + self.assertTrue(luts, f"no palette in {sorted(state)}") + self.assertEqual(len(luts), len(indices)) + # The dense weight is emptied out in favour of the palette. + originals = [ + v for k, v in state.items() if k.endswith("weight.original") + ] + self.assertTrue(all(v.numel() == 0 for v in originals)) + + def test_palette_size_follows_n_bits(self): + """A k-bit palette holds 2**k entries; guards a silently ignored preset.""" + for case in self.CASES: + with self.subTest(case.name): + state = self._export(case).state_dict + lut = next(v for k, v in state.items() if k.endswith(".lut")) + self.assertEqual(lut.shape[-2], case.palette_entries) + + def test_export_emits_lut_to_dense(self): + for case in self.CASES: + with self.subTest(case.name): + self.assertGreater( + _op_count(self._export(case), "lut_to_dense"), 0 + ) + + def test_lowers_through_to_executorch(self): + for case in self.CASES: + with self.subTest(case.name): + self.assertGreater(self._pte_size(case), 0) + + def test_narrower_palettes_produce_smaller_programs(self): + sizes = {case.name: self._pte_size(case) for case in self.CASES} + self.assertLess(sizes["w4"], sizes["w6"], sizes) + self.assertLess(sizes["w6"], sizes["w8"], sizes) + + +class PalettizerApiTest(unittest.TestCase): + def test_finalize_before_prepare_is_rejected(self): + palettizer = CoreAIPalettizer(_model()) + with self.assertRaisesRegex(RuntimeError, "prepare"): + palettizer.finalize() + + def test_calibration_is_not_required(self): + """Plain k-means clusters weights, so it needs no calibration data. + + The quantizer's calibration collects activation ranges and is + mandatory; here it only supplies sensitivity weights. + """ + palettizer = CoreAIPalettizer( + _model(), PalettizationTest.CASES[0].config() + ) + palettizer.prepare((torch.randn(2, 64),)) + self.assertIsInstance(palettizer.finalize(), nn.Module) + + +if __name__ == "__main__": + unittest.main() From ee2ab2795a75265eb7e77c92e55293aa100a35c2 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 12 Aug 2026 13:41:52 -0700 Subject: [PATCH 2/2] up --- .../apple/coreai/palettizer/palettizer.py | 24 ++++-- backends/apple/coreai/test/test_palettizer.py | 84 +++++++++++++++++-- 2 files changed, 91 insertions(+), 17 deletions(-) diff --git a/backends/apple/coreai/palettizer/palettizer.py b/backends/apple/coreai/palettizer/palettizer.py index 26bbb2b5cfd..9d433155a94 100644 --- a/backends/apple/coreai/palettizer/palettizer.py +++ b/backends/apple/coreai/palettizer/palettizer.py @@ -72,22 +72,27 @@ def prepare( ) return self._prepared + def _require_prepared(self, method: str) -> None: + if self._prepared is None: + raise RuntimeError(f"Call prepare() before {method}().") + def calibration_mode(self, *, loss_fn: Callable, sensitivity_path=None): """Context manager for sensitivity-weighted clustering (SqueezeLLM). Optional, and unlike the quantizer's calibration this is not needed for a plain run: k-means clusters the weights, which requires no data. Supplying a loss function collects squared gradients as per-element - sensitivities so clustering favors the weights that matter most. + sensitivities so clustering favors the weights that matter most. The + block must call ``step(output, target)`` at least once. + + There is no ``training_mode`` counterpart: ``KMeansPalettizer`` does not + support training-time compression. """ + self._require_prepared("calibration_mode") return self._palettizer.calibration_mode( loss_fn=loss_fn, sensitivity_path=sensitivity_path ) - def training_mode(self): - """Context manager for palettization-aware training.""" - return self._palettizer.training_mode() - def finalize(self, mmap_dir=None) -> nn.Module: """Produce the export-ready palettized model. @@ -100,8 +105,11 @@ def finalize(self, mmap_dir=None) -> nn.Module: """ from coreai_opt.common import ExportBackend - if self._prepared is None: - raise RuntimeError("Call prepare() before finalize().") - return self._palettizer.finalize( + self._require_prepared("finalize") + finalized = self._palettizer.finalize( backend=ExportBackend.CoreAI, mmap_dir=mmap_dir ) + # coreai_opt clears the model's prepared marker, so a second call would + # otherwise pass this guard and fail deeper with a worse message. + self._prepared = None + return finalized diff --git a/backends/apple/coreai/test/test_palettizer.py b/backends/apple/coreai/test/test_palettizer.py index 52c532357a0..9d46e165aab 100644 --- a/backends/apple/coreai/test/test_palettizer.py +++ b/backends/apple/coreai/test/test_palettizer.py @@ -12,8 +12,10 @@ weight. """ +import tempfile import unittest from dataclasses import dataclass +from pathlib import Path import torch import torch.nn as nn @@ -28,6 +30,8 @@ def _model(): + """Deterministic, so two builds can be compared to each other.""" + torch.manual_seed(0) return nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 64)).eval() @@ -45,6 +49,14 @@ def _op_count(exported_program, prefix): ) +def _first_lut(finalized): + """The first palette in a finalized model, for comparing clusterings.""" + for name, tensor in finalized.state_dict().items(): + if name.endswith(".lut"): + return tensor + raise AssertionError("no palette in the finalized model") + + @dataclass(frozen=True) class PalettizeCase: """One palettization preset and the palette it should produce.""" @@ -110,22 +122,31 @@ def test_weights_become_a_lookup_table(self): originals = [ v for k, v in state.items() if k.endswith("weight.original") ] + self.assertEqual( + len(originals), + len(luts), + f"expected one emptied weight per palette in {sorted(state)}", + ) self.assertTrue(all(v.numel() == 0 for v in originals)) def test_palette_size_follows_n_bits(self): - """A k-bit palette holds 2**k entries; guards a silently ignored preset.""" + """A k-bit palette holds 2**k entries; guards a silently ignored preset. + + Every palette is checked, not just the first: the model has two linears, + and a preset reaching only one of them would still be wrong. + """ for case in self.CASES: with self.subTest(case.name): state = self._export(case).state_dict - lut = next(v for k, v in state.items() if k.endswith(".lut")) - self.assertEqual(lut.shape[-2], case.palette_entries) + luts = [v for k, v in state.items() if k.endswith(".lut")] + self.assertTrue(luts, "no palette found") + for lut in luts: + self.assertEqual(lut.shape[-2], case.palette_entries) def test_export_emits_lut_to_dense(self): for case in self.CASES: with self.subTest(case.name): - self.assertGreater( - _op_count(self._export(case), "lut_to_dense"), 0 - ) + self.assertGreater(_op_count(self._export(case), "lut_to_dense"), 0) def test_lowers_through_to_executorch(self): for case in self.CASES: @@ -144,18 +165,63 @@ def test_finalize_before_prepare_is_rejected(self): with self.assertRaisesRegex(RuntimeError, "prepare"): palettizer.finalize() + def test_calibration_before_prepare_is_rejected(self): + palettizer = CoreAIPalettizer(_model()) + with self.assertRaisesRegex(RuntimeError, "before calibration_mode"): + palettizer.calibration_mode(loss_fn=lambda *_: torch.tensor(0.0)) + + def test_finalize_twice_is_rejected(self): + """coreai_opt clears its own prepared marker, so the guard must too.""" + palettizer = CoreAIPalettizer(_model(), PalettizationTest.CASES[0].config()) + palettizer.prepare((torch.randn(2, 64),)) + palettizer.finalize() + with self.assertRaisesRegex(RuntimeError, "prepare"): + palettizer.finalize() + + def test_no_training_mode(self): + """KMeansPalettizer does not support training-time compression. + + Forwarding to it would return a context manager that raises only on + entry, advertising a capability that cannot work. + """ + self.assertFalse(hasattr(CoreAIPalettizer, "training_mode")) + def test_calibration_is_not_required(self): """Plain k-means clusters weights, so it needs no calibration data. The quantizer's calibration collects activation ranges and is mandatory; here it only supplies sensitivity weights. """ - palettizer = CoreAIPalettizer( - _model(), PalettizationTest.CASES[0].config() - ) + palettizer = CoreAIPalettizer(_model(), PalettizationTest.CASES[0].config()) palettizer.prepare((torch.randn(2, 64),)) self.assertIsInstance(palettizer.finalize(), nn.Module) + def test_sensitivity_calibration_changes_the_palette(self): + """Weighted k-means must actually use the collected sensitivities.""" + sample = (torch.randn(2, 64),) + config = PalettizationTest.CASES[0].config() + + plain = CoreAIPalettizer(_model(), config) + plain.prepare(sample) + plain_lut = _first_lut(plain.finalize()) + + with tempfile.TemporaryDirectory() as d: + path = str(Path(d) / "sensitivity") + weighted = CoreAIPalettizer(_model(), config) + prepared = weighted.prepare(sample) + with weighted.calibration_mode( + loss_fn=lambda out, target: (out - target).pow(2).mean(), + sensitivity_path=path, + ) as calibrate: + calibrate.step(prepared(*sample), torch.randn(2, 64)) + weighted_lut = _first_lut(weighted.finalize()) + + self.assertEqual(plain_lut.shape, weighted_lut.shape) + self.assertFalse( + torch.equal(plain_lut, weighted_lut), + "sensitivity-weighted clustering produced an identical palette", + ) + if __name__ == "__main__": unittest.main()