CLAUDE.md

# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## What this is

A research/ablation codebase for **compressing FLUX.2 [klein] distilled 4B** (a 4-step, CFG-free
rectified-flow MM-DiT text-to-image model) into a **smaller, faster model**. Dev rig (since
2026-06-10) is 1× **RTX PRO 4500 Blackwell 32 GB** (sm_120, SDPA; previously an A100-80GB —
see the env section, several old assumptions changed); designed to lift onto a B200.

**ACTIVE TRACK — W4A8 SVDQuant (quantization).** Our own fake-quant implementation of SVDQuant
(4-bit weights + 8-bit activations, with a high-precision low-rank branch that absorbs outliers):
calibrate activation stats → per-Linear `smooth → (whitened) SVD low-rank + iterative refine → 4-bit
residual` → eval against the frozen teacher. Code in `flux2distill/svdquant.py`; pipeline =
`scripts/11–13` (drive the grid one cell at a time with `scripts/run_cell.sh`).
**Best result (2026-06-10, new box/new eval axis): `r64 plain+refine SMOOTH=0 = 0.0297 eval-loss
@ 3.43×`** (smoothed twin re-evals at 0.0325; the surgery frontier sat at 0.23–0.32). **Key findings:**
(1) the upgrades are NOT the monotonic win the old "L-shape" ablation implied — **refine** is the
reliable lever; **whitening alone is non-monotonic in rank** (overfits the 300-img Gram; hurts r16 &
r64, helps r32). (2) ⚠️ **SmoothQuant is HARMFUL at W4A8 — CONFIRMED at all ranks (2026-06-10):**
`SMOOTH=0` beats α=0.5 by −8.6% (r64/r32) to −14.1% (r16); 8-bit activations don't need migration,
and smoothing just widens the 4-bit-weight distribution (worst-layer wrecon 0.15→0.26). **SMOOTH=0 is
the default now** (α sweep: U-shaped, off best, α=0.1 worst). (3) **W4A4 works ONLY with per-group
activations** (`AGROUP=64`, the paper's actual spec): per-token A4 is catastrophic (0.51); g64 acts
→ **0.0610 @ r128 no-smooth** (~2× the A8 champion). Same clean recipe at every bit-width:
plain+refine, no smooth, per-group W+A. NB the 2026-06-01 grid numbers (best
0.0446) are the OLD eval axis — not comparable to new-box numbers (see RESULTS.md 2026-06-10).
Full grid + math in `report/QUANT_REPORT.md` (+ `report/QUANT_REPORT.pdf` with all per-cell montages).

**★ NOW (2026-06-13) — NVFP4 + a DEPLOYABLE fused model on real Blackwell kernels.** Added NVFP4
(Blackwell-native FP4: E2M1 + group-16 + FP8 scales) to `svdquant.py` (`w_fmt`/`a_fmt`). **NVFP4 beats
INT4 on BOTH quality and speed.** Quality (same axis): NVFP4 W4A4 g16 r128 = **0.0303** (≈ the INT4
W4A8 champ, but full 4-bit acts; ~2× better than INT4 W4A4); NVFP4-W + FP8-A r128 = **0.0169** = new
overall champion. Speed (real Nunchaku kernels, built from source for sm_120): NVFP4 W4A4 = **2.49–2.75×**
per-layer vs bf16; **INT4 is BROKEN on Blackwell** (1677 ms/step, slower than bf16 — no INT4 tensor
cores). Wrote our own bf16→Nunchaku NVFP4 exporter (`flux2distill/nunchaku_export.py`, no
DeepCompressor; convention validated to 2.99% in `scripts/26`) and shipped a **deployable fused
klein-4B**: `outputs/nvfp4/deploy/klein4b_nvfp4_fused.safetensors` (2.9 GB), correct images on the real
FP4 kernel, **1.74×@512 / 1.90×@1024 end-to-end, −24% VRAM**. Full write-up + figures:
**`report/NVFP4_REPORT.md`** (`report/figures/GRID_overview.png`). Setup/footguns:
**`docs/CUDA_SETUP_RUNBOOK.md`**. Next-step kernel optimization plan: **`docs/SPEEDUP_IDEAS.md`**.
Scripts 20–28 (profile/bench/convert), `scripts/make_nvfp4_figures.py`. Nunchaku source build +
recipe in `/workspace/build_nunchaku/do_build.sh`.

**★★ NOW (2026-06-14) — NVFP4 HEAD-TO-HEAD (matched, paired) + the BFL-official comparison, closed.**
Ran the deployable model vs the field on **N=512 MJHQ-30k, 512px, 4 steps, guidance 1.0, seed=idx**,
with **image-space fidelity-to-teacher** metrics (this is a *closer-to-teacher* claim, not "prettier").
**Master result: the SVDQuant low-rank branch helps at NVFP4 W4A4** — r0→r128 = **LPIPS −19.7%,
PSNR +1.27 dB, FID-vs-teacher −14.7%** (PickScore/CLIP flat → no semantic loss), and the **real
Nunchaku FP4 kernel reproduces the fake-quant** (LPIPS 0.167 vs 0.173 → the gain is real on the
deployed model). Models compared: **A** teacher bf16 · **C** ours NVFP4 W4A4 r128 (real Nunchaku
kernel; 0.352 s/img · 42.6 ms/step · 12.6 GB) · **D** plain-NVFP4 r0 (our controlled "plain NVFP4"
stand-in, identical recipe minus the low-rank branch) · **E** BFL **official FP8** (W8A8, the real
public baseline; closest to teacher at 0.080/23.0 dB as 8-bit should be). FID-vs-real is ~flat
(~88–90 incl. teacher — it tracks klein-vs-MJHQ style, not quant; reported but not discriminating).
Numbers: `RESULTS.md` (2026-06-14) + `outputs/eval/h2h/metrics.json`; speed/VRAM
`outputs/nvfp4/benchmark_headtohead.json`; full write-up + figures **`report/HEADTOHEAD_klein4b_nvfp4.md`**
(`report/figures/h2h_*.png`, incl. the teacher|ours|r0|BFL-fp8 montage on the text/hand/eggs probes).
Pipeline (one process per model, per-run logs): `scripts/run_h2h.sh` (A/D/fake-q-r128/C/E gen) →
`scripts/34_metrics.py` → `scripts/42_h2h_figures.py`; sensitive probes `scripts/run_probes.sh`
(uses `PROMPTS_JSON=outputs/eval/probes.json`, an env override now on scripts 32/35/41).

**⛔ BFL official NVFP4 (`black-forest-labs/FLUX.2-klein-4b-nvfp4`) is NOT runnable on this box — DO
NOT re-attempt the manual-dequant path.** It is the SAME klein-4b weights (per-tensor scale matches the
teacher to 0.1% on multiple layers), NVIDIA-ModelOpt NVFP4 W4A4, but stored in the **cutlass/TensorRT
tensor-core SWIZZLED layout** (packed-U8 weights AND FP8 group-16 block-scales jointly interleaved).
Plain `modelopt.NVFP4QTensor.dequantize` → **cosine ≈ 0** vs teacher (fused *and* non-fused layers);
modelopt defers to `tensorrt_llm` (`cutlass_fp4_scale_to_modelopt_fp4_scale`) to de-swizzle, not viable
on sm_120 (no FLUX TRT runtime). The diffusers `from_single_file` + `NVIDIAModelOptConfig` path also
fails (FLUX.2 single-file converter isn't modelopt-aware: `torch.chunk`s scalar scale tensors, drops
every `input_scale`). → Reproducing BFL's NVFP4 quality independently needs THEIR runtime; this is
itself a finding (`report/HEADTOHEAD_klein4b_nvfp4.md §5`). The BFL-NVFP4 *attempt* is `scripts/40_load_bfl.py`.
**BFL's FP8 IS plain (not swizzled)** — `weight.float()·weight_scale` → cos 0.9997 vs teacher — so it
loads faithfully via `scripts/41_gen_bfl_fp8.py` (dequant + diffusers key-remap + static-fp8 activation
hooks = true W8A8). For a "plain NVFP4" comparator, use our own rank-0 (`scripts/32_gen_eval.py fq:0`),
the clean one-variable ablation.

**SHELVED TRACK — block surgery.** The prior approach (depth-prune single-stream blocks → replace
with cheaper "surrogate" blocks → freeze the rest and recover via short distillation) topped out at
~1.15–1.26× and was quality-bounded. Kept for record: `block_surgery_plan.md`,
`block_surgery_todo.md`, `report/REPORT.{md,pdf}`, scripts `03/05/07/08/09`, and
`flux2distill/{surgery,surrogate,calibration}.py`. Quantization is a better, orthogonal axis.

**Read these first (real decision history + results):** `RESULTS.md` (both tracks on one axis +
all metrics tables), **`report/QUANT_REPORT_2026-06-10.{md,pdf}` (the LATEST campaign: SMOOTH=0
verdict, α U-sweep, W4A4 per-token→per-group fix, converged recipe, 20-cell montage appendix —
built by `scripts/build_report_2026_06_10.py`)**, `report/QUANT_REPORT.md` / `.pdf` (2026-05-31
A100-era: methodology + math: whitening, Cholesky→eigh, refinement, the 4×3 grid), `plan.md`
(active quant design), `TODO.md` (quant backlog: the 2000-img calib re-sweep is the next experiment),
`README.md` (run order). `block_surgery_*.md` + `report/REPORT.pdf` hold the shelved-track deep-dive.
`init-plan.md` is the original proposal, superseded.

## Remote bucket & where the storage weight is (`hf://buckets/Mercity/FluxDistill`)

`/workspace` is mirrored to the Xet bucket **`Mercity/FluxDistill`** (`hf sync`). The full bucket is
**~310–340 GB**, but **the bytes are NOT spread out — they are almost entirely large `.pt` files
under `outputs/`.** Know this before any `hf sync`/`hf download`, or you will pull hundreds of GB by
accident (a naive `hf sync … ./` downloads the whole thing).

Storage map (approx, 2026-06-14):

| Path | Size | What / why it's heavy |
| --- | --- | --- |
| `outputs/**/quant_state.pt` | **~169 GB** (23 files) | the ~7.3–7.6 GB fake-quant state per grid cell — **dead weight after eval** (loss + montages are already saved); safe to delete |
| `outputs/train_*/**.pt` | **~108 GB** (19 files) | recovery-training checkpoints (`student_best.pt`/`student_final.pt`) |
| `models/klein-4b/` | ~23 GB | teacher weights (public, re-downloadable via `hf download`) |
| `outputs/**/*.safetensors` | ~10 GB | incl. the deployable fused NVFP4 model |
| `.venv/` | ~4 GB | NOTE: no `.venv` is used on the box since 2026-06-10; this is a stale dir in the bucket |
| `.cache/` | ~1 GB | HF/pip cache |
| `build_nunchaku/` | ~0.8 GB | Nunchaku source build tree |
| everything else (code, eval imgs/`.png`/`.jpg`, logs, json, `report/`, `docs/`, `scripts/`, `flux2distill/`) | **~1.3 GB** | the actual source + results you usually want |

**Takeaway:** `mjhq_ref` (eval reference images, ~0.1 GB) is NOT the weight — the **`.pt` checkpoints
are** (~278 GB combined). To pull a usable working copy *without* the giant checkpoints, exclude
`*.pt` (and the stale env/build dirs). The "code + models + eval artifacts, no checkpoints" recipe
(~34 GB) is:

```bash
hf sync hf://buckets/Mercity/FluxDistill ./ \
  --exclude "outputs/*.pt" \
  --exclude "outputs/eval/imgs/mjhq_ref/*" \
  --exclude "build_nunchaku/*" \
  --exclude ".venv/*" --exclude ".cache/*" \
  --exclude "*__pycache__*" --exclude "*.pyc"
```

Drop `--exclude "outputs/*.pt"` to also fetch checkpoints (adds ~278 GB — slow link, many hours).
`hf sync` matching is Python `fnmatch`, where `*` crosses `/`, so `outputs/*.pt` matches
`outputs/<exp>/quant_state.pt` at any depth. **Always `--dry-run` first and aggregate the plan by
size** (not by guessing patterns) to confirm what's actually heavy before committing to a download.

## Environment & commands

No build/lint/test suite — this is a script-driven research repo. Everything runs with
`PYTHONPATH=.` from the repo root.

**⚠️ BOX REBUILT 2026-06-10 (third box). What changed vs the docs/reports written before:**
1. **GPU: A100-80GB → RTX PRO 4500 Blackwell 32 GB** (sm_120), driver CUDA 13.0. Plain/RTN/refine
   builds (~18 GB) fit; eval peaks ~31 GB (fits only thanks to the 13_eval memory fixes below);
   the **whiten path (~28 GB Gram on-GPU) is borderline** — prefer no-whiten cells here, or
   implement the TODO Gram-offload first. Upside: Blackwell IS the real INT4/FP4 kernel target.
2. **NO `.venv` — system python 3.11** (deliberate choice). **torch 2.12.0+cu130** (the old cu126
   advice is obsolete — cu126 has no sm_120 kernels and won't run this GPU), diffusers 0.39.dev
   (git, for `Flux2*`), transformers 5.10.2, torchao 0.17. Pip cache: `/workspace/.cache/pip`.
3. Model weights re-downloaded 2026-06-10 to `models/klein-4b/` (~23 GB, public, no token —
   **`hf download`**, NOT the deprecated `huggingface-cli`). HF cache pinned to `/workspace/.cache`.
4. **The eval axis SHIFTED** (new kernels + transformers bump → different prompt embeddings /
   eval draws): the unchanged grid-best checkpoint re-evals at **0.0325**, not its recorded
   0.0446 (load verified key-complete). **Old-box numbers are NOT comparable to new-box numbers**
   — compare only within a box era; re-anchored baselines live in RESULTS.md (2026-06-10 section).
5. `scripts/13_eval_svdquant.py` got 32-GB memory fixes (CPU-side state load; text encoder
   parked on CPU during the loss step; numerics untouched) and `scripts/run_cell.sh` now takes
   optional `[SMOOTH] [ALPHA]` args 5–6 and exports `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.

**⚠️ THIS IS AN EPHEMERAL POD — ONLY `/workspace` PERSISTS (it's the repo, backed by the HF bucket).**
A restart wipes `models/`, the Python site-packages, AND the agent's `~/.claude/.../memory/` store.
**So agent "memory" does NOT survive here — record every durable fact in THIS file (CLAUDE.md) or the
repo docs (`README.md`/`RESULTS.md`/`TODO.md`/`report/`/`docs/`), which sync to
`hf://buckets/Mercity/FluxDistill`.** Do not rely on `/remember`-style memory on this box.
After a restart, `models/` and the Python stack must be rebuilt. To recover:
(a) `hf download black-forest-labs/FLUX.2-klein-4B --local-dir models/klein-4b` (~23 GB, public);
(b) reinstall `torch==2.12.0 --index-url https://download.pytorch.org/whl/cu130`, then
`transformers==5.10.2 torchao==0.17 accelerate safetensors` and diffusers from git (Flux2). **Gotcha:**
a leftover cu124 `torchvision`/`torchaudio` breaks the `Qwen3ForCausalLM` import (dead
`torchvision::nms` op) → `Flux2KleinPipeline` import fails; `pip uninstall -y torchvision torchaudio`
fixes it. Use `PIP_CACHE_DIR=/workspace/.cache/pip`. Builds are bit-deterministic on a fixed
box/stack: the 4 deleted W4A4 g64 champions rebuilt to exactly 0.0610/0.0620/0.0742/0.0759
(2026-06-13). NOTE: many `quant_state.pt` survive on disk (≈18 W4A8 + the 4 rebuilt W4A4), so
TODO.md's "only 3 SMOOTH=0 states survive" line is stale.

Caches: `data/monet_cache/` (400 latents — used for BOTH training/eval AND the 300-img quant
calibration) survived the rebuild. `data/monet_calib/` (the 2000-img full-res calib) does NOT
exist yet — only needed for the deferred 2k re-sweep (`scripts/11`). All `outputs/abl_c300_*`
grid cells incl. `quant_state.pt` survived. PDF-report deps (markdown+weasyprint + apt
libpango/cairo/poppler) are NOT reinstalled yet on this box.
NOTE: the shelved block-surgery `.pt` model states were deleted to reclaim space; their sample
images / logs / `selection.json` were kept (so block-surgery montages can be regenerated, weights cannot).

**ACTIVE pipeline — W4A8 SVDQuant** (scripts 11–13):

```bash
export PYTHONPATH=.   # system python (no .venv since 2026-06-10)
# Preferred: one grid cell (build + eval) at a time, each with its own logs:
bash scripts/run_cell.sh 64 plain_refine 0 3        # args: RANK variant WHITEN REFINE [SMOOTH] [ALPHA]
bash scripts/run_cell.sh 64 plain_refine_nosmooth 0 3 0   # SMOOTH=0 (s=1, no migration)
# (variant name is just the output-dir tag; WHITEN/REFINE/SMOOTH/ALPHA are the real knobs)

# Or call the build/eval scripts directly (set OUT explicitly to avoid name collisions):
RANK=64 ALPHA=0.5 WGROUP=64 N_CALIB=300 WHITEN=0 REFINE=3 CALIB_DIR=data/monet_cache \
  OUT=outputs/abl_c300_r64_plain_refine python3 -u scripts/12_build_svdquant.py
python3 -u scripts/13_eval_svdquant.py outputs/abl_c300_r64_plain_refine   # vel-loss + 8 montages

python3 scripts/make_quant_report_assets.py   # regen the analysis figures from the grid numbers
python3 scripts/build_report_pdf.py           # assemble report/QUANT_REPORT.pdf (+ per-cell montages)
```
- `12` env vars: `RANK ALPHA WBITS ABITS WGROUP N_CALIB WHITEN REFINE CALIB_DIR MB OUT`. **Always set
  `OUT`** when sweeping: the default dir tag is only `_{plain|whiten}` (ignores refine), so refine
  on/off would collide. The grid uses `outputs/abl_c300_r{R}_{plain|whiten|plain_refine|whiten_refine}`.
  Falls back to `data/monet_cache` if `CALIB_DIR` is absent (warns).
- `13` eval = SAME held-out first-16 metric as the surgery runs → directly comparable in `RESULTS.md`;
  it also renders **8 probe prompts** (storefront text, mountain lake, fisherman, neon street, chalkboard
  text, breakfast counting flat-lay, hand/fingers, spiderweb macro) as teacher|quant montages.
- **Run experiments ONE AT A TIME** with per-run logs + a Monitor (no batched bg loops); see the
  saved memory. Plain build ~4 min, whiten ~7, whiten+refine ~10 (300-calib); eval ~2 min. The build
  writes the 7.3 GB `quant_state.pt` *before* `quant_config.json`, so wait for the `DONE ->` line
  before eval (race). The fake-quant state file is dead weight after eval (loss+montages are saved) —
  safe to delete to reclaim space if quota is tight.

**SHELVED pipeline — block surgery** (scripts 01–10), run in order:

```bash
export PYTHONPATH=.
python3 scripts/01_inspect_model.py                 # introspect transformer (block names, param counts)
python3 scripts/02_teacher_smoke.py                 # sanity: teacher 4-step generation
python3 scripts/09_build_linattn.py 10 4            # BUILD student: drop_k=10, FFN on deepest 4 blocks
python3 scripts/04_gen_eval.py baseline             # teacher-vs-student image pairs across prompts
python3 scripts/06_cache_data.py 400                # cache N monet images @512 -> VAE latents
python3 scripts/10_bench.py                         # measure inference speedup vs teacher

# RECOVERY TRAINING (the core loop). Knobs are env vars + positional args:
STUDENT_DIR=outputs/student_linattn4 SCHED=cosine MB=8 ACCUM=2 \
  python3 -u scripts/08_train_recover.py 300 adamw 1e-4
#                                        ^steps ^opt(adamw|muon) ^base_lr
```

- Always run training with `python3 -u` (unbuffered) so the per-step log file updates live.
- `08_train_recover.py` env vars: `STUDENT_DIR` (which built student to recover), `SCHED`
  (`cosine`|`constant`), `MB`/`ACCUM` (micro-batch / grad-accum). It logs to
  `outputs/train_<tag>/train.log` and saves `student_best.pt`/`student_final.pt` + per-step samples.
- To watch a long run, tail the `eval_vel_loss=` lines in its `train.log`.
- `make_report_assets.py` regenerates the report graphs/montages from the logs.

There are two superseded scripts kept as record: `03_build_student.py` (v1 per-token surgery, produced
a collapsed model) and `07_train.py` (the flawed run that trained **all** weights and diverged). Use
`05_build_student_v2.py` / `09_build_linattn.py` for builds and `08_train_recover.py` for training.

## Architecture — ACTIVE quant track (`flux2distill/svdquant.py`)

The quant pipeline is **calibrate → decompose → quantize → eval**, all PTQ (no training):

**1. `SVDQuantLinear`** replaces an `nn.Linear`. Forward = `(x̂·L₁·L₂)[bf16 low-rank] +
Q8(x̂)·Q4(R)[4-bit residual] + b`, where `x̂ = x/smooth`. Buffers: `smooth` (in,), `lora_down`
(r,in), `lora_up` (out,r), `w_res` (out,in, fake-quant'd 4-bit). `compressed_bytes()` reports the
size a real low-bit packing would reach (the fake-quant buffers themselves are bf16).

**2. `from_linear` does the decomposition** (smooth → low-rank → 4-bit residual): the `s_j` smoothing
factor, then the low-rank fit, then group-wise 4-bit residual. Two SVD modes + refinement:
- **whitened SVD** (`whiten=True`, default): minimizes OUTPUT error `‖X̂(Ŵ−L)‖` via an eigen
  square-root `M` of the smoothed-act Gram `Ĝ=X̂ᵀX̂` (SVD of `Ŵ·M`, map back by `M⁻¹`). Uses
  `torch.linalg.eigh` + eigenvalue clamp, **NOT Cholesky** (the bf16 Gram is non-PD → Cholesky
  crashes). Per-layer try/except → plain-SVD fallback. `whiten=False` = plain SVD of `Ŵ` (base paper).
- **iterative refinement** (`refine_iters`, default 3, 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.

**3. `collect_act_stats`** registers forward-pre-hooks that accumulate per-channel absmax (for
`smooth`) AND the raw Gram `XᵀX` (for whitening) in ONE calibration pass — the smoothed Gram is
`Ĝ = Gᵣ⊘(s·sᵀ)` since smoothing is diagonal. Gram is ~18 GB fp32 on-GPU for the 100 layers.

**4. `apply_svdquant_from_stats`** quantizes every target Linear (the 100 inside the 25 blocks;
embedders/norms/`proj_out` stay bf16). **`quant_config.json` is the build↔load contract** — records
per-layer `specs` (shape/rank); `apply_svdquant_empty` rebuilds empty modules from it so the saved
state_dict round-trips via `load_state_dict`.

## Critical gotchas — quant track

- **System python since 2026-06-10** (no `.venv`); torch 2.12.0+cu130. **Old/new eval numbers
  don't mix** — the axis shifted with the box (see RESULTS.md 2026-06-10); compare within-box only.
- **32 GB VRAM ceiling.** Eval peaks ~31 GB even after the memory fixes; don't add residents to
  `scripts/13`. Whiten builds (~28 GB) are borderline — prefer plain cells or do the Gram-offload TODO.
- **Cholesky is wrong here.** The smoothed-act Gram is non-PD (bf16 accumulation + dead channels) →
  `_LinAlgError`. Use the symmetric eigen square-root with eigenvalue clamping (same math, robust).
- **Build/eval race.** `scripts/12` saves the 7.8 GB `quant_state.pt` then `quant_config.json`;
  launching `13` too early hits `FileNotFoundError` on the config. Wait for the `DONE ->` log line.
- **rank = quality knob, NOT speed.** Low-rank branch is ~r/in ≈ 0.5–2% of the GEMM; rank costs
  memory (bf16 branch), not compute. The SPEED knob is bit-width (W4A8 vs W4A4). No real speedup on
  this A100 — fake-quant measures quality; throughput needs Ada/Blackwell fused kernels.
- **SmoothQuant is NOT a free win at W4A8 — it hurts (α=0.5).** Migration helps low-bit *activations*;
  at A8 the acts are already easy, so smoothing only widens the 4-bit *weight* spread → worse. RTN
  (`SMOOTH=0`, s=1) beat smoothed rank-0 by 21% and beat smoothed SVD cells. `SMOOTH` env on `scripts/12`
  (default 1; `from_linear(..., smooth=False)` sets s=1). Also note: **α=0 ≠ no-smooth** (α=0 →
  `s=1/max|W|`, the all-on-activations extreme); true no-smooth is `SMOOTH=0`. Tune α / disable smoothing
  before trusting any W4A8 number. The α=0.5 grid is mis-tuned (open: re-run no-smooth/low-α).
- **Whitening is NOT a free win at low calib.** The 2026-06-01 full grid showed whitening-alone is
  non-monotonic in rank (hurts r16 & r64, helps r32) because it overfits the noisy 300-img Gram.
  **Refine** is the dependable upgrade; whitening only helps paired with refine or at moderate rank.
  Don't assume "more upgrades = better" — verify per (rank, calib-size). The open question (TODO ★):
  does a 2000-img calib stabilize whitening? Until then, **at high rank prefer plain+refine** (no Gram).
- **Set `OUT` per run.** The build's default dir tag ignores `REFINE`, so plain vs plain_refine (or
  whiten vs whiten_refine) silently overwrite each other. The grid encodes all knobs in the dir name.
- **`hf download`, not `huggingface-cli`** (deprecated, refuses to run) — and pin torch to **cu130**
  (this box: CUDA-13.0 driver, Blackwell sm_120; the older cu126 advice was for the destroyed A100 box).

## Architecture — SHELVED surgery track (the big picture, requires reading multiple files)

The whole system is a **3-stage loop over the `flux2distill/` package**: build a student (surgery) →
recover it (training) → evaluate. The non-obvious glue:

**1. The teacher is a `Flux2KleinPipeline`; we operate on `pipe.transformer`.** Its denoiser has
`transformer_blocks` (5 **double**-stream blocks) and `single_transformer_blocks` (20 **single**-stream
blocks), d=3072. **All surgery targets the single blocks** (they are ~78% of compute; the 5 double
blocks do fragile cross-modal binding and are left untouched). Single blocks are referenced by index
**S0..S19** (S0 = first/shallowest, S19 = last/deepest — see `plan.md`'s block convention).

**2. Surgery (`surgery.py`)** = pick least-useful single blocks and swap them for surrogates:
`importance_by_ablation` (leave-one-out: skip each block, measure final-latent change — the preferred
selector) → `select_blocks_by_importance` → `attach_surrogates` replaces `single_transformer_blocks[i]`
in place. A surrogate must satisfy the **single-block forward contract** (in the model loop it's
called with `encoder_hidden_states=None`, `split_hidden_states=False`, and must return a **single
tensor** — see how `LinearAttentionSurrogate.forward` mirrors `Flux2SingleTransformerBlock`).

**3. Surrogate types (`surrogate.py`)** — two kinds via `make_surrogate(kind=...)`:
- `lowrank`: per-token `x + B·σ(A·x)`. **Cannot mix tokens** → fine for warm-start but collapses under
  aggressive pruning (this is *the* key finding). Closed-form `lstsq_lowrank_init`.
- `linear_attention`: O(N) linear attention with **RoPE** (head_dim 128, reuses the model's rotary),
  a **depthwise-conv** local branch, a **focused** (learnable-exponent) feature map, and an **optional
  FFN** (`use_ffn`). This is the working surrogate. Output projection zero-inits to identity.

**4. `selection.json` is the contract between build and load.** Every build writes
`outputs/<student>/{student_state.pt, selection.json}`. `selection.json` records `surrogate_idx`,
`kind`, and per-surrogate params (`heads, head_dim, conv_kernel, ffn_hidden`, and **`ffn_idx`** = the
subset of dropped blocks that get an FFN). Loaders (`eval_utils.load_student`, `08_train_recover.py`)
**reconstruct the student structure from `selection.json` via `attach_surrogates`, then
`load_state_dict`** — so the saved state_dict only round-trips if you attach with the *same* config.

**5. Recovery training (`08_train_recover.py`) — the critical rule: FREEZE everything, train ONLY the
surrogates.** It loads the frozen teacher (`pipe.transformer`) and a separate student, sets
`requires_grad=False` on the whole student, then re-enables grad only on the surrogate modules
(casting *those* to fp32 = master weights; bf16 autocast compute). Loss = velocity matching to the
teacher (`losses.py`) + a light real-data flow-matching term. Optimizer split lives in
`model_utils.build_param_groups` (Muon on 2D weights, AdamW on the rest) but for surrogate-only
recovery a single AdamW @1e-4 is the validated default. The metric is a fixed held-out **velocity-
matching loss** logged as `eval_vel_loss=`.

## Critical gotchas / hard-won rules — SHELVED surgery track

- **Never train all weights.** Training the pretrained (kept) blocks at any meaningful LR diverges to
  noise. Freeze the base; only surrogates train. AdamW @1e-4 (adapter regime), **not** Muon's 0.02.
- **fp32-master surrogates + the pipeline → dtype clash.** The student's surrogates are fp32 but the
  pipeline runs other modules in bf16, so any `pipe(...)` generation with that student **must be
  wrapped in `torch.autocast("cuda", dtype=torch.bfloat16)`** (see `sample()` in `08`). Training
  forwards already are.
- **Cosine LR floors at 30% of base** (`MIN_LR_FRAC`), never decays to 0 (a dead tail breaks
  resume/extend; constant LR jitters at the floor).
- **Surrogate quality vs speed/drop-count is a real frontier.** Linear-attn recovers far better than
  per-token; the FFN improves quality but is heavy (kills the speedup from dropping more). More
  dropped blocks = faster + worse. See the frontier table in `TODO.md`/`report/`.
- `jinja2>=3.1` is required for the Qwen3 chat template (older versions crash `encode_prompt`).
- At 512px the token sequence is fixed: 512 text + 1024 image = **1536**; the VAE latent is 32-channel,
  packed to 128-ch / 1024 tokens. `pipe._encode_vae_image` + `pipe._pack_latents` is the encode path.