align / code /docs /MEASUREMENTS.md
viai957's picture
Add files using upload-large-folder tool
d909bd8 verified
|
Raw History Blame Contribute Delete
59.8 kB

Measurements: first GPU evidence pass (G0 control plane on NCCL, P0-M MoE dispatch)

Raw numbers from the first real-GPU run of the G0 work, with what each one decides in docs/ROADMAP.md. Nothing here is an estimate: every figure is copied from a file under .tmp/gpu_smoke/ (named per section). Scope is deliberately small: tiny models, 2 GPUs, one single-GPU benchmark. Nothing here is a 4-rank, 7B, or multi-minute result, and sections "Not measured" lists what that leaves open.

Environment

Item Value
Date / window 2026-10-12, 17:45 to 17:57 IST (host pl-gpt-machine)
GPUs used physical GPU 1 (GPU-bdbde387-01e3-14d4-9a9a-5d40e8cf3427, PCI 00:1C.0) and GPU 0 (GPU-c569cb57-edc9-3dd1-5e5e-6822ebcb56f0, PCI 00:1B.0). Sections (a), (b): CUDA_VISIBLE_DEVICES=1,0 (rank 0 -> GPU 1, rank 1 -> GPU 0). Sections (c), (d): GPU 1 alone.
GPU NVIDIA A10G, sm_86, 23028 MiB, SM clock 1710 MHz (max 1710), 35-42 C, 58-62 W idle, 133 W during the benchmark. PCIe gen4 x8; all four GPUs PHB to each other (no NVLink).
Driver / CUDA / torch / triton / NCCL 580.65.06 / 12.8 / 2.11.0+cu128 / 3.6.0 / 2.28.9+cuda12.9
Python /apps/development/rlhf/align/.venv/bin/python (3.12), PYTHONPATH=/apps/development/rlhf/align (the package is not installed in the venv)
Foreign load on the GPUs one idle nmo_trace_server.py (circuit-tracer UI, not ours) holds 12045 / 10853 / 12045 / 12045 MiB on GPUs 0-3 at 0% utilisation for the whole window. Only ~11 GiB per card was free; nothing here needed more than 4.4 GiB. Its presence did not change utilisation (0% before and after) but a benchmark taken next to an idle resident context is not a clean-box number.
Raw files .tmp/gpu_smoke/{bench_c.json,bench_c2.json,diag_d.json,runs/*/}; scripts bench_c.py, bench_c2.py, diag_d.py, worker.py, driver.py, compare.py, preflight.sh (all scratch, not committed)

Preflight (all passed; re-run before every GPU command, last values shown)

MemAvailable_GB=121  load1=1.06  diskq(nvme2n1)=0  D-state=0  date=2026-10-12T17:47:41+05:30
GPU  used/total MiB  util
0    12045 / 23028   0 %
1    10853 / 23028   0 %
2    12045 / 23028   0 %
3    12045 / 23028   0 %
PICK = GPUs 1,0   (least memory used; each >= 4 GB free, < 50% util)

Thresholds: MemAvailable >= 45 GB, load < 96, disk queue < 64, D-state <= 5. Heavy-job lock (.align_heavy_job.lock) was free before and was taken by rank 0 of each torchrun via apply_resource_guards(heavy=True). After the last run: the only GPU processes were the four foreign server pids, same memory as before.

(a) FSDP2 + adamw8bit + per-rank optimizer state: SIGTERM one rank, resume, compare

Setup. worker.py under torchrun --standalone --nproc_per_node=2, NCCL. Tiny dense align Mistral (4 layers, d=256, 8 q / 4 kv heads, vocab 1024, bf16, 2,885,888 parameters, 39 tensors), production wrapping (apply_activation_checkpointing then apply_fsdp: per-block FSDP2 + root, bf16 params, fp32 reduce), optimizer.name=adamw8bit (max_grad_norm=1.0, 1.4M local params, block-wise 8-bit moments on device), train.per_rank_optimizer_state=true, grad_accum=2, 20 micro-steps (10 optimizer steps), save_every_micro_steps=10, tripwires.armed=("grad_norm",) (as configs/sft.toml), ALIGN_PG_TIMEOUT_S=120. Each rank trains on its own rows. Peak device memory per rank: 0.034 GiB.

Interrupt. driver.py polls history.jsonl and sends a real external SIGTERM to rank 1's pid only once micro-step 15 is logged (a 0.4 s sleep per micro-step makes the timing reliable). It was sent at t = 13.08 s after launch.

Run What Result
a_full1, a_full2 uninterrupted, 20 micro-steps, same seed final params bit-identical (0 of 2,885,888 elements differ). Run-to-run nondeterminism of this setup on 2 GPUs is zero.
a_int SIGTERM to rank 1 at logged micro 15 both ranks logged coordinated stop (signal) at optimizer step 8 (micro 16) (the next optimizer-step boundary), both returned stop_exit_code = 143 (128+15), checkpoints micro_00000010 and micro_00000016; checkpoint_problem() returns None (complete) for both; each holds model.safetensors, resume.pt, optimizer_rank0.pt, optimizer_rank1.pt, tokenizer files. No HALT.json (a signal stop is not a verdict). torchrun itself exited 1, as the integration report measured.
a_res16 resume latest (= the stop checkpoint, micro 16), finish to 20 resumed_micro_step 16, finished at micro 20 / optimizer step 10, exit 0.
a_res10 resume the explicit cadence checkpoint micro_00000010, finish to 20 resumed_micro_step 10, finished at 20 / 10, exit 0.

Final-parameter comparison against the uninterrupted run (compare.py, bf16 params gathered with full_state_dict()):

Comparison max abs diff differing elements verdict
a_full1 vs a_full2 (run-to-run) 0.0 0 / 2,885,888 bit-exact
a_full1 vs a_res16 (resume from the SIGTERM stop checkpoint) 0.0 0 / 2,885,888 bit-exact
a_full1 vs a_res10 (resume from the cadence checkpoint) 0.0 0 / 2,885,888 bit-exact
sensitivity: a_full1@20 vs a_int@16 (different training stage) 0.0273 2,141,203 differs, so the comparison can see change
negative control: resume from 16 with optimizer_rank0.pt and optimizer_rank1.pt swapped 0.0273 2,333,055 differs, so a wrong optimizer shard is detected, i.e. the bit-exact result depends on the per-rank files being read correctly

Loss, lr and grad-norm in history.jsonl also match the uninterrupted run at every micro-step the interrupted and resumed runs share: a_int 16 of 16, a_res16 20 of 20, and the resumed portion (micros 11-20) of a_res10. Bit-exactness here is explained, not lucky: with 2 ranks an NCCL reduce-scatter sums two values (commutative, so order-free), params are bf16 with no master copy so the checkpoint round trip is lossless, 8-bit moments are saved in their quantised form, and RNG/stream/scheduler state are restored. A 4-rank NCCL ring sums in a fixed ring order too, but that is not shown here.

(b) NaN on one rank: skip cap HALTs both ranks, no NCCL hang

Same setup, train.max_consecutive_skipped_steps=3. The adapter multiplies the loss by NaN on one rank only. (The injection index is trainer.micro_step + 1, so --nan-from 5 first hits micro_step 4. The table uses the real micro_step.)

Run Injection Result
b_halt rank 1, every micro-step from 4 both ranks raised NonFiniteHaltError at micro_step 10 (optimizer step 1), identical message ("4 consecutive bad optimizer steps exceed train.max_consecutive_skipped_steps=3 (last cause: non-finite loss and skipped step)"), skip_stats = {consecutive 4, bad_total 4, boundaries 5}, optimizer_skipped_steps 4 on both, exit code 81 on both, HALT.json {"code": 81} written. Worker-side time from the start of main to the raise was 1.24 s (rank 0) and 1.39 s (rank 1); the whole launch took 7.8 s. No NCCL error, watchdog or timeout line in the log (the only NCCL matches are the version banner and a deprecation warning, see Findings). No checkpoint existed yet (first cadence save is micro 10), so none was written.
b_halt_rank0 rank 0, every micro-step from 4 identical outcome with the roles swapped: both ranks exit 81 at micro_step 10, 1.2 s.
b_transient rank 1, one micro-step (real micro_step 6) not a halt. Both ranks skip_stats = {consecutive 0, bad_total 1, boundaries 10}, optimizer_skipped_steps 1, 9 of 10 optimizer steps applied, finished at micro 20, exit 0, all parameters finite.

The grad gate's skip is global (NaN in one rank's gradient reaches both shards through the reduce-scatter, and the global-norm is non-finite on both), and the non-finite-loss flag is MAX-reduced with the stop flags in the one boundary collective, so both ranks always take the same branch.

(c) MoE expert dispatch: loop vs grouped_mm, one A10G (GPU 1)

Shape. E=64 experts, top-8, d_model 2048, expert FFN 1024 (SwiGLU: gate_up [64, 2048, 2048], down [64, 2048, 1024]), bf16, align.model.moe.moe_experts_forward. Median of 10 after 3 warmups, CUDA events, includes gather, both expert matmuls, SiLUup, routing-weight scale and index_add_ combine; the router itself is excluded from this table (timed in the breakdown below). "Balanced" = the real router on random activations; "skewed" = a Zipf(1.0) expert popularity, k distinct experts per token (max expert load 7.1x the mean, no empty expert). Backward uses a gradient equal to the output; weights and x require grad. FLOPs per layer forward = 2Tk(d2I + Id) = 412 GFLOP at T=4096, 1649 GFLOP at T=16384; backward counted as 2x forward.

T routing loop fwd ms loop fwd+bwd ms grouped fwd ms grouped fwd+bwd ms fwd+bwd speedup loop / grouped peak GiB fwd+bwd TFLOP/s loop / grouped
4096 balanced 13.132 39.684 12.388 35.325 1.123x 2.17 / 1.92 31.2 / 35.0
4096 skewed 14.557 42.345 13.732 37.518 1.129x 2.17 / 1.92 29.2 / 33.0
16384 balanced 49.139 136.851 46.440 128.696 1.063x 4.40 / 4.40 36.2 / 38.4
16384 skewed 47.846 135.295 45.353 127.205 1.064x 4.40 / 4.40 36.6 / 38.9

(min/max over the 10 repeats are within 1% of the median in bench_c.json.)

Is the fallback taken? Yes, confirmed three independent ways.

  1. TORCH_CPP_LOG_LEVEL=INFO on a 2-call probe prints, on every call: [I... GroupedMMUtils.h:128] fallback path for torch._grouped_mm, performance may not be optimal (probe_log.err). The same string is in libtorch_cuda.so next to a _grouped_mm_fallback symbol.
  2. The profiler on one F.grouped_mm(x_sorted, gate_up^T, offs) call (T=4096 balanced) shows 64 aten::mm and 64 aten::item CPU ops and 93 CUDA kernels: 64 cuBLAS GEMMs (36x ampere_bf16_s16816gemm_bf16_128x128_ldg8_f2f_stages_32x5_tn, 28x ampere_bf16_s1688gemm_bf16_128x128_ldg8_f2f_stages_32x1_tn), 28 memsets, and 1 Memcpy DtoH (the offs read-back that forces a host sync). No CUTLASS grouped kernel launches.
  3. Timing: grouped_mm is only 6-13% faster than the loop (above), which is what swapping one loop of 64 cuBLAS calls for another loop of 64 cuBLAS calls looks like.

Agreement of the two backends on the same inputs (bf16, T=4096 balanced; the other three rows are the same within noise): output y max abs diff 2.4e-3 on an absmax of 0.465 (one bf16 ulp at that magnitude, relative 5.3e-3); dx 3.9e-3 on 0.652 (6.0e-3); dgate_up, ddown, dweights max abs diff exactly 0.0 (bitwise identical weight and routing-weight gradients). Hypothesis, not measured further: the forward 1-ulp difference is the cuBLAS kernel choice (the s1688 vs s16816 split above) plus the nondeterministic index_add_ order.

Where the time goes (bench_c2.json, loop backend forward, balanced routing, ms):

Component T=4096 T=16384
route_top_k (fp32 router + softmax + top-k) 0.273 1.120
build_dispatch_plan (stable sort + scatter_add counts) 0.210 0.209
counts.tolist() host sync (idle stream) 0.032 0.032
gather index_select 0.552 2.238
64x gate_up linears 5.536 20.806
SiLU * up 0.726 2.817
64x down linears 3.056 10.962
scale + index_add_ combine 2.317 9.374
full loop forward (measured) 13.162 49.225
sum of the seven non-router components 12.429 46.438
GEMM share of the loop forward 65.3% 64.5%
GEMM throughput inside the layer 48.0 TFLOP/s 51.9 TFLOP/s
whole layer forward throughput 31.3 TFLOP/s 33.5 TFLOP/s

Dense-GEMM ceiling on this card (cuBLAS bf16, same GPU, same clocks): per-expert shapes run as 64 back-to-back F.linear: 512 rows/expert: gate_up 48.8, down 43.8 TFLOP/s; 2048 rows/expert: gate_up 52.9, down 50.2 TFLOP/s; one large 8192x4096 @ 4096x4096: 68.0 TFLOP/s.

(d) Tiny OLMoE on CUDA, both dispatch backends

diag_d.py (GPU 1): OLMoE architecture, 2 layers, d=128, expert width 64, 16 experts top-4, 4 heads, vocab 512, sdpa attention, identical weights loaded into a loop and a grouped_mm model, batch 4x32, CE loss, full backward.

dtype logits max abs diff (absmax) loss loop / grouped router counts identical worst grad rel diff params with all-zero grad finite pass
fp32 1.9e-7 (0.908) 6.259638 / 6.259639 yes (every layer) 6.7e-7 (layers.1.attn.k_proj.weight) 0 yes yes (tol 1e-4)
bf16 3.9e-3 (0.906) 6.259564 / 6.259566 yes (every layer) 1.0e-2 (layers.0.mlp.experts.down_proj, 5.7e-6 on 5.6e-4) 0 yes yes (tol 5e-2)

grouped_mm_unavailable_reason("cuda") returned None (the 2-expert probe runs on sm_86: the fallback works, it is just not fast). Every parameter, including router and stacked experts, receives a non-zero gradient on CUDA.

What each number decides

Number Decision it feeds
grouped_mm is the fallback on sm_86 (log line, 64 mm + 64 item + 1 DtoH, 93 kernels) Closes the "expect the slow fallback" half of P1-11 with evidence. docs/REFERENCE_GAP_ANALYSIS.md line for torch._grouped_mm ("A10G is firm degraded") is now measured, not inferred.
grouped_mm 1.06-1.13x faster than loop (fwd+bwd), same memory, weight grads bit-identical Either backend is correct on A10G; the choice is worth at most ~6-12%. Keep loop the default (oracle, no extra probe/alignment constraint); grouped_mm may be used where grouped_mm_shape_problem is None. Not a reason to build anything.
GEMMs are 65% of layer forward and run at 48-52 TFLOP/s in-layer, 44-53 TFLOP/s for the same shapes in isolation, vs 68 TFLOP/s for one large GEMM P1-12 (Triton grouped GEMM) fails its own entry condition ("only if P1-11 shows the loop is the bottleneck"): the loop's GEMM efficiency equals cuBLAS on those shapes, so a Triton kernel's best case is closing 49 -> ~68 TFLOP/s on 65% of the layer, at most ~18% of forward time (an upper bound: cuBLAS itself does not reach 68 on these shapes), with a high bar to beat cuBLAS on sm_86 with a 99 KB shared-memory limit. Recommend defer P1-12 and spend the effort on the glue below, if anything.
Non-GEMM glue = 35% of layer forward: combine (y*w + index_add_) 17.6% / 19.0%, gather 4.2% / 4.5%, SiLU*up 5.5% / 5.7%, routing + plan 3.7% / 2.7% (T=4096 / 16384) If a custom kernel is worth writing for MoE on A10G it is a fused gather and a fused scale+combine (+ SiLU*up), not the GEMM. Worth a P1 line item only after P1-15 shows MoE layers dominate step time.
Host sync in the loop (counts.tolist()) costs 0.032 ms on an idle stream Isolated cost is negligible. Its real cost (draining the launch queue once per layer and blocking CPU run-ahead) is not measured by this benchmark; P1-15's end-to-end MFU is the test. Do not remove it on this evidence.
Layer fwd+bwd at 4096 tokens = 35-40 ms (loop 39.7), 16384 tokens = 127-137 ms; 31-39 TFLOP/s whole-layer First anchor for the OLMoE-1B-7B throughput estimate (16 MoE layers, ~0.6 s/micro-batch of expert compute at 4096 tokens/GPU before attention, optimizer and comms), to be replaced by the P1-15 measured MFU. The analysis doc's 25 TFLOP/s assumption is below this single-GPU layer figure, as expected before comms.
Peak 2.2 GiB (4096) / 4.4 GiB (16384) for one layer incl. 0.75 GiB expert weights + 0.75 GiB grads Per-layer transient for the P1-15 memory arithmetic: peak minus the 1.5 GiB of weights and grads is ~0.7 GiB at 4096 tokens and ~2.9 GiB at 16384 (before attention, residual stream and activation checkpointing, which this benchmark does not include).
(d): loop and grouped_mm agree on CUDA end to end in fp32 (1e-7 / 7e-7) and bf16 (1 ulp / 1e-2 worst grad), identical routing, no dead parameter Clears the CUDA half of the P5 dispatch parity claim (previously CPU only). Backend choice does not need a model-level gate.
(a): resume from a SIGTERM stop checkpoint and from a cadence checkpoint is bit-exact vs an uninterrupted run on 2 NCCL ranks, FSDP2 + adamw8bit + per-rank optimizer files; negative control detects wrong shards Moves P0-1 / P0-8 / P0-9 / per-rank from "CPU-verified" to "NCCL 2-rank verified" for the lossless-resume and coordinated-stop claims. Leaves open: 4 ranks, 7B-scale save/restore timing, kill -9 mid-save.
(a): coordinated stop took 1 boundary (micro 15 signal -> micro 16 stop), both ranks, 143 on both The boundary-latency of a stop is grad_accum micro-steps. At accum 8 and ~minutes per micro-step on the 7B job this is the number to budget in the supervisor's SIGTERM grace period; it needs the 7B step time to turn into seconds.
(b): persistent NaN on either single rank halts both at the same micro-step in ~1.3 s with identical error and exit 81, no NCCL hang; a single transient NaN skips exactly one step and the run continues healthy Moves P0-2 from "CPU gloo" to "NCCL 2-rank verified": the skip-cap/non-finite agreement needs no extra collective to avoid the lone-raiser hang. The remaining P0-2 item (GradNormTripwire treating a NaN norm as blind) is unchanged.

Findings that were not asked for

  1. P0-6 (rank-local logging) is confirmed on real ranks. history.jsonl receives one row per rank per micro-step: 40 rows for 20 micro-steps at world 2, rows from the two ranks carry different per-rank losses, and under the rank-1 NaN injection the file shows a NaN loss row (rank 1) and a finite row (rank 0) for the same micro-step. Nothing interleaved or corrupted a line in these runs. A consumer reading "the last row" or counting rows gets rank-dependent answers. Reduce at log cadence and write from rank 0 only, as P0-6 says.
  2. TORCH_NCCL_TRACE_BUFFER_SIZE is deprecated on torch 2.11 (Warning: Environment variable TORCH_NCCL_TRACE_BUFFER_SIZE is deprecated; use TORCH_FR_BUFFER_SIZE instead, 4 lines per launch). configure_nccl_diagnostics sets the old name for the buffer size (it sets both names only for the dump-file prefix). Whether the old name still enables the recorder was not tested: nccl_trace/ was created and is empty because no collective timed out. P0-10's open question (does DUMP_ON_TIMEOUT need TORCH_NCCL_ENABLE_MONITORING) still needs a deliberate hang, not done here.
  3. The guards caught a real launch-environment fault. The first launch was refused with writable path '/home/vignesh/.opencode/tmp' resolves to device 64768, not the /apps device: the Claude session's TMPDIR points at the root filesystem. driver.py now unsets it so apply_resource_guards sets /apps/.../align/.tmp. Sections (c) and (d) were single-process scripts that did not go through apply_resource_guards (no heavy-job lock, inherited TMPDIR); they wrote nothing outside .tmp/gpu_smoke/.
  4. torchrun reaps the survivor. After rank 1 exited (143) torchrun logged Sending process ... closing signal SIGTERM for rank 0 and exited 1. The checkpoint was complete because save_checkpoint ends in a barrier, so a peer cannot exit before rank 0 finishes writing; in these runs rank 0 also completed its own post-train dump. Anything a stage script does after train() returns on rank 0 only (an eval, a copy) can be cut off by this; the existing scripts return immediately on a coordinated stop, which is the safe shape.
  5. The tiny runs do not learn (loss 6.96 to 7.02 over 20 steps on random tokens); they validate the control plane and optimizer plumbing (trunk_report.moved_fraction = 1.0, all 39 tensors change between micro 16 and 20), not optimisation quality.

Not measured / still open

  • 4 ranks, a 7B model, FSDP2 with MoE (the OLMoE RouterStats through FSDP2 hooks), expert sharding on dim 0, EP all-to-all (P1-13), any run longer than 10 seconds, NCCL behaviour under a real hang, kill -9 during a save, offload_adamw, the DPO/GRPO/PPO stage scripts, P0-12 (Triton FA2 gate), P0-14 (adamw8bit at 1.8B local elements; this run used 1.4M).
  • Timings are single-GPU, single-layer, random weights, with an idle foreign context resident on the card; they are not an end-to-end MFU and do not include FSDP gathers, attention, the optimizer or comms.
  • agree figures compare loop vs grouped_mm, not either backend vs the HF reference (that is the oracle tier, CPU).
  • Bit-exactness in (a) is for 2 ranks on this tiny model; a different kernel mix (flash/mem-efficient attention backward, larger embeddings) could introduce run-to-run noise that this setup did not exhibit (two uninterrupted runs were bit-identical).

Track M: is align an OLMoE repo on the REAL checkpoint? (allenai/OLMoE-1B-7B-0125, 2026-10-12)

Everything below was measured on the released weights, not on a tiny model. Raw files: .tmp/olmoe/ (parity/, parity_post/, parity_trained/, smoke_prefix_run/ = first 20-step pair, smoke/ = final pair, smoke_trim/, probe_*/, det_test.json, gradvar_before_fix.json); the code that produced them is scripts/olmoe_parity.py, scripts/olmoe_sft_smoke.py (subcommands prepare, train, probe, drive, analyze) and configs/olmoe_sft_smoke.toml; the pins are tests/test_olmoe_scripts.py. Nothing here is an estimate unless it says so. The run sequence matters: the first 20-step pair ("pre-fix") exposed a bug that was then fixed in align/model/moe.py (section f); the whole pair was then repeated on the fixed code ("final"). Numbers are from the final pair unless marked pre-fix.

