code/README.md

align

Consolidated post-training alignment codebase for Mistral-7B. SFT, DPO, PPO, and GRPO training with torch + triton + numpy + safetensors only.

Status (2026-10-12, HEAD 654c20c: Track M on real weights, Lean verifier)

Full CPU suite (not gpu, not oracle), HEAD 654c20c:
  1608 passed, 13 skipped

Oracle tier (legacy venv, transformers 5.15, subprocess only; -m "oracle and not gpu"):
  17 passed, 6 skipped in 32s   (OLMoE logits/router/aux-loss/checkpoint parity at ~1e-6)

Dependency firewall:
  2 passed, 0 violators, no whitelist additions

Dependency firewall: clean (`tests/test_dependency_firewall.py`).

Measured at 654c20c, on the released OLMoE-1B-7B-0125 weights and 4x A10G (details, raw files and caveats in docs/MEASUREMENTS.md):

• OLMoE parity: fp32 max abs logit diff 1.4e-5 to 3.0e-5 against transformers 5.15, top-1 74/74, router sets 100% on 16 layers (section e).
• Deterministic MoE combine: a bf16 index_add_ gave 5 of 5 different expert sets on identical inputs. Fixed (section f).
• 4-GPU FSDP2 SFT kill/resume: SIGTERM to one rank, supervisor relaunch, loss within the run-to-run floor (section g). On 2 ranks, resume is bit-exact (section a).
• FSDP throughput: 0.53 TFLOP/s per GPU at 1,024-token micro-batches. FSDP is comm-bound over PCIe. Ceiling 0.2 to 0.3 B tokens/day on this box (section h).
• Dense sampler: rank sharding fixed. Before, every rank trained the same batches (tests/test_training_sampler_sharding.py).
• Lean verifier (align/verify/lean/, see docs/LEAN_VERIFIER.md): 177/178 known-good accepted, 0/61 adversarial accepted. No OS sandbox yet.

Still unmeasured: kill -9 during a save, DPO/GRPO/PPO on multiple GPUs, grouped_mm under FSDP, expert parallelism, and a 7B dense model.

