Auto-research: SFT data mixture search on a fast proxy for Olmo 3.5

Systematically search SFT data mixtures on a fast proxy model, so the Olmo 3.5 hero SFT run starts from measured settings rather than guesses. Umbrella issue for the mixture-search programme; specific dataset additions are tracked separately (#1851, #1852).

**Scope (updated 2026-08-28): think-only.** The programme trains on Dolci-Think; instruct data is out. This retires the two axes that were previously ranked 1 and 2 in the search space (think:instruct ratio, and mixed-vs-staged think→instruct) — see *Search space* below. Instruct results are kept where they are *evidence* for the think-only choice, not as things to sweep.

**Timeline.** ~1 month of proxy work while pretraining finishes, then ~2 weeks for hero SFT, then ~2 weeks DPO + RL. The proxy work has to be done by the time the hero base lands.

**Deliverable.** Not "the best mixture" — a *mixture-vs-metric response surface with error bars*, so that when the hero base arrives with a different midtrain profile we can reason about how to adjust rather than reopening the search with two weeks left.

## Proxy model

OLMoE3 latent-KDA 1.2B. The full loop (tokenize → SFT → HF convert → eval) is validated end to end: #1843.

**Switch the base from midtrain to long-context.** Pipeline validation used the midtrain checkpoint only because it was the only stage with an HF export at the time. Production SFT happens *after* long-context extension — both dense `Olmo-3-1025-7B` and `Olmo-Hybrid-7B` were SFT'd from post-LC 65k bases, so midtrain made the MoE a stage behind everything it was compared against. Use:

```
oe-training-default/ai2-llm/checkpoints/jacobm/olmoe3/eval-comparisons/legacy-cx8-v2/latent-kda-l2/1.2B/long-context/
```

Its HF config declared `max_position_embeddings: 8192` against 65536 trained; fixed in place 2026-08-26 with Jacob's authorization (all 9 LC exports had it; backups + revert at `~/handoff/lc-config-fix-backup/`). Switching bases invalidates the existing baseline set, so run it as a deliberate arm, not a silent substitution.

The LC think baseline (0.5 epoch, seq 32768) is running as its own tracked arm: #1854.

**Regenerate the SFT config when switching bases.** `scripts/train/debug/kda_mt_sft.json` is derived from the *midtrain* checkpoint and silently drops `recompute_each_block: True` and `lm_head.loss_implementation: fused_linear`, which the LC checkpoint config carries and which dominate memory at long sequence. Rerun `make_kda_sft_config.py` against the LC checkpoint's `config.json` rather than reusing the file.

**Cost per run** (8x B300): 0.54B-token screen ≈ 0.9 h / 7.4 GPU-h; 5B-token confirmation ≈ 8.6 h / 68 GPU-h. Budget in *tokens*, not epochs: the throughput was measured on instruct data, and while the token-normalised cost carries over, an epoch count does not — Dolci-Think rows are ~13x longer.

Use **seq 32768** for think: it is the existing cache's native length (no re-tokenize) and cuts mid-trace truncation 30.6% → 1.95%. Sequence length 16384 measured ~47% *faster* per token than 8192 (31,556 vs 21,544 tok/s/device) at 91.7% of a B300, but those numbers were taken with the two levers above off — the regenerated LC config changes the fit calculus (measured in #1854).

## Prerequisites — do these before any mixture sweep

1. **Noise floor.** Rerun one identical config 3x and eval each. Seed-to-seed variance on this model is unmeasured; every "finding" below the floor is wasted GPU time. This is experiment #1.
2. **Freeze an eval battery with one sensitive task per axis, plus a held-out task** kept out of the selection loop entirely. Optimising ~50 runs against a small noisy battery will overfit to it. GSM8K is useless as a *selection* signal here (base 0.6308 == 1-epoch instruct SFT 0.6308) but valuable as a calibrated control — think moved it +6.29 where instruct moved it 0.0000 (#1849).
3. **Verify the tool-row transform before any tool axis.** The SFT transform drops **82% of `Dolci Instruct Tool Use` rows** (410/500) as prefix-unstable. That was measured on the *instruct* tool set; whether Dolci-Think's tool-bearing rows hit the same path is unverified. The defect is in the transform, not the dataset, so assume it applies until someone checks — a "tool" component that is silently 18% of its nominal weight invalidates any arm that carries it.
4. **Rank-stability check.** Take 3 deliberately different mixtures, run at 0.3 ep and 1.0 ep, and test whether the 0.3-ep ranking predicts the 1.0-ep ranking. If it does, the whole search runs ~3x cheaper. If it doesn't, that is worth knowing before spending the month on short runs.

## Search space, in priority order

1. **Composition within Dolci-Think.** How its constituent subsets are weighted against each other, plus the candidate additions tracked in #1851 / #1852. This is the mixture axis now that the think:instruct dial is gone.
2. **Token budget × composition interaction.** Gains arrive early but the tail is not free: on the MoE, 0.1 epoch captured 72% of the IFEval gain and the remaining 0.9 epochs added +12.5 more.
3. **Tool-subset proportion** — gated on prerequisite 3 above.
4. **LR.** Currently 2.5e-5, inherited from the Olmo Hybrid instruct recipe and never swept for this architecture (midtraining ran at 1.6e-4). Secondary unless something looks unstable.

**Retired 2026-08-28** (kept here so the change is auditable, not silently dropped):

- ~~think : instruct ratio~~ — moot under think-only. The dense report's finding that the two mixtures move disjoint axes (think → maths, instruct → instruction-following) is now background for *why* think-only, not a dial to sweep.
- ~~Mixed vs staged~~ — was ranked the highest-value single question (does one mixed run match Olmo 3's think-then-instruct two-stage recipe at equal tokens). There is no second stage to compare against now. Worth reviving only if instruct data comes back into scope.

## Budget and real constraints

Compute is **not** the limiting factor: 16 GPUs for a month at 40–60% utilisation is 4,600–6,900 GPU-h, enough for 190–280 full-epoch runs against a search that needs 40–80.

The actual constraints:
- **B300 only.** Measured peak 206.9 GiB reserved at seq 8192; H100 (74.5 GiB) and H200 (131 GiB) both too small. No fallback without an EP or FSDP code change.
- **Do not pass `--preemptible` on holmes.** It makes a job backfill-only on a strict-priority cluster, so it never schedules when full: 70 min of `FailedScheduling` vs 30 s without it. 27 scripts under `scripts/train/debug/` pass it by default. Pass an explicit `--timeout` instead, since allocated jobs get an 8 h min-runtime shield.
- **Storage is the likely real bottleneck.** 207 GB per checkpoint, ~2 TB per long run. Budget deletion into the loop rather than as cleanup.
- **Tokenisation cache hazards:** #1818 (cache key not reproducible across builds, silently discards tokenization) and #1815 (cache not invalidated by label-derivation changes) both bite a workflow that tokenises many mixtures. Confirm cache *hits* rather than assuming.
- Dolci-Think is ~24.7B tokens at ~10,951 tok/row (13x instruct's ~840); tokenise a subset unless doing multiple epochs.

## Open questions

- **Hero architecture and size are unconfirmed.** Every proxy finding rests on this lineage being representative. If Olmo 3.5 is dense or hybrid instead, the harness transfers but the proxy should change. Worth settling with the pretraining team early.
- **Multi-node is untested for this model.** Both 2x8 smoke attempts died to holmes cordons before running model code. Expected ~96% scaling on 800 Gbps/GPU InfiniBand, but unmeasured; the likelier hazard is the synchronous 207 GB checkpoint write from 16 ranks.

Related: #1843 (pipeline validation), #1849 (think vs instruct on the MoE), #1851 / #1852 (specific dataset additions), #1804 / #1815 / #1818 (tokenisation hazards).