Blame for 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:
```bash
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:
```bash
.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:
```bash
.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):
   ```bash
   cd /apps/development/rlhf/align
   bash scripts/setup_environment.sh  # uv pip install -e ".[dev]"
   ```

2. **Run SFT** (CPU smoke test):
   ```bash
   .venv/bin/python scripts/sft.py --smoke
   ```

3. **Run DPO** (CPU smoke test):
   ```bash
   .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):
   ```bash
   torchrun --nproc_per_node=3 scripts/grpo.py --config configs/grpo_stageb.toml
   ```

5. **Run PPO** (3 ranks, smoke test):
   ```bash
   torchrun --nproc_per_node=3 scripts/ppo.py --smoke
   ```

6. **Run tests**:
   ```bash
   .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.