File size: 59,810 Bytes
d909bd8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
# 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, SiLU*up, 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 = 2*T*k*(d*2I + I*d) = 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.