code/docs/ARCHITECTURE.md

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:

   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.