Environment

Item Value
Host / date pl-gpt-machine, 2026-10-12 18:20 to 22:10 IST. 186 GB RAM, no swap. Preflight before every heavy step (MemAvailable >= 45 GB, load < 96, disk queue < 64, D-state <= 5, GPUs < 2 GiB used) and the heavy-job flock for every run; another track's lean_ceiling_test held the flock for ~15 min mid-session and my next job was refused (exit 83) rather than overlapping, then waited.
GPUs 4x A10G, 23,028 MiB, all four idle (0 MiB) (the foreign 12 GiB nmo_trace_server of the first GPU pass was gone). PCIe, every pair PHB; NCCL chose its SHM host-memory transport, not P2P (Channel ... via SHM/direct/direct), and NCCL_P2P_LEVEL=SYS changed nothing. NCCL 2.28.9, torch 2.11.0+cu128.
Checkpoint allenai/OLMoE-1B-7B-0125 @ 9b0c1aa8, 6 safetensors shards, stored fp32 (27.68 GB, torch_dtype: float32; the brief's "~14 GB" is the bf16 size), 3,219 tensors, per-expert hub keys. Downloaded with huggingface_hub in the legacy venv (3 min 32 s; 26 GiB on /apps): safetensors + config + tokenizer only.
Oracle transformers 5.15.0 in /apps/development/rlhf/nano-mistral-o1/.venv, subprocess only (the legacy-side code is a source string; tests/test_dependency_firewall.py passes).
Data (smoke) the repo's cached PRM800K SFT shards (nano-mistral-o1/data/shards/sft/train, 11,739 Mistral-tokenized rows) decoded with the vendored Mistral tokenizer, split at prompt_len, re-tokenized with the OLMoE tokenizer; 320 rows of 640-968 tokens (mean 757; prompt mean 111), EOS-closed, prompt masked. Real data, not synthetic.

(e) Parity on the real checkpoint (CPU, align vs transformers 5.15, separate processes)

3 fixed prompts tokenized by the HF tokenizer (5, 39 and 30 tokens = 74 positions: prose, Python, prose with numbers); ids go to align. align loop + sdpa, transformers eager. Each engine saves logits and per-layer router logits; a third step compares them.

dtype max abs logit diff mean abs top-1 router top-8 set agreement, min over 16 layers verdict
fp32 1.4e-05 (sequential index_add_ combine, pre-fix) / 3.0e-05 (fixed-order combine, post-fix); |logit| up to 25.3 1.1e-06 to 1.4e-06 74/74 1.00000 on every layer; max router-logit diff 9e-06 PASS (gate: < 1e-3, 100%, >= 99.9%)
bf16 1.31 0.048 to 0.064 74/74 0.892 to 0.973 per layer recorded, not gated

Gate chosen: fp32, max abs < 1e-3 AND top-1 100% AND per-layer router set agreement >= 99.9% (olmoe_parity.evaluate, pinned both ways in tests). Why not the brief's alternative (bf16, top-1 100% with max abs < 0.1): the weights are stored fp32, so fp32 compute is the same function and any residual is accumulation order (measured 1.4e-05 to 3.0e-05, 30x to 70x under the gate); and the bf16 alternative is unsatisfiable by two correct implementations (measured max abs 1.31: bf16 logits of magnitude 16-32 have an ulp of 0.125; align runs its router in fp32 while transformers' bf16 router logits differ from it by up to 0.34, flipping 3-11% of top-8 sets). A bf16 tolerance that passes would be fitted to the answer. bf16 is run and recorded, never gated.

