docs: M2d — ragged-batching lever, 9× measured, step bottleneck → rollout
Records the M2d lever (batch the GRPO training-side forwards), the right-pad-is-free insight, both exact gates, the end-to-end no-OOM smoke, and the 9× throughput. The honest decomposition correction: M2c claimed the training forwards "dominate" the step; the clean per-component bench falsifies the strong form — they were ~2.5 s of the ~8.5 s step (~30%), worth the 9×, but the rollout (~6 s) was always the larger share. After M2d the step is ~95% rollout, so the next step-level lever is full B×G rollout batching (today only the G samples of each prompt decode in lockstep; the B prompts are still sequential). Same measure-first lesson, once more. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -576,3 +576,54 @@ still a real, correctness-gated improvement (cleaner code, less PCIe, ~10% decod
|
||||
headline is that the *next* decode lever is **ragged batched prefill of the per-sample forwards**,
|
||||
not the cache. The M2 decode engine is now M2a (single-seq) + M2b (batched) + M2c (device cache),
|
||||
all token-identical-gated; the post-training stack remains complete with its bottleneck mapped.
|
||||
|
||||
### M2d — batch the GRPO training-side forwards (landed; the lever M2c named, + a decomposition correction)
|
||||
|
||||
M2c named the next lever: **ragged batched prefill of the per-sample training-side forwards**. Those
|
||||
forwards are the two phases that, per step, run one single-sequence `forward` per sample: the
|
||||
`per_token_logp` **captures** (logπ_old policy + logπ_ref reference) and the inner **clipped-PG**
|
||||
forward/backwards. M2d packs all `N = B·G` ragged samples of a step into ONE `forward_batched`.
|
||||
|
||||
**The enabling property — right-padding is free under causal attention.** Pad each ragged completion
|
||||
on the RIGHT to the batch's `Lmax`. A real completion row sits at an earlier position than the
|
||||
trailing pad, and causal masking forbids attending forward, so its logits are **bit-identical** to
|
||||
the unpadded single-sequence forward; the pad rows are garbage but masked out (`target = -100`). This
|
||||
is exactly why training engines pad-and-mask rather than run ragged. Two new pieces:
|
||||
- `per_token_logp_batched` (`crates/xtrain-train/src/grpo_batch.rs`): right-pad → one
|
||||
`forward_batched(batch = N)` → slice each sample's logπ back to its real length.
|
||||
- `ops::clipped_pg_loss_batched` (`crates/xtrain-autodiff/src/ops.rs`): like the per-sample
|
||||
`clipped_pg_loss`, but takes **per-row** `advantage[t]` (the owning sample's `A`) and **per-row**
|
||||
`weight[t]` (the full normaliser; the caller passes `1/(N·n_s)`). It does NOT compute its own
|
||||
`1/n_tokens`, so folding `weight = 1/(N·n_s)` reproduces the looped `Σ_s (1/N)(1/n_s)…`
|
||||
**bit-for-bit** (the per-row CE backward is row-local). A `--micro` knob packs in chunks to bound
|
||||
the `[chunk·Lmax, vocab]` logits memory; the weight uses the GLOBAL `N`, so chunked
|
||||
grad-accumulation is exact. Both `train_grpo` and the bench call these shared helpers.
|
||||
|
||||
**Correctness gates (exact, not bf16-noisy):**
|
||||
- `xtrain-model::forward_batched_ragged_matches_looped` — forward_batched on right-padded ragged
|
||||
sequences == per-sequence single-seq forward on the real rows, **max|Δlogit| = 3.7e-7 (fp32) and
|
||||
0.0 (bf16)**, both composed + flash. Pins "right-pad is free".
|
||||
- `xtrain-autodiff::clipped_pg_loss_batched_matches_looped` — batched op == looped
|
||||
`Σ_s (1/N)·clipped_pg_loss_s`, **loss Δ=1.5e-8, grad max|Δ|=7.5e-9 (f32)**.
|
||||
Composed, these prove the batched GRPO step == the looped step. End-to-end: a short SFT (v12 base,
|
||||
150 steps, arith) → `train_grpo` 12 steps runs clean — **no OOM** (1B master + AdamW + batched
|
||||
activations fit with `micro=16`), mean-reward rises, the batched inner executes.
|
||||
|
||||
**Throughput (bench `bin/bench_grpo_batch`, v12 1.05B, N=48 ragged, micro=16, β=0, weight-independent):**
|
||||
|
||||
| phase (per step) | looped (single-seq) | batched (M2d) | speedup |
|
||||
|-------------------------|---------------------|---------------|---------|
|
||||
| capture `per_token_logp`| 622 ms | 71 ms | 8.7× |
|
||||
| inner clipped-PG fwd+bwd| 1907 ms | 208 ms | 9.2× |
|
||||
| **training forwards** | **2526 ms** | **280 ms** | **9.0×**|
|
||||
|
||||
**The decomposition correction (the honest finding).** M2c claimed "the per-sample training forwards
|
||||
now dominate the step." The clean per-component bench falsifies the strong form: the training
|
||||
forwards were **~2.5 s of the ~8.5 s step (~30%)** — substantial and worth the 9× win, but the
|
||||
**rollout (`generate_cached_batch`, ~6 s) was always the larger share.** After M2d cuts the training
|
||||
forwards to ~0.28 s, the step is **~95% rollout** — the long pole has swung back to the rollout. So
|
||||
M2d removes the training-forward overhang (a real, exactly-gated 9× on its component), and re-confirms
|
||||
the same measure-first lesson one more time: the next **step-level** lever is **full B×G rollout
|
||||
batching** — today only the `G` samples of each prompt decode in lockstep (M2b); the `B` prompts are
|
||||
still sequential. M2d closes the "ragged batched per-sample forwards" lever M2c named; the post-
|
||||
training stack stays complete, now with the step decomposition measured, not asserted.
|
||||
|
||||
Reference in New Issue
Block a user