Blame for code/docs/REFERENCE_GAP_ANALYSIS.md

# align: reference-repo gap analysis

Scope: `/apps/development/rlhf/align` at HEAD 1d1c8ee (CPU suite 1107 passed / 13 skipped), compared against the read-only
repos in `/apps/development/rlhf/reference-repos/`. Date of analysis: 2026-10-12.

Method and honesty rules. Every status below comes from static reading and grep. No GPU job, no 7B load, no model download
and no test run was done for this report. Where a number is my arithmetic or an assumption it is labelled. Items a
verification pass corrected are already folded in. "Not found" means a search of the stated place came back empty.
Paths are relative to `/apps/development/rlhf/align/` unless they start with `reference-repos/`.

Hardware envelope used throughout: one node, 4x A10G (SM86, 23 GB usable, no FP8, no NVLink; `nvidia-smi topo -m`
shows every pair as PHB, i.e. PCIe through the host bridge), 186 GB RAM, no swap, shared box. Host rules in
`/apps/CLAUDE.md` (45 GB MemAvailable preflight, one heavy job at a time, nothing on `/` or `/tmp`) apply to every
validation step named here.

---

## 1. Verdicts

### 1.1 Is align robust enough for real training today?

**Conditionally yes for what it already does; no for multi-day, pretraining-scale runs.**

What is strong (confirmed in code):
- Lossless resume of weights, fp32 master/moments, scheduler, batch-stream position and sampler state
  (`align/train/trainer.py:626-746`, `:888-1019`); `.partial` staging plus `os.rename` publication and startup partial
  cleanup (`trainer.py:696-728`, `:869-876`).
- Host-safety layer stronger than OLMo-core's: RAM/load/disk-queue preflight and exclusive flock
  (`align/utils/guard.py:377-478`), stage gate ladder (`align/train/stage_gates.py`), non-reentrant activation checkpointing,
  FSDP2 bf16/fp32-reduce (`align/train/fsdp.py:61-123`), chunked fused CE and DPO that never materialise full logits
  (`align/kernels/fused_ce.py:234`, `fused_dpo.py:526`), strong GRPO train/infer mismatch handling
  (`align/rl/grpo_loss.py:107-164`).

What blocks "real run" status (all are concrete, mostly small):
1. **No graceful stop.** No SIGTERM/SIGINT handler in the trainer; `scripts/sft.py:143-147` documents a torchrun SIGTERM
   landing mid-save. No wall-clock limit.
2. **Silent NaN spin.** A non-finite gradient is skipped (`offload_adamw.py:193-201`), the scheduler does not advance on
   skips (`trainer.py:447-454`), so the horizon guard never trips, and `GradNormTripwire` is "blind" on NaN
   (`probes.py:325-326`). A permanently-NaN run burns GPU hours until data runs out. No consecutive-skip cap exists outside PPO
   (`ppo/health.py:221-232`).
3. **Clipping off by default.** `max_grad_norm` defaults to `None` (`config.py:136`); no shipped sft/dpo/pretrain TOML sets
   it; `offload_adamw` ignores it (`offload_adamw.py:606-620`).
4. **Tripwires unarmed for SFT and pretrain.** `TripwireConfig.armed` defaults to `()` (`config.py:329-335`); only `dpo_*.toml`
   and `grpo_stageb.toml:68` arm anything. The code exists (`probes.py:165-431`) but is not active.
5. **Save is a full-model gather.** Every rank gathers to CPU, rank 0 writes ~14 GB (`train_module.py:128-152`,
   `trainer.py:676-681`); the README projects 62.6 GB host peak (35.5 GB with per-rank optimizer files). Per-rank
   optimizer state cannot be resharded (`trainer.py:925-947`).
6. **FSDP reduce-scatters on every micro-batch.** No `no_sync` / `set_requires_gradient_sync` anywhere (grep over align/,
   tests/, docs/ is empty); `train_batch` calls `backward()` each micro-step (`train_module.py:88-97`) with grad_accum 8
   (SFT) and 85 (`configs/grpo_stageb.toml:63-64`), over PCIe, in fp32. A throughput bug, unmeasured.
7. **Possibly rank-local logging and tripwires** (read from code, not executed): `Tracker` is constructed on every rank
   (`trainer.py:315-317`), logs `float(out.loss.detach())` unreduced (`:505-531`), and a HALT tripwire can raise on one rank
   only (`:551-589`), stranding peers until the NCCL timeout. Verify first, then fix.
8. **Loss normalisation is per-micro-batch mean**, equal-weighted across ranks (`train_module.py:88-97`,
   `sft_loss.py:97-116`). Biased with prompt-masked SFT or token-budget batching; becomes a hard interface problem for MoE.
9. **GPU evidence is thin.** README states adamw8bit is CPU-verified only (no GPU step time at 1.8B local elements, GPU/NCCL
   multi-rank unrun). The Triton FA2 backend is not selectable at HEAD: `.cache/kernel_gates/` is empty and
   `verify_equivalence_gate` raises without `triton_fa2_equivalence.json` there (`kernels/gate.py:40,138`); the file under
   `model/artifacts/` is not at the gate path. SDPA is the working backend.
10. Smaller: `latest_checkpoint` does not validate completeness and name-sorts (`trainer.py:878-886`); no fsync before the
    publishing rename; `pinned_checkpoints` is not persisted in `resume.pt`; no RNG state is saved; the optimizer has one
    param group.

Corrections to stale premises (verified): python3.12 headers now exist at `/usr/include/python3.12/Python.h`, `.cache/pyheaders`
is absent, so `kernels/_toolchain.py` is a no-op here. It also never blocked `torch.compile`; it only prepends headers to
`CPATH` and is called only from `kernels/ops.py:156` and `kernels/flash_attention.py:92`. `torch.compile` is simply unused
in align (grep hits are comments).

