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:
vh
2026-08-24 22:10:51 -07:00
parent 872c2c562f
commit 7b5fd91d3c
9 changed files with 1239 additions and 0 deletions
+214
View File
@@ -321,3 +321,217 @@ Also fold in:
`vllm/vllm-openai:v0.27.1` already on disk. **If it still no-ops, the harness
must emit merged weights** — and Eitri needs that requirement while he is
early, not after the run.
---
## 6. Round-1 aborted; throughput root-caused (measured 2026-08-24 22:00 PDT)
Run-01 launched, reached step 19 of 1,312 at ~35–46 s/it, and was **killed by
operator instruction** — not a crash, not an OOM. ETA was ~13.9 h at 8.6% MFU
and the operator elected to root-cause before spending the window.
Nothing was destroyed: the 609 MB encode cache, `order-manifest.jsonl`,
`truncation-report.json` and `resume-run-01.sh` are all preserved at
`/tank/erp-tune/run-01/`. **There are no checkpoints** — the first was due at
step 100, so brokkr's `lora_B` inert-adapter gate never ran. That question is
open and moves to the restart.
Model-agnostic lessons from this investigation are in
[`training-throughput-playbook.md`](training-throughput-playbook.md); the
probes are at [`scripts/training-probes/`](../../scripts/training-probes/).
What follows is Gemma-4-specific.
### 6.1 Where the step time goes
Real checkpoint, GPU0, `attn_implementation="sdpa"`, PEFT + gradient
checkpointing + the chunked CE, fwd+bwd, best-of-2 after warmup:
| shape | time | peak |
|---|---:|---:|
| 2 × 2,048 | 1.776 s | 53.2 GiB |
| 2 × 8,192 | 11.570 s | 62.3 GiB |
| 2 × 16,384 | **35.017 s** | 76.6 GiB |
Fitting `t(w) = A·w + B·w²` over all three (per-sequence `w`, batch 2):
A = 6.8715e-04 s/token B = 8.8509e-08 s/token²
| w | predicted | measured | linear | quadratic | quad share |
|---:|---:|---:|---:|---:|---:|
| 2,048 | 1.779 | 1.776 | 1.407 | 0.371 | 20.9% |
| 8,192 | 11.569 | 11.570 | 5.629 | 5.940 | 51.3% |
| 16,384 | 35.017 | 35.017 | 11.258 | 23.759 | **67.8%** |
**Two terms, three points, residuals under 3 ms across an 8× range.** No fixed
per-batch term was needed, which refutes the launch-bound hypothesis outright —
~3,840 expert-GEMM launches per forward are not the cost.
Independently, the profiler kernel table (device rows only — see playbook §3.4):
| device kernel | ms | of step |
|---|---:|---:|
| `fmha_cutlassB_bf16_aligned_128x64_k65536_sm80` (attn BWD) | 16,144.6 | 46.1% |
| `fmha_cutlassF_bf16_aligned_32x128_gmem_sm80` (attn FWD) | 6,691.2 | 19.1% |
| `cutlass_80_tensorop_bf16_s16816gemm` ×3 (dense GEMM) | 2,774.0 | 7.9% |
| elementwise / vectorized / unrolled ×9 | 4,787.4 | 13.7% |
| gather / Memcpy DtoD / dropout | 951.6 | 2.7% |
| **attention total** | **22,835.8** | **65.2%** |
**Scaling fit says 67.8% quadratic; kernel table says 65.2% attention. Two
independent methods, 2.6 points apart.**
### 6.2 ⚠ The attention kernels are Ampere, on a Blackwell card
`fmha_cutlass*_sm80` on sm_120. There is no Blackwell-tuned attention kernel in
this path at all, and the forward is additionally on `gmem` — the
global-memory fallback tier of the memory-efficient backend, selected when the
working set will not fit in shared memory.
This is the mechanism behind the 100%-SM / 27-TFLOPS / 304-TFLOPS-capable
reading: the chip is saturated running a kernel generation behind on the
dominant cost centre.
The candidate fix is a purpose-built kernel for this architecture's mixed
256/512 head-dim split — `zzhhjjj/gemma-triton-flash-attn`
(`register_triton_attention()`, then `_attn_implementation = "triton_gqa"`),
reported 9.23× over SDPA at N=16K D=256 SWA and 2.94× fwd+bwd at D=512.
`flex_attention` + `BlockMask` is the no-new-dependency alternative.
⚠ **Prefer a UNIFORM backend over a per-layer split.** vLLM special-cased this
exact mixed-head-dim architecture and measured mixed backends **8% slower** than
uniform. And `attn_implementation` is all-or-nothing at `from_pretrained` /
`set_attn_implementation` — per-layer routing requires a custom function
registered on `ALL_ATTENTION_FUNCTIONS` branching on `module.head_dim` /
`sliding_window`.
⚠ **FA2 is not available for this model**: it caps head_dim at 256 and the 5
global layers are at 512. FA3 is Hopper-only. Do not bet on FA4 on sm_120.
### 6.3 Masking is CORRECT — and padding is what costs
Band structure asserted directly against the real config at n=16,384:
sliding_attention max 1,024 allowed/row, saturates at row 1,023 PASS
Constraints were **not** silently dropped; the 25 sliding layers were genuinely
windowed. Run-01 was training the model we intended.
The same probe found the mechanism nobody had measured:
| 2D mask supplied | `full_attention` mask returned |
|---|---|
| `None` | **`None`** → `is_causal` fast path AVAILABLE |
| all-ones (no padding) | **`None`** → `is_causal` fast path AVAILABLE |
| right-padded (what `collate_mixed` emits) | 4D `16384²` → **fast path LOST** |
**Padding is what pins the 5 global layers to an explicit mask.** The 25
sliding layers get a 4D tensor either way — `sdpa_attention_forward` sets
`is_causal=True` only when `attention_mask is None`, and a 1024 window cannot
be expressed as `is_causal`.
Isolated, same width, only the mask differing:
2 × 16,384, no padding 35.244 s 26,048 loss targets
2 × 16,384, 50% pad on row 1 38.567 s 19,640 loss targets
**9.4% slower for 24% less work.**
### 6.4 The corpus is 29.9% padding — and bucketing is the biggest win available
Measured off the preserved encode cache in true `SequentialSampler` order:
records 20,982 (3,583 rp-dialogue / 12,003 prose-chunk / 5,396 actual-play)
seq len min/mean/max 142 / 2,752 / 16,384
micro-batches (mb=2) 10,491
real tokens 57,733,156
padded tokens 82,337,318
PADDING WASTE 29.9%
mb width p50/p90/p99 2,092 / 10,634 / 16,341
micro-batches at 16,384 3 of 10,491 (0.0%)
⚠ Note the last line against §6.1: **the 2 × 16,384 benchmark shape occurs in
three micro-batches out of 10,491.** Weighted over the real distribution the
quadratic share is ~51%, not 67.8%.
**Bucket-to-pair, shuffle-to-mix** (brokkr's design, validated on measured
lengths — form micro-batches within length buckets, then shuffle the resulting
*micro-batches* globally):
| bucket | waste | predicted step | zero-pad mb | roots/accum window |
|---:|---:|---:|---:|---:|
| current | 29.9% | 44.3 s → 16.13 h | 0.1% | 3.68 |
| **2** | **0.0%** | **28.6 s → 10.40 h** | **78.3%** | 3.56 |
| 8 | 0.0% | 28.6 s → 10.41 h | 65.3% | 3.54 |
| 32 | 0.1% | 28.6 s → 10.43 h | 41.9% | 3.55 |
| 128 | 0.6% | 28.8 s → 10.51 h | 14.7% | 3.55 |
| 512 | 2.4% | 29.7 s → 10.82 h | 4.1% | 3.61 |
**≥35.5% wall clock, no kernel work, no new dependency, peak memory unchanged.**
Two findings that changed the design:
- **Bucket size is not a diversity knob.** Roots per accumulation window are
flat at 3.54–3.61 across a 256× range. The global micro-batch shuffle does
all the mixing. Use the tightest bucket.
- **35.5% is a floor.** Zero-pad micro-batches go 0.1% → 78.3%, which puts the
5 global layers back on `is_causal` for most of the run (§6.3). The cost
model does not capture that. Direction certain, magnitude not yet measured at
representative shapes.
⚠ **Source-homogeneity is a real hazard here** — length correlates hard with
root (kvasir short, chunked RP windows long), so length-homogeneous batches are
root-homogeneous batches. The global micro-batch shuffle is what prevents an
accumulation window drawing its whole gradient from one source. It is
load-bearing, not decoration.
### 6.5 The chunked CE is fine — do not swap it
2 × 16,384 CE forward 374 ms of 35.329 s = 1.1%
2 × 4,096 CE forward 93 ms of 4.387 s = 2.1%
⚠ **Forward only** — the `torch.utils.checkpoint` recompute runs inside
`.backward()`, outside the timing window. Even at 3× it is ~3%.
`liger-kernel` fused linear CE is a ~1–3% lever on this shape. §2's finding
stands unchanged: chunking is what makes seq 16384 *reachable*, and it is not
what makes it slow.
### 6.6 MoE is ~8% — stop optimising it
Dense GEMM is 7.9% of the step, confirming the earlier decomposition bound of
~10% from the kernel side.
On `grouped_mm`: **the trace does not adjudicate it.** Run-01 was relaunched on
`eager`, so the profile shows the *default* path — 25,463 `aten::mm` dispatches
in one fwd+bwd, far more than the ~90 a grouped path would produce, so the
default is per-expert sequential. Whether the flag changes that when set is a
different measurement and was not run. At 7.9% it is not worth running.
### 6.7 Restart parameters for round 2
**Do not relaunch without the sampler change.** It is the only lever that wins
under every branch of the diagnosis.
1. **Implement bucket-to-pair + shuffle-to-mix** in the harness, tightest
bucket, global micro-batch shuffle. Expected ~16.1 h → ~10.4 h or better.
2. **Re-assert the mask band structure** after the sampler change —
`scripts/training-probes/step0_mask.py`, 30 s, no GPU. The sampler touches
batch composition, which is what drives mask construction.
3. **Resume with `/tank/erp-tune/resume-run-01.sh`, NEVER the original launch
command** — it begins `rm -rf /tank/erp-tune/run-01` and would destroy the
609 MB encode cache (2.5 min to reuse, ~4.3 h to rebuild). ⚠ A sampler change
alters record *order*, not encoding, so the cache stays valid — but bump
`encode_version` if anything upstream of `input_ids` changes.
4. **Run the `lora_B` inert-adapter gate at step 100.** It never ran in round 1.
Norm every `lora_B` tensor in the checkpoint: all-non-zero = real, all-zero =
INERT (kill the run), partial = module-selection problem. This is the one
failure that stays invisible until brokkr's acceptance gate reports
base-identical numbers.
5. **The corpus override is ONE RUN ONLY** (`operator-2026-08-25-rnd-run`). A
second run needs a second operator grant.
6. **Attention backend is round 2's second lever**, gated on an A/B on the
replica — not on argument. It can run while the tuned job trains.
⚠ GPU0 is currently **reserved and idle** by operator instruction; `sec` /
mog-sec remains down. The window is still open, so
`playbooks/ana-ml2-training-window-close.yaml` has NOT been run.