code/docs/ROADMAP.md

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.Parameters 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.