plan.md

# FLUX.2 [klein] 4B → W4A8 SVDQuant — Plan

> **Status:** Active track (supersedes the block-surgery track) · **Dev rig:** 1× A100-80GB
> (SDPA, fake-quant simulation) · **Teacher:** klein **distilled** 4-step 4B (Apache 2.0)
> **Deliverable:** a 4-bit-weight / 8-bit-activation klein-4B that holds teacher-at-4-steps
> quality, as a fast + small inference endpoint; Nunchaku-deployable later.

The block-surgery track (depth-prune single blocks → linear-attn surrogates → distill) is
**shelved** — it topped out at ~1.15–1.26× and was quality-bounded (see
`block_surgery_plan.md` / `block_surgery_todo.md` for the full history). Quantization is a
better, orthogonal axis: it attacks memory *and* compute, is post-training (no fragile
freeze-and-distill loop), and is proven on this exact model family (Nunchaku already ships
SVDQuant weights for klein-9B).

---

## 0. Method — SVDQuant (our own fake-quant)

SVDQuant (Li et al., ICLR 2025, arXiv 2411.05007 / Nunchaku / DeepCompressor). Per Linear
`y = x Wᵀ + b` (W: (out,in)):

1. **Smooth** — migrate per-channel activation outliers into the weights:
   `x̂ = x / s`, `Ŵ = W ⊙ s`, with `s_j = max|x_j|^α / max|W_:,j|^(1-α)` (SmoothQuant/AWQ
   form). Exactly invariant: `x̂ Ŵᵀ = x Wᵀ`. (`max|x_j|` from calibration.)
2. **Absorb** — keep the top-**rank** components as a **16-bit low-rank branch** `L` that
   soaks up the dominant (outlier) energy. Residual `R = Ŵ − L` is smooth → easy to low-bit.
   Two ways to pick `L` (both implemented; `WHITEN` env):
   - **plain** = SVD of `Ŵ`, minimizes weight error `‖Ŵ−L‖` (base SVDQuant paper).
   - **whitened** (DEFAULT) = activation-aware SVD minimizing OUTPUT error `‖X̂(Ŵ−L)‖` via the
     act-Gram `Ĝ=X̂ᵀX̂` (ASVD/SVD-LLM idea; **our add, not in the paper**). Robust eigen-√ of `Ĝ`,
     NOT Cholesky (bf16 Gram is non-PD). This is what makes rank a real quality knob.
3. **Refine** (DEFAULT 3 iters; SVDQuant §4.2) — re-fit `L` to `Ŵ−Q(R)` to absorb the 4-bit
   rounding error, keep the best iterate. Free at inference; build-time only.
4. **Quantize** — residual weights to **4-bit** (group-wise along in-dim), activations to
   **8-bit** (per-token dynamic). The low-rank branch keeps full-precision `x̂`.

Forward: `y = (x̂·L)[16-bit] + Q8(x̂)·Q4(R)ᵀ[low-bit] + b`. The low-rank branch is the
high-precision *error-correction* path; rank ≪ 3072 so it is ~free on compute (memory only).
**Full math + the Cholesky→eigh fix + ablation caveats: `report/QUANT_REPORT.md`.**

**RESULTS — full 4×3 grid (2026-06-01, 300-img calib; supersedes the earlier 100-calib L-shape).**
Every method × every rank, one fixed calib (`data/monet_cache`), so all 12 cells are comparable
(full table + per-cell metrics in `RESULTS.md`, report `report/QUANT_REPORT.{md,pdf}`). Best per rank:
r16 whiten+refine 0.0556 / r32 whiten+refine 0.0476 / r64 **plain+refine 0.0446** (@ 3.67/3.59/3.43×).
**Overall best = r64 plain+refine 0.0446 @ 3.43×**, all ~4–5× below the surgery frontier (0.23–0.32).
Headline finding: the upgrades are NOT monotonic — **refine** is the reliable lever; **whitening alone
is non-monotonic in rank** (overfits the 300-img Gram: hurts r16 & r64, helps r32) and only helps paired
with refine or at moderate rank. At high rank, plain+refine (no Gram) wins.

**⚠️ Bigger finding (2026-06-01 mechanism ablation): SmoothQuant (α=0.5) is HARMFUL at W4A8.** The RTN
floor (no smooth, no SVD, `SMOOTH=0`/s=1) = **0.0573** — beats SmoothQuant rank-0 (0.0729, −21%) AND the
smoothed SVD r16/r32 cells. SmoothQuant migrates outliers into the weights, which only helps low-bit
*activations*; at A8 the acts are already fine, so it just widens the 4-bit *weight* spread (worst-layer
wrecon 0.15→0.26). **The whole α=0.5 grid is mis-tuned.** → #1 next experiment: re-run no-smooth/low-α
(beats the 2000-img re-sweep in priority). See `TODO.md` / `RESULTS.md`.

