# align roadmap: robust training, OLMoE-class MoE, formal-proof track

Companion to `docs/REFERENCE_GAP_ANALYSIS.md` (evidence and citations live there; item IDs below map to its sections).
Paths are relative to `/apps/development/rlhf/align/`. Effort: S <=2 days, M ~1 week, L 2-4 weeks, XL >1 month.

## Ground rules for every item

**Dependency policy** (`docs/DEPENDENCY_POLICY.md`, enforced by `tests/test_dependency_firewall.py`): runtime imports are
torch, triton, numpy, safetensors and stdlib only. Reference code is ported by reading it, never imported. Oracles against
HF/transformers/mistral-common live behind the `oracle` marker in the legacy venv. Anything needing a banned package
(`pyarrow` is the single whitelisted file, `data/ingest_parquet.py`) goes in `extras/`. Every item below states "compliant"
only if it needs nothing outside that set; kernels additionally pass through `kernels/gate.py`.

**Host-safety preconditions** (`/apps/CLAUDE.md`, enforced in `utils/guard.py`) apply to every step marked GPU:
1. `free -g; uptime; awk '$3=="nvme2n1"{print $12}' /proc/diskstats` passes: MemAvailable >= 45 GB, 1-min load < 2x cores,
   disk queue < 64. `ps -eo stat | grep -c '^D'` not sustained above ~5.
2. `nvidia-smi` shows the needed GPUs free (they carried 13-19 GB of other jobs when the maps were written). Single-GPU
   validation pins one device via `CUDA_VISIBLE_DEVICES` before importing torch.
3. One heavy job at a time (heavy-job flock, `acquire_heavy_job_lock`). Never concurrent with a data build.
4. No writes to `/` or `/tmp`; scratch in `.tmp/`; at most 2 checkpoints kept; `--num-workers` <= 2 per rank; save every <= 100
   micro-steps for anything longer than 20 minutes; delete partials immediately.
5. Verify progress by `stat -c %y` on the log, not PIDs.
Items marked **GPU** need those; everything else is CPU-testable inside the existing CPU test loop (gloo for multi-rank).

**Definition of done for any item**: new tests green, `pytest tests/ -m "not gpu"` count does not regress, firewall test green,
README/ARCHITECTURE updated where behaviour changes, and a recorded deviation note if a documented contract changes.

---

## Status at 654c20c (2026-10-12)

Measured on the released OLMoE-1B-7B-0125 weights and on 4x A10G. Sources: `docs/MEASUREMENTS.md` (Track M sections (a) to (h)) and `docs/LEAN_VERIFIER.md`. Each line names its section.

- **OLMoE real-checkpoint parity (P1 MoE core, P5 rows):** fp32 max abs logit diff 1.4e-5 to 3.0e-5 against transformers 5.15, top-1 74/74, router top-8 sets 100% on all 16 layers (MEASUREMENTS (e)). The gate is fp32 < 1e-3. bf16 (1.31) is a noise floor, not a gate.
- **Deterministic MoE dispatch (P1-3):** the bf16 `index_add_` combine gave 5 of 5 different top-8 expert sets on identical inputs. The fixed-order combine gives 1 of 1 (MEASUREMENTS (f)).
- **4-GPU FSDP2 SFT kill/resume (P0-1, P0-8, P0-9):** a SIGTERM to one of four ranks stops all four at the same optimizer boundary. The supervisor relaunch resumes at micro 44 and loss stays within the run-to-run floor (max 3.3e-3 against a tolerance of 1.09e-2; MEASUREMENTS (g)). On 2 ranks, resume from a stop checkpoint and from a cadence checkpoint is bit-exact against an uninterrupted run, and a wrong optimizer shard is detected (MEASUREMENTS (a)).
- **FSDP cost, replacing the old constants (P1-15):** 277 padded tokens/s at 4 x 1,024-token micro-batches is **0.53 TFLOP/s per GPU**. FSDP is comm-bound over PCIe (13.55 of 14.8 s per micro-step). A micro-batch of 8,192 tokens gives 2,140 tokens/s. The 4-GPU ceiling for this box is about 2 to 3.4 k tokens/s, which is **0.2 to 0.3 B tokens/day**. This replaces the earlier 25 TFLOP/s assumption and the 1.1 B tokens/day figure (MEASUREMENTS (h)).
- **Dense sampler rank sharding (fixed in 654c20c):** `scripts/sft.py` and `scripts/pretrain.py` gave every rank the same batch stream. `build_training_sampler` now shards through `DistributedBatchSampler` (`tests/test_training_sampler_sharding.py`). No dense 7B run has been made since the fix.
- **Lean verifier, Track L (commit label "P1"; implements the P4-1 row):** pre-registered ceiling test accepts 177 of 178 known-good proofs (1 timeout) and 0 of 61 adversarial cases (`docs/LEAN_VERIFIER.md`). **Open:** no OS sandbox. The REPL runs as the same user with filesystem access. Run the pool as an unprivileged user before adversarial RL.
- **Still open before a multi-hour GPU run:** P0-3 (`offload_adamw` clipping), P0-5 (`no_sync`), P0-6 (rank-local logging; the `history.jsonl` collision is live), P0-7 (token-count normalisation), P0-14 (adamw8bit at 1.8B local elements), P1-1 (`ModelConfig.moe`), and `kill -9` during a save.

