TODO.md

TODO — FLUX.2 klein-4B quantization track

Backlog. Status: [ ] todo · [~] in progress · [x] done. (Block-surgery backlog archived in block_surgery_todo.md.)

★ 2026-06-13 — NVFP4 + deployable fused model — DONE

• NVFP4 in svdquant.py (E2M1 + group-16 + FP8 scales; w_fmt/a_fmt) + unit test (scripts/test_nvfp4.py). Beats INT4 ~2× at matched bits.
• NVFP4 sweep (was the "NVFP4 sim" TODO below): W4A4 r{32,64,128} = 0.0390/0.0364/0.0303; NVFP4-W+FP8-A r{64,128} = 0.0204/0.0169 (overall champion). scripts/run_nvfp4_cell.sh.
• Real kernel speed — built Nunchaku from source for sm_120 (/workspace/build_nunchaku, do_build.sh); NVFP4 W4A4 = 2.49–2.75× per-layer; INT4 broken on Blackwell (hardware). Profilers scripts/20–25.
• bf16→Nunchaku NVFP4 exporter (flux2distill/nunchaku_export.py, no DeepCompressor; validated 2.99% in scripts/26) → deployable fused klein-4B (scripts/27/28): outputs/nvfp4/deploy/klein4b_nvfp4_fused.safetensors (2.9 GB), correct images, 1.74×@512 / 1.90×@1024, −24% VRAM. Report report/NVFP4_REPORT.md; figures scripts/make_nvfp4_figures.py.
• Docs: docs/CUDA_SETUP_RUNBOOK.md (setup + footguns), docs/SPEEDUP_IDEAS.md (next steps).

★★ 2026-06-14 — NVFP4 head-to-head (N=512) — DONE

• Matched head-to-head A teacher / C ours-NVFP4-r128-real-kernel / D plain-NVFP4-r0 / E BFL-official-FP8, N=512 MJHQ-30k. Low-rank branch helps: LPIPS −19.7%, PSNR +1.27 dB, FID-teacher −14.7%; real kernel ≈ fake-q; no semantic loss. report/HEADTOHEAD_klein4b_nvfp4.md, RESULTS.md (2026-06-14), outputs/eval/h2h/metrics.json, report/figures/h2h_*.png. Pipeline scripts/run_h2h.sh → scripts/34_metrics.py → scripts/42_h2h_figures.py.
• BFL-official comparison (was DEFERRED in report/NVFP4_SVDQUANT_EVAL.md) — RESOLVED: their FP8 loads faithfully (W8A8, scripts/41_gen_bfl_fp8.py) = model E; their NVFP4 is BLOCKED — cutlass/TensorRT swizzled layout, needs their runtime (scripts/40_load_bfl.py + HEADTOHEAD §5). Do NOT re-attempt manual dequant. Use our r0 as the "plain NVFP4" stand-in.
• (open, low priority) Run BFL-official NVFP4 if a tensorrt_llm / TRT FLUX path becomes available on sm_120 — only path to a direct BFL-NVFP4 number.

★★ NEXT — kernel/pipeline speedup (full menu in docs/SPEEDUP_IDEAS.md)

Lift the ~1.9× toward the per-step 2.3× / per-layer 2.5× by attacking the non-transformer overhead and the attention. Priority order:

• CUDA graphs over the denoise loop (4B is launch-bound at small sizes) — cheapest big win.
• Quantize text-encoder (Qwen3) + VAE decoder — the fixed ~30% overhead cap.
• Fix the fused batch=1 limitation (packed-rotary broadcast over batch) — serving throughput.
• Output-quant fusion (qout: fold layer N+1's act-quant into N's epilogue).
• FP8 attention (QKᵀ/AV in E4M3) — the high-res ceiling; verify the GEMM uses Blackwell tcgen05/TMEM, not a Hopper wgmma path.
• GEMM tile/warp autotune per shape; TMA loads; rank Pareto + per-layer rank; real wcscales; FP6.
• Measure everything across {quant level × batch × res} vs bf16 (extend outputs/nvfp4/benchmark.json).

★ Deployment coverage — INT4 export for OLDER GPUs (Ada/Ampere/Turing, RTX 20/30/40)

