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.chunks 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:

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):

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:

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.