## P0: robustness blockers (do first; each is small)

Goal: a multi-hour 4-GPU SFT/DPO/GRPO or continued-pretrain run that cannot silently burn GPU time, loses little on a kill,
and whose logs can be trusted.

| ID | Item | Acceptance test | Files | Effort | GPU? |
|---|---|---|---|---|---|
| P0-1 | Coordinated cancel: SIGTERM/SIGINT handler sets a flag; flag all-reduced (MAX) on the gloo group every K micro-steps; loop exits at an optimizer-step boundary, writes a final synchronous checkpoint, restores handlers. Add `time_limit_s` and `extra_steps_after_cancel` | Multi-process gloo test: send SIGTERM to one rank mid-run, all ranks stop at the same step, checkpoint loads and resumes to bit-identical state; no hang | edit `train/trainer.py`, `utils/dist.py`; new `tests/test_trainer_cancel.py` | S | no |
| P0-2 | Non-finite protection: halt when consecutive skipped optimizer steps exceed N or the all-reduced logged loss is non-finite; make `GradNormTripwire` treat NaN as a fire, not "blind"; decision from a reduced value on every rank | Injected NaN-grad test: run halts after N skips with a clear error on all ranks, no NCCL-timeout hang; PPO `max_consecutive_skips` behaviour unchanged | edit `train/trainer.py`, `train/probes.py:325`, `config.py`; extend `tests/test_trainer*.py` | S | no |
| P0-3 | Clipping on for every optimizer: move global-norm clip into `TrainModule.optim_step` (covers `offload_adamw`), default 1.0 in sft/dpo/pretrain TOMLs, log pre-clip norm | Test that offload_adamw and adamw8bit clip identically on a synthetic grad; pre-clip norm in `history.jsonl`; stale comment `offload_adamw.py:632` removed | edit `train/train_module.py`, `train/offload_adamw.py`, `configs/{sft,dpo_*,pretrain}.toml` | S | no |
| P0-4 | Arm tripwires by default for SFT/pretrain (`grad_norm`, `length_drift` as applicable) and add the spike-skip gate (rolling 128-step mean + 6 sigma, skip not halt, spike-rate metric) | Config test: sft/pretrain TOML load with armed tripwires; spike injection skips one step and logs it | edit `config.py:329-335`, `train/probes.py`, optimizers' pre-update gate | S-M | no |
| P0-5 | FSDP `no_sync` / `set_requires_gradient_sync(False)` on non-final micro-batches | CPU gloo test: gradient after N micro-batches equals the sync-every-step result (fp32 tolerance) and reduce-scatter call count drops to 1 per optimizer step | edit `train/train_module.py`, `train/fsdp.py` | S | **GPU** to measure step-time gain (single 4-rank 20-step probe via the existing probe gate) |
| P0-6 | Verify, then fix, rank-local logging: reduce loss/token counts on the gloo group at log cadence, rank-0-only `Tracker` writes, tripwire decisions from reduced metrics | First a 2-rank gloo test that reproduces (or refutes) interleaved `history.jsonl` records; after the fix, exactly one record per log step with the dp-mean loss | edit `train/trainer.py:315-317,505-589`, `utils/tracking.py` | M | no |
| P0-7 | Token-count loss normalisation: TrainModule computes valid-token count of the whole rank batch first, passes `loss_div_factor` to the loss adapter, sum-reduction per micro-batch, FSDP divide factor handled | Test: split one batch into 1, 2, 4 micro-batches and across 1/2 ranks, loss and grads equal within fp32 tolerance (today they differ when valid tokens vary). Interface change documented because MoE aux losses depend on it | edit `train/train_module.py`, `train/sft_loss.py`, `train/dpo_loss.py`, `rl/grpo_module.py` (check contract) | M | no |
| P0-8 | Checkpoint hardening: fsync files and parent dir before the publishing rename; `latest_checkpoint` validates completeness (resume.pt + safetensors + manifest) and prefers highest step across the run's own dirs; persist `pinned_checkpoints` in `resume.pt`; marker-first delete in prune; sparse permanent tier (`permanent_every_micro_steps`, weights-only) | Kill-during-save and kill-during-prune tests leave a loadable newest checkpoint; stale higher-numbered foreign dir does not shadow; pins survive resume | edit `train/trainer.py:626-886`, `config.py` | S | no |
| P0-9 | Save RNG state (torch cpu/cuda, numpy, python) and split `load_optim_state` / `load_trainer_state` flags in resume | Resume reproduces dropout/sampling stream exactly on one rank; warm-start-without-optimizer path covered | edit `train/trainer.py:888-1019` | S | no |
| P0-10 | Distributed init: `cpu:gloo,cuda:nccl` composite backend, configurable timeout, `set_device` before init; GC control (disable auto GC, gen-1 collect every 1000 steps and after saves) | gloo test unchanged; unit test that GC is restored on exit; timeout reads from config | edit `utils/dist.py:32-41`, `train/trainer.py` | S | no (NCCL timeout behaviour only matters on GPU) |
| P0-11 | Launch supervisor `scripts/supervise.sh`/`.py`: relaunch torchrun after crash only if the host preflight passes AND `history.jsonl` mtime shows progress since the last attempt; bounded retries with backoff; GPU-idle and temperature check via nvidia-smi; reuses the heavy-job flock | Fake-crash script test: retries stop after the bound, never launches a second concurrent job, refuses on a failing preflight | new `scripts/supervise.py`, edit `utils/guard.py` | S | no (dry-run test) |
| P0-12 | Triton FA2 gate decision: either re-run `scripts/gate_flash_attention.py` on one free GPU so the gate lands in `.cache/kernel_gates`, or document SDPA as the only supported backend and drop the dead artifact path | `select_attn_implementation` result matches the documented state on a clean checkout | edit `kernels/gate.py` doc, `model/attention.py`; README | S | **GPU** (one A10G, minutes) |
| P0-13 | Sharded sync checkpoint (no rank-0 gather): each rank writes its own weight shard and its own optimizer shard with the OLMo write discipline (tmp + fdatasync + replace), a manifest records world size, mesh and per-tensor placements; keep HF-safetensors export as a separate streaming per-tensor step. Optimizer state for adamw8bit/offload written dequantise-on-save or with a documented per-rank layout plus a reshard tool | Round trip at world size 2 on CPU; load at world size 1 and 4 for weights via the reshard path; host-RAM peak at save is bounded by one shard (assert with the existing allocation-ledger style test); `roundtrip_gate` still passes | new `train/dcp_ckpt.py`, edit `train/trainer.py`, `train/train_module.py`, `model/checkpoint.py` | L | **GPU** for a 7B 4-rank save/load timing run under the preflight; start with the tiny config on CPU |
| P0-14 | GPU evidence pass for adamw8bit: single-GPU step time and parity at 1.8B local elements, then 4-rank 20-step run with the stage gates | Numbers recorded in README status; no regression vs CPU-verified claims | scripts only (`bench_7b_end_to_end.py`) | S | **GPU, full preflight, exclusive** |

