docs: M3 — DPO results (infra correct, held-out correctness flat, over-optimization collapse)

Implementation log (docs/18) + Phase-3 row (evolution.md): the two ops + gates,
pair-gen (gold chosen / sampled-wrong rejected), reference-logprob caching, the
training loop, and the honest finding — reward margin + pref-acc rise but
held-out arithmetic correctness stays ~5-8% (flat within std-error) and
over-optimizes to collapse (margin +34 → 0% format). DPO reweights, it does not
install the capability; motivates M4 GRPO (optimize the verifiable reward online).

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));
+        }),
+    )
+}

crates/xtrain-autodiff/tests/autograd.rs

@@ -1005,3 +1005,83 @@ fn transpose_var(x: &Var) -> Var {
         }),
     )
 }
+
+// seq_logprob (M3 DPO): Σ log p(target) over non-ignored rows. Grad-check with a
+// completion mask — rows 0,1 are -100 (prompt, contribute 0), rows 2..6 supervised.
+#[test]
+fn seq_logprob_bwd() {
+    require_gpu();
+    let (rows, cols) = (6usize, 9usize);
+    let x_h = fill(rows * cols, 202);
+    let targets: Vec<i32> = (0..rows)
+        .map(|r| if r < 2 { -100 } else { (r * 2 % cols) as i32 })
+        .collect();
+    let target = Tensor::from_slice(&targets, &[rows]).to_device(Device::Cuda(0));
+
+    let x = Var::leaf(cuda(&x_h, &[rows, cols]));
+    let lp = ops::seq_logprob(&x, &target);
+    lp.backward();
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+
+    // Numeric scalar = seq_logprob = −Σ per_row (per_row is 0 for ignored rows).
+    let tgt = targets.clone();
+    let lx = move |v: &[f32], s: &[usize]| {
+        let t = Tensor::from_slice(&tgt, &[rows]).to_device(Device::Cuda(0));
+        let (_, per_row) = cuda(v, s).cross_entropy(&t);
+        -per_row
+            .to_device(Device::Cpu)
+            .as_slice::<f32>()
+            .iter()
+            .sum::<f32>()
+    };
+    report(
+        "seq_logprob dX",
+        &grad_check(&x_h, &[rows, cols], &lx, dx.as_slice::<f32>(), cfg_nonlinear()),
+    );
+}
+
+// dpo_loss (M3): scalar DPO loss with the two policy logprobs as parents. Grad-check
+// each parent (finite diff of softplus(−Δ)) + the degenerate points the gate pins:
+// policy==reference ⇒ Δ=0, L=log2, grads ∓β/2; β=0 ⇒ grads 0.
+#[test]
+fn dpo_loss_bwd_and_degenerate() {
+    require_gpu();
+    let (ref_c, ref_r, beta) = (0.5f32, 0.9f32, 0.1f32);
+    let (pc0, pr0) = (1.2f32, 0.7f32);
+    let softplus = |z: f32| z.max(0.0) + (-(z.abs())).exp().ln_1p();
+
+    let pc = Var::leaf(cuda(&[pc0], &[1]));
+    let pr = Var::leaf(cuda(&[pr0], &[1]));
+    let l = ops::dpo_loss(&pc, &pr, ref_c, ref_r, beta);
+    l.backward();
+    let dpc = pc.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
+    let dpr = pr.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
+
+    let l_of_pc = move |v: &[f32], _s: &[usize]| softplus(-(beta * ((v[0] - ref_c) - (pr0 - ref_r))));
+    report("dpo_loss dpc", &grad_check(&[pc0], &[1], &l_of_pc, &[dpc], cfg_nonlinear()));
+    let l_of_pr = move |v: &[f32], _s: &[usize]| softplus(-(beta * ((pc0 - ref_c) - (v[0] - ref_r))));
+    report("dpo_loss dpr", &grad_check(&[pr0], &[1], &l_of_pr, &[dpr], cfg_nonlinear()));
+
+    // Degenerate 1: policy == reference ⇒ Δ=0 ⇒ L=log2, grads = (∓β/2).
+    let pc2 = Var::leaf(cuda(&[ref_c], &[1]));
+    let pr2 = Var::leaf(cuda(&[ref_r], &[1]));
+    let l2 = ops::dpo_loss(&pc2, &pr2, ref_c, ref_r, beta);
+    let lval = l2.value().to_device(Device::Cpu).as_slice::<f32>()[0];
+    l2.backward();
+    let d2c = pc2.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
+    let d2r = pr2.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
+    assert!((lval - 2f32.ln()).abs() < 1e-5, "L at Δ=0 must be log2, got {lval}");
+    assert!(
+        (d2c + beta * 0.5).abs() < 1e-5 && (d2r - beta * 0.5).abs() < 1e-5,
+        "grads at Δ=0 must be ∓β/2, got ({d2c},{d2r})"
+    );
+
+    // Degenerate 2: β=0 ⇒ grads 0.
+    let pc3 = Var::leaf(cuda(&[pc0], &[1]));
+    let pr3 = Var::leaf(cuda(&[pr0], &[1]));
+    let l3 = ops::dpo_loss(&pc3, &pr3, ref_c, ref_r, 0.0);
+    l3.backward();
+    let d3c = pc3.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
+    assert!(d3c.abs() < 1e-9, "β=0 ⇒ grad 0, got {d3c}");
+    println!("dpo_loss OK: grad-check (dpc,dpr) + degenerate (Δ=0→log2 & ∓β/2, β=0→0)");
+}