The gate has teeth on these weights (the same align code with one classic porting bug injected, compared with the same transformers tensors; identical before and after the combine fix):

Injected defect max abs logit diff top-1 min router set agreement passes
none 1.4e-05 to 3.0e-05 100% 100% yes
QK-norm per head instead of over the full projection 11.6 81% 5.4% no
Mistral RoPE base 1e6 instead of 1e4 7.2 95% 25.7% no
expert gate/up stacked in the wrong order 32.7 19% 0.0% no
Mixtral routing (top-k, then softmax over the kept experts) 16.8 47% 2.7% no

Pinned against the released files without loading weights (test_real_checkpoint_config_is_the_documented_architecture_and_its_index_maps_totally_onto_the_model): the hub config.json (rope spelled the 4.x way, fp32 dtype) reads as the 6.919B architecture; all 3,219 index keys are consumed by hf_to_align and every parameter of the model is produced with the right shape (checked on meta tensors built from the safetensors headers); index total_size = 4 x 6,919,161,856. The slow+oracle test (-m "slow and oracle") re-runs the whole parity table and the four mutants (passes, 107 s).

Round trip of a TRAINED checkpoint. The checkpoint align's Trainer wrote at micro 80 of the SIGTERM-resumed run (bf16 weights, config.json, tokenizer files, 13.84 GB, gathered from FSDP2 shards by rank 0) loads in transformers 5.15 and gives the same logits as align loading the same files: fp32 compute, max abs 1.43e-05, top-1 74/74, router sets 100% (parity_trained/). The "an align checkpoint loads in transformers" claim therefore holds at 7B-MoE scale through the real save path.

