Files
esh-pfi-infrastructure/scripts/training-probes/README.md
T
vh 7b5fd91d3c 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.
2026-08-24 22:10:51 -07:00

2.5 KiB
Raw Blame History

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.

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.