crates/xtrain-train/src/bin/gen_dpo_pairs.rs

@@ -0,0 +1,157 @@
+//! Generate DPO preference pairs for the verifiable arithmetic task (M3).
+//!
+//! Per the aligned decision: **chosen = the gold answer** (`sft_answer`, always
+//! correct), **rejected = a sampled-incorrect completion from the SFT model** — a
+//! format-valid but wrong boxed answer, i.e. a hard negative drawn from the model's
+//! own distribution. Since the SFT model is only ~8% correct (M1), a single GREEDY
+//! decode is wrong ~92% of the time, so we use the KV-cache greedy engine (M2a) and
+//! simply skip the ~8% of prompts where greedy happens to be correct (no usable
+//! negative). Fast (cached), deterministic, and one clean hard negative per prompt.
+//!
+//! Writes `<out>` as `question<TAB>chosen<TAB>rejected` (bare text, like the SFT
+//! TSV — `train_dpo` adds the `User:/Assistant:` frame). Problems are deduped.
+
+#[cfg(no_cuda)]
+fn main() {
+    eprintln!("gen_dpo_pairs: built without CUDA (no_cuda); run on a GPU host.");
+}
+
+#[cfg(not(no_cuda))]
+use std::collections::HashSet;
+#[cfg(not(no_cuda))]
+use std::io::Write;
+#[cfg(not(no_cuda))]
+use xtrain_cuda::device;
+#[cfg(not(no_cuda))]
+use xtrain_model::{Config, TinyTransformer, generate_greedy_cached};
+#[cfg(not(no_cuda))]
+use xtrain_tensor::Device;
+#[cfg(not(no_cuda))]
+use xtrain_train::task::{Op, GenConfig, check_answer, gen_problem, parse_boxed_answer};
+
+#[cfg(not(no_cuda))]
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+#[cfg(not(no_cuda))]
+fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
+    args.iter()
+        .position(|a| a == name)
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(default)
+}
+
+#[cfg(not(no_cuda))]
+fn flag_value(args: &[String], name: &str) -> Option<String> {
+    args.iter()
+        .position(|a| a == name)
+        .and_then(|i| args.get(i + 1))
+        .cloned()
+}
+
+/// Keep only the first answer "turn": cut at the first `<|endoftext|>` then the
+/// first newline (mirrors eval_arith).
+#[cfg(not(no_cuda))]
+fn first_answer_segment(continuation: &str) -> &str {
+    let s = continuation
+        .split("<|endoftext|>")
+        .next()
+        .unwrap_or(continuation);
+    s.split('\n').next().unwrap_or(s)
+}
+
+#[cfg(not(no_cuda))]
+fn main() {
+    use xserv_tokenizer::Tokenizer;
+
+    let args: Vec<String> = std::env::args().collect();
+    let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
+    let ckpt = positionals.first().expect("usage: gen_dpo_pairs <sft_ckpt> <tokenizer.json> [flags]");
+    let tok_path = positionals
+        .get(1)
+        .map(|s| s.as_str())
+        .unwrap_or("/opt/wjh/models/gpt2/tokenizer.json");
+
+    let n_heads = flag(&args, "--heads", 52usize);
+    let head_dim = flag(&args, "--head-dim", 32usize);
+    let n_layers = flag(&args, "--layers", 22usize);
+    let ffn = flag(&args, "--ffn", 6656usize);
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
+    let n_pairs: usize = flag(&args, "--n", 2000);
+    let seed: u64 = flag(&args, "--seed", 1234);
+    let max_add: i64 = flag(&args, "--max-add", 999);
+    let max_mul: i64 = flag(&args, "--max-mul", 99);
+    let max_new: usize = flag(&args, "--max-tokens", 32);
+    let out = flag_value(&args, "--out").expect("--out <file> is required");
+
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let tok = Tokenizer::from_file(std::path::Path::new(tok_path));
+    let cfg = Config::from_arch(tok.vocab_size(), n_heads, head_dim, n_layers, ffn)
+        .with_kv_heads(kv_heads);
+    let mut seed_init = 1u64;
+    let model = TinyTransformer::new(cfg, device, |shape| {
+        seed_init = seed_init.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed_init, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed_init, 0.04)
+        }
+    });
+    xtrain_train::checkpoint::load_into(std::path::Path::new(ckpt.as_str()), &model.params())
+        .expect("load SFT checkpoint");
+
+    let gcfg = GenConfig {
+        max_add,
+        max_mul,
+        ops: vec![Op::Add, Op::Sub, Op::Mul],
+    };
+    let mut rng = seed.max(1);
+    let mut keys = HashSet::new();
+    let mut writer = std::io::BufWriter::new(std::fs::File::create(&out).expect("create out"));
+    let (mut written, mut skipped, mut attempts) = (0usize, 0usize, 0usize);
+
+    while written < n_pairs {
+        attempts += 1;
+        if attempts > n_pairs * 4 {
+            eprintln!("gen_dpo_pairs: stopping early at {written} pairs after {attempts} attempts");
+            break;
+        }
+        let p = gen_problem(&mut rng, &gcfg);
+        if !keys.insert(p.key()) {
+            continue;
+        }
+        let prompt_text = format!("User: {}\nAssistant:", p.question());
+        let ids: Vec<i32> = tok.encode(&prompt_text).into_iter().map(|t| t as i32).collect();
+        let out_ids = generate_greedy_cached(&model, device, &ids, max_new);
+        let cont = tok.decode(&out_ids[ids.len()..].iter().map(|&t| t as u32).collect::<Vec<_>>());
+        let seg = first_answer_segment(&cont).trim();
+        // A valid hard negative: a well-formed boxed answer that is WRONG.
+        if parse_boxed_answer(seg).is_some() && !check_answer(seg, p.answer()) {
+            writeln!(writer, "{}\t{}\t{}", p.question(), p.sft_answer(), seg).expect("write");
+            written += 1;
+        } else {
+            skipped += 1; // greedy was correct (~8%) or malformed → no clean negative
+        }
+    }
+    writer.flush().expect("flush");
+    println!(
+        "wrote {written} DPO pairs to {out} (skipped {skipped} no-negative; {attempts} attempts; \
+         chosen=gold, rejected=greedy-incorrect)"
+    );
+}

