Flat Mode — Quantize Without a Sensitivity Plan

Who this page is for: you have a BF16 model and want a quantized MLX build, but you don’t have (and don’t want to build) this project’s own real per-tensor sensitivity plan — the stratified-calibration → KL-sensitivity → MILP-allocation pipeline that normally decides which bit width every individual tensor gets. Flat mode skips all of that. You pick one bit width, and every real trunk tensor in the model gets it — except a small, automatically-protected set of tensors that get a real hard floor instead, so quality-critical layers are never silently flattened down to your chosen width.

This is the primary path for anyone outside this project — nobody else has this project’s own private Hessian-sweep data, so --plan (a full custom multi-bit allocation) only works for someone who has run that whole separate pipeline themselves. Flat mode needs nothing but a source model.

What actually happens when you run it

Flat mode is implemented by one shared module, flat_plan.py, called from two otherwise-independent build scripts — YAQA (05_full_model_quantize.py) and GPTQ (08_gptq_apply_plan.py). Both take the exact same --flat-* flags and produce the exact same kind of plan; only the actual correction math that gets applied afterward differs (YAQA’s two-sided Hessian correction vs. GPTQ’s one-sided correction — see README.md for why this project built both).

Given a source model and a bit width, Flat mode:

1. Loads the model’s real module structure (no weight data, just the list of every quantizable nn.Linear/SwitchLinear tensor — the same leaf-type set the build scripts actually quantize).
2. Classifies each tensor into one of three groups:
   • Boundary — every tensor in the first and last --flat-n-protect real transformer layer indices (default: just layer 0 and the very last layer).
   • lm_head — the model’s own output projection, matched by name.
   • Trunk — everything else.
3. Assigns bits: trunk tensors get exactly --flat-bits. Boundary and lm_head tensors get max(--flat-bits, their own floor) — so asking for a higher --flat-bits than the floor never silently downgrades a protected tensor; it only ever raises it.
4. Writes a real plan.json to disk (same schema every other plan file in this project already uses) and points the build script’s own --plan at it — every existing, already-trusted call site in either script then runs completely unchanged. Flat mode never touches the actual correction logic, only which bit width gets requested.

Why boundary gets a real hard floor at all: in a normal MILP-built plan, boundary protection comes from a soft cost weighting the solver can freely trade against a real token budget. Flat mode has no solver and no budget to trade against — it’s one direct assignment — so it can’t reuse “no hard floor” as if it meant “no protection needed.” The default floor (6-bit) is the same one already given to lm_head, justified by real data from this project’s own Hessian sweep: the last boundary layer contains the single highest-danger tensor of all 497 scored in the model.

The flags (identical on both scripts)

--plan and --flat-bits are mutually exclusive in spirit: if you pass both, --flat-bits wins and its synthesized plan silently replaces whatever --plan pointed at.

Layer 0 (8 tensors -- a GatedDeltaNet layer):
  language_model.model.layers.0.linear_attn.in_proj_a
  language_model.model.layers.0.linear_attn.in_proj_b
  language_model.model.layers.0.linear_attn.in_proj_qkv
  language_model.model.layers.0.linear_attn.in_proj_z
  language_model.model.layers.0.linear_attn.out_proj
  language_model.model.layers.0.mlp.down_proj
  language_model.model.layers.0.mlp.gate_proj
  language_model.model.layers.0.mlp.up_proj

Layer 63 (7 tensors -- a standard attention layer):
  language_model.model.layers.63.mlp.down_proj
  language_model.model.layers.63.mlp.gate_proj
  language_model.model.layers.63.mlp.up_proj
  language_model.model.layers.63.self_attn.k_proj
  language_model.model.layers.63.self_attn.o_proj
  language_model.model.layers.63.self_attn.q_proj
  language_model.model.layers.63.self_attn.v_proj

Running it — YAQA

Quick test first (3 example tensors, a couple of minutes, needs your own real BF16 model):

"$PYBIN" 05_full_model_quantize.py --source /Users/yourname/Downloads/my-model \
    --flat-bits 5 --subset default --output /Users/yourname/Downloads/my-model-quicktest

What success looks like: a line like [flat-mode] Synthesized plan: 497 trunk tensors at Q5, 15 boundary tensor(s) floored to >= Q6, 1 lm_head tensor floored to >= Q6 (your own real counts will differ by model), then the model loads, a real stratified calibration batch builds across all 6 content domains, Part 5a PASS: real Hessian collection works end-to-end ... no NaN, and finally a real per-tensor error-reduction number for each tensor (90%+ is expected and is success).

./run_full_yaqa.sh /Users/yourname/Downloads/my-model-quantized \
    --source /Users/yourname/Downloads/my-model --flat-bits 5 --n-calibration 4

(--n-calibration 4 is this project’s own validated choice for real production builds, not a hard requirement — the bare code default is 1.) If interrupted for any reason, run the exact same command again; run_full_yaqa.sh finds the .yaqa_resume folder it already wrote next to your output and continues from there instead of starting over.

Watch a real run live: ./monitor.sh <output>.yaqa_resume/logs/batch_<N>.log, or open the live HTML dashboard it also serves (generate_dashboard.py).

Running it — GPTQ

GPTQ is the other, independently-built full-model path in this project — a simpler, one-sided correction (versus YAQA’s two-sided H_I+H_O correction), still real and still Hessian-weighted. Unlike YAQA, 08_gptq_apply_plan.py has its own built-in resumability and internal calibration batching, so there’s no separate wrapper script to call — you run it directly, for a quick test or a real full run alike:

"$PYBIN" 08_gptq_apply_plan.py --source /Users/yourname/Downloads/my-model \
    --flat-bits 5 --output /Users/yourname/Downloads/my-model-gptq-flat

If interrupted, re-running the same command resumes from <output>.gptq_resume automatically, the same idea as YAQA’s .yaqa_resume.

Before committing to a real multi-hour run, --inspect reports real architecture facts about --source (and --plan, if given) as JSON and exits immediately — no model loading, no GPU work, no calibration:

"$PYBIN" 08_gptq_apply_plan.py --source /Users/yourname/Downloads/my-model --inspect

Which one should I use?

This project’s own real, benchmarked production build uses YAQA for the trunk (497 tensors) and GPTQ scoped narrowly to just lm_head (the one tensor YAQA structurally can’t reach — its 248,320-wide output dimension alone needs 246.7GB of Hessian memory, more than this project’s own correction loop can hold). That combination — not a full-model GPTQ build — is what’s actually shipped and benchmarked (92.14% mean across MMLU/GSM8K/IFEval/BFCL/HumanEval, highest of six real models compared).

GPTQ’s own --flat-bits path on the full trunk is real and available — useful if you want a faster, simpler build and are comfortable trading away YAQA’s extra correction quality — but it’s not what this project’s own numbers above were measured on. See README.md’s “Why two implementations” section for the full real comparison.