An external reference for the trainer's MFU: the same model and precision policy in PyTorch — +12.4% at B=16, +0.99% at the matched B=32

[feitreim]

The trainer's 81,532 tokens/s and 36.20% MFU at B=16 had nothing to be a percentage of. This adds the thing they are a percentage of: the same 4.39B model in PyTorch, fed the same tokens off the same shard, measured over the same twelve steps with the first discarded, and divided by the same training_flops_per_token() over the same 2.25 PFLOP/s. No PyTorch FLOP counter is anywhere near the table.

The table — matched precision policy

optim's policy, reproduced term for term: bf16 masters for every matrix-shaped parameter (ParameterKind::{Embedding, Matrix, Head}), fp32 norms and router (::Norm, ::Router), fp32 AdamW moments, and the update arithmetic in fp32 on a widened master that rounds back to bf16 on commit. One B200, B=16, torch 2.13.0+cu130, wiki-val-00000.tok.

Tier	tokens/s	MFU	vs trainer	peak VRAM	CUDA graphs	compile
trainer (bin/train.rs, main)	81,532	36.20%	—	—	no	—
PyTorch eager	49,282	21.88%	−39.6%	132.6 GiB	no	30.9 s
torch.compile()	91,669	40.71%	+12.4%	95.5 GiB	no	38.0 s
torch.compile(mode="reduce-overhead")	91,289	40.54%	+12.0%	95.5 GiB	yes	23.1 s

Compile time is excluded from throughput — compiled tiers get three warmup steps before the timed window opens, eager gets one because its optimizer step is compiled — and reported here instead.

The line to answer is +12.4%. mode="reduce-overhead" is CUDA graphs, which the trainer does not have, and it is labelled as such everywhere — but it lands 0.4% behind plain torch.compile(), inside the noise of an 11-step window on a 357 ms step. Whatever the gap is made of, it is not graphs. The script asks inductor whether graphs were actually recorded rather than trusting the mode name: cuda_graphs=False for default, True for reduce-overhead.

mode="max-autotune" remains unrun. The escalation continues only while the trainer is ahead, and 81,532 is not ahead of 91,669.

The rematch at B=32

The table above pits a B=16 trainer against a B=16 baseline, and both halves of
that have since moved. #108-#114 cut the trainer's cost per unit of batch from
5.33 GiB to 3.39, #117 spent the room on B = 32, and main (c5a3b1a) now
reports 92,718 tokens/s / 41.17% MFU. So the baseline was re-run at the
matched batch — same policy, same shard, same twelve-steps-minus-one window,
all three tiers in one container.

Tier	tokens/s	MFU	vs trainer	peak VRAM	CUDA graphs	compile
trainer (bin/train.rs, c5a3b1a)	92,718	41.17%	—	174.05 GiB	no	—
PyTorch eager	—	—	—	out of memory	—	—
torch.compile()	93,634	41.58%	+0.99%	150.1 GiB	no	102.4 s
torch.compile(mode="reduce-overhead")	94,348	41.89%	+1.76%	150.0 GiB	yes	34.2 s

Same gates as before, and they all held at the new batch: parameter count
asserted term by term against the Rust model, 174 of 174 moments fp32, losses
11.17 → 8.48 finite on the real shard, and inductor asked whether it recorded
graphs rather than trusted by mode name.

The trainer did not cross

It closed 11.4 points of a 12.4-point gap and stopped 0.99% short. Against
#117's conservative same-container number (90,718) the gap is 3.2%. Both point
estimates put torch.compile() ahead, so by the standing escalation rule
mode="max-autotune" stays locked — it is still unrun, and what unlocks it
is a trainer that beats 93,634.

The 0.99% deserves its caveat rather than a victory lap in either direction: it
is below the 1.4% spread #117 measured between two containers running the same
trainer binary, and the trainer's number comes from that PR's container while
these come from today's. At matched batch the two are a coin-flip the coin has
not landed on. What is not a coin-flip is the direction of travel, and +12.4% →
+0.99% is the whole of the interesting result.