crates/xtrain-train/src/bin/train_dpo.rs

@@ -0,0 +1,233 @@
+//! DPO training on the verifiable arithmetic task (M3 / Stage P1).
+//!
+//! Loads the SFT checkpoint as the policy AND uses it as the frozen reference:
+//! reference logprobs `log πref(chosen)` / `log πref(rejected)` are **precomputed
+//! once** before any optimizer step (when policy == reference), then cached as
+//! constants — so only one model stays resident (the design's reference-logprob
+//! caching). Each step forwards the policy on the chosen and rejected completions,
+//! takes [`seq_logprob`] of each, and minimises [`dpo_loss`]; the two forwards
+//! share the policy params, so backward accumulates both branches' grads.
+//!
+//! Health metrics (per docs/18, the doc-13 "don't trust loss alone" lesson): the
+//! chosen−rejected **reward margin** and **preference accuracy** (margin > 0) — both
+//! should rise. The arithmetic-correctness payoff is measured separately by running
+//! `eval_arith` on the saved checkpoint.
+//!
+//!   train_dpo <tokenizer.json> <dpo.tsv> --init-ckpt <sft.ckpt> <arch flags> \
+//!       --beta 0.1 --steps 1000 --lr 5e-7 --ckpt <out.ckpt>
+
+#[cfg(no_cuda)]
+fn main() {
+    eprintln!("train_dpo: built without CUDA (no_cuda); run on a GPU host.");
+}
+
+#[cfg(not(no_cuda))]
+use xtrain_autodiff::ops;
+#[cfg(not(no_cuda))]
+use xtrain_cuda::device;
+#[cfg(not(no_cuda))]
+use xtrain_model::{Config, TinyTransformer, ids_tensor};
+#[cfg(not(no_cuda))]
+use xtrain_tensor::Device;
+
+#[cfg(not(no_cuda))]
+fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
+        })
+        .collect()
+}
+
+#[cfg(not(no_cuda))]
+fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
+    args.iter()
+        .position(|a| a == name)
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(default)
+}
+
+#[cfg(not(no_cuda))]
+fn flag_value(args: &[String], name: &str) -> Option<String> {
+    args.iter()
+        .position(|a| a == name)
+        .and_then(|i| args.get(i + 1))
+        .cloned()
+}
+
+/// Frame a (question, completion) the same way the SFT loader does
+/// (`User: …\nAssistant:` prompt + ` {completion}\n<|endoftext|>`), then return the
+/// next-token (input, target) pair: input = tokens[..L-1], target = labels[1..L]
+/// with the prompt positions masked to -100 (only completion tokens supervised).
+#[cfg(not(no_cuda))]
+fn frame(
+    tok: &xserv_tokenizer::Tokenizer,
+    question: &str,
+    completion: &str,
+) -> (Vec<i32>, Vec<i32>) {
+    let prompt = format!("User: {question}\nAssistant:");
+    let answer = format!(" {completion}\n<|endoftext|>");
+    let p_ids: Vec<i32> = tok.encode(&prompt).into_iter().map(|t| t as i32).collect();
+    let a_ids: Vec<i32> = tok.encode(&answer).into_iter().map(|t| t as i32).collect();
+    let mut tokens = p_ids.clone();
+    tokens.extend_from_slice(&a_ids);
+    let mut labels = vec![-100i32; p_ids.len()];
+    labels.extend_from_slice(&a_ids);
+    let l = tokens.len();
+    (tokens[..l - 1].to_vec(), labels[1..l].to_vec())
+}
+
+/// Sequence logprob `Σ log πθ(completion)` of a framed (input, target) pair.
+#[cfg(not(no_cuda))]
+fn seq_lp(
+    model: &TinyTransformer,
+    device: Device,
+    input: &[i32],
+    target: &[i32],
+) -> xtrain_autodiff::tape::Var {
+    let logits = model.forward(&ids_tensor(input, device));
+    ops::seq_logprob(&logits, &ids_tensor(target, device))
+}
+
+#[cfg(not(no_cuda))]
+fn scalar(v: &xtrain_autodiff::tape::Var) -> f32 {
+    v.value().to_device(Device::Cpu).as_slice::<f32>()[0]
+}
+
+#[cfg(not(no_cuda))]
+fn main() {
+    use xserv_tokenizer::Tokenizer;
+    use xtrain_optim::GpuAdamW;
+
+    let args: Vec<String> = std::env::args().collect();
+    let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
+    let tok_path = positionals.first().expect("usage: train_dpo <tokenizer.json> <dpo.tsv> [flags]");
+    let tsv_path = positionals.get(1).expect("usage: train_dpo <tokenizer.json> <dpo.tsv> [flags]");
+
+    let n_heads = flag(&args, "--heads", 52usize);
+    let head_dim = flag(&args, "--head-dim", 32usize);
+    let n_layers = flag(&args, "--layers", 22usize);
+    let ffn = flag(&args, "--ffn", 6656usize);
+    let kv_heads = flag(&args, "--kv-heads", n_heads);
+    let beta: f32 = flag(&args, "--beta", 0.1);
+    let steps: usize = flag(&args, "--steps", 1000);
+    let lr: f32 = flag(&args, "--lr", 5e-7);
+    let wd: f32 = flag(&args, "--wd", 0.0);
+    let clip: f32 = flag(&args, "--clip", 1.0);
+    let log_every: usize = flag(&args, "--log-every", 50);
+    let init_ckpt = flag_value(&args, "--init-ckpt").expect("--init-ckpt <sft.ckpt> is required");
+    let out_ckpt = flag_value(&args, "--ckpt").expect("--ckpt <out> is required");
+
+    // Load preference pairs: question<TAB>chosen<TAB>rejected.
+    let raw = std::fs::read_to_string(tsv_path).expect("read dpo tsv");
+    let pairs: Vec<(String, String, String)> = raw
+        .lines()
+        .filter(|l| !l.trim().is_empty())
+        .map(|l| {
+            let mut it = l.splitn(3, '\t');
+            let q = it.next().expect("question").to_string();
+            let c = it.next().expect("chosen").to_string();
+            let r = it.next().expect("rejected").to_string();
+            (q, c, r)
+        })
+        .collect();
+    assert!(!pairs.is_empty(), "no DPO pairs in {tsv_path}");
+
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    let tok = Tokenizer::from_file(std::path::Path::new(tok_path.as_str()));
+    let cfg = Config::from_arch(tok.vocab_size(), n_heads, head_dim, n_layers, ffn)
+        .with_kv_heads(kv_heads);
+    let mut seed_init = 1u64;
+    let model = TinyTransformer::new(cfg, device, |shape| {
+        seed_init = seed_init.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed_init, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed_init, 0.04)
+        }
+    });
+    xtrain_train::checkpoint::load_into(std::path::Path::new(&init_ckpt), &model.params())
+        .expect("load SFT checkpoint");
+    model.eval(); // DPO runs without dropout (deterministic logprobs)
+
+    // Pre-tokenize every pair once.
+    let framed: Vec<((Vec<i32>, Vec<i32>), (Vec<i32>, Vec<i32>))> = pairs
+        .iter()
+        .map(|(q, c, r)| (frame(&tok, q, c), frame(&tok, q, r)))
+        .collect();
+
+    // Reference logprobs: computed ONCE while policy == reference (SFT init), cached.
+    println!("precomputing reference logprobs for {} pairs…", framed.len());
+    let mut ref_c = Vec::with_capacity(framed.len());
+    let mut ref_r = Vec::with_capacity(framed.len());
+    for ((ci, ct), (ri, rt)) in &framed {
+        ref_c.push(scalar(&seq_lp(&model, device, ci, ct)));
+        ref_r.push(scalar(&seq_lp(&model, device, ri, rt)));
+    }
+
+    let params = model.params();
+    let mut opt = GpuAdamW::new(wd);
+    let n = framed.len();
+    // A fixed shuffle (LCG-strided) so steps sweep the dataset without bias.
+    let mut order: Vec<usize> = (0..n).collect();
+    let mut s = 0x9E3779B97F4A7C15u64;
+    for i in (1..n).rev() {
+        s = s.wrapping_mul(6364136223846793005).wrapping_add(1);
+        let j = (s >> 33) as usize % (i + 1);
+        order.swap(i, j);
+    }
+
+    let start = std::time::Instant::now();
+    let (mut win_loss, mut win_margin, mut win_acc) = (0f32, 0f32, 0usize);
+    for step in 0..steps {
+        let i = order[step % n];
+        let ((ci, ct), (ri, rt)) = &framed[i];
+        let lpc = seq_lp(&model, device, ci, ct);
+        let lpr = seq_lp(&model, device, ri, rt);
+        let (lpc_v, lpr_v) = (scalar(&lpc), scalar(&lpr));
+        let margin = (lpc_v - ref_c[i]) - (lpr_v - ref_r[i]); // implicit reward margin
+        let loss = ops::dpo_loss(&lpc, &lpr, ref_c[i], ref_r[i], beta);
+        win_loss += scalar(&loss);
+        win_margin += margin;
+        win_acc += (margin > 0.0) as usize;
+
+        loss.backward();
+        let _ = xtrain_train::clip::clip_grad_norm_gpu(&params, clip, 1.0);
+        opt.step(lr, &params);
+        for p in &params {
+            p.zero_grad();
+        }
+
+        if (step + 1) % log_every == 0 || step == steps - 1 {
+            let w = log_every.min(step + 1) as f32;
+            println!(
+                "step {:5}/{steps}: loss {:.4} | reward-margin {:+.4} | pref-acc {:.1}% | {:.1}s",
+                step + 1,
+                win_loss / w,
+                win_margin / w,
+                100.0 * win_acc as f32 / w,
+                start.elapsed().as_secs_f32(),
+            );
+            win_loss = 0.0;
+            win_margin = 0.0;
+            win_acc = 0;
+        }
+    }
+
+    xtrain_train::checkpoint::save(std::path::Path::new(&out_ckpt), &params).expect("save ckpt");
+    println!(
+        "DPO done: {} pairs, {steps} steps, beta {beta}, lr {lr:.1e} → {out_ckpt}",
+        framed.len()
+    );
+}

