| # 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. |
| |
