post-train: M3 — seq_logprob + dpo_loss autograd ops

Two new ops for DPO (M3), both reusing existing kernels (no new CUDA):

- seq_logprob(logits, target): Σ log πθ(target) over non-ignored (target≥0)
  positions — the per-sequence logprob DPO compares between policy and
  reference. = −Σ per_row of cross_entropy (ignored rows already 0, like SFT
  masking); backward = cross_entropy_backward(probs, target, −upstream) (sum,
  no mean division). Gate: finite-diff grad-check with a -100 completion mask.

- dpo_loss(lpθ_chosen, lpθ_rejected, lpref_chosen, lpref_rejected, β): scalar
  L = −log σ(Δ) = softplus(−Δ) with the two policy logprobs as parents (ref
  logprobs constant). Gate: grad-check both parents + degenerate points
  (policy==ref ⇒ Δ=0, L=log2, grads ∓β/2; β=0 ⇒ grads 0). Same formula as TRL.

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

crates/xtrain-autodiff/src/ops.rs

@@ -439,3 +439,81 @@ pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
         }),
     )
 }
+
+/// Per-sequence log-probability: `Σ log πθ(target)` over the non-ignored
+/// (`target ≥ 0`) positions — the quantity DPO (M3) compares between policy and
+/// reference. `target` is `[rows]` I32 carrying `-100` (ignore) at masked positions
+/// (e.g. the prompt) and the gold token id elsewhere; ignored positions contribute
+/// 0, exactly like the SFT cross-entropy masking. Returns a scalar `[1]` Var.
+///
+/// Reuses the CE forward (per-row `−log p(target)`) and backward, so no new kernel:
+/// `seq_logprob = −Σ per_row`, and `d(seq_logprob)/d(logits) = −(probs − onehot)`
+/// = `cross_entropy_backward(probs, target, −upstream)` (a SUM, so no mean
+/// division — contrast [`cross_entropy`], which divides by `valid_rows`).
+pub fn seq_logprob(x: &Var, target: &Tensor) -> Var {
+    let logit_dtype = x.value().dtype();
+    let (probs, per_row) = x.value().cross_entropy(target);
+    // per_row[r] = −log p(target_r), and is 0 for ignored rows (target < 0), so the
+    // sum already counts only the supervised (completion) positions.
+    let sum_neg_lp: f32 = per_row
+        .to_device(xtrain_tensor::Device::Cpu)
+        .as_slice::<f32>()
+        .iter()
+        .sum();
+    let out = Tensor::from_slice(&[-sum_neg_lp], &[1]).to_device(x.value().device());
+
+    let target = target.clone();
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(move |d, parents| {
+            let upstream = d.to_device(xtrain_tensor::Device::Cpu).as_slice::<f32>()[0];
+            // d(Σ log p)/d(logits) = −(probs − onehot); SUM, so no /valid_rows.
+            let dx = Tensor::cross_entropy_backward(&probs, &target, -upstream);
+            Var::push_grad(&parents[0], dx.to_dtype(logit_dtype));
+        }),
+    )
+}
+
+/// DPO loss (Rafailov et al., M3) for one preference pair, as a scalar `[1]` Var
+/// whose two parents are the POLICY sequence-logprobs of the chosen and rejected
+/// completions (from [`seq_logprob`]); the REFERENCE logprobs are constants
+/// (precomputed once from the frozen SFT model). With
+///   `Δ = β·[(lpθ_chosen − lpref_chosen) − (lpθ_rejected − lpref_rejected)]`
+/// the loss is `L = −log σ(Δ) = softplus(−Δ)`. Only the policy terms carry gradient:
+///   `∂L/∂lpθ_chosen = −β·(1−σ(Δ))`,  `∂L/∂lpθ_rejected = +β·(1−σ(Δ))`.
+/// Degenerate points the M3 gate pins: `πθ == πref` ⇒ `Δ = 0`, `L = log 2`, implicit
+/// reward 0; `β → 0` ⇒ gradient → 0. Same formula as TRL
+/// (`-logsigmoid(β·(pol_c − pol_r − (ref_c − ref_r)))`).
+pub fn dpo_loss(
+    lp_pol_chosen: &Var,
+    lp_pol_rejected: &Var,
+    lp_ref_chosen: f32,
+    lp_ref_rejected: f32,
+    beta: f32,
+) -> Var {
+    use xtrain_tensor::Device;
+    let scalar = |v: &Var| v.value().to_device(Device::Cpu).as_slice::<f32>()[0];
+    let pc = scalar(lp_pol_chosen);
+    let pr = scalar(lp_pol_rejected);
+    let delta = beta * ((pc - lp_ref_chosen) - (pr - lp_ref_rejected));
+    // L = softplus(−Δ) = log(1 + e^{−Δ}) (numerically stable).
+    let nd = -delta;
+    let l = nd.max(0.0) + (-(nd.abs())).exp().ln_1p();
+    let dev = lp_pol_chosen.value().device();
+    let out = Tensor::from_slice(&[l], &[1]).to_device(dev);
+
+    Var::from_op(
+        out,
+        vec![lp_pol_chosen.clone(), lp_pol_rejected.clone()],
+        Box::new(move |d, parents| {
+            let up = d.to_device(Device::Cpu).as_slice::<f32>()[0];
+            // s = σ(−Δ) = 1 − σ(Δ); ∂L/∂Δ = −s, and ∂Δ/∂pc = β, ∂Δ/∂pr = −β.
+            let s = 1.0 / (1.0 + delta.exp());
+            let g = up * beta * s;
+            let dev = parents[0].value().device();
+            Var::push_grad(&parents[0], Tensor::from_slice(&[-g], &[1]).to_device(dev));
+            Var::push_grad(&parents[1], Tensor::from_slice(&[g], &[1]).to_device(dev));
+        }),
+    )
+}