post-train: M2d — batch the GRPO training-side forwards (op + module + wiring)
After M2b/M2c made the rollout cheap, the GRPO step is dominated by the per-sample
single-sequence training-side forwards: the per_token_logp captures (policy +
reference) and the inner clipped-PG forward/backwards. M2d packs all N=B·G ragged
samples of a step into ONE forward_batched.
Enabling property — right-padding is free under causal attention: a real completion
row sits at an earlier position than the trailing pad, and causal masking forbids
attending forward, so its logits equal the unpadded single-sequence forward; pad
rows are masked out (target=-100).
- ops::clipped_pg_loss_batched: like clipped_pg_loss but takes per-row advantage[t]
(the owning sample's A) and per-row weight[t] (the full normaliser). It does NOT
compute its own 1/n_tokens, so the caller passing weight=1/(N·n_s) reproduces the
looped Σ_s (1/N)(1/n_s)·clipped_pg_loss_s bit-for-bit (per-row CE backward is
row-local).
- grpo_batch.rs (shared module): per_token_logp_batched (right-pad → one
forward_batched(N) → slice back to real length) + looped baselines +
inner_pg_step_{looped,batched}. A --micro knob chunks the pack to bound the
[chunk·Lmax, vocab] logits memory; weight uses the GLOBAL N so chunked
grad-accumulation stays exact.
- train_grpo restructured to collect-all-samples-then-batch; per-window phase timers
(rollout / capture / inner) to keep the step decomposition honest. Default micro =
B·G; bench-measured 9× on the training forwards.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
162
crates/xtrain-train/src/grpo_batch.rs
Normal file
162
crates/xtrain-train/src/grpo_batch.rs
Normal file
@@ -0,0 +1,162 @@
|
||||
//! Batched GRPO training-side forwards (post-training M2d). After M2b/M2c made the
|
||||
//! rollout cheap, the GRPO **step** is dominated by the per-sample full-sequence
|
||||
//! forwards: the `per_token_logp` captures (policy + reference) and the inner
|
||||
//! clipped-PG `forward`/`backward`s — each a single-sequence `forward` over a short
|
||||
//! ragged completion. This module packs the `N = B·G` ragged samples of a step into
|
||||
//! ONE `forward_batched`, amortising the per-launch overhead across N (the same win
|
||||
//! M2b gave the rollout).
|
||||
//!
|
||||
//! 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 is 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' own outputs are garbage but are masked out (`target = -100`).
|
||||
//!
|
||||
//! Both the looped (baseline) and batched paths live here so they share one source of
|
||||
//! truth — `bin/bench_grpo_batch` A/Bs them (timing + a closeness gate), and the
|
||||
//! per-row equivalence of the loss op is pinned by `clipped_pg_loss_batched_matches_looped`
|
||||
//! in `xtrain-autodiff/tests/autograd.rs`.
|
||||
|
||||
#![cfg(not(no_cuda))]
|
||||
|
||||
use xtrain_autodiff::ops;
|
||||
use xtrain_model::{TinyTransformer, ids_tensor};
|
||||
use xtrain_tensor::{Device, Tensor};
|
||||
|
||||
/// One framed completion of a GRPO step: the next-token `(input, target)` pair
|
||||
/// (prompt positions masked to `-100` in `target`), its group-relative `adv`, and the
|
||||
/// per-position rollout-time / reference logprobs the clipped-PG loss needs.
|
||||
pub struct PgSample {
|
||||
pub input: Vec<i32>,
|
||||
pub target: Vec<i32>,
|
||||
pub adv: f32,
|
||||
pub logp_old: Vec<f32>,
|
||||
pub logp_ref: Vec<f32>,
|
||||
}
|
||||
|
||||
// ------------------------------- looped (baseline) -------------------------------
|
||||
|
||||
/// Per-position `logπ(target_t)` of one framed `(input, target)` pair (= `−per_row`
|
||||
/// of cross_entropy; masked positions are 0). One single-sequence forward, no grad.
|
||||
pub fn per_token_logp(model: &TinyTransformer, device: Device, input: &[i32], target: &[i32]) -> Vec<f32> {
|
||||
let logits = model.forward(&ids_tensor(input, device)).value();
|
||||
let (_, per_row) = logits.cross_entropy(&ids_tensor(target, device));
|
||||
per_row
|
||||
.to_device(Device::Cpu)
|
||||
.as_slice::<f32>()
|
||||
.iter()
|
||||
.map(|p| -p)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// One inner clipped-PG epoch the looped way: per sample, a single-sequence forward +
|
||||
/// [`ops::clipped_pg_loss`] scaled by `1/N` + backward (grads accumulate on `model`'s
|
||||
/// params). Returns the summed scaled loss. Caller does clip + opt.step + zero_grad.
|
||||
pub fn inner_pg_step_looped(
|
||||
model: &TinyTransformer,
|
||||
device: Device,
|
||||
batch: &[PgSample],
|
||||
eps: f32,
|
||||
beta: f32,
|
||||
) -> f32 {
|
||||
let scale = 1.0 / batch.len() as f32;
|
||||
let mut total = 0f32;
|
||||
for s in batch {
|
||||
let logits = model.forward(&ids_tensor(&s.input, device));
|
||||
let loss = ops::clipped_pg_loss(&logits, &ids_tensor(&s.target, device), &s.logp_old, &s.logp_ref, s.adv, eps, beta);
|
||||
let scaled = ops::scale(&loss, scale);
|
||||
total += scaled.value().to_device(Device::Cpu).as_slice::<f32>()[0];
|
||||
scaled.backward();
|
||||
}
|
||||
total
|
||||
}
|
||||
|
||||
// ------------------------------- batched (M2d) -----------------------------------
|
||||
|
||||
/// Right-pad `m` ragged `i32` rows (each `< lmax` long) to `[m*lmax]` sequence-major,
|
||||
/// filling with `pad`. Used for both the id stream (pad = 0, arbitrary) and the target
|
||||
/// stream (pad = −100, ignored by cross_entropy).
|
||||
fn pack_i32(rows: &[&[i32]], lmax: usize, pad: i32) -> Vec<i32> {
|
||||
let mut flat = vec![pad; rows.len() * lmax];
|
||||
for (i, r) in rows.iter().enumerate() {
|
||||
flat[i * lmax..i * lmax + r.len()].copy_from_slice(r);
|
||||
}
|
||||
flat
|
||||
}
|
||||
|
||||
/// Batched [`per_token_logp`]: pack `samples` (each `(input, target)`) right-padded to
|
||||
/// `Lmax`, run ONE `forward_batched(batch = N)`, and slice each sample's `logπ` back to
|
||||
/// its real length. Equal to looping [`per_token_logp`] (right-pad is free under causal
|
||||
/// attention), to bf16 batch-reduction tolerance. `samples` are processed in chunks of
|
||||
/// `micro` (≥1) to bound the `[chunk*Lmax, vocab]` logits memory.
|
||||
pub fn per_token_logp_batched(
|
||||
model: &TinyTransformer,
|
||||
device: Device,
|
||||
samples: &[(Vec<i32>, Vec<i32>)],
|
||||
micro: usize,
|
||||
) -> Vec<Vec<f32>> {
|
||||
let mut out = Vec::with_capacity(samples.len());
|
||||
for chunk in samples.chunks(micro.max(1)) {
|
||||
let m = chunk.len();
|
||||
let lmax = chunk.iter().map(|(i, _)| i.len()).max().unwrap();
|
||||
let ins: Vec<&[i32]> = chunk.iter().map(|(i, _)| i.as_slice()).collect();
|
||||
let tgs: Vec<&[i32]> = chunk.iter().map(|(_, t)| t.as_slice()).collect();
|
||||
let ids = Tensor::from_slice(&pack_i32(&ins, lmax, 0), &[m * lmax]).to_device(device);
|
||||
let tgt = Tensor::from_slice(&pack_i32(&tgs, lmax, -100), &[m * lmax]).to_device(device);
|
||||
let logits = model.forward_batched(&ids, m).value();
|
||||
let (_, per_row) = logits.cross_entropy(&tgt);
|
||||
let pr = per_row.to_device(Device::Cpu).as_slice::<f32>().to_vec();
|
||||
for (i, (inp, _)) in chunk.iter().enumerate() {
|
||||
let b = i * lmax;
|
||||
out.push((0..inp.len()).map(|r| -pr[b + r]).collect());
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// One inner clipped-PG epoch, batched: pack the batch (in `micro`-sized chunks) and run
|
||||
/// ONE `forward_batched` + [`ops::clipped_pg_loss_batched`] + backward per chunk. The
|
||||
/// per-row `weight = 1/(N·n_s)` uses the GLOBAL `N = batch.len()` (not the chunk size),
|
||||
/// so chunked grad-accumulation reproduces the looped `Σ_s (1/N)(1/n_s)…` exactly.
|
||||
/// Returns the summed loss. Caller does clip + opt.step + zero_grad.
|
||||
pub fn inner_pg_step_batched(
|
||||
model: &TinyTransformer,
|
||||
device: Device,
|
||||
batch: &[PgSample],
|
||||
eps: f32,
|
||||
beta: f32,
|
||||
micro: usize,
|
||||
) -> f32 {
|
||||
let inv_n = 1.0 / batch.len() as f32;
|
||||
let mut total = 0f32;
|
||||
for chunk in batch.chunks(micro.max(1)) {
|
||||
let m = chunk.len();
|
||||
let lmax = chunk.iter().map(|s| s.input.len()).max().unwrap();
|
||||
let ins: Vec<&[i32]> = chunk.iter().map(|s| s.input.as_slice()).collect();
|
||||
let tgs: Vec<&[i32]> = chunk.iter().map(|s| s.target.as_slice()).collect();
|
||||
let ids = Tensor::from_slice(&pack_i32(&ins, lmax, 0), &[m * lmax]).to_device(device);
|
||||
let tgt = Tensor::from_slice(&pack_i32(&tgs, lmax, -100), &[m * lmax]).to_device(device);
|
||||
|
||||
let mut logp_old = vec![0f32; m * lmax];
|
||||
let mut logp_ref = vec![0f32; m * lmax];
|
||||
let mut advantage = vec![0f32; m * lmax];
|
||||
let mut weight = vec![0f32; m * lmax];
|
||||
for (i, s) in chunk.iter().enumerate() {
|
||||
let b = i * lmax;
|
||||
let li = s.input.len();
|
||||
logp_old[b..b + li].copy_from_slice(&s.logp_old);
|
||||
logp_ref[b..b + li].copy_from_slice(&s.logp_ref);
|
||||
let n_s = s.target.iter().filter(|&&t| t >= 0).count().max(1) as f32;
|
||||
let w = inv_n / n_s; // = 1/(N · n_s)
|
||||
for r in 0..lmax {
|
||||
advantage[b + r] = s.adv;
|
||||
weight[b + r] = w;
|
||||
}
|
||||
}
|
||||
let logits = model.forward_batched(&ids, m);
|
||||
let loss = ops::clipped_pg_loss_batched(&logits, &tgt, &logp_old, &logp_ref, &advantage, &weight, eps, beta);
|
||||
total += loss.value().to_device(Device::Cpu).as_slice::<f32>()[0];
|
||||
loss.backward();
|
||||
}
|
||||
total
|
||||
}
|
||||
Reference in New Issue
Block a user