Palettization, pruning, joint compression, and mixed precision¶
Part 9 · Core AI: compression and numeric formats · Reference 02
Release scope (2026-10-06 audit). Current compression behavior is pinned to
coreai-opt==0.3.0 (2026-09-25), implementation commit 189612be. Python is >=3.11,<3.14;
base requirements are torch>=2.8.0 and torchao>=0.15.0, with no upper bounds.
The source defines tested pairs Torch 2.8.0/2.9.1 with TorchAO 0.15.0, Torch 2.10.0 with
TorchAO 0.16.0, and Torch 2.11–2.14 with TorchAO 0.18.0. Import-time compatibility warnings
are not a guarantee that another pair works. A C++ compiler and Ninja are required for vendored
k-means JIT compilation. coremltools is optional: the coreml extra requires
coremltools>=8.3 and numpy>=2,<2.4.
The coreai extra pins coreai-core==1.0.0b3, coreai-torch==0.4.3 and includes
scikit-learn>=1.7.2. b3 publishes cp314 wheels, but this compression package still excludes
Python 3.14. coreai-models at db63a2d8 has a narrower export profile: torch==2.9.0,
torchao<0.18. Native fixtures here use Python 3.12, NumPy 2.4.6 and both Torch 2.11.0/AO 0.18.0
and Torch 2.9.0/AO 0.17.0. Core AI artifacts target OS 27. Source review and fixture outcomes are
recorded separately in the refresh evidence.
Published model benchmarks retain their original dates and environments.
See the shared evidence conventions. API citations and runtime checks attest their named source revision or fixture.
What this covers¶
Apply palettization, pruning, or combined compression through the supported configuration APIs. Validate each stage and the exported model; storage reduction alone does not establish runtime speed or retained quality.
What this does not cover¶
Related references: guide 01.
What you need¶
pip install coreai-opt # or: uv pip install coreai-opt
pip install 'coreai-opt[coreai]' # adds coreai-core==1.0.0b3, coreai-torch==0.4.3, scikit-learn
- A C++ compiler on the host, at runtime. This is not optional for palettization and it is not
a build-time-only dependency.
coremltoolswas dropped as a runtime dependency in commitedd4720; the 1-D k-means it used to provide is now a vendored C++ core JIT-compiled bytorch.utils.cpp_extension.load()on first use. §2.3 covers what that means in practice. - A representative sample of your real inputs, if you are doing anything with activations. For
weight-only palettization and pruning the
example_inputstuple only needs to be shape-correct, because both are data-free. For joint compression it must be real data. - An evaluation metric you trust, wired up before you start. Every decision in this guide is a trade. Apple's own compression-exploration agent skill refuses to run without one, and it is right to.
- Patience for k-means on a large model. Clustering a 600 M-parameter encoder is not instant.
§7.2 covers
num_workers, which is the one lever that matters and which session 325 never mentions.
Contents¶
- Lookup tables are a different idea, and the ANE is why
- The palettizer in eight lines
PalettizationSpec: six fields, and what each one costs- The three schemes, with diagrams
- ⚠️ The ANE rank-5 ceiling
- Sizing: what a bit-width actually buys
- Determinism, workers, and fast k-means mode
- Sensitivity-weighted k-means (SqueezeLLM)
lut_qspec: quantizing the palette itself- What
finalize()emits - Pruning: the technique nobody presented
- Program-level compression:
palettize_weightsandsparsify_weights - Joint compression
- Mixed precision
- Choosing per layer: the sweep, the Debugger, and SAM3
- Apple's PSNR acceptance gates
- Community-measured findings, labelled
- The worked examples the repo ships
- ⚠️ Silent failures, consolidated
- Numbers, attributed
- Quick reference
- Sources and evidence ledger
1. Lookup tables are a different idea, and the ANE is why¶
1.1 Two families, one API¶
coreai-opt ships three compression techniques, and the package's own enum is the cleanest
statement of the taxonomy:
✅ VERIFIED —
src/coreai_opt/common.py:110-134:
Those three are not variations on a theme. They throw away different information:
| What is stored | What is thrown away | Reconstruction | |
|---|---|---|---|
| Quantization | integer codes + scale (+ zero_point) |
precision, uniformly across the value range | x ≈ (q − zp) · scale — an arithmetic formula |
| Palettization | integer indices + a palette of 2^n_bits centroids |
everything except which cluster a weight belongs to | x ≈ LUT[index] — a table lookup |
| Pruning | a binary mask + the surviving weights | entire weights, set exactly to zero | x' = x · mask — nothing to reconstruct |
The crucial difference between the first two is where the compression error comes from. Quantization imposes a uniform grid whether or not your weights are uniformly distributed — and neural-network weights are famously not; they are heavy-tailed and clustered near zero. Every grid point in the tails is a grid point wasted. Palettization runs k-means on the weights you actually have and puts its centroids where the mass is.
This is why, at the same nominal bit-width, the two techniques do not produce the same error. See §17.1 for a community measurement of exactly how large that gap is at per-channel granularity.
The API surface is deliberately parallel, which session 325 says out loud:
✅ VERIFIED — WWDC26 session 325, "Dive into Core AI model authoring and optimization", 325:243–244: "As before, I construct a
KMeansPalettizersimilar to theQuantizer, and pass it the model and config. Then, I prepare and finalize."
That parallel is real and it is enforced by a shared base class:
✅ VERIFIED —
src/coreai_opt/base_model_compressor.py.Quantizer,KMeansPalettizerandMagnitudePrunerall derive from_BaseModelCompressorand all implement__init__(model, config)→prepare(...)→finalize(...).calibration_mode()andtraining_mode()are optional and raiseNotImplementedErroron compressors that do not support them, with the message"{cls} does not implement training_mode(). This compressor doesn't support training time compression."
The capability matrix is where they stop being parallel:
Quantizer |
KMeansPalettizer |
MagnitudePruner |
|
|---|---|---|---|
| Execution modes | GRAPH (default) and EAGER | EAGER only | EAGER only |
| Weights | ✅ | ✅ | ✅ |
| Activations | ✅ | ✗ (weight-only, enforced) | ✗ (weight-only, enforced) |
calibration_mode() |
✅ (activation ranges) | ✅ (different meaning — sensitivities, §8) | ✗ |
training_mode() / schedules |
✅ QAT | ✅ PATSchedule; frozen default strategy | own training loop + step() |
step() |
✅ (QAT schedule) | ✅ PATSchedule, inside training_mode() |
✅ (sparsity schedule) |
mmap_dir on finalize |
✅ (GRAPH/EAGER + CoreAI only) | ✅ (CoreAI only) | ✗ |
| Presets | w8 w4 w4_per_block |
w4 w6 w8 |
none |
✅ VERIFIED — the eager-only constraint for palettization is stated in the docs (
docs/src/palettization/config.md:119): "Palettization supports eager mode only." It is not a soft preference; there is no_graph/directory underpalettization/orpruning/, onlyquantization/_graph/.
The practical consequence of eager palettization is that palettization uses PyTorch's
__torch_function__ protocol to intercept ops rather than torch.export + torchao PT2E. That
means it tolerates dynamic control flow that would break graph mode, but it also means it gets
none of graph mode's conveniences: no Conv+BN folding, no shared-observer logic, and op names in
your config are module-qualified call sites (linear1.linear) rather than global aten node
names (linear_1). If you inspect a model to find names for a config, inspect it in the mode you
will compress in. Guide 01 §13 covers ModelInspector.
1.2 Why lookup tables suit the Neural Engine¶
Session 325's justification is one sentence, and it names iOS specifically:
✅ VERIFIED — WWDC26 session 325, 325:241–242: "For compression, I apply 4-bit palettization with per-channel scales to the two encoders. There is a preset available for this, but I use the lower-level representation here to showcase the APIs. This lookup-table-based compression, is well-suited for power efficiency on iOS."
(The "with per-channel scales" clause in that quote is contradicted by the shipped code. §5 is about that contradiction, which turns out to be the most useful thing in this guide.)
Apple does not explain why on stage. The reason is legible from the Neural Engine's documented
constraints, which are written down in Apple's own model-authoring agent skill:
✅ VERIFIED —
skills/skills/model-authoring/references/neural_engine_rules.md, vendored inapple/coreai-models: the Neural Engine's supported dtypes are fp16 / int8 / int16, and fp32 falls back to GPU/CPU. Max tensor rank is 5. Shapes must be fully static. The last axis must be contiguous and 64-byte aligned.
Put those together and the argument writes itself:
- The ANE is a fixed-function engine. It runs a closed set of hardware ops — convolution, normalisation, pooling, elementwise — and cannot execute arbitrary code. Any weight format it supports has to be one the hardware knows how to expand.
- The ANE runs fp16 compute. Whatever format your weights are stored in, they arrive at the MAC arrays as fp16. So the question is never "can the engine do 4-bit arithmetic" — it is "how cheaply can the engine turn 4 stored bits into an fp16 value".
- A table lookup is the cheapest possible answer.
LUT[index]is an indexed read from a tiny table that fits in local memory. There is no multiply, no add, no zero-point subtraction. A blockwise-affine dequantisation needs a multiply and (if asymmetric) a subtract per weight, plus a scale tensor streamed alongside the weights. - Inference at this scale is memory-bandwidth-bound, not compute-bound. Session 330 says so directly: ✅ VERIFIED — WWDC26 session 330, 330:21–23: "state-of-the-art machine learning models are getting larger. The inference stage is typically memory bandwidth bound, so compressing the weights becomes necessary both to better fit models into memory and to save memory bandwidth." Fewer bytes read is less DRAM traffic is less energy. That is what "power efficiency" means here — it is not a figure of speech.
There is corroborating evidence from a source that is not Apple's marketing: the community model zoo's own rule of thumb.
🟢 COMMUNITY-MEASURED —
notes/repos/john-rocky-models.md, single-author archive with self-declared uncontrolled benchmarks; treat as a strong hint, not a specification: "statically-compiled ANE execution requires palettized (LUT) weights — blockwise-linear int4 is a GPU-only format there. If you aren't explicitly targeting ANE, target GPU and move on."
If that is right — and it is consistent with everything Apple ships — then the choice between quantization and palettization is not primarily a quality choice. It is a compute-unit choice. Which is exactly how Apple's own export tooling treats it.
1.3 The strongest evidence: Apple's shipped presets split by platform¶
The apple/coreai-models LLM exporter has two preset families, and they do not overlap.
✅ VERIFIED —
python/src/coreai_models/export/presets.py:macOSDEFAULT_MACOS_COMPRESSION_PRESET = "4bit" DEFAULT_IOS_COMPRESSION_PRESET = "4bit_weight_palettized_group32"MACOS_PRESETS:"none","4bit"— where4bitis a quantization config (int4,symmetric_with_clipping,per_blockblock_size=32axis=1). iOSIOS_PRESETS:"none","4bit_weight_palettized_group8","4bit_weight_palettized_group32"— both palettization configs ({"n_bits": 4, "granularity": {"type": "per_grouped_channel", "axis": 0, "group_size": 8 or 32}}).
And the CLI enforces the split rather than merely defaulting to it:
✅ VERIFIED —
python/src/coreai_models/llm/export.py,_resolve_export_config(export.py:263-368). Passing the wrong family raises:RuntimeError("macOS quantization preset provided, but platform is iOS.")andRuntimeError("iOS palettization preset provided, but platform is macOS."). The YAML loader (_load_compression_config_object, export.py:163-237) applies the same rule to hand-written configs: a file whose top-level key iskmeans_palettization_configrequires--platform iOS; one whose top key isquantization_configrequires--platform macOS.
That is not a style guide. That is Apple's production LLM exporter refusing to build a palettized macOS bundle or a quantized iOS bundle. Take the hint.
1.4 Where palettization sits in the pipeline¶
Identical to quantization — it is an optional stage that consumes an nn.Module and produces an
nn.Module, inserted before conversion:
┌──────────────────────── you are here ────────────────────────┐
nn.Module ──► re-author ──► KMeansPalettizer(model, config) │
(Part 10) .prepare(example_inputs) ← k-means runs HERE │
.finalize(backend=CoreAI) ← weights freed HERE │
└─────────────────────────────────────────────────────────────┘
──► torch.export(...).run_decompositions(coreai_torch.get_decomp_table())
──► cast_to_16_bit_precision(exported_program) ← compress FIRST, cast SECOND
──► coreai_torch.TorchConverter().add_exported_program(...).to_coreai()
──► rewrite already completed inside to_coreai() ──► ai_program.save_asset(Path("model.aimodel")) (Part 8)
The property that makes this workable is that every compressor output is itself a PyTorch model. You can score a palettized model in PyTorch, on your own eval harness, before you have converted anything. Apple's docs make the stronger point that for evaluation you should not finalize at all:
✅ VERIFIED —
KMeansPalettizer.finalizedocstring: "Only callfinalizewhen exporting to a target backend. For torch-based evaluation, use the model returned byprepare()directly rather than callingfinalize." Apple's compression-exploration agent skill repeats this as an operational rule: "Do not callfinalize()" during a sweep, "Calibration is not needed for weight-only compression." (skills/skills/model-compression-exploration/SKILL.md.)
That is not a stylistic preference either — see §10.2 for what finalize() does to your weights.
2. The palettizer in eight lines¶
2.1 The minimum viable palettization¶
import copy
import torch
import torch.nn as nn
from coreai_opt import ExportBackend
from coreai_opt.palettization import KMeansPalettizer, KMeansPalettizerConfig
model = MyModel().eval()
float_model = copy.deepcopy(model) # prepare() mutates in place — see §2.4
example_inputs = (torch.randn(1, 3, 224, 224),) # MUST be a tuple; shape-correct is enough
config = KMeansPalettizerConfig.presets.w4() # 4-bit, per-grouped-channel, axis 0, group 16
palettizer = KMeansPalettizer(model, config)
prepared = palettizer.prepare(example_inputs, num_workers=8) # k-means runs here
score(prepared) # evaluate in PyTorch, before export
finalized = palettizer.finalize(backend=ExportBackend.CoreAI) # default backend; destructive
✅ VERIFIED — exact signatures,
src/coreai_opt/palettization/kmeans/palettizer.py:class KMeansPalettizer(_BasePalettizer, _EagerCompressionComponentBuilderMixin): def __init__(self, model: nn.Module, config: KMeansPalettizerConfig | None = None) def prepare(self, example_inputs: tuple[torch.Tensor], sensitivity_path: str | None = None, num_workers: int = 1) -> nn.Module @contextmanager def calibration_mode(self, model=None, *, loss_fn: Callable, sensitivity_path: str | None = None) def finalize(self, model: nn.Module | None = None, backend: ExportBackend = ExportBackend.CoreAI, *, mmap_dir: str | PathLike[str] | None = None) -> nn.Module def save_sensitivities(self, path: str) -> None @classmethod def get_op_type_resolver(cls) -> Callable[[Callable], str | None]
Three things about that signature list are worth pausing on, because they differ from Quantizer:
num_workersis aprepare()parameter, not a config field. It is the single biggest wall-clock lever in the whole palettization workflow and session 325 never mentions it. §7.2.sensitivity_pathappears inprepare()andcalibration_mode(). These are the two ends of the same feature:calibration_mode()writes the file,prepare()reads it. §8.- 0.3.0 adds
training_mode()andstep(), plusPATScheduleon module configs andtraining_strategy_speconPalettizationSpec. The default_DefaultTrainingStrategyuses frozen one-shot centroids/indices: palettized weights receive no gradient; other parameters can adapt around them. It does not automatically learn centroids or perform differentiable palettization. Register a customTrainingStrategyto change that behavior.
PATSchedule(enable_fake_palettize=N) postpones activation. When enabled after warm-up, the
implementation re-clusters warmed-up weights. step() advances the schedule at your chosen cadence.
Source: palettization/spec/training_strategy.py, config/palettization_config.py, and
kmeans/palettizer.py at 189612be. Choose the training strategy explicitly when describing results.
2.2 The preset roster, and what w4() expands to¶
✅ VERIFIED —
src/coreai_opt/palettization/config/_presets/kmeans_palettizer_config.pyandmodule_kmeans_palettizer_config.py. The same three presets exist on bothKMeansPalettizerConfig.presetsandModuleKMeansPalettizerConfig.presets:
| Preset | Signature | Resulting PalettizationSpec |
|---|---|---|
w4 |
w4(*, axis: int = 0, group_size: int = 16) |
n_bits=4, PerGroupedChannelGranularity(axis=axis, group_size=group_size) |
w6 |
w6(*, axis: int = 0, group_size: int = 16) |
n_bits=6, PerGroupedChannelGranularity(axis=axis, group_size=group_size) |
w8 |
w8() |
n_bits=8, PerTensorGranularity() |
Applied to both the "weight" and "in_proj_weight" state keys (see §2.5).
Note what is not here. There is no w2, no w1, no w3, no vector-palettization preset, and no
preset that sets lut_qspec or enable_per_channel_scale. Everything below 4 bits, everything
involving cluster_dim, and everything involving joint compression, you build by hand.
Note also the asymmetry between w4/w6 and w8: the two low-bit presets default to
per-grouped-channel granularity, while w8 is per-tensor. That is a considered choice —
at 8 bits a single 256-entry palette for a whole tensor is usually fine and the LUT overhead of
per-channel palettes would be significant; at 4 bits a single 16-entry palette for a whole tensor
is usually not fine. §6 has the arithmetic.
presets.w4() expanded by hand, so you can see every field it sets:
from coreai_opt.palettization import (
KMeansPalettizerConfig, ModuleKMeansPalettizerConfig, PalettizationSpec,
)
from coreai_opt.palettization.spec import PerGroupedChannelGranularity
spec = PalettizationSpec(
n_bits=4, # palette = 2**4 = 16 entries
lut_qspec=None, # fp LUT (default)
granularity=PerGroupedChannelGranularity(axis=0, group_size=16),
cluster_dim=1, # scalar k-means
enable_per_channel_scale=False, # ← see §5 before changing this
)
config = KMeansPalettizerConfig(
global_config=ModuleKMeansPalettizerConfig(
op_state_spec={"weight": spec, "in_proj_weight": spec},
enable_fast_kmeans_mode=True, # default; §7.3
rounding_precision=4, # default; §7.3
),
)
That is equivalent to KMeansPalettizerConfig.presets.w4() in effect. Whether the preset
literally constructs it this way is an implementation detail; the spec fields and the two state
keys are ✅ verified from the preset source and from OpKMeansPalettizerConfig.get_default_state_spec.
2.3 The vendored _kmeans1d dependency, and why your build box needs a compiler¶
This is the part of the install that surprises people, so it gets its own section.
✅ VERIFIED — changelog fragment
changelog.d/31.changed, landed in commitedd4720: "Replace the coremltools-based 1D k-means used by palettization with a vendored C++ core that is JIT-compiled at runtime viatorch.utils.cpp_extension.coremltoolsis no longer a runtime dependency (it is now an optional dependency, installable via thecoremlextra). This requires a C++ compiler to be available on the host at runtime."✅ VERIFIED —
src/coreai_opt/deps/_kmeans1d/containscore.py+_core.cpp. The JIT load happens atcore.py:30-55viatorch.utils.cpp_extension.load()with flags["-std=c++11", "-O2", "-DNDEBUG"], plus-stdlib=libc++on darwin.ninja>=1.11andsetuptools>=42are declared as runtime dependencies inpyproject.toml:51-62precisely because of this.
The vendored library is MIT-licensed kmeans1d. The reason a specialised 1-D k-means exists at
all is worth knowing: clustering scalars is not the general k-means problem. In one dimension the
optimal clustering can be computed exactly by dynamic programming rather than approximated by
Lloyd's algorithm, which is both faster and deterministic. That determinism is a property you
inherit — and lose the moment you set cluster_dim > 1 and move to genuine multi-dimensional
k-means (§7.1).
Practical consequences:
- First palettization in a fresh environment pays a compile. It is a one-time cost per
environment, cached by
torch.utils.cpp_extension. In CI, expect the first palettization job to be slower than the rest and do not mistake it for a hang. - A container without build-essential / Xcode CLT will fail at
prepare(), not atpip install. If you are baking a Docker image for a compression pipeline, install a compiler. - Quantization-only workflows never hit this, which is why teams that quantize for months and then try palettization once get a confusing failure on a box they thought was working.
2.4 prepare() mutates in place, and eager mode makes that worse¶
Shared with quantization, but sharper here.
✅ VERIFIED — the compressor lifecycle doc (
docs/src/introduction/how_to_use_coreaiopt.md) and_BaseModelCompressor:prepare()"may modify in place — use the returned model". Re-preparing raisesRuntimeError("Model has already been prepared. Cannot re-prepare a prepared model.").
In eager mode — which is the only mode palettization has — compression is implemented with
PyTorch parametrizations registered on the module in place. The consequence is blunt: after
prepare(), reading .weight on either the original object or the returned object gives
you the fake-palettized value, because they are the same object. If you want the float weights for
a comparison, copy.deepcopy the model before prepare(), or save the tensors you care about.
float_model = copy.deepcopy(model) # do this FIRST
prepared = KMeansPalettizer(model, config).prepare(example_inputs)
# model is prepared is True (same object). float_model still has dense weights.
prepare() also handles train/eval mode for you on the way in and not on the way out:
✅ VERIFIED —
docs/src/introduction/how_to_use_coreaiopt.md:39,79: "You don't need to put the model in.eval()or.train()before callingprepare()— the API runs the trace internally in eval mode and restores the original mode when it returns." and "The finalized model inherits the current training mode, so call.eval()on it before running inference or downstream conversion."
2.5 The two default state keys, and the weight-only enforcement¶
✅ VERIFIED —
src/coreai_opt/palettization/config/palettization_config.py:
"in_proj_weight" is nn.MultiheadAttention's fused QKV parameter. If you write your own
op_state_spec and only list "weight", you will silently leave every nn.MultiheadAttention in
your model uncompressed — a large omission in a transformer. Copy both keys.
The weight-only nature is enforced by validation, not by convention:
✅ VERIFIED —
WeightOnlyModuleValidationMixin/WeightOnlyOpValidationMixinraiseValueErrorif you setop_input_spec,op_output_spec,module_input_specormodule_output_spec: "{cls} does not support {key}. This is a weight-only compression type that only supportsop_state_specandmodule_state_spec."
So the only way to combine palettized weights with quantized activations is joint compression (§13). There is no single config that expresses both.
2.6 Which ops are palettizable¶
✅ VERIFIED —
src/coreai_opt/palettization/kmeans/supported_ops_registry.py(_KMeansPalettizerSupportedOpsRegistry), registry key → torch function:
key function conv1d/conv2d/conv3dF.conv1d/F.conv2d/F.conv3dconv_transpose1d/conv_transpose2d/conv_transpose3dF.conv_transpose1d/2d/3dlinearF.linearmulti_head_attention_forwardF.multi_head_attention_forward
ConvTranspose1d/2d/3dsupport was added in 0.2.1 (CHANGELOG). The registry'sregisterdecorator validates at runtime that the registered class subclasses_PalettizationSupportMixin, raisingTypeErrorotherwise.
Read that list for what it excludes: F.embedding is not on it. Neither is any gather, any
normalisation parameter, any bias, and nothing you wrote by hand as raw tensor arithmetic.
This has two important consequences.
One: your norm and RoPE parameters stay full precision automatically. That is usually what you
want, and it means you do not need .without(nn.LayerNorm, nn.RMSNorm)-style exclusions the way you
do with quantization. A community formulation of the same observation:
🟢 COMMUNITY-MEASURED —
notes/repos/john-rocky-models.md: "k-means palettizesF.linear/F.convweights only, so RMSNorm/RoPE params stay full precision automatically."
Two: embedding tables are your problem. In a modern LLM the embedding / lm_head matrix is
frequently the single largest tensor in the model, and palettization cannot touch it. Apple's own
iOS exporter handles this by explicitly excluding embeddings from the palettization config and
quantizing them separately:
✅ VERIFIED —
python/src/coreai_models/export/presets.py: both iOS palettization presets excludetorch.nn.modules.sparse.Embeddingandcoreai_models.primitives.ios.embedding.LoadEmbeddings. And frommodels/README.md:54: "AlliOSpalettization presets quantize the Embedding to 8-bit per tensor by default." The implementation isquantize_per_tensorinprimitives/ios/quantization.py— symmetric,nbits=8only,scale = max|x| / 127, with the scale clamped to a minimum of1e-6. The exporter exposes--disable-embedding-quantization-iosto keep the embedding in float32.
That is the pattern to copy: palettize the projections, quantize the embedding, and account for the embedding separately in your bits-per-weight arithmetic. Note that the published BPW figures for Apple's iOS models carry an asterisk for exactly this reason — see §20.
3. PalettizationSpec: six fields, and what each one costs¶
QuantizationSpec has nine fields. PalettizationSpec has six in 0.3.0, including the training strategy.
The first five control the representation; the sixth controls training-time behavior.
✅ VERIFIED —
src/coreai_opt/palettization/spec/spec.py:86-93:And the factory (class PalettizationSpec(CompressionSpec): # pydantic BaseModel, frozen=True, extra="forbid" n_bits: Literal[1, 2, 3, 4, 6, 8] = 4 lut_qspec: QuantizationSpec | None = None granularity: PalettizationGranularity = PerTensorGranularity() cluster_dim: PositiveInt = 1 enable_per_channel_scale: bool = False training_strategy_spec: TrainingStrategySpec = DefaultTrainingSpec()default_weight_palettization_spec()) returns exactly those defaults:n_bits=4, lut_qspec=None, PerTensorGranularity(), cluster_dim=1, enable_per_channel_scale=False.
Note frozen=True, extra="forbid", inherited from CompressionSpec. You cannot mutate a spec after
construction and you cannot typo a field name — PalettizationSpec(nbits=4) raises rather than
silently doing nothing. To vary one field, use pydantic's model_copy:
base = PalettizationSpec(n_bits=4, granularity=PerGroupedChannelGranularity(axis=0, group_size=32))
six_bit = base.model_copy(update={"n_bits": 6})
3.1 n_bits — six legal values, and two conspicuous absences¶
✅ VERIFIED —
Literal[1, 2, 3, 4, 6, 8]. Docstring: "Number of bits used for palette indices. Determines palette size (2^n_bits entries). Must be one of {1, 2, 3, 4, 6, 8}."
There is no 5 and no 7. This trips people building sweep scripts who assume range(1, 9). The
same set appears at the program level (_VALID_N_BITS = {1,2,3,4,6,8} in
coreai_utils/passes/weight_palettization.py) and in sparsify_weights' palettize_nbits
argument, so it is a format-level constraint, not a validation quirk of one entry point.
The palette size is 2**n_bits:
n_bits |
palette entries | index storage | typical use |
|---|---|---|---|
| 1 | 2 | 1 bit/weight | research / binary nets. Expect catastrophic quality. |
| 2 | 4 | 2 bits/weight | ~8× compression. Apple's skill calls the quality "usually unacceptable" — §16. |
| 3 | 8 | 3 bits/weight | rare; useful only as a mixed-precision floor for insensitive layers |
| 4 | 16 | 4 bits/weight | the default and the one Apple ships on iOS |
| 6 | 64 | 6 bits/weight | the safety valve when 4 is not enough. SAM3's text encoder uses this. |
| 8 | 256 | 8 bits/weight | ~2× compression, near-lossless for most vision models |
3.2 granularity — two classes, not three¶
Quantization has three granularity classes. Palettization has two.
✅ VERIFIED —
coreai_opt.palettization.spec.__all__=PalettizationGranularity, PalettizationSpec, PerGroupedChannelGranularity, PerTensorGranularity, default_weight_palettization_spec, plusDefaultTrainingSpec,TrainingStrategy, andTrainingStrategySpecin 0.3.0. Registry keys (used in YAML):"per_tensor","per_grouped_channel".class PerGroupedChannelGranularity(PalettizationGranularity): axis: int | None = Field(default=None, ge=0, le=1) group_size: intnum_blocks_to_cluster = weight.shape[axis] // group_size. Raises the internal_IncompatibleGranularityErrorifaxisisNoneor out of range, or if the shape is not divisible bygroup_size.
Two things to notice.
axis is constrained to 0 or 1. Field(ge=0, le=1). You cannot palettize per-group along a
spatial axis of a conv kernel. Axis 0 is the output-channel axis for Conv* and Linear; axis 1 is
the input-channel axis. Apple's shipped iOS presets and Apple's SAM3 recipe both use axis=0.
There is no PerChannelGranularity — but you can express it. group_size=1 gives
shape[0] // 1 == shape[0] groups, i.e. one palette per channel. That is arithmetic from the
verified formula, not a separate API. And it is what the community found works best:
🟢 COMMUNITY-MEASURED —
notes/repos/john-rocky-models.md, single-author archive, uncontrolled benchmarks: "Per-channel (group_size=1) basically always wins; at per-channel, k-means beats quantization by ~15–19 dB at both 8-bit and 4-bit. Per-tensor palettization can be worse than per-channel quantization." ⚠️ Note the cost side, which that quote omits: one LUT per output channel is a lot of LUT. §6.3 does the arithmetic, and it is why Apple's own presets usegroup_size=8/16/32rather than1.
⚠️ Import hazard. PerTensorGranularity exists twice with the same spelling:
coreai_opt.quantization.spec.PerTensorGranularity and
coreai_opt.palettization.spec.PerTensorGranularity. They are different classes and they are not
interchangeable. If you import the wrong one into a PalettizationSpec, pydantic will reject it;
if you import the wrong one into lut_qspec (which takes a quantization spec), validation raises
ValueError. Prefer qualified imports in any file that touches both:
from coreai_opt.palettization.spec import PerTensorGranularity as PalettPerTensor
from coreai_opt.quantization.spec import PerTensorGranularity as QuantPerTensor
3.3 cluster_dim — the fractional-bit lever¶
✅ VERIFIED — docstring: "The dimension of centroids for each lookup table... When
cluster_dim > 1, it indicates 2-D clustering, and eachcluster_dimlength of weight vectors along the output channel are palettized using the same 2-D centroid." Arithmetic fromdocs/src/palettization/basics.md:16:n_bits=4, cluster_dim=2⇒ k-means on 2-D data with 16 centroids, LUT shape16×2, effective bits per weight = 4/2 = 2.
This is the only way to get below 1 bit per stored index without going to n_bits=1, and it is
covered in detail in §4.3. Two hard constraints:
✅ VERIFIED —
enable_fast_kmeans_mode=True(the default) withcluster_dim > 1raisesValueError("enable_fast_kmeans_mode is not supported when cluster_dim > 1. ..."). You must setenable_fast_kmeans_mode=False. ✅ VERIFIED — at the program level,palettize_weights(..., enable_per_channel_scale=True, cluster_dim>1)raisesValueError. The two features do not compose.
3.4 enable_per_channel_scale — the field with the trap¶
✅ VERIFIED — docstring: "When set to True, weights are normalized along the output channels using per-channel scales before being palettized."
The idea is sound and standard: if channel A's weights span ±0.01 and channel B's span ±3.0, a single shared palette wastes almost all its entries. Normalise each channel to a common range first, palettize the normalised weights, and store one fp16 scale per channel to undo it at runtime.
The default is False. Leave it False if you are targeting the Neural Engine. §5 is entirely
about why, and it is the most important section in this guide.
3.5 lut_qspec — quantizing the palette¶
✅ VERIFIED —
_SUPPORTED_LUT_DTYPES = {torch.int8, torch.uint8, torch.float8_e4m3fn, torch.float8_e5m2}.validate_lut_qspecraisesValueErrorif the LUT dtype is not in that set, iflut_qspec.granularityis not the quantizationPerTensorGranularity, or iflut_qspec.qformulation == MINVAL("Uselut_qspec.qformulation=ZPinstead."). FP8 dtypes additionally require symmetric quantization.
Covered in §9. Short version: quantizing the LUT is what unlocks the integer execution path, and it is mandatory for joint compression to be worth doing.
3.6 The two module-level fields that are not on the spec¶
Two knobs live on ModuleKMeansPalettizerConfig rather than on PalettizationSpec, because they
control the clustering process rather than the format:
✅ VERIFIED —
src/coreai_opt/palettization/config/palettization_config.py:class ModuleKMeansPalettizerConfig(WeightOnlyModuleValidationMixin, ModuleCompressionConfig[OpKMeansPalettizerConfig, PalettizationSpec]): # @final enable_fast_kmeans_mode: bool = True rounding_precision: PositiveInt = 4enable_fast_kmeans_modedocstring: "enables optimizations for faster K-means clustering by rounding the weights before clustering if data is in float16 range. If weight dtype is float32, weights are cast to float16 and then rounded."rounding_precisionis the number of decimal places, default 4.
§7.3 covers when to turn that off.
3.7 The config hierarchy is the same one you already know¶
KMeansPalettizerConfig uses the identical three-level machinery as QuantizerConfig — because
it is literally the same generic base class, CompressionConfig[...].
✅ VERIFIED —
src/coreai_opt/palettization/config/palettization_config.py:
Everything guide 01 §4–§5 says applies unchanged:
- Precedence:
module_name_configs>module_type_configs>global_config. module_name_configskeys are regexes matched withre.fullmatch."detector"does not match"detector.decoder.linear";"detector.*"does.module_type_configskeys must be fully-qualified class names ("torch.nn.modules.linear.Linear"), or the class object itself."torch.nn.Linear"raisesValueError(f"Expected fully-qualified name, got {module_type}").Nonemeans "leave this alone." Omitting a field applies defaults; explicitly passingNonedisables compression for that scope. This is the mechanism for "compress everything except X".only_for(...)/without(...)chainable helpers exist and returnSelf.only_forcannot be called twice.KMeansPalettizerConfigandModuleKMeansPalettizerConfigare@finaland raiseTypeErroron subclassing.- YAML:
KMeansPalettizerConfig.from_yaml(path)/.from_dict({...}), top-level keykmeans_palettization_config, withpalettization_specavailable as an anchor host.
The two mechanisms you will actually reach for:
# Compress everything EXCEPT the detector subtree (the SAM3 move — guide 01 §13)
config = KMeansPalettizerConfig.presets.w4(group_size=32)
config.set_module_name("detector.*", None)
# Or: compress ONLY the two encoders, nothing else
config = (KMeansPalettizerConfig.presets.w4(group_size=32)
.only_for("image_encoder.*", "text_encoder.*"))
✅ VERIFIED —
only_forsemantics (config/compression_config.py): it disablesglobal_configand deep-copies it onto each target. Calling it twice raisesValueError("only_for requires a non-disabled global_config to redistribute as per-module overrides. If you've already called only_for or set_global(None), pass all targets in one only_for(...) call instead of chaining.")
4. The three schemes, with diagrams¶
The coreai-opt docs give each scheme a diagram. Here they are as ASCII, with the tensor shapes
made explicit, because the shapes are what determine whether the ANE will take your model.
Throughout: a Linear weight W of shape [C_out, C_in] = [512, 1024], so 524,288 weights, and
n_bits = 4 (palette size 16).
4.1 Scheme 1 — scalar palettization, per tensor¶
PerTensorGranularity(), cluster_dim=1.
W [512, 1024] fp16 ONE palette for the whole tensor
┌──────────────────────────────┐
│ 0.31 -0.02 0.88 ... │ k-means(k=16) over all 524 288 scalars
│ 0.29 0.44 -0.51 ... │ ─────────────────────────────────────────►
│ ... │
└──────────────────────────────┘
LUT [1, 1, 16, 1] indices [512, 1024] × 4 bits
┌─────────────────┐ ┌──────────────────────┐
│ c0 -0.83 │ │ 9 4 14 ... │
│ c1 -0.61 │ │ 9 11 2 ... │
│ ... │ │ ... │
│ c15 0.79 │ └──────────────────────┘
└─────────────────┘
reconstruct: W'[i,j] = LUT[ indices[i,j] ]
Storage: 524288 × 4 / 8 = 262,144 B of indices, plus 16 × 2 = 32 B of LUT. Against 1,048,576 B
of fp16, that is 3.999× compression. The LUT is free, numerically speaking.
And that is the problem. Sixteen numbers have to represent every weight in the layer. If the distribution is bimodal, or if a handful of channels have a very different dynamic range from the rest, those sixteen centroids get pulled apart and everything suffers.
This is the presets.w8() scheme (with 256 entries, where it works well) and the scheme you should
not use at 4 bits without measuring.
4.2 Scheme 2 — scalar palettization, per grouped channel¶
PerGroupedChannelGranularity(axis=0, group_size=G), cluster_dim=1. This is presets.w4() and
presets.w6(), and it is what Apple ships for iOS.
With axis=0, group_size=128 on our [512, 1024] weight: 512 // 128 = 4 groups, so four
independent palettes, each fitted to a 128-row slab.
W [512, 1024] FOUR palettes, one per row-group
┌──────────────────────────────┐
│ rows 0..127 (group 0) │ ──► k-means(k=16) ──► LUT[0] : 16 centroids
├──────────────────────────────┤
│ rows 128..255 (group 1) │ ──► k-means(k=16) ──► LUT[1] : 16 centroids
├──────────────────────────────┤
│ rows 256..383 (group 2) │ ──► k-means(k=16) ──► LUT[2] : 16 centroids
├──────────────────────────────┤
│ rows 384..511 (group 3) │ ──► k-means(k=16) ──► LUT[3] : 16 centroids
└──────────────────────────────┘
LUT tensor [4, 1, 16, 1] indices [512, 1024] × 4 bits
reconstruct: W'[i,j] = LUT[ i // 128, 0, indices[i,j], 0 ]
Storage: 262,144 B of indices + 4 × 16 × 2 = 128 B of LUT. Still ~3.998× — the LUT overhead is
noise at this group size.
Now push it: group_size=16 gives 512 // 16 = 32 palettes → 32 × 16 × 2 = 1,024 B. Still noise.
group_size=1 gives 512 palettes → 512 × 16 × 2 = 16,384 B, which is 6.25% on top of the indices.
That is the point at which "per-channel always wins" starts costing you real bytes — and at 8 bits
it gets much worse, because each palette is 256 entries:
✅ VERIFIED — Apple's compression-exploration skill states this as a pitfall verbatim: "8-bit per-channel LUT stores 256 × fp16 entries per output channel" and "scale/ZP overhead 5-15% at 2-4 bit fine granularity", and "at
block_size=16+ int4 the effective width is ~5 bits." (skills/skills/model-compression-exploration/SKILL.md.)
⚠️ The divisibility rule bites here. shape[axis] % group_size must be 0. 512 % 128 == 0 is
fine; 1000 % 16 != 0 is not, and the failure mode is not an exception — see §19.1.
4.3 Scheme 3 — vector palettization¶
cluster_dim = D > 1. This is the scheme that gets you fractional bits per weight, and it is the
one nobody presented.
Instead of clustering scalars, you cluster contiguous D-length vectors along the output-channel axis, and each cluster centroid is itself a D-vector. One index now addresses D weights.
With n_bits=4, cluster_dim=2 on our [512, 1024] weight:
W [512, 1024] pair up along the output-channel axis
┌───────────────────┐
│ row 0 ─┐ │ each 2-vector (W[0,j], W[1,j]) is one data point
│ row 1 ─┘ pair │ ───► k-means in ℝ² with k = 2**4 = 16 centroids
│ row 2 ─┐ │
│ row 3 ─┘ pair │ LUT is 16 × 2 (16 centroids, each a 2-vector)
│ ... │
└───────────────────┘
indices: 512·1024 / 2 = 262 144 indices × 4 bits = 131 072 B
LUT : 16 × 2 × 2 B = 64 B
effective bits per weight = n_bits / cluster_dim = 4 / 2 = 2.0
~8× compression at a nominal 4-bit setting. That is the pitch: vector quantization exploits correlation between adjacent weights, which scalar k-means throws away entirely. If neighbouring output channels are correlated — and in a well-trained network they often are — a 2-D codebook can represent the joint distribution far more efficiently than two independent 1-D codebooks.
The price list is long, and it is why this scheme is rare in production:
| Cost | Detail |
|---|---|
| Non-deterministic | ✅ vector k-means uses numpy.random + torch.randint for centroid init. §7.1. |
| Fast mode unavailable | ✅ enable_fast_kmeans_mode must be False — a ValueError otherwise. |
| No per-channel scale | ✅ cluster_dim > 1 + enable_per_channel_scale=True → ValueError at the program level. |
| CoreML backend restricted | ✅ at most one of {cluster_dim>1, lut_qspec, enable_per_channel_scale}; two raises CoreMLExportError. |
| Slower to fit | multi-dimensional Lloyd's iterations, not the exact 1-D DP. |
| 🔴 Runtime support unquantified | see the GAP box below. |
🔴 GAP — is vector palettization fast on device? The
.aimodelformat supports it (the LUT shape contract has a dedicatedVECTOR_SIZEaxis, ✅ verified below), andfinalize()emitslut_to_densefor it. But no Apple benchmark, no Apple sample, and no community measurement in our corpus reports device latency or compute-unit placement forcluster_dim > 1. It is plausible that the wider LUT interacts badly with the ANE's rank and alignment rules (§5), and it is plausible that it is fine. We do not know. What would resolve it: running acluster_dim=2and acluster_dim=1artefact throughllm-benchmark-style timing withSpecializationOptions.from_preferred_compute_unit_kind(...)and comparing placement, or an Apple doc page onlut_to_denselowering. Safe default meanwhile: usecluster_dim=1. Reach for vector palettization only when scalar palettization at your target bit-width has already failed your quality gate, and always measure device latency before shipping it — a size win that halves your throughput is not a win.
4.4 The LUT tensor shape contract¶
All three schemes produce a LUT with the same rank. This is the single most useful fact in the palettization implementation, because it is what §5 is about.
✅ VERIFIED —
src/coreai_opt/palettization/spec/fake_palettize.py:139-151:
So:
| Scheme | NUM_LUT_AXIS_0 |
NUM_LUT_AXIS_1 |
NUM_PALETTES |
VECTOR_SIZE |
rank |
|---|---|---|---|---|---|
| per-tensor, scalar | 1 | 1 | 2**n_bits |
1 | 4 |
| per-grouped-channel axis 0, scalar | shape[0] // group_size |
1 | 2**n_bits |
1 | 4 |
| per-grouped-channel axis 1, scalar | 1 | shape[1] // group_size |
2**n_bits |
1 | 4 |
| any granularity, vector | as above | as above | 2**n_bits |
cluster_dim |
4 |
The LUT is always rank 4 at the PyTorch level. Hold onto that number.
The other buffers the fake-palettize module carries, for completeness:
✅ VERIFIED — buffers on
_FakePalettizeImplBase:lut,indices,per_channel_scale,quantized_lut,lut_quantization_scale,lut_quantization_zero_point, plusfake_palett_enabledandobserver_enabled(bothuint8tensors). A custom_load_from_state_dictre-registers these dynamically-created buffers on load — so a palettized model does round-trip throughstate_dict()/load_state_dict().
5. ⚠️ The ANE rank-5 ceiling¶
This section exists because of a single docstring inside Apple's own SAM3 export pipeline that contradicts what Apple said on stage. It is the highest-value finding in this guide, and it generalises well beyond palettization.
5.1 The rule¶
✅ VERIFIED —
skills/skills/model-authoring/references/neural_engine_rules.md(479 lines), vendored inapple/coreai-modelsand written by Apple: Max tensor rank is 5. Supported dtypes are fp16 / int8 / int16 (fp32 falls back to GPU/CPU). Shapes must be fully static. The last axis must be contiguous and 64-byte aligned; a singleton last axis costs 32× memory at fp16 and 64× at int8, so keep the last axis at ≥ 32 fp16 elements and prefer powers of two.
Five is not a soft target. A rank-6 intermediate is rejected — the graph does not run on the Neural Engine. It does not fail; it moves.
5.2 The contradiction¶
Session 325, describing the SAM3 encoders:
✅ VERIFIED — WWDC26 session 325, 325:241: "For compression, I apply 4-bit palettization with per-channel scales to the two encoders."
The shipped code that reproduces that demo does the opposite, and says why in a docstring:
✅ VERIFIED —
apple/coreai-models,models/sam3/pipeline.py:136-142,SegmentationExportConfigdocstring, verbatim:"Both encoders deliberately disable per-channel scale:
enable_per_channel_scale=Truelowers tomps.dequantize_lutops with rank-6 LUTs, which ANE rejects (max tensor rank 5), forcing the runtime to fall back to GPU. Keeping it off keeps the asset ANE-compatible at the cost of a small PyTorch-side quality regression."
And the code:
✅ VERIFIED —
models/sam3/pipeline.py:208-245:from coreai_opt import ExportBackend from coreai_opt.palettization import ( KMeansPalettizer, KMeansPalettizerConfig, ModuleKMeansPalettizerConfig, PalettizationSpec, ) from coreai_opt.palettization.spec import PerGroupedChannelGranularity def _make_pal_config(n_bits: int, group_size: int) -> KMeansPalettizerConfig: spec = PalettizationSpec( n_bits=n_bits, granularity=PerGroupedChannelGranularity(axis=0, group_size=group_size), ) return KMeansPalettizerConfig( global_config=ModuleKMeansPalettizerConfig(op_state_spec={"weight": spec}), ) img_pal_config = _make_pal_config(config.image_n_bits, config.image_group_size) # 4, 32 txt_pal_config = _make_pal_config(config.text_n_bits, config.text_group_size) # 6, 8 img_enc = ImageEncoderModule(sam3_lite.image_encoder); img_enc.eval() img_palettizer = KMeansPalettizer(img_enc, img_pal_config) img_enc = img_palettizer.prepare(example_inputs=(pixel_ref,)) img_enc = img_palettizer.finalize(backend=ExportBackend.CoreAI)enable_per_channel_scaleis never set — it takes its default ofFalse.
There are two readings of the discrepancy and the guide will not pretend to know which:
- The talk used "per-channel scales" loosely to mean the per-grouped-channel granularity
(
PerGroupedChannelGranularity(axis=0, ...)), which is a per-channel-ish thing and is what the code does. - The recipe changed after the talk was recorded, and the docstring is the post-mortem.
Either way, the shipped code and its docstring win over spoken narration, per the series'
precedence rules. Write enable_per_channel_scale=False — which is also the default, so write
nothing.
⚠️ SILENT FAILURE — per-channel scale moves your model off the Neural Engine, and nothing tells you. Turning on
enable_per_channel_scaledoes not raise, does not warn, does not fail conversion, and does not change your PyTorch numerics in any alarming way. In fact it makes them slightly better — that is the whole point of the feature. What changes is where the model runs: the ANE rejects the rank-6 LUT and the runtime falls back to the GPU. On a Mac you may not even notice; the GPU is fast. On an iPhone you have just traded the energy-efficient engine for the power-hungry one, in exchange for a fractional dB. Your battery-life regression will show up in field telemetry weeks later. How to detect it: run the artefact underSpecializationOptions.from_preferred_compute_unit_kind(ComputeUnitKind.ane())and check where ops actually land, or use the Core AI Debugger's compute-unit comparison (Part 10). "Preferred" means preferred — the runtime silently relocates ops it cannot place.
5.3 The rank arithmetic, and why you cannot reshape your way out¶
The observed fact is rank 6. The exact derivation is not published, but the budget is easy to see and worth reasoning through, because the same arithmetic governs every ANE decision you make.
Recall the verified PyTorch-side LUT shape from §4.4:
Rank 4 is comfortably inside the budget. But the LUT operand that reaches
mps.dequantize_lut is not that tensor — it is a tensor the lowering constructs, and it has to
carry enough structure to be broadcast against the weight it reconstructs. On the ANE the weight is
already living in a 4-D Conv2d layout ([C_out, C_in, 1, 1]) or a BC1S activation layout
([B, C, 1, S]). Add the grouped-LUT axes and you are at 5. Add another axis for a per-channel
scale plane and you are at 6.
🔴 GAP — the precise axis accounting for the rank-6 LUT. Apple states the outcome (rank 6, ANE rejects) but not the derivation, and
mps.dequantize_lutis an MPSGraph-dialect op with no published operand spec in our corpus. We do not know which axis is added, whether it is the scale plane itself or a broadcast axis introduced to align it, or whether the same thing happens withPerTensorGranularity+ per-channel scale. What would resolve it: an MLIR dump of the lowered program (ai_programinspection, Part 8) or an Apple doc page formps.dequantize_lut. Safe default meanwhile:enable_per_channel_scale=Falsefor any ANE target, regardless of granularity. It is the default and Apple's own shipped recipe agrees.
The workaround you know does not apply here. The standard ANE rank remedy is to reshape — drop singleton dimensions, fold axes together, and get back under 5. Apple's skill gives the canonical conversions:
✅ VERIFIED —
references/neural_engine_rules.md, BC1S conversion helpers:Note thex = x.permute(0, 2, 1).unsqueeze(2) # (B, S, D) → (B, D, 1, S) x = x.squeeze(2).permute(0, 2, 1) # back def gpu_to_bc1s(x): B, H, S, D = x.shape return x.permute(0, 1, 3, 2).reshape(B, H * D, 1, S) def bc1s_to_gpu(x, n_heads, head_dim): B, _, _, S = x.shape return x.reshape(B, n_heads, head_dim, S).permute(0, 1, 3, 2)reshape(B, H*D, 1, S)ingpu_to_bc1s: folding the head and head-dim axes together is exactly a rank reduction, from 4 to 4 but with the seq axis freed — and the same technique takes a rank-5 tensor to rank 4 when you have a spare singleton.
That technique works for tensors your graph creates. It is useless for the LUT, because:
- The rank-6 tensor is synthesised by the lowering, not by any line of your PyTorch. There is no
reshapeyou can insert, because the tensor does not exist until aftertorch.exportandTorchConverterhave run. - The automatic frontend module rewrite does not provide a supported remedy for a rank-6 LUT. Validate the lowered graph and runtime placement instead of relying on an optimizer to repair it.
So the rank-6 LUT has exactly one remedy: do not create it. Leave enable_per_channel_scale
off.
Where the reshape trick does earn its keep is everywhere else in an ANE model, and the interaction
with palettization is real: if you have already spent your rank budget on a 5-D KV cache
([n_layers, B, H_kv*D, 1, max_S] — ✅ verified from the skill's KV-cache table), you have zero
headroom left for anything the compression lowering wants to add. Rank is a budget shared across
your authoring choices and your compression choices, and most people only account for the first.
Practical rank checklist for an ANE-targeted, palettized model:
rank 4 activations in BC1S (B, C, 1, S)
rank 4 Conv2d weights (C_out, C_in, 1, 1)
rank 4 palettization LUT, PyTorch side (G0, G1, 2**n_bits, cluster_dim)
rank 5 KV cache, if you have one (n_layers, B, H_kv*D, 1, max_S)
─────
rank 5 HARD CEILING — anything above is rejected and the graph moves to GPU
rank 6 what enable_per_channel_scale produces ← the trap
5.4 The other lever the SAM3 recipe used instead¶
Having given up per-channel scale, Apple recovered the quality a different way: asymmetric bit-widths per component.
✅ VERIFIED —
apple/coreai-models,models/sam3/README.md, the three-function split with per-function compression:
Function Compression Inputs Outputs image_encode4-bit k-means palettization (gs=32) + fp16 pixel_valuesbackbone_featurestext_encode6-bit k-means palettization (gs=8) + fp16 input_idstext_featuresdetectfp16, no weight compression backbone_features,text_featurespred_masks,pred_boxes,pred_logits,presence_logits
The transcript describes this as "I apply 4-bit palettization ... to the two encoders" — a simplification. The shipped recipe uses three different compression levels for three components, chosen by where the parameters and the sensitivity are:
- Image + text encoders together are 96% of parameters (✅ session 325, 325:60). They get the aggressive treatment.
- The detector is 4% of parameters (✅ 325:158) and is compression-sensitive. It gets nothing. Session 325, 325:248: "The detector stays uncompressed. I know that it's sensitive to compression from our previous exercise."
- Within the encoders, the text encoder — much smaller, and feeding a cross-attention that the whole detection depends on — gets 6 bits and a 4× finer group size (8 vs 32).
That is mixed precision (§14) applied by hand, and it is the pattern the rest of this guide teaches how to derive rather than guess.
⚠️ The CLI flag that undoes it. ✅ VERIFIED —
models/sam3/README.md: the export script's--n-bitsand--group-sizeflags apply uniformly to both encoders, overriding the asymmetric defaults. Souv run models/sam3/export.py --n-bits 4does not "keep the defaults and set the image encoder to 4" — it silently drags the text encoder from 6/gs8 down to 4/gs32-or- whatever-you-passed. If you are reproducing the demo, pass neither flag.
The runnable commands, for reference:
uv run models/sam3/export.py # lite (iOS) export — the WWDC26 325 demo
uv run models/sam3/export.py --help
uv run models/sam3/export.py --full # plain HF Sam3Model, float32, 1008×1008
uv run models/sam3/export.py --full --dtype float16
✅ VERIFIED —
models/sam3/README.md. Note--image-sizedefaults to 336 for the lite path ("the resolution we recommend for iOS deployment") and 1008 for--full; session 325 flags the change at 325:247: "I changed the input image size from 1008 pixels to 336 to run on an iPhone." Note also that--dtype(float16|float32, defaultfloat32) only exists on the--fullpath. §17.3 is about that default.
6. Sizing: what a bit-width actually buys¶
Every compression decision is a size/quality trade, and you cannot evaluate the trade if you are guessing at the size. This section is arithmetic.
6.1 The formula¶
🟢 COMMUNITY-MEASURED / derived —
notes/repos/john-rocky-models.mdrecords this as the working formula used across that archive's exports. The individual terms are consistent with the verifiedcoreai-optstorage layout (§4.4), but the formula itself is a community synthesis, not an Apple-published API:
Apple's compression-exploration skill ships a script that computes the same things:
✅ SOURCE-INSPECTED —
compression_metrics.pyat model snapshotdb63a2d8:check_divisibility(model, axis, block_size) -> dict[str, tuple[int, int]],extract_layer_specs(prepared, *, quantizer=None) -> list[LayerSpec], andcompute_average_bitwidth(specs: list[LayerSpec]) -> float. Prefer keyword arguments for axis and block size. Package-levelcoreai_opt.inspection.bits_per_weightis the 0.3.0 analytical alternative for prepared eager quantization/palettization.
6.2 Worked: a [512, 1024] Linear at every setting¶
fp16 baseline: 524,288 × 2 = 1,048,576 B = 1.000 MiB.
| Scheme | index/weight bytes | LUT / scale bytes | total | ratio | effective bpw |
|---|---|---|---|---|---|
| fp16 | 1,048,576 | — | 1,048,576 | 1.00× | 16.00 |
P8, per-tensor (presets.w8()) |
524,288 | 512 | 524,800 | 2.00× | 8.01 |
| P4, per-tensor | 262,144 | 32 | 262,176 | 4.00× | 4.00 |
| P4, grouped axis 0, gs 32 (16 LUTs) | 262,144 | 512 | 262,656 | 3.99× | 4.01 |
| P4, grouped axis 0, gs 8 (64 LUTs) | 262,144 | 2,048 | 264,192 | 3.97× | 4.03 |
| P4, grouped axis 0, gs 1 (512 LUTs) | 262,144 | 16,384 | 278,528 | 3.76× | 4.25 |
| P6, grouped axis 0, gs 8 (64 LUTs) | 393,216 | 8,192 | 401,408 | 2.61× | 6.13 |
| P8, grouped axis 0, gs 8 (64 LUTs) | 524,288 | 32,768 | 557,056 | 1.88× | 8.50 |
| P8, grouped axis 0, gs 1 (512 LUTs) | 524,288 | 262,144 | 786,432 | 1.33× | 12.00 |
P4, cluster_dim=2, per-tensor |
131,072 | 64 | 131,136 | 8.00× | 2.00 |
| int8 per-channel quant (for contrast) | 524,288 | 1,024 | 525,312 | 2.00× | 8.02 |
| int4 per-block-32 quant (for contrast) | 262,144 | 32,768 | 294,912 | 3.56× | 4.50 |
Read the last three rows of the palettization block. 8-bit per-channel palettization is a 1.33×
compression, not a 2× — three quarters of the "savings" go straight back into 512 palettes of 256
fp16 entries each. That is the concrete form of the skill's warning, and it is the reason
presets.w8() is per-tensor while presets.w4() is per-grouped-channel.
Also read the int4-per-block row: 4.50 effective bits, because a per-block-32 int4 quantization
stores an fp16 scale per 32 weights (16 bits / 32 weights = 0.5 bits/weight). Palettization's LUT
overhead amortises over the whole group rather than per block, which is why 4-bit palettization at
gs=32 lands at 4.01 bpw and 4-bit block quantization at 4.50. That difference is real and it is in
palettization's favour.
6.3 The group-size decision, stated as a rule¶
For our C_in = 1024:
n_bits |
gs=1 | gs=8 | gs=16 | gs=32 | gs=128 |
|---|---|---|---|---|---|
| 4 | 0.250 | 0.031 | 0.016 | 0.008 | 0.002 |
| 6 | 1.000 | 0.125 | 0.062 | 0.031 | 0.008 |
| 8 | 4.000 | 0.500 | 0.250 | 0.125 | 0.031 |
The rule that falls out: the finer the group, the more the palette costs, and the cost grows
exponentially with n_bits. At 4 bits you can afford almost any group size. At 8 bits, per-channel
palettization is a bad deal in bytes even though it is a good deal in dB.
Which is exactly why Apple ships group_size=8 and group_size=32 variants at 4 bits and nothing
finer, and why the community's "per-channel always wins" finding (§17.1) needs the size column
attached to it before you act on it.
6.4 The embedding asterisk¶
Apple's published BPW figures for the iOS LLM models carry a footnote, and it matters:
✅ VERIFIED —
apple/coreai-models, quality tables collected frommodels/*/README.md: "'BPW' marked*includes the INT8-per-tensor embedding."
So 5.71* for a "mixed 4-bit/8-bit palettized" Qwen3 0.6B is not the bit-width of the
transformer weights — it is the whole-model average including a large int8 embedding table that
palettization never touched (§2.6). When you compute your own BPW, decide which convention you are
using and say so, because the two numbers can differ by more than a bit.
7. Determinism, workers, and fast k-means mode¶
7.1 ⚠️ Scalar palettization is reproducible. Vector palettization is not.¶
✅ VERIFIED —
docs/src/palettization/config.md:43-72. Withcluster_dim > 1, vector k-means usesnumpy.random+torch.randintfor centroid initialisation and is non-deterministic. Seeding works only withnum_workers=1, because "k-means runs in spawned worker processes that do not inherit the parent's RNG state." Scalar palettization (cluster_dim == 1) is deterministic.
The documented reproducibility recipe:
import numpy as np, torch
seed = 42
np.random.seed(seed); torch.manual_seed(seed)
model_1 = KMeansPalettizer(copy.deepcopy(model), config).prepare(example_inputs) # num_workers=1
np.random.seed(seed); torch.manual_seed(seed)
model_2 = KMeansPalettizer(copy.deepcopy(model), config).prepare(example_inputs) # identical
⚠️ SILENT FAILURE — a vector-palettized model is a different model every time you build it. Nothing warns you.
prepare()succeeds, the model scores fine, you ship. Then a rebuild from the identical commit produces a different artefact with different numerics, your regression test that compares against a golden hash fails, and you spend a day looking for a dependency change that is not there. Worse: if your CI evaluates the artefact and your release build is a different artefact, the quality number you shipped against was never measured on the bytes you shipped. Rule: ifcluster_dim > 1, pinnum_workers=1and seed both RNGs, or do not use it. The underlying reason scalar palettization escapes this is §2.3 — 1-D k-means is solved exactly by dynamic programming, with no random initialisation to seed.
7.2 num_workers — the wall-clock lever session 325 never mentions¶
✅ VERIFIED —
KMeansPalettizer.prepare(..., num_workers: int = 1), and the docstring: "1runs clustering sequentially. Values greater than1usetorch.multiprocessingto parallelize clustering across layers. It is recommended to use more than one worker process to parallelize the clustering, especially when multiple CPUs are available." Implementation (kmeans/palettizer.py): -num_workers < 1→ValueError(f"num_workers must be >= 1, got {num_workers}")-== 1→_calculate_centroids_sequential: one forward pass, tqdm bar"Palettizing layers (num_workers=1)", hooks tick per fake-palettize module -> 1→_calculate_centroids_parallel:torch.multiprocessing.get_context("spawn").Pool(...), worker count capped at the number of layers, tqdm"Palettizing layers (num_workers=N)". Source comment: "spawn (not fork) so workers don't inherit the parent's CUDA context or other process-global state." Each worker returns the mutated_KMeansFakePalettizeand the parent swaps it back intomodule.parametrizations[attr][idx].
Practical notes:
- The unit of parallelism is a layer. A model with 8 palettizable layers gets no benefit from
num_workers=32; the pool is capped. spawnmeans real process startup and real pickling. For a tiny model,num_workers=1will be faster. For SAM3's 600 M-parameter image encoder it is not close.spawnis also why seeding does not propagate (§7.1). The two facts are the same fact.- A sensible default is
min(os.cpu_count(), number_of_palettizable_layers), and for anything bigger than a toy model you should pass it explicitly:
7.3 enable_fast_kmeans_mode rounds your weights, by default¶
✅ VERIFIED —
ModuleKMeansPalettizerConfig.enable_fast_kmeans_mode: bool = True, docstring: "enables optimizations for faster K-means clustering by rounding the weights before clustering if data is in float16 range. If weight dtype is float32, weights are cast to float16 and then rounded."rounding_precision: PositiveInt = 4is the number of decimal places.
The mechanism is a deduplication trick: rounding to 4 decimal places collapses a large float32 tensor into far fewer distinct values, and 1-D k-means over a small set of distinct values with multiplicities is dramatically cheaper than over millions of unique floats.
The cost is that you have quietly perturbed the data your centroids are fitted to. For most models at 4 bits this is invisible — a 4-decimal rounding is far below the resolution of a 16-entry palette. It becomes relevant when:
- your weights are unusually small in magnitude (rounding to 4 decimals is a large relative error
on values around
1e-4); - you are at
n_bits=8and per-channel, where centroids are close together and the rounding is no longer negligible; - you are chasing the last dB before a gate.
Turn it off, or raise the precision:
ModuleKMeansPalettizerConfig(
op_state_spec={"weight": spec, "in_proj_weight": spec},
enable_fast_kmeans_mode=False, # exact clustering; slower
)
# ...or keep it on but round less aggressively:
ModuleKMeansPalettizerConfig(
op_state_spec={"weight": spec, "in_proj_weight": spec},
rounding_precision=6,
)
And remember it is mandatory to disable for cluster_dim > 1 (§3.3).
A related historical signal worth knowing: bfloat16 support was patched three separate times in this package.
✅ VERIFIED — commits
4df45c0"fix(palettization): cast bfloat16 sensitivities to float32",859d7c9"fix(palettization): cluster bfloat16 weights as float32", and (on the quantization side)f6baedf"fix(quantization): match the float_range bound dtype to the input".
Three fixes to one dtype path in a repo with roughly 35 commits total is a strong signal that the bf16 path is newer and less exercised than fp32/fp16. If your model is bf16 and something numerically strange happens, converting to fp32 before palettizing is a legitimate first diagnostic step.
8. Sensitivity-weighted k-means (SqueezeLLM)¶
Plain k-means minimises squared error in weight space. But not all weights matter equally: an error in a weight the network is highly sensitive to costs you far more output error than the same magnitude of error somewhere flat. Sensitivity-weighted k-means fixes the objective by weighting each weight by its importance.
✅ VERIFIED —
docs/src/palettization/andkmeans/palettizer.py. Based on SqueezeLLM: Dense-and-Sparse Quantization (https://arxiv.org/pdf/2306.07629). Squared gradients are used as per-element importance weights for weighted k-means. From theprepare()docstring: "sensitivity values indicate the importance of each weight element... k-means clustering will place centroids closer to more sensitive weight values", computed viacalibration_mode(loss_fn=...)which "uses the loss function to compute gradients via backpropagation, and the squared gradients are collected as sensitivity values".
8.1 The API¶
import torch.nn.functional as F
from coreai_opt.palettization import KMeansPalettizer, KMeansPalettizerConfig
palettizer = KMeansPalettizer(model, KMeansPalettizerConfig.presets.w4(group_size=32))
prepared_model = palettizer.prepare(example_inputs, num_workers=8)
# Collect sensitivities and RE-CLUSTER using them.
with palettizer.calibration_mode(loss_fn=F.cross_entropy,
sensitivity_path="sensitivities.pt") as skm:
for batch, target in calibration_dataloader:
output = prepared_model(batch)
skm.step(output, target) # computes loss + loss.backward() internally
Reuse on a later run, skipping the calibration pass entirely:
Or save them after the fact:
8.2 What calibration_mode() actually does — and it is a lot¶
This context manager is doing considerably more than "collect some statistics", and knowing the sequence saves you from several surprises.
✅ VERIFIED —
src/coreai_opt/palettization/kmeans/palettizer.py: 1. Saves a fullstate_dictcheckpoint to a temporary file (tempfile.NamedTemporaryFile(prefix="palettizer_calibration_", suffix=".pt")) and callszero_grad(). 2. Disables observers and fake palettization for the duration of sensitivity collection — so gradients are measured on the uncompressed weights. 3. Registersparam.register_hook(lambda grad: torch.square(grad))on everyrequires_gradparameter. 4. On exit:_construct_sensitivitiestakes-param.grad.cpu()per parameter;_normalize_sensitivitiescomputesval = 100 * -val, normalises by the max into[0, 1], then setsval[val == 0] = min(val[val != 0])and clips values below_SENSITIVITY_CLIP_THR = 1e-12. Source comment: "Clipping very small or zero sensitivity values stabilizes k-means, they can lead to divergence otherwise." 5. Restores the checkpoint, attaches the sensitivities to each_KMeansFakePalettize, re-enables observers, recomputes all centroids with the samenum_workerschosen atprepare()time, then restores fake-palettization on / observers off. 6. Ifstep()was never called:RuntimeError("calibration_mode requires at least one call to step(). No calibration data was processed.")
Consequences you should plan around:
- Exiting the context manager re-runs the entire clustering pass. If clustering took twenty
minutes at
prepare(), budget another twenty on exit. This is not a bug; it is the whole point. - It writes a full model checkpoint to your temp directory. For a multi-gigabyte model, make
sure
$TMPDIRhas room. ANo space left on devicehere is confusing because nothing in your code wrote a file. num_workersis inherited fromprepare(), not passed again. Choose it correctly the first time.- Gradients are taken on uncompressed weights, which is the right thing (you want the sensitivity of the original function) but means memory peaks at full precision.
8.3 The sensitivity file format¶
✅ VERIFIED — sensitivity dict keys are
f"{module_name}.parametrizations.{attr_name}.original", and shapes must match the parameter exactly (asserted). Loading usestorch.load(path, weights_only=True), hardened in commit367dfd5("Use weights_only=True when loading sensitivities").
Two implications. First, the key format is tied to the parametrization naming, so a sensitivity
file is only valid for a model with the same module structure — rename a submodule and the file
silently applies to nothing. Second, weights_only=True means you cannot stuff arbitrary metadata
into the file; keep provenance in a sidecar JSON.
8.4 When it is worth it¶
Sensitivity weighting costs you a calibration dataset, a backward pass over it, and a second full clustering pass. Spend that when:
- you are at 4 bits or below and short of your quality gate by a few dB;
- the model has known-sensitive layers you would rather not exclude outright (excluding them costs size; weighting them costs nothing at runtime);
- you have already tried the cheap levers — finer
group_size, one bit-width up — and want the quality without the bytes.
Skip it when you are at 6 or 8 bits and already passing, because you are spending an hour for tenths of a dB you do not need.
🔴 GAP — no published quality delta for sensitivity weighting in
coreai-opt. The SqueezeLLM paper reports large gains for LLM weight quantization, but neither Apple nor any measurement in our corpus reports a before/after PSNR or accuracy number forcoreai-opt's implementation on any model. Theedsrandresnet50example pages do not use it. What would resolve it: an ablation on one of the shipped examples, or an Apple example page. Safe default meanwhile: treat it as a tool to try when you are close to a gate, measure it on your own model, and do not build a schedule around an assumed gain.
9. lut_qspec: quantizing the palette itself¶
9.1 The idea¶
The LUT is small — 16 or 64 or 256 fp16 values per group. Compressing it saves almost nothing (§6.2). So why does the API let you quantize it?
Because it is not about size. It is about which execution path the runtime can take.
✅ VERIFIED —
docs/src/utils/joint_compression.md, verbatim: "A floating-point LUT causes operations to execute in floating-point regardless of the activation quantization, whereas anINT8LUT allows the runtime to use the faster W_INT8-A_INT8 execution path where available."
An fp16 palette means the reconstructed weight is an fp16 value, which means the matmul is an fp16 matmul, which means quantizing your activations to int8 bought you nothing but a pair of extra casts. An int8 palette means the reconstructed weight is an int8 value and the whole matmul can run on the integer path.
So: lut_qspec is not a palettization feature. It is the thing that makes joint compression
worth doing. If you are not quantizing activations, leave it None.
9.2 The constraints¶
✅ VERIFIED —
validate_lut_qspecinpalettization/spec/spec.py: -_SUPPORTED_LUT_DTYPES = {torch.int8, torch.uint8, torch.float8_e4m3fn, torch.float8_e5m2}— note no int4, no int2, no FP4. The palette is not where you save bits. -lut_qspec.granularitymust be the quantizationPerTensorGranularity. Anything else raisesValueError. (This is the import hazard from §3.2 — you need the class fromcoreai_opt.quantization.spec, not the identically-named one frompalettization.spec.) -lut_qspec.qformulationmust beZP;MINVALraises with "Uselut_qspec.qformulation=ZPinstead." - FP8 dtypes additionally require symmetric quantization (inherited fromQuantizationSpec's ownvalidate_qscheme_for_fp_quant).
import torch
from coreai_opt.palettization import PalettizationSpec
from coreai_opt.quantization.spec import QuantizationSpec, QuantizationScheme
from coreai_opt.quantization.spec import PerTensorGranularity as QuantPerTensor
lut_qspec = QuantizationSpec(
dtype=torch.int8,
qscheme=QuantizationScheme.SYMMETRIC,
granularity=QuantPerTensor(), # the ONLY legal granularity here
)
spec = PalettizationSpec(
n_bits=4,
lut_qspec=lut_qspec,
granularity=PerGroupedChannelGranularity(axis=0, group_size=32),
)
9.3 The CoreML restriction¶
✅ VERIFIED —
validate_coreml_palettization_compatibilityin_utils/export_utils.py: on the CoreML backend, at most one of{cluster_dim > 1, lut_qspec, enable_per_channel_scale}may be set. Two of them raisesCoreMLExportError("CoreML export does not support cluster_dim + lut_qspec on <ctx>. Use backend=ExportBackend.CoreAI instead."), because the combination "hits an unsupported CoreML/MIL op configuration (mismatched tensor ranks, orlut_to_densedivisibility errors)". Expanded in commits56c4a36and012f399.
Note the parenthetical: mismatched tensor ranks. The same class of problem as §5, showing up in a different backend. Rank is the recurring constraint in this whole subsystem.
The full CoreML restriction matrix, for reference:
✅ VERIFIED —
src/coreai_opt/_utils/export_utils.py:17-47:So on CoreML: no FP4/FP8 anywhere, no int2/uint2 weights, no per-channel or per-block activation quantization, no MINVAL formulation, and the palettization at-most-one rule above. EveryCOREML_SUPPORTED_WEIGHT_DTYPES = {torch.int8, torch.uint8, torch.int4, torch.uint4} COREML_SUPPORTED_ACTIVATION_DTYPES = {torch.int8, torch.uint8} COREML_SUPPORTED_LUT_DTYPES = {torch.int8, torch.uint8} COREML_SUPPORTED_ACTIVATION_GRANULARITIES = {PerTensorGranularity}CoreMLExportErrormessage ends with the string "Use backend=ExportBackend.CoreAI instead." (common.py:163-177).
10. What finalize() emits¶
10.1 The op-chaining table¶
finalize(backend=ExportBackend.CoreAI) replaces the fake-palettize parametrization with real
custom layers from coreai-torch, and which ops you get depends on which features you enabled.
✅ VERIFIED —
src/coreai_opt/palettization/kmeans/_prepare_for_export.py:225-233._prepare_for_mlir_export(model, mmap_dir=...)swaps incoreai_torch._compression.custom_layers.PalettizeModule/ScaledPalettizeModule, wrapped bycoreai_torch._compression.utils.wrap_for_parametrization. The op chains:
Configuration Emitted ops 1. Palettization only lut_to_dense2. Quantized LUT ( lut_qspec)lut_to_dense(int LUT)+constexpr_blockwise_shift_scale(lut_scale)3. Per-channel scale lut_to_dense+constexpr_blockwise_shift_scale(pcs)4. Both lut_to_dense(int LUT)+constexpr_blockwise_shift_scale(fused_scale)wherefused_scale = lut_scale * per_channel_scale
Row 3 is the one §5 is about — the extra constexpr_blockwise_shift_scale is where the additional
rank comes from. Row 4 is worth noticing for a different reason: the two scales are fused into
one tensor, so combining quantized-LUT and per-channel-scale does not cost you two scale planes.
It still costs you the rank.
✅ VERIFIED —
coreai_torchis lazy-imported (lazy_import_coreai_torch), socoreai-optis usable without it installed. If it is missing atfinalize(backend=CoreAI)you get anImportErrorfrom that helper and nothing earlier.
10.2 ⚠️ Eager-only KMeansPalettizer.finalize(CoreAI) frees dense weights¶
✅ VERIFIED —
kmeans/_prepare_for_export.py, verbatim: "The dense pre-palettization weight stored on the parametrization list is always replaced with a zero-size placeholder so its storage can be released." And from thefinalizedocstring: "Whenbackend=ExportBackend.CoreAI, finalize frees the original dense weights in place: on each parametrized weight,parametrizations[...].originalis replaced with a zero-size placeholder so its storage can be released." This is the eager-only k-means palettizer’s Core AI backend behavior, not a universal rule for every compressor, execution mode or backend.2⚠️ SILENT FAILURE — you cannot undo
finalize(), and there is no warning. The call returns a working model, so nothing looks wrong. But the float weights are gone from the process, and if the object you finalized was your only copy (rememberprepare()mutates in place, §2.4), your float model is gone too. The failure surfaces later, when you want to re-palettize at a different bit-width and discover you have to reload the checkpoint from disk — or, worse, when your sweep script finalizes in a loop and every iteration after the first is clustering already-palettized weights. Two rules that make this a non-event: (1)copy.deepcopybeforeprepare(), always; (2) do not callfinalize()during a sweep at all — Apple's own compression-exploration skill says exactly this, becauseprepare()alone gives you a scoreable model (§1.4).
The quantizer has the same behaviour in eager mode (guide 01 §2.5), so this is a property of the package, not of palettization.
10.3 mmap_dir¶
✅ VERIFIED —
finalize(..., mmap_dir=...)is CoreAI-only (ValueErrorotherwise), and "the files inmmap_dirmust remain in place for the lifetime of the returned model; removing them invalidates the mmap-backed weights." Both quantizer and palettizer require all tensors on CPU and an empty target directory;_validate_mmap_dir_constraintsruns before anything else infinalize().
The 0.3.0 source writes per-weight safetensors files and reloads them through mmap. Both quantizer modes and the eager palettizer require CPU tensors and an empty directory; preserve the files for the returned model's lifetime. §19.7 in guide 01 runs graph and eager quantizer fixtures.
10.4 The full export chain, end to end¶
Putting §1.4's diagram into code, with the ordering rule that matters:
from pathlib import Path
import torch
import coreai_torch
from coreai_opt import ExportBackend
from coreai_opt.casting import cast_to_16_bit_precision
from coreai_opt.palettization import KMeansPalettizer, KMeansPalettizerConfig
# 1. Compress (PyTorch level)
palettizer = KMeansPalettizer(model, KMeansPalettizerConfig.presets.w4(group_size=32))
palettizer.prepare(example_inputs, num_workers=8)
finalized = palettizer.finalize(backend=ExportBackend.CoreAI).eval() # .eval() — §2.4
# 2. Export
exported_program = torch.export.export(finalized, example_inputs).run_decompositions(
coreai_torch.get_decomp_table()
)
# 3. Cast — COMPRESS FIRST, CAST SECOND
cast_to_16_bit_precision(exported_program) # mutates in place
# 4. Convert
converter = coreai_torch.TorchConverter()
converter.add_exported_program(exported_program)
ai_program = converter.to_coreai()
ai_program.save_asset(Path("model.aimodel"))
✅ VERIFIED —
docs/src/introduction/integration_coreai.mdis the source of this sequence, and it explains why the compressed model survives conversion: "Under the hood,finalize()replaces coreai-opt's internal fake-quantize/fake-palettize ops with PyTorch custom ops whose definitions match the corresponding compression ops in the Core AI dialect. This allowscoreai-torchto recognize the ops and map them correctly in the Core AI representation." ✅ VERIFIED — the ordering rule,docs/src/utils/casting.md: "compress first, cast second. Any quantized int8 buffers are left untouched; any remaining FP32 weights move to FP16."
Guide 01 §14 covers cast_to_16_bit_precision in detail. The one-line version: it rewrites the
ExportedProgram graph in place, the FP pass is aggressive (everything except tensors exceeding the
fp16 range, ≈ ±65504) and the INT pass is conservative (skips constant-foldables, indexing feeds, and
anything not consumed by a compute-intensive op). It may change the dtypes of your model's inputs and
outputs, so your calling code may need updating.
11. Pruning: the technique nobody presented¶
Pruning received zero seconds of stage time across WWDC26. Session 325 does not name it. The
coreai-opt README does — "quantization, palettization (codebook-based compression), and
pruning" — and the repository ships a complete implementation with its own config hierarchy, two
schemes, two schedules, a supported-ops registry, export tests and a tutorial notebook.
That gap is worth taking seriously in both directions. Pruning is real and supported. It is also, for most readers, the wrong tool, and Apple says so more bluntly than vendors usually do.
11.1 The verdict, first¶
✅ VERIFIED —
docs/src/pruning/overview.md:7, verbatim:"Unless the original PyTorch model already has a large fraction of weights close to zero across all of its weight parameters, post-training pruning will almost always degrade accuracy. It is most useful as a quick way to evaluate the impact of sparsity on model size and inference latency before committing to a fine-tuning workflow."
Read that carefully, because it is doing two things. It tells you post-training pruning does not work — and it tells you what pruning is for in this package: a measurement instrument. You prune to find out what sparsity would buy you, so you can decide whether to spend the training budget to get there properly.
Compare that with quantization, where data-free post-training weight quantization to 8 bits is routinely near-lossless and takes seconds:
✅ VERIFIED —
docs/src/landing_page.md:43-47: data-free workflows are "typically seconds to minutes even for large models. Often works well for reducing the model down to 8 bits, or even 6 or 4 bits, with only a slight decrease in accuracy." Fine-tuning-based is "the most time-intensive workflow, but typically the only way to recover accuracy at the most aggressive compression ratios for weights (4 bits and below)."
So the honest decision rule is:
Do you have a training loop, a dataset, and time?
├── No ──► Do not prune. Quantize or palettize. You will get more size reduction,
│ more reliably, in less time, with a runtime that definitely exploits it.
└── Yes ──► Is your model already known to be over-parameterised for its task,
or is structured sparsity a hardware requirement you have verified?
├── No ──► Spend the training budget on QAT instead. Guide 01 §11.
└── Yes ──► Prune with a PolynomialDecaySchedule while fine-tuning. §11.6.
And there is a second, sharper reason for that ordering, which is the gap box in §11.8.
11.2 The API¶
✅ VERIFIED —
coreai_opt.pruning.__all__=MagnitudePruner, MagnitudePrunerConfig, ModuleMagnitudePrunerConfig, PruningSpec;coreai_opt.pruning.spec.__all__=ChannelStructured, PruneImplBase, PruningScheme, PruningSpec, Unstructured, default_weight_pruning_spec;coreai_opt.pruning.config.__all__=ConstantSparsitySchedule, MagnitudePrunerConfig, ModuleMagnitudePrunerConfig, OpMagnitudePrunerConfig, PolynomialDecaySchedule, SparsityScheduleBase.
class MagnitudePruner(_BasePruner, _EagerCompressionComponentBuilderMixin): def __init__(self, model: nn.Module, config: MagnitudePrunerConfig | None = None) def prepare(self, example_inputs: tuple[torch.Tensor]) -> nn.Module def step(self) -> None def finalize(self, model: nn.Module | None = None, backend: ExportBackend = ExportBackend.CoreAI) -> nn.Module
Note the absences, and note that they are different absences from palettization's:
- No
mmap_dironfinalize. Quantization and palettization both have it; pruning does not. - No
calibration_mode(). Magnitude pruning is data-free by construction — it ranks weights by|w|, and no forward pass is needed to do that. - No
training_mode(). But there is astep(), and pruning is explicitly designed to be used inside your own training loop. §11.6. - No presets.
MagnitudePrunerConfig.presetsdoes not exist. Every pruning config is hand-written. - Eager only, like palettization.
11.3 PruningSpec and the two schemes¶
✅ VERIFIED —
src/coreai_opt/pruning/spec/spec.py:Scheme registry keysclass PruningSpec(CompressionSpec): target_sparsity: float = Field(default=0.5, ge=0.0, le=1.0) pruning_scheme: PruningScheme = Unstructured() pruning_algo: type[PruneImplBase] = "default" # → _MagnitudePruneImpl"unstructured"/"channel_structured":
Three fields. pruning_algo is an extension point — PruneImplBase is exported publicly, so you
can register your own ranking criterion — but only one implementation ships.
Unstructured pruning zeroes individual weights, anywhere, chosen by magnitude:
W [4, 6], target_sparsity = 0.5 mask
┌────────────────────────────────┐ ┌────────────────────────┐
│ 0.81 -0.02 0.44 0.01 ... │ │ 1 0 1 0 ... │
│ -0.03 0.67 -0.55 0.02 ... │ ──► │ 0 1 1 0 ... │
│ 0.05 -0.71 0.02 0.63 ... │ │ 0 1 0 1 ... │
│ 0.90 0.01 -0.48 0.03 ... │ │ 1 0 1 0 ... │
└────────────────────────────────┘ └────────────────────────┘
keep the 12 largest |w| out of 24; zero the rest. W' = W * mask
Channel-structured pruning removes entire channels, ranked by L1 norm:
W [C_out=8, C_in, ...], axis=0, target_sparsity = 0.5
channel importance = |W|.sum(over every dim except axis)
c0: 41.2 c1: 3.8 c2: 27.9 c3: 5.1
c4: 33.6 c5: 2.2 c6: 19.4 c7: 8.7
prune floor(8 * 0.5) = 4 lowest → c5, c1, c3, c7 zeroed ENTIRELY
mask broadcast to the full weight shape.
✅ VERIFIED — mask arithmetic in
pruning/spec/prune.py: - unstructured:num_keep = numel - floor(numel * sparsity), thentopk(|w|.flatten(), num_keep); mask is1.0at kept indices. - channel-structured:num_prune = floor(num_channels * sparsity); channel importance is the L1 norm,weight.abs().sum(dim=all_dims_except_axis); keeps the topnum_channels - num_prune; the mask is broadcast/expanded to the full weight shape. -sparsity == 0.0→ones_like;sparsity >= 1.0→zeros_like. - The mask is a buffer namedmask.forwardrecomputes it only when_dirty, casts/resizes as needed, then returnsweight * self.mask. Mask dtype matches the weight — fixed in commit3b8d61a, the very first bug-fix commit in the repository.
Which to choose, in one line each:
Unstructuredgives you the best quality at a given sparsity, because it is free to keep any weight. It gives you no dense speedup — a zero in the middle of a dense matmul still gets multiplied. Its value is entirely in compressed storage (if the format encodes sparsity) and in any hardware that exploits fine-grained sparsity.ChannelStructuredremoves whole output (or input) channels, which is genuinely structural — a pruned model can in principle be rebuilt smaller and run faster on any hardware. It costs far more quality at the same sparsity, because you are forced to throw away large weights that happen to live in a weak channel.
11.4 ⚠️ Realized sparsity rounds down, and it can round to zero¶
✅ VERIFIED —
docs/src/pruning/config.md:34-38: forChannelStructured, realized sparsity rounds down to a multiple of1/num_channels. Worked example from the docs:num_channels=7, target_sparsity=0.5⇒floor(7 * 0.5) = 3channels pruned ⇒ 3/7 ≈ 43%, not 50%.⚠️ SILENT FAILURE — a channel-structured target you cannot reach produces a smaller sparsity and no warning. With
num_channels = 8andtarget_sparsity = 0.1,floor(8 * 0.1) = 0: not one channel is pruned, the layer is completely untouched, and the API reports nothing. Your "10% sparse" model is 0% sparse in every narrow layer and your size estimate is wrong by exactly the amount you were counting on. This compounds with the divisibility silent skip in §19.1 —coreai-opthas a consistent house style of degrading quietly rather than raising. Detection: afterprepare(), compute realized sparsity yourself:Do not trustfor name, module in prepared.named_modules(): for attr, plist in getattr(module, "parametrizations", {}).items(): w = getattr(module, attr) zeros = (w == 0).sum().item() print(f"{name}.{attr}: realized sparsity {zeros / w.numel():.3f}")target_sparsityas a description of what you got. It is a request.
11.5 Schedules¶
Pruning is the only compressor whose step() advances a sparsity schedule rather than a QAT
schedule. Two implementations ship.
✅ VERIFIED —
src/coreai_opt/pruning/config/sparsity_schedule.py:Formula:@SparsityScheduleBase.register("constant") class ConstantSparsitySchedule(SparsityScheduleBase): begin_step: NonNegativeInt = 0 # sparsity(s) = target_sparsity if s >= begin_step else 0.0 @SparsityScheduleBase.register("polynomial_decay") class PolynomialDecaySchedule(SparsityScheduleBase): begin_step: int = Field(default=0, ge=0) total_iters: PositiveInt # REQUIRED — no default power: PositiveFloat = 3.0 initial_sparsity: float = Field(default=0.0, ge=0.0, le=1.0) update_frequency: PositiveInt = 1Behaviour: belown_updates = max((total_iters - 1) // update_frequency + 1, 1) i = offset // update_frequency t = i / max(n_updates - 1, 1) sparsity = target + (initial - target) * (1 - t) ** powerbegin_step→initial_sparsity; at or beyondbegin_step + total_iters→target_sparsity; on an off-boundary step it returnsprev_sparsity— and raisesValueErrorifprev_sparsity is Nonewhenupdate_frequency > 1.
The polynomial-decay curve with power=3.0 (the default) is the standard "prune fast early, slow
down near the target" shape from the sparsity literature. Sketched, for initial=0.0,
target=0.7, total_iters=30:
sparsity
0.70 ┤ ╭──────────────────────
│ ╭─────╯
0.60 ┤ ╭─────╯
│ ╭────╯
0.50 ┤ ╭────╯
│ ╭───╯
0.40 ┤ ╭─╯
│ ╭╯
0.20 ┤╭╯
││
0.00 ┼╯──────┬──────┬──────┬──────┬──────┬──────► step
0 5 10 15 20 25 30
Most of the damage is done in the first third, leaving two thirds of the fine-tuning budget to recover from it. That is the entire point of the shape: a network can re-route around gradually removed weights, but not around a sudden 70% amputation.
Use ConstantSparsitySchedule when you want the measurement Apple describes in §11.1 — set
begin_step=0, prune immediately, evaluate, and learn what sparsity costs you before you commit.
11.6 The fine-tuning loop¶
✅ VERIFIED —
docs/src/pruning/overview.md:56-89:from coreai_opt.pruning import MagnitudePruner, MagnitudePrunerConfig, ModuleMagnitudePrunerConfig, PruningSpec from coreai_opt.pruning.config import PolynomialDecaySchedule config = MagnitudePrunerConfig( global_config=ModuleMagnitudePrunerConfig( op_state_spec={"weight": PruningSpec(target_sparsity=0.7)}, sparsity_schedule=PolynomialDecaySchedule(begin_step=0, total_iters=num_epochs, power=3.0), ), ) pruner = MagnitudePruner(model, config) prepared_model = pruner.prepare(example_inputs) for epoch in range(num_epochs): prepared_model.train() for batch, target in train_dataloader: optimizer.zero_grad() criterion(prepared_model(batch), target).backward() optimizer.step() pruner.step() # advance schedule + recompute masksstep()is a no-op when nothing has a schedule._build_scheduled_modules()runs atprepare()time and applies the step-0 state before the initialisation forward pass.
The critical detail is where pruner.step() sits. In this example it is once per epoch,
because total_iters=num_epochs. In the QAT example in guide 01 §11, quantizer.step() is once per
mini-batch, because the schedule integers are in mini-batches. The unit is whatever cadence you
call step() at, and the schedule's total_iters must be in the same unit. Mismatch them and your
schedule completes in the first epoch or never completes at all — with no error either way.
Also note: the masked weights are still parameters. They receive gradients and the optimizer updates them; the mask is re-applied on the next forward. That is what lets a weight come back if the schedule has not yet reached its target — and it is also why you must not naively compute "non-zero parameters" from the optimizer's state.
11.7 Configuration¶
✅ VERIFIED —
src/coreai_opt/pruning/config/magnitude_pruner_config.py:Defaults targetclass ModuleMagnitudePrunerConfig(WeightOnlyModuleValidationMixin, ModuleCompressionConfig[...]): sparsity_schedule: SparsityScheduleBase | None = None class MagnitudePrunerConfig(CompressionConfig[ModuleMagnitudePrunerConfig]): _CONFIG_KEY = "magnitude_pruning_config" # YAML top-level key _SPEC_KEY = "pruning_spec"{"weight", "in_proj_weight"}at 50% unstructured. There are no presets. Supported ops (pruning/supported_ops_registry.py):linear,conv1d/2d/3d,conv_transpose1d/2d/3d,multi_head_attention— allF.*functions.
The same three-level hierarchy, the same None-means-skip semantics, the same regex module-name
matching, and the same fully-qualified-type-name requirement as everything else in this package. The
YAML top-level key is magnitude_pruning_config.
A per-layer example — prune the deep layers harder than the shallow ones, and leave the first and last alone entirely:
from coreai_opt.pruning import MagnitudePrunerConfig, ModuleMagnitudePrunerConfig, PruningSpec
from coreai_opt.pruning.spec import Unstructured
from coreai_opt.pruning.config import PolynomialDecaySchedule
schedule = PolynomialDecaySchedule(begin_step=0, total_iters=40, power=3.0)
config = MagnitudePrunerConfig(
global_config=ModuleMagnitudePrunerConfig(
op_state_spec={"weight": PruningSpec(target_sparsity=0.5, pruning_scheme=Unstructured()),
"in_proj_weight": PruningSpec(target_sparsity=0.5)},
sparsity_schedule=schedule,
),
)
# Deep blocks tolerate more.
config.set_module_name(r"backbone\.blocks\.(1[0-9]|2[0-9])\..*", ModuleMagnitudePrunerConfig(
op_state_spec={"weight": PruningSpec(target_sparsity=0.7)},
sparsity_schedule=schedule,
))
# Boundary layers: leave alone. (See §17.2 for why this is not superstition.)
config.set_module_name(r"stem\..*", None)
config.set_module_name(r"head\..*", None)
11.8 🔴 The gap that should govern your decision¶
🔴 GAP — we cannot verify that Core AI's runtime does anything with unstructured sparsity.
MagnitudePruner.finalize(backend=CoreAI)exists andtests/export/contains a pruning-export test, so a pruned model converts. But: what op the masked weight lowers to, whether the.aimodelstores a sparse encoding or a dense tensor full of zeros, and whether the ANE or the GPU exploits sparsity for either size or latency, are stated nowhere — not in thecoreai-optdocs, not in thecoreai-modelsskills, not in any session, and not in any community measurement in our corpus. The program-levelsparsify_weightspass (§12.2) does name aweight_num_thresholdand offersn:mratios, which is the kind of thing you build for hardware that exploits structured sparsity — but that is inference, not evidence. What would resolve it: an.aimodelsize comparison between a 0%-sparse and a 70%-sparse version of the same model (does the file shrink?), plus a latency comparison on device. Both are a half-day of work and would be the most valuable single measurement anyone could contribute to this topic. Safe default meanwhile: assume unstructured pruning buys you nothing at runtime. Use it as Apple's docs describe it — a measurement instrument — and if you need size, use palettization or quantization, whose runtime benefit is documented. If you need structural speedup, useChannelStructuredand physically rebuild the smaller model yourself, which is a technique this package does not automate.
12. Program-level compression: palettize_weights and sparsify_weights¶
Everything so far operates on a PyTorch nn.Module. There is a second, entirely separate entry
point that operates on an already-converted Core AI program — no PyTorch model required, no
prepare()/finalize() lifecycle, just a function that takes an AIProgram and returns a
compressed one.
✅ VERIFIED —
src/coreai_opt/coreai_utils/__init__.py,__all__=CompressionGranularity, DType, palettize_weights, quantize_weights, sparsify_weights.from coreai_opt.coreai_utils import CompressionGranularity, DType from coreai_opt.coreai_utils import quantize_weights, palettize_weights, sparsify_weightsclass DType(_StrEnum): INT2, UINT2, INT4, UINT4, INT8, UINT8, FP4_E2M1FN, FP8_E4M3FN, FP8_E5M2, FP8_E8M0FNU def is_int(self) -> bool class QScheme(_StrEnum): SYMMETRIC, ASYMMETRIC class CompressionGranularity(_StrEnum): PER_TENSOR, PER_CHANNEL, PER_BLOCK, PER_GROUPED_CHANNEL✅ SOURCE-INSPECTED (0.3.0) —
QSchemeis defined incoreai_opt.coreai_utils.commonand is not re-exported bycoreai_utils.__init__. The tagged documentation's package-level import is stale. Import it fromcommon.
Only constants consumed by a specific op set are candidates:
✅ VERIFIED —
coreai_utils/passes/__init__.py:
Note coreai.gather_nd on that list — the program-level path can compress a gather, which is
exactly what the PyTorch-level palettizer cannot do (§2.6). That is a genuinely useful asymmetry if
your problem is an embedding table.
12.1 palettize_weights¶
✅ VERIFIED —
coreai_utils/passes/weight_palettization.py:63-76:def palettize_weights( coreai_program: AIProgram, lut_dtype: DType | None, n_bits: int = 4, granularity: CompressionGranularity = CompressionGranularity.PER_TENSOR, group_size: int = 32, cluster_dim: int = 1, enable_per_channel_scale: bool = False, weight_num_threshold: int = 1024, num_kmeans_workers: int = 4, enable_fast_kmeans_mode: bool = True, rounding_precision: int = 4, in_place: bool = False, ) -> AIProgram_VALID_LUT_DTYPES = {INT8, UINT8, FP8_E4M3FN, FP8_E5M2};_VALID_N_BITS = {1,2,3,4,6,8};granularity ∈ {PER_TENSOR, PER_CHANNEL, PER_GROUPED_CHANNEL}.enable_per_channel_scale=True+cluster_dim > 1→ValueError. ⚠️lut_dtypeis positional argument #2 and has no default. The docs example passes it as a keyword; if you pass positionally you will silently put yourn_bitsin thelut_dtypeslot.
Differences from the PyTorch path worth knowing:
KMeansPalettizer |
palettize_weights |
|
|---|---|---|
| Input | nn.Module |
AIProgram |
| Granularity vocabulary | per-tensor, per-grouped-channel | per-tensor, per-channel, per-grouped-channel |
| Worker default | num_workers=1 |
num_kmeans_workers=4 |
| Size floor | none | weight_num_threshold=1024 — smaller tensors skipped |
| Compresses embeddings | ✗ | ✓ (via coreai.gather_nd) |
| Sensitivity weighting | ✓ | ✗ |
| Evaluate before committing | ✓ (still a torch model) | ✗ (you must run the program) |
Entry point:
from pathlib import Path
from coreai.authoring import AIModelAsset
from coreai_opt.coreai_utils import CompressionGranularity, DType, palettize_weights
ai_asset = AIModelAsset.load(Path("model.aimodel"))
compressed = palettize_weights(
coreai_program=ai_asset.program,
lut_dtype=None, # keyword — see the warning above
n_bits=4,
granularity=CompressionGranularity.PER_GROUPED_CHANNEL,
group_size=32,
in_place=False,
)
compressed.save_asset(Path("model_palettized.aimodel"))
✅ VERIFIED — the shape of this entry point is from
docs/src/utils/coreai_compression.md:11-31(which showsquantize_weights; the palettization call is the same pattern). Note the docs useweight_num_threshold=2048in their "advanced" examples although the code default is1024— a small inconsistency, but if you are reproducing a documented result, pass the value explicitly.
When to use this path. Rarely, and deliberately. You lose the ability to score the compressed
model in PyTorch, you lose sensitivity weighting, and you are compressing a program that has already
been through the frontend module rewrite — so op fusion decisions were made against uncompressed weights. Use it
when you have an .aimodel and no PyTorch source, or when you need to compress an op the PyTorch
palettizer cannot reach.
12.2 sparsify_weights — the pruning features the PyTorch path does not have¶
✅ VERIFIED —
coreai_utils/passes/weight_sparsification.py:55-64:Rules: -def sparsify_weights( coreai_program: AIProgram, target_sparsity: float | None = 0.5, block_size: int | None = None, n_m_ratio: tuple[int, int] | None = None, quantize_dtype: DType | None = None, palettize_nbits: int | None = None, weight_num_threshold: int = 1024, in_place: bool = False, ) -> AIProgramtarget_sparsityXORn_m_ratio— setting both, or neither, raisesValueError. -quantize_dtypeXORpalettize_nbits— joint sparse+quant or sparse+palettized, never both. -quantize_dtype ∈ {INT8, UINT8, FP8_E4M3FN, FP8_E5M2};palettize_nbits ∈ {1,2,3,4,6,8}. -block_sizemust be> 1; block sparsity runs along the output-channel dimension, and only for linear/conv. -n_m_ratio=(n, m): "Out of everymelements, thenwith lowest magnitude are set to zero", along the input-channel axis, linear/conv only.
This is a strictly larger feature set than MagnitudePruner:
| Feature | MagnitudePruner (PyTorch) |
sparsify_weights (program) |
|---|---|---|
| Unstructured | ✓ | ✓ (target_sparsity) |
| Channel-structured | ✓ (ChannelStructured) |
— |
| Block sparsity (output-channel dim) | ✗ | ✓ (block_size) |
| N:M sparsity (input-channel axis) | ✗ | ✓ (n_m_ratio) |
| Sparse + quantized in one pass | ✗ | ✓ (quantize_dtype) |
| Sparse + palettized in one pass | ✗ | ✓ (palettize_nbits) |
| Schedules / fine-tuning | ✓ | ✗ |
| Size floor | none | weight_num_threshold=1024 |
n_m_ratio=(2, 4) — two zeros in every four consecutive input-channel weights — is the canonical
structured-sparsity pattern that GPU tensor cores exploit on other vendors' hardware. Its presence
here is suggestive.
🔴 GAP — nothing tells us whether Apple silicon exploits
n:mor block sparsity. The API accepts these patterns; no Apple documentation, session, skill or community measurement in our corpus says the ANE or the Apple GPU reads them any faster than dense weights, or that the.aimodelstores them any smaller. This is the same gap as §11.8, restated at the program level. What would resolve it: a size + latency comparison ofsparsify_weights(n_m_ratio=(2,4))against dense, on device. Safe default meanwhile: do not plan a schedule aroundn:msparsity. If you try it, measure the artefact size first — that costs five minutes and answers half the question.
Note the genuinely useful capability in the last two rows of that table: sparsify_weights will do
sparse + palettized in a single pass, which the PyTorch path cannot express at all (you would
have to prune, finalize, then palettize the result, and MagnitudePruner.finalize() produces a
model whose parametrizations the palettizer would then have to survive — untested territory).
13. Joint compression¶
Joint compression means palettized weights and quantized activations in one model. It is a first-class documented workflow with its own doc page, its own test, and three hard rules.
13.1 The three rules¶
✅ VERIFIED —
docs/src/utils/joint_compression.md:Rule 1 — the order is mandatory: palettize weights →
palettizer.finalize()→ quantize activations → calibrate →quantizer.finalize().Rule 2 — you must finalize the palettizer in between, and the reason is structural: "
quantizer.prepareusestorch.export, which cannot trace through the parametrizations."Rule 3 — quantize the LUT, or you have wasted your time: "A floating-point LUT causes operations to execute in floating-point regardless of the activation quantization, whereas an
INT8LUT allows the runtime to use the faster W_INT8-A_INT8 execution path where available."And the restriction that governs everything else: "Models compressed via the joint compression flow can currently only be finalized to the
Core AIbackend."
Rule 2 deserves a moment. Palettization is eager-only and implemented with parametrizations.
Activation quantization needs graph mode, which needs torch.export, which chokes on
parametrizations. So the only way to stack them is to collapse the palettization into real ops
first. That is why the flow has an apparently redundant finalize() in the middle, and why the
ordering cannot be reversed: quantizing activations first would leave a GraphModule that the
eager palettizer's __torch_function__ interception is not built to walk.
There is a small piece of plumbing that makes rule 2 work at all, and it is worth knowing about because it explains an otherwise puzzling line in the source:
✅ VERIFIED —
base_model_compressor.py:21,57-69registers a non-persistent marker buffer_COREAI_OPT_PREPARED_ATTR = "_coreai_opt_prepared"so that re-preparing a prepared model raises.KMeansPalettizer.finalize()deletes that marker (palettization/kmeans/palettizer.py:420-423) "so the palettized model can be handed toQuantizerfor joint compression." Without that one line, step 3 of the flow would raiseRuntimeError("Model has already been prepared.").
13.2 The complete flow¶
import copy
import torch
import torch.nn as nn
import coreai_opt as opt
from coreai_opt import ExportBackend
from coreai_opt.palettization import (
KMeansPalettizer, KMeansPalettizerConfig, ModuleKMeansPalettizerConfig, PalettizationSpec,
)
from coreai_opt.palettization.spec import PerGroupedChannelGranularity
from coreai_opt.quantization import Quantizer, QuantizerConfig, ModuleQuantizerConfig
from coreai_opt.quantization.spec import QuantizationSpec, QuantizationScheme
from coreai_opt.quantization.spec import PerTensorGranularity as QuantPerTensor
model = MyModel().eval()
float_model = copy.deepcopy(model)
example_inputs = (torch.randn(1, 3, 224, 224),) # REPRESENTATIVE data — this seeds qparams
# ── STEP 1: palettize the weights, with an INT8 LUT ────────────────────────────
lut_qspec = QuantizationSpec(
dtype=torch.int8,
qscheme=QuantizationScheme.SYMMETRIC,
granularity=QuantPerTensor(), # the only legal LUT granularity
)
palett_config = KMeansPalettizerConfig(
global_config=ModuleKMeansPalettizerConfig(
op_state_spec={
"weight": PalettizationSpec(n_bits=4, lut_qspec=lut_qspec),
"in_proj_weight": PalettizationSpec(n_bits=4, lut_qspec=lut_qspec),
},
),
)
palettizer = KMeansPalettizer(model, palett_config)
palettizer.prepare(example_inputs, num_workers=8)
# ── STEP 2: finalize the palettizer — MANDATORY, not optional ─────────────────
palettized_model = palettizer.finalize(backend=ExportBackend.CoreAI)
# ── STEP 3: quantize activations ONLY, on the palettized model ────────────────
act_spec = QuantizationSpec(dtype=torch.int8, qscheme=QuantizationScheme.SYMMETRIC)
quant_config = QuantizerConfig(
global_config=ModuleQuantizerConfig(
op_state_spec=None, # ← weights are ALREADY compressed. None == "leave alone".
op_input_spec={"*": act_spec},
op_output_spec={"*": act_spec},
),
)
quantizer = Quantizer(palettized_model, quant_config)
prepared_model = quantizer.prepare(example_inputs)
# ── STEP 4: calibrate the activation ranges ───────────────────────────────────
with quantizer.calibration_mode():
for batch in calibration_dataloader: # ~128 representative samples
prepared_model(batch)
# ── STEP 5: finalize to Core AI. This is the ONLY legal backend here. ─────────
final_model = quantizer.finalize(backend=ExportBackend.CoreAI)
✅ VERIFIED — this is
docs/src/utils/joint_compression.md's flow, with the config shapes checked against the class definitions. The critical line isop_state_spec=Nonein step 3: that is theNone-means-disable semantics from guide 01 §4.3, and getting it wrong means you quantize weights that are already palettized — a second lossy pass over data that has already been through k-means.
13.3 What it costs, measured by Apple¶
✅ APPLE-PUBLISHED —
docs/src/examples/edsr.md. Model:edsr_r16f64super-resolution, 1.5 M parameters; evaluated on B100, with 20 calibration and 80 evaluation samples:
Configuration PSNR Weight storage FP32 baseline 30.68 dB ~5.5 MB W_INT8 _ A_INT8 30.33 dB (−0.35) ~1.4 MB (4×) W_P4(INT8) _ A_INT8 joint 29.86 dB (−0.47 more) ~0.7 MB (8×)
Read the deltas. Going from fp32 to int8-weight/int8-activation costs 0.35 dB and buys 4×. Going from there to 4-bit palettized weights with an int8 LUT costs another 0.47 dB and buys a further 2×. Both are cheap by the standards of §16's gates, and the second step is the one that takes you from "fits comfortably" to "fits on a phone".
✅ APPLE-PUBLISHED — the reference tests, which double as acceptance thresholds:
tests/test_joint_compression.py::test_p4a8_compression_mnist_accuracy(marked@pytest.mark.slow, skipped withoutcoreai) asserts an MNIST baseline of > 97.0% and joint P4-A8 after calibration of > 90.0%, and assertspost_calib_acc == finalized_accforExportBackend._TORCH— i.e. finalize must be numerically exact for the torch backend. The MNIST model has 6 weight-bearing layers (conv1,conv2,conv_transpose1,conv_transpose2,dense1,dense2),batch_size=128,num_calibration_batches=17. There is a matching export test:tests/export/test_pt2e_mlir_export.py::test_mnist_p4a8_compression_export.
That MNIST assertion is a useful sanity target for your own pipeline. If your joint-compressed toy model does not land within a few points of its float baseline, something is misconfigured rather than fundamentally lossy.
13.4 Things that go wrong¶
- You forgot
lut_qspec. Everything runs; the model is smaller; the activation quantization bought you nothing because the fp LUT forces floating-point ops. Nothing warns you. Symptom: your measured latency does not improve at all. - You left
op_state_specat its default in step 3. Weights get quantized on top of being palettized. Symptom: quality drops far more than the EDSR numbers suggest it should. - You used
torch.randnforexample_inputsor for calibration. The initialisation forward pass inprepare()seeds the activation qparams, andcalibration_mode()refines them. Random data gives you ranges that have nothing to do with your model's real activations. Apple's docs say this three separate times, which is a strong hint about how often it happens. - You tried
finalize(backend=ExportBackend.CoreML). Not supported for joint compression, by documented restriction. - You expected
ExportBackend.CoreMLto accept the int8 LUT anyway. It does accept int8 LUTs (COREML_SUPPORTED_LUT_DTYPES = {int8, uint8}) — but the joint flow is Core-AI-only regardless, and the palettization at-most-one rule (§9.3) will catch you if you also wanted per-channel scale.
14. Mixed precision¶
Uniform compression is a simplifying assumption, and it is almost always the wrong one. Layers differ enormously in how much error they tolerate — by orders of magnitude, not percentages — and a single global bit-width is therefore either too aggressive for the sensitive layers or too timid for the rest. Usually both at once.
coreai-opt treats mixed precision as a first-class documented workflow with its own doc page
(docs/src/utils/mixed_precision.md) and its own worked example
(docs/src/examples/mixed_precision_palettization.md). Neither appears in any session.
14.1 The result that makes the case¶
✅ APPLE-PUBLISHED —
docs/src/examples/mixed_precision_palettization.md. ResNet50 palettized, evaluated on the full ImageNet validation set (50,000 images) on the mps backend:
Configuration BPW Size Top-1 FP16 baseline 16 48.64 MB 75.02% Uniform 4-bit per-tensor 4 12.16 MB 65.87% Greedy mixed precision (target 4) 3.95 12.03 MB 70.27%
Read those three rows carefully, because the comparison is not the one people expect.
The mixed-precision model is smaller than the uniform 4-bit model — 12.03 MB against 12.16 MB, 3.95 BPW against 4.00 — and it recovers 4.4 percentage points of top-1 accuracy. It is not trading size for quality. It is spending the same budget more intelligently.
The recipe it found:
✅ APPLE-PUBLISHED — same page. Distribution across 54 configured layers: 2 layers at 6-bit (
conv1,layer1.0.downsample.0), 50 layers at 4-bit, 2 layers at 2-bit (layer1.1.conv1,layer3.4.conv2).
Two layers get more bits, two get fewer, fifty stay put. That is the entire delta, and it is
worth 4.4 points. Note which two got more: conv1 is the stem — the very first convolution,
seeing raw pixels — and layer1.0.downsample.0 is a projection shortcut. Both are boundary
layers in the sense §17.2 describes.
And the shape of the curve:
✅ APPLE-PUBLISHED — same page, verbatim: the curve has an inflection at ≈ 4.0 realized BPW: "below it, every additional 0.5 BPW buys us 15-35 percentage points of accuracy; above it, gains drop to 1-2 points per 0.5 BPW."
top-1
75% ┤ ╭──────────────● fp16 (16 bpw)
│ ╭──────────╯
70% ┤ ● 3.95 ╯ ← mixed precision
│ ╭╯
65% ┤ ● 4.00 ╯ ← uniform 4-bit per-tensor
│ ╭╯
60% ┤ ╭╯
│ ╭╯ ▲
40% ┤ ╭─╯ │ inflection ≈ 4.0 BPW
│ ╭─╯ │ below: 15-35 pts per 0.5 BPW
20% ┤╯ │ above: 1-2 pts per 0.5 BPW
┼────┬────┬────┬────┬────┬────┬────► realized BPW
2 2.5 3 3.5 4 4.5 5
That inflection is the single most actionable fact in this section. It says:
- Above ~4 BPW, buying more bits is a bad deal. You are paying 12.5% more storage for one accuracy point.
- Below ~4 BPW, giving up bits is catastrophic. A layer pushed from 4 to 3.5 effective bits can cost you tens of points.
- Therefore the right mixed-precision strategy is not "shave everything a bit" — it is "keep almost everything at the inflection, and move a handful of layers in each direction." Which is exactly the recipe Apple found.
14.2 The workflow¶
✅ VERIFIED (shape) — the method is described consistently in
docs/src/utils/mixed_precision.mdand in the community's reading of the same material: per-layer bit-widths come from a layer-sensitivity sweep — compress one layer at a time, score by PSNR — then walk least-loss-first toward a target average bit-width.
The algorithm in words:
- Baseline. Score the uncompressed model on your metric. Record it.
- Sensitivity sweep. For each compressible layer L, and for each candidate bit-width b:
compress only L to b, leave everything else uncompressed, score. You now have a
(layer, bits) → quality_losstable. This isn_layers × n_bit_widthsevaluations, and it parallelises perfectly. - Greedy assignment. Start with every layer at its highest candidate bit-width. Repeatedly pick
the
(layer, next-lower-bits)move with the smallest quality loss per bit saved, apply it, and recompute the average bit-width. Stop when you hit your target. - Verify jointly. The sweep measured layers in isolation; errors compound. Score the final configuration end-to-end and expect it to be somewhat worse than the sum of the parts.
Step 4 is not optional and is the step people skip. A per-layer sensitivity table is a heuristic ranking, not an additive error model.
✅ SOURCE-INSPECTED (0.3.0) —
docs/src/utils/mixed_precision.mdexplicitly labels the sensitivity and greedy recipe algorithms as pseudocode. There is no publiccoreai_opt.mixed_precisionmodule or greedy-search function in the release package. Implement the sweep with per-module configs and prepared models (§14.3), and usecoreai_opt.inspection.bits_per_weightfor supported eager quantization/palettization models (§15.6). The BPW utility does not support graph quantization, pruning, or finalized models.
14.3 Writing the sweep yourself¶
Every primitive this needs is verified. The only thing you supply is a scoring function.
"""Per-layer sensitivity sweep + greedy bit assignment for palettization.
Uses only verified coreai-opt APIs:
- KMeansPalettizerConfig / ModuleKMeansPalettizerConfig / PalettizationSpec
- config.set_global(None) + config.set_module_name(name, cfg)
- KMeansPalettizer(...).prepare(...) [NO finalize() — Apple's skill says don't]
"""
from __future__ import annotations
import copy
import itertools
from dataclasses import dataclass
import torch
import torch.nn as nn
from coreai_opt.palettization import (
KMeansPalettizer, KMeansPalettizerConfig, ModuleKMeansPalettizerConfig, PalettizationSpec,
)
from coreai_opt.palettization.spec import PerGroupedChannelGranularity
CANDIDATE_BITS = (8, 6, 4, 2) # legal values are {1,2,3,4,6,8} — no 5, no 7
GROUP_SIZE = 32
def _spec(n_bits: int) -> PalettizationSpec:
return PalettizationSpec(
n_bits=n_bits,
granularity=PerGroupedChannelGranularity(axis=0, group_size=GROUP_SIZE),
# enable_per_channel_scale left False — see §5. Do not turn this on for ANE targets.
)
def _module_cfg(n_bits: int) -> ModuleKMeansPalettizerConfig:
s = _spec(n_bits)
return ModuleKMeansPalettizerConfig(op_state_spec={"weight": s, "in_proj_weight": s})
def _config_for(assignment: dict[str, int]) -> KMeansPalettizerConfig:
"""Build a config that palettizes ONLY the named modules, each at its own bit-width."""
cfg = KMeansPalettizerConfig(global_config=None) # None == compress nothing by default
for module_name, n_bits in assignment.items():
# module_name_configs keys are regexes matched with re.fullmatch — escape dots.
cfg.set_module_name(module_name.replace(".", r"\."), _module_cfg(n_bits))
return cfg
def palettizable_modules(model: nn.Module) -> list[str]:
"""Modules whose forward reaches a palettizable op (§2.6)."""
kinds = (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d,
nn.ConvTranspose1d, nn.ConvTranspose2d, nn.ConvTranspose3d,
nn.MultiheadAttention)
return [n for n, m in model.named_modules() if isinstance(m, kinds)]
def check_divisibility(model: nn.Module, names: list[str], group_size: int) -> list[str]:
"""Pre-flight the SILENT SKIP in §19.1: axis-0 size must be divisible by group_size."""
bad = []
for n in names:
mod = model.get_submodule(n)
w = getattr(mod, "weight", None)
if w is None:
w = getattr(mod, "in_proj_weight", None)
if w is not None and w.shape[0] % group_size != 0:
bad.append(f"{n}: shape[0]={w.shape[0]} not divisible by {group_size}")
return bad
@dataclass
class Cell:
module: str
n_bits: int
loss: float # baseline_score - score, higher is worse
numel: int
def sweep(model: nn.Module, example_inputs: tuple, score, baseline: float,
names: list[str], num_workers: int = 8) -> list[Cell]:
"""One layer at a time, one bit-width at a time. Embarrassingly parallel across cells."""
cells: list[Cell] = []
for name, n_bits in itertools.product(names, CANDIDATE_BITS):
work = copy.deepcopy(model) # prepare() mutates in place — §2.4
palettizer = KMeansPalettizer(work, _config_for({name: n_bits}))
prepared = palettizer.prepare(example_inputs, num_workers=num_workers)
# NOTE: no finalize(). The prepared model is scoreable and finalize() is destructive (§10.2).
numel = model.get_submodule(name).weight.numel()
cells.append(Cell(name, n_bits, baseline - score(prepared), numel))
del work, prepared, palettizer
return cells
def greedy_assign(cells: list[Cell], names: list[str], target_bpw: float) -> dict[str, int]:
"""Start everyone at the top bit-width; take the cheapest bit-per-loss move until at target."""
by_cell = {(c.module, c.n_bits): c for c in cells}
numel = {c.module: c.numel for c in cells}
total = sum(numel[n] for n in names)
assignment = {n: max(CANDIDATE_BITS) for n in names}
ladder = sorted(CANDIDATE_BITS, reverse=True)
def bpw(a: dict[str, int]) -> float:
return sum(numel[n] * b for n, b in a.items()) / total
while bpw(assignment) > target_bpw:
best, best_cost = None, float("inf")
for n, cur in assignment.items():
i = ladder.index(cur)
if i + 1 >= len(ladder):
continue
nxt = ladder[i + 1]
bits_saved = (cur - nxt) * numel[n]
delta_loss = by_cell[(n, nxt)].loss - by_cell[(n, cur)].loss
cost = delta_loss / max(bits_saved, 1) # loss per bit saved
if cost < best_cost:
best, best_cost = (n, nxt), cost
if best is None:
break # everyone is at the floor
assignment[best[0]] = best[1]
return assignment
Driving it:
baseline_score = score(model)
names = palettizable_modules(model)
problems = check_divisibility(model, names, GROUP_SIZE)
if problems:
raise SystemExit("These layers would be SILENTLY SKIPPED:\n " + "\n ".join(problems))
cells = sweep(model, example_inputs, score, baseline_score, names, num_workers=8)
assignment = greedy_assign(cells, names, target_bpw=4.0)
# STEP 4 — verify jointly. Isolated sensitivity is a ranking, not an error model.
final = copy.deepcopy(model)
final_prepared = KMeansPalettizer(final, _config_for(assignment)).prepare(example_inputs, num_workers=8)
print("joint score:", score(final_prepared), "target bpw:", 4.0)
The check_divisibility pre-flight is not decoration — it is the difference between a sweep whose
size numbers mean something and one whose "compressed" layers were quietly left dense.
14.4 Apple's shipped mixed-precision configs¶
The clearest evidence that Apple treats this as production practice rather than a demo: the LLM export registry references hand-written mixed-precision YAML for two Qwen3 models on iOS.
✅ VERIFIED —
apple/coreai-models,python/src/coreai_models/llm/model_registry.py, iOS variants:
short_name hf_id compression qwen3-0.6bQwen/Qwen3-0.6Bnone+compression_config="models/qwen3/qwen3_0_6b_mixed_4bit_8bit.yaml"qwen2.5-1.5b-instructQwen/Qwen2.5-1.5B-Instruct4bit_weight_palettized_group8qwen3-4bQwen/Qwen3-4Bnone+compression_config="models/qwen3/qwen3_4b_mixed_4bit_8bit.yaml"
And the quality it buys, published:
✅ APPLE-PUBLISHED — WikiText-2 perplexity via lm-evaluation-harness, from
models/*/README.md. BPW marked*includes the INT8-per-tensor embedding (§6.4):
Model Compression BPW Platform Perplexity Qwen3 0.6B none (float16) 16.00 iOS 26.16 Qwen3 0.6B Mixed 4-bit/8-bit palettized (YAML) 5.71* iOS 30.90 Qwen3 4B none (float16) 16.00 iOS 16.41 Qwen3 4B Mixed 4-bit/8-bit palettized (YAML) 4.89* iOS 18.80 Qwen3 4B none (float16) 16.00 macOS 16.41 Qwen3 4B 4-bit quantized 4.50 macOS 18.33 Qwen2.5 1.5B Instruct none 16.00 iOS 12.21 Qwen2.5 1.5B Instruct 4-bit palettized (gs 8) 4.63* iOS 14.64
Two observations that are easy to miss:
- The 0.6B model pays much more than the 4B model. 26.16 → 30.90 is a +18% perplexity regression at 5.71 BPW; 16.41 → 18.80 is +15% at a lower 4.89 BPW. Small models have less redundancy to give up. If your model is under a billion parameters, budget for a worse compression curve than the papers promise.
- Apple did not use uniform palettization for these two. The registry entry is literally
compression="none"plus a YAML file, meaning the preset system was insufficient and someone hand-wrote per-module bit-widths. That is the strongest possible signal about how well uniform 4-bit works on a small transformer.
14.5 Writing the YAML¶
Mixed precision is where YAML earns its place, because the file is the artefact you version and review rather than a Python function that produces one.
✅ VERIFIED —
KMeansPalettizerConfig._CONFIG_KEY = "kmeans_palettization_config",_SPEC_KEY = "palettization_spec". Loading:KMeansPalettizerConfig.from_yaml(path)usesyaml.safe_load; an empty file emitswarnings.warn("Empty YAML content detected, returning None …")and returnsNone; a non-mapping raisesValueError; unexpected top-level keys raiseRuntimeError. The only allowed top-level keys are_CONFIG_KEYand_SPEC_KEY— the latter exists purely to host YAML anchors. ✅ VERIFIED — the exporter's loader (coreai_models/llm/export.py:163-237) additionally requires exactly one coreai-opt top-level key after popping an optionalcoreai_models:block, which may contain only{"calibrate_activations"}; andkmeans_palettization_configrequires--platform iOS.
# mixed_4bit_8bit.yaml — anchors live under palettization_spec, config under kmeans_palettization_config
palettization_spec:
p8: &p8
n_bits: 8
granularity: {type: per_grouped_channel, axis: 0, group_size: 32}
p4: &p4
n_bits: 4
granularity: {type: per_grouped_channel, axis: 0, group_size: 32}
kmeans_palettization_config:
global_config:
op_state_spec:
weight: *p4
in_proj_weight: *p4
module_name_configs:
# Boundary layers keep more bits (§17.2).
'model\.layers\.0\..*':
op_state_spec: {weight: *p8, in_proj_weight: *p8}
'model\.layers\.(2[0-9])\..*':
op_state_spec: {weight: *p8, in_proj_weight: *p8}
# Leave the head alone entirely. null == "do not compress this scope".
'lm_head': null
from coreai_opt.palettization import KMeansPalettizerConfig
config = KMeansPalettizerConfig.from_yaml("mixed_4bit_8bit.yaml")
Three YAML-specific reminders, all verified in guide 01 §5.4 and unchanged here: keys under
module_name_configs are re.fullmatch regexes (so escape your dots, and remember that
model\.layers\.1\..* matches layer 1 but not layer 10 — model\.layers\.1[0-9]\..* does);
module_type_configs keys must be fully-qualified class paths; and null means disable,
which is not the same as omitting the entry.
⚠️ The registry YAMLs are source-tree-only. ✅ VERIFIED —
_resolve_registry_compression_configraisesSystemExitwhen running from a wheel: "Registry preset references <path>, but the YAML lives in the source tree which is unavailable in this install." If you want to read Apple's actual mixed-precision recipes, cloneapple/coreai-models;pip installwill not give them to you.
15. Choosing per layer: the sweep, the Debugger, and SAM3¶
Three different mechanisms in this stack answer the same question — which layers tolerate compression? — and they answer it at three different points in the workflow. Most people know about one of them.
| Mechanism | Where it runs | What it tells you | Cost |
|---|---|---|---|
| Layer-sensitivity sweep (§14) | PyTorch, before conversion | which layer × bit-width cells are cheap | n_layers × n_bits evaluations |
Apple's model-compression-exploration skill |
PyTorch, agent-driven | ~60 configs across the Pareto frontier | hours, parallelised |
| Core AI Debugger (Part 10) | after conversion, per-op | which op diverged, with PSNR per sync point | one comparison run |
They compose. The sweep tells you what to try; the Debugger tells you what actually broke.
15.1 The SAM3 diagnosis, and why it belongs here¶
Guide 01 §13 tells the SAM3 story from the quantization side: presets.w4() applied uniformly to an
850 M-parameter model, 3 GB → ~430 MB, and an occluded flower that stops being detected. The
diagnosis is worth repeating here because the fix is a mixed-precision fix.
✅ VERIFIED — WWDC26 session 325, 325:156–162, verbatim:
"I'm noticing that the vast majority of low-PSNR sync points are actually coming from the detector decoder. This tells me that the quantization scheme applied earlier has mildly corrupted the detector results. Since we previously identified that the detector block only accounts for 4% of model parameters, we're not getting much benefit from compressing it anyway. So, I'll return to the Jupyter notebook, and try changing the quantization scheme to ignore the detector." "Great! I can see that we have once again reached baseline quality where all flowers are detected and the model is only a fraction of the size! Core AI Debugger turned hours of manual tensor comparison into a visual diagnosis. I started with missing detections and reached a revised quantization scheme in minutes."
The mechanism is the one this guide has used repeatedly:
# The SAM3 move: compress everything, leave one subtree alone.
config = KMeansPalettizerConfig.presets.w4(group_size=32)
config.set_module_name(r"detector\..*", None) # None == "leave this alone"
Two lessons generalise beyond SAM3:
- Parameter share and sensitivity are usually anti-correlated. The detector is 4% of the parameters and 100% of the problem. Compressing it saves nothing measurable and costs everything. Before you compress a block, ask what fraction of your bytes it actually is — if the answer is under 5%, skip it and stop thinking about it.
- The final layers of a network are where errors become visible. Every upstream error has already been absorbed or amplified by the time it reaches the decoder, and the decoder's job is to make discrete decisions (is this a detection?) where small numeric shifts flip outcomes.
15.2 The three-function split also selects coreai-models’ ANE preference¶
Session 325 presents the three-entrypoint split (image_encode / text_encode / detect) as a
latency trick — run each at a different cadence, and get a 76%-faster second inference when only
the prompt changes.
✅ VERIFIED — WWDC26 session 325, 325:255–256: "here's the payoff of the three-function split. I swapped the prompt to butterfly and only re-ran the text encoder and the detector. As a result, the second inference is 76% faster, even after warmup."
Reading the optional coreai-models Swift package shows the split also selects that loader’s
Neural Engine preference. This is package policy, not a Core AI framework naming contract; direct
AIModel callers choose their own SpecializationOptions.1 It still bears
directly on compression when you adopt that loader because intended ANE residency affects the
appropriate format (§1.2). Part 10 covers the policy; the compression-side consequence is simply this:
if you split a model into per-cadence functions, you can and should give each function its own
compression config, which is precisely what the SAM3 recipe does (§5.4).
⚠️ The 76% number requires caller-side work Apple's own package does not do.
CoreAISegmentationEnginere-runsimage_encodeon every call and exposes no cache. If you want the figure, you cache the backbone features yourself. Part 7 covers this.
15.3 Apple's model-compression-exploration skill¶
Apple ships an agent skill that automates the sweep. Even if you never run it, its contents are the closest thing to an Apple-endorsed compression methodology that exists, and it is worth mining.
✅ VERIFIED —
skills/skills/model-compression-exploration/SKILL.mdinapple/coreai-models. It drives roughly 30 main-sweep + 30 refinementcoreai_optconfigurations.Step 3 timing probe: start with
QuantizerConfig.presets.w8()(graph mode default); ifQuantizer.preparefails, fall back toQuantizerConfig.presets.w8(execution_mode=ExecutionMode.EAGER)for the whole sweep.KMeansPalettizerConfig.presets.w6()= 6-bit, per-grouped-channel,group_size=16; palettization is eager-only.The three config groups: - 1a — channel-structured quantization (6 configs):
{int8, int4} × {symmetric, asymmetric, symmetric_with_clipping}- 1b — block-structured quantization (9 configs): int4 per-block ×{block_size 16, 32, 128} × {3 qschemes}- 2 — palettization (15 configs):{(8-bit per-tensor), (6-bit per-tensor), (6-bit gs 4/8/16), (4-bit gs 4/8/16)} × {enable_per_channel_scale True/False}Preset anchors:
presets.w8()/w4()(per-channel symmetric),presets.w4_per_block(block_size=32),KMeansPalettizerConfig.presets.w8()/w6()/w4().Refinement pass: filter out configs with PSNR < 10 dB or IoU < 0.1, pick the 95th and 75th percentile survivors as seeds, and generate 5 layer-skip variants each via
set_module_nameoverrides.Operational rules: "Do not call
finalize()"; "Calibration is not needed for weight-only compression."Pitfalls it names: per-block / per-grouped-channel silently skip layers whose weight dimension is not divisible (pre-check with
check_divisibility()); scale/ZP overhead is 5-15% at 2-4 bit fine granularity; an 8-bit per-channel LUT stores 256 × fp16 entries per output channel; atblock_size=16+ int4 the effective width is ~5 bits.Output record shape, verbatim:
{"group":"2", "config":{"name":"palette_grouped_gs4_6bit_pcs0_skip-Embedding","path":"path/to/config"}, "time_taken":1000, "output_quality_metrics":[{"name":"bbox","metric":"iou","value":0.7}, {"name":"logits","metric":"psnr","value":16}], "compression_metrics":{"average_bitwidth":5,"compression_ratio":1.7,"theoretical_model_size":402}}Report format: exactly 5 configs per group spanning the Pareto frontier; columns
Config | PSNR (dB) | Avg Bitwidth | Compression Ratio; a scatter plot saved ascompression_exploration.png.Parallelisation: one subagent per group appending to a shared
results.jsonl; the main agent polls with/loop 5m. Bundled scripts:scripts/compression_metrics.py(theoretical size, average bitwidth, divisibility, parametrize walk) andscripts/quality_metrics.py(PSNR / SNR / IoU plus a dispatcher).
Five things in that specification are worth stealing whether or not you use an agent:
- The config-name convention —
palette_grouped_gs4_6bit_pcs0_skip-Embeddingencodes scheme, group size, bit-width, per-channel-scale flag and skip-list in one string. Adopt something like it; sweep results are useless if you cannot tell what produced them six weeks later. - Two metrics, not one. The record carries both a task metric (
iouonbbox) and a numeric metric (psnronlogits). PSNR alone will happily tell you a model is fine when its detections have collapsed. - The PSNR < 10 dB / IoU < 0.1 filter. Below those, a config is not "worse" — it is broken, and including it distorts every percentile you compute.
- Pareto framing. "Five configs spanning the frontier" is the right deliverable. A single "best" config is a decision you have taken away from whoever ships the model.
- The
enable_per_channel_scale True/Falseaxis in group 2. Note that the skill does sweep it — meaning it is a legitimate quality lever. Just re-read §5 before you let a sweep pickpcs1for an ANE target: the sweep is scoring PyTorch quality, and it has no idea it just moved your model to the GPU.
That last point is the most important interaction in this guide. An automated sweep optimising
PSNR will choose enable_per_channel_scale=True every time, because it improves PSNR, and it will
never tell you what it cost. If you run this skill against an iOS target, constrain the search
space to pcs0 before you start.
15.4 The .aimodel is a build artefact, not a pure function of the recipe¶
One more caution before you treat any sweep result as durable.
🟢 COMMUNITY-MEASURED —
notes/repos/john-rocky-models.md, single-author, uncontrolled: "An.aimodelis a build artefact, not a pure function of the recipe" — the samecoreai.llm.export qwen3-0.6binvocation produced a 2.2× faster artefact on macOS 26 than on the 27 beta, attributed to loss of native quantized-Linear lowering (same code, same wheels). "Version-stamp and keep your artifacts." ⚠️ Beta-era measurement; the mechanism is identified but not Apple-confirmed. Treat it as a reason to record your toolchain version alongside every sweep result, not as a claim about any shipping OS.
15.5 Hand-off to the Debugger¶
The sweep gives you a ranking. The Debugger gives you a diagnosis. Use the second when the first does not explain what you are seeing — specifically, when a configuration the sweep predicted would be fine is not fine.
What Part 10 covers, and why it matters here:
✅ VERIFIED — WWDC26 session 325, 325:139–152. The Debugger builds a comparison session between a specialized Core AI model and a PyTorch reference loaded from an Intermediates File. The navigator fills with operation pairs — "These pairs are called sync points, places where the specialized model's output is expected to match the original PyTorch result. The debugger automatically identifies these points throughout the model." Each carries a similarity metric; "The default metric is a peak signal-to-noise ratio or PSNR, but this can be changed to whichever similarity indicator suits your model best." Green = similar, yellow = moderately diverged, red = significantly different. The workflow is "sort by similarity" then use the up-arrow key to walk the low-PSNR sync points "one-by-one to see if a pattern emerges", with the source viewer showing each op's PyTorch module hierarchy.
The reason that last detail matters for compression: the source viewer maps a diverging op back to
a PyTorch module path — which is exactly the string you need for set_module_name(...). The
Debugger's output is directly actionable as a config change, and that round trip is the whole
workflow:
sweep ──► a config that looks good on paper
│
▼
convert + run ──► quality is not what the sweep predicted
│
▼
Debugger: sort sync points by PSNR, walk the reds
│
▼
source viewer gives you "detector.decoder.layers.3.self_attn"
│
▼
config.set_module_name(r"detector\.decoder\..*", None) ──► re-measure
🔴 GAP — the Debugger's available similarity metrics are not enumerated. Session 325 says the metric "can be changed to whichever similarity indicator suits your model best" but never lists the options, and no doc page in our corpus does either. What would resolve it: the Debugger's metric picker in Xcode 27, or an Apple doc page. Safe default meanwhile: PSNR is the default and is what every Apple acceptance gate in §16 is expressed in; use it, and carry a task metric alongside it as §15.3's record format does.
15.6 Inspection and analytical BPW in 0.3.0¶
ModelInspector(model, example_inputs, execution_mode, compressor=None, dynamic_shapes=None,
export_with_no_grad=True) requires an explicit execution mode. format_summary(colorize=None)
accepts the colorization keyword. Use the mode you plan to compress in; graph and eager op names
differ. bits_per_weight(prepared) returns BitsPerWeightResult(bpw, per_module_map, total_bits,
total_weights) for prepared eager quantization/palettization. It counts bias, norms, skipped
parameters and buffers at full precision, including non-persistent buffers except its internal
prepared marker. It rejects unsupported/pruned models; it is not a file-size or runtime-speed
measurement. Source: inspection/model_inspector.py and inspection/bits_per_weight.py at 189612be.
16. Apple's PSNR acceptance gates¶
There is no Apple documentation page titled "how good does my compressed model have to be". There is something better: Apple wrote acceptance thresholds into an agent skill, because an agent needs a number to decide against. Those numbers are the closest thing to an official standard that exists, and this guide treats them as one.
16.1 The authoring gates¶
✅ VERIFIED —
skills/skills/model-authoring/SKILL.md:94-99, vendored inapple/coreai-models:
Comparison Gate Re-authored model vs source model (both PyTorch) > 70 dB ANE-layout vs GPU-layout (both PyTorch) > 70 dB Compiled Core AI model vs torch ≥ 40 dB After 4-bit palettization ≥ 35 dB
Four numbers, and the structure of the ladder is as informative as the values.
The first two gates are about work you did, and they are strict. Re-authoring a model for the
Neural Engine — swapping nn.Linear for nn.Conv2d(1×1), converting to BC1S layout, splitting
attention per head — must be numerically transparent. 70 dB is essentially "the same tensor with
fp16 rounding noise". If your re-authored model scores 55 dB against the original, you did not
re-author it; you changed it, and you have a bug. Same for ANE-layout vs GPU-layout: two layouts of
the same math must agree.
The third gate is about the toolchain, and it is much looser. 40 dB accepts real fp16 numerical drift through compilation, fusion and automatic rewriting. That is a 30 dB relaxation from the authoring gates, and it tells you Apple expects conversion to cost you real precision.
The fourth gate is about compression, and it is looser still. 35 dB after 4-bit palettization is the floor below which Apple's own skill says investigate.
The ladder in one line: your code must be exact, the toolchain may drift, compression may cost.
16.2 The sizing table that goes with them¶
✅ VERIFIED — same skill, palettization sizing guidance:
Bit-width Expected compression Expected PSNR Flag for investigation below 8-bit ≈ 2× > 55 dB < 50 dB 4-bit ≈ 4× ~ 40 dB < 35 dB 2-bit ≈ 8× 25–35 dB — "Usually unacceptable"
And a second, differently-scoped table from the deployment side of the same skills tree:
✅ VERIFIED —
skills/skills/model-deployment(PSNR acceptance table):
Scenario Expected PSNR Investigate if below float32 end-to-end > 70 dB 60 dB fp16 on-device > 50 dB 40 dB 4-bit palettized ~40 dB 30 dB
The two tables agree on the shape and differ slightly on the thresholds (35 vs 30 for 4-bit; 40 vs 50 for the compiled/fp16 tier). They are scoped differently — one is an authoring gate, one is a deployment expectation — and the difference is small enough not to matter in practice. Use the stricter number as your CI gate and the looser one as your "do not ship below this" floor.
16.3 Using them¶
import torch
def psnr(reference: torch.Tensor, candidate: torch.Tensor) -> float:
"""Peak signal-to-noise ratio in dB, using the reference's dynamic range as peak."""
reference = reference.detach().float()
candidate = candidate.detach().float()
mse = torch.mean((reference - candidate) ** 2)
if mse == 0:
return float("inf")
peak = reference.abs().max()
return float(20 * torch.log10(peak / torch.sqrt(mse)))
GATES = {
"reauthored_vs_source": 70.0, # strict: your code must be transparent
"ane_layout_vs_gpu": 70.0, # strict: two layouts, same math
"compiled_vs_torch": 40.0, # loose: toolchain drift is expected
"palettized_4bit": 35.0, # loosest: compression costs
}
def gate(name: str, value: float) -> None:
threshold = GATES[name]
verdict = "PASS" if value >= threshold else "INVESTIGATE"
print(f"{name:<24} {value:6.2f} dB (gate {threshold:.0f} dB) {verdict}")
Three practical notes on applying them:
- PSNR is peak-relative, so it is scale-sensitive. Comparing logits (unbounded, occasionally enormous) against comparing a normalised feature map will give you very different dB for the same relative error. Fix a convention and keep it.
- Gate per output, not on a mean. A model with four outputs where three score 80 dB and one scores 12 dB has an average of 63 dB and a broken output.
- Carry a task metric alongside. §15.3's record format does exactly this —
psnronlogitsandiouonbbox. For an autoregressive model, the equivalent of a task metric is greedy-token agreement over a real generation, not a single-step logit comparison; see §17.1.
16.4 The gate that is not a PSNR gate¶
For LLM decoders, dB is a poor proxy, and the community archive is emphatic about why:
🟢 COMMUNITY-MEASURED —
notes/repos/john-rocky-models.md, single-author, uncontrolled benchmarks. Its pass bar for a language model is "per-token cosine ≥ 0.999 on logits AND greedy argmax token-exact" across a real generation, with the reasoning: "Per-step matters: an AR loop can look fine at step 1 and drift by step 30." It also records the hardest-won lesson of that whole archive as "read the generations" — a model that passes every numeric gate can still emit broken grammar.
That is a community bar, not Apple's. But the logic is sound and generalises: for any model whose output feeds an iterative process, single-step numeric similarity understates compounding error.
17. Community-measured findings, labelled¶
Everything in this section comes from notes/repos/john-rocky-models.md, a single-author
community archive with self-declared uncontrolled benchmarks. Its numbers are frequently unique —
nobody else has published Core AI compression measurements at this scale — and they are frequently
useful. They are not Apple figures, they were not produced under a controlled protocol, and
several of them are internally in tension. Every one is labelled.
17.1 Dense int4 k-means: what the archive actually found¶
Finding 1 — palettization beats quantization by a lot, at fine granularity.
🟢 COMMUNITY-MEASURED: "Per-channel (group_size=1) basically always wins; at per-channel, k-means beats quantization by ~15–19 dB at both 8-bit and 4-bit. Per-tensor palettization can be worse than per-channel quantization." Mechanism the author gives: "k-means fits a per-group lookup table to the actual weight clusters → tracks non-uniform weight distributions far better than symmetric per-block int4 at the same bit width."
15–19 dB is a very large gap — more than the entire margin between Apple's 4-bit gate (35 dB) and
its 8-bit expectation (55 dB). Treat it as a strong hint that if you are quantizing weights per-
channel and struggling, palettization at the same nominal width is the first thing to try. Then read
§6.3 and price the LUTs before you commit to group_size=1.
Finding 2 — int4 is a cliff, not a slope.
🟢 COMMUNITY-MEASURED, stated repeatedly across the archive: "int4 is a cliff, not a slope" — "the failure is capacity, so no clever rounding rescues". And: "Across Gemma 4 E2B and Qwen3.5, linear int4 and k-means int4 both flip next-token argmax vs the HF reference; int8 k-means palettization reproduces HF top-1 exactly at ~half the fp16 size." "Finer groups are the main int4 lever (group32 → group8 helps), but still don't reach exact. Per-channel scale is marginal or harmful."
Note the last clause independently corroborates §5 from a completely different direction: per-channel scale was measured to be "marginal or harmful" on quality, which means the ANE cost it imposes buys nothing even before you count the compute-unit fallback.
Finding 3 — selective 4-bit works where uniform 4-bit does not.
🟢 COMMUNITY-MEASURED, the archive's summary stance for LLM decoders: "int8 k-means palettization is the floor that stays exact when applied across the whole transformer; whole-model int4 degrades. SELECTIVE 4-bit works: k-means int4 on the FFN + lm_head only (attention/embeddings kept ≥ int8/fp16) measured top-1 exact." Specifically for Gemma 4: "the gate/up MLP projections must be int8 for exactness — keeping them at 4-bit caps accuracy regardless of other layers." Recommended recipe: "int8 k-means, group 32, all projections; keep tied
lm_head+ 1-D conv (SSM) full precision."
This is mixed precision (§14) arrived at empirically, and it agrees structurally with Apple's own mixed 4/8-bit Qwen3 YAMLs (§14.4): neither party found uniform 4-bit acceptable on a transformer.
Finding 4 — but which int8, and the archive contradicts itself, on purpose.
🟢 COMMUNITY-MEASURED, and flagged by the archive itself as a tension: a separate measurement on MoE experts (ZAYA / LFM) found
sym8— symmetric linear int8 with per-K-block-32 scales — clean, and k-means int8 lossier, for top-k ≥ 4 routing. The same note records the reverse for top-1 routing, wherekm8(k-means int8, 256-entry codebook) "recovers fp16 quality" while sym8 fails. The archive's own synthesis: "int8 is the safe floor everywhere; which int8 (k-means vs symmetric-linear-block32) is tensor-role- and routing-dependent and must be gated per model."
⚠️ Do not flatten those two into one rule. Both were measured, at different dates, on different tensor roles. The defensible takeaway is the archive's own: int8 is the floor, and the choice of which int8 is a per-model measurement, not a default.
Finding 5 — dense-path int4 k-means is a speed lever with a known quality cost.
🟢 COMMUNITY-MEASURED —
dense-int4km-flagship-session-findings.md, dated 2026-07-01. All numbers from that author's own microbenchmarks:
Measurement Number Hardware lm_headint4km per-op vs fp16, vocab 248 K2.77× M4 Max Qwen3.6-35B decode, int4km dense path vs baseline 2.18× ratio M4 Max GPU LFM-8B on-device decode 1.23× sustained / 1.43× avg iPhone 17 Pro (A19) lm_headFP4-E2M1 vs fp16, vocab 248 K1.72× (and 1.01× vs int4km) M4 Max ⚠️ The author's own caveats, which matter more than the numbers: "Absolute tok/s from the quick Mac driver are ~10× too slow" — only the ratio is valid, not the throughput. On quality: "The dense-int4km lever itself is quality-safe (LFM-8B coherent). Flagship int4 degrades quality … the 2.18× is a speed win at a known int4 quality cost."
And one structural constraint from the same archive that bears directly on §2.6:
🟢 COMMUNITY-MEASURED: "k-means palettization is
F.linear-only", so an int4 head "needs a kernel path, not coreai-opt'sF.linearquantizer." Thelm_headis "the largest single tensor, high sensitivity, needs per-row (per-output-channel) scales for matmul efficiency" — e.g. 262,144 × 1536 for Gemma 4.
17.2 Boundary layers¶
🟢 COMMUNITY-MEASURED: "Boundary layers (first/last) are high-error — skipping them can add up to +9 dB; always ablate."
Independently corroborated by Apple's own mixed-precision result (§14.1), where the two layers
promoted to 6 bits were conv1 (the stem) and layer1.0.downsample.0 (a projection shortcut). Two
sources, two model families, same conclusion. Ablate your first and last compressible layers before
anything else; it is one experiment and it is frequently the largest single win.
17.3 The free 1.7× that everyone leaves on the table¶
This one is not about compression at all, and it is the cheapest win in this guide.
🟢 COMMUNITY-MEASURED — vision models, M4 Max, Python runtime, method stated by the author: load each official
.aimodelwithSpecializationOptions.from_preferred_compute_unit_kind(<unit>), synthetic inputs from the function descriptors, 3 warmup + 20 timed runs, median single-inference latency. Caveat the author attaches: "'Preferred' means the runtime may still place unsupported ops elsewhere."
Model Recipe Artefact GPU ANE CPU Winner clip-vit-base-patch32 fp32 static 577 MB 6.54 ms 5.43 ms 18.76 ms ANE clip-vit-base-patch32 fp16 ( --dtype float16)289 MB 6.31 ms 3.68 ms — ANE, 1.7× over GPU yolos-base fp32 static 488 MB 444.8 ms 456.7 ms 733.7 ms GPU (≈ tie) sam3 fp32 static 3.1 GB 559.9 ms 565.7 ms 2789.7 ms GPU (≈ tie) depth-anything-3 (small) fp32 static 101 MB 7.30 ms 6.84 ms 34.58 ms ANE The author's conclusion, verbatim: "every official CV recipe DEFAULTS to float32, and at fp32 the big ViTs land in a GPU/ANE tie on M4 Max. But the scripts expose
--dtype float16, and fp16 is what the ANE runs natively: CLIP at fp16 drops to 3.68 ms on ANE (1.7× faster than GPU, 1.5× faster than fp32-ANE) at half the artifact size. If you're deploying these recipes to ANE, pass--dtype float16." (First ANE load of a new variant pays a ~5 s one-time specialization for CLIP.)
The mechanism is the one from §5.1: fp32 is not an ANE dtype. Apple's own authoring skill says so — "fp16 only — no fp32 literals anywhere", and fp32 "falls back to GPU/CPU". Shipping an fp32 recipe to an ANE target is asking for the fallback by construction.
So before you spend a week on a palettization sweep:
✅ VERIFIED — the
--dtype {float16,float32}flag with defaultfloat32exists on the SAM3 export script's--fullpath (models/sam3/README.md). Whether every utility export script exposes it identically is 🔴 unverified — check--helpfor yours.
⚠️ Two honest caveats before you take this as a universal rule. First, it is one author's
measurement on one machine, and the two large models in that table (yolos-base, sam3) show
essentially no ANE advantage at all — so "ANE wins" is not general. Second, the archive separately
records that on iOS 27 beta, "CV graphs silently fall back to GPU even with an ANE preference
and a pure-ViT split backbone (fingerprint-identical outputs, zero ANE-compile wait)". Measure on
your target, not on a Mac.
17.4 Palettization composes with stateful export, with an ordering rule¶
🟢 COMMUNITY-MEASURED: "read the export spec (reference inputs / dynamic shapes / state names) from the ORIGINAL model first (the finalized palettized model loses that method), palettize, then drive
export_to_coreaiwith that spec." Verified by the author as top-1-exact for Gemma 4 (dual-KV) and Qwen3.5 (hybrid 4-state).
That is the same failure mode as §10.2 in a different costume: finalize() returns a model that has
lost something you needed, and the loss is silent. Read what you need off the float model first.
17.5 A beta bug worth knowing about¶
🟢 COMMUNITY-MEASURED, beta-era, reproduction dated 2026-06-11: "per-channel (axis-0) int8 Linear weights are broken on the macOS-27-beta MPSGraph GPU delegate — torch-level numerics are clean but the lowered matmul returns garbage (minimal head-only repro, multiple shapes, sym and clipping alike); use per-block-32 there."
This is a quantization bug, not a palettization one, and it is beta-specific. It is included here for one reason: it is a case where your PyTorch score is clean and the device output is garbage, which is exactly the failure class the Debugger exists for and exactly why §16's compiled vs torch gate is a separate gate from the authoring gates. Do not assume a passing torch-level PSNR means your artefact works.
18. The worked examples the repo ships¶
If you read one thing after this guide, read these. coreai-optimization ships four example pages
and four Jupyter notebooks, and the notebooks are the only place in the entire Core AI ecosystem
where you can watch a complete compression workflow run end to end on a model small enough to
iterate on. Remember the warning at the top of this guide: Core AI has zero Apple sample-code
Xcode projects.
These notebooks are the closest thing that exists to a first-party sample.
18.1 The four MNIST notebooks¶
✅ VERIFIED —
docs/src/tutorials/:They are executed in CI:mnist_quantization.ipynb mnist_palettization.ipynb mnist_palettization_and_activation_quantization.ipynb ← joint compression mnist_pruning.ipynb ← added in commit b1535b4make test-tutorialsrunsdocs/tests/test_tutorials.pyvia papermill (Makefile), onubuntu-latest, as part of stage 1 of.github/workflows/ci.yaml. So they are kept working, not merely kept.✅ VERIFIED — the model they use, from
tests/test_joint_compression.pyandtests/models/mnist.py: 6 weight-bearing layers —conv1,conv2,conv_transpose1,conv_transpose2,dense1,dense2— withbatch_size = 128andnum_calibration_batches = 17. Checked-in artefacts:tests/_test_artifacts/mnist/mnist_pretrained_1epoch_09032025.ptandmnist_example_input_11122025.pt. Acceptance: baseline > 97.0%, joint P4-A8 after calibration > 90.0%.
Suggested order, and what each one is actually for:
| # | Notebook | Read it to learn |
|---|---|---|
| 1 | mnist_quantization |
the lifecycle, GRAPH mode, calibration. Guide 01's material, runnable. |
| 2 | mnist_palettization |
this guide's §2–§7 with real numbers, in about five minutes of compute. |
| 3 | mnist_pruning |
the only runnable pruning example anywhere, including step() cadence. |
| 4 | mnist_palettization_and_activation_quantization |
§13's five-step flow, with the op_state_spec=None line in context. |
Notebook 3 deserves particular attention precisely because pruning got no stage time — it is the
only place you can see the schedule/step() interaction working rather than reasoning about it
from a docstring.
Note that the presence of conv_transpose1 / conv_transpose2 in the MNIST model is not
incidental: ConvTranspose* palettization was added in 0.2.1, and this model is how it is
tested.
To run them:
git clone https://github.com/apple/coreai-optimization
cd coreai-optimization
make env-tutorial # creates .venv-tutorial with the tutorial dependency group
source .venv-tutorial/bin/activate
jupyter lab docs/src/tutorials/
✅ VERIFIED —
Makefile:make env-tutorialbuilds.venv-tutorialwith--with-tutorial. Also useful:make env(dev + coreai + coreml groups),make test-fast(--marker 'not slow'),make api-list [MODULE=coreai_opt.palettization.spec.spec]to print the public surface of any module. ⚠️ Auvrule from Apple's ownAGENTS.md(which is literally titled# CLAUDE.md): "Always pass--no-synctouv run:uv run --no-sync --active ….uv runimplicitly syncs the active project to its default-groups before running, which re-resolves dependencies and can clobber a venv's group-pinned packages." If you are working across the torch-version venvs, this will save you an afternoon.
18.2 The four example pages¶
✅ VERIFIED —
docs/src/examples/:toy_models.md,resnet50.md,edsr.md,mixed_precision_palettization.md.
toy_models.md — a Conv2d → ReLU → Linear stack, used to demonstrate the per-tensor SNR
comparison utility. Its output table is the single best illustration of how error accumulates
through a network:
✅ APPLE-PUBLISHED —
docs/src/utils/activation_comparison.md:286-295, default INT8, graph mode:conv_weight -> activation_post_process_1 SNR = 47.17 dB conv_bias -> conv_bias SNR = inf dB linear_weight -> activation_post_process_4 SNR = 48.13 dB x -> activation_post_process_0 SNR = 43.20 dB conv2d -> conv2d SNR = 42.40 dB relu -> activation_post_process_2 SNR = 38.94 dB flatten -> activation_post_process_3 SNR = 38.94 dB linear -> activation_post_process_5 SNR = 35.74 dB
Read it top to bottom: 48 dB at the weights, 35.7 dB by the output. Twelve dB of accumulated
loss across three operations, monotonically decreasing. flatten costs exactly nothing (identical
to relu — it is a pure reshape), and the final linear costs 3.2 dB on its own. This is why §16's
gates get looser as you move down the stack, and why gating on a final-output PSNR alone hides where
the damage happened.
resnet50.md — quantization PTQ, and the clearest available demonstration that the choice of
qparams calculator matters:
✅ APPLE-PUBLISHED — 128 evaluation samples from imagenette:
Config Accuracy FP32 baseline 78.12% W_INT8(per-channel) _ A_INT8(per-tensor), moving_average, 896 calibration samples74.22% same but global_minmaxactivations75.78% W_FP8_E4M3 _ A_FP8_E4M3, global_minmax76.56% ⚠️ Note the small evaluation set — 128 samples. A 1.5-point difference on 128 images is inside the noise. Treat the ordering as informative and the magnitudes as indicative.
edsr.md — the joint-compression example, reproduced in §13.3. Super-resolution is an
unusually good compression testbed because PSNR is the task metric, not just a proxy.
mixed_precision_palettization.md — §14.1. The only place Apple publishes a full-ImageNet
(50,000 image) evaluation of a compression recipe.
18.3 The test suite as documentation¶
✅ VERIFIED —
tests/layout:tests/ test_smoke.py # minimal quant (eager+graph) + palettization smoke test_joint_compression.py # P4-A8 MNIST accuracy (slow; needs coreai) test_palettization_preset.py # what each preset expands to palettization/ # kmeans palettizer (+mnist), kmeans_fake_palettize, # kmeans1d, kmeans_parallel, support mixins, config, spec pruning/ # magnitude pruner (+mnist), config/spec, sparsity schedule export/ # eager & graph MIL/MLIR export, kmeans export, pruning export coreai_utils/ # quantize/palettize/sparsify weights, sparse utils models/ # resnet.py, simple.py, mnist.py fixtures/ # quantization.py, palettization.py, pruning.py, fp4.py, fp8.py
tests/palettization/test_kmeans_parallel.py is the reference for num_workers behaviour,
tests/pruning/test_sparsity_schedule.py for the schedule formula, and tests/fixtures/ for
minimal working config objects you can copy. When a docstring and your intuition disagree, the test
is the tiebreaker.
19. ⚠️ Silent failures, consolidated¶
coreai-opt has a consistent house style: when compression cannot be applied, it is skipped with
a log line rather than raised. That is a defensible design for a sweep tool and a dangerous one
for a build pipeline. Here is everything in this guide's scope that fails quietly.
19.1 A group size your weight is not divisible by leaves the layer uncompressed¶
The big one.
✅ VERIFIED —
PerGroupedChannelGranularityraises the internal_IncompatibleGranularityErrorwhenshape[axis] % group_size != 0. That error is caught inside_FakePalettizeImplBase.forward, which logs a warning, permanently disables the module, and removes the parametrization. Palettization has a_remove_disabled_fake_palett_modulespass analogous to quantization's_remove_disabled_fake_quant_nodes. The log line, from the transcript corpus's reading of the source:"Tensor incompatible with granularity: ... Skipping palettization."⚠️ SILENT FAILURE.
prepare()returns successfully. The model scores better than you expected, because some layers were never compressed. Your size estimate — computed from the config — is wrong. You ship a model that is bigger than your release notes say, and possibly bigger than the memory budget you sized the feature against. There is no exception, no non-zero exit, and the warning is one line in atqdm-heavy log you were not reading.The same hazard exists for
cluster_dim(a shape not divisible by the vector size) and, on the quantization side, forPerBlockGranularityblock sizes.Apple's own compression-exploration skill names this as a pitfall and prescribes the fix: "per-block / per-grouped-channel silently skip layers whose weight dim isn't divisible (pre-check with
check_divisibility())". The community archive independently names it: "Check divisibility before trusting a size."Pre-flight it — the
A count mismatch is your signal.check_divisibilityimplementation in §14.3 is 12 lines and catches every instance. Or verify after the fact:
19.2 enable_per_channel_scale=True moves your model off the Neural Engine¶
§5.2. Rank-6 LUT, ANE rejects, runtime falls back to GPU. No warning, no error, marginally better PyTorch numerics. Detected only by inspecting compute-unit placement on device.
19.3 Eager-only KMeansPalettizer.finalize(CoreAI) frees dense weights irreversibly¶
§10.2. The dense pre-palettization weight is replaced with a zero-size placeholder. Combined with
prepare()'s in-place mutation (§2.4), a careless script can leave you with no float model in the
process at all. copy.deepcopy before prepare(); do not finalize() during a sweep.
2
19.4 Vector palettization produces a different model every run¶
§7.1. Non-deterministic centroid initialisation, and seeding does not reach spawned workers. A golden-hash regression test will fail for reasons unrelated to your change, and — worse — the artefact you evaluated is not the artefact you shipped.
19.5 Channel-structured pruning can silently prune zero channels¶
§11.4. floor(num_channels × target_sparsity). With 8 channels and a 10% target, nothing is pruned
and nothing says so.
19.6 Joint compression without lut_qspec buys you no speed¶
§13.4. A floating-point LUT forces floating-point ops regardless of how thoroughly you quantized the activations. Everything runs; the latency does not move; nothing explains why.
19.7 Joint compression with the default op_state_spec double-compresses weights¶
§13.4. Forgetting op_state_spec=None in the quantization config means the already-palettized
weights get quantized on top. Quality drops far more than the published EDSR deltas suggest, and the
cause is a missing three-character value.
19.8 --n-bits on the SAM3 export script overrides both encoders¶
§5.4. The shipped defaults are deliberately asymmetric (image 4/gs32, text 6/gs8); the CLI flag applies uniformly and silently discards the asymmetry.
19.9 A sweep optimising PSNR will pick the setting that costs you the ANE¶
§15.3. enable_per_channel_scale is one of the axes Apple's own exploration skill sweeps, and it
improves PyTorch PSNR. The sweep has no visibility into compute-unit placement. Constrain the search
space before you run it against an iOS target.
19.10 Inherited from quantization, and still true here¶
From guide 01 §17, because they apply to every compressor in the package:
example_inputsmust be a tuple; a bare tensor raises, but an empty tuple raises differently.- Re-preparing a prepared model raises — except after
KMeansPalettizer.finalize(), which deliberately clears the marker so joint compression can work (§13.1). That asymmetry is intentional and easy to misread as a bug. - Configs are
@final; specs are frozen. Mutate withmodel_copy(update={...}). module_type_configskeys must be fully-qualified;"torch.nn.Linear"matches nothing and raises only if it has no dot —"torch.nn.Linear"does have dots, so it passes the syntactic check and then matches nothing. ⚠️ That one is genuinely silent.only_forcannot be chained.
20. Numbers, attributed¶
Every number in this guide, with its source and its caveats. Nothing here is "measured by us".
20.1 Apple-published¶
| Number | Context | Source | Caveats |
|---|---|---|---|
| 75.02% → 65.87% → 70.27% top-1 | ResNet50 fp16 / uniform 4-bit per-tensor / greedy mixed precision | docs/src/examples/mixed_precision_palettization.md |
full ImageNet val (50 k), mps backend |
| 48.64 / 12.16 / 12.03 MB, BPW 16 / 4 / 3.95 | same three ResNet50 configs | same | theoretical weight storage |
| 2 layers @ 6-bit, 50 @ 4-bit, 2 @ 2-bit | the greedy recipe (conv1, layer1.0.downsample.0 up; layer1.1.conv1, layer3.4.conv2 down) |
same | 54 configured layers |
| inflection at ≈ 4.0 BPW; below it 15–35 pts per 0.5 BPW, above it 1–2 | ResNet50 palettization curve | same | ResNet50 only; do not assume it transfers |
| 30.68 / 30.33 / 29.86 dB PSNR | EDSR fp32 / W_INT8-A_INT8 / W_P4(INT8)-A_INT8 joint | docs/src/examples/edsr.md |
edsr_r16f64, 1.5 M params, B100, 20 calib / 80 eval |
| ~5.5 / ~1.4 / ~0.7 MB (1× / 4× / 8×) | same three EDSR configs | same | weight storage |
| 78.12 / 74.22 / 75.78 / 76.56% | ResNet50 PTQ: fp32 / int8 moving_average / int8 global_minmax / FP8-E4M3 global_minmax |
docs/src/examples/resnet50.md |
⚠️ only 128 eval samples; 896 calibration samples |
| 47.17 → 35.74 dB SNR down the stack | toy Conv2d→ReLU→Linear, default INT8, graph mode |
docs/src/utils/activation_comparison.md:286-295 |
illustrative toy model |
| MNIST baseline > 97.0%, joint P4-A8 > 90.0% | test assertions | tests/test_joint_compression.py |
acceptance thresholds, not measurements |
| Qwen3 0.6B iOS: 26.16 → 30.90 perplexity at 5.71* BPW | mixed 4/8-bit palettized YAML | apple/coreai-models, models/*/README.md |
WikiText-2 via lm-evaluation-harness; * includes int8 embedding |
| Qwen3 4B iOS: 16.41 → 18.80 at 4.89* BPW | mixed 4/8-bit palettized YAML | same | same |
| Qwen3 4B macOS: 16.41 → 18.33 at 4.50 BPW | 4-bit quantized (not palettized) | same | different technique, different platform |
| Qwen2.5 1.5B iOS: 12.21 → 14.64 at 4.63* BPW | 4-bit palettized gs 8 | same | |
| 8-bit ≈ 2× / > 55 dB, 4-bit ≈ 4× / ~40 dB, 2-bit ≈ 8× / 25–35 dB | palettization sizing guidance | skills/skills/model-authoring/SKILL.md |
2-bit annotated "Usually unacceptable" |
| 70 / 70 / 40 / 35 dB | the four acceptance gates | model-authoring/SKILL.md:94-99 |
§16.1 |
| > 70 / > 50 / ~40 dB | deployment-side PSNR expectations | model-deployment skill |
§16.2; slightly different scoping |
| 96% of SAM3 params in the two encoders; 4% in the detector | SAM3 structure | WWDC26 session 325, 325:60 and 325:158 | 848 M params per models/sam3/README.md; talk says "850-million" |
| 76% faster second inference | SAM3 three-function split, prompt-only change | session 325, 325:255-256 | ⚠️ requires caller-side caching Apple's package does not do |
| 5-15% scale/ZP overhead at 2-4 bit fine granularity | sizing guidance | model-compression-exploration/SKILL.md |
|
~5 bits effective at block_size=16 + int4 |
same | same | quantization, included for contrast |
20.2 Community-measured — notes/repos/john-rocky-models.md¶
⚠️ Single-author archive with self-declared uncontrolled benchmarks. Not Apple figures. Several entries are internally in tension (§17.1, finding 4) and the archive says so itself.
| Number | Context | Hardware / date | Caveats |
|---|---|---|---|
| ~15–19 dB k-means over quantization | per-channel, at both 8-bit and 4-bit | not stated | the largest single claim in this table; unreplicated |
| +9 dB from skipping boundary layers | first/last compressible layers | not stated | corroborated in shape by Apple's ResNet50 recipe |
| 3.68 ms ANE vs 6.31 ms GPU (1.7×) | clip-vit-base-patch32, fp16, 289 MB | M4 Max, 3 warmup + 20 timed, median | fp32 variant: 5.43 ANE / 6.54 GPU, 577 MB |
| 444.8 GPU / 456.7 ANE ms | yolos-base fp32, 488 MB | M4 Max, same protocol | ≈ tie — ANE advantage is not general |
| 559.9 GPU / 565.7 ANE ms | sam3 fp32, 3.1 GB | M4 Max, same protocol | ≈ tie |
| 2.77× per-op | lm_head int4km vs fp16, vocab 248 K |
M4 Max microbenchmark | per-op, not end-to-end |
| 2.18× ratio | Qwen3.6-35B decode, dense int4km path | M4 Max GPU, 2026-07-01 | ⚠️ author: "absolute tok/s ~10× too slow; only the RATIO is valid" |
| 1.23× sustained / 1.43× avg | LFM-8B decode, dense int4km | iPhone 17 Pro (A19), thermally matched | |
| 1.72× fp4-E2M1 vs fp16; 1.01× vs int4km | lm_head, vocab 248 K |
M4 Max | fp4 ≈ int4km on speed |
| 2.2× artefact speed difference | same export recipe, macOS 26 vs 27 beta | beta-era | "an .aimodel is a build artefact, not a pure function of the recipe" |
| 7.0 GB fp32 → 3.5 GB fp16 → 1.9 GB int8 | Gemma 4 E2B core | — | sizes, not quality |
20.3 What is not measured anywhere¶
Stated plainly, because absence of a number is information:
- No latency or size number for
cluster_dim > 1on any hardware (§4.3). - No before/after quality delta for sensitivity-weighted k-means in
coreai-opt(§8.4). - No evidence that any Apple compute unit exploits unstructured, block or
n:msparsity for either size or speed (§11.8, §12.2). - No quality or latency number for any non-LLM model in
apple/coreai-models. ✅ Verified:Tools/benchmarkis actuallyllm-benchmarkand importsCoreAILanguageModels; there is no non-LLM benchmark tool in that repository. - No enumeration of the Core AI Debugger's available similarity metrics (§15.5).
21. Quick reference¶
21.1 Imports¶
# Palettization
from coreai_opt.palettization import (
KMeansPalettizer, KMeansPalettizerConfig, ModuleKMeansPalettizerConfig, PalettizationSpec,
)
from coreai_opt.palettization.spec import (
PalettizationGranularity, PerGroupedChannelGranularity,
PerTensorGranularity, # ⚠️ NOT the quantization one
default_weight_palettization_spec, DefaultTrainingSpec, TrainingStrategy, TrainingStrategySpec,
)
from coreai_opt.palettization.config import PATSchedule
# Pruning
from coreai_opt.pruning import (
MagnitudePruner, MagnitudePrunerConfig, ModuleMagnitudePrunerConfig, PruningSpec,
)
from coreai_opt.pruning.spec import (
ChannelStructured, PruneImplBase, PruningScheme, Unstructured, default_weight_pruning_spec,
)
from coreai_opt.pruning.config import (
ConstantSparsitySchedule, OpMagnitudePrunerConfig, PolynomialDecaySchedule, SparsityScheduleBase,
)
# Program level (no PyTorch model needed)
from coreai_opt.coreai_utils import (
CompressionGranularity, DType, palettize_weights, quantize_weights, sparsify_weights,
)
from coreai_opt.coreai_utils.common import QScheme # not in the package __all__; §12
# Shared
from coreai_opt import ExportBackend, CoreMLExportError
from coreai_opt.casting import cast_to_16_bit_precision
21.2 The decision table¶
| Situation | Reach for |
|---|---|
| Targeting iOS / Neural Engine | Palettization. presets.w4(group_size=32). enable_per_channel_scale=False. |
| Targeting macOS / GPU | Quantization. Guide 01. Apple's own exporter agrees. |
| First experiment on any model | presets.w8() (per-tensor, 2×) — establishes the ceiling cheaply |
| 4-bit is not good enough | ↑ to w6, or ↓ group_size, or mixed precision, or sensitivity weighting |
| Need weight and activation compression | Joint compression (§13). Palettize with lut_qspec first. |
| Need < 4 effective bits | cluster_dim=2 — but read §4.3's GAP box first |
| Model is small (< 1 B params) | Expect a worse curve than the papers. See §14.4. |
| Want speed on device | Try --dtype float16 before compressing anything (§17.3) |
| Want to know which layers to skip | Sensitivity sweep (§14.3), then the Debugger (§15.5) |
| Considering pruning | Read §11.1 and §11.8. Probably don't. |
| Need to compress an embedding table | Not palettization (§2.6). int8 per-tensor, or palettize_weights at the program level. |
21.3 Field cheat-sheet¶
PalettizationSpec(
n_bits=4, # {1, 2, 3, 4, 6, 8} — NO 5, NO 7
lut_qspec=None, # int8/uint8/fp8_e4m3fn/fp8_e5m2; per-tensor; ZP only
granularity=PerTensorGranularity(), # or PerGroupedChannelGranularity(axis={0,1}, group_size=N)
cluster_dim=1, # >1 ⇒ vector; requires enable_fast_kmeans_mode=False
enable_per_channel_scale=False, # ⚠️ True ⇒ rank-6 LUT ⇒ ANE rejects ⇒ GPU fallback
training_strategy_spec=DefaultTrainingSpec(), # frozen one-shot centroids/indices
)
ModuleKMeansPalettizerConfig(
op_state_spec={"weight": spec, "in_proj_weight": spec}, # BOTH keys
enable_fast_kmeans_mode=True, # rounds weights to `rounding_precision` decimals first
rounding_precision=4,
pat_schedule=PATSchedule(enable_fake_palettize=0), # step() controls activation
)
PruningSpec(
target_sparsity=0.5, # a REQUEST, not a guarantee (§11.4)
pruning_scheme=Unstructured(), # or ChannelStructured(axis=0)
pruning_algo="default", # → _MagnitudePruneImpl
)
PolynomialDecaySchedule(
begin_step=0, total_iters=N, # REQUIRED; units = your step() cadence
power=3.0, initial_sparsity=0.0, update_frequency=1,
)
21.4 Legal-value quick check¶
| Thing | Legal values |
|---|---|
n_bits |
1, 2, 3, 4, 6, 8 |
| LUT dtype | int8, uint8, float8_e4m3fn, float8_e5m2 |
| LUT granularity | quantization PerTensorGranularity only |
| LUT formulation | ZP only (MINVAL rejected) |
| Palettization granularity axis | 0 or 1 |
| Palettization ops | conv1d/2d/3d, conv_transpose1d/2d/3d, linear, multi_head_attention_forward |
| Pruning ops | linear, conv1d/2d/3d, conv_transpose1d/2d/3d, multi_head_attention |
| Program-level palettization granularity | PER_TENSOR, PER_CHANNEL, PER_GROUPED_CHANNEL |
sparsify_weights |
target_sparsity XOR n_m_ratio; quantize_dtype XOR palettize_nbits |
| Joint compression backend | ExportBackend.CoreAI only |
| CoreML palettization | at most one of {cluster_dim>1, lut_qspec, enable_per_channel_scale} |
21.5 Pre-flight checklist¶
□ copy.deepcopy(model) before prepare()
□ example_inputs is a TUPLE
□ shape[axis] % group_size == 0 for every target layer ← §19.1, the big one
□ enable_per_channel_scale is False (unless GPU-only) ← §5
□ op_state_spec lists BOTH "weight" and "in_proj_weight" ← §2.5
□ num_workers > 1 for anything bigger than a toy ← §7.2
□ cluster_dim == 1, or seeded with num_workers=1 ← §7.1
□ no finalize() during a sweep ← §10.2
□ embeddings accounted for separately ← §2.6, §6.4
□ a task metric alongside PSNR ← §15.3
□ the toolchain version recorded next to every result ← §15.4
22. Sources and evidence ledger¶
22.1 Primary — shipped source¶
apple/coreai-optimizationat release189612be, packagecoreai-opt0.3.0. Files read for this guide:palettization/{__init__,base_palettizer}.py,palettization/spec/{__init__,spec,granularity,fake_palettize}.py,palettization/config/{__init__,palettization_config}.py,palettization/config/_presets/{kmeans_palettizer_config,module_kmeans_palettizer_config}.py,palettization/kmeans/{palettizer,kmeans_fake_palettize,supported_ops_registry,_prepare_for_export}.py,pruning/{__init__,magnitude_pruner,supported_ops_registry}.py,pruning/spec/{__init__,spec,scheme,prune}.py,pruning/config/{__init__,sparsity_schedule,magnitude_pruner_config}.py,coreai_utils/{__init__,common}.py,coreai_utils/passes/{__init__,weight_quantization,weight_palettization,weight_sparsification}.py,base_model_compressor.py,common.py,config/compression_config.py,_utils/export_utils.py,deps/_kmeans1d/core.py,casting/__init__.py,pyproject.toml,Makefile,CHANGELOG.md,AGENTS.md,.github/workflows/ci.yaml,changelog.d/{31.changed,42.fixed,52.fixed}.- Docs in that repo:
docs/src/palettization/{basics,overview,config}.md,docs/src/pruning/{basics,overview,config}.md,docs/src/utils/{joint_compression,mixed_precision,casting,coreai_compression,activation_comparison}.md,docs/src/examples/{toy_models,resnet50,edsr,mixed_precision_palettization}.md,docs/src/introduction/{installation,how_to_use_coreaiopt,integration_coreai}.md,docs/src/landing_page.md,docs/src/tutorials/*.ipynb(listing only). apple/coreai-models:models/sam3/{pipeline.py,README.md,export.py},python/src/coreai_models/llm/{export.py,model_registry.py},python/src/coreai_models/export/presets.py,models/README.md,models/*/README.md.- Tests:
tests/test_joint_compression.py,tests/test_smoke.py,tests/models/mnist.py, directory listings fortests/{palettization,pruning,export,coreai_utils,fixtures}.
22.2 Primary — Apple's agent skills (vendored in apple/coreai-models)¶
These are Apple engineers' empirical rules, written for machine consumption. In this guide they are treated as stronger than session transcripts and weaker than source.
skills/skills/model-authoring/SKILL.md— the four PSNR gates (:94-99), the palettization sizing table (:149-153), the ANE/GPU at-a-glance table, the KV-cache conventions table.skills/skills/model-authoring/references/neural_engine_rules.md(479 lines) — max tensor rank 5, fp16/int8/int16 dtypes, 64-byte last-axis alignment, BC1S conversions, causal mask conventions, conv stride factorisation.skills/skills/model-authoring/references/gpu_rules.md(297 lines) — the contrast case.skills/skills/model-authoring/references/common_issues.md(176 lines).skills/skills/model-compression-exploration/SKILL.md— the ~60-config sweep, the three groups, the refinement pass, the output record shape, the pitfalls list.skills/skills/model-deployment— the deployment-side PSNR acceptance table.
22.3 Primary — WWDC26¶
- Session 325, "Dive into Core AI model authoring and optimization". Cited at 325:60 (96%
encoders), 325:64–68 (
coreai-optcapability claims), 325:139–152 (Debugger sync points and PSNR), 325:156–162 (the detector diagnosis), 325:158 (4% detector), 325:241–248 (the palettization passage), 325:249–262 (export and the 76% figure), 325:263–267 (closing recommendations). - Session 330, "…Metal tensors and TensorOps". Cited at 330:21–26 for the memory-bandwidth-bound framing only. Everything else about session 330 belongs to Part 11.
22.4 Community¶
notes/repos/john-rocky-models.md— a single-author community archive with self-declared uncontrolled benchmarks. Every citation in this guide is labelled 🟢 COMMUNITY-MEASURED with hardware and caveats where the archive states them. Itscompression.md,compression-reference.md,compute-units-and-authoring.md,apple-models-bench.mdanddense-int4km-flagship-session-findings.mdsections were used.
22.5 What this guide deliberately does not claim¶
For the record, and so that a future editor does not "fix" an absence into an invention:
- No
coreai_opt.mixed_precisionmodule. The package tree hasquantization/,palettization/,pruning/,casting/,inspection/,coreai_utils/,config/,deps/,_utils/. There is no mixed-precision subpackage in it. §14's sweep is code you write. - No built-in differentiable centroid training by default, and no
DKMPalettizer; 0.3.0 provides training strategies and PATSchedule (§2.1). - No pruning presets.
MagnitudePrunerConfig.presetsdoes not exist. - No
n_bits=5orn_bits=7. - No claim that the Core AI runtime exploits sparsity. §11.8 is a declared gap, not an omission.
- No
.coreaimodelor.aiassetfile extension, nocoreai-torch convertCLI, no "iOS 20 / macOS 17", and no on-device LoRA training API. All four are in circulation and all four are fabricated. Part 1 carries the known-bad-claims reference.
22.6 Open questions, collected¶
Every 🔴 GAP in this guide, in one place, so they can be closed by someone with a device:
| § | Unknown | What would resolve it |
|---|---|---|
| 2.1 | Resolved: learned centroids require a custom training strategy | 0.3.0 default freezes centroids/indices; see training_strategy.py |
| 4.3 | Device latency and compute-unit placement for cluster_dim > 1 |
timing two artefacts under from_preferred_compute_unit_kind |
| 5.3 | The precise axis accounting behind the rank-6 LUT | an MLIR dump of the lowered program, or an mps.dequantize_lut doc page |
| 8.4 | The quality delta from sensitivity-weighted k-means in coreai-opt |
an ablation on edsr or resnet50 |
| 10.3 | Resolved: per-weight safetensors and mmap reload | _utils/export_utils.py at 189612be; native GRAPH/EAGER fixtures |
| 11.8 | Whether Core AI exploits unstructured sparsity for size or speed | .aimodel size + on-device latency, 0% vs 70% sparse |
| 12.2 | Whether Apple silicon exploits n:m or block sparsity |
same measurement, with n_m_ratio=(2,4) |
| 14.2 | Resolved: greedy search is documentation pseudocode, not a public helper | docs/src/utils/mixed_precision.md and package tree at 189612be |
| 15.5 | The Debugger's full list of similarity metrics | the metric picker in Xcode 27 |
| 17.3 | Whether every utility export script exposes --dtype |
--help on each script in apple/coreai-models |
Part 9 · Reference 02. Previous: 01 — coreai-opt quantization. Part 10
covers the Core AI Debugger, hardware-specific authoring, and the rank/layout rules this guide's §5
depends on.
-
The classifier and preferences are implemented in the optional
apple/coreai-modelspackage’s pinnedModelStructure.swift. Core AI’s.defaultbehavior is documented separately in Managing model specialization and caching. ↩ -
The pinned k-means palettizer docstring limits this behavior to its Core AI backend:
KMeansPalettizer.finalize. ↩↩