# 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.