docs/18-post-training-rl-sft.md

@@ -410,3 +410,59 @@ short arithmetic-eval lengths the cache is overhead-bound and gives ~nothing —
 per-layer host round-trip is part of why short-seq is overhead-bound; M2b's device-side cache
 targets it.) This is the same measure-first lesson as T17 (process-per-GPU throughput-neutral):
 the win is real but only in the regime that actually stresses the bottleneck.
+
+### M3 — DPO (offline preference optimization, landed; honest negative result)
+
+The first real alignment method. Infra landed and gated; the empirical finding is that DPO
+**does not improve held-out arithmetic correctness on this task** — a genuine, on-theme negative
+result (the design doc's "RL is finicky" risk, made concrete).
+
+**Two new autograd ops (`xtrain-autodiff`, both reuse the CE kernel — no new CUDA):**
+- `seq_logprob(logits, target)` = `Σ log πθ(target)` over non-ignored positions (the per-
+  sequence logprob DPO compares). `= −Σ per_row` of cross_entropy (ignored rows already 0, like
+  SFT masking); backward = `cross_entropy_backward(probs, target, −upstream)` (SUM, no mean).
+  **Gate:** finite-diff grad-check with a `-100` completion mask.
+- `dpo_loss(lpθ_chosen, lpθ_rejected, lpref_chosen, lpref_rejected, β)` = `−log σ(Δ)` 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).
+
+**Pair construction (`gen_dpo_pairs`, aligned decision):** chosen = gold answer; rejected = the
+SFT model's own **greedy** (KV-cache engine, M2a) completion when it's a format-valid WRONG
+boxed answer — a hard negative in the model's distribution. Since SFT is ~8% correct (M1),
+greedy is wrong ~92% of the time, so this is fast and deterministic; ~8% of prompts are skipped
+(greedy correct). 1500 pairs generated (158 skipped) in ~8 min.
+
+**Training (`train_dpo`):** loads the SFT ckpt as policy AND frozen reference; **precomputes the
+reference logprobs once** (while policy == reference) and caches them — one resident model. Each
+step forwards the policy on chosen + rejected, `seq_logprob` each, minimises `dpo_loss`; the two
+forwards share params so backward accumulates both branches. Loss **starts at exactly log2**
+(Δ=0 at init) — a built-in correctness check that fired correctly. Tracks reward margin +
+preference accuracy.
+
+**Result (v12 1.05B, 1500 pairs, β=0.1; 100 held-out prompts, vs the SFT baseline format
+100/100, correct 8/100):**
+
+| run                       | reward margin | pref-acc | format | correct |
+|---------------------------|---------------|----------|--------|---------|
+| SFT (baseline)            | —             | —        | 100/100 | 8/100  |
+| DPO lr 5e-7 × 300         | +0.78         | ~82%     | 100/100 | 7/100  |
+| DPO lr 5e-7 × 800         | +1.25         | ~82%     | 100/100 | 5/100  |
+| DPO lr 1e-6 × 2000        | **+34.2**     | ~76%     | **0/100** | 0/100 |
+
+The reward margin and preference accuracy rise cleanly (the loss IS being optimized — the infra
+is correct), but the implicit reward **does not transfer to held-out correctness**: it stays
+~5–8% (all within the ~2.7% std-error of 100 prompts — statistically flat), and pushing harder
+**over-optimizes to collapse** (margin +34 = huge KL from the reference → the model emits
+garbage, `46 * 80 = CRAFTIE SERIES SERIES…`, format 0%).
+
+**The lesson (why):** chosen and rejected differ only in the final number tokens, so DPO raises
+`log p(correct) − log p(wrong)` for the *specific* training pairs — it **reweights the existing
+distribution, it does not install the capability**. The base model has no arithmetic algorithm,
+so preferring correct-vs-wrong final answers on seen pairs cannot generalize to unseen problems;
+and the only way to drive the margin far is to globally distort the distribution → incoherence.
+**DPO works when the chosen is already plausible under the policy; it cannot manufacture
+knowledge the model lacks.** This is the precise motivation for **M4 GRPO**: optimize the *actual
+verifiable reward* online (sample → check → reinforce what is genuinely correct), rather than a
+fixed-pair proxy — though GRPO faces the same 8%-correct sparsity, so whether it moves the metric
+is M4's open question. Gate met for M3 = the infra is correct (op grad-checks, log2-at-init,
+margin/acc rise); the correctness flatness is the reported finding, not a bug.

