init-plan.md

> **Status:** Proposal · **Compute target:** 1× B200, ~13–15h, ~$66–75 · **Deliverable:** a ~2B open base model + a downstream distill-down kit.
> 

# 1. Objective

Compress **FLUX.2 [klein] 4B** (rectified-flow MM-DiT) into a **~2B general-purpose text-to-image base model** that retains most of the teacher's capability, runs meaningfully faster, and serves as a strong foundation for downstream **narrow-domain fine-tuning** by third parties (e.g. wallpapers, single styles, product imagery). The model is recovered via knowledge distillation on a single B200 within a strict ~$70 / half-day budget. We deliberately do **not** reduce inference steps; we accept 4→8 steps in exchange for the smaller, faster backbone.

# 2. Background

- **Teacher:** FLUX.2 [klein] **base** 4B (undistilled — full training signal, best for fine-tuning). Architecture: `d=3072`, `24 heads`, **5 double-stream + 20 single-stream** MM-DiT blocks, shared-across-block AdaLN modulation, no bias.
- **Conditioning:** Qwen3-4B text encoder (intermediate-layer stack → 7680-dim conditioning). Frozen; not part of the 4B denoiser.
- **Latent space:** shared FLUX.2 VAE (128-ch latent) — **identical across the whole klein family (4B & 9B)**. This is what makes cross-size velocity distillation valid (see §7).
- **Capacity distribution:** the 20 single-stream blocks hold **~67%** of backbone params; the 5 double-stream blocks (cross-modal binding) hold ~33%. This asymmetry drives the architecture choice below.
- **License:** klein 4B is Apache 2.0 (commercial-friendly); klein 9B is non-commercial.

# 3. Target Architecture (~2B student)

Full-width, depth-pruned, with lightweight block replacement for the removed depth.

| Component | Spec | Rationale |
| --- | --- | --- |
| Width | `d=3072` (unchanged) | Same as teacher → preserves texture/detail capacity; feature matching needs no projector |
| Double-stream | **5 (all kept)** | Preserve cross-modal/instruction binding |
| Single-stream | **8 kept full** (of 20) | Depth-only prune; recovers far better than width prune |
| Removed depth | **12 single blocks → light BRLL surrogates** | Linear/low-rank residual-branch replacements, shortcut preserved; least-squares (training-free) init = warm start |
| Effective size | core ~2.04B + surrogates ~0.41B ≈ **~2.45B** |  |
| Inference | **~1.7× faster than 4B**, ~5 GB BF16 | Still runs on consumer GPUs |

# 4. Methodology

**Initialization**

- Prune from klein-**base**-4B; inherit surviving block weights.
- Replace the 12 removed single-stream blocks with BRLL-style surrogates, **least-squares initialized** from teacher activations (training-free warm start → less to recover).

**Distillation (online teacher)**

- **Velocity/output matching** against the 4B teacher (flow-matching target) — primary signal, runs on unlimited teacher-generated targets.
- **Feature matching** on retained blocks — free here because student/teacher widths match (no projector). Use **masked KD** (exclude DiT extreme-activation outliers).
- **Real-data flow-matching loss** on curated images for grounding.
- **No adversarial (LADD) phase in v1** — too costly/finicky for the budget; revisit for a v2 texture polish.

**Data pipeline**

- **300k curated images + captions**, diversity-first (broad styles/subjects/composition + in-image text coverage).
- **Precompute VAE latents offline** (~79 GB @512px BF16, streaming shards). VAE never runs in the training loop.
- **Recompute Qwen3-4B text embeddings online** (caching them would be ~2.4 TB — not worth it).
- **Multi-timestep sampling**: fresh random timestep per image per epoch (epochs = new velocity-field supervision, not repetition).
- **4 epochs × 300k = 1.2M views.** Views kept at full count (quality prioritized over the ~$17 saving).

**Optimizer & precision**

