| # Architecture |
|
|
| ## Package map (one line per module, ownership) |
|
|
| ### Config & core |
|
|
| - `align.config`: TOML config parser, struct-mode validation, CLI overrides. Owns config schema, training hyperparameters, stage-specific options. |
| - `align.registry`: Named object factories (schedulers, optimizers, batch sources). No dependencies beyond stdlib + torch. |
| - `align.exceptions`: Shared error types (InvalidCheckpoint, HorizonReached, etc.). |
|
|
| - `align.exit_codes`: Stdlib-only trainer <-> supervisor contract. Exit codes 80 (tripwire HALT), 81 (non-finite / skip-cap HALT), 82 (wall budget spent), 83 (preflight refused), `ExitCodeError` / `PreflightRefused`, the `<run_dir>/HALT.json` marker, and `run_with_exit_codes(main)` used by every stage script. |
|
|
| ### Tokenizer |
|
|
| - `align.tokenizer.core`: Pure-Python BPE (Mistral-7B-v0.3 compatible). Loads tokenizer.json, implements encode/decode, handles special tokens and byte fallback. |
| - `align.tokenizer.__init__`: Exports `Tokenizer` class. Falls back to vendored copy (`assets/tokenizer/mistral-7b-v0.3/`) when the HF snapshot is absent. |
|
|
| ### Model |
|
|
| - `align.model.hf_config`: Reads and writes `config.json` (model architecture, vocab size, rope theta). Typed via dataclass. |
| - `align.model.mistral`: `MistralForCausalLM` (full model), `MistralForValueHead` (critic), `MistralForProcessReward` (per-token scores). FSDP2-compatible modules. |
| - `align.model.attention`: Flash attention v2 wrapper (sdpa, eager fallback). Handles determinism and backward passes. |
| - `align.model.gates`: Token-gating layers (used in some loss functions). |
| - `align.model.moe`: MoE core (no aux loss added; statistics returned). fp32 `Router` (softmax over all experts, then top-k), stable-argsort dropless dispatch with a loop and a `grouped_mm` backend, stacked 3-D expert parameters, `RouterStats` (sufficient statistics of every aux loss), `load_balancing_loss` (Switch/OLMoE and HF conventions), `router_z_loss`, `all_reduce_totals`, `expert_load_metrics` (MaxVio, utilization, dead experts, entropy). |
| - `align.model.olmoe`: OLMoE architecture (`OlmoeConfig`, `OlmoeForCausalLM`): QK-norm over the full projection, 64-expert/top-8 sparse blocks, `forward(..., return_router_stats=True)` returns `(out, [RouterStats])`. Parity with transformers 5.15 to ~1e-6 (oracle tier). |
| - `align.model.olmoe_checkpoint`: HF config reader/writer (refuses scaled rope, bias, tied embeddings) and the HF <-> align key mapping (per-expert hub keys or fused, stacked `gate_up_proj [E,2I,d]` / `down_proj [E,d,I]`). |
| - `align.model.checkpoint`: Loads/saves checkpoints (safetensors, sharded for FSDP2). Handles tokenizer fallback, atomic writes, state_dict exports. |
| |
| ### Training infrastructure |
| |
| - `align.train.trainer`: `Trainer` class (main loop). Owns batch iteration, checkpoint save/resume, logging, tripwire gates. Does NOT own the loss computation; delegates to `TrainModule`. |
| - `align.train.train_module`: `TrainModule` abstract class. Subclasses implement `forward()` (loss), `backward()`, `post_backward_hook()`. Owns gradient clipping, LR scheduling application, optimizer step. |
| - `align.train.moe_module`: `MoETrainModule`, the MoE-aware TrainModule. Wraps the model so the loss adapter is unchanged, collects `RouterStats`, accumulates the Demons Alg. 1 window buffer, adds load-balance + z-loss, logs `moe/*` telemetry. `MoEAuxConfig` lives in `align.config` (`[moe_aux]`). |
| - `align.train.ckpt_format`: `WeightFormat` seam (`DenseFormat`, `OlmoeFormat`): which converter writes/reads the weight file and which tokenizer, if any, rides along. The Trainer's checkpoint orchestration is family-independent. |
| - `align.train.sft_loss`: SFT (cross-entropy). Owns label construction and vocabulary. |
| - `align.train.dpo_loss`: DPO (dpop + NLL). Owns preference pair handling, reference logprob attachment, loss symmetry. |
| - `align.train.dpo_stage.py`: `DpoStageConfig`, stage-specific gates and resumption rules. |
| - `align.train.offload_adamw`: AdamW with CPU offload of the optimizer state (fp32 master + moments on host, ~12 B/param). Handles FSDP2 DTensors, distributed grad norm, staging caps. Also hosts `make_optimizer`, the single place a config's `optimizer.name` becomes an optimizer (and the banner string). |
| - `align.train.adamw8bit` / `align.train.quant8`: torch-only block-wise 8-bit AdamW (`optimizer.name = "adamw8bit"`, what every 7.25B full-FT config names). uint8 dynamic-codebook moments + fp32 absmax per 2048-block, state on the device (2.004 B/param, zero host RSS), stochastic-rounding write-back into bf16 params, local-shard stepping under FSDP2, global pre-clip grad norm, bit-exact `state_dict` resume (refuses a different world size). |
| - `align.train.fsdp`: FSDP2 setup (fully sharded data parallel). Owns mixed precision policy (bfloat16 compute, float32 reduce), parameter sharding, process groups. |
| - `align.train.probes`: Telemetry (weight movement, token counts, entropy, skip metrics). `TrunkMovementProbe` tracks only `requires_grad` parameters. |
| - `align.train.stage_gates`: Data readiness gates (shard counts, manifest, coordinate reference caches). |
| - `align.train.stage_smoke`: Smoke-test harness (small model, synthetic data, mini runs). |
| - `align.train.ref_cache`: Reference policy logprob cache (DPO). Owns hash functions, atomic writes, cache key validation. |
|
|
| ### Reinforcement Learning (GRPO & rewards) |
|
|
| - `align.rl.grpo_module`: `GrpoTrainModule` (subclass of `TrainModule`). Owns GRPO loss, advantage centering, batch scaling, clip range, entropy baseline. |
| - `align.rl.grpo_loss`: GRPO loss function (grouped relative policy optimization). Owns batch weighting, reward integration, KL penalty. |
| - `align.rl.advantages`: `compute_generalized_advantages` (GAE). Shared with PPO. |
| - `align.rl.batch_source`: `GrpoBatchSource` (batch iterator, handles sampling seed, curriculum). Owns generation queuing and batch shuffling. |
| - `align.rl.rewards`: `CompositeReward` (outcome + KL + length penalties). Separate from loss to allow reward reweighting. |
| - `align.rl.tripwire`: Health monitors (`CollapseCollapseReward`, `EntropyTripwire`, `RatioTripwire`). Armed with thresholds, fire on deviation, halt the run. |
| - `align.rl.canary`: Early-stopping logic (threshold on reward, loss divergence). |
| - `align.rl.stage_config`: `RolloutShape`, `GrpoStageConfig`. Owns rollout backend selection, world-size validation. |
| - `align.rl.format_instrument`: Format string templates for logging and telemetry. |
| - `align.rl.ops`: Distributed reduce operations (all-reduce, broadcast). |
| - `align.rl.types`: Shared dataclasses (`RolloutBatch`, `AdvantageData`, `PolicyLogprobs`). |
|
|
| ### PPO (Process reward models) |
|
|
| - `align.ppo.trainer`: `PpoTrainModule` (subclass of `TrainModule`). Owns PPO loss (clipped surrogate + entropy), GAE, value head training. |
| - `align.ppo.loss`: PPO loss computation. Owns ratio clipping, advantage centering. |
| - `align.ppo.gae`: Generalized advantage estimation (shared with GRPO via `align.rl.advantages`). |
| - `align.ppo.value_head`: Value head architecture options (linear, MLP, per-token or outcome). |
| - `align.ppo.lora`: LoRA adapters for the PPO value head. Own implementation on `nn.Linear` (no peft). Handles FSDP2-sharded weight distribution. |
| - `align.ppo.rewards`: Process reward module (per-token scores). Integrates with the loss. |
| - `align.ppo.health`: PPO-specific health monitors (`RatioTripwire` with skip tracking). |
| - `align.ppo.rollout`: PPO rollout generator (samples from the policy, collects trajectories). |
| - `align.ppo.stage_config`: `PPOStageConfig`. Owns reward model selection, KL settings, value head options. |
|
|
| ### Inference |
|
|
| - `align.inference.engine`: `Engine` class. KV-cache sampler, greedy/top-k decoding, batch bucketing. Deterministic, no randomness. |
| - `align.inference.bucketed`: Batching logic for variable-length sequences (padding, masking). |
| - `align.inference.receipts`: Checkpoint identity (content hash, manifest hash). Used for golden regression and cache validation. |
|
|
| ### Rollout (policy sampling) |
|
|
| - `align.rollout.protocol`: Shared dataclasses (`RolloutRequest`, `RolloutResult`, `RolloutMetrics`). Interface between trainer and sampler. |
| - `align.rollout.client`: `make_rollout_client()` factory. Returns a sampler based on `config.rollout.backend`. |
| - `align.rollout.native`: Native spool sampler. Pure torch KV-cache, generation loop, reward scoring. Runs in-process or on a remote device (via DTensor send/recv). |
|
|
| ### Data pipeline |
|
|
| - `align.data.hub`: HF hub streaming (JSONL, parquet row groups). HTTP-based, no `datasets` dependency. |
| - `align.data.schema`: Shared dataclasses for prompts, pairs, outcomes (`Prompt`, `PreferencePair`, `OutcomeExample`). |
| - `align.data.encoding`: Text-to-token encoding. Owns tokenization, batching, padding. |
| - `align.data.build`: Dataset builders (SFT, DPO, GRPO, PPO). Owns corpus merging, filtering, curation. |
| - `align.data.sft_corpus`: SFT data sources (math reasoning, instruction following). |
| - `align.data.prm800k_pairs`: PRM800K preference pairs (step-level annotations). |
| - `align.data.argilla`: Argilla task data (preference annotations). |
| - `align.data.stageb_sources`: Stage B outcome data (MATH-500, PRM800K). |
| - `align.data.pretrain_corpus`: Pretraining data (raw text). |
| - `align.data.store`: Shard storage and manifest (local file-based data storage). |
| - `align.data.loader`: Batch iterator (handles streaming, sharding, resumption). Dense `scripts/sft.py` and `scripts/pretrain.py` build their sampler with `build_training_sampler`, which shards by rank via `DistributedBatchSampler` (fixed in 654c20c). Before the fix every rank received the same batch stream, so a 4-rank run averaged four identical gradients (`tests/test_training_sampler_sharding.py`). |
| - `align.data.outcome_agreement`: Outcome agreement metric (verifier consensus). |
| - `align.data.receiptss`: Data fingerprinting (used for cache validation). |
| - `align.data.evalsets`: Named evaluation datasets (golden regression, benchmarks). |
| - `align.data.ingest_parquet`: One-time parquet-to-shard converter (whitelisted for `pyarrow`). |
|
|
| ### Evaluation |
|
|
| - `align.eval.golden`: Golden regression checkpoints and their expected decode outputs. Owns anchor values (0.09 for SFT, 0.258 for DPO). |
| - `align.eval.verifier`: Integration wrapper around `align.verify` (math answer checking). Metric for outcome agreement. |
| - `align.eval.harness`: Evaluation main loop (golden regression, benchmark runs). |
| - `align.eval.metrics`: Utility functions (win rate, accuracy, kl divergence). |
| - `align.eval.roundtrip`: Checkpoint round-trip testing (save -> load -> forward pass). |
|
|
| ### Answer verification |
|
|
| - `align.verify.extract`: Boxed-answer extraction (finds `\boxed{...}` in LaTeX). Raises `NotImplementedError` if parsing fails. |
| - `align.verify.normalize`: Answer normalization (convert to rational/decimal/interval/set). |
| - `align.verify.equiv`: Equivalence checking (exact rational equality, interval overlap, set equality). |
| - `align.verify.expr`: Small-scale expression evaluator (rationalize `\frac{1}{2}` etc.). Intentionally limited scope. |
| - `align.verify.__init__`: Exports `is_readable()` (can this answer be checked?), `are_equivalent()` (do two answers match?). |
| - `align.verify.lean` (package): strict Lean 4 proof verifier, the formal-proof reward (ROADMAP P4-1). `strict.py` (statement and proof assembly, lexical and semantic checks), `repl_pool.py` (REPL worker pool with timeouts, memory watchdog, process-group kill), `reward.py` (reward hook, reward 0 for every failure class), `toolchain.py` (pinned Lean v4.9.0-rc1, Mathlib d1d1e4b7, child environment). Design and measurements: `docs/LEAN_VERIFIER.md`. Not an OS sandbox: the REPL runs as the same user (see that doc, "No OS sandbox"). |
|
|
| ### Utils |
|
|
| - `align.utils.dist`: Process group / rank utilities. Owns rank 0 sync, all-reduce, gloo group (control plane for rollout). |
| - `align.utils.guard`: Resource guards (heavy-job lock, preflight checks). Enforces the host rules; every refusal is a `PreflightRefused` (exit 83). |
| - `align.utils.supervisor` + `scripts/supervise.py`: Bounded-retry launch supervisor (stdlib only). See "Control plane" below. |
| - `align.utils.tracking`: Metrics logging (JSONL history, optional stdout). No wandb dependency. |
|
|
| ## Trainer / TrainModule split (state management) |
|
|
| **`Trainer` (align/train/trainer.py):** The main loop. |
|
|
| - **Owns:** batch iteration, checkpoint save/resume, terminal conditions (data exhausted, horizon, tripwire), logging. |
| - **Does NOT own:** loss computation, backward pass, optimizer step, LR scheduling application. |
| - **Interface to TrainModule:** calls `module.forward(batch) -> loss`, `module.backward()`, `module.post_backward_hook(...)`, `module.optim_step() -> applied_bool`. |
| - **State persisted in checkpoint.pt:** `stream` (batch resumption), `scheduler`, `optimizer_per_rank`, `micro_step`, `optimizer_step`, `seed`. |
|
|
| **`TrainModule` (align/train/train_module.py):** Computation, not state. |
| |
| - **Owns:** forward pass (model, loss function), backward pass (pytorch autograd), gradient clipping, LR schedule application, optimizer step. |
| - **Does NOT own:** batch iteration, checkpoint save/resume. |
| - **Subclasses:** `SftTrainModule`, `DpoTrainModule`, `GrpoTrainModule`, `PpoTrainModule`. |
| - **Resumption:** built fresh from config + model checkpoint at every launch. `Trainer.resume_from()` restores the scheduler and optimizer state into the new module. |
| |
| **Resume flow:** |
| 1. Load model weights through `Trainer.weight_format` (`align.train.ckpt_format`: `DenseFormat` by default, `OlmoeFormat` for an MoE module; the family is injected, never sniffed from key names). |
| 2. Build TrainModule from config (initializes optimizer, scheduler). |
| 3. Call `Trainer.resume_from(checkpoint_dir)` — restores `scheduler.state_dict()`, `optimizer_state_dict()`, `micro_step`, `seed`. |
| 4. Stream resumes from the saved `offset` and `consumed_total`. |
| 5. `TrainModule.load_extra_state(resume.pt["extra_state"])` restores module-private state (GRPO: the collapse tripwire's baseline window and latch). The base `TrainModule.extra_state()` returns `{}`; subclasses override the pair. Only tensors/dicts/lists/primitives (the file is read with `weights_only=True`). |
| 6. The resumed checkpoint is registered in `Trainer.checkpoints_written` (a per-process list), which launch scripts read for the rollout engine's weights and keep-N pruning reads for deletion. |
| 7. Optimizer state is either one gathered blob in `resume.pt` (default) or one `optimizer_rank{r}.pt` per rank (`train.per_rank_optimizer_state = true`); the checkpoint's own manifest decides how a resume reads it, and per-rank files refuse a different rank count. |
|
|
| This design avoids serializing the entire TrainModule (which includes the model, the optimizer, and all their state), and instead captures only the restorable pieces. |
|
|
| ## Control plane (G0): decisions that every rank must take together |
|
|
| One rule drives the design. A rank-local fact (a signal, a wall clock, a tripwire verdict, a non-finite loss) becomes an action only after a MAX all-reduce at a program point every rank reaches. The failure being prevented is always the same: one rank acts alone while its peers sit in a collective, so nothing crashes and the run hangs until the NCCL watchdog fires (the 5-hour stall of 2026-08-12). |
|
|
| | Mechanism | Where | Decision point | Failure it prevents | |
| |---|---|---|---| |
| | Coordinated stop (SIGTERM/SIGINT, `train.max_wall_seconds`) | `Trainer._boundary_control`, `_StopSignals` | Optimizer-step boundary, one MAX reduce of `[signalled, wall, nonfinite, skipped]` | Python's default SIGTERM killing a rank mid-save; a stop mid-accumulation (a checkpoint holds no gradients). A stop always writes a checkpoint and wins over `max_micro_steps`. | |
| | HALT tripwires and arm-time self-test | `Trainer._agree_or_raise` | Reduced verdict, then the same `HaltError` on every rank (rank-ordered gather of messages) | A lone raiser unwinding while peers wait in the next micro-step's all-gather. | |
| | Skip cap | `Trainer._record_step_health` | Computed from reduced flags: consecutive bad steps > `max_consecutive_skipped_steps` (3) or bad fraction > `max_skipped_fraction` (5%, evaluated after 100 boundaries) | A run that skips forever (schedule stalls, loss "healthy") training nothing. A bad step is one the grad gate skipped, or any micro-batch with a non-finite loss; a step that returns False by design (PPO critic warmup) does not count. | |
| | Checkpoint durability | `fsync_tree`, `_publish` | Files, then directories, then the rename, then the parent directory | A rename that reaches the journal before the data (zero-length files that pass `isdir`). | |
| | Checkpoint completeness | `checkpoint_problem`, `latest_checkpoint` | Weights header extent, `resume.pt` loop state, every per-rank optimizer file | Resuming (or the supervisor counting progress) from a truncated newest directory. | |
| | RNG, pins, skip stats in `resume.pt` | `capture_rng_state`, `_resume_payload` | Per-rank gather at save; restored when the world size matches | A resume that repeats or forgets dropout/sampling streams, forgives a bad run, or loses the first-alarm pin. | |
| | Process-group timeout + flight recorder | `utils.dist.init_distributed(run_dir=...)` | `ALIGN_PG_TIMEOUT_S` (default 1800); NCCL trace dump under `<run_dir>/nccl_trace` | A hang leaving no evidence; a dump landing in `/tmp`. | |
|
|
| **Exit codes and the halt marker (`align.exit_codes`).** HALT conditions are verdicts about the *run*, not the host, and must not be relaunched: 80 tripwire, 81 skip cap / non-finite, 82 wall budget spent (after the checkpoint), 83 preflight refused. `HaltError` (RuntimeError subclass, `exit_code` attribute) carries the code; every stage script runs `main` through `run_with_exit_codes`, and returns `Trainer.stop_exit_code` after a coordinated stop instead of falling into the post-train eval (82 for wall clock, 128+signum for a signal). Because `torchrun` collapses *every* worker exit code to 1 (measured on torch 2.11), the Trainer also writes `<run_dir>/HALT.json` (rank 0, tmp + fsync + rename) before raising; the supervisor treats a marker written after the attempt started as the verdict whatever the launcher exit code said. |
| |
| **Launch supervisor (`scripts/supervise.py`, `align.utils.supervisor`).** Relaunch iff: exit != 0, not a HALT code, no fresh HALT marker, retries < 2, progress made during the attempt (`history.jsonl` mtime advanced AND the newest *complete* `checkpoints/micro_N` increased, validated with `checkpoint_problem`), and the host preflight passes after the backoff. The child runs in its own process group, which is killed and confirmed empty before any relaunch. The preflight on the first launch belongs to the job itself. |
| |
| **Not covered by the control plane:** the PPO critic-only warmup deliberately returns an unapplied step (not a skip); `offload_adamw` does not clip (P0-3), so the default `max_grad_norm = 1.0` is a no-op there; the GPU-idle/temperature check of P0-11 is not implemented; NCCL status: on 2 A10G ranks, a SIGTERM to one rank stops both at the same boundary and resumes bit-exactly, and a NaN on one rank halts both without a hang (MEASUREMENTS (a), (b)). On 4 A10G ranks (OLMoE-1B-7B, FSDP2), a SIGTERM to one rank with supervisor relaunch resumes with loss within the run-to-run floor (MEASUREMENTS (g)). Still open: `kill -9` during a save, and NCCL behaviour under a real hang. |
| |
| ## Mixture of experts (Track M) |
| |
| **Layering.** `model/moe.py` and `model/olmoe.py` own the architecture and *return* statistics; they add no loss. `train/moe_module.py` owns the objective. `train/ckpt_format.py` owns the file format. Nothing in the Trainer knows what a router is. |
|
|
| **Routing.** fp32 router, softmax over all experts, then top-k (not top-k then softmax), no renormalisation for OLMoE-1B-7B. Dispatch is dropless: a stable argsort of the flat expert assignment gives contiguous per-expert segments; a Python loop or `F.grouped_mm` computes them. Experts are stacked 3-D parameters (`gate_up_proj [E,2I,d]`, `down_proj [E,d,I]`), sharded on dim 0 by FSDP2, so `E` must divide the world size. The combine is fixed-order (unsort, then sum in a fixed order, custom gather backward). An earlier bf16 `index_add_` combine was nondeterministic: 5 of 5 runs on identical inputs gave different top-8 expert sets. Fixed in `align/model/moe.py` (MEASUREMENTS, Track M (f)). |
|
|
| **Aux losses on the seam (`MoETrainModule`).** |
|
|
| 1. The loss adapter is unchanged. The module hands it a proxy model that adds `return_router_stats=True` (only when gradients are enabled, so a `no_grad` reference forward is not counted), unwraps `(out, stats)`, and keeps the stats. |
| 2. *Window buffer (Demons Alg. 1).* Per micro-batch one packed float64 all-reduce of detached `[counts, tokens, prob_sum, z_sum, entropy_sum]`; summed into a buffer that lives for one accumulation window and is reset at `optim_step()`/`zero_grads()`. |
| 3. *Global-scope load balance.* `f_i` comes from the buffer; `P_i` stays each rank's differentiable `prob_sum`. Micro-batch `m` is scaled by `world * m / n_buf`, which is exactly `world * accum / N_window` once `m = accum`, so after the Trainer's `/grad_accum` and the gradient reduction's mean the window's gradient is the gradient of one loss over the global batch. The approximation is confined to `f_i` of micro-batches 1..accum-1; at the last micro-batch `f_i` is exact. `lb_scope = "micro"` keeps the per-rank formulation for ablation. |
| 4. *Router z-loss* uses the same scale, with the global token count as denominator. |
| 5. *Defaults.* `[moe_aux]`: load balance 0.01 and z-loss 0.001 when the stage is `pretrain` or `proxy`; both 0 for every post-training stage (a pretrained router already balances the pretraining distribution, and narrow post-training data legitimately skews expert use). `None` means "stage default"; any number overrides. |
| 6. *Telemetry* (always on unless disabled): `moe/lb_loss`, `moe/z_loss` (exact global-window values), `moe/max_vio`, `moe/max_vio_worst`, `moe/utilization`, `moe/utilization_worst`, `moe/frac_dead_worst`, `moe/load_entropy`, `moe/router_entropy`, plus per-layer MaxVio/utilization with `per_layer_metrics`. Values are those of the last completed window (NaN before the first), so every history row has the same columns. |
| 7. The returned `LossOutput.loss` is the CE alone (comparable with dense runs); a non-finite aux term turns it into NaN so the Trainer's non-finite accounting still sees it. |
|
|
| **Limits.** A decoupled-backward adapter (DPO) cannot carry an aux loss (raises; telemetry still works). GRPO/PPO modules override `train_batch` and are not MoE-aware. Measured on 4x A10G with FSDP2 and real OLMoE-1B-7B weights (MEASUREMENTS, Track M (g)): SFT with the MoE telemetry on. Not measured: `grouped_mm` end to end under FSDP, expert parallelism (only its all-to-all volume was timed, (h)), and MoE under DPO, GRPO or PPO. Router fp32 optimizer state (P1-8), the collapse tripwire (P1-9) and `ModelConfig.moe` (P1-1) are not done. |
|
|
| **Checkpoints.** `OlmoeFormat` writes HF-loadable weights + `config.json` (bf16 weights, fp32 optimizer master is authoritative, as for dense). It copies a tokenizer only from an explicit `tokenizer_dir`: the dense path's pinned Mistral files would be wrong for a 50304-token model. A checkpoint written by align loads in transformers with equal logits (oracle tier). |
|
|
| ## Reward composition vs. loss (the seam) |
|
|
| Three separate computations, three separate containers: |
|
|
| ### DPO path (`align.train.dpo_loss.py`) |
| |
| - **Loss:** `dpo_loss(policy_logps, reference_logps, pairs) -> scalar loss`. |
| - Formula: `-log(sigmoid(beta * (log(pi/pi_ref) - log(pi_c/pi_ref_c))))` (dpop) + NLL on the chosen arm. |
| - Owns preference pair handling, reference logprob integration, label masking. |
| - No "reward" object; the beta scaling **is** the reward weighting. |
| - **Testing:** `tests/test_dpo_loss.py`, `tests/test_dpo_stage.py`. |
|
|
| ### GRPO path (`align.rl.grpo_module.py` + `align.rl.rewards.py`) |
| |
| - **Rewards:** `CompositeReward(outcome, kl_penalty, overlong_penalty) -> tensor of per-sample scores`. |
| - Owns reward composition, penalty scaling, validation. |
| - Does NOT apply loss weighting; that is the module's job. |
| - **Loss:** `grpo_loss(policy_logps, reference_logps, advantages, ...) -> scalar loss`. |
| - Formula: `-(rho * advantages).mean() - beta * kl_penalty`. |
| - Clips advantages to a range, centers around zero, scales by world size (distributed gradient averaging correction). |
| - Owns advantage centering, clipping, batch scaling. |
| - **Module:** `GrpoTrainModule.forward(batch)` calls both, combines them. |
| - **Testing:** `tests/test_grpo_module.py`, `tests/test_grpo_loss.py`, `tests/test_rewards.py`. |
|
|
| ### PPO path (`align.ppo.trainer.py` + `align.ppo.loss.py`) |
|
|
| - **Rewards:** Per-token value targets (via GAE on per-token scores) or outcome rewards. The value head predicts these. |
| - **Loss:** PPO surrogate loss (clipped ratio * advantage) + value head MSE + entropy bonus. |
| - Owns ratio clipping (default 0.2), advantage centering. |
| - **Module:** `PpoTrainModule.forward(batch)` computes advantages via GAE, then the loss. |
| - **Testing:** `tests/test_ppo.py`, `tests/test_ppo_loss.py`. |
|
|
| **The seam:** rewards are computed **outside** the loss, so they can be reweighted (e.g., curriculum) without recomputing logprobs. Inside the loss, rewards are summarized (summed, averaged, gated) to a single scalar or per-token target. |
|
|
| ## Three-distributions rule for PPO |
|
|
| PPO requires three policy distributions to be in sync: |
|
|
| 1. **Reference policy:** frozen weights (SFT checkpoint or a prior PPO checkpoint). |
| 2. **Training policy:** weights being updated (the main model). |
| 3. **Rollout policy:** the one generating data (should match training policy). |
|
|
| **Invariant:** `rollout_policy == training_policy`. If they diverge, the KL penalty is computed against stale data. |
|
|
| **Implementation:** |
| - The rollout samples from the latest `training_policy` weights. |
| - The reference policy is loaded once and frozen. |
| - If the rollout is on a remote device (native backend), weights are synced after each checkpoint via `align.rollout.checkpoint_weight_loader`. |
|
|
| **PPO loss math depends on this:** the ratio `pi / pi_ref` uses the current step's logprobs; a stale reference introduces a train-infer mismatch. |
|
|
| ## Rho-correction contract for GRPO |
|
|
| In GRPO, the advantage is computed as `A = r - beta * log(pi / pi_ref)`. But `pi` and `pi_ref` can be stale or inexact due to: |
|
|
| 1. **Batch async:** the rollout samples from step N, training happens on step N+1. |
| 2. **Precision loss:** logits are bf16 under FSDP2 mixed precision; log_softmax is inexact. |
| 3. **Implementation divergence:** the rollout uses a custom sampler; training uses the model. |
| |
| The `rho` correction factor measures this: `rho = exp(log_pi_train - log_pi_rollout)`. If they match, rho = 1. If the training policy is slightly better, rho > 1 (advantage is inflated, we correct down). If training is worse, rho < 1. |
| |
| **Contract:** The rollout reports logprobs as `log p(action | state)` (temperature-scaled if T > 1). Training recomputes them the same way. If logprobs match, advantages are exact; if they diverge, rho absorbs the error. |
| |
| **Known divergences that rho cannot fix:** |
| - Determinism: if the sampler uses torch.multinomial and training uses log_softmax, bit-level differences are expected. |
| - Sequence-level effects: if the sampler reorders prompts, state dependencies shift. |
| - Model architecture: if the sampler runs an old checkpoint, rho correction works only if the two policies are close in KL. |
|
|
| **Testing:** `tests/test_grpo_module.py::test_rho_correction_recovers_exact_advantage_when_logprobs_match`. |
|
|
| ## What moved to `extras/` and why |
|
|
| - **`extras/vllm_rollout/`:** vLLM worker for high-throughput sampling (100x faster than native torch sampler). |
| - **Why it is optional:** introduces ~20 heavy dependencies (vllm, ray, pydantic, etc.); most runs don't need this throughput. |
| - **How it works:** separate venv (`.venv-rollout`), separate process, talks to trainer via `align.rollout.protocol.py` (shared dataclasses only). |
| - **When to use:** multi-node rollout, >1000 tokens/s target. |
| |
| - **`align/data/ingest_parquet.py`:** one-time converter from parquet (GSM8K, MATH, argilla) to the repo's shard format. |
| - **Why it is optional:** `pyarrow` is heavy (C++ build, wheels may fail). Most development uses pre-ingested shards. |
| - **How it works:** whitelisted in firewall (`path`, not import), lazy import (`importorskip`), does not affect runtime. |
| - **When to use:** when onboarding a new dataset source. |
|
|
| ## Parity oracle design |
|
|
| Testing against the legacy stack (`nano-mistral-o1` with transformers 5.15, math_verify) requires: |
| |
| 1. **Two venvs:** legacy (with transformers) in `/apps/development/rlhf/nano-mistral-o1/.venv`, and align (torch-only) in `/apps/development/rlhf/align/.venv`. |
| 2. **Oracle tests:** marked with `@pytest.mark.oracle` in the test file. |
| 3. **Running oracle tests:** |
| ```bash |
| cd /apps/development/rlhf/nano-mistral-o1 |
| PYTHONPATH=/apps/development/rlhf/align .venv/bin/python -m pytest /apps/development/rlhf/align/tests/ -m oracle |
| ``` |
| This runs the align test file **in the legacy venv context**, so transformers and math_verify are available. |
|
|
| 4. **Test structure:** oracle tests have two branches: |
| - **Legacy code path:** `oracle = transformers.AutoTokenizer.from_pretrained(...)`. |
| - **Align code path:** `align_result = align.tokenizer.Tokenizer.from_snapshot(...)`. |
| - **Assert:** `oracle == align_result` or `oracle.logits ≈ align_result.logits`. |
|
|
| 5. **Coverage:** Parity is checked on: |
| - **Tokenizer:** 365 texts x2 directions (encode, decode) + 176 id lists x2 + vocab maps. Oracle: 100% identical. |
| - **Verifier:** 400 constructed pairs (gold standard answers vs. predicted). Oracle: 99.75% agreement (1 disagreement where align is correct and legacy is wrong). |
| - **Model forward:** Tiny models, sdpa vs. eager attention, bfloat16. Oracle: identical logits. |
| - **Checkpoint load:** Transformers-written checkpoint loads via align. Oracle: weights match exactly. |
|
|
| 6. **Why oracle is separate:** transformers introduces 100+ dependencies (datasets, huggingface_hub, tokenizers, tqdm, ...). Align must not import any of them at runtime. Oracle is run on-demand, isolated, and never shipped. |
| |
| ## Summary |
| |
| The architecture cleanly separates concerns: |
| |
| - **Config** → **Model** → **Trainer** → **TrainModule** (computation) + **Batch source** (data). |
| - **Loss** is pluggable (SFT, DPO, GRPO, PPO all have different formulas). |
| - **Reward** is separate from loss (enables curriculum, reward analysis). |
| - **Checkpoint** is atomic (temp file + os.replace, never partial). |
| - **Resume** is lossless (every micro-step is recoverable; no data is dropped mid-accumulation). |
| - **Parity** is tested via oracle (legacy venv, no import overhead on align). |
| |
| This design makes it possible to: |
| - Swap loss functions without touching batch sources. |
| - Rerun old checkpoints and get the same output (deterministic inference engine). |
| - Add new reward signals without recomputing logprobs. |
| - Keep the dependency footprint minimal (torch + triton + numpy + safetensors). |
| |
| ## Expert parallelism |
| |
| `align/model/moe_ep.py` (opt-in; default MoE path is bitwise unchanged, hooks in `moe.py` are `ep=None`). |
|
|
| **Why.** Measured on this box: FSDP over SHM NCCL (~4 GB/s, no P2P) spends 13.55 s of a 14.8 s OLMoE micro-step in |
| collectives, because FSDP all-gathers all 64 experts of a layer (805 MB bf16, forward and again in backward) to use the |
| ~1/64 each token hits. EP keeps the weights where they are and moves the tokens. |
|
|
| **Pipeline per layer.** Local routing (router replicated, so `RouterStats` and the load-balancing loss are identical to |
| the non-EP model; `MoEAux` all-reduces over the data-parallel group as before) -> local stable sort by global expert id -> |
| equal-split all-to-all of the `[ep, E/ep]` count matrix (one host sync) -> row all-to-all (dropless, no capacity or |
| padding) -> fixed regroup to expert-major, source-rank-ascending order -> local experts through the same |
| `run_experts_backend` -> inverse permutation -> reverse all-to-all -> the unchanged weight-scale and fixed-order slot sum. |
| No collective accumulates anything, so no reduction order is inherited from NCCL. The receive order equals the order a |
| single process sees when it stable-sorts the rank-ordered concatenation of all tokens, which is why equality is bitwise |
| in fp64. Backward is the reverse all-to-all with swapped splits (`_AllToAllRows`, exact: a row move has no arithmetic). |
|
|
| **Layout.** World `W = R x ep`. EP groups are `ep` adjacent ranks; the `R` replicas of one expert shard form a replicate |
| group. Non-expert parameters (attention, norms, embeddings, lm_head, router) are reduced or FSDP-sharded over all `W` |
| ranks exactly as today. Expert shards stay OUT of that FSDP unit. On the 4-GPU box `W = ep = 4`, `R = 1`: experts are never |
| gathered. Order of operations: build model -> `apply_expert_parallel` -> optimizer -> (FSDP the rest). |
| |
| **Gradient scale (the silent one).** Replicated parameters get `(1/W) * sum` over ranks. An expert shard's gradient |
| arrives already SUMMED over its EP group (the dispatch backward delivers every contribution). `finalize_expert_grads` |
| sums over the replicate group (when `R > 1`) and multiplies by `1/W`; skipping it trains the experts at `ep` times the |
| learning rate with no error. The FSDP-of-experts alternative (`fully_shard(block.mlp.experts, mesh=replicate_mesh)`) |
| must set `set_gradient_divide_factor(W)` to land on the same scale; it is specified, not implemented, here. |
| `MoETrainModule(ep_groups=...)` calls it in `optim_step`. |
|
|
| **Clipping.** Grad-norm clipping must see all expert shards: `global_grad_norm` all-reduces the shard squares over the EP |
| group. An optimizer that clips by its own norm is refused (`max_grad_norm` must be None; pass `ep_max_grad_norm`), |
| otherwise each rank clips by a different coefficient and replicated parameters drift apart. |
|
|
| **Deadlocks closed.** A rank that receives zero rows must still issue every forward and backward all-to-all |
| (`_no_rows` ties the empty output to the local weights, which also yields exact-zero weight grads instead of `None`); |
| `combine_experts` and `_DispatchGather.backward` take the width explicitly so zero-row tensors pass. A desynchronised |
| collective stream is reported by a split-size sanity check rather than an allocator error. |
|
|
| **Checkpoints.** Expert tensors are sharded on dim 0. `MoETrainModule.full_model_state_dict` gathers them (collective over |
| the EP group; like the base it returns the state on rank 0), `load_model_state_dict` slices, so the file is the one a |
| non-EP run writes. |
|
|
| **Bytes per micro-step per rank, OLMoE-1B-7B (d=2048, I=1024, E=64, k=8, 16 layers, bf16, ep=W=4), analytic:** |
|
|
| | tokens/micro-step/rank | EP all-to-all (4 passes) | FSDP expert all-gather + reduce-scatter (3 collectives) | ratio | |
| |---|---|---|---| |
| | 1024 | 1.50 GiB | 27.0 GiB | 18.0x | |
| | 4096 | 6.00 GiB | 27.0 GiB | 4.5x | |
|
|
| EP traffic scales with tokens, FSDP traffic with parameters; the crossover is ~18.4k tokens per rank per micro-step. |
| Counted-vs-analytic is checked in `tests/test_moe_ep.py` (exact on balanced routing; within 10% on a real OLMoE-shaped layer |
| with a random router). The number to trust for wall time is the on-GPU timing in the GPU validation plan. |
|
|
| **Limits.** Validated with gloo on CPU only; NCCL all-to-all, bf16 wire dtype and DTensor norm paths are specified, not run. |
| `E` must divide by `ep`. Load imbalance (one rank receiving most rows) is dropless-correct but makes the slowest rank set the |
| step time; EP does not fix routing skew. |
|
|