Cost: each engine loads in ~6 s (page cache hot) and runs the 3 prompts in ~2.9 s on 8 threads; peak RSS 51 GB (align) / 50 GB (transformers) at fp32, mostly mapped checkpoint pages.

(f) A bug only the real weights could show: the MoE combine was nondeterministic (FIXED in align/model/moe.py)

Symptom. In the first 4-GPU pair the pre-clip global grad norm of optimizer step 1 was 5.44, 8.09 and 11.65 in three runs with identical weights, rows and seed (the losses at micro 1-4 agreed to 1e-3). Repeating the exact step-1 accumulation window four times inside ONE process on frozen weights (probe --only gradvar) gave 5.41, 10.21, 5.42, 13.28 with rank-mean losses 0.4216 / 0.4211 / 0.4214 / 0.4221: forward and backward were not reproducible.

Isolation (.tmp/olmoe/det_test.py: the real layers 0-3 + embeddings on one A10G, bf16, one batch, 5 repeats, hash of every layer's top-8 expert ids):

variant distinct losses / 5 distinct router top-8 hashes / 5 distinct layer-0 grad hashes / 5
as shipped (sdpa, index_add_ combine) 5 5 5
as shipped, eager attention 5 5 5
fixed-order combine, sdpa 1 1 5 (attention-backward noise, norm 42.2606 to 42.2635)
fixed-order combine, eager attention 1 1 1

Cause. moe_experts_forward summed a token's k=8 expert outputs with out.index_add_(0, token_idx, y) and gathered tokens with x.index_select(0, token_idx), whose autograd backward is another index_add_ of k rows. On CUDA both are atomic scatters; in bf16 every partial sum is rounded, so the order the SMs happen to run in changes the hidden state by an ulp, and the router, sitting on near-ties, then selects different experts for the same tokens. CPU index_add_ is sequential and (measured) accumulates at higher precision; the oracle tier compares fp32 on CPU; so neither could see it.

Fix. combine_experts (un-sort to (token, slot) order with the plan's inverse permutation, then one sum(dim=1): fp32 accumulation, one rounding, fixed order) and _DispatchGather (a custom autograd function whose backward is the same fixed-order reduction). Equal to the old code in fp32 up to summation order; more accurate than 8 sequential bf16 adds in bf16; no measurable speed change (compute-only 16-layer fwd+bwd, 1,024 tokens: 0.621 s before and after). After the fix:

  • one MoE layer at OLMoE shape (E=64, k=8, T=4096, bf16) is bit-identical over 8 repeats on the A10G (test_gpu_moe_layer_is_bit_reproducible_across_repeats_in_bf16; the old code fails it);
  • the FSDP step-1 window gives 14.42, 14.38, 14.35, 14.34 with bit-identical losses (0.42198020964860916 x4); the 0.6% residual is the SDPA mem-efficient backward (gradient noise, no routing effect);
  • the old values (5.4 to 13.3) were not noise around 14.4: bf16 atomics under-reported the step-1 gradient norm by up to 2.7x;
  • two full 20-step runs now agree on the step-1 norm to 0.2% (14.33, 14.30), and on per-step mean losses to <= 0.0016 (the pre-fix pair: <= 0.0096).

Regression tests (tests/test_olmoe_scripts.py): fixed-order exactness against an fp64 sum rounded once, an emulated-atomic bf16 combine that makes the same property fail on CPU, the slot-order permutation test through both backends, structural pins that moe_experts_forward goes through the helpers and that the gather has the custom backward node, and the GPU repeat test. Six mutants of moe.py (inline index_add_, sequential bf16 reduce, missing un-sort, wrong inverse, plain index_select, inline combine) are all killed. Oracle tier after the fix (-m "oracle and not gpu"): 17 passed, 6 skipped, the README baseline; full CPU suite 1600 passed, 13 skipped (other tracks' tests included).

Still open: grad-norm excursions. Even after the fix, three of 20 optimizer-step grad norms differ between the two final runs by 1.2x to 2.4x with per-micro losses equal to 1e-3 (step 11: 1.32 vs 1.77, step 12: 1.42 vs 3.35, step 16: 1.70 vs 1.45); the other 17 agree to <= 4%. They are not caused by the resume (step 11 was logged before the stop; step 16 differs in the other direction). Replaying step 12's window four times from the saved micro-44 weights gives 3.350 / 3.346 / 3.347 / 3.341: from fixed weights the norm is reproducible to 0.3%, so the uninterrupted run reached the same window with weights that differ by bf16-level noise and got 1.42. The gradient norm is a discontinuous function of the weights (consistent with router top-k flips), not the output of a nondeterministic kernel. Control: the following window (micro 49-52) from the same weights gives 1.3189 / 1.3187 / 1.3190 / 1.3188 (+-0.02%).

(g) 4-GPU post-training smoke: FSDP2, bf16, adamw8bit, per-rank optimizer state, LBL off, telemetry on

torchrun --nproc_per_node=4 scripts/olmoe_sft_smoke.py train --config configs/olmoe_sft_smoke.toml: OlmoeForCausalLM (6.919B params, loop dispatch, sdpa), non-reentrant activation checkpointing on all 16 blocks, fully_shard per block + root (bf16 params, fp32 reduce), MoETrainModule with lb_weight = z_weight = 0 and per-layer telemetry on, fused chunked SFT cross-entropy, micro-batch 1 x seq 1024 (every micro-batch padded to 1024) x accum 4 x 4 ranks = 16,384 padded tokens (12,080 real) per optimizer step, lr 1e-5 (10% warmup, cosine), clip 1.0, adamw8bit (3.23 GiB of state per rank), 20 optimizer steps, checkpoint at step 10. Ranks draw distinct rows (rank stride over one seeded permutation) and load the model one at a time.

Quantity Measured (final pair; the pre-fix pair agrees to the digit unless noted)
Micro-step time, median of 77 intervals, 4 ranks in lockstep 14.79 s (min 14.60; max 16.89 = a micro-step containing an optimizer step); optimizer step = 59.2 s
Throughput 277 padded tokens/s (4 GPUs), 204 real tokens/s, 175 loss-tokens/s; 61 optimizer steps/h of 16 k padded tokens (the 20-step run is 20 minutes of training)
Achieved model FLOPs 277 x 6 x 1.282B = 2.1 TFLOP/s total, 0.53 TFLOP/s per GPU (~0.8% of an A10G's 70 TFLOP/s bf16; the roadmap assumed 20-25 TFLOP/s per card)
GPU memory per rank allocator peak 14.83 GiB, reserved peak 17.79 GiB, nvidia-smi peak 18.5 GiB (18,932 MiB on rank 0) of 22.5 GiB (23,028 MiB): 4.0 GiB headroom on rank 0 (rank peaks 18,932 / 18,900 / 18,667 / 18,667 MiB; the allocator peaks are identical on all ranks). [Review 2026-10-12: this said 3.6 GiB, a MiB/GiB unit slip: 23,028 - 18,932 = 4,096 MiB.] After wrap + optimizer init 6.45 GiB (3.22 params + 3.23 8-bit state). The load peaks at 13.1 GiB (the full bf16 model on the card before sharding).
Mean loss per optimizer step (16 micro-batches each) 0.464, 0.377, 0.361, 0.271, 0.247, 0.326, 0.169, 0.304, 0.361, 0.210, 0.247, 0.220, 0.219, 0.366, 0.254, 0.299, 0.410, 0.188, 0.288, 0.178 (first-5 mean 0.344, last-5 mean 0.273). Low from the start (the base model has seen this kind of math) and noisy (16 rows/step): it goes down; it is not a convergence claim.
Pre-clip grad norm 14.33 (step 1), 3.66, 2.45, 1.98, 1.90, 2.26, 1.30, 1.71, 1.80, 1.28, 1.32, 1.42, 1.26, 1.53, 1.32, 1.70, 1.74, 1.20, 1.43, 0.99 (step 20). Above the 1.0 clip on 19 of 20 steps. Identical on all four ranks. GradNormTripwire ceiling is 500.
Skips / halts 0 / none. Trunk-movement probe passed; every parameter, including all expert tensors, moved.

Router telemetry with load balancing OFF (exact window values at the optimizer-step boundary; 64 experts, top-8; MaxVio = (max load - mean)/mean; utilization = sum_i min(f_i, 1/E)):

layer 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
MaxVio, step 1 4.75 5.33 2.95 2.18 6.23 6.32 5.57 5.54 6.15 6.21 5.48 6.28 5.00 4.24 3.82 6.18
MaxVio, step 20 4.95 5.28 3.01 2.43 6.10 5.96 5.52 5.32 5.81 6.00 5.20 5.91 4.84 4.27 4.56 4.95
utilization, step 1 .757 .769 .724 .707 .661 .617 .563 .570 .631 .601 .644 .587 .653 .651 .639 .635
utilization, step 20 .756 .754 .719 .716 .667 .598 .580 .575 .632 .607 .647 .592 .657 .637 .659 .607

Layer-mean MaxVio 5.14 at step 1, 5.01 at step 20 (range 4.72 to 5.14 over the run; per-layer drift 0.3 to 1.3), worst layer 6.3 / 6.1, utilization 0.65 (worst layer 0.56 / 0.58), worst-layer dead-expert fraction (< 0.1x mean load) 15.6% -> 12.5% (up to 10 of 64 experts), router entropy 3.80 -> 3.79 nats (log 64 = 4.16), the load-balance loss VALUE (not applied) 1.50 -> 1.49 (1.0 = balanced), the z-loss value (not applied) 8.6 -> 8.5. The pretrained OLMoE router is already strongly skewed on this narrow math distribution (the busiest expert gets ~6x its fair share in most layers), and 20 steps of LBL-off SFT at lr 1e-5 do not move it (no collapse, no drift beyond noise). That is evidence that LBL-off is safe for a short SFT, not for a long one, and says nothing about LBL-on (no A/B was run). Each telemetry point is one optimizer-step window all-reduced over the 4 ranks (16 distinct rows, ~12k real tokens), so a single point carries sampling noise from which problems happened to be in the window: the step-1 vs step-20 comparison is two such samples, and the 4.72-5.14 range over all 20 windows is the honest noise band.

Checkpoint save and host memory (the 45 GB floor question). A checkpoint is 27.7 GB: model.safetensors 13,838,724,744 B (bf16, per-expert hub layout, tokenizer files copied in) + 4 x 3,466,699,433 B of per-rank 8-bit optimizer state.

Before malloc_trim (pre-fix pair) After (final pair and smoke_trim/)
Save wall time 230-233 s 228-241 s idle; 347 s once while a 3-minute CPU test suite ran on the same host (optimizer files done after ~2 min, model.safetensors ~2 min later)
Rank-0 RSS during a save 2.1 -> 30.3 GB (+28.1 GB = the gathered 13.8 GB bf16 state + the 13.8 GB per-expert clones align_to_hf makes) 2.07 -> 30.2 GB (unchanged: the transient is real)
Ranks 1-3 steady RSS 14.3 GB each, 13.4 GiB of it one anonymous [heap] (the freed bf16 state dict from the rank-serial load, stranded in glibc's brk heap; rank 0 kept ~11 GB after each save) 2.0 GB each; rank 0 2.06 GB after every save
Sum of the four ranks' RSS peak 73.5 GiB (pre-fix ref), 89 GiB (resume) peak 46.1 GiB (a save), 90 GiB (resume): the resume reads four full CPU state dicts and is not helped
System MemAvailable 141 idle -> 97 training -> 67.8 GB min (during a save) 134.6 idle/training -> 96.6 GB min at saves; 70.1-75.9 GB min around the resume and the final save of the SIGTERM run (another tenant held ~35 GB then)
Model load, per rank 26.1-26.6 GB peak RSS (13.8 GB bf16 + the 13.8 GB stacked copy hf_to_align builds), 11-65 s; four ranks at once would be ~106 GB same peak, then RSS falls to 0.74-0.98 GB (trimmed: true)
keep_checkpoints = 2 not honoured across a resume: the SIGTERM run left 3 checkpoints (40, 44, 80 = 83 GB), because pruning only knows checkpoints written by the current process same in the recorded run. FIXED in review (Trainer.resume_from now adopts the lower-numbered checkpoints on disk; higher-numbered zombies are still ignored; tests/test_control_plane.py::test_resume_adopts_older_checkpoints_...). Not re-run on GPUs.

trim_host_heap (malloc_trim(0)) is applied after the load and after every save, outside the save's RSS window; pinned by an AST test and 3 mutants. Added in review, NOT re-measured on GPUs: a trim after resume_from_latest() too. In the recorded SIGTERM run ranks 1-3 sat at 13.7 GB RSS from the resume until the first save (smoke/sup/ckpt_rss_rank1.jsonl: rss_start_gb 13.73 at micro 80), the same stranded-heap shape as after the initial load. It does not change the 90 GiB summed peak DURING the resume read.

SIGTERM, supervisor, resume (scripts/supervise.py -> torchrun; SIGTERM to rank 1's pid only, 0.2 s after the step-10 checkpoint was complete on disk):

Event Observed (final pair; the pre-fix pair is identical in structure)
Stop all four ranks logged coordinated stop (signal) at optimizer step 11 (micro 44): the next boundary, one accumulation window (4 micro-steps, 59 s) later; stop checkpoint micro_00000044 written (236 s)
Supervisor attempt 1: child exit 1 (torchrun collapses the 143s), progressed: true, newest complete checkpoint 44; backoff 10 s; preflight ok; attempt 2 launched, resumed at micro 44 / optimizer step 11, ran micro 45-80, exit 0, "success" after 2 attempts. 2,072 s total.
Cost of the interruption last micro-44 log -> first micro-45 log = 331 s (236 s stop save + 10 s backoff + ~85 s relaunch, serial load and resume) against 14.8 s for a normal micro-step
Loss continuity vs the uninterrupted reference (same code, rows, seed) micro 1-44 (same computation, two runs): max abs dloss 2.9e-03, mean 7.0e-04 (the run-to-run floor). Micro 45-80 (resumed): max 3.3e-03, mean 1.1e-03. Rule: resumed max <= 2 x floor + 5e-3 = 1.09e-02 -> PASS with 3.3x room. Per-step mean losses of the two runs differ by <= 0.0016. (Pre-fix pair: floor 1.2e-02, resumed 4.7e-03; the combine fix tightened the floor 4x.) Review correction: the original text here claimed 'a wrong optimizer shard or data position moves this loss by 0.1 to 1'. That was never measured (no mutant resume was run on the GPUs) and is not credible for a lost optimizer state at lr 1e-5 over 36 micro-steps, so the 1e-2 tolerance is a continuity check of unknown sensitivity. What actually establishes that this was a real resume is now checked exactly by olmoe_sft_smoke.py analyze (see 'Review of Track M' below).
Gradient norm at the resume see "grad-norm excursions" in (f): the first resumed step (12) logged 3.35 against 1.42 in the reference, and replaying that window from the checkpoint weights reproduces 3.35

(h) FSDP cost vs tokens, compute, and expert parallelism

The micro-step is 14.8 s for ~0.6 s of compute. Three independent probes on the idle box (olmoe_sft_smoke.py probe --only collectives|compute|scan; scan is the real FSDP2 model, fwd+bwd, no optimizer):

Probe Result
FSDP micro-step time vs tokens L=256: 14.60 s; 1,024: 14.72-14.77 s; 2,048: 14.92 s; 4,096: 15.02 s; 8,192: 15.28 s. Peak allocator memory 10.96 / 11.04 / 11.16 / 11.40 / 12.12 GiB (+3.23 GiB for the 8-bit state in training: 15.4 GiB at 8,192 by that arithmetic. [Review: the training run's 14.83 GiB at 1,024 tokens is 0.56 GiB ABOVE this probe's 11.04 + 3.23 = 14.27 (gradient-accumulation and optimizer transients the scan does not have), so the honest 8,192 figure is ~15.9 GiB allocated, ~19.6 GiB nvidia-smi. It still fits, with ~2.9 GiB to spare; it is no longer 'comfortable'.]). Time is flat in tokens: +0.08 s per extra 1,024 tokens; ~14.5 s is fixed.
NCCL on this box (one block = 419.6 M params; bf16 all-gather of the block, fp32 reduce-scatter; un-overlapped) all-gather 0.209 s (4.0 GB/s algbw), reduce-scatter 0.430 s (3.9 GB/s); the same at 1/16, 1/4 and full block size and with NCCL_P2P_LEVEL=SYS: it is the shared-memory transport's bandwidth, not latency. Per micro-step: 16 x (2 all-gathers [forward + activation-checkpoint recompute] + 1 reduce-scatter) = 13.55 s = 92% of the measured 14.8 s.
Compute only (truncated random-init model on one GPU, no FSDP, activation checkpointing on): t = fixed + n_layers x per-layer, 2- and 4-layer fits loop dispatch: 1,024 tokens 35.9 ms/layer + 47 ms -> 0.62 s for 16 layers; 4,096 tokens 86.3 ms/layer + 128 ms -> 1.51 s. grouped_mm (the torch fallback on sm_86): 0.44 s (-30%) and 1.42 s (-5.6%).
Expert-parallel all-to-all (bf16 activations, T x k x d per rank; 6 per layer per micro-step = forward + recompute + backward) 1,024 tokens: 9.1 ms each (3.7 GB/s), 0.87 s per micro-step; 4,096 tokens: 36.4 ms each, 3.50 s per micro-step

A 128k-token optimizer step (32,768 tokens per rank). FSDP columns use the measured micro-step times above; compute and EP columns are arithmetic on measured per-op times, not runs (the 8,192-token compute 2.7 s and all-to-all 7.0 s are linear extrapolations of the 1,024/4,096 points; EP adds ~0.5 s per step for the dense-parameter all-reduce, assumed non-overlapped):

tokens/micro-batch micro-steps/step FSDP step (measured t_micro x n) compute alone FSDP / compute EP-4 step (est.) EP / compute FSDP tok/s EP tok/s (est.) compute-bound ceiling
1,024 32 32 x 14.72 = 471 s 32 x 0.62 = 19.8 s 24x 32 x (0.62+0.87)+0.5 = 48 s 2.4x 278 ~2,700 6,600
4,096 8 8 x 15.02 = 120 s 8 x 1.51 = 12.1 s 10x 8 x (1.51+3.50)+0.5 = 41 s 3.4x 1,090 ~3,200 10,800
8,192 4 4 x 15.28 = 61 s 4 x ~2.7 (extrapolated) = ~11 s 5.6x 4 x (2.7+7.0)+0.5 = 39 s 3.6x 2,140 ~3,400 ~12,000

What each number decides

Number Decision it feeds
fp32 logits 1.4e-05 to 3.0e-05, top-1 74/74, router sets 100% on all 16 layers; the trained checkpoint round-trips through transformers at 1.4e-05; four injected bugs fail by 5-6 orders P5 "HF logit parity" is met on the released weights; align is an OLMoE repo. Gate = fp32 < 1e-3 / 100% / 99.9%. bf16 (1.31) is a recorded noise floor and must not be used as a gate.
bf16 router set agreement 89-97% between two correct implementations Under bf16, router replay (R3) and any flip-rate threshold must be measured against align's own sampler, never against transformers: the HF-vs-align flip rate (3-11% per layer) is the baseline disagreement of bf16 routers, far above the roadmap's 0.1% skip-R3 heuristic.
Combine nondeterminism: 5/5 distinct top-8 sets for identical inputs before the fix, 1/1 after Fixed. Any MoE flip-rate or mismatch-KL probe taken on the pre-fix code measured atomics, not engine mismatch: re-baseline. Residual: SDPA backward gives ~0.6% gradient noise (no routing effect), and the gradient norm is discontinuous in the weights (excursions in (f)).
Pre-clip grad norm 14.3 at step 1, > 1.0 on 19 of 20 steps With Adam the clip does not shrink the update, it re-weights early steps against late ones (step 1 is down-weighted 14x). Log a clipped fraction next to any lr claim and compare against an unclipped baseline before attributing an SFT effect to the lr.
FSDP micro-step 14.8 s, flat in tokens; 92% is NCCL over host shared memory at ~4 GB/s Roadmap gate 3 (FSDP vs EP): exceeded. FSDP-only is 24x / 10x / 5.6x compute at 128k tokens per step with 1,024 / 4,096 / 8,192-token micro-batches; the criterion was 3x. Build EP (P1-13) if ~1,000-2,000 tokens/s is not enough.
EP all-to-all 3.5 s vs FSDP 13.55 s per micro-step at 4,096 tokens (16x less at 1,024; 2x less at 8,192) EP is worth 3x at 4,096-token micro-batches and 1.6x at 8,192, and is still comm-bound (3.4-3.6x compute) on this PCIe-SHM box. The free lever is the micro-batch: 1,024 -> 8,192 tokens takes plain FSDP from 278 to 2,140 tokens/s. The 4-GPU ceiling for OLMoE post-training here is ~2-3.4 k tokens/s = 0.2-0.3 B tokens/day, against the roadmap's 1.1 B/day (which assumed 25 TFLOP/s per card): re-state the OLMoE budget before pricing an SFT epoch or a proxy.
277 tokens/s = 0.53 TFLOP/s per GPU at the shipped micro-batch; 4.2% of the 6,600 tokens/s compute-bound ceiling Replaces the 25 TFLOP/s MFU assumption (P1-15).
grouped_mm -30% (1,024 tokens) / -5.6% (4,096) of compute, but compute is 4% of the FSDP micro-step Roadmap gate 2 (grouped GEMM) stays closed for FSDP-only training (worth <= 0.2 s of 14.8 s). It becomes a 5-30% lever only after the comm wall is removed (EP or larger micro-batches). Keep loop as the default and the oracle.
Peak 14.83 GiB allocated / 18.5 GiB nvidia-smi of 22.5; +0.08 s per extra 1,024 tokens; 8,192-token micro-batches fit (~15.9 GiB allocated; see the correction under (h)) Headroom is 4.0 GiB at 1,024 tokens and the model fits at 8,192. OLMoE-1B-7B full-FT with adamw8bit is bandwidth-limited, not memory-limited.
adamw8bit + FSDP2 on 3-D expert tensors (sharded on dim 0, 16 experts per rank): 20 steps, no skips, all parameters moved Closes the "adamw8bit / FSDP2 with MoE RouterStats on GPU" part of P0-14 and of the docs/ARCHITECTURE.md limits: router stats ride out of non-reentrant-checkpointed blocks without double counting (telemetry stable across 80 micro-steps).
MaxVio ~5 (worst 6.3), utilization 0.65, 12-16% dead experts at step 1, unchanged at step 20 with LBL off Roadmap gate 7 (LBL off unless the load tripwire fires): not fired. No MaxVio threshold is pre-registered yet; these are the first numbers to base one on. Because the base router is already this skewed on math data, an LBL-on arm would pull away from the pretrained solution: run that A/B before enabling it.
Save 228-241 s idle; rank-0 peak RSS 30.2 GB (+28 GB); MemAvailable min 67.8 GB before the heap fix, 96.6 GB after; four ranks 90 GiB during a resume Roadmap gate 4 (checkpoint RSS): the criterion was a per-rank save peak within ~10 GB of the 45 GB floor. It is not met (30.2 GB; 15 GB under), so reshardable checkpoints are not forced into P0 by this measurement. The job's own footprint was the bigger risk: 40 GiB of dead heap on ranks 1-3, now fixed (malloc_trim). Remaining: the 28 GB save transient (needs a per-tensor streaming writer) and the 90 GiB resume (four full CPU state dicts: needs a shard-local read). With both, one OLMoE job coexists with a second ~60 GB job; today it coexists with ~35 GB.
SIGTERM to one rank -> all four stop at the same boundary -> supervisor relaunch -> loss within 3.3e-3 of the uninterrupted run (floor 2.9e-3) Moves P0-1 / P0-8 / P0-9 from "2-rank NCCL" to "4-rank NCCL, 6.9B model, real save/resume" verified. Budget for the SIGTERM grace period: boundary latency 59 s + 236 s for the stop checkpoint; the whole interruption cost 5.5 minutes.

Findings that were not asked for (Track M)

  1. scripts/sft.py / scripts/pretrain.py (dense) give every rank the same batch stream: build_training_sampler returns one unsharded LengthGroupedSampler and the Trainer does no rank slicing (only dpo_stage wraps in DistributedBatchSampler). By code reading, a 4-GPU dense SFT trains 4 identical copies of each batch: the reduced gradient equals one rank's, the configured 24-row "global batch" is 6 rows, and the horizon arithmetic is 4x too long. Not verified by a run (no dense 7B weights on this box). The OLMoE smoke shards by rank (RankSampler, disjointness pinned).
  2. keep_checkpoints leaks one checkpoint per resume (13.8 + 13.9 GB): pruning is driven by a per-process list. Fixed in review (see the table above).
  3. The history.jsonl collision (P0-6) is still live in the Trainer; the smoke gives ranks 1-3 their own rank{r}/history.jsonl and rank 0 the one the supervisor watches.
  4. The checkpoint is fp32 on disk (27.7 GB), twice the brief's "~14 GB"; load_hf_safetensors(dtype=bf16) casts shard by shard so a bf16 load never holds both.
  5. /apps is 94-95% full. One SIGTERM/resume pair of 20-step runs writes ~130 GB of checkpoints; all checkpoints from these runs were deleted after analysis (histories, telemetry and logs kept).
  6. Host-memory fix beyond the brief: malloc_trim after load and save recovered ~40 GiB of dead glibc heap on a no-swap box (above). It applies to any rank-serial 7B load in this repo, dense included.
  7. The shared heavy-job lock did its job against another track's long job (refused with exit 83, then waited).
  8. torch could not create .cache/torch/kernels (a warning on every MoE entropy call); the empty directory now exists. Untracked: .cache/.

Not measured / still open (Track M)

  • LBL-on vs LBL-off A/B; a router flip rate between a sampler and the trainer (needs the MoE decode path, P3-1); grouped_mm end to end under FSDP; EP itself (only its all-to-all volume was timed); a 128k-token step (largest micro-batch run: 8,192 tokens, forward+backward only); runs longer than 80 micro-steps; kill -9 mid-save; a second large job alongside (the 45 GB floor under real contention); GRPO / DPO on OLMoE.
  • Throughput numbers are for loop dispatch with sdpa; grouped_mm and other attention variants were timed compute-only.
  • The grad-norm excursions of (f) are characterised (reproducible from fixed weights, not from a kernel) but not traced to a specific router flip.

Review of Track M (independent reviewer, 2026-10-12)

What was re-checked, how, and what changed. Raw: .tmp/review/ (parity1/, parity1.log).

Claim Check Result
fp32 parity on the real checkpoint Re-ran olmoe_parity.run (CPU, separate processes, current working tree incl. the fixed-order combine) on one prompt that is not in the original three: "The sum of the first n integers is n(n+1)/2; for n = 100 that gives 5050, and 17 * 23 = 391." (32 tokens, |logit| up to 23.1) max abs 1.57e-05, mean 1.38e-06, top-1 32/32, router top-8 set agreement 1.00000 on all 16 layers (min overlap 8/8), max router-logit diff 9.4e-06. Gate PASS. Consistent with the recorded 1.4e-05 / 3.0e-05. Both workers took the heavy-job lock; peak RSS ~51 GB; MemAvailable was observed at >= 72 GB while it ran.
Recorded parity tables Recomputed from parity/, parity_post/ report_*.json and mutants_fp32.json All numbers in section (e) match the raw files (fp32 1.383e-05 / 2.956e-05, bf16 1.306, 0.892-0.973, the four mutants).
LBL off, telemetry on launch_rank*.json of every attempt of both runs; history.jsonl keys lb_weight = z_weight = 0.0 on every rank of every launch; moe/l{0..15}/max_vio, utilization, moe/lb_loss, moe/z_loss logged every row; optimizer banner = adamw8bit on-device; moe_backend = loop.
The resume was a real resume sup/launch_rank*_*.json, raw history.jsonl per rank, supervisor.jsonl Every rank's 2nd launch: resumed_from = .../checkpoints/micro_00000044, resumed_micro_step = 44; first launches None / 0. Each rank's raw history is micro 1..80 exactly once (80 rows, no duplicates, no restart). Supervisor: attempt 1 exit 1 progressed ckpt 44, attempt 2 exit 0 ckpt 80.
Same data position batch row ids of all 4 ranks x 80 micro-steps, ref vs sup 320/320 identical, and lr identical at micro 40/44/48/80.
Per-rank optimizer files used per_rank_optimizer_state = true in the config (pinned by test); trainer writes optimizer_rank{r}.pt and checkpoint_problem requires one per rank (align/train/trainer.py); each rank's ckpt_rss_rank{r}.jsonl shows rank 0 alone at a 30.2 GB peak and ranks 1-3 at 5.25 GB (a 3.5 GB local write, no gather); the 4 x 3,466,699,433 B sizes match 6.919 B x 2.004 B/param / 4 Consistent. The checkpoint directories themselves were deleted to free 130 GB, so the file names could not be re-listed; this is inference from the trainer code, the config and the RSS records.
Throughput and memory arithmetic 4 x 1024 / 14.79 s = 277 padded tok/s; x 757/1024 = 204 real; x 646/1024 = 175 loss tokens; 277 x 6 x 1.282 B = 2.1 TFLOP/s = 0.53 per GPU; 3600 / 59.2 = 61 steps/h; step-time median over 77 intervals re-derived from history.jsonl (14.795 s ref / 14.797 s sup); alloc peak 14.828 / reserved 17.79 GiB identical on ranks 0-3; nvidia-smi peaks 18,932 / 18,900 / 18,667 / 18,667 MiB; FSDP/EP/128k table arithmetic redone All consistent, except the headroom (3.6 GiB, should be 4.0) and the 8,192-token memory extrapolation (15.4 should be ~15.9): both corrected above.
Telemetry and loss numbers Recomputed from rank-0 and per-rank histories: MaxVio 5.1396 -> 5.0053, worst 6.32 -> 6.10, dead fraction 0.156 -> 0.125, entropy 3.8035 -> 3.7915, lb 1.503 -> 1.4875, z 8.589 -> 8.512; per-step mean losses and first-5 / last-5 means (0.344 / 0.2726); grad-norm list; clip active on 19 of 20 All match.
Resume loss tolerance Re-ran analyze on the recorded pair floor 2.93e-03, resumed 3.30e-03, tolerance 1.086e-02, PASS: matches the doc.

Defects found and fixed (each with a regression test; mutants that remove the fix fail)

  1. The resume verdict could PASS on a resume that never happened (scripts/olmoe_sft_smoke.py). analyze de-duplicated history by micro-step (last wins) and compared loss curves; a relaunch that ignored the checkpoint and retrained from micro 0 yields a de-duplicated curve equal to the reference, so the verdict was PASS. Now analyze also requires check_resume_provenance (second launch loaded micro_{stop} and resumed at the stop step on every rank; raw history advances by exactly one per row) and compare_batches (every rank saw the same row at every micro-step as the uninterrupted run), and exits 1 otherwise. The recorded run satisfies all of it. Tests: test_analyze_passes_a_genuine_resume_and_fails_a_restart_from_zero_..., test_provenance_flags_each_way_a_resume_can_be_fake, test_batch_compare_is_exact_..., test_analyze_fails_when_a_rank_saw_a_different_row_..., test_the_recorded_track_m_run_satisfies_the_provenance_gates; 5 script mutants killed.
  2. The claimed sensitivity of the loss-curve gate was never measured (see the correction in the resume table). The docstring and this document no longer say a wrong optimizer shard moves the loss by 0.1-1. Still open: a mutant resume on the GPUs (zeroed 8-bit moments; base weights instead of the checkpoint) to find out what the 1e-2 tolerance actually detects.
  3. No heap trim after the resume read (13.7 GB stranded on ranks 1-3 until the first save). Trim added; AST test test_heap_trim_is_also_wired_after_the_resume_read; mutant killed. Not re-measured on GPUs.
  4. keep_checkpoints leaked one checkpoint per resume (align/train/trainer.py, a file the track did not own): resume_from adopts the lower-numbered checkpoint directories; higher-numbered zombies are still ignored. Adopted steps go in FRONT of the resumed one because scripts/ppo.py::_latest_weights reads checkpoints_written[-1] as the newest (the first version of the fix appended them after and broke tests/test_ppo.py::test_resume_restores_the_latest_weights_pointer, caught by the full suite). tests/test_trainer.py::test_resume_registers_the_resumed_checkpoint_in_checkpoints_written used to pin exactly [8] and now pins [4, 8]. Test tests/test_control_plane.py::test_resume_adopts_older_checkpoints_so_keep_n_does_not_leak_one_per_resume; 2 mutants killed (no adoption; adopting zombies too).
  5. Doc arithmetic: headroom 3.6 -> 4.0 GiB (MiB/GiB slip) and the 8,192-token allocator extrapolation 15.4 -> ~15.9 GiB. Telemetry-noise caveat added.

Observations, not changed

  • Resume on step 12 showed grad norm 3.35 against 1.42 in the reference even though the batch rows are identical (320/320) and the lr is identical: the difference is in the weights. The track's explanation (replaying the window from the saved micro-44 weights reproduces 3.35 to 0.3%) is consistent with that; it rules out a data-position bug and a kernel nondeterminism, and leaves "weights differ by bf16-level noise and the gradient is discontinuous in them". It does not by itself exclude a lossy checkpoint round trip (bf16 weights are authoritative under stochastic rounding, so there should be none), which only the mutant/identity test in item 2 would settle.
  • README.md, docs/ARCHITECTURE.md and docs/ROADMAP.md status lines and the 25 TFLOP/s / 1.1 B tokens/day constants are stale relative to section (g)/(h); left to the owner, as the track reported.
  • Shared batch stream in dense scripts/sft.py / pretrain.py: confirmed by reading, still unfixed. build_training_sampler returns one unsharded LengthGroupedSampler; Trainer's ResumableBatchStream does list(self.sampler) with no rank slice (only dpo_stage wraps in DistributedBatchSampler). Every rank therefore trains on the same batch indices, FSDP averages four identical gradients, and len(sampler) (the LR horizon) is 4x the per-rank work. Not run (no dense 7B weights); not fixed here because the horizon arithmetic in the entry gates (sampler_derived_horizon) and global_batch_rows semantics must change together, which is the owner's call. Severity: high for any multi-GPU dense run.