Replicate Olmo 3 Think SFT on the KDA proxy: 2 epochs of Dolci-Think + paper-comparable evals

Claude Code · bridge-cse session (Abhishek's #1853/#1859 anchor work) — agent-written

Successor to #1859 (1-epoch baseline); feeds #1853. Goal: replicate the Olmo 3 paper's Think SFT recipe on the OLMoE3 latent-KDA 1.2B proxy and compare evals as directly as possible.

## Recipe (primary source: the released checkpoint's own config)

`/weka/oe-adapt-default/jacobm/checkpoints/olmo3-7b-think-sft/config.json` says: **`max_duration: 2 epochs`**, global batch **1,048,576 tokens** (matches ours), **LR 5e-5** (ours was 2.5e-5), seq **32768**.

We match epochs and batch — the recipe's transferable, token-denominated parts. **Two documented divergences:**

1. **LR stays 2.5e-5, not 5e-5.** LR does not transfer across architecture and scale: 5e-5 was tuned for 7B dense, has no evidence on this 1.2B-active KDA MoE, and our 2.5e-5 does (#1854 probe + two clean runs). Keeping it also makes epoch count the *only* delta vs #1859, so 1ep-vs-2ep stays unconfounded.
2. **Seq 65536, not 32768.** Per #1859's H005 result, training at 32768 on a 65536-native base collapses RULER@65536 (0.0883 vs base 0.3751) while 65536 is throughput-neutral. (Plausibly true of the paper model too — RULER is absent from its released tables.) Data: same Dolci-Think cache at 65536, olmo123 template, `allenai/olmo-3-tokenizer-instruct-dev` — matching the paper's think-tokenization workaround per `docs/olmo3.md`.

## Training launch

- Branch `think/lc-baseline`, `scripts/train/debug/oc_sft_olmoe3_kda_think.sh`, MODE=train, `SEQ=65536 NNODES=2` (LR default 2.5e-5), 2 epochs ≈ 47,214 steps, **~30h wall / ~477 GPU-h** on 2×8 B300 (measured: 14.9h/epoch).
- **Fresh run — do NOT resume #1859's step23607 to "save an epoch"**: its LR schedule already decayed to zero over one epoch, so a continuation is not the same trajectory as a 2-epoch schedule.
- 30-step gate first (memory/rendezvous), then train. `PREEMPTIBLE=0` on holmes; explicit `--timeout`; base = LC `cx8-samebatch/step63802`; config `kda_lc_sft.json` (already regenerated from this base — do not use `kda_mt_sft.json`).
- Convert with `-s 65536`; verify `max_position_embeddings` in the export before any eval.

## Comparison targets (Olmo-3-7B-Think-SFT, paper conditions)

| benchmark | 7B Think SFT | 7B Think DPO | 7B Think final |
|---|---|---|---|
| MATH | 94.4 | 92.4 | 95.1 |
| IFEval | 77.9 | 75.9 | 88.2 |
| IFBench | 30.0 | 28.3 | 41.6 |
| PopQA | 20.8 | 24.7 | 23.7 |
| MMLU | 74.9 | 74.8 | 77.8 |
| AIME 2024 | 69.6 | 74.6 | 71.6 |

Our reference rows: anchor step11768 and 1-epoch step23607 (#1859; durable at `/weka/oe-adapt-default/abhishekr/checkpoints/olmoe3-kda-1.2b-dolci-think-sft-65536/`). 1-epoch short-context battery is in flight and will be appended here.

## Eval program (the risky half — gate before spending)

1. **Measure termination rate on the new checkpoint first.** The #1854-era checkpoint ran 41% of open-ended prompts to the cap; `max_tokens=None` tasks cost ~15x estimates on such a model (measured: MMLU:cot projected ~193 GPU-h). No uncapped task launches before this number exists.
2. **32K generation budgets** to match paper conditions. AIME is feasible for the first time: a 32K budget fits the 65536 context (measured penalty for capping at 16K: 15–21pp — do not cap lower).
3. **Code evals need a sandbox configured** (`CodeExecutionScorer3s requires a sandbox`) or they return `Success` with 0.0000 on both arms. Treat unconfigured code scores as missing, not zero.
4. Cost model, measured on this stack: **1.82M decode tokens/GPU-h**; prompt length matters independently (RULER@32768 runs 43 tok/s vs 880 for short-prompt tasks).
5. popqa metric variant differs across harness versions (`popqa_contains` vs `exact_match_first_line`) — pick one, state it, never cross-compare.

## Ops traps for whoever runs this

- Eval launcher: `~/handoff/xarch-evals/eval_kda_mc.sh ` with `OLMO_EVAL_DIR=/root/repos/olmo-eval-moe`. Task names: `ifeval_mt_wildchat_unused_withRewrite`, `ifeval_mt_ood_wildchat_unused_withRewrite`, `ifeval_ood`; plain `ifeval` exists only at olmo-eval-launch commit b03d314 (checkout ready at `/root/repos/olmo-eval-launch-ifeval`).
- Read Beaker status from the most recently created job, not `jobs[0]`; `beaker job events` before diagnosing queue delays.
- RULER grid (8192–65536) on the final checkpoint is mandatory — inside-window scores alone read a cliff as a win.