MoE + Dolci-Think SFT baseline: seq 16384, reasoning-general math gains where instruct SFT gave 0.0000

> # ⚠️ THIS EXPERIMENT USED THE **MIDTRAIN** CHECKPOINT, NOT THE LONG-CONTEXT ONE
>
> Everything below is SFT from the **8192-native midtrain** base:
> `.../olmo-ddp/midtraining/mt-1p2b-kda-ev2-neg-nope-gated-latentmoe-l2-paper-cx8-samebatch-lr1p6e-4-r1/step63802`
>
> A **long-context sibling of the same lineage and step count exists**, native 65536:
> ```
> /weka/oe-training-default/ai2-llm/checkpoints/jacobm/olmoe3/olmo-ddp/long_context/legacy-cx8-v2/latent-kda-l2/1.2B/cx8-samebatch/step63802
> ```
> (verified: `train_module.max_sequence_length = 65536`, `rank_microbatch_size = 65536`)
>
> **Read this issue as an instruct-vs-think data-mixture experiment on the midtrain base**, which is
> what it is and what it is internally consistent for. Do not read it as the recommended baseline
> checkpoint. Midtrain was the consistent choice at the time — #1843 and every baseline number here
> use it, including the instruct-SFT arm that supplies the `+0.0000` control the whole result rests on
> — but a follow-up on the LC base is the better artifact to build on.
>
> **The LC HF export declares the wrong context — confirmed by its author.** The export at
> `eval-comparisons/legacy-cx8-v2/latent-kda-l2/1.2B/long-context/hf` declares
> `max_position_embeddings: 8192` for a model trained at **65536**. Jacob, asked directly: *"that was
> just a mistake if so, the conversion code should set the config to 65k at all times, the model was
> trained to 65k."*
>
> (This warning previously said the 8192 was a deliberate choice for stage-comparable eval instance
> sets. That was a plausible story fitted to the evidence, not something the artifacts could settle,
> and it was wrong. Only the author could resolve it.)
>
> **Scope:** all **9** long-context exports under `eval-comparisons/` are affected —
> `latent-kda-l2` and `kda-nonlatent`, every size 275m/480m/810m/1.2B, plus
> `480m/long-context-lr40pct`. The 17 pretrain and midtrain exports declare 8192 and are **correct**,
> having been trained at 8192. Uniform-8192 batch conversion, right for every stage except LC.
>
> **No re-conversion needed.** The value appears in exactly two fields —
> `config.json:max_position_embeddings` and `tokenizer_config.json:model_max_length`. `tokenizer.json`
> sets no truncation or padding, `generation_config.json` has no length field, and the model is NoPE
> (`use_rope: false`, so `modeling_olmo3moe.py:1126` leaves `rotary_embs = None`) — no weights or
> buffers depend on it. Two JSON edits, versus ~2 h per 35 GB re-export.
>
> **Why it matters:** the value governs vLLM's `max_model_len`. Evaluating an unfixed export caps
> context at 8192 and silently drops over-length prompts — the same mechanism that dropped 512
> ifeval_mt instances in this issue. Fix the config or pass `hf_overrides` before evaluating. For
> *training* it is irrelevant; the olmo-core DCP config carries 65536 correctly. These are Jacob's
> artifacts and patching is his call.
>
> **Note also:** at seq 16384 the LC base changes nothing about truncation (still 30.6% of think rows
> cut). Truncation only improves if you also train at 32k/64k — and 16384 already sits at 245.4 GiB of
> 267.7 on a B300, so 65536 is a memory-architecture problem (4x microbatch, needs
> `recompute_each_block`, likely multi-node), not a config flag. "LC base at seq 16384" and "train at
> 64k" are separable experiments; the former costs nothing extra.

---

Tracking issue for making **OLMoE3 latent-KDA MoE + `allenai/Dolci-Think-SFT`** the standing SFT baseline for future iterations. Follow-on from #1843 (which covered the same model on Dolci-**Instruct**).

Model: 18.5B total / 1.3B active, 20 layers (15 KimiDeltaAttention + 4 NoPE full-attention), 512 routed experts top-16 + 1 shared, d_model 1280. Base = Jacob's midtrain `step63802`.

Plan, full run log and failure ledger: `/weka/oe-adapt-default/abhishekr/handoff/moe-think-plan.md`

## Headline: it works, and the gain is reasoning-general

Anchor run at **1,806,696,448 tokens** — chosen to exactly match dense-think, dense-instruct, hybrid-instruct and MoE-instruct, so the numbers sit in one matrix. seq 16384, LR 2.5e-5, 1x8 B300.

| task | n | base | + think | Δ | paired McNemar |
|---|---|---|---|---|---|
| **gsm_symbolic** | 5000 | 0.5606 | 0.6476 | **+8.70** | z=12.29, p=1e-34 |
| **gsm8k** | 1319 | 0.6308 | 0.6907 | **+5.99** | z=4.92, p=8.5e-07 |
| **math500** | 500 | 0.3100 | 0.3520 | **+4.20** | z=2.01, p=0.044 |
| **ifeval_mt** (restricted) | 1262 | 0.2686 | 0.4128 | **+14.4** | ~7.6 unpaired |

**Instruct SFT moved gsm8k by exactly +0.0000 at the same token budget.** That flatness is what makes these deltas attributable to the think mixture rather than to SFT in general.

**Pre-registered test passes.** GSM-Symbolic exists to separate real arithmetic reasoning from template familiarity on GSM8K. Criterion fixed before the base number was measured: symbolic Δ ≥ gsm8k Δ ⇒ reasoning-general. Result: **+8.70 > +5.99**, the same ordering the dense report found (+7.2 > +5.8). Not contamination or memorisation.