docs/evolution.md

@@ -99,6 +99,8 @@ Phase 1/2 把**预训练全栈**学完后，Phase 3 转向**后训练 infra**（
 
 **M2a（KV-cache 增量解码引擎，单序列，已落地）**：两个 forward-only 原语 + 裸 Tensor 逐 token block forward，各自隔离闸门。`rope_at`（绝对位置 RoPE，新 kernel，不动训练 `rope` → 训练路径零风险）逐位等于全序列 rope 的对应行；`decode_attention`（单 query × cached-K/V，由现成 strided-gemm + 普通 softmax 组合，**零新 kernel**）等于全 causal attention 末行（max|Δ| 6e-8）。引擎 `generate_greedy_cached` 镜像 `block_forward` 在 Tensor 层（无 autograd tape，推理不需梯度），靠**公开 `params()` 稳定顺序**拿权重（零 model 可见性改动）。**核心闸门 = token-identical**：与朴素全重算贪心逐 token 一致（小 GQA 单测 + v12 1.05B 上 cached eval 与 naive **逐字节相同**：format 100/100, correct 8/100）。**吞吐 baseline（v12, batch1, F32，profile-first 实测）= cache 收益随序列长度而定**：max_new 32 ≈ 持平（108 vs 111，短序列 launch 开销 bound）、128 **~1.9×**（69 vs 133）、256 naive **OOM** vs cached 129 tok/s。cached 吞吐**近恒定**（O(1)/token + 恒定显存），naive **衰减**（O(t)/token，O(seq²) 图 → OOM）。⇒ 短 eval prompt overhead-bound、cache 几乎无收益，真正受益的是**长 rollout**（DPO 造对 / GRPO completion）——与 T17（process-per-GPU 吞吐中性）同一条 measure-first 教训：收益真实，但只在真正压到瓶颈的 regime 里。M2a 的 per-layer 主机往返是短序列 overhead-bound 的一部分原因，M2b（device 端 cache + 批量 ragged）针对它。
 
