google-ai-edge/litert-samples/accuracy-safe-quantization

Accuracy-safe quantization

A quantization is done when three things hold, in this order:

1. it exports and the file shrinks by what the recipe predicts,
2. output parity with the float source holds on a task-level check, not

just a smoke test,

1. the quantized model still passes the deployment check on the target

runtime and device.

Quantization rewrites the graph, so step 3 is a fresh obligation every time: re-run the same CompiledModel verification you used to accept the float conversion (see the gpu-clean-conversion skill), then the on-device numerical check.

All recipes below are ai-edge-quantizer (pip install ai-edge-quantizer; 0.9.0 as of 2026-10, which pins ai-edge-litert 2.2.0 on the host), plain Python, no build step. Worked examples live in this repo under models/bonsai/bonsai_image_4b/converted/ and models/qwen/qwen3_tts/converted/.

Choosing a lane

Start with the lightest recipe that meets the size budget, and move down only on evidence:

Full-integer static quantization (static_wi8_ai8 — quantized activations, calibration data required) is a different lane aimed at NPU/AOT targets and is not covered here.

Recipes

Recipes layer by regex: broad rule first, narrow overrides after — that is how one file mixes lanes (bonsai's DiT puts everything at int8 channelwise, then overrides .*TransformerBlock_.* to int4 blockwise).

fp16 float-casting:

from ai_edge_quantizer import quantizer, recipe_manager
from ai_edge_quantizer.recipe import AlgorithmName, qtyping
rm = recipe_manager.RecipeManager()
rm.add_quantization_config(
    regex=".*", operation_name=qtyping.TFLOperationName.ALL_SUPPORTED,
    op_config=qtyping.OpQuantizationConfig(
        weight_tensor_config=qtyping.TensorQuantizationConfig(
            num_bits=16, dtype=qtyping.TensorDataType.FLOAT),
        compute_precision=qtyping.ComputePrecision.FLOAT),
    algorithm_key=AlgorithmName.FLOAT_CASTING)
quantizer.Quantizer("model_fp32.tflite", rm.get_quantization_recipe()) \
    .quantize().export_model("model_fp16.tflite")

Dynamic-range int (swap bits / granularity / algorithm per the table):

from ai_edge_quantizer.qtyping import QuantGranularity as G
from ai_edge_quantizer.qtyping import TFLOperationName as OP
rm = recipe_manager.RecipeManager()
rm.add_dynamic_config(regex=".*", operation_name=OP.FULLY_CONNECTED,
                      num_bits=4, granularity=G.BLOCKWISE_32,
                      algorithm_key=AlgorithmName.OCTAV)
rm.add_dynamic_config(regex=".*", operation_name=OP.EMBEDDING_LOOKUP,
                      num_bits=8, granularity=G.CHANNELWISE)

Weight-only uses the same signature via rm.add_weight_only_config(...) — models/bonsai/bonsai_image_4b/converted/quantize_weight_only.py wraps it as a reusable CLI.

ai_edge_quantizer.recipe also ships these as presets (dynamic_wi8_afp32(), dynamic_wi4b32_afp32(), weight_only_wi8_afp32(), …). The litert-torch LLM exporter accepts a preset name as its quantization_recipe argument, and a custom recipe can be registered by assigning a callable onto the module — the qwen3_tts talker recipe (models/qwen/qwen3_tts/converted/export_talker.py) registers BOCTAV4 (blockwise-32 OCTAV int4 + int8 embeddings) that way.

Verify after every step

• Size first. fp16 ≈ ½, int8 ≈ ¼, int4 blockwise ≈ ⅐ of fp32 (block

scales add overhead). If the file did not shrink as predicted, the regex did not match — fix that before measuring anything.

• Parity against the float reference. Same inputs through the float and

quantized models; correlation on outputs plus the task-level check (argmax match, token-for-token greedy decode, IoU).

• A smoke gate is a floor, not a parity verdict. An LLM can pass most of

a handful of chat prompts and still score near zero on a real benchmark. Before publishing an int4 decoder, run a task benchmark at real length (e.g. GSM8K-style, n≥100) against the float baseline.

• Long generations, specifically. Granularity problems do not show up

in short outputs.

• On the target device. Host emulation of int kernels is pessimistic —

int8 graphs have scored visibly worse on host CPU than the same graphs on the device GPU delegate. Never reject a recipe on desktop numbers alone; never accept one without device numbers.

When it breaks or degrades

Some models are genuinely 4-bit sensitive — small reasoning-distilled decoders (~1–2 B) often fail int4 quality gates that instruct-tuned peers and larger models pass. When int4 fails on quality, ship int8 as the quality row rather than forcing it; int4 becomes a speed reference.

Watch for

• Embeddings stay int8 even in int4 recipes — both shipped LLM-lane

recipes in this repo do this deliberately.

• Bytes are not speed. int4's latency win depends on the backend's

kernel efficiency: the same model can gain ~1.5× on one device and barely 1.1× on another. Measure on the target; don't project from file size.

• Check whether the container is exact. Ternary weights land in int4

blockwise as exactly {-7, 0, +7} — zero rounding error. When the weight distribution matches the container, parity is free; verify it rather than budgeting for loss that isn't there. For exact-container cases use min-max, not OCTAV — OCTAV's clipping optimization can move a grid that min-max reproduces exactly.

• int2 has recipes before it has a consumer: ai-edge-quantizer 0.9.0

ships dynamic_wi2b32_afp32 and its siblings (2.5 bits/weight at block-32 with an fp16 scale per block), and as of 2026-08 the CPU runtime refused the tensor type at prepare — a hard load failure, not degradation. Re-test the load on the runtime you ship before spending time there.

• Auxiliary tables cast to fp16 need the same discipline. Casting

host-side embedding/projection tables halves them; verify generated outputs are unchanged before shipping (qwen3_tts did, and it held).

• Pin the toolchain. Quantized-graph compatibility moves with the

runtime; a graph exported from a dev checkout can fail GPU kernel initialization on a release runtime. Record ai-edge-quantizer / litert-torch versions in the recipe README next to the numbers.

Output layout

Quantization extends the model recipe from gpu-clean-conversion; it does not get its own tree:

models/<family>/<model>/converted/
  export_*.py              float export (existing)
  quantize_*.py            one script per quantized variant
  verify_*.py              parity checks, reused for every variant
  README.md                recipe, sizes, parity numbers, gate results,
                           device, toolchain versions

Keep each variant separately re-runnable. State which variant is the quality row and which is the speed row when they differ. Weights are not committed.