Exit criterion for P0: kill -TERM, kill -9 at a random point, injected NaN, and injected loss spike each end in the documented
state in automated tests; one real 4-GPU run of at least 2 hours survives a deliberate SIGTERM and resumes with matching loss.

### G0 status (2026-10-12): what landed, in commit-ready form

Gate G0 is the set of control-plane fixes that must hold before any multi-hour GPU run. Everything below is verified on CPU (single process, plus 2-process gloo for the collective halves); **nothing here has run on NCCL, four ranks or a 7B model.** Items are listed by roadmap ID. "Done" means the acceptance test in the table exists, passes, and has a mutant that removes the fix and fails it.

| ID | Status | What landed (commit message) | Not done |
|---|---|---|---|
| P0-1 | **done** | `trainer: coordinated stop` - SIGTERM/SIGINT set a flag; `train.max_wall_seconds` is a second source; both are MAX-reduced with the skip/non-finite flags in ONE collective at every optimizer-step boundary, never mid-window. A stop always writes a checkpoint, wins over `max_micro_steps`, restores handlers, and a second SIGINT escapes. Stage scripts exit `Trainer.stop_exit_code` (82 wall budget, 128+n signal) instead of running the post-train eval of a half-trained model. 2-rank gloo tests: SIGTERM to rank 1 only, SIGTERM from a parent, wall clock on one rank. | `extra_steps_after_cancel`; the flag is reduced every boundary, not every K micro-steps (cheaper than the plan, no downside seen). |
| P0-2 | **done** | `trainer: skip cap + non-finite loss + rank-reduced HALT` - a "bad" step is one the grad gate skipped or one with any non-finite micro-batch loss (a loss-only NaN with finite grads used to train on). Halts (`NonFiniteHaltError`, exit 81) when consecutive bad steps exceed `max_consecutive_skipped_steps` (3) or the bad fraction exceeds `max_skipped_fraction` (5%, evaluated after 100 boundaries). Tripwire HALTs and the arm-time self-test are reduced first, so every rank raises the same `HaltError` (exit 80) instead of one rank leaving its peers in a collective. Counters persist in `resume.pt`. PPO's critic-only warmup (a step that returns False by design) does not count. | `GradNormTripwire` still treats a NaN norm as "blind" (the skip cap covers it); GPU/NCCL no-hang check. |
| P0-3 | partial | `config: max_grad_norm defaults to 1.0` (OLMoE / OLMo 2 value; `adamw8bit` clips with the global norm and logs the PRE-clip norm, so the grad-norm tripwire is not blinded); stated explicitly in sft/pretrain/dpo TOMLs and `DpoStageConfig`. | `offload_adamw` still does not clip (the default is a no-op there); the clip is not yet in `TrainModule.optim_step`; stale comment at `offload_adamw.py:632`. |
| P0-4 | partial | `configs: arm grad_norm tripwire for sft and pretrain` (the only tripwire those stages can feed; `logps_chosen`/`length_drift` need preference pairs and stay off - arming a guard you cannot feed is how a run halts at step 1 or runs blind). | The rolling-mean + 6 sigma spike-skip gate and its spike-rate metric. |
| P0-8 | partial | `checkpoint: durable publish + completeness` - fsync every file then every directory, rename, fsync the parent; `latest_checkpoint` skips (and warns about) a directory whose weights header, `resume.pt` loop state, or any per-rank optimizer file is missing/truncated, while keeping it on disk; `pinned_checkpoints` persist in `resume.pt` and survive resume + prune. | Marker-first delete in prune; the sparse permanent tier; highest-step-across-dirs rule. |
| P0-9 | partial | `checkpoint: persist RNG` - python, numpy, torch CPU and the current CUDA device are gathered per rank at save and restored on resume (skipped with a warning at a different world size). | Separate `load_optim_state` / `load_trainer_state` flags. CUDA RNG path tested with a fake only. |
| P0-10 | partial | `dist: explicit process-group timeout + NCCL flight recorder` - `ALIGN_PG_TIMEOUT_S` (default 1800 s, parsed strictly), `TORCH_NCCL_TRACE_BUFFER_SIZE`, `TORCH_NCCL_DUMP_ON_TIMEOUT` and the dump-prefix variables (both generations of the name) set before group creation with the dump under `<run_dir>/nccl_trace`, never `/tmp`. Every stage script passes `run_dir`. | `cpu:gloo,cuda:nccl` composite backend; GC control; `TORCH_NCCL_ENABLE_MONITORING` is not set (the header says DUMP_ON_TIMEOUT needs it; its 2.11 default is unverified - needs a 2-rank NCCL run). |
| P0-11 | partial | `supervise: bounded-retry launch supervisor` (`scripts/supervise.py`) - relaunch only if exit != 0, not a HALT code, no fresh `HALT.json`, retries < 2, `history.jsonl` mtime AND the newest *complete* checkpoint advanced, and the host preflight passes after the backoff; child in its own process group, killed and confirmed empty before a relaunch. `align.exit_codes` defines 80/81/82/83 and the halt marker: **measured, `torchrun` exits 1 whatever its workers exit with**, so the marker file, not the exit code, is what reaches the supervisor under torchrun. Guard refusals are `PreflightRefused` (83). | nvidia-smi GPU-idle/temperature check; a crash before the first checkpoint of a fresh run stops the supervisor (deliberately strict). |
| per-rank | done | `dpo, grpo: plumb train.per_rank_optimizer_state` into `DpoStageConfig` / GRPO `TrainCfg` and the scripts' `_run_config`; ON in `dpo_argilla.toml`, `dpo_steppairs.toml`, `grpo_stageb.toml` (sft/pretrain already); GRPO `[train]` loader is now struct-mode (a typo'd key raises). Resume then requires the same world size (the trainer refuses otherwise). | NCCL / 4-rank / 7B unrun. The GRPO `[rollout]` table still filters silently through `__dataclass_fields__`. |
| health knobs | done | `max_wall_seconds`, `max_consecutive_skipped_steps`, `max_skipped_fraction` exist on `TrainConfig`, `DpoStageConfig`, GRPO `TrainCfg`, PPO `TrainCfg` and are forwarded by each script (a field that is not copied trains at the default while the TOML says otherwise). | - |

