Upload DIMBA_REBUILD_SPEC.md with huggingface_hub

DIMBA_REBUILD_SPEC.md

@@ -0,0 +1,84 @@
+# DIMBA rebuild spec — full fixes + plan
+**3 Jul 2026.** Everything learned from the $358 budget28 run, its post-mortem (docs/ROOT_CAUSE_BUDGET28.md), the repair session (~$31), and the repo audit — as a clean-rebuild checklist. Ordered by importance within each section.
+
+---
+
+## 1. Objective / loss fixes (the proven root cause)
+
+**1.1 min-SNR floor: 0.5, not 1e-3 — as the DEFAULT.**
+`weight = clamp(snr_flow, min=0.5, max=5.0)`. The old floor (1e-3) removed all gradient pressure at high t, so the model converged to a noise-independent blend exactly where every generation starts (t=1.0). In the current repo the fix exists but is opt-in (`snr_floor` param, trainer.py:117) — every call site except `repair_finetune.py` still trains broken. In the rebuild: make the fixed behavior the only behavior. *Validated: this + 1.2 took an overfit probe from unigram soup to locally-coherent text.*
+
+**1.2 CE anchor must fade with noise.**
+`ce_loss = (ce_per_sample * (1 - t)).mean()`. Unweighted CE at t=1.0 has the marginal token distribution as its optimum — it literally *teaches* unigram output at pure noise. Default-on.
+
+**1.3 Timestep sampling: 50/50 uniform + logit-normal.**
+Logit-normal alone undersamples t≈1, so the dead zone gets almost no training coverage even with a fixed loss. Cheap, keep it.
+
+**1.4 SFT loss: response-CE at uniform t + snr-floored latent MSE anchor on the diffuse mask** (commit 1d6ee9e in the old repo). Without the anchor, SFT unpins the flow-ODE geometry at high t that 1.1–1.3 bought.
+
+**1.5 Classifier-free guidance (NEW — likely the missing "commitment" piece).**
+Train with prompt dropout (~10% of batches get a null/empty prompt), sample with guidance scale ~2–4: `v = v_uncond + s·(v_cond − v_uncond)`. The SFT-repaired model's exact failure was *no mode commitment* — prompt-relevant texture but never an answer. CFG is the standard diffusion mechanism for forcing commitment to the conditioning signal, and DIMBA never had it. Cheap to add, high expected value.
+
+**1.6 Self-conditioning (NEW).**
+On 50% of training steps, run one no-grad forward, feed the resulting x0̂ back as an extra input channel for the graded step. Standard in continuous text diffusion (Analog Bits, CDCD) precisely because it dramatically improves sample commitment/quality at small scale. If DIMBA lacks it, add it.
+
+---
+
+## 2. Distillation pipeline fixes
+
+**2.1 Keep KD on through stage 3.** `kd_weight` annealed 1.0→0.3 over the run, never 0. Teacher forward on the same batch, KL from teacher logits into the student's x0-logits. The original run had the teacher connected for 1,000 of ~855k steps — 99.9% of the $358 was from-scratch diffusion pretraining at 135M, which cannot work. Currently `kd_weight=0.0` is still the default in every preset (train_4090.py:174,180; train_h100.py:196,292). *(Untested at scale — today's optional KD run is the cheap test.)*
+
+**2.2 Stage 1/2 alignment: run to convergence, not a fixed 500 steps.**
+≥5–10k steps each, stop on flattening of the alignment loss. Stage 2 ended at 0.556 last time — visibly unconverged — and it was the only bridge across the AR→diffusion gap.
+
+**2.3 BLOCKING coherence gates between every phase.**
+Generation-based, not loss-based: sample N=16 fixed prompts with the production sampler, score the outputs with teacher PPL, hard-assert a threshold or abort the pipeline. The warnings already exist as log lines (train_4090.py:1149, 1399: "SFT cannot fix an incoherent base", "GRPO cannot bootstrap…") — they fired last time and the run proceeded anyway. Losses looked healthy at every stage of a run that produced an unusable model; only generations tell the truth.
+
+**2.4 Plateau early-stop in stage 3.**
+Stop when loss improvement < ε over a sliding window. Budget28's stage-3 loss converged at ~1.6B of 28B tokens; ~93% of stage-3 compute was spent flat. Would have refunded ~$300.
+
+**2.5 Per-t-bucket loss logging.**
+Log the *unweighted* diffusion loss separately for t-buckets (e.g. [0, .3), [.3, .7), [.7, .9), [.9, 1.0]). The aggregate min-SNR-weighted loss is dominated by mid/low-t and mathematically cannot show a dead high-t region — that's why babysitting missed it for 28B tokens.
+
+**2.6 Partial-noise recovery ladder as a standard periodic eval.**
+The diagnostic that found the root cause: noise a held-out row to t_start ∈ {0.3, 0.7, 0.9, 1.0}, integrate the ODE to 0, report token recovery. Run every ~5k steps. A cliff between 0.9 and 1.0 = objective broken; catch it in the first $5, not after $358.
+
+**2.7 Free-generation probes in the training log** (repair_finetune.py already does this — 2 prompts every EVAL_EVERY steps, decoded into the log). Make it standard in every trainer.
+
+---
+
+## 3. Code hygiene
+
+**3.1 One shared checkpoint loader.** The 49152→49154 (block-CoT think tokens) vocab mismatch is hand-patched with `strict=False` in 9+ scripts; `generate.py` can silently construct a *random* model when keys mismatch. The loader should: read config from the checkpoint, adopt the checkpoint's actual embedding rows, load strict, and hard-fail on any missing key.
+**3.2** Fix `trust_remote_code` dataset loads in the GRPO data path (Bespoke-Stratos-17k errors — non-fatal but noisy and skips data).
+**3.3** Kill hardcoded paths / dead code flagged in the audit (`/root/dimba_repo_audit.md`).
+**3.4** Sampler entry point should assert the checkpoint's training family (flow vs DDPM) instead of letting you sample a flow model with a DDIM schedule.
+
+---
+
+## 4. Strategy: the actual rebuild path
+
+**Recommended: Plan A — pretrained Mamba init, one gap at a time.**
+
+- **Phase 0 — infra ($0):** implement §1–§3 in the clean repo. Contract tests: overfit-8-sentences must pass free generation; partial-noise ladder must show no cliff.
+- **Phase 1 — base model ($0):** take `state-spaces/mamba2-130m` (pretrained AR Mamba-2). Bidirectionalize (init forward+backward directions from the AR weights). Keep its tokenizer. *(Alternative Plan B: MOHAWK-distill SmolLM→AR Mamba yourself, 3–5B tokens ≈ $15–25, verify it generates coherent text, then proceed — pick this only if you need the SmolLM tokenizer/quality.)*
+- **Phase 2 — diffusion conversion, validation slice (~$10):** 1–2B tokens with the full fixed objective (§1), KD from the original AR Mamba (§2.1), all gates and evals (§2.3–2.7) on. **Hard gate:** free generations must show phrase-level coherence and the ladder must be cliff-free, or stop and rethink — do not scale a run that fails here.
+- **Phase 3 — diffusion conversion, full (~$60–120):** 10–20B tokens only after Phase 2 passes. Early-stop on plateau.
+- **Phase 4 — SFT (~$10–15):** fixed SFT loss (§1.4) + CFG prompt-dropout data. Gate: fixed prompts must get *answered*, even badly.
+- **Phase 5 — RL (~$10, optional):** GRPO only if Phase 4's policy produces parseable answers at >0 rate; otherwise there is no reward gradient (proven twice now).
+- Total: **~$100–170**, roughly a third of budget28, with 5 abort points instead of 0.
+
+**Seriously consider before committing: switch to masked *discrete* diffusion (MDLM / LLaDA-style) instead of continuous-latent flow.** Continuous-latent text diffusion at 135M is the hardest known regime for this task — the t=1.0 "posterior mean blend" pathology and the commitment problem are both *inherent* to the continuous parameterization, which is exactly where both of our failures live. Masked discrete diffusion demonstrably produces coherent text at small scale, keeps the bidirectional Mamba backbone, and makes §1.1/1.2 moot (no SNR weighting, CE-native objective). Bigger rebuild, but it removes the two failure modes instead of patching them.
+
+**Scale honesty:** budget28 + the recovery session validated one negative jointly: *135M continuous diffusion + ≤$400 does not produce a QA-capable model*, even with a perfect denoiser as the starting point. If the goal is a model that answers questions, either the discrete-diffusion switch or ~350M+ params should be part of the plan.
+
+---
+
+## 5. Assets that survive (private HF repo `devnull37/d1-135m-28b`)
+
+| File | Use in rebuild |
+|---|---|
+| `repaired_final.pt` | Best continuous-diffusion checkpoint (perfect denoiser t≤0.9 on unseen text). Init for any continuation of Plan-continuous. |
+| `budget28_distill_final.pt` | The collapsed original — negative-control / ablation baseline. |
+| `sft_repaired_final.pt`, `budget28_sft_final.pt`, `budget28_grpo_final.pt` | Forensics only. |
+| `ROOT_CAUSE_BUDGET28.md` | The full experiment chain (overfit contract → clamp → partial-noise ladder) — reuse these diagnostics as the Phase-0 contract tests. |