Download code/docs/MEASUREMENTS.md from viai957/align: direct link, hf CLI and curl.
- Browser
- Download file 59.8 kB
-
https://huggingface.co/viai957/align/resolve/main/code/docs/MEASUREMENTS.md
- Command line
-
hf download hf://viai957/align/code/docs/MEASUREMENTS.md
-
curl -L -o MEASUREMENTS.md https://huggingface.co/viai957/align/resolve/main/code/docs/MEASUREMENTS.md
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.
TORCH_CPP_LOG_LEVEL=INFOon a 2-call probe prints, on every call:[I... GroupedMMUtils.h:128] fallback path fortorch._grouped_mm, performance may not be optimal(probe_log.err). The same string is inlibtorch_cuda.sonext to a_grouped_mm_fallbacksymbol.- The profiler on one
F.grouped_mm(x_sorted, gate_up^T, offs)call (T=4096 balanced) shows 64aten::mmand 64aten::itemCPU ops and 93 CUDA kernels: 64 cuBLAS GEMMs (36xampere_bf16_s16816gemm_bf16_128x128_ldg8_f2f_stages_32x5_tn, 28xampere_bf16_s1688gemm_bf16_128x128_ldg8_f2f_stages_32x1_tn), 28 memsets, and 1Memcpy DtoH(theoffsread-back that forces a host sync). No CUTLASS grouped kernel launches. - 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
- P0-6 (rank-local logging) is confirmed on real ranks.
history.jsonlreceives 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. TORCH_NCCL_TRACE_BUFFER_SIZEis 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_diagnosticssets 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 (doesDUMP_ON_TIMEOUTneedTORCH_NCCL_ENABLE_MONITORING) still needs a deliberate hang, not done here.- 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'sTMPDIRpoints at the root filesystem.driver.pynow unsets it soapply_resource_guardssets/apps/.../align/.tmp. Sections (c) and (d) were single-process scripts that did not go throughapply_resource_guards(no heavy-job lock, inheritedTMPDIR); they wrote nothing outside.tmp/gpu_smoke/. torchrunreaps the survivor. After rank 1 exited (143) torchrun loggedSending process ... closing signal SIGTERMfor rank 0 and exited 1. The checkpoint was complete becausesave_checkpointends 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 aftertrain()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.- 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
RouterStatsthrough 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 -9during 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.
agreefigures 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)
scripts/sft.py/scripts/pretrain.py(dense) give every rank the same batch stream:build_training_samplerreturns one unshardedLengthGroupedSamplerand the Trainer does no rank slicing (onlydpo_stagewraps inDistributedBatchSampler). 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).keep_checkpointsleaks one checkpoint per resume (13.8 + 13.9 GB): pruning is driven by a per-process list. Fixed in review (see the table above).- The
history.jsonlcollision (P0-6) is still live in the Trainer; the smoke gives ranks 1-3 their ownrank{r}/history.jsonland rank 0 the one the supervisor watches. - 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. /appsis 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).- Host-memory fix beyond the brief:
malloc_trimafter 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. - The shared heavy-job lock did its job against another track's long job (refused with exit 83, then waited).
torchcould 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_mmend 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 -9mid-save; a second large job alongside (the 45 GB floor under real contention); GRPO / DPO on OLMoE. - Throughput numbers are for
loopdispatch withsdpa;grouped_mmand 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)
- The resume verdict could PASS on a resume that never happened (
scripts/olmoe_sft_smoke.py).analyzede-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. Nowanalyzealso requirescheck_resume_provenance(second launch loadedmicro_{stop}and resumed at the stop step on every rank; raw history advances by exactly one per row) andcompare_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. - 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.
- 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. keep_checkpointsleaked one checkpoint per resume (align/train/trainer.py, a file the track did not own):resume_fromadopts the lower-numbered checkpoint directories; higher-numbered zombies are still ignored. Adopted steps go in FRONT of the resumed one becausescripts/ppo.py::_latest_weightsreadscheckpoints_written[-1]as the newest (the first version of the fix appended them after and broketests/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_writtenused to pin exactly[8]and now pins[4, 8]. Testtests/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).- 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.mdanddocs/ROADMAP.mdstatus 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_samplerreturns one unshardedLengthGroupedSampler;Trainer'sResumableBatchStreamdoeslist(self.sampler)with no rank slice (onlydpo_stagewraps inDistributedBatchSampler). Every rank therefore trains on the same batch indices, FSDP averages four identical gradients, andlen(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) andglobal_batch_rowssemantics must change together, which is the owner's call. Severity: high for any multi-GPU dense run.