Track M, CPU-complete (P5 rows): dropless MoE core + OLMoE architecture + HF mapping with transformers-5.15 logit parity (oracle tier); `MoETrainModule` aux losses with the global-batch window buffer, router z-loss and `moe/*` telemetry in `history.jsonl` (`docs/ARCHITECTURE.md`, "Mixture of experts"); a tiny OLMoE steps, logs, checkpoints and resumes bit-exactly through the shared Trainer with per-rank optimizer files (`tests/test_olmoe_trainer_smoke.py`). P1-3 (CPU part), P1-4 (aux losses; the token-count normalisation of P0-7 is still open), P1-6 and P1-9 telemetry are covered; P1-1 (`ModelConfig.moe`), P1-7 (flat expert layout, sharded loader), P1-8 (fp32 router optimizer state), P1-9 collapse tripwire, P1-11/12/13 (GPU measurements) are open.

**Open before the first multi-hour GPU run:** P0-14 (adamw8bit GPU step time at 1.8B local elements), `offload_adamw` clipping (P0-3), token-count loss normalisation (P0-7), `no_sync` (P0-5), rank-local logging (P0-6), and `kill -9` during a save on real ranks. SIGTERM on real ranks is done (see Status at 654c20c).

---

## P1: MoE core (dense-first, torch reference, tiny sizes)

Goal: a mini-B / mini-A OLMoE-shaped model trains correctly on CPU and on one-to-four A10Gs, with every fast path gated
against a torch reference. Order matters; P0-7 (loss normalisation) is a hard prerequisite for P1-3.

