post-train: M4 — clipped_pg_loss + scale_rows (GRPO policy-gradient op)

The GRPO (M4) token-level loss op + the one primitive it needs:

- scale_rows(x[r,c], s[r]): per-row scale (new ~5-line CUDA kernel). The
  clipped-PG backward scales each completion token's row of (probs − onehot) by
  its own per-token coefficient, which cross_entropy_backward's single scalar
  scale can't express.
- clipped_pg_loss(logits, target, logp_old, logp_ref, A, eps, beta): per-token
  ρ_t = exp(logπθ_t − logp_old_t), L = −mean min(ρA, clip(ρ,1±ε)A) + β·mean KL
  (k3 estimator), masked to completion tokens. Backward reuses the CE machinery
  (probs − onehot) + scale_rows. Gates: grad-check the active PG path + the A=0
  (KL-only) path; degenerate value checks ε→∞ ⇒ vanilla PG, β=0 ⇒ no KL.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-autodiff/src/ops.rs

@@ -517,3 +517,83 @@ pub fn dpo_loss(
         }),
     )
 }
+
+/// GRPO clipped policy-gradient loss (M4) for ONE completion, a scalar `[1]` Var
+/// with the policy logits as the single parent. Per non-ignored (completion) token
+/// `t` (`target[t] ≥ 0`):
+///   `logπθ_t = log softmax(logits[t])[target_t]` (`= −per_row[t]` of cross_entropy)
+///   `ρ_t = exp(logπθ_t − logp_old[t])`
+///   `pg_t = min(ρ_t·A, clip(ρ_t, 1−ε, 1+ε)·A)`
+///   `kl_t = exp(logp_ref[t] − logπθ_t) − (logp_ref[t] − logπθ_t) − 1`  (k3 estimator)
+/// `L = −mean_t pg_t + β·mean_t kl_t` over the `N` completion tokens.
+///
+/// `advantage` `A` is the group-relative advantage (constant per completion in
+/// GRPO); `logp_old`/`logp_ref` are per-position constants (old policy at rollout
+/// time / frozen reference). Backward reuses the CE machinery + the per-row
+/// `scale_rows`: `dL/dlogits[t,:] = g_t·(onehot − probs)[t,:]` with
+///   `g_t = −(1/N)A·ρ_t·[unclipped active] + (β/N)(1 − exp(logp_ref_t − logπθ_t))`.
+/// Degenerate points the gate pins: `A=0` ⇒ only the KL term; `ε→∞` ⇒ vanilla PG
+/// (no clip); `β=0` ⇒ no KL term.
+#[allow(clippy::too_many_arguments)]
+pub fn clipped_pg_loss(
+    logits: &Var,
+    target: &Tensor,
+    logp_old: &[f32],
+    logp_ref: &[f32],
+    advantage: f32,
+    eps: f32,
+    beta: f32,
+) -> Var {
+    use xtrain_tensor::Device;
+    let logit_dtype = logits.value().dtype();
+    let (probs, per_row) = logits.value().cross_entropy(target);
+    let rows = per_row.shape()[0];
+    let per_row_h = per_row.to_device(Device::Cpu).as_slice::<f32>().to_vec();
+    let target_h = target.to_device(Device::Cpu).as_slice::<i32>().to_vec();
+    assert_eq!(logp_old.len(), rows, "logp_old must have one entry per position");
+    assert_eq!(logp_ref.len(), rows, "logp_ref must have one entry per position");
+
+    let mut s = vec![0f32; rows]; // per-row scale for cross_entropy_backward(·,·,1.0)
+    let (mut pg_sum, mut kl_sum, mut n) = (0f32, 0f32, 0f32);
+    for t in 0..rows {
+        if target_h[t] < 0 {
+            continue; // masked (prompt) position — no contribution, no gradient
+        }
+        n += 1.0;
+        let lp = -per_row_h[t]; // logπθ_t
+        let ratio = (lp - logp_old[t]).exp();
+        let clipped = ratio.clamp(1.0 - eps, 1.0 + eps);
+        let (unclipped_term, clipped_term) = (ratio * advantage, clipped * advantage);
+        pg_sum += unclipped_term.min(clipped_term);
+        let active = unclipped_term <= clipped_term; // min picks unclipped ⇒ grad flows
+        let d = logp_ref[t] - lp;
+        kl_sum += d.exp() - d - 1.0;
+        let pg_grad = if active { -advantage * ratio } else { 0.0 };
+        let kl_grad = beta * (1.0 - d.exp());
+        s[t] = -(pg_grad + kl_grad); // dL/dlogits = g·(onehot−probs) = −g·(probs−onehot)
+    }
+    let inv_n = if n > 0.0 { 1.0 / n } else { 1.0 };
+    for v in &mut s {
+        *v *= inv_n;
+    }
+    let loss_val = -pg_sum * inv_n + beta * kl_sum * inv_n;
+    let dev = logits.value().device();
+    let out = Tensor::from_slice(&[loss_val], &[1]).to_device(dev);
+    let s_dev = Tensor::from_slice(&s, &[rows]).to_device(dev);
+
+    let target = target.clone();
+    Var::from_op(
+        out,
+        vec![logits.clone()],
+        Box::new(move |d, parents| {
+            let up = d.to_device(Device::Cpu).as_slice::<f32>()[0];
+            // (probs − onehot), masked rows already 0; per-row scale by s; × upstream.
+            let ce = Tensor::cross_entropy_backward(&probs, &target, 1.0);
+            let mut dx = ce.scale_rows(&s_dev);
+            if up != 1.0 {
+                dx = dx.scale(up);
+            }
+            Var::push_grad(&parents[0], dx.to_dtype(logit_dtype));
+        }),
+    )
+}