- **Muon** on 2D hidden weights (attention/MLP) + **AdamW** on embedders, `proj_out`, norms, modulation. Tune LR + weight decay only.
- **BF16** compute, FP32 master weights + FP32 moments. No quantization in training.

**Resolution**

- Bulk training at **512px**; optional short 768/1024 polish phase if memory/time headroom remains.

# 5. Compute & Infrastructure (1× B200, 192 GB)

**Memory budget (fixed ~45 GB)**

| Item | GB |
| --- | --- |
| FP32 master weights | ~9.8 |
| BF16 weights (compute) | ~4.9 |
| BF16 grads | ~4.9 |
| Muon momentum (1 moment, FP32) | ~9.0 |
| AdamW m,v (FP32, ~8% of params) | ~1.6 |
| Teacher 4B (BF16, fwd only) | ~6.8 |
| Qwen3-4B encoder (BF16, fwd only) | ~8.0 |
| **Total fixed** | **~45** |

→ ~137 GB free for activations. Muon saves ~8 GB vs full AdamW.

**Batch sizing**

- Activations ≈ **1.63 GB/sample** @512px (FlashAttention, no checkpointing).
- **Max micro-batch ≈ 84**; run **micro-batch 64**, gradient-accumulate to **effective ~256** (Muon's large-batch sweet spot). No gradient checkpointing needed.
- ~18,750 micro-steps (≈4,700 optimizer steps at accum-4).

**Throughput optimizations**

- `torch.compile` (CUDA-graphs mode) — required to hit target MFU.
- **FlashAttention-4** via PyTorch FlexAttention (BF16, Blackwell-native).
- Precomputed latents + aggressive prefetch (keeps GPU fed, protects MFU).

**Time & cost** — 1.2M views × ~40 TFLOP/view:

| MFU | Wall-clock | Cost (@~$5/h) |
| --- | --- | --- |
| 35% | ~17 h | ~$85 |
| 40% | ~15 h | ~$75 |
| 45% | ~13 h | ~$66 |

Plus a one-time ~30–60 min offline VAE-latent precompute (can run on cheaper hardware).

# 6. Expected Outcomes & Quality Profile

A full-width 2B with a small, defensible gap to the 4B.

- **Robust (holds up well):** global structure, color, lighting, coherence; simple/moderate prompt adherence (encoder + double-stream preserved); texture/fine detail (full width retained).
- **Degrades (worst first):** dense in-image text/typography; complex multi-object composition, counting, precise spatial/attribute binding; small faces/hands and micro-detail.
- **Forgiving for the use case:** downstream narrow fine-tuning re-concentrates the 2B's capacity on one domain, so lost broad generality matters far less for end users.

# 7. Downstream "Distill-Down" Kit

Ship the 2B base **plus** a kit so users can specialize it via a bigger teacher rather than fine-tuning the 2B directly — best for **small downstream datasets**.

- **Mechanism:** fine-tune the larger teacher (4B or 9B) on the narrow domain → use it as a data amplifier (generate synthetic in-domain samples) → distill into the 2B via **velocity matching**.
- **Cross-size validity:** the entire klein family shares one VAE/latent space, so 9B→2B and 4B→2B velocity matching is well-defined. Each model uses its own text encoder; **the student stays on Qwen3-4B** (do not bolt the 9B's Qwen3-8B encoder onto the 2B — dimension mismatch).
- **Asset to ship:** because the 2B descends from the 4B, document the **block-correspondence map** (which student block ↔ which teacher block) + projector recipe for optional feature matching — the hard part users can't easily derive themselves.
- **License note:** 4B teacher path = commercial-OK (Apache 2.0); 9B teacher path = non-commercial only.

# 8. Risks & Open Questions

- **Muon on warm-start recovery is under-validated** (strong evidence is LLM/from-scratch + one flow-matching benchmark). → A/B against AdamW on the first training chunk before committing.
- **1.2M views is a partial recovery in absolute terms.** → Monitor the loss curve; be willing to extend if a specific capability (e.g. text) stays weak.
- **Surrogate weight is a tradeoff:** light = cheap warm start, near-2B FLOPs; heavy = better quality but more cost and less speedup. Default light.
- **Coverage holes:** quality gaps will track wherever the 300k set is thin. Curate for breadth, not volume.
- **MFU sensitivity:** a sloppy loop at ~25% MFU ≈ 1.6× longer/costlier; the §5 optimizations are how we hit 40–45%.

# 9. Implementation & Engineering Setup

What we actually build, in dependency order.

**A. Data & latent caching (offline, one-time)**

- Curate + caption 300k images; assign stable IDs.
- Batch **VAE-encode all images → BF16 latents written to streaming shards** (WebDataset / tensordict). The VAE runs only here, never in the training loop.
- Dataloader streams `(latent, caption)` with aggressive prefetch; per step samples timestep `t` and noise, forms `x_t = (1−t)·noise + t·latent`.
- Text embeddings: run **Qwen3-4B online**, intercepting intermediate layers (→ 7680-dim); encoder frozen. (Optionally batch-precompute per shard if I/O-bound.)

**B. Model surgery (prune + BRLL surrogates)**

- Load klein-**base**-4B transformer (diffusers Flux2 / BFL reference impl).
- Depth-select: keep 5 double + 8 of 20 single (importance / CKA similarity or learnable mask); inherit weights for kept blocks.
- Replace the 12 removed single blocks with **BRLL surrogates** — linear/low-rank residual branch with the shortcut preserved; **least-squares-init from teacher activations** on a calibration batch (training-free warm start).
- Sanity: shape + forward-parity checks. Full width `d=3072` → no projector needed for feature matching.

**C. Distillation training loop (online teacher)**

- Teacher 4B loaded **frozen** (eval, `no_grad`, BF16, compiled).
- Per step: `latents + prompts → text emb → sample t, noise → x_t`; **teacher forward → target velocity** `v_teacher`; student forward → `v_student`.
- Loss = `w1·velocity-MSE` (primary) + `w2·feature-match` (masked KD on retained blocks) + `w3·real-data flow-matching`. Tune weights.
- Param-group split: **Muon** on 2D hidden weights, **AdamW** on embedders / `proj_out` / norms / modulation.
- BF16 autocast, FP32 master weights + moments; grad-accumulate to effective ~256; `torch.compile` the student; **FlashAttention-4** via FlexAttention.

**D. Infra, logging & eval**

- Single B200, micro-batch 64, resumable checkpointing.
- Log per-component losses, throughput / MFU, and periodic sample grids for eyeballing.
- Eval harness: FID/CLIP + targeted prompt probes (text rendering, composition/counting, texture) vs the 4B; downstream fine-tune smoke test.

**E. Distill-down kit (reuses C)**

- Same loop with the teacher **swappable (4B or 9B)**; student stays on the Qwen3-4B encoder.
- Flow: domain fine-tune the teacher → generate synthetic in-domain samples → velocity-match into the 2B.
- Package: block-correspondence map + projector recipe + scripts + minimal config.

# 10. Milestones

1. **Data + latents** — curate 300k diverse pairs; precompute VAE latents to shards.
2. **Architecture** — prune to 5D+8S full-width; least-squares-init BRLL surrogates.
3. **Recovery run** — Muon+AdamW, online teacher, velocity + feature + FM losses, 1.2M views @512px, BF16.
4. **Eval** — FID/CLIP + targeted probes (text, composition, texture) vs 4B; downstream fine-tune smoke test.
5. **(Optional) Polish** — short 768/1024 phase; optional LADD for v2.
6. **Release** — 2B base weights + distill-down kit (correspondence map, scripts) + model card.

# 11. Budget Summary

| Line | Estimate |
| --- | --- |
| Recovery training (1× B200, 1.2M views) | **~$66–75** |
| VAE-latent precompute (one-time) | <$5 |
| **Total per run** | **~$70** |

Experiments to be tracked individually under this page per the standard experiment format.