• Extend flux2distill/nunchaku_export.py to precision='int4' (int4 codes, fp16/bf16 per-group-64 scales via pack_scale not pack_micro_scale, no wtscale/wcscales) and produce real deployable checkpoints klein4b_int4_{w4a8,w4a4}.safetensors — the artifacts a 4090/30-series user actually loads (Blackwell has NO INT4 tensor cores; INT4 is the older-GPU format). Quality is verifiable HERE (the INT4 kernel runs on sm_120, just emulated/slow); REAL speed needs an Ada GPU. NB: the saved fake-quant quant_state.pt champions are bf16-stored quality refs, NOT runnable-fast on any GPU — they're the source recipe; the INT4 export is what makes them deployable on older cards.

Build-out (lead) — DONE

• flux2distill/svdquant.py — SVDQuantLinear (smooth + SVD low-rank + 4-bit residual), fake-quant primitives, act-stat+Gram hooks, build/load surgery.
• Whitened (activation-aware) SVD — min OUTPUT error ‖X̂(Ŵ−L)‖ via eigen-√ of the smoothed-act Gram (robust to non-PD; per-layer plain-SVD fallback). Makes rank a real knob.
• Iterative low-rank refinement (SVDQuant §4.2) — re-fit L to absorb 4-bit quant error, keep best iterate. REFINE env (default 3). −13% at r16 vs plain.
• scripts/11–13 (calib download / build / eval); 2000-img calib cached at data/monet_calib (100-img subset is plenty, ~3min/build).
• Initial rank sweep (100-calib, whitened+refine) — 0.0574/0.0494/0.0454. SUPERSEDED by the full 300-calib grid below (which revises the per-knob story — see RESULTS.md). Kept as history.

ACTIVE campaign (2026-06-01) — full method×rank GRID, fixed calib — ✅ DONE

Closed the "L-shape, not a grid" gap: ran the complete 2×2 method grid at EVERY rank, all on one fixed 300-img calib (data/monet_cache latents) so every cell is comparable.

• Grid: ranks {16,32,64} × {plain, plain+whiten, plain+refine, plain+whiten+refine} = 12 builds+evals, fixed 300-calib, α=0.5, group=64, W4A8. Dirs outputs/abl_c300_r{R}_{variant}, 8-probe montages each. Run one-at-a-time w/ per-run logs + Monitor. Full grid in RESULTS.md.
• plain+refine (the missing cell) — now run at all 3 ranks. Refine's solo effect: helps everywhere EXCEPT r16 (overfits weight error). r64 plain+refine = 0.0446 is the grid best.
• Collate into RESULTS.md grid (done) + report assets/montages (in progress).

Grid takeaways (overturn the old "each upgrade compounds" claim):