Last pass in one line: Gate G0 (control plane) and the CPU half of Track M landed. A SIGTERM or wall-clock budget on any rank stops every rank at the same optimizer-step boundary after a checkpoint; HALTs, skip caps and non-finite losses are rank-reduced and exit with distinct codes (80-83) that survive torchrun through <run_dir>/HALT.json; checkpoints are fsynced, completeness-checked and carry RNG/pins/skip stats; scripts/supervise.py relaunches only after a non-HALT crash that made progress on a healthy host; clipping (1.0) and a grad_norm tripwire are on for SFT/pretrain; per-rank optimizer files are on for DPO and GRPO. Mixture of experts: a dropless OLMoE (transformers parity), MoETrainModule with global-batch load-balance and router z-loss (Demons window buffer, off for post-training, 0.01/0.001 for pretrain/proxy) and moe/* telemetry in history.jsonl, and a tiny OLMoE that trains, checkpoints and resumes bit-exactly through the shared Trainer. See docs/ARCHITECTURE.md ("Control plane", "Mixture of experts") and docs/ROADMAP.md ("G0 status").

Not done, in order of consequence: offload_adamw ignores max_grad_norm (P0-3); token-count loss normalisation and no_sync (P0-5/7); rank-local logging (P0-6); ModelConfig.moe so a TOML can build an OLMoE (today it is constructed in code); fp32 router optimizer state and the router-collapse tripwire; the adamw8bit GPU step time at 1.8B local elements (P0-14).

What this is

A pure-alignment training stack for Mistral-7B-v0.3 without HF transformers, omegaconf, wandb, bitsandbytes, or vLLM at runtime. Every module is imported and tested in isolation on CPU. All state is preserved through checkpoints: resume is lossless, and trajectory continuity is gated.

Dependency policy (three lines)

1. Allowed at runtime: torch, triton, numpy, safetensors, and align itself. No others.
2. Removed from transformers stack: tokenizer (vendored as pure-Python BPE), config reader (stdlib toml + own codec), optimizer (in-repo torch block-wise 8-bit AdamW, optimizer.name = "adamw8bit"; CPU-offloaded AdamW stays as the RunConfig default), LoRA (own nn.Linear adapter).
3. Banned: vLLM, datasets, peft, deepspeed, math_verify. They are isolated (extras/) or replaced (align.verify, align.data.hub).

See docs/DEPENDENCY_POLICY.md for the replacement table and vendored asset inventory.

Training stage ladder

SFT (scripts/sft.py)
├─ Input: base model + SFT corpus
├─ Output: SFT checkpoint
│
├─ DPO (scripts/dpo.py)
│  ├─ Input: SFT checkpoint + preference pairs + reference cache
│  ├─ Output: DPO checkpoint
│  └─ Golden anchor: 0.09 win rate vs SFT on primary contrast
│
├─ PPO (scripts/ppo.py)
│  ├─ Input: SFT checkpoint + reward model
│  ├─ Reward: process reward (per-token) or outcome
│  ├─ Output: PPO checkpoint
│  └─ Status: draft; `ppo_stage.toml` has unresolved [TBD] fields
│
└─ GRPO (scripts/grpo.py)
   ├─ Input: SFT checkpoint + reward model
   ├─ Variant: grouped relative policy optimization (batch-constrained scaling)
   ├─ Output: GRPO checkpoint
   ├─ Reward shapes: outcome agreement, KL penalty, length (soft/hard)
   └─ Stage B: 32 prompts x 16 generations x 768 tokens, 3-rank FSDP2

Each stage reads config from `configs/<name>.toml` (struct mode: unknown keys raise).
Resume from any checkpoint via `--checkpoint /path/to/step_N.pt`.

Config system

• Format: .toml, not YAML. Parsed by stdlib tomllib, validated via dataclasses.
• Struct mode: unknown keys raise KeyError at parse time (no silent misspellings).
• CLI overrides: --config.key.path value sets nested fields (e.g., --config.optimizer.lr 2e-4).
• Predefined configs:
  • configs/sft.toml — full-FT or LoRA SFT on 32 prompts
  • configs/dpo.toml — DPO on preference pairs, aligned to SFT golden anchor
  • configs/grpo_stageb.toml — GRPO stage B (32 prompts, 16 generations per prompt)
  • configs/ppo_stage.toml — PPO (DRAFT: unresolved [TBD] fields, not launchable)
  • configs/pretrain.toml — pretraining on raw text (causal LM loss)

Each config names the base model snapshot (model_snapshot), data sources, optimizer, LR schedule, and training hyperparameters.

Rollout backends

GRPO and PPO can sample from the policy via two routes (set rollout.backend in config or script):

Backend: "spool" (default)

• Pure PyTorch: sampling, KV-cache, batch bucketing, reward scoring all in-process.
• No dependencies beyond torch, triton, numpy.
• Throughput: ~200 tokens/s per GPU (A10G, bf16, seq_len=768).
• Used in: stage-B and PPO smoke tests.

Backend: "native"

• Same as spool, but the native sampler is a standalone torch module that can be synced to a separate device (vLLM-like architecture).
• Requires NMO_ROLLOUT_GPU env var to point to a free CUDA device >= world size.
• Untested end-to-end; opt-in via rollout.backend = "native" in config.
• scripts/grpo.py wires it via make_rollout_client; scripts/ppo.py uses spool only.

Backend: "vllm" (extras/)

• High-throughput vLLM worker (100x faster sampling, but requires extras/vllm_rollout/.venv-rollout).
• Optional: not in the main venv. Isolated behind align/rollout/protocol.py.
• Not present on this box; would need a separate venv setup.

Host rules

No GPU jobs, no loading the 7B model, no processes in /tmp, no pkill patterns that self-match. Checkpoints write to /apps/development/rlhf/align/.tmp/ and are atomic (temp file + os.replace). The box crashed once (load 619, 2GB free RAM, no swap); follow /apps/CLAUDE.md rules before any heavy job.

Heavy-job lock: align/.align_heavy_job.lock is per-project. It does NOT exclude legacy nmo or mprm jobs running concurrently.

Tests

Run all CPU tests:

cd /apps/development/rlhf/align
.venv/bin/python -m pytest tests/ -q -m "not gpu and not oracle and not slow"

Subset by file or test name:

.venv/bin/python -m pytest tests/test_trainer.py -q
.venv/bin/python -m pytest tests/test_trainer.py::test_scheduled_lr_is_actually_applied_to_the_optimizer -v

Test markers

• -m "not gpu": skip GPU-only tests (inference gates, 7B model loads, etc.). ~60 tests deselected.
• -m "not oracle": skip parity tests that require the legacy transformers/math_verify stack. Those run under the legacy venv: cd /apps/development/rlhf/nano-mistral-o1 && .venv/bin/python -m pytest /apps/development/rlhf/align/tests/ -m oracle -k test_....
• -m "not slow": skip long-running tests (e.g., end-to-end 7B bench).

Golden regression gate (correctness threshold)

Every completed run writes a golden result to align/eval/golden.py. The regression test (tests/test_golden_regression.py) compares the current policy's greedy decode (via align.inference.Engine) against the recorded output. This is the ultimate correctness gate: if the model output differs, the regression fails and the run is NOT VALID.

The anchor values are deterministic and independent of hardware. A checkpoint's identity (model weights, tokenizer, config) determines its golden output; a 1-bit difference in weights changes the decode.

Anchor checkpoints (recorded output + model snapshot ID):

• SFT: 0.09 win rate vs base on primary contrast (recorded via bench_7b_end_to_end.py).
• DPO: 0.258 on the same contrast.
• PPO: TBD (not yet trained).
• GRPO: TBD (not yet trained).

Run the regression gate:

.venv/bin/python -m pytest tests/test_golden_regression.py -v

If it fails, the discrepancy is printed: expected output vs. observed output. Reconcile before merging.

Quick start

1. Set up the venv (if not done):

   cd /apps/development/rlhf/align
   bash scripts/setup_environment.sh  # uv pip install -e ".[dev]"
   

2. Run SFT (CPU smoke test):

   .venv/bin/python scripts/sft.py --smoke
   

3. Run DPO (CPU smoke test):

   .venv/bin/python scripts/precompute_ref_logprobs.py \
     --checkpoint /path/to/sft.pt \
     --output-dir /apps/development/rlhf/align/.tmp/ref_cache_smoke
   
   .venv/bin/python scripts/dpo.py --smoke
   

4. Run GRPO (stage B, 3 ranks):

   torchrun --nproc_per_node=3 scripts/grpo.py --config configs/grpo_stageb.toml
   

5. Run PPO (3 ranks, smoke test):

   torchrun --nproc_per_node=3 scripts/ppo.py --smoke
   

6. Run tests:

   .venv/bin/python -m pytest tests/ -q -m "not gpu and not oracle and not slow"
   

Unresolved items (as of 2026-10-12)

See docs/DEPENDENCY_POLICY.md for the status section. Key open work:

• Base model snapshot missing on this box: /apps/voxtral/.cache/huggingface/hub/models--mistralai--Mistral-7B-v0.3/.... The tokenizer is vendored, but weights are not. This blocks SFT from the base and the 7B golden regression tests.
• PPO config unfinished: configs/ppo_stage.toml has [TBD pre-reg freeze] and [TBD prm.checkpoint] placeholders. The process reward model does not exist and must be trained.
• Data is not migrated: configs point at /apps/development/rlhf/nano-mistral-o1/data/. The alignment repo depends on the legacy tree.
• High-impact correctness findings: see the integration reports (.tmp/integration_report.md, .tmp/integration_report_2.md). Fixed: scheduler LR applied and logged as applied, grad_norm tripwires live (pre-clip), adamw8bit constructible, GRPO tripwire state checkpointed, DPO ref-cache key bound to checkpoint CONTENT, PPO gamma / lm_head / all-skipped-window / resume-pointer holes closed.
• adamw8bit GPU step time is unmeasured. No GPU step time at 1.8B local elements has been measured, and its checkpoint-save peak (3x the 13.5GB state under the rank-0 gather) is an estimate from the code path. train.per_rank_optimizer_state = true cuts that 3x; it is now validated on 2 CPU/gloo ranks over real FSDP2 shards (bit-exact resume for adamw8bit and offload_adamw, tests/test_trainer_dist.py) and is ON in configs/sft.toml and configs/pretrain.toml. It is also plumbed for DPO and GRPO (ROADMAP, per-rank row). The 2- and 4-rank NCCL runs (docs/MEASUREMENTS.md (a) and (g)) are SFT only; DPO and GRPO multi-rank on GPU are unrun.
• GRPO resume is not wired in scripts/grpo.py. The tripwire state is saved and restored by the Trainer, but a relaunch into an existing run dir still starts from the SFT substrate. Wiring it needs the prompt cursor (and the extra draws of dead-group retries) checkpointed, or a resume re-draws the earliest prompt blocks.
• Existing DPO ref caches must be re-precomputed (the key now carries a content hash of the checkpoint).
• GRPO log-softmax runs in bf16 under FSDP2 (train/infer rho misses the rounding); fixing it needs a GPU memory check at the 2-row micro shape.

Architecture and internals

See docs/ARCHITECTURE.md for:

• Package map (one line per module, what it owns, what it exports).
• Trainer / TrainModule split (how state is managed, how resume works).
• Reward composition vs. loss (DPO dpop+NLL, PPO per-token value head, GRPO batch rewards).
• The three-distributions rule for PPO (policy, reference, reward model).
• Rho-correction contract for GRPO (logprob mismatch between train and inference).
• What moved to extras/ and why (vLLM rollout, parquet ingest).
• Parity oracle design (tokenizer + verifier run under legacy venv, model forward pass checked).

Citation

This codebase is a consolidation of Anthropic alignment infrastructure: nano-mistral-o1, mistral-prm-ppo, and supporting libraries. It is CPU-tested, and GPU runs are recorded in docs/MEASUREMENTS.md (OLMoE-1B-7B SFT on 2 and 4 A10G with FSDP2). The dependency firewall prevents accidental regressions to monolithic frameworks.