| ID | Item | Acceptance test | Files | Effort | GPU? |
|---|---|---|---|---|---|
| P1-1 | Config surface: `ModelConfig.moe` sub-dataclass (num_experts, top_k, expert_hidden, router dtype, gating, normalize flag, lb/z weights, capacity_factor optional, shared expert optional), `qk_norm`, `norm_style`; struct-mode rejects unknown keys; `hf_config` writes/reads the fields and refuses silent reload as dense | Round-trip config test; loading an MoE config into the dense path raises | edit `config.py`, `model/hf_config.py`, `model/mistral.py:70` validate | M | no |
| P1-2 | QK-norm in `Attention` (RMSNorm on q and k per head) with parity test vs a numpy reference; parameter names and HF key tables updated | Existing dense tests untouched (flag off); QK-norm forward/backward matches reference | edit `model/mistral.py:211`, `model/checkpoint.py` tables | M | no |
| P1-3 | `MoEFeedForward`: fp32 router (softmax over all experts, top-k, no renorm by default), stable-argsort dropless dispatch, torch loop-over-experts oracle, index_add combine; capacity-factor path as parity oracle (capacity >= T*k equals dropless exactly); aux-loss channel through `DecoderBlock` that survives non-reentrant AC and does not count recompute passes | Dropless == capacity(inf) == per-expert loop; grads match AC on/off to bf16 tolerance; empty-expert and padding-token cases | new `model/moe.py`; edit `model/mistral.py:367-395`, `train/fsdp.py:90-123` | M-L | no |
| P1-4 | Aux losses: Switch load-balance loss (E/k formulation, weight 0.01, divided by n_layers) and router z-loss (0.001), both normalised by the same token count as CE (uses P0-7); adapters add them to `LossOutput` with metrics | Formula test (balanced routing gives 1.0); invariance to micro-batch split | new `model/moe_loss.py`; edit `train/sft_loss.py`, `train/train_module.py:44-69` | S-M | no |
| P1-5 | LM-head z-loss (1e-5) in chunked fused CE, adding `mult*lse^2` to the objective and its gradient; re-gate | Matches eager reference; gate artifact regenerated | edit `kernels/fused_ce.py`, `kernels/gate.py` (parameterise by kernel name first) | S | no (reference) / **GPU** for the gate artifact |
| P1-6 | Init: trunc-normal(0, 0.02, 3 sigma) for router, experts (incl. w2) in fp32 then cast; explicit expert init; test that no parameter is left uninitialised (catches bare 3-D `nn.Parameter`s `apply()` skips); depth-scaled option | Test enumerates all params and asserts non-garbage statistics | edit `model/mistral.py:559-566`, `model/moe.py` | S | no |
| P1-7 | Checkpoint: flat expert layout `(E*h, d)`, router keys, HF-style per-expert export, strict `_remap` tables for MoE, buffer handling in sharded loader (`train_module.py:173-178` currently rejects buffers), `roundtrip_gate` for expert layout | Save/load/export round trip at world size 1 and 2; unknown key still raises | edit `model/checkpoint.py:66-98,281`, `train/train_module.py:154-199` | M | no |
| P1-8 | Optimizer groups: per-name overrides (wd, quantised vs fp32 state) with strict glob matching; router and norms in fp32 state (`adamw8bit.py` quantises any tensor with >= 4096 local elements, including the 131k-element router) | Group test; router state stays fp32; single-group behaviour unchanged | edit `train/adamw8bit.py:257`, `train/offload_adamw.py`, `config.py` | S | no |
| P1-9 | MoE telemetry and tripwire: per-layer expert load max/mean, fraction of experts below 0.1x mean, router entropy, router lse^2, drop fraction, per-expert grad-norm > 0 check for experts that got tokens; tripwire halts on router collapse; counters tagged to exclude AC recompute | Synthetic collapsed router fires; healthy one does not | new `train/moe_probes.py`; edit `train/probes.py` | S | no |
| P1-10 | Mini-B end-to-end CPU smoke (d768 L12 E32 k4) and mini-A single-GPU smoke: loss decreases, aux losses finite, resume lossless, checkpoint round trip | CPU smoke in suite (tiny); GPU smoke records tokens/s and peak memory | `configs/moe_mini_b.toml`, `configs/moe_mini_a.toml`, `scripts/moe_smoke.py` | M | mini-A **GPU** (one device, preflight) |
| P1-11 | Grouped-GEMM decision: micro-benchmark torch `_grouped_mm` vs loop on one A10G (expect the slow fallback; sm90/sm100 CUTLASS only) | Benchmark table recorded in the analysis doc | `scripts/bench_grouped_mm.py` | S | **GPU** (one device, minutes) |
| P1-12 | Triton grouped GEMM (fwd, dX, dW) for SM86 with autotune pruned to num_stages <= 2 (99 KB smem), template from `reference-repos/DeepSpeed/deepspeed/moe/group_gemm_triton.py`, gated against the loop oracle for fwd and both grads, empty experts included. Only if P1-11 shows the loop is the bottleneck | Gate artifact passes on SM86; throughput >= loop; no host sync | new `kernels/grouped_gemm.py`, gate entry | L | **GPU** |
| P1-13 | Parallelism decision for experts: benchmark FSDP-only experts vs EP all-to-all on 4 A10G over PCIe at 4096 tokens/micro-batch. Build the EP mesh path (HSDP substrate, `all_to_all_single` custom autograd, split-size exchange) only if FSDP-only is comm-bound in the measurement | Measured per-layer ms for both; decision recorded; if EP built: parity vs single-rank on gloo (uneven splits verified in the 2-process probe) | `scripts/bench_moe_parallel.py`; later `train/ep.py` | M (bench) + L (EP) | **GPU, exclusive, full preflight** |
| P1-14 | Sliced upcycling tool (Mistral-7B-v0.3 FFN into E=16 slices of 896) as a state-dict transform. Exact-equivalence test uses forced combine weights 1.0 with k=E (softmax weights sum to 1, so plain routing does not reproduce dense logits) | CPU test on a tiny dense model: equivalence within summation-order tolerance; load result into the P1-1 model | new `scripts/upcycle_dense_to_moe.py` | M | no (conversion); **GPU** only for a 7B parity check, streaming per tensor |
| P1-15 | OLMoE-1B-7B bring-up smoke: micro-batch 1x4096, adamw8bit, 4 ranks, a few hundred steps; record memory (expect ~13.5 GiB/card, my arithmetic) and real MFU | Fits; no OOM; MFU recorded. Measured: 0.53 TFLOP/s per GPU (MEASUREMENTS (h)); the 25 TFLOP/s assumption is retired | `configs/olmoe_1b7b.toml` | S | **GPU, all four free, exclusive** |