• Refine = reliable workhorse (best variant at every rank uses it; only r16-plain regresses).
• Whitening ALONE is unreliable at 300-calib — non-monotonic in rank (hurts r16 & r64, helps r32); it overfits the noisy 300-img Gram. Only earns its keep at r32 or when paired with refine.
• Strong whiten×refine interaction (refine in the whitened metric fixes whitening's overfit).
• At high rank, skip whitening: r64 plain+refine (0.0446, no Gram, simpler) ≥ r64 whiten+refine (0.0451).

Deferred follow-ups (queued — see plan.md §5)

• ★ NOW THE KEY EXPERIMENT — re-sweep at 2000-image calib. The grid shows whitening misbehaving at 300-calib (overfit Gram). Re-run the grid (or at least the whiten / whiten+refine cells) at 2000-img full-res calib (scripts/11 → data/monet_calib, needs image download) to test: does whitening become reliably beneficial with richer activation stats? This decides whether to keep whitening at all. Reuse the cached-stats idea below to make it cheap.
• Eval probes 4 → 8 (richer visual comparison): kept the original 4 (storefront text, mountain lake, fisherman, neon street) + added chalkboard-text, multi-object breakfast flat-lay, hand/face five-fingers, dewy-spiderweb macro. In scripts/13_eval_svdquant.py.

★★ NOW THE #1 EXPERIMENT — SmoothQuant is HURTING W4A8 (confirmed 2026-06-01)

Mechanism ablation (rank-0, 300-calib) found removing SmoothQuant IMPROVES eval-loss: RTN (no smooth, SMOOTH=0, s=1) = 0.0573 vs SmoothQuant rank-0 (α=0.5) = 0.0729 (−21%), and RTN even beats the smoothed SVD r16/r32 cells (0.0620/0.0586). The whole α=0.5 grid is mis-tuned — the SVD branch was partly papering over smoothing damage (worst-layer wrecon 0.15→0.26 with smoothing).

• RTN floor (SMOOTH=0) and SmoothQuant rank-0 baselines built+evaled (RESULTS.md). New SMOOTH env on scripts/12 (from_linear(smooth=False) → s=1). NB α=0 ≠ no-smooth.
• ★ DONE (2026-06-10, new box): SMOOTH=0 at r{16,32,64} plain+refine — wins at every rank (−14.1/−8.6/−8.6%), new best r64 ns = 0.0297 (new axis; vs smoothed twin re-eval 0.0325). SMOOTH=0 is now the default. NB: box rebuilt → eval axis shifted (0.0446→0.0325 for the same checkpoint); compare within-box only. Full table in RESULTS.md.
• ★ DONE (2026-06-10): α sweep {off, 0.1, 0.25, 0.5} × r{32,64} — replicated U-shape, identical ordering both ranks: off < 0.25 < 0.5 ≪ 0.1. No α beats off; α=0.1 (the weight-equalizing extreme) is the WORST (+20% vs off-adjacent wrecon). SMOOTH=0 is the permanent W4A8 default. Full table in RESULTS.md.
• DONE (2026-06-10): W4A4 3-cell ablation — r64 nosmooth 0.5103 (catastrophic, 17× the A8 twin) / r64 α=0.5 0.3885 (smoothing FLIPS to +24% helpful at A4) / r128 α=0.5 0.3060 (rank = activation shield, +21%). Naive per-token A4 not viable. Full table in RESULTS.md.
• DONE (2026-06-10): W4A4 α-up sweep — r64 a75 0.2819; r128 a75 0.2080 (W4A4 best); r128 a100 0.2397 (turns — weight damage outruns act relief). α*≈0.75 for per-token A4; the optimal α tracks the bottleneck. Full table in RESULTS.md.
• DONE (2026-06-10): per-group activation quant (AGROUP=64) — the W4A4 fix: r64-ns 0.0742 (−85% vs per-token), r128-ns 0.0610 = W4A4 best (~2× the A8 champion). Smoothing dead weight again at both ranks → recipe converges to plain+refine, no smooth, per-group W+A. Implemented as a_group through svdquant.py + scripts/12+13 (config-recorded, back-compat: old configs load as per-token). Full 2×2 + qualitative in RESULTS.md.
• NVFP4 sim (explicit, user-requested): FP4 E2M1 element grid + group-16 + FP8-quantized scales in fake_quant_act/fake_quant_weight (~30 lines). This is the Blackwell-native deployable format (this box's GPU runs it in silicon). Pair with the AGROUP=16 ladder cell.
• (corrected 2026-06-10) the jasperai/monet calib data is NOT Monet paintings — it's a diverse photographic set (260/400 captions mention text/signs; 262 mention people/hands). So calib content-narrowness is NOT a confound; whitening's 300-calib instability remains attributable to Gram sample noise (the 2k re-sweep still tests that). A probe-aligned eval batch (text-heavy crops scored directly by the metric) is still a nice-to-have, but the "calib lacks text" premise was wrong.
• Act group-size ladder at A4 — unit test shows ~−25% act-error per halving (g64 0.129 → g32 0.099 → g16 0.077): run AGROUP=16 r128 nosmooth as the sim-only upper bound. NB g64 = Ada-INT4 deployable; g16 is only deployable as NVFP4 on Blackwell → pairs with the NVFP4 sim (E2M1 grid + FP8 scales, ~30 lines in fake_quant_act/fake_quant_weight).
• (deferred) whiten+refine at A4 (act-aware L as act shield; Gram ~18 GB → tight on 32 GB, may need CPU-Gram offload).
• (deferred) the 2000-img calib re-sweep below now ranks BELOW fixing α — do α first.
• Weight group size {32, 64, 128} at rank 32 — finer = better, ~free in fake-quant.
• More refine iters {5, 8} — all layers hit the iter=3 cap, so the knob isn't saturated.
• W4A4 — the Nunchaku sweet spot; map the aggressive corner vs W4A8.

⚠ Model states NOT saved — rebuild before any use

Due to the ~250 GB volume quota, every cell's 7.4 GB fake-quant quant_state.pt was deleted after its eval during the 2026-06-10 campaign (losses/configs/montages all kept; builds verified bit-deterministic on this box). Currently on disk: ONLY the 3 W4A8 SMOOTH=0 states (incl. the W4A8 champion r64 = 0.0297). The W4A4 champion and all other cells' weights are gone.

• Rebuild the W4A4 champion when needed: ABITS=4 AGROUP=64 bash scripts/run_cell.sh 128 w4a4g64_nosmooth 0 3 0 (~15 min; expect eval_vel_loss=0.0610 to reproduce).
• Any other deleted cell rebuilds the same way from its recorded knobs (dir name encodes them; quant_config.json in each dir is the contract). NB determinism is same-box/stack — if the box is rebuilt again, regenerated models follow the same recipe but may not be bit-identical and the eval axis may move (re-anchor baselines first, as done 2026-06-10).
• (optional) hf sync the champion states to the HF bucket for durability (needs HF_TOKEN).

Efficiency / infra

• One-time stats cache — absmax+Gram depend only on (teacher, calib), not rank/α/group; cache once, reuse across the whole sweep (one calib pass instead of per-build). Biggest sweep speedup; makes the ablation grid + α/group sweeps cheap. Also lets the Gram offload to CPU between builds → makes a 24 GB card viable for the whiten path.
• Batch the SVD across same-shape layers — currently 100 sequential per-layer SVDs (the dominant cost; a single small SVD never saturates the GPU). The 100 Linears have only a few distinct shapes (qkv/mlp repeat across blocks) → stack same-shape layers, do batched svd/eigh/matmul (works for whiten + refine too). Est. 3–5× on decompose. Code change + re-validate against locked numbers.
• GPU: drop the A100-80GB. Workload peaks ~18 GB (plain/RTN/refine) / ~28 GB (whiten); the SVD loop is latency-bound, not throughput-bound → A100 compute/VRAM mostly idle. Move to a RTX 5090 (32 GB, Blackwell) — cheaper, fits everything incl. the Gram, AND it's the real INT4/FP4 kernel target (A100 sm_80 INT4 is weak, can't measure real speed). 4090 (24 GB) ok for no-whiten cells; 48 GB (A6000/L40) if you want whiten headroom with zero code changes.
• Per-layer rel-err report — which layers are hardest to quantize (higher-rank candidates).
• Ablate calib size (100 vs 2k) and σ coverage (uniform vs the 4 schedule sigmas) on eval-loss.

Deployment (deferred — needs real kernels, not this A100)

• Real low-bit number on Ada/Blackwell (RTX 40/50) via Nunchaku fused INT4/FP4 — our A100 (sm_80) under-represents it; fake-quant here is quality-only. Expect ~3× latency, ~3–3.5× mem.
• Nunchaku-ready export — align our decomposition layout/packing so a checkpoint loads in Nunchaku (we don't build with DeepCompressor). Verify numerics match.
• B200 scale-up + text-encoder caching + larger batch (remaining fixed-overhead wins).

Knob summary (what each does) — updated from the 2026-06-01 full grid

• rank = quality knob. ~Free vs a bf16 baseline, but NOT free on real 4-bit kernels: the fused Nunchaku branch still runs 16-bit GEMMs whose share of the (now 3–4× faster) layer is ~5–10% latency at r32 (SVDQuant §4.3/E.5, measured), scaling ~linearly → r64 ≈ 10–20%, r128 ≈ 20–40%. Plus MEMORY (bf16 branch). Pick method per rank (r64→plain+refine no-smooth is the W4A8 champion).
• bit-width = the SPEED knob (W4A8 vs W4A4). NOT rank.
• α / SMOOTH = activation↔weight outlier split (smoothing). SMOOTH=0 (off) is the validated default at W4A8 (2026-06-10: beats α=0.5 at every rank); small-α sweep still open.
• group = weight-quant granularity. Untuned (64 default).
• whiten = output-error vs weight-error SVD (eigen-√, robust to non-PD). UNRELIABLE at 300-calib — non-monotonic in rank (overfits the Gram). Only use paired w/ refine or at r32.
• refine = error-feedback iters (3, all layers still hit the cap → not saturated). The RELIABLE upgrade — helps everywhere except r16-plain. Default ON.

Notes

• Continuity: eval uses the SAME held-out first-16 of data/monet_cache as 08_train_recover.py, so quant numbers sit on the same axis as the surgery frontier in RESULTS.md.
• Quantization is orthogonal to the (shelved) block surgery — in principle they stack, but we lead with quant alone since it's the bigger, cleaner lever.