**Why W4A8 (not the paper's W4A4):** 8-bit activations are far more forgiving — the right
conservative lead for a **4-step distilled** model where per-step quant error doesn't
average out over many steps. We lead with W4A8 + rank 32, then sweep (rank, α, A4). W4A4
(the Nunchaku sweet spot) is a later comparison.

**Fake-quant, not real kernels.** On the A100 we *simulate* (quantize→dequantize in bf16)
to map the **quality** frontier via the existing held-out velocity-loss metric. We are not
training with DeepCompressor (effectively unmaintained) or running Nunchaku kernels here.
The decomposition is kernel-agnostic, so a checkpoint built here is the same math a fused
W4A8 kernel would run — making it Nunchaku-ready later (an open task, deferred).

---

## 1. What we quantize

- **All `nn.Linear` inside the 25 transformer blocks** (`transformer_blocks.*` +
  `single_transformer_blocks.*`) — uniform W4A8, **no per-block precision games**
  (mixing precisions is slower/messier and we have no evidence the 5 cross-modal double
  blocks need more; revisit only if a fragility shows up).
- **Kept in bf16:** norms (not Linear), embedders (`x_embedder`, `context_embedder`,
  time/text embedders), final `proj_out`. Tiny share of compute; quantizing the output
  projection tends to cost quality for ~nothing.
- Knobs (`flux2distill/config.py::QuantConfig`): `rank=32`, `alpha=0.5`, `w_bits=4`,
  `a_bits=8`, `w_group=64`.

## 2. Calibration data

- **5–7k full-res `jasperai/monet` images via `url`** (NOT the thumbnail `image` field) →
  VAE latents + captions, cached to `data/monet_calib/` by `scripts/11_cache_calib.py`.
  Download is slow but worth it — more images = better activation statistics.
- Calibration forward = real latents noised across sigmas with their captions (the **same**
  forward the eval metric uses), collecting per-input-channel activation abs-max.

## 3. Metric & continuity

- Primary: **held-out velocity-matching loss vs the teacher** on the fixed first-16 of
  `data/monet_cache` — identical construction to `08_train_recover.py`, so the quant number
  is **directly comparable** to the block-surgery frontier (teacher-vs-student loss; a good
  W4A8 quant should land *well below* the surgery numbers ~0.23–0.32, since quant perturbs
  far less than dropping whole blocks). Also report velocity rel-L2 and teacher-vs-quant
  image montages.

---

## 4. Pipeline (scripts)

```bash
source .venv/bin/activate; export PYTHONPATH=.    # system python has NO torch
python3 scripts/11_cache_calib.py 2000 data/monet_calib   # full-res monet via url (already cached)
RANK=32 ALPHA=0.5 WGROUP=64 N_CALIB=100 WHITEN=1 REFINE=3 \
  python3 -u scripts/12_build_svdquant.py                 # calibrate(+Gram)->whiten/SVD/refine/4-bit
python3 -u scripts/13_eval_svdquant.py outputs/svdquant_r32_a0.5_w4a8_whiten   # vel-loss + montages
```

`flux2distill/svdquant.py` holds the module (`SVDQuantLinear`), the fake-quant primitives,
`collect_act_stats` (absmax+Gram hooks), and the build/load surgery (`apply_svdquant_from_stats` /
`apply_svdquant_empty`). `quant_config.json` is the build↔load contract (records `specs`
= per-layer shape/rank so empty modules rebuild before `load_state_dict`). Build env knobs:
`RANK ALPHA WBITS ABITS WGROUP N_CALIB WHITEN REFINE CALIB_DIR MB OUT`.

## 5. Roadmap (see `TODO.md` for the live checklist)

1. ✅ **W4A8 baseline + rank sweep** {16,32,64}, whitened+refine — DONE (results in §0 / RESULTS.md).
2. ✅ **Whitening + iterative refinement** implemented and measured.
3. 🔄 **ACTIVE (2026-06-01): full method×rank GRID at fixed calib.** ranks {16,32,64} ×
   {plain, plain+whiten, plain+refine, plain+whiten+refine} = 12 builds, all at **300-calib
   from `data/monet_cache`** (closes the L-shape: every cell comparable; runs the missing
   `plain+refine`). Also expanded the visual eval to **8 probe prompts**. See `TODO.md`.
4. **DEFERRED — 2000-image calib re-sweep.** After the 300-calib grid picks the method, re-run
   ranks {16,32,64} at the chosen method on the **full-res 2000-image** calib (`scripts/11` →
   `data/monet_calib`) to confirm the choice holds with richer activation statistics. (Needs the
   image download; do NOT run until the 300-calib grid is read.)
5. **Sweep** α, group size, more refine iters; then **W4A4** for the aggressive corner.
6. **One-time stats cache** (absmax+Gram independent of rank/α/group) → cheap sweeps.
7. **Nunchaku-ready export** + **real low-bit kernels** on Ada/Blackwell → the throughput number.

## 6. Open items

- Real speedup needs fused low-bit kernels (Nunchaku, Ada/Blackwell — NOT our A100 sm_80). On
  this box we only have the quality story; throughput/VRAM land on a Blackwell / B200 run.
- ~~The method ablation is an L-shape, not a grid~~ — **CLOSED (2026-06-01):** the full 4×3 grid ran.
  New open item it surfaced: **whitening is unreliable at 300-calib** (non-monotonic in rank) → the
  deferred 2000-img re-sweep (§5, item 4) tests whether richer calib stabilizes it.
- Whether the 5 double (cross-modal) blocks tolerate W4A8 as well as the singles — only revisit
  with per-block precision if eval/montages show binding failures.
- α for W4A8: 8-bit activations may want *less* smoothing (less weight burden) — sweep.