### 1.2 Can align serve as an OLMoE repo today?

**No.** There is zero MoE code: no router, expert FFN, aux or z-loss, grouped GEMM, dispatch, expert-aware checkpoint keys,
or MoE config fields (greps for moe/expert/router/top_k/all_to_all/z_loss over `align/`, `scripts/`, `configs/` find nothing
relevant). The model is dense Mistral only (`align/model/mistral.py:354-366` `MLP`; `model/hf_config.py:49` hard-codes
`MistralForCausalLM`; key tables `model/checkpoint.py:66-75` are dense `model.layers.N.*`). Also absent: QK-norm (OLMoE
uses it; `grep qk_norm|q_norm|k_norm align/` is empty), from-scratch sharded init, token-addressed data order, data mixtures,
doc-masked attention, WSD schedule, per-group weight decay, LM-head z-loss, in-loop eval.

What carries over unchanged: RMSNorm, RoPE, GQA attention backends, fused chunked CE (needs a z-loss term), FSDP2 per-block
sharding, activation checkpointing, the optimizer family, TrainModule/Trainer, resume machinery, config struct-mode, gates.
The MoE core is roughly 10 work items (ROADMAP P1) on top of the P0 robustness fixes.

Hardware reality for OLMoE-1B-7B itself (section 3): it fits on four free A10Gs at micro-batch 1x4096 (about 13.5 GiB per
card, my arithmetic) but a 5T-token run is about 12 years on this node. "OLMoE-class" here means the architecture and recipe at
reduced scale, not the published run.

### 1.3 Distance to a Leanstral-like build

**Far, and the target is not defined from this side.** The Leanstral 1.5 recipe is pending a separate research report at
`/apps/development/rlhf/research/notes/` (not consulted; nothing here assumes its architecture, size, context length, tokenizer
or training stages). Whatever that report says, the code side needs these regardless:

- **Lean verifier environment.** Nothing exists in align (`align/verify/*` is math-answer checking only). The only reference
  is `reference-repos/miles/examples/experimental/formal_math/single_round/reward_fn.py:18-41` (single-round, binary reward,
  Kimina server). It has known traps: it slices the response from its first `:=` (`:63-69`), and whether `sorry`, extra axioms or
  statement edits are rejected is delegated to a client library whose source is not in the repos. These must be tested
  adversarially before any reward is trusted. No multi-round proof repair exists in any reference.
- **Long context.** Engine KV is contiguous per slot, no paging, no prefix sharing (`inference/engine.py:11-17`); on a 24 GB
  card next to ~14.5 GB of bf16 weights an 8k-token completion allows about 8 concurrent sequences (derived estimate).
  Packed training has no document masking.
- **RL infrastructure.** Synchronous rollout (`rl/batch_source.py:246-247`), checkpoint-as-weight-sync, single rollout GPU, no
  external-process fault tolerance, no multi-turn token masks, no dynamic-sampling top-up, no pass@k eval.
- **Tokenizer and chat format.** Only the Mistral v0.3 BPE and a single-turn `[INST]` template; no tool-call/multi-turn
  serialisation; no other tokenizer versions (`tokenizer/core.py:22-24` raises on other pipelines).
- A Lean toolchain is an external binary; per `docs/DEPENDENCY_POLICY.md` it must be reached by stdlib subprocess/HTTP, with
  concurrency bounded by the host preflight.

---

## 2. Capability matrix

Legend. **Status**: have / partial / missing / n-a. **Torch-only?**: yes = torch/numpy/safetensors/stdlib; triton = needs
Triton (allowed); dep = needs a banned dependency. **A10G**: yes / degraded / no. **Effort**: S (<=2 days), M (~1 week),
L (2-4 weeks), XL (>1 month). **For**: R = robust-training, O = olmoe, L = leanstral-like, N = nice-to-have. "Best reference"
paths are under `reference-repos/`.

### 2.1 Training robustness