Out of scope for P1: aux-loss-free bias balancing and shared experts (not in the OLMoE-1B-7B recipe; add later as flags),
tensor/context/pipeline parallelism, FP8.

---

## P2: data and pretraining pipeline

Goal: a resumable, world-size-invariant pretraining run on a blended corpus, with held-out eval in the loop.

| ID | Item | Acceptance test | Files | Effort | GPU? |
|---|---|---|---|---|---|
| P2-1 | Token-addressed fixed-length dataset: flat token files + doc-ends array (datatrove layout), uint32 permutation memmap per (seed, epoch, size), dp rank takes `indices[:, rank::W]`, resume offset = tokens_processed // global_batch_tokens, dataset fingerprint (sha256 over files, sizes, vocab) compared on resume | Same global batch sequence at world size 1, 2, 4; resume at a different world size continues the same order; fingerprint mismatch refuses | new `data/token_dataset.py`, `data/token_loader.py`; edit `train/trainer.py` ResumableBatchStream | L | no |
| P2-2 | Parallel, resumable tokenise-and-shard build with per-shard receipts; bounded RAM per process, no more than 2 workers during any training job | Killed build resumes from the last committed shard; peak RSS asserted under a cap; refuses to start while a training job holds the heavy lock | edit `data/pretrain_corpus.py`, `data/store.py`; new `scripts/build_pretrain.py` | L | no (CPU job, still under the load preflight) |
| P2-3 | Data mixture: per-source target ratios, max repetition, Hamilton apportionment, ratio-difference report; deterministic cursor checkpointed with the model | Realised ratios within one instance of target; same seed gives same order | new `data/mixture.py` | M | no |
| P2-4 | Quality layer: Gopher/C4 rule filters, exact + MinHash-LSH dedup (numpy signatures, disk-sorted, union-find), n-gram decontamination against align's eval sets, degenerate-instance filter, vocab-range check in collate | Each filter has a labelled golden set; dedup RSS bounded; decontam removes seeded leaks | new `data/filters.py`, `data/dedup.py`, `data/decontam.py` | M-L | no |
| P2-5 | Doc-boundary attention masking: first via per-token segment ids in the SDPA path/`flex_attention` on SM86 (validate JIT under the vendored toolchain), then optionally in align's Triton kernel (block-skip assumption at `kernels/flash_attention.py:114-136` must be reworked and re-gated) | Packed batch equals the same documents run separately; no cross-document leakage in a gradient test | edit `model/attention.py`, `model/mistral.py`, optional `kernels/flash_attention.py` | M-L | **GPU** (SM86 JIT, block sizes unmeasured) |
| P2-6 | Schedules and groups: WSD, linear, alpha_f floor, stateless token-denominated schedule that supports extending the horizon on resume | Schedule values match OLMo-core formulae at sampled steps; resume with changed horizon honoured | edit `train/trainer.py:157-212`, `config.py` | S | no |
| P2-7 | In-loop held-out NLL eval at cadence plus eval-on-finish; spike-recovery workflow (`steps_to_skip` ranges, `hard_stop`) | Eval cadence test; skipped ranges keep token accounting | edit `train/trainer.py`; new `eval/lm_eval.py` | M | no / **GPU** for the first real run |
| P2-8 | From-scratch init flow: meta-device build, post-FSDP sharded materialisation with deterministic per-parameter seeding | Same weights regardless of world size (CPU gloo test at tiny size) | edit `model/mistral.py:515-530`, `train/fsdp.py`, `train/stage_smoke.py` | M | **GPU** for 7B-class |
| P2-9 | Generic byte-level BPE tokenizer with a `unicodedata`-based pre-tokenizer (OLMoE vocab 50,304 class), oracle test against the reference tokenizer in the legacy venv; confirm token dtype handling above 32,768 | Byte-identical on a held-out text corpus | new `tokenizer/bytelevel.py` | L | no |
| P2-10 | Async or bounded-staging checkpoint save (only after P0-13): shard-by-shard staging capped by the 45 GB rule | Peak host RAM during save under the cap; no NCCL watchdog abort | edit `train/dcp_ckpt.py` | M | **GPU** |
| P2-11 | Optional: `torch.compile` per block after AC and before FSDP; profiler callback; MFU/allocator counters in `history.jsonl`; GC and `device_peak_flops` override for A10G (OLMo falls back to an A100 table) | Compile on/off parity; profile trace written under `.tmp/` on one rank | edit `model/mistral.py`, `train/trainer.py` | M | **GPU** |

