vh
7b5fd91d3c
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.
54 lines
2.5 KiB
Markdown
54 lines
2.5 KiB
Markdown
# Training throughput probes
|
||
|
||
Instruments for finding where a training step's time actually went. Written
|
||
2026-08-24 during the Gemma-4 26B-A4B ERP/RP tune investigation; the lessons
|
||
they produced live in
|
||
[`docs/pfi/training-throughput-playbook.md`](../../docs/pfi/training-throughput-playbook.md).
|
||
|
||
**These are diagnostic instruments, not production code.** They hard-code paths
|
||
for that run. Adapt the constants at the top; keep the measurement design.
|
||
|
||
## The probes
|
||
|
||
| script | settles | GPU | runtime |
|
||
|---|---|---|---|
|
||
| `step0_mask.py` | mask band structure; which layers keep the `is_causal` fast path | no | ~30 s |
|
||
| `step2_padding.py` | padding waste, length distribution, CE chunk sizing | no | ~2 min |
|
||
| `step_bucket.py` | bucketing gain, bucket-size sweep, source diversity | no | ~3 min |
|
||
| `step1_profile.py` | scaling fit, padding penalty, CE wall clock, kernel table | **yes** | ~15 min |
|
||
|
||
Run in that order. Only the last needs the real checkpoint, and it wants an
|
||
idle card — it loads ~48 GiB and peaks near 77 GiB at `2 × 16,384`.
|
||
|
||
## Design rules worth preserving when you adapt these
|
||
|
||
**`step1_profile.py` reuses the harness's own `discover_target_modules` and
|
||
replicates its `compute_loss` byte-for-byte** rather than re-implementing the
|
||
step. A probe that reimplements the training step measures the probe. If you
|
||
port this, keep the import from the real harness.
|
||
|
||
**`step0_mask.py` needs no weights and no GPU** — SDPA backend selection and
|
||
mask construction depend on shapes, dtype and mask presence, not on weight
|
||
values. That is what makes the correctness assertion cheap enough to run before
|
||
every job.
|
||
|
||
**The scaling test takes three points, not two.** Two points over three
|
||
plausible terms (quadratic, linear, fixed-per-batch) is underdetermined; see
|
||
playbook §1.1 for the hour that cost.
|
||
|
||
**`step_bucket.py` sweeps bucket size deliberately.** The first version
|
||
re-sorted within each bucket, which silently collapsed every bucket size to a
|
||
full global sort and made the sweep a no-op. If you change the pairing logic,
|
||
check that the sweep still varies something.
|
||
|
||
## Raw evidence
|
||
|
||
`step1-profile-output-2026-08-24.txt` is the unedited output of the run the
|
||
playbook's numbers come from — scaling points, padding penalty, CE timing, and
|
||
the full `key_averages()` kernel table. Kept so the claims can be re-derived
|
||
rather than taken on faith.
|
||
|
||
⚠ That table **double-counts**: `key_averages()` lists both the ATen op and the
|
||
CUDA kernel it launched, each carrying the same self device time. Sum device
|
||
kernel rows only. See playbook §3.4.
|