YAQA_MLX_Vegalaboratories_LTD

A real, staged port of YAQA (Yet Another Quantization Algorithm — model-preserving adaptive rounding, Cornell-RelaxML) from its original PyTorch/CUDA implementation to Apple’s MLX framework.

Author: Hakim Ghelab, VegaLaboratories LTD

What this is

YAQA is a whole-model-aware quantization rounding method: instead of correcting each weight using only local, per-layer information (as GPTQ does), it collects a real Hessian sensitivity sketch — both an input-side map (H_I) and an output-side map (H_O) — from the entire model’s real backward pass, so the correction accounts for how each weight actually affects the model’s real output. GPTQ only ever has the equivalent of H_I; YAQA’s LDLQ_2hess correction algorithm is two-sided by construction and needs both. No MLX port of YAQA existed before this work (checked directly via GitHub repo/code search, 2026-09-03 — see PORT_LEDGER.md).

Status

Staged 5-step plan — all 5 steps complete, independently verified end to end on real layers of a real 27B model (Qwen3.5-27B hybrid architecture — standard attention + GatedDeltaNet linear-attention layers), and shipped as a real, loadable production build (363/363 tensors corrected — 362 YAQA + 1 scoped GPTQ for lm_head — 6.015 real bits/weight), including a rebuild with the MTP sidecar’s own real YAQA correction applied. Since then: a new MILP-aware Hessian mechanism that actively protects the tensors the Hessian signal flags as most at risk, and this project’s first real end-to-end benchmark numbers (MMLU / GSM8K / IFEval / BFCL / HumanEval, temperature 0) across 6 real built models. See the timeline above for the full, dated story — every stop traces back to a real log, commit, or doc in this repo.

• Step 1 — the gradient-capture mechanism MLX needs (mx.custom_function) works as required.
• Step 2 — Sketch B’s real Hessian formula, hand-derived from the real source and cross-verified two independent ways. (Sketch B specifically, not Sketch A — Sketch A’s cheap path structurally cannot supply H_O at all; see PORT_LEDGER.md Step 2.)
• Step 3 — the real LDLQ_2hess block-sequential rounding algorithm, cross-verified against an independent NumPy reference.
• Step 4 — the complete real pipeline (Hessian collection, Hadamard incoherence-processing rotation, block-Cholesky factors, the real quantizer) run end to end on a real layer of a real 27B model: YAQA’s correction reduces the real, Hessian-weighted error metric by 75% versus naive rounding.
• Step 5 — generalized to many real tensors at once. Real finding: 52% of this model’s real tensors (every MLP projection, lm_head) have no valid Hadamard factor — a genuine limitation of the real upstream algorithm, not a porting bug. Second real finding (2026-09-04): the reference repo’s actual deployment format is a trellis-VQ codebook with runtime SU/SV/Hadamard rotation via custom CUDA kernels, not affine quantization. Third real finding (same day): the rotation-aware forward pass IS derivable and buildable on real MLX primitives (mx.hadamard_transform, mx.quantized_matmul) — derived by hand, verified numerically, built as yaqa_core.YaqaQuantizedLinear. Production decision: --incoherence none (unrotated, real mlx_lm-loadable today) is the default; --incoherence hadamard stays as a real, verified research branch pending a loader that doesn’t exist yet. Fourth real finding: the original “wrap all 363 tensors, one pass” architecture needs 557GB simultaneously — does not fit in 128GB. Fixed via memory-bounded batching (--hessian-batch-budget-gb) plus running each batch as a separate process (run_full_yaqa.sh) after a real Metal resource-handle limit was found to survive mx.clear_cache(). The real full run has since completed (2026-09-07): 363/363 tensors corrected, real production build saved and loadable, three post-build regressions found by direct comparison against the source model and fixed. The MTP sidecar’s own real YAQA correction (2026-09-08) followed the same path — all 7 MTP tensors corrected, up to 9,675x better than naive on the worst tensor — and the production build was relaunched with it applied. See PORT_LEDGER.md for the full story and 00_DOCS/HTML/yaqa_pipeline_comparison.html for a visual comparison.

See PORT_LEDGER.md for the full, honest engineering log — every real bug found and fixed, every real command and its real output, and a plain-English explanation of what any of this means, written for a non-technical reader.

Requirements

• Python 3.11 (this project uses 3.11.15) with a real MLX install. Verified working versions (2026-09-03):
  • mlx 0.32.2
  • mlx-lm 0.31.3
  • numpy 2.4.6
• Apple Silicon Mac (MLX is Metal-based; there is no CPU-only or CUDA path).
• For the real-model scripts (04_real_layer_test.py, 05_full_model_quantize.py): a local copy of the real BF16 source model this project targets, plus enough RAM to hold the model and, for Step 5, every target tensor’s Hessian simultaneously during collection (this project’s own machine: 128GB RAM, confirmed comfortable for the largest real Hessian in the current plan, ~1.2GB for a 17408×17408 H_I).
• The synthetic-only scripts (01_gradient_collection_test.py through 03_rounding_update_test.py, test_regression.py) need no real model and run in seconds to a couple of minutes.