+**M3（DPO，离线偏好优化，已落地 + 诚实负结果）**：两个复用 CE kernel 的新算子（零新 CUDA）——`seq_logprob`（Σ log πθ over 非 mask 位，反向 = CE_backward 取负求和；grad-check + mask）、`dpo_loss`（−log σ(Δ)，双 policy logprob 父节点；grad-check + 退化 Δ=0→log2/∓β·½、β=0→0）。造对（`gen_dpo_pairs`）= chosen=gold、rejected=SFT 自己 greedy（用 M2a 引擎）的格式合法**错误**答案（8% greedy 答对的跳过）。训练（`train_dpo`）把 SFT ckpt 同时作 policy 和冻结 reference，**一次性预算 reference logprob 并缓存**（单模型驻留），每步 policy forward chosen+rejected → seq_logprob → dpo_loss，两 forward 共享 param 累积梯度；**loss 起步恰好 log2**（Δ=0 内置校验）。**结果（v12, 1500 对, β0.1；100 留出题 vs SFT 8/100）**：reward-margin 与 pref-acc 干净上升（loss 被正确优化、infra 对），但**不转化为 held-out 正确率**——lr5e-7×300→7%、×800→5%、lr1e-6×2000→margin+34 **崩溃**（0% 格式、输出垃圾），三档都在 100 题 ~2.7% 标准误内 = 统计持平。**教训**：chosen/rejected 只差最终数字 token，DPO 提升的是**特定训练对的 token 偏好、reweight 现有分布,不 install 能力**；base 模型没有算术算法,偏好优化不泛化,推狠了只是全局扭曲分布→不连贯。**DPO 在 chosen 本就 plausible 时有效,不能凭空造模型没有的知识**——这正是 M4 GRPO 的动机:在线优化**真实可验证 reward**(采样→check→强化真正对的)而非固定对的 proxy(但 GRPO 同样面对 8% 稀疏,能否抬动指标是 M4 的 open question)。与 v8/T17 同源的诚实账：跑通+闸门齐全,负结果如实记。
+
 ## 四、perf 杠杆台账（详见 [known-issues.md](known-issues.md)）
 
 - **已修**：KI-1 单序列 launch-bound（T10）· KI-5 per-op cudaMalloc 串行（T11）· KI-2 bf16/OOM（T12）· KI-3 激活重计算（T13，解锁 dim1024，v8 用上）。