xtrain/crates/xtrain-train/src/grpo_batch.rs

//! 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
}