Install

This project uses a uv-managed Python tool environment. If you already have mlx/mlx-lm installed some other way, any Python 3.11 environment with the versions above will work — just substitute your own interpreter path below.

# This project's own real interpreter (adjust if you're not on the same machine):
PYBIN="/Users/hghelab/.local/share/uv/tools/mlx-optiq/bin/python3"

# Or set up your own, from the real pinned versions in requirements.txt:
uv venv --python 3.11
uv pip install -r requirements.txt

Running

cd scripts/yaqa_port

# Fast, synthetic-only regression suite (no real model needed, ~1-2 min):
"$PYBIN" test_regression.py

# Steps 1-3 individually (all synthetic, no real model needed):
"$PYBIN" 01_gradient_collection_test.py
"$PYBIN" 02_sketch_b_synthetic_test.py
"$PYBIN" 03_rounding_update_test.py

# Step 4: one real layer, end to end (needs the real source model on disk):
"$PYBIN" 04_real_layer_test.py
"$PYBIN" 04_real_layer_test.py --from-cache          # skip the model load/backward pass on repeat runs
"$PYBIN" 04_real_layer_test.py --sigma-reg 0.5        # override the damping constant

# Step 5: many real tensors at once (needs the real source model + a real per-tensor plan on disk):
"$PYBIN" 05_full_model_quantize.py --subset default   # 3 representative real tensors, comparison only, no save
"$PYBIN" 05_full_model_quantize.py --subset "language_model.model.layers.63.mlp.up_proj"  # one tensor by name

Step 5c: real output model

--incoherence controls whether Hadamard rotation is carried into the DEPLOYED tensor (not the correction step, which uses it either way where a valid factor exists):

• none (production default, 2026-09-04): deploys unrotated for every tensor. Reuses 08_gptq_apply_plan.py’s real save/sidecar/MTP mechanism completely unchanged, real mlx_lm-loadable today, no runtime overhead. Keeps YAQA’s actual advantage over GPTQ (two-sided H_I+H_O correction) on every tensor.
• hadamard (research branch): deploys rotation-aware via yaqa_core.YaqaQuantizedLinear (real, derived + verified end to end — test_regression.py test [8]) for the 173/363 tensors that support it. Requires a real loader that does not exist yet — the saved model is NOT loadable with plain mlx_lm.load() until that’s built.

Real memory constraint, found 2026-09-04: all 363 real plan targets’ Hessians held simultaneously need 557GB (310GB even excluding lm_head, which is excluded from YAQA’s two-sided correction separately — its 248,320-wide output dimension alone needs 246.7GB, and even the real reference algorithm doesn’t run its per-layer correction loop on this tensor). lm_head is NOT left uncorrected, though: 06_lm_head_gptq.py gives it a real, scoped one-sided GPTQ correction instead (its input side is small and safe; only its output side was ever the problem) — see the Files table below and PORT_LEDGER.md for the full reasoning. The original “wrap everything, one pass” design for the other 362 tensors does not fit in 128GB either. Fixed by splitting targets into memory-bounded batches (--hessian-batch-budget-gb, default 8GB). A SECOND real constraint found immediately after: running multiple batches inside one long-lived process hits a hard Metal resource-handle limit ([metal::malloc] Resource limit (499000) exceeded) that mx.clear_cache() does NOT reset (tested directly, twice). Fixed by running each batch as its own separate process — use run_full_yaqa.sh, not 05_full_model_quantize.py directly, for any run with more than a couple of tensors:

# The real full run (production, mlx_lm-loadable):
./run_full_yaqa.sh /path/to/output --n-calibration 4 --incoherence none

# Watch it live in another terminal (auto-refreshing, color-coded):
./monitor.sh <output_dir>.yaqa_resume/logs/batch_<N>.log

run_full_yaqa.sh queries the real batch count, runs each batch as a fresh process (--batch-only N --resume-dir <output>.yaqa_resume), then does one final --assemble-from-resume pass that loads every batch’s cached result and performs the real save. STATUS (2026-09-04): the batching + separate-process fix is being validated now on the 3-tensor subset — do not treat it as confirmed working until PORT_LEDGER.md says so.

Manual/advanced flags on 05_full_model_quantize.py (normally only used by run_full_yaqa.sh itself, not directly): --resume-dir <dir>, --batch-only N, --print-num-batches, --assemble-from-resume.

04_real_layer_test.py and 05_full_model_quantize.py both have real, hardcoded paths near the top of the file (SOURCE_MODEL, PLAN_PATH, TARGET_LAYER) pointing at this project’s own real model and plan — edit those constants to point at your own real model/plan if running this elsewhere.

Files

© 2026 Hakim Ghelab, VegaLaboratories LTD. All rights reserved.