docs(gemma4-erp-tune): root-cause the 8.6% MFU — attention on Ampere kernels, 29.9% padding
Run-01 was killed at step 19 by operator instruction to root-cause before spending a ~13.9h window. Two independent methods now agree on where the step time went, and neither was the hypothesis the consult panel converged on. Scaling fit (3 points, 2 params, residuals <3ms over an 8x range): A = 6.87e-4 s/token, B = 8.85e-8 s/token^2 quadratic share 20.9% @ w=2048 -> 67.8% @ w=16384 No fixed term was needed, which refutes launch-bound outright. Profiler kernel table (device rows only): attention 22,835.8 ms 65.2% fmha_cutlass*_sm80 dense GEMM 2,774.0 ms 7.9% other 5,739.0 ms 16.4% The attention kernels are sm80 — Ampere-generation CUTLASS running on an sm_120 Blackwell card, with the forward on the gmem fallback tier. That is the mechanism behind 100% SM utilisation at 27 of 304 available TFLOPS. Correctness cleared separately: the sliding mask asserts at max 1024 allowed/row, so the 25 windowed layers were genuinely windowed. The same probe found that right-padding is what pins the 5 global layers to an explicit 4D mask and off the is_causal fast path — measured at 9.4% slower for 24% less loss work at fixed width. The largest available win is not the attention kernel. The corpus is 29.9% padding, and bucket-to-pair + shuffle-to-mix takes it to 0.0% for >=35.5% wall clock, no new dependency, unchanged peak memory. Bucket size turned out not to be a diversity knob — roots per accumulation window are flat across a 256x range, so the global micro-batch shuffle does that work alone and the bucket should be tight. Adds docs/pfi/training-throughput-playbook.md as the durable model-agnostic home (sibling to the quantization playbook), the four probes under scripts/training-probes/ with raw output kept for re-derivation, and a §6 to the sizing doc carrying the Gemma-4-specific numbers and round-2 restart parameters. Measured negatives recorded so they are not re-chased: grouped_mm (0.9% slower, and MoE is only 7.9% of the step), CUDA graphs / torch.compile over the expert loop (no fixed cost to amortise), liger fused CE (~1-3% lever), FA4 on sm_120. Round-1 state preserved: 609MB encode cache, order manifest, truncation report, resume script. No checkpoints — it died at step 19 and the first was due at 100, so the lora_B inert-adapter gate never ran and moves to the restart.
This commit is contained in:
@@ -0,0 +1,359 @@
|
||||
# Training throughput playbook — how to find where the step time went
|
||||
|
||||
_Sibling to [`model-quantization-playbook.md`](model-quantization-playbook.md).
|
||||
That one is for making a model small; this one is for making a training run
|
||||
fast. Same contract: **model-agnostic lessons live here, model-specific ones
|
||||
stay in the per-model artifact and link up.**_
|
||||
|
||||
First written 2026-08-24 out of the Gemma-4 26B-A4B ERP/RP tune, which ran at
|
||||
8.6% MFU and cost a four-model frontier panel and most of a night to explain.
|
||||
The worked example in §7 is that run. The lessons above it are not about
|
||||
Gemma-4.
|
||||
|
||||
> **Read this before hypothesising about kernels.** The single most expensive
|
||||
> failure in that investigation was not a wrong hypothesis. It was *four
|
||||
> people, including four frontier models, reasoning confidently from
|
||||
> arithmetic instead of spending ten minutes on a measurement that settled
|
||||
> it.* Two of the panel's conclusions were retracted by their own authors
|
||||
> within the hour. Every retraction was a derivation; every survivor was a
|
||||
> measurement.
|
||||
|
||||
---
|
||||
|
||||
## 1. The 10-minute triage — do this FIRST, always
|
||||
|
||||
Before you profile, before you read a modelling file, before you ask anyone:
|
||||
**measure the step's scaling curve.** Three sequence lengths, fixed batch,
|
||||
fwd+bwd, best-of-2 after a warmup.
|
||||
|
||||
t(w) = A·w + B·w² w = per-sequence length
|
||||
|
||||
Fit two parameters to three points. The residuals tell you which regime you
|
||||
are in, and the regime tells you which lever exists:
|
||||
|
||||
| observed `t(4w)/t(w)` | regime | the lever |
|
||||
|---|---|---|
|
||||
| ~4× | **linear** — per-token work dominates | fewer tokens; fused elementwise |
|
||||
| ~16× | **quadratic** — attention dominates | attention backend / kernel |
|
||||
| ~1× | **launch-bound** — fixed per-batch cost | CUDA graphs, `torch.compile`, bigger batch |
|
||||
|
||||
**If the two-term fit closes with residuals under ~1%, launch-bound is
|
||||
refuted.** You did not need a constant term, so there is not a meaningful one.
|
||||
This is the cheapest possible refutation of the most seductive wrong answer,
|
||||
and it costs one extra data point.
|
||||
|
||||
### ⚠ 1.1 ⭐⭐ Three points minimum. A two-point fit with three plausible terms is UNDETERMINED
|
||||
|
||||
This is the lesson that cost the most. A two-point fit over {quadratic,
|
||||
linear, fixed} has infinitely many solutions, and which one you land on is
|
||||
decided by whichever per-step number you happened to quote. In the worked
|
||||
example a peer produced **two confident, opposite conclusions from the same
|
||||
method inside an hour** — "attention is ~5 s of 35" and then "attention is
|
||||
21–33 s of 35" — because the inputs drifted between attempts.
|
||||
|
||||
Three points, two parameters, and check the residuals. If they do not close,
|
||||
you have a third term and you need a fourth point.
|
||||
|
||||
### ⚠ 1.2 ⭐⭐ Benchmark the shape you RUN, not the worst case you can construct
|
||||
|
||||
The quadratic share is **strongly shape-dependent** — in the worked example it
|
||||
ran 20.9% at w=2,048, 51.3% at w=8,192, 67.8% at w=16,384. A synthetic
|
||||
`max_seq_len` benchmark therefore measures the shape where attention looks
|
||||
worst, and generalising from it overstates the attention prize by ~1.3×.
|
||||
|
||||
Get the real distribution off the encode cache and weight by it:
|
||||
|
||||
E[t] = A·E[w] + B·E[w²]
|
||||
|
||||
**`E[w²]` is not `E[w]²`.** For a bimodal length distribution they can differ
|
||||
by 2× or more, and a quadratic term is dominated by the rare long batches that
|
||||
an `E[w]²` shortcut averages away. In the worked example `E[n²]/E[n]²` was
|
||||
**2.08**.
|
||||
|
||||
Sanity check the weighted prediction against the observed `s/it` before you
|
||||
trust any of it.
|
||||
|
||||
---
|
||||
|
||||
## 2. The reference probe set
|
||||
|
||||
Committed at [`scripts/training-probes/`](../../scripts/training-probes/).
|
||||
Run them in this order; each is minutes and none needs the real checkpoint
|
||||
except the profiler.
|
||||
|
||||
| probe | what it settles | needs GPU? |
|
||||
|---|---|---|
|
||||
| `step0_mask.py` | mask band structure + which layers keep the fast path | no |
|
||||
| `step2_padding.py` | padding waste, length distribution, CE chunk sizing | no |
|
||||
| `step_bucket.py` | bucketing gain, bucket-size sweep, root diversity | no |
|
||||
| `step1_profile.py` | scaling fit, padding penalty, CE wall clock, kernel table | yes |
|
||||
|
||||
`step1_profile.py` loads the real model but reuses the harness's own
|
||||
`discover_target_modules` and `compute_loss`, so it measures the thing that
|
||||
actually runs rather than a re-implementation. **Keep that property when you
|
||||
adapt it** — a probe that reimplements the training step measures the probe.
|
||||
|
||||
---
|
||||
|
||||
## 3. The recurring landmines
|
||||
|
||||
### 3.1 ⭐⭐ Right-padding is a compute tax AND a backend tax
|
||||
|
||||
Everyone knows padding wastes tokens. The second effect is the one that gets
|
||||
missed: **an explicit padding mask can knock fast-path-eligible layers off
|
||||
`is_causal`.**
|
||||
|
||||
`scaled_dot_product_attention` takes `is_causal=True` **or** an `attn_mask`,
|
||||
never both usefully. HF sets `is_causal=True` only when `attention_mask is
|
||||
None`. Right-pad a batch and you hand it a 2D mask, it materialises a 4D
|
||||
tensor, and every layer that could have taken the clean causal route now takes
|
||||
a masked dense one.
|
||||
|
||||
Measured, same width, same `n`, only the mask differing:
|
||||
|
||||
no padding 35.244 s 26,048 loss targets
|
||||
50% pad on one row 38.567 s 19,640 loss targets
|
||||
|
||||
**9.4% slower for 24% less work.** Verify this on your own stack with
|
||||
`step0_mask.py` — it prints whether `create_causal_mask` returns `None` or a
|
||||
tensor for each mask case.
|
||||
|
||||
### 3.2 ⭐⭐ Length-bucket to PAIR, shuffle micro-batches to MIX — and the bucket should be TIGHT
|
||||
|
||||
Naive length-bucketing has a real hazard: length correlates with data source,
|
||||
so length-homogeneous batches are **source-homogeneous batches**, and an
|
||||
accumulation window can end up drawing its entire gradient from one root.
|
||||
|
||||
The fix costs nothing: **form micro-batches within length buckets, then
|
||||
shuffle the resulting micro-batches globally.** Padding efficiency is a
|
||||
property of the pairing alone, so all of the saving survives the shuffle.
|
||||
|
||||
**The non-obvious part — bucket size is not a diversity knob.** Measured
|
||||
across a 256× range of bucket sizes, roots per accumulation window stayed flat
|
||||
at 3.54–3.61 (against 3.68 for a pure shuffle). The *global micro-batch
|
||||
shuffle* does all of the mixing; the bucket contributes nothing to diversity
|
||||
and only costs padding. So use the tightest bucket you can — which in the
|
||||
limit is a full length sort.
|
||||
|
||||
| bucket | padding waste | zero-pad micro-batches | roots/window |
|
||||
|---|---|---|---|
|
||||
| current (shuffle) | 29.9% | 0.1% | 3.68 |
|
||||
| 2 | 0.0% | **78.3%** | 3.56 |
|
||||
| 32 | 0.1% | 41.9% | 3.55 |
|
||||
| 512 | 2.4% | 4.1% | 3.61 |
|
||||
|
||||
Note the `zero-pad micro-batches` column — that is §3.1 compounding. A tight
|
||||
bucket does not merely cut tokens, it puts most batches back on the causal
|
||||
fast path.
|
||||
|
||||
**Peak memory does not rise.** `padded = batch × max(len)`, so one long record
|
||||
forces a full-width batch regardless of its partner. Bucketing pairs long
|
||||
records *with each other*, which roughly halves the number of worst-case
|
||||
batches.
|
||||
|
||||
### 3.3 ⭐⭐ Check the kernel GENERATION, not just the backend name
|
||||
|
||||
The backend name (`EFFICIENT_ATTENTION`, `FLASH_ATTENTION`, …) is not the whole
|
||||
story. Read the actual kernel symbols out of the profiler:
|
||||
|
||||
fmha_cutlassF_bf16_aligned_32x128_gmem_sm80
|
||||
fmha_cutlassB_bf16_aligned_128x64_k65536_sm80
|
||||
^^^^
|
||||
|
||||
`sm80` is **Ampere**. Those were running on an sm_120 Blackwell card, on the
|
||||
dominant cost centre of the step. A backend can be "selected correctly" and
|
||||
still be a generation behind, and nothing in the config surface tells you.
|
||||
|
||||
Also read the variant suffix: `gmem` on the forward kernel is the
|
||||
**global-memory fallback tier** of the memory-efficient path, chosen when the
|
||||
working set will not fit in shared memory. Wrong backend *and* that backend's
|
||||
slow path.
|
||||
|
||||
### 3.4 ⭐ `key_averages()` double-counts — use device-kernel rows only
|
||||
|
||||
`torch.profiler`'s `key_averages()` table lists both the ATen op and the CUDA
|
||||
kernel it launched, each carrying the same `self_device_time_total`. Summing
|
||||
the whole table gives you roughly **2× the real step time**.
|
||||
|
||||
The tell is exact equality between an `aten::` row and a kernel row:
|
||||
|
||||
aten::_efficient_attention_backward 30 16144.6
|
||||
fmha_cutlassB_bf16_aligned_128x64_k65536 30 16144.6
|
||||
|
||||
Filter to device kernels (`void …`, `fmha_…`, `cutlass::…`, `Memcpy…`) and
|
||||
sanity-check the total against the measured wall clock. In the worked example
|
||||
the filtered total came to 89.5% of the step, which is the right shape; the
|
||||
unfiltered total came to 202%.
|
||||
|
||||
### 3.5 ⭐ Time the loss forward AND account for its backward recompute separately
|
||||
|
||||
If the loss head is gradient-checkpointed, a CUDA-event window around the
|
||||
forward loop measures **half the story at best** — the recompute happens inside
|
||||
`.backward()`, outside your window.
|
||||
|
||||
State the caveat explicitly when you report the number. In the worked example
|
||||
the CE forward measured 374 ms of a 35.3 s step (1.1%); even at 3× for
|
||||
recompute-plus-backward it is ~3%, which was enough to kill a proposed
|
||||
dependency swap — but "1.1%" alone would have been an unearned claim.
|
||||
|
||||
### 3.6 ⭐⭐ Assert mask band structure directly; never infer it from performance
|
||||
|
||||
`transformers` can **silently skip mask creation** and pass
|
||||
`attention_mask=None` when a mask function is not registered. If that fires on
|
||||
a sliding-window model, the windowed layers do full causal attention — not a
|
||||
speed bug, **a different model from the one you will serve**.
|
||||
|
||||
There is a tempting alibi: "if constraints were dropped we would be on the
|
||||
fast path and fast; we are slow, therefore correct." It is decent evidence and
|
||||
it is not an assertion. Materialise the mask once and count allowed positions
|
||||
per row:
|
||||
|
||||
sliding_attention max 1,024 allowed/row, saturates at row 1,023 PASS
|
||||
|
||||
Thirty seconds, on CPU, no weights. Do it before every run that changes the
|
||||
masking path, and before believing any optimisation result.
|
||||
|
||||
### 3.7 ⭐ "Bit-identical output from a different backend" — ask *could this have disagreed?*
|
||||
|
||||
A backend flag that produces `max_abs_diff == 0.0` against the reference is
|
||||
either (a) legitimately the same GEMMs behind a different launcher, or (b) a
|
||||
flag that never took. **Argument cannot separate these** — in the worked
|
||||
example three frontier models split 2–1 on it and the majority was not
|
||||
obviously right.
|
||||
|
||||
Do not resolve it by vote. **Count kernel launches.** A per-expert loop leaves
|
||||
`n_experts` dispatches per layer visible; a grouped path leaves one. That is
|
||||
unambiguous and falls out of a trace you are running anyway.
|
||||
|
||||
Related trap: **a trace of the default path does not test the flag.** If the
|
||||
run was relaunched without the flag set, the profile tells you what the default
|
||||
does and nothing about the flag. Say so rather than over-claiming.
|
||||
|
||||
### 3.8 ⭐ MFU is a denominator argument waiting to happen — report the decomposition instead
|
||||
|
||||
MFU invites an unwinnable fight about what counts as a FLOP (active vs dense
|
||||
params for MoE, whether checkpoint recompute counts, whether frozen-base
|
||||
skipped GEMMs count). That fight consumed an hour of the worked example and
|
||||
produced nothing.
|
||||
|
||||
Report these **beside** MFU, not instead of it:
|
||||
|
||||
- tokens/s, and **real (unpadded) tokens/s** separately
|
||||
- achieved hardware FLOPs straight from the profiler
|
||||
- the time decomposition (attention / GEMM / elementwise / other)
|
||||
|
||||
Then the denominator stops mattering.
|
||||
|
||||
**The reading that actually diagnosed it** was not an MFU number at all:
|
||||
|
||||
> 100% SM utilisation at 279–292 W, running 27 TFLOPS, on a card that does
|
||||
> 304 TFLOPS on a dense GEMM at the same power.
|
||||
|
||||
**SM-busy, tensor-core-idle.** The chip is fully occupied doing work that is
|
||||
not matrix multiplication. No FLOP-counting convention changes that, and it
|
||||
points straight at the kernel table.
|
||||
|
||||
### 3.9 Frozen-base LoRA is ~4ND, not ~6ND — and the arithmetic intensity does NOT drop
|
||||
|
||||
A claim that circulated and was wrong: "frozen-base LoRA has structurally lower
|
||||
arithmetic intensity, so a dense-GEMM ceiling is unreachable in principle."
|
||||
|
||||
The correct accounting: forward is 2ND, input-gradient backward through the
|
||||
frozen weights is 2ND, and only the weight-gradient (~2ND) is skipped. So
|
||||
**~4ND against ~6ND — two-thirds of the work, at the same arithmetic intensity
|
||||
per remaining GEMM.** You do fewer GEMMs; the ones you do are exactly as dense.
|
||||
|
||||
Gradient checkpointing is a separate, real ~⅓ recompute tax. Account for it
|
||||
separately rather than folding it into an intensity story.
|
||||
|
||||
---
|
||||
|
||||
## 4. Panel / consult discipline for perf work
|
||||
|
||||
Perf investigations are unusually good at generating confident wrong answers,
|
||||
because the arithmetic is easy and the ground truth is expensive. Specific
|
||||
guards, learned the hard way:
|
||||
|
||||
- **Every arm's claim gets a measurement or an expiry date.** In the worked
|
||||
example the panel produced four self-retractions in ninety minutes. The
|
||||
measurements produced zero.
|
||||
- **Treat cross-arm agreement as weak evidence.** Ask arms to attack a
|
||||
hypothesis rather than extend it; agreement among similarly-primed readers of
|
||||
the same artifact is not independent confirmation.
|
||||
- **A dispute about what a specific dispatcher does is a question of fact.**
|
||||
Do not put it to a panel. Instrument it.
|
||||
- **When an arm says "you missed X," check what they read.** If your settled
|
||||
artifact was not on their reading list, the "miss" is usually
|
||||
restatement-of-a-settled-prior, not a genuine gap.
|
||||
|
||||
---
|
||||
|
||||
## 5. Superseded claims — do not follow these
|
||||
|
||||
| claim | status | replaced by |
|
||||
|---|---|---|
|
||||
| "Explicit mask → `EFFICIENT_ATTENTION`" is over-specific; Blackwell defaults to `CUDNN_ATTENTION` | **WRONG** (2026-08-24) | Measured: sm_120 selects `fmha_cutlass*_sm80`, i.e. `EFFICIENT_ATTENTION`. The original claim was right. |
|
||||
| Attention's quadratic share is ~5 s of a 35 s step | **WRONG** (2026-08-24) | Measured 22.8 s / 65.2% at w=16,384; 67.8% by independent scaling fit |
|
||||
| Frozen-base LoRA has structurally lower arithmetic intensity | **WRONG** (2026-08-24) | ~4ND vs 6ND at unchanged intensity — see §3.9 |
|
||||
| The chunked CE is a 2–5× under-estimated cost centre | **WRONG** (2026-08-24) | Measured 1.1% of step forward, ≲3% with recompute |
|
||||
| `attn_implementation="flash_attention_2"` is the per-layer lever | **NOT A FLAG** (2026-08-24) | All-or-nothing at `from_pretrained`; per-layer needs a custom fn on `ALL_ATTENTION_FUNCTIONS`. FA2 also caps head_dim at 256. |
|
||||
| Bucket size ~256 is needed to preserve source diversity | **UNNECESSARY** (2026-08-24) | Diversity is flat in bucket size; the global micro-batch shuffle does that work — see §3.2 |
|
||||
|
||||
## 6. Measured negatives — don't re-chase
|
||||
|
||||
- **Fused MoE kernel (`grouped_mm`) as the throughput fix.** Measured 0.9%
|
||||
*slower* than the Python loop and bit-identical. Independently, dense GEMM is
|
||||
only 7.9% of the step, so the whole category is capped near 10%.
|
||||
- **CUDA graphs / `torch.compile` over the expert loop.** The two-term scaling
|
||||
fit closed without a constant term, so there is no meaningful fixed per-batch
|
||||
cost to amortise. ~3,840 expert-GEMM launches per forward are not what you
|
||||
are paying for.
|
||||
- **`liger-kernel` fused linear CE.** Real and correct, but a ~1–3% lever on
|
||||
this shape. Not a project.
|
||||
- **FlashAttention-4 on sm_120.** Public reports are sour — one measurement of
|
||||
1.07× over FA2, and an sm_120 patch people could not get working that fell
|
||||
back to torch SDPA. Do not bet a round on it.
|
||||
|
||||
---
|
||||
|
||||
## 7. Worked example — Gemma-4 26B-A4B ERP/RP tune, 2026-08-24
|
||||
|
||||
Model-specific detail lives in
|
||||
[`gemma4-erp-tune-sizing.md`](gemma4-erp-tune-sizing.md) §6. The short version,
|
||||
because the *shape* of the investigation is the transferable part:
|
||||
|
||||
**Symptom.** 8.6% MFU, ~35–46 s/it, 1,312 steps, ~13.9 h ETA.
|
||||
|
||||
**What the panel produced.** Four frontier arms plus an orchestrator, over
|
||||
ninety minutes: a sliding-window hypothesis, a retraction of it, a retraction
|
||||
of the retraction, a correctness scare that resolved itself, two mutually
|
||||
contradictory readings of one dispatcher, and four self-corrections.
|
||||
|
||||
**What settled it, in about twenty minutes of GPU time:**
|
||||
|
||||
scaling fit (3 points, 2 params, residuals <3 ms over an 8× range)
|
||||
A = 6.87e-4 s/token B = 8.85e-8 s/token²
|
||||
quadratic share: 20.9% @ w=2,048 → 67.8% @ w=16,384
|
||||
|
||||
kernel table (device rows only)
|
||||
attention 22,835.8 ms 65.2% fmha_cutlass*_sm80
|
||||
dense GEMM 2,774.0 ms 7.9%
|
||||
other 5,739.0 ms 16.4%
|
||||
|
||||
Two independent methods, 2.6 points apart. **Attention was the answer, on
|
||||
Ampere-generation kernels, with the forward on a global-memory fallback tier.**
|
||||
|
||||
**The largest actionable win was not the attention kernel.** It was a sampler
|
||||
change — bucket-to-pair, shuffle-to-mix — worth 29.9% of tokens and ~35.5% of
|
||||
wall clock, with no new dependency, no kernel work, and unchanged peak memory.
|
||||
It also wins under *every* branch of the diagnosis, which is why it was
|
||||
recommended while the rest was still unresolved.
|
||||
|
||||
**The transferable ordering:**
|
||||
|
||||
1. Assert correctness (mask band structure). Everything downstream assumes it.
|
||||
2. Scaling curve. Names the regime in ten minutes.
|
||||
3. Kernel table. Names the cost centre.
|
||||
4. Data-side levers first (padding, bucketing) — they need no dependency and
|
||||
they multiply into every other cost.
|
||||
5. Kernel/backend levers last, gated on 2 and 3.
|
||||
Reference in New Issue
Block a user