What changed shape

• Eager no longer fits. It reached 174.48 GiB of the card's 178.35 and died
  allocating 6.14 more. Its footprint was already the largest row in the B=16
  table and the batch doubled underneath it. An OOM is a measurement, so it gets
  a row.
• reduce-overhead is now 0.76% ahead of default, where at B=16 it was
  0.4% behind. Both margins are sub-1% on an 11-step window, so this is a sign
  flip inside the noise rather than graphs suddenly paying. It stays excluded
  from the headline for the reason it always was: the trainer has no CUDA graphs
  and cuda-oxide cannot give it any. cuda_graphs=False for default, True
  here, both verified.
• Compile time is order-dependent, not tier-dependent. 102.4 s for default
  against 34.2 s for reduce-overhead is inductor's cold cache, not a property
  of the modes — default ran first in the container, exactly as it did at B=16
  (38.0 s vs 23.1 s).

Memory

95.5 → 150.1 GiB over 16 → 32 is 3.41 GiB per unit of batch, against the
trainer's measured 3.39 — the two scale alike, which is the sanity check worth
having. So PyTorch's 24 GiB of headroom at B=32 is all intercept: 40.9 GiB fixed
against the trainer's 65.5. Part of that is real and already named below — bf16
gradients on the 50 matrix parameters save ~8.8 GiB the trainer spends on fp32.
The rest is not a like-for-like subtraction, because max_memory_allocated()
counts what the caching allocator handed out and not the CUDA context or cuBLAS
workspaces beside it. The slopes are measured by the same instrument; the
intercepts are not.

One container, or two

Two, and the harness cannot make it one. torch_baseline is the only Modal
function on torch_image; every trainer entrypoint runs on the kernel image;
neither image contains the other's toolchain. Pairing them in a single container
means one image carrying both torch and the whole cuda-oxide backend — the
multi-gigabyte tax this baseline was split out to avoid in the first place. The
trainer's 92,718 is #117's sweep container and 90,718 is its train_ab
container; the two PyTorch rows are today's.

modal run modal_app.py::pytorch_baseline --tiers ",default,reduce-overhead" --batches 32

Policy verification

The script refuses to report a number unless the moments are fp32, and prints what it actually holds rather than what it meant to:

dtypes {"parameters": {"bfloat16": 50, "float32": 37},
        "gradients":  {"bfloat16": 50, "float32": 37},
        "moments":    {"float32": 174}}

87 parameters: 50 bf16 masters (embedding, lm-head, qkv, o_proj, gate_up, down), 37 fp32 (24 block norms + 12 routers + the final norm), and 174 of 174 moments fp32. Parameter count is asserted against the Rust model's term by term first:

{"embedding": 154389504, "lm_head": 154927104, "final_norm": 3072,
 "blocks": 4077232128, "expected_total": 4386551808, "actual_total": 4386551808}

Losses run 11.18 → 8.42 over the window on the real shard, finite, checked.

The stock autocast policy, kept for the record

The first pass measured fp32 masters under bf16 autocast with AdamW(fused=True) — the default a practitioner reaches for, carrying a handicap the trainer does not: a weight cast per matmul and twice the parameter and gradient bytes. Those numbers stay here rather than being replaced.

Tier	B	tokens/s	MFU	vs trainer @ B=16	peak VRAM	CUDA graphs
PyTorch eager	16	44,635	19.82%	−45.3%	148.7 GiB	no
PyTorch eager	12	46,668	20.72%	−42.8%	125.8 GiB	no
torch.compile()	16	88,840	39.45%	+9.0%	111.6 GiB	no
torch.compile(mode="reduce-overhead")	16	89,541	39.76%	+9.8%	111.5 GiB	yes

