block_surgery_todo.md

TODO — surrogate strengthening & roadmap

Backlog of knobs/experiments for the FLUX.2 klein 4B→2B distillation. Ordered roughly by ROI for image distillation. Status: [ ] todo · [] in progress · [x] done.

Surrogate capability knobs (linear-attention)

• RoPE positional encoding on q,k (head_dim=128) — DONE. Helped text/structure legibility.
• Depthwise-conv local branch on v — DONE (FLatten-style). In the upgraded run.
• Warm-start = mimic the teacher block — DONE (RoPE-aware seq fit; residual rel-err 1.0→~0.65). NOTE: one-shot fit on teacher inputs → limited end-to-end benefit (sequential mismatch). Upgraded run (RoPE+conv+warmstart, drop6) floored at 0.231 vs simple linattn 0.253 (9% lower). Confound: same run also doubled data (400) + batch (16) + cosine-30% LR — not surrogate-only.
• Sequential/progressive warm-start (fit blk12, re-capture, fit blk13, …) to fix the mismatch.
• Feature map φ: elu+1 → focused with learnable exponent (FLatten) — DONE.
• Make it a real block: added an FFN (ffn_hidden=1024) — DONE. Improves warm-start (rel-err → 0.47–0.59) BUT is heavy: at drop-8 it cut the speedup to 1.20× (~= drop-6). Lesson: FFN buys quality at fixed drop, but kills the speed-via-more-drops strategy.
• Gating (GLA-style) data-dependent forget gates — selectivity. Small cost.
• Width (heads × head_dim) — raw capacity dial. Cost linear in width.
• Stack depth (2 linattn ops) — alternative to FFN for capacity.
• 2D depthwise conv on the image-token grid (vs current 1D over the sequence) — better spatial locality; needs text/image token split + 32×32 reshape.

Surgery / compression frontier

MEASURED (eval velocity-loss floor, 512/4-step, 300 steps): teacher 4B 3.876B 1.00x — per-token drop6 3.158B 1.19x 0.308 linattn drop6 3.177B 1.15x 0.253 (RoPE+conv+warmstart: 0.231) linattn+FFN drop8 2.995B 1.20x 0.269

• Speed needs LIGHT surrogates + more drops, NOT FFN. drop-8+FFN proved adding capability while dropping more nets ~same speed (1.20×) + worse quality. For ~1.45×: drop 10–12 with a light surrogate (linattn+RoPE+conv+warmstart, NO FFN / tiny FFN).
• Re-run importance ablation per drop_k; never drop the un-prunable 5 double blocks (the ~31% FLOP floor) — revisit only if a much bigger speedup is needed.

Training recipe

• Freeze base, train surrogates only; AdamW≈Muon (tie); fp32 master; grad-clip.
• Cosine LR floor = 30% of base (was 0/15%) so the LR visibly decays and settles — DONE.
• Size STEPS to the data (avoid dead cosine tail).
• Trajectory velocity matching on the 4 schedule sigmas (vs current σ~U(0,1)).
• Feature matching on retained blocks (masked KD).

Scale-up (B200 target)

• FlashAttention-4 / FlexAttention, torch.compile, larger batch, 300k data + offline latent shards.

Session outcome 2026-05-31 — full deep-dive report at report/REPORT.pdf

Block IDs: single-stream S0..S19 (S0 = first/shallowest after the 5 double blocks; S19 = last/ deepest before output). "Deepest N" = highest IDs. Lower importance = safer to drop.

Frontier (measured, eval velocity-loss floor, 512/4-step):

  per-token  drop6  (drop S12-17)                     3.158B  1.19x  0.308
  linattn    drop6  (drop S12-17) simple              3.177B  1.15x  0.253
  linattn    drop6  (drop S12-17) +RoPE+conv+warmstart 3.177B 1.15x  0.231  <- best quality
  linattn    drop8  (drop S10-17) +focused+FFN(all 8) 2.995B  1.20x  0.269  <- best perceptual
  linattn    drop10 (drop S8-17)  FFN on S14-17       2.737B  1.26x  0.322  <- most compression
  v1 per-token drop12 (drop S0-11)                    2.441B ~1.45x  COLLAPSED

• Measured compute split: single blocks 78%, double 21% (per-block 1.08x). Double-block FFN shrink ≈ 2% speedup → REJECTED; double blocks = fragile cross-modal core, don't cut.
• Recommended next single-block deliverable: drop-7, FFN on the deepest 3 (incl. last, S17-ish). Projected ~3.08B, ~1.17x wall, loss ~0.25 (the drop-6/drop-8 midpoint).
• The real ~2x lever (separate project): step reduction 4→2 via step-distillation (DMD/consistency).
• Free one-time wins later: cache/batch the Qwen3 text-encode; B200 + FlashAttention + fp8 + 300k data.