---

## Experiment P2: Goedel-Prover-V2-8B SFT with the Lean verifier in the loop (in progress)

Not the data-pipeline P2 above; the label is the one the experiment brief uses. User-approved, and it runs ahead of the P4 gate that says to wait for the Leanstral report.

- **Model:** `Goedel-LM/Goedel-Prover-V2-8B`, Qwen3ForCausalLM, 36 layers, hidden 4096, 32 heads / 8 KV heads, head_dim 128, intermediate 12288, vocab 151936, rope_theta 1e6, RMS eps 1e-6, untied embeddings, bf16. Snapshot: `/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a`.
- **Reference code:** `/apps/development/rlhf/reference-repos/Goedel-Prover-V2`. Prompt templates, code extraction and correction prompts come from there (`src/utils.py` DeepSeekCoTHandler, `src/inference.py`). Benchmarks are `dataset/minif2f.jsonl` (244 problems) and `dataset/MOBench.jsonl` (360 problems). Sampling: temperature 1.0, top_p 0.95.
- **User decisions:** (1) generation uses vLLM, in `extras/` (isolated `.venv-rollout`, vllm 0.26.0). (2) Completions are capped at **16,384 new tokens** for the main loop. A **40,960-token** setting is kept for reference baselines. Every number reported from this experiment must state its cap.
- **Verifier:** `align/verify/lean` with the pins above (Lean v4.9.0-rc1, Mathlib d1d1e4b7). Eight workers at about 0.8 GB PSS each. Any Lean run with more than 2 workers takes the heavy-job flock.
- **Data contract:** SFT input is built from the sampled `prompt_ids` and `completion_ids` with no re-tokenisation, and the loss is on completion tokens only. The model code is `align/model/qwen3.py` and `align/model/qwen3_checkpoint.py`.
- **Status:** in progress. No results are recorded in this document yet.

---

## P3: MoE-aware RL post-training

Goal: GRPO that stays stable when the policy is MoE and rollouts are long. Prerequisite: P1 complete, P0 robustness in place.

| ID | Item | Acceptance test | Files | Effort | GPU? |
|---|---|---|---|---|---|
| P3-1 | MoE decode path in `Engine` (router + expert loop; padding tokens excluded from routing counts); engine emits per-token top-k expert ids | Engine logprobs equal trainer forward on the same tokens (existing LOGPROB_BASIS discipline) for the MoE model | edit `inference/engine.py`, `rollout/native.py` | M | **GPU** (small model) |
| P3-2 | Routing replay: router accepts forced `expert_indices`, scores re-gathered at recorded indices so gate gradients survive; indices stored as int16/int32 per sample | With replay on, trainer routing equals rollout routing exactly; mismatch metric near zero; storage matches (T-1) x L x k x bytes | edit `model/moe.py`, `rl/assembly.py`, `rl/types.py` | L | no (CPU parity) / GPU for real runs |
| P3-3 | GSPO sequence-level ratio, OPSM mask, behind pre-registered ablation flags (A3 conformance still refuses unregistered changes); note the length-weighting interaction with Dr.GRPO's constant denominator | Unit tests on synthetic ratios; A3 test updated; default behaviour unchanged | edit `rl/grpo_loss.py:182-262` | S | no |
| P3-4 | Multiple minibatches per rollout batch (so clip, OPSM and rho actually bind); decide with data whether routing replay is needed at that point | Ratio is 1 on the first minibatch and not on later ones; clip fraction logged | edit `rl/batch_source.py:296-307`, `rl/grpo_module.py` | M | no |
| P3-5 | Rollout correctness guards: refuse top_p < 1 / top_k > 0 in RL configs; truncation fraction metric and optional drop/mask of `finish_reason != stop`; bf16 log-softmax fix or measured justification; mismatch-KL metric | Config refusal test; metrics present | edit `rl/stage_config.py`, `rl/assembly.py`, `rl/grpo_module.py` | S | no |
| P3-6 | In-memory trainer-to-engine weight sync over torch.distributed (layer-by-layer `full_tensor` + copy into existing engine parameters), replacing the 13.5 GB checkpoint write per sync; layer-grouped for MoE | Engine output equal before/after sync to the checkpoint path; peak host RAM under cap | edit `rollout/native.py:66-90`, `rl/batch_source.py` | M | **GPU, full preflight** |
| P3-7 | Bounded-staleness async pipeline: `policy_version` on request/result, bounded queue with drop rule, producer overlapped with training; exactly-once result accounting | Staleness never exceeds K in a deterministic fake-engine test; resume loses no samples | edit `rollout/protocol.py`, `rl/batch_source.py` | M | **GPU** to measure overlap benefit |
| P3-8 | Eval and curriculum: unbiased pass@k, stop-rate and length stats, named eval sets; oversampling top-up to a constant live-group count with per-prompt pass-rate memory | pass@k matches the closed form on synthetic data; batch size constant under 50% dead groups | edit `rl/canary.py`, `rl/batch_source.py` | M | no |
| P3-9 | MoE RL telemetry into the GRPO tripwire framework (router entropy, expert load); decide whether to keep aux loss during RL by ablation, not by assumption (trl keeps it, open-instruct GRPO apparently does not) | Tripwire fires on injected router collapse | edit `rl/tripwire.py` | S | no |

