gahow/xtrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Compare commits

base: gahow:main
Branches Tags
gahow:main
..
compare: gahow:1eef10afd9b24b7a7501e98330dd86c8740cbdda
Branches Tags
gahow:main

2 Commits
main ... 1eef10afd9

Author SHA1 Message Date
Gahow Wang
1eef10afd9 docs: T21 — record DDP-dropout wiring gap + fix (known-issues / evolution / dropout doc) 
- known-issues.md: new "DDP-dropout wiring" Fixed entry (gap + fix +
  regression test), with the meta-lesson that op/single-GPU unit tests can
  miss launcher-level integration gaps — only the V9-PILOT end-to-end run on
  the real launcher path exposed it.
- 17-dropout.md: annotate the DDP-combination note with the T18 wiring gap
  and its T21 fix.
- evolution.md: T21 row (Infra) recording the fix + meta-lesson.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 21:08:35 +08:00
Gahow Wang
1b58bd8626 test: T21 — DDP-dropout regression (live under DDP + p=0 bit-identical) 
Adds ddp_dropout_is_live_and_p0_bit_identical, run via the real launcher
path (train_rank, world=2). It would have caught the original bug:

- GATE A: a p=0 DDP run is BIT-IDENTICAL (loss trace + final params) to the
  no-dropout path — the regression guard for --dropout 0 (default).
- GATE B: a p=0.2 DDP run's loss trace DIFFERS measurably (>1e-3) from p=0.
  On the pre-T21 code the model stays in eval mode, so p=0.2 would be an
  identity and the trace would be bit-identical to p=0 — this gate fails.
- GATE C: model.is_training() == true after the run (direct proof that
  train_rank called model.train() and it survived the final-step eval).
- p>0 run is finite (no NaN/Inf).

eval_every < steps so a periodic eval fires mid-run (flipping to eval mode),
exercising the per-step model.train() restore discipline the pilot called out.
Run with --test-threads=1 like the other DDP tests (shared-GPU deadlock).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 21:08:27 +08:00
40 changed files with 139 additions and 5856 deletions
Expand all files Collapse all files

45
README.md
View File

@@ -6,16 +6,14 @@ inference side). A learning project: hand-write the entire training-systems stac
gradient checkpointing), then use it to run a multi-version **scaling study** that maps
the data-vs-capacity frontier for a tiny model.
> **Status: complete — three phases.**
> **Status: complete — two phases.**
> **Phase 1** = the from-scratch full stack (T1–T13) + an 8-version scaling study (v0–v8):
> hand-write the whole training-systems stack, then map the data-vs-capacity frontier.
> **Phase 2** = systems-stack depth (T14–T18): hand-write the five deferred training-stack
> features — fused flash-attention, real GQA, gradient accumulation, process-per-GPU DDP,
> dropout. **Phase 3** = one Chinchilla-style double-axis run (v9): dim1280 true-GQA +
> 6.01B FineWeb tokens, validating the v8 conclusion that data and capacity must scale
> together. Trains Qwen3-compatible LMs whose weights load into **xserv**; deterministic
> gates stay byte-identical, while large BF16 checkpoints are served and checked for
> prompt-level drift. This README is the capstone; per-topic detail lives in [`docs/`](docs/).
> dropout. Trains a Qwen3-compatible LM whose weights load into **xserv** and generate
> **token-identical** output — the closed loop held byte-for-byte across both phases. This
> README is the capstone; per-topic detail lives in [`docs/`](docs/).
---
@@ -36,8 +34,7 @@ borrows, the rest hand-written CUDA + Rust:
Every op's backward is verified against **finite differences** and against **PyTorch**
(forward + per-parameter grads, batch > 1). Trained weights export to HF-safetensors and
load into xserv (Qwen3, BF16); deterministic fixtures produce token-identical greedy output,
and large checkpoints are validated end-to-end in the serving path.
load into xserv (Qwen3, BF16) producing token-identical greedy output — the closed loop.
## The build journey — Phase 1 (T1–T13) + Phase 2 (T14–T18)
@@ -109,7 +106,7 @@ Each is opt-in, kept the default path **bit-identical**, and held a **hard corre
residual ~5×@8; with all 8 GPUs at 95–99% util, the residual is the **NCCL all-reduce + PCIe
topology wall**, not context serialization. The third profile-first falsification (see below).
## The scaling study — v0 → v10
## The scaling study — v0 → v8
Same Qwen3-style architecture throughout; we scaled **dim** and **data** and read out val
loss (full per-run detail in [`docs/runs/`](docs/runs/)).
@@ -122,13 +119,11 @@ loss (full per-run detail in [`docs/runs/`](docs/runs/)).
| v6 | FineWeb-edu 1.02ep | 768 / 127M | 3.07\* | **corpus swap → graduates to real text** |
| v7 | FineWeb-edu 1.45ep | 768 / 127M | 3.01\* | same subset, more epochs → near-ceiling |
| **v8** | FineWeb-edu 1.05ep | **1024 / 226M** | **2.98\*** | **capacity → helps** |
| **v9** | FineWeb-edu 6.01B / ~1ep | **1280 / 357M + GQA** | **2.89\*** | **data + capacity → helps** |
| **v10** | FineWeb-edu 6.76B / ~1ep | **1280 / 357M + GQA** | **2.88\*** | **data-only top-up → small gain** |
\* FineWeb-edu val is a different (harder) distribution — **not comparable** to the
TinyStories val of v0–v5. Judge v6+ by sample quality + transfer, not the number.
### Four findings
### Three findings
1. **Data volume saturates.** TinyStories at dim768: 3.5× more tokens (v4→v5) bought only
−5% val, curve flat. The narrow synthetic corpus is exhausted at this model size.
@@ -137,18 +132,10 @@ TinyStories val of v0–v5. Judge v6+ by sample quality + transfer, not the numb
historical/scientific expository prose. (Cost: TinyStories transfer val 1.11 → 2.75.)
3. **Capacity helps.** v8 (dim1024, ~1 epoch) beats both v6 (dim768, same epoch, by 0.085)
and v7 (dim768, *more* data, by 0.035) → the dim768 runs were partly capacity-limited.
4. **Double-axis scale helps.** v9 scales both axes (dim1280/core357M + 6.01B FineWeb tokens)
and beats v8 by another 0.095 val loss (~3.2%). The direction is validated, but the gain is
still incremental and greedy decoding still repeats.
5. **Moving validation tails must stop.** v10 added one more FineWeb shard and got moving-tail
val 2.8816, but appending data moves the held-out tail. A fixed eval v1 was created from the
shard010 tail: v6/v7/v8/v9/v10 = 3.2328 / 3.1850 / 3.1515 / 2.9278 / 2.8814. Future runs
should report this fixed eval first.
**Meta-finding:** every lever is now in the **~3% or smaller** regime. Single-axis moves were
exhausted by v8; v9 confirms Chinchilla-style double-axis scale works; v10 shows a data-only
top-up mostly adapts to the new shard. The next useful run should change model/context, not just
append another shard.
**Meta-finding:** every *single*-axis lever (data volume, corpus breadth, capacity) is now
worth only **~3%**. Per the Chinchilla lesson, further gains require scaling **data and
capacity together** — single-axis moves are exhausted.
## Efficiency — throughput & MFU
@@ -179,18 +166,18 @@ versions — a fixed-MFU estimate is off by up to ~100× for the early launch-bo
the line: flash == composed SDPA (grads/PyTorch), GQA group=1 bit-identical to MHA, gradient
accumulation `accum=1` bit-identical, dropout p=0 bit-identical *and* dropout × recompute
bit-exact, the default path unchanged on every feature, and the **xserv closed-loop md5
byte-identical (`b04fc9f9`) throughout the deterministic gates**.
- **The closed loop matters.** Exporting to xserv and checking generated continuations caught
real bugs and proved the whole stack end-to-end.
byte-identical (`b04fc9f9`) throughout both phases**.
- **The closed loop matters.** Exporting to xserv and checking token-identical greedy output
caught real bugs and proved the whole stack end-to-end.
## Running it
Everything trains on a remote 8× RTX 5090 box; model artifacts live in a registry
(`tiny-models/v0…v10`). Serve any trained version in xserv:
(`tiny-models/v0…v8`). Serve any trained version in xserv:
```bash
# on the GPU box
cargo run -p xserv-model --release --bin xserv-cli -- <registry>/v10-fineweb-edu-dim1280-gqa-data6765 --max-tokens 100
cargo run -p xserv-model --release --bin xserv-cli -- <registry>/v8-fineweb-edu-dim1024 --max-tokens 100
# then type a prompt, e.g. In science,
```
@@ -205,6 +192,6 @@ cargo test --workspace # autograd grad-checks, PyTorch parity, DDP, e
## Doc index
- [`docs/evolution.md`](docs/evolution.md) — per-milestone changes across algorithm / architecture / infra / dataset.
- [`docs/runs/README.md`](docs/runs/README.md) — the v0–v10 comparison; [`docs/runs/0N-*.md`](docs/runs/) — per-run detail.
- [`docs/runs/README.md`](docs/runs/README.md) — the v0–v8 comparison; [`docs/runs/0N-*.md`](docs/runs/) — per-run detail.
- [`docs/00-*` … `17-*`](docs/) — per-phase design docs (build chain → tensor → autograd → transformer → training → perf → distributed → export → batched → allocator → bf16 → recompute → flash-attention → GQA → grad-accum → process-per-GPU → dropout).
- [`docs/known-issues.md`](docs/known-issues.md) — perf backlog (KI-1/2/3/5 fixed; process-per-GPU CLOSED = measured no-op; KI-4 = accepted modeling tradeoff).

263
crates/xtrain-autodiff/src/ops.rs
View File

@@ -398,8 +398,7 @@ pub fn repeat_kv(kv: &Var, nh: usize, batch: usize) -> Var {
}
/// Cross-entropy mean loss over logits `x:[rows,cols]` with one I32 target per
/// row. Negative targets are ignored, which is useful for assistant-only SFT
/// masks. Returns a scalar [`Var`]. Backward: `dx = (probs - onehot)/valid_rows`,
/// row. Returns a scalar [`Var`]. Backward: `dx = (probs - onehot)/rows`,
/// scaled by the upstream scalar grad.
pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
// CE math is fp32 (cross_entropy upcasts bf16 logits internally + caches fp32
@@ -408,22 +407,10 @@ pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
// fp32 logits buffer) is a real activation-memory saving at large vocab.
let logit_dtype = x.value().dtype();
let (probs, per_row) = x.value().cross_entropy(target);
let cols = x.value().shape()[1] as i32;
let target_host = target.to_device(xtrain_tensor::Device::Cpu);
let valid_rows = target_host
.as_slice::<i32>()
.iter()
.filter(|&&t| {
if t >= cols {
panic!("cross_entropy target {t} out of vocab range {cols}");
}
t >= 0
})
.count()
.max(1);
let rows = x.value().shape()[0];
// Mean loss as a host scalar wrapped back into a [1] tensor.
let mean = per_row.to_device(xtrain_tensor::Device::Cpu);
let mean_val: f32 = mean.as_slice::<f32>().iter().sum::<f32>() / valid_rows as f32;
let mean_val: f32 = mean.as_slice::<f32>().iter().sum::<f32>() / rows as f32;
let loss = Tensor::from_slice(&[mean_val], &[1]).to_device(x.value().device());
let target = target.clone();
@@ -433,251 +420,9 @@ pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
Box::new(move |d, parents| {
// `d` is the scalar upstream grad (1.0 when this is the loss root).
let upstream = d.to_device(xtrain_tensor::Device::Cpu).as_slice::<f32>()[0];
let scale = upstream / valid_rows as f32;
let scale = upstream / rows as f32;
let dx = Tensor::cross_entropy_backward(&probs, &target, scale);
Var::push_grad(&parents[0], dx.to_dtype(logit_dtype));
}),
)
}
/// 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));
}),
)
}
/// 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));
}),
)
}
/// Batched GRPO clipped-PG loss over `N` ragged completions packed into ONE
/// `forward_batched` (M2d): `logits` is `[R, vocab]` with `R = N·Lmax` rows in
/// sequence-major order (sample 0's `Lmax` rows, then sample 1's, …), each ragged
/// completion right-padded to the batch's `Lmax`. Prompt AND pad rows are masked
/// (`target < 0`), so they contribute nothing and carry no gradient — the
/// **right-pad-is-free-under-causal-attention** property (a real completion row
/// never attends to the trailing pad rows, so its logits equal the unpadded
/// single-sequence forward's).
///
/// Unlike the per-sample [`clipped_pg_loss`] (which folds a single scalar
/// `advantage` and a global `1/N_tokens` normaliser), this op takes **per-row**
/// `advantage[t]` (the owning sample's group-relative `A`) and **per-row**
/// `weight[t]` (the full normaliser, e.g. `1/(N_samples · n_s)` for sample `s`'s
/// completion rows, `0` at masked rows). It does NOT compute its own `inv_n`. With
/// `weight[t] = 1/(N_samples·n_s)` and `advantage[t] = A_s` this is **bit-equivalent
/// to the looped path** `Σ_s scale·(1/n_s)·clipped_pg_loss_s` (`scale = 1/N_samples`):
/// the per-row backward is local (`cross_entropy_backward` is row-wise), so the
/// batched row-`t` gradient equals the looped sample-`s` row-`t` gradient, and the
/// scalar loss equals the looped weighted sum. (`tests/autograd.rs`:
/// `clipped_pg_loss_batched_matches_looped`.) Degenerate points match
/// [`clipped_pg_loss`] (`A=0` ⇒ KL only; `ε→∞` ⇒ vanilla PG; `β=0` ⇒ no KL).
#[allow(clippy::too_many_arguments)]
pub fn clipped_pg_loss_batched(
logits: &Var,
target: &Tensor,
logp_old: &[f32],
logp_ref: &[f32],
advantage: &[f32],
weight: &[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 row");
assert_eq!(logp_ref.len(), rows, "logp_ref must have one entry per row");
assert_eq!(advantage.len(), rows, "advantage must have one entry per row");
assert_eq!(weight.len(), rows, "weight must have one entry per row");
let mut s = vec![0f32; rows]; // per-row scale for cross_entropy_backward(·,·,1.0)
let mut loss_val = 0f32;
for t in 0..rows {
if target_h[t] < 0 {
continue; // masked (prompt or pad) row — no contribution, no gradient
}
let (a, w) = (advantage[t], weight[t]);
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 * a, clipped * a);
let pg_t = unclipped_term.min(clipped_term);
let active = unclipped_term <= clipped_term; // min picks unclipped ⇒ grad flows
let d = logp_ref[t] - lp;
let kl_t = d.exp() - d - 1.0;
let pg_grad = if active { -a * ratio } else { 0.0 };
let kl_grad = beta * (1.0 - d.exp());
// The full per-row normaliser is folded into s (no global inv_n here).
s[t] = -(pg_grad + kl_grad) * w;
loss_val += (-pg_t + beta * kl_t) * w;
}
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];
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));
}),
)
}

263
crates/xtrain-autodiff/tests/autograd.rs
View File

@@ -1005,266 +1005,3 @@ 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)");
}
// clipped_pg_loss (M4 GRPO): per-token clipped PG + k3 KL, one completion. Grad-check
// the active (in-trust-region) path + the A=0 (KL-only) path, plus value-level
// degenerate checks (ε→∞ ⇒ vanilla PG, β=0 ⇒ no KL).
#[test]
fn clipped_pg_loss_bwd_and_degenerate() {
require_gpu();
let (rows, cols) = (6usize, 10usize);
let x_h = fill(rows * cols, 303);
// rows 0,1 masked (prompt); 2..6 supervised (completion).
let targets: Vec<i32> = (0..rows)
.map(|r| if r < 2 { -100 } else { (r * 2 % cols) as i32 })
.collect();
let mk_target = || Tensor::from_slice(&targets, &[rows]).to_device(Device::Cuda(0));
// logp_old = logπθ at the base logits ⇒ ρ≈1 (in trust region → active path).
let (_, per_row0) = cuda(&x_h, &[rows, cols]).cross_entropy(&mk_target());
let logp_old: Vec<f32> = per_row0
.to_device(Device::Cpu)
.as_slice::<f32>()
.iter()
.map(|p| -p)
.collect();
let logp_ref: Vec<f32> = logp_old.iter().map(|l| l - 0.3).collect(); // exercise KL
let (eps, beta) = (0.2f32, 0.1f32);
// Host replica of the forward loss as a function of per-row CE values.
let host_loss = {
let (tg, lo, lr) = (targets.clone(), logp_old.clone(), logp_ref.clone());
move |per_row_h: &[f32], a: f32, e: f32, b: f32| -> f32 {
let (mut pg, mut kl, mut n) = (0f32, 0f32, 0f32);
for t in 0..per_row_h.len() {
if tg[t] < 0 {
continue;
}
n += 1.0;
let lp = -per_row_h[t];
let ratio = (lp - lo[t]).exp();
let clipped = ratio.clamp(1.0 - e, 1.0 + e);
pg += (ratio * a).min(clipped * a);
let d = lr[t] - lp;
kl += d.exp() - d - 1.0;
}
let inv = if n > 0.0 { 1.0 / n } else { 1.0 };
-pg * inv + b * kl * inv
}
};
let per_row_of = |v: &[f32], s: &[usize]| {
let (_, pr) = cuda(v, s).cross_entropy(&mk_target());
pr.to_device(Device::Cpu).as_slice::<f32>().to_vec()
};
// (1) grad-check the active PG path (A>0, ρ≈1).
let adv = 0.7f32;
let x = Var::leaf(cuda(&x_h, &[rows, cols]));
let loss = ops::clipped_pg_loss(&x, &mk_target(), &logp_old, &logp_ref, adv, eps, beta);
loss.backward();
let dx = x.grad().unwrap().to_device(Device::Cpu);
let hl = host_loss.clone();
let lx = move |v: &[f32], s: &[usize]| hl(&per_row_of(v, s), adv, eps, beta);
report(
"clipped_pg dX (active)",
&grad_check(&x_h, &[rows, cols], &lx, dx.as_slice::<f32>(), cfg_nonlinear()),
);
// (2) grad-check the A=0 path (loss = β·mean KL; PG gradient must vanish).
let x0 = Var::leaf(cuda(&x_h, &[rows, cols]));
let loss0 = ops::clipped_pg_loss(&x0, &mk_target(), &logp_old, &logp_ref, 0.0, eps, beta);
loss0.backward();
let dx0 = x0.grad().unwrap().to_device(Device::Cpu);
let hl0 = host_loss.clone();
let lx0 = move |v: &[f32], s: &[usize]| hl0(&per_row_of(v, s), 0.0, eps, beta);
report(
"clipped_pg dX (A=0, KL only)",
&grad_check(&x_h, &[rows, cols], &lx0, dx0.as_slice::<f32>(), cfg_nonlinear()),
);
// (3) ε→∞ ⇒ vanilla PG (no clip): loss value == −mean(ρA) + β·mean KL.
let big = 1e9f32;
let lv = ops::clipped_pg_loss(&Var::leaf(cuda(&x_h, &[rows, cols])), &mk_target(), &logp_old, &logp_ref, adv, big, beta);
let got = lv.value().to_device(Device::Cpu).as_slice::<f32>()[0];
let pr0 = per_row_of(&x_h, &[rows, cols]);
let want = host_loss(&pr0, adv, big, beta);
assert!((got - want).abs() < 1e-4, "ε→∞ vanilla loss mismatch: {got} vs {want}");
// (4) β=0 ⇒ no KL term (loss == −mean pg only).
let lvb = ops::clipped_pg_loss(&Var::leaf(cuda(&x_h, &[rows, cols])), &mk_target(), &logp_old, &logp_ref, adv, eps, 0.0);
let gotb = lvb.value().to_device(Device::Cpu).as_slice::<f32>()[0];
let wantb = host_loss(&pr0, adv, eps, 0.0);
assert!((gotb - wantb).abs() < 1e-5, "β=0 loss mismatch: {gotb} vs {wantb}");
println!("clipped_pg_loss OK: grad-check (active + A=0) + degenerate (ε→∞ vanilla, β=0 no KL)");
}
// clipped_pg_loss_batched (M2d): N ragged completions packed + right-padded into ONE
// forward must equal the looped per-sample path Σ_s (1/N)·clipped_pg_loss_s. The
// per-row CE backward is row-local, so folding weight = 1/(N·n_s) into the batched
// op reproduces the looped gradient and weighted-sum loss bit-for-bit (f32 path).
#[test]
fn clipped_pg_loss_batched_matches_looped() {
require_gpu();
let (n, lmax, cols) = (3usize, 5usize, 10usize);
let rows = n * lmax;
let x_h = fill(rows * cols, 909);
// Per sample: row 0 = prompt (-100); rows 1..real_len = completion; rest = pad
// (-100). Different real_len ⇒ n_s = {2, 3, 1} completion rows.
let real_len = [3usize, 4, 2];
let adv_s = [0.7f32, -0.5, 0.3];
let mut targets = vec![-100i32; rows];
for s in 0..n {
for r in 1..real_len[s] {
let t = s * lmax + r;
targets[t] = ((t * 3) % cols) as i32;
}
}
let mk_target = || Tensor::from_slice(&targets, &[rows]).to_device(Device::Cuda(0));
// logp_old ≈ logπθ at base logits (ρ≈1), logp_ref offset to exercise the KL term.
let (_, per_row0) = cuda(&x_h, &[rows, cols]).cross_entropy(&mk_target());
let logp_old: Vec<f32> = per_row0
.to_device(Device::Cpu)
.as_slice::<f32>()
.iter()
.map(|p| -p)
.collect();
let logp_ref: Vec<f32> = logp_old.iter().map(|l| l - 0.3).collect();
let (eps, beta) = (0.2f32, 0.1f32);
// Per-row advantage (sample's A) + per-row weight 1/(N·n_s) (full normaliser).
let n_of = |s: usize| (0..lmax).filter(|&r| targets[s * lmax + r] >= 0).count() as f32;
let mut advantage = vec![0f32; rows];
let mut weight = vec![0f32; rows];
for s in 0..n {
let w = (1.0 / n as f32) * (1.0 / n_of(s));
for r in 0..lmax {
advantage[s * lmax + r] = adv_s[s];
weight[s * lmax + r] = w;
}
}
// Batched: one packed [R, vocab] forward + one backward.
let xb = Var::leaf(cuda(&x_h, &[rows, cols]));
let lb = ops::clipped_pg_loss_batched(
&xb, &mk_target(), &logp_old, &logp_ref, &advantage, &weight, eps, beta,
);
lb.backward();
let gb = xb.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>().to_vec();
let lb_val = lb.value().to_device(Device::Cpu).as_slice::<f32>()[0];
// Looped reference: per-sample slice → clipped_pg_loss → scale(1/N) → backward.
let mut g_ref = vec![0f32; rows * cols];
let mut loss_ref = 0f32;
for s in 0..n {
let r0 = s * lmax;
let xs_h = x_h[r0 * cols..(r0 + lmax) * cols].to_vec();
let tgt_s: Vec<i32> = targets[r0..r0 + lmax].to_vec();
let lo_s = logp_old[r0..r0 + lmax].to_vec();
let lr_s = logp_ref[r0..r0 + lmax].to_vec();
let xs = Var::leaf(cuda(&xs_h, &[lmax, cols]));
let tgt = Tensor::from_slice(&tgt_s, &[lmax]).to_device(Device::Cuda(0));
let ls = ops::clipped_pg_loss(&xs, &tgt, &lo_s, &lr_s, adv_s[s], eps, beta);
let scaled = ops::scale(&ls, 1.0 / n as f32);
scaled.backward();
let gs = xs.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>().to_vec();
g_ref[r0 * cols..(r0 + lmax) * cols].copy_from_slice(&gs);
loss_ref += scaled.value().to_device(Device::Cpu).as_slice::<f32>()[0];
}
let max_g = gb
.iter()
.zip(&g_ref)
.map(|(a, b)| (a - b).abs())
.fold(0.0f32, f32::max);
assert!(
(lb_val - loss_ref).abs() < 1e-5,
"batched loss {lb_val} vs looped {loss_ref}"
);
assert!(max_g < 1e-5, "batched grad vs looped: max|Δ| = {max_g}");
println!(
"clipped_pg_loss_batched OK: loss Δ={:.2e}, grad max|Δ|={:.2e} (== looped Σ_s 1/N·pg_s)",
(lb_val - loss_ref).abs(),
max_g
);
}

45
crates/xtrain-cuda/src/ffi.rs
View File