| Capability | align status (path) | Best reference | Torch-only? | A10G | Effort | For |
|---|---|---|---|---|---|---|
| Coordinated SIGTERM/cancel, step-boundary final save, wall-clock limit | missing (`train/trainer.py:395-484` has no cancel path; `scripts/sft.py:143-147`) | `OLMo-core/src/olmo_core/train/trainer.py:641-667,1234-1250,1369-1392`; `OLMo/olmo/train.py:1057-1108` | yes | yes | S | R O L |
| Non-finite loss halt + consecutive-skip cap | partial (skip exists `offload_adamw.py:193-201`; cap only in PPO `ppo/health.py:221-232`) | `OLMo-core/.../trainer.py:1430-1441` | yes | yes | S | R O |
| Global grad-norm clipping on for every optimizer, pre-clip norm logged | partial (`adamw8bit.py:418-443` implements; `config.py:136` default None; offload ignores) | `OLMo-core/.../transformer/train_module.py:514-524,611-638` | yes | yes | S | R O |
| Tripwires armed by default for SFT/pretrain | partial (code `train/probes.py:165-431`; `config.py:329-335` armed=()) | n/a (align's own design) | yes | yes | S | R |
| Loss/grad spike skip-and-continue (rolling mean+sigma) and spike-rate metric | partial (HALT-only `probes.py:282-340`) | `OLMo-core/src/olmo_core/optim/skip_step_optimizer.py:12-116`; `train/callbacks/stability_monitor.py` | yes | yes | M | R O |
| Token-count loss normalisation across micro-batches/ranks | partial (`train/train_module.py:88-97`, `sft_loss.py:97-116`) | `OLMo-core/.../transformer/train_module.py:354-422`; `nn/moe/loss.py:29-134` | yes | yes | M | R O L |
| FSDP `no_sync` on non-final micro-batches | missing (grep empty; `train_module.py:88-97`) | `fairseq2/src/fairseq2/trainer.py:640-650`; `nn/fsdp/unified.py:113-122` | yes | yes | S | R |
| Rank-reduced metrics, rank-0 writer, tripwire decisions from reduced values | partial / verify (`trainer.py:315-317,505-531,551-589`) | `OLMo-core/.../train/trainer.py:1037-1105,1394-1441` | yes | yes | M | R O |
| Skip-ranges, hard stop, loss-based early stop | partial (`max_micro_steps` only `trainer.py:461-463`) | `OLMo-core/.../trainer.py:254-259,1487-1498`; `OLMo/olmo/train.py:1062-1125` | yes | yes | S | R O |
| GC control (disable auto GC, periodic gen-1) | missing (manual only `scripts/sft.py:160`) | `OLMo-core/.../callbacks/garbage_collector.py:12-53` | yes | yes | S | R N |
| Throughput/MFU and CUDA allocator counters in history | partial (`probes.py:528`, tokens/s only) | `OLMo-core/.../callbacks/speed_monitor.py`, `gpu_memory_monitor.py` | yes | yes | S | R N |
| In-loop held-out NLL/perplexity eval | missing (stage-exit gate only `scripts/sft.py:176-185`) | `OLMo-core/.../callbacks/evaluator_callback.py:45-300` | yes | yes | M | R O L |
| Launch supervisor (bounded retry gated on preflight and progress, GPU idle/temp check) | partial (`utils/guard.py:81-566`; no retry loop) | `OLMo-core/src/scripts/lambda/healthchecks.sh` (launch_with_retries.sh is a Slurm wrapper, unbounded retry, weaker model) | yes | yes | S | R |
| RNG state in checkpoints; trainer/optimizer-state load split | partial (none saved; `trainer.py:888-1019`) | `OLMo-core/.../trainer.py:803-848`; `train/utils.py:31-95` | yes | yes | S | R O |
| Atomic publish + retention | have (`trainer.py:626-728,856-886`); gaps: no fsync before rename, no completeness check in `latest_checkpoint`, pins not persisted | `OLMo-core/.../distributed/checkpoint/filesystem.py:105-160` | yes | yes | S | R |
| Sparse permanent checkpoint tier (rewind target) | partial (keep 2, `config.py:166`) | `OLMo-core/.../callbacks/checkpointer.py:63-110` | yes | yes | S | R O |
| LR schedules: WSD, linear, alpha_f floor, stateless token-denominated | partial (cosine only, stateful `trainer.py:172-212`) | `OLMo-core/src/olmo_core/optim/scheduler.py:160-503` | yes | yes | S | O R |
| Param-group overrides (wd, quantisation floor) | missing (`adamw8bit.py:257-263` one group) | `OLMo-core/.../optim/config.py:44-229` | yes | yes | S | O |
| Weight EMA | missing | `TinyRecursiveModels/models/ema.py` | yes | yes | S | N R |
| Weight init: dense trunk N(0,0.02) exists; no depth-scaled, no meta/sharded init, no MoE init | partial (`model/mistral.py:559-566`, default `init_weights=True`) | `OLMo-core/.../nn/transformer/model.py:265-340`, `init.py:182-254` | yes | yes | M | O |
| Low-memory optimizer for ~7B (8-bit moments, host offload) | partial, CPU-verified only (`train/adamw8bit.py`, `offload_adamw.py`; README status) | not in OLMo-core | yes | degraded | M | O |
| Anomaly mode / fp16 scaler | missing | `fairseq2/src/fairseq2/trainer.py:513-665` | yes | yes | S | N |
| Optimizers beyond AdamW family (Muon, Lion, schedule-free) | missing | `OLMo-core/src/olmo_core/optim/{muon,lion}.py` | yes | yes | S | N |

### 2.2 Distributed and checkpointing

| Capability | align status (path) | Best reference | Torch-only? | A10G | Effort | For |
|---|---|---|---|---|---|---|
| Sharded topology-agnostic checkpoint (DCP-style), no rank-0 gather, bounded-memory unshard | partial (gather-based `train_module.py:128-152`, `trainer.py:676-746`; per-rank optimizer files refuse other world size `:925-947`; no DCP path for optimizer state) | `OLMo-core/.../distributed/checkpoint/__init__.py:71-393,465-`; `fairseq2/src/fairseq2/checkpoint/manager.py` | yes | yes | L | R O L |
| Async checkpoint save | missing | `OLMo-core/.../train/checkpoint.py:136-181`; `callbacks/checkpointer.py:168-325` | yes | degraded (host staging copy vs 45 GB rule) | M | R O |
| Data-reader state in checkpoint | have (`trainer.py` resume ~:960-985 refuses missing sampler state; `rl/batch_source.py:425,462`) | `fairseq2/.../checkpoint/manager.py:378-402` | yes | yes | n/a | R |
| FSDP2 per-block, bf16 params, fp32 reduce | have (`train/fsdp.py:61-86`; embeddings/lm_head in root unit) | `OLMo-core/.../nn/transformer/model.py:859-935` | yes | yes | n/a | R O |
| Distributed init: composite `cpu:gloo,cuda:nccl`, configurable timeout, set_device first | partial (`utils/dist.py:32-41`, `:108-121`) | `OLMo-core/src/olmo_core/train/__init__.py:85-148`; `distributed/utils.py:72-216` | yes | yes | S | R O |
| Activation checkpointing: full block only, non-reentrant | partial (`fsdp.py:90-123`) | `OLMo-core/.../nn/transformer/model.py:726-838` | yes | yes | M | O R |
| `torch.compile` per block | missing (comments only; no longer toolchain-blocked on this box) | `OLMo-core/.../model.py:840-857` | yes | yes (unmeasured) | M | O N |
| HSDP / DDP / 2-D mesh | missing | `OLMo-core/.../distributed/parallel/__init__.py:110-253` | yes | degraded | M | N (substrate for EP) |
| Expert parallelism, all-to-all | missing | `OLMo-core/.../nn/moe/parallel_mlp.py:206-871`; `ops/moe.py:224-281`; `torchfeather/torchfeather/distributed/expert_parallel.py:74-263` | yes | degraded (PCIe) | L-XL | O |
| Tensor / context / pipeline parallel | missing | `OLMo-core/.../distributed/parallel/{tensor,context,pipeline}_parallel.py` | yes (ring CP needs ring_flash_attn: dep) | degraded | L-XL | N |
| Float8 training | n-a | `OLMo-core/src/olmo_core/float8/__init__.py` | dep | **no** (SM86 has no FP8) | XL | none |
| Multi-node, elastic resume | missing (`trainer.py:937-947`) | `OLMo-core` DCP | yes | n/a | L | none now |
| Profiler callback | missing | `OLMo-core/.../callbacks/profiler.py` | yes | yes | S | N O |
| Per-layer grad/activation stats | partial (`probes.py` TrunkMovementProbe) | `OLMo-core/.../callbacks/gap_monitor.py` | yes | yes | M | O N |

### 2.3 MoE (OLMoE recipe)

| Capability | align status (path) | Best reference | Torch-only? | A10G | Effort | For |
|---|---|---|---|---|---|---|
| MoE block, ModelConfig fields, param naming, HF key tables | missing (`mistral.py:354-396`; `hf_config.py:49,68-91`; `checkpoint.py:66-75`) | `OLMo-core/.../nn/transformer/block.py:512-596`; `nn/moe/moe.py:123-251`; `nn/transformer/config.py:1079-1100` | yes | yes | L (M understated: config, hf_config, key tables, AC return contract, QK-norm all change) | O L |
| QK-norm in attention | missing (`mistral.py:211-226`, grep empty) | `OLMoE/configs/OLMoE-1B-7B-0924.yml` (`attention_layer_norm: true`); `OLMo-core/.../config.py:1087` | yes | yes | M | O |
| Router: fp32 linear gate, softmax top-k, no renorm (OLMoE-1B-7B) | missing | `OLMo-core/.../nn/moe/router.py:416-455`; `torchfeather/torchfeather/model/moe/moe.py:129-192` | yes | yes | S | O L |
| Load-balancing aux loss (weight 0.01, /n_layers) | missing | `OLMo-core/.../nn/moe/loss.py:29-111`, `moe.py:147-151` | yes | yes | S | O |
| Router z-loss (0.001) | missing | `OLMo-core/.../nn/moe/loss.py:114-138` | yes | yes | S | O |
| Aux-loss gradient plumbing through FSDP2, AC, grad accum | missing (`train_module.py:44-69,88-97`) | `OLMo-core/src/olmo_core/ops/__init__.py:4-28` | yes | yes | M | O R |
| Aux-loss-free bias balancing (DeepSeek style) | missing; not in the OLMoE recipe | `OLMo-core/.../router.py:303-328`; `torchfeather/.../components/optimizer.py:186-241` | yes | yes | M | L N |
| Shared experts | missing; not in OLMoE-1B-7B | `OLMo-core/.../moe.py:241-249` | yes | yes | S | N |
| Dropless dispatch (stable argsort, gather, scatter) | missing | `torchfeather/.../moe.py:202-407`; `OLMo-core/.../kernels/moe.py` | yes (triton optional) | yes | M | O |
| Capacity-factor routing | missing; parity oracle for dropless | `OLMo-core/.../parallel_mlp.py:350-631` | yes | yes | S | N |
| Grouped GEMM via `torch._grouped_mm` | missing. **A10G is firm "degraded"**: `libtorch_cuda` carries CUTLASS grouped kernels only for sm90/sm100 plus a fallback; `DeepSpeed/deepspeed/moe/group_gemm_triton.py:1-20` states sm80/sm86 falls back to a per-group `at::mm` loop with a device-to-host sync on `offs`. Inductor template gated at capability >=9.0 (`torch/_inductor/kernel/mm_grouped.py:225`) | `torchfeather/.../moe.py:51-86` | yes | degraded | S (correctness only) | O |
| Triton grouped GEMM for SM86 (fwd, dX, dW) | missing | `DeepSpeed/deepspeed/moe/group_gemm_triton.py:56-355` (written for sm80/86, autograd, fits 99 KB smem); `Liger-Kernel/src/liger_kernel/ops/fused_moe*.py` (prune autotune to stages<=2) | triton | yes (unrun) | L | O |
| Expert/router init (trunc-normal 0.02 incl. router and w2) and "no uninitialised param" test | partial (`mistral.py:562-566` visits Linear/Embedding only; bare 3-D expert Parameters would be missed) | `OLMo-core/.../init.py:182-254`; `OLMo-MoE/olmo/model.py:709-741` | yes | yes | S | O |
| Expert-aware checkpoint layout + round-trip gate | missing (`checkpoint.py:281` is the place) | `OLMo-core/src/olmo_core/nn/hf/convert.py:50-53,203-278` | yes | yes | M | O R |
| Sparse upcycling dense to MoE | missing | `OLMoE/scripts/sparsify_ckpt_unsharded.py:27-66` (no upcycling code in OLMo-core) | yes | yes | M | O L |
| MoE telemetry and router-collapse tripwire | missing | `OLMo-core/.../router.py:370-414` | yes | yes | S | R O |
| LM-head z-loss (1e-5) inside fused CE | missing (grep empty in `kernels/fused_ce.py`) | `OLMo-core/.../functional/cross_entropy_loss.py:12-50`; `Liger-Kernel/.../ops/cross_entropy.py:223` | yes | yes | S | O R |
| EP x FSDP2 composition | partial (`train_module.py:154-199`: non-`(Shard(0),)` placements take the generic `distribute_tensor` branch, correct but puts each whole tensor on GPU; not wrong) | `OLMo-core/.../nn/moe/mlp.py:71-112` | yes | degraded | L | O |
| Capacity of megablocks fork dependency | n-a: OLMoE 2024 run needed a megablocks fork (source not in reference-repos); OLMo-core has none | `OLMo-core` | yes | yes | S | O |

### 2.4 Data

| Capability | align status (path) | Best reference | Torch-only? | A10G | Effort | For |
|---|---|---|---|---|---|---|
| Token-addressed, world-size-invariant data order (uint32 permutation memmap, offset = tokens/global_batch) | partial (`data/loader.py:52-168` ragged batches; `trainer.py:118` materialises `list(self.sampler)`; resume at other world size reorders silently) | `OLMo-core/.../data/data_loader.py:441-533,619-777` | yes | yes | L | R O |
| Fixed-length memmap token dataset with doc lengths | partial (`data/store.py:304-372`; `pretrain_corpus.py:110-157` EOS-joined, no doc ids) | `OLMo-core/.../data/numpy_dataset.py:451-690`; `datatrove/.../pipeline/tokens/tokenizer.py` | yes | yes | M | O |
| Data mixtures (target ratio, max repetition, apportionment) | missing | `OLMo-core/.../data/source_mixture.py:30-456`; `dolma/src/mixer.rs` | yes | yes | M | O L R |
| Parallel resumable tokenise/shard build | partial (single process, serial; `pretrain_corpus.py:64` 500M ceiling; `store.py:316-319` loads every row's extras at open) | `datatrove/.../tokens/*`; `dolma/python/dolma/tokenizer/*` | yes | yes | L | R O |
| Exact + MinHash dedup | missing | `datatrove/.../dedup/minhash.py:41-598`; `exact_dedup.py` | yes (numpy, disk-sorted) | yes (CPU) | M | R O |
| Heuristic quality, language, URL/PII filters | missing (`data/argilla.py:82` is preference-pair only) | `datatrove/.../filters/gopher_quality_filter.py`; `DCLM/baselines/baselines_configs/dclm_baseline_refinedweb.yaml` | yes (fastText classifier is dep) | yes | M | R O L |
| Text-level benchmark decontamination | partial (id-level `data/build.py:199-229`) | `datatrove/.../decont/n_grams.py:28-194` | yes | yes | S | R L |
| Degenerate-instance filter, vocab-range check in collate | missing (only all-masked guard `sft_loss.py:84-93`) | `OLMo-core/.../data/utils.py:642-725`; `collator.py:31-40` | yes | yes | S | R O |
| Doc-boundary attention masking | missing (`kernels/flash_attention.py:824-832` causal + key padding only) | `attention-gym/attn_gym/masks/document_mask.py:34-63`; torch `flex_attention` | yes | degraded (JIT and 99 KB smem unmeasured) | M-L | R O L |
| Generic tokenizer (byte-level BPE with Unicode classes, other vocab sizes) | missing (`tokenizer/core.py:22-24`) | `mistral-common/.../tekken.py:109`; OLMoE uses `allenai/gpt-neox-olmo-dolma-v1_5` (`OLMoE/README.md:92-96`) | yes (stdlib has no `\p{L}`; use `unicodedata`) | yes | L | O L |
| Chat template: multi-turn, tools, V3 tool-call tokens | partial (`data/schema.py:264-273` single `[INST]`) | `mistral-common/.../instruct.py:473-740` (port, deps pydantic/tiktoken banned) | yes | yes | M | L R |
| Pretraining from scratch pipeline | missing (`scripts/pretrain.py` is continued pretrain from snapshot; config points at another repo's data dir `configs/pretrain.toml:52-55`) | `OLMo-core/src/scripts/train/OLMoE-1B-7B.py` | yes | degraded | M | O |

### 2.5 Kernels

| Capability | align status (path) | Best reference | Torch-only? | A10G | Effort | For |
|---|---|---|---|---|---|---|
| Fused linear + chunked CE | have (`kernels/fused_ce.py:180-234`) | `Liger-Kernel/.../ops/fused_linear_cross_entropy.py:51` | yes | yes | n/a | R O L |
| Fused DPO | have (`kernels/fused_dpo.py:375,526`; torch-only, no triton import) | n/a | yes | yes | n/a | R |
| Triton FA2 fwd+bwd | present but **not selectable at HEAD** (gate absent from `.cache/kernel_gates`; square, no-KV-cache blocks only `mistral.py:321-340`); SDPA already ships FA2 kernels | `kernels/_fa2_reference_upload.py` | triton | yes | S (rerun gate on GPU) | R |
| Fused RMSNorm/RoPE/SwiGLU with backward | partial (forward-only, flag-gated off `kernels/ops.py:39-49,163-213`) | `Liger-Kernel/.../ops/{rms_norm,rope,swiglu}.py` | triton | yes | S each | N |
| Generalised equivalence gate (parameterised by kernel name) | partial (`kernels/gate.py:62-208`, resolver is FA2-specific) | n/a | yes | yes | S | R O L |
| Doc-masked / varlen flash attention | missing | `attention-gym`; torch `varlen_attn` | triton | degraded | L | O L |
| FP8 GEMM, MLA decode, wgmma/TMA kernels | n-a | `DeepGEMM`, `FlashMLA`, `TileKernels` (SM90/SM100; TileLang) | dep | **no** | XL | none |
| MoE permute/gather/scatter Triton kernels | missing (torch index ops are the oracle) | `OLMo-core/.../kernels/moe.py:31-192`; `Liger-Kernel/.../fused_moe_kernels.py` | triton | yes | M | O (only if profiling demands) |

### 2.6 Inference and rollout

| Capability | align status (path) | Best reference | Torch-only? | A10G | Effort | For |
|---|---|---|---|---|---|---|
| Continuous batching, contiguous per-slot KV | have (`inference/engine.py:11-17,588-645`) | `nano-vllm/nanovllm/engine/scheduler.py` | yes | yes | n/a | R |
| Paged KV cache | missing (by design `engine.py:11-17`) | `vllm-upstream/vllm/v1/attention/ops/triton_unified_attention.py` (no flash-attn dep; nano-vllm needs flash_attn) | triton | yes | L | R L |
| Prefix sharing across the G samples of a prompt | missing (`rollout/native.py` prefills each prompt G times) | `nano-vllm/.../block_manager.py:36-120` (use blake2b, not xxhash) | yes | yes | S (copy KV rows) to M (block hash) | R L |
| Chunked prefill, token budget, preemption | missing | `vllm-upstream/vllm/v1/core/sched/scheduler.py:562-790` | yes | yes | M | R |
| CUDA-graph decode, fused sampler, top-k/min-p/penalties | missing | `nano-vllm/.../model_runner.py:196-257` | yes | yes | M | N |
| Speculative decoding (n-gram proposer + exact rejection sampler) | missing (`engine.py:48-51`) | `vllm-upstream/.../rejection_sampler.py`; `DeepSpec/deepspec/utils/sampling.py:30-45` | yes | degraded | L | N |
| MoE inference path (router + expert loop/grouped) | missing | `vllm-upstream/.../models/olmoe.py:63-118` | triton | degraded | XL | O |
| Per-request failure isolation, restartable rollout | partial (spool has heartbeat `rollout/protocol.py:167-227`; native raises) | `prime-rl/src/prime_rl/orchestrator/dispatcher.py:1-22` | yes | yes | S-M | R L |
| Multi-engine / time-sliced rollout | partial (one engine GPU, rank 0 only) | `open-instruct/open_instruct/data_loader.py:395`; `miles/miles/rollout/sleep_rollout.py` | yes | degraded | M | R L |

### 2.7 RL post-training

| Capability | align status (path) | Best reference | Torch-only? | A10G | Effort | For |
|---|---|---|---|---|---|---|
| Train/infer logprob mismatch correction (rho) | have (`rl/grpo_loss.py:107-164`) | `open-instruct/open_instruct/grpo_utils.py:410`; `miles/.../corrections.py` | yes | yes | n/a | R O |
| bf16 log-softmax gap under FSDP2 (README:200) | partial | n/a | yes | yes | S | R O |
| Sampling-support replay (top-p/top-k) | missing (minimum: refuse top_p<1 in RL config; `rollout/protocol.py:44`) | `prime-rl/src/prime_rl/trainer/rl/loss.py:74-91` | yes | yes | S (refusal) / M | R |
| Async / bounded-staleness pipelined RL | partial (submit/collect split exists `rollout/protocol.py:99-108` but called back to back `rl/batch_source.py:246-247`) | `prime-rl/.../orchestrator/dispatcher.py:400-441`; `open-instruct/.../data_loader.py:413,994-1003` | yes | degraded | M | R O L |
| In-memory trainer-to-engine weight sync | partial (checkpoint is the sync, ~13.5 GB write per sync) | `open-instruct/open_instruct/vllm_utils.py:1356-1457` | yes | degraded (PCIe) | M | R O |
| Multiple minibatches per rollout | missing (`rl/batch_source.py:296-307`; ratio is 1 and clip never binds) | `open-instruct/.../grpo_utils.py:102`; `trl/.../grpo_config.py:183` | yes | yes | M | R O |
| GSPO sequence-level ratio, OPSM | missing / partial (`grpo_loss.py:131-142` TV mask is different) | `trl/trl/trainer/grpo_trainer.py:3007-3033,3146-3157` | yes | yes | S | O R |
| MoE routing replay (R3) | missing (needs MoE model and decode path) | `miles/miles/utils/replay_base.py`; `prime-rl/.../models/layers/moe.py:179-227` | yes | yes | L | O |
| Truncation mask/metric (finish_reason unused) | partial (`rl/rewards.py:22-36` soft only) | `open-instruct/.../data_loader.py:1187-1204` | yes | degraded (KV-bound) | S | R L |
| Dynamic sampling top-up, pass-rate memory | partial (`rl/batch_source.py:270-285` retries, shrinking batch) | `open-instruct/.../data_loader.py:508-516,1006-1030` | yes | yes | M | R L |
| Eval during RL: pass@k, stop rate, multiple sets | partial (`rl/canary.py:26-92` greedy pass@1 only) | `open-instruct/.../grpo_utils.py:34-67,785-870` | yes | degraded | S | R L |
| Multi-turn rollout with observation-token loss mask | missing (`rl/assembly.py:126-150`) | `open-instruct/.../vllm_utils.py:949-1166`; `miles/.../multi_turn.py:25-70` | yes | yes | L | L N |
| Environment abstraction (reset/step/state, pool) | missing | `open-instruct/open_instruct/environments/base.py:43-135` (imports ray/pydantic: re-implement) | yes | yes | M | L N |
| Sandboxed code-execution verifier | missing | `open-instruct/open_instruct/code_utils/code_utils.py:96-340` (reliability_guard is not a security boundary) | yes | yes | M | N L |
| **Lean 4 verifier environment** | **missing** | `miles/examples/experimental/formal_math/single_round/{reward_fn,kimina_wrapper}.py` (single round, traps in 1.3) | yes (stdlib subprocess/HTTP) | yes | L | L |
| Rejection sampling / expert iteration pipeline | missing (parts exist: Engine, OutcomeVerifier, SFT) | `open-instruct/open_instruct/rejection_sampling/rejection_sampling.py:231-301` (RM-argmax, not verifier-gated) | yes | yes | S | L R |
| Reward model (Bradley-Terry) training | missing (`ppo/value_head.py:149` head only; README:194) | `trl/trl/trainer/reward_trainer.py:743-760` | yes | yes | S | N |
| LoRA in GRPO | partial (`ppo/lora.py` PPO only) | `prime-rl` docs LoRA | yes | yes | M | R N |
| PPO stage launchable | partial (README: draft, `[TBD]` fields) | n/a | yes | degraded | L | R |
| Recurrent/ACT reasoning cores (HRM, TRM) | missing; puzzle-scale, no evidence of LLM transfer; repos can not resume exactly | `HRM/models/hrm/hrm_act_v1.py:186-207` | yes | yes | M | N |

---

## 3. OLMoE on 4x A10G: sizing and token budget

### 3.1 Reference configuration (read from code)

OLMoE-1B-7B-0924 (`reference-repos/OLMoE/configs/OLMoE-1B-7B-0924.yml:11-45`; cross-check
`reference-repos/OLMo-core/src/olmo_core/nn/transformer/config.py:1079-1100`): 16 layers, d_model 2048, 16 heads (MHA),
64 experts, top-8, expert hidden 1024 (SwiGLU), dropless, RMSNorm, QK-norm, RoPE, seq 4096, vocab 50,280 padded to 50,304,
untied embeddings, no biases, init N(0,0.02) cut at 3 sigma, load-balance weight 0.01, router z-loss 0.001. Recipe: AdamW
lr 4e-4, wd 0.1, betas (0.9, 0.95), cosine with warmup (2,500 steps), alpha_f 0.1, global batch 1024 x 4096 = 4.19M tokens,
clip 1.0, bf16 with FSDP full shard.

Re-derived counts: 6.919B total, 1.282B active (matches `OLMoE/README.md:4`). Expert params 402.65M/layer, attention
16.78M/layer, embeddings 2 x 103.0M.

Note: the published run had `decay_embeddings: true` (`OLMoE-1B-7B-0924.yml:52-53`), so a single param group with wd 0.1 matches
it. The embedding-wd=0 override belongs to the later OLMo-core recipe (`OLMoE-1B-7B.py:33-41`). The "5.1T tokens" figure is
not in these repos; only the scheduler horizon t_max = 5e12 (`yml:64-66`) is.

### 3.2 Memory (my arithmetic; nothing measured)

Per parameter with align's stack: bf16 weight 2 B + bf16 grad 2 B + 8-bit AdamW moments 2.004 B = 6.004 B, FSDP over 4 ranks.

| Model | Total / active | Static per rank | Verdict |
|---|---|---|---|
| OLMoE-1B-7B (vocab 50,304) | 6.919B / 1.282B | 9.67 GiB | fits micro-batch 1x4096, about 13.5 GiB with one gathered block (0.84 GB, x2 with prefetch), root unit 0.41 GB, AC residents 0.27 GB, ~1 GB block transient, ~0.5 GB context |
| Same with Mistral vocab 32,768 | 6.847B / 1.210B | 9.57 GiB | fits |
| "mid": d1536 L16 E64 k8 h768 | 3.877B / 0.706B | 5.42 GiB | fits easily |
| **mini-A**: d1024 L16 E64 k8 h512 | 1.746B / 0.337B | 2.44 GiB | recommended ablation size |
| **mini-B**: d768 L12 E32 k4 h384 | 0.419B / 0.121B | 0.59 GiB | CI / smoke size |
| Mistral-7B v0.3 sliced upcycle, E=16 x h=896 | 7.250B / depends on k | 10.14 GiB | same footprint as the dense stack align already budgets |
| Mixtral-style 8 x full 14336 FFN | 46.7B | 65 GiB | impossible |

Ceiling: with ~5 GiB reserved, static <= ~17 GiB/rank gives about 12B total; sensible limit is 8-10B. The recipe's fp32 fused
AdamW (16 B/param, ~27.6 GiB/rank) does not fit; host-offload AdamW (12 B/param, ~83 GB host) violates the 45 GB preflight when
anything else runs. So adamw8bit is the only plausible optimizer, and it has never been exercised on MoE shapes or on GPU.
Flat expert layout `(E*h, d)` is preferable to `(E, d, h)` so `Shard(0)` works for any world size (GRPO stage B uses 3 ranks;
64 experts do not divide by 3).

### 3.3 Token budget in GPU-days

Assumptions (stated, not measured): sustained 25 TFLOP/s per A10G in bf16 (sensitivity 15-35); FLOPs per token = 6 x active
params; MoE dispatch and PCIe communication overhead ignored, so these are optimistic lower bounds on time. No MFU
measurement exists in the repo (`probes.py` reports tokens/s only).

| Run | FLOPs | GPU-days at 25 TFLOP/s | Wall-clock on 4 GPUs |
|---|---|---|---|
| OLMoE-1B-7B, 5e12 tokens (scheduler horizon) | 6 x 1.282e9 x 5e12 = 3.85e22 | ~17,800 (12,700 at 35; 29,700 at 15) | ~4,450 days, about 12 years (8.7 y at 35 TFLOP/s) |
| OLMoE-1B-7B throughput | 7.7 GFLOP/token | n/a | 100 TFLOP/s / 7.7e9 = ~13k tokens/s = ~1.1B tokens/day |
| OLMoE-1B-7B, 20B tokens | 1.5e20 | ~71 | ~18 days |
| mini-A (0.337B active), throughput | 2.02 GFLOP/token | n/a | ~49k tokens/s = ~4.3B tokens/day |
| mini-A, 20B tokens | 4.0e19 | ~19 | ~4.7 days |
| mini-B (0.121B active), 5B tokens | 3.6e18 | ~1.7 | ~10 hours |

Conclusions. A faithful OLMoE-1B-7B pretraining run is out of reach on this node by a factor of several hundred; the
published 5T-token horizon would take over a decade. The realistic goal is architecture-and-recipe reproduction at mini-A
scale (1-20B tokens, days not years), plus an OLMoE-1B-7B bring-up smoke test (a few hundred steps) to prove it fits and to
measure real MFU. Upcycling a pretrained dense checkpoint, or using pre-tokenised public shards read via memmap, is the only
way to reach useful quality; tokenising web-scale data on this box is not viable (pure-Python BPE ~75-165k tokens/s/core
measured on CPU, ~1 core-hour per 300-500M tokens).

Communication reality check (assumed ~12 GB/s effective PCIe, unmeasured). Per layer per 4096-token micro-batch, EP
all-to-all moves about 403 MB fwd+bwd (~34 ms/layer, ~0.54 s over 16 layers). FSDP-style expert weight gathering moves
about 2.42 GB/layer (~201 ms/layer, ~3.2 s over 16 layers) against ~0.9-2.1 s of compute. So FSDP-only experts are
comm-bound at 4096 tokens per micro-batch and EP is ~6x cheaper in bytes, but dropless EP costs a host sync per layer and
cannot be compiled. The published OLMoE-1B-7B run used no EP (`OLMo-MoE/olmo/config.py:1361-1363`). Benchmark both before
committing; do not assume.

Correctness caveat on upcycling tests. The OLMoE script makes E identical copies plus a fresh router with softmax top-k and no
renormalisation, so it is not function-preserving. For the sliced variant, "k=E reproduces dense logits" holds only if the
combine weights are forced to 1.0 (router bypass), because softmax weights sum to 1 (about 1/E each at init). The slicing
itself is exact for SwiGLU.

---

## 4. Excluded repos and paths

Excluded from review and why.

- **Speech/ASR/TTS/duplex/voice** (out of scope for a text LLM stack): Easy-Turn, Emilia-pipe, FD-Bench, Full-Duplex-Bench,
  Game-Time-Benchmark, HumDial-FDBench, IndicConformerASR, IndicVoices, MTR-DuplexBench, NeMo (LLM parts not reviewed),
  NeMo-text-processing, OIWER-Orthographically-Informed-Benchmarking-for-ASR, SID-bench, URO-Bench, VoiceBench,
  superintelligenceco-voicebench, vistaar, delayed-streams-modeling, hibiki, hibiki-zero, icefall, indic-punct,
  indicConformer-finetuning, indicfdb, llm_wer, moshi, moshi-finetune, omnilingual-asr, open_asr_leaderboard, personaplex,
  pocket-tts, pyannote-audio, sherpa-onnx, silero-vad, tts_longeval, unmute, whisper-timestamped, whisperX, whisper_streaming.
- **Interpretability UI**: neuronpedia. (interp-engine and circuit-tracer are library-level and were not reviewed here.)
- **Unreviewed, unclassified**: Dhruva-Evaluation-Suite (modality unstated), astra-brief, clef-head (no README).
- **Reviewed for patterns only, not adoptable**: 3FS (needs RDMA NICs and FoundationDB; only the "avoid FUSE for small random
  reads" lesson applies), nnsight and TransformerLens (analysis tooling, depend on transformers; TransformerLens has an
  OLMoE-1B-7B-0924 adapter that could serve as an oracle in the legacy venv), NeMo-Curator (RAPIDS and Ray; datatrove is the
  right template), AgileRL (evolutionary HPO irrelevant; GSPO/packing covered elsewhere), HRM and TinyRecursiveModels
  (small puzzle models, model-only checkpoints that cannot resume; negative pattern), DeepSpeed (banned dependency; Triton
  grouped GEMM and router math are pattern sources only).
- **Hopper/Blackwell-only, not A10G-feasible**: DeepGEMM sm90/sm100 kernels and Mega MoE, FlashMLA, TileKernels (TileLang, SM90/SM100),
  FP8/MXFP, prime-rl DeepGEMM/DeepEP/FA3/FA4, Liger Blackwell tiling, torch `_scaled_grouped_mm`, Inductor grouped_mm template.
- **Infrastructure with no single-node value**: Beaker launch/callbacks, remote-store I/O, Ray placement, SLURM/k8s, prefill-decode
  disaggregation, wandb/comet/slack reporting, HF-hub conversion callbacks (need transformers).
- **Not read or not found**: the megablocks fork behind the 2024 OLMoE run (source absent); prime-rl `verifiers`, `prime-envs`,
  `renderers` submodules (empty directories); Lean environments anywhere other than miles `formal_math/single_round`;
  `nmo-research/docs/0[234]_*.md` (not opened by the capability maps; claims rest on code); the Leanstral 1.5 report.
- **Not verified by execution anywhere**: SM86 behaviour of `flex_attention`, `varlen_attn`, Triton grouped GEMM autotune;
  all throughput and PCIe figures; the all-ranks-write-`history.jsonl` observation; adamw8bit on GPU.
- **Documentation errors noticed in align itself**: `docs/ARCHITECTURE.md:156-166,195-212` misdescribes the GRPO loss/rho (code
  is DAPO clip with rho as a multiplicative weight and constant Dr.GRPO denominator, `rl/grpo_loss.py:182-199`);
  `README.md:198` says GRPO resume is unwired while `rl/batch_source.py` has state_dict/load_state_dict; comment at
  `train/offload_adamw.py:632` ("Not a field of OptimizerConfig today") is stale since `config.py:136` exists.