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);.partialstaging plusos.renamepublication 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):
- No graceful stop. No SIGTERM/SIGINT handler in the trainer;
scripts/sft.py:143-147documents a torchrun SIGTERM landing mid-save. No wall-clock limit. - 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, andGradNormTripwireis "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). - Clipping off by default.
max_grad_normdefaults toNone(config.py:136); no shipped sft/dpo/pretrain TOML sets it;offload_adamwignores it (offload_adamw.py:606-620). - Tripwires unarmed for SFT and pretrain.
TripwireConfig.armeddefaults to()(config.py:329-335); onlydpo_*.tomlandgrpo_stageb.toml:68arm anything. The code exists (probes.py:165-431) but is not active. - 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). - FSDP reduce-scatters on every micro-batch. No
no_sync/set_requires_gradient_syncanywhere (grep over align/, tests/, docs/ is empty);train_batchcallsbackward()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. - Possibly rank-local logging and tripwires (read from code, not executed):
Trackeris constructed on every rank (trainer.py:315-317), logsfloat(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. - 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. - 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 andverify_equivalence_gateraises withouttriton_fa2_equivalence.jsonthere (kernels/gate.py:40,138); the file undermodel/artifacts/is not at the gate path. SDPA is the working backend. - Smaller:
latest_checkpointdoes not validate completeness and name-sorts (trainer.py:878-886); no fsync before the publishing rename;pinned_checkpointsis not persisted inresume.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 isreference-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 whethersorry, 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-24raises on other pipelines). - A Lean toolchain is an external binary; per
docs/DEPENDENCY_POLICY.mdit 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,rendererssubmodules (empty directories); Lean environments anywhere other than milesformal_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.jsonlobservation; adamw8bit on GPU. - Documentation errors noticed in align itself:
docs/ARCHITECTURE.md:156-166,195-212misdescribes 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:198says GRPO resume is unwired whilerl/batch_source.pyhas state_dict/load_state_dict; comment attrain/offload_adamw.py:632("Not a field of OptimizerConfig today") is stale sinceconfig.py:136exists.