@@ -139,51 +139,6 @@ unsafe extern "C" {
period: i32,
s: CudaStream,
);
// RoPE at an absolute position offset (KV-cache decode, forward only): row
// `tok`'s position is `pos0 + tok` (no modulo). For a single decode token
// (tokens == 1) the one row sits at absolute position `pos0`.
pub fn launch_rope_at_f32(
x: *const f32,
y: *mut f32,
tokens: i32,
heads: i32,
head_dim: i32,
theta: f32,
pos0: i32,
s: CudaStream,
);
// RoPE with a per-row absolute position (batched KV-cache decode, M2b): row
// `tok`'s position is `positions[tok]`. Forward only.
pub fn launch_rope_pos_f32(
x: *const f32,
positions: *const i32,
y: *mut f32,
tokens: i32,
heads: i32,
head_dim: i32,
theta: f32,
s: CudaStream,
);
// Concatenate along the sequence dim: a:[bh,ta,hd], b:[bh,tb,hd] →
// out:[bh,ta+tb,hd] (device-side KV-cache append, M2c).
pub fn launch_cat_seq_f32(
a: *const f32,
b: *const f32,
out: *mut f32,
bh: i32,
ta_hd: i32,
tb_hd: i32,
s: CudaStream,
);
// Per-row scale: y[r,c] = x[r,c] * s[r] (GRPO policy-gradient backward).
pub fn launch_scale_rows_f32(
x: *const f32,
s: *const f32,
y: *mut f32,
rows: i32,
cols: i32,
stream: CudaStream,
);
pub fn launch_rope_dx_f32(
dy: *const f32,
dx: *mut f32,

23
crates/xtrain-distributed/src/bin/train_ddp.rs
View File

@@ -88,7 +88,6 @@ fn main() {
let val_tokens: usize = flag(&args, "--val-tokens", 0);
let eval_every: usize = flag(&args, "--eval-every", 0);
let eval_batches: usize = flag(&args, "--eval-batches", 64);
let sft_tsv = args.iter().any(|a| a == "--sft-tsv");
// Dropout (Phase T18/T21): residual-path dropout prob, active at training time
// only (inverted scaling), identity at eval/sampling/export. Default 0 = off
// (forward graph bit-identical to the no-dropout path). Mirrors bin/train; the
@@ -110,11 +109,6 @@ fn main() {
.position(|a| a == "--ckpt")
.and_then(|i| args.get(i + 1))
.map(PathBuf::from);
let init_ckpt: Option<PathBuf> = args
.iter()
.position(|a| a == "--init-ckpt")
.and_then(|i| args.get(i + 1))
.map(PathBuf::from);
// Use every visible GPU as a rank (CUDA_VISIBLE_DEVICES selects the set;
// device ordinals are 0..count within it).
@@ -135,19 +129,12 @@ fn main() {
);
// Reuse the cached token-id stream (v1's u16 cache); never re-tokenize 2GB.
let corpus = if sft_tsv {
Corpus::load_sft_tsv_cached(&tok_path, &corpus_path)
} else {
Corpus::load_cached(&tok_path, &corpus_path)
};
let corpus = Corpus::load_cached(&tok_path, &corpus_path);
println!(
"corpus: {} tokens, vocab {}",
corpus.len(),
corpus.vocab_size
);
if sft_tsv {
println!("SFT TSV: ON (assistant-only loss via ignore-index labels)");
}
let vocab = corpus.vocab_size;
// Hold out a tail slice for validation (rank 0 evaluates on it).
let (train_corpus, valid) = if val_tokens > 0 {
@@ -213,10 +200,6 @@ fn main() {
if dropout > 0.0 {
println!("dropout: ON (p={dropout}, residual-path, train-only inverted scaling)");
}
if let Some(path) = &init_ckpt {
println!("init checkpoint: {}", path.display());
}
let init_ckpt_for_ranks = init_ckpt.clone();
let results = launch(
&devices,
&train_corpus,
@@ -233,10 +216,6 @@ fn main() {
if flash {
m = m.with_flash(true);
}
if let Some(path) = &init_ckpt_for_ranks {
xtrain_train::checkpoint::load_into(path, &m.params())
.expect("load init checkpoint");
}
m
},
);

16
crates/xtrain-distributed/src/bin/train_ddp_mp.rs
View File

@@ -10,9 +10,7 @@
//!
//! Versus `train_ddp` (thread-per-GPU, kept as the regression baseline) the ONLY
//! difference is the launch model + cross-process UniqueId bootstrap. CLI flags
//! mirror `train_ddp` (incl. `--dropout` — same T21 wiring: `cfg.dropout` set here
//! and `train_rank` re-asserts `model.train()` each step), so it doubles as the
//! before→after throughput driver.
//! are identical, so it doubles as the before→after throughput driver.
//!
//! Run on dash5 (pick idle GPUs — dash5 is shared):
//! export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
@@ -110,11 +108,6 @@ fn main() {
let val_tokens: usize = flag(&args, "--val-tokens", 0);
let eval_every: usize = flag(&args, "--eval-every", 0);
let eval_batches: usize = flag(&args, "--eval-batches", 64);
// Dropout (Phase T18/T21): residual-path dropout prob, active at training time
// only (inverted scaling), identity at eval/sampling/export. Default 0 = off
// (bit-identical to the no-dropout path). Mirrors bin/train_ddp; propagates into
// cfg.dropout (below) and relies on T21's per-step model.train() in train_rank.
let dropout: f32 = flag(&args, "--dropout", 0.0f32);
let opts = ModelOpts {
bf16: args.iter().any(|a| a == "--bf16"),
recompute: args.iter().any(|a| a == "--recompute"),
@@ -143,9 +136,7 @@ fn main() {
(corpus, None)
};
let mut cfg =
Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
cfg.dropout = dropout;
let cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
if env.rank == 0 {
println!(
@@ -171,9 +162,6 @@ fn main() {
if opts.flash {
println!("flash-attention: ON (fused SDPA kernel, no materialized scores)");
}
if dropout > 0.0 {
println!("dropout: ON (p={dropout}, residual-path, train-only inverted scaling)");
}
}
let dcfg = DdpConfig {

209
crates/xtrain-distributed/tests/ddp_correctness.rs
View File

@@ -27,7 +27,6 @@ fn synth_corpus(vocab: usize, n_tokens: usize) -> Corpus {
.collect();
Corpus {
tokens,
labels: None,
vocab_size: vocab,
}
}
@@ -39,24 +38,17 @@ fn test_config(vocab: usize) -> Config {
cfg
}
/// Run `cfg`/`dcfg` as a DDP job over `devices` (the same launcher path as
/// production — `DdpContext::init` + `train_rank` per rank) and return rank 0's
/// (loss trace, final params on host, final `is_training()` flag). `cfg` carries
/// the dropout prob; `dcfg` carries the loop knobs. Caller asserts.
///
/// `world == 1` is the deterministic path: `all_reduce_average_grads` short-circuits
/// (no NCCL collective), so the run is bit-reproducible — used for the bit-identity
/// gate. `world >= 2` exercises the real cross-rank NCCL all-reduce, which is not
/// bit-reproducible run-to-run on this PCIe box (KI-5), so those gates use the same
/// ULP/relative tolerances as the rest of this file.
fn run_ddp(
devices: &[u32],
/// Run `cfg`/`dcfg` as a 2-rank DDP job (the same launcher path as production) and
/// return rank 0's (loss trace, final params on host, final `is_training()` flag).
/// `cfg` carries the dropout prob; `dcfg` carries the loop knobs. Caller asserts.
fn run_ddp2(
cfg: Config,
corpus: &Corpus,
valid: Option<&Corpus>,
dcfg: &DdpConfig,
) -> (Vec<f32>, Vec<Vec<f32>>, bool) {
let world = devices.len();
let world = 2usize;
let devices = [0u32, 1u32];
let id = get_unique_id();
let results: Vec<(Vec<f32>, Vec<Vec<f32>>, bool)> = std::thread::scope(|s| {
let handles: Vec<_> = devices
@@ -435,31 +427,21 @@ fn ddp_throughput_scaling() {
}
}
/// T21 regression: prove dropout is actually LIVE under DDP (with `p>0`), and that
/// `p=0` is bit-identical to the no-dropout path. Guards the V9-PILOT launcher-
/// wiring gap — `train_ddp` had no `--dropout` flag and `train_rank` never called
/// `model.train()`, so under DDP every forward ran in the default eval mode and
/// dropout was a silent identity regardless of config. Op/single-GPU tests never
/// exercised dropout-under-DDP, so it slipped through; this test runs the REAL
/// launcher path (`DdpContext::init` + `train_rank`).
/// T21 regression: prove dropout is actually LIVE under DDP, and that p=0 is
/// bit-identical to the no-dropout path. Guards the V9-PILOT launcher-wiring gap —
/// `train_ddp` had no `--dropout` flag and `train_rank` never called `model.train()`,
/// so under DDP every forward ran in the default eval mode and dropout was a silent
/// identity regardless of config. Op/single-GPU tests never exercised dropout-under-
/// DDP, so it slipped through; this test runs the real launcher path (`train_rank`).
///
/// On the pre-T21 code, both load-bearing gates FAIL: GATE B (p>0 trace would be
/// bit-identical to p=0 — model stuck in eval mode → dropout is identity) and GATE C
/// (`is_training()` would be false after the run).
///
/// p=0 regression (GATE A) is checked at `world=1`, ONE step, where the NCCL
/// all-reduce short-circuits: the p=0 FORWARD is byte-identical to no-dropout so the
/// loss is BIT-IDENTICAL (== 0.0), and the post-step params match within the engine's
/// atomicAdd backward-reduction ULP floor (< 1e-7, dropout-independent — the
/// fresh-train md5 caveat). The cross-rank NCCL all-reduce (`world>=2`) is not
/// bit-reproducible run-to-run on this PCIe box (KI-5, observed ≤~2.4e-7), so the
/// `world=2` p=0-vs-no-dropout check (GATE A2) uses the same KI-5 ULP tolerance as the
/// rest of this file. GATE B's live-dropout signal (>1e-3) sits ~4 orders of magnitude
/// above every noise floor here, so it carries the load.
/// With dropout fixed across the 4 sub-runs, all three checks below would FAIL on the
/// pre-T21 code: (a) the p>0 trace would be bit-identical to p=0 (model stuck in eval
/// mode → identity), and (c) `is_training()` would be false after the run.
#[test]
fn ddp_dropout_is_live_and_p0_bit_identical() {
if device::device_count().unwrap_or(0) < 2 {
eprintln!("skip: need >= 2 GPUs");
let world = 2usize;
if device::device_count().unwrap_or(0) < world as i32 {
eprintln!("skip: need >= {world} GPUs");
return;
}
@@ -492,92 +474,62 @@ fn ddp_dropout_is_live_and_p0_bit_identical() {
ckpt_path: None,
};
// --- GATE A: p=0 == no-dropout at world=1, ONE step (the deterministic scope). ---
// The regression guard for `--dropout 0`. ops::dropout(p=0) returns x.clone() (a
// graph no-op) regardless of training mode, so the p=0 FORWARD graph is byte-for-
// byte the no-dropout forward → loss[0] must be BIT-IDENTICAL (the load-bearing
// claim, asserted == 0.0). At world=1 the NCCL all-reduce short-circuits, and one
// step has no optimizer-state compounding; the only residual non-determinism is
// the engine's atomicAdd backward-reduction ORDER (the documented fresh-train md5
// caveat — dropout-INDEPENDENT, present with or without the dropout op), which
// moves the post-step params by a single grad ULP. So params are checked against
// that tight reduction floor (< 1e-7), the same nature as the cross-rank KI-5
// tolerance used elsewhere in this file — not a dropout signal. GATE B (live) has
// a >1e-3 signal, ~4 orders of magnitude above this floor, so it carries the load.
let d1 = [0u32];
let dcfg_1step = DdpConfig {
steps: 1,
eval_every: 0,
..base_dcfg.clone()
};
let cfg_nodrop = test_config(vocab); // cfg.dropout defaults to 0.0
assert_eq!(cfg_nodrop.dropout, 0.0, "baseline cfg must have dropout 0");
let mut cfg_p0 = test_config(vocab);
cfg_p0.dropout = 0.0; // explicitly set p=0 — must not perturb anything
let (loss_nd1, params_nd1, _) = run_ddp(&d1, cfg_nodrop, &corpus, None, &dcfg_1step);
let (loss_p01, params_p01, _) = run_ddp(&d1, cfg_p0, &corpus, None, &dcfg_1step);
let max_loss_diff_1 = (loss_nd1[0] - loss_p01[0]).abs();
let max_param_diff_1 = params_nd1
.iter()
.zip(&params_p01)
.flat_map(|(a, b)| a.iter().zip(b).map(|(x, y)| (x - y).abs()))
.fold(0.0f32, f32::max);
// (1) p=0 config — the no-dropout baseline. cfg.dropout defaults to 0.0.
let cfg_p0 = test_config(vocab);
assert_eq!(cfg_p0.dropout, 0.0, "baseline cfg must have dropout 0");
let (loss_p0, params_p0, _) = run_ddp2(cfg_p0, &corpus, Some(&valid), &base_dcfg);
// (2) Same config, dropout disabled by p=0 but explicitly set — must be the
// SAME run (sanity: setting dropout=0 doesn't perturb anything).
let mut cfg_p0b = test_config(vocab);
cfg_p0b.dropout = 0.0;
let (loss_p0b, params_p0b, _) = run_ddp2(cfg_p0b, &corpus, Some(&valid), &base_dcfg);
// (3) Same config + data + seed, but dropout p=0.2 ON.
let mut cfg_p = test_config(vocab);
cfg_p.dropout = 0.2;
let (loss_p, _params_p, train_flag_p) = run_ddp2(cfg_p, &corpus, Some(&valid), &base_dcfg);
// GATE A — p=0 is bit-identical to the no-dropout path (regression guard).
// ops::dropout(p=0) is a clone no-op regardless of training mode, so these two
// runs must agree to the last bit on BOTH the loss trace and the final params.
let mut max_loss_diff = 0.0f32;
for (a, b) in loss_p0.iter().zip(&loss_p0b) {
max_loss_diff = max_loss_diff.max((a - b).abs());
}
let mut max_param_diff = 0.0f32;
for (a, b) in params_p0.iter().zip(&params_p0b) {
for (x, y) in a.iter().zip(b) {
max_param_diff = max_param_diff.max((x - y).abs());
}
}
println!(
"T21 GATE A (world=1, 1 step, p=0 vs no-dropout): |loss diff| = {max_loss_diff_1:.3e} \
(bit-identical forward), max |param diff| = {max_param_diff_1:.3e} (atomicAdd floor)"
"T21 GATE A (p=0 bit-identical): max |loss diff| = {max_loss_diff:.3e}, \
max |param diff| = {max_param_diff:.3e}"
);
assert_eq!(
max_loss_diff_1, 0.0,
"world=1 p=0 forward loss not bit-identical to no-dropout path"
max_loss_diff, 0.0,
"p=0 DDP loss trace not bit-identical to no-dropout path"
);
assert!(
max_param_diff_1 < 1e-7,
"world=1 p=0 post-step params diverged from no-dropout beyond the atomicAdd \
reduction floor: {max_param_diff_1:.3e}"
);
// --- world=2 runs: real cross-rank NCCL all-reduce (the production path). ---
let d2 = [0u32, 1u32];
let mut cfg_p0_w2 = test_config(vocab);
cfg_p0_w2.dropout = 0.0;
let mut cfg_p_w2 = test_config(vocab);
cfg_p_w2.dropout = 0.2;
let (loss_p0_2, _params_p0_2, _) = run_ddp(&d2, cfg_p0_w2, &corpus, Some(&valid), &base_dcfg);
let (loss_p_2, _params_p_2, _) = run_ddp(&d2, cfg_p_w2, &corpus, Some(&valid), &base_dcfg);
// GATE A2 — under DDP (world=2), p=0 matches a separate no-dropout baseline within
// NCCL's run-to-run ULP noise (KI-5; the all-reduce is not bit-reproducible). This
// confirms enabling dropout=0 doesn't perturb the DDP path beyond that noise floor.
let (loss_nd_2, _, _) = run_ddp(&d2, test_config(vocab), &corpus, Some(&valid), &base_dcfg);
let max_loss_diff_2 = loss_nd_2
.iter()
.zip(&loss_p0_2)
.map(|(a, b)| (a - b).abs())
.fold(0.0f32, f32::max);
println!(
"T21 GATE A2 (world=2 p=0 vs no-dropout, KI-5 noise): max |loss diff| = {max_loss_diff_2:.3e}"
);
assert!(
max_loss_diff_2 < 1e-6,
"world=2 p=0 diverged from no-dropout beyond NCCL noise: {max_loss_diff_2:.3e}"
assert_eq!(
max_param_diff, 0.0,
"p=0 DDP final params not bit-identical to no-dropout path"
);
// GATE B — dropout is LIVE with p>0 under DDP. If model.train() were not wired
// (the pre-T21 bug), the model would stay in eval mode and the p=0.2 forward would
// be IDENTITY → loss trace bit-identical to p=0 (diff at the ~1e-7 NCCL noise
// floor). A difference orders of magnitude above that proves dropout masks are
// actually applied during the training forward — and that they survive the mid-run
// eval flips (model.train() is re-asserted each step). Inverted scaling + masking
// perturbs every step, so the gap is large (>1e-3 ≫ KI-5 noise ~2.4e-7).
let max_live_diff = loss_p0_2
.iter()
.zip(&loss_p_2)
.map(|(a, b)| (a - b).abs())
.fold(0.0f32, f32::max);
// (the pre-T21 bug), the model would stay in eval mode and the p=0.2 forward
// would be IDENTITY → loss trace bit-identical to p=0. A real, sizeable
// difference proves dropout masks are actually applied during the training
// forward (and survive the mid-run eval flips, since model.train() is re-asserted
// each step). Inverted scaling + masking perturbs every step, so the gap is large.
let mut max_live_diff = 0.0f32;
for (a, b) in loss_p0.iter().zip(&loss_p) {
max_live_diff = max_live_diff.max((a - b).abs());
}
println!(
"T21 GATE B (dropout live, world=2): p0[last]={:.6} p0.2[last]={:.6} max |loss diff| = {max_live_diff:.3e}",
loss_p0_2.last().unwrap(),
loss_p_2.last().unwrap()
"T21 GATE B (dropout live): p0[last]={:.6} p0.2[last]={:.6} max |loss diff| = {max_live_diff:.3e}",
loss_p0.last().unwrap(),
loss_p.last().unwrap()
);
assert!(
max_live_diff > 1e-3,
@@ -585,30 +537,17 @@ fn ddp_dropout_is_live_and_p0_bit_identical() {
(model.train() not wired): max |loss diff| {max_live_diff:.3e}"
);
// No NaN/Inf in the p>0 run (dropout converges normally under DDP).
// GATE C — train_rank leaves the model in TRAINING mode (direct proof that
// model.train() was called and survives the final-step eval). On the pre-T21
// code this would be false (model never left the default eval mode).
assert!(
loss_p_2.iter().all(|l| l.is_finite()),
"p=0.2 DDP loss has non-finite values"
train_flag_p,
"model not in training mode after DDP run — model.train() not wired in train_rank"
);
// GATE C — train_rank actually sets TRAINING mode (direct, complementary proof of
// model.train() being wired). Use a dedicated short run with eval_every=0 so no
// eval fires: a model that finishes a training step in training mode proves
// train_rank called model.train(). (With eval enabled, eval_loss → model.eval()
// runs LAST on the final step and legitimately leaves the model in eval mode —
// same as the single-GPU loop — so is_training() after an eval-enabled run reflects
// the final eval, not the training-mode wiring. GATE B already proves dropout
// survives the mid-run eval flips via the per-step model.train() restore.) On the
// pre-T21 code is_training() stays false (model never left the default eval mode).
let dcfg_noeval = DdpConfig {
steps: 2,
eval_every: 0,
..base_dcfg.clone()
};
let (_, _, train_flag) = run_ddp(&d1, cfg_p_w2, &corpus, None, &dcfg_noeval);
// No NaN/Inf in the p>0 run (dropout converges normally under DDP).
assert!(
train_flag,
"model not in training mode after a no-eval DDP run — model.train() not wired in train_rank"
loss_p.iter().all(|l| l.is_finite()),
"p=0.2 DDP loss has non-finite values"
);
println!("T21 GATE C (train_rank sets training mode): is_training() == true ✅");
}

133
crates/xtrain-distributed/tests/ddp_proc.rs
View File

@@ -10,14 +10,6 @@
//! (a) multi-process loss matches single-GPU within `<1e-3`,
//! (b) cross-rank params agree within `<1e-6` (KI-5 ULP tolerance),
//! (c) multi-process loss matches the thread-per-GPU `launch` path within `<1e-3`.
//!
//! T21-for-proc regression `proc_per_gpu_dropout_is_live_and_p0_matches_no_dropout`
//! (below) additionally proves that `--dropout` propagates through the process-per-
//! GPU launcher — the analogue of the thread-per-GPU T21 fix. Pre-fix
//! `train_ddp_mp` had no `--dropout` flag, so `cfg.dropout` stayed 0 regardless of
//! what the user passed, silently disabling dropout under process-per-GPU. The
//! GATE B loss-trace signal (>1e-3 gap between p=0 and p=0.2) sits orders of
//! magnitude above the KI-5 cross-rank noise floor and catches that gap directly.
#![cfg(not(no_cuda))]
@@ -45,7 +37,6 @@ fn synth_corpus() -> Corpus {
.collect();
Corpus {
tokens,
labels: None,
vocab_size: VOCAB,
}
}
@@ -82,20 +73,8 @@ fn dcfg(batch_size: usize) -> DdpConfig {
// The dump dir is passed launcher→worker via this env key (separate from the
// XTRAIN_* keys the launcher sets); workers write `rank{N}.dump` there.
const ENV_DUMP_DIR: &str = "XTRAIN_TEST_DUMP_DIR";
// Optional launcher→worker channel for `cfg.dropout`. Absent = 0.0 = the existing
// correctness test's contract (no perturbation). The T21-for-proc regression test
// below sets it before each `launch_processes` call to prove the process-per-GPU
// path actually plumbs `--dropout` into every worker's model.
const ENV_DROPOUT: &str = "XTRAIN_TEST_DROPOUT";
const GLOBAL_BATCH: usize = 8;
fn worker_dropout() -> f32 {
std::env::var(ENV_DROPOUT)
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(0.0)
}
// ── Worker entry: runs when this test binary is re-execed by launch_processes ─
fn run_as_worker_if_needed() {
@@ -107,13 +86,7 @@ fn run_as_worker_if_needed() {
// production `run_worker` wrapper is exercised by `bin/train_ddp_mp` on dash5.
let ctx = DdpContext::init(env.rank, env.world, env.id, env.local_rank);
let device = Device::Cuda(env.local_rank);
// Mirrors bin/train_ddp_mp's `cfg.dropout = dropout` wiring — the T21-for-proc
// regression: if this line were missing (the pre-fix launcher's exact gap),
// `cfg.dropout` would stay 0 and the GATE B test below would find a bit-
// identical p=0 / p=0.2 loss trace and FAIL.
let mut cfg = test_config();
cfg.dropout = worker_dropout();
let model = build_model(cfg, device);
let model = build_model(test_config(), device);
let res = train_rank(
&ctx,
&model,
@@ -229,16 +202,8 @@ fn proc_per_gpu_matches_single_gpu_and_thread_path() {
let dump_dir = std::env::temp_dir().join(format!("xtrain_t17_{}", std::process::id()));
std::fs::create_dir_all(&dump_dir).unwrap();
// SAFETY: single-threaded test (forced by --test-threads=1) sets this env
// before spawning workers; no concurrent env access. ENV_DROPOUT is cleared
// defensively — libtest orders `--test-threads=1` runs alphabetically, so the
// sibling `proc_per_gpu_dropout_is_live_...` test (starts with 'd') runs BEFORE
// this one (starts with 'm'). If it happened to leak `ENV_DROPOUT=0.2` in this
// process's env, the workers here would inherit it (Command inherits parent
// env by default) and build with dropout=0.2 while the single-GPU baseline
// (run_single_gpu → test_config → dropout=0) stays at 0 — GATE (a) would blow up.
// Explicit remove here severs that ordering coupling.
// before spawning workers; no concurrent env access.
unsafe {
std::env::remove_var(ENV_DROPOUT);
std::env::set_var(ENV_DUMP_DIR, &dump_dir);
}
// Re-exec the test binary but run ONLY this test, single-threaded, so the
@@ -307,100 +272,6 @@ fn proc_per_gpu_matches_single_gpu_and_thread_path() {
let _ = std::fs::remove_dir_all(&dump_dir);
}
/// T21-for-proc regression: prove that `--dropout` actually reaches the model
/// under process-per-GPU. The pre-fix `bin/train_ddp_mp` had no `--dropout` flag
/// and never set `cfg.dropout`, so the launcher's worker built its model with
/// dropout stuck at 0 — silent identity, regardless of what the user passed. The
/// thread-per-GPU T21 fix caught the analogous gap; this test caps the same gap
/// on the proc-per-GPU path with the same GATE-B pattern (loss trajectory of a
/// p=0.2 run differs from p=0 by a large margin, well above the NCCL noise floor).
///
/// Both runs share the corpus, the initial params (via `build_model`'s deterministic
/// LCG), and every other config knob; the ONLY difference is `cfg.dropout`. If the
/// worker didn't plumb the env-provided dropout into `cfg.dropout` (the exact pre-
/// fix regression), both traces would be bit-identical and this test would FAIL.
/// The `>1e-3` threshold sits orders of magnitude above the KI-5 cross-rank ULP
/// noise floor (~1e-7 on this PCIe box), so it's a hard signal for "dropout is
/// active" rather than a noise measurement. Mirrors
/// `ddp_dropout_is_live_and_p0_bit_identical` in ddp_correctness.rs for T21's
/// thread-per-GPU fix.
#[test]
fn proc_per_gpu_dropout_is_live_and_p0_matches_no_dropout() {
run_as_worker_if_needed();
let world = 2usize;
if device::device_count().unwrap_or(0) < world as i32 {
eprintln!("skip: need >= {world} GPUs");
return;
}
let base_dump_dir = std::env::temp_dir().join(format!("xtrain_t21mp_{}", std::process::id()));
std::fs::create_dir_all(&base_dump_dir).unwrap();
let worker_args = [
"--exact".to_string(),
"proc_per_gpu_dropout_is_live_and_p0_matches_no_dropout".to_string(),
"--test-threads=1".to_string(),
"--nocapture".to_string(),
];
// Helper: launch `world` workers with a specific dropout prob (via env), read
// rank 0's loss trace, clean up. Uses a subdir per run so the two invocations
// do not clobber each other's dumps.
let mut launch_with_dropout = |p: f32, tag: &str| -> Vec<f32> {
let dump_dir = base_dump_dir.join(tag);
std::fs::create_dir_all(&dump_dir).unwrap();
// SAFETY: single-threaded test (forced by --test-threads=1); no concurrent env access.
unsafe {
std::env::set_var(ENV_DUMP_DIR, &dump_dir);
std::env::set_var(ENV_DROPOUT, format!("{p}"));
}
launch_processes(world, &worker_args).expect("worker processes failed");
let (losses, _) = read_dump(dump_dir.to_str().unwrap(), 0);
losses
};
let loss_p0 = launch_with_dropout(0.0, "p0");
let loss_p1 = launch_with_dropout(0.2, "p02");
// GATE B — dropout is LIVE under process-per-GPU with p>0. If the worker
// didn't set `cfg.dropout` (the pre-fix gap), the two traces would match to
// the ~1e-7 NCCL noise floor. Anything above ~1e-3 is unambiguous evidence
// that dropout masks are actually applied in every worker's forward.
let max_live_diff = loss_p0
.iter()
.zip(&loss_p1)
.map(|(a, b)| (a - b).abs())
.fold(0.0f32, f32::max);
println!(
"T21-proc GATE B (dropout live under proc-per-GPU): p0[last]={:.6} p0.2[last]={:.6} max |loss diff| = {max_live_diff:.3e}",
loss_p0.last().unwrap(),
loss_p1.last().unwrap()
);
assert!(
max_live_diff > 1e-3,
"p=0.2 proc-per-GPU loss matches p=0 — dropout NOT plumbed through the \
process-per-GPU launcher (cfg.dropout stayed 0 in the worker): max |loss diff| {max_live_diff:.3e}"
);
// No NaN/Inf in the p>0 run.
assert!(
loss_p1.iter().all(|l| l.is_finite()),
"p=0.2 proc-per-GPU loss has non-finite values"
);
// Clear the launcher→worker env keys so we don't leak state to anything that
// runs later in this process. `proc_per_gpu_matches_single_gpu_and_thread_path`
// clears ENV_DROPOUT defensively too, but keeping the invariant "each test
// leaves the env as it found it" costs nothing.
// SAFETY: single-threaded test (forced by --test-threads=1); no concurrent env access.
unsafe {
std::env::remove_var(ENV_DROPOUT);
std::env::remove_var(ENV_DUMP_DIR);
}
let _ = std::fs::remove_dir_all(&base_dump_dir);
}
fn max_rel(a: &[f32], b: &[f32]) -> f32 {
a.iter()
.zip(b)

436
crates/xtrain-model/src/decode.rs
View File

@@ -1,436 +0,0 @@
//! KV-cache incremental-decode engine (post-training M2a, single sequence).
//!
//! The naive sampler ([`crate::TinyTransformer`] via `train::sample::generate`)
//! re-runs the full forward over the whole growing prefix every step — O(t²) and
//! a fresh autograd graph per token. This is the inference engine that replaces it:
//! a per-layer **K/V cache** + a **single-token incremental forward** that processes
//! one new token at a time, attending to the cached keys/values.
//!
//! Built on three primitives, all gated by their own correctness tests:
//! - [`Tensor::rope_at`](xtrain_tensor::Tensor::rope_at): RoPE at the token's
//! absolute position (not row-in-tile), so cached post-RoPE K matches the full
//! forward (bit-identical, `integration::rope_at_matches_full_rope_row`).
//! - [`Tensor::decode_attention`](xtrain_tensor::Tensor::decode_attention): the
//! single-query × cached-K/V SDPA, equal to the full causal attention's last row
//! (`integration::decode_attention_matches_full_attention_last_row`).
//! - this module's per-token block forward, mirroring `model::block_forward` at the
//! raw-Tensor level (no autograd tape — inference needs no gradients).
//!
//! Correctness gate (the M2 centerpiece): KV-cache greedy decode is **token-
//! identical** to the naive full-recompute greedy (`tests/decode_kv.rs`).
//!
//! Prefill is just the first `prompt.len()` decode steps (one token at a time) —
//! one code path, at the cost of a non-batched prefill (M2b adds batched prefill +
//! ragged batch decode). The cache is host-accumulated (token-major f32) and the
//! K/V tensor is rebuilt per step; the host round-trip is small (`num_kv·head_dim`
//! floats/token/layer) and is the honest M2a baseline — M2b moves it device-side.
#![cfg(not(no_cuda))]
use crate::TinyTransformer;
use xtrain_tensor::{DType, Device, Tensor};
/// Per-layer K/V cache: token-major host accumulation. For each layer, `k[li]` and
/// `v[li]` hold `[T, num_kv, head_dim]` (f32, flattened), grown by one token's
/// `num_kv·head_dim` values per decode step. Stored f32 (an exact upcast of the
/// bf16 projection output); rebuilt to the compute dtype when forming the K/V
/// tensor, so bf16 values round-trip bit-for-bit.
struct KVCache {
k: Vec<Option<Tensor>>,
v: Vec<Option<Tensor>>,
}
impl KVCache {
fn new(n_layers: usize) -> Self {
Self {
k: (0..n_layers).map(|_| None).collect(),
v: (0..n_layers).map(|_| None).collect(),
}
}
/// Append one token's K/V (`[bh,1,hd]`, compute dtype) to layer `li`, growing the
/// device-resident `[bh,T,hd]` cache via `cat_seq` (no host round-trip, M2c).
fn append(&mut self, li: usize, k_bh: Tensor, v_bh: Tensor) {
self.k[li] = Some(match self.k[li].take() {
Some(c) => c.cat_seq(&k_bh),
None => k_bh,
});
self.v[li] = Some(match self.v[li].take() {
Some(c) => c.cat_seq(&v_bh),
None => v_bh,
});
}
}
/// Linear `x @ W` in the compute dtype — mirrors `model::linear` (bf16 casts the
/// fp32-master weight to bf16 on the fly; the activation stream is already bf16).
fn linear_t(cdt: DType, x: &Tensor, w: &Tensor) -> Tensor {
match cdt {
DType::F32 => x.matmul(w),
DType::BF16 => x.matmul(&w.to_dtype(DType::BF16)),
_ => unreachable!("compute dtype must be F32/BF16"),
}
}
/// A norm/QK-norm gamma in the compute dtype — mirrors `model::norm_gamma`.
fn gamma_t(cdt: DType, g: &Tensor) -> Tensor {
match cdt {
DType::F32 => g.clone(),
DType::BF16 => g.to_dtype(DType::BF16),
_ => unreachable!("compute dtype must be F32/BF16"),
}
}
/// Greedy KV-cache decode: continue `prompt` by `max_new` tokens, argmax each step.
/// Returns the full token sequence (prompt + generated), matching the naive
/// `sample::generate` interface for `temperature == 0`. Token-identical to the
/// naive full-recompute greedy (gated by `tests/decode_kv.rs`).
pub fn generate_greedy_cached(
model: &TinyTransformer,
device: Device,
prompt: &[i32],
max_new: usize,
) -> Vec<i32> {
let mut rng = 0u64;
generate_cached(model, device, prompt, max_new, 0.0, &mut rng)
}
/// KV-cache decode with temperature sampling (`temperature == 0` → greedy argmax,
/// matching [`generate_greedy_cached`]; otherwise sample from `softmax(logits/T)`).
/// The KV-cache rollout the GRPO loop uses: each step allocates only a single-row
/// `[1, vocab]` logits buffer (vs the naive sampler's `[seq, vocab]`), so it is far
/// lighter on memory + the allocator — the naive sampler fragments the caching
/// allocator over a long rollout, which is the M4 "rollout is the long pole" wall.
pub fn generate_cached(
model: &TinyTransformer,
device: Device,
prompt: &[i32],
max_new: usize,
temperature: f32,
rng_state: &mut u64,
) -> Vec<i32> {
assert!(!prompt.is_empty(), "prompt must be non-empty");
let cfg = model.config();
let cdt = model.compute_dtype();
let n_layers = cfg.n_layers;
// params() is a stable, documented order (see TinyTransformer::params):
// [0] = embed [vocab, dim]
// [1 + li*11 .. +11] = layer li's 11 leaves, in block_params order:
// attn_norm, wq, wk, wv, q_norm, k_norm, wo, ffn_norm, w_gate, w_up, w_down
// [1 + n_layers*11] = final_norm [dim]
// [1 + n_layers*11 + 1] = lm_head [dim, vocab]
let params: Vec<Tensor> = model.params().iter().map(|p| p.value()).collect();
assert_eq!(
params.len(),
1 + n_layers * 11 + 2,
"unexpected param layout for decode"
);
let embed = &params[0];
let final_norm = &params[1 + n_layers * 11];
let lm_head = &params[1 + n_layers * 11 + 1];
let mut cache = KVCache::new(n_layers);
let mut tokens = prompt.to_vec();
// Prefill: feed each prompt token in order; the last step's logits are the
// distribution for the first generated token.
let mut logits = Vec::new();
for (pos, &tok) in prompt.iter().enumerate() {
logits = decode_step(&params, cfg, cdt, device, &mut cache, tok, pos, embed, final_norm, lm_head);
}
for _ in 0..max_new {
let next = if temperature <= 0.0 {
argmax(&logits) as i32
} else {
sample_temperature(&logits, temperature, rng_state) as i32
};
tokens.push(next);
let pos = tokens.len() - 1; // absolute position of the token just appended
logits = decode_step(&params, cfg, cdt, device, &mut cache, next, pos, embed, final_norm, lm_head);
}
tokens
}
/// Sample a token from `softmax(logits / temperature)` (numerically stable). Same
/// LCG + inverse-CDF scheme as the naive `sample::sample_temperature`.
fn sample_temperature(row: &[f32], temperature: f32, rng_state: &mut u64) -> usize {
let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let exps: Vec<f32> = row.iter().map(|&x| ((x - max) / temperature).exp()).collect();
let sum: f32 = exps.iter().sum();
*rng_state = rng_state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
let r = ((*rng_state >> 32) as f32 / u32::MAX as f32) * sum;
let mut acc = 0.0;
for (i, &e) in exps.iter().enumerate() {
acc += e;
if acc >= r {
return i;
}
}
exps.len() - 1
}
/// One incremental decode step for token `tok` at absolute position `pos`: append
/// its K/V to the cache and return the next-token logits as host f32 `[vocab]`.
#[allow(clippy::too_many_arguments)]
fn decode_step(
params: &[Tensor],
cfg: &crate::Config,
cdt: DType,
device: Device,
cache: &mut KVCache,
tok: i32,
pos: usize,
embed: &Tensor,
final_norm: &Tensor,
lm_head: &Tensor,
) -> Vec<f32> {
let (nh, hd, num_kv) = (cfg.n_heads, cfg.head_dim, cfg.num_kv_heads);
let dim = cfg.dim;
let scale = 1.0 / (hd as f32).sqrt();
let (theta, eps) = (cfg.rope_theta, cfg.eps);
let n_layers = cfg.n_layers;
// Embedding (fp32 table) → activation stream in the compute dtype.
let ids = Tensor::from_slice(&[tok], &[1]).to_device(device);
let mut h = embed.embedding(&ids); // [1, dim] f32
if cdt == DType::BF16 {
h = h.to_dtype(DType::BF16);
}
for li in 0..n_layers {
let base = 1 + li * 11;
let (attn_norm, wq, wk, wv) =
(&params[base], &params[base + 1], &params[base + 2], &params[base + 3]);
let (q_norm, k_norm, wo) = (&params[base + 4], &params[base + 5], &params[base + 6]);
let (ffn_norm, w_gate, w_up, w_down) =
(&params[base + 7], &params[base + 8], &params[base + 9], &params[base + 10]);
// --- Attention sub-block (pre-norm + cached-KV attention + residual) ---
let normed = h.rms_norm(&gamma_t(cdt, attn_norm), eps).0; // [1, dim]
// Q: project → per-head QK-norm → RoPE at absolute position `pos`.
let q = linear_t(cdt, &normed, wq).reshape(&[1, nh, hd]); // [1, nh, hd]
let q = q.reshape(&[nh, hd]).rms_norm(&gamma_t(cdt, q_norm), eps).0;
let q = q.reshape(&[1, nh, hd]).rope_at(theta, pos);
let q_bh = q.reshape(&[nh, 1, hd]); // seq=1 ⇒ the head-transpose is a no-op on data
// K: same as Q (QK-norm + RoPE). V: project only. Append each as [num_kv,1,hd]
// (bh-major) into the device cache; no host round-trip, no transpose (M2c).
let k = linear_t(cdt, &normed, wk).reshape(&[1, num_kv, hd]);
let k = k.reshape(&[num_kv, hd]).rms_norm(&gamma_t(cdt, k_norm), eps).0;
let k_bh = k.reshape(&[1, num_kv, hd]).rope_at(theta, pos).reshape(&[num_kv, 1, hd]);
let v_bh = linear_t(cdt, &normed, wv).reshape(&[num_kv, 1, hd]);
cache.append(li, k_bh, v_bh);
// repeat_kv the cached [num_kv,T,hd] to [nh,T,hd] for the SDPA.
let expand = |c: &Tensor| if num_kv == nh { c.clone() } else { c.repeat_kv(nh, 1) };
let k_full = expand(cache.k[li].as_ref().unwrap());
let v_full = expand(cache.v[li].as_ref().unwrap());
let attn = q_bh.decode_attention(&k_full, &v_full, scale); // [nh, hd]
let attn = attn.reshape(&[1, dim]); // concat heads (nh·hd == dim)
let attn_out = linear_t(cdt, &attn, wo); // [1, dim]
h = h.add(&attn_out);
// --- MLP sub-block (pre-norm + SwiGLU + residual) ---
let normed = h.rms_norm(&gamma_t(cdt, ffn_norm), eps).0;
let gate = linear_t(cdt, &normed, w_gate);
let up = linear_t(cdt, &normed, w_up);
let act = gate.silu().mul(&up); // swiglu = silu(gate) ∘ up
let down = linear_t(cdt, &act, w_down);
h = h.add(&down);
}
let h = h.rms_norm(&gamma_t(cdt, final_norm), eps).0;
let logits = linear_t(cdt, &h, lm_head); // [1, vocab]
logits
.to_dtype(DType::F32)
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec()
}
fn argmax(row: &[f32]) -> usize {
row.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
.unwrap()
.0
}
// ===================================================================
// M2b — batched KV-cache decode (G samples of one prompt, in lockstep)
// ===================================================================
/// Batched K/V cache: `G` sequences advancing together. Per layer, a device-resident
/// `[G·num_kv, T, head_dim]` grown one token per step via `cat_seq` (M2c — no host
/// round-trip). Same as M2a's device cache with a G dimension in `bh`.
struct BatchKVCache {
k: Vec<Option<Tensor>>,
v: Vec<Option<Tensor>>,
}
impl BatchKVCache {
fn new(n_layers: usize) -> Self {
Self {
k: (0..n_layers).map(|_| None).collect(),
v: (0..n_layers).map(|_| None).collect(),
}
}
fn append(&mut self, li: usize, k_bh: Tensor, v_bh: Tensor) {
self.k[li] = Some(match self.k[li].take() {
Some(c) => c.cat_seq(&k_bh),
None => k_bh,
});
self.v[li] = Some(match self.v[li].take() {
Some(c) => c.cat_seq(&v_bh),
None => v_bh,
});
}
}
/// Batched KV-cache decode: roll out `n_samples` (G) completions of the SAME
/// `prompt` in lockstep — all G share the prompt, so they advance at one common
/// decode position each step (uniform RoPE via `rope_pos`). Returns G full token
/// sequences (prompt + sampled continuation). The G-way batching amortises the
/// per-step kernel launches across G (the rollout long-pole). Token-identical per
/// row to G independent single-sequence decodes (gated by `tests/decode_batch.rs`).
///
/// `temperature == 0` ⇒ greedy (all G identical); `> 0` ⇒ independent samples
/// (per-row draw from one shared `rng_state`). No finished-mask: all G generate
/// `max_new` tokens; the caller cuts each at `<|endoftext|>` (a perf-only early
/// stop is the M2b+ follow-up). Ragged (different-length prompts) is also deferred.
pub fn generate_cached_batch(
model: &TinyTransformer,
device: Device,
prompt: &[i32],
n_samples: usize,
max_new: usize,
temperature: f32,
rng_state: &mut u64,
) -> Vec<Vec<i32>> {
assert!(!prompt.is_empty(), "prompt must be non-empty");
assert!(n_samples > 0, "n_samples must be > 0");
let cfg = model.config();
let cdt = model.compute_dtype();
let n_layers = cfg.n_layers;
let params: Vec<Tensor> = model.params().iter().map(|p| p.value()).collect();
let embed = &params[0];
let final_norm = &params[1 + n_layers * 11];
let lm_head = &params[1 + n_layers * 11 + 1];
let g = n_samples;
let mut cache = BatchKVCache::new(n_layers);
let mut seqs: Vec<Vec<i32>> = vec![prompt.to_vec(); g];
// Prefill: feed each prompt token (identical across G) at its position.
let mut logits = Vec::new(); // [G, vocab] flattened
for (pos, &tok) in prompt.iter().enumerate() {
let toks = vec![tok; g];
logits = decode_step_batch(&params, cfg, cdt, device, &mut cache, &toks, pos, embed, final_norm, lm_head);
}
let vocab = cfg.vocab;
for _ in 0..max_new {
let mut next = Vec::with_capacity(g);
for row in 0..g {
let lg = &logits[row * vocab..(row + 1) * vocab];
let t = if temperature <= 0.0 {
argmax(lg) as i32
} else {
sample_temperature(lg, temperature, rng_state) as i32
};
next.push(t);
seqs[row].push(t);
}
let pos = seqs[0].len() - 1; // all G are at the same position
logits = decode_step_batch(&params, cfg, cdt, device, &mut cache, &next, pos, embed, final_norm, lm_head);
}
seqs
}
/// One batched decode step: `toks` is one current token per sequence (`[G]`), all at
/// absolute position `pos`. Appends each sequence's K/V and returns logits `[G·vocab]`.
#[allow(clippy::too_many_arguments)]
fn decode_step_batch(
params: &[Tensor],
cfg: &crate::Config,
cdt: DType,
device: Device,
cache: &mut BatchKVCache,
toks: &[i32],
pos: usize,
embed: &Tensor,
final_norm: &Tensor,
lm_head: &Tensor,
) -> Vec<f32> {
let (nh, hd, num_kv) = (cfg.n_heads, cfg.head_dim, cfg.num_kv_heads);
let dim = cfg.dim;
let g = toks.len();
let scale = 1.0 / (hd as f32).sqrt();
let (theta, eps) = (cfg.rope_theta, cfg.eps);
let n_layers = cfg.n_layers;
// Uniform per-row position (all G at the same decode step).
let positions = Tensor::from_slice(&vec![pos as i32; g], &[g]).to_device(device);
let ids = Tensor::from_slice(toks, &[g]).to_device(device);
let mut h = embed.embedding(&ids); // [G, dim] f32
if cdt == DType::BF16 {
h = h.to_dtype(DType::BF16);
}
for li in 0..n_layers {
let base = 1 + li * 11;
let (attn_norm, wq, wk, wv) =
(&params[base], &params[base + 1], &params[base + 2], &params[base + 3]);
let (q_norm, k_norm, wo) = (&params[base + 4], &params[base + 5], &params[base + 6]);
let (ffn_norm, w_gate, w_up, w_down) =
(&params[base + 7], &params[base + 8], &params[base + 9], &params[base + 10]);
let normed = h.rms_norm(&gamma_t(cdt, attn_norm), eps).0; // [G, dim]
// Q: project → per-head QK-norm → RoPE at `pos` for every row.
let q = linear_t(cdt, &normed, wq).reshape(&[g, nh, hd]);
let q = q.reshape(&[g * nh, hd]).rms_norm(&gamma_t(cdt, q_norm), eps).0;
let q = q.reshape(&[g, nh, hd]).rope_pos(&positions, theta);
let q_bh = q.reshape(&[g * nh, 1, hd]); // bh = G·nh
// K/V appended as [G·num_kv,1,hd] (bh-major) into the device cache (M2c).
let k = linear_t(cdt, &normed, wk).reshape(&[g, num_kv, hd]);
let k = k.reshape(&[g * num_kv, hd]).rms_norm(&gamma_t(cdt, k_norm), eps).0;
let k_bh = k
.reshape(&[g, num_kv, hd])
.rope_pos(&positions, theta)
.reshape(&[g * num_kv, 1, hd]);
let v_bh = linear_t(cdt, &normed, wv).reshape(&[g * num_kv, 1, hd]);
cache.append(li, k_bh, v_bh);
// repeat_kv the cached [G·num_kv,T,hd] to [G·nh,T,hd] for the SDPA.
let expand = |c: &Tensor| if num_kv == nh { c.clone() } else { c.repeat_kv(nh, g) };
let k_full = expand(cache.k[li].as_ref().unwrap());
let v_full = expand(cache.v[li].as_ref().unwrap());
let attn = q_bh.decode_attention(&k_full, &v_full, scale); // [G·nh, hd]
let attn = attn.reshape(&[g, dim]); // concat heads per sequence
let attn_out = linear_t(cdt, &attn, wo);
h = h.add(&attn_out);
let normed = h.rms_norm(&gamma_t(cdt, ffn_norm), eps).0;
let gate = linear_t(cdt, &normed, w_gate);
let up = linear_t(cdt, &normed, w_up);
let act = gate.silu().mul(&up);
let down = linear_t(cdt, &act, w_down);
h = h.add(&down);
}
let h = h.rms_norm(&gamma_t(cdt, final_norm), eps).0;
linear_t(cdt, &h, lm_head)
.to_dtype(DType::F32)
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec()
}

5
crates/xtrain-model/src/lib.rs
View File

@@ -25,8 +25,3 @@ pub use config::Config;
mod model;
#[cfg(not(no_cuda))]
pub use model::{TinyTransformer, batched_ids_tensor, ids_tensor, param_to_host};
#[cfg(not(no_cuda))]
pub mod decode;
#[cfg(not(no_cuda))]
pub use decode::{generate_cached, generate_cached_batch, generate_greedy_cached};

61
crates/xtrain-model/tests/dropout.rs
View File

@@ -66,18 +66,10 @@ fn tiny_cfg(dropout: f32) -> Config {
fn batch_data(cfg: &Config, device: Device) -> (xtrain_tensor::Tensor, xtrain_tensor::Tensor) {
let (batch, seq) = (3usize, 6usize);
let seqs: Vec<Vec<i32>> = (0..batch)
.map(|b| {
(0..seq)
.map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32)
.collect()
})
.map(|b| (0..seq).map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32).collect())
.collect();
let tgts: Vec<Vec<i32>> = (0..batch)
.map(|b| {
(0..seq)
.map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
.collect()
})
.map(|b| (0..seq).map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32).collect())
.collect();
(
batched_ids_tensor(&seqs, device),
@@ -102,11 +94,7 @@ fn fwd_bwd(
let loss = m.loss_batched(ids, tgt, batch);
let loss_val = host(&loss.value())[0];
loss.backward();
let grads: Vec<Vec<f32>> = m
.params()
.iter()
.map(|p| host(&p.grad().unwrap()))
.collect();
let grads: Vec<Vec<f32>> = m.params().iter().map(|p| host(&p.grad().unwrap())).collect();
(logits, loss_val, grads)
}
@@ -198,9 +186,7 @@ fn recompute_with_dropout(dtype: DType, grad_tol: f32) {
// Both models: same init, train mode, p=0.2. step_seed starts at 0 and bumps
// to 1 on the first training forward in BOTH, so they draw the same masks.
let off = build(cfg, device)
.with_compute_dtype(dtype)
.with_training(true);
let off = build(cfg, device).with_compute_dtype(dtype).with_training(true);
let on = build(cfg, device)
.with_compute_dtype(dtype)
.with_recompute(true)
@@ -208,19 +194,11 @@ fn recompute_with_dropout(dtype: DType, grad_tol: f32) {
let off_loss = off.loss_batched(&ids, &tgt, batch);
off_loss.backward();
let off_grads: Vec<Vec<f32>> = off
.params()
.iter()
.map(|p| host(&p.grad().unwrap()))
.collect();
let off_grads: Vec<Vec<f32>> = off.params().iter().map(|p| host(&p.grad().unwrap())).collect();
let on_loss = on.loss_batched(&ids, &tgt, batch);
on_loss.backward();
let on_grads: Vec<Vec<f32>> = on
.params()
.iter()
.map(|p| host(&p.grad().unwrap()))
.collect();
let on_grads: Vec<Vec<f32>> = on.params().iter().map(|p| host(&p.grad().unwrap())).collect();
let mut max_rel = 0.0f32;
for (a, b) in off_grads.iter().flatten().zip(on_grads.iter().flatten()) {
@@ -262,18 +240,10 @@ fn flash_plus_dropout_grad_check_fp32() {
cfg.dropout = 0.2;
let seq = 40usize;
let seqs: Vec<Vec<i32>> = (0..batch)
.map(|b| {
(0..seq)
.map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32)
.collect()
})
.map(|b| (0..seq).map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32).collect())
.collect();
let tgts: Vec<Vec<i32>> = (0..batch)
.map(|b| {
(0..seq)
.map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
.collect()
})
.map(|b| (0..seq).map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32).collect())
.collect();
let ids = batched_ids_tensor(&seqs, device);
let tgt = batched_ids_tensor(&tgts, device);
@@ -307,16 +277,7 @@ fn flash_plus_dropout_grad_check_fp32() {
);
// Same tolerances as the flash-vs-composed gate (flash.rs run_fp32): flash
// differs from composed only by reduction order; dropout masks are identical.
assert!(
logit_rel < 1e-3,
"[F32] flash+dropout logits diverged: {logit_rel:.2e}"
);
assert!(
loss_rel < 1e-3,
"[F32] flash+dropout loss diverged: {loss_rel:.2e}"
);
assert!(
grad_rel < 2e-2,
"[F32] flash+dropout grads diverged: {grad_rel:.3e}"
);
assert!(logit_rel < 1e-3, "[F32] flash+dropout logits diverged: {logit_rel:.2e}");
assert!(loss_rel < 1e-3, "[F32] flash+dropout loss diverged: {loss_rel:.2e}");
assert!(grad_rel < 2e-2, "[F32] flash+dropout grads diverged: {grad_rel:.3e}");
}

97
crates/xtrain-model/tests/ragged_batch.rs
View File

@@ -1,97 +0,0 @@
// M2d gate: does forward_batched on RIGHT-PADDED ragged sequences reproduce the
// per-sequence single-seq forward on the real (non-pad) rows? The batched GRPO
// training-side forwards depend on this "right-pad is free under causal attention"
// property — a real completion row is at an earlier position than the trailing pad,
// and causal masking forbids attending forward, so its logits should be unchanged.
//
// Tested in fp32 (exact) over both SDPA cores (composed + fused flash), since the
// bench uses flash and a kernel could in principle leak the pad keys into the online
// softmax.
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, ids_tensor};
use xtrain_tensor::{DType, Device, Tensor};
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()
}
fn build(cfg: Config, device: Device, dtype: DType, flash: bool) -> TinyTransformer {
let mut seed = 1u64;
let m = TinyTransformer::new(cfg, device, |shape| {
seed = seed.wrapping_add(1);
let n: usize = shape.iter().product();
if shape.len() == 1 {
fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
} else {
fill(n, seed, 0.08)
}
});
m.with_compute_dtype(dtype).with_flash(flash)
}
fn host(t: &Tensor) -> Vec<f32> {
t.to_dtype(DType::F32).to_device(Device::Cpu).as_slice::<f32>().to_vec()
}
#[test]
fn forward_batched_ragged_matches_looped() {
if device::device_count().unwrap_or(0) == 0 {
eprintln!("no CUDA device; skipping");
return;
}
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let mut cfg = Config::tiny();
cfg.vocab = 32;
cfg.n_layers = 2;
let vocab = cfg.vocab;
// Ragged lengths incl. one crossing the flash tile (>32) and short ones.
let lens = [6usize, 40, 9, 4];
let lmax = *lens.iter().max().unwrap();
let n = lens.len();
let seqs: Vec<Vec<i32>> = lens
.iter()
.enumerate()
.map(|(b, &l)| (0..l).map(|i| ((b * 7 + i * 3 + 1) % vocab) as i32).collect())
.collect();
for (dtype, tol) in [(DType::F32, 2e-3f32), (DType::BF16, 3e-1f32)] {
for flash in [false, true] {
let m = build(cfg, device, dtype, flash);
// Looped: each sequence on its own (the ground truth).
let looped: Vec<Vec<f32>> = seqs.iter().map(|s| host(&m.forward(&ids_tensor(s, device)).value())).collect();
// Batched: right-pad each to lmax (pad id 0), one forward_batched(batch = n).
let mut flat = vec![0i32; n * lmax];
for (i, s) in seqs.iter().enumerate() {
flat[i * lmax..i * lmax + s.len()].copy_from_slice(s);
}
let ids = Tensor::from_slice(&flat, &[n * lmax]).to_device(device);
let batched = host(&m.forward_batched(&ids, n).value()); // [n*lmax, vocab]
let mut dmax = 0f32;
for (i, s) in seqs.iter().enumerate() {
for r in 0..s.len() {
for c in 0..vocab {
let a = looped[i][r * vocab + c];
let b = batched[(i * lmax + r) * vocab + c];
dmax = dmax.max((a - b).abs());
}
}
}
println!("dtype={dtype:?} flash={flash}: ragged right-pad vs looped, max|Δlogit| (real rows) = {dmax:.3e}");
assert!(dmax < tol, "dtype={dtype:?} flash={flash}: right-pad NOT free under causal — max|Δ| = {dmax}");
}
}
println!("forward_batched_ragged_matches_looped OK: right-pad is free under causal (fp32+bf16, composed + flash)");
}

196
crates/xtrain-tensor/src/tensor.rs
View File

@@ -790,107 +790,6 @@ impl Tensor {
out
}
/// RoPE at an absolute position offset (KV-cache decode, forward only).
/// `self`:[tokens,heads,head_dim]; row `r`'s position is `pos0 + r` (no
/// modulo). For a single new decode token pass `tokens == 1` → the one row is
/// rotated at absolute position `pos0`. Mirrors [`rope`](Self::rope)'s dtype
/// handling (bf16 → f32 → bf16); no backward (inference path).
#[cfg(not(no_cuda))]
pub fn rope_at(&self, theta: f32, pos0: usize) -> Self {
assert_eq!(self.ndim(), 3, "rope_at requires [tokens,heads,head_dim]");
let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
assert_eq!(head_dim % 2, 0, "head_dim must be even");
if self.dtype == DType::BF16 {
return self
.to_dtype(DType::F32)
.rope_at(theta, pos0)
.to_dtype(DType::BF16);
}
let out = Tensor::zeros(&self.shape, DType::F32, self.device());
unsafe {
xtrain_cuda::ffi::launch_rope_at_f32(
self.data_ptr() as *const f32,
out.data_ptr() as *mut f32,
tokens as i32,
heads as i32,
head_dim as i32,
theta,
pos0 as i32,
std::ptr::null_mut(),
);
}
out
}
/// RoPE with a PER-ROW absolute position (batched KV-cache decode, M2b).
/// `self`:[tokens,heads,head_dim]; row `t`'s position is `positions[t]` (an
/// I32 `[tokens]` tensor). For G-way batched decode all G rows share one decode
/// position; for ragged batches each row carries its own. Mirrors `rope_at`'s
/// dtype handling; forward only.
#[cfg(not(no_cuda))]
pub fn rope_pos(&self, positions: &Tensor, theta: f32) -> Self {
assert_eq!(self.ndim(), 3, "rope_pos requires [tokens,heads,head_dim]");
let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
assert_eq!(head_dim % 2, 0, "head_dim must be even");
assert_eq!(positions.dtype, DType::I32, "positions must be I32");
assert_eq!(positions.numel(), tokens, "one position per token");
if self.dtype == DType::BF16 {
return self
.to_dtype(DType::F32)
.rope_pos(positions, theta)
.to_dtype(DType::BF16);
}
let out = Tensor::zeros(&self.shape, DType::F32, self.device());
unsafe {
xtrain_cuda::ffi::launch_rope_pos_f32(
self.data_ptr() as *const f32,
positions.data_ptr() as *const i32,
out.data_ptr() as *mut f32,
tokens as i32,
heads as i32,
head_dim as i32,
theta,
std::ptr::null_mut(),
);
}
out
}
/// Concatenate along the sequence (middle) dim: `self`:[bh,ta,hd] ++
/// `other`:[bh,tb,hd] → `[bh,ta+tb,hd]`. The device-side KV-cache append (M2c):
/// the cache stays on the GPU and grows by one token per decode step, removing
/// the M2a/M2b host round-trip. Mirrors the bf16 cast handling of the other
/// structural kernels.
#[cfg(not(no_cuda))]
pub fn cat_seq(&self, other: &Tensor) -> Self {
assert_eq!(self.ndim(), 3, "cat_seq requires [bh,t,hd]");
assert_eq!(other.ndim(), 3, "cat_seq requires [bh,t,hd]");
assert_eq!(self.dtype, other.dtype, "cat_seq dtype mismatch");
let (bh, ta, hd) = (self.shape[0], self.shape[1], self.shape[2]);
let (bh2, tb, hd2) = (other.shape[0], other.shape[1], other.shape[2]);
assert_eq!(bh, bh2, "cat_seq bh mismatch");
assert_eq!(hd, hd2, "cat_seq head_dim mismatch");
if self.dtype == DType::BF16 {
return self
.to_dtype(DType::F32)
.cat_seq(&other.to_dtype(DType::F32))
.to_dtype(DType::BF16);
}
let out = Tensor::zeros(&[bh, ta + tb, hd], DType::F32, self.device());
unsafe {
xtrain_cuda::ffi::launch_cat_seq_f32(
self.data_ptr() as *const f32,
other.data_ptr() as *const f32,
out.data_ptr() as *mut f32,
bh as i32,
(ta * hd) as i32,
(tb * hd) as i32,
std::ptr::null_mut(),
);
}
out
}
/// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an
/// orthogonal map, so it needs no cached forward values, only `theta`/`period`.
#[cfg(not(no_cuda))]
@@ -1010,31 +909,6 @@ impl Tensor {
dx
}
/// Per-row scale: `out[r,c] = self[r,c] * s[r]`. `self`:[rows,cols] F32,
/// `s`:[rows] F32. Used by the GRPO (M4) policy-gradient backward, where each
/// completion token's row of `(probs − onehot)` is scaled by its own per-token
/// coefficient (the per-token clipped-PG + KL gradient). Forward-only.
#[cfg(not(no_cuda))]
pub fn scale_rows(&self, s: &Tensor) -> Self {
assert_eq!(self.ndim(), 2, "scale_rows requires a 2D tensor");
assert_eq!(self.dtype, DType::F32, "scale_rows is F32");
assert_eq!(s.dtype, DType::F32, "scale vector is F32");
let (rows, cols) = (self.shape[0], self.shape[1]);
assert_eq!(s.numel(), rows, "scale vector must have one entry per row");
let out = Tensor::zeros(&self.shape, DType::F32, self.device());
unsafe {
xtrain_cuda::ffi::launch_scale_rows_f32(
self.data_ptr() as *const f32,
s.data_ptr() as *const f32,
out.data_ptr() as *mut f32,
rows as i32,
cols as i32,
std::ptr::null_mut(),
);
}
out
}
// --- Structural / model ops (the T5 kernels) ---
/// Reshape to `new_shape` (must keep `numel`). Pure metadata change on a
@@ -1202,76 +1076,6 @@ impl Tensor {
(out, probs)
}
/// Decode-time (incremental) attention: a SINGLE query position against a
/// cached K/V of length `t` (KV-cache decode, forward only). `self` = Q
/// `[bh,1,head_dim]`; `k`,`v` = `[bh,t,head_dim]`, already repeat_kv-expanded
/// to `bh` heads. Returns out `[bh,head_dim]` (= `[bh,1,head_dim]` flattened).
///
/// No causal mask is needed — the one query sits at the end, so every cached
/// key (positions `0..t`) is visible. This is exactly the LAST query row of the
/// full causal [`attention`](Self::attention), so KV-cache greedy decode is
/// token-identical to full recompute. Softmax is computed in f32 (matching the
/// causal path) with `scale` folded in before the exponentials.
#[cfg(not(no_cuda))]
pub fn decode_attention(&self, k: &Tensor, v: &Tensor, scale: f32) -> Self {
assert_eq!(self.ndim(), 3, "decode_attention Q must be [bh,1,head_dim]");
assert_eq!(self.shape[1], 1, "decode_attention Q seq must be 1");
assert_eq!(k.ndim(), 3, "decode_attention K must be [bh,t,head_dim]");
assert_eq!(k.shape(), v.shape(), "K/V shape mismatch");
assert_eq!(self.dtype, k.dtype, "Q/K dtype mismatch");
assert_eq!(self.dtype, v.dtype, "Q/V dtype mismatch");
let (bh, hd) = (self.shape[0], self.shape[2]);
assert_eq!(k.shape[0], bh, "Q/K batch-head mismatch");
assert_eq!(k.shape[2], hd, "Q/K head_dim mismatch");
let t = k.shape[1]; // cached length
let dt = self.dtype;
let dev = self.device();
// scores[bh,1,t] = Q[bh,1,hd] · Kᵀ[bh,hd,t] (per-head batched GEMM).
// [bh,1,t] is stored identically to [bh,t]; allocate 2D so the rowwise
// softmax can run without a reshape.
let scores = Tensor::zeros(&[bh, t], dt, dev);
strided_batched_gemm(
dt,
false,
true,
1,
t,
hd,
self.data_ptr(),
hd,
k.data_ptr(),
t * hd,
scores.data_ptr(),
t,
bh,
);
// probs = softmax(scale · scores) over the t keys (f32, like the causal path).
let probs = scores
.to_dtype(DType::F32)
.scale(scale)
.softmax()
.to_dtype(dt);
// out[bh,1,hd] = probs[bh,1,t] · V[bh,t,hd].
let out = Tensor::zeros(&[bh, hd], dt, dev);
strided_batched_gemm(
dt,
false,
false,
1,
hd,
t,
probs.data_ptr(),
t,
v.data_ptr(),
t * hd,
out.data_ptr(),
hd,
bh,
);
out
}
/// Backward of [`attention`](Self::attention). Inputs: forward `q`,`k`,`v`,
/// the cached `probs`, the upstream `dout` (all batched `[bh,seq,*]`), and the
/// same `scale`. Returns `(dq, dk, dv)`.

167
crates/xtrain-tensor/tests/integration.rs
View File

@@ -56,170 +56,3 @@ fn elementwise_scale_kernel() {
r.len()
);
}
/// (c) `rope_at` (KV-cache decode RoPE at an absolute position) is bit-identical
/// to the full-sequence `rope`'s corresponding row. This is the invariant the
/// decode KV-cache relies on: a single new token RoPE'd at position `t` must equal
/// what the full-sequence forward would have produced at row `t` (so cached
/// post-RoPE K matches the full-recompute path → token-identical decode).
#[test]
fn rope_at_matches_full_rope_row() {
assert!(
device::device_count().expect("device count") > 0,
"no CUDA device"
);
device::set_device(0).unwrap();
let (n, heads, hd) = (7usize, 3usize, 8usize);
let theta = 10000.0f32;
// Deterministic pseudo-random fill in [-1, 1).
let host: Vec<f32> = (0..n * heads * hd)
.map(|i| ((i * 37 % 101) as f32 / 50.0) - 1.0)
.collect();
// Full-sequence rope (period = n → row r gets position r).
let full = Tensor::from_slice(&host, &[n, heads, hd]).to_device(Device::Cuda(0));
let roped_full = full
.rope(theta, n)
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec();
let row_len = heads * hd;
for t in 0..n {
let row = &host[t * row_len..(t + 1) * row_len];
let roped_row = Tensor::from_slice(row, &[1, heads, hd])
.to_device(Device::Cuda(0))
.rope_at(theta, t)
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec();
let expect = &roped_full[t * row_len..(t + 1) * row_len];
assert_eq!(
roped_row.as_slice(),
expect,
"rope_at(pos0={t}) != full rope row {t}"
);
}
println!("rope_at OK: bit-identical to full rope across {n} positions");
}
/// (d) `decode_attention` (single query vs cached K/V, no mask) equals the LAST
/// query row of the full causal `attention`. This is the core decode-engine
/// invariant: the incremental path must reproduce what the full-recompute forward
/// computes for the final position, so KV-cache greedy decode is token-identical.
/// Tolerance is fp rounding (different softmax kernel + reduction order), not bits.
#[test]
fn decode_attention_matches_full_attention_last_row() {
assert!(
device::device_count().expect("device count") > 0,
"no CUDA device"
);
device::set_device(0).unwrap();
let (bh, t, hd) = (6usize, 5usize, 8usize);
let scale = 1.0 / (hd as f32).sqrt();
let n = bh * t * hd;
let qh: Vec<f32> = (0..n).map(|i| ((i * 31 % 97) as f32 / 48.0) - 1.0).collect();
let kh: Vec<f32> = (0..n).map(|i| ((i * 53 % 89) as f32 / 44.0) - 1.0).collect();
let vh: Vec<f32> = (0..n).map(|i| ((i * 17 % 83) as f32 / 41.0) - 1.0).collect();
let q = Tensor::from_slice(&qh, &[bh, t, hd]).to_device(Device::Cuda(0));
let k = Tensor::from_slice(&kh, &[bh, t, hd]).to_device(Device::Cuda(0));
let v = Tensor::from_slice(&vh, &[bh, t, hd]).to_device(Device::Cuda(0));
// Reference: full causal attention, take each head's last query row.
let (full, _) = q.attention(&k, &v, scale);
let full_h = full.to_device(Device::Cpu).as_slice::<f32>().to_vec();
// Decode: build Q_last [bh,1,hd] from each head's last row, attend to all K/V.
let mut ql = vec![0f32; bh * hd];
for b in 0..bh {
let src = (b * t + (t - 1)) * hd;
ql[b * hd..(b + 1) * hd].copy_from_slice(&qh[src..src + hd]);
}
let q_last = Tensor::from_slice(&ql, &[bh, 1, hd]).to_device(Device::Cuda(0));
let dec = q_last
.decode_attention(&k, &v, scale)
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec();
assert_eq!(dec.len(), bh * hd, "decode out shape");
let mut max_abs = 0f32;
for b in 0..bh {
for d in 0..hd {
let got = dec[b * hd + d];
let exp = full_h[(b * t + (t - 1)) * hd + d];
max_abs = max_abs.max((got - exp).abs());
}
}
assert!(
max_abs < 1e-4,
"decode_attention vs full last-row max abs diff {max_abs} exceeds 1e-4"
);
println!("decode_attention OK: matches full causal last row (bh={bh}, t={t}, max|Δ|={max_abs:.2e})");
}
/// (e) `rope_pos` (per-row positions, M2b batched decode): with positions
/// [0,1,…,n-1] it is bit-identical to the full-sequence `rope` (period=n); with a
/// uniform position P every row matches `rope_at(·, P)` of that single row. This is
/// the primitive the batched decode uses (G rows sharing one decode position).
#[test]
fn rope_pos_matches_rope_and_rope_at() {
assert!(device::device_count().expect("device count") > 0, "no CUDA device");
device::set_device(0).unwrap();
let (n, heads, hd) = (7usize, 3usize, 8usize);
let theta = 10000.0f32;
let host: Vec<f32> = (0..n * heads * hd).map(|i| ((i * 37 % 101) as f32 / 50.0) - 1.0).collect();
let x = Tensor::from_slice(&host, &[n, heads, hd]).to_device(Device::Cuda(0));
// positions [0,1,…,n-1] ⇒ identical to the full-sequence rope.
let seq_pos: Vec<i32> = (0..n as i32).collect();
let pos_t = Tensor::from_slice(&seq_pos, &[n]).to_device(Device::Cuda(0));
let got = x.rope_pos(&pos_t, theta).to_device(Device::Cpu).as_slice::<f32>().to_vec();
let want = x.rope(theta, n).to_device(Device::Cpu).as_slice::<f32>().to_vec();
assert_eq!(got, want, "rope_pos [0..n] != full rope");
// uniform position P ⇒ each row matches rope_at(single row, P).
let p = 5i32;
let uni = Tensor::from_slice(&vec![p; n], &[n]).to_device(Device::Cuda(0));
let got_u = x.rope_pos(&uni, theta).to_device(Device::Cpu).as_slice::<f32>().to_vec();
let row_len = heads * hd;
for t in 0..n {
let row = &host[t * row_len..(t + 1) * row_len];
let want_row = Tensor::from_slice(row, &[1, heads, hd])
.to_device(Device::Cuda(0))
.rope_at(theta, p as usize)
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec();
assert_eq!(&got_u[t * row_len..(t + 1) * row_len], want_row.as_slice(), "uniform pos row {t}");
}
println!("rope_pos OK: == full rope for [0..n] and == rope_at(P) per row for uniform P");
}
/// (f) `cat_seq` (device-side KV-cache append, M2c): concatenating [bh,ta,hd] ++
/// [bh,tb,hd] along the seq dim equals the host-side interleaved concat (per bh row,
/// a's block then b's block). This is the device append that removes the M2a/M2b
/// host round-trip.
#[test]
fn cat_seq_matches_host_concat() {
assert!(device::device_count().expect("device count") > 0, "no CUDA device");
device::set_device(0).unwrap();
let (bh, ta, tb, hd) = (4usize, 3usize, 2usize, 5usize);
let ah: Vec<f32> = (0..bh * ta * hd).map(|i| i as f32 * 0.1).collect();
let bhost: Vec<f32> = (0..bh * tb * hd).map(|i| -(i as f32) - 1.0).collect();
let a = Tensor::from_slice(&ah, &[bh, ta, hd]).to_device(Device::Cuda(0));
let b = Tensor::from_slice(&bhost, &[bh, tb, hd]).to_device(Device::Cuda(0));
let got = a.cat_seq(&b).to_device(Device::Cpu).as_slice::<f32>().to_vec();
// Host reference: per bh row, a's ta*hd then b's tb*hd.
let mut want = vec![0f32; bh * (ta + tb) * hd];
for r in 0..bh {
let (oa, ob, oo) = (r * ta * hd, r * tb * hd, r * (ta + tb) * hd);
want[oo..oo + ta * hd].copy_from_slice(&ah[oa..oa + ta * hd]);
want[oo + ta * hd..oo + (ta + tb) * hd].copy_from_slice(&bhost[ob..ob + tb * hd]);
}
assert_eq!(got, want, "cat_seq != host interleaved concat");
println!("cat_seq OK: [bh={bh},{ta}+{tb},{hd}] == host concat");
}

268
crates/xtrain-train/src/bin/bench_grpo_batch.rs
View File

@@ -1,268 +0,0 @@
//! Micro-benchmark + closeness gate for the M2d batched GRPO training-side forwards.
//!
//! After M2b/M2c the GRPO *step* is no longer rollout-bound — it is the `N = B·G`
//! per-sample full-sequence forwards (the `per_token_logp` captures + the inner
//! clipped-PG forward/backwards). This bin isolates exactly that, weight-independently
//! (step wall-clock depends on shapes + launch counts, not on what the weights are), by
//! synthesising `N` realistic ragged samples and A/B-timing the looped vs batched path
//! for BOTH phases — plus asserting they agree numerically (the looped-vs-batched
//! closeness gate; per-row bit-equivalence of the loss op is pinned by the autograd
//! test `clipped_pg_loss_batched_matches_looped`).
//!
//! bench_grpo_batch <tokenizer.json> --init-ckpt <base.ckpt> <arch flags> \
//! --n 48 --plen 12 --clen 24 --micro 16 --reps 3
#[cfg(no_cuda)]
fn main() {
eprintln!("bench_grpo_batch: built without CUDA (no_cuda); run on a GPU host.");
}
#[cfg(not(no_cuda))]
use xtrain_cuda::device;
#[cfg(not(no_cuda))]
use xtrain_model::{Config, TinyTransformer};
#[cfg(not(no_cuda))]
use xtrain_tensor::{DType, Device, Tensor};
#[cfg(not(no_cuda))]
use xtrain_train::grpo_batch::{PgSample, inner_pg_step_batched, inner_pg_step_looped, per_token_logp, per_token_logp_batched};
#[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()
}
#[cfg(not(no_cuda))]
fn load_model(cfg: Config, device: Device, ckpt: &str) -> TinyTransformer {
let mut seed = 1u64;
let m = TinyTransformer::new(cfg, device, |shape| {
seed = seed.wrapping_add(1);
let n: usize = shape.iter().product();
if shape.len() == 1 {
fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
} else {
fill(n, seed, 0.04)
}
})
.with_compute_dtype(DType::BF16)
.with_flash(true);
xtrain_train::checkpoint::load_into(std::path::Path::new(ckpt), &m.params()).expect("load ckpt");
m.eval();
m
}
#[cfg(not(no_cuda))]
fn elapsed_ms<F: FnMut()>(reps: usize, mut f: F) -> f32 {
let start = std::time::Instant::now();
for _ in 0..reps {
f();
}
start.elapsed().as_secs_f32() * 1e3 / reps as f32
}
/// Per-position argmax of the model over each ragged `input` (one `forward_batched`
/// per `micro`-chunk). Used to teacher-force WELL-CONDITIONED targets (the top-1 token,
/// high prob) so the closeness gate's logp isn't the ~−20 of a random token — where
/// `−log p` amplifies bf16 noise. This matches real GRPO (targets are model samples).
#[cfg(not(no_cuda))]
fn model_argmax(model: &TinyTransformer, device: Device, inputs: &[Vec<i32>], vocab: usize, micro: usize) -> Vec<Vec<i32>> {
let mut out = Vec::with_capacity(inputs.len());
for chunk in inputs.chunks(micro.max(1)) {
let m = chunk.len();
let lmax = chunk.iter().map(|s| s.len()).max().unwrap();
let mut flat = vec![0i32; m * lmax];
for (i, s) in chunk.iter().enumerate() {
flat[i * lmax..i * lmax + s.len()].copy_from_slice(s);
}
let ids = Tensor::from_slice(&flat, &[m * lmax]).to_device(device);
let logits = model.forward_batched(&ids, m).value().to_dtype(DType::F32).to_device(Device::Cpu);
let v = logits.as_slice::<f32>();
for (i, s) in chunk.iter().enumerate() {
let mut row = Vec::with_capacity(s.len());
for r in 0..s.len() {
let base = (i * lmax + r) * vocab;
let mut best = 0usize;
for c in 1..vocab {
if v[base + c] > v[base + best] {
best = c;
}
}
row.push(best as i32);
}
out.push(row);
}
}
out
}
#[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 tok_path = positionals.first().expect("usage: bench_grpo_batch <tokenizer.json> [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 n: usize = flag(&args, "--n", 48); // B·G samples per step
let plen: usize = flag(&args, "--plen", 12); // prompt tokens
let clen: usize = flag(&args, "--clen", 24); // max completion tokens
let micro: usize = flag(&args, "--micro", 16);
let reps: usize = flag(&args, "--reps", 3);
let (eps, beta) = (flag(&args, "--eps", 0.2f32), flag(&args, "--beta", 0.0f32));
let init_ckpt = flag_value(&args, "--init-ckpt").expect("--init-ckpt <base.ckpt> 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.as_str()));
let vocab = tok.vocab_size();
let cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
let policy = load_model(cfg, device, &init_ckpt);
let params = policy.params();
// --- Synthesise N ragged samples (frame-shaped: prompt masked, ragged completion).
// Token IDs are random-but-valid; only the SHAPES drive the forward cost.
let mut rng = 0xC0FFEEu64;
let mut next = || {
rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
(rng >> 33) as usize
};
let mut io: Vec<(Vec<i32>, Vec<i32>)> = Vec::with_capacity(n);
let mut advs: Vec<f32> = Vec::with_capacity(n);
for _ in 0..n {
let pl = plen.saturating_sub(2) + next() % 5; // jitter prompt length a little
let cl = 4 + next() % clen.max(1); // completion 4..=clen
let total = pl + cl;
let toks: Vec<i32> = (0..total).map(|_| (next() % vocab) as i32).collect();
let mut labels = vec![-100i32; pl]; // prompt masked
labels.extend_from_slice(&toks[pl..]);
let l = toks.len();
io.push((toks[..l - 1].to_vec(), labels[1..l].to_vec())); // target masked at [..pl-1]
advs.push(if next() % 2 == 0 { 0.7 } else { -0.7 });
}
let toklens: Vec<usize> = io.iter().map(|(i, _)| i.len()).collect();
let (lmin, lmax) = (*toklens.iter().min().unwrap(), *toklens.iter().max().unwrap());
println!("samples N={n}, seq len {lmin}..{lmax} (ragged), micro={micro}, β={beta}\n");
// Replace random completion targets with the model's own argmax (teacher forcing):
// well-conditioned logp (top-1, not the ~−20 of a random token where bf16 noise
// blows up via −log p). The completion target positions are where the skeleton is
// ≥0; prompt positions stay masked (−100).
let inputs: Vec<Vec<i32>> = io.iter().map(|(i, _)| i.clone()).collect();
let preds = model_argmax(&policy, device, &inputs, vocab, micro);
for (s, (_, target)) in io.iter_mut().enumerate() {
for j in 0..target.len() {
if target[j] >= 0 {
target[j] = preds[s][j];
}
}
}
// ---------------- Phase 1: capture (per_token_logp) ----------------
let logp_loop: Vec<Vec<f32>> = io.iter().map(|(i, t)| per_token_logp(&policy, device, i, t)).collect();
let logp_batch = per_token_logp_batched(&policy, device, &io, micro);
let cap_dmax = logp_loop
.iter()
.zip(&logp_batch)
.flat_map(|(a, b)| a.iter().zip(b).map(|(x, y)| (x - y).abs()))
.fold(0.0f32, f32::max);
let t_cap_loop = elapsed_ms(reps, || {
let _: Vec<Vec<f32>> = io.iter().map(|(i, t)| per_token_logp(&policy, device, i, t)).collect();
});
let t_cap_batch = elapsed_ms(reps, || {
let _ = per_token_logp_batched(&policy, device, &io, micro);
});
// Build PgSamples from the (matching) capture; ref = old − 0.3 to exercise KL.
let batch: Vec<PgSample> = io
.iter()
.zip(&advs)
.zip(&logp_batch)
.map(|(((input, target), &adv), lp)| PgSample {
input: input.clone(),
target: target.clone(),
adv,
logp_old: lp.clone(),
logp_ref: lp.iter().map(|v| v - 0.3).collect(),
})
.collect();
// ---------------- Phase 2: inner clipped-PG (forward + backward) ----------------
// Representative grad snapshots: layer-0 wq (params[2]) + final_norm.
let wq0 = &params[2];
let fnorm = &params[1 + n_layers * 11];
let snap = |v: &xtrain_autodiff::Var| -> Vec<f32> {
v.grad().map(|g| g.to_device(Device::Cpu).as_slice::<f32>().to_vec()).unwrap_or_default()
};
let zero = |ps: &[xtrain_autodiff::Var]| ps.iter().for_each(|p| p.zero_grad());
zero(&params);
inner_pg_step_looped(&policy, device, &batch, eps, beta);
let (gq_loop, gn_loop) = (snap(wq0), snap(fnorm));
zero(&params);
inner_pg_step_batched(&policy, device, &batch, eps, beta, micro);
let (gq_batch, gn_batch) = (snap(wq0), snap(fnorm));
zero(&params);
let reldiff = |a: &[f32], b: &[f32]| -> f32 {
let num = a.iter().zip(b).map(|(x, y)| (x - y).abs()).fold(0.0f32, f32::max);
let den = a.iter().map(|x| x.abs()).fold(0.0f32, f32::max).max(1e-12);
num / den
};
let gq_rel = reldiff(&gq_loop, &gq_batch);
let gn_rel = reldiff(&gn_loop, &gn_batch);
// Time only forward+backward — the lever. opt.step + grad-clip are identical in
// both paths (one call over `params` after the per-sample loop), so they would
// only add a constant; excluding them also dodges the unrelated 1B-Adam-state
// memory wall (the M4 finding) that this diagnostic doesn't need to reproduce.
let t_inner_loop = elapsed_ms(reps, || {
inner_pg_step_looped(&policy, device, &batch, eps, beta);
zero(&params);
});
let t_inner_batch = elapsed_ms(reps, || {
inner_pg_step_batched(&policy, device, &batch, eps, beta, micro);
zero(&params);
});
// ---------------- Report ----------------
let spd = |a: f32, b: f32| if b > 0.0 { a / b } else { 0.0 };
println!("=== closeness gate (looped vs batched) ===");
println!(" capture per_token_logp : max|Δ| = {cap_dmax:.3e}");
println!(" inner grad wq[0] : rel|Δ| = {gq_rel:.3e}");
println!(" inner grad final_norm : rel|Δ| = {gn_rel:.3e}");
println!("\n=== timing (mean of {reps} reps, ms/phase) ===");
println!(" capture : looped {t_cap_loop:8.1} batched {t_cap_batch:8.1} ({:.2}× )", spd(t_cap_loop, t_cap_batch));
println!(" inner : looped {t_inner_loop:8.1} batched {t_inner_batch:8.1} ({:.2}× )", spd(t_inner_loop, t_inner_batch));
let (step_loop, step_batch) = (t_cap_loop + t_inner_loop, t_cap_batch + t_inner_batch);
println!(" STEP : looped {step_loop:8.1} batched {step_batch:8.1} ({:.2}× )", spd(step_loop, step_batch));
// The RIGOROUS correctness gates live in the test suite (exact, not bf16-noisy):
// - xtrain-model forward_batched_ragged_matches_looped (forward+pad == looped)
// - xtrain-autodiff clipped_pg_loss_batched_matches_looped (op == looped, f32)
// This is a smoke check at the 1B/bf16 scale: single-seq vs batched GEMM differ in
// batch-reduction order, so a loose band, with well-conditioned (argmax) targets.
assert!(cap_dmax < 0.2, "capture closeness smoke FAILED: max|Δlogp| = {cap_dmax}");
assert!(gq_rel < 0.2 && gn_rel < 0.2, "inner grad closeness smoke FAILED: wq {gq_rel}, fn {gn_rel}");
println!("\nSMOKE PASS (bf16 band): batched ≈ looped; rigorous gates are the two tests above.");
}

209
crates/xtrain-train/src/bin/eval_arith.rs
View File

@@ -1,209 +0,0 @@
//! Verifiable-task eval (post-training, M1+). Load a checkpoint, greedily generate an
//! answer for each held-out arithmetic prompt, parse the `\boxed{}` answer, and report
//! the exact-match pass-rate against the gold file. Two signals are printed:
//! **format** (fraction that emitted any boxed integer) and **correctness** (fraction
//! whose boxed answer matches gold). This is the M1 format-baseline metric and the
//! reusable verifiable-eval harness for M3 (DPO) / M4 (GRPO).
//!
//! eval_arith <ckpt> <tokenizer.json> --heads 52 --head-dim 32 --kv-heads 13 \
//! --layers 22 --ffn 6656 \
//! --prompts-file <dir>/arith_eval_prompts.txt \
//! --gold-file <dir>/arith_eval_gold.txt --max-tokens 48 --show 8
#[cfg(no_cuda)]
fn main() {
eprintln!("eval_arith: built without CUDA (no_cuda); run on a GPU host (dash5).");
}
#[cfg(not(no_cuda))]
use std::path::PathBuf;
#[cfg(not(no_cuda))]
use xtrain_cuda::device;
#[cfg(not(no_cuda))]
use xtrain_model::{Config, TinyTransformer};
#[cfg(not(no_cuda))]
use xtrain_tensor::Device;
#[cfg(not(no_cuda))]
use xtrain_train::sample::generate;
#[cfg(not(no_cuda))]
use xtrain_train::task::{check_answer, parse_boxed_answer};
// Same deterministic LCG init scheme as bin/train.rs / bin/greedy_sample.rs (the
// values are overwritten by the loaded checkpoint; init just shapes the tensors).
#[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()
}
#[cfg(not(no_cuda))]
fn decode_escapes(s: &str) -> String {
s.replace("\\n", "\n").replace("\\t", "\t")
}
/// The model keeps generating past the answer (no EOS stop in the sampler), so keep
/// only the first answer "turn": cut at the first `<|endoftext|>` and then at the
/// first newline. The arithmetic answer is a single line, so this isolates it.
#[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()
.map(|s| PathBuf::from(s.as_str()))
.expect("usage: eval_arith <ckpt> <tokenizer.json> [flags]");
let tok_path = positionals
.get(1)
.map(|s| PathBuf::from(s.as_str()))
.unwrap_or_else(|| PathBuf::from("/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 max_new = flag(&args, "--max-tokens", 48usize);
let n_show = flag(&args, "--show", 8usize);
let prompts_file = flag_value(&args, "--prompts-file").expect("--prompts-file is required");
let gold_file = flag_value(&args, "--gold-file").expect("--gold-file is required");
// M2: decode through the KV-cache incremental engine instead of the naive
// full-recompute sampler. Token-identical to the naive path (gated by
// tests/decode_kv.rs); this flag also lets us A/B the two for the speedup.
let use_cached = args.iter().any(|a| a == "--cached");
// Prompts: skip the `#` header / blank lines and decode escaped newlines so the
// count and order line up with the gold file.
let prompts: Vec<String> = std::fs::read_to_string(&prompts_file)
.unwrap_or_else(|e| panic!("read prompts {prompts_file}: {e}"))
.lines()
.map(str::trim)
.filter(|l| !l.is_empty() && !l.starts_with('#'))
.map(decode_escapes)
.collect();
let golds: Vec<i64> = std::fs::read_to_string(&gold_file)
.unwrap_or_else(|e| panic!("read gold {gold_file}: {e}"))
.lines()
.map(str::trim)
.filter(|l| !l.is_empty())
.map(|l| l.parse::<i64>().expect("gold line not an integer"))
.collect();
assert_eq!(
prompts.len(),
golds.len(),
"prompt/gold count mismatch ({} vs {})",
prompts.len(),
golds.len()
);
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let tok = Tokenizer::from_file(&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 = 1u64;
let model = TinyTransformer::new(cfg, device, |shape| {
seed = seed.wrapping_add(1);
let n: usize = shape.iter().product();
if shape.len() == 1 {
fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
} else {
fill(n, seed, 0.04)
}
});
xtrain_train::checkpoint::load_into(&ckpt, &model.params()).expect("load checkpoint");
println!(
"eval_arith: ckpt {} | {} prompts | max_new {} | decode={}",
ckpt.display(),
prompts.len(),
max_new,
if use_cached { "kv-cache" } else { "naive" }
);
let (mut n_boxed, mut n_correct) = (0usize, 0usize);
let mut shown = 0usize;
let mut gen_tokens = 0usize;
let t0 = std::time::Instant::now();
for (prompt, &gold) in prompts.iter().zip(&golds) {
let ids: Vec<i32> = tok.encode(prompt).into_iter().map(|t| t as i32).collect();
let out = if use_cached {
xtrain_model::generate_greedy_cached(&model, device, &ids, max_new)
} else {
let mut rng = 7u64;
generate(&model, device, &ids, max_new, 0.0, &mut rng)
};
gen_tokens += out.len() - ids.len();
let cont = tok.decode(&out[ids.len()..].iter().map(|&t| t as u32).collect::<Vec<_>>());
let seg = first_answer_segment(&cont);
if parse_boxed_answer(seg).is_some() {
n_boxed += 1;
}
let ok = check_answer(seg, gold);
if ok {
n_correct += 1;
}
if shown < n_show {
let q = prompt.replace('\n', " ");
println!(" [{}] gold={gold} got={seg:?} {}", q, if ok { "OK" } else { "x" });
shown += 1;
}
}
let elapsed = t0.elapsed().as_secs_f64();
let n = prompts.len() as f64;
println!(
"RESULT format(boxed)={}/{} ({:.1}%) | correct={}/{} ({:.1}%)",
n_boxed,
prompts.len(),
100.0 * n_boxed as f64 / n,
n_correct,
prompts.len(),
100.0 * n_correct as f64 / n,
);
println!(
"TIMING decode={} | {:.2}s | {} gen tokens | {:.1} tok/s",
if use_cached { "kv-cache" } else { "naive" },
elapsed,
gen_tokens,
gen_tokens as f64 / elapsed,
);
}

106
crates/xtrain-train/src/bin/gen_arith_task.rs
View File

@@ -1,106 +0,0 @@
//! Generate the M1 verifiable-arithmetic post-training dataset. Pure host tool (no
//! CUDA): writes
//! <out>/arith_sft.tsv user<TAB>assistant rows for `train --sft-tsv`
//! <out>/arith_eval_prompts.txt greedy_sample `--prompts-file` format (held out)
//! <out>/arith_eval_gold.txt parallel gold integers for the checker
//!
//! Eval problems are deduped against train (no leakage). The SFT rows carry just the
//! user/assistant content; `data::load_sft_tsv_cached` adds the `User:/Assistant:`
//! frame + `<|endoftext|>` and masks the prompt, so the eval prompt lines here
//! reconstruct exactly that frame (`User: <q>\nAssistant:`, literal `\n` decoded by
//! greedy_sample).
//!
//! Example:
//! cargo run -p xtrain-train --release --bin gen_arith_task -- \
//! --n 20000 --eval 500 --seed 1 --out-dir /dashscope-tmp/wjh/xtrain_post/arith
use std::collections::HashSet;
use std::fs::{self, File};
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use xtrain_train::task::{GenConfig, Op, gen_problem, unique_space};
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(|v| v.parse().ok())
.unwrap_or(default)
}
fn main() {
let args: Vec<String> = std::env::args().collect();
let n_train: usize = flag(&args, "--n", 20000);
let n_eval: usize = flag(&args, "--eval", 500);
let seed: u64 = flag(&args, "--seed", 1);
let max_add: i64 = flag(&args, "--max-add", 999);
let max_mul: i64 = flag(&args, "--max-mul", 99);
let out_dir: PathBuf = args
.iter()
.position(|a| a == "--out-dir")
.and_then(|i| args.get(i + 1))
.map(PathBuf::from)
.expect("--out-dir <dir> is required");
fs::create_dir_all(&out_dir).expect("create out dir");
let cfg = GenConfig {
max_add,
max_mul,
ops: vec![Op::Add, Op::Sub, Op::Mul],
};
// Guard: train + eval are deduped (and eval is held out from train), so the
// request must fit comfortably inside the unique key space. Cap at 80% to keep
// dedup fast and the disjoint-eval loop terminating.
let space = unique_space(&cfg);
let need = (n_train + n_eval) as u64;
assert!(
need * 5 <= space * 4,
"requested {need} unique problems but the space is only {space} \
(max_add={max_add}, max_mul={max_mul}); raise --max-add/--max-mul or lower --n/--eval"
);
let mut rng = seed.max(1);
// Train: dedup so the same problem is not repeated and so eval can be held out.
let mut train_keys = HashSet::new();
let mut tsv = BufWriter::new(File::create(out_dir.join("arith_sft.tsv")).expect("create tsv"));
while train_keys.len() < n_train {
let p = gen_problem(&mut rng, &cfg);
if !train_keys.insert(p.key()) {
continue;
}
writeln!(tsv, "{}\t{}", p.question(), p.sft_answer()).expect("write tsv");
}
tsv.flush().expect("flush tsv");
// Eval: disjoint from train (skip any key seen in train) and from itself.
let mut prompts =
BufWriter::new(File::create(out_dir.join("arith_eval_prompts.txt")).expect("create eval"));
let mut golds =
BufWriter::new(File::create(out_dir.join("arith_eval_gold.txt")).expect("create gold"));
writeln!(prompts, "# verifiable arithmetic eval prompts (held out from arith_sft.tsv)")
.expect("write header");
let mut eval_keys = HashSet::new();
while eval_keys.len() < n_eval {
let p = gen_problem(&mut rng, &cfg);
if train_keys.contains(&p.key()) || !eval_keys.insert(p.key()) {
continue;
}
writeln!(prompts, "User: {}\\nAssistant:", p.question()).expect("write prompt");
writeln!(golds, "{}", p.answer()).expect("write gold");
}
prompts.flush().expect("flush prompts");
golds.flush().expect("flush golds");
println!(
"wrote {} train rows + {} eval prompts to {} (ops=+,-,* max_add={} max_mul={} seed={})",
train_keys.len(),
eval_keys.len(),
out_dir.display(),
max_add,
max_mul,
seed
);
}

157
crates/xtrain-train/src/bin/gen_dpo_pairs.rs
View File

@@ -1,157 +0,0 @@
//! 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)"
);
}

70
crates/xtrain-train/src/bin/greedy_sample.rs
View File

@@ -7,8 +7,7 @@
//! export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
//! cargo run -p xtrain-train --release --bin greedy_sample -- \
//! /tmp/xtrain_v4.ckpt /opt/wjh/models/gpt2/tokenizer.json \
//! --heads 24 --head-dim 32 --layers 18 --ffn 2048 \
//! --prompts-file scripts/chat_alpha_fixed_prompts.txt --max-tokens 120
//! --heads 24 --head-dim 32 --layers 18 --ffn 2048
#[cfg(no_cuda)]
fn main() {
@@ -53,60 +52,6 @@ fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
.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()
}
#[cfg(not(no_cuda))]
fn flag_values(args: &[String], name: &str) -> Vec<String> {
args.iter()
.enumerate()
.filter_map(|(i, a)| {
if a == name {
args.get(i + 1).cloned()
} else {
None
}
})
.collect()
}
#[cfg(not(no_cuda))]
fn decode_prompt_escapes(s: &str) -> String {
s.replace("\\n", "\n").replace("\\t", "\t")
}
#[cfg(not(no_cuda))]
fn load_prompts(args: &[String]) -> Vec<String> {
let mut prompts = Vec::new();
if let Some(path) = flag_value(args, "--prompts-file") {
let text = std::fs::read_to_string(&path)
.unwrap_or_else(|e| panic!("failed to read prompts file {path}: {e}"));
prompts.extend(
text.lines()
.map(str::trim)
.filter(|line| !line.is_empty() && !line.starts_with('#'))
.map(decode_prompt_escapes),
);
}
prompts.extend(
flag_values(args, "--prompt")
.into_iter()
.map(|p| decode_prompt_escapes(&p)),
);
if prompts.is_empty() {
prompts = ["Once upon a time", "One day", "The little"]
.into_iter()
.map(String::from)
.collect();
}
prompts
}
#[cfg(not(no_cuda))]
fn main() {
use xserv_tokenizer::Tokenizer;
@@ -130,8 +75,6 @@ fn main() {
// GQA (Phase T15): num K/V heads (must match the ckpt; default = --heads).
let kv_heads = flag(&args, "--kv-heads", n_heads);
let max_new = flag(&args, "--max-tokens", 40usize);
let temperature = flag(&args, "--temperature", 0.0f32);
let prompts = load_prompts(&args);
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
@@ -163,16 +106,11 @@ fn main() {
});
xtrain_train::checkpoint::load_into(&ckpt, &model.params()).expect("load checkpoint");
println!(
"decode: prompts={} max_new={} temperature={}",
prompts.len(),
max_new,
temperature
);
let prompts = ["Once upon a time", "One day", "The little"];
for p in prompts {
let ids: Vec<i32> = tok.encode(&p).into_iter().map(|t| t as i32).collect();
let ids: Vec<i32> = tok.encode(p).into_iter().map(|t| t as i32).collect();
let mut rng = 7u64;
let out = generate(&model, device, &ids, max_new, temperature, &mut rng);
let out = generate(&model, device, &ids, max_new, 0.0, &mut rng);
let text = tok.decode(&out.iter().map(|&t| t as u32).collect::<Vec<_>>());
println!("[{p}] → {text}");
}

19
crates/xtrain-train/src/bin/train.rs
View File

@@ -115,7 +115,6 @@ fn main() {
let val_tokens: usize = flag(&args, "--val-tokens", 0);
let eval_every: usize = flag(&args, "--eval-every", 0);
let eval_batches: usize = flag(&args, "--eval-batches", 64);
let sft_tsv = args.iter().any(|a| a == "--sft-tsv");
// Dropout (Phase T18): residual-path dropout prob, active at training time
// only (inverted scaling), identity at eval/sampling/export. Default 0 = off
// (forward graph bit-identical to the no-dropout path).
@@ -137,11 +136,6 @@ fn main() {
.cloned()
.unwrap_or_else(|| "/tmp/xtrain_tinystories.ckpt".to_string()),
);
let init_ckpt: Option<PathBuf> = args
.iter()
.position(|a| a == "--init-ckpt")
.and_then(|i| args.get(i + 1))
.map(PathBuf::from);
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
@@ -152,19 +146,12 @@ fn main() {
tok_path.display(),
corpus_path.display()
);
let corpus = if sft_tsv {
Corpus::load_sft_tsv_cached(&tok_path, &corpus_path)
} else {
Corpus::load_cached(&tok_path, &corpus_path)
};
let corpus = Corpus::load_cached(&tok_path, &corpus_path);
println!(
"corpus: {} tokens, vocab {}",
corpus.len(),
corpus.vocab_size
);
if sft_tsv {
println!("SFT TSV: ON (assistant-only loss via ignore-index labels)");
}
let vocab = corpus.vocab_size;
// Hold out a tail slice for validation (if requested and the corpus is big).
let (train_corpus, valid) = if val_tokens > 0 {
@@ -219,10 +206,6 @@ fn main() {
if dropout > 0.0 {
println!("dropout: ON (p={dropout}, residual-path, train-only inverted scaling)");
}
if let Some(path) = &init_ckpt {
xtrain_train::checkpoint::load_into(path, &model.params()).expect("load init checkpoint");
println!("init checkpoint: loaded {}", path.display());
}
// Eval-only mode: load a checkpoint and score it on the held-out val set, then
// exit. Used to put an EXISTING model (e.g. v0) and a new one on the same

233
crates/xtrain-train/src/bin/train_dpo.rs
View File

@@ -1,233 +0,0 @@
//! 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()
);
}

294
crates/xtrain-train/src/bin/train_grpo.rs
View File

@@ -1,294 +0,0 @@
//! GRPO training on the verifiable arithmetic task (M4 / Stage P3) — online,
//! critic-free RL. The centerpiece: generation INSIDE the training loop.
//!
//! Each step: sample B prompts (fresh problems), roll out G completions per prompt
//! (temperature sampling via the naive sampler — batched/cached rollout is the M2b/
//! M4-perf follow-up), score each with the rule-based checker (reward ∈ {0,1}),
//! compute the **group-relative advantage** `A_i = (r_i − mean) / (std + ε)` (no
//! critic), then K inner clipped-PG epochs minimising [`clipped_pg_loss`] with a KL
//! leash to the frozen reference (πref = the SFT checkpoint). Reward = pure 0/1
//! correctness; the KL term (β) is what keeps format/coherence (the M3 collapse
//! lesson — here it is an explicit leash, not just a hope).
//!
//! Health signal (the falsifiable "it learns"): **mean rollout reward must rise**
//! (the RL analogue of T5's overfit-27/27). Held-out correctness is measured by
//! eval_arith on the saved checkpoint.
//!
//! train_grpo <tokenizer.json> --init-ckpt <sft.ckpt> <arch flags> \
//! --steps 200 --group 6 --prompts 8 --temp 1.0 --beta 0.04 --eps 0.2 \
//! --lr 1e-6 --max-add 20 --max-mul 9 --ckpt <out.ckpt>
#[cfg(no_cuda)]
fn main() {
eprintln!("train_grpo: built without CUDA (no_cuda); run on a GPU host.");
}
#[cfg(not(no_cuda))]
use xtrain_cuda::device;
#[cfg(not(no_cuda))]
use xtrain_model::{Config, TinyTransformer, generate_cached_batch};
#[cfg(not(no_cuda))]
use xtrain_tensor::{DType, Device};
#[cfg(not(no_cuda))]
use xtrain_train::grpo_batch::{PgSample, inner_pg_step_batched, per_token_logp_batched};
#[cfg(not(no_cuda))]
use xtrain_train::task::{check_answer, gen_problem, GenConfig, Op};
#[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()
}
#[cfg(not(no_cuda))]
fn first_answer_segment(c: &str) -> &str {
let s = c.split("<|endoftext|>").next().unwrap_or(c);
s.split('\n').next().unwrap_or(s)
}
/// Build a model from the SFT checkpoint (bf16 compute to fit two 1B models). The
/// policy enables activation recompute (T13) so its backward fits alongside the
/// frozen reference + the Adam state; the reference only forwards (no backward).
#[cfg(not(no_cuda))]
fn load_model(cfg: Config, device: Device, ckpt: &str, recompute: bool) -> TinyTransformer {
let mut seed = 1u64;
let m = TinyTransformer::new(cfg, device, |shape| {
seed = seed.wrapping_add(1);
let n: usize = shape.iter().product();
if shape.len() == 1 {
fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
} else {
fill(n, seed, 0.04)
}
})
.with_compute_dtype(DType::BF16)
.with_recompute(recompute)
.with_flash(true);
xtrain_train::checkpoint::load_into(std::path::Path::new(ckpt), &m.params()).expect("load ckpt");
m.eval();
m
}
/// Frame (question, completion) like the SFT loader and return the next-token
/// (input, target) pair (prompt masked to -100). Same as train_dpo.
#[cfg(not(no_cuda))]
fn frame(tok: &xserv_tokenizer::Tokenizer, question: &str, completion: &str) -> (Vec<i32>, Vec<i32>) {
let p_ids: Vec<i32> = tok
.encode(&format!("User: {question}\nAssistant:"))
.into_iter()
.map(|t| t as i32)
.collect();
let a_ids: Vec<i32> = tok
.encode(&format!(" {completion}\n<|endoftext|>"))
.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())
}
#[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_grpo <tokenizer.json> [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 steps: usize = flag(&args, "--steps", 200);
let group: usize = flag(&args, "--group", 6);
let n_prompts: usize = flag(&args, "--prompts", 8);
let inner: usize = flag(&args, "--inner", 1);
// M2d: pack the step's N=B·G ragged samples into forward_batched chunks of this
// many samples (bounds the [chunk·Lmax, vocab] logits memory). Default = whole batch.
let micro: usize = flag(&args, "--micro", n_prompts * group.max(1));
let temp: f32 = flag(&args, "--temp", 1.0);
let beta: f32 = flag(&args, "--beta", 0.04);
let eps: f32 = flag(&args, "--eps", 0.2);
let lr: f32 = flag(&args, "--lr", 1e-6);
let clip: f32 = flag(&args, "--clip", 1.0);
let max_new: usize = flag(&args, "--max-tokens", 24);
let max_add: i64 = flag(&args, "--max-add", 20);
let max_mul: i64 = flag(&args, "--max-mul", 9);
let seed: u64 = flag(&args, "--seed", 20260630);
let log_every: usize = flag(&args, "--log-every", 20);
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");
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 policy = load_model(cfg, device, &init_ckpt, false); // flash keeps attn memory bounded
// Frozen πref for the KL leash — only resident when β>0 (a second 1B model is the
// memory long-pole; β=0 is pure PG and skips it, the gated degenerate).
let reference = if beta > 0.0 {
Some(load_model(cfg, device, &init_ckpt, false))
} else {
None
};
let gcfg = GenConfig {
max_add,
max_mul,
ops: vec![Op::Add, Op::Sub, Op::Mul],
};
let params = policy.params();
let mut opt = GpuAdamW::new(0.0);
let mut rng = seed.max(1);
let start = std::time::Instant::now();
let (mut win_reward, mut win_solved, mut win_n) = (0f32, 0usize, 0usize);
// Per-window phase timers (ms): rollout / capture / inner — to keep the step
// decomposition honest (M2d cut the training-side forwards 9×, so the question is
// what now dominates the step).
let (mut t_roll, mut t_cap, mut t_inner) = (0f32, 0f32, 0f32);
for step in 0..steps {
// ---- Rollout: B prompts × G completions, scored, group-advantage ----
// Collect ALL the step's framed samples first (input, target, adv), so the
// training-side forwards can be batched across the whole step (M2d) instead of
// run one ragged sequence at a time.
let t0 = std::time::Instant::now();
let mut raw: Vec<(Vec<i32>, Vec<i32>, f32)> = Vec::new();
for _ in 0..n_prompts {
let p = gen_problem(&mut rng, &gcfg);
let prompt_ids: Vec<i32> = tok
.encode(&format!("User: {}\nAssistant:", p.question()))
.into_iter()
.map(|t| t as i32)
.collect();
// M2b batched rollout: the G samples of this prompt decode in lockstep
// (one forward per step over the whole group → G× fewer kernel launches
// than G sequential single-seq rollouts; the M4 rollout long-pole fix).
let mut comps: Vec<(String, f32)> = Vec::with_capacity(group);
let outs = generate_cached_batch(&policy, device, &prompt_ids, group, max_new, temp, &mut rng);
for out in &outs {
let cont = tok.decode(&out[prompt_ids.len()..].iter().map(|&t| t as u32).collect::<Vec<_>>());
let seg = first_answer_segment(&cont).trim().to_string();
let r = if check_answer(&seg, p.answer()) { 1.0 } else { 0.0 };
comps.push((seg, r));
}
let mean = comps.iter().map(|c| c.1).sum::<f32>() / group as f32;
let var = comps.iter().map(|c| (c.1 - mean).powi(2)).sum::<f32>() / group as f32;
let std = var.sqrt();
win_reward += mean * group as f32;
win_solved += comps.iter().filter(|c| c.1 > 0.5).count();
win_n += group;
// A whole group with no reward variance gives zero advantage → skip
// (no learning signal, and avoids dividing by ~0).
if std < 1e-6 {
continue;
}
for (seg, r) in &comps {
let adv = (r - mean) / (std + 1e-4);
let (input, target) = frame(&tok, &p.question(), seg);
raw.push((input, target, adv));
}
}
t_roll += t0.elapsed().as_secs_f32() * 1e3;
// ---- Batched capture (M2d): logπ_old (policy) + logπ_ref (frozen) over ALL
// samples in forward_batched chunks, instead of one forward per sample. ----
if !raw.is_empty() {
let t1 = std::time::Instant::now();
let io: Vec<(Vec<i32>, Vec<i32>)> = raw.iter().map(|(i, t, _)| (i.clone(), t.clone())).collect();
let logp_old = per_token_logp_batched(&policy, device, &io, micro);
// β=0 ⇒ KL term drops ⇒ logp_ref unused; pass zeros (no reference model).
let logp_ref = match &reference {
Some(r) => per_token_logp_batched(r, device, &io, micro),
None => raw.iter().map(|(i, _, _)| vec![0.0; i.len()]).collect(),
};
let batch: Vec<PgSample> = raw
.iter()
.zip(logp_old)
.zip(logp_ref)
.map(|(((input, target, adv), lo), lr)| PgSample {
input: input.clone(),
target: target.clone(),
adv: *adv,
logp_old: lo,
logp_ref: lr,
})
.collect();
t_cap += t1.elapsed().as_secs_f32() * 1e3;
// ---- K inner clipped-PG epochs, batched over the captured samples ----
let t2 = std::time::Instant::now();
for _ in 0..inner {
inner_pg_step_batched(&policy, device, &batch, eps, beta, micro);
let _ = xtrain_train::clip::clip_grad_norm_gpu(&params, clip, 1.0);
opt.step(lr, &params);
for p in &params {
p.zero_grad();
}
}
t_inner += t2.elapsed().as_secs_f32() * 1e3;
}
if (step + 1) % log_every == 0 || step == steps - 1 {
let w = log_every.min(step + 1) as f32; // steps in this window
println!(
"step {:5}/{steps}: mean-reward {:.3} | solved {}/{} | {:.0}s | ms/step roll {:.0} cap {:.0} inner {:.0}",
step + 1,
win_reward / win_n.max(1) as f32,
win_solved,
win_n,
start.elapsed().as_secs_f32(),
t_roll / w,
t_cap / w,
t_inner / w,
);
win_reward = 0.0;
win_solved = 0;
win_n = 0;
t_roll = 0.0;
t_cap = 0.0;
t_inner = 0.0;
// Periodic save so a later OOM (naive rollout fragments the allocator —
// the long-pole the design doc flagged) still leaves an evaluatable ckpt.
xtrain_train::checkpoint::save(std::path::Path::new(&out_ckpt), &params).expect("save");
}
}
xtrain_train::checkpoint::save(std::path::Path::new(&out_ckpt), &params).expect("save ckpt");
println!("GRPO done: {steps} steps, G={group}, B={n_prompts}, beta {beta}, lr {lr:.1e} → {out_ckpt}");
}

184
crates/xtrain-train/src/data.rs
View File

@@ -15,7 +15,6 @@ use xserv_tokenizer::Tokenizer;
/// A tokenized corpus: one flat stream of token ids, plus the vocab size.
pub struct Corpus {
pub tokens: Vec<i32>,
pub labels: Option<Vec<i32>>,
pub vocab_size: usize,
}
@@ -34,7 +33,6 @@ impl Corpus {
let ids: Vec<i32> = tok.encode(text).into_iter().map(|t| t as i32).collect();
Self {
tokens: ids,
labels: None,
vocab_size: tok.vocab_size(),
}
}
@@ -54,11 +52,7 @@ impl Corpus {
tokens.len(),
cache.display()
);
return Self {
tokens,
labels: None,
vocab_size,
};
return Self { tokens, vocab_size };
}
let me = Self::load(tokenizer_path, corpus_path);
write_u16_cache(&cache, &me.tokens);
@@ -70,103 +64,22 @@ impl Corpus {
me
}
/// Load assistant-only SFT data from a two-column TSV:
///
/// ```text
/// user<TAB>assistant
/// ```
///
/// Literal `\n` and `\t` escapes are decoded. Each row is formatted as
/// `User: ...\nAssistant:` + answer + `<|endoftext|>`. Labels are `-100`
/// for prompt tokens and the token id itself for answer/EOS tokens, so the
/// cross-entropy op ignores prompt rows while still training the assistant
/// answer and stop token.
pub fn load_sft_tsv_cached(tokenizer_path: &Path, corpus_path: &Path) -> Self {
let token_cache = cache_path(corpus_path);
let label_cache = label_cache_path(corpus_path);
let vocab_size = Tokenizer::from_file(tokenizer_path).vocab_size();
if token_cache.exists() && label_cache.exists() {
let tokens = read_u16_cache(&token_cache);
let labels = read_i32_cache(&label_cache);
assert_eq!(
tokens.len(),
labels.len(),
"SFT cache token/label length mismatch"
);
println!(
"corpus: read {} cached SFT tokens from {} (+ labels {})",
tokens.len(),
token_cache.display(),
label_cache.display()
);
return Self {
tokens,
labels: Some(labels),
vocab_size,
};
}
let tok = Tokenizer::from_file(tokenizer_path);
let text = std::fs::read_to_string(corpus_path)
.unwrap_or_else(|e| panic!("failed to read SFT corpus {}: {e}", corpus_path.display()));
let mut tokens = Vec::new();
let mut labels = Vec::new();
for (lineno, line) in text.lines().enumerate() {
if line.trim().is_empty() {
continue;
}
let (user, assistant) = line
.split_once('\t')
.unwrap_or_else(|| panic!("SFT TSV line {} missing tab", lineno + 1));
let user = decode_tsv_escapes(user);
let assistant = decode_tsv_escapes(assistant);
let prompt = format!("User: {user}\nAssistant:");
let answer = format!(" {assistant}\n<|endoftext|>");
let prompt_ids: Vec<i32> = tok.encode(&prompt).into_iter().map(|t| t as i32).collect();
let answer_ids: Vec<i32> = tok.encode(&answer).into_iter().map(|t| t as i32).collect();
let (row_tokens, row_labels) = sft_row(&prompt_ids, &answer_ids);
tokens.extend(row_tokens);
labels.extend(row_labels);
}
assert_eq!(tokens.len(), labels.len(), "SFT tokens/labels mismatch");
write_u16_cache(&token_cache, &tokens);
write_i32_cache(&label_cache, &labels);
println!(
"corpus: tokenized {} SFT tokens → cached to {} (+ labels {})",
tokens.len(),
token_cache.display(),
label_cache.display()
);
Self {
tokens,
labels: Some(labels),
vocab_size: tok.vocab_size(),
}
}
/// Split off the last `n` tokens as a held-out validation corpus, leaving the
/// rest as the train corpus. Returns `(train, valid)`. Used for periodic val
/// loss during training without leaking the eval window into training.
pub fn split_tail(self, n: usize) -> (Self, Self) {
let n = n.min(self.tokens.len() / 10); // never hand off more than 10%
let cut = self.tokens.len() - n;
let valid_tokens = self.tokens[cut..].to_vec();
let valid_labels = self.labels.as_ref().map(|labels| labels[cut..].to_vec());
let valid = self.tokens[cut..].to_vec();
let mut train = self.tokens;
train.truncate(cut);
let train_labels = self.labels.map(|mut labels| {
labels.truncate(cut);
labels
});
(
Self {
tokens: train,
labels: train_labels,
vocab_size: self.vocab_size,
},
Self {
tokens: valid_tokens,
labels: valid_labels,
tokens: valid,
vocab_size: self.vocab_size,
},
)
@@ -188,27 +101,11 @@ impl Corpus {
pub fn sample(&self, seq: usize, rng_state: &mut u64) -> (Vec<i32>, Vec<i32>) {
assert!(self.tokens.len() > seq + 1, "corpus shorter than a window");
let max_start = self.tokens.len() - seq - 1;
let mut start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
if let Some(labels) = &self.labels {
for _ in 0..16 {
if labels[start + 1..start + seq + 1].iter().any(|&t| t >= 0) {
break;
}
start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
}
}
let start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
let input = self.tokens[start..start + seq].to_vec();
let target = self.target_window(start, seq);
let target = self.tokens[start + 1..start + seq + 1].to_vec();
(input, target)
}
/// Deterministic target labels for an input window starting at `start`.
pub fn target_window(&self, start: usize, seq: usize) -> Vec<i32> {
match &self.labels {
Some(labels) => labels[start + 1..start + seq + 1].to_vec(),
None => self.tokens[start + 1..start + seq + 1].to_vec(),
}
}
}
/// Drop a leading partial line (before the first newline) and everything after
@@ -230,12 +127,6 @@ fn cache_path(corpus_path: &Path) -> PathBuf {
PathBuf::from(s)
}
fn label_cache_path(corpus_path: &Path) -> PathBuf {
let mut s = corpus_path.as_os_str().to_os_string();
s.push(".labels.i32.bin");
PathBuf::from(s)
}
/// Read a flat little-endian `[u16]` cache into an `i32` id stream.
fn read_u16_cache(path: &Path) -> Vec<i32> {
let mut r = BufReader::new(
@@ -249,18 +140,6 @@ fn read_u16_cache(path: &Path) -> Vec<i32> {
.collect()
}
fn read_i32_cache(path: &Path) -> Vec<i32> {
let mut r = BufReader::new(
std::fs::File::open(path).unwrap_or_else(|e| panic!("open cache {}: {e}", path.display())),
);
let mut buf = Vec::new();
r.read_to_end(&mut buf).expect("read cache");
assert!(buf.len() % 4 == 0, "corrupt i32 cache (odd byte count)");
buf.chunks_exact(4)
.map(|b| i32::from_le_bytes([b[0], b[1], b[2], b[3]]))
.collect()
}
/// Write an id stream as a flat little-endian `[u16]` cache. Ids must fit in u16
/// (GPT-2 vocab = 50257 < 65536); asserts otherwise.
fn write_u16_cache(path: &Path, tokens: &[i32]) {
@@ -275,35 +154,6 @@ fn write_u16_cache(path: &Path, tokens: &[i32]) {
w.flush().expect("flush cache");
}
fn write_i32_cache(path: &Path, labels: &[i32]) {
let mut w = BufWriter::new(
std::fs::File::create(path)
.unwrap_or_else(|e| panic!("create cache {}: {e}", path.display())),
);
for &t in labels {
w.write_all(&t.to_le_bytes()).expect("write cache");
}
w.flush().expect("flush cache");
}
fn decode_tsv_escapes(s: &str) -> String {
s.replace("\\n", "\n").replace("\\t", "\t")
}
/// Build one SFT example's `(tokens, labels)` from already-tokenized prompt/answer
/// ids: prompt tokens are masked to the ignore-index (`-100`, which `cross_entropy`
/// skips) so only the answer + EOS tokens contribute to the loss. Pure (no tokenizer
/// / no CUDA) so the assistant-only masking is unit-testable directly.
fn sft_row(prompt_ids: &[i32], answer_ids: &[i32]) -> (Vec<i32>, Vec<i32>) {
let mut tokens = Vec::with_capacity(prompt_ids.len() + answer_ids.len());
tokens.extend_from_slice(prompt_ids);
tokens.extend_from_slice(answer_ids);
let mut labels = Vec::with_capacity(prompt_ids.len() + answer_ids.len());
labels.extend(std::iter::repeat(-100).take(prompt_ids.len()));
labels.extend_from_slice(answer_ids);
(tokens, labels)
}
/// Tiny LCG (same constants as the model tests' deterministic fill) so dataset
/// sampling is reproducible from a single u64 seed.
fn next_rand(state: &mut u64) -> u64 {
@@ -312,27 +162,3 @@ fn next_rand(state: &mut u64) -> u64 {
.wrapping_add(1442695040888963407);
*state >> 16
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn sft_row_masks_prompt_supervises_answer() {
let prompt = [5, 6, 7];
let answer = [8, 9]; // includes the EOS token in real use
let (tokens, labels) = sft_row(&prompt, &answer);
// Tokens are prompt then answer, in order.
assert_eq!(tokens, vec![5, 6, 7, 8, 9]);
// Prompt positions are ignore-index (-100); answer positions are supervised.
assert_eq!(labels, vec![-100, -100, -100, 8, 9]);
assert_eq!(tokens.len(), labels.len());
}
#[test]
fn sft_row_handles_empty_answer() {
let (tokens, labels) = sft_row(&[1, 2], &[]);
assert_eq!(tokens, vec![1, 2]);
assert_eq!(labels, vec![-100, -100]);
}
}

162
crates/xtrain-train/src/grpo_batch.rs
View File

@@ -1,162 +0,0 @@
//! 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
}

3
crates/xtrain-train/src/lib.rs
View File

@@ -10,13 +10,10 @@
pub mod clip;
pub mod data;
pub mod schedule;
pub mod task;
#[cfg(not(no_cuda))]
pub mod checkpoint;
#[cfg(not(no_cuda))]
pub mod grpo_batch;
#[cfg(not(no_cuda))]
pub mod sample;
#[cfg(not(no_cuda))]
mod train_loop;

240
crates/xtrain-train/src/task.rs
View File

@@ -1,240 +0,0 @@
//! Verifiable arithmetic task (post-training, M1). A tiny two-operand integer
//! arithmetic task with a deterministic, rule-based checker: the assistant must end
//! its answer with `\boxed{N}`, and the reward is exact-match on `N`.
//!
//! This single module is the shared task spec for the whole post-training stack —
//! M1 SFT-data generation, M3 DPO preference-pair construction, and M4 GRPO reward
//! scoring all parse/score through here, so the task lives in exactly one place.
//!
//! Host-only (no CUDA): generation + parsing + checking are pure, so this compiles
//! and unit-tests on a GPU-less host.
use std::fmt;
/// The supported binary operations.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Op {
Add,
Sub,
Mul,
}
impl Op {
pub fn symbol(self) -> char {
match self {
Op::Add => '+',
Op::Sub => '-',
Op::Mul => '*',
}
}
}
impl fmt::Display for Op {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.symbol())
}
}
/// A single two-operand arithmetic problem.
#[derive(Clone, Copy, Debug)]
pub struct Problem {
pub a: i64,
pub b: i64,
pub op: Op,
}
impl Problem {
/// The exact integer answer (the verifiable gold label).
pub fn answer(self) -> i64 {
match self.op {
Op::Add => self.a + self.b,
Op::Sub => self.a - self.b,
Op::Mul => self.a * self.b,
}
}
/// The user-turn question text. No template wrapping — the SFT loader
/// (`data::load_sft_tsv_cached`) adds the `User:/Assistant:` frame.
pub fn question(self) -> String {
format!("What is {} {} {}?", self.a, self.op, self.b)
}
/// The assistant-turn SFT target: restate the equation and end with the boxed
/// answer. This teaches the answer FORMAT (the checker only reads `\boxed{}`);
/// arithmetic correctness is what DPO (M3) / GRPO (M4) later improve.
pub fn sft_answer(self) -> String {
format!("{} {} {} = \\boxed{{{}}}.", self.a, self.op, self.b, self.answer())
}
/// A stable dedup key, so eval problems can be held out from train.
pub fn key(self) -> (i64, char, i64) {
(self.a, self.op.symbol(), self.b)
}
}
/// Operand-range configuration for problem sampling. Multiplication uses a smaller
/// range (`max_mul`) so products stay modest; add/sub use `max_add`.
#[derive(Clone)]
pub struct GenConfig {
pub max_add: i64,
pub max_mul: i64,
pub ops: Vec<Op>,
}
impl Default for GenConfig {
fn default() -> Self {
Self {
max_add: 999,
max_mul: 99,
ops: vec![Op::Add, Op::Sub, Op::Mul],
}
}
}
/// Number of distinct problems this config can produce (the key space). Used to
/// guard the dedup generator against requesting more unique problems than exist —
/// otherwise train/eval dedup loops near saturation get pathologically slow or, for
/// a disjoint eval, never terminate.
pub fn unique_space(cfg: &GenConfig) -> u64 {
cfg.ops
.iter()
.map(|op| {
let max = if *op == Op::Mul { cfg.max_mul } else { cfg.max_add };
((max as u64) + 1).pow(2) // ordered (a, b) pairs in [0, max]
})
.sum()
}
/// Sample one problem deterministically from the LCG state `rng`. Operands are drawn
/// in `[0, max]` per the op; subtraction may yield a negative answer (the checker /
/// parser handle a leading `-`).
pub fn gen_problem(rng: &mut u64, cfg: &GenConfig) -> Problem {
let op = cfg.ops[(next_rand(rng) as usize) % cfg.ops.len()];
let max = if op == Op::Mul { cfg.max_mul } else { cfg.max_add };
let a = rand_range(rng, max);
let b = rand_range(rng, max);
Problem { a, b, op }
}
/// Parse the integer inside the LAST `\boxed{...}` in `text`. Returns `None` if there
/// is no well-formed boxed integer (no box, empty, or non-integer contents). "Last"
/// so a model that emits intermediate boxes still scores on its final answer.
pub fn parse_boxed_answer(text: &str) -> Option<i64> {
const TAG: &str = "\\boxed{";
let mut found = None;
let mut rest = text;
while let Some(i) = rest.find(TAG) {
let after = &rest[i + TAG.len()..];
match after.find('}') {
Some(j) => {
if let Ok(n) = after[..j].trim().parse::<i64>() {
found = Some(n);
}
rest = &after[j + 1..];
}
None => break,
}
}
found
}
/// Verifiable reward: does the completion's boxed answer exactly match `gold`?
pub fn check_answer(completion: &str, gold: i64) -> bool {
parse_boxed_answer(completion) == Some(gold)
}
/// `[0, max]` inclusive draw from the LCG.
fn rand_range(rng: &mut u64, max: i64) -> i64 {
debug_assert!(max >= 0);
(next_rand(rng) % (max as u64 + 1)) as i64
}
/// Same LCG constants as the dataset sampler (`data::next_rand`), kept local so the
/// task module stays dependency-free and host-only.
fn next_rand(state: &mut u64) -> u64 {
*state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
*state >> 1
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn answer_question_and_sft_target() {
let p = Problem {
a: 12,
b: 13,
op: Op::Mul,
};
assert_eq!(p.answer(), 156);
assert_eq!(p.question(), "What is 12 * 13?");
assert_eq!(p.sft_answer(), "12 * 13 = \\boxed{156}.");
let s = Problem {
a: 3,
b: 8,
op: Op::Sub,
};
assert_eq!(s.answer(), -5);
}
#[test]
fn parse_takes_last_boxed_and_handles_edges() {
assert_eq!(parse_boxed_answer("\\boxed{3} then \\boxed{156}."), Some(156));
assert_eq!(parse_boxed_answer("\\boxed{-7}"), Some(-7));
assert_eq!(parse_boxed_answer("\\boxed{ 42 }"), Some(42));
assert_eq!(parse_boxed_answer("no box here"), None);
assert_eq!(parse_boxed_answer("\\boxed{abc}"), None);
assert_eq!(parse_boxed_answer("\\boxed{unterminated"), None);
}
#[test]
fn check_is_exact_match() {
assert!(check_answer("the result is \\boxed{156}.", 156));
assert!(!check_answer("the result is \\boxed{155}.", 156));
assert!(!check_answer("no boxed answer at all", 156));
}
#[test]
fn sft_target_is_always_self_consistent() {
// The SFT target's boxed answer must always check against the problem's own
// gold — across all ops/operands. This is the M1 data invariant.
let cfg = GenConfig::default();
let mut rng = 12345u64;
for _ in 0..2000 {
let p = gen_problem(&mut rng, &cfg);
assert!(
check_answer(&p.sft_answer(), p.answer()),
"self-inconsistent SFT target for {p:?}"
);
}
}
#[test]
fn unique_space_counts_ordered_pairs_per_op() {
// add+sub+mul each contribute (max+1)^2 ordered pairs.
let cfg = GenConfig {
max_add: 9,
max_mul: 4,
ops: vec![Op::Add, Op::Sub, Op::Mul],
};
assert_eq!(unique_space(&cfg), 100 + 100 + 25);
// The shipped default is comfortably large (millions), so 20k requests are
// a tiny fraction and dedup stays fast.
assert!(unique_space(&GenConfig::default()) > 1_000_000);
}
#[test]
fn generation_is_deterministic_from_seed() {
let cfg = GenConfig::default();
let (mut r1, mut r2) = (7u64, 7u64);
for _ in 0..200 {
assert_eq!(
gen_problem(&mut r1, &cfg).key(),
gen_problem(&mut r2, &cfg).key()
);
}
}
}

2
crates/xtrain-train/src/train_loop.rs
View File

@@ -207,7 +207,7 @@ pub fn eval_loss(
break;
}
let input: Vec<i32> = valid.tokens[s..s + seq].to_vec();
let target = valid.target_window(s, seq);
let target: Vec<i32> = valid.tokens[s + 1..s + seq + 1].to_vec();
let ids = ids_tensor(&input, device);
let targets = ids_tensor(&target, device);
let loss = model.loss(&ids, &targets);

83
crates/xtrain-train/tests/decode_batch.rs
View File

@@ -1,83 +0,0 @@
// M2b batched KV-cache decode — the token-identical gate.
//
// Batched decode rolls out G samples of one prompt in lockstep (one common decode
// position each step, uniform RoPE via rope_pos, KV cache carrying a G dimension).
// Under GREEDY decoding all G rows are deterministic and must each equal the
// single-sequence greedy decode (generate_greedy_cached, itself gated token-
// identical to the naive sampler). This pins that the G-way batching indexes each
// sequence's K/V correctly (no cross-row contamination) and reproduces M2a exactly.
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_model::{generate_cached_batch, generate_greedy_cached, Config, TinyTransformer};
use xtrain_tensor::{DType, Device};
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()
}
fn build(cfg: Config, device: Device) -> TinyTransformer {
let mut seed = 1u64;
TinyTransformer::new(cfg, device, |shape| {
seed = seed.wrapping_add(1);
let n: usize = shape.iter().product();
if shape.len() == 1 {
fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
} else {
fill(n, seed, 0.08)
}
})
.with_compute_dtype(DType::F32)
}
#[test]
fn batched_greedy_decode_matches_single_seq() {
assert!(
device::device_count().expect("device count") > 0,
"no CUDA device"
);
device::set_device(0).unwrap();
let device = Device::Cuda(0);
// Real GQA (8 query / 2 kv heads → group 4) so repeat_kv(nh, batch=G) is exercised.
let cfg = Config::from_arch(48, 8, 16, 4, 256).with_kv_heads(2);
let model = build(cfg, device);
let prompt: Vec<i32> = vec![3, 9, 1, 14, 5];
let max_new = 24usize;
let g = 5usize;
let single = generate_greedy_cached(&model, device, &prompt, max_new);
let mut rng = 0u64;
let batched = generate_cached_batch(&model, device, &prompt, g, max_new, 0.0, &mut rng);
assert_eq!(batched.len(), g, "expected {g} sample rows");
for (row, seq) in batched.iter().enumerate() {
assert_eq!(
seq.len(),
single.len(),
"row {row} length {} vs single {}",
seq.len(),
single.len()
);
if seq != &single {
let first = seq.iter().zip(&single).position(|(a, b)| a != b).unwrap();
panic!(
"batched row {row} diverges from single-seq at index {first}: {:?} vs {:?}",
seq[first], single[first]
);
}
}
println!(
"batched decode OK: all {g} greedy rows token-identical to single-seq over {max_new} tokens"
);
}

94
crates/xtrain-train/tests/decode_kv.rs
View File

@@ -1,94 +0,0 @@
// M2a KV-cache decode engine — the token-identical correctness gate.
//
// The centerpiece M2 invariant: greedy decode through the KV-cache incremental
// engine (`xtrain_model::generate_greedy_cached`) must be TOKEN-IDENTICAL to the
// naive full-recompute greedy (`xtrain_train::sample::generate` at temperature 0),
// which re-runs the whole forward over the growing prefix each step. Same tokens ⇒
// the cache + decode-time attention + RoPE-at-position reproduce the full forward.
//
// Numerics note: a randomly-initialised model has near-uniform logits, so argmax
// can be fragile to ~1e-6 differences. This unit gate therefore runs in F32 (the
// tightest path, and the dtype the eval harness actually uses) on a small model.
// The headline gate on the trained v12 checkpoint (peaked logits → robust argmax)
// is run on the GPU box and recorded in docs/18.
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, generate_greedy_cached};
use xtrain_tensor::{DType, Device};
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()
}
fn build(cfg: Config, device: Device, dtype: DType) -> TinyTransformer {
let mut seed = 1u64;
let m = TinyTransformer::new(cfg, device, |shape| {
seed = seed.wrapping_add(1);
let n: usize = shape.iter().product();
if shape.len() == 1 {
fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
} else {
fill(n, seed, 0.08)
}
});
m.with_compute_dtype(dtype)
}
// A real GQA config (8 query / 2 kv heads → group 4) to exercise repeat_kv in the
// decode path; head_dim 16, dim 128, 4 layers.
fn gqa_cfg() -> Config {
Config::from_arch(48, 8, 16, 4, 256).with_kv_heads(2)
}
#[test]
fn kv_cache_decode_is_token_identical_to_naive_f32() {
assert!(
device::device_count().expect("device count") > 0,
"no CUDA device"
);
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let model = build(gqa_cfg(), device, DType::F32);
let prompt: Vec<i32> = vec![1, 5, 9, 13, 2, 7];
let max_new = 24usize;
let mut rng = 7u64;
let naive = xtrain_train::sample::generate(&model, device, &prompt, max_new, 0.0, &mut rng);
let cached = generate_greedy_cached(&model, device, &prompt, max_new);
assert_eq!(
naive.len(),
cached.len(),
"length mismatch: naive {} vs cached {}",
naive.len(),
cached.len()
);
if naive != cached {
// Report the first divergence for debugging.
let first = naive
.iter()
.zip(&cached)
.position(|(a, b)| a != b)
.unwrap();
panic!(
"token divergence at index {first}: naive={:?} cached={:?}\nnaive ={naive:?}\ncached ={cached:?}",
naive[first], cached[first]
);
}
println!(
"KV-cache decode token-identical to naive over {} generated tokens (F32, GQA 8/2)",
max_new
);
}

1
crates/xtrain-train/tests/grad_accum.rs
View File

@@ -216,7 +216,6 @@ fn synth_corpus(vocab: usize, n_tokens: usize) -> Corpus {
tokens: (0..n_tokens)
.map(|i| (i * 7 + 3) as i32 % vocab as i32)
.collect(),
labels: None,
vocab_size: vocab,
}
}

101
csrc/ops/nn.cu
View File

@@ -242,96 +242,6 @@ void launch_rope_f32(const float* x, float* y, int tokens, int heads,
rope_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta, period);
}
// RoPE at an absolute position offset (KV-cache decode-time, forward only). Same
// rotate_half as rope_k, but row `tok`'s position is `pos0 + tok` (no modulo) —
// a single new decode token sits at absolute position pos0. The training rope_k
// (position = tok % period) is left untouched, so this adds no training-path risk.
__global__ void rope_at_k(const float* x, float* y, int heads, int head_dim,
float theta, int pos0) {
int tok = blockIdx.x;
int head = blockIdx.y;
int half = head_dim / 2;
int i = threadIdx.x;
if (i >= half) return;
int pos = pos0 + tok;
float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
float angle = (float)pos * freq;
float c = cosf(angle), sn = sinf(angle);
int base = (tok * heads + head) * head_dim;
float x0 = x[base + i], x1 = x[base + i + half];
y[base + i] = x0 * c - x1 * sn;
y[base + i + half] = x1 * c + x0 * sn;
}
void launch_rope_at_f32(const float* x, float* y, int tokens, int heads,
int head_dim, float theta, int pos0, void* s) {
dim3 grid(tokens, heads);
int blk = head_dim / 2;
rope_at_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta, pos0);
}
// RoPE with a PER-ROW absolute position (batched KV-cache decode, M2b): row `tok`'s
// position is `positions[tok]` (an i32 per token). For G-way batched decode all G
// rows share one decode position; for ragged batches each row carries its own.
// Forward only; the training rope_k is untouched.
__global__ void rope_pos_k(const float* x, const int* positions, float* y,
int heads, int head_dim, float theta) {
int tok = blockIdx.x;
int head = blockIdx.y;
int half = head_dim / 2;
int i = threadIdx.x;
if (i >= half) return;
int pos = positions[tok];
float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
float angle = (float)pos * freq;
float c = cosf(angle), sn = sinf(angle);
int base = (tok * heads + head) * head_dim;
float x0 = x[base + i], x1 = x[base + i + half];
y[base + i] = x0 * c - x1 * sn;
y[base + i + half] = x1 * c + x0 * sn;
}
void launch_rope_pos_f32(const float* x, const int* positions, float* y,
int tokens, int heads, int head_dim, float theta, void* s) {
dim3 grid(tokens, heads);
int blk = head_dim / 2;
rope_pos_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, positions, y, heads, head_dim, theta);
}
// Concatenate along the sequence (middle) dim: a:[bh,ta,hd], b:[bh,tb,hd] →
// out:[bh,ta+tb,hd] with out[:, :ta]=a, out[:, ta:]=b. The device-side KV-cache
// append (M2c): keeps K/V on the GPU and grows by one token per step, removing the
// host round-trip the M2a/M2b host cache paid. One block per bh row.
__global__ void cat_seq_k(const float* a, const float* b, float* out,
int ta_hd, int tb_hd) {
int i = blockIdx.x; // bh row
int o_hd = ta_hd + tb_hd;
const float* ar = a + (long)i * ta_hd;
const float* br = b + (long)i * tb_hd;
float* outr = out + (long)i * o_hd;
for (int j = threadIdx.x; j < ta_hd; j += blockDim.x) outr[j] = ar[j];
for (int j = threadIdx.x; j < tb_hd; j += blockDim.x) outr[ta_hd + j] = br[j];
}
void launch_cat_seq_f32(const float* a, const float* b, float* out,
int bh, int ta_hd, int tb_hd, void* s) {
cat_seq_k<<<bh, 256, 0, (cudaStream_t)s>>>(a, b, out, ta_hd, tb_hd);
}
// Per-row scale: y[r,c] = x[r,c] * s[r]. One block per row. Used by the GRPO
// (M4) policy-gradient backward, where each completion token's row of
// (probs − onehot) is scaled by its own per-token coefficient.
__global__ void scale_rows_k(const float* x, const float* s, float* y,
int rows, int cols) {
int r = blockIdx.x;
float sr = s[r];
for (int c = threadIdx.x; c < cols; c += blockDim.x)
y[r * cols + c] = x[r * cols + c] * sr;
}
void launch_scale_rows_f32(const float* x, const float* s, float* y,
int rows, int cols, void* st) {
int blk = cols < 1024 ? cols : 1024;
if (blk < 32) blk = 32;
scale_rows_k<<<rows, blk, 0, (cudaStream_t)st>>>(x, s, y, rows, cols);
}
__global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim,
float theta, int period) {
int tok = blockIdx.x;
@@ -428,7 +338,7 @@ __global__ void cross_entropy_fwd_k(const float* x, const int* target,
for (int c = threadIdx.x; c < cols; c += blockDim.x) pr[c] *= inv;
if (threadIdx.x == 0) {
int t = target[r];
loss[r] = t < 0 ? 0.0f : -logf(pr[t]);
loss[r] = -logf(pr[t]);
}
}
void launch_cross_entropy_fwd_f32(const float* x, const int* target,
@@ -444,13 +354,8 @@ __global__ void cross_entropy_dx_k(const float* probs, const int* target,
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= rows * cols) return;
int r = i / cols, c = i % cols;
int t = target[r];
if (t < 0) {
dx[i] = 0.0f;
} else {
float g = probs[i] - (c == t ? 1.0f : 0.0f);
dx[i] = g * scale;
}
float g = probs[i] - (c == target[r] ? 1.0f : 0.0f);
dx[i] = g * scale;
}
void launch_cross_entropy_dx_f32(const float* probs, const int* target,
float* dx, int rows, int cols, float scale, void* s) {

629
docs/18-post-training-rl-sft.md
View File

@@ -1,629 +0,0 @@
# Phase: Post-Training Infra — SFT / DPO / Reward Model / GRPO — Design Document
> Status: **DESIGN — decisions locked, pending go-ahead to implement.** Nothing
> implemented yet. This doc proposes the scope, the staged build, the new infra pieces,
> and the correctness gates for a standard post-training stack on top of the xtrain
> training framework. Decisions D1–D4 are resolved (see "Resolved decisions"):
> **DPO → GRPO (reward model optional) · rule-based/verifiable reward · KV-cache decode
> engine built up front · a verifiable task as the optimization/eval target.**
## Goal
Build a **standard, from-scratch post-training infrastructure** — the systems layer that
turns a pretrained base LM into an aligned chat model — and use it to run chat
alignment. The deliverable that matters here is the **infra and the lessons**, not the
end-to-end chat quality (see the project's learning-axis framing). Each stage should
teach exactly one new post-training systems concept and ship with a hard correctness
gate, matching the Phase-1/Phase-2 culture (grad-checks, PyTorch parity, bit-identical
default paths, profile-first).
Concretely we want to be able to answer, with our own code:
- How does **offline preference optimization (DPO)** differ from SFT in the training
loop — what is the reference model, why two forwards, what is the loss?
- How does a **reward model** turn preferences into a scalar signal?
- How does **online RL (GRPO)** actually run — the rollout engine, reward scoring,
group-relative advantage, the clipped policy-gradient update, the KL leash?
- Where are the **memory and throughput** pressure points that make post-training infra
different from pretraining infra (multiple models resident, generation in the loop)?
## Baseline: what already exists vs. what is missing
What the framework already gives us (verified in code, reused as-is):
| capability | where | reuse for post-training |
|---|---|---|
| batched forward → logits `[B*S, vocab]` | `model.rs::forward_batched` | logprob extraction for DPO/RM/GRPO |
| cross-entropy with **ignore-index −100** | `ops.rs::cross_entropy`, `nn.cu` | assistant-only / completion-only masking |
| assistant-only **SFT** (TSV, masked labels) | `data.rs::load_sft_tsv_cached` (commit `fbf4ac2`) | SFT chat baseline = DPO init + reference |
| bf16 mixed precision, fp32 master | `with_compute_dtype` | policy + frozen reference both bf16 compute |
| recompute / flash / grad-accum | `with_recompute` / `with_flash` / `--accum-steps` | bound activation memory with 2–3 models resident |
| DDP (thread + process-per-GPU) | `xtrain-distributed` | data-parallel post-training |
| AdamW + clip + LR sched + checkpoint | `xtrain-optim`, `checkpoint.rs`, `schedule.rs` | unchanged optimizer path |
| single-seq greedy/temperature sampling | `sample.rs::generate` | **slow** rollout fallback (no KV cache) |
What is **missing** and must be built (these are the actual lessons):
1. **Per-sequence completion logprob** — a way to read `Σ log πθ(y_t | x, y_<t)` over the
completion tokens of a sequence. CE gives a *mean* scalar; DPO/GRPO need a *per-sequence
masked sum*. New op or thin wrapper over the CE per-row machinery.
2. **Frozen reference model** held in memory alongside the trainable policy (no grad, no
optimizer), or its logprobs precomputed and cached.
3. **Pairwise preference loss** (DPO) and **Bradley-Terry ranking loss** (RM).
4. **Reward head** — a `[dim,1]` scalar head reading the last non-pad position (RM only).
5. **Rollout / generation engine** — batched autoregressive sampling. Current `generate`
is single-sequence and re-runs the full forward each step (no KV cache). Online RL needs
batched rollouts; a real **KV-cache incremental-decode engine** is the centerpiece infra
build.
6. **GRPO machinery** — group sampling, group-relative advantage, clipped PG loss, KL
penalty, the actor-learner loop.
## The post-training landscape — where the infra lives
```
data models in memory new systems concept
SFT (prompt, answer) policy loss masking (have it)
DPO (prompt, chosen, reject) policy + ref(frozen) dual forward, pairwise logσ loss
RM (prompt, chosen, reject) reward model scalar head, ranking loss
PPO prompts + reward source policy+ref+RM+critic rollout + GAE + clipped PG (4 models)
GRPO prompts + reward source policy+ref(+RM) rollout + group baseline + clipped PG
```
The pedagogical ladder is **SFT → DPO → (RM) → GRPO**. DPO is the cheapest "real" alignment
method (no generation, no reward model, reuses the training loop almost verbatim) and is the
right first rung. GRPO is chosen over PPO as the online-RL rung because it **drops the value
critic** (group-relative advantage replaces the learned baseline) — that removes a whole
model and the GAE machinery while still teaching the complete online-RL loop. PPO is noted
as an optional later extension, not a primary target.
## Proposed scope & sequencing (recommended path)
> ✅ **DECISION D1 (scope/sequencing) — LOCKED: P0 → P1(DPO) → P3(GRPO), P2(reward
> model) optional.** With D3 locked to "KV-cache engine up front", the engine becomes a
> foundational milestone that both DPO pair-generation and GRPO rollouts sit on. Effective
> build order: **P0 → KV-cache decode engine → P1(DPO) → P3(GRPO) → P2(optional)** (see
> "Milestones").
### Stage P0 — SFT chat baseline (light; mostly reuse)
Goal: a clean SFT checkpoint to serve as **both the DPO/GRPO init and the frozen
reference**. With D4 = verifiable task, P0 SFT teaches the **task format** (e.g. arithmetic
prompts → a parseable answer such as `\boxed{N}`) so the model emits checker-readable
completions; the same template is reused by rollout and eval. The current SFT (commit
`fbf4ac2`) already does single-turn assistant-only masking; P0 only adds what alignment
needs:
- a fixed **chat template** (the `User:/Assistant:` + `<|endoftext|>` format already used,
promoted to a documented constant shared by SFT data prep, rollout, and eval),
- optional **multi-turn masking** (supervise every assistant turn, mask user turns),
- optional **sequence packing** (concatenate examples to fill `seq`, reset attention/RoPE
per example — note `forward_batched` already isolates sequences, so packing = careful
index bookkeeping, not new attention code).
Gate: masking unit test (only assistant tokens contribute to loss); packing does not leak
loss across example boundaries. **Hypothesis:** a documented chat template + multi-turn mask
gives a reproducible SFT reference without changing the training numerics for single-turn data
(bit-identical to `fbf4ac2` on single-turn input).
### Stage P1 — DPO (offline preference optimization) ⭐ first real method
New infra:
1. **Preference data — constructed from the verifiable checker (D4).** On a verifiable task
there is no off-the-shelf preference set, so we build pairs: sample several completions
per prompt from the P0 SFT model (using the KV-cache engine built in the prior milestone),
score each with the rule-based checker, take a **correct** completion as `chosen` and an
**incorrect** one as `rejected`. This is a one-time offline data-prep step; DPO training
itself is then static. Tokenize each as `template(prompt) + completion + EOS`; build a
completion mask (prompt = masked).
2. **`seq_logprob(logits, target_ids, mask) → [B]`**: per-sequence sum of
`log softmax(logits)[target]` over masked positions. Implement by reusing the CE per-row
path (CE per-row = `−log πθ(target)`), summing `−per_row` over the mask. Add a grad-checked
op so the backward is exact.
3. **Frozen reference** `πref`: load the SFT checkpoint into a second model in **eval/no-grad**
bf16. Its logprobs are **constants** in the loss. Optimization to teach: **precompute and
cache reference logprobs** once over the dataset → the reference model need not stay
resident during training (one model in memory, like SFT).
4. **DPO loss** (Rafailov et al.): with
`Δ = β[(logπθ(yw|x) − logπref(yw|x)) − (logπθ(yl|x) − logπref(yl|x))]`,
`L = −log σ(Δ)`. Only `πθ` terms carry gradient.
Memory: policy (fp32 master + Adam m/v + bf16 + grads) + reference (bf16 only, or cached
logprobs → zero). Recompute + accum keep activations bounded; 1B fits 32 GB comfortably.
Correctness gates:
- `seq_logprob` finite-difference grad-check (tiny model).
- DPO-loss + grad **PyTorch parity** (the project's standard gate).
- **Degenerate checks**: `πθ == πref` at init ⇒ `Δ = 0`, `L = log 2`, implicit reward 0;
`β → 0` ⇒ gradient → 0.
- **Health metric**: chosen−rejected **reward margin** rises over training; accuracy
(margin > 0) increases. Reported, not just loss (the doc-13 lesson: val/loss alone is not a
sufficient signal).
Application: chat alignment via DPO on English preference pairs. This is the **offline
chat-alignment deliverable**.
### Stage P2 — Reward model (Bradley-Terry) — OPTIONAL
> ✅ **DECISION D2 (reward source) — LOCKED: rule-based / verifiable reward first.** GRPO
> brings up on the deterministic checker; a learned reward model is **deferred/optional** (only
> if we later want general-chat GRPO). So this whole stage is optional and not on the critical
> path.
New infra: a **scalar reward head** (`[dim,1]`) reading the hidden state at the last
non-pad position; **ranking loss** `−log σ(r(x,yw) − r(x,yl))`. Reuses the preference data
and the dual-sequence forward from P1.
Gates: ranking-loss grad-check; held-out **pairwise accuracy** (`r_w > r_l`); a frozen RM
loads/serves the scalar correctly.
### Stage P3 — GRPO (online RL, critic-free) ⭐ the deep infra lesson
This is the centerpiece. It introduces **generation inside the training loop**.
**(a) Rollout / generation engine — built up front (its own milestone).**
> ✅ **DECISION D3 (rollout depth) — LOCKED: build the KV-cache incremental-decode engine
> up front**, as a foundational milestone *before* DPO/GRPO, rather than starting naive. It is
> then the shared substrate for DPO pair-generation and GRPO rollouts. Tradeoff accepted:
> front-loads the single hardest build and delays the first alignment result, in exchange for
> a real generation engine and a clean, isolated infra lesson.
The engine: per-layer **K/V cache**, **single-token incremental forward** (process the prompt
once to fill the cache, then decode one token at a time), **batched ragged decode** (B prompts
× G samples; sequences hit EOS at different lengths → finished-mask / left-padding /
compaction). The current attention assumes a full causal window over `seq`; incremental decode
needs a **decode-time attention path** — query length 1 against cached K/V of length `t`, with
RoPE position = `t`. This reuses the composed SDPA shapes (one-row query), so it can land as a
distinct code path without disturbing the training attention (flash/GQA/composed unchanged).
Hard gate (the centerpiece correctness lesson): **KV-cache decode == full-recompute decode,
token-identical** greedy output — the same byte-/token-identical discipline the project uses
for the xserv export closed loop. A throughput baseline (decode tokens/s, cache-fill vs.
per-token decode) is recorded here, before any rollout optimization (profile-first).
**(b) Reward scoring.** Rule-based verifiable reward first (e.g., exact-match on a synthetic
arithmetic/format task) or RM from P2. Returns a scalar per completion.
**(c) Group-relative advantage.** Sample `G` completions per prompt; advantage
`A_i = (r_i − mean(r_group)) / (std(r_group) + ε)`. No critic, no GAE.
**(d) Clipped policy-gradient loss with KL leash.** Per completion token,
`ρ_t = exp(logπθ_t − logπθ_old_t)` (old = policy at rollout time), token loss
`−min(ρ_t A, clip(ρ_t, 1±ε) A) + βKL(πθ‖πref)`, masked to completion tokens. KL via the k3
estimator.
**(e) Actor-learner loop.** sample prompt batch → rollout G each → score → advantage →
capture `πθ_old` logprobs → K inner epochs of clipped PG updates → repeat. Reference `πref`
fixed throughout.
Memory: policy + reference (+ RM if learned). Each 1B; recompute + accum bound activations.
Throughput note: rollout (generation) will dominate wall-clock — a baseline must be recorded
(tokens/s of generation vs. update) **before** any rollout optimization, per the project's
profile-first rule.
Correctness gates:
- PG-loss finite-diff grad-check.
- **Degenerate checks**: `G = 1` ⇒ advantage 0 ⇒ no PG signal, only KL; `ε → ∞` ⇒ vanilla PG;
`β = 0` ⇒ no KL term.
- (KV-cache decode token-identical to full-recompute is gated in the engine milestone, a
prerequisite of GRPO.)
- **Synthetic RL overfit**: on a tiny verifiable task with a known optimum, mean reward must
rise to the optimum (the RL analogue of T5's "overfit 27/27" — a hard, falsifiable signal
that the loop is correct, independent of fuzzy chat quality).
## Evaluation
- **Offline (DPO/RM)**: reward margin, preference accuracy, KL drift from reference, plus the
fixed chat-prompt generation suite (`scripts/chat_alpha_fixed_prompts.txt`) judged before/
after — reusing and extending the doc-13 recommendation for a generation-based eval harness
(exact-match math, code syntax, stop-token, refusal appropriateness, corruption).
- **Online (GRPO)**: mean reward curve, KL-to-reference, response length, the verifiable-task
pass rate, and the same fixed-prompt suite.
- **Selection by generation eval, not loss** — the recurring doc-13/v11 lesson: lower
post-training loss did not mean better generations.
## Memory & throughput budget (8× RTX 5090, 1.05B model, indicative)
- Params (bf16) ~2.1 GB; fp32 master ~4.2 GB; AdamW m/v ~8.4 GB; grads ~2.1 GB → policy
optimizer state alone ~17 GB before activations. Recompute + grad-accum keep activations
small; this is why post-training reuses the Phase-1/2 memory levers unchanged.
- DPO: + reference (bf16 ~2.1 GB, or 0 if logprobs cached). Fits.
- GRPO: + reference (~2.1 GB) (+ RM ~2.1 GB if learned). Fits; rollout activations are the new
variable. **Generation, not the update, is expected to be the throughput bottleneck** — to be
measured, not assumed.
## Correctness-gate philosophy (unchanged from Phase 1/2)
Every stage ships: (1) a finite-difference grad-check on the new loss/op, (2) PyTorch parity
on loss + grads where applicable, (3) explicit degenerate-case bit/again checks (β→0, G=1,
ε→∞, ref==policy), (4) a falsifiable "it actually learns" signal (reward margin up / synthetic
RL overfit), and (5) **no change to the default training path** when post-training flags are
off. New CUDA kernels (if any, e.g. decode-time attention) get the same fwd/bwd-vs-reference
gates as flash/GQA.
## Risks & tradeoffs
- **Rollout engine is the long pole.** A correct KV-cache incremental-decode path is a real
build (decode-time attention, ragged batch). Mitigation: naive rollout first; KV-cache as an
isolated, separately-gated sub-phase.
- **RL is finicky.** KL leash, advantage normalization, clip range, reward hacking. Mitigation:
synthetic verifiable task with a known optimum as the bring-up gate before any real chat reward.
- **Reward-model noise** can mislead GRPO. Mitigation: rule-based reward first.
- **Tokenizer (KI-4)** — gpt2 50257 vocab is kept for the xserv closed loop; unchanged here.
- **Two/three resident models** raise memory; bounded by recompute/accum and (for DPO) reference
logprob caching.
## Resolved decisions (aligned 2026-06-29)
- **D1 — Scope & sequencing → DPO → GRPO, reward model optional.**
- **D2 — Online-RL reward source → rule-based / verifiable reward first** (RM deferred/optional).
- **D3 — Rollout engine depth → build the KV-cache incremental-decode engine up front** (not
naive-first), as a foundational milestone before DPO/GRPO.
- **D4 — Alignment task / eval target → a verifiable task** (arithmetic/format/GSM8K-style) with
a deterministic exact-match reward, for a clean, falsifiable RL signal.
## Milestones (locked order)
1. **M1 — P0 SFT task baseline.** Chat template + assistant-only masking on the verifiable
task; produces the reference + init checkpoint. Gate: masking unit test; single-turn
bit-identical to `fbf4ac2`.
2. **M2 — KV-cache decode engine** (D3, up front). Per-layer K/V cache + incremental
decode-time attention + batched ragged decode. Gate: **token-identical to full-recompute
greedy**; record decode throughput baseline.
3. **M3 — P1 DPO.** Verifiable-checker pair construction (via M2) → `seq_logprob` op
(grad-check) → DPO loss (PyTorch parity; ref==policy and β→0 degenerate checks) → DPO
training loop → run + reward-margin / preference-accuracy curve.
4. **M4 — P3 GRPO.** Group rollout (M2) + rule-based reward + group-relative advantage +
clipped PG with KL leash. Gate: PG grad-check; G=1/ε→∞/β=0 degenerate checks; **synthetic
verifiable-task RL-overfit** (mean reward → known optimum) → verifiable-task GRPO run.
5. **M5 (optional) — P2 reward model.** Scalar head + ranking loss + pairwise-accuracy gate;
enables GRPO-with-RM for general chat.
> Each milestone is one design+gate cycle; results get appended here (like the run docs) and a
> row in `docs/evolution.md` (algorithm/infra dimensions) when it lands.
## Implementation log
### M1 — SFT task baseline (landed)
The verifiable task and its data pipeline are implemented and verified host-side (no CUDA
needed); the SFT run + eval ran on dash5 (1×5090). **Result: SFT moves answer-format
adherence 0% → 100%, with arithmetic correctness 8% — exactly the intended split (SFT buys
the format; correctness is M3/M4's job).**
**Verifiable task (the spec, in one Rust module — `crates/xtrain-train/src/task.rs`):**
- Two-operand integer arithmetic, ops `+ − ×`; operands `[0,999]` for `+/−`, `[0,99]` for `×`
(modest products); subtraction may be negative. (Ranges enlarged from the first cut to keep
the unique-key space ≫ requested rows — see the saturation guard below.)
- User turn: `What is A op B?`. SFT target: `A op B = \boxed{N}.` — teaches the answer FORMAT;
the checker reads only `\boxed{}`, so arithmetic *correctness* is what M3/M4 improve.
- Rule-based reward: `parse_boxed_answer` (takes the LAST `\boxed{int}`) + `check_answer`
(exact match vs. gold). This is the single shared checker reused by M3 (pair construction)
and M4 (GRPO reward).
- Why this task: trivial deterministic checker, freely scalable difficulty, and it directly
probes the base model's known arithmetic weakness (v12 SFT failed `12 * 13`).
**Data generator (`crates/xtrain-train/src/bin/gen_arith_task.rs`, pure host bin):**
writes `arith_sft.tsv` (`user<TAB>assistant` for `--sft-tsv`), `arith_eval_prompts.txt`
(`greedy_sample --prompts-file` format), and `arith_eval_gold.txt` (parallel gold ints).
Train rows are deduped; eval is held out from train (no leakage). A **saturation guard**
(`unique_space()` + `assert need·5 ≤ space·4`) rejects requests that approach the unique-key
space, since deduped train + disjoint eval near saturation get pathologically slow (or, for
the disjoint-eval loop, never terminate). With the shipped defaults the space is ~2.01M keys,
so a 20 000 + 500 request is a tiny fraction (gen runs in ~0.2 s).
**Scorer (`crates/xtrain-train/src/bin/eval_arith.rs`):** loads a checkpoint, greedily
generates a continuation per held-out prompt, isolates the first answer segment (cut at the
first `<|endoftext|>` then first newline), and reports two signals via the shared checker —
**format** (fraction emitting any `\boxed{int}`) and **correctness** (exact-match vs. gold).
This is the reusable verifiable-eval harness for M3 (DPO) / M4 (GRPO). It uses the *naive*
no-KV-cache sampler (full forward per token), so even 100 prompts is slow — concrete
motivation for M2 (the KV-cache decode engine).
**Masking made testable:** the assistant-only label masking in `load_sft_tsv_cached` was
extracted into a pure `sft_row(prompt_ids, answer_ids)` helper (behavior-preserving — the
single-turn path is bit-identical to `fbf4ac2`).
**Gate (verified locally in `no_cuda` mode):** `cargo test -p xtrain-train --lib` → 9/9 pass,
including `sft_row` masks prompt→`-100` / supervises answer, the SFT-target self-consistency
invariant (always checker-correct over 2000 samples), parser edge cases, and seed determinism.
A 200/50 generation run confirmed clean 2-column TSV, correct gold (incl. negatives), and 0
train/eval leakage.
**Run (dash5, 1×5090, from the v12 1.05B base):**
1. dataset: `gen_arith_task --n 20000 --eval 500 --seed 1 --out-dir <dir>` → 20 000 train +
500 held-out eval, 0 leakage.
2. SFT: `train <tok> <dir>/arith_sft.tsv --sft-tsv --init-ckpt <v12-base.ckpt> --heads 52
--head-dim 32 --kv-heads 13 --layers 22 --ffn 6656 --bf16 --recompute --flash --seq 256
--batch 16 --steps 250 --max-lr 1e-4 --min-lr 1e-5 --ckpt arith_sft_v12.ckpt` → the P0
reference/init checkpoint. Train loss 4.68 → ~0.34, best val 0.386, no OOM, ~4.3K tok/s.
3. eval: `eval_arith <ckpt> <tok> <arch> --prompts-file <dir>/arith_eval_prompts.txt
--gold-file <dir>/arith_eval_gold.txt --max-tokens 32`, base vs. SFT, on 100 held-out prompts.
**M1 result (100 held-out prompts, greedy, max_new 32):**
| checkpoint | format (`\boxed{}`) | correct (exact-match) |
|---------------------|----------------------|-----------------------|
| v12 base (pre-SFT) | 0 / 100 (0%) | 0 / 100 (0%) |
| arith SFT | **100 / 100 (100%)** | 8 / 100 (8%) |
The base model never emits the format — it answers `"I don't know."` / restates the question
and stops. SFT moves format **0% → 100%**: every completion cleanly restates the equation and
boxes an integer (`46 * 80 = \boxed{3380}.`). Correctness is only **8%**: the format is fully
learned but the *arithmetic* is the base model's own weak capability — e.g. it boxes 3380 for
gold 3680, −10 for gold 5; it does get some right (`895 − 353 = \boxed{542}.` ✓). That residual
gap is exactly what the verifiable reward in M3 (DPO) / M4 (GRPO) is built to close.
**Gate met:** format 0% → 100% confirms the assistant-only SFT path is wired end-to-end; the
held-out correct > 0 confirms the checker + eval harness score real matches (not just format).
M1 delivers the format floor + the reusable task spec / checker / eval harness — not arithmetic
skill, which is downstream by design.
### M2a — KV-cache incremental-decode engine (single sequence, landed)
The decode engine (D3, built up front) that replaces the naive sampler — which re-runs the
full forward over the growing prefix every step (O(t²), a fresh autograd graph per token). Two
forward-only primitives + a raw-Tensor per-token block forward, each gated in isolation.
**Primitives (`xtrain-tensor`, both forward-only):**
- `Tensor::rope_at(theta, pos0)` — RoPE at a token's *absolute* position (`pos = pos0 + row`,
no modulo), vs the training `rope` (`pos = row % period`) which is left untouched (new CUDA
kernel `rope_at_k` → no training-path risk). Cached K is stored post-RoPE, so it must match
what the full forward produced at that position. **Gate:** bit-identical to the full-sequence
rope's row `t` (`integration::rope_at_matches_full_rope_row`).
- `Tensor::decode_attention(k, v, scale)` — single-query × cached-K/V SDPA (`[bh,1,hd]` vs
`[bh,t,hd]`, no causal mask: the one query sees all cached keys). Composed from the existing
strided batched GEMM + plain softmax — **no new kernel**. **Gate:** equals the full causal
attention's last query row, max |Δ| 6e-8 (`integration::decode_attention_matches_…`).
**Engine (`xtrain-model/src/decode.rs`, `generate_greedy_cached`):** per-layer K/V cache +
single-token incremental forward. Prefill = the first `prompt.len()` decode steps (one code
path). Mirrors `model::block_forward` at the raw-Tensor level (no autograd tape — inference
needs no grads), pulling weights via the public `params()` stable order (no model-internal
visibility changes). The cache is host-accumulated token-major f32, rebuilt per step — the
honest M2a baseline; M2b moves it device-side + adds batched ragged decode.
**Gate (the M2 centerpiece — token-identical):** KV-cache greedy decode is byte-for-byte the
same token sequence as the naive full-recompute greedy. Verified two ways:
- `xtrain-train/tests/decode_kv.rs` — small GQA model (8 query / 2 kv heads), F32, 24 generated
tokens, exact token-equality. (Unit gate runs F32: a random model's near-uniform logits make
argmax fragile to ~1e-6, so the tightest path is used; the trained model below has peaked
logits → robust.)
- v12 1.05B SFT checkpoint: `eval_arith --cached` produces the **identical** eval outcome to the
naive run (format 100/100, correct 8/100) and byte-identical completions.
**Throughput baseline (v12 1.05B, batch 1, F32, profile-first — measured, not assumed):** the
cache win is **sequence-length-dependent**, which is the honest systems finding here:
| max_new | naive | kv-cache | note |
|---------|-------|----------|------|
| 32 | 108 tok/s | 111 tok/s | ~1.0× — both **launch/overhead-bound** at short seq |
| 128 | 69 tok/s | **133 tok/s** | **~1.9×** — naive's O(t²) recompute starts to bite |
| 256 | **OOM** | 129 tok/s | naive rebuilds the O(seq²) graph every step → OOM |
Cached throughput stays ~constant (O(1)/token compute + constant memory); naive **decays**
(108→69 tok/s, O(t)/token) and eventually **OOMs** (the full autograd graph per step). So at the
short arithmetic-eval lengths the cache is overhead-bound and gives ~nothing — it matters for
**long rollouts** (DPO pair-generation, GRPO completions), exactly where M3/M4 use it. (M2a's
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.
### M4 — GRPO (online RL, critic-free, landed; infra + two honest systems walls)
The centerpiece: generation INSIDE the training loop. Infra built and gated; the run surfaces
two concrete systems findings (the memory long-pole + the rollout long-pole, both flagged in the
design doc's Risks) and the same capability wall as M3.
**Task made learnable first (per the aligned decision "easier task → then M4"):** the v12 SFT
model scores ~8% on the hard task *and* on easy problems — it learned format, not arithmetic. So
the easy task (operands ≤20, ops `+ − ×`) was re-SFT'd from the v12 base → **held-out 18.7%**
(100% format), a baseline with reward variance for GRPO. Note: even easy arithmetic plateaus at
~19% held-out (250 vs 600 SFT steps identical) — a 1B web-text model does not generalize the
add/sub algorithm from ~550 examples; it memorizes train (982 total problems, 550 seen).
**New op (`xtrain-autodiff`, reuses the CE kernel + one new primitive):**
- `clipped_pg_loss(logits, target, logp_old, logp_ref, A, ε, β)` — per completion token
`ρ_t = exp(logπθ_t − logp_old_t)`, `L = −mean min(ρA, clip(ρ,1±ε)A) + β·mean KL` (k3), masked
to completion tokens. Backward reuses `(probs − onehot)` + `scale_rows` (a new ~5-line per-row
scale kernel — the per-token coefficient varies, which CE-backward's single scalar can't
express). **Gate:** grad-check the active PG path + the A=0 (KL-only) path; degenerate value
checks ε→∞ ⇒ vanilla PG, β=0 ⇒ no KL.
**Loop (`train_grpo`):** per step — sample B prompts, roll out G completions each, score (reward
0/1), group-relative advantage `A=(r−mean)/(std+ε)` (no critic; all-correct/all-wrong groups
skipped — zero advantage), capture `logπθ_old`/`logπref` per token, K inner clipped-PG epochs.
Rollout uses the M2 KV-cache engine with **temperature sampling** (added in M4): single-row
`[1,vocab]` logits per step vs the naive sampler's `[seq,vocab]`.
**Systems wall #1 — memory (the design doc's "two/three resident models"):** KL-leash GRPO needs
policy + frozen reference, two 1.05B fp32-master models + AdamW m/v ≈ 21 GB fixed + training
activations → unreliably OOMs on a 32 GB 5090 (fragmentation tips it over). To get a completing
run, `β=0` (pure PG) drops the reference model (−4.2 GB). So the *principled* KL-leash version is
memory-bound at this model size on this hardware — a real, reported constraint, not a bug.
**Systems wall #2 — rollout (the design doc's "rollout is the long pole"):** the naive sampler's
growing `[seq,vocab]` allocations fragment the caching allocator over a long rollout → OOM. The
cached temperature rollout (single-row logits) is lighter; but single-sequence cached decode is
slow (the M2a host-round-trip), so rollout still dominates wall-clock (~16 s/step at G=6·B=6).
Batched ragged decode (M2b) is the real fix and is deferred to where it is load-bearing.
**Result (easy task, β=0, G=6·B=6, 40 steps, lr 5e-7; 150 held-out, vs SFT 28/150 = 18.7%):**
mean rollout reward fluctuates ~0.58–0.81 (noisy, inflated by train-set overlap in the sampled
problems); **format stays 100/100** (no collapse even without the KL leash, at this gentle lr);
**held-out 30/150 = 20.0%** — `+1.3 pp`, within the ~3% std-error of 150 prompts, i.e.
**statistically flat**, the same wall as M3 DPO.
**The consistent M3+M4 lesson:** on a task where the base model lacks the underlying capability,
**neither offline preference optimization (DPO) nor online RL (GRPO) moves held-out correctness**
— each optimizes its objective (margin / reward) on the *training distribution* it can reach
(here inflated by memorization), but cannot install a *generalizable* algorithm the model never
had. RL reinforces what the model already does; it does not teach arithmetic. Gate met for M4 =
the infra is correct (PG/KL grad-checks + degenerate checks, the loop runs, reward signal + KL
leash wired, format held); the held-out flatness + the two memory/throughput walls are the
reported findings. The honest end-state of the post-training arc: **a complete, correctness-gated
SFT → KV-cache → DPO → GRPO stack** — the infrastructure learned in full, with measured, honest
limits on what alignment can do for a capability the base model lacks.
### M2b — batched KV-cache decode (landed; completes the M2 engine, fixes the rollout long-pole)
Built after M4 (where the rollout long-pole bit hardest): decode the **G samples of one prompt in
lockstep** — one forward per step over the whole group → G× fewer kernel launches, the deferred
fix from M2a.
**One new primitive:** `rope_pos(x, positions[])` — RoPE with a *per-row* absolute position (new
forward-only kernel), since the G batched rows share one decode position (M2a's `rope_at` does
`pos0 + row`, wrong for a batch at a single position). **Gate:** bit-identical to the full rope
for positions `[0..n]`, and to `rope_at(P)` per row for a uniform `P`.
**Engine (`generate_cached_batch`):** `BatchKVCache` carries a G dimension (`[T, G·num_kv, hd]`
host-accumulated → `[G·num_kv, T, hd]`); the batched `decode_step` threads G through embed /
projections / QK-norm / `rope_pos` / cache. Two M2a pieces drop in unchanged: `decode_attention`
is already batch-agnostic (`bh = G·nh`), and `repeat_kv(nh, batch=G)` broadcasts per group. No
finished-mask (all G generate `max_new`; the caller cuts at EOS) and no ragged-length prompts yet
— both perf-only follow-ups.
**Gate (token-identical):** all G **greedy** rows are byte-identical to the single-sequence decode
(`tests/decode_batch.rs`, 8 query / 2 kv heads → exercises the `repeat_kv` batching) — pins that
G-way batching indexes each sequence's K/V with no cross-row contamination.
**Throughput (v12 1.05B, G=6·B=6, easy task, rollout wired into `train_grpo`):** ~8.5 s/step vs
~14–16 s/step for the single-seq cached rollout — **~1.7×**, rollout-inclusive. Short of the full
G× because (a) the per-token-logp forwards + the PG update also cost, and (b) the M2a per-layer
**host round-trip** is still there (now G× the data in one transfer, not removed). The full
device-side cache (no host round-trip) is the remaining decode-engine optimization. Batching also
**stabilises memory**: one batched forward per step vs G separate allocations that fragmented the
caching allocator (the M4 OOM). So M2b closes the decode-engine milestone (M2a single-seq + M2b
batched) and turns the rollout long-pole from "OOM/unbounded" into a bounded ~1.7× win — measured,
with the device-cache as the named next lever.
### M2c — device-side KV cache (landed; the bottleneck moved, a profile-first finding)
The named M2b follow-up: keep K/V on the GPU (`[bh,T,hd]`, an `Option<Tensor>` per layer) and
grow it by one token per step via a new `cat_seq` kernel (concat along the seq dim) — removing the
M2a/M2b per-layer **host round-trip** (`to_cpu`/`from_slice`/re-upload) *and* the `transpose_3d01`.
Both single-seq and batched decode refactored to it (cleaner than the host `Vec` + rebuild).
**Gates hold:** `cat_seq == host concat`; `decode_kv` single-seq + `decode_batch` G-way both still
**token-identical**; GQA training path unaffected.
**The finding (why this is a measure-first lesson, not a speedup story):** removing the host
round-trip buys **~10%** on *pure* single-seq decode (133 → 147 tok/s @128) but **does not move the
GRPO step** (~8.5 s/step, unchanged). Because after M2b batching, the rollout is no longer the
step's bottleneck — the per-sample **`per_token_logp` captures** (2 forwards/sample) and the
**PG-update** forwards+backwards (`model.forward`, full-sequence, per sample) now dominate. So the
long pole **shifted** from the rollout to the training-side forwards (cf. T11/T17/M2a: profile
before optimizing — the bottleneck you fixed is not the one that remains). The device cache is
still a real, correctness-gated improvement (cleaner code, less PCIe, ~10% decode); the honest
headline is that the *next* decode lever is **ragged batched prefill of the per-sample forwards**,
not the cache. The M2 decode engine is now M2a (single-seq) + M2b (batched) + M2c (device cache),
all token-identical-gated; the post-training stack remains complete with its bottleneck mapped.
### M2d — batch the GRPO training-side forwards (landed; the lever M2c named, + a decomposition correction)
M2c named the next lever: **ragged batched prefill of the per-sample training-side forwards**. Those
forwards are the two phases that, per step, run one single-sequence `forward` per sample: the
`per_token_logp` **captures** (logπ_old policy + logπ_ref reference) and the inner **clipped-PG**
forward/backwards. M2d packs all `N = B·G` ragged samples of a step into ONE `forward_batched`.
**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 sits 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 are garbage but masked out (`target = -100`). This
is exactly why training engines pad-and-mask rather than run ragged. Two new pieces:
- `per_token_logp_batched` (`crates/xtrain-train/src/grpo_batch.rs`): right-pad → one
`forward_batched(batch = N)` → slice each sample's logπ back to its real length.
- `ops::clipped_pg_loss_batched` (`crates/xtrain-autodiff/src/ops.rs`): like the per-sample
`clipped_pg_loss`, but takes **per-row** `advantage[t]` (the owning sample's `A`) and **per-row**
`weight[t]` (the full normaliser; the caller passes `1/(N·n_s)`). It does NOT compute its own
`1/n_tokens`, so folding `weight = 1/(N·n_s)` reproduces the looped `Σ_s (1/N)(1/n_s)…`
**bit-for-bit** (the per-row CE backward is row-local). A `--micro` knob packs in chunks to bound
the `[chunk·Lmax, vocab]` logits memory; the weight uses the GLOBAL `N`, so chunked
grad-accumulation is exact. Both `train_grpo` and the bench call these shared helpers.
**Correctness gates (exact, not bf16-noisy):**
- `xtrain-model::forward_batched_ragged_matches_looped` — forward_batched on right-padded ragged
sequences == per-sequence single-seq forward on the real rows, **max|Δlogit| = 3.7e-7 (fp32) and
0.0 (bf16)**, both composed + flash. Pins "right-pad is free".
- `xtrain-autodiff::clipped_pg_loss_batched_matches_looped` — batched op == looped
`Σ_s (1/N)·clipped_pg_loss_s`, **loss Δ=1.5e-8, grad max|Δ|=7.5e-9 (f32)**.
Composed, these prove the batched GRPO step == the looped step. End-to-end: a short SFT (v12 base,
150 steps, arith) → `train_grpo` 12 steps runs clean — **no OOM** (1B master + AdamW + batched
activations fit with `micro=16`), mean-reward rises, the batched inner executes.
**Throughput (bench `bin/bench_grpo_batch`, v12 1.05B, N=48 ragged, micro=16, β=0, weight-independent):**
| phase (per step) | looped (single-seq) | batched (M2d) | speedup |
|-------------------------|---------------------|---------------|---------|
| capture `per_token_logp`| 622 ms | 71 ms | 8.7× |
| inner clipped-PG fwd+bwd| 1907 ms | 208 ms | 9.2× |
| **training forwards** | **2526 ms** | **280 ms** | **9.0×**|
**The decomposition correction (the honest finding).** M2c claimed "the per-sample training forwards
now dominate the step." The clean per-component bench falsifies the strong form: the training
forwards were **~2.5 s of the ~8.5 s step (~30%)** — substantial and worth the 9× win, but the
**rollout (`generate_cached_batch`, ~6 s) was always the larger share.** After M2d cuts the training
forwards to ~0.28 s, the step is **~95% rollout** — the long pole has swung back to the rollout. So
M2d removes the training-forward overhang (a real, exactly-gated 9× on its component), and re-confirms
the same measure-first lesson one more time: the next **step-level** lever is **full B×G rollout
batching** — today only the `G` samples of each prompt decode in lockstep (M2b); the `B` prompts are
still sequential. M2d closes the "ragged batched per-sample forwards" lever M2c named; the post-
training stack stays complete, now with the step decomposition measured, not asserted.

42
docs/evolution.md
View File

@@ -33,9 +33,9 @@
---
## 二、Scaling runs（v0–v10）—— 主要动「模型架构」与「数据集」
## 二、Scaling runs（v0–v8）—— 主要动「模型架构」与「数据集」
架构始终是 **Qwen3-style**（RoPE + RMSNorm + QK-norm + SwiGLU，gpt2 50257 词表），逐版放大 dim/层/头（v8 起首次拨容量轴到 dim1024，v9 进入 dim1280+真 GQA 双轴点，v10 固定架构只补数据轴）；其余维度逐版变化如下：
架构始终是 **Qwen3-style**（RoPE + RMSNorm + QK-norm + SwiGLU，gpt2 50257 词表），逐版放大 dim/层/头（v8 起首次拨容量轴到 dim1024）；其余维度逐版变化如下：
| ver | 模型架构（dim/层/头·hd · 核心/总参） | 数据集（语料 · 实训 token · epoch） | 算法/精度 | Infra（GPU · 吞吐） | 结果（val · 备注） |
|---|---|---|---|---|---|
@@ -48,27 +48,22 @@
| v6 | dim768/18L（同 v4/v5） | **FineWeb-edu** 真实网页 · 2.29B · 1.02ep | bf16 | 8 GPU · 218K | val **3.07**：⚠️**FineWeb 留出集，与 v0–v5 不可比**（真实网页熵高，~3.0 是预期）；判据=采样质量+transfer。第一版脱离 TinyStories，**语言种类质变**（小故事→真实说明文）；transfer→TinyStories val 2.75(v5 native 1.11)，纯通用数据对窄分布有代价；val 末步仍单调降=未饱和 |
| v7 | dim768/18L（同 v4/v5/v6） | **同 v6 的 FineWeb-edu 子集**（非新数据）· 3.28B · **1.45ep** | bf16 | 8 GPU · 218K | val **3.01**（与 v6 可比）：⚠️**同子集多 epoch 近天花板**——唯一变量=epoch(1.02→1.45)，多喂 ~1B token val 仅 ↓0.05 且 ~step44000 后走平、采样无质变。与 v5 的 TinyStories 数据量饱和同类（重复老数据边际薄）；真·更多数据要**新 shards** |
| v8 | **dim1024**/18L/**32h** · **226M/329M**（+78% 容量，ffn 2730） | **同 v6/v7 的 FineWeb-edu 子集**（非新数据）· 2.36B · **1.05ep** | bf16 **+ 激活重计算(T13)** | 8 GPU · 129K（重算税） | val **2.98**（与 v6/v7 可比）：⭐**容量轴 A/B——容量有用**：唯一变量=dim768→dim1024，同 ~1ep v6 3.07→**2.98**（↓0.085），且 v8(1.05ep) < v7(1.45ep 更多老数据) 3.01 ⇒ 放大容量 > 重复老数据 ⇒ v6/v7 部分 capacity-limited。⚠️但增益仅 ~3%、val 末步**仍在降未饱和** ⇒ **单轴(数据/容量)单步都已 ~3%/lever = 全面边际递减，要双轴一起 scale(Chinchilla)** |
| v9 | **dim1280**/18L/**40h/10kv GQA** · **357M/486M**（ffn 4096） | **FineWeb-edu 扩展 shards 000-009** · **6.01B** · **~1.00ep** | bf16 + recompute + **flash + grad-accum + true GQA** | 8 GPU · **78.6K**（21.25h） | val **2.8854**（与 v6-v8 可比）：✅**双轴 Chinchilla 点有效**——容量从 v8 226M→357M，同时数据从 2.255B 子集→6.013B token，best val 比 v8 再降 **0.0947 (~3.2%)**。采样写真实说明文更稳一些，但 greedy 重复仍明显；收益仍是稳健增量而非质变 |
| v10 | **同 v9** | **FineWeb-edu 扩展 shards 000-010** · **6.765B** · **~1.00ep** | bf16 + recompute + flash + grad-accum + true GQA | 8 GPU · **79.0K**（23.86h） | moving-tail val **2.8816**；固定 eval v1 上 v9 **2.9278**→v10 **2.8814**。结论：补 shard010 对新分布有效，但只补数据轴不解决 greedy 重复；后续应固定 eval set，并优先试更大模型+长 context |
> 实训 token = steps×batch×seq（非数据集大小）。v0–v5 的 val 是同一 1M-token TinyStories 留出集。v6 起换 FineWeb-edu，
> 且 v9/v10 追加新 shards 会移动默认 tail-heldout；严格横比改用 fixed eval v1（shard010 tail 1M）：
> v6/v7/v8/v9/v10 = **3.2328 / 3.1850 / 3.1515 / 2.9278 / 2.8814**。
> 实训 token = steps×batch×seq（非数据集大小）。val 同一 1M-token TinyStories 留出集（v0–v5 可比；v6 起换 FineWeb-edu 留出集，分布不同、与 v0–v5 不可比；v6/v7/v8 同一 FineWeb 留出集、三版彼此可比 3.07/3.01/2.98）。
---
## 三、各维度的累积演进（轴向看一条线怎么走的）
- **算法**：手写 autograd(tape)+扇出累加 → AdamW/LR-sched/grad-clip → +QK-norm(Qwen3) → batched forward → bf16 混合精度(fp32 master) → 激活重计算(T13) → 融合 flash-attention(T14，online softmax + flash 式 bwd) → 梯度累积(T16，复用 tape SUM，等效大 batch 而显存随 micro) → dropout(T18，counter-based 设备 RNG + inverted scaling，train/eval 切换)。
- **模型架构**：固定 Qwen3-style；dim **32→256→384→512→768→1024→1280**（v8 首拨容量轴，v9 进入 dim1280）；核心参数 **41K→357M**（总 3.26M→486M）。+QK-norm(T9，Qwen3 兼容) → **真 GQA(T15，`num_kv_heads<num_heads`，repeat_kv broadcast + 组内梯度求和；默认=nh→MHA 逐位回归；v9 用 40 query / 10 kv）**——架构补齐到现代 LLM 标配（MHA/GQA/MQA 一条 `num_kv_heads` 轴），两条 SDPA(composed/flash) 共用同一 broadcast，导出真 `num_key_value_heads` 且 xserv 闭环。
- **模型架构**：固定 Qwen3-style；dim **32→256→384→512→768→1024**（v8 首拨容量轴，头数 24→32）；核心参数 **41K→226M**（总 3.26M→329M）。+QK-norm(T9，Qwen3 兼容) → **真 GQA(T15，`num_kv_heads<num_heads`，repeat_kv broadcast + 组内梯度求和；默认=nh→MHA 逐位回归)**——架构补齐到现代 LLM 标配（MHA/GQA/MQA 一条 `num_kv_heads` 轴），两条 SDPA(composed/flash) 共用同一 broadcast，导出真 `num_key_value_heads` 且 xserv 闭环。
- **Infra**：单卡 fp32 → cuBLAS/GPU-optim(T7) → NCCL DDP(T8) → batched forward(T10) → caching allocator(T11) → bf16(T12) → 激活重计算(T13，解锁 dim1024) → flash-attention(T14，不物化 N×N，attention 显存收益随 seq 增长) → 梯度累积(T16，DDP 只在累积边界通信，显存随 micro 不随有效 batch) → process-per-GPU(T17，torchrun 式独立进程/CUDA context，复用 T8 train_rank 零改动)。吞吐 **3.3K→217K tok/s**（dim768 bf16），dim1024+重算 ~129K（重算税）；MFU **0.4%→17%**（每次提升都对应一块 perf 基建，详见 known-issues + MFU 分析）。T13/T14/T16 是三条**显存杠杆**（重计算压激活峰值、flash 不物化 N×N attention scores、梯度累积解耦有效 batch 与激活显存），可叠加放大有效 batch。**T17 实测=负结果记账**：process-per-GPU 在本尺度对吞吐**中性**（thread ~5.27× vs proc ~5.31×@8，差<1% 噪声），8 卡全 95–99% util ⇒ 残留非线性是 NCCL/PCIe 通信墙、**非**单 context 串行——把 KI-5/T11 doc 长挂的「process-per-GPU 是残留串行的解」猜想实测钉死推翻（方法论同 T11 证伪「分桶 all-reduce」）。
- **数据集**：TinyStories 3MB 切片 → 全量 TinyStories（epoch 0.01→5.33，**至饱和**）→ **v6 毕业到 FineWeb-edu 真实网页**（2.255B 语料，1.02ep）→ **v7 同子集多 epoch（1.45ep，近顶）→ v8 同子集换大模型**（dim1024，1.05ep）→ **v9 扩新 FineWeb shards 到 6.013B token 并同步放大模型** → **v10 补 shard010 到 6.765B token（只拨数据轴）**。tokenizer 全程 gpt2 BPE（复用 xserv-tokenizer；保闭环优先，KI-4 接受）。
- **数据集**：TinyStories 3MB 切片 → 全量 TinyStories（epoch 0.01→5.33，**至饱和**）→ **v6 毕业到 FineWeb-edu 真实网页**（2.255B 语料，1.02ep）→ **v7 同子集多 epoch（1.45ep，近顶）→ v8 同子集换大模型**（dim1024，1.05ep）。tokenizer 全程 gpt2 BPE（复用 xserv-tokenizer；v6 刻意不换 tokenizer 以隔离「数据来源」变量，KI-4 留后续版本）。
- **v5→v6 数据轴的质变**：v0–v5 都吃合成幼儿故事（TinyStories，低熵、词汇受控），v5 证明同尺寸模型在它上面已饱和；v6 第一版换成**真实教育类网页文本**（FineWeb-edu），语言种类发生质变——采样从「只会写小故事」变成「能写历史/科学/说明文」。
- ⚠️ **同子集多 epoch 也有天花板（v6→v7）**：v6 的 FineWeb val 才训 1.02ep、末步仍单调降，曾被读作「还没喂够」；v7 把**同一 2.255B 子集**喂到 1.45ep（多 ~1B token），FineWeb val 仅 ↓0.05（3.07→3.01）且 ~step44000 后走平、采样无质变 ⇒ **该子集在 dim768 已近天花板**。这与 v5 的 TinyStories 数据量饱和是**同一类现象**：**「重复喂老数据」边际都薄，无论是 v5 的同语料多 epoch 还是 v7 的同子集多 epoch**。真正抬天花板的是 v6「换更广的新语料」那一步——**杠杆在「更多样的新 token」，不在「同数据多读几遍」**。后续要继续降 val，必须补**新 FineWeb shards**（更多样、不重复），不是同子集加 epoch。
- ⚠️ **val 可比性**：v0–v5 的 val 是同一 TinyStories 1M 留出集（彼此可比）；**v6 起换 FineWeb-edu 留出集，分布不同、val 不能和 v0–v5（~1.1）比大小**——真实网页熵高，~3.0 是预期而非回退。v9/v10 追加 shards 后默认 tail-heldout 会移动，不能再只看 moving-tail best。为后续建立 fixed eval v1（shard010 tail 1M）：v6/v7/v8/v9/v10 = **3.2328 / 3.1850 / 3.1515 / 2.9278 / 2.8814**。
- ⚠️ **val 可比性**：v0–v5 的 val 是同一 TinyStories 1M 留出集（彼此可比）；**v6 起换 FineWeb-edu 留出集，分布不同、val 不能和 v0–v5（~1.1）比大小**——真实网页熵高，~3.0 是预期而非回退；**v6/v7/v8 同一 FineWeb 留出集、三版彼此可比**（3.07→3.01→2.98）。v6 的判据还有采样质量 + **transfer eval**（v6→TinyStories val 2.75 vs v5 native 1.11，量化「纯通用数据对窄分布的代价」）。
- ⭐ **容量轴有用，但也只有 ~3%（v8）**：v6/v7 在 dim768 上「吃不动更多数据」，v8 用最干净的 A/B 回答了「是数据见够还是容量不够」——**冻结数据子集、纯把 dim768→dim1024（core 127M→226M，+78%）**，同 ~1 epoch 下 FineWeb val **3.07→2.98（↓0.085）**，且 v8（1.05ep）还低于 v7（1.45ep 更多老数据）的 3.01。⇒ **容量有用，v6/v7 部分是 capacity-limited（不全是数据见够）**；放大容量比「给小模型多喂老数据」更值。**但增益只有 ~3%**，与数据轴单步杠杆同量级。
- ✅ **双轴一起 scale 有效（v9）**：v9 把 v8 的提案落地：模型 core 226M→357M，数据 2.255B 子集→6.013B token（实训 6.012B），best FineWeb val **2.9801→2.8854**，再降 **0.0947 (~3.2%)**。这确认 Chinchilla 式双轴方向正确；但收益仍是 ~3% 级稳健增量，greedy 重复仍在，说明小尺度下“更好 val”尚未完全转化成肉眼质变。
- 📌 **只补数据轴边际有限（v10）**：v10 保持 v9 架构，仅补 shard010 到 6.765B token。fixed eval v1 上 v9 2.9278→v10 2.8814，说明新 shard 分布被学到；但 moving-tail best 只从 2.8854→2.8816，且 greedy 复读不变。下一步更值得改模型/context，而不是继续一片片补数据。
- 🧭 **元结论：单轴单步都已 ~3%/lever = 全面边际递减，要双轴一起 scale（Chinchilla 小尺度复现）**：把三条轴并起来看——数据量轴（v5/v7 同子集多 epoch，饱和，~1.6–5%/步）、数据广度轴（v6 换语料，是一次性换分布红利）、容量轴（v8，有用但 ~3%）——**到 v8，任何单轴的单步杠杆都收敛到 ~3%/lever**。而 v8 容量 +78% 却只配同样的 2.36B token、val 末步仍在降 ⇒ 数据立刻成新瓶颈。⇒ **要继续进步，容量与数据必须匹配地一起 scale，而不是单独猛拨一根轴**——这正是 Chinchilla 在这个 toy 尺度上的复现。
## 三·五、Phase 2 系统栈深度综合（T14–T18 五条特性按四维收束）
@@ -86,29 +81,6 @@ scaling 科学线（v0–v8）收官后，项目重启回到本职「学训练
> 📌 两条 integration 发现（非回归，pre-existing，记账）：① **DDP 三个测试并行会争 2 卡 deadlock** → 文档/测试用 `--test-threads=1`（或标 serial）跑。② **fresh-train md5 run-to-run 不定**——反向 atomicAdd 归约序非确定 → 有效的确定性闸门是**导出（export）重确定性**（同 ckpt 重导 safetensors md5 逐位一致），**不是** fresh-train 复现。
## 三·六、Phase 3 后训练栈（SFT → KV-cache → DPO → GRPO，详见 [18-post-training-rl-sft.md](18-post-training-rl-sft.md)）
Phase 1/2 把**预训练全栈**学完后，Phase 3 转向**后训练 infra**（对齐方向）。锁定路线 DPO→GRPO（reward model 可选）、**rule-based 可验证 reward 优先**、**KV-cache 增量解码引擎前置自建**、任务取**可验证算术**（确定性 exact-match，给 RL 干净可证伪信号）。里程碑 M1（SFT baseline）→ M2（KV-cache 解码引擎，token-identical 闸门）→ M3（DPO）→ M4（GRPO）→ M5（可选 RM）。按维度落点：
- **算法**：后训练损失族——SFT（assistant-only masking，已有）→ DPO（`seq_logprob` 算子 + Bradley-Terry/σ(Δ) 偏好损失，frozen reference）→ GRPO（group-relative advantage，无 critic + clipped PG + KL leash）。每条沿用 Phase 1/2 闸门规矩：新损失/算子有限差分 grad-check + PyTorch parity + 退化检查（β→0 / G=1 / ε→∞ / ref==policy）+ 一条可证伪「真在学」信号（reward margin↑ / 合成 RL overfit）。
- **Infra**：**KV-cache 增量解码引擎（M2，前置）**是这一阶段的硬核——per-layer K/V cache + 单 token 增量 forward（prompt 灌一次 cache 后逐 token 解码）+ ragged 批量解码。硬闸门 = **解码逐 token 等价于全重算 greedy**（同 xserv 导出闭环的逐位纪律），并先记解码吞吐 baseline（profile-first）。它是 DPO 造对 + GRPO rollout 的共享底座。
- **数据集**：可验证任务自带数据生成器——两操作数整数算术（`+ − ×`），rule-based checker 读 `\boxed{}` 做 exact-match，是 M1 SFT 数据 + M3 造对 + M4 GRPO reward 的单一共享 spec。
- **模型架构**：复用 v12 1.05B 基座，不动架构。
**M1（SFT task baseline，已落地）**：可验证算术任务 + 数据生成器 + 评分器一套，host-side 9/9 单测过（masking、SFT-target 自洽 2000 样、parser 边界、种子确定性）。dash5 单卡从 v12 基座 SFT（loss 4.68→~0.34，best val 0.386）。**100 留出题 eval：格式 `\boxed{}` 习得率 base 0% → SFT 100%；算术正确率 8%。**——SFT 只买**格式**（0%→100% 干净落地），算术正确性是 base 模型本身弱项（如 `46*80` 框成 3380），正是 M3/M4 的可验证 reward 要去补的残差。一条诚实账：M1 用的是**朴素无 KV-cache 采样器**（每 token 全量 forward），100 题已经很慢——这正是 M2 解码引擎前置的动机。
**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 同源的诚实账：跑通+闸门齐全,负结果如实记。
**M4（GRPO,在线 critic-free RL,已落地 + 两道诚实系统墙 + 一致负结果）**：新算子 `clipped_pg_loss`（per-token ρ + clip + k3 KL,反向用新增 `scale_rows` per-row 缩放 kernel;grad-check active+A=0 路径 + 退化 ε→∞ vanilla/β=0 无KL)。环 `train_grpo`:采 B prompt × rollout G → checker reward 0/1 → group-relative advantage `(r−mean)/(std+ε)`(无 critic,全对/全错组跳过)→ 存 πθ_old/πref per-token → K 内层 clipped-PG。rollout 用 **M2 引擎 + 新加的 temperature 采样**（单行 logits 比 naive `[seq,vocab]` 轻)。**先把任务改简单**:v12 SFT 在硬/易题都 ~8-9%(只会格式不会算术)→ 在 easy(操作数≤20)上从 v12 base 重训 SFT → held-out **18.7%**;但 250/600 步同样 18.7% = 1B web-text 模型从 ~550 例**不泛化加减法、只记 train**。**两道系统墙(设计文档 Risks 预言)**:① 显存——KL-leash 要 policy+reference 两个 1B fp32-master+Adam≈21GB,加激活在 32GB 5090 上不稳定 OOM → 只能 `β=0`(去掉 reference)跑完;② rollout 长杆——naive 采样增长序列撑碎 allocator,cached 采样更轻但单序列慢仍主导墙钟(~16s/step)。**结果**(easy, β=0, G6·B6, 40步, lr5e-7;150 留出 vs SFT 18.7%):reward 噪声 ~0.58-0.81(被 train 重叠抬),**format 100/100 不崩**(温和 lr 下 β=0 也没崩),**held-out 20.0%**(+1.3pp,~3% 标准误内 = 统计持平)。**M3+M4 一致教训**:模型缺底层能力时,离线偏好(DPO)和在线 RL(GRPO)**都不抬 held-out**——各自在能触及的训练分布上优化目标(被记忆抬高),装不进可泛化算法;**RL 强化模型已会的,不教算术**。**后训练弧诚实终态 = 一套完整、闸门齐全的 SFT → KV-cache → DPO → GRPO 栈**,infra 学全,并测得对齐对"base 缺失能力"能做什么的诚实边界。
**M2b（批量 KV-cache 解码,已落地,补全 M2 引擎 + 修 rollout 长杆）**：M4 后补的 rollout 长杆修复——一个 prompt 的 **G 个样本同步解码**(每步一次 forward 跑整组 → G× 更少 kernel 启动)。一个新原语 `rope_pos`(逐 row 绝对位置 kernel,G 行共享一个解码位置;闸门 = `[0..n]` 逐位等于全 rope、统一 P 逐行等于 `rope_at(P)`,bit-identical)。引擎 `generate_cached_batch`:`BatchKVCache` 带 G 维,批量 `decode_step` 把 G 贯穿 embed/proj/QK-norm/`rope_pos`/cache;**M2a 两件零改动复用**——`decode_attention` 本就 batch-agnostic(bh=G·nh)、`repeat_kv(nh,batch=G)` 按组广播。闸门 = G 个贪心行逐字节等于单序列(`tests/decode_batch.rs`,8q/2kv 头练 repeat_kv 批量)。**吞吐**(v12, G6·B6, 接进 train_grpo):**~8.5s/step vs 单序列 ~14-16s/step ≈ 1.7×**(rollout-inclusive;未到满 G× 因 per_token_logp + PG 更新也占时间、M2a 主机往返还在);且**显存更稳**(一次批量 forward vs G 次分配撑碎 allocator 的 M4 OOM)。⇒ M2 引擎闭环(M2a 单序列 + M2b 批量),rollout 长杆从"OOM/无界"变成有界 ~1.7× 收益,device 端 cache 是点名的下一杠杆。
**M2c（device 端 KV cache,已落地,瓶颈转移的 profile-first 发现）**：K/V 留 device 为 `[bh,T,hd]`(每层 `Option<Tensor>`),每步用新 `cat_seq` kernel(沿 seq 拼接)append 一个 token——去掉 M2a/M2b 每层**主机往返** + `transpose_3d01`,单序列和批量都重构到它(比 host Vec+rebuild 干净)。闸门全保:`cat_seq`==host concat、decode_kv 单序列 + decode_batch 批量仍 **token-identical**、GQA 训练路径不受影响。**发现(measure-first 的点,不是加速故事)**:去掉主机往返让**纯单序列解码 +10%**(133→147 tok/s@128),但 **GRPO step 不动**(~8.5s/step)——因为 M2b 批量化后 rollout 已不是 step 瓶颈,**per-sample `per_token_logp` 捕获(2×/样本)+ PG 更新 forward/backward(全序列 `model.forward`)成了主导**。长杆从 rollout **转移**到训练侧 forward(同 T11/T17/M2a:profile 后再动手——你修的不是剩下的瓶颈)。device cache 仍是真实、闸门齐全的改进(更干净、少 PCIe、解码 +10%),但下一杠杆是 **per-sample forward 的 ragged 批量**而非 cache。M2 引擎现 = M2a(单序列)+ M2b(批量)+ M2c(device cache),全 token-identical-gated;后训练栈完整、瓶颈已测绘。
**M2d（批量 GRPO 训练侧 forward,已落地,M2c 点名的杠杆 + 一处 decomposition 纠正）**：M2c 点名的下一杠杆——把每步 `N=B·G` 个 ragged 样本的训练侧 forward(`per_token_logp` 捕获 + inner clipped-PG fwd/bwd)打包进**一次 `forward_batched`**。**使能性质 = causal 下右 padding 免费**:真 completion 行位置早于尾部 pad,causal 禁止前向 attend,故真行 logits 与单序列 forward **逐位相同**,pad 行垃圾被 `target=-100` 屏蔽——这正是训练引擎 pad-and-mask 而非跑 ragged 的原因。两件新东西:`per_token_logp_batched`(右 pad → 一次 `forward_batched(N)` → 按真长切片)、`ops::clipped_pg_loss_batched`(per-row `advantage[t]` + per-row `weight[t]`,caller 传 `1/(N·n_s)`,op 不再自算 `1/n_tokens` → 折进 weight 即与 looped `Σ_s (1/N)(1/n_s)…` **逐位等价**;`--micro` 分块界定 `[chunk·Lmax,vocab]` logits 显存,weight 用全局 N 故分块梯度累积精确)。**两道精确闸门**:`forward_batched_ragged_matches_looped`(右 pad 批量 forward == 单序列,fp32 max|Δ|=3.7e-7、bf16 **0.0**,composed+flash)+ `clipped_pg_loss_batched_matches_looped`(批量 op == looped,loss Δ=1.5e-8/grad 7.5e-9,f32),复合即证端到端等价;端到端短 SFT→`train_grpo` 12 步**不 OOM**(1B master+AdamW+批量激活 micro=16 容得下)、批量 inner 执行。**吞吐(bench,v12 1.05B,N=48,micro16,权重无关)**:capture 622→71ms(8.7×)、inner 1907→208ms(9.2×)、**训练侧 forward 合计 2526→280ms(9.0×)**。**Decomposition 纠正(诚实发现)**:M2c 说"训练侧 forward 主导 step",干净分量 bench 证伪强形式——训练侧 forward 是 **~8.5s step 里的 ~2.5s(~30%)**,可观、值这 9×,但 **rollout(`generate_cached_batch` ~6s)一直是更大头**;M2d 把训练侧砍到 ~0.28s 后,step **~95% 是 rollout**,长杆又摆回 rollout。⇒ M2d 拔掉训练侧 forward 这块 overhang(分量级精确 9×),再次印证 measure-first:**step 级下一杠杆 = 全 B×G rollout 批量**(今天只有每 prompt 的 G 同步、B 个 prompt 仍串行)。后训练栈保持完整,step decomposition 现为**实测**而非断言。
## 四、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 用上）。

152
docs/runs/09-v9-fineweb-edu-dim1280-gqa.md
View File

@@ -1,152 +0,0 @@
# Scaling Run v9: Chinchilla 双轴 — dim1280/18L true GQA(core 356.9M) + FineWeb-edu 6.01B token + Phase-2 stack — Design Document
## Goal
v8 给出的元结论是：单独拨容量轴有用，但只有约 3% 的边际；单独重复旧数据也只有约 1.6% 的边际。要继续明显超过
v8，必须把 **模型容量 + 新 token** 一起放大，而不是只拨一根轴。
v9 就是这个双轴点：
1. **模型轴**：dim1024/core 226M -> **dim1280/core 356.9M**，同时启用真 GQA（40 query heads / 10 kv heads）。
2. **数据轴**：v6-v8 的 2.255B FineWeb 子集 -> **6.013B token**，追加了新 FineWeb-edu shards 003-009。
3. **系统栈**：使用 Phase-2 现代路径：`--flash + --accum-steps + bf16 + recompute + DDP`。dropout 设为 0，按标准预训练。
> v9 的 val 仍是 FineWeb-edu 分布，不能和 v0-v5 的 TinyStories val 直接比。注意：v9 扩展 cache 后默认
> tail-heldout 已经从 v6-v8 的旧 tail 移到新 shards 末尾；严格横比后续以 fixed eval v1 为准。
## Data
| 项 | 值 |
|----|----|
| 来源 | FineWeb-edu `sample/10BT`，原 shards 000-002 + 新 shards 003-009 |
| token cache | `data/fineweb-edu.txt.u16.bin` |
| 总 token | **6,013,639,492** |
| held-out val | 末尾 **1,000,000** token |
| train corpus | 6,012,639,492 token |
| 训练消费 token | **6,012,600,320** = 91745 steps x effective batch 256 x seq 256 |
| epoch | ~1.00 |
P3-DATA 目标本来是约 7B token；shard 010 下载 `curl rc=18` 中断，所以最终停在 6.01B。对 core 356.9M 来说，
D/N 约 **16.8 token/param**，低于理想 Chinchilla 20，但已经远高于 v8 的约 10.4，是一个干净的双轴 scale 点。
## Architecture
| 项 | v8 | **v9** |
|----|----|----|
| dim | 1024 | **1280** |
| layers | 18 | 18 |
| query heads x head_dim | 32 x 32 | **40 x 32** |
| kv heads | 32 (MHA) | **10 (true GQA, group=4)** |
| ffn | 2730 | **4096** |
| core params | 226.50M | **356.89M** |
| total params | 329.42M | **485.55M** |
| export tensors | 201 | **201** |
`config.json` writes real `num_key_value_heads = 10`, so xserv loads v9 as true GQA rather than MHA.
## Training
| 项 | 值 |
|----|----|
| optimizer | hand-written AdamW, wd=0.1 |
| schedule | warmup -> cosine, max_lr 6e-4 -> min_lr 6e-5 |
| grad clip | global norm 1.0 |
| steps | **91745** |
| effective global batch | **256** (`--batch 128 --accum-steps 2`) |
| seq_len | 256 |
| precision | bf16 mixed precision, fp32 master |
| memory stack | activation recompute + flash-attention + gradient accumulation |
| world size | 8 x RTX 5090 |
| wall clock | **21h15m** |
| steady throughput | **~78.6K tok/s** |
| peak observed memory | ~17GB / GPU |
Command:
```sh
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 cargo run -p xtrain-distributed --release --bin train_ddp -- \
/opt/wjh/models/gpt2/tokenizer.json data/fineweb-edu.txt \
--heads 40 --head-dim 32 --kv-heads 10 --layers 18 --ffn 4096 \
--steps 91745 --batch 128 --accum-steps 2 --seq 256 \
--max-lr 6e-4 --min-lr 6e-5 --val-tokens 1000000 --eval-every 1000 \
--eval-batches 64 --bf16 --recompute --flash --dropout 0.0 \
--ckpt /dashscope-tmp/wjh/xtrain_v9.ckpt
```
## Results
- train loss: **11.1550 -> 2.9340**
- first val: step 1000 = **5.1517**
- best val: step 91000 = **2.8854**
- final val: step 91745 = **2.8873**
- exit code: **0**
FineWeb val curve milestones:
| step | 1000 | 10000 | 20000 | 30000 | 40000 | 50000 | 60000 | 70000 | 80000 | 90000 | 91000 | final |
|------|------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|
| val | 5.1517 | 3.4820 | 3.2953 | 3.2026 | 3.1422 | 3.0844 | 3.0148 | 2.9616 | 2.9160 | 2.8915 | **2.8854** | 2.8873 |
The curve kept improving into the last 1K-step window, then the final eval bounced slightly from 2.8854 to 2.8873. This is close to
the floor for this run, but not a clear overfit failure.
## Comparison
| | v6 | v7 | v8 | **v9** |
|---|---|---|---|---|
| model | dim768/core127M | dim768/core127M | dim1024/core226M | **dim1280/core357M + GQA** |
| data | 2.29B | 3.28B same subset | 2.36B same subset | **6.01B expanded shards** |
| best val | 3.0652 | 3.0149 | 2.9801 | **2.8854** |
On the run-local moving tail, v9 beats v8 by **0.0947** val loss (~3.2% relative), essentially the same size as the
v6->v8 capacity gain but now on top of it. A later fixed eval v1 check still supports the same direction
(v8 3.1515 -> v9 2.9278 on shard010-tail holdout), while making the moving-tail caveat explicit. This confirms
the v8 prediction: **双轴 scale 有效**. It is still an incremental gain, not a qualitative jump.
## Samples
xserv greedy samples (`--max-tokens 60`) are more coherent than the v8 examples on some prompts, but repetition remains:
```text
[The history of] the United States is the story of the people, the places, and the events that have shaped the nation...
[In science,] the term "scientific method" is used to describe the process of gathering information and testing it...
[The most important] thing is to be aware of the symptoms and to seek medical attention...
[Water is] a natural resource that is essential for human life...
```
The model writes real explanatory English and the domain mix is FineWeb-like. Greedy decoding still falls into repeated clauses on
some prompts (`scientific method`, symptoms, and earlier fixed prompts), so the val gain is more visible in the metric than in a
dramatic sample-quality leap.
## xserv validation
Registry path:
```text
/opt/wjh/projects/tiny-models/v9-fineweb-edu-dim1280-gqa
```
Files:
- `config.json`
- `model.safetensors` (BF16, 201 tensors, 927MB)
- `tokenizer.json`
- `xtrain.ckpt` (fp32 master checkpoint, 1.9GB)
- `RUN.md`
xserv loads v9 as:
```text
Model: qwen3, layers=18, hidden=1280, heads=40/10 kv, vocab=50257
Loaded 201 tensors
Ready (KV cache, dtype=bf16).
```
Token-match check against xtrain greedy (`max-tokens 40`):
- `Once upon a time`: xtrain and xserv matched through the checked continuation.
- `One day`: diverged after "large, dark," (`very tall man` vs `metallic object`) from BF16 greedy tie sensitivity.
- `The little`: same repetitive pattern, with a short BF16 path divergence.
This is the same class of BF16-vs-f32 greedy drift seen in v8; the important integration result is that xserv successfully loads
true GQA (`kv_heads=10 < heads=40`) and generates from the exported weights.

200
docs/runs/10-v10-fineweb-edu-dim1280-gqa-data6765.md
View File

@@ -1,200 +0,0 @@
# Scaling Run v10: Data-axis follow-up — dim1280/18L true GQA + FineWeb-edu 6.765B token — Design Document
## Goal
v9 证明了双轴 scale（更大模型 + 更多新 token）有效：best val 从 v8 的 2.9801 降到 2.8854。
但 v9 的数据量只有 6.013B token，D/N 约 16.8，低于 Chinchilla 经验里的 20。v10 的目标很窄：
1. **只补数据轴**：补上 v9 中断的 FineWeb-edu shard010，把 cache 从 6.013B 推到 6.765B。
2. **架构不变**：完全复用 v9 dim1280 / 18L / 40q-10kv GQA / ffn4096。
3. **验证边际**：看 D/N 从 16.8 到 18.95 是否还能显著降低 val。
## Data
| 项 | 值 |
|----|----|
| 来源 | FineWeb-edu `sample/10BT`，shards 000-010 |
| token cache | `data/fineweb-edu.txt.u16.bin` |
| 总 token | **6,765,333,808** |
| held-out val | 末尾 **1,000,000** token |
| train corpus | 6,764,333,808 token |
| 训练消费 token | **6,764,298,240** = 103215 steps x effective batch 256 x seq 256 |
| epoch | ~1.00 |
Important caveat: xtrain 当前训练入口用“全 cache 的末尾 1M token”做 held-out。追加 shard010 后，v10 的 val tail
和 v9 的 val tail 不再是同一个切片。因此 v9 原报告的 2.8854 与 v10 原报告的 2.8816 不能被当作严格同一
验证集上的横比。
为了解决这个问题，本轮创建了固定 eval set：
```text
/dashscope-tmp/wjh/xtrain_fixed_eval_v1/fineweb-fixed-eval-v1.txt.u16.bin
```
它包含 shard010 末尾 11M token；前 10M token 只是为了复用现有 `split_tail(val_tokens=1M)`，真正 eval 的是最后
1M token。该 fixed eval v1 对 v6-v9 都是未见数据；对 v10 也是训练时 held-out。
## Architecture
v10 与 v9 完全相同：
| 项 | 值 |
|----|----|
| dim | 1280 |
| layers | 18 |
| query heads x head_dim | 40 x 32 |
| kv heads | 10 (true GQA, group=4) |
| ffn | 4096 |
| core params | 356.89M |
| total params | 485.55M |
| export tensors | 201 |
## Training
| 项 | 值 |
|----|----|
| optimizer | hand-written AdamW, wd=0.1 |
| schedule | warmup -> cosine, max_lr 6e-4 -> min_lr 6e-5 |
| grad clip | global norm 1.0 |
| steps | **103215** |
| effective global batch | **256** (`--batch 128 --accum-steps 2`) |
| seq_len | 256 |
| precision | bf16 mixed precision, fp32 master |
| memory stack | activation recompute + flash-attention + gradient accumulation |
| world size | 8 x RTX 5090 |
| wall clock | **23h51m** |
| steady throughput | **~79.0K tok/s** |
| peak observed memory | ~17GB / GPU |
Command:
```sh
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 cargo run -p xtrain-distributed --release --bin train_ddp -- \
/opt/wjh/models/gpt2/tokenizer.json data/fineweb-edu.txt \
--heads 40 --head-dim 32 --kv-heads 10 --layers 18 --ffn 4096 \
--steps 103215 --batch 128 --accum-steps 2 --seq 256 \
--max-lr 6e-4 --min-lr 6e-5 --val-tokens 1000000 --eval-every 1000 \
--eval-batches 64 --bf16 --recompute --flash --dropout 0.0 \
--ckpt /dashscope-tmp/wjh/xtrain_v10.ckpt
```
## Results
- train loss: **11.1575 -> 2.9000**
- first val: step 999 = **5.3048**
- best val: step 103214 = **2.8816**
- final val: step 103214 = **2.8816**
- exit code: **0**
FineWeb moving-tail val milestones:
| step | 999 | 9999 | 19999 | 29999 | 39999 | 49999 | 59999 | 69999 | 79999 | 89999 | 99999 | final |
|------|-----|------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|
| val | 5.3048 | 3.5622 | 3.3282 | 3.2450 | 3.1886 | 3.1342 | 3.0714 | 3.0202 | 2.9724 | 2.9236 | 2.8950 | **2.8816** |
The curve still improved at the final eval. There is no overfit signal in this run.
## Fixed Eval V1
Fixed eval v1 (`shard010 tail 1M`, seq256, 64 eval batches):
| version | fixed eval v1 |
|---------|---------------|
| v6 | 3.2328 |
| v7 | 3.1850 |
| v8 | 3.1515 |
| v9 | 2.9278 |
| **v10** | **2.8814** |
This is the cleanest cross-version result in the v10 round. It says:
- v9's double-axis gain transfers to a shard010 holdout: v8 3.1515 -> v9 2.9278.
- v10 further improves on the new shard010 distribution: v9 2.9278 -> v10 2.8814.
- The apparent v9 moving-tail 2.8854 -> v10 moving-tail 2.8816 delta is tiny and not strict apples-to-apples.
## Decoding
Greedy decoding still repeats. Fixed prompts from xtrain:
```text
[Once upon a time] there was a king who had a daughter. She was beautiful and beautiful...
[The little] The little boy was a little boy. The little boy was a little boy...
[One day] I was walking down the street and I saw a man with a dog...
```
Temperature 0.8 is more varied and less immediately looped, but coherence remains weak:
```text
[Once upon a time] I was a kid who did not go to the beach to swim...
[The little] ones are not as loud as the adults...
[One day] I was on the edge of the water, and I saw something I had never seen before...
```
xserv loads the exported v10 true-GQA weights and generates FineWeb-like explanatory prose, but repeated sentence frames remain:
```text
[The history of] the city of San Francisco is a story of the growth of the city...
[In science,] the term "observation" is used to describe the act of observing something...
[Water is] the most important element in the human body...
```
Conclusion: decoding remains a separate bottleneck. The current xtrain sampler only supports greedy and temperature sampling; top-p and
repetition penalty exist in xserv's chat path, but not in the raw xtrain sampler or `xserv-cli` path used for weight validation. A clean
next step is to add a raw generation tool with `temperature/top-p/repetition-penalty` so decoding experiments do not depend on chat
templates.
## xserv Validation
Registry path:
```text
/opt/wjh/projects/tiny-models/v10-fineweb-edu-dim1280-gqa-data6765
```
Files:
- `config.json`
- `model.safetensors` (BF16, 201 tensors, 927MB)
- `tokenizer.json`
- `xtrain.ckpt` (fp32 master checkpoint, 1.9GB)
xserv loads v10 as true GQA:
```text
Model: qwen3, layers=18, hidden=1280, heads=40/10 kv, vocab=50257
Loaded 201 tensors
Ready (KV cache, dtype=bf16).
```
## v11 Feasibility: Bigger Model + Longer Context
A v11 smoke test prioritized the user's chosen direction: larger model plus longer context.
Candidate:
| item | value |
|------|-------|
| dim / layers | 1536 / 20 |
| heads / kv_heads | 48 / 12 |
| ffn | 6144 |
| core / total params | 684.26M / 838.65M |
| stack | bf16 + recompute + flash + accum + 8 GPU DDP |
Smoke results:
| seq | batch / accum | effective batch | peak mem | tok/s | result |
|-----|---------------|-----------------|----------|-------|--------|
| 512 | 64 / 4 | 256 | **30530 MiB** | **44.7K** | 50 steps OK |
| 1024 | 32 / 8 | 256 | **30530 MiB** | **31.0K** | 20 steps OK |
Both fit, but the memory margin is thin on 32GB RTX 5090. Expected one-epoch wall clock on 6.76B tokens:
- seq512: roughly **42h**
- seq1024: roughly **61h**
Recommendation: make v11 a controlled run, not a blind launch. Use fixed eval v1, keep data fixed, and choose either:
1. **v11a practical**: dim1536/20L, seq512, batch64/accum4. Faster, still doubles context over v10.
2. **v11b long-context**: dim1536/20L, seq1024, batch32/accum8. More aligned with "long context", but ~2.5 days and tight memory.
For scientific clarity, v11 should not append more data before training; use the current 6.765B train cache while preserving fixed eval v1.

251
docs/runs/12-v12-1b-longctx-chat-alpha.md
View File

@@ -1,251 +0,0 @@
# Scaling Run v12: 1B-class long-context base → chat-alpha-v2 — Design Document
## Goal
v11 proved that a larger `dim1536/20L` model can train at `seq1024` on dash5, and it improved the fixed-eval-data-v1, long-context (`seq1024`) score to **2.7467**. It also proved the current bottleneck: greedy generation still repeats, and broad real-data SFT on top of v11 regressed chat quality despite lower SFT validation loss.
v12 therefore separates the next phase into two gates:
1. **Base gate**: train a stronger English base model around 1B total params with `seq1024`, using the existing FineWeb-edu 6.765B-token cache and fixed eval data v1.
2. **Chat gate**: only after the base gate is healthy, run assistant-only English SFT and judge it with fixed prompt generation, not SFT loss alone.
Success means the model is serviceable enough for a small chat-alpha: stable base loss, lower fixed eval than v11, less repetitive fixed generation, and SFT that improves instruction behavior without destroying arithmetic/refusal/debug prompts.
## Baseline: What v11 Taught Us
| item | v11 |
|------|-----|
| arch | dim1536 / 20L / 48q-12kv GQA / ffn6144 |
| params | 684.26M core / 838.65M total |
| data | FineWeb-edu 6.765B token, 1 epoch |
| context | seq1024 |
| throughput | ~30.96K tok/s on 8 x RTX 5090 |
| fixed eval data v1, seq1024 | **2.7467** |
| issue | greedy repetition remains; direct real SFT regressed generation quality |
SFT result from v11:
| model | train result | generation result |
|-------|--------------|-------------------|
| `v11-chat-alpha-sft-v2-anchor` | synthetic assistant-only anchor | current best narrow chat-alpha |
| `v11-chat-alpha-real-sft-v1` | SFT val 1.4272 | bad hallucination, math failure |
| `v11-chat-alpha-real-mix-v1` | SFT val 2.0543 | better than direct real-SFT, still worse than anchor |
Conclusion: SFT data quality matters, but v11's base is still too weak for broad real SFT to become a general chat model.
## Architecture
v12 target: slightly above 1B total params while staying close to the proven v11 shape and keeping GQA group size 4.
| item | value |
|------|-------|
| dim | **1664** |
| layers | **22** |
| query heads x head_dim | **52 x 32** |
| kv heads | **13** |
| GQA group | 4 |
| ffn | **6656** |
| core params | **883.4M** |
| embed + lm_head | **167.3M** |
| total params | **1.0506B** |
Why this shape:
- It is a controlled step from v11 rather than a new architecture family.
- `52/13` preserves true GQA with group 4.
- Total params are near the requested 1B target.
- `dim1664` is less aggressive than `dim1792/22L` and has a better chance to fit `seq1024` on 32GB 5090s.
## Data
Base pretraining stays English-oriented and uses the current token cache. Pass the `.txt` stem to xtrain; `Corpus::load_cached` appends `.u16.bin` internally.
```text
/opt/wjh/projects/xtrain/data/fineweb-edu.txt
cache = /opt/wjh/projects/xtrain/data/fineweb-edu.txt.u16.bin
tokens = 6,765,333,808
```
Training uses the last 1M tokens as moving-tail validation. Every cross-version v12 claim must also run fixed eval data v1 with the long-context `seq1024` setting, matching the v11 `eval_v11_seq1024.log` score of **2.7467**. This is distinct from the older v10 table that used the same fixed eval data with `seq256`.
```text
/dashscope-tmp/wjh/xtrain_fixed_eval_v1/fineweb-fixed-eval-v1.txt
cache = /dashscope-tmp/wjh/xtrain_fixed_eval_v1/fineweb-fixed-eval-v1.txt.u16.bin
```
No new FineWeb shards are added in this phase. The experiment is model/context scale, not another data-axis change.
## Training Plan
Primary v12 run:
| item | value |
|------|-------|
| world | 8 x RTX 5090 on dash5 |
| precision | bf16 mixed precision, fp32 master |
| memory stack | recompute + flash + grad accumulation |
| seq | **1024** |
| micro global batch | **16** (2 sequences/rank) |
| accum | **15** |
| effective global batch | **240** |
| tokens/step | **245,760** |
| full steps | **27,524** |
| max_lr → min_lr | **4e-4 → 4e-5** |
| eval | moving-tail 1M every 500 steps; fixed eval data v1 at seq1024 after checkpoints |
| smoke throughput | **~24.5K tok/s** |
| estimated full wall clock | **~76-78h** |
The reduced micro-batch is intentional: v11 `seq1024` with global batch 32 already sat near the 5090 memory limit. v12 has larger weights; an initial `batch24/accum10` smoke OOMed after step 0, while `batch16/accum15` passed a 10-step smoke at ~29.4GB/GPU and preserved the same 245,760 tokens/step.
Command wrapper:
```sh
scripts/run_v12_phase.sh start-pilot
scripts/run_v12_phase.sh start-full
scripts/run_v12_phase.sh status
scripts/run_v12_phase.sh eval-fixed
scripts/run_v12_phase.sh export
scripts/run_v12_phase.sh sample
```
## Gates
### Gate 0: build and smoke
Run:
```sh
scripts/run_v12_phase.sh smoke
```
Pass criteria:
- no CUDA OOM
- no NaN loss
- first 30 steps decrease from initialization
- peak memory leaves enough margin for eval
### Gate 1: pilot
Run:
```sh
scripts/run_v12_phase.sh start-pilot
```
Default pilot is 300 steps with held-out eval every 100 steps.
Pass criteria:
- train loss decreases smoothly
- grad norm does not spike persistently
- moving-tail eval is finite and improving
- checkpoint can be reloaded by `eval-fixed`
### Gate 2: full base
Run only after the pilot passes:
```sh
scripts/run_v12_phase.sh start-full
```
Pass criteria:
- fixed eval data v1 at `seq1024` beats v11's **2.7467**
- generation samples improve or at least do not regress on repetition
- checkpoint exports and xserv loads the true GQA config
### Gate 3: chat-alpha SFT
After a healthy v12 base:
1. Use assistant-only SFT (`--sft-tsv`) with English-only data.
2. Start from narrow anchors first, then mix in Smol-SmolTalk.
3. Judge with fixed generation prompts before calling it useful.
The primary high-quality source remains `HuggingFaceTB/smol-smoltalk` filtered to English single-turn examples, with local anchors preserved to keep deterministic behavior.
## Evaluation
Base metrics:
- moving-tail val during training
- fixed eval data v1 at `seq1024`
- xtrain fixed prompt samples from `scripts/chat_alpha_fixed_prompts.txt`
- xserv exported-model smoke
Chat metrics:
- fixed prompt answers for SFT explanation, SFT data provenance, arithmetic, refusal, repetition-debug checklist, summary, and simple code generation
- compare against `v11-chat-alpha-sft-v2-anchor`
- reject models that lower SFT validation loss but hallucinate more in fixed prompts
## Artifacts
Expected paths:
```text
/dashscope-tmp/wjh/xtrain_v12/
/dashscope-tmp/wjh/xtrain_v12/xtrain_v12_pilot.ckpt
/dashscope-tmp/wjh/xtrain_v12/xtrain_v12.ckpt
/opt/wjh/projects/tiny-models/v12-fineweb-edu-1b-longctx
```
## Results
### Gate 0/1: smoke + pilot
- `batch24/accum10` smoke OOMed after step 0.
- `batch16/accum15` smoke passed 10 steps: train loss **11.2347 -> 7.9459**, ~24.5K tok/s, ~29.4GB/GPU.
- 300-step pilot passed: train loss **11.2296 -> 5.4832**, val **6.5810 -> 5.9642 -> 5.5888**, exit code 0.
- Pilot checkpoint reload matched final val: fixed eval data v1 at seq1024 = **5.5891**.
- Fixed chat prompts still repeat heavily, as expected for a 300-step base; use them as a regression baseline, not as chat quality.
### Gate 2: full base
Full run completed on dash5:
| item | result |
|------|--------|
| wall clock | **81h01m** |
| throughput | **~24.55K tok/s** |
| train loss | **11.2294 -> 2.6696** |
| moving-tail best val | **2.7411** |
| moving-tail final val | **2.7412** |
| fixed eval data v1, seq1024 reload | **2.7410** |
| exit code | **0** |
Validation milestones:
| step | 499 | 999 | 1499 | 1999 | 2499 | 21999 | 23999 | 25999 | 26999 | 27499 | final |
|------|-----|-----|------|------|------|-------|-------|-------|-------|-------|-------|
| val | 5.3029 | 4.4079 | 3.9287 | 3.6964 | 3.5555 | 2.7805 | 2.7637 | 2.7468 | 2.7443 | **2.7411** | 2.7412 |
Compared with v11's fixed eval data v1 at seq1024 (**2.7467**), v12 reaches **2.7410** after reload. This is a real but very small gain
(~0.006 absolute), despite the parameter increase from 838.65M to 1.0506B total and the slower 24.55K tok/s throughput. The result says the
larger 1B-class base is viable and marginally better, but this scale step did not produce a qualitative base-model jump.
Generation:
- Raw FineWeb-style prompts are better than the pilot checkpoint and can produce plausible explanatory prose.
- Greedy repetition remains visible, especially on story-like prompts.
- Chat prompts are not reliable without SFT: SFT data provenance is hallucinated, arithmetic still fails, and the model repeats template-like text.
- xserv loads the export correctly as true GQA: `layers=22, hidden=1664, heads=52/13 kv`.
Exported model:
```text
/opt/wjh/projects/tiny-models/v12-fineweb-edu-1b-longctx
```
Files:
- `config.json`
- `model.safetensors` (2.0GB)
- `tokenizer.json`
- `xtrain.ckpt` (4.0GB)
Conclusion: v12 passes the base gate and is a better SFT starting point than v11 by metric, but the gain is narrow. The next step should be
assistant-only chat SFT from v12 with conservative anchors first, then a small Smol-SmolTalk mix. Do not expect the base checkpoint itself to
serve as a usable chat model.

180
docs/runs/13-v12-chat-sft-quality.md
View File

@@ -1,180 +0,0 @@
# v12 Chat SFT Quality Check
Date: 2026-06-29
## Goal
Turn the completed v12 1.05B base checkpoint into a usable chat-alpha model with
SFT, then judge whether it is stable enough to call a high-quality chat model.
Base checkpoint:
```text
/dashscope-tmp/wjh/xtrain_v12/xtrain_v12.ckpt
```
Architecture:
```text
dim=1664 layers=22 heads=52 kv_heads=13 head_dim=32 ffn=6656
total params=1.0506B
```
## Stage A: Synthetic SFT
Data:
```text
/dashscope-tmp/wjh/xtrain_sft_alpha_v2/chat_alpha_v2_sft.tsv
211,257 examples, about 14.96M SFT tokens
```
Run:
```text
/dashscope-tmp/wjh/xtrain_sft_v12_alpha_v2/chat_alpha_v12_v2.ckpt
```
Metrics:
```text
train loss: 3.5730 -> 0.0426
eval: step39 0.1078, step79 0.0582, step119 0.0466, step159 0.0423,
step199 0.0403, step239 0.0390, step279 0.0389, step319 0.0378
best/final val loss: 0.0378
```
Export:
```text
/opt/wjh/projects/tiny-models/v12-chat-alpha-sft-v2
```
Quality notes:
- Learns the User/Assistant format and usually stops correctly.
- Too narrow and template-heavy.
- Fails basic math and code prompts in fixed greedy evaluation.
## Stage B: Anchor SFT
Data:
```text
/dashscope-tmp/wjh/xtrain_sft_alpha_v2_anchor/chat_alpha_v2_anchor.tsv
32,020 examples, about 1.73M SFT tokens
```
Run:
```text
/dashscope-tmp/wjh/xtrain_sft_v12_anchor/chat_alpha_v12_anchor.ckpt
```
Metrics:
```text
train loss: 1.7777 -> 0.1165
eval: step19 0.3447, step39 0.1449, step59 0.1217, step79 0.1158
best/final val loss: 0.1158
```
Export:
```text
/opt/wjh/projects/tiny-models/v12-chat-alpha-sft-v2-anchor
```
Generation artifacts:
```text
/dashscope-tmp/wjh/xtrain_sft_v12_anchor/generation/anchor_xserv_greedy.txt
/dashscope-tmp/wjh/xtrain_sft_v12_anchor/generation/anchor_diagnostic_greedy.txt
```
Quality notes:
- Better project-context answers and summaries than synthetic-only.
- Still unreliable on basic multiplication, yes/no facts, translation, and code.
- Overuses "cannot verify" style answers outside appropriate uncertainty cases.
## Stage C: Real-Mix Repair
Data:
```text
/dashscope-tmp/wjh/xtrain_sft_real_mix_v1/smol_smoltalk_real_mix.tsv
96,287 examples, about 25.3M SFT tokens
```
Run:
```text
/dashscope-tmp/wjh/xtrain_sft_v12_real_mix_repair/chat_alpha_v12_real_mix_repair.ckpt
```
Training setup:
```text
init=/dashscope-tmp/wjh/xtrain_sft_v12_anchor/chat_alpha_v12_anchor.ckpt
steps=200
seq=512
batch=32
accum=8
effective batch=256
lr=1e-6 -> 2e-7
```
Metrics:
```text
train loss: 2.7391 -> 2.0384
eval: step49 2.1964, step99 2.0383, step149 1.9801, step199 1.9570
best/final val loss: 1.9570
```
Export:
```text
/opt/wjh/projects/tiny-models/v12-chat-alpha-real-mix-repair
```
Generation artifacts:
```text
/dashscope-tmp/wjh/xtrain_sft_v12_real_mix_repair/generation/real_mix_repair_xserv_greedy.txt
/dashscope-tmp/wjh/xtrain_sft_v12_real_mix_repair/generation/real_mix_repair_diagnostic_greedy.txt
/dashscope-tmp/wjh/xtrain_sft_v12_real_mix_repair/generation/real_mix_repair_diagnostic_greedy_reppenalty1.txt
```
Quality notes:
- Loss improved cleanly and the model kept chat formatting.
- Fixed prompt math `17% of 240` improved in the standard suite.
- General diagnostic math still fails, e.g. `12 * 13`.
- Code generation remains unusable for simple Python function prompts.
- Some outputs contain corrupted or off-topic fragments.
- Reducing repeat penalty from 1.15 to 1.0 did not fix the failures.
## Verdict
The SFT pipeline works, and v12 can be turned into a chat-shaped model that follows
the prompt format and stops correctly. However, none of the three SFT variants is a
stable high-quality chat model yet.
The limiting issue is no longer infrastructure. It is data and objective quality:
the current synthetic/anchor data is too narrow, while the current real-mix data
adds breadth but also noisy or low-quality behavior. Validation loss alone is not a
sufficient selection signal for chat quality.
## Recommended Next Step
Build a smaller, higher-precision SFT curriculum before another large run:
1. Keep the anchor data, but reduce over-refusal templates.
2. Add verified small instruction sets for math, code, translation, summarization,
and closed-book common facts.
3. Add an automatic fixed-prompt eval harness that scores exact-match math, simple
code syntax, refusal appropriateness, stop-token behavior, and corruption.
4. Train a short curriculum from the v12 base or v12 anchor checkpoint, then pick
by generation eval rather than SFT loss alone.

17
docs/runs/README.md
View File

@@ -22,10 +22,6 @@ val loss 一栏给的是各版**各自训练 run 报告的 best val**（held-out
**v8 改测容量轴**：同 v6/v7 子集、纯把 dim768→dim1024（core 127M→226M），FineWeb val 3.07/3.01→**2.98** ⇒
**容量有用**（v6/v7 部分 capacity-limited）；但增益仅 ~3%、val 末步仍在降未饱和 ⇒ **到 v8，数据轴与容量轴的
单步杠杆都收敛到 ~3%/lever = 全面边际递减，要双轴一起 scale**（Chinchilla，详见 [08-v8](08-v8-fineweb-edu-dim1024.md)）。
**v9 兑现双轴**：dim1024→dim1280（core 226M→357M）并把 FineWeb token 从 2.255B 子集扩到 6.013B，
best val **2.8854**，相比 v8 再降 0.0947（~3.2%）。结论：双轴 scale 有效，但仍是稳健增量而非质变。
**v10 只补数据轴**：同 v9 架构，只补 shard010 到 6.765B token，moving-tail best/final val **2.8816**。
注意追加 shard 会移动 held-out tail；固定 eval v1 上 v6→v10 为 **3.2328 / 3.1850 / 3.1515 / 2.9278 / 2.8814**。
⚠️ **v6 起换了保留集（语料）**：v0–v5 的 val 都是 **TinyStories** 1M 留出集（彼此可比）；v6 换成纯
**FineWeb-edu**（真实网页文本），它的 val（3.07）是**另一把尺子上的另一个分布**，**不能**和 v0–v5 的
@@ -43,15 +39,18 @@ best val **2.8854**，相比 v8 再降 0.0947（~3.2%）。结论：双轴 scale
| [v6-fineweb-edu-dim768](06-v6-fineweb-edu-dim768.md) | **FineWeb-edu** 真实网页 (2.255B 语料) | ~2.29B | ~1.02 | 768 / 18 / 24·32 / 2048 (**同 v4/v5**) | 127.43M | 204.63M | **3.0652** ⚠️*(FineWeb val,与上不可比)* | **第一版脱离 TinyStories**，唯一变量=数据来源 + 8 卡 DDP bf16；~1.9h/8 卡 ~218K tok/s。**val 是另一分布**(真实网页熵高,~3.0 是预期非回退)，判据=采样质量+transfer。FineWeb val 末步仍单调降=未饱和；**transfer**: v6→TinyStories val **2.75**(v5 native 1.11)，纯通用数据对窄分布有代价。采样: v6 写真实说明文 vs v5 一律掉进小故事 |
| [v7-fineweb-edu-dim768](07-v7-fineweb-edu-dim768.md) | **同 v6 的 2.255B FineWeb-edu 子集**(非新数据) | ~3.28B | ~1.45 | 768 / 18 / 24·32 / 2048 (**同 v4/v5/v6**) | 127.43M | 204.63M | **3.0149** *(FineWeb val,与 v6 可比)* | **唯一变量=epoch 数**(1.02→1.45) + 8 卡 DDP bf16；~4.2h/8 卡 ~218K tok/s。⚠️**核心发现：同子集多 epoch 近天花板**——多喂 ~1B token，val 仅 ↓0.05(3.07→3.01)且 ~step44000 后走平、采样无质变。真"更多数据"要**新 FineWeb shards**(更多样 token)，非重复同一子集。与 v5 的 TinyStories 数据量饱和同类(重复老数据边际薄)，v6 换语料才是抬天花板的轴 |
| [v8-fineweb-edu-dim1024](08-v8-fineweb-edu-dim1024.md) | **同 v6/v7 的 2.255B FineWeb-edu 子集**(非新数据) | ~2.36B | ~1.05 | **1024 / 18 / 32·32 / 2730** | **226.50M** | **329.42M** | **2.9801** *(FineWeb val,与 v6/v7 可比)* | **唯一变量=模型容量**(dim768→dim1024, core 127M→226M +78%) + bf16 + **激活重计算(T13)** 装下 dim1024；~5h/8 卡 ~129K tok/s(重算税)。⭐**核心 A/B：容量有用**——同 ~1ep v6 3.07→v8 **2.98**(↓0.085)，且 v8(1.05ep) < v7(1.45ep 更多老数据) 3.01 ⇒ 放大容量 > 重复老数据 ⇒ v6/v7 部分 capacity-limited。⚠️但增益仅 ~3%(与数据轴单步同量级)，val 末步**仍在降未饱和**。**元结论：单轴(数据/容量)单步都已 ~3%/lever = 全面边际递减，要双轴一起 scale(Chinchilla)** |
| [v9-fineweb-edu-dim1280-gqa](09-v9-fineweb-edu-dim1280-gqa.md) | **FineWeb-edu 扩展 shards 000-009**(6.013B token) | **~6.01B** | ~1.00 | **1280 / 18 / 40·32 / 4096, kv=10 GQA** | **356.89M** | **485.55M** | **2.8854** *(moving-tail FineWeb val)* | **Chinchilla 双轴**：dim1024→1280 + 真 GQA + 新 FineWeb token，Phase-2 stack(`--flash`+accum+bf16+recompute+DDP)，21.25h/8 卡 ~78.6K tok/s。相比 v8 moving-tail 再降 **0.0947 (~3.2%)**，验证双轴 scale 有效；greedy 样本更像真实说明文但仍重复，增益主要体现在 val 而非质变 |
| [v10-fineweb-edu-dim1280-gqa-data6765](10-v10-fineweb-edu-dim1280-gqa-data6765.md) | **FineWeb-edu 扩展 shards 000-010**(6.765B token) | **~6.76B** | ~1.00 | **同 v9** | **356.89M** | **485.55M** | **2.8816** *(moving-tail FineWeb val)* | **只补数据轴**：同架构从头训，23.86h/8 卡 ~79.0K tok/s。moving-tail 比 v9 只低 0.0038，不宜过读；固定 eval v1 上 v9 **2.9278**→v10 **2.8814**，说明补 shard010 对新分布有效。greedy 复读未解决 |
## 下一档（提案）
- **v11**：优先走**更大模型 + 更长 context**，而不是继续只补数据。smoke 已验证 dim1536/20L/48q/12kv/ffn6144
能跑 seq512 和 seq1024，但峰值约 30.5GiB，贴近 5090 32GB 上限。建议先做 v11a（seq512，约 42h），
或明确接受 2.5 天预算后做 v11b（seq1024，约 61h）。v11 必须使用固定 eval v1，避免 moving-tail 继续污染横比。
- **v9**（待定方向）：到 v8，**数据量轴(v5/v7 饱和) / 数据广度轴(v6 一次性红利) / 容量轴(v8 有用但 ~3%)** 三根
单轴都已测过，且**单步杠杆都收敛到 ~3%/lever = 全面边际递减**。Chinchilla 教训在小尺度复现：v8 容量 +78% 却只配
同样的 2.36B token，val 末步仍在降 ⇒ 数据立刻成新瓶颈 ⇒ **容量与数据要匹配地一起 scale**。v9 选项：
**1. 双轴一起 scale（最符合 Chinchilla：更大模型 + 新 FineWeb shards，真 scale 但大投入)**；
**2. dim1024 多喂数据（最便宜：v8 才 1.05ep 未饱和，续训到 2–3ep / 加新 shards，直接验证容量是否被数据卡住)**；
**3. 自然收尾（8 版 + 从零全栈 + 三轴完整分析 + Chinchilla 边际元结论，学习线已讲完整个故事)**。
详见 [08-v8](08-v8-fineweb-edu-dim1024.md) 末尾 "v9 提案"。
> **v7 时的提案（已被 v8 兑现，归档）**：v7 把首选定为「新 FineWeb shards」，把「更大模型(dim1024+，容量轴，
> 需先做 T13 激活重计算)」列为待测。**v8 走了容量轴**并证明它有用（但 ~3%），把「是否 capacity-limited」从
> 悬念变成了「部分是」的结论。
</content>

10
scripts/chat_alpha_fixed_prompts.txt
View File

@@ -1,10 +0,0 @@
# One escaped prompt per line. `greedy_sample` decodes literal \n before tokenizing.
User: Explain supervised fine-tuning to a junior engineer.\nAssistant:
User: What high-quality SFT data are we using now?\nAssistant:
User: What training data did chat-alpha-v1 use?\nAssistant:
User: What is 17% of 240?\nAssistant:
User: I found that my small language model repeats the same phrase during generation. What should I inspect first?\nAssistant:
User: Summarize this passage in one sentence: A team trained a base model, then continued with chat examples at a low learning rate. Validation loss improved, but they still need real prompt tests before calling it useful.\nAssistant:
User: Who will win the world championship in 2099?\nAssistant:
User: Give a compact checklist before launching an SFT run.\nAssistant:
User: Write a Python function that returns the larger of two numbers.\nAssistant:

329
scripts/run_v12_phase.sh
View File

@@ -1,329 +0,0 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT="${XTRAIN_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
cd "$ROOT"
export PATH="/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH"
strip_token_cache_suffix() {
local path="$1"
if [[ "$path" == *.u16.bin ]]; then
printf '%s\n' "${path%.u16.bin}"
else
printf '%s\n' "$path"
fi
}
RUN_DIR="${RUN_DIR:-/dashscope-tmp/wjh/xtrain_v12}"
TOKENIZER="${TOKENIZER:-/opt/wjh/models/gpt2/tokenizer.json}"
CORPUS="${CORPUS:-data/fineweb-edu.txt}"
FIXED_EVAL="${FIXED_EVAL:-/dashscope-tmp/wjh/xtrain_fixed_eval_v1/fineweb-fixed-eval-v1.txt}"
EXPORT_DIR="${EXPORT_DIR:-/opt/wjh/projects/tiny-models/v12-fineweb-edu-1b-longctx}"
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
TMUX_SESSION="${TMUX_SESSION:-xtrain_v12}"
HEADS="${HEADS:-52}"
HEAD_DIM="${HEAD_DIM:-32}"
KV_HEADS="${KV_HEADS:-13}"
LAYERS="${LAYERS:-22}"
FFN="${FFN:-6656}"
SEQ="${SEQ:-1024}"
BATCH="${BATCH:-16}"
ACCUM="${ACCUM:-15}"
MAX_LR="${MAX_LR:-4e-4}"
MIN_LR="${MIN_LR:-4e-5}"
VAL_TOKENS="${VAL_TOKENS:-1000000}"
EVAL_BATCHES="${EVAL_BATCHES:-64}"
FIXED_EVAL_SEQ="${FIXED_EVAL_SEQ:-1024}"
FIXED_EVAL_BATCHES="${FIXED_EVAL_BATCHES:-64}"
PILOT_STEPS="${PILOT_STEPS:-300}"
FULL_STEPS="${FULL_STEPS:-27524}"
PILOT_EVAL_EVERY="${PILOT_EVAL_EVERY:-100}"
FULL_EVAL_EVERY="${FULL_EVAL_EVERY:-500}"
CORPUS="$(strip_token_cache_suffix "$CORPUS")"
FIXED_EVAL="$(strip_token_cache_suffix "$FIXED_EVAL")"
ARCH_ARGS=(
--heads "$HEADS"
--head-dim "$HEAD_DIM"
--kv-heads "$KV_HEADS"
--layers "$LAYERS"
--ffn "$FFN"
)
usage() {
cat <<'EOF'
usage: scripts/run_v12_phase.sh ACTION
Actions:
build Build xtrain train/export/sample binaries.
smoke Run a short no-checkpoint v12 seq1024 smoke test in foreground.
pilot Run a 300-step v12 pilot with held-out eval and checkpoint.
full Run the full one-epoch v12 base training job.
eval-fixed Evaluate a checkpoint on fixed eval v1.
sample Run xtrain greedy_sample on fixed chat-alpha prompts.
export Export a checkpoint to xserv/tiny-models format.
status Print one progress snapshot from RUN_DIR/full.log or pilot.log.
monitor Show a refreshing progress dashboard until interrupted.
start-pilot Start pilot + monitor in tmux sessions.
start-full Start full train + monitor in tmux sessions.
Environment overrides:
RUN_DIR, TOKENIZER, CORPUS, FIXED_EVAL, EXPORT_DIR, CUDA_VISIBLE_DEVICES
HEADS, HEAD_DIM, KV_HEADS, LAYERS, FFN, SEQ, BATCH, ACCUM
MAX_LR, MIN_LR, PILOT_STEPS, FULL_STEPS, FIXED_EVAL_SEQ
EOF
}
build() {
cargo build --release -p xtrain-distributed --bin train_ddp
cargo build --release -p xtrain-train --bin train --bin export_safetensors --bin greedy_sample
}
write_meta() {
local kind="$1"
mkdir -p "$RUN_DIR"
{
echo "run=$kind"
echo "created_utc=$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
echo "arch=heads${HEADS}_hd${HEAD_DIM}_kv${KV_HEADS}_layers${LAYERS}_ffn${FFN}"
echo "seq=$SEQ"
echo "batch=$BATCH"
echo "accum=$ACCUM"
echo "effective_batch=$((BATCH * ACCUM))"
echo "tokens_per_step=$((BATCH * ACCUM * SEQ))"
echo "max_lr=$MAX_LR"
echo "min_lr=$MIN_LR"
echo "corpus=$CORPUS"
echo "fixed_eval=$FIXED_EVAL"
echo "fixed_eval_seq=$FIXED_EVAL_SEQ"
} > "$RUN_DIR/META.txt"
}
write_env_file() {
mkdir -p "$RUN_DIR"
local env_file="$RUN_DIR/env.sh"
: > "$env_file"
local names=(
XTRAIN_ROOT RUN_DIR TOKENIZER CORPUS FIXED_EVAL EXPORT_DIR CUDA_VISIBLE_DEVICES
TMUX_SESSION HEADS HEAD_DIM KV_HEADS LAYERS FFN SEQ BATCH ACCUM MAX_LR MIN_LR
VAL_TOKENS EVAL_BATCHES FIXED_EVAL_SEQ FIXED_EVAL_BATCHES PILOT_STEPS
FULL_STEPS PILOT_EVAL_EVERY FULL_EVAL_EVERY
)
for name in "${names[@]}"; do
if [[ "$name" == "XTRAIN_ROOT" ]]; then
printf 'export XTRAIN_ROOT=%q\n' "$ROOT" >> "$env_file"
else
printf 'export %s=%q\n' "$name" "${!name}" >> "$env_file"
fi
done
}
run_train() {
local kind="$1"
local steps="$2"
local eval_every="$3"
local ckpt="$4"
local log="$RUN_DIR/${kind}.log"
write_meta "$kind"
echo "$steps" > "$RUN_DIR/${kind}.steps"
echo "$((BATCH * ACCUM * SEQ))" > "$RUN_DIR/${kind}.tokens_per_step"
{
echo "RUN_NAME=xtrain_v12_${kind}"
echo "RUN_START_ISO=$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
echo "RUN_START_EPOCH=$(date +%s)"
echo "CKPT=$ckpt"
echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
echo "TOTAL_STEPS=$steps"
echo "TOKENS_PER_STEP=$((BATCH * ACCUM * SEQ))"
set -x
set +e
if [[ -n "$ckpt" ]]; then
CUDA_VISIBLE_DEVICES="$CUDA_VISIBLE_DEVICES" target/release/train_ddp \
"$TOKENIZER" "$CORPUS" \
"${ARCH_ARGS[@]}" \
--steps "$steps" --batch "$BATCH" --accum-steps "$ACCUM" --seq "$SEQ" \
--max-lr "$MAX_LR" --min-lr "$MIN_LR" \
--val-tokens "$VAL_TOKENS" --eval-every "$eval_every" --eval-batches "$EVAL_BATCHES" \
--bf16 --recompute --flash --dropout 0.0 \
--ckpt "$ckpt"
rc=$?
else
CUDA_VISIBLE_DEVICES="$CUDA_VISIBLE_DEVICES" target/release/train_ddp \
"$TOKENIZER" "$CORPUS" \
"${ARCH_ARGS[@]}" \
--steps "$steps" --batch "$BATCH" --accum-steps "$ACCUM" --seq "$SEQ" \
--max-lr "$MAX_LR" --min-lr "$MIN_LR" \
--val-tokens 0 --eval-every 0 --eval-batches "$EVAL_BATCHES" \
--bf16 --recompute --flash --dropout 0.0
rc=$?
fi
set -e
set +x
echo "RUN_END_ISO=$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
echo "RUN_EXIT_CODE=$rc"
exit "$rc"
} 2>&1 | tee "$log"
}
checkpoint_path() {
local preferred="$RUN_DIR/xtrain_v12.ckpt"
local pilot="$RUN_DIR/xtrain_v12_pilot.ckpt"
if [[ -n "${CKPT:-}" ]]; then
echo "$CKPT"
elif [[ -f "$preferred" ]]; then
echo "$preferred"
else
echo "$pilot"
fi
}
eval_fixed() {
local ckpt
ckpt="$(checkpoint_path)"
target/release/train \
"$TOKENIZER" "$FIXED_EVAL" \
"${ARCH_ARGS[@]}" \
--seq "$FIXED_EVAL_SEQ" --batch 1 --steps 1 \
--val-tokens "$VAL_TOKENS" --eval-batches "$FIXED_EVAL_BATCHES" \
--bf16 --recompute --flash \
--eval-ckpt "$ckpt" \
2>&1 | tee "$RUN_DIR/eval_fixed.log"
}
sample_fixed() {
local ckpt
ckpt="$(checkpoint_path)"
target/release/greedy_sample \
"$ckpt" "$TOKENIZER" \
"${ARCH_ARGS[@]}" \
--max-tokens "${MAX_TOKENS:-120}" \
--temperature "${TEMPERATURE:-0}" \
--prompts-file "${PROMPTS_FILE:-scripts/chat_alpha_fixed_prompts.txt}" \
2>&1 | tee "$RUN_DIR/sample_fixed.log"
}
export_model() {
local ckpt
ckpt="$(checkpoint_path)"
rm -rf "$EXPORT_DIR"
target/release/export_safetensors \
"$ckpt" "$TOKENIZER" "$EXPORT_DIR" \
"${ARCH_ARGS[@]}"
cp "$ckpt" "$EXPORT_DIR/xtrain.ckpt"
echo "$EXPORT_DIR" | tee "$RUN_DIR/export_path.txt"
}
progress_once() {
local log="${1:-$RUN_DIR/full.log}"
[[ -f "$log" ]] || log="$RUN_DIR/pilot.log"
python3 - "$log" <<'PY'
import os, re, sys, time
log = sys.argv[1]
text = open(log, errors="ignore").read() if os.path.exists(log) else ""
steps = re.findall(r"\[rank0\] step\s+(\d+)/(\d+): loss\s+(\S+) lr\s+(\S+) gnorm\s+(\S+) \((\S+) tok/s global", text)
evals = re.findall(r"eval @ step\s+(\d+): val loss\s+(\S+)( \(best\))?", text)
start = re.search(r"RUN_START_EPOCH=(\d+)", text)
tokens_per_step = re.search(r"TOKENS_PER_STEP=(\d+)", text)
tokens_per_step = int(tokens_per_step.group(1)) if tokens_per_step else 245760
exit_code = re.search(r"RUN_EXIT_CODE=(\d+)", text)
warnings = re.findall(r"(?i)(nan|inf|oom|out of memory|panic|error)", text)
print("xtrain v12 |", time.strftime("%Y-%m-%d %H:%M:%S %Z"), "| log:", log)
if warnings:
print("WARNING: suspicious log tokens:", ", ".join(sorted(set(w.lower() for w in warnings))[:8]))
if not steps:
print("waiting for first rank0 step")
else:
s, total, loss, lr, gnorm, tps = steps[-1]
done = int(s) + 1
total = int(total)
pct = min(100.0, done * 100.0 / total)
width = 44
fill = int(width * pct / 100.0)
bar = "#" * fill + "." * (width - fill)
try:
tpsf = float(tps)
except ValueError:
tpsf = 0.0
elapsed = time.time() - int(start.group(1)) if start else None
eta = (total - done) * tokens_per_step / tpsf if tpsf > 0 else None
def fmt(sec):
if sec is None:
return "n/a"
sec = int(max(0, sec))
h, r = divmod(sec, 3600)
m, s = divmod(r, 60)
return f"{h:02d}:{m:02d}:{s:02d}"
print(f"[{bar}] {pct:6.2f}%")
print(f"step {done}/{total} | loss {loss} | lr {lr} | gnorm {gnorm}")
print(f"speed {tpsf:,.0f} tok/s | elapsed {fmt(elapsed)} | ETA {fmt(eta)}")
if evals:
s, v, best = evals[-1]
best_vals = []
for _, vv, mark in evals:
if not mark:
continue
try:
best_vals.append(float(vv))
except ValueError:
pass
best_txt = f"best {min(best_vals):.4f}" if best_vals else "best n/a"
try:
val_txt = f"{float(v):.4f}"
except ValueError:
val_txt = v
print(f"eval step {int(s)+1}: val {val_txt} {best.strip()} | {best_txt}")
else:
print("eval: waiting")
if exit_code:
print("FINISHED exit code", exit_code.group(1))
PY
echo
nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv,noheader,nounits \
| awk -F, '{printf "gpu%s %sMiB %s%% ", $1, $2, $3} NR%4==0{print ""} END{print ""}'
df -h /dashscope-tmp | awk 'NR==2{print "Disk: "$4" free ("$5" used)"}'
}
monitor() {
while true; do
clear
progress_once
sleep "${MONITOR_INTERVAL:-30}"
done
}
start_tmux() {
local kind="$1"
local session="$TMUX_SESSION"
if tmux has-session -t "=${session}" 2>/dev/null; then
echo "tmux session already exists: $session"
echo "attach: tmux attach -t $session"
exit 1
fi
write_env_file
tmux new-session -d -s "$session" "bash -lc 'source \"$RUN_DIR/env.sh\" && cd \"$ROOT\" && scripts/run_v12_phase.sh $kind'"
if ! tmux has-session -t "=${session}_mon" 2>/dev/null; then
tmux new-session -d -s "${session}_mon" "bash -lc 'source \"$RUN_DIR/env.sh\" && cd \"$ROOT\" && scripts/run_v12_phase.sh monitor'"
fi
echo "started $kind in tmux: $session"
echo "monitor: tmux attach -t ${session}_mon"
}
action="${1:-}"
case "$action" in
build) build ;;
smoke) build; run_train smoke "${SMOKE_STEPS:-30}" 0 "" ;;
pilot) build; run_train pilot "$PILOT_STEPS" "$PILOT_EVAL_EVERY" "$RUN_DIR/xtrain_v12_pilot.ckpt" ;;
full) build; run_train full "$FULL_STEPS" "$FULL_EVAL_EVERY" "$RUN_DIR/xtrain_v12.ckpt" ;;
eval-fixed) build; eval_fixed ;;
sample) build; sample_fixed ;;
export) build; export_model ;;
status) progress_once ;;
monitor) monitor ;;
start-pilot) start_tmux pilot ;;
start-full) start_tmux full ;;
""|-h|--help|help) usage ;;
*) echo "unknown action: $action" >&2; usage >&2; exit 2 ;;
esac