SPEED.md

Training speed / MFU — state & open levers

Honest current numbers on 8× B300 (Blackwell Ultra, sm_103, CUDA 13 / torch cu130), so you don't re-walk dead ends. MFU is measured against the achievable bf16 matmul roofline = ~1500 TFLOP/s (a pure n=16384 bf16 matmul on an idle B300, power-capped ~1650 MHz/1070 W). Nominal spec is ~2500 but the cu130 Blackwell kernels only reach ~1500 in a pure matmul, so 1500 is the honest denominator — a real workload cannot beat it. src/mxf/mfu.py computes MFU (numerator counts only real non-pad tokens, so padding lowers it honestly).

Pretrain (scripts/pretrain.py) — the main training loop

config	TFLOP/s	MFU
bs64, no compile	520	35%
bs64 + torch.compile	817	54%
bs96 + compile + fixed-shape (round L to mult-of-64)	860	57% ← best clean single-GPU
sequence packing (2048-block, compile)	708	47% ← regressed, don't use
8-GPU DDP (the actual big run): no compile	~540	~36%

What's been tried / known

• torch.compile fuses the 152k-vocab cross-entropy — the single biggest win (35→54%). Do keep it.
• Length-bucketing + rounding padded L to a multiple of 64 bounds compile to ~2–3 static shapes (otherwise it recompiles every batch — a dynamic-shape trap).
• Sequence packing REGRESSES on this box (57→47%): no flash-attn build for sm_103, so packed attention uses a dense 4D mask and sdpa computes the full block² score matrix, wasting ~90% of it. Correct (bit-exact no-leak proof exists) but slower. It's behind --pack-len (default 0 = off). It'd help only with a block-sparse attention backend.
• ⚠️ THE BIG OPEN BUG: torch.compile + DDP crashes (single-GPU+compile ✓, 8-GPU no-compile ✓, 8-GPU+compile ✗ — all ranks, likely the layer-1 injection forward-hook's graph-break fighting DDP's reducer). So the 8-GPU run currently runs no-compile at ~36% MFU instead of the ~57% compile path. Fixing this is the highest-leverage speed win (compile-after-DDP-wrap, static_graph=True, torch._dynamo DDPOptimizer settings, or moving the injection out of a Python hook into an in-graph module so there's no graph-break).
• No gradient checkpointing (it silently corrupts grads: the inject-hook context exits before the checkpointed backward recompute). Don't enable it without moving injection in-graph. This caps batch at ~bs96–128.

Untried levers (candidates for hill-climbing)

1. Fix compile+DDP → recovers 57% at 8-GPU (biggest single win). Try in-graph injection to kill the graph-break.
2. FP8 (torchao float8 / TE) — Blackwell's design point, ~2× the bf16 roofline. Real code + correctness surface (LoRA + injection), and changes the roofline denominator.
3. FlexAttention block-sparse — would make sequence packing actually help (long fixed blocks → high arithmetic intensity, zero padding); HF-Qwen3 flex integration is the fiddly part.
4. Target-only loss / chunked CE — lm_head is applied to all positions incl. prompt+pad; only target positions need loss. compile already fuses most of it, but computing the head on target rows only saves memory (→ bigger batch) + FLOPs.
5. Reuse-KV for the shared 98-token prompt prefix — but injection differs per row, so prefill activations differ; not shareable as-is.

Embed / clustering (scripts/embed_cluster_acts.py) — forward-only, layer-27 early-exit

Streaming out-of-core; the 8-way sharded corpus download keeps GPUs fed. Forward-only, short (64-tok) seqs → lower arithmetic intensity than training. This is the 200M-doc prefill pass. Live MFU measured from doc throughput; expect it to be memory/IO-influenced (short seq + mfs writes).

Injection (src/mxf/inject.py) — the thing that makes compile/DDP/checkpointing hard

Norm-matched additive: h_p += coeff·‖h_p‖·v/‖v‖ at the layer-1 marker, .detach()ed. It's a register_forward_hook → graph-break under compile. Making it an in-graph module op (subclass the layer, do the injection inside forward with per-row marker positions/vecs passed as tensors) would likely unlock compile+DDP and gradient checkpointing at once — probably the highest-value refactor for speed.

Roofline reproducer

import torch, time
n=16384; a=torch.randn(n,n,device="cuda",dtype=torch.bfloat16); b=a.clone()
for _ in range(50): a@b
torch.cuda.synchronize(); t=time.time()
for _ in range(50): c=a@b
torch.cuda.synchronize(); print(2*n**3*50/(time.time()-t)/1e12, "TFLOP/s")  # ~1500 on B300/cu130