---

## P4: formal-proof track (placeholder; design waits for the Leanstral 1.5 research report)

Do not start model-specific work until `/apps/development/rlhf/research/notes/` contains the Leanstral 1.5 report. The items below
are generic enablers that hold for any Lean-4 prover regardless of what that report says.

| ID | Item | Acceptance test | Files | Effort | GPU? |
|---|---|---|---|---|---|
| P4-1 | Lean verifier environment via stdlib `subprocess`/HTTP (no kimina client): assemble statement + proof without slicing at the first `:=`, per-call timeout, memory cap per child, concurrency bounded by the host preflight, failure taxonomy (`no_code`, `lean_<status>`, `timeout`, `crash`) mapped to reward 0 | Adversarial suite must be rejected: `sorry`, added axioms, `native_decide`/unsafe tricks, statement edits, helper-lemma-before-theorem handled correctly, hung process, huge output | new `align/verify/lean/` package (`strict.py`, `repl_pool.py`, `reward.py`, `toolchain.py`), `tests/test_lean_verifier.py` (Lean binary tests marked and skipped when absent). Landed in 654c20c, see Status | L | no GPU; needs a Lean toolchain installed under `/apps` |
| P4-2 | Generic sandboxed-subprocess harness (spawn, timeout, kill, classify) shared by Lean and code execution | Fuzz test: no zombie, bounded RSS | new `align/verify/sandbox.py` | M | no |
| P4-3 | Environment abstraction (reset/step/state, exactly-once episode result: episode, failure, or cancellation) | Fault-injection test | new `align/rl/env.py` | M | no |
| P4-4 | Multi-turn rollout with per-token loss mask (observation tokens masked, token ids concatenated never re-tokenised) | Mask test; GRPO loss ignores masked tokens | edit `rollout/protocol.py`, `rl/assembly.py:126-150`, `inference/engine.py` | L | **GPU** for engine continuation |
| P4-5 | Rejection-sampling / expert-iteration script: generate, verifier-filter (reward == 1), dedup, write SFT shard with per-prompt pass rate | Selection invariants on a fixture | new `scripts/expert_iter.py` | S | no / GPU for generation |
| P4-6 | Long-context inference: prefix sharing for the G samples (copy KV rows first), then paged KV and chunked prefill if 8k+ completions are required; KV capacity budget stated per card (about 64k tokens across all slots with 7B bf16, derived estimate) | Same logprobs as contiguous path; capacity gain measured | edit `inference/engine.py` | M then L | **GPU** |
| P4-7 | Multi-turn chat template and tool-call serialisation (port of the V3 path from mistral-common, oracle-tested), including training that the control-token embedding rows are not zero in the base weights (check `[TOOL_CALLS]` and friends before SFT) | Token-identical to mistral-common in the legacy venv | edit `data/schema.py`, `tokenizer/core.py` | M | no |
| P4-8 | Document masking and long-context training validated for the target's context length; context parallelism only if the report requires it and only via an in-repo kernel (ring_flash_attn is banned) | Decided after the report | n/a | n/a | n/a |

---

## Suggested order of attack

1. P0-1, P0-2, P0-3, P0-4, P0-8, P0-9, P0-10 in one sitting (all S, CPU-only, high value).
2. P0-5 and P0-6 (verify first), then P0-7 (interface change, needed by MoE).
3. P0-12, P0-14 as one exclusive GPU session after the preflight; then P0-13 (the largest P0 item).
4. P1-1 through P1-10 on CPU; P1-11 and P1-15 benchmarks in one GPU session; only then decide P1-12 and P1-13.
5. P2 items 1, 2, 3, 6, 7 as the minimum for a resumable blended pretrain run; the rest as needed.
6. P3 and P4 when the target model and the Leanstral report are known.