| # 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 |
| ```python |
| 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 |
| ``` |
| |