**Cross-architecture transfer**, on gsm8k code verified byte-identical across `479f90e..moe-parity`:

| | dense 7B | MoE 18.5B/1.3B-active |
|---|---|---|
| gsm8k | +5.8 | **+5.99** |
| gsm_symbolic | +7.2 | **+8.70** |

## seq 16384 instead of the native 8192

At 8192, **49.1%** of Dolci-Think rows are cut mid-trace; at 16384, **30.6%**. Measured from the tokenised cache's doc offsets. Truncated rows are unterminated reasoning traces, i.e. supervision to never stop.

16384 turned out to be **~47% faster per token** than 8192 (31,556 vs 21,544 tok/s/device on B300) — KDA is linear attention and only 4 of 20 blocks are full attention at d_model 1280, so longer sequences amortise better. 245.4 GiB reserved of 267.7 at 8 ranks.

Caveat worth stating: this *narrows* the context gap to dense/hybrid (32768) rather than opening a new one, but the within-family instruct-vs-think contrast now varies context as well as data. Instruct rows are ~840 tokens and sequences are packed, so seq length is near-irrelevant for that arm.

## Saturation curve — the actionable finding

| step | tokens | gsm8k | Δ |
|---|---|---|---|
| base | — | 0.6308 | — |
| 345 | 0.36B | 0.6748 | +4.40 |
| 690 | 0.72B | 0.6710 | +4.02 |
| 1380 | 1.45B | 0.6937 | +6.29 |
| 1723 | 1.81B | 0.6907 | +5.99 |

**~3/4 of the gsm8k effect is present by step 345 — 1.9% of one think epoch.** 1380→1723 is flat within noise. Meanwhile ifeval_mt is still climbing (+11.3 at 1380 → +14.4 at 1723). So math transfer is cheap and early-saturating; instruction-following is the axis that keeps paying.

## LR is not a sensitive knob

Two arms at the **full anchor length** (not a short screen — `build_scheduler` derives warmup and annealing from `num_training_steps`, so a 300-step probe fully anneals inside its window and systematically favours the larger LR):

| arm | last-100-step mean CE | load imbalance |
|---|---|---|
| 2.5e-5 | 0.8681 | 3.476 |
| 5e-5 | 0.8577 | 3.376 |

Indistinguishable. No expert collapse in either — imbalance fell monotonically in both.

## Methodology notes that changed conclusions

1. **ifeval_mt denominators differ by context length.** A 16384 checkpoint scored **1582** instances vs the 8192 baselines' **1262**, because vLLM accepts prompts the shorter-context models rejected. Comparing raw numbers would have been invalid *and looked entirely normal*. Restricted metric via `handoff/rescore/restrict_to_baseline.py`, validated against a known baseline first. (192 prompts exceed even 16384, so "full set" n=1582 is still not the full 1774.)
2. **Run the paired test.** math500 unpaired reads z=1.72 / 1.39 — "no effect". Paired McNemar: z=2.50 / 2.01, significant. At n=500, 80% of items agree between arms and contribute variance without signal. Reported unpaired, this task would have been written up as a null and that would have been wrong.
3. **math500 is underpowered** (n=500, SE on a difference ≈ 3.0pt): a null there is uninformative, not evidence of no effect. Pre-registered before the base landed.
4. **Converter defaults to `--max-sequence-length 8192`.** Converting a 16384 checkpoint with the default exports `max_position_embeddings=8192` and silently caps vLLM — worst on exactly the long-generation tasks the larger context exists to serve. All five exports verified at 16384 before evaluation.
5. **The eval launcher's default `OLMO_EVAL_DIR` is pre-#296**, i.e. the buggy IFEval loose scorer. All numbers here are at `abhishekr/moe-parity` (post-fix). Baselines were re-scored offline to match; the fix costs ~1.5-1.7pt uniformly and leaves deltas unchanged.

## In progress

**Phase 2** — the actual baseline run: 0.5 epoch (9,124 steps / 9.57B tokens, 5x the anchor), 1x8, seq 16384, LR 5e-5, checkpoints every 1,825 steps. Currently ~85%. Will extend the curve to test whether the early plateau holds at 5x the budget.

## Known-unmeasured: 2x8 multi-node

Two smoke attempts were cancelled by cordoned nodes **before running any model code**, so nothing is known about 16-rank DDP or the 207 GB synchronous checkpoint write at that scale. Not a negative result — an absence of measurement. Worth retrying when holmes is healthy; it would halve wall-clock (5.3h vs 10.5h per half-epoch).

## Repo-level issues surfaced along the way

- **96 scripts (27 under `scripts/train/debug/`) pass `--preemptible`.** On holmes that makes a job unallocated ⇒ backfill-only ⇒ it never schedules when the cluster is busy. A job that failed to schedule for 70 minutes at priority `high` ran in **30 seconds** without the flag. This is the default failure mode for real work on holmes and reads like a capacity problem.
- `--max_retries` does not cover cordon cancellations (`canceledCode: 12`), and multi-node additionally requires all replicas up within 10 minutes — so one bad node kills the group with no retry.
- Cache-key non-reproducibility (#1818) blocked reusing an existing tokenisation; see comment there.

## Artefacts

- Anchor arms: W&B `nwq5613d` (2.5e-5), `ly87xpq5` (5e-5); Phase 2: `6ks90lcj`
- Checkpoints: `deletable_checkpoint_states/r708wdtx` (anchor A, 5 steps + HF), `givs31pn` (arm B), `s3zl5u0c` (Phase 2)
- Launch script: `scripts/train/debug/oc_sft_olmoe3_kda_think.sh` (branch `think/moe-baseline`)
- Tooling: `handoff/rescore/{restrict_to_baseline,mcnemar_generic}.py`