Matching the policy is worth +3.2% and 16.1 GiB against that — the whole of the handicap and not much more. The trainer's reference at B=12 (79,200 / 35.17%) was measured against the fp32-master eager tier only; the matched policy was run at B=16, the recommended config.

Where the policies still cannot meet

• Gradients — favors PyTorch. The trainer keeps fp32 gradients for every parameter; its weight-gradient GEMMs write fp32 directly beside a bf16 master. Autograd produces gradients in the parameter's dtype, so the 50 bf16 masters get bf16 gradients (the 37 fp32 parameters do match). Half the gradient bytes to write, read and reduce. It costs accuracy rather than speed — most visibly in the embedding's scatter-add, which accumulates in bf16 where the trainer accumulates in fp32. Closing it means a custom autograd function behind every matmul: a heroic port, not a baseline.
• Optimizer fusion. torch.optim.AdamW sizes moments after the parameter, so a bf16 master would get bf16 moments — not the policy. MasterAdamW keeps them fp32 and transcribes optim::adamw_step directly (epsilon on the corrected second moment, decoupled decay, copy_ into bf16 as MasterStorage::commit's round-to-nearest). Its step is torch.compiled because PyTorch ships no fused kernel for a mixed parameter/moment dtype, and that is not a flourish: the same arithmetic uncompiled is eleven passes over 4.39B parameters, 77 ms a step, dragging torch.compile() from 91,669 to 75,490 — a number that measures this file rather than the trainer's kernels. Every row above has a fused optimizer one way or the other.
• The classifier — favors the trainer. Its fused classifier never materializes fp32 logits over the padded vocabulary; F.cross_entropy does, in both memory and bandwidth.
• Everything else matches. Norms, router, moments and the update math are fp32 on both sides. Under bf16 masters the embedding lookup is already bf16, so the residual-stream cast the fp32 policy needed is a no-op — the one concession PyTorch used to get is gone.

Fairness of the model itself

Transcribed from crates/nn/ and gpu/model/src/lib.rs rather than approximated: bias-free linears, fused [D, 3, D] QKV, RMSNorm at 1e-5 with fp32 statistics, interleaved-pair RoPE at theta 10000 computed in fp32, causal SDPA on the flash backend (runtime-probed, cuDNN fallback that says so), SwiGLU experts as [E, D, 2FF] / [E, FF, D] batched GEMMs, AdamW at lr 3e-4 / wd 0.1 / betas (0.9, 0.999) / eps 1e-8, aux-loss decaying 1e-2 linearly to zero at 10,000 steps.

The lm-head is padded to VP = 50,432 because the Rust pads it, because the MFU formula's D * VP counts it, and because 50,257 would cost cuBLAS its aligned epilogue; the loss still normalizes over VOCAB alone, as the fused classifier does with its VOCAB/VP pair.

Routing is the part most easily strawmanned, so it is copied most literally from nn::MoeFfn and moe_bin_assign: top-2 of 8 on renormalized softmax probabilities, capacity slots handed out in flattened (token, rank) order, overflow parked on one scratch row that no expert reads and no gradient leaves. At a capacity factor of one both sides compute E * C expert rows whatever the routing looks like — so load balance cannot move either number, and random tokens would have measured the same thing. The real shard is used anyway, because only that also gives a loss curve worth reading.

Shape of the change

gpu/model/baselines/pytorch_baseline.py is new and standalone. modal_app.py gains a second image and a pytorch_baseline entrypoint; the baseline shares nothing with the kernel toolchain — no rustc, no LLVM, no codegen backend — so adding torch to the image every kernel run pulls would have been a multi-gigabyte tax on runs that never import it. Nothing the trainer builds or runs is touched.

modal run modal_app.py::pytorch_baseline                             # eager, B=16
modal run modal_app.py::pytorch_baseline --tiers "default" --batches 16
modal run modal_app.py::pytorch_baseline --tiers "reduce-overhead" --batches 16
modal run modal_app.py::pytorch_baseline --masters fp32              # stock autocast