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

Compare commits

base: gahow:2e68942032dd677dfb87e2a6df04d92fe1af1e62
Branches Tags
gahow:main
...
compare: gahow:main
Branches Tags
gahow:main

142 Commits
2e68942032 ... main

Author SHA1 Message Date
Gahow Wang
6465a2d5ce test: T21-for-proc — clear ENV_DROPOUT across tests to sever ordering coupling 
libtest with --test-threads=1 (the documented invariant for this file's DDP
tests) runs tests alphabetically. The new
proc_per_gpu_dropout_is_live_and_p0_matches_no_dropout ('d') runs BEFORE
proc_per_gpu_matches_single_gpu_and_thread_path ('m'). It sets ENV_DROPOUT=0.2
via std::env::set_var; if left in place, the correctness test's spawned workers
would inherit it (Command inherits parent env by default) and build with
cfg.dropout=0.2 while its single-GPU baseline (run_single_gpu → test_config →
dropout=0) stays at 0 — GATE (a) `max_rel_single < 1e-3` would blow up by
orders of magnitude.

Two defenses:
- correctness test remove_var(ENV_DROPOUT) before spawn (belt): even if the
  dropout test forgot to clean up, this test starts from a clean env.
- dropout test remove_var(ENV_DROPOUT, ENV_DUMP_DIR) at exit (suspenders):
  keep the invariant "each test leaves the env as it found it" so any future
  test added after these two starts clean too.

Same --test-threads=1 SAFETY comment applies (no concurrent env access).
 2026-07-01 14:09:42 +08:00
Gahow Wang
33a1aee9ec test: T21-for-proc — dropout-live regression under process-per-GPU 
Analogue of the ddp_dropout_is_live_and_p0_bit_identical test (T21, thread-per-
GPU) for the process-per-GPU launcher. Runs launch_processes twice on the same
corpus / init / config with the ONLY difference being cfg.dropout (passed
launcher→worker via a new XTRAIN_TEST_DROPOUT env — worker re-execs cannot
inherit argv changes), reads rank 0's loss trajectory from both runs, and
asserts GATE B: max |loss diff| > 1e-3.

The threshold sits ~4 orders of magnitude above this box's KI-5 cross-rank NCCL
noise floor (~1e-7), so it is an unambiguous "dropout mask is applied" signal,
not a noise measurement. Pre-fix (missing cfg.dropout = ... in the worker /
launcher, exactly the gap the paired launcher commit closes) both traces are
bit-identical and this test FAILs.

Also wires ENV_DROPOUT into the shared worker entry so the existing correctness
test's contract is unchanged (absent env → 0.0 → same synth run as before).
p0/ and p02/ subdirs isolate the two invocations' dumps.
 2026-07-01 13:51:31 +08:00
Gahow Wang
86de6bfb51 distributed: T21-for-proc — wire --dropout into the process-per-GPU launcher 
T21 fixed --dropout under thread-per-GPU (train_ddp): added the flag, set
cfg.dropout, and made train_rank re-assert model.train() each step so the
training forward stays live across periodic eval flips. The process-per-GPU
launcher (train_ddp_mp) was left out: it never parsed --dropout, so cfg.dropout
stayed at Config::from_arch's 0.0 default, and the worker's model built with
dropout permanently disabled — silently, regardless of what the user passed.

The gap is the exact same launcher-wiring class the V9-PILOT caught: op-level
+ single-GPU tests pass, the DDP-thread T21 regression test passes, but the
proc-per-GPU launcher path was never exercised end-to-end with dropout>0.

Mirror bin/train_ddp exactly: parse --dropout (default 0, bit-identical
default), set cfg.dropout before build_model, print an ON banner on rank 0.
train_rank's per-step model.train() from T21 is reused unchanged (proc-per-GPU
uses the same train_rank).

Follow-up test that exercises this wiring end-to-end (GATE B loss-trace
divergence between p=0 and p=0.2 under process-per-GPU) lands in the next
commit.
 2026-07-01 13:51:17 +08:00
Gahow Wang
4379868f2d docs: M2d — ragged-batching lever, 9× measured, step bottleneck → rollout 
Records the M2d lever (batch the GRPO training-side forwards), the right-pad-is-free
insight, both exact gates, the end-to-end no-OOM smoke, and the 9× throughput.

The honest decomposition correction: M2c claimed the training forwards "dominate" the
step; the clean per-component bench falsifies the strong form — they were ~2.5 s of
the ~8.5 s step (~30%), worth the 9×, but the rollout (~6 s) was always the larger
share. After M2d the step is ~95% rollout, so the next step-level lever is full B×G
rollout batching (today only the G samples of each prompt decode in lockstep; the B
prompts are still sequential). Same measure-first lesson, once more.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 23:03:28 +08:00
Gahow Wang
0e82b2438e test: M2d — ragged-forward + batched-op equivalence gates + throughput bench 
Two exact correctness gates (composed = the end-to-end batched GRPO step == looped):
- xtrain-model forward_batched_ragged_matches_looped: forward_batched on RIGHT-padded
  ragged sequences == per-sequence single-seq forward on the real rows. fp32
  max|Δlogit| = 3.7e-7, bf16 = 0.0, both composed + flash SDPA. Pins "right-pad is
  free under causal".
- 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).

bench_grpo_batch: weight-independent micro-bench of the per-sample training forwards
(loads v12 base as policy, N realistic ragged samples, teacher-forced argmax targets
so the closeness smoke isn't −log-amplified by random low-prob tokens). Measured on
dash5 (v12 1.05B, N=48, micro=16): capture 622→71 ms (8.7×), inner 1907→208 ms
(9.2×), training forwards 2526→280 ms (9.0×).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 23:03:09 +08:00
Gahow Wang
c2ebf62ae1 post-train: M2d — batch the GRPO training-side forwards (op + module + wiring) 
After M2b/M2c made the rollout cheap, the GRPO step is dominated by the per-sample
single-sequence training-side forwards: the per_token_logp captures (policy +
reference) and the inner clipped-PG forward/backwards. M2d packs all N=B·G ragged
samples of a step into ONE forward_batched.

Enabling property — right-padding is free under causal attention: a real completion
row sits at an earlier position than the trailing pad, and causal masking forbids
attending forward, so its logits equal the unpadded single-sequence forward; pad
rows are masked out (target=-100).

- ops::clipped_pg_loss_batched: like clipped_pg_loss but takes per-row advantage[t]
  (the owning sample's A) and per-row weight[t] (the full normaliser). It does NOT
  compute its own 1/n_tokens, so the caller passing weight=1/(N·n_s) reproduces the
  looped Σ_s (1/N)(1/n_s)·clipped_pg_loss_s bit-for-bit (per-row CE backward is
  row-local).
- grpo_batch.rs (shared module): per_token_logp_batched (right-pad → one
  forward_batched(N) → slice back to real length) + looped baselines +
  inner_pg_step_{looped,batched}. A --micro knob chunks the pack to bound the
  [chunk·Lmax, vocab] logits memory; weight uses the GLOBAL N so chunked
  grad-accumulation stays exact.
- train_grpo restructured to collect-all-samples-then-batch; per-window phase timers
  (rollout / capture / inner) to keep the step decomposition honest. Default micro =
  B·G; bench-measured 9× on the training forwards.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 23:02:56 +08:00
Gahow Wang
41d46208a6 docs: M2c — device KV cache + the bottleneck-shift finding 
Implementation log (docs/18) + Phase-3 row (evolution.md): cat_seq device cache,
gates hold (token-identical), and the profile-first finding — ~10% single-seq
decode but no GRPO-step change because the long pole shifted to the per-sample
logp/PG forwards after M2b batching. Names ragged batched prefill as the next
decode lever.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 17:39:10 +08:00
Gahow Wang
3a3425960c post-train: M2c — device-side KV cache (cat_seq), profile-first bottleneck shift 
Device-resident KV cache: keep K/V on the GPU as [bh,T,hd], grow by one token
per step via a new cat_seq kernel (concat along seq) — removes 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; cache is Option<Tensor>
per layer (cleaner than the host Vec + rebuild).

Gates all hold: cat_seq == host concat; decode_kv single-seq + decode_batch
G-way both still TOKEN-IDENTICAL; GQA training path unaffected.

Honest measurement (the point): 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 + the PG-update
forwards/backwards (model.forward, full-seq) now dominate. Measure-first lesson
(cf. T11/T17/M2a): the long pole shifted to the training-side forwards; the next
decode lever (ragged batched prefill) targets those, not the cache.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 17:38:16 +08:00
Gahow Wang
0f76c0fdb0 docs: M2b — batched decode results (token-identical + ~1.7x rollout, device-cache next) 
Implementation log (docs/18) + Phase-3 row (evolution.md): rope_pos primitive +
gate, the batched engine (decode_attention/repeat_kv reused), the token-
identical batch gate, and the measured ~1.7x rollout-inclusive step speedup +
memory stabilization. Closes the M2 decode engine (M2a single-seq + M2b
batched); names the device-side cache as the remaining lever.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 17:20:01 +08:00
Gahow Wang
361c5290fa post-train: M4 — use M2b batched rollout in GRPO (~1.7× step) 
train_grpo rolls out a prompt's G samples with one generate_cached_batch call
instead of G sequential generate_cached calls. Measured on v12 1.05B (G=6, B=6,
easy task): ~8.5 s/step vs ~14-16 s/step single-seq cached — ~1.7× (rollout-
inclusive; short of G× because per_token_logp + the PG update also cost, and the
M2a host round-trip remains). Also more stable memory: one batched forward per
step vs G allocations that fragment the caching allocator.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 17:18:54 +08:00
Gahow Wang
2c9b58cb3b post-train: M2b — batched KV-cache decode (G-way, token-identical) 
The rollout long-pole fix deferred from M2a: decode the G samples of one prompt
in lockstep (one forward per step over the group → G× fewer kernel launches).

- rope_pos(x, positions[]): RoPE with a per-row absolute position (new forward-
  only kernel) — G rows share one decode position. Gate: == full rope for
  [0..n], == rope_at(P) per row for uniform P (bit-identical).
- generate_cached_batch: BatchKVCache [T, G·num_kv, hd] + batched decode_step.
  decode_attention is already batch-agnostic (bh = G·nh); repeat_kv(nh, batch=G)
  broadcasts per group. No finished-mask / ragged prompts yet (perf-only / next).
- Gate (tests/decode_batch.rs): all G greedy rows token-identical to the single-
  sequence decode (8 query / 2 kv heads → exercises repeat_kv batching).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 17:18:54 +08:00
Gahow Wang
096e45b845 docs: M4 — GRPO results (infra + memory/rollout walls + capability-wall negative result) 
Implementation log (docs/18) + Phase-3 row (evolution.md): the clipped_pg_loss
op + gates, the actor-learner loop, the easy-task SFT baseline (held-out 18.7%,
plateaus → no generalization), the two systems walls the design doc flagged
(two 1B models OOM the 32GB box → β=0; naive rollout fragments the allocator →
cached temperature sampling, rollout still the long pole), and the result:
format holds, held-out 20.0% (+1.3pp, statistically flat) — the same wall as
DPO. Closes the SFT→KV-cache→DPO→GRPO post-training arc with honest limits.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 17:01:22 +08:00
Gahow Wang
7fb3b32fd9 post-train: M4 — GRPO actor-learner loop + cached temperature rollout 
train_grpo: the online, critic-free RL loop — per step sample B prompts, roll
out G completions each, score with the rule-based checker (reward 0/1), compute
group-relative advantage A=(r−mean)/(std+ε), then K inner clipped_pg_loss
epochs with a KL leash to the frozen reference. Reward = pure 0/1 correctness
(KL is the format protector, the M3 collapse lesson). Tracks mean rollout reward
(the falsifiable "it learns" signal). Periodic checkpoint save.

decode: generate_cached adds temperature sampling to the KV-cache engine (M2) —
single-row [1,vocab] logits per step vs the naive sampler's [seq,vocab], far
lighter on the caching allocator (the naive sampler fragments it over a long
rollout). generate_greedy_cached now routes through it (temp 0); decode_kv
token-identical gate still passes.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 16:59:05 +08:00
Gahow Wang
aaa77082ef post-train: M4 — clipped_pg_loss + scale_rows (GRPO policy-gradient op) 
The GRPO (M4) token-level loss op + the one primitive it needs:

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

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 14:07:02 +08:00
Gahow Wang
99090465bf docs: M3 — DPO results (infra correct, held-out correctness flat, over-optimization collapse) 
Implementation log (docs/18) + Phase-3 row (evolution.md): the two ops + gates,
pair-gen (gold chosen / sampled-wrong rejected), reference-logprob caching, the
training loop, and the honest finding — reward margin + pref-acc rise but
held-out arithmetic correctness stays ~5-8% (flat within std-error) and
over-optimizes to collapse (margin +34 → 0% format). DPO reweights, it does not
install the capability; motivates M4 GRPO (optimize the verifiable reward online).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 12:38:06 +08:00
Gahow Wang
2f827fd6d8 post-train: M3 — DPO pair-gen + training loop (verifiable arithmetic) 
gen_dpo_pairs: 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 from the model's distribution. ~8% of prompts skipped (greedy
correct). Writes question<TAB>chosen<TAB>rejected (bare, SFT-framed at train).

train_dpo: loads the SFT ckpt as policy AND frozen reference; precomputes the
reference logprobs ONCE (policy==ref) 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.
Tracks reward margin + preference accuracy (the doc-13 "don't trust loss alone"
health signal). Loss starts at exactly log2 (Δ=0 at init) — a built-in check.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 12:37:01 +08:00
Gahow Wang
f3c764ce95 post-train: M3 — seq_logprob + dpo_loss autograd ops 
Two new ops for DPO (M3), both reusing existing kernels (no new CUDA):

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

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

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 12:11:01 +08:00
Gahow Wang
b39e6e7110 docs: M2a — KV-cache decode engine results (token-identical + length-dependent speedup) 
Implementation log (docs/18) + Phase-3 row (evolution.md): the two decode
primitives and their gates, the engine design (host-cache baseline), the
token-identical centerpiece gate, and the measured throughput baseline showing
the cache win is sequence-length-dependent (~1.0x@32, ~1.9x@128, naive OOM@256).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 12:01:10 +08:00
Gahow Wang
eff26a0898 post-train: M2a — KV-cache incremental decode engine (token-identical) 
Single-sequence KV-cache decode (xtrain-model/src/decode.rs): per-layer K/V
cache + single-token incremental forward (prefill = first prompt.len() decode
steps, one code path). Mirrors model::block_forward at the raw-Tensor level (no
autograd tape — inference needs no grads), using rope_at + decode_attention.
Cache is host-accumulated token-major f32, rebuilt per step (the honest M2a
baseline; M2b moves it device-side + batched ragged).

Gate (the M2 centerpiece): KV-cache greedy decode is TOKEN-IDENTICAL to the
naive full-recompute greedy — tests/decode_kv.rs (small GQA model, F32, 24
tokens) and corroborated on the v12 1.05B SFT checkpoint (cached eval =
naive eval byte-for-byte: format 100/100, correct 8/100).

eval_arith --cached A/Bs the two paths + reports decode tok/s. Measured on v12
(1.05B, batch 1, F32): the cache win is sequence-length-dependent —
  max_new=32   naive 108 vs cached 111 tok/s  (~1.0x; overhead-bound)
  max_new=128  naive  69 vs cached 133 tok/s  (~1.9x)
  max_new=256  naive OOM     vs cached 129 tok/s
Cached throughput stays ~constant (O(1)/token) while naive decays (O(t)/token,
O(seq^2) graph → OOM at length). Short eval prompts are overhead-bound, so the
cache matters for long rollouts (DPO/GRPO), not the arithmetic eval itself.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 12:00:03 +08:00
Gahow Wang
c88e2ab88c post-train: M2 — decode primitives (rope_at + decode_attention) 
Two forward-only Tensor primitives the KV-cache decode engine is built on,
each gated by an isolated correctness test:

- rope_at(theta, pos0): RoPE at an absolute position (pos = pos0 + row, no
  modulo) for a single decode token, vs the training rope_k (pos = row %
  period) left untouched. New forward-only CUDA kernel, no training-path risk.
  Gate: bit-identical to the full-sequence rope's corresponding row.
- decode_attention(k, v, scale): single-query × cached-K/V SDPA, composed from
  the existing strided batched GEMM + plain (non-causal) softmax — no new
  kernel. Gate: equals the full causal attention's last query row (max |Δ| 6e-8).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 12:00:03 +08:00
Gahow Wang
1574e21d89 post-train: M1 — verifiable-arith eval scorer + SFT format-baseline result 
eval_arith: load ckpt, greedy-generate per held-out prompt, parse \boxed{}
via the shared task checker, report format(boxed) + correctness pass-rates.
Reused as the verifiable-eval harness for M3 (DPO) / M4 (GRPO).

M1 result (100 held-out prompts, v12 1.05B base): SFT moves answer-format
adherence 0% -> 100%, arithmetic correctness 8% -- the intended split (SFT
buys the format; correctness is the verifiable-reward job of M3/M4). Logged
in docs/18 implementation log + a Phase-3 row in docs/evolution.md.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-30 11:13:19 +08:00
Gahow Wang
cb64604496 post-train: M1 fix — enlarge arith key space + saturation guard 
The default operand ranges (max_add=99, max_mul=12) gave only ~20k unique
problems, so 'gen_arith_task --n 20000 --eval 500' (a) made train dedup
pathologically slow near saturation and (b) made the disjoint-eval loop never
terminate. A background run stalled after ~10k train rows with no eval files.

Fix (root cause, not a workaround):
- enlarge default ranges to max_add=999, max_mul=99 (~2.01M key space) so 20k+
  requests are a tiny fraction and dedup stays trivial;
- add unique_space() + a generator guard that errors clearly when n+eval exceeds
  80% of the key space, instead of looping forever.

Verified: cargo test 10/10; full 20000/500 gen now 0.2s, all 3 files, 0
train/eval leakage; guard panics on an oversized (--max-add 99) request.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-29 23:28:25 +08:00
Gahow Wang
9c70e99ae4 post-train: M1 — verifiable arithmetic task + SFT data generator 
First post-training milestone (docs/18). Lands the verifiable task + its data
pipeline, all verified host-side (no CUDA); the SFT run itself reuses the
existing --sft-tsv path on the GPU box.

- task.rs: the shared task spec — two-operand integer arithmetic, answer in
  \boxed{N}, with parse_boxed_answer + check_answer (exact-match rule-based
  reward). One module reused by M1 (SFT data), M3 (DPO pairs), M4 (GRPO reward).
- gen_arith_task bin: writes arith_sft.tsv (--sft-tsv format) + held-out
  arith_eval_prompts.txt (greedy_sample format) + arith_eval_gold.txt; train
  deduped, eval disjoint from train.
- data.rs: extract assistant-only masking into a pure, testable sft_row()
  (behavior-preserving; single-turn bit-identical to fbf4ac2).

Gate (verified locally, no_cuda): cargo test -p xtrain-train --lib = 9/9 pass
(masking, SFT-target self-consistency over 2000 samples, parser edges, seed
determinism); a 200/50 gen run = clean 2-col TSV, correct gold incl. negatives,
0 train/eval leakage. SFT training run + format-eval pending on dash5.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-29 22:52:25 +08:00
Gahow Wang
ab32168dcc docs: post-training stack design — SFT → KV-cache → DPO → GRPO (docs/18) 
Design doc for a from-scratch post-training infra on top of xtrain. Ladder:
SFT (have it) → DPO → reward model (optional) → GRPO, each rung one new
post-training systems concept + a hard correctness gate (grad-check, PyTorch
parity, degenerate checks, a falsifiable 'it learns' signal).

Decisions aligned with the user (D1-D4):
- D1 scope: DPO → GRPO, reward model optional.
- D2 reward: rule-based / verifiable first; learned RM deferred.
- D3 rollout: build the KV-cache incremental-decode engine UP FRONT (not
  naive-first) as the foundational milestone before DPO/GRPO.
- D4 task: a verifiable task (arithmetic/format) with deterministic exact-match
  reward, for a clean RL signal.

Locked milestone order: M1 SFT task baseline → M2 KV-cache decode engine
(token-identical gate) → M3 DPO → M4 GRPO → M5 optional reward model. Status:
design only, no implementation yet.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-29 22:44:25 +08:00
Gahow Wang
7a1fba95b5 docs: v12 — 1.05B long-ctx base + chat-alpha SFT quality check 
- run 12: dim1664/22L true-GQA 1.05B base, seq1024, 6.765B FineWeb tokens,
  81h on 8x5090. Fixed eval v1 @seq1024 = 2.7410 vs v11 2.7467 — a real but
  marginal gain; v11->v12 is a capacity-only step on fixed data, so the ~0.2%
  return confirms the 1B base is now data-limited.
- run 13: three SFT stages from the v12 base (synthetic / anchor /
  real-mix-repair). The pipeline works and produces a chat-shaped model that
  follows the format and stops, but none of the variants is a stable
  high-quality chat model — bottleneck is SFT data quality + selection signal
  (val loss decouples from generation quality), not infra.
- scripts/run_v12_phase.sh wrapper + chat_alpha_fixed_prompts.txt eval set.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-29 16:19:12 +08:00
Gahow Wang
fbf4ac2917 sft: assistant-only SFT (ignore-index CE) + chat-prompt greedy eval 
Enable assistant-only supervised fine-tuning and a fixed chat-prompt eval path
used by the v12 SFT runs:

- cross_entropy ignores negative targets (-100 ignore-index), normalizing by
  valid rows instead of all rows; CUDA fwd/bwd skip t<0 (ops.rs, nn.cu).
- Corpus gains optional labels + load_sft_tsv_cached: two-column TSV is
  formatted as 'User: .. \nAssistant:' + answer + <|endoftext|>, prompt tokens
  masked to -100 while answer+EOS are supervised; i32 label cache alongside the
  u16 token cache; sample() retries windows that are fully masked; eval uses
  target_window so masking applies to val loss too (data.rs, train_loop.rs).
- train + train_ddp: --sft-tsv selects the TSV loader, --init-ckpt continues
  training from a base checkpoint.
- greedy_sample: --prompts-file/--prompt/--temperature for fixed chat-prompt
  generation eval.

Test fixtures updated for the new Corpus.labels field; dropout.rs carries
incidental rustfmt. Not rebuilt locally (no CUDA toolchain on this checkout);
correctness rests on the documented v12 base+SFT runs on the GPU box.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-29 16:19:02 +08:00
Gahow Wang
5c27493a90 docs: backfill v9/v10 scaling runs + reframe README to v0–v10 / three phases 
Add per-run design+result docs for the two Chinchilla-axis runs that were
done but never committed:
- v9 (dim1280 true-GQA, core 357M, 6.01B FineWeb tokens): double-axis scale,
  best moving-tail val 2.8854 (~3.2% below v8) — direction validated, gain
  still incremental, greedy repetition remains.
- v10 (same arch, data-only top-up to 6.765B): moving-tail 2.8816; fixed
  eval v1 v6→v10 = 3.2328/3.1850/3.1515/2.9278/2.8814.

Extend the comparison tables in docs/runs/README.md and docs/evolution.md to
v10, and reframe README to v0–v10 with Phase 3 = the v9 double-axis run. No
code changes.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-29 16:18:48 +08:00
Gahow Wang
a1370446fe 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:22:49 +08:00
Gahow Wang
980605474b 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 (DdpContext::init + train_rank). It would have caught the original bug:

- GATE A (world=1, ONE step — the deterministic scope): the p=0 FORWARD is
  byte-identical to no-dropout (ops::dropout(p=0) is a graph no-op) so the
  step loss is BIT-IDENTICAL (== 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), so the
  post-step params are checked against that tight ULP floor (< 1e-7).
- GATE A2 (world=2): p=0 matches a separate no-dropout baseline within NCCL's
  run-to-run ULP noise (< 1e-6, KI-5 — the all-reduce is not bit-reproducible
  on this PCIe box). Enabling dropout=0 doesn't perturb the DDP path beyond it.
- GATE B (world=2): a p=0.2 run's loss trace DIFFERS by > 1e-3 from p=0 —
  orders of magnitude above every noise floor here (~3e-2 observed). On the
  pre-T21 code the model stays in eval mode, so p=0.2 would be an identity and
  the trace would match p=0 at the noise floor — this gate fails. (Verified by
  simulating the bug: with model.train() removed, GATE B drops to 2.4e-7.)
- GATE C: a dedicated no-eval run ends with model.is_training() == true,
  direct proof that train_rank called model.train().
- 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:22:49 +08:00
Gahow Wang
81f3cf59e5 distributed: T21 — wire dropout into the DDP path (--dropout + model.train()) 
V9-PILOT caught a launcher-level integration gap: T18 wired dropout into
the single-GPU bin/train, but the DDP path never did. train_ddp had no
--dropout flag and never set cfg.dropout, and ddp.rs::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.

Fix, mirroring the single-GPU train/eval discipline:
- train_ddp.rs: add a --dropout <p> flag (default 0 = off, matching the
  prior behavior) and set cfg.dropout from it; log it when on.
- ddp.rs::train_rank: call model.train() at the start of each step (before
  the micro-batch loop). eval_loss() flips the model to eval mode and does
  not restore it, so re-asserting train() each step keeps dropout live
  across eval boundaries.

--dropout 0 (default) is bit-identical to the prior DDP path: cfg.dropout
stays 0 and ops::dropout(p=0) is a clone no-op regardless of training mode.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 21:08:17 +08:00
Gahow Wang
db70abe450 docs: T20 — Phase-2 systems-depth capstone (reframe README to two phases) 
Re-conclude xtrain as TWO phases now that Phase-2 (T14–T18) is merged on main:

README.md
- Status header: "complete (T1–T13) + scaling v0–v8" → "complete — two phases"
  (Phase 1 = from-scratch stack T1–T13 + v0–v8 scaling study; Phase 2 = the five
  deferred systems-stack features T14–T18).
- Crate table: note the Phase-2 additions (fused flash-attn + repeat_kv + dropout
  in autodiff; GQA + dropout in model; grad-accum in train; process-per-GPU
  launcher in distributed).
- Build-journey section retitled Phase 1 + Phase 2; replaced the run-on T14–T18
  prose with a structured "## Phase 2" summary (5 features + honest results:
  flash = mem-not-walltime win, GQA group-sum backward, grad-accum −74% mem,
  dropout × recompute bit-exact, T17 throughput-neutral falsification).
- Engineering lessons: T17 added as the THIRD profile-first falsification;
  reinforced honest-correctness with the Phase-2 hard gates + md5 b04fc9f9.
- Doc index: doc range …14-* → …17-*; KI status line (process-per-GPU CLOSED,
  KI-4 accepted tradeoff).

docs/evolution.md
- New "三·五、Phase 2 systems-depth synthesis": ties the 5 features into the
  per-axis (算法/架构/Infra/数据) narrative + the two integration notes.

docs/known-issues.md
- KI-4 reframed as a deliberately-accepted modeling tradeoff (保 xserv closed
  loop; T19 DROPPED), not "open".
- New integration notes: (a) DDP tests need --test-threads=1 (parallel deadlock);
  (b) fresh-train md5 is non-deterministic (atomicAdd reduction order) → the valid
  determinism gate is export re-determinism, not fresh-train reproduction.
- (process-per-GPU item was already CLOSED=measured no-op in T17.)

Docs-only; no code touched.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 18:11:47 +08:00
Gahow Wang
71b0a1621f docs: T17 process-per-GPU results — measured throughput-neutral 
Records the key empirical finding: process-per-GPU is statistically identical
to thread-per-GPU at this scale (thread 5.27x vs proc 5.31x @8, <1% noise; all
8 GPUs 95-99% util). The residual ~5.3x@8 non-linearity is the NCCL/PCIe
communication wall, NOT single-CUDA-context launch/cuBLAS serialization as the
old KI-5/T11 note speculated — measurement falsifies that hypothesis (same
methodology as T11 falsifying "bucket the all-reduce"). Correctness all green:
proc==thread loss 1.5e-7, cross-rank 1.2e-7, full regression + xserv md5
b04fc9f9 identical. Closes the process-per-GPU backlog item (measured no-op);
default training path unchanged. evolution.md Infra row + README T17 row +
known-issues entry.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 18:03:14 +08:00
Gahow Wang
4abb17383a test: process-per-GPU DDP correctness (ddp_proc.rs) 
Self-launching test: worker mode (XTRAIN_RANK set) trains on synthetic corpus
and dumps loss+params; launcher mode runs single-GPU baseline + thread-per-GPU
launch + spawns 2 worker processes, then asserts (a) proc loss == single-GPU
<1e-3, (b) cross-rank params <1e-6 (KI-5 ULP), (c) proc loss == thread-per-GPU
<1e-3. Run with --test-threads=1 (distributed harness property).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 17:48:52 +08:00
Gahow Wang
a188c8a277 distributed: train_ddp_mp bin (process-per-GPU launcher/worker) 
Dual-mode binary self-detecting via XTRAIN_RANK: launcher spawns one worker
per visible GPU forwarding full argv; worker rebuilds config from argv and runs
run_worker. CLI flags identical to train_ddp (thread-per-GPU, kept), so it
doubles as the before->after throughput driver. thread-per-GPU path untouched.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 17:48:52 +08:00
Gahow Wang
ffd548b80b distributed: process-per-GPU launcher + worker (proc.rs) 
torchrun-style process-per-GPU: launch_processes spawns one worker process per
GPU (re-exec current_exe with XTRAIN_{RANK,WORLD,LOCAL_RANK,NCCL_ID} env),
mints the ncclUniqueId once in the launcher and hex-injects it via env (no
shared FS/TCP, race-free). worker_env/run_worker read the env, bind the device
(own CUDA context), DdpContext::init + build_model + train_rank reused from T8
UNCHANGED. hex_encode/decode_unique_id are host-testable pure fns.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 17:48:43 +08:00
Gahow Wang
c470c627a7 docs: Phase T17 — process-per-GPU DDP design 
torchrun-style: launcher spawns N worker processes, each with its own CUDA
context; cross-process ncclUniqueId distributed via launcher-minted hex env
injection (race-free, no shared FS / TCP); train_rank + grad all-reduce reused
unchanged. Keeps thread-per-GPU path as regression baseline. ZeRO-1 dropped
(user scope decision).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 17:44:38 +08:00
Gahow Wang
2ff4573a31 docs: T15 GQA results + evolution row (模型架构) + README build-journey row 
Backfill docs/14-gqa.md gate table (dash5 numbers); add T15 evolution row +
cumulative 模型架构 line; README build-journey T15 row + Phase 2 prose + doc
index range (00..14).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 01:44:58 +08:00
Gahow Wang
39df0b40c1 gqa: fix kv-proj shape test param indices (embed,attn_norm precede wq) 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 01:38:42 +08:00
Gahow Wang
830d06ad01 gqa: real grouped-query attention (repeat_kv op + both SDPA paths + wiring + tests) 
- repeat_kv CUDA kernel: fwd head-block gather, bwd DETERMINISTIC group-sum (each
  kv head sums its group of query-head grads; no atomics) + Tensor/ops node.
- Config gains num_kv_heads (default = n_heads → MHA); wk/wv project to kv_dim;
  attention() repeat_kv-broadcasts K/V to nh heads before the UNCHANGED composed
  & flash SDPA → GQA on both paths. group=1 is identity → MHA bit-identical.
- --kv-heads flag on train/train_ddp/export_safetensors/greedy_sample; export
  writes real num_key_value_heads (xserv repeat_kv grouping aligned).
- Tests: repeat_kv grad-check (group>1 grad-sum + group=1 identity); model gqa.rs
  (GQA flash==composed fp32/bf16, group=1 bit-identical to MHA, kv-proj shape);
  parity_dump+parity.py GQA path (repeat_interleave) via XTRAIN_PARITY_KV_HEADS.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 01:37:37 +08:00
Gahow Wang
62b1cb5dc7 docs: Phase T15 — GQA design (repeat_kv broadcast op + backward grad-sum) 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 01:30:34 +08:00
Gahow Wang
4b6d3e0a79 test: flash+dropout cross-feature grad-check (Phase-2 integration) 
Add flash_plus_dropout_grad_check_fp32 to xtrain-model dropout tests: the two
orthogonal Phase-2 features (T14 flash-attn, T18 dropout) in the same model must
still grad-check. Both models run train-mode p=0.2 (identical masks, seed is
flash-independent) so the only delta is the SDPA reduction order — checked against
the flash-vs-composed tolerance.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 00:43:54 +08:00
Gahow Wang
c36cdf74d1 Merge t18-dropout into main 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

# Conflicts:
#	README.md
#	crates/xtrain-autodiff/tests/autograd.rs
#	crates/xtrain-model/src/model.rs
#	crates/xtrain-train/src/bin/train.rs
#	crates/xtrain-train/src/train_loop.rs
#	docs/evolution.md
 2026-06-18 00:41:41 +08:00
Gahow Wang
f26db882e5 Merge t16-grad-accum into main 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

# Conflicts:
#	README.md
#	docs/evolution.md
 2026-06-18 00:37:11 +08:00
Gahow Wang
9e958cb0f9 Merge t14-flash-attention into main 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 00:35:46 +08:00
Gahow Wang
80fafa1914 docs: T18 evolution row + README build-journey row (dropout) 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 00:06:06 +08:00
Gahow Wang
e625aa05dd dropout: wire into model (residual sites) + train/eval switch + flag (T18) 
Config.dropout (default 0). TinyTransformer gets a Cell<bool> training switch
(train()/eval()/with_training, default eval = safe) + a Cell<u64> step_seed bumped
once per training forward. forward_batched derives a per-layer block_seed (pure fn
of step_seed×layer) and block_forward derives two per-site seeds, inserting
ops::dropout at the attn and ffn sub-block outputs (before each residual). The
seed is a pure function of (step_seed, layer, site) so the checkpoint (T13)
recompute re-derives the same masks → grads stay exact. p=0 or eval → no dropout
node → graph bit-identical to pre-T18.

train_loop: model.train() per step (restored after eval flips to eval); eval_loss
runs model.eval(). bin/train: --dropout flag → cfg.dropout. Export/sampling run in
eval (default), so exported weights are dropout-free (xserv closed loop unaffected).

Model-level tests (dropout.rs): p=0 bit-identical to no-dropout (logits/loss/grads);
eval(p>0) == p=0 identity; train differs from eval + finite; recompute-with-dropout
grads match non-recompute (fp32 + bf16).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 00:05:32 +08:00
Gahow Wang
5eb27783f8 dropout: autodiff op + fixed-seed grad-check (T18) 
ops::dropout(x,p,seed): fwd runs Tensor::dropout, caches the mask in the backward
closure, bwd pushes dx=d⊙mask. p==0 returns x.clone() (no node) so the default
graph is unchanged. Tests in autograd.rs: fixed-seed finite-diff grad-check (mask
held constant across the ± perturbation — dropout is a fixed elementwise linear
map of x); E[out]≈input + keep-rate≈1-p over a seed sweep; p=0 kernel identity.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 00:05:32 +08:00
Gahow Wang
1fdd0c5002 dropout: device RNG kernel + Tensor fwd/bwd (T18) 
csrc/ops/dropout.cu: counter-based RNG (splitmix64 over seed^index) → fp32
uniform → Bernoulli(keep=1-p); fwd writes out=x⊙mask + an fp32 mask buffer
(per-element 1/(1-p) or 0); bwd applies the same mask (dx=d⊙mask). fp32 + bf16
activation variants (mask fp32 in both; uniform is dtype-independent so masks
match across precisions). Stateless → re-run with same seed = same mask (T13
recompute-safe). Registered in build.rs + FFI decls.

Tensor::dropout(p,seed)->(out,mask) and Tensor::dropout_backward(d,mask) wrap the
launches (contiguous F32/BF16, default stream, per-op sync via the kernels).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 00:05:18 +08:00
Gahow Wang
6b8c1e4e0f docs: Phase T18 — dropout design (device RNG + mask) 
Counter-based (stateless) RNG → Bernoulli(keep=1-p) mask, inverted 1/(1-p)
scaling at train, identity at eval. New autodiff `dropout` op (fwd generates +
applies mask, bwd applies the SAME cached mask). Wired at the two residual-path
sites (attn / ffn outputs); attention-probs dropout deliberately skipped (fused
SDPA doesn't materialise probs). Documents the RNG choice, per-site deterministic
seed (so T13 recompute reproduces the same mask), train/eval switch, p=0
bit-identity, and the acceptance gates.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-18 00:05:08 +08:00
Gahow Wang
8bd7db16e1 docs: T16 grad-accum results — evolution row + README build-journey 
dash5-verified gate numbers: accum=N bit-close to N× big batch (loss
8.5e-8 / grad 3.8e-5), accum=1 bit-identical (0.0), DDP+accum matches
single-GPU (5.7e-7), memory flat (same effective batch 64: 27.7GB big →
7.2GB accum, −74%), xserv closed loop md5-identical + token-identical.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:52:32 +08:00
Gahow Wang
b06b553f99 test: drop unused Var import in grad_accum 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:49:04 +08:00
Gahow Wang
abe5ceb913 test: grad-accum equivalence + accum=1 bit-identity + DDP+accum 
- grad_accum.rs: accum=N×B grads bit-close to a single N·B big batch;
  accum_steps=1 bit-identical (max|Δ|==0) to no-accum; real train() loop
  with accum tracks a big-batch baseline over 20 AdamW steps.
- ddp_correctness.rs: world=2 + accum=2 matches a single-GPU big batch of
  the same effective size (loss + cross-rank + vs-baseline).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:45:40 +08:00
Gahow Wang
7a03b0054a train+ddp: micro-batch gradient accumulation (--accum-steps) 
Accumulate grads over N micro-batches, then one AdamW step + zero_grad,
for an effective batch of N×micro at one micro-batch's activation cost.
Each micro-loss is scaled by 1/N before backward (the tape SUM-accumulates
the scaled grads) so the boundary grad equals a single step over an N×
batch. accum==1 skips the scale → bit-identical to the pre-T16 path.

DDP: the cross-rank all-reduce fires ONLY at the accumulation boundary
(intermediate micro-steps are local-only, no NCCL); the /world average is
orthogonal to the per-micro 1/N, so the boundary grad is the effective
global-batch mean. New --accum-steps flag in both train binaries; effective
batch is printed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:45:33 +08:00
Gahow Wang
d01fec6639 docs: Phase T16 — gradient accumulation design 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:41:17 +08:00
Gahow Wang
9064ced4c2 docs: T14 flash-attention results + evolution/README rows 
Fill in the design doc's measured results (grad-check, flash==composed,
PyTorch parity, peak mem -16%/-23%, tok/s tradeoff), add the T14 row to
evolution.md (算法/Infra) and the README build-journey table.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:34:10 +08:00
Gahow Wang
d217f4fbd3 perf: spread flash bwd dK/dV atomics across all threads 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:27:33 +08:00
Gahow Wang
4d7b69f8d4 perf: cache softmax weights in shared mem (drop hd× redundant expf) 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:24:56 +08:00
Gahow Wang
9b05f4f93f test: flash==composed bf16 uses robust mean/p99 metric (repo convention) 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:19:08 +08:00
Gahow Wang
c0f0b67510 test: eps=2e-3 for flash dQ/dK finite-diff (cuts f32 rounding term) 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:17:44 +08:00
Gahow Wang
80602099dc test: scale Q/K in flash grad-check for well-conditioned grads 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:17:04 +08:00
Gahow Wang
f38beb0346 test: flash finite-diff grad-check uses single-tile clean regime 
Match the trusted composed grad-check dims (seq=5<FA_TILE); the multi-tile
online-softmax path is gated by flash_bwd_matches_composed_bwd (seq=40),
sharper than finite-diff on the near-zero grads a long softmax produces.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:16:20 +08:00
Gahow Wang
01fb22d114 test: flash bwd vs composed bwd (sharper than finite-diff) 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:12:30 +08:00
Gahow Wang
5f3b81ac96 test+bins: flash grad-check, flash==composed, PyTorch parity, --flash flag 
autograd: flash_attention_batched_bwd (dQ/dK/dV finite-diff, seq>tile)
+ flash_matches_composed_fwd. model/tests/flash.rs: flash==composed
on-vs-off (logits/loss/every param grad), fp32 + bf16. parity_dump:
XTRAIN_PARITY_FLASH dumps the flash path for the same parity.py oracle
(PyTorch SDPA parity at B>1). train + train_ddp get the --flash flag.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:10:39 +08:00
Gahow Wang
0e20821633 autodiff+model: flash-attention op + --flash opt-in wiring 
ops::flash_attention autograd node (fwd caches O(N) logsumexp instead of
O(N²) probs; bwd via Tensor::flash_attention_backward). Model gets a
use_flash bool + with_flash(bool) builder; the SDPA core in attention()
picks ops::flash_attention vs ops::attention. flash threads through
block_forward so the recompute (T13) segment also runs flash. Default
off = composed path, graph unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:10:32 +08:00
Gahow Wang
326a6fadfe cuda: fused flash-attention kernel (fwd + flash-style bwd) 
csrc/ops/flash_attention.cu: a single fused fwd kernel (one block per
query row, streams KV in tiles of 32, online softmax — running max/sum
+ rescaled V accumulator, causal mask inlined, never materializes the
[bh,S,S] scores) writing out[bh,S,hd] + the per-row logsumexp L (O(N),
saved for backward). flash-style bwd: recompute scores from Q/K/V + L,
collapse the softmax Jacobian with D[i]=ΣdO·O, dQ owned per row, dK/dV
atomicAdd across rows. Tensor::flash_attention / flash_attention_backward
wrap them (bf16 upcasts Q/K/V→f32 for the kernel, same fp32-softmax
policy as composed).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:10:25 +08:00
Gahow Wang
65a2264227 docs: Phase T14 — fused flash-attention design 
Design doc for the hand-written single fused flash-attention kernel:
online softmax tiled over KV, NEVER materializing the [bh,S,S] score
matrix; flash-style backward (recompute scores from saved logsumexp +
D=ΣdO·O, dQ/dK/dV). Opt-in --flash; composed T10 path stays default.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 23:10:16 +08:00
Gahow Wang
31cc2bf745 docs: capstone README — full-stack + scaling study (v0-v8) writeup 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 16:17:26 +08:00
Gahow Wang
511f35d40c docs: run v8 — dim1024 capacity helps (val 2.98) 
v8 = capacity-axis A/B: freeze the v6/v7 2.255B FineWeb-edu subset, scale
dim768→dim1024 (core 127M→226M, +78%) via bf16 + T13 activation recompute.
8-GPU DDP, 2.36B tok (1.05 ep), ~129K tok/s (recompute tax), ~5h.

Result (same FineWeb val, v6/v7/v8 comparable): v6 3.0652 / v7 3.0149 /
v8 2.9801. Capacity helps — v8 (1.05ep) beats v6 at the same ~1ep by 0.085
AND beats v7 (smaller model, 1.45ep more old data) by 0.035 ⇒ v6/v7 were
partly capacity-limited, scaling capacity > repeating old data. But the gain
is only ~3% (same magnitude as the data-axis single-step lever), and v8's
val was still descending at the end (not saturated).

Meta-finding: every single-axis lever (data-volume v5/v7, breadth v6,
capacity v8) is now ~3%/lever ⇒ broad diminishing returns; to progress,
scale capacity AND data together (Chinchilla, reproduced at toy scale).

- docs/runs/08-v8-fineweb-edu-dim1024.md: full capacity experiment + v7-vs-v8 samples
- docs/runs/README.md: +v8 row, v9 proposal
- docs/evolution.md: +T13 infra row, +v8 scaling row, capacity-axis & diminishing-returns notes

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 15:12:01 +08:00
Gahow Wang
0150263055 perf: KI-3 fixed — dim1024 batch32 fits, mem 31.1→14.6GB, tok/s 39.7K→31.5K 
Per-block activation recompute (T13) measured on dash5 (1× RTX 5090 32GB, bf16,
batch32 seq256, steady-state):

- Correctness (exact, hard gate): recompute on-vs-off grads are BIT-IDENTICAL —
  fp32 AND bf16: loss / logits / every param grad max rel = 0.00e0 (not "within
  tol", exactly equal). Full suite green with recompute on/off; DDP loss-match
  5.67e-7; DDP+recompute 2-rank descends 11.079→6.010.
- dim768 (18L/24h ffn2048, core 127M): peak mem 31144→14562 MiB (−53%), tok/s
  39.7K→31.5K (−20%, the extra-forward tradeoff, in the predicted 20–35% band).
- dim1024 (18L/32h ffn2730, core 226M): recompute OFF OOMs (hits 32100/32607
  MiB → OutOfMemory); recompute ON fits at 16596 MiB, ~23K tok/s, converges.
  → KI-3 payoff achieved: dim1024 batch32 unblocked, v8 can proceed.

Fill docs/12 bench table; mark KI-3 FIXED in docs/known-issues.md.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 09:50:29 +08:00
Gahow Wang
69c5f07359 docs: Phase T13 — activation recompute 
Design doc for per-block gradient checkpointing (KI-3): the no-tape forward +
recompute-on-backward design, the `checkpoint` primitive, per-block wrapping,
the exactness/correctness argument (same kernels + inputs → identical grads),
composition with bf16+DDP+batched, and the verification plan (on-vs-off grad
gate + memory/throughput before→after, dim1024-fits). Bench table left as TBD
to fill after the dash5 run.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 09:45:16 +08:00
Gahow Wang
f202351be5 model: per-block activation recompute (--recompute) 
Wrap each transformer block's forward in the checkpoint primitive when
recompute is enabled (Phase T13 / KI-3). To make the block forward a pure
segment fn (no `&self` borrow, so it can re-run in the backward closure),
extract the block body + its helpers (linear / norm_gamma / attention /
swiglu_mlp) into free functions parameterised by (cfg, compute_dtype) and add
`Block::block_params()` (the 11 leaves in the params() per-block order). The
non-recompute path calls `block_forward` directly — identical graph to before.

- `TinyTransformer::with_recompute(bool)` builder (opt-in; default off keeps the
  unchanged tape / bit-identical numerics).
- `--recompute` flag wired into bin/train and bin/train_ddp (DDP: each rank
  checkpoints independently).

Correctness gate: tests/recompute.rs builds two identical models (recompute
on/off), runs the same batched loss+backward, and asserts the forward logits,
the loss, and EVERY parameter grad match within tight fp tol — parameterised
over fp32 and bf16 (T12 composition).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 09:42:42 +08:00
Gahow Wang
c396b39483 autodiff: checkpoint primitive (recompute-on-backward) 
Add `xtrain_autodiff::checkpoint::checkpoint(segment_fn, input, params)`, a
higher-order autograd node (à la torch.utils.checkpoint) for activation
recomputation (Phase T13 / KI-3):

- forward: run `segment_fn` on detached leaves so its internal ops are NOT
  recorded on the outer tape; keep only the output value (the local sub-tape —
  and thus the segment's intermediate activations — drops immediately). The
  checkpoint node's parents are [input, ..params].
- backward: re-run `segment_fn` from the saved input + (unchanged) param values
  into a fresh local tape, seed the recomputed output with the upstream grad,
  backprop, then push the recovered input/param grads to the real parents. Local
  tape drops at the end → recomputed activations freed.

Exact by construction (same deterministic kernels, same inputs) → grads match
the non-checkpointed path. Composes with bf16 (T12, same path on recompute) and
DDP (T8, per-rank).

Supporting change: `Var::backward_seeded(seed)` — backward from an explicit
non-scalar upstream grad (the segment output is generally not a scalar);
`backward()` is now the scalar wrapper that seeds ones.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 09:42:31 +08:00
Gahow Wang
9c557f0609 docs: run v7 — FineWeb subset near-ceiling at dim768 (val 3.01) 
v7 = same arch as v4/v5/v6 (dim768/18L, bf16, 8-GPU DDP global 256),
trained the SAME 2.255B-token FineWeb-edu subset to 1.45 epoch (vs v6's
1.02), best FineWeb val 3.0149 (v6 3.0652). Exported + archived to
registry v7-fineweb-edu-dim768, serves in xserv (coherent expository
English, ~v6 quality).

Key finding: more epochs of the SAME subset gave only ~0.05 val drop and
the curve flattened (~step 44000) with no sampling quality gain → the
2.255B FineWeb subset is near its ceiling at dim768. Same class as v5's
TinyStories data-volume saturation: repeating old data has thin margins;
true further gains need FRESH shards (more diverse tokens), as v6's
corpus-swap (which raised the ceiling) showed.

Adds docs/runs/07-v7-*.md; updates docs/runs/README.md (+v7 row, intro
saturation note, v8 proposal) and docs/evolution.md (+v7 row, dataset-axis
ceiling note).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-17 03:55:47 +08:00
Gahow Wang
b4bb426d48 docs: run v6 — FineWeb-edu graduation (val 3.07, new distribution) 
第一版脱离 TinyStories：纯 FineWeb-edu 真实网页文本(2.255B 语料)，架构同
v4/v5(dim768/18L, core 127.43M)，8 卡 DDP bf16，2.29B tok/1.02ep，~1.9h
@218K tok/s。train 11.03→3.14，best/final FineWeb val 3.0652。

方法论：FineWeb val(3.07) 与 v0–v5 的 TinyStories val(~1.1) 不可比——真实
网页熵高，~3.0 是预期非回退；判据是采样质量 + transfer eval。

- 新增 docs/runs/06-v6-fineweb-edu-dim768.md：数据管线(scripts/fineweb_to_txt.py)
  / 架构(同 v4/v5,隔离数据变量) / 超参 / 结果(val 单调降无走平=未饱和) /
  方法论说明 / transfer eval(v6→TinyStories val 2.75 vs v5 native 1.11,纯通用
  数据对窄分布有代价) / v5-vs-v6 同提示词采样对比(v6 写真实说明文 vs v5 一律
  掉进小故事)
- README 对比表加 v6 行(val 单独标注分布) + 换轴说明 + v7 提案
- evolution.md scaling 表 v6 行定稿 + 数据轴 TinyStories→FineWeb-edu 毕业说明

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 22:21:43 +08:00
Gahow Wang
88bec270af docs: evolution overview — per-milestone changes across algorithm/arch/infra/dataset axes 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 19:30:52 +08:00
Gahow Wang
7e5ea9976b data: FineWeb-edu parquet->txt prep script (Scaling v6) 
v6 broadens data from TinyStories to FineWeb-edu (HuggingFaceFW/fineweb-edu
sample/10BT) while freezing the v4/v5 arch. scripts/fineweb_to_txt.py streams
the parquet text column row-group by row-group and joins docs with
<|endoftext|> so xtrain's existing Corpus loader (gpt2 BPE, u16 cache) handles
it unchanged. Corpus .txt/.parquet/.u16.bin stay dash5-only (gitignored).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 19:04:45 +08:00
Gahow Wang
579365f4a0 docs: run v5 — TinyStories saturation at dim768 (val 1.11) 
设计文档 05-v5-tinystories-dim768.md（中文，xserv 风格）：数据 2.49B tok/5.33ep、
架构同 v4(净测数据变量)、bf16 8 卡 global 256、train 11.07→1.06 best val 1.1102。
核心发现「数据天花板」：v4(1.54ep)1.169→v5(5.33ep)1.110 仅 ↓5% 且末段 val 走平
⇒ TinyStories 在 dim768/127M-core 近饱和，v6 该换轴（更大模型/更广语料，非更多 TinyStories）。
xserv BF16 服务 3/3 prompt 逐 token 一致。

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 17:56:25 +08:00
Gahow Wang
8a1e29543b run: v5 archive + export (dim768, bf16, 5.33ep, val 1.11) 
v0–v5 对比表加 v5 行 + tokens-trained / epoch 两列，让 TinyStories 数据饱和可见
（v4→v5 同 arch 数据 ×3.5 仅 val ↓5% 且末段走平）。下一档提案改为 v6 换轴。
导出 201 tensors + RUN.md 存入 dash5 registry v5-tinystories-dim768（checkpoint/safetensors 不入库）。

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 17:56:25 +08:00
Gahow Wang
5b7dde1736 test: bf16 test reads f32-cast logits (forward now returns bf16) 
The `keep bf16 logits` change made forward_batched return bf16 logits
in bf16 mode; the bf16 test's host read must cast to f32 first.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 14:29:24 +08:00
Gahow Wang
320c1ae4fb perf: KI-2 FIXED — dim768 bf16 fits batch 32, tok/s 31.5K→40.8K 
bf16 mixed precision (fp32 master) solves the v4 dim768 fp32 batch-32
OOM and speeds up the now-compute-bound dim768 GEMMs (dash5 1× RTX
5090 32GB, dim768/18L/24h×32 ffn2048 seq256, steady-state):

  config          batch  peak mem   tok/s   fits 32GB
  fp32            16      27.2 GB    31.5K   yes
  bf16            16      19.3 GB    35.5K   yes   (-29% mem / +13% tok/s)
  fp32            32      —          —       OOM
  bf16            32      31.1 GB    40.8K   yes   (+29% vs fp32-b16)

Verified on dash5: fp32 suite green at tight tol + xserv export md5
bit-identical to registry; bf16 looser-tol (loss 1.2e-4, logits p99
6.8e-3, grad 1.0e-2) + 150-step convergence tracks fp32 (3.984 vs
3.988); 2-GPU bf16 DDP at per-rank batch 32 trains cleanly.

Mark KI-2 FIXED; fill docs/11 results.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 14:28:20 +08:00
Gahow Wang
48922cb628 perf: keep bf16 logits (no persistent fp32 logits buffer) 
At vocab 50257 the logits tensor [B*S, vocab] is ~1.6GB fp32 at batch
32 — held across the whole backward. Keep it bf16: cross_entropy
upcasts the bf16 logits to fp32 internally (transient) + caches fp32
probs, and its backward casts dx back to bf16 to chain into the
bf16 lm_head matmul backward. The sampler casts bf16 logits→f32 before
the host argmax/softmax. Halves the persistent logits activation.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 14:20:48 +08:00
Gahow Wang
30db62d8f2 docs: Phase T12 — bf16 mixed precision design 
docs/11-bf16-mixed-precision.md: the AMP split (bf16 linears +
activations, fp32 master / norms / softmax / RoPE / CE, no loss
scaling), the cast-op bridge, module layout, and the dual
verification gate (fp32 unchanged + bf16 looser-tol + convergence +
mem/throughput). Memory/throughput before->after to be filled from
the dash5 bench.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 14:15:02 +08:00
Gahow Wang
0a2a4dcaa8 train: --bf16 flag (fp32-master AMP) + bf16 correctness test 
- TinyTransformer::with_compute_dtype(BF16): embedding stays fp32
  master then casts to bf16; each linear casts its fp32 weight to bf16
  on the fly; logits cast back to fp32 for cross-entropy. Default F32
  reproduces the v0-v4 forward graph bit-for-bit.
- --bf16 flag on bin/train and bin/train_ddp (off by default).
- tests/bf16.rs: same fp32 master weights run fp32 vs bf16; assert
  loss/logits/grads within a loose bf16 tol, no NaN, and grads are
  fp32 (master untouched).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 14:14:55 +08:00
Gahow Wang
b0086b5214 autodiff: bf16 mixed-precision path (fp32 master via cast op) 
Tensor ops dispatch on dtype: fp32 branch unchanged (bit-identical),
bf16 branch routes matmul/attention through GemmEx and elementwise
through the bf16 kernels. Norm/softmax/RoPE/cross-entropy upcast to
fp32 around the existing fp32 kernels (standard AMP: reductions/loss
fp32, matmuls bf16). Transposes route bf16 through fp32 (pure layout).

New autodiff `cast` op is the AMP bridge: forward downcasts a fp32
master leaf to bf16 for the matmul; backward upcasts the bf16 grad
back to fp32. So the fp32 leaf accumulates an fp32 grad and AdamW /
clip / DDP all-reduce stay fp32 and completely unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 14:14:48 +08:00
Gahow Wang
d05115ddf3 cuda: bf16 cuBLAS GemmEx (16BF in/out, fp32 accum) + cast kernels 
Add the bf16 compute primitives for T12 mixed precision:
- DType::BF16 (half::bf16 as TensorDType), 2 bytes.
- cublasGemmEx / cublasGemmStridedBatchedEx FFI + CUDA_R_16BF /
  CUBLAS_COMPUTE_32F constants (values per xserv gemm.rs).
- cublas::gemm_ex / gemm_ex_strided_batched: same row-major⟺col-major
  transpose algebra as sgemm, bf16 in/out, fp32 accumulation.
- csrc/ops/cast.cu: f32<->bf16 cast + bf16 elementwise (add/mul/scale/
  silu(+dx)/add_bias/sum_rows), each load->fp32->compute->store bf16.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 14:14:39 +08:00
Gahow Wang
511ceebbb3 docs: KI-2 trigger — dim768 fp32 batch-32 OOM 
v4 surfaced the concrete bf16 trigger: dim768 fp32 OOMs at per-rank batch 32
(global 256) in 32GB, forcing per-rank 16 (global 128). bf16 (halve activation
mem) would restore the batch-256 sweet spot. Record it on KI-2; keep KI-2 as
the backlog item it is (still deferred).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 13:14:42 +08:00
Gahow Wang
ff79fee3c5 docs: run v4 — TinyStories, dim768, val 1.17 
Design doc docs/runs/04-v4-tinystories-dim768.md (data 720.9M tok ~1.54ep /
arch dim768/18L core 127.4M vs v3 / hparams 22000 steps, global batch 128
per-rank 16, seq 256, lr 6e-4->6e-5 warmup 1100 + cosine, clip 1.0, world=8
DDP fp32 / results train 11.07->1.14, best val 1.1690, ~145K tok/s 8-GPU /
v3->v4 improvement: val 1.30->1.17 + side-by-side samples). Notes that this run
validated T11's caching allocator at dim768 multi-GPU and that dim768 fp32
batch-32 OOM is the bf16 trigger. Update docs/runs/README.md comparison table
to v0/v1/v2/v3/v4 and the next-rung proposal to v5.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 13:14:37 +08:00
Gahow Wang
734e119db3 run: v4 archive + export (dim768, 8-GPU DDP, val 1.17) 
v4 scaling run finished: dim768/18L, core 127.43M (total 204.63M), trained
720.9M tokens (~1.54 epoch) on 8x RTX 5090 DDP fp32, ~145K tok/s, ~84 min,
best val 1.1690. Checkpoint archived to registry
(~/projects/tiny-models/v4-tinystories-dim768/) and exported to xserv HF Qwen3
safetensors (201 tensors, BF16); xserv serves it and matches xtrain greedy
token-for-token on all 3 fixed prompts (40 tok).

Add `greedy_sample` bin: load a trained ckpt with its arch flags and print
xtrain's own greedy continuations for the fixed run prompts, so they can be
diffed against xserv's greedy on the exported weights (the per-run token-match
check). Same model/config/init scheme as bin/train.rs + bin/export_safetensors.rs.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 13:14:28 +08:00
Gahow Wang
f85bd4d276 perf: KI-5 FIXED — single-GPU 40K->93K tok/s, DDP scaling 1.3x->5x@8 
Device caching/pool allocator removes the per-op cudaMalloc serialization that
was the real DDP bottleneck (and a single-GPU cost). Measured on dash5 (8x RTX
5090, dim384/12L, per-rank batch 32, seq 256, steady-state tok/s):

  single-GPU: 40226 -> 92638 tok/s  (~2.3x)
  DDP scaling (global batch 32*world):
    world  before        after
      1    39801 1.00x    92385 1.00x
      2    47229 1.19x   146821 1.59x
      4    52854 1.33x   269867 2.92x
      8    48996 1.23x   461270 4.99x

8-GPU absolute throughput 49K -> 461K tok/s (9.4x); nvidia-smi shows all 8 GPUs
at 95-99% util during the run (KI-5 saw only 1-2/8 busy). Loss trajectories are
bit-identical before/after (10.9026->4.8453). xserv closed loop green: re-export
of the v3 ckpt is md5-identical to the registry safetensors and xserv serves it.

Mark KI-5 FIXED in docs/known-issues.md with before/after table; fill in the
design doc's measured numbers. Residual ~5x@8 (not perfectly linear) is the
~7% all-reduce + 8-GPU PCIe/launch overhead; process-per-GPU is the next lever
if v4 needs higher linearity.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 11:15:02 +08:00
Gahow Wang
4c3f332f64 docs: Phase T11 — caching allocator 
Design doc for the device caching/pool allocator (KI-5 re-diagnosis recap, size
classes, per-device + thread-safety, Drop->return, transparency/correctness
argument, why skip-memset uninit is deferred, dual verification gates). Before/
after numbers filled after dash5 measurement.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 11:04:11 +08:00
Gahow Wang
b7104e2cb7 test: loosen flaky DDP cross-rank assertion to <1e-6; scale to world=8 
The cross-rank `max|p0-p1| == 0.0` check is flaky on this PCIe-only box: NCCL's
all-reduce is not bit-reproducible run-to-run across ranks (algorithm/chunk
choice is unstable), so cross-rank params can differ by a few ULP (observed
<=1.2e-7) even with identical init + averaged grads. The load-bearing gate is the
loss-trajectory match (~5.7e-7); a tight <1e-6 tolerance is the honest invariant.

Also extend ddp_throughput_scaling to include world=8 for the KI-5 before/after
scaling table.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 11:04:11 +08:00
Gahow Wang
28801fbfe5 cuda: device caching allocator (pool GpuBuffer alloc) 
Every tape op allocates its output via Tensor::zeros -> GpuBuffer::alloc ->
cudaMalloc, a synchronous process-serialized driver call. Under the single-
process thread-per-GPU DDP model the rank threads' hundreds of per-step allocs
serialize through the driver (KI-5 root cause); it costs single-GPU too.

Add a per-device, size-classed caching pool: GpuBuffer::alloc serves from a
free-list (request rounded up to a size class so repeating training shapes
reuse buffers), only cudaMalloc on a miss; Drop returns the buffer to the pool
instead of cudaFree. Thread-safe via a global registry keyed by device id with
each device's free-list behind its own Mutex (registry lock held only to clone
out the per-device Arc<Mutex<_>>, so rank threads don't contend across devices).
The buffer records its alloc-time device so Drop returns to the right pool.

Transparent: physical capacity may be rounded up, but len()/memset/copy bounds
all use the requested length, so the rounded tail is never read and numerics are
unchanged. zeros() still memsets (reused buffers hold stale bytes).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 11:04:02 +08:00
Gahow Wang
d422c68704 docs: KI-5 — correct cross-rank divergence attribution (pre-existing flaky) 
The ~1-ULP cross-rank param divergence is NOT caused by coalescing: the
original ungrouped all-reduce is itself run-to-run nondeterministic on
this box (6 reruns: cross-rank diff {0, 0, 5.96e-8, 5.96e-8, 1.19e-7,
1.19e-7}), so the T8 test's `max|p0-p1| == 0.0` assertion is flaky here
(passes ~1/3 of runs) independent of T11. Diffs are ≤1.19e-7 (a few ULP,
numerically benign; loss-match stays ~6e-7). Noted as a follow-up to
loosen the assertion to a tight tolerance; coalescing was reverted purely
because it gives ~0 scaling benefit.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 09:42:13 +08:00
Gahow Wang
84092fb28d docs: KI-5 re-diagnosis — all-reduce is NOT the DDP bottleneck (T11) 
T11 set out to coalesce/overlap the gradient all-reduce per the original
KI-5 hypothesis. Profiling on dash5 (8× RTX 5090, dim384, per-rank batch
32, seq 256) falsifies that hypothesis:

  - grad all-reduce is only ~6-7% of each step;
  - per-rank fwd+bwd inflates ~linearly with world (136→780 ms for the
    SAME per-rank workload) and dominates;
  - coalescing the ~150 per-tensor all-reduces into one grouped/flat
    launch gives ~0 scaling gain AND breaks cross-rank bit-identity
    (max|p0-p1| 0.0 → 1.49e-8), violating the T8 correctness gate — so
    the coalescing commit (b8b5821) was reverted.

Real bottleneck (NOCOMM=1 still inflates; util shows 1-2 of 8 GPUs busy
at a time; CPU not starved; per-thread default stream doesn't help):
single-process thread-per-GPU ranks serialize on the single CUDA
context's per-op cudaMalloc / driver calls. Fix direction (out of T11
scope): a caching/pool allocator, or process-per-GPU. Recorded in
docs/known-issues.md with the measured table; KI-5 stays Open.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 09:40:45 +08:00
Gahow Wang
88c2c15768 Revert "dist: coalesce grads into buckets for all-reduce (KI-5)" 
This reverts commit b8b58212dc.
 2026-06-16 09:39:38 +08:00
Gahow Wang
b8b58212dc dist: coalesce grads into buckets for all-reduce (KI-5) 
Replace the per-parameter eager all-reduce (~150 tiny serial NCCL calls
for dim512, DDP's dominant cost after T10's batched forward) with a
coalesced bucketed all-reduce: pack grads into a few large contiguous
scratch buffers, all-reduce each bucket once (fused via ncclGroupStart/
End), fold the 1/world average into one per-bucket scale, unpack back.

The packed buffer is the concatenation of the grad tensors, so NCCL's
element-wise sum over a bucket equals the per-tensor sums — bit-identical
to the un-bucketed path; only launch/latency overhead is removed. DDP
cross-rank param identity + loss-match are preserved.

Adds xtrain_cuda::device::copy_d2d (cudaMemcpy D2D) for the pack/unpack.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 09:09:44 +08:00
Gahow Wang
a78502e0f0 docs: run v3 — TinyStories, dim512, val 1.30 
Per-run design doc docs/runs/03-v3-tinystories-dim512.md (data 245.8M tok full
TinyStories ~0.53 epoch / arch dim512 16L core 67.13M vs total 118.59M, what
changed vs v2 / hyperparams 30000 steps batch 32 seq 256 lr 6e-4→6e-5 warmup
1500 + cosine clip 1.0 single-GPU batched via T10 / results train 10.91→1.40
best val 1.3027 ~26K tok/s / improvement vs v2 1.71→1.30 with side-by-side
samples). Notes v3 validated T10 batched forward at scale and avoided KI-5 by
staying single-GPU; v4 proposal + open levers (KI-2/3/4/5, data ladder).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 03:37:45 +08:00
Gahow Wang
64b2a8c09e run: v3 archive + export (dim512, single-GPU batched, val 1.30) 
v3 trained (30000 steps × batch 32 × seq 256 = 245.8M tok, ~0.53 epoch),
single-GPU batched via T10 (~26K tok/s, ~2.65h). Archived to registry
~/projects/tiny-models/v3-tinystories-dim512/ (xtrain.ckpt + config.json +
model.safetensors BF16 179 tensors + tokenizer.json + RUN.md) and served in
xserv (loads 16L/dim512 qwen3, 2/3 prompts token-match xtrain greedy; 3rd
diverges on BF16 drift as in v1/v2).

best/final val 1.3027 (beats ~1.4 target). val ladder on the same held-out
1M-token set: v0 3.80 / v1 2.58 / v2 1.71 / v3 1.30. T10 (batched forward)
validated at scale (KI-1 root cause = launch-bound, not all-reduce); single-GPU
avoids KI-5. Update docs/runs/README.md comparison table.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 03:37:36 +08:00
Gahow Wang
9a25616a30 docs: Phase T10 — batched forward 
docs/09-batched-forward.md: the launch-bound diagnosis recap, the
[B*S,dim]-flatten + fused batched-attention design (RoPE per-seq position +
causal masking inline in softmax), the attention forward/backward via
strided-batched GEMM, autograd implications, the looped-split/merge dead-end
post-mortem (1127 tok/s, host round-trips), verification methods + before→after
throughput, and the v3 recommendation (per-rank batch 16-32, single/small world
until KI-5 bucketed all-reduce lands).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 00:44:50 +08:00
Gahow Wang
4ccab0fb42 perf: KI-1 fixed — GPU util 0-15%→37-54%, tok/s 1653→25627 (15.5x) 
Mark KI-1 (single-sequence launch-bound, the root cause of "DDP weak scaling")
FIXED by the T10 batched forward. dim384/12L, batch 16, seq 256, 1 GPU,
back-to-back A/B:

  before (single-seq): ~1653 tok/s, GPU util 0-15%, ~3 GB
  after  (batched):    25627 tok/s (batch16) / 40263 (batch32),
                       util 37% mean / 54% peak, ~10 GB
  → single-GPU ~15.5x (batch16) / ~24x (batch32); util 0-15% → 37-54%.

A single GPU at batch 32 (40K tok/s) now beats the old 4-GPU setup (3163) ~12x.
The v3 falsification history (larger batch doesn't help a single-seq design) is
kept. DDP residual weak scaling is a NEW, higher-level bottleneck batching
exposes (eager all-reduce of all params each step) → recorded as KI-5
(bucketed/overlapped all-reduce), out of T10 scope.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 00:44:43 +08:00
Gahow Wang
25b032445d train: real batched step (drop loop+SUM) 
Feed a real batch of B sequences as ONE batched forward/backward, replacing the
"loop B times + let the tape SUM grads + clip ×1/B" hack. CE mean over B*S rows
is already the batch-mean loss, so backward yields the batch-mean gradient
directly → clip pre-scale = 1.0.

DDP stays equivalent: each rank runs one batched forward over its b_local =
B_global/world sequences (local-mean grad Σ_local/b_local); all_reduce_average
(sum across ranks /world) = Σ_global/B_global = global batch-mean → clip
pre-scale 1.0. The ddp_correctness single-GPU baseline batches the same way.
DDP loss matches single-GPU 5.7e-7, cross-rank params bit-identical (0.0).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 00:44:33 +08:00
Gahow Wang
5353b38402 model: batched forward [B,S] 
forward_batched(ids[B*S], batch)/loss_batched: run B equal-length sequences as
ONE forward over flattened [B*S] ids, so every linear is one big [B*S,dim] GEMM.
Attention reshapes to [B*nh,S,hd], runs the fused batched causal SDPA (per-seq
mask + RoPE period=S, no cross-sequence attention), writes back [B*S,dim]. The
old per-(batch,head) loop + host-round-tripping split/merge_heads + the additive
causal_mask leaf are gone. forward(ids[seq]) is now forward_batched(ids,1), so
the sampler / inference path (batch=1) is unchanged.

+batched_ids_tensor helper. New batched.rs test: batched forward == looped
single-sequence (logits identical 0.0, grads 6.4e-4, loss identical). PyTorch
parity now exercises B>1 (B=2,S=4): loss 5e-8, logits 6.9e-6, all 25 param
grads within rtol — verifying per-seq RoPE position + per-seq causal masking.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 00:44:25 +08:00
Gahow Wang
7821bd9c34 autograd: batch dim for ops (flatten linears, batched attention) 
Add the batched-forward primitives. Linears/norms/elementwise/embedding/CE
already act on flat [rows,dim], so they work unchanged on [B*S,dim]; only
attention + RoPE need sequence awareness:

- RoPE: kernel takes a `period` (= seq len) so position = row % period, i.e.
  per-sequence position on a flattened batch (period == tokens = single seq).
- Fused batched causal attention: new `Tensor::attention`/`attention_backward`
  + ops node, running QKᵀ and PV as cublasSgemmStridedBatched over the B*nh
  (sequence,head) blocks (new sgemm_strided_batched binding) and a causal
  softmax kernel (scale + per-row causal mask inline) — the whole attention is
  3 launches regardless of B*nh, no per-head/per-seq loop, no host round-trip.
- transpose_4d12 ([B,S,nh,hd] <-> [B,nh,S,hd]) to lay out the batched heads.

grad-checks: new batched-rope, transpose_4d12, batched-attention dQ/dK/dV all
pass finite-diff (attn dK 1.5e-2, dQ 7.5e-3, dV 2.9e-4; rest tighter) alongside
the existing 12.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-16 00:44:15 +08:00
Gahow Wang
d2a585c5cb docs: KI-1 re-diagnosed in v3 — larger batch does NOT fix DDP weak scaling 
v3 tested the documented mitigation (raise global_batch to amortize the
per-step all-reduce). Isolated back-to-back A/B on 4× RTX 5090, dim384/12L,
seq256:

  global_batch 32 (8/rank)  → 3163 tok/s
  global_batch 256 (64/rank)→ 3200 tok/s   (8× batch, +1.2%, within noise)

8× larger batch = 1/8 the all-reduces per token, yet no speedup → all-reduce
is NOT the bottleneck. GPU util 0–15%, mem ~2–3 GB/32 GB → the workload is
launch-bound: the single-sequence model design (each sequence its own tiny
forward/backward, per-op kernel launches) starves the GPU, and batching only
adds proportionally more serial launches. Real fix is batched (multi-sequence)
forward so GEMMs fill the GPU — a T4/T5 autograd/model change, not a batch knob.
Bucketed/overlapped all-reduce stays deferred (no value until launch-bound is
fixed). KI-1 kept Open with the corrected root cause.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 23:20:26 +08:00
Gahow Wang
bf679f6f1f docs: run v2 — TinyStories, dim384/12L, DDP 4-card (val 1.71) 
Scaling run v2 design doc + comparison-table update. v2 = dim384/12L/12h
SwiGLU ffn1536 (core 28.32M, total 66.92M), trained 4500 steps / ~36.9M
tokens on full TinyStories (reused v1 u16 cache) via NCCL DDP across 4
RTX 5090s. Best val 1.7055 (train 10.89→1.72), a clear jump over v1 2.58
and v0 3.80. Exported to xserv (135 BF16 tensors) and archived in the
dash5 registry; xserv greedy token-matches xtrain on 2/3 fixed prompts
(3rd diverges late under BF16 drift). Records the DDP weak-scaling caveat
(global batch too small → all-reduce dominates) → links docs/known-issues
KI-1; v3 proposal applies KI-1's fix (much larger global batch).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 22:38:31 +08:00
Gahow Wang
c87a0bc44e docs: known-issues / perf backlog — KI-1 DDP weak scaling at small global batch 
Surfaced by v2 (world=4, global_batch=32): ~3593 tok/s, no speedup vs v1
single-GPU. Root cause + proposed fixes recorded; also consolidates deferred
T7 items (bf16, activation recompute) and the large-vocab modeling note.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 19:56:58 +08:00
Gahow Wang
7090b475fb train: bring DDP trainer to parity with bin/train (val + checkpoint + cache + arch) 
The T8 DDP path now matches the single-GPU `bin/train`: CLI-tunable arch
(scaling-ladder rung), the cached token-id stream (`load_cached`), held-out
val-loss eval + best-val checkpointing, and LR warmup→cosine. Rank 0 owns the
val corpus and runs the no-grad eval / writes the best checkpoint (params are
bit-identical across ranks). The eval/checkpoint logic is reused from
`xtrain-train` (`eval_loss`, `checkpoint::save`) rather than duplicated.

- DdpConfig gains eval_every / eval_batches / ckpt_path.
- train_rank takes `valid: Option<&Corpus>` and returns DdpResult
  (losses + evals + best_val); launch threads the val corpus to rank 0 only.
- bin/train_ddp reworked to the bin/train CLI (positional tokenizer/corpus +
  --dim/--heads/--head-dim/--layers/--ffn/--steps/--batch/--seq/--max-lr/
  --val-tokens/--eval-every/--ckpt), reusing the u16 cache.
- DDP correctness test updated to the new signatures (semantics unchanged).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 19:34:40 +08:00
Gahow Wang
264660527f docs: run v1 — TinyStories full, dim256 
docs/runs/01-v1-tinystories-dim256.md + docs/runs/README.md comparison table.
v1: full TinyStories train (468.3M tok, u16-cached) + dim256/8L (core 8.39M).
Same-held-out-set val loss v0 3.8050 → v1 2.5847 (−1.22); v1 samples coherent
stories vs v0's "mommy's mommy's mommy" loop; exports + serves token-identical
in xserv. Single RTX 5090, ~25.9 min, ~3310 tok/s.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 19:09:46 +08:00
Gahow Wang
ec8114ecbc train: --eval-ckpt eval-only mode (v0-vs-v1 same-set val loss) 
Expose eval_loss() and add a --eval-ckpt <path> branch to bin/train: load an
existing checkpoint into a model of the given arch and score it on the held-out
val split, then exit. Lets v0 and v1 be measured on the identical validation set
(the acceptance metric) without a separate eval binary.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 18:44:40 +08:00
Gahow Wang
e44e50ef78 data: full TinyStories + tokenized-id cache, val loss, CLI arch 
- Corpus::load_cached: tokenize the (large) corpus ONCE, cache the id stream to
  <corpus>.u16.bin (gpt2 vocab 50257 < 65536 → exact u16), read cache on reruns.
- Corpus::split_tail: hold out a tail slice as a validation corpus.
- train(): take an optional valid corpus + eval_every/eval_batches; periodic
  deterministic val-loss eval that checkpoints the BEST val model; returns
  TrainResult{train_losses, evals, best_val}. T6 fixed-cadence path preserved.
- bin/train + bin/export_safetensors: read architecture (--heads/--head-dim/
  --layers/--ffn) + opt knobs (--steps/--batch/--seq/--max-lr/--val-tokens/
  --eval-every) from CLI flags; defaults reproduce the v0-baseline tiny config.
- gitignore the multi-GB corpus + *.u16.bin caches + *.ckpt (dash5-only).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 18:34:48 +08:00
Gahow Wang
15f1e526c7 train: parameterize model size (scaling ladder) 
Add Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn) so the model
size is a tunable rung instead of a hardcoded tiny config, and Config::core_params()
(num_params minus the two vocab×dim tables) — the figure the ladder is sized
against (the 50257-vocab embed+lm_head adds a fixed ~25M that is not capacity).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 18:34:39 +08:00
Gahow Wang
8981cf7982 docs: T9 verification results (xserv == xtrain, dash5) 
Capture the closed-loop run: train (loss 10.84->3.59) -> export (47 tensors,
BF16) -> xserv dump-logits + greedy. Top-1 + top-11 token order identical,
logits within ~1e-2 (BF16-vs-f32 drift), greedy generation token-for-token
identical across two prompts.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:37:46 +08:00
Gahow Wang
e246c3bec2 export: dump_logits bin for xserv-vs-xtrain comparison 
xtrain-side top-k next-token logit dump (f32 forward, same model/config/ckpt
as the exporter) mirroring xserv's dump-logits, so the closed-loop check can
compare both sides numerically for the same prompt + weights.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:36:41 +08:00
Gahow Wang
18c2229b4b docs: Phase T9 — export to xserv 
Architecture diff table (xtrain TinyTransformer vs xserv qwen3.rs), the
QK-norm structural decision + BF16 acceptance criterion, the tensor-name +
layout mapping table, and the dash5 closed-loop verification recipe.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:33:32 +08:00
Gahow Wang
1c76573cb4 export: safetensors + config.json for xserv qwen3 
New bin export_safetensors: load an xtrain checkpoint, map every param to its
HF Qwen3 tensor name, transpose 2D projection weights [in,out]->[out,in]
(1D norms + [vocab,dim] embed/lm_head kept), cast to BF16 (xserv's qwen3
forward is BF16-only), and write config.json + model.safetensors + a copy of
the gpt2 tokenizer.json. Sized exactly like bin/train.rs. safetensors 0.5 to
match xserv. GPU body gated behind not(no_cuda).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:33:26 +08:00
Gahow Wang
7a4f69e430 model: add per-head QK-norm (Qwen3-compat) for xserv export 
xserv's Qwen3 forward unconditionally applies per-head RMSNorm to Q and K
(q_norm/k_norm, shape [head_dim]) before RoPE — even gamma=1 is a real RMS
divide, not identity. xtrain never had this, so an exact xserv<->xtrain loop
was structurally impossible. Add it (reusing the 2D rms_norm op on the
[seq*nh, hd] head rows, inserted between reshape and rope to mirror
qwen3.rs's order) so the trained model is genuinely Qwen3-compatible.

params() inserts q_norm,k_norm after wv; num_params() counts them; the
PyTorch parity refs (parity.py / adamw_parity.py) + their name lists add the
same step so the dumps stay self-consistent.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:33:19 +08:00
Gahow Wang
ad82e8bf92 dist: lengthen scaling bench so NCCL init amortizes 
30-step bench charged the one-time NCCL init + 4 model builds (present at world=4,
absent at world=1) against the wall clock, understating steady-state scaling
(in-loop tok/s already showed ~53k at 4 GPUs). Bump to 150 steps.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:18:23 +08:00
Gahow Wang
818f76a18f dist: drop unused import; relax DDP-vs-single-GPU param tolerance 
dash5 verify: loss trajectory matches single-GPU to max_rel 1.16e-7 and
cross-rank params are bit-identical (0.0), but DDP-vs-single-GPU per-param rel
diff is ~2.8e-3 after 20 AdamW steps — expected, since the two differ only in
gradient summation order (fp add isn't associative) and that rounding compounds.
Bump check (c) 1e-3 -> 1e-2 (a/b stay tight). Also remove an unused DType import.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:17:31 +08:00
Gahow Wang
0131f05b26 docs: Phase T8 — distributed data parallel 
Design doc for the NCCL DDP path: comm bootstrap (rank-0 UniqueId + grouped
CommInitRank), thread-per-GPU launch model (Var is !Send), all-reduce-then-
local-step scheme (in-place fp32 AllReduce on .grad() + /world, each rank steps
its own GpuAdamW), why params stay consistent (NCCL bit-identical reduce + same
init/state), batch sharding math vs single-GPU, verification plan + scaling
table. Lists TP/PP/ZeRO/bf16-comm as out-of-scope follow-ups.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:15:49 +08:00
Gahow Wang
cf5e3987df dist: multi-rank launcher + ddp acceptance test 
bin/train_ddp: spawn one thread per visible GPU (CUDA_VISIBLE_DEVICES selects
the set), NCCL all-reduce gradients each step, train the tiny transformer on
TinyStories; doubles as the throughput driver (prints global tok/s). no_cuda
build keeps a stub main.

tests/ddp_correctness: (1) 2-rank DDP vs single-GPU over the same synthetic data
-> loss trajectory max_rel < 1e-3, cross-rank params bit-identical (==0.0), DDP
vs single-GPU params rel < 1e-3; (2) 1/2/4-GPU throughput table on a fixed
per-GPU workload. Gated #[cfg(not(no_cuda))], auto-skips with < 2 GPUs.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:15:41 +08:00
Gahow Wang
163f567c80 dist: ddp all-reduce + sharded batch 
DDP training step (train_rank) on top of DdpContext: each rank advances the
SAME RNG, draws the whole global batch, and runs forward+backward only on its
shard (i % world == rank) so the union over ranks is the single-GPU batch in the
same order. After backward, all-reduce-average the device grads, then finish the
mean with clip(pre_scale = 1/b_local) -> Sigma_global/B_global, identical to the
single-GPU clip(1/B). Each rank then runs its own GpuAdamW.step; same init +
same averaged grad + same optimizer state keep params bit-identical across ranks.

Adds a deterministic build_model (same LCG init as bin/train) shared by ranks +
baseline, a per-step loss all-reduce for the reported global-mean loss, and the
thread-per-GPU launch() helper (thread::scope; Var graph is !Send so each rank
builds its model thread-locally, only UniqueId/config/&Corpus cross threads).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:15:29 +08:00
Gahow Wang
e27df50ca9 dist: nccl ffi + comm bootstrap 
New crate xtrain-distributed (mirrors xserv-distributed): hand-written NCCL
FFI (GetUniqueId / CommInitRank / AllReduce / CommDestroy / Group{Start,End},
ncclUniqueId passed by value per the NCCL ABI) and a safe DdpContext wrapper —
rank 0 mints the UniqueId, every rank inits its communicator under a group, and
all_reduce_average_grads in-place AllReduce(sum)s each param's .grad() device
buffer then scales by 1/world (reuses T7's scale_inplace kernel). AllReduce runs
on the null stream so it orders with the model's kernels (no extra barrier).

build.rs follows the per-crate convention: no nvcc -> no_cuda cfg (crate
compiles to empty, cargo check passes host-side); with nvcc, links -lnccl
-lcudart like xserv-distributed's build.rs.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:14:56 +08:00
Gahow Wang
5e8add2a41 docs: Phase T7 — performance 
Design doc for the T7 fp32-preserving speedups: cuBLAS matmul fwd/bwd
(row-major⟺col-major layout), GPU AdamW + GPU grad-norm (no per-step
param/grad roundtrip), drop per-op sync + device memset. Includes the
verification table (regression suite green + tok/s 2770→8220 ~3x), the
deferred bf16/recompute follow-up rationale, and the T8 all-reduce note.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 17:00:29 +08:00
Gahow Wang
a842e432b5 perf: streams / drop per-op sync 
Default-stream kernels run in order and every host read goes through a
stream-ordered cudaMemcpy (to_device), so the per-op cudaDeviceSynchronize
after each kernel was pure overhead — remove all 21 in tensor.rs. Host
data is still correctly ordered by the D2H memcpy that reads it.

Also zero op-output buffers with cudaMemset (device-side, async) instead of
a blocking H2D memcpy of a host zero buffer on every allocation — that
copy was itself a hidden per-op sync point.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:56:17 +08:00
Gahow Wang
8070c1949a perf: make xtrain-cuda a regular dep of xtrain-optim (GPU AdamW) 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:53:52 +08:00
Gahow Wang
b0e397ca81 perf: GPU AdamW + grad-norm 
Eliminate the per-step GPU↔host roundtrip of every parameter/gradient.

- optim.cu: adamw_step (m/v on device, in-place param update), sumsq_accum
  (block-reduced global grad sum-of-squares), scale_inplace.
- GpuAdamW: device m/v state per param; step launches the kernel reading
  each param's .grad() and rewriting the param buffer in place — no host
  roundtrip. Host AdamW kept as the torch-parity reference.
- clip_grad_norm_gpu: device sum-of-squares reduction (only the scalar norm
  comes back), in-place rescale of grads by pre_scale·clip_factor.
- train_loop: use GpuAdamW + clip_grad_norm_gpu.
- test: GPU AdamW vs host reference parity (max abs err < 1e-6).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:53:09 +08:00
Gahow Wang
0e5c7d22e2 perf: cuBLAS matmul fwd/bwd 
Route Tensor::matmul and matmul_backward through cuBLAS Sgemm instead of
the hand-written tiled kernel. fp32 → same GEMM up to rounding order, so
the T3 cuBLAS tolerance and downstream grad-checks are preserved.

- cublas.rs: thread-local persistent handle + row-major sgemm helper with
  transpose flags (col-major⟺row-major as the T3 oracle does).
- matmul_backward: dA/dB via cuBLAS OP_T, dropping the two transpose
  kernels + their allocations the T3 version ran.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:48:35 +08:00
Gahow Wang
5df1d4d57b test: resolve real_training corpus default via CARGO_MANIFEST_DIR 
cargo runs tests with cwd = crate dir, so the bare relative default
data/tinystories-valid-3mb.txt didn't resolve. Anchor it to the repo root via
CARGO_MANIFEST_DIR so the test runs out of the box (still overridable with
XTRAIN_CORPUS).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:41:12 +08:00
Gahow Wang
2f8118fda9 test: tighten AdamW parity (f32 reference, 10 steps, allclose tol) 
The loss trajectory already matched torch.optim.AdamW (worst relerr ~2e-4),
but the float64 torch reference diverged per-weight from the f32 GPU training
after the model memorised the batch (flat region: weights underdetermined,
loss identical). Fixes: run the torch reference in float32 (match engine
precision), shorten to 10 steps (weights still well-determined), and compare
final params with an allclose-style rtol+atol metric (a pure relative metric is
misleading on near-zero weights).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:34:18 +08:00
Gahow Wang
29b4d30b6c docs: Phase T6 — training loop 
Design doc for the T6 training stack: Goal / Module Layout / Key Design
Decisions (AdamW math + decoupled WD, LR schedule, global-norm grad clip with
batch averaging, checkpoint format, data pipeline + xserv tokenizer reuse,
sampler) / 验证方法 (AdamW parity, checkpoint round-trip, real training, host
unit tests).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:30:14 +08:00
Gahow Wang
22b7434b23 test: AdamW PyTorch parity + checkpoint round-trip + real training 
Acceptance tests (GPU-gated not(no_cuda), run on dash5):
- adamw_parity_dump.rs + adamw_parity.py: build the tiny model with fixed init,
  run N AdamW steps on a fixed batch, dump the loss trajectory + final params;
  the Python side rebuilds the identical model and runs torch.optim.AdamW with
  matched lr/wd/betas/eps, comparing trajectory + final params within rtol.
- checkpoint_roundtrip.rs: train a few steps, save, load into a fresh model with
  a DIFFERENT init, assert identical logits/loss on a fixed input.
- real_training.rs (#[ignore], --release): train on TinyStories for a bounded
  budget; assert loss drops substantially and print greedy samples.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:30:06 +08:00
Gahow Wang
77a82bfeee train: loop + checkpoint save/load + sampler + train binary 
Training loop (train_loop.rs): sample batch_size sequences, forward loss +
backward (tape SUMs grads), clip_grad_norm with ×1/batch averaging, AdamW step
with scheduled lr, zero_grad; logs loss/lr/gnorm/tok-s and checkpoints
periodically; returns the loss trace.

Checkpoint (checkpoint.rs): flat little-endian dump of params() in order
(magic/version/count + per-param ndim/dims/f32 data); load_into validates and
overwrites a matching model's params via set_value (exact f32 round-trip).

Sampler (sample.rs): autoregressive greedy / temperature generation — re-runs
forward on the growing prefix (model is single-sequence, RoPE pos=row).

bin/train.rs: end-to-end entry — load tokenizer+corpus, train a tiny 4-layer
model for a bounded budget, checkpoint, print samples. no_cuda stub keeps it
buildable on a GPU-less host.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:29:58 +08:00
Gahow Wang
7d84a64f5c data: gpt2 bpe via xserv-tokenizer + TinyStories corpus + lr schedule + grad clip 
New xtrain-train crate scaffold. Data pipeline reuses xserv's from-scratch
GPT-2/Qwen BPE via a path-dep (../../../xserv/crates/xserv-tokenizer, resolves
on both ~/projects and dash5 /opt/wjh/projects): Corpus::load tokenizes the
corpus into one id stream and samples fixed-length (input, target) next-token
windows (LCG-seeded, reproducible). Trims a range-downloaded file to whole
stories (<|endoftext|> boundaries).

Also the host-only training math: LrSchedule (linear warmup + cosine decay)
and global L2 grad-norm + clip scale, each with a local unit test.

Corpus: data/tinystories-valid-3mb.txt — first ~3MB of TinyStories-valid
(fetched on dash5 via hf-mirror.com; HF direct unreachable). Substitution
noted: a real TinyStories subset, not the full set.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:29:32 +08:00
Gahow Wang
f22429f5b8 optim: hand-written AdamW (decoupled weight decay + bias correction) 
New xtrain-optim crate. AdamW with per-param m/v moments keyed by params()
index, global bias correction, and decoupled weight decay (matches
torch.optim.AdamW). Split into a pure-host step_host (flat f32 buffers,
unit-testable on a GPU-less host) and a step(&[Var]) wrapper that round-trips
each param value/grad through the GPU tensor (gated not(no_cuda)). Per-step lr
argument leaves room for an LR schedule.

Host unit test checks the update against an independent reference recurrence
over 20 steps and the pure-decay (g=0) boundary.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:28:23 +08:00
Gahow Wang
8565565647 docs: Phase T5 — tiny transformer 
Goal / Module Layout / Key Design Decisions (multi-head layout via
reshape+transpose_3d01+split/merge_heads, embedding gather/scatter-add,
x@W convention, causal mask, params API, overfit methodology) / 验证方法
with the dash5 results (grad-checks, overfit 2.82->0.004, PyTorch parity).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:09:30 +08:00
Gahow Wang
603c85e1e0 model: silence torch parity warning (read loss before backward) 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:09:30 +08:00
Gahow Wang
3366f30c4d model: PyTorch parity harness (weight dump + equivalent torch model) 
parity_dump.rs (#[ignore] fixture generator) dumps the model's exact
weights, ids, forward logits, loss, and per-param grads after one
backward. parity.py rebuilds the IDENTICAL model in PyTorch (same x@W
convention, RoPE rotate_half pos=row, RMSNorm, SwiGLU, causal SDPA),
runs fwd+bwd, and compares logits + every grad within rtol.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:07:30 +08:00
Gahow Wang
e3912c2380 model: tiny RoPE+RMSNorm+SwiGLU transformer + overfit test 
New crate xtrain-model: a from-scratch decoder built entirely from the
autodiff op set.
- Config (tiny: dim=32, 2 layers, 2 heads, head_dim=16, ffn=64).
- TinyTransformer: embedding -> N x {pre-RMSNorm -> multi-head causal
  attention (RoPE, additive causal mask, per-head SDPA) -> residual;
  pre-RMSNorm -> SwiGLU MLP -> residual} -> final RMSNorm -> LM head.
  x@W weight convention (engine GEMM is plain A@B); dim=n_heads*head_dim.
- params()/zero_grad-able leaves for the optimizer; param_to_host export.
- overfit test: char-level bring-up (embedded text -> vocab -> shifted
  targets), minimal hand-written GD (p -= lr*grad) memorises one fixed
  batch -> loss ~0 + greedy argmax matches targets. End-to-end fwd+bwd
  correctness signal. Gated #![cfg(not(no_cuda))].

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:05:20 +08:00
Gahow Wang
0acfa5df11 ops: grad-check the T5 structural ops 
Finite-diff grad-checks (same L=sum(W∘out) harness as autograd.rs) for
embedding (incl. repeated ids), reshape, transpose_3d01, transpose_2d,
and split/merge_heads round-trip. Gated #![cfg(not(no_cuda))].

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:05:20 +08:00
Gahow Wang
7fb1a29057 ops: embedding/reshape/transpose/split-merge-heads fwd+bwd 
Phase T5 structural ops on top of the T4 set, needed to assemble the
tiny transformer:
- embedding: gather rows by I32 ids (CUDA kernel) / scatter-add backward
  (atomic, so repeated ids accumulate). csrc/ops/model.cu + ffi.
- reshape: contiguous metadata-only view (Tensor::reshape), no kernel.
- transpose_3d01: [a,b,c]->[b,a,c] for the multi-head layout (kernel).
- autograd nodes: embedding/reshape/transpose_3d01/transpose_2d, plus
  split_heads (->Vec<Var>) / merge_heads for per-head attention.
- tape: Var::zero_grad + set_value so a hand-written GD step can update
  params and clear grads between steps.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 16:05:09 +08:00
Gahow Wang
777f3c7949 docs: Phase T4 — autograd engine 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:53:55 +08:00
Gahow Wang
e7ce504b1f ops: differentiable autograd nodes + per-op grad-check tests 
ops.rs wraps each Tensor op as a Var node with its backward closure (forward
caches captured by move). swiglu = mul(silu(gate), up); attention is composed
(matmul+scale+softmax+matmul), no fused kernel. tests/autograd.rs grad-checks
every op via the L=sum(W∘out) template, plus a fan-out grad-accumulation test
(dL/dx=4x) and an end-to-end composed-attention grad-check (dQ/dK/dV). Adds
xtrain-cuda dev-dep for device selection in tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:53:55 +08:00
123 changed files with 45037 additions and 167 deletions
Expand all files Collapse all files

8
.gitignore vendored
View File

@@ -9,3 +9,11 @@
# Claude Code runtime state # Claude Code runtime state
/.claude/ /.claude/
# Large scaling-run corpora + tokenized id caches live on dash5 only, never in
# git (the small data/tinystories-valid-3mb.txt is committed as a fixture).
/data/tinystories-train.txt
/data/fineweb-edu.txt
/data/*.parquet
*.u16.bin
*.ckpt

163
Cargo.lock generated
View File

@@ -2,6 +2,15 @@
# It is not intended for manual editing. # It is not intended for manual editing.
version = 4 version = 4
[[package]]
name = "aho-corasick"
version = "1.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301"
dependencies = [
"memchr",
]
[[package]] [[package]]
name = "cc" name = "cc"
version = "1.2.64" version = "1.2.64"
@@ -41,6 +50,18 @@ dependencies = [
"zerocopy", "zerocopy",
] ]
[[package]]
name = "itoa"
version = "1.0.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
[[package]]
name = "memchr"
version = "2.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4"
[[package]] [[package]]
name = "proc-macro2" name = "proc-macro2"
version = "1.0.106" version = "1.0.106"
@@ -59,6 +80,88 @@ dependencies = [
"proc-macro2", "proc-macro2",
] ]
[[package]]
name = "regex"
version = "1.12.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f1292b7759ae1cb9ec195452d1390a074f0cd8541ab7a5a8c31cd6db45d4a6ba"
dependencies = [
"aho-corasick",
"memchr",
"regex-automata",
"regex-syntax",
]
[[package]]
name = "regex-automata"
version = "0.4.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f"
dependencies = [
"aho-corasick",
"memchr",
"regex-syntax",
]
[[package]]
name = "regex-syntax"
version = "0.8.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
[[package]]
name = "safetensors"
version = "0.5.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cc0cdb7198d738a111f6df8fef42cb175412c311d0c4ac9126ff4e550ad1a0e8"
dependencies = [
"serde",
"serde_json",
]
[[package]]
name = "serde"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
dependencies = [
"serde_core",
"serde_derive",
]
[[package]]
name = "serde_core"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
dependencies = [
"serde_derive",
]
[[package]]
name = "serde_derive"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "serde_json"
version = "1.0.150"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9"
dependencies = [
"itoa",
"memchr",
"serde",
"serde_core",
"zmij",
]
[[package]] [[package]]
name = "shlex" name = "shlex"
version = "2.0.1" version = "2.0.1"
@@ -88,10 +191,20 @@ version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
[[package]]
name = "xserv-tokenizer"
version = "0.1.0"
dependencies = [
"regex",
"serde",
"serde_json",
]
[[package]] [[package]]
name = "xtrain-autodiff" name = "xtrain-autodiff"
version = "0.1.0" version = "0.1.0"
dependencies = [ dependencies = [
"xtrain-cuda",
"xtrain-tensor", "xtrain-tensor",
] ]
@@ -102,6 +215,36 @@ dependencies = [
"cc", "cc",
] ]
[[package]]
name = "xtrain-distributed"
version = "0.1.0"
dependencies = [
"xtrain-autodiff",
"xtrain-cuda",
"xtrain-model",
"xtrain-optim",
"xtrain-tensor",
"xtrain-train",
]
[[package]]
name = "xtrain-model"
version = "0.1.0"
dependencies = [
"xtrain-autodiff",
"xtrain-cuda",
"xtrain-tensor",
]
[[package]]
name = "xtrain-optim"
version = "0.1.0"
dependencies = [
"xtrain-autodiff",
"xtrain-cuda",
"xtrain-tensor",
]
[[package]] [[package]]
name = "xtrain-tensor" name = "xtrain-tensor"
version = "0.1.0" version = "0.1.0"
@@ -112,6 +255,20 @@ dependencies = [
"xtrain-cuda", "xtrain-cuda",
] ]
[[package]]
name = "xtrain-train"
version = "0.1.0"
dependencies = [
"half",
"safetensors",
"xserv-tokenizer",
"xtrain-autodiff",
"xtrain-cuda",
"xtrain-model",
"xtrain-optim",
"xtrain-tensor",
]
[[package]] [[package]]
name = "zerocopy" name = "zerocopy"
version = "0.8.52" version = "0.8.52"
@@ -131,3 +288,9 @@ dependencies = [
"quote", "quote",
"syn", "syn",
] ]
[[package]]
name = "zmij"
version = "1.0.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"

4
Cargo.toml
View File

@@ -4,6 +4,10 @@ members = [
"crates/xtrain-cuda", "crates/xtrain-cuda",
"crates/xtrain-tensor", "crates/xtrain-tensor",
"crates/xtrain-autodiff", "crates/xtrain-autodiff",
"crates/xtrain-model",
"crates/xtrain-optim",
"crates/xtrain-train",
"crates/xtrain-distributed",
] ]
[workspace.package] [workspace.package]

232
README.md
View File

@@ -1,50 +1,210 @@
# xtrain # xtrain
A from-scratch **Rust + CUDA** LLM **training** engine — the sibling of A from-scratch **Rust + CUDA** LLM **training** engine — the sibling of **xserv** (the
[xserv](https://github.com/) (the inference side). GPU-first. inference side). A learning project: hand-write the entire training-systems stack
(autograd → backward → optimizer → training loop → distributed → mixed precision →
gradient checkpointing), then use it to run a multi-version **scaling study** that maps
the data-vs-capacity frontier for a tiny model.
The goal is to learn the full training-systems stack by hand: autograd / backward > **Status: complete — three phases.**
passes / optimizers (AdamW) / the training loop / distributed logic. Heavy lifting > **Phase 1** = the from-scratch full stack (T1–T13) + an 8-version scaling study (v0–v8):
is borrowed where it makes sense (GEMM → cuBLAS after a hand-written version, > hand-write the whole training-systems stack, then map the data-vs-capacity frontier.
multi-GPU comms → NCCL, tokenizer → reused from xserv), but the core is written > **Phase 2** = systems-stack depth (T14–T18): hand-write the five deferred training-stack
from scratch. The target architecture is a tiny modern transformer > features — fused flash-attention, real GQA, gradient accumulation, process-per-GPU DDP,
(RoPE + RMSNorm + SwiGLU, ~1–30M params) whose forward aligns with xserv's Qwen3, > dropout. **Phase 3** = one Chinchilla-style double-axis run (v9): dim1280 true-GQA +
so the backward passes map one-to-one onto xserv's existing forward kernels and > 6.01B FineWeb tokens, validating the v8 conclusion that data and capacity must scale
trained weights can flow back into xserv. > 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/).
## Status ---
Bootstrapping (P0). This repo currently contains only the project skeleton and a ## What got built (from scratch, by hand)
working Rust↔CUDA build chain, verified by a trivial vector-add CUDA kernel.
## Layout 7 crates, no ML framework — only cuBLAS / NCCL / safetensors as deliberate "heavy-lifting"
borrows, the rest hand-written CUDA + Rust:
``` | crate | what's hand-written |
xtrain/ |---|---|
├── Cargo.toml # workspace | `xtrain-cuda` | CUDA Runtime FFI, RAII `GpuBuffer`, **caching/pool allocator**, cuBLAS (sgemm + bf16 GemmEx) bindings |
├── csrc/ # CUDA sources (.cu) | `xtrain-tensor` | tensor (dtype/shape/strides/storage), elementwise + transpose + embedding kernels |
│ └── test/vecadd.cu # trivial element-wise vector-add (smoke test) | `xtrain-autodiff` | **tape autograd engine** (grad accumulation), per-op backward, finite-diff grad-check, **checkpoint** (recompute) primitive, **fused flash-attention** (online-softmax) fwd/bwd, **`repeat_kv`** broadcast (GQA), **`dropout`** (counter-based device RNG + mask) |
└── crates/ | `xtrain-model` | tiny **Qwen3-style** transformer (RoPE + RMSNorm + QK-norm + SwiGLU), batched forward, **GQA** (`num_kv_heads<num_heads`), residual/MLP **dropout** |
└── xtrain-cuda/ # CUDA Runtime FFI + build.rs (nvcc → sm_120) | `xtrain-optim` | hand-written **AdamW** (host + GPU kernels) |
├── build.rs # compiles csrc/*.cu via the `cc` crate, links cudart | `xtrain-train` | training loop, LR schedule, grad clip, **gradient accumulation**, checkpoint, BPE corpus + cache, samplers, safetensors export |
├── src/ # ffi / error / device / memory | `xtrain-distributed` | **NCCL DDP** (thread-per-GPU + torchrun-style process-per-GPU launcher / cross-process `ncclUniqueId`, all-reduce) |
└── tests/ # vecadd smoke test
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.
## The build journey — Phase 1 (T1–T13) + Phase 2 (T14–T18)
Each phase: design doc + implementation + tests + a scoped commit (see [`docs/`](docs/) and
[`docs/evolution.md`](docs/evolution.md) for the per-axis changelog). **Phase 1 (T1–T13)**
hand-built the stack and fixed the four real bottlenecks; **Phase 2 (T14–T18)** went back to
hand-write five deferred training-stack features — see the Phase-2 summary below the table.
| phase | what | result |
|---|---|---|
| T1–T2 | Rust↔CUDA build chain · tensor abstraction | vector-add verified · roundtrip |
| T3–T4 | hand GEMM fwd/bwd + finite-diff · **tape autograd** + 11 op backwards | grads vs cuBLAS 1e-7 / finite-diff |
| T5 | tiny transformer (RoPE+RMSNorm+SwiGLU) | overfit + **PyTorch parity** |
| T6 | AdamW + training loop + checkpoint · GPT-2 BPE + TinyStories | first **coherent English** |
| T7 | cuBLAS + GPU optimizer + drop syncs | ~3× (2.7K→8.5K tok/s) |
| T8 | NCCL DDP | multi-GPU (weak scaling, then) |
| T9 | + per-head **QK-norm** (Qwen3-compat) + safetensors export | **xserv closed loop, token-identical** |
| **T10** | **batched multi-sequence forward** (fixes KI-1) | **single-GPU 15–24×**; MFU 0.4%→14% |
| **T11** | **device caching allocator** (fixes KI-5) | single-GPU 2.3×; **8-GPU 461K tok/s** |
| **T12** | **bf16 mixed precision** (fp32 master, fixes KI-2) | dim768 OOM solved; −29% mem |
| **T13** | **activation recompute** / checkpointing (fixes KI-3) | dim1024 fits; grads bit-identical |
| **T14** | **fused flash-attention** kernel (online softmax, no materialized N×N; opt-in `--flash`) | peak mem −16%@1k / −23%@2k seq; flash==composed (grads/PyTorch) |
| **T15** | **grouped-query attention** (`num_kv_heads<num_heads`; `repeat_kv` broadcast feeds both SDPA paths; backward sums each kv head's group; `--kv-heads`) | repeat_kv grad-check + **group=1 bit-identical to MHA**; GQA flash==composed; PyTorch GQA B>1; **xserv closed loop with real `num_key_value_heads`** token-identical |
| **T16** | **gradient accumulation** (`--accum-steps`; DDP all-reduces only at the boundary) | equiv to N× big batch (grad 3.8e-5); same effective-64 batch 27.7GB→7.2GB (−74%) |
| **T17** | **process-per-GPU** DDP (torchrun-style: 1 worker process / CUDA context per GPU; launcher mints `ncclUniqueId` → hex env injection; `train_rank` reused unchanged; thread-per-GPU path kept) | proc==thread loss 1.5e-7, cross-rank 1.2e-7, xserv md5 identical · **measured no-op on throughput**: thread 5.27× vs proc 5.31×@8 (8 GPUs 95–99% util) → residual non-linearity is NCCL/PCIe, *not* CUDA-context serialization (falsifies the old KI-5 hypothesis) |
| **T18** | **dropout** (hand counter-based device RNG + mask, inverted scaling, train/eval switch) | fixed-seed grad-check; **p=0 bit-identical**; recompute-safe |
The four performance fixes (T10–T13) each removed a real bottleneck — see
[`docs/known-issues.md`](docs/known-issues.md) — which is where **Phase 1** closed.
## Phase 2 — systems-stack depth (T14–T18)
Phase 1 fixed bottlenecks; Phase 2 went back to hand-write the five training-stack features
that had been **explicitly deferred** earlier (project's actual goal = learn the whole stack).
Each is opt-in, kept the default path **bit-identical**, and held a **hard correctness gate**:
- **T14 · fused flash-attention** ([`docs/13-flash-attention.md`](docs/13-flash-attention.md)) —
a single hand-written kernel: **online (streaming) softmax, tiled over KV, never materializes
the `N×N` scores**; flash-style backward recomputes scores + the `D=ΣdO·O` Jacobian
simplification for dQ/dK/dV. Opt-in `--flash`, default off. **The win is memory, not
wall-clock**: peak activation **−16%@seq1024 / −23%@seq2048** (grows with seq, since the
`N×N` never lands), but **~2.3× slower** at head-dim 64 (a hand kernel can't beat cuBLAS
tensor-cores on a small head). Gate: flash == composed (loss rel `0.0`, grad `4.4e-5`),
PyTorch B>1 `7.9e-6`.
- **T15 · real GQA** ([`docs/14-gqa.md`](docs/14-gqa.md)) — `num_kv_heads < num_heads` via a new
`repeat_kv` **broadcast op** that copies K/V `group = nh/num_kv` times to feed **both**
(composed + flash) SDPA paths **unchanged**; its **backward is a deterministic group-sum**
(no atomics) collapsing each kv head's query-head group. Gate: `repeat_kv` grad-check +
**group=1 bit-identical to MHA** (regression guard); **xserv closed loop with real
`num_key_value_heads`** token-identical.
- **T16 · gradient accumulation** ([`docs/15-grad-accum.md`](docs/15-grad-accum.md)) — N
micro-steps scaled by `1/N` accumulate on the tape, then one AdamW step; DDP **all-reduces
only at the accumulation boundary**. Decouples effective batch from activation memory: same
effective batch 64, big-batch **27.7GB (OOM)** → accum 4×16 **7.2GB (−74%)**. Gate: `accum=N`
≡ one N× batch (grad `3.8e-5`); `accum=1` bit-identical.
- **T18 · dropout** ([`docs/17-dropout.md`](docs/17-dropout.md)) — a **stateless counter-based
device RNG** (Philox-style bit-mix) → Bernoulli mask, inverted `1/(1−p)` scaling in train,
identity in eval; wired at the two residual sites (attn-out, mlp-out). Stateless RNG is what
makes it **compose bit-exactly with T13 activation recompute** — the backward re-run
regenerates the *same* mask from `(seed, index)`. Gate: fixed-seed grad-check; **p=0
bit-identical**.
- **T17 · process-per-GPU** ([`docs/16-process-per-gpu.md`](docs/16-process-per-gpu.md)) — a
torchrun-style launcher: one worker process + CUDA context per GPU, the launcher mints one
`ncclUniqueId` and **hex-injects it into each child's env** (no shared FS/TCP, no race); the
worker reuses the T8 `train_rank` **unchanged**. Built and **correct** (proc vs thread loss
`1.5e-7`, cross-rank `1.2e-7`, xserv md5 identical) — but **measured throughput-neutral**:
8-GPU thread **491K (5.27×)** vs proc **493K (5.31×)**, `<1%`. This **falsifies** the
long-standing KI-5/T11 hypothesis that thread-per-GPU's shared CUDA context caused the
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
Same Qwen3-style architecture throughout; we scaled **dim** and **data** and read out val
loss (full per-run detail in [`docs/runs/`](docs/runs/)).
| ver | data (trained tok / epoch) | dim / core params | val loss | axis explored |
|---|---|---|---|---|
| v0–v3 | TinyStories (↑) | 32→512 / 41K→67M | 3.80 → 1.30 | bring-up |
| v4 | TinyStories 1.54ep | 768 / 127M | 1.17 | — |
| v5 | TinyStories 5.33ep | 768 / 127M | **1.11** | **data volume → saturates** |
| 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
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.
2. **Corpus > more-of-the-same.** Swapping TinyStories → FineWeb-edu (v5→v6) was a
*qualitative* jump: the model went from only-writes-kid-stories to writing genuine
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.
## Efficiency — throughput & MFU
The throughput story is the perf-infra report card (RTX 5090, bf16/fp32):
| | v1 | v2 | v3 | v4 | v5 |
|---|---|---|---|---|---|
| tok/s | 3.3K (1 GPU) | 3.6K (4 GPU) | 26K (1 GPU) | 145K (8 GPU) | 217K (8 GPU) |
| **MFU** | 0.4% | 0.2% | 14% | 17% | 13% |
| enabled by | — | DDP (weak) | **batched (T10)** | **alloc (T11)** | **bf16 (T12)** |
v1/v2 ran at **<0.5% MFU** — the single-sequence design left the GPU idle (launch-bound).
**Batched forward (T10) was the single biggest unlock** (~35× MFU jump). 6ND is an accurate
FLOPs count, but predicting *time* needs the *realized* MFU, which varied ~40× across
versions — a fixed-MFU estimate is off by up to ~100× for the early launch-bound runs.
## Engineering lessons
- **Profile before optimizing.** *Three* "known" fixes were *falsified by measurement*: (1)
"bigger batch fixes DDP scaling" (real cause: single-seq launch-bound → T10); (2) "bucket the
all-reduce" (real cause: per-op `cudaMalloc` serialization → T11 caching allocator); and (3)
"process-per-GPU would fix the residual ~5×@8" (T17 — built the torchrun-style launcher and
measured it **throughput-neutral**: the residual is the NCCL/PCIe communication wall, not
shared-context serialization). All three would have been no-ops; each got measured and either
reverted or recorded as a deliberate negative result instead of shipped on faith.
- **Honest correctness.** QK-norm was *added* to match xserv's Qwen3 (not faked); every change
kept a hard correctness gate, and **no tolerance was ever loosened to go green**. Phase 2 held
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.
## 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:
```bash
# on the GPU box
cargo run -p xserv-model --release --bin xserv-cli -- <registry>/v10-fineweb-edu-dim1280-gqa-data6765 --max-tokens 100
# then type a prompt, e.g. In science,
``` ```
The build mirrors xserv's approach: `build.rs` invokes `nvcc` (via the `cc` crate) Build/test the engine itself (CUDA compiles + runs on the GPU box; host-side `cargo check`
to compile `csrc/*.cu` targeting `sm_120` (RTX 5090) and links them into the Rust works anywhere via the `no_cuda` cfg):
crate over hand-written `extern "C"` FFI.
## Building & testing
CUDA compilation and execution happen on a GPU box (dash5, 8× RTX 5090, sm_120):
```sh ```sh
export PATH=/usr/local/cuda/bin:$HOME/.cargo/bin:$PATH export PATH=/usr/local/cuda/bin:$HOME/.cargo/bin:$PATH
cargo build cargo test --workspace # autograd grad-checks, PyTorch parity, DDP, etc.
cargo test -p xtrain-cuda -- --nocapture # runs the vecadd smoke test
``` ```
On a machine without `nvcc`/GPU, `build.rs` detects the missing toolchain, skips ## Doc index
CUDA compilation, and sets a `no_cuda` cfg — so host-side `cargo check` still
works (the GPU smoke test is compiled out). - [`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/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).

4
crates/xtrain-autodiff/Cargo.toml
View File

@@ -5,3 +5,7 @@ edition.workspace = true
[dependencies] [dependencies]
xtrain-tensor = { path = "../xtrain-tensor" } xtrain-tensor = { path = "../xtrain-tensor" }
[dev-dependencies]
# Acceptance tests need device selection (set_device) to drive the GPU.
xtrain-cuda = { path = "../xtrain-cuda" }

103
crates/xtrain-autodiff/src/checkpoint.rs Normal file
View File

@@ -0,0 +1,103 @@
//! Activation recomputation / gradient checkpointing (Phase T13, KI-3).
//!
//! A higher-order autograd primitive — the analogue of `torch.utils.checkpoint`.
//! It runs a *segment* of the model (a transformer block, here) WITHOUT recording
//! the segment's internal ops on the surrounding tape, so the segment's
//! intermediate activations are freed right after the forward instead of being
//! kept alive until backward. When the segment's output-grad arrives in backward,
//! the segment forward is **re-run** from the saved input (into a throwaway local
//! tape), the recomputed output is seeded with the upstream grad, and the gradient
//! is backpropagated through the local tape to recover the input-grad and the
//! parameter-grads — which are then pushed to the real tape's parents. The local
//! tape is dropped at the end of the closure, freeing the recomputed activations.
//!
//! ## Why it is exact (the hard gate)
//! The recompute runs the *same* `segment_fn` from the *same* input value and the
//! *same* parameter values (parameters are leaves that persist across forward and
//! backward; only their grad slot changes). The forward kernels are deterministic,
//! so the recomputed output equals the original output bit-for-bit, and the local
//! backward is the ordinary analytic backward of that segment. Therefore the input-
//! and parameter-grads are identical to those a non-checkpointed forward would
//! produce — checkpointing trades compute (one extra forward per segment) for
//! memory, never correctness.
//!
//! ## Composition
//! - **bf16 (T12):** `segment_fn` is the unchanged block forward, so the recompute
//! runs the same bf16 path; the `cast` op's grad upcast still bridges bf16→fp32.
//! - **DDP (T8):** each rank checkpoints its own forward/backward independently;
//! the param-grads recovered here feed the same per-rank `.grad()` slots that the
//! all-reduce averages — no change to the distributed path.
//! - **batched (T10):** the segment input/output carry the `[batch*seq, …]` batch
//! dim transparently; `checkpoint` is shape-agnostic.
#![cfg(not(no_cuda))]
use crate::tape::Var;
use std::rc::Rc;
/// Run `segment_fn(input, params)` with activation recomputation.
///
/// `segment_fn(x, p)` must build the segment's forward graph from a single input
/// `x` and the parameter slice `p`, returning the single segment output. It is
/// called once now (forward, result detached from the outer tape) and once per
/// backward (recompute). It MUST be deterministic and depend only on `x` and `p`
/// (this is what makes the recompute exact).
///
/// `params` are the segment's learnable leaves; their grads are accumulated into
/// the SAME leaves the optimizer reads (so DDP / AdamW are unchanged).
///
/// Returns the segment output as a `Var` on the outer tape whose backward triggers
/// the recompute. Equivalent — grad-for-grad — to calling `segment_fn(input,
/// params)` directly, but without keeping the segment's internal activations alive.
pub fn checkpoint<F>(segment_fn: F, input: &Var, params: &[Var]) -> Var
where
F: Fn(&Var, &[Var]) -> Var + 'static,
{
let segment_fn = Rc::new(segment_fn);
// --- Forward (no taping of internals) ---
// Detach the input and params into fresh leaves so `segment_fn` builds a LOCAL
// tape disconnected from the outer graph. We only keep the output's value; the
// local `Var`s (and thus the segment's intermediate activations) are dropped
// when this scope ends.
let out_value = {
let x_det = Var::leaf(input.value());
let params_det: Vec<Var> = params.iter().map(|p| Var::leaf(p.value())).collect();
let out_local = segment_fn(&x_det, &params_det);
out_local.value()
};
// Parents on the OUTER tape: the segment input, then the params (so their grads
// land in the leaves the optimizer reads).
let mut parents = Vec::with_capacity(1 + params.len());
parents.push(input.clone());
parents.extend(params.iter().cloned());
let segment_fn = segment_fn.clone();
Var::from_op(
out_value,
parents,
Box::new(move |dout, parents| {
// --- Backward (recompute) ---
// Rebuild fresh leaves from the CURRENT input/param values (params are
// unchanged since forward; input is the saved segment input), re-run the
// forward to rebuild the local tape, seed the recomputed output with the
// upstream grad, and backprop through the local tape.
let x_det = Var::leaf(parents[0].value());
let params_det: Vec<Var> = parents[1..].iter().map(|p| Var::leaf(p.value())).collect();
let out_local = segment_fn(&x_det, &params_det);
out_local.backward_seeded(dout.clone());
// Push the recovered grads to the real parents (engine SUMs on fan-out).
if let Some(dx) = x_det.grad() {
Var::push_grad(&parents[0], dx);
}
for (det, parent) in params_det.iter().zip(&parents[1..]) {
if let Some(dp) = det.grad() {
Var::push_grad(parent, dp);
}
}
// `out_local` / the local tape drop here → recomputed activations freed.
}),
)
}

2
crates/xtrain-autodiff/src/lib.rs
View File

@@ -18,6 +18,8 @@ pub use finite_diff::{GradCheckConfig, GradCheckResult, ParamFn, grad_check};
// kernels via xtrain-tensor, so they are gated behind `not(no_cuda)` (the // kernels via xtrain-tensor, so they are gated behind `not(no_cuda)` (the
// per-crate convention); the grad_check harness above stays host-only. // per-crate convention); the grad_check harness above stays host-only.
#[cfg(not(no_cuda))] #[cfg(not(no_cuda))]
pub mod checkpoint;
#[cfg(not(no_cuda))]
pub mod ops; pub mod ops;
#[cfg(not(no_cuda))] #[cfg(not(no_cuda))]
pub mod tape; pub mod tape;

683
crates/xtrain-autodiff/src/ops.rs Normal file
View File

@@ -0,0 +1,683 @@
//! Differentiable ops as autograd nodes (Phase T4).
//!
//! Each function runs the forward [`Tensor`] kernel, then builds a [`Var`] whose
//! backward closure computes the analytic gradient (see
//! `docs/03-autograd-engine.md` for the math) and pushes it to each parent via
//! [`Var::push_grad`] (which SUMs — correct under fan-out). Forward outputs that
//! the backward needs (softmax `y`, rms `inv_rms`, cross-entropy `probs`) are
//! cached by moving them into the closure.
//!
//! Attention is NOT a node here: it is composed from `matmul` + `scale` +
//! `softmax` in user code, and its backward falls out of theirs.
#![cfg(not(no_cuda))]
use crate::tape::Var;
use xtrain_tensor::{DType, Tensor};
/// dtype cast as an autograd node (Phase T12 — the AMP bridge between fp32 master
/// weights / fp32 reductions and the bf16 compute stream). Forward casts `x` to
/// `target`; **backward casts the upstream grad back to `x`'s dtype**. So a fp32
/// master-weight leaf fed through `cast(w, BF16)` into a bf16 matmul accumulates
/// an **fp32** grad — AdamW / clip / DDP all-reduce stay fp32, untouched.
pub fn cast(x: &Var, target: DType) -> Var {
let src = x.value().dtype();
if src == target {
return x.clone();
}
let out = x.value().to_dtype(target);
Var::from_op(
out,
vec![x.clone()],
Box::new(move |d, parents| {
Var::push_grad(&parents[0], d.to_dtype(src));
}),
)
}
/// `C = A @ B` (2D). Backward: `dA = dC @ Bᵀ`, `dB = Aᵀ @ dC`.
pub fn matmul(a: &Var, b: &Var) -> Var {
let out = a.value().matmul(&b.value());
Var::from_op(
out,
vec![a.clone(), b.clone()],
Box::new(|dc, parents| {
let a = parents[0].value();
let b = parents[1].value();
let (da, db) = Tensor::matmul_backward(&a, &b, dc);
Var::push_grad(&parents[0], da);
Var::push_grad(&parents[1], db);
}),
)
}
/// Elementwise `out = a + b` (same shape). Backward: grad flows unchanged to both.
pub fn add(a: &Var, b: &Var) -> Var {
let out = a.value().add(&b.value());
Var::from_op(
out,
vec![a.clone(), b.clone()],
Box::new(|d, parents| {
Var::push_grad(&parents[0], d.clone());
Var::push_grad(&parents[1], d.clone());
}),
)
}
/// Elementwise `out = a * b` (Hadamard). Backward: `da = d∘b`, `db = d∘a`.
pub fn mul(a: &Var, b: &Var) -> Var {
let out = a.value().mul(&b.value());
Var::from_op(
out,
vec![a.clone(), b.clone()],
Box::new(|d, parents| {
let a = parents[0].value();
let b = parents[1].value();
Var::push_grad(&parents[0], d.mul(&b));
Var::push_grad(&parents[1], d.mul(&a));
}),
)
}
/// Broadcast bias add: `out[r,c] = x[r,c] + bias[c]`. Backward: `dx = d`,
/// `dbias[c] = sum_r d[r,c]` (sum over the broadcast dim).
pub fn add_bias(x: &Var, bias: &Var) -> Var {
let out = x.value().add_bias(&bias.value());
Var::from_op(
out,
vec![x.clone(), bias.clone()],
Box::new(|d, parents| {
Var::push_grad(&parents[0], d.clone());
Var::push_grad(&parents[1], d.sum_rows());
}),
)
}
/// Scale by a constant: `out = x * alpha`. Backward: `dx = d * alpha`.
pub fn scale(x: &Var, alpha: f32) -> Var {
let out = x.value().scale(alpha);
Var::from_op(
out,
vec![x.clone()],
Box::new(move |d, parents| {
Var::push_grad(&parents[0], d.scale(alpha));
}),
)
}
/// RMSNorm: `y = x * rsqrt(mean(x²)+eps) * gamma`. Caches `inv_rms` for backward.
pub fn rms_norm(x: &Var, gamma: &Var, eps: f32) -> Var {
let (y, inv_rms) = x.value().rms_norm(&gamma.value(), eps);
Var::from_op(
y,
vec![x.clone(), gamma.clone()],
Box::new(move |dy, parents| {
let x = parents[0].value();
let gamma = parents[1].value();
let (dx, dgamma) = Tensor::rms_norm_backward(&x, &gamma, dy, &inv_rms);
Var::push_grad(&parents[0], dx);
Var::push_grad(&parents[1], dgamma);
}),
)
}
/// SiLU: `y = x * sigmoid(x)`. Backward uses the forward `x`.
pub fn silu(x: &Var) -> Var {
let out = x.value().silu();
Var::from_op(
out,
vec![x.clone()],
Box::new(|dy, parents| {
let x = parents[0].value();
Var::push_grad(&parents[0], Tensor::silu_backward(&x, dy));
}),
)
}
/// SwiGLU (SiLU-gated GLU): `out = silu(gate) ∘ up`. Composed from `silu` + `mul`
/// so its backward comes from theirs — no dedicated kernel needed.
pub fn swiglu(gate: &Var, up: &Var) -> Var {
mul(&silu(gate), up)
}
/// Dropout (Phase T18). With probability `p` zero each element, scale the kept
/// ones by `1/(1-p)` (inverted dropout — `E[out] == x`). The keep/drop mask is
/// drawn by a counter-based RNG from `(seed, element index)`, so it is fully
/// determined by `seed` (same `seed` ⇒ same mask: stable across the T13 recompute
/// re-run, and held fixed across the ± perturbation of a finite-diff grad-check).
/// Forward caches the per-element scale `mask`; **backward applies the same mask**
/// (`dx = d ⊙ mask`), making dropout a fixed elementwise linear map of `x`.
///
/// `p == 0` is a no-op: returns `x.clone()` (no node added) so the default graph
/// is bit-identical to the no-dropout path. eval-time identity is handled by the
/// caller simply not invoking dropout (the model's train/eval switch).
pub fn dropout(x: &Var, p: f32, seed: u64) -> Var {
if p == 0.0 {
return x.clone();
}
let (out, mask) = x.value().dropout(p, seed);
Var::from_op(
out,
vec![x.clone()],
Box::new(move |d, parents| {
Var::push_grad(&parents[0], Tensor::dropout_backward(d, &mask));
}),
)
}
/// RoPE (rotate_half) over `x:[tokens,heads,head_dim]` with per-sequence position
/// `row % period` (`period` = sequence length; `period == tokens` for a single
/// sequence). Orthogonal map, so the backward is the inverse rotation of `dy` — no
/// cached forward values needed.
pub fn rope(x: &Var, theta: f32, period: usize) -> Var {
let out = x.value().rope(theta, period);
Var::from_op(
out,
vec![x.clone()],
Box::new(move |dy, parents| {
Var::push_grad(&parents[0], Tensor::rope_backward(dy, theta, period));
}),
)
}
/// Row-wise softmax. Caches the output `y` for the Jacobian backward.
pub fn softmax(x: &Var) -> Var {
let y = x.value().softmax();
let y_cache = y.clone();
Var::from_op(
y,
vec![x.clone()],
Box::new(move |dy, parents| {
Var::push_grad(&parents[0], Tensor::softmax_backward(&y_cache, dy));
}),
)
}
/// Token embedding gather: `out[s,:] = table[ids[s], :]`. `table`:[vocab,dim]
/// (a learnable [`Var`]), `ids`:[seq] I32 (a constant index, not a `Var`).
/// Backward scatter-adds the upstream grad back into the table rows.
pub fn embedding(table: &Var, ids: &Tensor) -> Var {
let out = table.value().embedding(ids);
let vocab = table.value().shape()[0];
let ids = ids.clone();
Var::from_op(
out,
vec![table.clone()],
Box::new(move |dout, parents| {
let dtable = Tensor::embedding_backward(dout, &ids, vocab);
Var::push_grad(&parents[0], dtable);
}),
)
}
/// Reshape (contiguous, metadata-only). Backward reshapes the grad back to the
/// input shape. Used for the multi-head layout swap `[seq, h*hd] <-> [seq, h, hd]`.
pub fn reshape(x: &Var, new_shape: &[usize]) -> Var {
let in_shape: Vec<usize> = x.value().shape().to_vec();
let out = x.value().reshape(new_shape);
Var::from_op(
out,
vec![x.clone()],
Box::new(move |d, parents| {
Var::push_grad(&parents[0], d.reshape(&in_shape));
}),
)
}
/// 3D axis-(0,1) transpose `[a,b,c] -> [b,a,c]`. Self-inverse structure: the
/// backward is the same transpose applied to the grad.
pub fn transpose_3d01(x: &Var) -> Var {
let out = x.value().transpose_3d01();
Var::from_op(
out,
vec![x.clone()],
Box::new(|d, parents| {
Var::push_grad(&parents[0], d.transpose_3d01());
}),
)
}
/// 4D axis-(1,2) transpose `[a,b,c,d] -> [a,c,b,d]`. Self-inverse structure: the
/// backward is the same transpose applied to the grad. Lays out the batched
/// multi-head attention `[B,S,nh,hd] <-> [B,nh,S,hd]`.
pub fn transpose_4d12(x: &Var) -> Var {
let out = x.value().transpose_4d12();
Var::from_op(
out,
vec![x.clone()],
Box::new(|d, parents| {
Var::push_grad(&parents[0], d.transpose_4d12());
}),
)
}
/// 2D transpose `[r,c] -> [c,r]` as an autograd node (backward transposes the
/// grad back). Used for `Kᵀ` in attention scores.
pub fn transpose_2d(x: &Var) -> Var {
let out = x.value().transpose_2d();
Var::from_op(
out,
vec![x.clone()],
Box::new(|d, parents| {
Var::push_grad(&parents[0], d.transpose_2d());
}),
)
}
/// Split a `[heads, seq, head_dim]` tensor into one `[seq, head_dim]` [`Var`] per
/// head. Each head block is contiguous in this layout, so the forward copies the
/// head block into its own contiguous tensor; the backward scatters each head's
/// grad back into a zero `[heads, seq, head_dim]` grad (the engine then SUMs the
/// `heads` contributions on the shared parent — fan-out).
pub fn split_heads(x: &Var) -> Vec<Var> {
let v = x.value();
assert_eq!(v.ndim(), 3, "split_heads requires [heads,seq,head_dim]");
let (heads, seq, hd) = (v.shape()[0], v.shape()[1], v.shape()[2]);
let dev = v.device();
let flat_host = v.to_device(xtrain_tensor::Device::Cpu);
let flat = flat_host.as_slice::<f32>();
(0..heads)
.map(|h| {
let base = h * seq * hd;
let block = Tensor::from_slice(&flat[base..base + seq * hd], &[seq, hd]).to_device(dev);
Var::from_op(
block,
vec![x.clone()],
Box::new(move |d, parents| {
let mut host = vec![0.0f32; heads * seq * hd];
let dvals = d.to_device(xtrain_tensor::Device::Cpu);
let base = h * seq * hd;
host[base..base + seq * hd].copy_from_slice(dvals.as_slice::<f32>());
let g = Tensor::from_slice(&host, &[heads, seq, hd]).to_device(dev);
Var::push_grad(&parents[0], g);
}),
)
})
.collect()
}
/// Inverse of [`split_heads`]: stack per-head `[seq, head_dim]` outputs into a
/// `[heads, seq, head_dim]` tensor. Backward hands each head its own slice of the
/// grad.
pub fn merge_heads(heads_v: &[Var]) -> Var {
let heads = heads_v.len();
let v0 = heads_v[0].value();
let (seq, hd) = (v0.shape()[0], v0.shape()[1]);
let dev = v0.device();
let mut host = vec![0.0f32; heads * seq * hd];
for (h, hv) in heads_v.iter().enumerate() {
let block = hv.value().to_device(xtrain_tensor::Device::Cpu);
let base = h * seq * hd;
host[base..base + seq * hd].copy_from_slice(block.as_slice::<f32>());
}
let out = Tensor::from_slice(&host, &[heads, seq, hd]).to_device(dev);
Var::from_op(
out,
heads_v.to_vec(),
Box::new(move |d, parents| {
let dhost = d.to_device(xtrain_tensor::Device::Cpu);
let dflat = dhost.as_slice::<f32>();
for (h, parent) in parents.iter().enumerate() {
let base = h * seq * hd;
let g =
Tensor::from_slice(&dflat[base..base + seq * hd], &[seq, hd]).to_device(dev);
Var::push_grad(parent, g);
}
}),
)
}
/// Batched causal scaled-dot-product attention. `q`,`k`,`v` are each
/// `[bh, seq, head_dim]` (bh = batch·n_heads). Returns `[bh, seq, head_dim]`.
/// One fused op (2 batched GEMMs + 1 causal-softmax kernel forward; 4 batched
/// GEMMs + 1 softmax-backward kernel in backward) — replaces the per-(batch,head)
/// matmul/softmax loop, so attention is a handful of launches regardless of bh.
/// Caches the softmax `probs` for backward.
pub fn attention(q: &Var, k: &Var, v: &Var, scale: f32) -> Var {
let (out, probs) = q.value().attention(&k.value(), &v.value(), scale);
Var::from_op(
out,
vec![q.clone(), k.clone(), v.clone()],
Box::new(move |dout, parents| {
let q = parents[0].value();
let k = parents[1].value();
let v = parents[2].value();
let (dq, dk, dv) = Tensor::attention_backward(&q, &k, &v, &probs, dout, scale);
Var::push_grad(&parents[0], dq);
Var::push_grad(&parents[1], dk);
Var::push_grad(&parents[2], dv);
}),
)
}
/// Fused FLASH causal scaled-dot-product attention (Phase T14). Same interface as
/// [`attention`] (`q`,`k`,`v` each `[bh, seq, head_dim]`), but the forward is a
/// SINGLE fused kernel with an online softmax over KV tiles — the `[bh,seq,seq]`
/// score matrix is NEVER materialized, and backward caches only the per-row
/// logsumexp (O(N)) instead of the whole probs (O(N²)). Mathematically the same
/// SDPA, so it matches the composed [`attention`] within fp/bf16 tolerance.
/// Opt-in via the model's `--flash` flag; the composed path stays the default.
pub fn flash_attention(q: &Var, k: &Var, v: &Var, scale: f32) -> Var {
let (out, lse) = q.value().flash_attention(&k.value(), &v.value(), scale);
let out_cache = out.clone();
Var::from_op(
out,
vec![q.clone(), k.clone(), v.clone()],
Box::new(move |dout, parents| {
let q = parents[0].value();
let k = parents[1].value();
let v = parents[2].value();
let (dq, dk, dv) =
Tensor::flash_attention_backward(&q, &k, &v, &out_cache, &lse, dout, scale);
Var::push_grad(&parents[0], dq);
Var::push_grad(&parents[1], dk);
Var::push_grad(&parents[2], dv);
}),
)
}
/// GQA repeat_kv head broadcast (Phase T15). `kv`:[batch·num_kv, seq, head_dim]
/// (a K or V tensor) → `[batch·nh, seq, head_dim]`, each KV head broadcast to its
/// `group = nh/num_kv` query heads (qh ← kv head qh/group, contiguous groups —
/// matches xserv's repeat_kv). Feeds the unchanged composed/flash SDPA so GQA is
/// "free" for both. Backward SUMS the `group` query heads sharing each KV head back
/// onto it (the multi-group grad accumulation). `nh == num_kv` (group 1) is identity
/// → bit-identical to the MHA path. `batch` lets the op recover num_kv from kv's bh.
pub fn repeat_kv(kv: &Var, nh: usize, batch: usize) -> Var {
let bh_kv = kv.value().shape()[0];
let num_kv = bh_kv / batch;
let out = kv.value().repeat_kv(nh, batch);
Var::from_op(
out,
vec![kv.clone()],
Box::new(move |dout, parents| {
let din = Tensor::repeat_kv_backward(dout, num_kv, batch);
Var::push_grad(&parents[0], din);
}),
)
}
/// 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`,
/// 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
// probs). The grad must match the logits' dtype so it chains into a bf16
// lm_head matmul backward — cast dx back. Keeping logits bf16 (no persistent
// 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);
// 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 loss = Tensor::from_slice(&[mean_val], &[1]).to_device(x.value().device());
let target = target.clone();
Var::from_op(
loss,
vec![x.clone()],
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 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));
}),
)
}

27
crates/xtrain-autodiff/src/tape.rs
View File

@@ -68,6 +68,19 @@ impl Var {
self.0.borrow().grad.clone() self.0.borrow().grad.clone()
} }
/// Clear the accumulated gradient. Call on every parameter between training
/// steps so the next `backward` accumulates from zero (grads SUM otherwise).
pub fn zero_grad(&self) {
self.0.borrow_mut().grad = None;
}
/// Overwrite this node's value tensor in place. Used by the optimizer to
/// apply a parameter update (`p ← p − lr·grad`) while keeping the leaf's
/// identity stable across steps.
pub fn set_value(&self, value: Tensor) {
self.0.borrow_mut().value = value;
}
/// Pointer identity, used to dedup nodes during the topological sort. /// Pointer identity, used to dedup nodes during the topological sort.
fn id(&self) -> *const RefCell<VarNode> { fn id(&self) -> *const RefCell<VarNode> {
Rc::as_ptr(&self.0) Rc::as_ptr(&self.0)
@@ -95,14 +108,24 @@ impl Var {
"backward() expects a scalar loss; got shape {:?}", "backward() expects a scalar loss; got shape {:?}",
self.value().shape() self.value().shape()
); );
self.backward_seeded(ones_like(&self.value()));
}
/// Reverse-mode backward from this node seeded with an explicit upstream grad
/// `seed` (same shape as this node's value), instead of the scalar `dL/dL = 1`.
///
/// This is the entry point for **activation recomputation** (Phase T13): a
/// checkpointed segment re-runs its forward into a fresh local tape, then
/// backprops the upstream output-grad through it via this method (the segment
/// output is generally NOT a scalar). For a scalar root, [`backward`] is the
/// thin wrapper that seeds ones.
pub fn backward_seeded(&self, seed: Tensor) {
// 1. Topological order (post-order DFS), parents before children. // 1. Topological order (post-order DFS), parents before children.
let mut topo: Vec<Var> = Vec::new(); let mut topo: Vec<Var> = Vec::new();
let mut visited: Vec<*const RefCell<VarNode>> = Vec::new(); let mut visited: Vec<*const RefCell<VarNode>> = Vec::new();
build_topo(self, &mut topo, &mut visited); build_topo(self, &mut topo, &mut visited);
// 2. Seed the loss gradient with ones. // 2. Seed this node's gradient with the supplied upstream grad.
let seed = ones_like(&self.value());
self.accumulate(seed); self.accumulate(seed);
// 3. Walk in reverse: each node hands its grad to its parents' closures. // 3. Walk in reverse: each node hands its grad to its parents' closures.

1270
crates/xtrain-autodiff/tests/autograd.rs Normal file
View File

File diff suppressed because it is too large Load Diff

220
crates/xtrain-autodiff/tests/structural.rs Normal file
View File

@@ -0,0 +1,220 @@
// GPU grad-checks for the Phase T5 structural ops added on top of the T4 set:
// embedding (gather fwd / scatter-add bwd), reshape, transpose_3d01,
// transpose_2d, and split/merge_heads. Same harness as autograd.rs:
// L = sum(W ∘ out), W fixed random ⇒ upstream dOut = W; run backward(), then
// grad-check each leaf's .grad() against central finite differences.
//
// Gated behind `not(no_cuda)`: compiles out on a GPU-less host, runs on dash5.
#![cfg(not(no_cuda))]
use xtrain_autodiff::ops;
use xtrain_autodiff::tape::Var;
use xtrain_autodiff::{GradCheckConfig, grad_check};
use xtrain_cuda::device;
use xtrain_tensor::{Device, Tensor};
fn fill(n: usize, seed: u64) -> 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
})
.collect()
}
fn require_gpu() {
assert!(
device::device_count().expect("device count") > 0,
"no CUDA device"
);
device::set_device(0).unwrap();
}
fn cuda(data: &[f32], shape: &[usize]) -> Tensor {
Tensor::from_slice(data, shape).to_device(Device::Cuda(0))
}
fn weighted_sum(out: &Tensor, w: &[f32]) -> f32 {
out.to_device(Device::Cpu)
.as_slice::<f32>()
.iter()
.zip(w)
.map(|(o, w)| o * w)
.sum()
}
// Structural ops are exactly linear in their input → a large eps just sharpens
// f32 resolution (same as add/mul/transpose in autograd.rs).
fn cfg_linear() -> GradCheckConfig {
GradCheckConfig {
eps: 1e-2,
rel_tol: 2e-2,
atol: 1e-3,
}
}
fn report(name: &str, res: &xtrain_autodiff::GradCheckResult) {
println!(
"{name}: max_rel_err = {:.3e} (worst num={:.5} ana={:.5} @ {})",
res.max_rel_err, res.worst_numeric, res.worst_analytic, res.worst_index
);
assert!(res.passed, "{name} grad-check failed: {res:?}");
}
// L = sum(W ∘ out): a constant-W leaf mul + sum-to-scalar reduction.
fn scalar_loss(out: &Var, w: &[f32]) -> Var {
let wt = Var::leaf(cuda(w, out.value().shape()));
sum_all(&ops::mul(out, &wt))
}
fn sum_all(x: &Var) -> Var {
let xv = x.value();
let total: f32 = xv.to_device(Device::Cpu).as_slice::<f32>().iter().sum();
let scalar = Tensor::from_slice(&[total], &[1]).to_device(xv.device());
let shape: Vec<usize> = xv.shape().to_vec();
Var::from_op(
scalar,
vec![x.clone()],
Box::new(move |d, parents| {
let dval = d.to_device(Device::Cpu).as_slice::<f32>()[0];
let ones = vec![dval; shape.iter().product()];
let g = Tensor::from_slice(&ones, &shape).to_device(Device::Cuda(0));
Var::push_grad(&parents[0], g);
}),
)
}
// ---- embedding (gather fwd / scatter-add bwd) ----
// Includes a repeated id so the atomic scatter-add accumulation is exercised.
#[test]
fn embedding_bwd() {
require_gpu();
let (vocab, dim) = (5, 7);
let ids_host: Vec<i32> = vec![0, 3, 1, 3, 2, 0]; // 0 and 3 repeat
let seq = ids_host.len();
let table_h = fill(vocab * dim, 201);
let w = fill(seq * dim, 202);
let ids = Tensor::from_slice(&ids_host, &[seq]).to_device(Device::Cuda(0));
let table = Var::leaf(cuda(&table_h, &[vocab, dim]));
let out = ops::embedding(&table, &ids);
scalar_loss(&out, &w).backward();
let dtable = table.grad().unwrap().to_device(Device::Cpu);
let idf = ids_host.clone();
let wf = w.clone();
let lt = move |v: &[f32], s: &[usize]| {
let ids = Tensor::from_slice(&idf, &[seq]).to_device(Device::Cuda(0));
weighted_sum(&cuda(v, s).embedding(&ids), &wf)
};
report(
"embedding dTable",
&grad_check(
&table_h,
&[vocab, dim],
&lt,
dtable.as_slice::<f32>(),
cfg_linear(),
),
);
}
// ---- reshape ----
#[test]
fn reshape_bwd() {
require_gpu();
let (rows, cols) = (6, 8);
let x_h = fill(rows * cols, 211);
let w = fill(rows * cols, 212);
let x = Var::leaf(cuda(&x_h, &[rows, cols]));
let out = ops::reshape(&x, &[rows * 2, cols / 2]);
scalar_loss(&out, &w).backward();
let dx = x.grad().unwrap().to_device(Device::Cpu);
let wf = w.clone();
let lx =
move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).reshape(&[rows * 2, cols / 2]), &wf);
report(
"reshape dX",
&grad_check(&x_h, &[rows, cols], &lx, dx.as_slice::<f32>(), cfg_linear()),
);
}
// ---- transpose_3d01 ([a,b,c] -> [b,a,c]) ----
#[test]
fn transpose_3d01_bwd() {
require_gpu();
let (a, b, c) = (3, 4, 5);
let x_h = fill(a * b * c, 221);
let w = fill(a * b * c, 222);
let x = Var::leaf(cuda(&x_h, &[a, b, c]));
let out = ops::transpose_3d01(&x);
scalar_loss(&out, &w).backward();
let dx = x.grad().unwrap().to_device(Device::Cpu);
let wf = w.clone();
let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).transpose_3d01(), &wf);
report(
"transpose_3d01 dX",
&grad_check(&x_h, &[a, b, c], &lx, dx.as_slice::<f32>(), cfg_linear()),
);
}
// ---- transpose_2d ----
#[test]
fn transpose_2d_bwd() {
require_gpu();
let (r, c) = (5, 7);
let x_h = fill(r * c, 231);
let w = fill(r * c, 232);
let x = Var::leaf(cuda(&x_h, &[r, c]));
let out = ops::transpose_2d(&x);
scalar_loss(&out, &w).backward();
let dx = x.grad().unwrap().to_device(Device::Cpu);
let wf = w.clone();
let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).transpose_2d(), &wf);
report(
"transpose_2d dX",
&grad_check(&x_h, &[r, c], &lx, dx.as_slice::<f32>(), cfg_linear()),
);
}
// ---- split_heads + merge_heads round-trip (identity reshuffle of [nh,seq,hd]) ----
// out = merge_heads(split_heads(x)) must equal x, and its grad must be dOut=W
// reshuffled identically — i.e. dx grad-checks against the identity composition.
#[test]
fn split_merge_heads_bwd() {
require_gpu();
let (nh, seq, hd) = (3, 4, 5);
let x_h = fill(nh * seq * hd, 241);
let w = fill(nh * seq * hd, 242);
let x = Var::leaf(cuda(&x_h, &[nh, seq, hd]));
let heads = ops::split_heads(&x);
let out = ops::merge_heads(&heads); // back to [nh,seq,hd]
scalar_loss(&out, &w).backward();
let dx = x.grad().unwrap().to_device(Device::Cpu);
// forward is identity, so grad-check the identity map.
let wf = w.clone();
let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s), &wf);
report(
"split/merge_heads dX",
&grad_check(
&x_h,
&[nh, seq, hd],
&lx,
dx.as_slice::<f32>(),
cfg_linear(),
),
);
}

7
crates/xtrain-cuda/build.rs
View File

@@ -33,6 +33,13 @@ fn main() {
.file("../../csrc/ops/elementwise.cu") .file("../../csrc/ops/elementwise.cu")
.file("../../csrc/ops/gemm.cu") .file("../../csrc/ops/gemm.cu")
.file("../../csrc/ops/nn.cu") .file("../../csrc/ops/nn.cu")
.file("../../csrc/ops/model.cu")
.file("../../csrc/ops/optim.cu")
.file("../../csrc/ops/attention.cu")
.file("../../csrc/ops/flash_attention.cu")
.file("../../csrc/ops/repeat_kv.cu")
.file("../../csrc/ops/cast.cu")
.file("../../csrc/ops/dropout.cu")
.compile("xtrain_cuda_kernels"); .compile("xtrain_cuda_kernels");
} }

290
crates/xtrain-cuda/src/cublas.rs Normal file
View File

@@ -0,0 +1,290 @@
//! cuBLAS GEMM backend (Phase T7).
//!
//! The hand-written tiled kernel (csrc/ops/gemm.cu) is kept as the T3 learning
//! artifact + correctness oracle's counterpart, but the forward + both backward
//! matmuls now route through cuBLAS `Sgemm` — fp32, so the result is numerically
//! the same GEMM (only the rounding order changes), which is why the T3 tolerance
//! against cuBLAS holds unchanged.
//!
//! **Layout.** cuBLAS is column-major; our tensors are row-major. A row-major
//! `[r,c]` matrix handed to cuBLAS with leading dim `c` is read as its transpose
//! (col-major `[c,r]`). To get a row-major result `C[m,n] = opA(A)·opB(B)` we
//! compute the col-major transpose `Cᵀ[n,m] = opB(B)ᵀ·opA(A)ᵀ`; the bytes of
//! col-major `Cᵀ` are exactly row-major `C`. See [`sgemm`] for the index algebra.
//!
//! **Handle.** cuBLAS handle creation is expensive (T3's oracle made one per
//! call). We cache one handle per thread for the lifetime of the process.
#![cfg(not(no_cuda))]
use crate::ffi::{self, CublasHandle};
use std::cell::RefCell;
use std::ffi::c_void;
thread_local! {
static HANDLE: RefCell<Option<CublasHandle>> = const { RefCell::new(None) };
}
/// Run `f` with the thread's cached cuBLAS handle, creating it on first use.
fn with_handle<R>(f: impl FnOnce(CublasHandle) -> R) -> R {
HANDLE.with(|h| {
let mut slot = h.borrow_mut();
if slot.is_none() {
let mut handle: CublasHandle = std::ptr::null_mut();
let status = unsafe { ffi::cublasCreate_v2(&mut handle) };
assert_eq!(status, 0, "cublasCreate failed: {status}");
*slot = Some(handle);
}
f(slot.unwrap())
})
}
/// Row-major single-precision GEMM: `C[m,n] = opA(A) · opB(B)` with
/// `C = alpha·(…) + beta·C`. `A`/`B`/`C` are device pointers to row-major fp32
/// matrices; `trans_a`/`trans_b` request the transpose of the *logical* operand.
///
/// `m,n,k` are the dims of the math (`opA(A)` is `[m,k]`, `opB(B)` is `[k,n]`).
/// The stored, untransposed shapes are: `A` is `[m,k]` (or `[k,m]` if `trans_a`),
/// `B` is `[k,n]` (or `[n,k]` if `trans_b`). Their row-major leading dims are the
/// stored column counts, derived below.
///
/// We ask cuBLAS for col-major `Cᵀ[n,m] = opB(B)ᵀ · opA(A)ᵀ`. Since a row-major
/// `[r,c]` buffer is col-major `[c,r]`, a row-major operand already *is* its own
/// transpose to cuBLAS — so `opB(B)ᵀ` over the row-major bytes of `B` is obtained
/// by passing `B` with the OPPOSITE op flag of what `opB` would suggest. Working
/// it through: first cuBLAS arg = `B` with op `trans_b ? N : T`, second = `A` with
/// op `trans_a ? N : T`, sizes (m=n, n=m, k=k).
#[allow(clippy::too_many_arguments)]
pub fn sgemm(
trans_a: bool,
trans_b: bool,
m: usize,
n: usize,
k: usize,
alpha: f32,
a: *const f32,
b: *const f32,
beta: f32,
c: *mut f32,
) {
// Leading dims = stored (row-major) column count of each untransposed matrix.
let lda = if trans_a { m } else { k }; // A stored [m,k] or [k,m]
let ldb = if trans_b { k } else { n }; // B stored [k,n] or [n,k]
let ldc = n; // Cᵀ is [n,m] col-major with ld n (== row-major C[m,n])
let op_b = if trans_b {
ffi::CUBLAS_OP_T
} else {
ffi::CUBLAS_OP_N
};
let op_a = if trans_a {
ffi::CUBLAS_OP_T
} else {
ffi::CUBLAS_OP_N
};
with_handle(|handle| {
let status = unsafe {
ffi::cublasSgemm_v2(
handle, op_b, op_a, n as i32, // rows of Cᵀ
m as i32, // cols of Cᵀ
k as i32, &alpha, b, ldb as i32, a, lda as i32, &beta, c, ldc as i32,
)
};
assert_eq!(status, 0, "cublasSgemm failed: {status}");
});
}
/// Strided-batched row-major SGEMM: for each `i` in `0..batch`,
/// `C_i[m,n] = alpha·opA(A_i)·opB(B_i) + beta·C_i`, where `A_i`/`B_i`/`C_i` are
/// consecutive matrices laid `stride_*` elements apart in one contiguous buffer.
/// Same row-major⟺col-major trick as [`sgemm`] (compute col-major `Cᵀ`), applied
/// per batch element. Used for the batched attention `QKᵀ` / `PV` GEMMs (and their
/// backwards), so the whole attention runs as 2 batched-GEMM launches, not a
/// per-(batch,head) Python loop. `A`/`B`/`C` are device pointers to the first
/// matrix; strides are in ELEMENTS.
#[allow(clippy::too_many_arguments)]
pub fn sgemm_strided_batched(
trans_a: bool,
trans_b: bool,
m: usize,
n: usize,
k: usize,
alpha: f32,
a: *const f32,
stride_a: usize,
b: *const f32,
stride_b: usize,
beta: f32,
c: *mut f32,
stride_c: usize,
batch: usize,
) {
let lda = if trans_a { m } else { k };
let ldb = if trans_b { k } else { n };
let ldc = n;
let op_a = if trans_a {
ffi::CUBLAS_OP_T
} else {
ffi::CUBLAS_OP_N
};
let op_b = if trans_b {
ffi::CUBLAS_OP_T
} else {
ffi::CUBLAS_OP_N
};
with_handle(|handle| {
let status = unsafe {
ffi::cublasSgemmStridedBatched(
handle,
op_b,
op_a,
n as i32,
m as i32,
k as i32,
&alpha,
b,
ldb as i32,
stride_b as i64,
a,
lda as i32,
stride_a as i64,
&beta,
c,
ldc as i32,
stride_c as i64,
batch as i32,
)
};
assert_eq!(status, 0, "cublasSgemmStridedBatched failed: {status}");
});
}
/// bf16 row-major GEMM `C[m,n] = opA(A)·opB(B)` via `cublasGemmEx`: bf16 in/out,
/// **fp32 accumulation** (`CUBLAS_COMPUTE_32F`) — the standard AMP matmul (Phase
/// T12). `a`/`b`/`c` are device pointers to row-major **bf16** matrices; the
/// row-major⟺col-major transpose algebra is identical to [`sgemm`] (we compute
/// the col-major `Cᵀ`). `alpha`/`beta` are fp32 host scalars (compute is fp32).
#[allow(clippy::too_many_arguments)]
pub fn gemm_ex(
trans_a: bool,
trans_b: bool,
m: usize,
n: usize,
k: usize,
alpha: f32,
a: *const c_void,
b: *const c_void,
beta: f32,
c: *mut c_void,
) {
let lda = if trans_a { m } else { k };
let ldb = if trans_b { k } else { n };
let ldc = n;
let op_a = if trans_a {
ffi::CUBLAS_OP_T
} else {
ffi::CUBLAS_OP_N
};
let op_b = if trans_b {
ffi::CUBLAS_OP_T
} else {
ffi::CUBLAS_OP_N
};
let bf16 = ffi::CUDA_R_16BF;
with_handle(|handle| {
let status = unsafe {
ffi::cublasGemmEx(
handle,
op_b,
op_a,
n as i32,
m as i32,
k as i32,
&alpha as *const f32 as *const c_void,
b,
bf16,
ldb as i32,
a,
bf16,
lda as i32,
&beta as *const f32 as *const c_void,
c,
bf16,
ldc as i32,
ffi::CUBLAS_COMPUTE_32F,
ffi::CUBLAS_GEMM_DEFAULT,
)
};
assert_eq!(status, 0, "cublasGemmEx failed: {status}");
});
}
/// Strided-batched bf16 GEMM (Phase T12) — the [`gemm_ex`] analogue of
/// [`sgemm_strided_batched`] for the batched attention GEMMs. bf16 in/out, fp32
/// accumulation; strides are in ELEMENTS.
#[allow(clippy::too_many_arguments)]
pub fn gemm_ex_strided_batched(
trans_a: bool,
trans_b: bool,
m: usize,
n: usize,
k: usize,
alpha: f32,
a: *const c_void,
stride_a: usize,
b: *const c_void,
stride_b: usize,
beta: f32,
c: *mut c_void,
stride_c: usize,
batch: usize,
) {
let lda = if trans_a { m } else { k };
let ldb = if trans_b { k } else { n };
let ldc = n;
let op_a = if trans_a {
ffi::CUBLAS_OP_T
} else {
ffi::CUBLAS_OP_N
};
let op_b = if trans_b {
ffi::CUBLAS_OP_T
} else {
ffi::CUBLAS_OP_N
};
let bf16 = ffi::CUDA_R_16BF;
with_handle(|handle| {
let status = unsafe {
ffi::cublasGemmStridedBatchedEx(
handle,
op_b,
op_a,
n as i32,
m as i32,
k as i32,
&alpha as *const f32 as *const c_void,
b,
bf16,
ldb as i32,
stride_b as i64,
a,
bf16,
lda as i32,
stride_a as i64,
&beta as *const f32 as *const c_void,
c,
bf16,
ldc as i32,
stride_c as i64,
batch as i32,
ffi::CUBLAS_COMPUTE_32F,
ffi::CUBLAS_GEMM_DEFAULT,
)
};
assert_eq!(status, 0, "cublasGemmStridedBatchedEx failed: {status}");
});
}

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

@@ -13,12 +13,14 @@ unsafe extern "C" {
// --- Device --- // --- Device ---
pub fn cudaGetDeviceCount(count: *mut i32) -> i32; pub fn cudaGetDeviceCount(count: *mut i32) -> i32;
pub fn cudaSetDevice(device: i32) -> i32; pub fn cudaSetDevice(device: i32) -> i32;
pub fn cudaGetDevice(device: *mut i32) -> i32;
pub fn cudaDeviceSynchronize() -> i32; pub fn cudaDeviceSynchronize() -> i32;
// --- Memory --- // --- Memory ---
pub fn cudaMalloc(devptr: *mut *mut u8, size: usize) -> i32; pub fn cudaMalloc(devptr: *mut *mut u8, size: usize) -> i32;
pub fn cudaFree(devptr: *mut u8) -> i32; pub fn cudaFree(devptr: *mut u8) -> i32;
pub fn cudaMemcpy(dst: *mut u8, src: *const u8, count: usize, kind: i32) -> i32; pub fn cudaMemcpy(dst: *mut u8, src: *const u8, count: usize, kind: i32) -> i32;
pub fn cudaMemset(devptr: *mut u8, value: i32, count: usize) -> i32;
// --- Error --- // --- Error ---
pub fn cudaGetErrorString(error: i32) -> *const c_char; pub fn cudaGetErrorString(error: i32) -> *const c_char;
@@ -124,7 +126,9 @@ unsafe extern "C" {
pub fn launch_silu_f32(x: *const f32, y: *mut f32, n: i32, s: CudaStream); pub fn launch_silu_f32(x: *const f32, y: *mut f32, n: i32, s: CudaStream);
pub fn launch_silu_dx_f32(x: *const f32, dy: *const f32, dx: *mut f32, n: i32, s: CudaStream); pub fn launch_silu_dx_f32(x: *const f32, dy: *const f32, dx: *mut f32, n: i32, s: CudaStream);
// RoPE (rotate_half), x:[tokens,heads,head_dim], position = token index. // RoPE (rotate_half), x:[tokens,heads,head_dim], position = (token index %
// period). `period` = sequence length, so a flattened batch of sequences gets
// per-sequence positions; period == tokens reproduces the single-sequence case.
pub fn launch_rope_f32( pub fn launch_rope_f32(
x: *const f32, x: *const f32,
y: *mut f32, y: *mut f32,
@@ -132,8 +136,54 @@ unsafe extern "C" {
heads: i32, heads: i32,
head_dim: i32, head_dim: i32,
theta: f32, theta: f32,
period: i32,
s: CudaStream, 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( pub fn launch_rope_dx_f32(
dy: *const f32, dy: *const f32,
dx: *mut f32, dx: *mut f32,
@@ -141,6 +191,7 @@ unsafe extern "C" {
heads: i32, heads: i32,
head_dim: i32, head_dim: i32,
theta: f32, theta: f32,
period: i32,
s: CudaStream, s: CudaStream,
); );
@@ -177,8 +228,181 @@ unsafe extern "C" {
); );
} }
// cuBLAS — used ONLY as a correctness reference for the hand-written GEMM in // Structural ops for the tiny transformer (csrc/ops/model.cu): token embedding
// tests. Declared (and linked, see build.rs) only when CUDA is compiled in. // (gather fwd / scatter-add bwd) and a 3D axis-(0,1) transpose for the multi-head
// attention layout. F32 values, I32 ids, row-major contiguous.
#[cfg(not(no_cuda))]
unsafe extern "C" {
// Embedding: out[s,:] = table[ids[s], :]. table:[vocab,dim], ids:[seq] (I32).
pub fn launch_embedding_fwd_f32(
table: *const f32,
ids: *const i32,
out: *mut f32,
seq: i32,
dim: i32,
s: CudaStream,
);
// Scatter-add: dtable[ids[s],:] += dout[s,:] (dtable pre-zeroed; atomic).
pub fn launch_embedding_bwd_f32(
dout: *const f32,
ids: *const i32,
dtable: *mut f32,
seq: i32,
dim: i32,
s: CudaStream,
);
// 3D axis-(0,1) transpose: in:[a,b,c] -> out:[b,a,c]. out[j,i,k]=in[i,j,k].
pub fn launch_transpose_3d01_f32(
input: *const f32,
out: *mut f32,
a: i32,
b: i32,
c: i32,
s: CudaStream,
);
// 4D axis-(1,2) transpose: in:[a,b,c,d] -> out:[a,c,b,d]. out[i,k,j,l]=in[i,j,k,l].
pub fn launch_transpose_4d12_f32(
input: *const f32,
out: *mut f32,
a: i32,
b: i32,
c: i32,
d: i32,
s: CudaStream,
);
}
// Batched attention helper (csrc/ops/attention.cu): causal row-wise softmax over
// score rows [rows, seq] with query position = (row % seq); scales logits by
// `scale` (= 1/sqrt(head_dim)) and masks future columns to probability 0.
#[cfg(not(no_cuda))]
unsafe extern "C" {
pub fn launch_softmax_causal_f32(
x: *const f32,
y: *mut f32,
rows: i32,
seq: i32,
scale: f32,
s: CudaStream,
);
}
// Fused flash-attention (csrc/ops/flash_attention.cu, Phase T14). A SINGLE kernel
// each for forward/backward that streams over KV tiles with an online softmax and
// NEVER materializes the [bh,S,S] score matrix. Q/K/V/out are [bh,S,hd] row-major
// F32; the forward saves only the per-row logsumexp `l` ([bh*S], O(N)) for backward.
#[cfg(not(no_cuda))]
unsafe extern "C" {
// Forward: o[bh,S,hd] = softmax(causal(Q·Kᵀ·scale))·V, online over KV tiles.
// Also writes l[bh*S] = per-row logsumexp (saved for backward, not the scores).
#[allow(clippy::too_many_arguments)]
pub fn launch_flash_attention_fwd_f32(
q: *const f32,
k: *const f32,
v: *const f32,
o: *mut f32,
l: *mut f32,
bh: i32,
seq: i32,
hd: i32,
scale: f32,
s: CudaStream,
);
// Per-row D[i]=Σ_d dO[i,d]·O[i,d] over `rows`=bh*S rows of width `hd`. Must run
// before the backward kernel (which takes the precomputed D, not O).
pub fn launch_flash_attention_rowdot_f32(
d_o: *const f32,
o: *const f32,
d_d: *mut f32,
rows: i32,
hd: i32,
s: CudaStream,
);
// Backward: recomputes scores from Q/K/V + saved logsumexp `l` (NO cached probs)
// and the precomputed `d_d` (= D), produces dq/dk/dv. dq/dk/dv must be PRE-ZEROED
// (dk/dv are accumulated across query rows via atomicAdd).
#[allow(clippy::too_many_arguments)]
pub fn launch_flash_attention_bwd_f32(
q: *const f32,
k: *const f32,
v: *const f32,
d_o: *const f32,
l: *const f32,
d_d: *mut f32,
dq: *mut f32,
dk: *mut f32,
dv: *mut f32,
bh: i32,
seq: i32,
hd: i32,
scale: f32,
s: CudaStream,
);
}
// GQA repeat_kv head broadcast (csrc/ops/repeat_kv.cu, Phase T15). Expands a K/V
// tensor from [batch·num_kv, S, hd] to the full [batch·nh, S, hd] so the SDPA
// (composed or flash, both untouched) sees a full set of heads. Forward gathers
// (out head qh ← kv head qh/group, group = nh/num_kv); backward sums the `group`
// query heads sharing each kv head (deterministic, no atomics). All F32.
#[cfg(not(no_cuda))]
unsafe extern "C" {
// Forward: out[b·nh+qh] = in[b·num_kv + qh/group], per [S,hd] head block.
pub fn launch_repeat_kv_fwd_f32(
input: *const f32,
out: *mut f32,
batch: i32,
nh: i32,
num_kv: i32,
seq: i32,
hd: i32,
s: CudaStream,
);
// Backward: din[b·num_kv+kvh] = Σ_{r<group} dout[b·nh + kvh·group + r].
pub fn launch_repeat_kv_bwd_f32(
dout: *const f32,
din: *mut f32,
batch: i32,
nh: i32,
num_kv: i32,
seq: i32,
hd: i32,
s: CudaStream,
);
}
// GPU-side optimizer kernels (csrc/ops/optim.cu): AdamW step (m/v on device) and
// the global grad-norm reduction + in-place rescale (Phase T7).
#[cfg(not(no_cuda))]
unsafe extern "C" {
// One in-place AdamW step over a parameter tensor of `n` elements. `bc1`/`bc2`
// are the bias-correction denominators 1-beta^t.
#[allow(clippy::too_many_arguments)]
pub fn launch_adamw_step_f32(
p: *mut f32,
g: *const f32,
m: *mut f32,
v: *mut f32,
lr: f32,
b1: f32,
b2: f32,
eps: f32,
wd: f32,
bc1: f32,
bc2: f32,
n: i32,
s: CudaStream,
);
// acc += sum_i g[i]^2 (acc is one f32 on device, pre-zeroed). atomicAdd.
pub fn launch_sumsq_accum_f32(g: *const f32, acc: *mut f32, n: i32, s: CudaStream);
// In-place scalar scale: x[i] *= factor.
pub fn launch_scale_inplace_f32(x: *mut f32, factor: f32, n: i32, s: CudaStream);
}
// cuBLAS — the production GEMM backend (Phase T7) and the correctness oracle the
// T3 GEMM tests still compare against. Declared (and linked, see build.rs) only
// when CUDA is compiled in.
#[cfg(not(no_cuda))] #[cfg(not(no_cuda))]
pub type CublasHandle = *mut c_void; pub type CublasHandle = *mut c_void;
@@ -202,7 +426,198 @@ unsafe extern "C" {
c: *mut f32, c: *mut f32,
ldc: i32, ldc: i32,
) -> i32; ) -> i32;
#[allow(clippy::too_many_arguments)]
pub fn cublasSgemmStridedBatched(
handle: CublasHandle,
transa: i32,
transb: i32,
m: i32,
n: i32,
k: i32,
alpha: *const f32,
a: *const f32,
lda: i32,
stride_a: i64,
b: *const f32,
ldb: i32,
stride_b: i64,
beta: *const f32,
c: *mut f32,
ldc: i32,
stride_c: i64,
batch_count: i32,
) -> i32;
} }
#[cfg(not(no_cuda))] #[cfg(not(no_cuda))]
pub const CUBLAS_OP_N: i32 = 0; pub const CUBLAS_OP_N: i32 = 0;
#[cfg(not(no_cuda))]
pub const CUBLAS_OP_T: i32 = 1;
// --- bf16 mixed precision (Phase T12) ---
//
// cudaDataType / cublasComputeType enum values (same as xserv's gemm.rs). The
// bf16 GEMM uses bf16 in/out with fp32 accumulation (CUBLAS_COMPUTE_32F).
#[cfg(not(no_cuda))]
pub const CUDA_R_32F: i32 = 0;
#[cfg(not(no_cuda))]
pub const CUDA_R_16BF: i32 = 14;
#[cfg(not(no_cuda))]
pub const CUBLAS_COMPUTE_32F: i32 = 68;
/// CUBLAS_GEMM_DEFAULT — let cuBLAS pick the algorithm.
#[cfg(not(no_cuda))]
pub const CUBLAS_GEMM_DEFAULT: i32 = -1;
#[cfg(not(no_cuda))]
unsafe extern "C" {
// General GEMM with explicit in/out + compute types (bf16 path). `alpha`/
// `beta` are fp32 host scalars (compute type is fp32). Pointers are void* so
// the same FFI serves bf16 / fp32.
#[allow(clippy::too_many_arguments)]
pub fn cublasGemmEx(
handle: CublasHandle,
transa: i32,
transb: i32,
m: i32,
n: i32,
k: i32,
alpha: *const std::ffi::c_void,
a: *const std::ffi::c_void,
a_type: i32,
lda: i32,
b: *const std::ffi::c_void,
b_type: i32,
ldb: i32,
beta: *const std::ffi::c_void,
c: *mut std::ffi::c_void,
c_type: i32,
ldc: i32,
compute_type: i32,
algo: i32,
) -> i32;
#[allow(clippy::too_many_arguments)]
pub fn cublasGemmStridedBatchedEx(
handle: CublasHandle,
transa: i32,
transb: i32,
m: i32,
n: i32,
k: i32,
alpha: *const std::ffi::c_void,
a: *const std::ffi::c_void,
a_type: i32,
lda: i32,
stride_a: i64,
b: *const std::ffi::c_void,
b_type: i32,
ldb: i32,
stride_b: i64,
beta: *const std::ffi::c_void,
c: *mut std::ffi::c_void,
c_type: i32,
ldc: i32,
stride_c: i64,
batch_count: i32,
compute_type: i32,
algo: i32,
) -> i32;
}
// bf16 cast + elementwise kernels (csrc/ops/cast.cu). Pointers are void* (bf16
// buffers); f32 sides are typed. The activation stream flows bf16; the math
// accumulates in fp32 inside each kernel.
#[cfg(not(no_cuda))]
unsafe extern "C" {
pub fn launch_cast_f32_to_bf16(input: *const f32, out: *mut c_void, n: i32, s: CudaStream);
pub fn launch_cast_bf16_to_f32(input: *const c_void, out: *mut f32, n: i32, s: CudaStream);
pub fn launch_add_bf16(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
s: CudaStream,
);
pub fn launch_mul_bf16(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
s: CudaStream,
);
pub fn launch_scale_bf16(
input: *const c_void,
out: *mut c_void,
alpha: f32,
n: i32,
s: CudaStream,
);
pub fn launch_silu_bf16(x: *const c_void, y: *mut c_void, n: i32, s: CudaStream);
pub fn launch_silu_dx_bf16(
x: *const c_void,
dy: *const c_void,
dx: *mut c_void,
n: i32,
s: CudaStream,
);
pub fn launch_add_bias_bf16(
x: *const c_void,
bias: *const c_void,
out: *mut c_void,
rows: i32,
cols: i32,
s: CudaStream,
);
pub fn launch_sum_rows_bf16(
dout: *const c_void,
dbias: *mut c_void,
rows: i32,
cols: i32,
s: CudaStream,
);
}
// Dropout (Phase T18, csrc/ops/dropout.cu). A counter-based (stateless) RNG: the
// keep/drop decision for element `i` is `hash(seed, i)` — no global state, so a
// re-run with the same `seed` reproduces the same mask (compatible with T13
// activation recomputation). Forward writes `out = x ⊙ mask` and the fp32 `mask`
// buffer (mask[i] = (1/(1-p)) if kept else 0, the inverted-dropout scale);
// backward applies the SAME mask: dx = d ⊙ mask. fp32 + bf16 activation variants
// (mask is fp32 in both; the uniform is computed in fp32, dtype-independent).
#[cfg(not(no_cuda))]
unsafe extern "C" {
pub fn launch_dropout_fwd_f32(
x: *const f32,
out: *mut f32,
mask: *mut f32,
p: f32,
scale: f32,
seed: u64,
n: i32,
s: CudaStream,
);
pub fn launch_dropout_bwd_f32(
d: *const f32,
mask: *const f32,
dx: *mut f32,
n: i32,
s: CudaStream,
);
pub fn launch_dropout_fwd_bf16(
x: *const c_void,
out: *mut c_void,
mask: *mut f32,
p: f32,
scale: f32,
seed: u64,
n: i32,
s: CudaStream,
);
pub fn launch_dropout_bwd_bf16(
d: *const c_void,
mask: *const f32,
dx: *mut c_void,
n: i32,
s: CudaStream,
);
}

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

@@ -1,7 +1,10 @@
#[cfg(not(no_cuda))]
pub mod cublas;
pub mod device; pub mod device;
pub mod error; pub mod error;
pub mod ffi; pub mod ffi;
pub mod memory; pub mod memory;
mod pool;
pub use error::{CudaError, Result}; pub use error::{CudaError, Result};
pub use memory::GpuBuffer; pub use memory::GpuBuffer;

40
crates/xtrain-cuda/src/memory.rs
View File

@@ -1,18 +1,37 @@
use crate::error::{self, Result}; use crate::error::{self, Result};
use crate::ffi; use crate::ffi;
use crate::pool;
/// RAII wrapper around a GPU memory allocation. Dropping frees the memory. /// RAII wrapper around a GPU memory allocation. Dropping returns the buffer to
/// the per-device caching pool (see [`crate::pool`]) for reuse instead of
/// calling `cudaFree`.
///
/// `len` is the logical (requested) length used for all copy/memset bounds and
/// exposed via [`GpuBuffer::len`]; `cap` is the physical size class the pool
/// rounded up to (>= `len`), used only to bucket the buffer for reuse. The
/// extra `cap - len` bytes are never exposed to callers, so pooling is
/// numerically transparent. `device` records which device pool to return to.
pub struct GpuBuffer { pub struct GpuBuffer {
ptr: *mut u8, ptr: *mut u8,
len: usize, len: usize,
cap: usize,
device: i32,
} }
impl GpuBuffer { impl GpuBuffer {
/// Allocate at least `len` bytes on the calling thread's current device,
/// reusing a pooled buffer when one of the matching size class is free.
/// The contents are **uninitialized** (a reused buffer holds stale bytes);
/// callers that need zeros must memset (see [`crate::Storage::zeros`]).
pub fn alloc(len: usize) -> Result<Self> { pub fn alloc(len: usize) -> Result<Self> {
assert!(len > 0, "cannot allocate 0 bytes on GPU"); assert!(len > 0, "cannot allocate 0 bytes on GPU");
let mut ptr = std::ptr::null_mut(); let a = pool::acquire(len)?;
error::check(unsafe { ffi::cudaMalloc(&mut ptr, len) })?; Ok(Self {
Ok(Self { ptr, len }) ptr: a.ptr,
len,
cap: a.cap,
device: a.device,
})
} }
pub fn len(&self) -> usize { pub fn len(&self) -> usize {
@@ -46,13 +65,20 @@ impl GpuBuffer {
ffi::cudaMemcpy(dst.as_mut_ptr(), self.ptr, dst.len(), ffi::CUDA_MEMCPY_D2H) ffi::cudaMemcpy(dst.as_mut_ptr(), self.ptr, dst.len(), ffi::CUDA_MEMCPY_D2H)
}) })
} }
/// Set every byte of the buffer to `value` on the device (no host copy).
/// Used to zero op-output buffers without a blocking H2D memcpy of zeros.
pub fn memset(&mut self, value: u8) -> Result<()> {
error::check(unsafe { ffi::cudaMemset(self.ptr, value as i32, self.len) })
}
} }
impl Drop for GpuBuffer { impl Drop for GpuBuffer {
fn drop(&mut self) { fn drop(&mut self) {
if !self.ptr.is_null() { // Return to the device pool for reuse (no cudaFree). The pool retains
unsafe { ffi::cudaFree(self.ptr) }; // the raw pointer for the process lifetime; on process exit the OS
} // reclaims the device context, so this is not a leak.
pool::release(self.ptr, self.device, self.cap);
} }
} }

124
crates/xtrain-cuda/src/pool.rs Normal file
View File

@@ -0,0 +1,124 @@
//! Device caching / pool allocator (Phase T11, KI-5).
//!
//! Every tape op allocates its output buffer via [`crate::GpuBuffer::alloc`],
//! which used to call `cudaMalloc` + (for `zeros`) `cudaMemset` on *every* op.
//! `cudaMalloc`/`cudaFree` are synchronous, process-serialized driver calls; in
//! the single-process thread-per-GPU DDP model the rank threads' hundreds of
//! per-step allocations queue through the driver and serialize (KI-5). The cost
//! hurts single-GPU too.
//!
//! Fix: cache freed device buffers in a per-device, size-classed free list and
//! reuse them. Training has repeating shapes, so after warm-up the steady-state
//! `cudaMalloc` count per step is ~0. The pool is **transparent**: a `GpuBuffer`
//! handed out from the pool exposes exactly the bytes the caller requested (the
//! physical allocation may be rounded up to its size class, but `len()` and all
//! copy/memset bounds use the requested length), so numerics are unchanged.
//!
//! Thread-safety: DDP runs thread-per-GPU in one process. The pool is a global
//! registry keyed by device id; each device's free list lives behind its own
//! `Mutex`. A buffer remembers which device it was allocated on (the thread's
//! current CUDA device at `alloc` time) so `Drop` returns it to the right pool.
use crate::error::{self, Result};
use crate::ffi;
use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};
/// Allocation granularity. Requests are rounded *up* to a size class so that
/// op outputs of the same shape (the common case in training) land in the same
/// free list and are reused across steps.
///
/// Small allocations round up to a multiple of `MIN_CLASS`; larger ones round
/// up to the next power of two. Powers of two keep the number of distinct
/// classes bounded (so the free lists stay shallow) while wasting at most ~2×
/// per buffer — fine for fixed-shape training, and freed memory is reused, not
/// leaked.
const MIN_CLASS: usize = 512;
/// Below this threshold, round up to a multiple of `MIN_CLASS` (fine-grained);
/// at or above it, round up to the next power of two.
const POW2_THRESHOLD: usize = 1 << 20; // 1 MiB
/// Round a byte length up to its size class (the physical allocation size).
fn size_class(len: usize) -> usize {
debug_assert!(len > 0);
if len <= POW2_THRESHOLD {
len.div_ceil(MIN_CLASS) * MIN_CLASS
} else {
len.next_power_of_two()
}
}
/// Per-device free list: size class -> stack of cached raw device pointers.
#[derive(Default)]
struct DevicePool {
free: HashMap<usize, Vec<*mut u8>>,
}
// The raw pointers are device addresses, only ever dereferenced by the GPU.
// They are guarded by a `Mutex` and moved between threads as plain handles.
unsafe impl Send for DevicePool {}
type SharedPool = Arc<Mutex<DevicePool>>;
fn registry() -> &'static Mutex<HashMap<i32, SharedPool>> {
static REGISTRY: OnceLock<Mutex<HashMap<i32, SharedPool>>> = OnceLock::new();
REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
}
/// The CUDA device the calling thread is currently set to. DDP sets this once
/// per rank-thread, so it identifies which pool to use.
fn current_device() -> Result<i32> {
let mut dev = 0i32;
error::check(unsafe { ffi::cudaGetDevice(&mut dev) })?;
Ok(dev)
}
/// Run `f` with the (locked) pool for `device`, creating it on first use. The
/// registry mutex is held only long enough to clone out this device's
/// `Arc<Mutex<DevicePool>>`, so different devices' threads don't contend on the
/// per-device free list — true per-rank concurrency.
fn with_device_pool<R>(device: i32, f: impl FnOnce(&mut DevicePool) -> R) -> R {
let pool = {
let mut reg = registry().lock().unwrap();
reg.entry(device).or_default().clone()
};
let mut guard = pool.lock().unwrap();
f(&mut guard)
}
/// Allocation served by the pool: a raw device pointer plus the device it lives
/// on and the size class (capacity) of the physical buffer.
pub(crate) struct PoolAlloc {
pub ptr: *mut u8,
pub device: i32,
pub cap: usize,
}
/// Acquire a buffer of at least `len` bytes for the calling thread's current
/// device. Reuses a cached buffer of the matching size class if one is free,
/// otherwise `cudaMalloc`s a fresh one of the size-class capacity.
pub(crate) fn acquire(len: usize) -> Result<PoolAlloc> {
let cap = size_class(len);
let device = current_device()?;
let cached = with_device_pool(device, |pool| {
pool.free.get_mut(&cap).and_then(|stack| stack.pop())
});
if let Some(ptr) = cached {
return Ok(PoolAlloc { ptr, device, cap });
}
let mut ptr = std::ptr::null_mut();
error::check(unsafe { ffi::cudaMalloc(&mut ptr, cap) })?;
Ok(PoolAlloc { ptr, device, cap })
}
/// Return a buffer to its device's free list for reuse. Does NOT `cudaFree`.
pub(crate) fn release(ptr: *mut u8, device: i32, cap: usize) {
if ptr.is_null() {
return;
}
with_device_pool(device, |pool| {
pool.free.entry(cap).or_default().push(ptr);
});
}

13
crates/xtrain-distributed/Cargo.toml Normal file
View File

@@ -0,0 +1,13 @@
[package]
name = "xtrain-distributed"
version.workspace = true
edition.workspace = true
license.workspace = true
[dependencies]
xtrain-cuda = { path = "../xtrain-cuda" }
xtrain-tensor = { path = "../xtrain-tensor" }
xtrain-autodiff = { path = "../xtrain-autodiff" }
xtrain-model = { path = "../xtrain-model" }
xtrain-optim = { path = "../xtrain-optim" }
xtrain-train = { path = "../xtrain-train" }

33
crates/xtrain-distributed/build.rs Normal file
View File

@@ -0,0 +1,33 @@
use std::env;
use std::path::Path;
use std::process::Command;
// Mirror the per-crate convention (see xtrain-cuda/build.rs): with no nvcc/GPU
// locally, emit `no_cuda` so the NCCL FFI + DDP code compiles (but is not linked
// or run). On dash5, link NCCL exactly like xserv-distributed's build.rs.
fn main() {
println!("cargo:rustc-check-cfg=cfg(no_cuda)");
let cuda_path = env::var("CUDA_HOME")
.or_else(|_| env::var("CUDA_PATH"))
.unwrap_or_else(|_| "/usr/local/cuda".to_string());
if !nvcc_available(&cuda_path) {
println!("cargo:warning=nvcc not found — skipping NCCL link (host-only build).");
println!("cargo:rustc-cfg=no_cuda");
return;
}
println!("cargo:rustc-link-search=native={cuda_path}/lib64");
// NCCL is installed as a system library on dash5.
println!("cargo:rustc-link-search=native=/usr/lib/x86_64-linux-gnu");
println!("cargo:rustc-link-lib=dylib=nccl");
println!("cargo:rustc-link-lib=dylib=cudart");
}
fn nvcc_available(cuda_path: &str) -> bool {
if Command::new("nvcc").arg("--version").output().is_ok() {
return true;
}
Path::new(&format!("{cuda_path}/bin/nvcc")).exists()
}

256
crates/xtrain-distributed/src/bin/train_ddp.rs Normal file
View File

@@ -0,0 +1,256 @@
//! Multi-rank DDP training launcher (Phase T8 / Scaling v2): spawn one thread per
//! GPU, NCCL all-reduce the gradients each step, and train the tiny transformer on
//! TinyStories. At parity with the single-GPU `bin/train`: CLI-tunable arch
//! (scaling-ladder rung), the cached token-id stream, held-out val-loss eval, LR
//! warmup→cosine, grad clip, and best-val checkpointing. Doubles as the throughput
//! driver — run it with 1/2/4 GPUs and read the global tok/s line.
//!
//! Run on dash5 (pick idle GPUs — dash5 is shared):
//! export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
//! CUDA_VISIBLE_DEVICES=1,2 cargo run -p xtrain-distributed --release \
//! --bin train_ddp -- /opt/wjh/models/gpt2/tokenizer.json \
//! data/tinystories-train.txt \
//! --dim 384 --heads 12 --head-dim 32 --layers 12 --ffn 1536 \
//! --steps 6000 --batch 32 --seq 256 --max-lr 6e-4 \
//! --val-tokens 1000000 --eval-every 500 --ckpt /tmp/xtrain_v2.ckpt
//!
//! Positional: <tokenizer.json> <corpus.txt>. Everything else is a flag with a
//! sane default. The launcher uses every GPU visible to it (CUDA_VISIBLE_DEVICES
//! selects them), so rank devices are always 0..N within the visible set.
#[cfg(no_cuda)]
fn main() {
eprintln!("train_ddp: built without CUDA (no_cuda); run on a GPU host (dash5).");
}
#[cfg(not(no_cuda))]
use std::path::PathBuf;
// A flag like `--dim 384`: scan argv for `name`, parse the following token.
#[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 main() {
use xtrain_cuda::device;
use xtrain_distributed::{DdpConfig, build_model, launch};
use xtrain_model::Config;
use xtrain_train::data::Corpus;
use xtrain_train::schedule::LrSchedule;
let args: Vec<String> = std::env::args().collect();
// First two non-flag positionals: tokenizer.json, corpus.txt.
let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
let tok_path = positionals
.first()
.map(|s| PathBuf::from(s.as_str()))
.unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
let corpus_path = positionals
.get(1)
.map(|s| PathBuf::from(s.as_str()))
.unwrap_or_else(|| PathBuf::from("data/tinystories-valid-3mb.txt"));
// Architecture (scaling-ladder rung). Defaults = v0-baseline tiny config.
let n_heads = flag(&args, "--heads", 2usize);
let head_dim = flag(&args, "--head-dim", 16usize);
let n_layers = flag(&args, "--layers", 4usize);
let ffn = flag(&args, "--ffn", 64usize);
// GQA (Phase T15): num K/V heads (must divide --heads). Default = --heads (MHA).
let kv_heads = flag(&args, "--kv-heads", n_heads);
// `--dim` is informational; dim is always n_heads*head_dim. Warn on mismatch.
let dim_flag = flag(&args, "--dim", 0usize);
if dim_flag != 0 && dim_flag != n_heads * head_dim {
eprintln!(
"warning: --dim {dim_flag} != heads*head_dim {}; using {}",
n_heads * head_dim,
n_heads * head_dim
);
}
// Optimization knobs (mirror bin/train).
let steps: usize = flag(&args, "--steps", 100);
let batch: usize = flag(&args, "--batch", 16);
// Micro-batch gradient accumulation (Phase T16): effective global batch =
// accum_steps × batch, all-reducing only at the accumulation boundary. Default
// 1 = no accumulation (bit-identical to the pre-T16 DDP path).
let accum_steps: usize = flag(&args, "--accum-steps", 1).max(1);
let seq_len: usize = flag(&args, "--seq", 64);
let max_lr: f32 = flag(&args, "--max-lr", 3e-3);
let min_lr: f32 = flag(&args, "--min-lr", max_lr * 0.1);
let weight_decay: f32 = flag(&args, "--wd", 0.1);
let max_grad_norm: f32 = flag(&args, "--clip", 1.0);
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
// train_rank loop calls model.train() each step so dropout is actually live
// under DDP (T21 wired this — the launcher previously never set training mode).
let dropout: f32 = flag(&args, "--dropout", 0.0f32);
// bf16 mixed precision (Phase T12): fp32 master weights, bf16 linears +
// activations. Opt-in; default fp32 reproduces v0–v4 numerics.
let bf16 = args.iter().any(|a| a == "--bf16");
// Activation recomputation (Phase T13): per-block gradient checkpointing — each
// rank checkpoints its own forward/backward; exact grads, lower peak activation
// memory (lets dim1024 batch32 fit). Opt-in; default off.
let recompute = args.iter().any(|a| a == "--recompute");
// Fused flash-attention (Phase T14): single fused SDPA kernel, online softmax,
// no materialized [bh,S,S] scores. Opt-in; default off keeps the composed path.
let flash = args.iter().any(|a| a == "--flash");
let ckpt: Option<PathBuf> = args
.iter()
.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).
let count = device::device_count().expect("device_count") as u32;
assert!(count > 0, "no CUDA device visible");
let devices: Vec<u32> = (0..count).collect();
assert_eq!(
batch % devices.len(),
0,
"global batch {batch} not divisible by world {}",
devices.len()
);
println!(
"DDP: world={} devices={:?} | steps={steps} seq={seq_len} global_batch={batch}",
devices.len(),
devices
);
// 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)
};
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 {
let (t, v) = corpus.split_tail(val_tokens);
println!("split: {} train tokens / {} val tokens", t.len(), v.len());
(t, Some(v))
} else {
(corpus, None)
};
let mut cfg =
Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
cfg.dropout = dropout;
println!(
"model: dim {} layers {} heads {} kv_heads {} head_dim {} ffn {} → core {:.3}M params \
(+ embed/lm {:.2}M = {:.2}M total)",
cfg.dim,
cfg.n_layers,
cfg.n_heads,
cfg.num_kv_heads,
cfg.head_dim,
cfg.ffn_hidden,
cfg.core_params() as f32 / 1e6,
(cfg.num_params() - cfg.core_params()) as f32 / 1e6,
cfg.num_params() as f32 / 1e6,
);
let dcfg = DdpConfig {
seq_len,
batch_size: batch,
accum_steps,
steps,
schedule: LrSchedule {
max_lr,
min_lr,
warmup: (steps / 20).max(5),
total: steps,
},
weight_decay,
max_grad_norm,
log_every: 50,
seed: 42,
eval_every,
eval_batches,
ckpt_path: ckpt.clone(),
};
println!(
"training: {steps} steps, seq {seq_len}, global batch {batch} × accum {accum_steps} = \
effective global batch {}, lr {max_lr:.1e}→{min_lr:.1e}, eval every {eval_every}",
batch * accum_steps
);
if bf16 {
println!("bf16 mixed precision: ON (fp32 master weights)");
}
if recompute {
println!("activation recompute: ON (per-block gradient checkpointing)");
}
if 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)");
}
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,
valid.as_ref(),
&dcfg,
move |device| {
let mut m = build_model(cfg, device);
if bf16 {
m = m.with_compute_dtype(xtrain_tensor::DType::BF16);
}
if recompute {
m = m.with_recompute(true);
}
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
},
);
let r0 = &results[0];
let start = r0.losses.first().copied().unwrap_or(0.0);
let end = r0.losses.last().copied().unwrap_or(0.0);
println!("train loss: start {start:.4} → end {end:.4}");
if let Some(best) = r0.best_val {
println!("best val loss: {best:.4}");
}
if let Some((s, v)) = r0.evals.last() {
println!("final val loss (step {s}): {v:.4}");
}
if let Some(path) = &ckpt {
println!("best-val checkpoint → {}", path.display());
}
}

215
crates/xtrain-distributed/src/bin/train_ddp_mp.rs Normal file
View File

@@ -0,0 +1,215 @@
//! Process-per-GPU DDP launcher / worker (Phase T17, torchrun-style).
//!
//! ONE binary, two modes (it self-detects via `XTRAIN_RANK`):
//! - **launcher** (env unset): mints the NCCL `ncclUniqueId`, then spawns one
//! WORKER process per visible GPU, re-execing this same binary with the same
//! argv plus `XTRAIN_{RANK,WORLD,LOCAL_RANK,NCCL_ID}` env, and waits for them.
//! - **worker** (`XTRAIN_RANK` set): binds its GPU (→ its own CUDA context),
//! inits NCCL with the launcher-supplied id, builds its model, runs
//! `train_rank` — the T8 training step reused UNCHANGED.
//!
//! 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.
//!
//! Run on dash5 (pick idle GPUs — dash5 is shared):
//! export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
//! CUDA_VISIBLE_DEVICES=0,1,2,3 cargo run -p xtrain-distributed --release \
//! --bin train_ddp_mp -- /opt/wjh/models/gpt2/tokenizer.json \
//! data/tinystories-valid-3mb.txt \
//! --dim 384 --heads 12 --head-dim 32 --layers 12 --ffn 1536 \
//! --steps 200 --batch 128 --seq 256
#[cfg(no_cuda)]
fn main() {
eprintln!("train_ddp_mp: built without CUDA (no_cuda); run on a GPU host (dash5).");
}
#[cfg(not(no_cuda))]
use std::path::PathBuf;
// A flag like `--dim 384`: scan argv for `name`, parse the following token.
#[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 main() {
use xtrain_cuda::device;
use xtrain_distributed::DdpConfig;
use xtrain_distributed::proc::{ModelOpts, launch_processes, run_worker, worker_env};
use xtrain_model::Config;
use xtrain_train::data::Corpus;
use xtrain_train::schedule::LrSchedule;
let args: Vec<String> = std::env::args().collect();
// ── Launcher mode: no XTRAIN_RANK in env → spawn one worker per visible GPU.
let env = worker_env();
if env.is_none() {
let count = device::device_count().expect("device_count");
assert!(count > 0, "no CUDA device visible");
let world = count as usize;
// Forward the full argv (minus argv[0]) to each worker verbatim.
let extra: Vec<String> = args[1..].to_vec();
println!("DDP (process-per-GPU): launching {world} worker processes (one per visible GPU)");
match launch_processes(world, &extra) {
Ok(()) => {}
Err(e) => {
eprintln!("launcher: {e}");
std::process::exit(1);
}
}
return;
}
let env = env.unwrap();
// ── Worker mode: build config from the forwarded argv, then train this rank.
// First two non-flag positionals: tokenizer.json, corpus.txt.
let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
let tok_path = positionals
.first()
.map(|s| PathBuf::from(s.as_str()))
.unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
let corpus_path = positionals
.get(1)
.map(|s| PathBuf::from(s.as_str()))
.unwrap_or_else(|| PathBuf::from("data/tinystories-valid-3mb.txt"));
// Architecture (scaling-ladder rung). Defaults = v0-baseline tiny config.
let n_heads = flag(&args, "--heads", 2usize);
let head_dim = flag(&args, "--head-dim", 16usize);
let n_layers = flag(&args, "--layers", 4usize);
let ffn = flag(&args, "--ffn", 64usize);
let kv_heads = flag(&args, "--kv-heads", n_heads);
let dim_flag = flag(&args, "--dim", 0usize);
if dim_flag != 0 && dim_flag != n_heads * head_dim {
eprintln!(
"warning: --dim {dim_flag} != heads*head_dim {}; using {}",
n_heads * head_dim,
n_heads * head_dim
);
}
// Optimization knobs (mirror train_ddp).
let steps: usize = flag(&args, "--steps", 100);
let batch: usize = flag(&args, "--batch", 16);
let accum_steps: usize = flag(&args, "--accum-steps", 1).max(1);
let seq_len: usize = flag(&args, "--seq", 64);
let max_lr: f32 = flag(&args, "--max-lr", 3e-3);
let min_lr: f32 = flag(&args, "--min-lr", max_lr * 0.1);
let weight_decay: f32 = flag(&args, "--wd", 0.1);
let max_grad_norm: f32 = flag(&args, "--clip", 1.0);
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"),
flash: args.iter().any(|a| a == "--flash"),
};
let ckpt: Option<PathBuf> = args
.iter()
.position(|a| a == "--ckpt")
.and_then(|i| args.get(i + 1))
.map(PathBuf::from);
assert_eq!(
batch % env.world,
0,
"global batch {batch} not divisible by world {}",
env.world
);
// Each worker loads the corpus independently (read-only u16 cache hit → cheap).
let corpus = Corpus::load_cached(&tok_path, &corpus_path);
let vocab = corpus.vocab_size;
let (train_corpus, valid): (Corpus, Option<Corpus>) = if val_tokens > 0 {
let (t, v) = corpus.split_tail(val_tokens);
(t, Some(v))
} else {
(corpus, None)
};
let mut cfg =
Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
cfg.dropout = dropout;
if env.rank == 0 {
println!(
"model: dim {} layers {} heads {} kv_heads {} head_dim {} ffn {} → core {:.3}M params \
(+ embed/lm {:.2}M = {:.2}M total) | world={} mode=process-per-GPU",
cfg.dim,
cfg.n_layers,
cfg.n_heads,
cfg.num_kv_heads,
cfg.head_dim,
cfg.ffn_hidden,
cfg.core_params() as f32 / 1e6,
(cfg.num_params() - cfg.core_params()) as f32 / 1e6,
cfg.num_params() as f32 / 1e6,
env.world,
);
if opts.bf16 {
println!("bf16 mixed precision: ON (fp32 master weights)");
}
if opts.recompute {
println!("activation recompute: ON (per-block gradient checkpointing)");
}
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 {
seq_len,
batch_size: batch,
accum_steps,
steps,
schedule: LrSchedule {
max_lr,
min_lr,
warmup: (steps / 20).max(5),
total: steps,
},
weight_decay,
max_grad_norm,
log_every: 50,
seed: 42,
eval_every,
eval_batches,
ckpt_path: ckpt.clone(),
};
let res = run_worker(&env, cfg, opts, &train_corpus, valid.as_ref(), &dcfg);
if env.rank == 0 {
let start = res.losses.first().copied().unwrap_or(0.0);
let end = res.losses.last().copied().unwrap_or(0.0);
println!("train loss: start {start:.4} → end {end:.4}");
if let Some(best) = res.best_val {
println!("best val loss: {best:.4}");
}
if let Some((s, v)) = res.evals.last() {
println!("final val loss (step {s}): {v:.4}");
}
if let Some(path) = &ckpt {
println!("best-val checkpoint → {}", path.display());
}
}
}

307
crates/xtrain-distributed/src/ddp.rs Normal file
View File

@@ -0,0 +1,307 @@
//! The DDP training step + a single-process, thread-per-GPU launcher (Phase T8).
//!
//! Each rank owns one GPU and one thread. Per step it processes a DISJOINT shard
//! of the global batch, all-reduce-averages the gradients, then runs its own
//! `GpuAdamW.step`. Identical init + identical optimizer state across ranks keep
//! the parameters consistent — verified by the cross-rank param-identity check in
//! the tests.
//!
//! Sampling matches single-GPU bit-for-bit: every rank advances the SAME RNG and
//! draws all `B_global` sequences of a step, but only runs forward+backward on
//! the ones assigned to it (`global index % world == rank`). The union over ranks
//! is exactly the single-GPU batch in the same order, so the all-reduced grad sum
//! equals the single-GPU summed grad.
use std::path::PathBuf;
use std::thread;
use std::time::Instant;
use xtrain_autodiff::tape::Var;
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
use xtrain_optim::GpuAdamW;
use xtrain_tensor::Device;
use xtrain_train::checkpoint;
use xtrain_train::clip::clip_grad_norm_gpu;
use xtrain_train::data::Corpus;
use xtrain_train::eval_loss;
use xtrain_train::schedule::LrSchedule;
use crate::{DdpContext, get_unique_id};
/// Per-rank DDP training config. `batch_size` is the GLOBAL batch (split across
/// ranks); the rest mirror `xtrain_train::TrainConfig`.
#[derive(Clone)]
pub struct DdpConfig {
pub seq_len: usize,
/// Global batch size; must be divisible by the world size.
pub batch_size: usize,
/// Micro-batch gradient accumulation (Phase T16): each optimizer step
/// accumulates grads over `accum_steps` micro-batches, giving an EFFECTIVE
/// global batch of `accum_steps × batch_size`. The cross-rank all-reduce
/// fires ONLY at the accumulation boundary (after the last micro-step) —
/// intermediate micro-steps skip the NCCL collective entirely. `1` = no
/// accumulation (bit-identical to the pre-T16 DDP path).
pub accum_steps: usize,
pub steps: usize,
pub schedule: LrSchedule,
pub weight_decay: f32,
pub max_grad_norm: f32,
pub log_every: usize,
pub seed: u64,
/// Evaluate held-out val loss every `eval_every` steps (0 = never). Only rank
/// 0 holds the `valid` corpus and runs the eval (no grad), mirroring
/// `xtrain_train::TrainConfig`. The best-val model is checkpointed by rank 0
/// (every rank's params are identical, so rank 0's are the model's).
pub eval_every: usize,
pub eval_batches: usize,
/// Best-val checkpoint path (written by rank 0 when val improves). When unset,
/// or when `eval_every == 0`, no checkpoint is written.
pub ckpt_path: Option<PathBuf>,
}
/// Outcome of a DDP run on this rank: per-step mean-loss trace plus, when
/// `eval_every > 0`, the (step, val_loss) eval points and the best val loss
/// (eval/best are only populated on rank 0, which owns the `valid` corpus).
pub struct DdpResult {
pub losses: Vec<f32>,
pub evals: Vec<(usize, f32)>,
pub best_val: Option<f32>,
}
/// Run `cfg.steps` DDP steps on this rank's `model`/`corpus`, using `ctx` for the
/// gradient all-reduce. Returns this rank's per-step mean-loss trace (the mean
/// over the GLOBAL batch — every rank computes the same value because losses are
/// all-reduced alongside the grads) plus eval/best-val (rank 0 only). The
/// optimizer step is identical on every rank, so the parameters stay in lockstep.
///
/// `valid` is the held-out corpus for periodic val-loss eval. Only rank 0 needs
/// it (it runs the no-grad eval and writes the best-val checkpoint); pass `None`
/// on the other ranks (or when `cfg.eval_every == 0`).
pub fn train_rank(
ctx: &DdpContext,
model: &TinyTransformer,
device: Device,
corpus: &Corpus,
valid: Option<&Corpus>,
cfg: &DdpConfig,
) -> DdpResult {
assert_eq!(
cfg.batch_size % ctx.world,
0,
"global batch {} not divisible by world {}",
cfg.batch_size,
ctx.world
);
let params = model.params();
let mut opt = GpuAdamW::new(cfg.weight_decay);
let mut rng = cfg.seed;
let mut losses = Vec::with_capacity(cfg.steps);
let mut evals = Vec::new();
let mut best_val: Option<f32> = None;
// Each rank runs ONE batched forward over its b_local = batch_size/world
// sequences → backward grad = local mean (Σ_local / b_local). all_reduce_average
// (sum across ranks, /world) then gives Σ_global/(world·b_local) = Σ_global/
// B_global — already the global-batch mean — so the clip pre-scale is 1.0.
let batch_local = cfg.batch_size / ctx.world;
let accum = cfg.accum_steps.max(1);
let start = Instant::now();
let mut tokens_seen: u64 = 0;
// Rank 0 owns the held-out eval + best-val checkpoint (params are identical
// across ranks, so rank 0's are the model). Other ranks never touch `valid`.
let do_eval = ctx.rank == 0 && cfg.eval_every > 0 && valid.is_some();
for step in 0..cfg.steps {
let lr = cfg.schedule.lr(step);
// Accumulate grads over `accum` micro-batches, then ONE optimizer step
// (Phase T16). Per micro-batch: draw the whole micro global batch from the
// shared RNG (same on every rank), keep only this rank's shard (global index
// % world == rank), run it as ONE batched forward/backward. Each micro-loss
// is scaled by 1/accum before backward (the tape SUM-accumulates the scaled
// grads across the `accum` micro-backwards) so the boundary grad equals a
// single step over an `accum × batch_size` global batch. `accum == 1` skips
// the scale → bit-identical to the pre-T16 DDP path. The cross-rank
// all-reduce fires ONLY after the last micro-step (intermediate micro-steps
// are local-only, no NCCL).
let mut local_sum = 0.0f32; // Σ over micro of (local_mean · b_local)
// Training mode → dropout active (T18; no-op when cfg.dropout == 0). Set
// each step so it is restored after a periodic eval flips the model to eval
// mode (eval_loss calls model.eval() and does not restore). Mirrors the
// single-GPU loop's train/eval discipline — without this, DDP forwards run
// in the default eval (identity) mode and --dropout is silently ignored
// (the T21 launcher-wiring gap the V9-PILOT caught). Each micro-step's
// forward bumps the per-step seed → fresh masks.
model.train();
for _ in 0..accum {
let mut inputs = Vec::with_capacity(batch_local);
let mut targets_v = Vec::with_capacity(batch_local);
for i in 0..cfg.batch_size {
let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
if i % ctx.world == ctx.rank {
inputs.push(input);
targets_v.push(target);
}
}
let ids = batched_ids_tensor(&inputs, device);
let targets = batched_ids_tensor(&targets_v, device);
let loss = model.loss_batched(&ids, &targets, batch_local);
local_sum += read_scalar(&loss) * batch_local as f32; // local mean·b_local
if accum == 1 {
loss.backward();
} else {
xtrain_autodiff::ops::scale(&loss, 1.0 / accum as f32).backward();
}
tokens_seen += (batch_local * cfg.seq_len) as u64;
}
// Accumulation boundary: ONE AllReduce(sum) + /world over the accumulated
// grads → every rank holds the effective-batch (accum·B_global) mean grad
// (the per-micro 1/accum scaling is already baked into each backward; the
// /world here is orthogonal to accum). Intermediate micro-steps issued NO
// NCCL — only this single boundary collective per optimizer step.
ctx.all_reduce_average_grads(&params);
// Reported loss = effective-batch mean: AllReduce(sum) the per-rank local
// sums across ranks, /(accum·B_global).
let step_loss = all_reduce_loss(ctx, local_sum) / (accum * cfg.batch_size) as f32;
losses.push(step_loss);
// Grads are already the effective-batch mean — just clip (pre-scale 1.0).
let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, 1.0);
opt.step(lr, &params);
for p in &params {
p.zero_grad();
}
if ctx.rank == 0 && (step % cfg.log_every == 0 || step == cfg.steps - 1) {
let elapsed = start.elapsed().as_secs_f32();
// Global tok/s = per-rank tok/s × world (each rank does 1/world of it).
let tps = (tokens_seen as f32 / elapsed.max(1e-6)) * ctx.world as f32;
println!(
"[rank0] step {step:5}/{}: loss {step_loss:.4} lr {lr:.2e} gnorm {gnorm:.3} \
({tps:.0} tok/s global, {} ranks)",
cfg.steps, ctx.world
);
}
// Periodic held-out eval + best-val checkpoint (rank 0 only). Mirrors the
// single-GPU `xtrain_train::train` loop, reusing its `eval_loss` /
// `checkpoint::save` so single-GPU and DDP share one eval/ckpt path. Other
// ranks have nothing to do here (params are identical across ranks).
if do_eval && ((step + 1) % cfg.eval_every == 0 || step == cfg.steps - 1) {
let v = valid.unwrap();
let vl = eval_loss(model, device, v, cfg.seq_len, cfg.eval_batches);
evals.push((step, vl));
let improved = best_val.map(|b| vl < b).unwrap_or(true);
println!(
" [rank0] eval @ step {step}: val loss {vl:.4}{}",
if improved { " (best)" } else { "" }
);
if improved {
best_val = Some(vl);
if let Some(path) = &cfg.ckpt_path {
checkpoint::save(path, &params).expect("best checkpoint save");
}
}
}
}
DdpResult {
losses,
evals,
best_val,
}
}
/// Spawn `world` rank threads (one per GPU in `devices`), init NCCL, build an
/// identical model per rank via `make_model`, and run `train_rank`. Returns each
/// rank's `DdpResult` (loss traces are identical; eval/best-val are on rank 0).
/// The launcher owns the thread-per-GPU model: rank 0 mints the `UniqueId`, every
/// thread `cudaSetDevice`s its GPU, builds its `Var` graph locally (the graph is
/// `!Send`), and joins at the end.
///
/// `valid` is the held-out corpus for rank 0's periodic eval (only used when
/// `cfg.eval_every > 0`). `make_model(device)` must be deterministic — same params
/// on every rank — for the parameters to stay consistent.
pub fn launch<F>(
devices: &[u32],
corpus: &Corpus,
valid: Option<&Corpus>,
cfg: &DdpConfig,
make_model: F,
) -> Vec<DdpResult>
where
F: Fn(Device) -> TinyTransformer + Send + Sync,
{
let world = devices.len();
let id = get_unique_id();
thread::scope(|s| {
let handles: Vec<_> = devices
.iter()
.enumerate()
.map(|(rank, &dev)| {
let make_model = &make_model;
let cfg = cfg.clone();
s.spawn(move || {
let ctx = DdpContext::init(rank, world, id, dev);
let device = Device::Cuda(dev);
let model = make_model(device);
// Only rank 0 holds the val corpus for eval.
let v = if rank == 0 { valid } else { None };
train_rank(&ctx, &model, device, corpus, v, &cfg)
})
})
.collect();
handles.into_iter().map(|h| h.join().unwrap()).collect()
})
}
/// AllReduce(sum) a single host scalar across ranks by round-tripping it through a
/// one-element device buffer. Used only for the logged/returned loss, so the cost
/// (one tiny collective per step) is negligible. Returns the summed value.
fn all_reduce_loss(ctx: &DdpContext, local: f32) -> f32 {
use xtrain_tensor::Tensor;
if ctx.world == 1 {
return local;
}
let device = Device::Cuda(ctx.device);
let t = Tensor::from_slice(&[local], &[1]).to_device(device);
ctx.all_reduce_sum_f32_ptr(t.data_ptr() as *mut std::ffi::c_void, 1);
xtrain_cuda::device::synchronize().expect("loss all-reduce sync");
t.to_device(Device::Cpu).as_slice::<f32>()[0]
}
fn read_scalar(v: &Var) -> f32 {
v.value().to_device(Device::Cpu).as_slice::<f32>()[0]
}
/// Build a `TinyTransformer` on `device` with the SAME deterministic init the
/// single-GPU `bin/train` uses (LCG fill, gammas ~1). Used by both the launcher
/// and the correctness test so every rank — and the single-GPU baseline — start
/// from bit-identical parameters. `cfg` must be identical on every call.
pub fn build_model(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.04)
}
})
}
// Deterministic LCG fill in [-scale, scale) — same scheme as bin/train's `fill`.
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()
}

76
crates/xtrain-distributed/src/ffi.rs Normal file
View File

@@ -0,0 +1,76 @@
//! Minimal NCCL FFI bindings (hand-written, like the CUDA bindings in
//! xtrain-cuda). Only the collectives data-parallel training needs:
//! unique-id creation, communicator init/destroy, and AllReduce. Mirrors
//! xserv-distributed's FFI.
use std::ffi::c_void;
use std::os::raw::c_char;
use xtrain_cuda::ffi::CudaStream;
/// Opaque NCCL communicator handle (`ncclComm_t`).
pub type NcclComm = *mut c_void;
/// `ncclUniqueId` is a 128-byte opaque blob shared from rank 0 to every rank.
#[repr(C)]
#[derive(Clone, Copy)]
pub struct NcclUniqueId {
pub internal: [c_char; 128],
}
impl Default for NcclUniqueId {
fn default() -> Self {
Self { internal: [0; 128] }
}
}
// ncclDataType_t (subset) — DDP all-reduces fp32 gradients.
pub const NCCL_FLOAT32: i32 = 7;
// ncclRedOp_t
pub const NCCL_SUM: i32 = 0;
// ncclResult_t
pub const NCCL_SUCCESS: i32 = 0;
unsafe extern "C" {
pub fn ncclGetUniqueId(uid: *mut NcclUniqueId) -> i32;
// ncclUniqueId is passed BY VALUE (a 128-byte struct) per the NCCL ABI.
pub fn ncclCommInitRank(
comm: *mut NcclComm,
nranks: i32,
commid: NcclUniqueId,
rank: i32,
) -> i32;
pub fn ncclCommDestroy(comm: NcclComm) -> i32;
pub fn ncclAllReduce(
sendbuff: *const c_void,
recvbuff: *mut c_void,
count: usize,
datatype: i32,
op: i32,
comm: NcclComm,
stream: CudaStream,
) -> i32;
pub fn ncclGroupStart() -> i32;
pub fn ncclGroupEnd() -> i32;
pub fn ncclGetErrorString(result: i32) -> *const c_char;
}
pub fn err_string(result: i32) -> String {
unsafe {
let p = ncclGetErrorString(result);
if p.is_null() {
return format!("nccl error {result}");
}
std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned()
}
}
pub fn check(result: i32, what: &str) {
assert_eq!(
result,
NCCL_SUCCESS,
"{what} failed: {}",
err_string(result)
);
}

158
crates/xtrain-distributed/src/lib.rs Normal file
View File

@@ -0,0 +1,158 @@
//! Distributed data-parallel (DDP) primitives for xtrain (Phase T8).
//!
//! Launch model: **one OS thread per GPU** (same as xserv-distributed). Each
//! rank thread binds its device, builds its own model (xtrain's `Var` graph is
//! `Rc`-based and not `Send`, so it must be constructed thread-locally — only the
//! `UniqueId` and scalar config cross the thread boundary), processes a disjoint
//! shard of the global batch, then AllReduces every parameter's `.grad()` device
//! buffer in place, averages by world size, and runs its own `GpuAdamW.step`.
//! Identical init + identical optimizer state across ranks keeps the parameters
//! consistent without ever re-syncing the weights.
//!
//! NCCL is issued on the legacy null stream — every xtrain kernel launches on the
//! null stream (`std::ptr::null_mut()`), so the AllReduce stays correctly ordered
//! after the producing backward kernels and before the consuming optimizer step,
//! with no extra synchronization.
#![cfg(not(no_cuda))]
pub mod ddp;
pub mod ffi;
pub mod proc;
pub use ddp::{DdpConfig, DdpResult, build_model, launch, train_rank};
pub use proc::{
ModelOpts, WorkerEnv, build_worker_model, hex_decode_unique_id, hex_encode_unique_id,
launch_processes, run_worker, worker_env,
};
use std::ffi::c_void;
use ffi::{NcclComm, NcclUniqueId};
use xtrain_autodiff::tape::Var;
use xtrain_cuda::device;
pub use ffi::NcclUniqueId as UniqueId;
/// Generate a unique id on one rank (rank 0) and share the raw bytes to every
/// other rank out-of-band — across threads it is just a `Copy` struct moved into
/// each rank closure; across processes it would be written to a file/env.
pub fn get_unique_id() -> NcclUniqueId {
let mut id = NcclUniqueId::default();
ffi::check(unsafe { ffi::ncclGetUniqueId(&mut id) }, "ncclGetUniqueId");
id
}
/// Per-rank data-parallel context: the NCCL communicator plus this rank's
/// identity. AllReduce is in-place on the null stream.
pub struct DdpContext {
pub rank: usize,
pub world: usize,
pub device: u32,
comm: NcclComm,
}
// The communicator is owned by exactly one rank thread.
unsafe impl Send for DdpContext {}
impl DdpContext {
/// Initialize this rank. Must run on the thread that will own this rank's GPU
/// work; binds the thread to `device` first. All ranks call this concurrently
/// with the same `id` and `world` — the group wrapper lets the concurrent
/// inits rendezvous without deadlock.
pub fn init(rank: usize, world: usize, id: NcclUniqueId, device: u32) -> Self {
device::set_device(device).expect("set_device");
let mut comm: NcclComm = std::ptr::null_mut();
ffi::check(unsafe { ffi::ncclGroupStart() }, "ncclGroupStart(init)");
ffi::check(
unsafe { ffi::ncclCommInitRank(&mut comm, world as i32, id, rank as i32) },
"ncclCommInitRank",
);
ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
Self {
rank,
world,
device,
comm,
}
}
/// In-place AllReduce(sum) over `count` F32 elements at a raw device pointer,
/// issued on the null stream (so it orders with this rank's kernels). The
/// reduction is asynchronous; a later sync (the caller's, or the next null-
/// stream kernel) completes it.
///
/// # Safety
/// `ptr` must point to at least `count` valid F32 device elements on this
/// rank's device. The reduction is in-place (send == recv).
pub fn all_reduce_sum_f32_ptr(&self, ptr: *mut c_void, count: usize) {
if self.world == 1 {
return; // nothing to reduce
}
ffi::check(
unsafe {
ffi::ncclAllReduce(
ptr as *const c_void,
ptr,
count,
ffi::NCCL_FLOAT32,
ffi::NCCL_SUM,
self.comm,
std::ptr::null_mut(),
)
},
"ncclAllReduce",
);
}
/// AllReduce every parameter's `.grad()` across ranks and divide by `world`,
/// the one collective DDP needs per step.
///
/// Each rank ran forward+backward on its own shard of `b` sequences, so
/// `.grad()` holds the SUM over that shard (the tape's fan-out rule). After
/// `AllReduce(sum)` every rank holds `Σ_global` (the sum over all `world·b`
/// sequences); dividing by `world` leaves `Σ_global / world`. The DDP train
/// loop's clip pass then applies the remaining `1/b` (`pre_scale = 1/b_local`),
/// giving `Σ_global / (world·b) = Σ_global / B_global` — bit-for-bit the same
/// mean gradient the single-GPU loop computes from a batch of `B_global`.
/// Params without a grad are skipped.
///
/// A single-process group barrier is unnecessary: the all-reduces serialize
/// on the comm, and the in-place scale runs on the same null stream after.
pub fn all_reduce_average_grads(&self, params: &[Var]) {
if self.world == 1 {
return;
}
// 1. Sum every grad across ranks (in place, on the null stream).
for p in params {
if let Some(g) = p.grad() {
let n = g.numel();
self.all_reduce_sum_f32_ptr(g.data_ptr() as *mut c_void, n);
}
}
// 2. Average: scale each summed grad by 1/world (null-stream kernel,
// ordered after the AllReduce that produced it).
let inv_world = 1.0 / self.world as f32;
for p in params {
if let Some(g) = p.grad() {
unsafe {
xtrain_cuda::ffi::launch_scale_inplace_f32(
g.data_ptr() as *mut f32,
inv_world,
g.numel() as i32,
std::ptr::null_mut(),
);
}
}
}
device::synchronize().expect("grad all-reduce sync failed");
}
}
impl Drop for DdpContext {
fn drop(&mut self) {
if !self.comm.is_null() {
unsafe { ffi::ncclCommDestroy(self.comm) };
}
}
}

200
crates/xtrain-distributed/src/proc.rs Normal file
View File

@@ -0,0 +1,200 @@
//! Process-per-GPU DDP launcher + worker (Phase T17, torchrun-style).
//!
//! T8's DDP is single-process, thread-per-GPU: N rank threads share ONE CUDA
//! primary context, so much of the driver work (kernel launch, cuBLAS handle,
//! stream queueing) serializes at the context level — the residual ~5×@8
//! non-linearity left after T11's allocator fix (see docs/10 / KI-5).
//!
//! Process-per-GPU gives each rank its OWN OS process and OWN CUDA context, so
//! those driver calls no longer queue in a shared context. Only the LAUNCH model
//! and the cross-process NCCL bootstrap change; the training step
//! (`train_rank` → grad all-reduce → local AdamW) and the consistency argument
//! are reused from T8 UNCHANGED.
//!
//! UniqueId rendezvous: the LAUNCHER (the common parent of every worker) mints
//! the `ncclUniqueId` once, hex-encodes it, and injects it into each worker's env
//! at spawn time. No shared file / TCP server / polling — the id is atomically
//! present before the child exists, so there is no "id not ready yet" race. This
//! is the simplest single-node mechanism (see docs/16).
use std::path::PathBuf;
use std::process::{Command, Stdio};
use xtrain_model::{Config, TinyTransformer};
use xtrain_tensor::{DType, Device};
use xtrain_train::data::Corpus;
use crate::ddp::{DdpConfig, DdpResult, build_model, train_rank};
use crate::ffi::NcclUniqueId;
use crate::{DdpContext, get_unique_id};
// Env keys the launcher sets on every spawned worker (torchrun-style: a worker
// detects its role by the presence of `XTRAIN_RANK`).
pub const ENV_RANK: &str = "XTRAIN_RANK";
pub const ENV_WORLD: &str = "XTRAIN_WORLD";
pub const ENV_LOCAL_RANK: &str = "XTRAIN_LOCAL_RANK";
pub const ENV_NCCL_ID: &str = "XTRAIN_NCCL_ID";
/// Hex-encode the 128-byte `ncclUniqueId` for env transport (128 B → 256 chars,
/// well under any env-var length limit). `c_char` is signed on this target, so
/// reinterpret the bytes as `u8` first.
pub fn hex_encode_unique_id(id: &NcclUniqueId) -> String {
let mut s = String::with_capacity(256);
for &b in &id.internal {
s.push_str(&format!("{:02x}", b as u8));
}
s
}
/// Inverse of [`hex_encode_unique_id`]: parse 256 hex chars back into the
/// 128-byte opaque blob. Panics on malformed input (the launcher always writes a
/// well-formed value, so a bad value means a corrupted env).
pub fn hex_decode_unique_id(hex: &str) -> NcclUniqueId {
assert_eq!(
hex.len(),
256,
"NCCL id hex must be 256 chars, got {}",
hex.len()
);
let mut id = NcclUniqueId::default();
for (i, slot) in id.internal.iter_mut().enumerate() {
let byte = u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16).expect("NCCL id hex byte parse");
*slot = byte as std::os::raw::c_char;
}
id
}
/// Spawn `world` worker processes (re-exec of the current binary with the same
/// argv), each pinned to one GPU via `XTRAIN_LOCAL_RANK`, and wait for all of
/// them. The launcher mints the `ncclUniqueId` and injects it (hex) into every
/// worker's env, so the cross-process NCCL bootstrap needs no shared file/TCP.
///
/// Returns `Ok(())` iff every worker exits 0; otherwise an error naming the first
/// failing rank (so the caller — `main` / a test — can propagate a non-zero exit).
/// `extra_args` is forwarded to each worker verbatim (so all training hyper-params
/// pass straight through); the workers inherit the launcher's env (incl.
/// `CUDA_VISIBLE_DEVICES`) plus the four `XTRAIN_*` keys.
pub fn launch_processes(world: usize, extra_args: &[String]) -> Result<(), String> {
let exe = std::env::current_exe().map_err(|e| format!("current_exe: {e}"))?;
let id = get_unique_id();
let id_hex = hex_encode_unique_id(&id);
let mut children = Vec::with_capacity(world);
for rank in 0..world {
let child = Command::new(&exe)
.args(extra_args)
.env(ENV_RANK, rank.to_string())
.env(ENV_WORLD, world.to_string())
// Single node: local rank == global rank == device ordinal within the
// visible set. (Multi-node would split these; see docs/16 follow-up.)
.env(ENV_LOCAL_RANK, rank.to_string())
.env(ENV_NCCL_ID, &id_hex)
// Workers inherit stdout/stderr so rank 0's training log surfaces.
.stdout(Stdio::inherit())
.stderr(Stdio::inherit())
.spawn()
.map_err(|e| format!("spawn worker rank {rank}: {e}"))?;
children.push((rank, child));
}
let mut first_err: Option<String> = None;
for (rank, mut child) in children {
let status = child
.wait()
.map_err(|e| format!("wait worker rank {rank}: {e}"))?;
if !status.success() && first_err.is_none() {
first_err = Some(format!("worker rank {rank} exited with {status}"));
}
}
match first_err {
Some(e) => Err(e),
None => Ok(()),
}
}
/// The four `XTRAIN_*` values a worker reads from its env. Present iff this
/// process was spawned by [`launch_processes`].
pub struct WorkerEnv {
pub rank: usize,
pub world: usize,
pub local_rank: u32,
pub id: NcclUniqueId,
}
/// Read the worker env if this process is a spawned worker (i.e. `XTRAIN_RANK`
/// is set), else `None` (this process is the launcher).
pub fn worker_env() -> Option<WorkerEnv> {
let rank: usize = std::env::var(ENV_RANK).ok()?.parse().ok()?;
let world: usize = std::env::var(ENV_WORLD)
.expect("XTRAIN_WORLD set with XTRAIN_RANK")
.parse()
.expect("XTRAIN_WORLD parse");
let local_rank: u32 = std::env::var(ENV_LOCAL_RANK)
.expect("XTRAIN_LOCAL_RANK set with XTRAIN_RANK")
.parse()
.expect("XTRAIN_LOCAL_RANK parse");
let id_hex = std::env::var(ENV_NCCL_ID).expect("XTRAIN_NCCL_ID set with XTRAIN_RANK");
let id = hex_decode_unique_id(&id_hex);
Some(WorkerEnv {
rank,
world,
local_rank,
id,
})
}
/// Per-worker model construction knobs (the opt-in feature flags the launcher
/// forwards). Mirrors the closure `train_ddp` passes to the thread-per-GPU
/// `launch`, but here it runs once in this worker's own process/context.
#[derive(Clone, Copy, Default)]
pub struct ModelOpts {
pub bf16: bool,
pub recompute: bool,
pub flash: bool,
}
/// Run this worker: bind its GPU (→ its own CUDA context), init NCCL with the
/// launcher-supplied id, build its model with the deterministic init (same as
/// every rank + the single-GPU baseline), and run `train_rank`. Reuses the T8
/// training step verbatim — the only difference from thread-per-GPU is how this
/// rank was started and how it got the `UniqueId`.
///
/// `valid` is the held-out corpus for rank 0's periodic eval (pass `None` on
/// other ranks or when `cfg.eval_every == 0`).
pub fn run_worker(
env: &WorkerEnv,
cfg: Config,
opts: ModelOpts,
corpus: &Corpus,
valid: Option<&Corpus>,
dcfg: &DdpConfig,
) -> DdpResult {
// Binding the device here establishes this process's own CUDA primary context.
let ctx = DdpContext::init(env.rank, env.world, env.id, env.local_rank);
let device = Device::Cuda(env.local_rank);
let model = build_worker_model(cfg, opts, device);
let v = if env.rank == 0 { valid } else { None };
train_rank(&ctx, &model, device, corpus, v, dcfg)
}
/// Build the worker's model with the deterministic `build_model` init + the
/// opt-in feature flags. Shared by `run_worker` and the test worker.
pub fn build_worker_model(cfg: Config, opts: ModelOpts, device: Device) -> TinyTransformer {
let mut m = build_model(cfg, device);
if opts.bf16 {
m = m.with_compute_dtype(DType::BF16);
}
if opts.recompute {
m = m.with_recompute(true);
}
if opts.flash {
m = m.with_flash(true);
}
m
}
/// Convenience: the directory tests/bins can stash per-rank result dumps in
/// (a worker writes its loss/params there; the launching test reads them back).
pub fn rank_dump_path(dir: &std::path::Path, rank: usize) -> PathBuf {
dir.join(format!("rank{rank}.dump"))
}

614
crates/xtrain-distributed/tests/ddp_correctness.rs Normal file
View File

@@ -0,0 +1,614 @@
//! DDP acceptance (Phase T8). Gated to a GPU host; skips when fewer than 2 GPUs.
//!
//! 1. **Correctness**: K steps single-GPU (world=1, global batch B) vs 2-rank DDP
//! (B/2 of the SAME data in the same order each) → loss trajectories match
//! within tight fp tolerance (it's just gradient averaging), and the two
//! ranks' parameters are identical after the run.
//! 2. **Throughput**: 1 / 2 / 4 GPU global tok/s on the SAME per-GPU workload →
//! near-linear scaling. Prints the table (run with `--nocapture`).
#![cfg(not(no_cuda))]
use std::time::Instant;
use xtrain_cuda::device;
use xtrain_distributed::{DdpConfig, DdpContext, build_model, get_unique_id, launch, train_rank};
use xtrain_model::{Config, batched_ids_tensor};
use xtrain_optim::GpuAdamW;
use xtrain_tensor::Device;
use xtrain_train::clip::clip_grad_norm_gpu;
use xtrain_train::data::Corpus;
use xtrain_train::schedule::LrSchedule;
// A self-contained synthetic corpus so the test needs no tokenizer/data files.
fn synth_corpus(vocab: usize, n_tokens: usize) -> Corpus {
let tokens: Vec<i32> = (0..n_tokens)
.map(|i| (i * 7 + 3) as i32 % vocab as i32)
.collect();
Corpus {
tokens,
labels: None,
vocab_size: vocab,
}
}
fn test_config(vocab: usize) -> Config {
let mut cfg = Config::tiny();
cfg.vocab = vocab;
cfg.n_layers = 2;
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],
cfg: Config,
corpus: &Corpus,
valid: Option<&Corpus>,
dcfg: &DdpConfig,
) -> (Vec<f32>, Vec<Vec<f32>>, bool) {
let world = devices.len();
let id = get_unique_id();
let results: Vec<(Vec<f32>, Vec<Vec<f32>>, bool)> = std::thread::scope(|s| {
let handles: Vec<_> = devices
.iter()
.enumerate()
.map(|(rank, &dev)| {
let dcfg = dcfg.clone();
let corpus = &corpus;
s.spawn(move || {
let ctx = DdpContext::init(rank, world, id, dev);
let device = Device::Cuda(dev);
let model = build_model(cfg, device);
// Only rank 0 holds the val corpus (mirrors launch()).
let v = if rank == 0 { valid } else { None };
let res = train_rank(&ctx, &model, device, corpus, v, &dcfg);
let host = model
.params()
.iter()
.map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
.collect::<Vec<_>>();
(res.losses, host, model.is_training())
})
})
.collect();
handles.into_iter().map(|h| h.join().unwrap()).collect()
});
results.into_iter().next().unwrap()
}
// Single-GPU baseline: the SAME loop as the DDP rank but world=1, so the global
// batch is processed on one device. Returns (loss trace, final params on host).
fn run_single_gpu(cfg: Config, corpus: &Corpus, dcfg: &DdpConfig) -> (Vec<f32>, Vec<Vec<f32>>) {
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let model = build_model(cfg, device);
let params = model.params();
let mut opt = GpuAdamW::new(dcfg.weight_decay);
let mut rng = dcfg.seed;
let mut losses = Vec::new();
for step in 0..dcfg.steps {
let lr = dcfg.schedule.lr(step);
// Sample the whole global batch and run it as ONE batched forward/backward
// (matches the T10 DDP path: backward yields the global-batch mean grad).
let mut inputs = Vec::with_capacity(dcfg.batch_size);
let mut targets_v = Vec::with_capacity(dcfg.batch_size);
for _ in 0..dcfg.batch_size {
let (input, target) = corpus.sample(dcfg.seq_len, &mut rng);
inputs.push(input);
targets_v.push(target);
}
let ids = batched_ids_tensor(&inputs, device);
let targets = batched_ids_tensor(&targets_v, device);
let loss = model.loss_batched(&ids, &targets, dcfg.batch_size);
losses.push(loss.value().to_device(Device::Cpu).as_slice::<f32>()[0]);
loss.backward();
clip_grad_norm_gpu(&params, dcfg.max_grad_norm, 1.0);
opt.step(lr, &params);
for p in &params {
p.zero_grad();
}
}
let host = params
.iter()
.map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
.collect();
(losses, host)
}
#[test]
fn ddp_matches_single_gpu_and_params_consistent() {
let world = 2usize;
if device::device_count().unwrap_or(0) < world as i32 {
eprintln!("skip: need >= {world} GPUs");
return;
}
let vocab = 64usize;
let cfg = test_config(vocab);
let corpus = synth_corpus(vocab, 4096);
let steps = 20usize;
let dcfg = DdpConfig {
seq_len: 32,
batch_size: 8, // global; 4 per rank with world=2
accum_steps: 1,
steps,
schedule: LrSchedule {
max_lr: 3e-3,
min_lr: 3e-4,
warmup: 3,
total: steps,
},
weight_decay: 0.1,
max_grad_norm: 1.0,
log_every: 1_000_000, // silence per-step logging in the test
seed: 7,
eval_every: 0,
eval_batches: 0,
ckpt_path: None,
};
// Single-GPU baseline (world=1) over the global batch.
let (single_losses, single_params) = run_single_gpu(cfg, &corpus, &dcfg);
// 2-rank DDP over the SAME corpus/config; returns per-rank (losses, params).
let devices = [0u32, 1u32];
let id = get_unique_id();
let results: Vec<(Vec<f32>, Vec<Vec<f32>>)> = std::thread::scope(|s| {
let handles: Vec<_> = devices
.iter()
.enumerate()
.map(|(rank, &dev)| {
let dcfg = dcfg.clone();
let corpus = &corpus;
s.spawn(move || {
let ctx = DdpContext::init(rank, world, id, dev);
let device = Device::Cuda(dev);
let model = build_model(cfg, device);
let res = train_rank(&ctx, &model, device, corpus, None, &dcfg);
let host = model
.params()
.iter()
.map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
.collect::<Vec<_>>();
(res.losses, host)
})
})
.collect();
handles.into_iter().map(|h| h.join().unwrap()).collect()
});
let (ddp_losses, ddp_p0) = &results[0];
let (_, ddp_p1) = &results[1];
// (a) DDP loss trajectory matches single-GPU within tight tolerance.
let mut max_rel = 0.0f32;
for (s, d) in single_losses.iter().zip(ddp_losses) {
let rel = (s - d).abs() / s.abs().max(1e-6);
max_rel = max_rel.max(rel);
}
println!(
"DDP vs single-GPU loss: single[last]={:.6} ddp[last]={:.6} max_rel={max_rel:.2e}",
single_losses.last().unwrap(),
ddp_losses.last().unwrap()
);
assert!(
max_rel < 1e-3,
"DDP loss trajectory diverged from single-GPU: max_rel {max_rel:.3e}"
);
// (b) Cross-rank parameter identity (same init + same averaged grad + same
// optimizer state ⇒ identical params).
let mut max_pdiff = 0.0f32;
for (a, b) in ddp_p0.iter().zip(ddp_p1) {
for (x, y) in a.iter().zip(b) {
max_pdiff = max_pdiff.max((x - y).abs());
}
}
println!("cross-rank max |param diff| = {max_pdiff:.3e}");
// On this PCIe-only box, NCCL's all-reduce is not bit-reproducible run-to-run
// across ranks (algorithm/chunk choice is unstable), so cross-rank params can
// differ by a few ULP (observed ≤1.2e-7) even with identical init + averaged
// grads. The load-bearing gate is the loss-trajectory match (a, ~5.7e-7); a
// tight tolerance here, not bit-identity, is the honest invariant (KI-5).
assert!(
max_pdiff < 1e-6,
"ranks' params drifted apart: {max_pdiff:.3e}"
);
// (c) DDP final params match single-GPU final params within fp tolerance.
// Looser than (a)/(b): DDP and single-GPU differ only in the gradient SUMMATION
// ORDER (single-GPU sums B sequences in tape order; DDP sums per-rank shards
// then NCCL-sums across ranks). fp addition isn't associative, so that tiny
// per-step rounding compounds over the AdamW steps — a few e-3 relative on
// individual params is expected and benign. The loss-trajectory match (a, ~1e-7)
// and tight cross-rank agreement (b, <1e-6) are the load-bearing checks.
let mut max_sdiff = 0.0f32;
for (a, b) in ddp_p0.iter().zip(&single_params) {
for (x, y) in a.iter().zip(b) {
max_sdiff = max_sdiff.max((x - y).abs() / y.abs().max(1e-6));
}
}
println!("DDP vs single-GPU max rel |param diff| = {max_sdiff:.3e}");
assert!(max_sdiff < 1e-2, "DDP params diverged from single-GPU");
}
#[test]
fn ddp_with_accum_matches_single_gpu_big_batch() {
// T16: DDP + gradient accumulation must match a single-GPU big-batch baseline
// of the SAME effective batch. world=2, accum=2, per-rank micro-batch 2 →
// effective global batch = world·accum·b_local = 2·2·2 = 8. Compared against a
// single-GPU run with batch 8, accum 1 (the big-batch baseline). The all-reduce
// fires only at the accumulation boundary (once per optimizer step, not per
// micro-step) — enforced by the train_rank implementation; the load-bearing
// gate here is that loss + final params still match the big-batch baseline.
let world = 2usize;
if device::device_count().unwrap_or(0) < world as i32 {
eprintln!("skip: need >= {world} GPUs");
return;
}
let vocab = 64usize;
let cfg = test_config(vocab);
let corpus = synth_corpus(vocab, 4096);
let steps = 20usize;
let effective_batch = 8usize; // world(2) · accum(2) · b_local(2)
let sched = LrSchedule {
max_lr: 3e-3,
min_lr: 3e-4,
warmup: 3,
total: steps,
};
// Single-GPU big-batch baseline: world=1, accum=1, batch = effective_batch.
let baseline_cfg = DdpConfig {
seq_len: 32,
batch_size: effective_batch,
accum_steps: 1,
steps,
schedule: sched,
weight_decay: 0.1,
max_grad_norm: 1.0,
log_every: 1_000_000,
seed: 7,
eval_every: 0,
eval_batches: 0,
ckpt_path: None,
};
let (single_losses, single_params) = run_single_gpu(cfg, &corpus, &baseline_cfg);
// DDP + accumulation: world=2, accum=2 → per-rank micro-batch = batch/world = 2.
let ddp_cfg = DdpConfig {
batch_size: effective_batch / 2, // per-step global batch; ×accum = effective
accum_steps: 2,
..baseline_cfg
};
let devices = [0u32, 1u32];
let id = get_unique_id();
let results: Vec<(Vec<f32>, Vec<Vec<f32>>)> = std::thread::scope(|s| {
let handles: Vec<_> = devices
.iter()
.enumerate()
.map(|(rank, &dev)| {
let ddp_cfg = ddp_cfg.clone();
let corpus = &corpus;
s.spawn(move || {
let ctx = DdpContext::init(rank, world, id, dev);
let device = Device::Cuda(dev);
let model = build_model(cfg, device);
let res = train_rank(&ctx, &model, device, corpus, None, &ddp_cfg);
let host = model
.params()
.iter()
.map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
.collect::<Vec<_>>();
(res.losses, host)
})
})
.collect();
handles.into_iter().map(|h| h.join().unwrap()).collect()
});
let (ddp_losses, ddp_p0) = &results[0];
let (_, ddp_p1) = &results[1];
// (a) Loss trajectory matches the single-GPU big-batch baseline.
let mut max_rel = 0.0f32;
for (s, d) in single_losses.iter().zip(ddp_losses) {
max_rel = max_rel.max((s - d).abs() / s.abs().max(1e-6));
}
println!(
"DDP+accum(w2·a2·b2) vs single-GPU big-batch(8): single[last]={:.6} ddp[last]={:.6} max_rel={max_rel:.2e}",
single_losses.last().unwrap(),
ddp_losses.last().unwrap()
);
assert!(
max_rel < 1e-3,
"DDP+accum loss diverged from big-batch baseline: {max_rel:.3e}"
);
// (b) Cross-rank parameter agreement (same KI-5 ULP tolerance as the base test).
let mut max_pdiff = 0.0f32;
for (a, b) in ddp_p0.iter().zip(ddp_p1) {
for (x, y) in a.iter().zip(b) {
max_pdiff = max_pdiff.max((x - y).abs());
}
}
println!("DDP+accum cross-rank max |param diff| = {max_pdiff:.3e}");
assert!(
max_pdiff < 1e-6,
"ranks' params drifted apart: {max_pdiff:.3e}"
);
// (c) Final params match single-GPU big-batch within fp tolerance.
let mut max_sdiff = 0.0f32;
for (a, b) in ddp_p0.iter().zip(&single_params) {
for (x, y) in a.iter().zip(b) {
max_sdiff = max_sdiff.max((x - y).abs() / y.abs().max(1e-6));
}
}
println!("DDP+accum vs single-GPU big-batch max rel |param diff| = {max_sdiff:.3e}");
assert!(
max_sdiff < 1e-2,
"DDP+accum params diverged from big-batch baseline"
);
}
#[test]
fn ddp_throughput_scaling() {
let max_gpus = device::device_count().unwrap_or(0) as usize;
if max_gpus < 1 {
eprintln!("skip: no GPU");
return;
}
// Same PER-GPU workload at each world size (batch scales with world), so the
// per-rank cost is fixed and global tok/s should scale ~linearly. Use enough
// steps that the one-time NCCL init + model-build overhead (which is larger at
// world=4 and absent at world=1) amortizes — otherwise the wall-clock ratio
// understates steady-state scaling.
let per_gpu_batch = 8usize;
let vocab = 256usize;
let cfg = test_config(vocab);
let corpus = synth_corpus(vocab, 8192);
let steps = 150usize;
let seq_len = 64usize;
let worlds: Vec<usize> = [1, 2, 4, 8]
.into_iter()
.filter(|&w| w <= max_gpus)
.collect();
println!("\n=== DDP throughput scaling (per-GPU batch {per_gpu_batch}, seq {seq_len}) ===");
println!(
"{:>6} | {:>14} | {:>8}",
"GPUs", "tok/s (global)", "speedup"
);
let mut base = 0.0f64;
for &world in &worlds {
let devices: Vec<u32> = (0..world as u32).collect();
let dcfg = DdpConfig {
seq_len,
batch_size: per_gpu_batch * world,
accum_steps: 1,
steps,
schedule: LrSchedule {
max_lr: 1e-3,
min_lr: 1e-3,
warmup: 1,
total: steps,
},
weight_decay: 0.0,
max_grad_norm: 1.0,
log_every: 1_000_000,
seed: 1,
eval_every: 0,
eval_batches: 0,
ckpt_path: None,
};
let total_tokens = (steps * dcfg.batch_size * seq_len) as f64;
let t = Instant::now();
let _ = launch(&devices, &corpus, None, &dcfg, move |device| {
build_model(cfg, device)
});
let secs = t.elapsed().as_secs_f64();
let tps = total_tokens / secs;
if world == 1 {
base = tps;
}
println!(
"{:>6} | {:>14.0} | {:>7.2}x",
world,
tps,
tps / base.max(1e-9)
);
}
}
/// 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`).
///
/// 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.
#[test]
fn ddp_dropout_is_live_and_p0_bit_identical() {
if device::device_count().unwrap_or(0) < 2 {
eprintln!("skip: need >= 2 GPUs");
return;
}
let vocab = 64usize;
let corpus = synth_corpus(vocab, 4096);
let steps = 20usize;
// eval_every < steps so a periodic eval fires MID-run (flipping the model to
// eval mode via eval_loss → model.eval()). The per-step model.train() must
// restore training mode so dropout stays live across the eval boundary — this is
// exactly the train/eval discipline the pilot called out. A held-out slice gives
// rank 0 something to eval on.
let valid = synth_corpus(vocab, 512);
let base_dcfg = DdpConfig {
seq_len: 32,
batch_size: 8, // global; 4 per rank with world=2
accum_steps: 1,
steps,
schedule: LrSchedule {
max_lr: 3e-3,
min_lr: 3e-4,
warmup: 3,
total: steps,
},
weight_decay: 0.1,
max_grad_norm: 1.0,
log_every: 1_000_000, // silence per-step logging
seed: 7,
eval_every: 7, // fires at steps 6, 13, 19 — flips to eval mode mid-run
eval_batches: 4,
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);
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)"
);
assert_eq!(
max_loss_diff_1, 0.0,
"world=1 p=0 forward loss 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}"
);
// 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);
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()
);
assert!(
max_live_diff > 1e-3,
"p=0.2 DDP loss trace matches p=0 — dropout is NOT live under DDP \
(model.train() not wired): max |loss diff| {max_live_diff:.3e}"
);
// No NaN/Inf in the p>0 run (dropout converges normally under DDP).
assert!(
loss_p_2.iter().all(|l| l.is_finite()),
"p=0.2 DDP loss has non-finite values"
);
// 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);
assert!(
train_flag,
"model not in training mode after a no-eval DDP run — model.train() not wired in train_rank"
);
println!("T21 GATE C (train_rank sets training mode): is_training() == true ✅");
}

409
crates/xtrain-distributed/tests/ddp_proc.rs Normal file
View File

@@ -0,0 +1,409 @@
//! Process-per-GPU DDP acceptance (Phase T17). Gated to a GPU host; skips when
//! fewer than 2 GPUs. Run with `--test-threads=1` (distributed tests deadlock if
//! they contend for the same GPUs in parallel — known harness property).
//!
//! Self-launching: the test binary detects WORKER mode via `XTRAIN_RANK` (set by
//! `launch_processes`). In worker mode it runs `run_worker` on a synthetic corpus
//! and dumps its per-step loss trace + final params to a per-rank file; in normal
//! mode it is the launcher — it runs the single-GPU baseline, spawns N worker
//! processes (re-execing itself), reads their dumps back, and asserts:
//! (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))]
use std::io::Write;
use std::path::Path;
use xtrain_cuda::device;
use xtrain_distributed::proc::{launch_processes, rank_dump_path, worker_env};
use xtrain_distributed::{DdpConfig, DdpContext, build_model, train_rank};
use xtrain_model::{Config, batched_ids_tensor};
use xtrain_optim::GpuAdamW;
use xtrain_tensor::Device;
use xtrain_train::clip::clip_grad_norm_gpu;
use xtrain_train::data::Corpus;
use xtrain_train::schedule::LrSchedule;
// ── Shared fixture (identical on launcher + every worker, so they agree) ──────
const VOCAB: usize = 64;
const STEPS: usize = 20;
fn synth_corpus() -> Corpus {
let tokens: Vec<i32> = (0..4096)
.map(|i| (i * 7 + 3) as i32 % VOCAB as i32)
.collect();
Corpus {
tokens,
labels: None,
vocab_size: VOCAB,
}
}
fn test_config() -> Config {
let mut cfg = Config::tiny();
cfg.vocab = VOCAB;
cfg.n_layers = 2;
cfg
}
fn dcfg(batch_size: usize) -> DdpConfig {
DdpConfig {
seq_len: 32,
batch_size,
accum_steps: 1,
steps: STEPS,
schedule: LrSchedule {
max_lr: 3e-3,
min_lr: 3e-4,
warmup: 3,
total: STEPS,
},
weight_decay: 0.1,
max_grad_norm: 1.0,
log_every: 1_000_000,
seed: 7,
eval_every: 0,
eval_batches: 0,
ckpt_path: None,
}
}
// 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() {
let Some(env) = worker_env() else { return };
let dump_dir = std::env::var(ENV_DUMP_DIR).expect("dump dir env");
// This is the worker body `run_worker` performs in production (init ctx →
// build deterministic model → train_rank). We train ONCE inline so we can dump
// both the loss trace AND the final params for the launcher to check; the
// 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 res = train_rank(
&ctx,
&model,
device,
&synth_corpus(),
None,
&dcfg(GLOBAL_BATCH),
);
let params: Vec<Vec<f32>> = model
.params()
.iter()
.map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
.collect();
write_dump(&dump_dir, env.rank, &res.losses, &params);
std::process::exit(0);
}
fn write_dump(dir: &str, rank: usize, losses: &[f32], params: &[Vec<f32>]) {
let path = rank_dump_path(Path::new(dir), rank);
let mut f = std::fs::File::create(&path).expect("create dump");
// Line 1: losses (space-separated). Following lines: one param tensor each.
let loss_line: Vec<String> = losses.iter().map(|x| format!("{x:.8e}")).collect();
writeln!(f, "{}", loss_line.join(" ")).unwrap();
for p in params {
let line: Vec<String> = p.iter().map(|x| format!("{x:.8e}")).collect();
writeln!(f, "{}", line.join(" ")).unwrap();
}
}
fn read_dump(dir: &str, rank: usize) -> (Vec<f32>, Vec<Vec<f32>>) {
let path = rank_dump_path(Path::new(dir), rank);
let text = std::fs::read_to_string(&path).expect("read dump");
let mut lines = text.lines();
let losses: Vec<f32> = lines
.next()
.unwrap()
.split_whitespace()
.map(|s| s.parse().unwrap())
.collect();
let params: Vec<Vec<f32>> = lines
.map(|l| l.split_whitespace().map(|s| s.parse().unwrap()).collect())
.collect();
(losses, params)
}
// ── Single-GPU baseline (same loop as the DDP rank, world=1) ──────────────────
fn run_single_gpu(cfg: Config, corpus: &Corpus, d: &DdpConfig) -> (Vec<f32>, Vec<Vec<f32>>) {
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let model = build_model(cfg, device);
let params = model.params();
let mut opt = GpuAdamW::new(d.weight_decay);
let mut rng = d.seed;
let mut losses = Vec::new();
for step in 0..d.steps {
let lr = d.schedule.lr(step);
let mut inputs = Vec::with_capacity(d.batch_size);
let mut targets_v = Vec::with_capacity(d.batch_size);
for _ in 0..d.batch_size {
let (input, target) = corpus.sample(d.seq_len, &mut rng);
inputs.push(input);
targets_v.push(target);
}
let ids = batched_ids_tensor(&inputs, device);
let targets = batched_ids_tensor(&targets_v, device);
let loss = model.loss_batched(&ids, &targets, d.batch_size);
losses.push(loss.value().to_device(Device::Cpu).as_slice::<f32>()[0]);
loss.backward();
clip_grad_norm_gpu(&params, d.max_grad_norm, 1.0);
opt.step(lr, &params);
for p in &params {
p.zero_grad();
}
}
let host = params
.iter()
.map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
.collect();
(losses, host)
}
// ── The test (launcher mode) ──────────────────────────────────────────────────
#[test]
fn proc_per_gpu_matches_single_gpu_and_thread_path() {
// If this process was spawned as a worker, do the worker job and exit before
// the test framework runs anything else.
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 cfg = test_config();
let corpus = synth_corpus();
let d = dcfg(GLOBAL_BATCH);
// (1) Single-GPU baseline over the global batch.
let (single_losses, single_params) = run_single_gpu(cfg, &corpus, &d);
// (2) Thread-per-GPU path (T8 `launch`) — the regression baseline to match.
let thread_results =
xtrain_distributed::launch(&[0u32, 1u32], &corpus, None, &d, move |device| {
build_model(cfg, device)
});
let thread_losses = &thread_results[0].losses;
// (3) Process-per-GPU: spawn 2 worker processes (re-exec of this test binary),
// each dumps its loss trace + final params to a temp dir.
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.
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
// worker process does the worker job and exits without touching other tests.
let worker_args = [
"--exact".to_string(),
"proc_per_gpu_matches_single_gpu_and_thread_path".to_string(),
"--test-threads=1".to_string(),
"--nocapture".to_string(),
];
launch_processes(world, &worker_args).expect("worker processes failed");
let (proc_losses0, proc_p0) = read_dump(dump_dir.to_str().unwrap(), 0);
let (_proc_losses1, proc_p1) = read_dump(dump_dir.to_str().unwrap(), 1);
// (a) process-per-GPU loss matches single-GPU.
let max_rel_single = max_rel(&single_losses, &proc_losses0);
println!(
"proc-per-GPU vs single-GPU loss: single[last]={:.6} proc[last]={:.6} max_rel={max_rel_single:.2e}",
single_losses.last().unwrap(),
proc_losses0.last().unwrap()
);
assert!(
max_rel_single < 1e-3,
"proc-per-GPU loss diverged from single-GPU: {max_rel_single:.3e}"
);
// (c) process-per-GPU loss matches the thread-per-GPU path.
let max_rel_thread = max_rel(thread_losses, &proc_losses0);
println!(
"proc-per-GPU vs thread-per-GPU loss: thread[last]={:.6} proc[last]={:.6} max_rel={max_rel_thread:.2e}",
thread_losses.last().unwrap(),
proc_losses0.last().unwrap()
);
assert!(
max_rel_thread < 1e-3,
"proc-per-GPU loss diverged from thread-per-GPU: {max_rel_thread:.3e}"
);
// (b) cross-rank parameter agreement (KI-5 ULP tolerance).
let mut max_pdiff = 0.0f32;
for (a, b) in proc_p0.iter().zip(&proc_p1) {
for (x, y) in a.iter().zip(b) {
max_pdiff = max_pdiff.max((x - y).abs());
}
}
println!("proc-per-GPU cross-rank max |param diff| = {max_pdiff:.3e}");
assert!(
max_pdiff < 1e-6,
"ranks' params drifted apart: {max_pdiff:.3e}"
);
// Bonus sanity: proc-per-GPU final params vs single-GPU within fp tolerance.
let mut max_sdiff = 0.0f32;
for (a, b) in proc_p0.iter().zip(&single_params) {
for (x, y) in a.iter().zip(b) {
max_sdiff = max_sdiff.max((x - y).abs() / y.abs().max(1e-6));
}
}
println!("proc-per-GPU vs single-GPU max rel |param diff| = {max_sdiff:.3e}");
assert!(
max_sdiff < 1e-2,
"proc-per-GPU params diverged from single-GPU"
);
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)
.map(|(s, d)| (s - d).abs() / s.abs().max(1e-6))
.fold(0.0f32, f32::max)
}

12
crates/xtrain-model/Cargo.toml Normal file
View File

@@ -0,0 +1,12 @@
[package]
name = "xtrain-model"
version.workspace = true
edition.workspace = true
[dependencies]
xtrain-tensor = { path = "../xtrain-tensor" }
xtrain-autodiff = { path = "../xtrain-autodiff" }
[dev-dependencies]
# Acceptance tests drive the GPU (device selection) directly.
xtrain-cuda = { path = "../xtrain-cuda" }

26
crates/xtrain-model/build.rs Normal file
View File

@@ -0,0 +1,26 @@
use std::env;
use std::path::Path;
use std::process::Command;
// Same per-crate convention as the other crates: this crate's tiny-transformer
// forward/backward calls GPU ops (via xtrain-autodiff / xtrain-tensor), so it
// gates GPU code + tests behind `not(no_cuda)`. cfg does not propagate across
// crates, so each crate re-detects nvcc. No CUDA is compiled here.
fn main() {
println!("cargo:rustc-check-cfg=cfg(no_cuda)");
let cuda_path = env::var("CUDA_HOME")
.or_else(|_| env::var("CUDA_PATH"))
.unwrap_or_else(|_| "/usr/local/cuda".to_string());
if !nvcc_available(&cuda_path) {
println!("cargo:rustc-cfg=no_cuda");
}
}
fn nvcc_available(cuda_path: &str) -> bool {
if Command::new("nvcc").arg("--version").output().is_ok() {
return true;
}
Path::new(&format!("{cuda_path}/bin/nvcc")).exists()
}

125
crates/xtrain-model/src/config.rs Normal file
View File

@@ -0,0 +1,125 @@
//! Tiny-transformer hyperparameters. Host-only (no GPU), always compiled.
/// Architecture config for [`crate::TinyTransformer`]. Keep it tiny — T5 is a
/// correctness bring-up, not a real training run.
#[derive(Debug, Clone, Copy)]
pub struct Config {
/// Vocabulary size (char-level in the bring-up).
pub vocab: usize,
/// Model / residual width. Must equal `n_heads * head_dim`.
pub dim: usize,
/// Number of decoder blocks.
pub n_layers: usize,
/// Number of attention (query) heads.
pub n_heads: usize,
/// Number of key/value heads (Phase T15, GQA). Each KV head is shared by a
/// group of `n_heads / num_kv_heads` query heads (repeat_kv). Must divide
/// `n_heads`. `num_kv_heads == n_heads` (the default) = MHA, bit-identical to
/// the pre-T15 path; `num_kv_heads < n_heads` = real grouped-query attention,
/// shrinking the K/V projections to `num_kv_heads * head_dim` and exported as a
/// real `num_key_value_heads`.
pub num_kv_heads: usize,
/// Per-head dimension (`dim / n_heads`).
pub head_dim: usize,
/// SwiGLU hidden width (gate/up project to this, down projects back).
pub ffn_hidden: usize,
/// RMSNorm epsilon.
pub eps: f32,
/// RoPE base frequency (theta).
pub rope_theta: f32,
/// Dropout probability `p` (Phase T18). Applied at the attention/MLP sub-block
/// outputs (before each residual add) at TRAINING time, with inverted scaling
/// `1/(1-p)`; disabled (identity) at eval. Default `0.0` = no dropout, and the
/// forward graph is then bit-identical to the pre-T18 path.
pub dropout: f32,
}
impl Config {
/// A minimal config used by the bring-up / overfit test.
pub fn tiny() -> Self {
let n_heads = 2;
let head_dim = 16;
Config {
vocab: 0, // set by the caller from the char vocab
dim: n_heads * head_dim,
n_layers: 2,
n_heads,
num_kv_heads: n_heads, // default = MHA
head_dim,
ffn_hidden: 64,
eps: 1e-5,
rope_theta: 10000.0,
dropout: 0.0,
}
}
/// Build a config from the architecture knobs, deriving `dim = n_heads *
/// head_dim`. The scaling-run entry (`bin/train`) passes these from CLI so the
/// model size is a tunable ladder rung (v1 = dim256/8L, v2/v3 scale further),
/// instead of a hardcoded tiny config. `eps`/`rope_theta` keep the engine
/// defaults (also what the xserv export reconciles against).
pub fn from_arch(
vocab: usize,
n_heads: usize,
head_dim: usize,
n_layers: usize,
ffn_hidden: usize,
) -> Self {
Config {
vocab,
dim: n_heads * head_dim,
n_layers,
n_heads,
num_kv_heads: n_heads, // default = MHA; set via with_kv_heads for GQA
head_dim,
ffn_hidden,
eps: 1e-5,
rope_theta: 10000.0,
dropout: 0.0,
}
}
/// Set the number of K/V heads (Phase T15, GQA). Builder-style so existing
/// `from_arch` call sites stay MHA unless they opt in. Asserts `num_kv_heads`
/// divides `n_heads`.
pub fn with_kv_heads(mut self, num_kv_heads: usize) -> Self {
assert!(num_kv_heads > 0, "num_kv_heads must be > 0");
assert_eq!(
self.n_heads % num_kv_heads,
0,
"n_heads {} not divisible by num_kv_heads {num_kv_heads}",
self.n_heads
);
self.num_kv_heads = num_kv_heads;
self
}
/// KV projection width (`num_kv_heads * head_dim`). For GQA this is smaller than
/// `dim`; for MHA it equals `dim`.
pub fn kv_dim(&self) -> usize {
self.num_kv_heads * self.head_dim
}
/// Transformer-core parameter count: everything except the token embedding and
/// the LM head (the two `vocab × dim` tables). This is the figure the scaling
/// ladder is sized against — the 50257-vocab embed+lm_head adds a fixed ~25M on
/// top that does not reflect model capacity. `num_params() = core + 2·vocab·dim`.
pub fn core_params(&self) -> usize {
self.num_params() - 2 * self.vocab * self.dim
}
/// Total learnable parameter count (for logging / sanity).
pub fn num_params(&self) -> usize {
let per_layer = 2 * self.dim // 2 rmsnorm gammas
+ 2 * self.head_dim // q/k per-head norm gammas
+ self.dim * self.dim // q proj [dim,dim]
+ 2 * self.dim * self.kv_dim() // k/v proj [dim,kv_dim] (GQA: smaller)
+ self.dim * self.dim // out proj
+ 2 * self.dim * self.ffn_hidden // gate/up proj
+ self.ffn_hidden * self.dim; // down proj
self.vocab * self.dim // embedding
+ self.n_layers * per_layer
+ self.dim // final norm
+ self.dim * self.vocab // lm head
}
}

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

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

32
crates/xtrain-model/src/lib.rs Normal file
View File

@@ -0,0 +1,32 @@
//! Tiny modern-architecture transformer (Phase T5).
//!
//! A from-scratch decoder built entirely from the [`xtrain_autodiff`] op set:
//! token embedding → `n_layers` × {pre-RMSNorm → multi-head causal attention
//! (per-head QK-norm + RoPE) → residual; pre-RMSNorm → SwiGLU MLP → residual} →
//! final RMSNorm → LM-head matmul. The forward builds an autograd graph; calling
//! `.backward()` on the cross-entropy loss fills every parameter's `.grad()`.
//! Per-head QK-norm (Qwen3-style) makes the architecture xserv-compatible (T9).
//!
//! Conventions (matching the engine, not HuggingFace):
//! - Linear weights are `[in, out]` and applied as `x @ W` (no transpose), since
//! the engine's GEMM is plain `A @ B`.
//! - `dim == n_heads * head_dim` (no separate attention projection size).
//! - RoPE position = token row index (the kernel's built-in convention).
//! - Causal masking is an additive `[seq,seq]` constant (−1e9 above the diagonal)
//! added to the attention scores before softmax.
//!
//! Everything GPU-facing is gated behind `not(no_cuda)`; on a GPU-less host the
//! crate still `cargo check`s (only [`Config`] is visible there).
mod config;
pub use config::Config;
#[cfg(not(no_cuda))]
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};

529
crates/xtrain-model/src/model.rs Normal file
View File

@@ -0,0 +1,529 @@
//! The tiny transformer forward graph + parameter container (Phase T5).
#![cfg(not(no_cuda))]
use std::cell::Cell;
use crate::config::Config;
use xtrain_autodiff::ops;
use xtrain_autodiff::tape::Var;
use xtrain_tensor::{DType, Device, Tensor};
/// One decoder block's learnable tensors.
struct Block {
attn_norm: Var, // [dim]
wq: Var, // [dim, dim]
wk: Var, // [dim, kv_dim] — kv_dim = num_kv_heads·head_dim (GQA; = dim for MHA)
wv: Var, // [dim, kv_dim]
q_norm: Var, // [head_dim] — per-head QK-norm (Qwen3-style)
k_norm: Var, // [head_dim]
wo: Var, // [dim, dim]
ffn_norm: Var, // [dim]
w_gate: Var, // [dim, ffn_hidden]
w_up: Var, // [dim, ffn_hidden]
w_down: Var, // [ffn_hidden, dim]
}
/// A tiny RoPE+RMSNorm+SwiGLU decoder. Holds every parameter as a leaf [`Var`];
/// `forward` builds an autograd graph over them.
pub struct TinyTransformer {
cfg: Config,
embed: Var, // [vocab, dim]
blocks: Vec<Block>,
final_norm: Var, // [dim]
lm_head: Var, // [dim, vocab]
/// Compute dtype for the forward graph (Phase T12). `F32` (default) = the
/// original path, bit-identical to T10/T11. `BF16` = mixed precision: the
/// parameter leaves stay fp32 (master), but each linear's weight is cast to
/// bf16 on the fly and the activation stream flows bf16 (see
/// `docs/11-bf16-mixed-precision.md`). The cast op's backward upcasts the bf16
/// weight grad back to fp32, so AdamW/clip/DDP stay fp32 and unchanged.
compute_dtype: DType,
/// Activation recomputation / gradient checkpointing (Phase T13, KI-3). When
/// `true`, each transformer block's forward runs through
/// [`xtrain_autodiff::checkpoint`]: the block's internal activations are NOT
/// kept on the tape during forward (only the block input is), and the block
/// forward is re-run during backward to recover them. Trades ~one extra forward
/// per block for a large drop in peak activation memory → lets dim1024 batch32
/// fit. Default `false` = the unchanged path (every activation stored), so
/// existing numerics are bit-identical; recompute is mathematically exact, so
/// grads match the non-checkpointed path within fp tolerance.
recompute: bool,
/// Fused flash-attention (Phase T14). When `true`, the SDPA core runs through
/// the hand-written single fused kernel ([`ops::flash_attention`]): online
/// softmax over KV tiles, the `[bh,seq,seq]` score matrix NEVER materialized,
/// backward caches only the O(N) logsumexp. Default `false` = the composed T10
/// path (`cublasSgemmStridedBatched`×2 + causal-softmax kernel, O(N²) probs),
/// so the default graph is unchanged. Mathematically the same SDPA → grads/loss
/// match the composed path within fp/bf16 tolerance. Opt-in via `--flash`.
use_flash: bool,
/// Training mode for dropout (Phase T18). `true` → the attn/MLP sub-block
/// outputs pass through `ops::dropout` (with `cfg.dropout` and a per-step,
/// per-site seed); `false` (default) → dropout is identity (eval/sampling/
/// export). `Cell` so `train()`/`eval()` flip it through `&self` (the forward
/// takes `&self`). When `cfg.dropout == 0` this flag is irrelevant — the graph
/// is bit-identical to the no-dropout path either way.
training: Cell<bool>,
/// Per-step dropout RNG seed (Phase T18). Bumped once at the start of each
/// TRAINING forward so every step draws fresh masks; combined with the layer
/// index + a per-site constant to give each dropout site its own seed. The RNG
/// is counter-based, so re-running a checkpointed block's forward in backward
/// (T13) reproduces the same seed → the same mask (recompute stays exact).
step_seed: Cell<u64>,
}
impl TinyTransformer {
/// Build a model with parameters initialised from `init(shape) -> host data`.
/// The caller controls initialisation (deterministic for tests / PyTorch
/// parity). `init` receives the logical shape and returns row-major data.
pub fn new(cfg: Config, device: Device, mut init: impl FnMut(&[usize]) -> Vec<f32>) -> Self {
let leaf = |data: Vec<f32>, shape: &[usize]| -> Var {
Var::leaf(Tensor::from_slice(&data, shape).to_device(device))
};
let mut mk = |shape: &[usize]| -> Var {
let data = init(shape);
assert_eq!(data.len(), shape.iter().product::<usize>(), "init size");
leaf(data, shape)
};
let embed = mk(&[cfg.vocab, cfg.dim]);
let blocks = (0..cfg.n_layers)
.map(|_| Block {
attn_norm: mk(&[cfg.dim]),
wq: mk(&[cfg.dim, cfg.dim]),
// GQA (T15): K/V project to num_kv_heads·head_dim (= dim when MHA).
wk: mk(&[cfg.dim, cfg.kv_dim()]),
wv: mk(&[cfg.dim, cfg.kv_dim()]),
q_norm: mk(&[cfg.head_dim]),
k_norm: mk(&[cfg.head_dim]),
wo: mk(&[cfg.dim, cfg.dim]),
ffn_norm: mk(&[cfg.dim]),
w_gate: mk(&[cfg.dim, cfg.ffn_hidden]),
w_up: mk(&[cfg.dim, cfg.ffn_hidden]),
w_down: mk(&[cfg.ffn_hidden, cfg.dim]),
})
.collect();
let final_norm = mk(&[cfg.dim]);
let lm_head = mk(&[cfg.dim, cfg.vocab]);
Self {
cfg,
embed,
blocks,
final_norm,
lm_head,
compute_dtype: DType::F32,
recompute: false,
use_flash: false,
training: Cell::new(false),
step_seed: Cell::new(0),
}
}
pub fn config(&self) -> &Config {
&self.cfg
}
/// Set the forward compute dtype (Phase T12). `BF16` enables mixed precision
/// (fp32 master weights, bf16 linears + activations); `F32` (the default) is
/// the unchanged full-precision path. Builder-style so existing call sites
/// that don't opt in keep the fp32 numerics bit-for-bit.
pub fn with_compute_dtype(mut self, dtype: DType) -> Self {
assert!(
matches!(dtype, DType::F32 | DType::BF16),
"compute_dtype must be F32 or BF16"
);
self.compute_dtype = dtype;
self
}
pub fn compute_dtype(&self) -> DType {
self.compute_dtype
}
/// Enable per-block activation recomputation / gradient checkpointing (Phase
/// T13). Builder-style and opt-in; default off keeps the unchanged tape (every
/// activation stored). On, each block's forward is wrapped in
/// [`xtrain_autodiff::checkpoint`] — exact grads, lower peak activation memory.
pub fn with_recompute(mut self, recompute: bool) -> Self {
self.recompute = recompute;
self
}
pub fn recompute(&self) -> bool {
self.recompute
}
/// Enable the fused flash-attention SDPA core (Phase T14). Builder-style and
/// opt-in; default off keeps the composed T10 path (so the default graph is
/// unchanged). On, the SDPA runs through [`ops::flash_attention`] — same SDPA
/// math, online softmax, no materialized `[bh,seq,seq]` scores.
pub fn with_flash(mut self, use_flash: bool) -> Self {
self.use_flash = use_flash;
self
}
pub fn use_flash(&self) -> bool {
self.use_flash
}
/// Switch to training mode (Phase T18): dropout (if `cfg.dropout > 0`) is
/// active in subsequent forwards. The training loop calls this before stepping.
pub fn train(&self) {
self.training.set(true);
}
/// Switch to eval mode (Phase T18): dropout is identity. Held-out eval,
/// autoregressive sampling, and weight export all run in this mode (default).
pub fn eval(&self) {
self.training.set(false);
}
pub fn is_training(&self) -> bool {
self.training.get()
}
/// Builder-style train/eval toggle (Phase T18) — handy for tests that want a
/// model fixed in one mode. Equivalent to [`train`](Self::train) /
/// [`eval`](Self::eval) but chains off `new(..)`.
pub fn with_training(self, training: bool) -> Self {
self.training.set(training);
self
}
/// All learnable parameters, in a stable order. The optimizer (a hand-written
/// GD step in T5, AdamW in T6) iterates this; each holds its `.grad()` after
/// `backward()`.
pub fn params(&self) -> Vec<Var> {
let mut ps = vec![self.embed.clone()];
for b in &self.blocks {
ps.extend([
b.attn_norm.clone(),
b.wq.clone(),
b.wk.clone(),
b.wv.clone(),
b.q_norm.clone(),
b.k_norm.clone(),
b.wo.clone(),
b.ffn_norm.clone(),
b.w_gate.clone(),
b.w_up.clone(),
b.w_down.clone(),
]);
}
ps.push(self.final_norm.clone());
ps.push(self.lm_head.clone());
ps
}
/// Forward over a single sequence of token `ids` (`[seq]` I32 on this
/// model's device). Returns the logits [`Var`] of shape `[seq, vocab]`. This
/// is the batch-1 special case of [`forward_batched`](Self::forward_batched)
/// (used by the autoregressive sampler / inference path).
pub fn forward(&self, ids: &Tensor) -> Var {
self.forward_batched(ids, 1)
}
/// Batched forward over `batch` sequences of equal length `seq`, flattened to
/// `[batch*seq]` I32 ids in sequence-major order (sequence 0's `seq` tokens,
/// then sequence 1's, …). Returns logits `[batch*seq, vocab]` in the SAME flat
/// layout. The whole graph runs on the flattened tokens so every linear
/// projection is ONE big `[batch*seq, dim] × [dim, out]` GEMM (the
/// GPU-filling win); only attention is sequence-aware (per-sequence causal
/// mask + RoPE position, NO cross-sequence attention).
pub fn forward_batched(&self, ids: &Tensor, batch: usize) -> Var {
let total = ids.shape()[0];
assert_eq!(
total % batch,
0,
"ids len {total} not divisible by batch {batch}"
);
let seq = total / batch;
// Dropout (T18) is active only in training mode with p>0; otherwise it is
// identity (`ops::dropout` no-ops at p==0). Bump the per-step seed ONCE per
// training forward so each step draws fresh masks (counter-based RNG, so a
// checkpointed block's recompute reproduces the same seed → same mask).
let dropout_p = if self.training.get() {
self.cfg.dropout
} else {
0.0
};
if dropout_p > 0.0 {
self.step_seed.set(self.step_seed.get().wrapping_add(1));
}
let base_seed = self.step_seed.get();
// Embedding gathers from the fp32 master table; in bf16 mode cast the
// activation stream to bf16 here (norms are cast to bf16 gammas too).
let mut h = ops::embedding(&self.embed, ids); // [batch*seq, dim], fp32
if self.compute_dtype == DType::BF16 {
h = ops::cast(&h, DType::BF16);
}
for (li, b) in self.blocks.iter().enumerate() {
// Per-layer dropout seed: a deterministic function of (base_seed,
// layer index) — NOT a mutable counter — so the checkpoint recompute
// (which re-derives it from the captured base_seed/li) gets the same
// masks. The block derives its two per-site seeds from this.
let block_seed = base_seed
.wrapping_mul(0x100000001B3)
.wrapping_add(li as u64);
h = if self.recompute {
// Activation recomputation (T13): run the whole block forward inside
// `checkpoint` so its internal activations aren't kept on the tape;
// the block forward is re-run in backward to recover the grads. The
// segment fn captures only `Copy` config (no borrow of `self`) and
// receives the block's params via the slice, in `block_params` order.
// `flash` is captured too → the recompute segment also runs flash;
// `dropout_p`/`block_seed` are captured so the recompute re-derives
// the same per-site dropout masks (counter-based RNG, exact).
let (cfg, cdt, flash) = (self.cfg, self.compute_dtype, self.use_flash);
let seg = move |x: &Var, p: &[Var]| {
block_forward(cfg, cdt, flash, batch, seq, dropout_p, block_seed, x, p)
};
xtrain_autodiff::checkpoint::checkpoint(seg, &h, &b.block_params())
} else {
block_forward(
self.cfg,
self.compute_dtype,
self.use_flash,
batch,
seq,
dropout_p,
block_seed,
&h,
&b.block_params(),
)
};
}
let h = ops::rms_norm(
&h,
&norm_gamma(self.compute_dtype, &self.final_norm),
self.cfg.eps,
);
// lm_head matmul in compute dtype. Logits stay bf16 in bf16 mode — the
// cross_entropy op upcasts to fp32 internally (no persistent fp32 logits
// buffer, a real saving at vocab 50257), and its backward casts dx back.
linear(self.compute_dtype, &h, &self.lm_head) // [batch*seq, vocab]
}
/// Cross-entropy mean loss of `forward(ids)` against `targets` (`[seq]` I32).
pub fn loss(&self, ids: &Tensor, targets: &Tensor) -> Var {
let logits = self.forward(ids);
ops::cross_entropy(&logits, targets)
}
/// Batched cross-entropy mean loss: `forward_batched(ids, batch)` against
/// flat `targets` (`[batch*seq]` I32, same sequence-major layout). The CE mean
/// is over all `batch*seq` rows — identical to averaging the per-sequence
/// losses, so the loss value matches the looped single-sequence path.
pub fn loss_batched(&self, ids: &Tensor, targets: &Tensor, batch: usize) -> Var {
let logits = self.forward_batched(ids, batch);
ops::cross_entropy(&logits, targets)
}
}
impl Block {
/// The block's learnable leaves, in the fixed order the segment forward
/// (`block_forward`) indexes them — matches the per-block slice in
/// [`TinyTransformer::params`]. This is the param order `checkpoint` passes to
/// the recompute closure.
fn block_params(&self) -> Vec<Var> {
vec![
self.attn_norm.clone(),
self.wq.clone(),
self.wk.clone(),
self.wv.clone(),
self.q_norm.clone(),
self.k_norm.clone(),
self.wo.clone(),
self.ffn_norm.clone(),
self.w_gate.clone(),
self.w_up.clone(),
self.w_down.clone(),
]
}
}
/// Project `x` (activation, in the compute dtype) by weight `w` (an fp32 master
/// leaf). In bf16 mode the weight is cast to bf16 via the autograd `cast` op (whose
/// backward upcasts the grad to fp32); in fp32 mode this is just `matmul(x, w)`.
fn linear(cdt: DType, x: &Var, w: &Var) -> Var {
match cdt {
DType::F32 => ops::matmul(x, w),
DType::BF16 => ops::matmul(x, &ops::cast(w, DType::BF16)),
_ => unreachable!(),
}
}
/// A norm/QK-norm gamma in the compute dtype. fp32 master leaf → bf16 (cast op,
/// grad upcast) in bf16 mode; identity in fp32 mode.
fn norm_gamma(cdt: DType, gamma: &Var) -> Var {
match cdt {
DType::F32 => gamma.clone(),
DType::BF16 => ops::cast(gamma, DType::BF16),
_ => unreachable!(),
}
}
/// One transformer block's forward: pre-norm + multi-head causal attention +
/// (T18) dropout + residual, then pre-norm + SwiGLU MLP + dropout + residual.
/// Attention runs the composed or fused-flash (T14) SDPA per `flash`. Pure in
/// `(cfg, cdt, flash, batch, seq, dropout_p, block_seed, input, params)` (no
/// `&self`, all `Copy`) so it can be the segment fn of
/// [`xtrain_autodiff::checkpoint`] for activation recomputation (T13) — the
/// recompute re-derives the same per-site seeds, so the dropout masks are
/// reproduced bit-for-bit. `dropout_p == 0` makes `ops::dropout` a no-op (the
/// graph is then identical to the pre-T18 path). `params` is the block's leaves in
/// [`Block::block_params`] order.
#[allow(clippy::too_many_arguments)]
fn block_forward(
cfg: Config,
cdt: DType,
flash: bool,
batch: usize,
seq: usize,
dropout_p: f32,
block_seed: u64,
h: &Var,
p: &[Var],
) -> Var {
let (attn_norm, wq, wk, wv) = (&p[0], &p[1], &p[2], &p[3]);
let (q_norm, k_norm, wo) = (&p[4], &p[5], &p[6]);
let (ffn_norm, w_gate, w_up, w_down) = (&p[7], &p[8], &p[9], &p[10]);
// Per-site dropout seeds (XOR a site constant into the block seed) so the two
// residual-path dropouts draw independent masks within the same step/layer.
let attn_seed = block_seed ^ 0x0A7700;
let ffn_seed = block_seed ^ 0x0FF700;
// --- Attention sub-block (pre-norm + dropout + residual) ---
let normed = ops::rms_norm(h, &norm_gamma(cdt, attn_norm), cfg.eps);
let attn = attention(
cfg, cdt, flash, batch, seq, &normed, wq, wk, wv, q_norm, k_norm, wo,
);
let attn = ops::dropout(&attn, dropout_p, attn_seed);
let h = ops::add(h, &attn);
// --- MLP sub-block (pre-norm + dropout + residual) ---
let normed = ops::rms_norm(&h, &norm_gamma(cdt, ffn_norm), cfg.eps);
let mlp = swiglu_mlp(cdt, &normed, w_gate, w_up, w_down);
let mlp = ops::dropout(&mlp, dropout_p, ffn_seed);
ops::add(&h, &mlp)
}
/// Multi-head causal self-attention over a flattened batch. `x`:[batch*seq,dim]
/// (already normed), laid out sequence-major. The Q/K/V/O projections are big
/// `[batch*seq, dim]` GEMMs; the scaled-dot-product attention itself runs as a
/// fused BATCHED op over the `batch·n_heads` (sequence,head) blocks — each attends
/// within its own `[seq,seq]` causal window (NO cross-sequence attention), with
/// RoPE positions reset per sequence (`period = seq`). Causal masking is applied
/// inside the fused op's softmax kernel (no additive `[seq,seq]` mask tensor).
#[allow(clippy::too_many_arguments)]
fn attention(
cfg: Config,
cdt: DType,
flash: bool,
batch: usize,
seq: usize,
x: &Var,
wq: &Var,
wk: &Var,
wv: &Var,
q_norm: &Var,
k_norm: &Var,
wo: &Var,
) -> Var {
let (nh, hd) = (cfg.n_heads, cfg.head_dim);
let num_kv = cfg.num_kv_heads; // GQA (T15): K/V have fewer heads than Q
let total = batch * seq;
let scale = 1.0 / (hd as f32).sqrt();
// Project, qk-norm + RoPE, then lay out as a batched [B*heads, seq, hd] tensor.
// `heads` = nh for Q, num_kv for K/V (GQA; equal for MHA).
// [B*S,dim] @ [dim,heads*hd] = [B*S, heads*hd]
// reshape [B*S, heads, hd]
// qk-norm per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
// rope [B*S, heads, hd] with per-sequence position (period = seq)
// reshape [B, S, heads, hd] → transpose(1,2) → [B, heads, S, hd] → [B*heads, S, hd]
let to_bh = |proj: Var, heads: usize, norm: Option<&Var>| -> Var {
let r = ops::reshape(&proj, &[total, heads, hd]);
let r = match norm {
// Per-head RMSNorm: flatten the (B*S,heads) head rows, norm over hd,
// restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
Some(gamma) => {
let flat = ops::reshape(&r, &[total * heads, hd]);
let normed = ops::rms_norm(&flat, &norm_gamma(cdt, gamma), cfg.eps);
let r = ops::reshape(&normed, &[total, heads, hd]);
ops::rope(&r, cfg.rope_theta, seq)
}
None => r,
};
let r = ops::reshape(&r, &[batch, seq, heads, hd]);
let t = ops::transpose_4d12(&r); // [B, heads, S, hd]
ops::reshape(&t, &[batch * heads, seq, hd]) // [B*heads, S, hd]
};
let q = to_bh(linear(cdt, x, wq), nh, Some(q_norm));
// K/V are laid out with num_kv heads, then repeat_kv-broadcast to nh heads so
// the SDPA below (composed or flash, both unchanged) sees a full head set. The
// broadcast's backward sums each KV head's group of query-head grads (GQA). For
// MHA (num_kv == nh) repeat_kv is identity → bit-identical to the pre-T15 path.
let k = to_bh(linear(cdt, x, wk), num_kv, Some(k_norm));
let v = to_bh(linear(cdt, x, wv), num_kv, None);
let (k, v) = if num_kv == nh {
(k, v)
} else {
(ops::repeat_kv(&k, nh, batch), ops::repeat_kv(&v, nh, batch))
};
// Causal SDPA over all B*nh (sequence,head) blocks. `flash` (T14) picks the
// single fused flash kernel (online softmax, no materialized [bh,S,S] scores);
// otherwise the composed T10 path (2 batched GEMMs + 1 causal-softmax kernel).
let out = if flash {
ops::flash_attention(&q, &k, &v, scale) // [B*nh, S, hd]
} else {
ops::attention(&q, &k, &v, scale) // [B*nh, S, hd]
};
// Back to [B*S, dim]: [B*nh,S,hd] → [B,nh,S,hd] → transpose(1,2) →
// [B,S,nh,hd] → [B*S, dim].
let out = ops::reshape(&out, &[batch, nh, seq, hd]);
let out = ops::transpose_4d12(&out); // [B, S, nh, hd]
let concat = ops::reshape(&out, &[total, nh * hd]); // [B*S, dim]
linear(cdt, &concat, wo) // out projection
}
/// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[batch*seq,dim].
fn swiglu_mlp(cdt: DType, x: &Var, w_gate: &Var, w_up: &Var, w_down: &Var) -> Var {
let gate = linear(cdt, x, w_gate); // [seq, ffn_hidden]
let up = linear(cdt, x, w_up); // [seq, ffn_hidden]
let act = ops::swiglu(&gate, &up); // silu(gate) ∘ up
linear(cdt, &act, w_down) // [seq, dim]
}
/// Materialise a parameter's value back to a host `Vec<f32>` (for the GD step
/// and PyTorch parity export).
pub fn param_to_host(v: &Var) -> Vec<f32> {
v.value().to_device(Device::Cpu).as_slice::<f32>().to_vec()
}
/// Build an I32 id tensor on `device` from token ids.
pub fn ids_tensor(ids: &[i32], device: Device) -> Tensor {
Tensor::from_slice(ids, &[ids.len()]).to_device(device)
}
/// Flatten `batch` equal-length sequences into one `[batch*seq]` I32 tensor in
/// sequence-major order (the layout `forward_batched` expects). Each row of
/// `seqs` is one sequence; all must have the same length.
pub fn batched_ids_tensor(seqs: &[Vec<i32>], device: Device) -> Tensor {
assert!(!seqs.is_empty(), "empty batch");
let seq = seqs[0].len();
let mut flat = Vec::with_capacity(seqs.len() * seq);
for s in seqs {
assert_eq!(s.len(), seq, "ragged batch: sequences must be equal length");
flat.extend_from_slice(s);
}
Tensor::from_slice(&flat, &[flat.len()]).to_device(device)
}

142
crates/xtrain-model/tests/batched.rs Normal file
View File

@@ -0,0 +1,142 @@
// T10 batched-forward equivalence: a batched forward over B sequences must equal
// the old single-sequence path (run each sequence on its own, concatenate the
// logits) — both for the forward logits AND every parameter's gradient.
//
// This is THE on-GPU correctness gate for batching (no PyTorch needed): if the
// per-sequence RoPE position, per-sequence causal masking, or any flattened op
// were wrong, the batched logits/grads would drift from the looped reference.
//
// Forward equivalence: batched logits[b*S+i] == single-seq-b logits[i].
// Gradient equivalence: the batched loss is the mean over all B*S rows, i.e.
// (1/B)·Σ_b mean_i(loss_b); summing the B single-sequence losses and scaling by
// 1/B gives the SAME scalar, so their summed grads (tape fan-out) ×1/B match the
// batched grads. We check that.
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor, ids_tensor};
use xtrain_tensor::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)
}
})
}
fn host(t: &xtrain_tensor::Tensor) -> Vec<f32> {
t.to_device(Device::Cpu).as_slice::<f32>().to_vec()
}
#[test]
fn batched_matches_looped_single_sequence() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let mut cfg = Config::tiny();
cfg.vocab = 16;
let batch = 3usize;
let seq = 5usize;
// B distinct sequences (sequence-major), within vocab.
let seqs: Vec<Vec<i32>> = (0..batch)
.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()
})
.collect();
// --- Batched forward: ONE pass over [B*S]. ---
let bmodel = build(cfg, device);
let bids = batched_ids_tensor(&seqs, device);
let blogits = host(&bmodel.forward_batched(&bids, batch).value());
// --- Looped reference: each sequence on its own, concatenate logits. ---
let smodel = build(cfg, device);
let mut slogits = Vec::with_capacity(batch * seq * cfg.vocab);
for s in &seqs {
let ids = ids_tensor(s, device);
slogits.extend(host(&smodel.forward(&ids).value()));
}
// Forward equivalence (fp GEMM rounding only differs in summation order).
let max_rel = blogits
.iter()
.zip(&slogits)
.map(|(b, s)| (b - s).abs() / s.abs().max(1e-4))
.fold(0.0f32, f32::max);
println!("batched vs looped: logits max rel err = {max_rel:.3e}");
assert!(max_rel < 1e-3, "batched logits diverged: {max_rel:.3e}");
// --- Gradient equivalence. ---
// Batched: loss = mean over B*S rows; one backward.
let bparams = bmodel.params();
let btgt = batched_ids_tensor(&tgts, device);
let bloss = bmodel.loss_batched(&bids, &btgt, batch);
let bloss_val = host(&bloss.value())[0];
bloss.backward();
// Looped: Σ_b loss_b (each a per-sequence mean), then grad ×(1/B) == batched.
let sparams = smodel.params();
let mut sloss_sum = 0.0f32;
for (s, t) in seqs.iter().zip(&tgts) {
let ids = ids_tensor(s, device);
let tg = ids_tensor(t, device);
let l = smodel.loss(&ids, &tg);
sloss_sum += host(&l.value())[0];
l.backward();
}
println!(
"batched loss = {bloss_val:.6} looped mean = {:.6}",
sloss_sum / batch as f32
);
assert!(
(bloss_val - sloss_sum / batch as f32).abs() < 1e-4,
"batched loss != looped mean"
);
let mut max_grad_rel = 0.0f32;
for (bp, sp) in bparams.iter().zip(&sparams) {
let bg = host(&bp.grad().expect("batched grad"));
let sg = host(&sp.grad().expect("looped grad"));
for (g_b, g_s) in bg.iter().zip(&sg) {
// looped grad is the SUM over B sequences; ×(1/B) recovers the mean.
let g_s = g_s / batch as f32;
let rel = (g_b - g_s).abs() / g_s.abs().max(1e-4);
max_grad_rel = max_grad_rel.max(rel);
}
}
println!("batched vs looped: grad max rel err = {max_grad_rel:.3e}");
assert!(
max_grad_rel < 5e-3,
"batched grads diverged: {max_grad_rel:.3e}"
);
}

151
crates/xtrain-model/tests/bf16.rs Normal file
View File

@@ -0,0 +1,151 @@
// T12 bf16 mixed-precision correctness gate (on-GPU, no PyTorch).
//
// The SAME model (identical fp32 master weights) run in fp32 vs bf16 compute
// mode must agree within a LOOSE bf16 tolerance (bf16 = 7-bit mantissa ≈ 2-3
// decimal digits → ~1e-2 relative error is expected and acceptable), both for
// the forward loss/logits AND every parameter's gradient. We also assert no
// NaN/Inf leaks and that the fp32 grads are fp32 (the cast op upcast the bf16
// weight grad back to the fp32 master, so AdamW/clip/DDP stay fp32).
//
// This is the "bf16 within looser tol vs fp32 reference" gate; the short-run
// convergence comparison is the train_loop-level bench on dash5.
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
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)
}
})
}
fn host(t: &xtrain_tensor::Tensor) -> Vec<f32> {
t.to_device(Device::Cpu).as_slice::<f32>().to_vec()
}
#[test]
fn bf16_matches_fp32_within_loose_tol() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
// A few layers / heads so the bf16 rounding accumulates through the depth
// the real model has (not just a single matmul).
let mut cfg = Config::tiny();
cfg.vocab = 32;
cfg.n_layers = 3;
let batch = 2usize;
let seq = 8usize;
let seqs: Vec<Vec<i32>> = (0..batch)
.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()
})
.collect();
let ids = batched_ids_tensor(&seqs, device);
let tgt = batched_ids_tensor(&tgts, device);
// fp32 reference.
let fp32 = build(cfg, device);
let f_logits = host(&fp32.forward_batched(&ids, batch).value());
let f_loss = fp32.loss_batched(&ids, &tgt, batch);
let f_loss_val = host(&f_loss.value())[0];
f_loss.backward();
let f_params = fp32.params();
// bf16 — SAME init (build re-runs the same deterministic fill). The forward
// now returns bf16 logits (CE upcasts internally); cast to f32 to read.
let bf16 = build(cfg, device).with_compute_dtype(DType::BF16);
let b_logits = host(
&bf16
.forward_batched(&ids, batch)
.value()
.to_dtype(DType::F32),
);
let b_loss = bf16.loss_batched(&ids, &tgt, batch);
let b_loss_val = host(&b_loss.value())[0];
b_loss.backward();
let b_params = bf16.params();
// No NaN/Inf in the bf16 forward.
assert!(
b_logits.iter().all(|v| v.is_finite()) && b_loss_val.is_finite(),
"bf16 forward produced non-finite values"
);
// Forward loss within loose bf16 tol.
let loss_rel = (b_loss_val - f_loss_val).abs() / f_loss_val.abs().max(1e-4);
println!("bf16 vs fp32: loss {b_loss_val:.5} vs {f_loss_val:.5} (rel {loss_rel:.3e})");
assert!(
loss_rel < 2e-2,
"bf16 loss too far from fp32: {loss_rel:.3e}"
);
// Logits: bf16 has ~2-3 decimal digits → compare on a robust (median-style)
// basis, requiring the bulk to be within ~3e-2 and the mean error small.
let n = f_logits.len();
let mut rels: Vec<f32> = f_logits
.iter()
.zip(&b_logits)
.map(|(f, b)| (b - f).abs() / f.abs().max(1.0))
.collect();
rels.sort_by(|a, b| a.partial_cmp(b).unwrap());
let p99 = rels[(n as f32 * 0.99) as usize];
let mean: f32 = rels.iter().sum::<f32>() / n as f32;
println!("bf16 vs fp32 logits: mean rel {mean:.3e}, p99 rel {p99:.3e}");
assert!(mean < 1e-2, "bf16 logits mean rel err too high: {mean:.3e}");
assert!(p99 < 5e-2, "bf16 logits p99 rel err too high: {p99:.3e}");
// Gradients: fp32 master grads must be fp32 (cast op upcast), finite, and
// within loose bf16 tol of the fp32 reference (mean over each param tensor).
let mut worst_param_mean = 0.0f32;
for (fp, bp) in f_params.iter().zip(&b_params) {
let bg = bp.grad().expect("bf16 grad");
assert_eq!(bg.dtype(), DType::F32, "bf16-mode grad must be fp32 master");
let fg = host(&fp.grad().expect("fp32 grad"));
let bg = host(&bg);
assert!(bg.iter().all(|v| v.is_finite()), "bf16 grad has non-finite");
// Scale-relative mean error over the tensor (robust to a few small entries).
let scale = fg.iter().map(|v| v.abs()).fold(0.0f32, f32::max).max(1e-6);
let mean_err: f32 =
fg.iter().zip(&bg).map(|(f, b)| (f - b).abs()).sum::<f32>() / fg.len() as f32 / scale;
worst_param_mean = worst_param_mean.max(mean_err);
}
println!("bf16 vs fp32 grads: worst per-tensor scaled-mean err = {worst_param_mean:.3e}");
assert!(
worst_param_mean < 3e-2,
"bf16 grads too far from fp32: {worst_param_mean:.3e}"
);
}

322
crates/xtrain-model/tests/dropout.rs Normal file
View File

@@ -0,0 +1,322 @@
// T18 dropout model-level gates.
//
// 1. p=0 bit-identical: a model built with cfg.dropout=0 (in either train or
// eval mode) produces logits/loss/grads bit-for-bit identical to the same
// model with no dropout field touched — the default forward graph is
// unchanged (the regression guard).
// 2. eval identity: with p>0 but eval mode, the forward equals the p=0 forward
// bit-for-bit (dropout is OFF at eval).
// 3. train vs eval differ: with p>0 and train mode, the forward differs from
// eval (dropout actually does something) and grads are still finite.
// 4. recompute compatibility: with p>0 + train + recompute, grads match the
// non-recompute path (the counter-based seed reproduces the same mask on the
// backward re-run — T13 stays exact even with dropout in the block).
//
// (The fixed-seed grad-check of the dropout op and the E[out]≈x / keep-rate check
// live in xtrain-autodiff/tests/autograd.rs; p>0 training convergence is the
// dash5 short run noted in docs/17-dropout.md.)
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
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)
}
})
}
fn host(t: &xtrain_tensor::Tensor) -> Vec<f32> {
t.to_dtype(DType::F32)
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec()
}
fn tiny_cfg(dropout: f32) -> Config {
let mut cfg = Config::tiny();
cfg.vocab = 16;
cfg.n_layers = 4;
cfg.dropout = dropout;
cfg
}
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()
})
.collect();
let tgts: Vec<Vec<i32>> = (0..batch)
.map(|b| {
(0..seq)
.map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
.collect()
})
.collect();
(
batched_ids_tensor(&seqs, device),
batched_ids_tensor(&tgts, device),
)
}
fn require_gpu() -> Device {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
Device::Cuda(0)
}
// Run forward+backward, return (logits, loss, per-param grads).
fn fwd_bwd(
m: &TinyTransformer,
ids: &xtrain_tensor::Tensor,
tgt: &xtrain_tensor::Tensor,
batch: usize,
) -> (Vec<f32>, f32, Vec<Vec<f32>>) {
let logits = host(&m.forward_batched(ids, batch).value());
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();
(logits, loss_val, grads)
}
// --- Gate 3: p=0 is bit-identical to the no-dropout path (default graph). ---
#[test]
fn dropout_p0_bit_identical() {
let device = require_gpu();
let batch = 3;
// Reference: cfg.dropout default (0.0), never touched train/eval.
let cfg0 = tiny_cfg(0.0);
let (ids, tgt) = batch_data(&cfg0, device);
let ref_m = build(cfg0, device);
let (ref_logits, ref_loss, ref_grads) = fwd_bwd(&ref_m, &ids, &tgt, batch);
// p=0 in TRAINING mode: the seed bump is gated on p>0, the op no-ops at p==0,
// so the graph must be byte-identical.
let p0_train = build(tiny_cfg(0.0), device);
p0_train.train();
let (lt, lst, gt) = fwd_bwd(&p0_train, &ids, &tgt, batch);
assert_eq!(ref_logits, lt, "p=0 train logits not bit-identical");
assert_eq!(ref_loss, lst, "p=0 train loss not bit-identical");
for (i, (a, b)) in ref_grads.iter().zip(&gt).enumerate() {
assert_eq!(a, b, "p=0 train grad[{i}] not bit-identical");
}
println!("p=0 (train) vs no-dropout: logits/loss/grads bit-identical ✅");
}
// --- Gate 2: eval is exact identity (p>0 but eval mode == p=0). ---
#[test]
fn dropout_eval_is_identity() {
let device = require_gpu();
let batch = 3;
let cfg = tiny_cfg(0.2);
let (ids, tgt) = batch_data(&cfg, device);
// p=0 reference and a p=0.2 model held in eval — outputs must match bit-for-bit.
let ref_m = build(tiny_cfg(0.0), device);
let (ref_logits, ref_loss, ref_grads) = fwd_bwd(&ref_m, &ids, &tgt, batch);
let eval_m = build(cfg, device);
eval_m.eval(); // explicit; also the default
let (el, els, eg) = fwd_bwd(&eval_m, &ids, &tgt, batch);
assert_eq!(ref_logits, el, "eval (p>0) logits not identity");
assert_eq!(ref_loss, els, "eval (p>0) loss not identity");
for (i, (a, b)) in ref_grads.iter().zip(&eg).enumerate() {
assert_eq!(a, b, "eval (p>0) grad[{i}] not identity");
}
println!("eval (p=0.2) == no-dropout: bit-identical (eval is identity) ✅");
}
// --- Gate (train vs eval differ): with p>0 + train, dropout actually fires. ---
#[test]
fn dropout_train_differs_from_eval() {
let device = require_gpu();
let batch = 3;
let cfg = tiny_cfg(0.3);
let (ids, _tgt) = batch_data(&cfg, device);
let m = build(cfg, device);
m.eval();
let eval_logits = host(&m.forward_batched(&ids, batch).value());
m.train();
let train_logits = host(&m.forward_batched(&ids, batch).value());
let max_diff = eval_logits
.iter()
.zip(&train_logits)
.map(|(a, b)| (a - b).abs())
.fold(0.0f32, f32::max);
assert!(
max_diff > 1e-4 && train_logits.iter().all(|v| v.is_finite()),
"train logits should differ from eval (dropout active) and be finite; max_diff={max_diff}"
);
println!("train vs eval logits max diff {max_diff:.4e} (dropout active in train) ✅");
}
// --- Gate 4: p>0 + recompute grads match non-recompute (T13 stays exact). ---
// The counter-based seed is a pure function of (step_seed, layer, site); the
// checkpoint backward re-runs block_forward and re-derives the SAME seeds, so the
// recomputed dropout masks match the forward — grads stay bit-identical.
fn recompute_with_dropout(dtype: DType, grad_tol: f32) {
let device = require_gpu();
let batch = 3;
let cfg = tiny_cfg(0.2);
let (ids, tgt) = batch_data(&cfg, device);
// 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 on = build(cfg, device)
.with_compute_dtype(dtype)
.with_recompute(true)
.with_training(true);
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 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 mut max_rel = 0.0f32;
for (a, b) in off_grads.iter().flatten().zip(on_grads.iter().flatten()) {
max_rel = max_rel.max((a - b).abs() / a.abs().max(1e-3));
}
println!("[{dtype:?}] dropout p=0.2 recompute on/off grad max rel = {max_rel:.3e}");
assert!(
max_rel < grad_tol,
"[{dtype:?}] recompute grads diverged with dropout: {max_rel:.3e}"
);
}
#[test]
fn dropout_recompute_matches_fp32() {
recompute_with_dropout(DType::F32, 1e-4);
}
#[test]
fn dropout_recompute_matches_bf16() {
recompute_with_dropout(DType::BF16, 5e-3);
}
// --- Cross-feature gate (Phase-2 integration): flash (T14) + dropout (T18)
// together in the SAME model still grad-checks. Build two identical models, both
// in train mode with p=0.2 (so dropout fires), one with `--flash` on, one off.
// The dropout site seeds are a pure function of (step_seed, layer, site) and are
// INDEPENDENT of flash, so both models draw the SAME masks on their first training
// forward → the only difference is the SDPA reduction order. Assert logits/loss/
// grads match within the flash-vs-composed tolerance and are finite. This is the
// orthogonality check for the two Phase-2 features landing together.
#[test]
fn flash_plus_dropout_grad_check_fp32() {
let device = require_gpu();
let batch = 3;
// seq=40 > FA_TILE=32 exercises flash's online-softmax tile-rescale path.
let mut cfg = Config::tiny();
cfg.vocab = 16;
cfg.n_layers = 4;
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()
})
.collect();
let tgts: Vec<Vec<i32>> = (0..batch)
.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);
// Both: same init, train mode (dropout active), same step_seed progression →
// identical masks; one composed SDPA, one flash SDPA.
let off = build(cfg, device).with_training(true);
let on = build(cfg, device).with_flash(true).with_training(true);
let (off_logits, off_loss, off_grads) = fwd_bwd(&off, &ids, &tgt, batch);
let (on_logits, on_loss, on_grads) = fwd_bwd(&on, &ids, &tgt, batch);
assert!(
on_logits.iter().all(|v| v.is_finite()) && on_grads.iter().flatten().all(|v| v.is_finite()),
"flash+dropout produced non-finite logits/grads"
);
let logit_rel = off_logits
.iter()
.zip(&on_logits)
.map(|(a, b)| (a - b).abs() / a.abs().max(1e-4))
.fold(0.0f32, f32::max);
let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
let mut grad_rel = 0.0f32;
for (a, b) in off_grads.iter().flatten().zip(on_grads.iter().flatten()) {
grad_rel = grad_rel.max((a - b).abs() / a.abs().max(1e-3));
}
println!(
"[F32] flash+dropout vs composed+dropout: loss rel {loss_rel:.2e}, \
logits max rel {logit_rel:.2e}, grad max rel {grad_rel:.3e}"
);
// 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}"
);
}

209
crates/xtrain-model/tests/flash.rs Normal file
View File

@@ -0,0 +1,209 @@
// T14 flash-attention correctness gate: the fused flash SDPA core must match the
// composed T10 path (cublasSgemmStridedBatched×2 + causal-softmax kernel) in
// forward logits, loss, AND every parameter gradient — flash is the SAME SDPA
// math (online softmax never materializes the [bh,S,S] scores), so it differs
// from composed only by reduction order (in-kernel fp32 FMA vs cuBLAS, and the
// dK/dV atomicAdd order in backward). This test makes that a closed on-GPU loop:
//
// build two identical models (same init), one with `--flash` on, one off, run
// the SAME batched loss + backward on both, and assert
// 1. the forward logits match within tolerance
// 2. the loss matches
// 3. EVERY parameter's grad matches within tolerance
//
// Parameterised over fp32 AND bf16 (T12). bf16 just adds the bf16 rounding band on
// top — flash's bf16 path upcasts Q/K/V to fp32 for the kernel exactly like the
// composed path's fp32 softmax, so the two are still the same softmax numerics.
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
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, 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: &xtrain_tensor::Tensor) -> Vec<f32> {
t.to_dtype(DType::F32)
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec()
}
// fp32: same SDPA math, differs only by reduction order → tight per-element check.
fn run_fp32(logit_tol: f32, grad_tol: f32) {
let (off_logits, off_loss, off_grads, on_logits, on_loss, on_grads) = run_both(DType::F32);
let logit_rel = off_logits
.iter()
.zip(&on_logits)
.map(|(a, b)| (a - b).abs() / a.abs().max(1e-4))
.fold(0.0f32, f32::max);
let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
println!(
"[F32] flash on/off: loss {off_loss:.6}/{on_loss:.6} (rel {loss_rel:.2e}), \
logits max rel {logit_rel:.2e}"
);
assert!(
logit_rel < logit_tol,
"[F32] logits diverged: {logit_rel:.2e}"
);
assert!(loss_rel < logit_tol, "[F32] loss diverged: {loss_rel:.2e}");
let mut max_grad_rel = 0.0f32;
for (off_g, on_g) in off_grads.iter().zip(&on_grads) {
for (a, b) in off_g.iter().zip(on_g) {
max_grad_rel = max_grad_rel.max((a - b).abs() / a.abs().max(1e-3));
}
}
println!("[F32] flash on/off: grad max rel err = {max_grad_rel:.3e}");
assert!(
max_grad_rel < grad_tol,
"[F32] flash grads diverged from composed: {max_grad_rel:.3e}"
);
}
// bf16: ~2-3 decimal digits → robust comparison (mean + p99 with abs().max(1.0)
// for logits, per-tensor scale-relative mean for grads), the same convention as
// the repo's bf16.rs gate (per-element max-rel blows up on near-zero bf16 logits).
fn run_bf16() {
let (off_logits, off_loss, off_grads, on_logits, on_loss, on_grads) = run_both(DType::BF16);
let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
println!("[BF16] flash on/off: loss {off_loss:.5}/{on_loss:.5} (rel {loss_rel:.3e})");
assert!(loss_rel < 2e-2, "[BF16] loss diverged: {loss_rel:.3e}");
let n = off_logits.len();
let mut rels: Vec<f32> = off_logits
.iter()
.zip(&on_logits)
.map(|(f, b)| (b - f).abs() / f.abs().max(1.0))
.collect();
rels.sort_by(|a, b| a.partial_cmp(b).unwrap());
let p99 = rels[(n as f32 * 0.99) as usize];
let mean: f32 = rels.iter().sum::<f32>() / n as f32;
println!("[BF16] flash on/off logits: mean rel {mean:.3e}, p99 rel {p99:.3e}");
assert!(mean < 1e-2, "[BF16] logits mean rel too high: {mean:.3e}");
assert!(p99 < 5e-2, "[BF16] logits p99 rel too high: {p99:.3e}");
let mut worst = 0.0f32;
for (off_g, on_g) in off_grads.iter().zip(&on_grads) {
let scale = off_g
.iter()
.map(|v| v.abs())
.fold(0.0f32, f32::max)
.max(1e-6);
let mean_err: f32 = off_g
.iter()
.zip(on_g)
.map(|(f, b)| (f - b).abs())
.sum::<f32>()
/ off_g.len() as f32
/ scale;
worst = worst.max(mean_err);
}
println!("[BF16] flash on/off grads: worst per-tensor scaled-mean err = {worst:.3e}");
assert!(worst < 3e-2, "[BF16] flash grads diverged: {worst:.3e}");
}
#[allow(clippy::type_complexity)]
fn run_both(dtype: DType) -> (Vec<f32>, f32, Vec<Vec<f32>>, Vec<f32>, f32, Vec<Vec<f32>>) {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
// seq=40 > FA_TILE=32 so the online-softmax tile-rescale path is exercised.
let mut cfg = Config::tiny();
cfg.vocab = 16;
cfg.n_layers = 4;
let batch = 3usize;
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()
})
.collect();
let tgts: Vec<Vec<i32>> = (0..batch)
.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);
// --- flash OFF (composed reference) ---
let off = build(cfg, device, dtype, false);
let off_logits = host(&off.forward_batched(&ids, batch).value());
let off_loss = off.loss_batched(&ids, &tgt, batch);
let off_loss_val = host(&off_loss.value())[0];
off_loss.backward();
let off_grads: Vec<Vec<f32>> = off
.params()
.iter()
.map(|p| host(&p.grad().expect("off grad")))
.collect();
// --- flash ON ---
let on = build(cfg, device, dtype, true);
let on_logits = host(&on.forward_batched(&ids, batch).value());
let on_loss = on.loss_batched(&ids, &tgt, batch);
let on_loss_val = host(&on_loss.value())[0];
on_loss.backward();
let on_grads: Vec<Vec<f32>> = on
.params()
.iter()
.map(|p| host(&p.grad().expect("on grad")))
.collect();
(
off_logits,
off_loss_val,
off_grads,
on_logits,
on_loss_val,
on_grads,
)
}
#[test]
fn flash_matches_composed_fp32() {
// fp32: same SDPA math, differs only by reduction order (in-kernel fp32 FMA vs
// cuBLAS, dK/dV atomicAdd order). Tight per-element check, not bit-exact.
run_fp32(1e-3, 2e-2);
}
#[test]
fn flash_matches_composed_bf16() {
// bf16 (T12 composition): bf16 rounding band on the fp32-softmax core; robust
// (mean/p99/scaled-mean) comparison per the repo's bf16 convention.
run_bf16();
}

269
crates/xtrain-model/tests/gqa.rs Normal file
View File

@@ -0,0 +1,269 @@
// T15 GQA correctness gate. Real grouped-query attention (num_kv_heads <
// num_heads): K/V project to num_kv_heads·head_dim and are repeat_kv-broadcast to
// the full head set before the SDPA. This test pins three things:
//
// 1. GQA flash == GQA composed (forward logits + loss + EVERY param grad) — the
// repeat_kv broadcast feeds both SDPA paths unchanged, so they must agree; in
// particular the wk/wv grads (which flow back through repeat_kv's group-sum)
// must match. Parameterised over fp32 (tight) and bf16 (rounding band).
// 2. group==1 (num_kv_heads == n_heads) is BIT-IDENTICAL to the pre-T15 MHA path
// (a model with num_kv_heads explicitly == n_heads vs the default config):
// forward logits + every grad |Δ|=0. The regression guard.
// 3. wk/wv really shrank to [dim, kv_dim] under GQA (shape check).
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
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, 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: &xtrain_tensor::Tensor) -> Vec<f32> {
t.to_dtype(DType::F32)
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec()
}
// A real GQA config: 8 query heads, 2 kv heads → group 4. seq=40 > FA_TILE=32 so
// the flash online-softmax tile path is exercised too.
fn gqa_cfg() -> Config {
let mut cfg = Config::tiny();
cfg.vocab = 16;
cfg.n_layers = 3;
// tiny() is 2 heads; rebuild with 8 query / 2 kv heads keeping head_dim=16.
Config::from_arch(cfg.vocab, 8, cfg.head_dim, cfg.n_layers, cfg.ffn_hidden).with_kv_heads(2)
}
fn ids_targets(cfg: &Config, batch: usize, seq: usize) -> (Vec<Vec<i32>>, Vec<Vec<i32>>) {
let seqs = (0..batch)
.map(|b| {
(0..seq)
.map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32)
.collect()
})
.collect();
let tgts = (0..batch)
.map(|b| {
(0..seq)
.map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
.collect()
})
.collect();
(seqs, tgts)
}
#[allow(clippy::type_complexity)]
fn run_both(
cfg: Config,
dtype: DType,
) -> (Vec<f32>, f32, Vec<Vec<f32>>, Vec<f32>, f32, Vec<Vec<f32>>) {
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let (batch, seq) = (3usize, 40usize);
let (seqs, tgts) = ids_targets(&cfg, batch, seq);
let ids = batched_ids_tensor(&seqs, device);
let tgt = batched_ids_tensor(&tgts, device);
let off = build(cfg, device, dtype, false);
let off_logits = host(&off.forward_batched(&ids, batch).value());
let off_loss = off.loss_batched(&ids, &tgt, batch);
let off_loss_val = host(&off_loss.value())[0];
off_loss.backward();
let off_grads: Vec<Vec<f32>> = off
.params()
.iter()
.map(|p| host(&p.grad().expect("off grad")))
.collect();
let on = build(cfg, device, dtype, true);
let on_logits = host(&on.forward_batched(&ids, batch).value());
let on_loss = on.loss_batched(&ids, &tgt, batch);
let on_loss_val = host(&on_loss.value())[0];
on_loss.backward();
let on_grads: Vec<Vec<f32>> = on
.params()
.iter()
.map(|p| host(&p.grad().expect("on grad")))
.collect();
(
off_logits,
off_loss_val,
off_grads,
on_logits,
on_loss_val,
on_grads,
)
}
// GQA flash vs composed: same SDPA math on the same repeat_kv-broadcast K/V → fp32
// agrees to reduction-order, bf16 to its rounding band.
#[test]
fn gqa_flash_matches_composed_fp32() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
let cfg = gqa_cfg();
assert!(cfg.num_kv_heads < cfg.n_heads, "test must be real GQA");
let (off_l, off_loss, off_g, on_l, on_loss, on_g) = run_both(cfg, DType::F32);
let logit_rel = off_l
.iter()
.zip(&on_l)
.map(|(a, b)| (a - b).abs() / a.abs().max(1e-4))
.fold(0.0f32, f32::max);
let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
println!(
"[GQA F32] flash on/off: loss {off_loss:.6}/{on_loss:.6} (rel {loss_rel:.2e}), \
logits max rel {logit_rel:.2e}"
);
assert!(
logit_rel < 1e-3,
"[GQA F32] logits diverged: {logit_rel:.2e}"
);
assert!(loss_rel < 1e-3, "[GQA F32] loss diverged: {loss_rel:.2e}");
let mut worst = 0.0f32;
for (a_g, b_g) in off_g.iter().zip(&on_g) {
for (a, b) in a_g.iter().zip(b_g) {
worst = worst.max((a - b).abs() / a.abs().max(1e-3));
}
}
println!("[GQA F32] flash on/off grad max rel = {worst:.3e}");
assert!(worst < 2e-2, "[GQA F32] grads diverged: {worst:.3e}");
}
#[test]
fn gqa_flash_matches_composed_bf16() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
let (off_l, off_loss, off_g, on_l, on_loss, on_g) = run_both(gqa_cfg(), DType::BF16);
let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
println!("[GQA BF16] flash on/off: loss {off_loss:.5}/{on_loss:.5} (rel {loss_rel:.3e})");
assert!(loss_rel < 2e-2, "[GQA BF16] loss diverged: {loss_rel:.3e}");
let n = off_l.len();
let mut rels: Vec<f32> = off_l
.iter()
.zip(&on_l)
.map(|(f, b)| (b - f).abs() / f.abs().max(1.0))
.collect();
rels.sort_by(|a, b| a.partial_cmp(b).unwrap());
let mean: f32 = rels.iter().sum::<f32>() / n as f32;
let p99 = rels[(n as f32 * 0.99) as usize];
println!("[GQA BF16] logits: mean rel {mean:.3e}, p99 rel {p99:.3e}");
assert!(
mean < 1e-2,
"[GQA BF16] logits mean rel too high: {mean:.3e}"
);
assert!(p99 < 5e-2, "[GQA BF16] logits p99 rel too high: {p99:.3e}");
let mut worst = 0.0f32;
for (a_g, b_g) in off_g.iter().zip(&on_g) {
let scale = a_g.iter().map(|v| v.abs()).fold(0.0f32, f32::max).max(1e-6);
let mean_err: f32 =
a_g.iter().zip(b_g).map(|(f, b)| (f - b).abs()).sum::<f32>() / a_g.len() as f32 / scale;
worst = worst.max(mean_err);
}
println!("[GQA BF16] grads: worst per-tensor scaled-mean err = {worst:.3e}");
assert!(worst < 3e-2, "[GQA BF16] grads diverged: {worst:.3e}");
}
// REGRESSION GUARD: num_kv_heads == n_heads (group 1) must be BIT-IDENTICAL to the
// pre-T15 MHA path. Build one model with the default config (num_kv_heads ==
// n_heads, the untouched path: repeat_kv not even invoked) and one that explicitly
// sets num_kv_heads = n_heads, then assert forward logits + every grad match to the
// bit. (Same composed path, so this is exact equality, not a tolerance.)
#[test]
fn gqa_group1_bit_identical_to_mha() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let mut base = Config::tiny();
base.vocab = 16;
base.n_layers = 3;
let base = Config::from_arch(base.vocab, 4, base.head_dim, base.n_layers, base.ffn_hidden);
// `explicit` sets num_kv_heads = n_heads (already the default, but exercises the
// with_kv_heads path); they are the same config → must produce identical output.
let explicit = base.with_kv_heads(base.n_heads);
assert_eq!(base.num_kv_heads, explicit.num_kv_heads);
let (batch, seq) = (2usize, 8usize);
let (seqs, tgts) = ids_targets(&base, batch, seq);
let ids = batched_ids_tensor(&seqs, device);
let tgt = batched_ids_tensor(&tgts, device);
let run = |cfg: Config| -> (Vec<f32>, f32, Vec<Vec<f32>>) {
let m = build(cfg, device, DType::F32, false);
let logits = host(&m.forward_batched(&ids, batch).value());
let loss = m.loss_batched(&ids, &tgt, batch);
let loss_v = host(&loss.value())[0];
loss.backward();
let grads = m
.params()
.iter()
.map(|p| host(&p.grad().unwrap()))
.collect();
(logits, loss_v, grads)
};
let (la, sa, ga) = run(base);
let (lb, sb, gb) = run(explicit);
assert_eq!(la, lb, "group-1 logits must be bit-identical to MHA");
assert_eq!(sa, sb, "group-1 loss must be bit-identical to MHA");
for (a, b) in ga.iter().zip(&gb) {
assert_eq!(a, b, "group-1 grad must be bit-identical to MHA");
}
println!("[GQA group1] bit-identical to MHA: logits + loss + all grads |Δ|=0");
}
// Under GQA, wk/wv must be [dim, kv_dim] (= num_kv_heads·head_dim), wq stays [dim,dim].
#[test]
fn gqa_kv_proj_shape() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let cfg = gqa_cfg();
let m = build(cfg, device, DType::F32, false);
let p = m.params();
// params order: embed[0], then block 0 = [attn_norm[1], wq[2], wk[3], wv[4],
// q_norm[5], k_norm[6], wo[7], ...]
let wq = p[2].value().shape().to_vec();
let wk = p[3].value().shape().to_vec();
let wv = p[4].value().shape().to_vec();
assert_eq!(wq, vec![cfg.dim, cfg.dim], "wq must be [dim,dim]");
assert_eq!(wk, vec![cfg.dim, cfg.kv_dim()], "wk must be [dim,kv_dim]");
assert_eq!(wv, vec![cfg.dim, cfg.kv_dim()], "wv must be [dim,kv_dim]");
println!(
"[GQA shapes] wq {:?} wk {:?} wv {:?} (kv_dim {})",
wq,
wk,
wv,
cfg.kv_dim()
);
}

133
crates/xtrain-model/tests/overfit.rs Normal file
View File

@@ -0,0 +1,133 @@
// End-to-end acceptance for the Phase T5 tiny transformer: overfit one fixed
// char-level batch with a hand-written gradient-descent step and assert the loss
// collapses toward 0. This is THE signal that the whole fwd+bwd graph (embedding,
// RMSNorm, RoPE, multi-head attention, SwiGLU, LM head, cross-entropy) is wired
// correctly — a single buggy backward would stall the loss.
//
// The optimizer here is deliberately minimal (`p ← p − lr·grad`); AdamW / LR
// schedule / real data are T6. Gated behind `not(no_cuda)` (runs on dash5).
#![cfg(not(no_cuda))]
use xtrain_autodiff::tape::Var;
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, ids_tensor};
use xtrain_tensor::Device;
// Deterministic LCG fill in [-scale, scale).
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 require_gpu() {
assert!(
device::device_count().expect("device count") > 0,
"no CUDA device"
);
device::set_device(0).unwrap();
}
// One GD step over every parameter: p ← p − lr·grad, then zero the grad.
fn gd_step(params: &[Var], lr: f32) {
for p in params {
if let Some(g) = p.grad() {
let updated = p.value().add(&g.scale(-lr));
p.set_value(updated);
}
p.zero_grad();
}
}
#[test]
fn overfit_tiny_batch() {
require_gpu();
let device = Device::Cuda(0);
// --- Char-level bring-up: tiny embedded text → vocab → (input, target). ---
let text = "hello tiny transformer world";
let mut vocab_chars: Vec<char> = text.chars().collect();
vocab_chars.sort_unstable();
vocab_chars.dedup();
let vocab = vocab_chars.len();
let stoi = |c: char| vocab_chars.iter().position(|&x| x == c).unwrap() as i32;
let tokens: Vec<i32> = text.chars().map(stoi).collect();
// Next-token prediction: input = tokens[..n-1], target = tokens[1..].
let input: Vec<i32> = tokens[..tokens.len() - 1].to_vec();
let target: Vec<i32> = tokens[1..].to_vec();
let ids = ids_tensor(&input, device);
let targets = ids_tensor(&target, device);
// --- Tiny model with small-scale deterministic init. ---
let mut cfg = Config::tiny();
cfg.vocab = vocab;
let mut seed = 1u64;
let model = TinyTransformer::new(cfg, device, |shape| {
seed = seed.wrapping_add(1);
let n: usize = shape.iter().product();
// RMSNorm gammas ([dim]) init to ~1; everything else small random.
if shape.len() == 1 {
fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
} else {
fill(n, seed, 0.08)
}
});
let params = model.params();
println!(
"overfit: vocab={vocab} seq={} params={}",
input.len(),
cfg.num_params()
);
let read_loss = |l: &Var| -> f32 { l.value().to_device(Device::Cpu).as_slice::<f32>()[0] };
let lr = 0.3f32;
let steps = 200;
let start = read_loss(&model.loss(&ids, &targets));
let mut last = start;
for step in 0..steps {
let loss = model.loss(&ids, &targets);
last = read_loss(&loss);
if step % 20 == 0 || step == steps - 1 {
println!("step {step:3}: loss = {last:.6}");
}
loss.backward();
gd_step(&params, lr);
}
println!("overfit: start loss = {start:.6} → final loss = {last:.6} ({steps} steps)");
// A correct fwd+bwd memorises this tiny fixed batch: loss → ~0.
assert!(
last < 0.05,
"overfit failed to drive loss to ~0: start {start:.4} final {last:.4}"
);
assert!(last < start, "loss did not decrease");
// Sanity: greedy argmax should reproduce the target sequence after overfit.
let logits = model.forward(&ids).value().to_device(Device::Cpu);
let lg = logits.as_slice::<f32>();
let mut correct = 0;
for (r, &t) in target.iter().enumerate() {
let row = &lg[r * vocab..(r + 1) * vocab];
let argmax = row
.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
.unwrap()
.0 as i32;
if argmax == t {
correct += 1;
}
}
println!("overfit: greedy match {correct}/{}", target.len());
assert_eq!(correct, target.len() as i32, "did not memorise the batch");
}

198
crates/xtrain-model/tests/parity.py Normal file
View File

@@ -0,0 +1,198 @@
#!/usr/bin/env python3
"""PyTorch parity check for the xtrain tiny transformer (Phase T5).
Loads the weights/ids dumped by tests/parity_dump.rs, rebuilds the IDENTICAL
model in PyTorch (same x@W convention, same RoPE rotate_half + position=row,
same RMSNorm, SwiGLU, causal mask, per-head SDPA), runs forward + one backward,
and compares the forward logits and every parameter's gradient against the Rust
values within a relative tolerance.
Usage: python3 parity.py /tmp/xtrain_parity
"""
import sys
import os
import math
import torch
DIR = sys.argv[1] if len(sys.argv) > 1 else "/tmp/xtrain_parity"
def read_vec(name):
path = os.path.join(DIR, name)
shape = None
vals = []
with open(path) as f:
for line in f:
line = line.strip()
if line.startswith("# shape"):
shape = [int(x) for x in line.split()[2].split(",") if x]
elif line:
vals.append(float(line))
t = torch.tensor(vals, dtype=torch.float64)
if shape:
t = t.reshape(shape)
return t
def read_cfg():
cfg = {}
with open(os.path.join(DIR, "config.txt")) as f:
for line in f:
k, v = line.split()
cfg[k] = v
return cfg
def read_ids(name):
with open(os.path.join(DIR, name)) as f:
return [int(x) for x in f.read().split()]
cfg = read_cfg()
DIM = int(cfg["dim"])
NL = int(cfg["n_layers"])
NH = int(cfg["n_heads"])
# GQA (T15): num_kv_heads <= n_heads; each kv head shared by group query heads.
# Default to NH (MHA) for fixtures dumped before the field existed.
NKV = int(cfg.get("num_kv_heads", str(NH)))
GROUP = NH // NKV
HD = int(cfg["head_dim"])
EPS = float(cfg["eps"])
THETA = float(cfg["rope_theta"])
# Batched: B sequences of length SEQ, flattened sequence-major to [B*SEQ] ids.
B = int(cfg.get("batch", "1"))
SEQ = int(cfg["seq"])
ids = read_ids("ids.txt")
targets = read_ids("targets.txt")
assert len(ids) == B * SEQ, f"ids {len(ids)} != B*SEQ {B*SEQ}"
# Load params as leaf tensors requiring grad (float64 for a clean reference).
P = {}
def load(name):
t = read_vec(f"w_{name}.txt").clone().requires_grad_(True)
P[name] = t
return t
def rms_norm(x, gamma):
# y = x / sqrt(mean(x^2)+eps) * gamma (no mean subtraction)
ms = x.pow(2).mean(dim=-1, keepdim=True)
return x * torch.rsqrt(ms + EPS) * gamma
def rope(x): # x: [B*SEQ, nh, hd], position = (row % SEQ) — resets per sequence
half = HD // 2
out = torch.empty_like(x)
i = torch.arange(half, dtype=torch.float64)
freq = THETA ** (-(2.0 * i) / HD) # [half]
# Position within each sequence: rows 0..SEQ for seq 0, 0..SEQ for seq 1, ...
pos = (torch.arange(B * SEQ, dtype=torch.float64) % SEQ).reshape(B * SEQ, 1)
ang = pos * freq # [B*SEQ, half]
c = torch.cos(ang).reshape(B * SEQ, 1, half)
s = torch.sin(ang).reshape(B * SEQ, 1, half)
x0 = x[..., :half]
x1 = x[..., half:]
out[..., :half] = x0 * c - x1 * s
out[..., half:] = x1 * c + x0 * s
return out
emb = load("embed")
final_norm = load("final_norm")
lm_head = load("lm_head")
layers = []
for l in range(NL):
layers.append({p: load(f"l{l}_{p}") for p in
["attn_norm", "wq", "wk", "wv", "q_norm", "k_norm", "wo",
"ffn_norm", "w_gate", "w_up", "w_down"]})
idx = torch.tensor(ids, dtype=torch.long)
# Per-sequence causal mask (broadcast over the batch); NO cross-sequence attention.
mask = torch.triu(torch.full((SEQ, SEQ), -1.0e9, dtype=torch.float64), diagonal=1)
h = emb[idx] # [B*SEQ, dim] (everything stays flattened, matching the Rust path)
for L in layers:
# Attention
x = rms_norm(h, L["attn_norm"])
q = (x @ L["wq"]).reshape(B * SEQ, NH, HD)
# GQA: K/V project to NKV heads, then repeat each kv head GROUP times to NH.
k = (x @ L["wk"]).reshape(B * SEQ, NKV, HD)
v = (x @ L["wv"]).reshape(B * SEQ, NKV, HD)
# Per-head QK-norm (Qwen3-style), before RoPE.
q = rms_norm(q, L["q_norm"])
k = rms_norm(k, L["k_norm"])
q = rope(q) # [B*SEQ, nh, hd]
k = rope(k) # [B*SEQ, nkv, hd]
# Reshape to [B, *, SEQ, HD]; broadcast kv heads to NH (repeat_interleave along
# the head axis: kv head kvh → query heads [kvh*GROUP, (kvh+1)*GROUP), matching
# xtrain repeat_kv + xserv repeat_kv).
q = q.reshape(B, SEQ, NH, HD).transpose(1, 2) # [B, nh, seq, hd]
k = k.reshape(B, SEQ, NKV, HD).transpose(1, 2) # [B, nkv, seq, hd]
v = v.reshape(B, SEQ, NKV, HD).transpose(1, 2)
if GROUP > 1:
k = k.repeat_interleave(GROUP, dim=1) # [B, nh, seq, hd]
v = v.repeat_interleave(GROUP, dim=1)
scale = 1.0 / math.sqrt(HD)
scores = (q @ k.transpose(-1, -2)) * scale + mask # [B, nh, seq, seq]
probs = torch.softmax(scores, dim=-1)
out = probs @ v # [B, nh, seq, hd]
out = out.transpose(1, 2).reshape(B * SEQ, DIM) # [B*SEQ, dim]
attn = out @ L["wo"]
h = h + attn
# MLP
x = rms_norm(h, L["ffn_norm"])
gate = x @ L["w_gate"]
up = x @ L["w_up"]
act = torch.nn.functional.silu(gate) * up
mlp = act @ L["w_down"]
h = h + mlp
h = rms_norm(h, final_norm)
logits = h @ lm_head # [B*SEQ, vocab]
loss = torch.nn.functional.cross_entropy(
logits, torch.tensor(targets, dtype=torch.long), reduction="mean")
loss_val = loss.item()
loss.backward()
# ---- Compare ----
def relerr(a, b):
a = a.double()
b = b.double()
denom = b.abs().clamp(min=1e-6)
return ((a - b).abs() / denom).max().item()
ref_logits = read_vec("logits.txt")
ref_loss = read_vec("loss.txt").item()
print(f"loss: rust={ref_loss:.6e} torch={loss_val:.6e} "
f"relerr={abs(loss_val-ref_loss)/max(abs(ref_loss),1e-6):.2e}")
le = relerr(logits.detach(), ref_logits)
print(f"logits: max relerr = {le:.2e}")
RTOL = 2e-2
worst = le
worst_name = "logits"
fails = []
if le > RTOL:
fails.append(("logits", le))
for name, t in P.items():
ref_g = read_vec(f"g_{name}.txt")
ge = relerr(t.grad, ref_g)
if ge > worst:
worst, worst_name = ge, f"grad[{name}]"
if ge > RTOL:
fails.append((f"grad[{name}]", ge))
print(f"params checked: {len(P)} worst = {worst_name} @ {worst:.2e} (rtol={RTOL})")
if fails:
print("FAIL:")
for n, e in fails:
print(f" {n}: relerr={e:.3e}")
sys.exit(1)
print("PARITY OK: forward logits + all param grads within rtol")

190
crates/xtrain-model/tests/parity_dump.rs Normal file
View File

@@ -0,0 +1,190 @@
// PyTorch parity, step 1 of 2: dump the Rust tiny-transformer's exact weights,
// inputs, forward logits, loss, and per-parameter gradients (after one backward)
// to a directory, so an equivalent PyTorch model (tests/parity.py) can be built
// from the SAME weights and the forward + grads compared within rtol.
//
// Run: XTRAIN_PARITY_DIR=/tmp/xtrain_parity cargo test -p xtrain-model \
// --test parity_dump -- --nocapture --ignored
// then: python3 crates/xtrain-model/tests/parity.py /tmp/xtrain_parity
//
// Marked #[ignore] (it's a fixture generator, not a pass/fail assertion) and
// gated #![cfg(not(no_cuda))].
#![cfg(not(no_cuda))]
use std::fs;
use std::io::Write;
use std::path::PathBuf;
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, ids_tensor, param_to_host};
use xtrain_tensor::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 write_vec(dir: &PathBuf, name: &str, data: &[f32], shape: &[usize]) {
let mut f = fs::File::create(dir.join(name)).unwrap();
let shape_str: Vec<String> = shape.iter().map(|d| d.to_string()).collect();
writeln!(f, "# shape {}", shape_str.join(",")).unwrap();
for v in data {
writeln!(f, "{v:.8e}").unwrap();
}
}
#[test]
#[ignore = "fixture generator for PyTorch parity; run with --ignored"]
fn dump_for_parity() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let dir = PathBuf::from(
std::env::var("XTRAIN_PARITY_DIR").unwrap_or_else(|_| "/tmp/xtrain_parity".to_string()),
);
fs::create_dir_all(&dir).unwrap();
// Fixed config + ids (independent of any text, for reproducibility). B>1 so
// the batched forward is exercised: 2 sequences of length 4, flattened
// sequence-major to [B*S]=8 ids. Per-sequence RoPE position (resets at the
// sequence boundary) + per-sequence causal masking (no cross-sequence
// attention) are both checked against PyTorch.
// Default: tiny MHA (2 heads). With XTRAIN_PARITY_KV_HEADS=k set, dump a real
// GQA config (8 query heads / k kv heads) so parity.py checks GQA at B>1 — the
// kv-projection shapes + the repeat_kv group-sum backward against PyTorch.
let mut cfg = Config::tiny();
cfg.vocab = 12;
if let Ok(kv) = std::env::var("XTRAIN_PARITY_KV_HEADS") {
let kv: usize = kv.parse().expect("XTRAIN_PARITY_KV_HEADS");
cfg = Config::from_arch(cfg.vocab, 8, cfg.head_dim, cfg.n_layers, cfg.ffn_hidden)
.with_kv_heads(kv);
println!(
"parity: GQA config (n_heads {} kv_heads {})",
cfg.n_heads, cfg.num_kv_heads
);
}
let batch = 2usize;
let seq = 4usize;
let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6]; // [B*S], sequence-major
let targets: Vec<i32> = vec![1, 4, 1, 5, 9, 2, 6, 0];
// Same deterministic init as the overfit test.
let mut seed = 1u64;
let mut 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.08)
}
});
// T14: with XTRAIN_PARITY_FLASH set, dump from the fused flash-attention path.
// flash is the SAME SDPA math, so the SAME parity.py PyTorch oracle is the
// reference for both paths — running this once per path checks flash against
// PyTorch at B>1 (forward logits + every parameter grad).
if std::env::var("XTRAIN_PARITY_FLASH").is_ok() {
model = model.with_flash(true);
println!("parity: FLASH attention path");
}
// config + ids
{
let mut f = fs::File::create(dir.join("config.txt")).unwrap();
writeln!(f, "vocab {}", cfg.vocab).unwrap();
writeln!(f, "dim {}", cfg.dim).unwrap();
writeln!(f, "n_layers {}", cfg.n_layers).unwrap();
writeln!(f, "n_heads {}", cfg.n_heads).unwrap();
writeln!(f, "num_kv_heads {}", cfg.num_kv_heads).unwrap();
writeln!(f, "head_dim {}", cfg.head_dim).unwrap();
writeln!(f, "ffn_hidden {}", cfg.ffn_hidden).unwrap();
writeln!(f, "eps {:e}", cfg.eps).unwrap();
writeln!(f, "rope_theta {:e}", cfg.rope_theta).unwrap();
writeln!(f, "batch {batch}").unwrap();
writeln!(f, "seq {seq}").unwrap();
}
{
let mut f = fs::File::create(dir.join("ids.txt")).unwrap();
for v in &ids {
writeln!(f, "{v}").unwrap();
}
let mut f = fs::File::create(dir.join("targets.txt")).unwrap();
for v in &targets {
writeln!(f, "{v}").unwrap();
}
}
// Stable param order, named to match parity.py.
let names = param_names(&cfg);
let params = model.params();
assert_eq!(names.len(), params.len(), "param name/count mismatch");
for (name, p) in names.iter().zip(&params) {
let shape = p.value().shape().to_vec();
write_vec(&dir, &format!("w_{name}.txt"), &param_to_host(p), &shape);
}
// Batched forward logits + loss (B sequences as one forward), then backward
// → per-param grads.
let ids_t = ids_tensor(&ids, device);
let targets_t = ids_tensor(&targets, device);
let logits = model.forward_batched(&ids_t, batch);
write_vec(
&dir,
"logits.txt",
&param_to_host(&logits),
logits.value().shape(),
);
let loss = model.loss_batched(&ids_t, &targets_t, batch);
let loss_val = param_to_host(&loss)[0];
{
let mut f = fs::File::create(dir.join("loss.txt")).unwrap();
writeln!(f, "{loss_val:.8e}").unwrap();
}
loss.backward();
for (name, p) in names.iter().zip(&params) {
let g = p.grad().expect("param has no grad");
let gh = g.to_device(Device::Cpu);
write_vec(
&dir,
&format!("g_{name}.txt"),
gh.as_slice::<f32>(),
g.shape(),
);
}
println!("parity: dumped to {} (loss={loss_val:.6e})", dir.display());
}
fn param_names(cfg: &Config) -> Vec<String> {
let mut names = vec!["embed".to_string()];
for l in 0..cfg.n_layers {
for p in [
"attn_norm",
"wq",
"wk",
"wv",
"q_norm",
"k_norm",
"wo",
"ffn_norm",
"w_gate",
"w_up",
"w_down",
] {
names.push(format!("l{l}_{p}"));
}
}
names.push("final_norm".to_string());
names.push("lm_head".to_string());
names
}

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

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

161
crates/xtrain-model/tests/recompute.rs Normal file
View File

@@ -0,0 +1,161 @@
// T13 activation-recomputation correctness gate (the HARD gate).
//
// Gradient checkpointing is mathematically EXACT: the backward re-runs the same
// `segment_fn` from the same saved input and the same (unchanged) parameter
// values, so the recomputed activations equal the originals and the recovered
// grads equal the non-checkpointed grads — checkpointing trades compute for
// memory, never correctness. This test makes that a closed loop on-GPU:
//
// build two identical models (same init), one with `--recompute` on, one off,
// run the SAME batched loss + backward on both, and assert
// 1. the forward logits match (recompute doesn't touch forward output)
// 2. the loss matches
// 3. EVERY parameter's grad matches within a tight fp tolerance.
//
// Composition is covered by parameterising over fp32 AND bf16 (T12): the
// recompute path is the unchanged block forward, so it runs the same dtype path.
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
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, 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.08)
}
});
m.with_compute_dtype(dtype).with_recompute(recompute)
}
/// Upcast to fp32 then read to host — logits are bf16 in bf16 mode (grads are
/// always fp32 master, but this is uniform and harmless for fp32 tensors).
fn host(t: &xtrain_tensor::Tensor) -> Vec<f32> {
t.to_dtype(DType::F32)
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec()
}
fn run(dtype: DType, logit_tol: f32, grad_tol: f32) {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
// A few layers so checkpointing actually wraps multiple blocks.
let mut cfg = Config::tiny();
cfg.vocab = 16;
cfg.n_layers = 4;
let batch = 3usize;
let seq = 6usize;
let seqs: Vec<Vec<i32>> = (0..batch)
.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()
})
.collect();
let ids = batched_ids_tensor(&seqs, device);
let tgt = batched_ids_tensor(&tgts, device);
// --- recompute OFF (reference) ---
let off = build(cfg, device, dtype, false);
let off_logits = host(&off.forward_batched(&ids, batch).value());
let off_loss = off.loss_batched(&ids, &tgt, batch);
let off_loss_val = host(&off_loss.value())[0];
off_loss.backward();
let off_grads: Vec<Vec<f32>> = off
.params()
.iter()
.map(|p| host(&p.grad().expect("off grad")))
.collect();
// --- recompute ON ---
let on = build(cfg, device, dtype, true);
let on_logits = host(&on.forward_batched(&ids, batch).value());
let on_loss = on.loss_batched(&ids, &tgt, batch);
let on_loss_val = host(&on_loss.value())[0];
on_loss.backward();
let on_grads: Vec<Vec<f32>> = on
.params()
.iter()
.map(|p| host(&p.grad().expect("on grad")))
.collect();
// 1. Forward logits — recompute must not change the forward output.
let logit_rel = off_logits
.iter()
.zip(&on_logits)
.map(|(a, b)| (a - b).abs() / a.abs().max(1e-4))
.fold(0.0f32, f32::max);
// 2. Loss.
let loss_rel = (off_loss_val - on_loss_val).abs() / off_loss_val.abs().max(1e-4);
println!(
"[{dtype:?}] recompute on/off: loss {off_loss_val:.6}/{on_loss_val:.6} (rel {loss_rel:.2e}), \
logits max rel {logit_rel:.2e}"
);
assert!(
logit_rel < logit_tol,
"[{dtype:?}] logits diverged: {logit_rel:.2e}"
);
assert!(
loss_rel < logit_tol,
"[{dtype:?}] loss diverged: {loss_rel:.2e}"
);
// 3. Every parameter grad — the load-bearing gate.
let mut max_grad_rel = 0.0f32;
for (off_g, on_g) in off_grads.iter().zip(&on_grads) {
for (a, b) in off_g.iter().zip(on_g) {
let rel = (a - b).abs() / a.abs().max(1e-3);
max_grad_rel = max_grad_rel.max(rel);
}
}
println!("[{dtype:?}] recompute on/off: grad max rel err = {max_grad_rel:.3e}");
assert!(
max_grad_rel < grad_tol,
"[{dtype:?}] recompute grads diverged from non-recompute: {max_grad_rel:.3e}"
);
}
#[test]
fn recompute_matches_non_recompute_fp32() {
// fp32: recompute runs the identical deterministic kernels → grads match to
// (near) bit-exact; allow a hair for any nondeterministic GPU reduction.
run(DType::F32, 1e-5, 1e-4);
}
#[test]
fn recompute_matches_non_recompute_bf16() {
// bf16 (T12 composition): same bf16 path on recompute. The recompute is still
// exact w.r.t. the bf16 forward, so on/off match tightly (looser tol only for
// bf16 rounding, not for any recompute discrepancy).
run(DType::BF16, 5e-3, 5e-3);
}

10
crates/xtrain-optim/Cargo.toml Normal file
View File

@@ -0,0 +1,10 @@
[package]
name = "xtrain-optim"
version.workspace = true
edition.workspace = true
[dependencies]
xtrain-tensor = { path = "../xtrain-tensor" }
xtrain-autodiff = { path = "../xtrain-autodiff" }
# GPU AdamW (Phase T7) launches kernels + syncs the device directly.
xtrain-cuda = { path = "../xtrain-cuda" }

26
crates/xtrain-optim/build.rs Normal file
View File

@@ -0,0 +1,26 @@
use std::env;
use std::path::Path;
use std::process::Command;
// Per-crate convention (see the other crates): the AdamW *math* is host-only and
// always compiles, but `AdamW::step(&[Var])` round-trips parameter values/grads
// through GPU tensors, so that call site is gated behind `not(no_cuda)`. cfg does
// not propagate across crates, so this crate re-detects nvcc. No CUDA is compiled.
fn main() {
println!("cargo:rustc-check-cfg=cfg(no_cuda)");
let cuda_path = env::var("CUDA_HOME")
.or_else(|_| env::var("CUDA_PATH"))
.unwrap_or_else(|_| "/usr/local/cuda".to_string());
if !nvcc_available(&cuda_path) {
println!("cargo:rustc-cfg=no_cuda");
}
}
fn nvcc_available(cuda_path: &str) -> bool {
if Command::new("nvcc").arg("--version").output().is_ok() {
return true;
}
Path::new(&format!("{cuda_path}/bin/nvcc")).exists()
}

247
crates/xtrain-optim/src/lib.rs Normal file
View File

@@ -0,0 +1,247 @@
//! Hand-written AdamW optimizer (Phase T6).
//!
//! AdamW = Adam with **decoupled** weight decay (Loshchilov & Hutter, 2019): the
//! weight-decay term is applied directly to the parameter, NOT folded into the
//! gradient (so it does not interact with the adaptive `v` denominator). This
//! matches `torch.optim.AdamW`.
//!
//! Update for parameter `θ` at step `t` (1-indexed), with gradient `g`:
//! ```text
//! m ← β1·m + (1−β1)·g
//! v ← β2·v + (1−β2)·g²
//! m̂ ← m / (1 − β1ᵗ) (bias correction)
//! v̂ ← v / (1 − β2ᵗ)
//! θ ← θ − lr·( m̂ / (√v̂ + ε) + wd·θ )
//! ```
//! The `lr·wd·θ` term is the decoupled decay. Note PyTorch applies decay as
//! `θ ← θ·(1 − lr·wd)` then the Adam step; both are algebraically the same
//! first-order update — we fold decay into the single subtraction above, which
//! is what PyTorch's default (`maximize=False`, no `amsgrad`) computes.
//!
//! The math operates on flat host `f32` buffers ([`AdamW::step_host`]) so it is
//! unit-testable on a GPU-less host; [`AdamW::step`] is a thin wrapper that
//! round-trips each parameter's value/grad through the GPU tensor and is gated
//! behind `not(no_cuda)`.
/// Per-parameter optimizer state: the first (`m`) and second (`v`) moment
/// estimates, one f32 per element, kept flat (matching the parameter layout).
struct ParamState {
m: Vec<f32>,
v: Vec<f32>,
}
/// Decoupled-weight-decay Adam. One instance owns the moment state for a fixed
/// list of parameters, keyed by their index in the slice passed to `step`
/// (the model's stable `params()` order).
pub struct AdamW {
pub lr: f32,
beta1: f32,
beta2: f32,
eps: f32,
weight_decay: f32,
/// Global step count (shared across all params for bias correction).
t: u64,
/// Lazily sized to the parameter list on the first `step`.
state: Vec<ParamState>,
}
impl AdamW {
/// PyTorch-default hyperparameters except `lr`/`weight_decay`, which you set
/// (β1=0.9, β2=0.999, ε=1e-8).
pub fn new(lr: f32, weight_decay: f32) -> Self {
Self::with_betas(lr, weight_decay, 0.9, 0.999, 1e-8)
}
pub fn with_betas(lr: f32, weight_decay: f32, beta1: f32, beta2: f32, eps: f32) -> Self {
Self {
lr,
beta1,
beta2,
eps,
weight_decay,
t: 0,
state: Vec::new(),
}
}
/// Current global step (number of `step` calls so far).
pub fn step_count(&self) -> u64 {
self.t
}
/// Pure-host AdamW step over flat parameter/gradient buffers. `params[i]` is
/// updated in place using `grads[i]`; both are the i-th parameter's elements
/// in the model's stable order. Lazily allocates moment state on first call.
///
/// This is the testable core — no GPU, no autograd. `lr` is passed per call
/// so a schedule can vary it each step.
pub fn step_host(&mut self, lr: f32, params: &mut [Vec<f32>], grads: &[Vec<f32>]) {
assert_eq!(params.len(), grads.len(), "param/grad count mismatch");
if self.state.is_empty() {
self.state = params
.iter()
.map(|p| ParamState {
m: vec![0.0; p.len()],
v: vec![0.0; p.len()],
})
.collect();
}
assert_eq!(self.state.len(), params.len(), "param count changed");
self.t += 1;
let bc1 = 1.0 - self.beta1.powi(self.t as i32);
let bc2 = 1.0 - self.beta2.powi(self.t as i32);
for (i, (p, g)) in params.iter_mut().zip(grads).enumerate() {
assert_eq!(p.len(), g.len(), "param/grad len mismatch at {i}");
let st = &mut self.state[i];
for j in 0..p.len() {
let gj = g[j];
st.m[j] = self.beta1 * st.m[j] + (1.0 - self.beta1) * gj;
st.v[j] = self.beta2 * st.v[j] + (1.0 - self.beta2) * gj * gj;
let mhat = st.m[j] / bc1;
let vhat = st.v[j] / bc2;
// Decoupled weight decay: decay term uses the *current* param,
// matching PyTorch's `p ← p − lr·wd·p` applied alongside the step.
p[j] -= lr * (mhat / (vhat.sqrt() + self.eps) + self.weight_decay * p[j]);
}
}
}
}
#[cfg(not(no_cuda))]
mod gpu {
use super::AdamW;
use xtrain_autodiff::tape::Var;
use xtrain_tensor::{DType, Device, Tensor};
impl AdamW {
/// Apply one AdamW step to every parameter `Var`, using `lr` for this step
/// (so an LR schedule can vary it). Pulls each param's value and `.grad()`
/// to the host, runs [`AdamW::step_host`], and writes the updated value
/// back with `set_value`. A param with no grad is fed a zero grad, so the
/// Adam term vanishes and only decoupled weight decay applies (the model's
/// params all receive grads each step, so this is just a safety default).
///
/// Does NOT zero grads — the caller does that (matching the GD-step
/// template in the T5 overfit test).
///
/// This is the host-roundtrip reference path; training uses
/// [`GpuAdamW`] (kernel, m/v on device). Both are checked against the
/// torch parity in tests.
pub fn step(&mut self, lr: f32, params: &[Var]) {
let device = params[0].value().device();
let shapes: Vec<Vec<usize>> =
params.iter().map(|p| p.value().shape().to_vec()).collect();
let mut host_params: Vec<Vec<f32>> = params
.iter()
.map(|p| p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec())
.collect();
let host_grads: Vec<Vec<f32>> = params
.iter()
.zip(&host_params)
.map(|(p, hp)| match p.grad() {
Some(g) => g.to_device(Device::Cpu).as_slice::<f32>().to_vec(),
None => vec![0.0; hp.len()], // no grad → no update this step
})
.collect();
self.step_host(lr, &mut host_params, &host_grads);
for ((p, data), shape) in params.iter().zip(&host_params).zip(&shapes) {
let t = Tensor::from_slice(data, shape).to_device(device);
p.set_value(t);
}
}
}
/// GPU AdamW (Phase T7): the optimizer state (m/v moments) lives on the device
/// as one tensor pair per parameter, and the update runs as a CUDA kernel that
/// reads each param's `.grad()` and rewrites the param buffer in place — no
/// per-step GPU↔host roundtrip of params/grads. Same math as
/// [`AdamW::step_host`] (the parity reference).
pub struct GpuAdamW {
beta1: f32,
beta2: f32,
eps: f32,
weight_decay: f32,
t: u64,
/// Per-parameter (m, v) device buffers, sized lazily on first step.
state: Vec<(Tensor, Tensor)>,
}
impl GpuAdamW {
/// PyTorch-default betas/eps; you set lr (per-step) + weight decay.
pub fn new(weight_decay: f32) -> Self {
Self {
beta1: 0.9,
beta2: 0.999,
eps: 1e-8,
weight_decay,
t: 0,
state: Vec::new(),
}
}
pub fn step_count(&self) -> u64 {
self.t
}
/// One in-place AdamW step over every parameter `Var` at learning rate
/// `lr`. Updates the param value buffer and the device m/v state via the
/// `adamw_step_f32` kernel. Params are mutated in place, so the leaf `Var`
/// identities stay stable across steps (no `set_value`). Does NOT zero
/// grads — the caller does. A param without a grad is skipped this step.
pub fn step(&mut self, lr: f32, params: &[Var]) {
let device = params[0].value().device();
if self.state.is_empty() {
self.state = params
.iter()
.map(|p| {
let shape = p.value().shape().to_vec();
(
Tensor::zeros(&shape, DType::F32, device),
Tensor::zeros(&shape, DType::F32, device),
)
})
.collect();
}
assert_eq!(self.state.len(), params.len(), "param count changed");
self.t += 1;
let bc1 = 1.0 - self.beta1.powi(self.t as i32);
let bc2 = 1.0 - self.beta2.powi(self.t as i32);
for (p, (m, v)) in params.iter().zip(&self.state) {
let g = match p.grad() {
Some(g) => g,
None => continue,
};
let pv = p.value();
let n = pv.numel() as i32;
unsafe {
xtrain_cuda::ffi::launch_adamw_step_f32(
pv.data_ptr() as *mut f32,
g.data_ptr() as *const f32,
m.data_ptr() as *mut f32,
v.data_ptr() as *mut f32,
lr,
self.beta1,
self.beta2,
self.eps,
self.weight_decay,
bc1,
bc2,
n,
std::ptr::null_mut(),
);
}
}
xtrain_cuda::device::synchronize().expect("adamw step sync failed");
}
}
}
#[cfg(not(no_cuda))]
pub use gpu::GpuAdamW;

76
crates/xtrain-optim/tests/adamw_gpu.rs Normal file
View File

@@ -0,0 +1,76 @@
// GPU AdamW parity (Phase T7): the device-side AdamW kernel (m/v on device, no
// host roundtrip) must produce the same update as the host reference
// `AdamW::step_host` given identical params + grads across several steps with a
// varying lr. This is the new correctness gate for the GPU optimizer; the host
// path itself is already pinned to PyTorch by xtrain-train's adamw_parity test.
//
// Gated #![cfg(not(no_cuda))] (runs on dash5; needs a GPU to link + launch).
#![cfg(not(no_cuda))]
use xtrain_autodiff::tape::Var;
use xtrain_cuda::device;
use xtrain_optim::{AdamW, GpuAdamW};
use xtrain_tensor::{Device, Tensor};
fn grad(step: usize, idx: usize, j: usize) -> f32 {
let s = (step * 13 + idx * 7 + j * 3) as f32;
(s * 0.123).sin() * 0.5
}
#[test]
fn gpu_adamw_matches_host() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let dev = Device::Cuda(0);
let wd = 0.1f32;
// Two params of different sizes (exercises per-param device state).
let shapes: Vec<Vec<usize>> = vec![vec![2, 2], vec![3]];
let init: Vec<Vec<f32>> = vec![vec![0.5, -1.0, 2.0, 0.0], vec![1.5, -0.25, 0.75]];
// GPU side: leaf Vars on device.
let params: Vec<Var> = init
.iter()
.zip(&shapes)
.map(|(d, s)| Var::leaf(Tensor::from_slice(d, s).to_device(dev)))
.collect();
let mut gpu_opt = GpuAdamW::new(wd);
// Host reference.
let mut host_params = init.clone();
let mut host_opt = AdamW::new(0.0, wd);
for step in 0..15 {
let lr = 0.01 + 0.001 * step as f32; // varying lr
let grads: Vec<Vec<f32>> = shapes
.iter()
.enumerate()
.map(|(idx, s)| {
let n: usize = s.iter().product();
(0..n).map(|j| grad(step, idx, j)).collect()
})
.collect();
// Push grads onto the GPU Vars, run the device step, then clear.
for (p, (g, s)) in params.iter().zip(grads.iter().zip(&shapes)) {
p.zero_grad();
Var::push_grad(p, Tensor::from_slice(g, s).to_device(dev));
}
gpu_opt.step(lr, &params);
for p in &params {
p.zero_grad();
}
host_opt.step_host(lr, &mut host_params, &grads);
}
let mut max_err = 0.0f32;
for (p, hp) in params.iter().zip(&host_params) {
let got = p.value().to_device(Device::Cpu).as_slice::<f32>().to_vec();
for (a, b) in got.iter().zip(hp) {
max_err = max_err.max((a - b).abs());
}
}
println!("gpu vs host AdamW: max abs err = {max_err:.3e}");
assert!(max_err < 1e-6, "GPU AdamW diverged from host: {max_err:e}");
}

99
crates/xtrain-optim/tests/adamw_host.rs Normal file
View File

@@ -0,0 +1,99 @@
// Host-only unit test for the AdamW *math* (no GPU). Verifies the update against
// an independent, hand-rolled reference implementation of the same recurrence for
// several steps with non-trivial weight decay — catching bias-correction and
// decoupled-decay mistakes. The rigorous vs-PyTorch parity (end-to-end on a real
// model) lives in xtrain-train; this is the fast local guard on the formula.
use xtrain_optim::AdamW;
// Independent reference: the textbook AdamW recurrence, kept separate from the
// implementation so a shared bug can't hide.
struct RefAdamW {
b1: f32,
b2: f32,
eps: f32,
wd: f32,
t: i32,
m: Vec<f32>,
v: Vec<f32>,
}
impl RefAdamW {
fn new(n: usize, wd: f32) -> Self {
Self {
b1: 0.9,
b2: 0.999,
eps: 1e-8,
wd,
t: 0,
m: vec![0.0; n],
v: vec![0.0; n],
}
}
fn step(&mut self, lr: f32, p: &mut [f32], g: &[f32]) {
self.t += 1;
let bc1 = 1.0 - self.b1.powi(self.t);
let bc2 = 1.0 - self.b2.powi(self.t);
for i in 0..p.len() {
self.m[i] = self.b1 * self.m[i] + (1.0 - self.b1) * g[i];
self.v[i] = self.b2 * self.v[i] + (1.0 - self.b2) * g[i] * g[i];
let mhat = self.m[i] / bc1;
let vhat = self.v[i] / bc2;
p[i] -= lr * (mhat / (vhat.sqrt() + self.eps) + self.wd * p[i]);
}
}
}
#[test]
fn adamw_matches_reference_recurrence() {
let lr = 0.01;
let wd = 0.1;
let mut opt = AdamW::new(lr, wd);
// Two parameters of different sizes (exercises per-param state keying).
let mut p_impl = vec![vec![0.5f32, -1.0, 2.0, 0.0], vec![1.5f32, -0.25]];
let mut p_ref = p_impl.clone();
let mut r0 = RefAdamW::new(4, wd);
let mut r1 = RefAdamW::new(2, wd);
// Deterministic pseudo-grads that change every step.
let grad = |step: usize, idx: usize, j: usize| -> f32 {
let s = (step * 13 + idx * 7 + j * 3) as f32;
(s * 0.123).sin() * 0.5
};
for step in 0..20 {
let grads = vec![
(0..4).map(|j| grad(step, 0, j)).collect::<Vec<_>>(),
(0..2).map(|j| grad(step, 1, j)).collect::<Vec<_>>(),
];
opt.step_host(lr, &mut p_impl, &grads);
r0.step(lr, &mut p_ref[0], &grads[0]);
r1.step(lr, &mut p_ref[1], &grads[1]);
}
assert_eq!(opt.step_count(), 20);
for (pi, pr) in p_impl.iter().zip(&p_ref) {
for (a, b) in pi.iter().zip(pr) {
assert!((a - b).abs() < 1e-6, "impl {a} != ref {b}");
}
}
}
#[test]
fn zero_grad_only_decays() {
// With g=0 and wd>0, the step must reduce to pure decoupled decay:
// θ ← θ − lr·wd·θ (Adam term is 0/eps = 0).
let lr = 0.1;
let wd = 0.5;
let mut opt = AdamW::new(lr, wd);
let mut p = vec![vec![2.0f32]];
let g = vec![vec![0.0f32]];
opt.step_host(lr, &mut p, &g);
let expected = 2.0 - lr * wd * 2.0;
assert!(
(p[0][0] - expected).abs() < 1e-6,
"{} != {expected}",
p[0][0]
);
}

22
crates/xtrain-tensor/src/dtype.rs
View File

@@ -1,12 +1,16 @@
//! Tensor data types. //! Tensor data types.
//! //!
//! T2 only needs `F32`, but the enum + `TensorDType` trait are structured so //! T2 only needs `F32`; `BF16` was added in T12 for mixed-precision training
//! half-precision types (F16/BF16) can be added later (T7 mixed precision) //! (bf16 linears / activations, fp32 master weights — see
//! without touching call sites. //! `docs/11-bf16-mixed-precision.md`). The enum + `TensorDType` trait keep call
//! sites dtype-polymorphic.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DType { pub enum DType {
F32, F32,
/// bfloat16: 1 sign / 8 exponent / 7 mantissa. Same exponent range as f32
/// (so no loss scaling needed), ~2-3 decimal digits. The T12 AMP compute type.
BF16,
/// 32-bit signed integers. Used for cross-entropy targets (token ids). /// 32-bit signed integers. Used for cross-entropy targets (token ids).
I32, I32,
} }
@@ -15,6 +19,7 @@ impl DType {
pub fn size_bytes(self) -> usize { pub fn size_bytes(self) -> usize {
match self { match self {
DType::F32 => 4, DType::F32 => 4,
DType::BF16 => 2,
DType::I32 => 4, DType::I32 => 4,
} }
} }
@@ -22,6 +27,7 @@ impl DType {
pub fn name(self) -> &'static str { pub fn name(self) -> &'static str {
match self { match self {
DType::F32 => "f32", DType::F32 => "f32",
DType::BF16 => "bf16",
DType::I32 => "i32", DType::I32 => "i32",
} }
} }
@@ -50,6 +56,16 @@ impl TensorDType for f32 {
} }
} }
impl TensorDType for half::bf16 {
const DTYPE: DType = DType::BF16;
fn to_f64(self) -> f64 {
self.to_f64()
}
fn from_f64(v: f64) -> Self {
half::bf16::from_f64(v)
}
}
impl TensorDType for i32 { impl TensorDType for i32 {
const DTYPE: DType = DType::I32; const DTYPE: DType = DType::I32;
fn to_f64(self) -> f64 { fn to_f64(self) -> f64 {

6
crates/xtrain-tensor/src/storage.rs
View File

@@ -99,9 +99,11 @@ impl Storage {
match device { match device {
Device::Cpu => Ok(Storage::cpu(vec![0u8; len_bytes])), Device::Cpu => Ok(Storage::cpu(vec![0u8; len_bytes])),
Device::Cuda(dev) => { Device::Cuda(dev) => {
// No device memset in T2: stage zeros from the host. // Device-side memset (Phase T7): avoids a blocking H2D memcpy of a
// host zero buffer on every op-output allocation. cudaMemset is
// async on the default stream, so it doesn't serialize the stream.
let mut buf = GpuBuffer::alloc(len_bytes)?; let mut buf = GpuBuffer::alloc(len_bytes)?;
buf.copy_from_host(&vec![0u8; len_bytes])?; buf.memset(0)?;
Ok(Storage::cuda(buf, dev)) Ok(Storage::cuda(buf, dev))
} }
} }

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

File diff suppressed because it is too large Load Diff

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

@@ -56,3 +56,170 @@ fn elementwise_scale_kernel() {
r.len() 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");
}

31
crates/xtrain-train/Cargo.toml Normal file
View File

@@ -0,0 +1,31 @@
[package]
name = "xtrain-train"
version.workspace = true
edition.workspace = true
[dependencies]
xtrain-tensor = { path = "../xtrain-tensor" }
xtrain-autodiff = { path = "../xtrain-autodiff" }
xtrain-model = { path = "../xtrain-model" }
xtrain-optim = { path = "../xtrain-optim" }
xtrain-cuda = { path = "../xtrain-cuda" }
# Reuse xserv's from-scratch GPT-2/Qwen BPE (project decision). This relative
# path resolves on both ~/projects (local) and /opt/wjh/projects (dash5). The
# crate inherits xserv's workspace for its own deps (serde/regex) — Cargo reads
# the target package's workspace, not ours.
xserv-tokenizer = { path = "../../../xserv/crates/xserv-tokenizer" }
# T9 export to xserv: HF Qwen3 safetensors + BF16 weight cast.
half.workspace = true
safetensors = "0.5"
[[bin]]
name = "train"
path = "src/bin/train.rs"
[[bin]]
name = "export_safetensors"
path = "src/bin/export_safetensors.rs"
[[bin]]
name = "dump_logits"
path = "src/bin/dump_logits.rs"

26
crates/xtrain-train/build.rs Normal file
View File

@@ -0,0 +1,26 @@
use std::env;
use std::path::Path;
use std::process::Command;
// Per-crate convention: the training loop / sampler / checkpoint all drive GPU
// ops through the model + tensor layers, so the bulk of this crate is gated
// behind `not(no_cuda)`. The LR schedule and the grad-clip *math* are host-only
// and always compile. cfg does not propagate across crates, so re-detect nvcc.
fn main() {
println!("cargo:rustc-check-cfg=cfg(no_cuda)");
let cuda_path = env::var("CUDA_HOME")
.or_else(|_| env::var("CUDA_PATH"))
.unwrap_or_else(|_| "/usr/local/cuda".to_string());
if !nvcc_available(&cuda_path) {
println!("cargo:rustc-cfg=no_cuda");
}
}
fn nvcc_available(cuda_path: &str) -> bool {
if Command::new("nvcc").arg("--version").output().is_ok() {
return true;
}
Path::new(&format!("{cuda_path}/bin/nvcc")).exists()
}

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

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

98
crates/xtrain-train/src/bin/dump_logits.rs Normal file
View File

@@ -0,0 +1,98 @@
//! Phase T9 verification helper — dump xtrain's OWN top-k next-token logits for a
//! prompt, so they can be compared against xserv's `dump-logits` on the exported
//! model (the closed-loop acceptance check). f32 forward, same model/config/ckpt
//! as bin/train.rs + bin/export_safetensors.rs.
//!
//! export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
//! cargo run -p xtrain-train --release --bin dump_logits -- \
//! /tmp/xtrain_tinystories.ckpt /opt/wjh/models/gpt2/tokenizer.json "Once upon a time"
#[cfg(no_cuda)]
fn main() {
eprintln!("dump_logits: 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, 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 main() {
use xserv_tokenizer::Tokenizer;
let args: Vec<String> = std::env::args().collect();
let ckpt = args
.get(1)
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from("/tmp/xtrain_tinystories.ckpt"));
let tok_path = args
.get(2)
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
let prompt = args
.get(3)
.cloned()
.unwrap_or_else(|| "Once upon a time".to_string());
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let dev = Device::Cuda(0);
let tok = Tokenizer::from_file(&tok_path);
let mut cfg = Config::tiny();
cfg.vocab = tok.vocab_size();
cfg.n_layers = 4;
let mut seed = 1u64;
let model = TinyTransformer::new(cfg, dev, |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");
let ids: Vec<i32> = tok.encode(&prompt).into_iter().map(|t| t as i32).collect();
eprintln!("Prompt: {prompt}");
eprintln!("Token IDs: {ids:?}");
let logits = model
.forward(&ids_tensor(&ids, dev))
.value()
.to_device(Device::Cpu);
let lg = logits.as_slice::<f32>();
let vocab = cfg.vocab;
let last = &lg[(ids.len() - 1) * vocab..ids.len() * vocab];
let mut idx: Vec<(usize, f32)> = last.iter().copied().enumerate().collect();
idx.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
println!("Top-20 logits (last position):");
for (rank, (id, val)) in idx.iter().take(20).enumerate() {
let t = tok.decode(&[*id as u32]);
println!(" [{rank:>2}] id={id:>6} logit={val:>10.4} token={t:?}");
}
}

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

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

277
crates/xtrain-train/src/bin/export_safetensors.rs Normal file
View File

@@ -0,0 +1,277 @@
//! Phase T9 — export a trained xtrain checkpoint into the format xserv loads:
//! an HF Qwen3-style `config.json` + `model.safetensors` (+ a copy of the GPT-2
//! `tokenizer.json`), so xserv's `Qwen3` loader can serve the same weights.
//!
//! xtrain's `TinyTransformer` is (after T9) architecturally a tiny Qwen3:
//! RoPE (rotate_half, pos=row) + RMSNorm + per-head QK-norm + SwiGLU + separate
//! lm_head, MHA (n_kv_heads = n_heads). The only deltas to xserv are mechanical:
//! - tensor NAMES → HF Qwen3 names (`model.layers.{i}.self_attn.q_proj.weight` …)
//! - 2D proj LAYOUT → xtrain stores `[in,out]` (computes `x@W`); xserv/HF want
//! `[out,in]` (computes `x@Wᵀ`) → transpose every 2D projection weight.
//! 1D norms and the `[vocab,dim]` embedding/lm_head rows are unchanged.
//! - DTYPE → xserv's Qwen3 forward is BF16-only, so weights are written as BF16.
//!
//! See `docs/08-export-xserv.md` for the full architecture diff + mapping table.
//!
//! Run on dash5 (needs a GPU to materialise the checkpoint params). The model
//! architecture must match the checkpoint — pass the same arch flags used to
//! train (defaults reproduce the v0-baseline tiny config):
//! export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
//! cargo run -p xtrain-train --release --bin export_safetensors -- \
//! /tmp/xtrain_v1.ckpt /opt/wjh/models/gpt2/tokenizer.json /tmp/xtrain_export \
//! --heads 8 --head-dim 32 --layers 8 --ffn 1024
#[cfg(no_cuda)]
fn main() {
eprintln!("export_safetensors: built without CUDA (no_cuda); run on a GPU host (dash5).");
}
#[cfg(not(no_cuda))]
use std::path::{Path, PathBuf};
#[cfg(not(no_cuda))]
use half::bf16;
#[cfg(not(no_cuda))]
use xtrain_autodiff::tape::Var;
#[cfg(not(no_cuda))]
use xtrain_cuda::device;
#[cfg(not(no_cuda))]
use xtrain_model::{Config, TinyTransformer, param_to_host};
#[cfg(not(no_cuda))]
use xtrain_tensor::Device;
// A flag like `--layers 8`: scan argv for `name`, parse the following token.
#[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)
}
// Same deterministic init scheme as bin/train.rs, so a freshly-built model has
// the right shapes before `load_into` overwrites the values from the checkpoint.
#[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()
}
/// A param ready to serialize: HF name + the (possibly transposed) row-major
/// data + its shape. Stored as BF16 (xserv's Qwen3 forward is BF16-only).
#[cfg(not(no_cuda))]
struct Export {
name: String,
data: Vec<bf16>,
shape: Vec<usize>,
}
/// 1D norm / embedding row-table: keep layout, just cast to BF16.
#[cfg(not(no_cuda))]
fn keep(name: &str, v: &Var) -> Export {
let host = param_to_host(v);
let shape = v.value().shape().to_vec();
Export {
name: name.to_string(),
data: host.iter().map(|&x| bf16::from_f32(x)).collect(),
shape,
}
}
/// 2D projection weight: xtrain `[in,out]` (x@W) → HF `[out,in]` (x@Wᵀ). Transpose
/// the row-major matrix and cast to BF16.
#[cfg(not(no_cuda))]
fn transpose(name: &str, v: &Var) -> Export {
let host = param_to_host(v);
let shape = v.value().shape().to_vec();
assert_eq!(shape.len(), 2, "transpose expects a 2D weight: {name}");
let (rows, cols) = (shape[0], shape[1]); // [in, out]
let mut out = vec![bf16::ZERO; rows * cols];
for r in 0..rows {
for c in 0..cols {
// out[c, r] = in[r, c]
out[c * rows + r] = bf16::from_f32(host[r * cols + c]);
}
}
Export {
name: name.to_string(),
data: out,
shape: vec![cols, rows], // [out, in]
}
}
/// Assemble every export tensor in HF Qwen3 naming, reading the xtrain params in
/// their stable `params()` order:
/// embed → per block [attn_norm, wq, wk, wv, q_norm, k_norm, wo, ffn_norm,
/// w_gate, w_up, w_down] → final_norm → lm_head
#[cfg(not(no_cuda))]
fn build_exports(model: &TinyTransformer) -> Vec<Export> {
let cfg = model.config();
let p = model.params();
let mut it = p.iter();
let mut next = || it.next().expect("params() ran short");
let mut ex = Vec::new();
ex.push(keep("model.embed_tokens.weight", next())); // [vocab, dim]
for l in 0..cfg.n_layers {
let b = format!("model.layers.{l}");
ex.push(keep(&format!("{b}.input_layernorm.weight"), next()));
ex.push(transpose(&format!("{b}.self_attn.q_proj.weight"), next()));
ex.push(transpose(&format!("{b}.self_attn.k_proj.weight"), next()));
ex.push(transpose(&format!("{b}.self_attn.v_proj.weight"), next()));
ex.push(keep(&format!("{b}.self_attn.q_norm.weight"), next()));
ex.push(keep(&format!("{b}.self_attn.k_norm.weight"), next()));
ex.push(transpose(&format!("{b}.self_attn.o_proj.weight"), next()));
ex.push(keep(
&format!("{b}.post_attention_layernorm.weight"),
next(),
));
ex.push(transpose(&format!("{b}.mlp.gate_proj.weight"), next()));
ex.push(transpose(&format!("{b}.mlp.up_proj.weight"), next()));
ex.push(transpose(&format!("{b}.mlp.down_proj.weight"), next()));
}
ex.push(keep("model.norm.weight", next())); // [dim]
ex.push(transpose("lm_head.weight", next())); // [dim,vocab] → [vocab,dim]
assert!(it.next().is_none(), "params() had extra tensors");
ex
}
/// config.json matching xserv's `ModelConfig` for a Qwen3 with xtrain's dims and
/// reconciled fields (eps, rope theta, head_dim, n_kv_heads = n_heads, untied).
#[cfg(not(no_cuda))]
fn config_json(cfg: &Config) -> String {
format!(
r#"{{
"architectures": ["Qwen3ForCausalLM"],
"model_type": "qwen3",
"vocab_size": {vocab},
"hidden_size": {dim},
"intermediate_size": {ffn},
"num_hidden_layers": {layers},
"num_attention_heads": {heads},
"num_key_value_heads": {kv_heads},
"head_dim": {head_dim},
"max_position_embeddings": 2048,
"rms_norm_eps": {eps},
"rope_theta": {theta},
"tie_word_embeddings": false,
"attention_bias": false,
"hidden_act": "silu"
}}
"#,
vocab = cfg.vocab,
dim = cfg.dim,
ffn = cfg.ffn_hidden,
layers = cfg.n_layers,
heads = cfg.n_heads,
kv_heads = cfg.num_kv_heads, // GQA (T15): real num_key_value_heads (= n_heads for MHA)
head_dim = cfg.head_dim,
eps = cfg.eps,
theta = cfg.rope_theta,
)
}
#[cfg(not(no_cuda))]
fn main() {
use safetensors::tensor::{Dtype, TensorView};
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()))
.unwrap_or_else(|| PathBuf::from("/tmp/xtrain_tinystories.ckpt"));
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 out_dir = positionals
.get(2)
.map(|s| PathBuf::from(s.as_str()))
.unwrap_or_else(|| PathBuf::from("/tmp/xtrain_export"));
// Architecture must match the checkpoint. Defaults = v0-baseline tiny config.
let n_heads = flag(&args, "--heads", 2usize);
let head_dim = flag(&args, "--head-dim", 16usize);
let n_layers = flag(&args, "--layers", 4usize);
let ffn = flag(&args, "--ffn", 64usize);
// GQA (Phase T15): num K/V heads (must match the trained ckpt; default = --heads).
let kv_heads = flag(&args, "--kv-heads", n_heads);
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let dev = Device::Cuda(0);
// Size the model from the arch flags + gpt2 vocab; must match the checkpoint.
let tok = Tokenizer::from_file(&tok_path);
let vocab = tok.vocab_size();
let cfg = Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
println!(
"export: ckpt {} → {} (vocab {}, dim {}, layers {}, heads {}, kv_heads {}, head_dim {})",
ckpt.display(),
out_dir.display(),
cfg.vocab,
cfg.dim,
cfg.n_layers,
cfg.n_heads,
cfg.num_kv_heads,
cfg.head_dim,
);
let mut seed = 1u64;
let model = TinyTransformer::new(cfg, dev, |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");
let exports = build_exports(&model);
println!("export: {} tensors", exports.len());
// Serialize to safetensors. Each TensorView borrows the raw BF16 bytes.
let views: Vec<(String, TensorView)> = exports
.iter()
.map(|e| {
let bytes = unsafe {
std::slice::from_raw_parts(e.data.as_ptr() as *const u8, e.data.len() * 2)
};
let view = TensorView::new(Dtype::BF16, e.shape.clone(), bytes)
.unwrap_or_else(|err| panic!("bad tensor view {}: {err}", e.name));
(e.name.clone(), view)
})
.collect();
std::fs::create_dir_all(&out_dir).expect("mkdir out_dir");
let st = safetensors::tensor::serialize(views.iter().map(|(n, v)| (n.as_str(), v)), &None)
.expect("serialize safetensors");
std::fs::write(out_dir.join("model.safetensors"), st).expect("write model.safetensors");
std::fs::write(out_dir.join("config.json"), config_json(&cfg)).expect("write config.json");
copy_tokenizer(&tok_path, &out_dir);
println!(
"export: wrote config.json + model.safetensors + tokenizer.json to {}",
out_dir.display()
);
}
/// Place the tokenizer beside the weights so xserv loads it from the model dir.
#[cfg(not(no_cuda))]
fn copy_tokenizer(tok_path: &Path, out_dir: &Path) {
std::fs::copy(tok_path, out_dir.join("tokenizer.json")).expect("copy tokenizer.json");
}

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

@@ -0,0 +1,106 @@
//! 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 Normal file
View File

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

179
crates/xtrain-train/src/bin/greedy_sample.rs Normal file
View File

@@ -0,0 +1,179 @@
//! Greedy-generation helper for run verification — load a trained checkpoint with
//! its arch flags and print xtrain's OWN greedy continuations for the fixed run
//! prompts, so they can be diffed against xserv's greedy on the exported weights
//! (the token-match check each scaling run reports). f32 forward, same
//! model/config/ckpt + init scheme as bin/train.rs and bin/export_safetensors.rs.
//!
//! 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
#[cfg(no_cuda)]
fn main() {
eprintln!("greedy_sample: 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;
// Same deterministic LCG init scheme as bin/train.rs / bin/export_safetensors.rs.
#[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()
}
// A flag like `--layers 18`: scan argv for `name`, parse the following token.
#[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 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;
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()))
.unwrap_or_else(|| PathBuf::from("/tmp/xtrain_tinystories.ckpt"));
let tok_path = positionals
.get(1)
.map(|s| PathBuf::from(s.as_str()))
.unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
// Architecture must match the checkpoint. Defaults = v0-baseline tiny config.
let n_heads = flag(&args, "--heads", 2usize);
let head_dim = flag(&args, "--head-dim", 16usize);
let n_layers = flag(&args, "--layers", 4usize);
let ffn = flag(&args, "--ffn", 64usize);
// 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();
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);
println!(
"greedy_sample: ckpt {} (vocab {}, dim {}, layers {}, heads {}, kv_heads {}, head_dim {})",
ckpt.display(),
cfg.vocab,
cfg.dim,
cfg.n_layers,
cfg.n_heads,
cfg.num_kv_heads,
cfg.head_dim,
);
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!(
"decode: prompts={} max_new={} temperature={}",
prompts.len(),
max_new,
temperature
);
for p in prompts {
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 text = tok.decode(&out.iter().map(|&t| t as u32).collect::<Vec<_>>());
println!("[{p}] → {text}");
}
}

309
crates/xtrain-train/src/bin/train.rs Normal file
View File

@@ -0,0 +1,309 @@
//! End-to-end training entry point: load the GPT-2 BPE + a TinyStories corpus,
//! train the tiny transformer with hand-written AdamW for a BOUNDED budget,
//! evaluate held-out val loss, checkpoint the best, and print a few samples.
//!
//! The MODEL SIZE is a CLI-tunable scaling-ladder rung (v0 baseline = the
//! defaults; v1 = dim256/8L/8h via flags), not a hardcoded tiny config.
//!
//! Run on dash5 (needs a GPU + the corpus + tokenizer.json):
//! export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
//! cargo run -p xtrain-train --release --bin train -- \
//! /opt/wjh/models/gpt2/tokenizer.json data/tinystories-train.txt \
//! --dim 256 --heads 8 --head-dim 32 --layers 8 --ffn 1024 \
//! --steps 3000 --batch 16 --seq 128 --max-lr 6e-4 \
//! --val-tokens 200000 --eval-every 250 --ckpt /tmp/xtrain_v1.ckpt
//!
//! Positional: <tokenizer.json> <corpus.txt>. Everything else is a flag with a
//! sane default (defaults reproduce the v0-baseline tiny config).
// On a GPU-less host (no_cuda) the whole training body is unavailable; keep a
// stub `main` so the crate still builds for `cargo check`.
#[cfg(no_cuda)]
fn main() {
eprintln!("xtrain train: built without CUDA (no_cuda); run on a GPU host (dash5).");
}
#[cfg(not(no_cuda))]
use std::path::{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::DType;
#[cfg(not(no_cuda))]
use xtrain_tensor::Device;
#[cfg(not(no_cuda))]
use xtrain_train::data::Corpus;
#[cfg(not(no_cuda))]
use xtrain_train::sample::generate;
#[cfg(not(no_cuda))]
use xtrain_train::schedule::LrSchedule;
#[cfg(not(no_cuda))]
use xtrain_train::{TrainConfig, train};
// Deterministic LCG fill in [-scale, scale) — same init scheme as the T5 tests.
#[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()
}
// A flag like `--dim 256`: scan argv for `name`, parse the following token.
#[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 main() {
let args: Vec<String> = std::env::args().collect();
// First two non-flag positionals: tokenizer.json, corpus.txt.
let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
let tok_path = positionals
.first()
.map(|s| PathBuf::from(s.as_str()))
.unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
let corpus_path = positionals
.get(1)
.map(|s| PathBuf::from(s.as_str()))
.unwrap_or_else(|| PathBuf::from("data/tinystories-valid-3mb.txt"));
// Architecture (scaling-ladder rung). Defaults = v0-baseline tiny config.
let n_heads = flag(&args, "--heads", 2usize);
let head_dim = flag(&args, "--head-dim", 16usize);
let n_layers = flag(&args, "--layers", 4usize);
let ffn = flag(&args, "--ffn", 64usize);
// GQA (Phase T15): num K/V heads (must divide --heads). Default = --heads (MHA).
let kv_heads = flag(&args, "--kv-heads", n_heads);
// `--dim` is informational; dim is always n_heads*head_dim. Warn on mismatch.
let dim_flag = flag(&args, "--dim", 0usize);
if dim_flag != 0 && dim_flag != n_heads * head_dim {
eprintln!(
"warning: --dim {dim_flag} != heads*head_dim {}; using {}",
n_heads * head_dim,
n_heads * head_dim
);
}
// Optimization knobs.
let steps: usize = flag(&args, "--steps", 2000);
let batch_size: usize = flag(&args, "--batch", 8);
// Micro-batch gradient accumulation (Phase T16): effective batch =
// accum_steps × batch, at one micro-batch's activation-memory cost. Default 1
// = no accumulation (bit-identical to the pre-T16 path).
let accum_steps: usize = flag(&args, "--accum-steps", 1).max(1);
let seq_len: usize = flag(&args, "--seq", 64);
let max_lr: f32 = flag(&args, "--max-lr", 3e-3);
let min_lr: f32 = flag(&args, "--min-lr", max_lr * 0.1);
let weight_decay: f32 = flag(&args, "--wd", 0.1);
let max_grad_norm: f32 = flag(&args, "--clip", 1.0);
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).
let dropout: f32 = flag(&args, "--dropout", 0.0f32);
// bf16 mixed precision (Phase T12): fp32 master weights, bf16 linears +
// activations. Opt-in; default fp32 reproduces v0–v4 numerics.
let bf16 = args.iter().any(|a| a == "--bf16");
// Activation recomputation (Phase T13): per-block gradient checkpointing —
// exact grads, lower peak activation memory (lets dim1024 batch32 fit). Opt-in;
// default off stores every activation (unchanged numerics).
let recompute = args.iter().any(|a| a == "--recompute");
// Fused flash-attention (Phase T14): single fused SDPA kernel, online softmax,
// no materialized [bh,S,S] scores. Opt-in; default off keeps the composed path.
let flash = args.iter().any(|a| a == "--flash");
let ckpt: PathBuf = PathBuf::from(
args.iter()
.position(|a| a == "--ckpt")
.and_then(|i| args.get(i + 1))
.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();
let device = Device::Cuda(0);
println!(
"loading tokenizer {} + corpus {} (cached id stream)",
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)
};
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 {
let (t, v) = corpus.split_tail(val_tokens);
println!("split: {} train tokens / {} val tokens", t.len(), v.len());
(t, Some(v))
} else {
(corpus, None)
};
let mut cfg =
Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn).with_kv_heads(kv_heads);
cfg.dropout = dropout;
println!(
"model: dim {} layers {} heads {} kv_heads {} head_dim {} ffn {} → core {:.3}M params \
(+ embed/lm {:.2}M = {:.2}M total)",
cfg.dim,
cfg.n_layers,
cfg.n_heads,
cfg.num_kv_heads,
cfg.head_dim,
cfg.ffn_hidden,
cfg.core_params() as f32 / 1e6,
(cfg.num_params() - cfg.core_params()) as f32 / 1e6,
cfg.num_params() as f32 / 1e6,
);
let mut seed = 1u64;
let mut model = TinyTransformer::new(cfg, device, |shape| {
seed = seed.wrapping_add(1);
let n: usize = shape.iter().product();
if shape.len() == 1 {
// RMSNorm gammas → ~1.
fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
} else {
// Small fan-in-ish scale; keeps early logits tame.
fill(n, seed, 0.04)
}
});
if bf16 {
model = model.with_compute_dtype(DType::BF16);
println!("bf16 mixed precision: ON (fp32 master weights)");
}
if recompute {
model = model.with_recompute(true);
println!("activation recompute: ON (per-block gradient checkpointing)");
}
if flash {
model = model.with_flash(true);
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)");
}
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
// metric — the v0-vs-v1 val-loss comparison. The arch flags must match the ckpt.
if let Some(p) = args.iter().position(|a| a == "--eval-ckpt") {
let ckpt_path = PathBuf::from(args.get(p + 1).expect("--eval-ckpt <path>"));
xtrain_train::checkpoint::load_into(&ckpt_path, &model.params())
.expect("load eval checkpoint");
let v = valid.expect("--eval-ckpt needs --val-tokens > 0");
let vl = xtrain_train::eval_loss(&model, device, &v, seq_len, eval_batches);
println!("eval-only: {} → val loss {vl:.4}", ckpt_path.display());
sample_some(&model, device, &tok_path);
return;
}
let tcfg = TrainConfig {
seq_len,
batch_size,
accum_steps,
steps,
schedule: LrSchedule {
max_lr,
min_lr,
warmup: (steps / 20).max(20),
total: steps,
},
weight_decay,
max_grad_norm,
log_every: 50,
ckpt_path: Some(ckpt.clone()),
ckpt_every: 500,
eval_every,
eval_batches,
seed: 42,
};
println!(
"training: {} steps, seq {}, batch {} × accum {} = effective batch {}, \
lr {:.1e}→{:.1e}, eval every {}",
tcfg.steps,
tcfg.seq_len,
tcfg.batch_size,
tcfg.accum_steps,
tcfg.batch_size * tcfg.accum_steps,
tcfg.schedule.max_lr,
tcfg.schedule.min_lr,
tcfg.eval_every
);
let result = train(&model, device, &train_corpus, valid.as_ref(), &tcfg);
let start = result.train_losses.first().copied().unwrap_or(0.0);
let end = result.train_losses.last().copied().unwrap_or(0.0);
println!("train loss: start {start:.4} → end {end:.4}");
if let Some(best) = result.best_val {
println!("best val loss: {best:.4}");
}
if let Some((s, v)) = result.evals.last() {
println!("final val loss (step {s}): {v:.4}");
}
sample_some(&model, device, &tok_path);
}
#[cfg(not(no_cuda))]
fn sample_some(model: &TinyTransformer, device: Device, tok_path: &Path) {
use xserv_tokenizer::Tokenizer;
let tok = Tokenizer::from_file(tok_path);
let prompts = ["Once upon a time", "The little", "One day"];
println!("\n--- samples (greedy) ---");
for p in prompts {
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, 40, 0.0, &mut rng);
let text = tok.decode(&out.iter().map(|&t| t as u32).collect::<Vec<_>>());
println!("[{p}] → {text}");
}
println!("\n--- samples (temperature 0.8) ---");
for p in prompts {
let ids: Vec<i32> = tok.encode(p).into_iter().map(|t| t as i32).collect();
let mut rng = 13u64;
let out = generate(model, device, &ids, 40, 0.8, &mut rng);
let text = tok.decode(&out.iter().map(|&t| t as u32).collect::<Vec<_>>());
println!("[{p}] → {text}");
}
}

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

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

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

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

90
crates/xtrain-train/src/checkpoint.rs Normal file
View File

@@ -0,0 +1,90 @@
//! Checkpoint save/load. Dumps the model's `params()` (in their stable order) to
//! a flat binary file and reloads them into a model with matching architecture.
//!
//! Format (little-endian):
//! ```text
//! magic : u32 = 0x58545254 ("XTRT")
//! version : u32 = 1
//! n_params: u32
//! repeat n_params times:
//! ndim : u32
//! dims : [u32; ndim]
//! data : [f32; prod(dims)]
//! ```
//! Architecture/config is NOT stored here — the caller rebuilds the model from
//! the same `Config` and `load_into`s the params (the round-trip and resume both
//! know their config). Gated behind `not(no_cuda)` (it round-trips GPU tensors).
#![cfg(not(no_cuda))]
use std::fs::File;
use std::io::{BufReader, BufWriter, Read, Write};
use std::path::Path;
use xtrain_autodiff::tape::Var;
use xtrain_tensor::{Device, Tensor};
const MAGIC: u32 = 0x5854_5254;
const VERSION: u32 = 1;
/// Write every parameter (value) to `path` in `params()` order.
pub fn save(path: &Path, params: &[Var]) -> std::io::Result<()> {
let mut w = BufWriter::new(File::create(path)?);
w.write_all(&MAGIC.to_le_bytes())?;
w.write_all(&VERSION.to_le_bytes())?;
w.write_all(&(params.len() as u32).to_le_bytes())?;
for p in params {
let v = p.value().to_device(Device::Cpu);
let shape = v.shape();
w.write_all(&(shape.len() as u32).to_le_bytes())?;
for &d in shape {
w.write_all(&(d as u32).to_le_bytes())?;
}
for &x in v.as_slice::<f32>() {
w.write_all(&x.to_le_bytes())?;
}
}
w.flush()
}
/// Read a checkpoint and overwrite each parameter's value in `params` (in order).
/// Shapes must match the saved ones. Tensors are placed on each param's device.
pub fn load_into(path: &Path, params: &[Var]) -> std::io::Result<()> {
let mut r = BufReader::new(File::open(path)?);
assert_eq!(read_u32(&mut r)?, MAGIC, "bad checkpoint magic");
assert_eq!(read_u32(&mut r)?, VERSION, "unsupported checkpoint version");
let n = read_u32(&mut r)? as usize;
assert_eq!(n, params.len(), "checkpoint param count != model");
for p in params {
let ndim = read_u32(&mut r)? as usize;
let mut dims = Vec::with_capacity(ndim);
for _ in 0..ndim {
dims.push(read_u32(&mut r)? as usize);
}
let numel: usize = dims.iter().product();
let mut data = vec![0.0f32; numel];
for slot in data.iter_mut() {
*slot = read_f32(&mut r)?;
}
let device = p.value().device();
assert_eq!(
p.value().shape(),
dims.as_slice(),
"checkpoint shape mismatch"
);
p.set_value(Tensor::from_slice(&data, &dims).to_device(device));
}
Ok(())
}
fn read_u32<R: Read>(r: &mut R) -> std::io::Result<u32> {
let mut b = [0u8; 4];
r.read_exact(&mut b)?;
Ok(u32::from_le_bytes(b))
}
fn read_f32<R: Read>(r: &mut R) -> std::io::Result<f32> {
let mut b = [0u8; 4];
r.read_exact(&mut b)?;
Ok(f32::from_le_bytes(b))
}

146
crates/xtrain-train/src/clip.rs Normal file
View File

@@ -0,0 +1,146 @@
//! Global-norm gradient clipping. The norm is computed across *all* parameter
//! gradients jointly (the same as `torch.nn.utils.clip_grad_norm_`): if the total
//! L2 norm exceeds `max_norm`, every gradient is scaled by `max_norm / total`.
//!
//! The norm math is host-only and testable; [`clip_grad_norm`] wraps it over the
//! parameter `Var`s (GPU round-trip), gated behind `not(no_cuda)`.
/// Compute the global L2 norm over a set of flat gradient buffers.
pub fn global_l2_norm(grads: &[Vec<f32>]) -> f32 {
let mut sumsq = 0.0f64;
for g in grads {
for &x in g {
sumsq += (x as f64) * (x as f64);
}
}
sumsq.sqrt() as f32
}
/// The scale factor to apply for clipping to `max_norm` (1.0 if already under).
pub fn clip_scale(total_norm: f32, max_norm: f32) -> f32 {
if total_norm > max_norm && total_norm > 0.0 {
max_norm / total_norm
} else {
1.0
}
}
#[cfg(not(no_cuda))]
mod gpu {
use super::{clip_scale, global_l2_norm};
use xtrain_autodiff::tape::Var;
use xtrain_tensor::{Device, Tensor};
/// First multiply every parameter's `.grad()` by `pre_scale` (use `1/batch`
/// to turn accumulated summed grads into a batch mean; `1.0` for no-op), then
/// clip the result to a joint global L2 norm of `max_norm`, writing the final
/// grads back via the tape's grad slot. Returns the post-pre_scale total norm
/// (the value the clip threshold is compared against, handy for logging).
/// Parameters without a grad contribute 0.
pub fn clip_grad_norm(params: &[Var], max_norm: f32, pre_scale: f32) -> f32 {
let device = params[0].value().device();
let grads: Vec<Option<Vec<f32>>> = params
.iter()
.map(|p| {
p.grad()
.map(|g| g.to_device(Device::Cpu).as_slice::<f32>().to_vec())
})
.collect();
// Norm is measured on the (pre_scale-applied) grads — what the optimizer
// will actually see if no clipping is needed.
let scaled_present: Vec<Vec<f32>> = grads
.iter()
.flatten()
.map(|g| g.iter().map(|x| x * pre_scale).collect())
.collect();
let total = global_l2_norm(&scaled_present);
let factor = pre_scale * clip_scale(total, max_norm);
if (factor - 1.0).abs() < f32::EPSILON {
return total; // pre_scale==1 and under threshold → grads untouched
}
for (p, g) in params.iter().zip(&grads) {
if let Some(g) = g {
let shape = p.grad().unwrap().shape().to_vec();
let scaled: Vec<f32> = g.iter().map(|x| x * factor).collect();
let t = Tensor::from_slice(&scaled, &shape).to_device(device);
p.zero_grad();
Var::push_grad(p, t);
}
}
total
}
}
#[cfg(not(no_cuda))]
mod gpu_norm {
use super::clip_scale;
use xtrain_autodiff::tape::Var;
use xtrain_tensor::{DType, Device, Tensor};
/// GPU-side global-norm grad clip (Phase T7): compute the joint L2 norm of all
/// `pre_scale`-applied grads with a device reduction, then rescale every grad
/// in place by `pre_scale·clip_factor` — no per-step grad roundtrip to host
/// (only the single scalar norm comes back). Returns the post-pre_scale total
/// norm. Params without a grad contribute 0 and are skipped on rescale.
pub fn clip_grad_norm_gpu(params: &[Var], max_norm: f32, pre_scale: f32) -> f32 {
let device = params[0].value().device();
// sum-of-squares of the RAW grads accumulated on device.
let acc = Tensor::zeros(&[1], DType::F32, device);
for p in params {
if let Some(g) = p.grad() {
unsafe {
xtrain_cuda::ffi::launch_sumsq_accum_f32(
g.data_ptr() as *const f32,
acc.data_ptr() as *mut f32,
g.numel() as i32,
std::ptr::null_mut(),
);
}
}
}
xtrain_cuda::device::synchronize().expect("grad-norm reduce sync failed");
let raw_sumsq = acc.to_device(Device::Cpu).as_slice::<f32>()[0];
// Norm of the pre_scale-applied grads = pre_scale · sqrt(raw_sumsq).
let total = pre_scale * raw_sumsq.max(0.0).sqrt();
let factor = pre_scale * clip_scale(total, max_norm);
if (factor - 1.0).abs() >= f32::EPSILON {
for p in params {
if let Some(g) = p.grad() {
unsafe {
xtrain_cuda::ffi::launch_scale_inplace_f32(
g.data_ptr() as *mut f32,
factor,
g.numel() as i32,
std::ptr::null_mut(),
);
}
}
}
xtrain_cuda::device::synchronize().expect("grad rescale sync failed");
}
total
}
}
#[cfg(not(no_cuda))]
pub use gpu::clip_grad_norm;
#[cfg(not(no_cuda))]
pub use gpu_norm::clip_grad_norm_gpu;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn norm_and_scale() {
// grads = [3,4] → norm 5.
let g = vec![vec![3.0f32], vec![4.0]];
assert!((global_l2_norm(&g) - 5.0).abs() < 1e-6);
// Clip to 2.5 → scale 0.5.
assert!((clip_scale(5.0, 2.5) - 0.5).abs() < 1e-6);
// Already under → no scaling.
assert!((clip_scale(5.0, 10.0) - 1.0).abs() < 1e-6);
}
}

338
crates/xtrain-train/src/data.rs Normal file
View File

@@ -0,0 +1,338 @@
//! Data pipeline: load the GPT-2 BPE (reusing xserv's from-scratch tokenizer),
//! tokenize a text corpus into one flat token stream, and sample fixed-length
//! `(input, target)` windows for next-token prediction. Host-only (no GPU).
//!
//! For the scaling runs the corpus is large (full TinyStories ≈ 2 GB / ~470 M
//! tokens), and the from-scratch BPE is slow, so [`Corpus::load_cached`]
//! tokenizes ONCE and caches the id stream to a `<corpus>.u16.bin` next to the
//! text (GPT-2 vocab = 50257 < 65536, so u16 is exact). Subsequent runs mmap-read
//! the cache instead of re-tokenizing.
use std::io::{BufReader, BufWriter, Read, Write};
use std::path::{Path, PathBuf};
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,
}
impl Corpus {
/// Load `tokenizer.json` (GPT-2 BPE) and tokenize the UTF-8 text at
/// `corpus_path` into a single id stream. TinyStories separates stories with
/// `<|endoftext|>`; the GPT-2 tokenizer emits that as a single special token,
/// so document boundaries are preserved in the stream.
pub fn load(tokenizer_path: &Path, corpus_path: &Path) -> Self {
let tok = Tokenizer::from_file(tokenizer_path);
let text = std::fs::read_to_string(corpus_path)
.unwrap_or_else(|e| panic!("failed to read corpus {}: {e}", corpus_path.display()));
// The range-fetched corpus may start/end mid-story; drop a leading partial
// line and a trailing partial story so we only train on whole sentences.
let text = trim_to_whole_stories(&text);
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(),
}
}
/// Like [`load`](Self::load) but caches the tokenized id stream to
/// `<corpus_path>.u16.bin`. On the first run it tokenizes the (large) corpus
/// and writes the cache; on later runs it reads the cache directly, skipping
/// the slow BPE. The cache is just a flat little-endian `[u16]` (no header) —
/// it is keyed only by path, so delete it if the corpus or tokenizer changes.
pub fn load_cached(tokenizer_path: &Path, corpus_path: &Path) -> Self {
let cache = cache_path(corpus_path);
let vocab_size = Tokenizer::from_file(tokenizer_path).vocab_size();
if cache.exists() {
let tokens = read_u16_cache(&cache);
println!(
"corpus: read {} cached tokens from {}",
tokens.len(),
cache.display()
);
return Self {
tokens,
labels: None,
vocab_size,
};
}
let me = Self::load(tokenizer_path, corpus_path);
write_u16_cache(&cache, &me.tokens);
println!(
"corpus: tokenized {} tokens → cached to {}",
me.tokens.len(),
cache.display()
);
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 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,
vocab_size: self.vocab_size,
},
)
}
/// Total number of tokens.
pub fn len(&self) -> usize {
self.tokens.len()
}
pub fn is_empty(&self) -> bool {
self.tokens.is_empty()
}
/// Sample one `(input, target)` pair of length `seq` for next-token
/// prediction: a window `[s, s+seq+1)` → input `[s, s+seq)`, target shifted
/// by one. `rng_state` is advanced (a tiny LCG, so sampling is reproducible
/// from a seed without pulling in an RNG crate).
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 input = self.tokens[start..start + seq].to_vec();
let target = self.target_window(start, seq);
(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
/// the last `<|endoftext|>` marker, so a byte-range download still yields only
/// complete stories. Falls back to the raw text if no marker is present.
fn trim_to_whole_stories(text: &str) -> &str {
let start = text.find('\n').map(|i| i + 1).unwrap_or(0);
let body = &text[start..];
match body.rfind("<|endoftext|>") {
Some(end) => &body[..end + "<|endoftext|>".len()],
None => body,
}
}
/// `<corpus_path>.u16.bin` — the token-id cache beside the corpus text.
fn cache_path(corpus_path: &Path) -> PathBuf {
let mut s = corpus_path.as_os_str().to_os_string();
s.push(".u16.bin");
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(
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() % 2 == 0, "corrupt u16 cache (odd byte count)");
buf.chunks_exact(2)
.map(|b| u16::from_le_bytes([b[0], b[1]]) as 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]) {
let mut w = BufWriter::new(
std::fs::File::create(path)
.unwrap_or_else(|e| panic!("create cache {}: {e}", path.display())),
);
for &t in tokens {
assert!((0..=u16::MAX as i32).contains(&t), "token id {t} > u16");
w.write_all(&(t as u16).to_le_bytes()).expect("write cache");
}
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 {
*state = state
.wrapping_mul(6364136223846793005)
.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 Normal file
View File

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

25
crates/xtrain-train/src/lib.rs Normal file
View File

@@ -0,0 +1,25 @@
//! Training stack (Phase T6): LR schedule, global-norm grad clipping, checkpoint
//! save/load, the GPT-2 BPE data pipeline (reusing xserv's tokenizer), an
//! autoregressive sampler, and the training loop that wires them onto the T5
//! `TinyTransformer` + the hand-written AdamW (`xtrain-optim`).
//!
//! Host-only pieces (LR schedule, grad-norm math) always compile so the crate
//! `cargo check`s on a GPU-less host; everything that touches GPU tensors is
//! gated behind `not(no_cuda)`.
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;
#[cfg(not(no_cuda))]
pub use train_loop::{TrainConfig, TrainResult, eval_loss, train};

80
crates/xtrain-train/src/sample.rs Normal file
View File

@@ -0,0 +1,80 @@
//! Autoregressive text sampling from the trained model. The model is
//! single-sequence with RoPE position = row index, so generation re-runs the
//! forward on the growing prefix each step and reads the last row's logits — the
//! simplest correct approach (no KV cache; that is an inference/perf concern).
//!
//! Greedy when `temperature == 0`, else temperature sampling over the softmax.
#![cfg(not(no_cuda))]
use xtrain_model::{TinyTransformer, ids_tensor};
use xtrain_tensor::Device;
/// Generate `max_new` tokens continuing `prompt`. `temperature == 0` → greedy
/// argmax; otherwise sample from softmax(logits / temperature). `rng_state` is a
/// reproducible LCG seed (only used when temperature > 0).
pub fn generate(
model: &TinyTransformer,
device: Device,
prompt: &[i32],
max_new: usize,
temperature: f32,
rng_state: &mut u64,
) -> Vec<i32> {
let vocab = model.config().vocab;
let mut ids: Vec<i32> = prompt.to_vec();
for _ in 0..max_new {
let ids_t = ids_tensor(&ids, device);
// In bf16 mode the logits are bf16; cast to f32 (on device) before reading.
let logits = model.forward(&ids_t).value();
let logits = logits
.to_dtype(xtrain_tensor::DType::F32)
.to_device(Device::Cpu);
let lg = logits.as_slice::<f32>();
// Last row = next-token distribution for the current prefix.
let last = &lg[(ids.len() - 1) * vocab..ids.len() * vocab];
let next = if temperature <= 0.0 {
argmax(last)
} else {
sample_temperature(last, temperature, rng_state)
};
ids.push(next as i32);
}
ids
}
fn argmax(row: &[f32]) -> usize {
row.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
.unwrap()
.0
}
fn sample_temperature(row: &[f32], temperature: f32, rng_state: &mut u64) -> usize {
// Softmax with temperature (numerically stable).
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();
let r = (next_rand(rng_state) 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
}
fn next_rand(state: &mut u64) -> u32 {
*state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
(*state >> 32) as u32
}

58
crates/xtrain-train/src/schedule.rs Normal file
View File

@@ -0,0 +1,58 @@
//! Learning-rate schedule: linear warmup → cosine decay to a floor. Host-only
//! and pure (just arithmetic over the step index), so it unit-tests locally.
/// Warmup-then-cosine LR schedule.
///
/// - steps `0..warmup`: linear ramp `0 → max_lr`
/// - steps `warmup..total`: cosine decay `max_lr → min_lr`
/// - steps `>= total`: clamped at `min_lr`
#[derive(Clone, Copy, Debug)]
pub struct LrSchedule {
pub max_lr: f32,
pub min_lr: f32,
pub warmup: usize,
pub total: usize,
}
impl LrSchedule {
/// LR for a 0-indexed step.
pub fn lr(&self, step: usize) -> f32 {
if step < self.warmup {
// Linear warmup; +1 so step 0 already takes a (tiny) nonzero LR.
return self.max_lr * (step + 1) as f32 / self.warmup.max(1) as f32;
}
if step >= self.total {
return self.min_lr;
}
let progress = (step - self.warmup) as f32 / (self.total - self.warmup).max(1) as f32;
let cosine = 0.5 * (1.0 + (std::f32::consts::PI * progress).cos()); // 1 → 0
self.min_lr + (self.max_lr - self.min_lr) * cosine
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn warmup_then_cosine_shape() {
let s = LrSchedule {
max_lr: 1.0,
min_lr: 0.1,
warmup: 10,
total: 100,
};
// Warmup ramps up and reaches the peak at the end of warmup.
assert!(s.lr(0) < s.lr(5));
assert!(s.lr(5) < s.lr(9));
assert!((s.lr(9) - 1.0).abs() < 1e-6);
// Just after warmup, near the peak; midway, near the midpoint.
assert!(s.lr(10) > 0.95);
let mid = s.lr(55); // progress ~0.5 → cosine ~0.5
assert!((mid - (0.1 + 0.9 * 0.5)).abs() < 0.02);
// Decays monotonically to the floor by the end.
assert!(s.lr(99) < s.lr(55));
assert!(s.lr(100) <= s.min_lr + 1e-6);
assert!(s.lr(1_000) <= s.min_lr + 1e-6);
}
}

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

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

226
crates/xtrain-train/src/train_loop.rs Normal file
View File

@@ -0,0 +1,226 @@
//! The training loop: sample a batch of sequences → ONE batched forward `loss` →
//! backward → grad clip → AdamW step → zero grads; with an LR schedule, periodic
//! loss logging, and periodic checkpointing.
//!
//! Since T10 the model is batched (`loss_batched`): `batch_size` sequences are
//! flattened to `[batch*seq]` and run as a SINGLE forward/backward, so the linear
//! projections become big `[batch*seq, dim]` GEMMs that fill the GPU. The
//! cross-entropy mean is over all `batch*seq` rows — already the batch-mean loss,
//! so backward yields the batch-mean gradient directly (clip pre-scale = 1.0; no
//! more "loop B times + SUM + ×1/batch" hack).
#![cfg(not(no_cuda))]
use std::path::PathBuf;
use std::time::Instant;
use xtrain_model::{TinyTransformer, batched_ids_tensor, ids_tensor};
use xtrain_optim::GpuAdamW;
use xtrain_tensor::Device;
use crate::checkpoint;
use crate::clip::clip_grad_norm_gpu;
use crate::data::Corpus;
use crate::schedule::LrSchedule;
/// Knobs for a training run.
pub struct TrainConfig {
pub seq_len: usize,
pub batch_size: usize,
/// Micro-batch gradient accumulation (Phase T16): each optimizer step
/// accumulates grads over `accum_steps` micro-batches of `batch_size`
/// sequences, giving an EFFECTIVE batch of `accum_steps × batch_size` at the
/// activation-memory cost of a single micro-batch. `1` = no accumulation
/// (bit-identical to the pre-T16 path).
pub accum_steps: usize,
pub steps: usize,
pub schedule: LrSchedule,
pub weight_decay: f32,
pub max_grad_norm: f32,
pub log_every: usize,
/// Optional checkpoint path written every `ckpt_every` steps (and at the end).
/// When `eval_every > 0`, the checkpoint instead tracks the BEST val loss.
pub ckpt_path: Option<PathBuf>,
pub ckpt_every: usize,
/// Evaluate held-out val loss every `eval_every` steps (0 = never). Each eval
/// averages cross-entropy over `eval_batches` fixed windows of the val corpus.
pub eval_every: usize,
pub eval_batches: usize,
/// Seed for reproducible sequence sampling.
pub seed: u64,
}
/// Outcome of a run: per-step train losses and (step, val_loss) eval points.
pub struct TrainResult {
pub train_losses: Vec<f32>,
pub evals: Vec<(usize, f32)>,
pub best_val: Option<f32>,
}
/// Train `model` on `corpus` for `cfg.steps` AdamW steps. Returns the per-step
/// train-loss trace plus any (step, val_loss) eval points. Logs progress, and —
/// when `valid` is given and `cfg.eval_every > 0` — evaluates held-out val loss
/// periodically and checkpoints the BEST val model (else checkpoints on a fixed
/// cadence, as in T6). Logs progress.
pub fn train(
model: &TinyTransformer,
device: Device,
corpus: &Corpus,
valid: Option<&Corpus>,
cfg: &TrainConfig,
) -> TrainResult {
let params = model.params();
let mut opt = GpuAdamW::new(cfg.weight_decay);
let mut rng = cfg.seed;
let mut losses = Vec::with_capacity(cfg.steps);
let mut evals = Vec::new();
let mut best_val: Option<f32> = None;
let start = Instant::now();
let mut tokens_seen: u64 = 0;
// Best-val checkpointing only kicks in when we actually evaluate.
let track_best = valid.is_some() && cfg.eval_every > 0;
let accum = cfg.accum_steps.max(1);
for step in 0..cfg.steps {
let lr = cfg.schedule.lr(step);
// Accumulate grads over `accum` micro-batches of `batch_size` sequences,
// then take ONE optimizer step (Phase T16). Each micro-batch is ONE batched
// forward/backward; its loss is the CE mean over batch*seq rows, so backward
// yields that micro-batch's mean grad. To make the SUM over `accum` micro-
// batches equal a single step over an `accum × batch` batch, each micro-loss
// is scaled by 1/accum before backward (the tape SUM-accumulates the scaled
// grads). `accum == 1` skips the scale entirely → bit-identical to pre-T16.
let mut step_loss_sum = 0.0f32;
// Training mode → dropout active (T18; no-op when cfg.dropout == 0). Set
// each step so it is restored after a periodic eval flips to eval mode.
// Each micro-step's forward bumps the per-step seed → fresh masks.
model.train();
for _ in 0..accum {
let mut inputs = Vec::with_capacity(cfg.batch_size);
let mut targets_v = Vec::with_capacity(cfg.batch_size);
for _ in 0..cfg.batch_size {
let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
inputs.push(input);
targets_v.push(target);
}
let ids = batched_ids_tensor(&inputs, device);
let targets = batched_ids_tensor(&targets_v, device);
let loss = model.loss_batched(&ids, &targets, cfg.batch_size);
step_loss_sum += read_scalar(&loss);
if accum == 1 {
loss.backward();
} else {
xtrain_autodiff::ops::scale(&loss, 1.0 / accum as f32).backward();
}
tokens_seen += (cfg.batch_size * cfg.seq_len) as u64;
}
// Reported loss = mean over the effective batch = mean of the raw micro
// losses (each is itself a micro-batch mean of equal size).
let step_loss = step_loss_sum / accum as f32;
losses.push(step_loss);
// Backward already produced the effective-batch mean gradient — just clip.
let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, 1.0);
opt.step(lr, &params);
for p in &params {
p.zero_grad();
}
if step % cfg.log_every == 0 || step == cfg.steps - 1 {
let elapsed = start.elapsed().as_secs_f32();
let tps = tokens_seen as f32 / elapsed.max(1e-6);
println!(
"step {step:5}/{}: loss {step_loss:.4} lr {lr:.2e} gnorm {gnorm:.3} \
({tps:.0} tok/s)",
cfg.steps
);
}
// Periodic held-out eval (deterministic windows, no grad).
if let Some(v) = valid {
if cfg.eval_every > 0 && ((step + 1) % cfg.eval_every == 0 || step == cfg.steps - 1) {
let vl = eval_loss(model, device, v, cfg.seq_len, cfg.eval_batches);
evals.push((step, vl));
let improved = best_val.map(|b| vl < b).unwrap_or(true);
println!(
" eval @ step {step}: val loss {vl:.4}{}",
if improved { " (best)" } else { "" }
);
if improved {
best_val = Some(vl);
if let Some(path) = &cfg.ckpt_path {
checkpoint::save(path, &params).expect("best checkpoint save");
}
}
}
}
// Fixed-cadence checkpointing (only when not tracking best val).
if !track_best {
if let Some(path) = &cfg.ckpt_path {
if cfg.ckpt_every > 0 && (step + 1) % cfg.ckpt_every == 0 {
checkpoint::save(path, &params).expect("checkpoint save");
}
}
}
}
// Without periodic eval, still persist the final params (T6 behaviour). With
// best-val tracking the checkpoint already holds the best model — don't clobber.
if !track_best {
if let Some(path) = &cfg.ckpt_path {
checkpoint::save(path, &params).expect("final checkpoint save");
println!("saved checkpoint → {}", path.display());
}
}
TrainResult {
train_losses: losses,
evals,
best_val,
}
}
/// Mean cross-entropy over `batches` deterministic, non-overlapping windows of
/// the validation corpus (no backward — eval only). Deterministic so val loss is
/// comparable across steps and runs (and across models — the v0-vs-v1 metric).
pub fn eval_loss(
model: &TinyTransformer,
device: Device,
valid: &Corpus,
seq: usize,
batches: usize,
) -> f32 {
if valid.len() <= seq + 1 {
return f32::NAN;
}
// Eval mode → dropout is identity (T18).
model.eval();
let n_win = (valid.len() - 1) / seq; // disjoint windows that fit
let batches = batches.max(1).min(n_win.max(1));
let stride = (n_win / batches).max(1);
let mut sum = 0.0f32;
let mut count = 0usize;
for i in 0..batches {
let s = (i * stride) * seq;
if s + seq + 1 > valid.len() {
break;
}
let input: Vec<i32> = valid.tokens[s..s + seq].to_vec();
let target = valid.target_window(s, seq);
let ids = ids_tensor(&input, device);
let targets = ids_tensor(&target, device);
let loss = model.loss(&ids, &targets);
sum += read_scalar(&loss);
count += 1;
}
if count == 0 {
f32::NAN
} else {
sum / count as f32
}
}
fn read_scalar(v: &xtrain_autodiff::tape::Var) -> f32 {
v.value().to_device(Device::Cpu).as_slice::<f32>()[0]
}

203
crates/xtrain-train/tests/adamw_parity.py Normal file
View File

@@ -0,0 +1,203 @@
#!/usr/bin/env python3
"""AdamW-vs-PyTorch parity (Phase T6).
Loads the model dumped by tests/adamw_parity_dump.rs (config, ids, initial
params, the loss trajectory, and final params), rebuilds the IDENTICAL tiny
transformer in PyTorch from the same initial weights, and runs the SAME number
of `torch.optim.AdamW` steps with matched hyperparameters (lr, weight_decay,
betas, eps) on the same fixed batch. It then compares:
* the per-step loss trajectory (Rust AdamW vs torch AdamW), and
* the final parameters,
within a relative tolerance. A correct hand-written AdamW (bias correction +
decoupled weight decay) tracks torch's optimizer step-for-step.
Usage: python3 adamw_parity.py /tmp/xtrain_adamw
"""
import sys
import os
import math
import torch
DIR = sys.argv[1] if len(sys.argv) > 1 else "/tmp/xtrain_adamw"
def read_vec(name):
path = os.path.join(DIR, name)
shape = None
vals = []
with open(path) as f:
for line in f:
line = line.strip()
if line.startswith("# shape"):
shape = [int(x) for x in line.split()[2].split(",") if x]
elif line:
vals.append(float(line))
# float32 to match the engine's precision: this is an optimizer-trajectory
# parity over many steps, so we compare f32 training against an f32 reference
# (a float64 reference would diverge purely from precision over the steps).
t = torch.tensor(vals, dtype=torch.float32)
if shape:
t = t.reshape(shape)
return t
def read_cfg():
cfg = {}
with open(os.path.join(DIR, "config.txt")) as f:
for line in f:
k, v = line.split()
cfg[k] = v
return cfg
def read_ids(name):
with open(os.path.join(DIR, name)) as f:
return [int(x) for x in f.read().split()]
cfg = read_cfg()
DIM = int(cfg["dim"])
NL = int(cfg["n_layers"])
NH = int(cfg["n_heads"])
HD = int(cfg["head_dim"])
EPS = float(cfg["eps"])
THETA = float(cfg["rope_theta"])
LR = float(cfg["lr"])
WD = float(cfg["wd"])
N_STEPS = int(cfg["n_steps"])
ids = read_ids("ids.txt")
targets = read_ids("targets.txt")
SEQ = len(ids)
NAMES = ["embed"]
for l in range(NL):
for p in ["attn_norm", "wq", "wk", "wv", "q_norm", "k_norm", "wo",
"ffn_norm", "w_gate", "w_up", "w_down"]:
NAMES.append(f"l{l}_{p}")
NAMES += ["final_norm", "lm_head"]
# Load the IDENTICAL initial weights as leaf params (float32 reference).
P = {n: read_vec(f"w0_{n}.txt").clone().requires_grad_(True) for n in NAMES}
def rms_norm(x, gamma):
ms = x.pow(2).mean(dim=-1, keepdim=True)
return x * torch.rsqrt(ms + EPS) * gamma
def rope(x): # x: [seq, nh, hd], position = token index
half = HD // 2
out = torch.empty_like(x)
i = torch.arange(half, dtype=torch.float32)
freq = THETA ** (-(2.0 * i) / HD)
pos = torch.arange(SEQ, dtype=torch.float32).reshape(SEQ, 1)
ang = pos * freq
c = torch.cos(ang).reshape(SEQ, 1, half)
s = torch.sin(ang).reshape(SEQ, 1, half)
x0, x1 = x[..., :half], x[..., half:]
out[..., :half] = x0 * c - x1 * s
out[..., half:] = x1 * c + x0 * s
return out
idx = torch.tensor(ids, dtype=torch.long)
tgt = torch.tensor(targets, dtype=torch.long)
mask = torch.triu(torch.full((SEQ, SEQ), -1.0e9, dtype=torch.float32), diagonal=1)
def forward():
h = P["embed"][idx]
for l in range(NL):
x = rms_norm(h, P[f"l{l}_attn_norm"])
q = (x @ P[f"l{l}_wq"]).reshape(SEQ, NH, HD)
k = (x @ P[f"l{l}_wk"]).reshape(SEQ, NH, HD)
v = (x @ P[f"l{l}_wv"]).reshape(SEQ, NH, HD)
# Per-head QK-norm (Qwen3-style), before RoPE.
q = rms_norm(q, P[f"l{l}_q_norm"])
k = rms_norm(k, P[f"l{l}_k_norm"])
q = rope(q).transpose(0, 1)
k = rope(k).transpose(0, 1)
v = v.transpose(0, 1)
scale = 1.0 / math.sqrt(HD)
scores = (q @ k.transpose(-1, -2)) * scale + mask
probs = torch.softmax(scores, dim=-1)
out = (probs @ v).transpose(0, 1).reshape(SEQ, DIM)
h = h + out @ P[f"l{l}_wo"]
x = rms_norm(h, P[f"l{l}_ffn_norm"])
act = torch.nn.functional.silu(x @ P[f"l{l}_w_gate"]) * (x @ P[f"l{l}_w_up"])
h = h + act @ P[f"l{l}_w_down"]
h = rms_norm(h, P["final_norm"])
return h @ P["lm_head"]
# Match the Rust optimizer: torch.optim.AdamW with the same lr/wd/betas/eps.
opt = torch.optim.AdamW(list(P.values()), lr=LR, betas=(0.9, 0.999),
eps=1e-8, weight_decay=WD)
torch_losses = []
for _ in range(N_STEPS):
opt.zero_grad()
logits = forward()
loss = torch.nn.functional.cross_entropy(logits, tgt, reduction="mean")
torch_losses.append(loss.detach().item())
loss.backward()
opt.step()
def relerr(a, b):
a, b = a.double(), b.double()
denom = b.abs().clamp(min=1e-6)
return ((a - b).abs() / denom).max().item()
# allclose-style: a per-element error is acceptable if it is within rtol *or*
# atol (absolute). Weights span very small magnitudes, so a pure relative metric
# is misleading on near-zero entries; this matches torch.allclose's semantics.
def max_mismatch(a, b, rtol, atol):
a, b = a.double(), b.double()
err = (a - b).abs()
tol = atol + rtol * b.abs()
over = err - tol # > 0 only where it exceeds the combined tolerance
return over.max().item()
rust_losses = read_vec("losses.txt")
print("step rust_loss torch_loss relerr")
worst_loss = 0.0
for i in range(N_STEPS):
rl, tl = rust_losses[i].item(), torch_losses[i]
e = abs(rl - tl) / max(abs(tl), 1e-6)
worst_loss = max(worst_loss, e)
if i < 5 or i == N_STEPS - 1:
print(f"{i:4d} {rl:.6e} {tl:.6e} {e:.2e}")
print(f"loss trajectory: worst relerr = {worst_loss:.2e}")
RTOL = 2e-2
ATOL = 1e-3
worst_over, worst_name, worst_rel = 0.0, "", 0.0
fails = []
for n in NAMES:
ref = read_vec(f"wN_{n}.txt")
over = max_mismatch(P[n].detach(), ref, RTOL, ATOL)
rel = relerr(P[n].detach(), ref)
if over > worst_over:
worst_over, worst_name, worst_rel = over, n, rel
if over > 0.0:
fails.append((n, rel, over))
print(
f"final params: {len(NAMES)} checked, worst = {worst_name} "
f"(relerr {worst_rel:.2e}, tol-overflow {worst_over:.2e}) "
f"[rtol={RTOL}, atol={ATOL}]"
)
if worst_loss > RTOL or fails:
print("FAIL:")
if worst_loss > RTOL:
print(f" loss trajectory relerr {worst_loss:.3e} > {RTOL}")
for n, rel, over in fails:
print(f" param[{n}]: relerr={rel:.3e} tol-overflow={over:.3e}")
sys.exit(1)
print("ADAMW PARITY OK: loss trajectory + final params match torch.optim.AdamW (rtol/atol)")

173
crates/xtrain-train/tests/adamw_parity_dump.rs Normal file
View File

@@ -0,0 +1,173 @@
// AdamW-vs-PyTorch parity, step 1 of 2: build the tiny transformer with a fixed
// deterministic init, then run N steps of the hand-written AdamW on a FIXED
// (input, target) batch — recording the loss at each step and the final
// parameters. tests/adamw_parity.py rebuilds the identical model + torch.optim
// .AdamW with matched hyperparameters and compares the loss trajectory and final
// params within rtol. This is the rigorous correctness check for the optimizer.
//
// Run: XTRAIN_ADAMW_DIR=/tmp/xtrain_adamw cargo test -p xtrain-train \
// --test adamw_parity_dump -- --nocapture --ignored
// then: python3 crates/xtrain-train/tests/adamw_parity.py /tmp/xtrain_adamw
//
// Marked #[ignore] (fixture generator) and gated #![cfg(not(no_cuda))].
#![cfg(not(no_cuda))]
use std::fs;
use std::io::Write;
use std::path::PathBuf;
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, ids_tensor, param_to_host};
use xtrain_optim::AdamW;
use xtrain_tensor::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 write_vec(dir: &PathBuf, name: &str, data: &[f32], shape: &[usize]) {
let mut f = fs::File::create(dir.join(name)).unwrap();
let shape_str: Vec<String> = shape.iter().map(|d| d.to_string()).collect();
writeln!(f, "# shape {}", shape_str.join(",")).unwrap();
for v in data {
writeln!(f, "{v:.8e}").unwrap();
}
}
const LR: f32 = 0.01;
const WD: f32 = 0.1;
// Kept short on purpose: AdamW correctness shows in the per-step loss trajectory
// and the parameter values *while the loss is still well-determined*. Run it long
// enough to memorise the tiny batch and the model enters a flat, overparameterised
// region where many weight configs give the same loss — there f32(GPU) vs the
// torch reference diverge per-weight (large *relative* error on tiny weights)
// while the loss stays identical. 10 steps keeps both signals sharp.
const N_STEPS: usize = 10;
#[test]
#[ignore = "fixture generator for AdamW PyTorch parity; run with --ignored"]
fn dump_adamw_trajectory() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let dir = PathBuf::from(
std::env::var("XTRAIN_ADAMW_DIR").unwrap_or_else(|_| "/tmp/xtrain_adamw".to_string()),
);
fs::create_dir_all(&dir).unwrap();
let mut cfg = Config::tiny();
cfg.vocab = 12;
let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6];
let targets: Vec<i32> = vec![1, 4, 1, 5, 9, 2, 6, 0];
// Same deterministic init the parity dump uses (so the torch side can reuse it).
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.08)
}
});
// Dump config + ids + initial params (named for adamw_parity.py).
{
let mut f = fs::File::create(dir.join("config.txt")).unwrap();
writeln!(f, "vocab {}", cfg.vocab).unwrap();
writeln!(f, "dim {}", cfg.dim).unwrap();
writeln!(f, "n_layers {}", cfg.n_layers).unwrap();
writeln!(f, "n_heads {}", cfg.n_heads).unwrap();
writeln!(f, "head_dim {}", cfg.head_dim).unwrap();
writeln!(f, "ffn_hidden {}", cfg.ffn_hidden).unwrap();
writeln!(f, "eps {:e}", cfg.eps).unwrap();
writeln!(f, "rope_theta {:e}", cfg.rope_theta).unwrap();
writeln!(f, "lr {LR:e}").unwrap();
writeln!(f, "wd {WD:e}").unwrap();
writeln!(f, "n_steps {N_STEPS}").unwrap();
let mut g = fs::File::create(dir.join("ids.txt")).unwrap();
for v in &ids {
writeln!(g, "{v}").unwrap();
}
let mut g = fs::File::create(dir.join("targets.txt")).unwrap();
for v in &targets {
writeln!(g, "{v}").unwrap();
}
}
let names = param_names(&cfg);
let params = model.params();
for (name, p) in names.iter().zip(&params) {
let shape = p.value().shape().to_vec();
write_vec(&dir, &format!("w0_{name}.txt"), &param_to_host(p), &shape);
}
// Train N steps of AdamW with a CONSTANT lr (no schedule) on the fixed batch.
let ids_t = ids_tensor(&ids, device);
let targets_t = ids_tensor(&targets, device);
let mut opt = AdamW::new(LR, WD);
let mut losses = Vec::with_capacity(N_STEPS);
for _ in 0..N_STEPS {
let loss = model.loss(&ids_t, &targets_t);
losses.push(param_to_host(&loss)[0]);
loss.backward();
opt.step(LR, &params);
for p in &params {
p.zero_grad();
}
}
{
let mut f = fs::File::create(dir.join("losses.txt")).unwrap();
for l in &losses {
writeln!(f, "{l:.8e}").unwrap();
}
}
for (name, p) in names.iter().zip(&params) {
let shape = p.value().shape().to_vec();
write_vec(&dir, &format!("wN_{name}.txt"), &param_to_host(p), &shape);
}
println!(
"adamw parity: dumped to {} (loss {:.6e} → {:.6e} over {N_STEPS} steps)",
dir.display(),
losses.first().unwrap(),
losses.last().unwrap()
);
}
fn param_names(cfg: &Config) -> Vec<String> {
let mut names = vec!["embed".to_string()];
for l in 0..cfg.n_layers {
for p in [
"attn_norm",
"wq",
"wk",
"wv",
"q_norm",
"k_norm",
"wo",
"ffn_norm",
"w_gate",
"w_up",
"w_down",
] {
names.push(format!("l{l}_{p}"));
}
}
names.push("final_norm".to_string());
names.push("lm_head".to_string());
names
}

113
crates/xtrain-train/tests/checkpoint_roundtrip.rs Normal file
View File

@@ -0,0 +1,113 @@
// Checkpoint round-trip acceptance (Phase T6): train a few AdamW steps on a fixed
// batch, save the params, build a FRESH model (different init), load the
// checkpoint into it, and assert it produces identical logits + loss on a fixed
// input. This verifies the on-disk format dumps/reloads `params()` in order with
// exact f32 fidelity. Gated #![cfg(not(no_cuda))] (runs on dash5).
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, ids_tensor, param_to_host};
use xtrain_optim::AdamW;
use xtrain_tensor::Device;
use xtrain_train::checkpoint;
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 make_model(device: Device, vocab: usize, init_seed: u64) -> TinyTransformer {
let mut cfg = Config::tiny();
cfg.vocab = vocab;
let mut seed = init_seed;
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)
}
})
}
#[test]
fn checkpoint_roundtrip_identical_logits() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let vocab = 12;
let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6];
let targets: Vec<i32> = vec![1, 4, 1, 5, 9, 2, 6, 0];
let ids_t = ids_tensor(&ids, device);
let targets_t = ids_tensor(&targets, device);
// --- Train a few steps so the params are non-trivial (not the init). ---
let model = make_model(device, vocab, 1);
let params = model.params();
let mut opt = AdamW::new(0.01, 0.1);
for _ in 0..5 {
let loss = model.loss(&ids_t, &targets_t);
loss.backward();
opt.step(0.01, &params);
for p in &params {
p.zero_grad();
}
}
let path = std::env::temp_dir().join(format!("xtrain_ckpt_{}.bin", std::process::id()));
checkpoint::save(&path, &params).unwrap();
let ref_logits = param_to_host(&model.forward(&ids_t));
let ref_loss = param_to_host(&model.loss(&ids_t, &targets_t))[0];
// --- Fresh model with a DIFFERENT init; loading must overwrite it exactly. ---
let fresh = make_model(device, vocab, 999);
let fresh_params = fresh.params();
// Sanity: before load, the fresh model disagrees.
let pre = param_to_host(&fresh.forward(&ids_t));
let pre_diff: f32 = pre
.iter()
.zip(&ref_logits)
.map(|(a, b)| (a - b).abs())
.fold(0.0, f32::max);
assert!(
pre_diff > 1e-4,
"fresh model unexpectedly matched before load"
);
checkpoint::load_into(&path, &fresh_params).unwrap();
let got_logits = param_to_host(&fresh.forward(&ids_t));
let got_loss = param_to_host(&fresh.loss(&ids_t, &targets_t))[0];
let _ = std::fs::remove_file(&path);
// Exact f32 round-trip → bit-for-bit identical forward (same kernels, same
// inputs). Allow only float noise from re-running the forward.
let max_logit_diff: f32 = got_logits
.iter()
.zip(&ref_logits)
.map(|(a, b)| (a - b).abs())
.fold(0.0, f32::max);
println!(
"checkpoint round-trip: max logit diff = {max_logit_diff:.3e}, loss {ref_loss:.6} vs {got_loss:.6}"
);
assert!(
max_logit_diff < 1e-5,
"logits differ after reload: {max_logit_diff:e}"
);
assert!(
(got_loss - ref_loss).abs() < 1e-5,
"loss differs after reload: {ref_loss} vs {got_loss}"
);
}

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

@@ -0,0 +1,83 @@
// 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 Normal file
View File

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

295
crates/xtrain-train/tests/grad_accum.rs Normal file
View File

@@ -0,0 +1,295 @@
// T16 gradient-accumulation correctness gates.
//
// Gradient accumulation is mathematically EXACT: accumulating the grads of N
// micro-batches of B sequences (each micro-loss scaled by 1/N before backward,
// the tape SUM-accumulating) equals a single step over one N·B-sequence batch.
// This file makes that a closed loop on-GPU, plus the accum_steps=1 bit-identity
// regression guard.
//
// 1. accum_equiv_big_batch: same init, same N·B sequences in the same order.
// Path A = ONE batched loss over all N·B (the big-batch baseline). Path B =
// N micro-backwards of B each, scale(1/N), tape SUM. Assert loss and EVERY
// parameter grad match within fp tolerance (only the summation order differs,
// like the T8 DDP-vs-single-GPU and T13 recompute gates).
// 2. accum1_bit_identical: accum_steps=1 must reproduce the no-accum path
// bit-for-bit (the implementation skips the ×1/1 scale entirely) — every
// parameter grad max|Δ| == 0.0.
// 3. accum_train_converges: drive the real `train()` loop with accum and assert
// the per-step effective-batch loss trace tracks a big-batch baseline (errors
// stay bounded over many AdamW steps, not just one).
#![cfg(not(no_cuda))]
use xtrain_autodiff::ops;
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
use xtrain_tensor::Device;
use xtrain_train::data::Corpus;
use xtrain_train::schedule::LrSchedule;
use xtrain_train::{TrainConfig, train};
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)
}
})
}
fn host(t: &xtrain_tensor::Tensor) -> Vec<f32> {
t.to_device(Device::Cpu).as_slice::<f32>().to_vec()
}
// `n` deterministic (seq, target) pairs for the equivalence tests.
fn make_seqs(n: usize, seq: usize, vocab: usize) -> (Vec<Vec<i32>>, Vec<Vec<i32>>) {
let seqs = (0..n)
.map(|b| {
(0..seq)
.map(|i| ((b * 7 + i * 3 + 1) % vocab) as i32)
.collect()
})
.collect();
let tgts = (0..n)
.map(|b| {
(0..seq)
.map(|i| ((b * 5 + i * 2 + 2) % vocab) as i32)
.collect()
})
.collect();
(seqs, tgts)
}
// Run one big-batch forward/backward over all `seqs` and return the grads.
fn big_batch_grads(
model: &TinyTransformer,
device: Device,
seqs: &[Vec<i32>],
tgts: &[Vec<i32>],
) -> (f32, Vec<Vec<f32>>) {
let n = seqs.len();
let ids = batched_ids_tensor(seqs, device);
let tgt = batched_ids_tensor(tgts, device);
let loss = model.loss_batched(&ids, &tgt, n);
let loss_val = host(&loss.value())[0];
loss.backward();
let grads = model
.params()
.iter()
.map(|p| host(&p.grad().expect("grad")))
.collect();
(loss_val, grads)
}
// Accumulate over `accum` micro-batches of `b` sequences (drawn in order from the
// flat `seqs`/`tgts`), scaling each micro-loss by 1/accum before backward; the
// tape SUM-accumulates. Returns the mean of the raw micro losses + accumulated grads.
fn accum_grads(
model: &TinyTransformer,
device: Device,
seqs: &[Vec<i32>],
tgts: &[Vec<i32>],
accum: usize,
b: usize,
scale: bool,
) -> (f32, Vec<Vec<f32>>) {
let mut loss_sum = 0.0f32;
for m in 0..accum {
let s = &seqs[m * b..(m + 1) * b];
let t = &tgts[m * b..(m + 1) * b];
let ids = batched_ids_tensor(s, device);
let tgt = batched_ids_tensor(t, device);
let loss = model.loss_batched(&ids, &tgt, b);
loss_sum += host(&loss.value())[0];
if scale {
ops::scale(&loss, 1.0 / accum as f32).backward();
} else {
loss.backward(); // accum==1 bit-identity path
}
}
let grads = model
.params()
.iter()
.map(|p| host(&p.grad().expect("grad")))
.collect();
(loss_sum / accum as f32, grads)
}
#[test]
fn accum_equiv_big_batch() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let mut cfg = Config::tiny();
cfg.vocab = 16;
cfg.n_layers = 3;
let b = 2usize; // micro-batch
let accum = 4usize; // → effective batch 8
let seq = 6usize;
let (seqs, tgts) = make_seqs(b * accum, seq, cfg.vocab);
// Big-batch baseline (accum_steps=1, batch = b·accum).
let big = build(cfg, device);
let (big_loss, big_grads) = big_batch_grads(&big, device, &seqs, &tgts);
// Accumulated (accum micro-batches of b, scale 1/accum).
let acc = build(cfg, device);
let (acc_loss, acc_grads) = accum_grads(&acc, device, &seqs, &tgts, accum, b, true);
let loss_rel = (big_loss - acc_loss).abs() / big_loss.abs().max(1e-4);
let mut max_grad_rel = 0.0f32;
for (bg, ag) in big_grads.iter().zip(&acc_grads) {
for (x, y) in bg.iter().zip(ag) {
max_grad_rel = max_grad_rel.max((x - y).abs() / x.abs().max(1e-3));
}
}
println!(
"accum=={accum}×b{b} vs big-batch{}: loss {big_loss:.6}/{acc_loss:.6} (rel {loss_rel:.2e}), \
grad max rel {max_grad_rel:.3e}",
b * accum
);
// fp summation order differs (big batch sums b·accum rows once; accum sums per
// micro then across micros) → tight fp tol, same convention as T13 recompute.
assert!(loss_rel < 1e-5, "loss diverged: {loss_rel:.2e}");
assert!(
max_grad_rel < 1e-4,
"accum grads diverged from big batch: {max_grad_rel:.3e}"
);
}
#[test]
fn accum1_bit_identical() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let mut cfg = Config::tiny();
cfg.vocab = 16;
cfg.n_layers = 3;
let b = 4usize;
let seq = 6usize;
let (seqs, tgts) = make_seqs(b, seq, cfg.vocab);
// No-accum reference: one batched loss + backward (the pre-T16 path).
let reference = build(cfg, device);
let (_, ref_grads) = big_batch_grads(&reference, device, &seqs, &tgts);
// accum_steps=1 path: the loop runs ONE micro-batch and (by design) skips the
// ×1/1 scale → must be byte-for-byte identical to the reference backward.
let accum1 = build(cfg, device);
let (_, a1_grads) = accum_grads(&accum1, device, &seqs, &tgts, 1, b, false);
let mut max_abs = 0.0f32;
for (r, a) in ref_grads.iter().zip(&a1_grads) {
for (x, y) in r.iter().zip(a) {
max_abs = max_abs.max((x - y).abs());
}
}
println!("accum_steps=1 vs no-accum: grad max |Δ| = {max_abs:.3e}");
assert_eq!(
max_abs, 0.0,
"accum_steps=1 not bit-identical to no-accum: {max_abs:.3e}"
);
}
// A self-contained synthetic corpus (no tokenizer / data file needed).
fn synth_corpus(vocab: usize, n_tokens: usize) -> Corpus {
Corpus {
tokens: (0..n_tokens)
.map(|i| (i * 7 + 3) as i32 % vocab as i32)
.collect(),
labels: None,
vocab_size: vocab,
}
}
#[test]
fn accum_train_converges() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let vocab = 64usize;
let mut cfg = Config::tiny();
cfg.vocab = vocab;
cfg.n_layers = 2;
let corpus = synth_corpus(vocab, 4096);
let steps = 20usize;
let seq = 32usize;
// Same per-step RNG stream + effective batch 8 either way: the big-batch run
// (accum=1, batch=8) and the accumulated run (accum=4, batch=2) draw the SAME
// 8 sequences per step in the same order, so the per-step loss/grads — and thus
// the whole AdamW trajectory — track within fp tolerance.
let sched = LrSchedule {
max_lr: 3e-3,
min_lr: 3e-4,
warmup: 3,
total: steps,
};
let base = |batch, accum| TrainConfig {
seq_len: seq,
batch_size: batch,
accum_steps: accum,
steps,
schedule: sched.clone(),
weight_decay: 0.1,
max_grad_norm: 1.0,
log_every: 1_000_000,
ckpt_path: None,
ckpt_every: 0,
eval_every: 0,
eval_batches: 0,
seed: 7,
};
let big_model = build(cfg, device);
let big = train(&big_model, device, &corpus, None, &base(8, 1)).train_losses;
let acc_model = build(cfg, device);
let acc = train(&acc_model, device, &corpus, None, &base(2, 4)).train_losses;
let mut max_rel = 0.0f32;
for (x, y) in big.iter().zip(&acc) {
max_rel = max_rel.max((x - y).abs() / x.abs().max(1e-6));
}
// Final params should also stay close (errors don't blow up over the run).
let mut max_pdiff = 0.0f32;
for (p, q) in big_model.params().iter().zip(&acc_model.params()) {
for (x, y) in host(&p.value()).iter().zip(host(&q.value())) {
max_pdiff = max_pdiff.max((x - y).abs() / x.abs().max(1e-6));
}
}
println!(
"accum(4×2) vs big(8) over {steps} steps: loss[last] {:.6}/{:.6} max_rel {max_rel:.2e}, \
final param max rel {max_pdiff:.2e}",
big.last().unwrap(),
acc.last().unwrap()
);
assert!(
max_rel < 1e-3,
"accum loss trajectory diverged: {max_rel:.3e}"
);
assert!(
max_pdiff < 1e-2,
"accum final params diverged: {max_pdiff:.3e}"
);
}

130
crates/xtrain-train/tests/real_training.rs Normal file
View File

@@ -0,0 +1,130 @@
// Real-training acceptance (Phase T6): train the tiny transformer on the
// TinyStories corpus (tokenized with the reused GPT-2 BPE) for a BOUNDED budget
// and assert the loss decreases substantially — the end-to-end signal that the
// whole stack (data pipeline, AdamW, LR schedule, grad clip) learns. Prints the
// loss curve and a couple of greedy samples.
//
// Needs the corpus + tokenizer present, so it is #[ignore] (run with --ignored)
// and gated #![cfg(not(no_cuda))]. Paths are overridable via env vars.
//
// Run: cargo test -p xtrain-train --release --test real_training \
// -- --ignored --nocapture
#![cfg(not(no_cuda))]
use std::path::PathBuf;
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer};
use xtrain_tensor::Device;
use xtrain_train::data::Corpus;
use xtrain_train::sample::generate;
use xtrain_train::schedule::LrSchedule;
use xtrain_train::{TrainConfig, train};
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()
}
#[test]
#[ignore = "real training; needs corpus + tokenizer; run with --ignored --release"]
fn trains_on_tinystories() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let tok_path = PathBuf::from(
std::env::var("XTRAIN_TOKENIZER")
.unwrap_or_else(|_| "/opt/wjh/models/gpt2/tokenizer.json".into()),
);
// Default resolves relative to the repo root (cargo runs tests with cwd =
// crate dir, so `../../data/...` from crates/xtrain-train); override with
// XTRAIN_CORPUS for any other location.
let corpus_path = PathBuf::from(std::env::var("XTRAIN_CORPUS").unwrap_or_else(|_| {
format!(
"{}/../../data/tinystories-valid-3mb.txt",
env!("CARGO_MANIFEST_DIR")
)
}));
let corpus = Corpus::load(&tok_path, &corpus_path);
println!(
"corpus: {} tokens, vocab {}",
corpus.len(),
corpus.vocab_size
);
let mut cfg = Config::tiny();
cfg.vocab = corpus.vocab_size;
cfg.n_layers = 4;
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)
}
});
let steps = std::env::var("XTRAIN_STEPS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(800usize);
let tcfg = TrainConfig {
seq_len: 64,
batch_size: 8,
accum_steps: 1,
steps,
schedule: LrSchedule {
max_lr: 3e-3,
min_lr: 3e-4,
warmup: (steps / 20).max(20),
total: steps,
},
weight_decay: 0.1,
max_grad_norm: 1.0,
log_every: 50,
ckpt_path: None,
ckpt_every: 0,
eval_every: 0,
eval_batches: 0,
seed: 42,
};
let losses = train(&model, device, &corpus, None, &tcfg).train_losses;
// Average the first/last few steps to smooth per-step noise.
let head: f32 =
losses[..10.min(losses.len())].iter().sum::<f32>() / 10.0_f32.min(losses.len() as f32);
let tail_n = 10.min(losses.len());
let tail: f32 = losses[losses.len() - tail_n..].iter().sum::<f32>() / tail_n as f32;
println!("loss: start(avg10) {head:.4} → end(avg10) {tail:.4}");
// A couple of greedy samples (should show English structure, not gibberish).
use xserv_tokenizer::Tokenizer;
let tok = Tokenizer::from_file(&tok_path);
for p in ["Once upon a time", "The little"] {
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, 40, 0.0, &mut rng);
let text = tok.decode(&out.iter().map(|&t| t as u32).collect::<Vec<_>>());
println!("sample [{p}] → {text}");
}
// Bounded run: expect a substantial drop (not full convergence).
assert!(
tail < head - 0.5,
"loss did not decrease substantially: {head:.4} → {tail:.4}"
);
assert!(tail < 6.5, "final loss implausibly high: {tail:.4}");
}

93
csrc/ops/attention.cu Normal file
View File

@@ -0,0 +1,93 @@
// Batched scaled-dot-product attention helpers (Phase T10).
//
// The QKᵀ and PV matmuls run as cublasSgemmStridedBatched in Rust; the only
// kernel attention needs of its own is a CAUSAL row-wise softmax over the score
// rows. Scores are [B*nh, S, S] flattened to rows of length S; for a flat row r
// the query position within its sequence is `r % S`, so columns j > r%S are
// future positions and get probability 0 (no additive -1e9 mask tensor needed).
//
// The forward also folds in the 1/sqrt(head_dim) scale (applied to logits before
// the max/exp) so we don't need a separate scale pass. Backward is the ordinary
// softmax Jacobian (csrc/ops/nn.cu launch_softmax_dx_f32): masked entries have
// y=0, so their contribution vanishes — no causal-specific backward needed.
//
// All F32, row-major, contiguous. Reduction helpers mirror nn.cu (inlined so the
// file is self-contained, matching the csrc/ layout).
#include <math.h>
extern "C" {
__device__ __forceinline__ float att_warp_sum(float v) {
#pragma unroll
for (int off = 16; off > 0; off >>= 1)
v += __shfl_down_sync(0xffffffff, v, off);
return v;
}
__device__ __forceinline__ float att_warp_max(float v) {
#pragma unroll
for (int off = 16; off > 0; off >>= 1)
v = fmaxf(v, __shfl_down_sync(0xffffffff, v, off));
return v;
}
__device__ __forceinline__ float att_block_sum(float v) {
__shared__ float sh[32];
int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
int nwarps = (blockDim.x + 31) >> 5;
v = att_warp_sum(v);
if (lane == 0) sh[warp] = v;
__syncthreads();
v = (threadIdx.x < nwarps) ? sh[threadIdx.x] : 0.0f;
if (warp == 0) v = att_warp_sum(v);
__shared__ float bc;
if (threadIdx.x == 0) bc = v;
__syncthreads();
return bc;
}
__device__ __forceinline__ float att_block_max(float v) {
__shared__ float sh[32];
int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
int nwarps = (blockDim.x + 31) >> 5;
v = att_warp_max(v);
if (lane == 0) sh[warp] = v;
__syncthreads();
v = (threadIdx.x < nwarps) ? sh[threadIdx.x] : -INFINITY;
if (warp == 0) v = att_warp_max(v);
__shared__ float bc;
if (threadIdx.x == 0) bc = v;
__syncthreads();
return bc;
}
// One block per score row. rows = B*nh*S total; the query position within its
// sequence is (blockIdx.x % seq). Logits are scaled by `scale` (= 1/sqrt(hd))
// before softmax; columns j > qpos are masked to probability 0.
__global__ void softmax_causal_k(const float* x, float* y, int seq, float scale) {
int r = blockIdx.x;
int qpos = r % seq;
const float* xr = x + (size_t)r * seq;
float* yr = y + (size_t)r * seq;
int valid = qpos + 1; // attend to columns [0, qpos]
float m = -INFINITY;
for (int c = threadIdx.x; c < valid; c += blockDim.x)
m = fmaxf(m, xr[c] * scale);
m = att_block_max(m);
float sum = 0.0f;
for (int c = threadIdx.x; c < valid; c += blockDim.x) {
float e = expf(xr[c] * scale - m);
yr[c] = e;
sum += e;
}
sum = att_block_sum(sum);
float inv = 1.0f / sum;
for (int c = threadIdx.x; c < seq; c += blockDim.x)
yr[c] = (c < valid) ? yr[c] * inv : 0.0f;
}
void launch_softmax_causal_f32(const float* x, float* y, int rows, int seq,
float scale, void* s) {
int blk = seq < 1024 ? seq : 1024;
if (blk < 32) blk = 32;
softmax_causal_k<<<rows, blk, 0, (cudaStream_t)s>>>(x, y, seq, scale);
}
} // extern "C"

141
csrc/ops/cast.cu Normal file
View File

@@ -0,0 +1,141 @@
// bf16 mixed-precision kernels (Phase T12, KI-2).
//
// Two groups:
// 1. f32 <-> bf16 cast — the bridge between fp32 master weights / fp32
// reductions and the bf16 compute/activation stream.
// 2. bf16 elementwise ops (add / mul / silu / scale + their backwards) — the
// residual-stream ops that flow bf16 activations. Each loads bf16 -> float,
// computes in fp32, stores bf16 (so the math accumulates in fp32 while the
// stored activation is half-size). Matmuls go through cuBLAS GemmEx
// (cublas.rs); norm / softmax / rope / cross-entropy stay fp32 (the Rust
// wrappers upcast around the existing fp32 kernels).
//
// bf16 is __nv_bfloat16; __float2bfloat16 / __bfloat162float round-trip via fp32.
#include <cuda_bf16.h>
extern "C" {
// --- f32 <-> bf16 cast ---
__global__ void cast_f32_to_bf16_k(const float* in, __nv_bfloat16* out, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) out[i] = __float2bfloat16(in[i]);
}
void launch_cast_f32_to_bf16(const float* in, void* out, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
cast_f32_to_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
in, (__nv_bfloat16*)out, n);
}
__global__ void cast_bf16_to_f32_k(const __nv_bfloat16* in, float* out, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) out[i] = __bfloat162float(in[i]);
}
void launch_cast_bf16_to_f32(const void* in, float* out, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
cast_bf16_to_f32_k<<<grid, blk, 0, (cudaStream_t)s>>>(
(const __nv_bfloat16*)in, out, n);
}
// --- bf16 elementwise (load->fp32->compute->store bf16) ---
__global__ void add_bf16_k(const __nv_bfloat16* a, const __nv_bfloat16* b,
__nv_bfloat16* out, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n)
out[i] = __float2bfloat16(__bfloat162float(a[i]) + __bfloat162float(b[i]));
}
void launch_add_bf16(const void* a, const void* b, void* out, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
add_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
(const __nv_bfloat16*)a, (const __nv_bfloat16*)b, (__nv_bfloat16*)out, n);
}
__global__ void mul_bf16_k(const __nv_bfloat16* a, const __nv_bfloat16* b,
__nv_bfloat16* out, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n)
out[i] = __float2bfloat16(__bfloat162float(a[i]) * __bfloat162float(b[i]));
}
void launch_mul_bf16(const void* a, const void* b, void* out, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
mul_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
(const __nv_bfloat16*)a, (const __nv_bfloat16*)b, (__nv_bfloat16*)out, n);
}
__global__ void scale_bf16_k(const __nv_bfloat16* in, __nv_bfloat16* out,
float alpha, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) out[i] = __float2bfloat16(__bfloat162float(in[i]) * alpha);
}
void launch_scale_bf16(const void* in, void* out, float alpha, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
scale_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, alpha, n);
}
// SiLU: y = x*sigmoid(x). Backward: dx = dy * (sig + x*sig*(1-sig)).
__global__ void silu_bf16_k(const __nv_bfloat16* x, __nv_bfloat16* y, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float v = __bfloat162float(x[i]);
float sig = 1.0f / (1.0f + expf(-v));
y[i] = __float2bfloat16(v * sig);
}
}
void launch_silu_bf16(const void* x, void* y, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
silu_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
(const __nv_bfloat16*)x, (__nv_bfloat16*)y, n);
}
__global__ void silu_dx_bf16_k(const __nv_bfloat16* x, const __nv_bfloat16* dy,
__nv_bfloat16* dx, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float v = __bfloat162float(x[i]);
float sig = 1.0f / (1.0f + expf(-v));
float g = sig + v * sig * (1.0f - sig);
dx[i] = __float2bfloat16(__bfloat162float(dy[i]) * g);
}
}
void launch_silu_dx_bf16(const void* x, const void* dy, void* dx, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
silu_dx_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
(const __nv_bfloat16*)x, (const __nv_bfloat16*)dy, (__nv_bfloat16*)dx, n);
}
// Broadcast bias add: out[r,c] = x[r,c] + bias[c]. x:[rows,cols], bias:[cols].
__global__ void add_bias_bf16_k(const __nv_bfloat16* x, const __nv_bfloat16* bias,
__nv_bfloat16* out, int rows, int cols) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < rows * cols)
out[i] = __float2bfloat16(__bfloat162float(x[i]) +
__bfloat162float(bias[i % cols]));
}
void launch_add_bias_bf16(const void* x, const void* bias, void* out, int rows,
int cols, void* s) {
int blk = 256, grid = (rows * cols + blk - 1) / blk;
add_bias_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
(const __nv_bfloat16*)x, (const __nv_bfloat16*)bias, (__nv_bfloat16*)out,
rows, cols);
}
// Column-sum over rows: dbias[c] = sum_r dout[r,c] (bias backward), fp32 accum.
__global__ void sum_rows_bf16_k(const __nv_bfloat16* dout, __nv_bfloat16* dbias,
int rows, int cols) {
int c = blockIdx.x * blockDim.x + threadIdx.x;
if (c < cols) {
float acc = 0.0f;
for (int r = 0; r < rows; ++r) acc += __bfloat162float(dout[r * cols + c]);
dbias[c] = __float2bfloat16(acc);
}
}
void launch_sum_rows_bf16(const void* dout, void* dbias, int rows, int cols, void* s) {
int blk = 256, grid = (cols + blk - 1) / blk;
sum_rows_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
(const __nv_bfloat16*)dout, (__nv_bfloat16*)dbias, rows, cols);
}
} // extern "C"

109
csrc/ops/dropout.cu Normal file
View File

@@ -0,0 +1,109 @@
// Dropout kernels (Phase T18).
//
// A counter-based (stateless) RNG: the keep/drop decision for element `i` is a
// pure function of (seed, i) — no global RNG state is advanced. This is what
// makes dropout compatible with activation recomputation (T13): when a
// checkpointed block re-runs its forward in backward, the SAME seed regenerates
// the SAME mask, so the recomputed activations / grads stay bit-identical to the
// forward (no mask drift).
//
// Inverted dropout: at training time kept elements are scaled by 1/(1-p) so the
// expectation E[out] == x (no inference-time rescale needed; eval is identity,
// handled in Rust by simply not calling dropout).
//
// key = seed ^ (i * GOLDEN)
// h = splitmix64(key) // a few rounds of xorshift/multiply
// u = (h >> 40) / 2^24 in [0,1) // 24-bit uniform
// keep = u >= p // Bernoulli(keep = 1-p)
// out = keep ? x * scale : 0 // scale = 1/(1-p)
// mask = keep ? scale : 0 // cached for backward (dx = d * mask)
//
// fp32 + bf16 variants: bf16 loads/stores half-size activations but the uniform
// `u` is always computed in fp32, so the mask distribution is identical across
// dtypes (drop decisions don't depend on bf16 rounding). The mask buffer is fp32
// in both cases (it stores `scale` or 0 — exactly representable, tiny relative to
// the activation, reused only elementwise in backward).
#include <cuda_bf16.h>
#include <stdint.h>
extern "C" {
// splitmix64: cheap, well-mixed counter hash. Maps a 64-bit counter to a 64-bit
// pseudo-random output; we only need the high bits for a uniform.
__device__ __forceinline__ uint64_t splitmix64(uint64_t x) {
x += 0x9E3779B97F4A7C15ULL;
x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ULL;
x = (x ^ (x >> 27)) * 0x94D049BB133111EBULL;
return x ^ (x >> 31);
}
// Uniform [0,1) for element i under `seed`, computed in fp32 (dtype-independent).
__device__ __forceinline__ float dropout_uniform(uint64_t seed, int i) {
uint64_t key = seed ^ ((uint64_t)i * 0x9E3779B97F4A7C15ULL);
uint64_t h = splitmix64(key);
// Top 24 bits → [0,1) with 2^-24 resolution.
return (float)(h >> 40) * (1.0f / 16777216.0f); // 1/2^24
}
// fp32 forward: out[i] = keep ? x[i]*scale : 0 ; mask[i] = keep ? scale : 0.
__global__ void dropout_fwd_f32_k(const float* x, float* out, float* mask,
float p, float scale, uint64_t seed, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float keep = (dropout_uniform(seed, i) >= p) ? scale : 0.0f;
mask[i] = keep;
out[i] = x[i] * keep;
}
}
void launch_dropout_fwd_f32(const float* x, float* out, float* mask, float p,
float scale, uint64_t seed, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
dropout_fwd_f32_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, out, mask, p, scale,
seed, n);
}
// Backward applies the SAME cached mask elementwise: dx[i] = d[i] * mask[i].
__global__ void dropout_bwd_f32_k(const float* d, const float* mask, float* dx,
int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) dx[i] = d[i] * mask[i];
}
void launch_dropout_bwd_f32(const float* d, const float* mask, float* dx, int n,
void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
dropout_bwd_f32_k<<<grid, blk, 0, (cudaStream_t)s>>>(d, mask, dx, n);
}
// bf16 forward: activation is bf16; mask is fp32 (stores `scale` or 0). Uniform
// is fp32, so the mask matches the fp32 path bit-for-bit (same drop decisions).
__global__ void dropout_fwd_bf16_k(const __nv_bfloat16* x, __nv_bfloat16* out,
float* mask, float p, float scale,
uint64_t seed, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float keep = (dropout_uniform(seed, i) >= p) ? scale : 0.0f;
mask[i] = keep;
out[i] = __float2bfloat16(__bfloat162float(x[i]) * keep);
}
}
void launch_dropout_fwd_bf16(const void* x, void* out, float* mask, float p,
float scale, uint64_t seed, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
dropout_fwd_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
(const __nv_bfloat16*)x, (__nv_bfloat16*)out, mask, p, scale, seed, n);
}
__global__ void dropout_bwd_bf16_k(const __nv_bfloat16* d, const float* mask,
__nv_bfloat16* dx, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) dx[i] = __float2bfloat16(__bfloat162float(d[i]) * mask[i]);
}
void launch_dropout_bwd_bf16(const void* d, const float* mask, void* dx, int n,
void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
dropout_bwd_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
(const __nv_bfloat16*)d, mask, (__nv_bfloat16*)dx, n);
}
} // extern "C"

281
csrc/ops/flash_attention.cu Normal file
View File

@@ -0,0 +1,281 @@
// Hand-written fused flash-attention (Phase T14).
//
// The T10 composed SDPA path is 3 launches that MATERIALIZE the [bh,S,S] score
// matrix: cublasSgemmStridedBatched (Q·Kᵀ) → causal-softmax kernel (writes the
// whole probs) → cublasSgemmStridedBatched (P·V), and backward caches that whole
// probs. flash-attention NEVER materializes N×N: a single fused kernel streams
// over KV tiles with an ONLINE softmax (running max/sum + rescaled V accumulator),
// so peak attention activation drops from O(S²) to O(S·hd) (= the output itself).
//
// Layout (matches the T10 op): Q/K/V/out are [bh, S, hd] row-major contiguous,
// bh = batch·n_heads. The query's position within its sequence is the row index
// within its [S,hd] block (so the flat row's qpos = (row % S) is automatic here —
// we index per (bh, row)). CAUSAL: a query at position i attends to keys j ≤ i.
// `scale` (= 1/sqrt(hd)) is folded into the logits before the max/exp.
//
// All F32, contiguous. (bf16 callers upcast Q/K/V → f32 on the Rust side and
// downcast the f32 out, mirroring the composed path's fp32 softmax policy, so the
// kernel only ever sees fp32.) Reduction helpers are inlined (self-contained file,
// matching the csrc/ layout).
//
// Parallelisation: grid = bh*S, one block per query row; blockDim.x threads
// cooperate. Forward keeps m (running max), l (running sum), acc[hd] (rescaled
// V accumulator) in shared memory, streams KV in tiles of BK. Backward recomputes
// scores from Q/K/V + the saved logsumexp L[bh,S] (NO cached probs), uses
// D[i]=Σ dOᵢ·Oᵢ to collapse the softmax Jacobian, and atomicAdds dK/dV (which are
// accumulated across query rows).
#include <math.h>
extern "C" {
__device__ __forceinline__ float fa_warp_sum(float v) {
#pragma unroll
for (int off = 16; off > 0; off >>= 1)
v += __shfl_down_sync(0xffffffff, v, off);
return v;
}
__device__ __forceinline__ float fa_warp_max(float v) {
#pragma unroll
for (int off = 16; off > 0; off >>= 1)
v = fmaxf(v, __shfl_down_sync(0xffffffff, v, off));
return v;
}
__device__ __forceinline__ float fa_block_sum(float v) {
__shared__ float sh[32];
int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
int nwarps = (blockDim.x + 31) >> 5;
v = fa_warp_sum(v);
if (lane == 0) sh[warp] = v;
__syncthreads();
v = (threadIdx.x < nwarps) ? sh[threadIdx.x] : 0.0f;
if (warp == 0) v = fa_warp_sum(v);
__shared__ float bc;
if (threadIdx.x == 0) bc = v;
__syncthreads();
return bc;
}
__device__ __forceinline__ float fa_block_max(float v) {
__shared__ float sh[32];
int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
int nwarps = (blockDim.x + 31) >> 5;
v = fa_warp_max(v);
if (lane == 0) sh[warp] = v;
__syncthreads();
v = (threadIdx.x < nwarps) ? sh[threadIdx.x] : -INFINITY;
if (warp == 0) v = fa_warp_max(v);
__shared__ float bc;
if (threadIdx.x == 0) bc = v;
__syncthreads();
return bc;
}
#define FA_TILE 32 // KV tile width (columns streamed per step)
// One block per (bh-row, query-position). Computes out[bh, i, :] and L[bh, i] via
// an online softmax that streams the keys in tiles of FA_TILE — the [S,S] score
// row is never stored, only the per-tile partials flow through shared memory.
__global__ void flash_attn_fwd_k(const float* Q, const float* K, const float* V,
float* O, float* L, int seq, int hd, float scale) {
int row = blockIdx.x; // global query row over bh*S
int b = row / seq; // which (batch,head) block
int i = row % seq; // query position within the sequence (causal limit)
int t = threadIdx.x;
int nthreads = blockDim.x;
const float* q = Q + (size_t)row * hd;
const float* kb = K + (size_t)b * seq * hd; // this block's keys [seq,hd]
const float* vb = V + (size_t)b * seq * hd; // this block's values[seq,hd]
// Q row in shared memory (reused every tile); acc accumulator over hd.
extern __shared__ float smem[];
float* sq = smem; // [hd]
float* acc = smem + hd; // [hd]
for (int d = t; d < hd; d += nthreads) {
sq[d] = q[d];
acc[d] = 0.0f;
}
__shared__ float m_run, l_run;
if (t == 0) { m_run = -INFINITY; l_run = 0.0f; }
__syncthreads();
int valid = i + 1; // causal: attend to keys [0, i]
for (int j0 = 0; j0 < valid; j0 += FA_TILE) {
int tile = min(FA_TILE, valid - j0);
// Each thread computes whole logits for a strided subset of the tile's
// columns: s = scale * (q · k_j). hd is small (≤128) so the per-thread
// dot loop is cheap; this avoids a block-reduce per column.
__shared__ float s_tile[FA_TILE];
for (int c = t; c < tile; c += nthreads) {
const float* kj = kb + (size_t)(j0 + c) * hd;
float dot = 0.0f;
for (int d = 0; d < hd; ++d) dot += sq[d] * kj[d];
s_tile[c] = dot * scale;
}
__syncthreads();
// Tile max, then online rescale of (m, l, acc).
float tmax = -INFINITY;
for (int c = t; c < tile; c += nthreads) tmax = fmaxf(tmax, s_tile[c]);
tmax = fa_block_max(tmax);
__shared__ float m_new, corr;
if (t == 0) {
float mn = fmaxf(m_run, tmax);
corr = (m_run == -INFINITY) ? 0.0f : expf(m_run - mn); // rescale old state
m_new = mn;
}
__syncthreads();
// Overwrite s_tile with the softmax weights p = exp(s - m_new) ONCE per
// column (instead of recomputing expf inside the per-dim V loop, which
// would cost hd× the transcendentals). Sum them for l.
float lsum = 0.0f;
for (int c = t; c < tile; c += nthreads) {
float p = expf(s_tile[c] - m_new);
s_tile[c] = p;
lsum += p;
}
lsum = fa_block_sum(lsum);
// Rescale old accumulator + add this tile's p·V (p cached in s_tile).
// Each thread owns a strided subset of hd; loops over the tile columns.
for (int d = t; d < hd; d += nthreads) {
float a = acc[d] * corr;
for (int c = 0; c < tile; ++c)
a += s_tile[c] * vb[(size_t)(j0 + c) * hd + d];
acc[d] = a;
}
if (t == 0) {
l_run = l_run * corr + lsum;
m_run = m_new;
}
__syncthreads();
}
// out = acc / l ; L = m + log(l) (logsumexp, saved for backward).
float inv = 1.0f / l_run;
for (int d = t; d < hd; d += nthreads) O[(size_t)row * hd + d] = acc[d] * inv;
if (t == 0) L[row] = m_run + logf(l_run);
}
void launch_flash_attention_fwd_f32(const float* q, const float* k, const float* v,
float* o, float* l, int bh, int seq, int hd,
float scale, void* s) {
int blk = hd < 1024 ? hd : 1024;
if (blk < 32) blk = 32;
size_t shmem = (size_t)2 * hd * sizeof(float); // sq[hd] + acc[hd]
flash_attn_fwd_k<<<bh * seq, blk, shmem, (cudaStream_t)s>>>(q, k, v, o, l, seq, hd, scale);
}
// Per-row D[i] = Σ_d dO[i,d] · O[i,d]. One block per row (bh*S rows). Used to
// collapse the softmax Jacobian in backward (Σ_j P_ij dP_ij = dOᵢ·Oᵢ).
__global__ void flash_attn_rowdot_k(const float* dO, const float* O, float* D, int hd) {
int row = blockIdx.x;
int t = threadIdx.x;
const float* d = dO + (size_t)row * hd;
const float* o = O + (size_t)row * hd;
float v = 0.0f;
for (int c = t; c < hd; c += blockDim.x) v += d[c] * o[c];
v = fa_block_sum(v);
if (t == 0) D[row] = v;
}
// Backward: one block per query row i. Recomputes scores from Q/K/V + the saved
// logsumexp L (NO cached probs), streams KV in tiles. dQ accumulates locally (this
// row owns it). dK/dV are accumulated ACROSS query rows so they atomicAdd into the
// shared global buffers (pre-zeroed by the caller).
// p_ij = exp(Qᵢ·Kⱼ·scale - L[i]) ; dp_ij = dOᵢ·Vⱼ ;
// ds_ij = p_ij·(dp_ij - D[i])·scale
// dQᵢ += Σ_j ds_ij·Kⱼ ; dKⱼ += ds_ij·Qᵢ ; dVⱼ += p_ij·dOᵢ
__global__ void flash_attn_bwd_k(const float* Q, const float* K, const float* V,
const float* dO, const float* L, const float* D,
float* dQ, float* dK, float* dV,
int seq, int hd, float scale) {
int row = blockIdx.x;
int b = row / seq;
int i = row % seq;
int t = threadIdx.x;
int nthreads = blockDim.x;
const float* q = Q + (size_t)row * hd;
const float* doi = dO + (size_t)row * hd;
const float* kb = K + (size_t)b * seq * hd;
const float* vb = V + (size_t)b * seq * hd;
float* dkb = dK + (size_t)b * seq * hd;
float* dvb = dV + (size_t)b * seq * hd;
extern __shared__ float smem[];
float* sq = smem; // [hd] Qᵢ
float* sdo = smem + hd; // [hd] dOᵢ
float* dqa = smem + 2*hd; // [hd] dQᵢ accumulator
for (int d = t; d < hd; d += nthreads) {
sq[d] = q[d];
sdo[d] = doi[d];
dqa[d] = 0.0f;
}
__shared__ float Li, Di;
if (t == 0) { Li = L[row]; Di = D[row]; }
__syncthreads();
int valid = i + 1;
for (int j0 = 0; j0 < valid; j0 += FA_TILE) {
int tile = min(FA_TILE, valid - j0);
// Phase 1: per-column ds[c] and p[c] (the column owner does the dots).
__shared__ float s_ds[FA_TILE];
__shared__ float s_p[FA_TILE];
for (int c = t; c < tile; c += nthreads) {
const float* kj = kb + (size_t)(j0 + c) * hd;
const float* vj = vb + (size_t)(j0 + c) * hd;
float sdot = 0.0f, dpdot = 0.0f;
for (int d = 0; d < hd; ++d) {
sdot += sq[d] * kj[d];
dpdot += sdo[d] * vj[d];
}
float p = expf(sdot * scale - Li);
s_p[c] = p;
s_ds[c] = p * (dpdot - Di) * scale;
}
__syncthreads();
// Phase 2: dV_j += p·dOᵢ ; dK_j += ds·Qᵢ (accumulated across rows → atomic).
// Spread the tile×hd atomics over ALL threads (was serial in the column
// owner) — flatten (c,d) so every thread issues a balanced share.
for (int idx = t; idx < tile * hd; idx += nthreads) {
int c = idx / hd, d = idx % hd;
size_t off = (size_t)(j0 + c) * hd + d;
atomicAdd(&dvb[off], s_p[c] * sdo[d]);
atomicAdd(&dkb[off], s_ds[c] * sq[d]);
}
// dQᵢ += Σ_c ds[c] · K_{j0+c} (this row owns dQ — no atomic).
for (int d = t; d < hd; d += nthreads) {
float a = 0.0f;
for (int c = 0; c < tile; ++c)
a += s_ds[c] * kb[(size_t)(j0 + c) * hd + d];
dqa[d] += a;
}
__syncthreads();
}
for (int d = t; d < hd; d += nthreads) dQ[(size_t)row * hd + d] = dqa[d];
}
void launch_flash_attention_bwd_f32(const float* q, const float* k, const float* v,
const float* d_o, const float* l, float* d_d,
float* dq, float* dk, float* dv,
int bh, int seq, int hd, float scale, void* s) {
int blk = hd < 1024 ? hd : 1024;
if (blk < 32) blk = 32;
// d_d is the pre-computed D[i]=Σ dOᵢ·Oᵢ (the Rust wrapper runs rowdot first,
// since it holds the forward O). dq/dk/dv are pre-zeroed by the caller.
flash_attn_bwd_k<<<bh * seq, blk, (size_t)3 * hd * sizeof(float), (cudaStream_t)s>>>(
q, k, v, d_o, l, d_d, dq, dk, dv, seq, hd, scale);
}
// Standalone D = rowdot(dO, O) launcher (the Rust wrapper calls this before bwd).
void launch_flash_attention_rowdot_f32(const float* d_o, const float* o, float* d_d,
int rows, int hd, void* s) {
int blk = hd < 1024 ? hd : 1024;
if (blk < 32) blk = 32;
flash_attn_rowdot_k<<<rows, blk, 0, (cudaStream_t)s>>>(d_o, o, d_d, hd);
}
} // extern "C"

88
csrc/ops/model.cu Normal file
View File

@@ -0,0 +1,88 @@
// Structural ops the tiny transformer (Phase T5) needs on top of the T4 op set:
// token embedding (gather forward / scatter-add backward) and a 3D axis-(0,1)
// transpose used to lay out multi-head attention ([seq,heads,hd] <-> [heads,seq,hd]).
//
// reshape is a pure metadata change (no data movement) and so has no kernel — it
// lives entirely in the Rust Tensor layer. All kernels here are F32 row-major
// contiguous; ids are I32. Each launcher matches the existing csrc/ style.
extern "C" {
// =====================================================================
// Embedding: gather rows of a table by integer ids.
// table:[vocab, dim], ids:[seq] (I32) -> out[s,:] = table[ids[s], :]
// Backward (scatter-add): dtable[ids[s], :] += dout[s, :]. Multiple positions
// may map to the same id, so the accumulation must be atomic.
// =====================================================================
__global__ void embedding_fwd_k(const float* table, const int* ids, float* out,
int seq, int dim) {
int i = blockIdx.x * blockDim.x + threadIdx.x; // over seq*dim
if (i >= seq * dim) return;
int s = i / dim, c = i % dim;
out[i] = table[ids[s] * dim + c];
}
void launch_embedding_fwd_f32(const float* table, const int* ids, float* out,
int seq, int dim, void* s) {
int n = seq * dim, blk = 256, grid = (n + blk - 1) / blk;
embedding_fwd_k<<<grid, blk, 0, (cudaStream_t)s>>>(table, ids, out, seq, dim);
}
// dtable is assumed pre-zeroed (Tensor::zeros). Scatter-add with atomics so
// repeated ids accumulate correctly.
__global__ void embedding_bwd_k(const float* dout, const int* ids, float* dtable,
int seq, int dim) {
int i = blockIdx.x * blockDim.x + threadIdx.x; // over seq*dim
if (i >= seq * dim) return;
int s = i / dim, c = i % dim;
atomicAdd(&dtable[ids[s] * dim + c], dout[i]);
}
void launch_embedding_bwd_f32(const float* dout, const int* ids, float* dtable,
int seq, int dim, void* s) {
int n = seq * dim, blk = 256, grid = (n + blk - 1) / blk;
embedding_bwd_k<<<grid, blk, 0, (cudaStream_t)s>>>(dout, ids, dtable, seq, dim);
}
// =====================================================================
// 3D axis-(0,1) transpose: in:[a,b,c] -> out:[b,a,c] (last dim contiguous).
// out[j, i, k] = in[i, j, k]
// Its own backward is the same op with (a,b) swapped, so one kernel suffices.
// =====================================================================
__global__ void transpose_3d01_k(const float* in, float* out, int a, int b, int c) {
int idx = blockIdx.x * blockDim.x + threadIdx.x; // over a*b*c
if (idx >= a * b * c) return;
int k = idx % c;
int j = (idx / c) % b;
int i = idx / (b * c);
// out index: ((j*a) + i)*c + k
out[(j * a + i) * c + k] = in[idx];
}
void launch_transpose_3d01_f32(const float* in, float* out, int a, int b, int c, void* s) {
int n = a * b * c, blk = 256, grid = (n + blk - 1) / blk;
transpose_3d01_k<<<grid, blk, 0, (cudaStream_t)s>>>(in, out, a, b, c);
}
// =====================================================================
// 4D axis-(1,2) transpose: in:[a,b,c,d] -> out:[a,c,b,d]. out[i,k,j,l]=in[i,j,k,l].
// Lays out batched multi-head attention: [B,S,nh,hd] <-> [B,nh,S,hd], so a
// flattened [B*nh, S, hd] view feeds the strided-batched-GEMM attention. Its own
// backward is the same op (swap b,c), so one kernel suffices.
// =====================================================================
__global__ void transpose_4d12_k(const float* in, float* out, int a, int b, int c, int d) {
int idx = blockIdx.x * blockDim.x + threadIdx.x; // over a*b*c*d
if (idx >= a * b * c * d) return;
int l = idx % d;
int k = (idx / d) % c;
int j = (idx / (d * c)) % b;
int i = idx / (d * c * b);
// out[i,k,j,l] at ((i*c + k)*b + j)*d + l
out[(((i * c + k) * b) + j) * d + l] = in[idx];
}
void launch_transpose_4d12_f32(const float* in, float* out, int a, int b, int c, int d, void* s) {
int n = a * b * c * d, blk = 256, grid = (n + blk - 1) / blk;
transpose_4d12_k<<<grid, blk, 0, (cudaStream_t)s>>>(in, out, a, b, c, d);
}
} // extern "C"

125
csrc/ops/nn.cu
View File

@@ -215,14 +215,20 @@ void launch_silu_dx_f32(const float* x, const float* dy, float* dx, int n, void*
// dx[i+h] = dy[i+h]*cos - dy[i]*sin // dx[i+h] = dy[i+h]*cos - dy[i]*sin
// ===================================================================== // =====================================================================
__global__ void rope_k(const float* x, float* y, int heads, int head_dim, float theta) { // `period` is the sequence length: a flattened batch lays B sequences end to end
// along the `tokens` axis, so each token's RoPE position is its index WITHIN its
// own sequence, `tok % period`. With period == tokens (single sequence) this is
// the original position = row.
__global__ void rope_k(const float* x, float* y, int heads, int head_dim,
float theta, int period) {
int tok = blockIdx.x; int tok = blockIdx.x;
int head = blockIdx.y; int head = blockIdx.y;
int half = head_dim / 2; int half = head_dim / 2;
int i = threadIdx.x; int i = threadIdx.x;
if (i >= half) return; if (i >= half) return;
int pos = tok % period;
float freq = powf(theta, -(float)(2 * i) / (float)head_dim); float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
float angle = (float)tok * freq; float angle = (float)pos * freq;
float c = cosf(angle), sn = sinf(angle); float c = cosf(angle), sn = sinf(angle);
int base = (tok * heads + head) * head_dim; int base = (tok * heads + head) * head_dim;
float x0 = x[base + i], x1 = x[base + i + half]; float x0 = x[base + i], x1 = x[base + i + half];
@@ -230,20 +236,112 @@ __global__ void rope_k(const float* x, float* y, int heads, int head_dim, float
y[base + i + half] = x1 * c + x0 * sn; y[base + i + half] = x1 * c + x0 * sn;
} }
void launch_rope_f32(const float* x, float* y, int tokens, int heads, void launch_rope_f32(const float* x, float* y, int tokens, int heads,
int head_dim, float theta, void* s) { int head_dim, float theta, int period, void* s) {
dim3 grid(tokens, heads); dim3 grid(tokens, heads);
int blk = head_dim / 2; int blk = head_dim / 2;
rope_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta); rope_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta, period);
} }
__global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, float theta) { // 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 tok = blockIdx.x;
int head = blockIdx.y; int head = blockIdx.y;
int half = head_dim / 2; int half = head_dim / 2;
int i = threadIdx.x; int i = threadIdx.x;
if (i >= half) return; if (i >= half) return;
int pos = pos0 + tok;
float freq = powf(theta, -(float)(2 * i) / (float)head_dim); float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
float angle = (float)tok * freq; 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;
int head = blockIdx.y;
int half = head_dim / 2;
int i = threadIdx.x;
if (i >= half) return;
int pos = tok % period;
float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
float angle = (float)pos * freq;
float c = cosf(angle), sn = sinf(angle); float c = cosf(angle), sn = sinf(angle);
int base = (tok * heads + head) * head_dim; int base = (tok * heads + head) * head_dim;
float d0 = dy[base + i], d1 = dy[base + i + half]; float d0 = dy[base + i], d1 = dy[base + i + half];
@@ -251,10 +349,10 @@ __global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, f
dx[base + i + half] = d1 * c - d0 * sn; dx[base + i + half] = d1 * c - d0 * sn;
} }
void launch_rope_dx_f32(const float* dy, float* dx, int tokens, int heads, void launch_rope_dx_f32(const float* dy, float* dx, int tokens, int heads,
int head_dim, float theta, void* s) { int head_dim, float theta, int period, void* s) {
dim3 grid(tokens, heads); dim3 grid(tokens, heads);
int blk = head_dim / 2; int blk = head_dim / 2;
rope_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(dy, dx, heads, head_dim, theta); rope_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(dy, dx, heads, head_dim, theta, period);
} }
// ===================================================================== // =====================================================================
@@ -330,7 +428,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; for (int c = threadIdx.x; c < cols; c += blockDim.x) pr[c] *= inv;
if (threadIdx.x == 0) { if (threadIdx.x == 0) {
int t = target[r]; int t = target[r];
loss[r] = -logf(pr[t]); loss[r] = t < 0 ? 0.0f : -logf(pr[t]);
} }
} }
void launch_cross_entropy_fwd_f32(const float* x, const int* target, void launch_cross_entropy_fwd_f32(const float* x, const int* target,
@@ -346,8 +444,13 @@ __global__ void cross_entropy_dx_k(const float* probs, const int* target,
int i = blockIdx.x * blockDim.x + threadIdx.x; int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= rows * cols) return; if (i >= rows * cols) return;
int r = i / cols, c = i % cols; int r = i / cols, c = i % cols;
float g = probs[i] - (c == target[r] ? 1.0f : 0.0f); int t = target[r];
dx[i] = g * scale; if (t < 0) {
dx[i] = 0.0f;
} else {
float g = probs[i] - (c == t ? 1.0f : 0.0f);
dx[i] = g * scale;
}
} }
void launch_cross_entropy_dx_f32(const float* probs, const int* target, void launch_cross_entropy_dx_f32(const float* probs, const int* target,
float* dx, int rows, int cols, float scale, void* s) { float* dx, int rows, int cols, float scale, void* s) {

86
csrc/ops/optim.cu Normal file
View File

@@ -0,0 +1,86 @@
// GPU-side optimizer kernels (Phase T7): AdamW parameter update and the
// global grad-norm reduction + rescale. These eliminate the per-step GPU↔host
// roundtrip of every parameter/gradient that the T6 host AdamW + host clip did.
//
// All F32, row-major, contiguous. The math mirrors xtrain-optim::AdamW::step_host
// (the reference); bias correction is passed in as bc1/bc2 = 1 - beta^t.
#include <math.h>
extern "C" {
// One AdamW step over a single parameter tensor of `n` elements, in place.
// m ← b1·m + (1-b1)·g
// v ← b2·v + (1-b2)·g²
// p ← p − lr·( (m/bc1) / (sqrt(v/bc2) + eps) + wd·p )
// `m`/`v` are this parameter's moment buffers (persisted on device across steps).
__global__ void adamw_step_f32(
float* p, const float* g, float* m, float* v,
float lr, float b1, float b2, float eps, float wd,
float bc1, float bc2, int n
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= n) return;
float gi = g[idx];
float mi = b1 * m[idx] + (1.0f - b1) * gi;
float vi = b2 * v[idx] + (1.0f - b2) * gi * gi;
m[idx] = mi;
v[idx] = vi;
float mhat = mi / bc1;
float vhat = vi / bc2;
p[idx] -= lr * (mhat / (sqrtf(vhat) + eps) + wd * p[idx]);
}
void launch_adamw_step_f32(
float* p, const float* g, float* m, float* v,
float lr, float b1, float b2, float eps, float wd,
float bc1, float bc2, int n, void* stream
) {
int block = 256;
int grid = (n + block - 1) / block;
adamw_step_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
p, g, m, v, lr, b1, b2, eps, wd, bc1, bc2, n);
}
// Accumulate sum-of-squares of one gradient tensor into *acc (a single f32 on
// device, pre-zeroed by the caller). Block-reduces then one atomicAdd per block.
__global__ void sumsq_accum_f32(const float* g, float* acc, int n) {
__shared__ float shared[32];
int tid = blockIdx.x * blockDim.x + threadIdx.x;
float v = (tid < n) ? g[tid] * g[tid] : 0.0f;
// block reduce
int lane = threadIdx.x & 31;
int warp = threadIdx.x >> 5;
int nwarps = (blockDim.x + 31) >> 5;
#pragma unroll
for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off);
if (lane == 0) shared[warp] = v;
__syncthreads();
v = (threadIdx.x < nwarps) ? shared[threadIdx.x] : 0.0f;
if (warp == 0) {
#pragma unroll
for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off);
if (lane == 0) atomicAdd(acc, v);
}
}
void launch_sumsq_accum_f32(const float* g, float* acc, int n, void* stream) {
int block = 256;
int grid = (n + block - 1) / block;
sumsq_accum_f32<<<grid, block, 0, (cudaStream_t)stream>>>(g, acc, n);
}
// Scale one tensor in place by a scalar (used to apply pre_scale·clip_factor to
// each gradient). Same as scale_f32 but in place.
__global__ void scale_inplace_f32(float* x, float factor, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) x[idx] *= factor;
}
void launch_scale_inplace_f32(float* x, float factor, int n, void* stream) {
int block = 256;
int grid = (n + block - 1) / block;
scale_inplace_f32<<<grid, block, 0, (cudaStream_t)stream>>>(x, factor, n);
}
}

84
csrc/ops/repeat_kv.cu Normal file
View File

@@ -0,0 +1,84 @@
// repeat_kv: the grouped-query-attention (GQA) head broadcast (Phase T15).
//
// GQA projects K/V to fewer heads than Q (num_kv_heads < num_heads); each KV head
// is SHARED by a group of `group = num_heads / num_kv_heads` query heads. Before
// the SDPA (composed or fused-flash, both untouched), we expand the KV tensor from
// [B·num_kv, S, hd] to the full [B·nh, S, hd] so the existing per-(batch,head)
// attention sees a full set of heads. GQA is then "free" for both SDPA paths.
//
// Layout: K/V are [bh_kv, S, hd] = [B·num_kv, S, hd] row-major, contiguous; the
// output is [bh_q, S, hd] = [B·nh, S, hd]. The head ordering matches xserv's
// repeat_kv (crates/xserv-model/src/qwen3.rs): query head qh reads kv head
// qh/group, query heads CONTIGUOUS within a group (dst = kvh*group + r). So:
//
// out[b·nh + qh, :, :] = in[b·num_kv + qh/group, :, :]
//
// Forward is a gather (each output row copies one input row). Backward is its
// transpose: a kv head receives the SUM of the `group` query heads that share it
// din[b·num_kv + kvh, e] = Σ_{r∈[0,group)} dout[b·nh + kvh·group + r, e]
// — the multi-group-to-one grad accumulation GQA's correctness hinges on. Each
// input element is owned by exactly one thread that serially sums its `group`
// source rows: race-free, NO atomics, run-to-run DETERMINISTIC. group==1 makes
// both directions a plain copy (identity → bit-identical to the MHA path).
//
// All F32, contiguous. (bf16 callers upcast → f32 on the Rust side and downcast
// the f32 result, mirroring the rest of the attention stack's fp32 policy.)
#include <math.h>
extern "C" {
// Forward gather. grid-stride over the bh_q·S·hd output elements; each output
// element copies from its kv-head source row. b = (out_bh / nh), qh = out_bh % nh,
// kv source bh = b·num_kv + qh/group.
__global__ void repeat_kv_fwd_k(const float* in, float* out, int nh, int num_kv,
int group, int seq, int hd) {
long row_elems = (long)seq * hd; // S·hd per head block
// One block per (batch, query-head); threads cover the S·hd block.
int out_bh = blockIdx.x; // over B·nh
int b = out_bh / nh;
int qh = out_bh % nh;
int kvh = qh / group;
const float* src = in + ((long)b * num_kv + kvh) * row_elems;
float* dst = out + (long)out_bh * row_elems;
for (long e = threadIdx.x; e < row_elems; e += blockDim.x) dst[e] = src[e];
}
void launch_repeat_kv_fwd_f32(const float* in, float* out, int batch, int nh,
int num_kv, int seq, int hd, void* s) {
int group = nh / num_kv;
int blk = (seq * hd) < 256 ? (seq * hd) : 256;
if (blk < 32) blk = 32;
repeat_kv_fwd_k<<<batch * nh, blk, 0, (cudaStream_t)s>>>(in, out, nh, num_kv,
group, seq, hd);
}
// Backward sum. One block per (batch, kv-head); threads cover the S·hd block.
// Each owned input element sums the `group` contiguous query-head source rows.
__global__ void repeat_kv_bwd_k(const float* dout, float* din, int nh, int num_kv,
int group, int seq, int hd) {
long row_elems = (long)seq * hd;
int in_bh = blockIdx.x; // over B·num_kv
int b = in_bh / num_kv;
int kvh = in_bh % num_kv;
int qh0 = kvh * group; // first query head sharing this kv head
float* dst = din + (long)in_bh * row_elems;
const float* base = dout + ((long)b * nh + qh0) * row_elems;
for (long e = threadIdx.x; e < row_elems; e += blockDim.x) {
float acc = 0.0f;
for (int r = 0; r < group; ++r) acc += base[(long)r * row_elems + e];
dst[e] = acc;
}
}
void launch_repeat_kv_bwd_f32(const float* dout, float* din, int batch, int nh,
int num_kv, int seq, int hd, void* s) {
int group = nh / num_kv;
int blk = (seq * hd) < 256 ? (seq * hd) : 256;
if (blk < 32) blk = 32;
repeat_kv_bwd_k<<<batch * num_kv, blk, 0, (cudaStream_t)s>>>(dout, din, nh,
num_kv, group,
seq, hd);
}
} // extern "C"

22606
data/tinystories-valid-3mb.txt Normal file
View File

File diff suppressed because it is too large Load Diff

136
docs/03-autograd-engine.md Normal file
View File

@@ -0,0 +1,136 @@
# Phase: Autograd Engine + Op Backward — Design Document
## Goal
在 T3 的 `Tensor`（matmul/transpose/finite-diff harness）之上，交付 **tape-based 动态 autograd 引擎** + 一个 tiny 现代 transformer 所需算子的**前向 kernel + 解析 backward**，每个 backward 都用 T3 的有限差分 harness 对拍通过。
具体三件事：
1. **autograd 引擎**：define-by-run 的反向自动微分。`Var` 包一个 `Tensor` + 可选 grad；每个 op 在 tape 上记一个节点（父节点 + backward 闭包）；`backward()` 按逆拓扑序遍历，把梯度推给父节点。**关键正确性点：梯度累加**——一个张量被多个 op 消费（扇出）时，各路梯度必须**求和**（T3 没有累加路径，在这里实现）。
2. **算子节点**：`matmul` / `add` / `mul` / `add_bias`(broadcast) / `scale` / `rms_norm` / `silu` / `swiglu` / `rope` / `softmax` / `cross_entropy`，各带前向 CUDA kernel（需要时）+ 解析 backward。
3. **Attention 用组合**：`attn = matmul(softmax(matmul(Q,Kᵀ)·scale), V)`。一旦 matmul/softmax/scale 是 autograd 节点，attention 的 backward 自动成立——**不写 fused attention backward kernel**，只加一个端到端 grad-check 测试。
**明确不做**（留给 T5/T6）：组装 transformer / 训练 loop / 优化器 / embedding / KV-cache / GQA 重复。本 Phase 只到「算子 backward 逐个对拍通过」。
## Module Layout
```
csrc/ops/nn.cu # 所有 T4 算子的 fwd+bwd kernel + launch_*（含 inlined warp/block reduce）
crates/xtrain-cuda/
├── build.rs # 新增 nn.cu
└── src/ffi.rs # 新增 launch_* 声明（no_cuda 门控）
crates/xtrain-tensor/
├── src/dtype.rs # 新增 I32（cross-entropy target 用）
└── src/tensor.rs # add/mul/add_bias/sum_rows/rms_norm(+bwd)/silu(+bwd)/
# rope(+bwd)/softmax(+bwd)/cross_entropy(+bwd)（no_cuda 门控）
crates/xtrain-autodiff/ # 引擎落在这里（已含 grad_check harness，自然归宿）
├── build.rs # 新增：检测 nvcc → no_cuda cfg（cfg 不跨 crate 传播）
├── src/
│ ├── lib.rs # 导出 tape::Var + ops（no_cuda 门控）
│ ├── finite_diff.rs # T3 既有 harness（不动）
│ ├── tape.rs # Var / VarNode / backward / 梯度累加
│ └── ops.rs # 各算子的 Var 节点构造器
└── tests/autograd.rs # 每算子 grad-check + 扇出累加 + 组合 attention（#![cfg(not(no_cuda))]）
```
为什么引擎放 `xtrain-autodiff` 而不是新 crate：该 crate 本就是「自动微分」语义的归宿，且已持有 `grad_check`。前向 kernel/`Tensor` 方法仍按 T2/T3 约定落在 `xtrain-tensor`（与 `scale`/`matmul` 一致），引擎只是在其上叠 tape。
## Key Design Decisions
### Tape 设计：`Rc<RefCell<VarNode>>` + 逆拓扑遍历
```rust
pub struct VarNode {
value: Tensor, // 前向输出
grad: Option<Tensor>, // 反向累加的梯度
parents: Vec<Var>, // 计算来源
backward: Option<BackwardFn>, // None=叶子
}
pub struct Var(Rc<RefCell<VarNode>>);
type BackwardFn = Box<dyn Fn(&Tensor, &[Var])>;
```
- `Var` clone 只是 bump `Rc`，**clone 共享同一节点**——这正是「扇出」的识别方式（同一 `Rc::as_ptr` 在多处出现）。
- `backward()`：① post-order DFS 建拓扑序（按指针去重）；② 把 loss（必须是标量）的 grad 种子设为 1；③ 逆序遍历，每个节点把自己的 grad 传给父节点的 backward 闭包。
- 闭包签名 `Fn(&grad, &parents)`：给本节点已累加的 grad 和父节点列表，闭包算出各父的梯度贡献并 `push_grad` 回去。前向需要 cache 的中间量（softmax 的 `y`、rms 的 `inv_rms`、ce 的 `probs`）用 `move` 闭包捕获。
### 梯度累加（扇出求和）——本 Phase 的正确性核心
`push_grad(parent, g)` 一律走 `accumulate`：
```rust
fn accumulate(&self, g: Tensor) {
match self.grad.take() {
None => self.grad = Some(g), // 首次
Some(prev) => self.grad = Some(prev.add(&g)),// 扇出：SUM
}
}
```
任何节点（叶子或中间）都累加：中间节点需要完整 grad 才能继续链式；叶子的累加结果就是输出。一个张量喂多个消费者时，多路 `push_grad` 自动求和。`mul(&x, &x)` 这类「同一 `Var` 进同一节点两次」也正确：`parents=[x,x]`（同指针），两次 `push_grad` 累加，拓扑去重保证 x 只遍历一次但收齐两路。测试 `fanout_grad_accumulation` 专门验证：`y=x∘x + x∘x`，`dL/dx` 须 = 4x（四处 x 全部求和）。
### 各算子 backward 数学
记上游梯度为 `d`（=本节点输出的梯度）。
| op | forward | backward |
|----|---------|----------|
| `matmul` | `C=A@B` | `dA=d@Bᵀ`, `dB=Aᵀ@d`（复用 T3 `matmul_backward`）|
| `add` | `a+b` | `da=d`, `db=d` |
| `mul` | `a∘b` | `da=d∘b`, `db=d∘a` |
| `add_bias` | `x[r,c]+bias[c]` | `dx=d`, `dbias[c]=Σ_r d[r,c]`（沿广播维求和）|
| `scale` | `x·α` | `dx=d·α` |
| `silu` | `x·σ(x)` | `dx=d·(σ + x·σ·(1−σ))`, `σ=σ(x)` |
| `swiglu` | `silu(g)∘u` | 由 `silu`+`mul` 组合自动得 |
| `rope` | rotate_half 旋转 | RoPE 是正交变换，`dx` = 用**逆（转置）旋转**作用于 `d`（角度 +θ 的转置 ≡ −θ）|
| `softmax` | row-wise safe softmax → `y` | Jacobian：`dx[r,c]=y[r,c]·(d[r,c] − Σ_c' d·y)` |
| `cross_entropy` | mean NLL(softmax(x), tgt) | `dx = (probs − onehot)/rows`，再乘上游标量 grad |
**RMSNorm**（`y[r,c]=x[r,c]·ir·γ[c]`, `ir=rsqrt(mean(x²)+eps)`）：
设 `g[c]=d[r,c]·γ[c]`，`n=cols`，
```
dx[r,c] = ir·g[c] − x[r,c]·ir³/n·Σ_c'(g[c']·x[r,c'])
dγ[c] = Σ_r d[r,c]·x[r,c]·ir
```
前向 cache 每行 `inv_rms[r]`，backward 直接复用，避免重算 reduce。
**RoPE 反向推导**：前向是 2×2 旋转矩阵 `R(θ)`，正交 ⇒ `Rᵀ = R(−θ)`。故
```
dx[i] = d[i]·cos + d[i+h]·sin
dx[i+h] = d[i+h]·cos − d[i]·sin
```
position=0 时旋转是恒等，backward 也恒等（sanity check）。
**Softmax 反向推导**：`∂y_i/∂x_j = y_i(δ_ij − y_j)`，链式后
`dx_i = Σ_j d_j·y_i(δ_ij − y_j) = y_i(d_i − Σ_j d_j y_j)`，即每行减去 `Σ(d∘y)` 后乘 `y`。
**Cross-entropy 反向推导**：`L=−log softmax(x)[t]`，softmax+NLL 的经典结果 `∂L/∂x_c = softmax_c − [c=t]`；取 batch 平均 ⇒ 除以 rows。kernel 把 `scale=upstream/rows` 折进去。
### Attention 用组合，不写 fused kernel
```
Kᵀ = transpose(K)
scores = scale(matmul(Q, Kᵀ), 1/√d) # [s,s]
probs = softmax(scores)
out = matmul(probs, V) # [s,d]
```
每一步都是已有 autograd 节点，`backward()` 自动沿 matmul→softmax→scale→matmul 链回传，得到 `dQ/dK/dV`，无需手写 attention backward。测试 `attention_composed_bwd` 单头单 batch 端到端 grad-check Q/K/V 三者。（transpose 在测试里用一个临时 `Var::from_op` 节点包，因为引擎暂未把 transpose 列为 op——T5 若需要再补。）
### kernel 实现要点
- `nn.cu` 自带 inlined `warp/block_reduce_sum/max`（不引外部头文件，与现有 csrc/ 单文件风格一致）；block-reduce 末尾广播到全 block，便于 softmax/rms 的「标量广播」模式。
- 每个 op 各自 `cudaDeviceSynchronize()`（T3 约定，无 stream）。
- 全 F32、row-major、contiguous；cross-entropy target 用新增的 `DType::I32`。
## 验证方法
模板沿用 T3 `gemm.rs::run_bwd`：标量 loss `L = sum(W∘out)`，`W` 固定随机 ⇒ 上游 `dOut = W`；跑 op 的 `backward()` 拿 `.grad()`，对每个输入用 `grad_check` 与中心差分对拍。
- **每算子**一个 grad-check（线性/双线性 op 用大 eps=1e-2、rel_tol=2e-2；非线性 op 用 eps=1e-3、rel_tol=3e-2、atol=1e-3 压住近零梯度）。
- **扇出累加**：`fanout_grad_accumulation`，验证 `dL/dx=4x`。
- **组合 attention**：`attention_composed_bwd`，端到端 grad-check `dQ/dK/dV`。
- 全部 `#![cfg(not(no_cuda))]` 门控；本地只 `cargo check`/`fmt`，构建+实跑在 dash5（8× RTX 5090, sm_120），capture 每 op 的 pass + max rel-err。

157
docs/04-tiny-transformer.md Normal file
View File

@@ -0,0 +1,157 @@
# Phase T5: Tiny Transformer (fwd+bwd) — Design Document
## Goal
在 T4 的 autograd 引擎（`Var` tape + 11 个算子）之上，**组装一个 tiny 现代架构 decoder**（RoPE + RMSNorm + SwiGLU + 多头因果 attention），跑通 **char-level bring-up**，并以「**把一个小 batch overfit 到 loss≈0**」证明整条 forward+backward 计算图正确。
明确范围（T5 只做这些）：
1. **补 T4 留的缺口算子**为正式可微节点：`embedding`（按 token id gather / scatter-add 反向）、`reshape`、`transpose`（2D + 3D 轴(0,1)），并加 `split_heads`/`merge_heads` 处理多头布局。
2. **tiny transformer**：token embedding → `n_layers` × {pre-RMSNorm → 多头因果 attention → 残差；pre-RMSNorm → SwiGLU MLP → 残差} → final RMSNorm → LM head → cross-entropy。
3. **overfit 验证**：极简手写 GD step（`p -= lr·grad`）记住一个固定小 batch。
**不做**（留 T6）：真训练 loop、AdamW、LR schedule、grad clip、checkpoint、TinyStories/tokenizer。overfit 不需要优化器，一行 `p -= lr·grad` 足够。
## Module Layout
```
csrc/ops/model.cu # 新：embedding gather/scatter-add + 3D 轴(0,1) transpose kernel
# （reshape 是纯 metadata，无 kernel）
crates/xtrain-cuda/
├── build.rs # 新增 model.cu
└── src/ffi.rs # 新增 launch_embedding_fwd/bwd + launch_transpose_3d01（no_cuda 门控）
crates/xtrain-tensor/
└── src/tensor.rs # 新增 reshape / embedding(+bwd) / transpose_3d01（no_cuda 门控）
crates/xtrain-autodiff/
├── src/tape.rs # 新增 Var::zero_grad / set_value（供手写 GD step）
├── src/ops.rs # 新增节点 embedding / reshape / transpose_3d01 / transpose_2d
│ # / split_heads(->Vec<Var>) / merge_heads
└── tests/structural.rs # 新增：上述结构算子各自 grad-check
crates/xtrain-model/ # 新 crate：模型本体
├── build.rs # 检测 nvcc → no_cuda cfg（逐 crate）
├── src/
│ ├── lib.rs # 导出 Config / TinyTransformer / 辅助
│ ├── config.rs # 超参（host-only，no_cuda 也编）
│ └── model.rs # TinyTransformer：参数容器 + forward 图（no_cuda 门控）
└── tests/
├── overfit.rs # 端到端：char-level + 手写 GD → loss≈0
├── parity_dump.rs # PyTorch 对拍 fixture（dump 权重/logits/grad）
└── parity.py # 等价 PyTorch 模型，对比 forward + 每参数 grad
```
为什么新开 `xtrain-model` crate（而非塞进 autodiff）：对齐 xserv 的「模型层独立于算子层」分层；autodiff 只管 `Var`/`ops`，model 在其上拼网络，职责清晰。
## Key Design Decisions
### 约定（对齐引擎，不照搬 HuggingFace）
- **线性层权重 `[in, out]`，按 `x @ W`**。引擎的 GEMM 是裸 `A @ B`，用 `[in,out]` 省掉每次投影的 transpose（Qwen 的 `[out,in]` + `x@Wᵀ` 是推理侧权重布局的产物，训练侧自定义即可，回流 xserv 时在 T9 转置导出）。
- **`dim = n_heads · head_dim`**（无独立 attention 投影维度），tiny 配置 `dim=32, n_layers=2, n_heads=2, head_dim=16, ffn_hidden=64`。
- **RoPE position = token 行号**（kernel 内建约定）；`rope` 只接受 `[tokens, heads, head_dim]` 布局。
- **因果 mask = 加性常量 `[seq,seq]`**：对角线及以下为 0，以上为 −1e9，softmax 前 `add` 进 scores。引擎没有 masking 算子，用一个常量 leaf + `add` 即可（该 leaf 无下游、其 grad 不被读取，无害）。
### 多头布局：reshape + transpose_3d01 + split/merge_heads
引擎的 matmul/softmax 都是 2D 的，多头 attention 因此**逐头做 2D SDPA**。布局流水线：
```
proj = x @ Wq # [seq, dim]
reshape → [seq, nh, hd]
rope (kernel 要求正是 [tokens, heads, head_dim])
transpose_3d01 → [nh, seq, hd] # 让每个 head 的 [seq,hd] 块连续
split_heads → nh × [seq, hd] # 每头一个 Var
# 逐头：
scores = (q_i @ k_iᵀ)·(1/√hd) + mask # [seq,seq]
probs = softmax(scores)
out_i = probs @ v_i # [seq, hd]
merge_heads(out) → [nh, seq, hd]
transpose_3d01 → [seq, nh, hd]
reshape → [seq, dim]
@ Wo # 输出投影
```
- `reshape` 是**纯 metadata**（连续张量改 shape/strides，共享 storage，无 kernel、无数据搬运）；`[seq, nh·hd] ↔ [seq, nh, hd]` 正好是它。反向 reshape 回去。
- `transpose_3d01`（`[a,b,c]→[b,a,c]`）有 kernel，**自反**：反向是对 grad 再做一次同样的 transpose。
- `split_heads`：`[nh,seq,hd]` 在该布局下每个 head 块连续，前向把每块拷成独立连续张量返回 `Vec<Var>`；反向把每头 grad 散回零初始化的 `[nh,seq,hd]`，由引擎的**扇出 SUM** 累加。`merge_heads` 是逆操作。
### embedding 前向 gather / 反向 scatter-add
```
out[s,:] = table[ids[s], :] # table:[vocab,dim], ids:[seq] I32
dtable[ids[s],:] += dout[s,:] # 反向：原子 scatter-add
```
- ids 是**常量**（不是 `Var`），只有 table 参与求导。
- 反向必须 **atomicAdd**：多个位置可能映射到同一 id，梯度要累加（grad-check 测试里特意放了重复 id `[0,3,1,3,2,0]`）。`dtable` 先 `zeros` 再原子累加。
### Block 结构（pre-norm + residual，Qwen3 风格）
```
h = embedding(ids)
for block:
h = h + attention(rms_norm(h, γ_attn)) # 残差
h = h + swiglu_mlp(rms_norm(h, γ_ffn)) # 残差
logits = rms_norm(h, γ_final) @ lm_head
loss = cross_entropy(logits, targets)
```
SwiGLU MLP：`down( silu(x@W_gate) ∘ (x@W_up) )`，复用 T4 的 `swiglu = mul(silu(g), u)`。
### 参数 API（为 T6 优化器准备）
- `TinyTransformer::params() -> Vec<Var>`：稳定顺序的全部可学习叶子（embedding / 各 block 的 9 个 / final_norm / lm_head）。
- `Var::set_value(t)`：原地更新参数值（GD/AdamW 用），保持叶子身份在多 step 间稳定。
- `Var::zero_grad()`：清梯度。**关键**：每个 forward 建新图但叶子复用，上一 step 的 grad 不清会被 SUM 累加 → 每 step 后必须 zero。
- `param_to_host(&Var)`：把参数搬回 host `Vec<f32>`（GD step / 对拍导出）。
手写 GD step（overfit 用）：
```rust
for p in params {
if let Some(g) = p.grad() { p.set_value(p.value().add(&g.scale(-lr))); }
p.zero_grad();
}
```
### overfit 方法学
一个固定的小 batch（char-level 文本 → 字符表 → `(input, shifted target)`），反复跑 forward→backward→GD。**只要 fwd+bwd 全对，模型会记住这个 batch，loss → ~0**；任何一个 backward 错了，loss 会停在某个台阶下不去。这是比单算子 grad-check 更强的端到端信号。验收同时检查 greedy argmax 是否完全复现 target 序列。
## 验证方法
全部 `#![cfg(not(no_cuda))]` 门控，本地只 `cargo check`/`fmt`，构建+实跑在 dash5（8× RTX 5090, sm_120）。
1. **结构算子 grad-check**（`tests/structural.rs`）：沿用 T4 的 `L = sum(W∘out)` 有限差分 harness，对 `embedding`（含重复 id）、`reshape`、`transpose_3d01`、`transpose_2d`、`split/merge_heads` 往返各做一遍。
2. **overfit**（`tests/overfit.rs`）：tiny 模型 + 手写 GD，断言 `loss → <0.05` 且 greedy argmax 全对。
3. **PyTorch 对拍**（`parity_dump.rs` + `parity.py`）：Rust dump 出权重/ids/logits/loss/每参数 grad（一次 backward），Python 用**完全等价**的 PyTorch 模型（同 `x@W` 约定、同 RoPE rotate_half pos=行号、同 RMSNorm/SwiGLU/因果 SDPA）跑 fwd+bwd，对比 forward logits + 21 个参数 grad 的相对误差（rtol=2e-2）。
### dash5 实测结果
```
# 结构算子 grad-check（max rel-err）
embedding dTable 3.5e-5 reshape dX 3.4e-4
transpose_2d dX 2.3e-5 transpose_3d01 dX 1.9e-4
split/merge_heads dX 3.9e-5 （5/5 通过）
# overfit（lr=0.3, 200 steps, vocab=16, seq=27, params=21664）
step 0: loss = 2.821415
step 20: loss = 0.341899
step 40: loss = 0.099285
...
step 199: loss = 0.004009
start 2.821415 → final 0.004009 ; greedy match 27/27 ✅
# PyTorch 对拍（rtol=2e-2）
loss: rust=2.505827e0 torch=2.505827e0 relerr=1.4e-8
logits: max relerr = 9.5e-5
params checked: 21 worst = grad[l0_wo] @ 1.1e-3 → PARITY OK ✅
# 无回归：T4 autograd 12/12 仍全绿
```
forward 与 PyTorch 对到 ~1e-4、每参数 grad 对到 ≤1.1e-3，overfit loss 从 2.82 跌到 0.004 且完全复现 target——三层证据（单算子 finite-diff、端到端 overfit、PyTorch 对拍）一致确认整条 fwd+bwd 正确。

114
docs/05-training-loop.md Normal file
View File

@@ -0,0 +1,114 @@
# Phase T6: Training Loop + AdamW + Real Training — Design Document
## Goal
在 T5 的 `TinyTransformer`（`params()` / `forward` / `loss` + `Var::{value,grad,set_value,zero_grad}`）之上，搭起**真正的训练栈**，并在**真实文本语料**上把 loss 训下来：
1. **手写 AdamW**：per-param 一/二阶矩（m、v），bias correction，**decoupled weight decay**，对拍 `torch.optim.AdamW` 数值一致。
2. **训练 loop**：语料 → 采样定长序列 → `forward(loss)` → `backward` → **global-norm grad clip** → **AdamW step** → `zero_grad`；**LR schedule**（warmup + cosine）；周期性 loss 日志；**checkpoint 存/取**。
3. **采样器**：greedy / temperature，训练中/后吐文本看「在不在学」。
4. **数据**：**复用 xserv 的 GPT-2 BPE** tokenizer（path-dep），语料 = **TinyStories** 子集。
**不做**（留后续 Phase）：性能（cuBLAS 切换 / bf16 / 激活重计算 = T7）、分布式（NCCL 数据并行 = T8）。本 Phase 只要**正确性 + 清晰的学习信号**，训练预算有界（几分钟 / 几千步，非完全收敛）。
## Module Layout
```
crates/xtrain-optim/ # 新 crate：优化器
├── build.rs # 检测 nvcc → no_cuda cfg（逐 crate）
├── src/lib.rs # AdamW：step_host(纯 host 数学) + step(&[Var]) GPU 包装
└── tests/adamw_host.rs # host 单测：对独立参考递推 + 纯 decay 边界（本地可跑，无 GPU）
crates/xtrain-train/ # 新 crate：训练基建 + 入口
├── build.rs # 检测 nvcc → no_cuda cfg
├── Cargo.toml # path-dep: ../../../xserv/crates/xserv-tokenizer（本地/dash5 都解析）
├── src/
│ ├── lib.rs # 模块导出（host-only 与 GPU 件分门控）
│ ├── schedule.rs # LrSchedule：warmup + cosine（host-only，可本地单测）
│ ├── clip.rs # global L2 norm + clip_scale（host 数学）+ clip_grad_norm(&[Var])（GPU 门控）
│ ├── data.rs # Corpus：load tokenizer+语料 → token 流 → sample(input,target) 窗口
│ ├── checkpoint.rs # save / load_into：按 params() 顺序 dump/reload（GPU 门控）
│ ├── sample.rs # generate：greedy / temperature 自回归采样（GPU 门控）
│ ├── train_loop.rs # TrainConfig + train()：把以上接到 model+AdamW（GPU 门控）
│ └── bin/train.rs # 真训练入口：load 数据 → train → checkpoint → 采样
└── tests/
├── adamw_parity_dump.rs # AdamW 对拍 fixture：固定 init 跑 N 步 AdamW，dump loss 轨迹 + 终参
├── adamw_parity.py # 等价 PyTorch 模型 + torch.optim.AdamW，对比轨迹 + 终参
├── checkpoint_roundtrip.rs # 训几步→save→载入新模型→logits/loss 逐位一致
└── real_training.rs # TinyStories 有界训练：loss 大幅下降 + 采样在学
data/tinystories-valid-3mb.txt # 语料子集（committed，~3MB，TinyStories-valid 前 3MB，整故事截断）
```
**为什么拆两个 crate**：对齐 xserv 的分层（优化器与训练编排分开）。`xtrain-optim` 只管参数更新数学；`xtrain-train` 管数据/调度/checkpoint/采样/loop。AdamW 数学独立可测，不依赖 model。
**host / GPU 门控约定**（沿用全仓）：纯算术（`LrSchedule`、grad-norm 数学、AdamW 的 `step_host`）**始终编译**，本地 `cargo check` + 单测即可验证；凡 round-trip GPU 张量的（`step(&[Var])`、`clip_grad_norm(&[Var])`、checkpoint、采样、loop）一律 `#[cfg(not(no_cuda))]`，链接+实跑在 dash5。每 crate 的 `build.rs` 各自检测 nvcc（cfg 不跨 crate 传播）。
## Key Design Decisions
### AdamW：手写数学 + decoupled weight decay
第 `t` 步（1-indexed），参数 `θ`、梯度 `g`：
```text
m ← β1·m + (1−β1)·g
v ← β2·v + (1−β2)·g²
m̂ ← m / (1 − β1ᵗ) (bias correction)
v̂ ← v / (1 − β2ᵗ)
θ ← θ − lr·( m̂ / (√v̂ + ε) + wd·θ )
```
- **decoupled weight decay**（Loshchilov & Hutter 2019）：`wd·θ` 直接作用在参数上，**不**并进梯度（不进入自适应 `√v̂` 分母）——这正是 `torch.optim.AdamW` 的定义，区别于「L2 正则把 `wd·θ` 加到 `g`」的 Adam。
- 默认超参对齐 PyTorch：β1=0.9，β2=0.999，ε=1e-8。
- **状态 keyed by 参数在 `params()` 中的下标**（稳定序），首次 `step` 惰性按各参数 numel 分配 `m,v`。`t` 全局共享（所有参数同一 bias correction，和 PyTorch 一致）。
**实现分层**：`step_host(lr, &mut [Vec<f32>], &[Vec<f32>])` 是纯 host f32 数学（无 GPU、无 autograd，本地单测）；`step(lr, &[Var])` 把每参数的 `value()`/`grad()` 拉到 host、调 `step_host`、`set_value` 写回。这条路子（host 算优化器）对 tiny 模型完全够用，且让 AdamW 数学**脱离 GPU 可严格对拍**——T6 是正确性 Phase，不做 GPU 优化器 kernel（那是性能向，超范围）。`lr` 每步传入，给 schedule 留口。
### LR schedule：warmup + cosine
`step ∈ [0,warmup)` 线性 `0→max_lr`；`[warmup,total)` cosine `max_lr→min_lr`；`≥total` 钳到 `min_lr`。纯函数（只吃 step 下标），本地单测形状。
### grad clip：global L2 norm（+ batch 平均）
跨**所有**参数梯度联合算 L2 norm（同 `torch.nn.utils.clip_grad_norm_`）：`total > max_norm` 则全体 `×(max_norm/total)`。
模型是**单序列**（无 batch 维），一个 `batch_size` 的「batch」靠**跑 `batch_size` 次 forward+backward**、让 tape 的 fan-out 规则**把梯度 SUM** 起来实现。为得到 batch 均值梯度，clip 这一趟 host pass 里**先 `×1/batch_size`** 再算 norm/裁剪——`clip_grad_norm(params, max_norm, pre_scale)` 把「平均」与「裁剪」融成一次 host 往返（省一趟拷贝）。batch 是 T7/边角关切，这里只求正确。
### checkpoint 格式
按 `params()` 顺序 dump 每个参数的 value 到扁平二进制：
```text
magic u32 = "XTRT" | version u32 | n_params u32
×n_params: ndim u32 | dims[ndim] u32 | data[Πdims] f32 (小端)
```
不存架构/config——调用方用同一 `Config` 重建模型再 `load_into`（round-trip 与 resume 都自知 config）。`load_into` 校验 magic/version/数量/逐参数 shape，按各参数 device 写回 `set_value`。f32 精确往返 → 重载后 forward 逐位一致（同 kernel 同输入）。
### 数据管线 + tokenizer 复用
- **tokenizer = 复用 xserv 的 from-scratch GPT-2/Qwen BPE**：`Cargo.toml` path-dep `../../../xserv/crates/xserv-tokenizer`，该相对路径在本地 `~/projects` 与 dash5 `/opt/wjh/projects` 都解析；Cargo 按目标 crate 自身的 workspace（xserv 的）解析它的 `serde/regex` 依赖，不需要 xtrain 复制 workspace dep。加载 `/opt/wjh/models/gpt2/tokenizer.json`。
- **语料 = TinyStories 子集**：dash5 经 `hf-mirror.com` 取 `TinyStories-valid.txt` 前 ~3MB（HF 直连不可达，proxy 脚本只起后台 SOCKS；hf-mirror 直连 200），committed 进 `data/`。`Corpus::load` 整篇 tokenize 成一条 token 流（TinyStories 用 `<|endoftext|>` 分故事，GPT-2 BPE 正好出成单个 special token，文档边界保留）；range 下载会掐头去尾，故先丢首个不完整行、截到最后一个 `<|endoftext|>`，只训整故事。`sample(seq)` 随机取窗口 `[s,s+seq+1)` → input `[s,s+seq)` / target 右移一位（next-token），LCG 种子可复现，不引 RNG crate。
### 采样器
模型单序列、RoPE pos=行号，故自回归生成**每步对增长前缀重跑 forward、取末行 logits**（最简正确法；KV cache 是推理/性能向，超范围）。`temperature==0` greedy argmax，否则按 `softmax(logits/T)` 采样。
### 训练 loop（`train`）
每步：采 `batch_size` 序列各自 forward `loss` + backward（tape SUM 梯度）→ `clip_grad_norm(×1/batch + 裁剪)` → `AdamW::step(lr)` → 全参数 `zero_grad`；按 `log_every` 打 `loss/lr/gnorm/tok-s`，按 `ckpt_every` 存 checkpoint，返回逐步 loss 轨迹。
## 验证方法（验收）
GPU 测试全部 `#[cfg(not(no_cuda))]` 门控，在 dash5 实跑 capture：
1. **AdamW 对拍 PyTorch**（严格正确性）：同一 tiny 模型 + 相同 init，Rust AdamW 与 `torch.optim.AdamW`（lr/wd/betas/eps 全对齐）各跑 N 步固定 batch → **loss 轨迹**与**终参**逐项 rtol 内一致。
- fixture：`cargo test -p xtrain-train --test adamw_parity_dump -- --ignored --nocapture`
- 对比：`python3 crates/xtrain-train/tests/adamw_parity.py /tmp/xtrain_adamw`
2. **checkpoint round-trip**：训几步 → save → 载入**全新 init 的模型** → 固定输入 logits/loss 逐位一致（且证明载入前新模型确实不同）。
- `cargo test -p xtrain-train --test checkpoint_roundtrip`
3. **真训练**（端到端学习信号）：TinyStories 上有界训练（几百~几千步）→ loss 大幅下降 + greedy 采样显出英文结构（非乱码）。
- `cargo test -p xtrain-train --release --test real_training -- --ignored --nocapture`
- 或 `cargo run -p xtrain-train --release --bin train -- <tokenizer.json> <corpus.txt> [steps] [ckpt]`
4. **host 单测**（本地即跑）：AdamW 数学对独立参考递推、LR schedule 形状、grad-norm/clip 数学。
- `cargo test -p xtrain-optim -p xtrain-train`

139
docs/06-performance.md Normal file
View File

@@ -0,0 +1,139 @@
# Phase T7: Performance — Design Document
## Goal
T6 把真训练打通了（TinyStories，loss 10.83→3.43，采样连贯），但吞吐只有 **~2800 tok/s**。
T7 的目标是**在不牺牲数值正确性的前提下把训练显著加速**——上来先把 fp32 路径里那些纯开销榨干，
再视情况上 bf16 / 激活重计算。
按 xtrain.md T7 note 的优先级，**保 fp32 数值不回归**的三步是 must-have：
1. **matmul fwd/bwd 走 cuBLAS** —— 前向 + 两路反向（`dA=dC·Bᵀ`、`dB=Aᵀ·dC`）全切 cuBLAS `Sgemm`。fp32，等价于换了求和顺序的同一个 GEMM，所以正确性自动保住。
2. **GPU 侧优化器 + grad clip** —— 干掉每步把全部参数/梯度 GPU↔host 往返的开销：AdamW 的 m/v 状态搬到 device、update 走 kernel；global grad-norm 用 device reduction，只把那一个标量取回 host。
3. **stream / 减 sync** —— 不再每个 op 之后都 `cudaDeviceSynchronize`：default stream 上 kernel 本就顺序执行，host 读数据又都走 stream-ordered 的 `cudaMemcpy`，per-op sync 是纯开销，全删。
**不做（本 Phase 范围外）**：分布式数据并行 / NCCL all-reduce（T8）、导出回流 xserv（T9）。
**降级出口**：bf16 混合精度（④）/ 激活重计算（⑤）改数值、牵动每个 kernel，且本 model 太小（dim=32）属 latency-bound、bf16 tensor-core 收益有限——按 xtrain.md 的 escape hatch，①②③ 交付并实测加速后，④⑤ 记为 follow-up，不把正确性留在半截状态。详见末节。
## Module Layout
```
csrc/ops/optim.cu # 新：GPU AdamW step + global grad sumsq reduce + in-place scale
crates/xtrain-cuda/
├── src/cublas.rs # 新：持久化 cuBLAS handle + row-major sgemm(含转置位)
├── src/ffi.rs # +CUBLAS_OP_T、+optim.cu 的三个 launch_*、+cudaMemset
├── src/memory.rs # GpuBuffer::memset（device 置零，免 H2D 零拷贝）
├── src/lib.rs # pub mod cublas（not(no_cuda)）
└── build.rs # +optim.cu
crates/xtrain-tensor/
├── src/tensor.rs # matmul/matmul_backward 改走 cublas::sgemm；删 21 处 per-op sync
└── src/storage.rs # device zeros 改用 memset
crates/xtrain-optim/
├── src/lib.rs # +GpuAdamW（m/v on device，in-place update）；host AdamW 留作参考
├── Cargo.toml # xtrain-cuda 升为常规依赖（GpuAdamW 要发 kernel）
└── tests/adamw_gpu.rs # 新：GPU AdamW 对 host 参考逐位一致
crates/xtrain-train/
├── src/clip.rs # +clip_grad_norm_gpu（device reduce + in-place rescale）；host 版留作参考
└── src/train_loop.rs # 改用 GpuAdamW + clip_grad_norm_gpu
```
## Key Design Decisions
### ① cuBLAS matmul（row-major ⟺ col-major）
cuBLAS 是 **列主序**，我们的张量是 **行主序**。一个行主序 `[r,c]`、leading dim = `c` 的矩阵交给 cuBLAS，
被读作它的转置（列主序 `[c,r]`）。要拿到行主序结果 `C[m,n] = opA(A)·opB(B)`，就让 cuBLAS 算它的列主序转置
`Cᵀ[n,m] = opB(B)ᵀ·opA(A)ᵀ`——`Cᵀ` 列主序的字节布局正好就是 `C` 行主序。
`cublas::sgemm(trans_a, trans_b, m, n, k, …)` 落地为：第一参 = `B`（op = `trans_b ? N : T`），第二参 = `A`（op = `trans_a ? N : T`），尺寸 `(m=n, n=m, k=k)`，`lda/ldb/ldc` = 各自**存储态行主序的列数**：
```rust
let lda = if trans_a { m } else { k }; // A 存 [m,k] 或 [k,m]
let ldb = if trans_b { k } else { n }; // B 存 [k,n] 或 [n,k]
let ldc = n; // Cᵀ 是 [n,m] 列主序 ld=n（== 行主序 C[m,n]）
```
`trans_a=trans_b=false` 这一支与 T3 测试里的 cuBLAS oracle **逐参数一致**（同样 OP_N、交换顺序、m=N/n=M/k=K），所以前向天然对得上。
**反向用 cuBLAS 的转置位省两个 transpose kernel**：T3 版 `matmul_backward` 是 `dc.matmul(b.transpose_2d())` + `a.transpose_2d().matmul(dc)`，要起两个 transpose kernel + 两次分配。T7 直接：
```text
dA[M,K] = dC[M,N] · Bᵀ → sgemm(trans_a=false, trans_b=true, m=M,n=K,k=N, a=dC, b=B)
dB[K,N] = Aᵀ · dC[M,N] → sgemm(trans_a=true, trans_b=false, m=K,n=N,k=M, a=A, b=dC)
```
**为什么不回归**：全程 fp32，cuBLAS 与手写 tiled kernel 算的是同一个 GEMM，只差求和顺序的 rounding。
所以 T3「fwd 对 cuBLAS / bwd 对 finite-diff」的容差不变，下游 autograd grad-check、PyTorch 对拍也不变。
**handle 持久化**：cuBLAS handle 创建很贵（T3 oracle 每次调用都 create/destroy）。改为 **每线程缓存一个 handle**，进程生命周期内复用（`thread_local! + RefCell<Option<CublasHandle>>`）。
### ② GPU AdamW + GPU grad-norm（去掉每步全参往返）
T6 的瓶颈之一：`AdamW::step` 把每个参数的 value + grad 全 D2H 拉回 host、host 上跑 AdamW、再 H2D 写回；`clip_grad_norm` 同理把全部 grad 拉回 host 算范数。3.26M 参数 × 每步两趟 = 大量 PCIe 往返 + 同步。
**GpuAdamW**：m/v 矩状态以「每参一对 device `Tensor`」常驻显存，update 是一个 in-place kernel——读参数的 `.grad()`、原地改写参数 buffer（参数 leaf 的 storage 是 `Arc` 共享，原地写对所有 clone 可见，leaf 身份跨步稳定，无需 `set_value`）：
```text
m ← β1·m + (1−β1)·g ; v ← β2·v + (1−β2)·g²
p ← p − lr·( (m/bc1) / (√(v/bc2) + ε) + wd·p ) bc1/bc2 = 1−βᵗ（host 传入）
```
数学与 host `AdamW::step_host` 逐字对应；host AdamW **原样保留**作 PyTorch 对拍的参考，新增 `adamw_gpu` 测试拿同一组 params/grads 把 GPU 结果对 host 参考**逐位比**（实测 max abs err = 0）。
**clip_grad_norm_gpu**：`sumsq_accum` kernel 对每个 grad 做 block-reduce 后 `atomicAdd` 到一个 device 标量；只把这**一个标量**取回 host 求 `sqrt`、算 clip factor，再用 `scale_inplace` kernel 原地把每个 grad 乘 `pre_scale·factor`。整步只回传 1 个 float，不再拉全部 grad。
### ③ stream / 减 sync
每个 tensor op 之前 `Tensor::zeros` 分配输出、之后 `cudaDeviceSynchronize`——两处都是隐藏开销：
- **per-op sync 全删（21 处）**：default stream 上 kernel 顺序执行；任何 host 读数据都走 `to_device(Cpu)` → 阻塞且 stream-ordered 的 `cudaMemcpy`，自然等齐前面的 kernel。所以 op 后那次显式 sync 对正确性纯属多余（只是把异步 kernel 错误提前暴露，可接受地推迟到下一次 sync/memcpy）。
- **device zeros 改 `cudaMemset`**：原来每个 op 输出都用「host 零 buffer + 阻塞 H2D memcpy」置零，那次 H2D 本身就是个 per-op 同步点 + 一次拷贝；换成 device 端 `cudaMemset`（default stream 上异步，不串行化 stream）。
once-per-step 的 sync（clip 取范数前、AdamW step 末尾）保留——量级是每步一次，非每 op。
> CUDA-graph capture 是 optional bonus，本 Phase 未做。
## 验证方法
**两道闸都要过**：
**A. 数值不回归（fp32 容差不变，全绿）**——dash5 实跑：
| 测试 | 闸 | 结果 |
|---|---|---|
| T3 GEMM（fwd vs cuBLAS / bwd vs finite-diff） | rel-err 容差不变 | 5/5 ok |
| T4 autograd grad-check（每 op finite-diff） | ≤2e-2 不变 | 12/12 ok |
| T5 结构 grad-check + overfit + PyTorch 对拍 | overfit 27/27、logits relerr、21 参梯度 rtol 不变 | overfit 2.821→0.004 (27/27)；parity logits relerr 1.5e-4、21 grads OK |
| T6 AdamW vs torch + checkpoint round-trip | 轨迹/终参 rtol 不变、逐位一致 | AdamW relerr 4.6e-6；ckpt logit diff 0.0 |
| **T7 GPU AdamW vs host 参考** | 逐位一致 | max abs err **0.0** |
**B. 吞吐提升**——同 model/config（dim 32、4 层、vocab 50257、seq 64、batch 8、~3.26M 参），60 步计时取稳态：
| 步骤 | tok/s | 备注 |
|---|---|---|
| baseline (T6) | ~2770 | 起点 |
| ① cuBLAS matmul | ~3310 | matmul 非主瓶颈（model 小、latency-bound） |
| ② GPU AdamW + grad-norm | ~4070 | 去掉每步全参 GPU↔host 往返 |
| ③ drop per-op sync + memset | **~8220** | 删 21 处 per-op sync 是大头 |
端到端（real_training 800 步，新快路）：**~8500 tok/s 稳态**，loss 10.81→3.90(avg10)，采样
`Once upon a time, there was a little girl named Lily. She was very happy to play with her mom.`——
收敛与 T6 fp32 同轨。
**净加速 ~3×，零数值回归。**
## ④⑤ Follow-up（本 Phase 未做，记给后续）
- **④ bf16 混合精度（fp32 master）**：matmul/激活走 bf16、optimizer 持 fp32 master 拷贝。本 model dim=32 太小、属 launch/latency-bound，bf16 tensor-core 算力收益有限，唯一够大的 `lm_head [64,32]@[32,50257]` 主要吃带宽；且 bf16 改数值、要单独加宽容差 + 重验收敛，风险/收益此规模下不划算。等模型放大或上 T8 多卡再做更值。
- **⑤ 激活重计算**：反向重算 block 激活省显存。当前单序列、显存不紧，优先级低。
两者按 escape hatch 推迟，①②③ 的 fp32 加速已完整交付且全测绿。
## T8 衔接（数据并行 all-reduce）
T7 之后**梯度常驻 device**（`.grad()` 是 device tensor），优化器 update 也全在 device——这正好对接 T8 的 NCCL 数据并行：
- 各 rank 本地 `backward` 后，梯度已在显存里，**直接对 `params` 的 `.grad()` 张量 all-reduce**（无需先拉回 host）。
- all-reduce 取 **均值**后，每 rank 各自跑 `GpuAdamW.step`——因为各 rank 梯度一致、优化器状态从相同 init 同步演化，参数自然保持一致（无需再同步参数）。
- grad clip 的 global-norm 在 all-reduce **之后**算：`clip_grad_norm_gpu` 的 `sumsq_accum` 已是 device reduction，多卡只需把那个标量再 all-reduce 一次（或对已 all-reduce 的梯度本地算，因梯度已一致，结果天然相同）。

157
docs/07-distributed.md Normal file
View File

@@ -0,0 +1,157 @@
# Phase T8: Distributed Data Parallel (DDP) — Design Document
## Goal
T7 把单卡训练榨到 **~8500 tok/s**（梯度常驻 device，AdamW + grad-norm 都在 device 侧）。
T8 的目标是**把训练 scale 到 dash5 的多张 RTX 5090**：用 **NCCL 数据并行**——每个 rank 跑 global batch
的一个分片，反向后把每个参数的梯度在显存里 **all-reduce 求和、按 world size 取均值**，然后各 rank 用
**自己**那份（已同步的）优化器跑 `GpuAdamW.step`。同 init + 同优化器超参/状态 ⇒ 参数永远逐位一致，
无需再同步权重。
验收（两条都要）：
1. **正确性**：DDP 必须对住单卡。单卡跑 K 步 global batch B，DDP 用 2 rank 各跑同一批数据同序的 B/2 →
loss 轨迹与单卡在紧容差内吻合（纯梯度平均，理论上近乎一致）；且一步后断言两个 rank 的参数完全相同。
2. **吞吐**：固定每卡 workload，测 1/2/4 卡的 tok/s → 近线性 scaling，给表。
**本 Phase 范围内只做 DDP**。tensor/pipeline 并行、sharded optimizer 是可选 bonus，**不做**（model 太小，
当前学习目标是把 DDP 这条最朴素也最核心的多卡链路吃透；切分式并行留待 model 放大）。导出回流 xserv 是 T9。
## Module Layout
```
crates/xtrain-distributed/ # 新 crate（镜像 xserv-distributed）
├── build.rs # nvcc 探测 → no_cuda cfg；有 nvcc 时链 -lnccl -lcudart（同 xserv build.rs）
├── Cargo.toml # 依赖 cuda/tensor/autodiff/model/optim/train
├── src/ffi.rs # NCCL FFI：GetUniqueId/CommInitRank/AllReduce/CommDestroy/Group{Start,End}
├── src/lib.rs # DdpContext（comm bootstrap + all_reduce_average_grads）+ get_unique_id
├── src/ddp.rs # DDP 训练 step（train_rank）+ thread-per-GPU launcher + 确定性 build_model
└── src/bin/train_ddp.rs # 多 rank 启动器 / 吞吐 driver（CUDA_VISIBLE_DEVICES 选卡）
crates/xtrain-distributed/tests/
└── ddp_correctness.rs # 2 卡 vs 单卡 loss 对拍 + 跨 rank 参数一致 + 1/2/4 卡吞吐表
Cargo.toml # workspace members += xtrain-distributed
docs/07-distributed.md # 本文
```
整个 crate 用 `#![cfg(not(no_cuda))]` 在 crate 根门控：本地（无 nvcc）crate 编译为空，不引用任何 NCCL 符号，
`cargo check` 直接过；dash5 上 `not(no_cuda)` 全量编译并链接 NCCL。bin 另带一个 `#[cfg(no_cuda)]` stub main。
## Key Design Decisions
### ① NCCL comm bootstrap（rank 0 发 UniqueId）
NCCL 的握手：**rank 0 调 `ncclGetUniqueId`** 生成一个 128 字节不透明 id，**带外**分发给所有 rank，每个 rank
用 `ncclCommInitRank(comm, world, id, rank)` 加入通信组。`ncclUniqueId` 按 NCCL ABI **按值传递**（128 字节 struct），
FFI 里 `#[repr(C)] struct NcclUniqueId { internal: [c_char;128] }`，`ncclCommInitRank` 的 `commid` 参数直接传值。
并发 init 必须用 `ncclGroupStart()/ncclGroupEnd()` 包住，否则多个线程同时 `CommInitRank` 会互相等待握手而死锁——
group 让它们 rendezvous。这套与 xserv-distributed 的 `TpContext::init` 完全一致。
### ② Launch model：thread-per-GPU（与 xserv 一致）
**单进程、每卡一个 OS 线程**，线程启动先 `cudaSetDevice(dev)` 绑定本卡。选这个模型的硬约束：xtrain 的
autograd `Var` 是 `Rc<RefCell<…>>`，**不是 `Send`**，整张计算图不能跨线程。所以**每个 rank 线程在闭包内本地
`build_model` 自己的 `Var` 图**——跨线程边界的只有 `UniqueId`（`Copy`）、标量 config、和 `&Corpus`（只读共享，`Sync`）。
`launch()` 用 `std::thread::scope`：rank 0 发 id → 每线程 `DdpContext::init` → 本地建模 → `train_rank` → join。
`DdpContext` 持有 `comm`（裸指针），`unsafe impl Send` 因为它被恰好一个 rank 线程独占。
为什么不用多进程（torchrun 式）：单机多卡，单进程多线程改动最小、无 IPC；UniqueId 跨线程就是 move 一个 `Copy`
struct，跨进程才需要写文件/env 分发。多进程留待真正跨节点。
### ③ all-reduce-then-local-step：梯度在显存里平均，各 rank 各自 step
T7 起**梯度常驻 device**（`Var::grad()` 是 device tensor）。DDP 的通信因此极简——**直接对每个参数的 `.grad()`
device buffer 做 in-place AllReduce**，零 host 往返：
```rust
// DdpContext::all_reduce_average_grads(params)
for p in params { if let Some(g) = p.grad() {
ncclAllReduce(g.data_ptr(), g.data_ptr(), g.numel(), NCCL_FLOAT32, NCCL_SUM, comm, null_stream);
}}
for p in params { if let Some(g) = p.grad() {
launch_scale_inplace_f32(g.data_ptr(), 1.0/world, g.numel(), null_stream); // 复用 T7 的 kernel
}}
```
AllReduce 用 **fp32 + Sum**（梯度就是 fp32），发在 **null/legacy stream** 上——xtrain 每个 kernel 都在 null stream
（`std::ptr::null_mut()`）launch，所以 AllReduce 自然排在产出梯度的 backward kernel 之后、消费它的 scale/optimizer
kernel 之前，**无需额外 barrier**。平均直接复用 T7 的 `launch_scale_inplace_f32`（grad-clip 用的同一个 kernel）。
之后**每个 rank 各自跑 `GpuAdamW.step`**——不再同步权重。
### ④ 为什么参数保持一致
三个充分条件：(a) **同 init**：所有 rank（含单卡 baseline）用同一个确定性 `build_model`（同 LCG 种子），起点逐位相同；
(b) **同梯度**：NCCL AllReduce 对参与的所有 rank **返回逐位相同的归约结果**（这是 NCCL 的保证），平均、clip-norm
（device reduction，确定性）、AdamW kernel 都是确定性的 → 每步喂给优化器的梯度逐位相同；(c) **同优化器状态**：
每个 rank 各自维护 m/v，但因为起点相同、每步输入相同、超参相同，状态演化也相同。
⇒ 各 rank 参数**逐位一致**（测试里断言 `max|p0-p1| == 0.0`），不需要任何权重再同步。
> 对比单卡：DDP 与单卡的最终参数只在 **fp 容差**内一致（`<1e-3` rel），不逐位——因为求和顺序不同
> （单卡 tape 顺序 SUM B 个；DDP 各 rank 先 SUM 分片再 NCCL 跨 rank SUM），fp 加法不结合。跨 rank 才逐位。
### ⑤ Batch sharding：对住单卡的均值
单卡 loop：`batch_size=B` 个序列各 forward+backward，tape **SUM** 出 `Σ_B`，再 `clip(pre_scale=1/B)` → batch 均值。
DDP 严格对齐（保证 all-reduce 后的和与单卡逐序列一致）：**每个 rank 推进同一个 RNG、抽出整批 B 个序列，但只对
分到自己的那些（`i % world == rank`）跑 forward+backward**。各 rank 的并集 == 单卡那一批同序，于是：
```
rank 本地: Σ_local（分片内 b = B/world 个的和，tape SUM）
AllReduce: Σ_global = Σ_ranks Σ_local （= 单卡的 Σ_B，只是求和顺序不同）
/world: Σ_global / world
clip pre_scale = 1/b_local = world/B: Σ_global/world · world/B = Σ_global / B ← 与单卡 clip(1/B) 同一个量
```
「按 world 取均值」放在通信原语里（语义清晰），剩下的 `1/b_local` 交给本就存在的 clip 前置缩放完成——
最终喂给 AdamW 的就是 **global batch 均值梯度**，与单卡完全对齐。loss 日志同理：本地 loss 和跨 rank AllReduce 后
除以 B，得 global 均值（每个 rank 拿到相同值）。
## 验证方法
测试 `tests/ddp_correctness.rs`（`#[cfg(not(no_cuda))]`，<2 卡自动 skip），用合成语料（无需 tokenizer/数据文件）：
### 正确性：`ddp_matches_single_gpu_and_params_consistent`
同一 config + 合成语料，跑 20 步：
- **(a) loss 对拍**：单卡 baseline（world=1，global batch 8）vs 2-rank DDP（各 4）→ 整条 loss 轨迹 `max_rel < 1e-3`。
- **(b) 跨 rank 参数一致**：断言 `max|p_rank0 - p_rank1| == 0.0`（逐位，NCCL 保证 + 确定性 step）。
- **(c) DDP vs 单卡参数**：`max rel|Δp| < 1e-3`（fp 求和顺序差，不逐位）。
### 吞吐：`ddp_throughput_scaling`
固定**每卡 batch=8**（batch 随 world 放大，每 rank 成本不变），测 world ∈ {1,2,4} 的 global tok/s，打印：
```
GPUs | tok/s (global) | speedup
1 | ... | 1.00x
2 | ... | ~2x
4 | ... | ~4x
```
近线性即过（tiny model latency-bound + 跨卡通信开销会让大 world 略低于理想）。
### dash5 实跑
```bash
export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
# 选空闲卡（dash5 共享，先看 nvidia-smi）：
CUDA_VISIBLE_DEVICES=<idle GPUs> \
cargo test -p xtrain-distributed --release -- --nocapture --test-threads=1
# 真训练 / 吞吐 driver：
CUDA_VISIBLE_DEVICES=0,1 cargo run -p xtrain-distributed --release --bin train_ddp -- \
100 64 16 /opt/wjh/models/gpt2/tokenizer.json data/tinystories-valid-3mb.txt
```
实测数字回填见 xtrain.md T8 note / commit。
## 不做（本 Phase 范围外，记为 follow-up）
- **Tensor / pipeline 并行**：要切权重、layer 内/跨 layer 通信，对 tiny model 收益小、改动大 → model 放大再做。
- **Sharded optimizer（ZeRO）**：每卡只存 1/world 的优化器状态。当前 model 优化器状态很小，不是瓶颈。
- **bf16 AllReduce / 通信压缩**：fp32 已够，且会动数值正确性（与 T7 同理由延后）。
- **多进程 / 跨节点 bootstrap**：单机多卡用线程模型足够；跨节点再上文件/env 分发 UniqueId。

176
docs/08-export-xserv.md Normal file
View File

@@ -0,0 +1,176 @@
# Phase T9: Export to xserv — Design Document
## Goal
闭环 xtrain ↔ xserv：把 xtrain 练出的权重导成 **xserv 的 Qwen3 loader 能直接加载并服务**的格式
（HF 命名的 `config.json` + `model.safetensors` + 复用的 gpt2 `tokenizer.json`），让推理侧 xserv
跑出**与 xtrain 自身一致**的生成结果。
验收信号：**同一份权重 + 同一 prompt，xserv 的贪心生成 token 序列对住 xtrain 的贪心生成**
（logits 在浮点容差内吻合）。这是整条 P0→P6 学习链的收口——训练栈练出来的东西，推理栈真的能用。
> 范围与诚实边界：**不改 xserv**。xserv 是独立项目，T9 只调整 xtrain 的导出去适配 xserv 既有 loader。
> 架构差异中只有一项是「结构性」的（QK-norm，见下），处理方式见 **Key Design Decision 1**。
## 关键第一步：架构 diff（xtrain TinyTransformer vs xserv qwen3.rs）
逐 op 读 `crates/xtrain-model/src/model.rs` 对 `~/projects/xserv/crates/xserv-model/src/qwen3.rs`
（forward 路径 `forward_with_cache`，即 `dump-logits` 走的 prefill）：
| 维度 | xtrain TinyTransformer | xserv Qwen3 | 是否兼容 / 处理 |
|------|------------------------|-------------|----------------|
| RMSNorm 公式 | `x*rsqrt(mean(x²)+eps)*γ`（无均值减） | 同 | ✅ 一致 |
| RMSNorm eps | `cfg.eps`（tiny=1e-5） | `rms_norm_eps`（config.json 提供） | ✅ 导出 eps 到 config |
| **QK-norm** | **T9 前没有** | **强制** per-head RMSNorm(Q,K)（`q_norm`/`k_norm`，shape `[head_dim]`） | **结构性差异** → 见 Decision 1 |
| RoPE 约定 | rotate_half，`freq=θ^(-2i/hd)`，pos=行号 | rotate_half，cos/sin cache `freq=1/θ^(2i/hd)`，pos=token idx | ✅ **逐式一致**（见 rope.cu）|
| RoPE θ | `cfg.rope_theta`（10000） | `rope_theta`（config 提供，默认 1e6） | ✅ 导出 θ 到 config |
| Attention scale | `1/√head_dim` | `1/√head_dim`（attention.rs / flash） | ✅ 一致 |
| Causal mask | 加性 `-1e9` 上三角 | causal flag（online softmax） | ✅ 同义 |
| GQA | MHA（无 kv 分组） | 支持 GQA（`num_kv_heads`） | ✅ 设 `num_key_value_heads = num_attention_heads`（退化为 MHA）|
| SwiGLU | `down(silu(gate)∘up)`，gate/up 独立 proj | 同（融合存 `gate_up_proj`，loader 内部 cat）| ✅ 一致，导出仍分开 gate/up（loader 自己 cat）|
| 偏置 | 无 | `attention_bias=false`，不读 bias | ✅ 一致 |
| final norm | `final_norm` 后接 lm_head | `model.norm` 后接 lm_head | ✅ 一致 |
| embedding tying | 独立 `lm_head` | 独立（`tie_word_embeddings=false`）| ✅ 一致 |
| 2D 权重 layout | `[in,out]`，算 `x@W` | `[out,in]`，算 `x@Wᵀ`（loader `.transpose(0,1)`）| ⚠️ 导出须转置 |
| **dtype** | **f32**（训练 + 自身推理） | **BF16 only**（kernel 全 assert BF16）| ⚠️ 导出转 BF16；数值上不可能 bit-exact，见 Decision 2 |
| vocab | gpt2 BPE，50257 | config 提供 | ✅ 导出 vocab |
**结论**：除 QK-norm 一项外，其余差异都是机械的（命名 / 转置 / dtype / 退化 GQA）。QK-norm 是
xserv 对 Qwen3 的**强制**步骤（`head_rmsnorm(q,q_norm)` / `head_rmsnorm(k,k_norm)` 无条件执行，
且 γ=1 也不是恒等——它仍按每个 head 向量的 RMS 做归一），xtrain 训练时从未施加 → 若不处理，
xserv 的前向数学与 xtrain 不同，闭环不可能成立。
## Key Design Decisions
### Decision 1：给 xtrain 加 per-head QK-norm（而不是伪造 match，也不是停在 blocker）
逃生舱里给了三条路：① 给 xtrain 补 QK-norm 重训以对齐 Qwen3；② 改 xserv（被禁止）；
③ 当作 blocker 报告。选 ①——因为它**真的能闭环**且改动**外科手术级**：
- xserv 的顺序是 `reshape → head_rmsnorm → transpose_for_rope → rope`，即 **QK-norm 在 RoPE 之前**，
作用在每个 `[head_dim]` 的 head 向量上。
- xtrain 复用既有的 2D `rms_norm` op：把 `[seq,nh,hd]` reshape 成 `[seq*nh, hd]`，用 `[hd]` 的 γ
做 rms_norm，再 reshape 回去，**插在 reshape 与 rope 之间**——与 xserv 逐步对齐。autograd 全自动
（rms_norm/reshape 都已有 backward），优化器/checkpoint/DDP 都按 `params()` 泛化迭代，自动兼容。
- 每 block 新增 `q_norm`/`k_norm` 两个 `[head_dim]` leaf；`params()` 顺序在 `wv` 之后、`wo` 之前插入
`q_norm,k_norm`；`num_params()` 加 `2*head_dim/layer`；PyTorch 对拍参考（parity.py / adamw_parity.py）
同步加 QK-norm，名单同步——改完仍全套自洽。
这样导出的权重是**真·Qwen3 兼容**的（训练时就带 QK-norm），不是凑出来的假象。
### Decision 2：BF16 是 xserv 的硬约束 → 闭环判据用「贪心 token 一致」而非 bit-exact
xserv 的 Qwen3 前向 **只支持 BF16**（embedding/rmsnorm/gemm/silu/rope kernel 全部
`assert_eq!(dtype, BF16)`，`dump-logits` 也按 bf16 读 logits）。xtrain 是 f32。所以：
- 导出时把所有权重 `bf16::from_f32` 转成 BF16 写入 safetensors。
- **数值上不可能 bit-exact**：BF16 仅 8 位尾数，权重舍入 + 前向全程 BF16 累加，相对 xtrain 的 f32
会有 ~1e-2 量级的 logits 漂移。因此**闭环判据定为「同 prompt 下贪心 argmax token 序列一致」**
（+ 报告 logits 的 top-1/分布吻合度），这是 BF16 推理对 f32 训练能给出的最强、且诚实的判据。
### Decision 3：导出复用 train.rs 的 config + checkpoint，零猜测
导出器 `bin/export_safetensors.rs` 用与 `bin/train.rs` **完全相同**的 `Config`（gpt2 vocab、
`n_layers=4`、`Config::tiny()` 其余字段）建空模型 → `checkpoint::load_into` 灌入训练权重 →
按 `params()` 稳定序映射。tokenizer.json 直接 copy 进导出目录，两侧用同一份 BPE。
## 张量名 + layout 映射表
xtrain `params()` 序（T9 后）：
`embed[vocab,dim]` → 每 block `[attn_norm, wq, wk, wv, q_norm, k_norm, wo, ffn_norm, w_gate, w_up, w_down]`
→ `final_norm[dim]` → `lm_head[dim,vocab]`。
| xtrain 参数 | shape | → HF Qwen3 名 | HF shape | 操作 |
|-------------|-------|---------------|----------|------|
| `embed` | `[vocab,dim]` | `model.embed_tokens.weight` | `[vocab,dim]` | keep（行索引两侧同）|
| `attn_norm` | `[dim]` | `model.layers.{i}.input_layernorm.weight` | `[dim]` | keep |
| `wq` | `[dim,dim]` | `…self_attn.q_proj.weight` | `[dim,dim]` | **transpose** |
| `wk` | `[dim,dim]` | `…self_attn.k_proj.weight` | `[dim,dim]` | **transpose** |
| `wv` | `[dim,dim]` | `…self_attn.v_proj.weight` | `[dim,dim]` | **transpose** |
| `q_norm` | `[head_dim]` | `…self_attn.q_norm.weight` | `[head_dim]` | keep |
| `k_norm` | `[head_dim]` | `…self_attn.k_norm.weight` | `[head_dim]` | keep |
| `wo` | `[dim,dim]` | `…self_attn.o_proj.weight` | `[dim,dim]` | **transpose** |
| `ffn_norm` | `[dim]` | `…post_attention_layernorm.weight` | `[dim]` | keep |
| `w_gate` | `[dim,ffn]` | `…mlp.gate_proj.weight` | `[ffn,dim]` | **transpose** |
| `w_up` | `[dim,ffn]` | `…mlp.up_proj.weight` | `[ffn,dim]` | **transpose** |
| `w_down` | `[ffn,dim]` | `…mlp.down_proj.weight` | `[dim,ffn]` | **transpose** |
| `final_norm` | `[dim]` | `model.norm.weight` | `[dim]` | keep |
| `lm_head` | `[dim,vocab]` | `lm_head.weight` | `[vocab,dim]` | **transpose** |
全部以 **BF16** dtype 写入。config.json 字段：`architectures=["Qwen3ForCausalLM"]`、`model_type="qwen3"`、
`vocab_size`、`hidden_size=dim`、`intermediate_size=ffn`、`num_hidden_layers`、`num_attention_heads`、
`num_key_value_heads=num_attention_heads`、`head_dim`、`rms_norm_eps=eps`、`rope_theta`、
`tie_word_embeddings=false`、`attention_bias=false`、`hidden_act="silu"`。
## Module Layout
```
crates/xtrain-model/src/model.rs # +q_norm/k_norm leaf；attention 插 per-head QK-norm；params() 序更新
crates/xtrain-model/src/config.rs # num_params() 计入 QK-norm γ
crates/xtrain-train/src/bin/export_safetensors.rs # 导出器（本 Phase 核心实现）
crates/xtrain-train/Cargo.toml # +half, +safetensors="0.5"（对齐 xserv）, +bin
crates/xtrain-model/tests/parity{.py,_dump.rs} # PyTorch 对拍同步加 QK-norm
crates/xtrain-train/tests/adamw_parity{.py,_dump.rs}# 同上
```
## 验证方法
1. **本地**：`cargo check --workspace` + `cargo fmt --all -- --check` 过（导出器 GPU 体 gated 在
`not(no_cuda)`，host 侧只 check）。
2. **dash5（闭环）**：
```bash
export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
# ① 训练一个小模型 → checkpoint（或复用已训的）
cargo run -p xtrain-train --release --bin train -- \
/opt/wjh/models/gpt2/tokenizer.json data/tinystories-valid-3mb.txt \
<steps> /tmp/xtrain_tinystories.ckpt
# ② 导出
cargo run -p xtrain-train --release --bin export_safetensors -- \
/tmp/xtrain_tinystories.ckpt /opt/wjh/models/gpt2/tokenizer.json /tmp/xtrain_export
# ③ xserv 加载 + dump logits（同 prompt）
# (在 /opt/wjh/projects/xserv) cargo run -p xserv-model --release --bin dump-logits -- /tmp/xtrain_export "<prompt>"
# ④ 对拍：xserv 贪心 token 序列 vs xtrain 自身贪心（sample.rs generate, temp=0）
```
判据：**贪心 token 序列一致**（BF16 推理 vs f32 训练，logits top-1 吻合；分布在 BF16 容差内）。
## 验证结果（dash5 实跑，capture）
**训练**（CUDA_VISIBLE_DEVICES=0，1200 步，gpt2 vocab，dim 32 / 4 层 / 2 头 / ffn 64，~5M 参，8700 tok/s）：
`loss 10.84 → 3.59`，贪心采样输出连贯英文（QK-norm 加入后训练/采样无回归）。
**导出**：`export: 47 tensors`（embed + 4×11 + final_norm + lm_head），写出 config.json（见上）+
model.safetensors（BF16，6.5 MB）+ tokenizer.json。
**① logits 数值对拍**（同 prompt `"Once upon a time"`，token ids `[7454, 2402, 257, 640]`）：
| rank | xtrain (f32) | xserv (BF16) | token |
|------|--------------|--------------|-------|
| 0 | 11.7711 | 11.7500 | `,` |
| 1 | 10.4724 | 10.5000 | ` there` |
| 2 | 6.6288 | 6.6562 | ` upon` |
| 3 | 6.5125 | 6.5000 | ` to` |
| … | … | … | … |
| 10 | 5.3614 | 5.3438 | ` she` |
top-1 一致（`,`，id 11）；top-11 token 排序完全一致；logit 绝对差 ~1e-2（~0.2–0.9% 相对），
正是 **BF16 推理 vs f32 训练** 的预期舍入漂移，无结构性误差。
**② 贪心生成逐 token 一致**（xserv `xserv-cli --max-tokens 40` vs xtrain `sample.rs generate temp=0`）：
```
prompt "Once upon a time":
xtrain: Once upon a time, there was a little girl named Lily. Timmy loved to play
with her mommy. One day, Timmy's mommy's mommy's mommy. "I'm sorry, I
xserv: Once upon a time, there was a little girl named Lily. Timmy loved to play
with her mommy. One day, Timmy's mommy's mommy's mommy. "I'm sorry, I ← 逐 token 相同
prompt "One day":
两侧均: One day, Timmy's mommy's mommy's mommy. "I'm sorry, I can't be careful and
be careful. I'm sorry, I can't have a good time. ← 逐 token 相同
```
**结论：闭环成立**。xtrain 练出的权重，经导出后由 xserv 加载并服务，贪心生成与 xtrain 自身**逐 token 一致**，
logits 在 BF16 容差内吻合。整条 P0→P6 学习链收口。
> 注：xtrain 采样每步重跑全量 forward（无 KV cache），xserv 走 KV-cache prefill+decode；两者都是对同一
> logits 的 greedy argmax，故序列一致。BF16 漂移未在 40 步内造成任何 argmax 翻转。

187
docs/09-batched-forward.md Normal file
View File

@@ -0,0 +1,187 @@
# Phase T10: Batched 多序列 Forward — Design Document
## Goal
修 **KI-1 的根因**。v2 暴露、v3 重诊断的结论（见 [docs/known-issues.md](known-issues.md)）：
吞吐瓶颈**不是** all-reduce，而是 **单序列模型设计 launch-bound**——T5 的模型一次只过**一条**
序列，每个 op 的 GEMM 都是 `[seq, dim]` 这种小矩阵，喂不饱 GPU；"batch" 是靠**循环 B 次 forward +
让 tape SUM 梯度**伪造的，于是 GPU util 只有 0–15%、显存占 ~8%，加大 global batch 也只是按比例
增加串行 kernel-launch（v3 实测 gbatch 32→256 仅 +1.2%）。
T10 的目标：给 model + autograd 加 **batch 维**，把一个 step 的 B 条序列**一次性**过模型，让线性层变成
**一个大 GEMM** 填满 GPU——这是 launch-bound 的根本解。**硬闸门是正确性**：所有既有 grad-check /
PyTorch 对拍 / overfit / DDP 跨 rank 一致**必须仍过**（PyTorch 对拍现在用 **B>1**），xserv 推理仍**逐
token 一致**；在此之上拿吞吐收益。
## 核心设计：linears 摊平为 `[B*S, dim]`，attention 批量化
一个关键观察：**transformer 里只有 attention 需要"序列感知"**，其余全是 per-token 的。
- **线性投影 / 逐元素 / norm / embedding / CE 天然是 2D `[rows, dim]` 上的运算**，不在乎 `rows` 是
`seq` 还是 `B*S`。所以把激活**摊平成 `[B*S, dim]`** 喂进去，这些 op **原样复用**：
q/k/v/o、gate/up/down、lm_head 全部变成**一个大 `[B*S, dim] × [dim, out]` GEMM**（填 GPU 的收益所在）；
embedding 对 `[B*S]` ids gather；rms_norm / qk_norm / swiglu / silu / 逐元素按行（per-token）；
cross_entropy 对 `[B*S, vocab]` vs `[B*S]` targets，**mean over B*S 行 == 各序列 loss 的均值**。
- **attention 是唯一序列感知的 op**：reshape 成 `[B, nh, S, hd]`，**每个序列内**做因果 SDPA（per-seq
因果 mask，**绝不跨序列 attend**），写回 `[B*S, dim]`。在 `(B, nh)` 维上批量。
- **RoPE 位置必须 per-sequence 复位**：位置 = 在**自己序列内**的下标（`row % S`），**不是**摊平后的全局
行号。这是正确性陷阱，显式处理（kernel 加 `period` 参数）。
```text
ids [B*S]
└ embedding → [B*S, dim]
每层 block：
rms_norm → [B*S, dim] (per-token)
attention:
x@Wq/Wk/Wv → [B*S, dim] (大 GEMM)
reshape → [B*S, nh, hd]
qk-norm + RoPE(period=S) → [B*S, nh, hd] (RoPE 位置 = row % S)
→ [B, nh, S, hd] → [B*nh, S, hd]
fused batched SDPA(因果) → [B*nh, S, hd] (2× 批量GEMM + 1× causal-softmax)
→ [B*S, dim]; @Wo → [B*S, dim] (大 GEMM)
+residual; rms_norm; SwiGLU MLP(大 GEMM); +residual
final rms_norm; @lm_head → [B*S, vocab]
cross_entropy([B*S, vocab], targets[B*S]) → 标量（B*S 行的 mean）
```
`forward(ids[seq])` 现在只是 `forward_batched(ids, batch=1)` 的特例 → **采样 / 推理路径 batch=1 不变**。
## Module Layout
```
csrc/ops/attention.cu # 新：causal 行 softmax（含 scale + per-seq 因果 mask）
csrc/ops/nn.cu # rope/rope_dx 加 period（位置 = tok % period）
csrc/ops/model.cu # 加 transpose_4d12（[B,S,nh,hd]<->[B,nh,S,hd]）
crates/xtrain-cuda/
├── src/cublas.rs # 加 sgemm_strided_batched（批量 GEMM，行主序同 sgemm 套路）
├── src/ffi.rs # +cublasSgemmStridedBatched / launch_softmax_causal / transpose_4d12；rope 加 period
└── build.rs # +attention.cu
crates/xtrain-tensor/src/tensor.rs # 加 attention/attention_backward（fused，2/4× 批量GEMM）、transpose_4d12；rope(+period)
crates/xtrain-autodiff/
├── src/ops.rs # 加 attention 节点、transpose_4d12 节点；rope(+period)
└── tests/autograd.rs # +rope_batched / transpose_4d12 / attention(batched) grad-check
crates/xtrain-model/
├── src/model.rs # forward_batched(ids,batch)/loss_batched；attention 走 fused 批量 op；删 causal_mask 叶
├── src/lib.rs # +batched_ids_tensor
├── tests/batched.rs # 新：batched == looped 单序列（logits + 梯度）
├── tests/parity_dump.rs + parity.py # PyTorch 对拍改 B=2,S=4（per-seq RoPE + per-seq mask）
crates/xtrain-train/src/train_loop.rs # 真 batch：一次 batched forward/backward 替代 loop+SUM；clip pre-scale 1.0
crates/xtrain-distributed/
├── src/ddp.rs # 每 rank 一次 batched forward（b_local 条）；clip pre-scale 1.0
└── tests/ddp_correctness.rs # 单卡基线也改 batched（对齐）
```
## Key Design Decisions
### ① 为什么"摊平 linears"是主要收益
GPU 填不满是因为 GEMM 太小（`[256, 384]` 这种）。摊平成 `[B*S, dim]` 后，B=16/S=256 时左矩阵是
`[4096, 384]`——同一个 cuBLAS GEMM kernel，但 M 大了 16×，一次 launch 干 16× 的活，启动开销被摊薄，
SM 占用率上去。**embedding / rms_norm / silu / CE 等本就按行**，摊平后只是行数变多，**零改动复用**。
### ② RoPE 位置 per-sequence 复位（正确性陷阱）
摊平后第 `r` 行是序列 `r/S` 的第 `r%S` 个 token。RoPE 角度 = `pos * freq`，**`pos` 必须是 `r%S`**，
否则第 2 条及之后的序列会用错位置、与单序列结果不一致、PyTorch 对拍直接挂。kernel 加 `period`
（= 序列长度 S）参数，`pos = tok % period`；`period == tokens`（单序列）时退化为原"位置=行号"。
反向同样按 `period` 取角度（RoPE 是正交映射，反向 = 逆旋转，不需缓存 forward）。
### ③ Attention：先摊平到 `[B*nh, S, hd]`，fused 批量 SDPA
q/k/v 投影后是 `[B*S, nh, hd]`（顺序 b,s,head,hd）。要喂批量 GEMM，需排成 `[B*nh, S, hd]`（每个 batch
元素 = 一个 (b, head) 的整条序列）：
```text
[B*S, nh, hd] → reshape [B, S, nh, hd] → transpose_4d12 → [B, nh, S, hd] → reshape [B*nh, S, hd]
```
`transpose_4d12`（轴 1,2 互换，新加的结构 op，自反式 backward）是关键。然后 **fused attention op**：
```text
scores[B*nh,S,S] = Q · Kᵀ (cublasSgemmStridedBatched)
P[B*nh,S,S] = softmax(causal(scores / √hd)) (launch_softmax_causal，行内做 mask+scale)
out[B*nh,S,hd] = P · V (cublasSgemmStridedBatched)
```
**整个 attention 不论 B*nh 多大，都是 3 次 kernel launch**（2 批量 GEMM + 1 softmax），没有 per-head /
per-seq 的 Python 循环，也没有 host 往返。
**因果 mask 内联在 softmax kernel**：第 `r` 行的 query 位置 = `r % S`，列 `j > r%S` 是未来 → 概率置 0，
不需要 `[S,S]` 的 `-1e9` 加性 mask 张量（连 `causal_mask` 叶子都删了）。`1/√hd` 的 scale 也折进
softmax（取 max/exp 前乘），省一个 scale pass。
> **踩坑记录（重要）**：T10 的第一版 attention 用"per-(batch,head) 循环 + `split_heads`/`merge_heads`"，
> 而 split/merge_heads 内部走 **host 往返**（`to_device(Cpu)` 再传回）。linears 确实摊平成大 GEMM 了，
> 但 B=16 时 attention 路径里成千上万次小 kernel + 几 MB 的 host memcpy **反而把吞吐拖到 1127 tok/s
> （比单序列还慢）**。教训：摊平 linears 是必要非充分——attention 的 host 往返 / launch 风暴必须一起干掉。
> fused 批量 SDPA 之后才到 **25.6K tok/s**。
### ④ Attention backward（手写，grad-check 兜底）
fused op 的反向是标准 attention 反传，全用批量 GEMM + 复用既有 `softmax_backward`：
```text
dP = dOut · Vᵀ ; dV = Pᵀ · dOut
dScores = softmax_jacobian(P, dP) · (1/√hd) (复用行 softmax 反向；mask 处 P=0 → 自动零梯度)
dQ = dScores · K ; dK = dScoresᵀ · Q
```
新加的 `attention` / `transpose_4d12` 都有 finite-diff grad-check（见 ⑥）。
### ⑤ 训练 loop：真 batch 替代 loop+SUM
旧：循环 B 次 `model.loss(单序列)` + `backward()`，靠 tape 把 B 份梯度 SUM，clip 时 ×`1/B` 还原 batch-mean。
新：`batched_ids_tensor` 把 B 条序列摊平成 `[B*S]`，**一次** `loss_batched` + **一次** `backward()`。
CE 已是 B*S 行的 mean = batch-mean loss，**backward 直接给出 batch-mean 梯度**，所以 **clip pre-scale = 1.0**
（不再有 loop+SUM+×1/B）。
**DDP 同理且保持等价**：每 rank 跑 `b_local = B_global/world` 条的**一次** batched forward → backward
梯度 = 本地 mean `Σ_local/b_local`；`all_reduce_average`（跨 rank 求和 /world）得
`Σ_global/(world·b_local) = Σ_global/B_global` = 全局 batch-mean → clip pre-scale 也 = 1.0。
单卡基线（`ddp_correctness` 里）同步改成对全局 batch 的一次 batched forward，二者对齐。
## 验证方法
**双闸门，都必须过。**
### 正确性（无回归）
- **算子级 finite-diff**（`xtrain-autodiff`，15 个）：新增 `rope_batched`（per-seq 位置）、`transpose_4d12`、
`attention(batched)` 的 dQ/dK/dV，连同既有 12 个全过。
实测：`attn(batched) dQ 7.5e-3 / dK 1.5e-2 / dV 2.9e-4`，`transpose_4d12 8.2e-5`，`rope_batched 4.5e-4`。
- **batched == looped 单序列**（`xtrain-model/tests/batched.rs`，新）：同一组权重，batched forward 对
"逐序列单独 forward 再拼接" 的 logits + 每参数梯度逐一对比。实测 **logits 完全一致(0.0)**，
梯度 max rel **6.4e-4**，loss 完全一致。
- **PyTorch 对拍 B>1**（`parity.py`，改 B=2/S=4）：等价 PyTorch 模型（per-seq RoPE `pos=row%S`、per-seq
因果 mask、QK-norm、SwiGLU）对拍 forward logits + 全部 25 个参数梯度。实测 **loss relerr 5e-8、logits
6.9e-6、梯度全在 rtol 2e-2 内**。
- **overfit**（27/27 token）、**checkpoint 逐位**、**AdamW 对 torch**、**DDP 跨 rank 参数 bit-identical
(0.0) + DDP loss 对单卡 5.7e-7**：全过。
- **xserv 闭环**：短训→导出→xserv serve，对 xtrain 贪心**仍逐 token 一致**；采样路径 batch=1 仍工作。
### 吞吐（收益）
单卡 dim384/12L/12h、batch 16、seq 256，back-to-back：
| 路径 | tok/s | GPU util | 显存 |
|---|---|---|---|
| **before**（单序列 launch-bound，KI-1 基线）| ~1653 | 0–15% | ~3 GB |
| T10 第一版（looped split/merge，host 往返）| 1127 | ~11% | 14.8 GB |
| **after**（fused 批量 attention，batch 16）| **25627** | **37% 均值 / 54% 峰** | 10.2 GB |
| **after**（batch 32）| **40263** | — | — |
→ 单卡相对 KI-1 基线 **~15.5×**（batch 16）/ **~24×**（batch 32），GPU util 0–15% → 37–54%。
**DDP 4 卡（dim384, per-rank batch 32, global 128）**：1 卡 40.3K → 4 卡 47.2K tok/s（global），
仅 **~1.17×**。这是 batched forward **新暴露**的下一层瓶颈：单卡 compute 快了 15–24×后，每步**对全部
67M 参数的 eager all-reduce + host 侧 optimizer/clip 同步**成了 DDP 的主导开销（不随卡数缩小）。
注意：**单卡 batch 32 = 40K tok/s 已经把 KI-1 时代的 4 卡 3163 tok/s 干翻 ~12×**——根因（单卡
launch-bound）已修。DDP 近线性需要 **bucketed / 与 backward overlap 的 all-reduce**（KI-1 修复项 2），
此前 all-reduce 非瓶颈做了没用，现在才有意义——列为 v3+ 的 follow-up（本 Phase 范围外）。
## 给 v3 的 note（已解锁）
batched forward 就位后，v3（dim512/16L、200–300M tok）建议：**per-rank batch 16–32、seq 256，4 卡
global batch 64–128**。按单卡 ~25K tok/s（dim384，dim512 略低）、4 卡放大，200–300M tok 的训练时间从
KI-1 时代估的 17–26h 压到**数小时级**。bucketed/overlapped all-reduce（KI-1 修复项 2）现在才有意义，
作为 v3 之后的进一步优化项。

160
docs/10-caching-allocator.md Normal file
View File

@@ -0,0 +1,160 @@
# Phase T11: Device Caching / Pool Allocator — Design Document
## Goal
修 **KI-5 的根因**。T10 修掉单卡 launch-bound（1653→40K tok/s）后，DDP 多卡仍只有 ~1.4× 的弱扩展。
T11 第一版拟修复（**分桶 all-reduce**）经 dash5 实测**证伪并 revert**：grad all-reduce 每步只占
**~6–7%**，融成一发对 1/2/4/8 卡几乎无差（见 [docs/known-issues.md](known-issues.md) KI-5 表）。
实测重新定位的根因：**每个 tape op 的输出都走 `Tensor::zeros` → `GpuBuffer::alloc` →
`cudaMalloc` + `cudaMemset`**。`cudaMalloc`/`cudaFree` 是**同步、进程级串行**的 driver 调用；在
**单进程 thread-per-GPU** 的 DDP 模型下，N 个 rank 线程每步几百次 alloc 在**单 CUDA context** 里排队
互相串行（`NOCOMM=1` 完全不通信时 fwd+bwd 仍 136→780ms 膨胀 ~6×，`nvidia-smi` 抽样 8 卡同一时刻
只有 1–2 张在忙、轮流跑）。**这笔 per-op alloc 开销单卡也吃**——训练定形状、每步重复 malloc/free
同样的几百个 buffer，纯属浪费。
T11 的修复：在 `xtrain-cuda`（`GpuBuffer`/`cudaMalloc`/`cudaFree` 所在）加一个 **device caching /
pool allocator**——freed 的显存**进 per-device 的 size-classed free-list 复用，不 `cudaFree`**；
`alloc` 优先从 free-list 取，miss 才 `cudaMalloc`。训练定形状 → 命中率极高，**warm-up 后每步
`cudaMalloc` ≈ 0**，消掉串行 driver 调用风暴。
**硬闸门是正确性**：allocator 必须**透明**——交出的字节、数值与改前**逐位一致**，所有既有 grad-check /
PyTorch 对拍 / overfit / DDP / xserv 闭环**必须仍过**。在此之上拿吞吐收益。
## Module Layout
```
crates/xtrain-cuda/src/
pool.rs ← 新增：global per-device free-list registry + size-class 逻辑
memory.rs ← GpuBuffer::alloc 从 pool 取；Drop 归还 pool（不 cudaFree）
ffi.rs ← 加 cudaGetDevice（Drop 要知道 buffer 属哪个 device pool）
lib.rs ← `mod pool;`
```
`xtrain-tensor` **零改动**：`Storage::zeros` 仍 `GpuBuffer::alloc` + `memset(0)`，签名不变。
pool 完全藏在 `GpuBuffer` 后面，上层无感。
## Key Design Decisions
### 1. Size class（按粒度向上取整 → 跨步可复用）
请求字节数向上取整到一个 **size class**，同形状的 op 输出落进同一 free-list、跨 step 复用：
```rust
const MIN_CLASS: usize = 512; // 小分配的对齐粒度
const POW2_THRESHOLD: usize = 1 << 20; // 1 MiB
fn size_class(len) =
if len <= 1 MiB { ceil(len / 512) * 512 } // 细粒度，浪费 ≤512B
else { len.next_power_of_two() } // 粗粒度，class 数有界
```
小分配按 512B 对齐（浪费极小）；大分配按 2 的幂取整（class 数有界 → free-list 浅，最多浪费 ~2×，
但**显存是复用不是泄漏**，定形状训练里大 buffer 的 class 也就那么几个）。
**关键透明性**：物理分配是 `cap`（取整后），但 `GpuBuffer::len()` 仍返回**请求的 `len`**：
- `memset(0)` 只 zero **逻辑 `len`** 字节（不是 `cap`）；
- 所有 copy（H2D/D2H）bounds 用 `len`，D2H 拷回 host 也只拷 `len` 字节；
- op kernel 只按 shape（= `len`）读写。
→ `cap - len` 的尾部字节**永不被任何人读到**，所以 round-up 对数值**完全透明**。
### 2. Per-device + 线程安全（DDP thread-per-GPU）
DDP 是单进程 thread-per-GPU——pool 必须跨 rank 线程安全，且**不能让不同 device 的线程互相串行**
（否则没解决问题）：
```
global REGISTRY: Mutex<HashMap<device_id, Arc<Mutex<DevicePool>>>>
DevicePool { free: HashMap<size_class, Vec<*mut u8>> }
```
- **两级锁**：registry 锁只在「按 device_id 取出（或首次插入）该 device 的 `Arc<Mutex<DevicePool>>`」
这一瞬持有，立刻 clone Arc 出来、释放 registry 锁，再锁**该 device 自己的** pool。
→ 不同 device 的 rank 线程**各锁各的 pool，真并发**，registry 锁只是极短的查表。
- buffer 在 **alloc 时**记下当前线程的 CUDA device（`cudaGetDevice`，DDP 每 rank 线程开头 set 一次），
存进 `GpuBuffer.device`；**Drop 时**按这个 device 归还，保证 ptr 回到它所属 context 的 pool
（即使 drop 发生在另一个 device 的线程上也对）。
### 3. Drop → 归还（不 cudaFree）
```rust
impl Drop for GpuBuffer {
fn drop(&mut self) { pool::release(self.ptr, self.device, self.cap); }
}
```
free-list **无界**（轻量、不做 eviction）——定形状训练的 working set 有界，每步复用同一批 buffer，
free-list 深度自然收敛，不会无限涨。pool 持有的 ptr 活到**进程退出**，届时 OS 回收整个 device
context，**不是泄漏**。
**双重释放/泄漏边界审查**：`GpuBuffer` 无 `Clone`，独占 ptr；`Storage` 用 `Arc<GpuBuffer>` 共享，
最后一个 Arc 落地时 buffer 恰好 drop 一次 → `release` 一次。`acquire` 从 free-list `pop` 一个 ptr
交给**唯一**一个新 `GpuBuffer`，无别名。故无双重释放、无别名。
### 4. memset：保留（正确性优先），不做 skip-memset uninit
`Storage::zeros` 复用的 buffer 持有**陈旧字节**，故**继续 `memset(0)`**（正确性）。
任务给的 OPTIONAL bonus（给「完全覆盖输出」的 op 加 `uninit`/skip-memset）**本次不做**，诚实理由：
- 真正串行的是 `cudaMalloc`，**已被 pool 消掉**；`cudaMemset` 在 default stream 上 async、开销小。
- 要 skip 必须逐 op 证明输出被**完全覆盖**——`matmul`(beta=0 全写)能跳，但 `embedding_bwd`(scatter-**add**)、
`sumsq_accum`/`sum_rows`(累加器)、`adamw`(读写 m/v) **必须**预 zero。审查面大、收益小、正确性风险高。
- **正确性是硬闸门**，不为一个已非瓶颈的 async memset 冒风险。留作后续（若 profile 显示 memset 成新瓶颈再做）。
## 验证方法（双闸门）
### 闸门一：正确性（透明，零回归）
allocator 不改任何数值。全回归套**必须仍绿**：
- T3 GEMM 对 cuBLAS；T4 各 op finite-diff grad-check（15 个）；
- T5 结构 + overfit(27/27) + PyTorch 对拍（B>1，logits/每参数 grad）；
- T6 AdamW 对 torch + checkpoint 逐位；
- T8 DDP loss 对单卡（~5.7e-7）+ 跨 rank 一致；T10 batched==looped；
- **xserv 闭环**：导出权重对 xtrain 贪心仍逐 token 一致。
### 闸门二：吞吐（收益）
- **单卡 tok/s before/after**（malloc 风暴消失应↑）+ GPU util；
- **DDP 1/2/4/8 卡 scaling before/after**（KI-5 调查的表）；
`ddp_throughput_scaling` 测试扩到 world=8。
**诚实原则**：若单卡提速但多卡仍受限 → 说明串行比 malloc 更深（如单 context 下 kernel launch /
cuBLAS handle 仍串行），如实报告，并说明 **process-per-GPU**（每 rank 独立 context，torchrun 式）
是否是剩余的修复方向（profile 确认，如前两次调查）。
## 顺手项
- **放宽 DDP flaky 断言**：`ddp_correctness` 的 cross-rank `max|p0−p1| == 0.0` → `< 1e-6`。
承重闸门是 loss-match（~5.7e-7）；本机 PCIe-only NCCL all-reduce run-to-run 跨 rank 非逐位可复现，
diff ≤1.2e-7（几 ULP，数值无害）。`== 0.0` 过严 flaky。
## Before → After（dash5, 8× RTX 5090, sm_120）
实测（`train_ddp`, dim384/12L/12h·hd32 ffn1536 core 28.3M, per-rank batch 32, seq 256,
steady-state tok/s；before = parent `d422c68`, after = pooled）：
**单卡（KI-5 假设：per-op malloc 单卡也吃）**
| | tok/s | GPU util |
|---|---|---|
| before | 40226 | 8 卡轮流忙，1–2/8 |
| after | **92638** | — |
→ 单卡 **~2.3×**，loss 轨迹逐位对住（10.9026→4.8453 before/after 一致）。
**DDP 1/2/4/8 卡 scaling（global batch = 32×world）**
| world | before tok/s | before speedup | after tok/s | after speedup |
|---|---|---|---|---|
| 1 | 39801 | 1.00× | 92385 | 1.00× |
| 2 | 47229 | 1.19× | 146821 | 1.59× |
| 4 | 52854 | 1.33× | 269867 | 2.92× |
| 8 | 48996 | 1.23× | **461270** | **4.99×** |
→ 8 卡绝对吞吐 **49K → 461K tok/s = 9.4×**；scaling 从「~1.3× 封顶」恢复到 **~5×@8**。
8 卡运行 `nvidia-smi` 抽样 **8 卡全部 95–99% util**（KI-5 时只有 1–2/8 在忙）——
per-op cudaMalloc 串行确是根因，pool 消掉后 GPU 变 compute-bound 喂满。
**残留**：5×@8 非完美线性（grad all-reduce ~7% + 8 卡 PCIe / launch 余量），但弱扩展的悬崖已消。
KI-5 标 **FIXED**。若 v4 要更高线性度，下一步才是 process-per-GPU（每 rank 独立 context）。

98
docs/11-bf16-mixed-precision.md Normal file
View File

@@ -0,0 +1,98 @@
# Phase T12: bf16 混合精度（fp32 master）— Design Document
> KI-2 的具体落地。v4（dim768, fp32）在单卡 32GB 下 per-rank batch 32（global 256）**OOM**，被迫降到 batch 16 训练。bf16 把激活显存减半（找回 batch-256 甜点区），并在 dim768 这个已 compute-bound 的规模上加速 tensor-core GEMM。附带收益：xserv 推理是 **BF16-only**，bf16 训练让闭环更贴。
## Goal
在**不动 fp32 路径任何数值**的前提下，新增一个 **opt-in 的 bf16 混合精度模式**（标准 AMP，fp32 master weights）：
1. **正确性硬闸门**：fp32 全套回归（T3 GEMM / T4 12 算子 grad-check / T5 结构+overfit+PyTorch 对拍 / T6 AdamW+checkpoint / T8 DDP / T10 batched / xserv 闭环）在**同样紧容差**下保持绿。bf16 是**加法、可选**的，绝不扰动 fp32。
2. **bf16 正确性**：bf16 前向/梯度在**更松的 bf16 容差**（≈2–3 位十进制有效数字 → ~1e-2 相对误差）内对住 fp32 参考；一段**短 bf16 训练收敛对住 fp32**（loss 曲线接近、无 NaN/发散）。
3. **显存+吞吐（payoff）**：dim768 bf16 能跑 per-rank batch 32（解 OOM）；测 dim768 bf16 vs fp32 的显存+tok/s。
## 什么是 bf16、什么是 fp32（标准 AMP split）
| 组件 | 精度 | 理由 |
|---|---|---|
| **master weights** + AdamW state(m/v) + 优化器更新 | **fp32** | 小步长更新需要 fp32 精度，否则被 bf16 的 8-bit 尾数吃掉 |
| **linear GEMM**（q/k/v/o、gate/up/down、lm_head） | **bf16 in/out + fp32 accum** | compute+memory 主体；tensor-core 走 `cublasGemmEx`（`CUDA_R_16BF` in/out，`CUBLAS_COMPUTE_32F` 累加） |
| **激活流**（残差流、attention Q/K/V/probs/out、MLP 中间） | **bf16** | 激活显存减半——这是解 OOM 的关键，不只是 GEMM 提速 |
| **RMSNorm / QK-norm** | **fp32**（bf16→fp32 算 reduction→bf16） | 求和/rsqrt 数值敏感 |
| **softmax（attention）/ RoPE / cross-entropy** | **fp32** | softmax 的 exp/求和、CE 的 log、RoPE 的 sin/cos 都数值敏感 |
| **梯度 → AdamW** | **fp32** | 见下「cast 算子」——grad 在 fp32 master leaf 上累加，AdamW/clip/DDP all-reduce 全程 fp32、**完全不改** |
**无 loss scaling**：bf16 是 8-bit 指数（与 fp32 同动态范围），不像 fp16（5-bit 指数易下溢）。所以梯度不会下溢到 0，**不需要** loss scaling。
## Module Layout（surgical：fp32 路径逐字节不动）
核心思路：**所有 op 按 `self.dtype()` 分派**。fp32 分支跑原 kernel（一字不改）；bf16 分支是新增代码。
### 1. `xtrain-tensor::dtype` — 加 `BF16`
- `DType::BF16`，`size_bytes()=2`，`half::bf16` 实现 `TensorDType`（`half` crate 已是依赖）。
### 2. `xtrain-cuda` — bf16 GEMM + cast kernel
- `ffi.rs`：声明 `cublasGemmEx` / `cublasGemmStridedBatchedEx`（void* 指针、`a_type/b_type/c_type/compute_type`），常量 `CUDA_R_16BF=14`、`CUDA_R_32F=0`、`CUBLAS_COMPUTE_32F=68`（数值同 xserv `gemm.rs`）。
- `cublas.rs`：`gemm_ex(...)` / `gemm_ex_strided_batched(...)`——和 `sgemm` 同样的 row-major⟺col-major 转置代数，只是走 `GemmEx`、in/out=bf16、accum=fp32。
- `csrc/ops/cast.cu` + ffi：`launch_cast_f32_to_bf16` / `launch_cast_bf16_to_f32`（逐元素 `__float2bfloat16` / `__bfloat162float`）。
### 3. `xtrain-tensor::tensor` — dtype-polymorphic ops
- `to_dtype(target)`：f32↔bf16 cast（CUDA），同 dtype 直接 clone。
- `matmul` / `matmul_backward` / `attention` / `attention_backward`：按 dtype 分派——fp32 走原 `sgemm`（**不动**），bf16 走 `gemm_ex`，输出同 dtype。
- 逐元素 op（`add`/`mul`/`silu`/`scale`…）+ `embedding`：允许 bf16 输入。逐元素 kernel 对 bf16 走「load→fp32→算→store bf16」（新增 bf16 kernel）或对 norm/softmax/CE 在 wrapper 里 upcast→fp32 kernel→downcast。**fp32 调用走原 f32 kernel 不变。**
### 4. `xtrain-autodiff::ops` — `cast` 算子 + bf16 透传
- **`cast(x, target_dtype)`**：前向 `x.to_dtype(target)`；**反向把 grad cast 回 `x` 的 dtype**。这是 AMP 的关键钩子：
- fp32 master weight leaf → `cast(w, BF16)` 喂给 matmul；matmul 的 bf16 grad 经 cast 反向**升回 fp32**，累加在 fp32 leaf 上。
- ⟹ `.grad()` 是 **fp32**，AdamW / `clip_grad_norm_gpu` / DDP `all_reduce_average_grads` **一行不改**，全程 fp32 master。
- 其它 op 的 backward 自然按张量 dtype 流转（softmax/rms_norm wrapper 内部 upcast→fp32→downcast，对外是 bf16）。
### 5. `xtrain-model::TinyTransformer` — `compute_dtype` 开关
- `new_amp(cfg, device, dtype, init)` 或 `forward_batched` 接 `compute_dtype: DType`（默认 `F32` = 原路径，逐字节同）。
- bf16 模式：embedding 输出 `cast→bf16` 进入残差流；每个 weight matmul 前 `cast(w, BF16)`；norm/softmax/rope 对 bf16 激活自动 fp32 内算；最后 logits `cast→fp32` 给 cross_entropy。**fp32 模式 `compute_dtype==F32` 时跳过所有 cast，graph 与 T10/T11 完全一致。**
### 6. `xtrain-train` / `xtrain-distributed` — `--bf16` flag
- `TrainConfig`/`DdpConfig` 加 `compute_dtype: DType`；`train.rs`/`train_ddp.rs` 加 `--bf16` flag。
- AdamW / clip / checkpoint / DDP all-reduce **不改**（master 永远 fp32，grad 永远 fp32）。
## Key Design Decisions
- **cast 算子承载 fp32 master ↔ bf16 compute 的桥**：不需要在优化器里维护一份独立的 bf16 weight 副本——fp32 leaf 即 master，前向临时 cast 出 bf16，反向 grad 自动升回 fp32。最小改动、零优化器侵入。
- **按 dtype 分派而非新类型**：fp32 路径走的还是同一个函数的 `F32` 分支 → 原 kernel、原 cuBLAS 调用、原 launch 顺序，数值逐字节不变（满足硬闸门）。
- **norm/softmax/CE 不写 bf16 reduction kernel**：wrapper 里 `to_dtype(F32)` → 复用现有 fp32 kernel → `to_dtype(BF16)`。多两个 cast launch，但**复用已验证的 fp32 数值**，且这些不是显存/算力大头。
- **无 loss scaling**：bf16 8-bit 指数，省掉 fp16 那套 scale/unscale/inf-check。
## 验证方法（双闸门）
### 闸门 ① fp32 不回归（hard gate）
全套现有测试在原紧容差下保持绿（bf16 是 opt-in，默认 dtype=F32）：
- `cargo test` 全 crate：grad-check（rel ≤2e-2）、structural、GEMM 对 cuBLAS（~1e-7）、batched==looped、overfit 27/27、AdamW GPU bit-exact + host 对 torch、checkpoint 逐位、DDP loss 对单卡 <1e-6、**PyTorch 对拍**（loss/logits/grad）。
- **xserv 闭环**：v4 ckpt（fp32 训）重导 safetensors md5 一致 + xserv 贪心逐 token 对住。
### 闸门 ② bf16 正确性 + 收敛
- **bf16 looser-tol 数值**：同一组随机权重/输入，bf16 forward logits 与 bf16 grad 对 fp32 参考在 **rel ~1e-2**（bf16 2–3 位有效数字）内。
- **短训练收敛**：dim768（或缩小代理）bf16 跑数百步，loss 曲线对住 fp32，无 NaN/发散，end loss 接近。
### 闸门 ③ 显存 + 吞吐（payoff）
- **dim768 bf16 能跑 per-rank batch 32**（v4 OOM 的触发点）。
- 测 dim768 **bf16 vs fp32** 的峰值显存 + steady-state tok/s（预期：显存↓、tok/s↑）。
## 实测结果（dash5, 1× RTX 5090 32GB, sm_120）
**闸门 ① fp32 不回归**：全套测试在原紧容差绿（autograd 15 / structural 5 / GEMM 5 / batched==looped / overfit 27/27 / AdamW GPU bit-exact + host 对 torch / checkpoint 逐位 / DDP 2）。**xserv 闭环**：v3 ckpt 用 T12 代码重导 `model.safetensors` 与 registry **md5 逐位一致**（`b04fc9f9a0c9af04c47d9ca649aea12e`）——export/fp32 数值零漂移。
**闸门 ② bf16 正确性 + 收敛**：
- **looser-tol（tests/bf16.rs）**：同 fp32 master 跑 fp32 vs bf16——loss rel `1.2e-4`、logits mean rel `2.0e-3` / p99 `6.8e-3`、grad worst scaled-mean `1.0e-2`，无 NaN，grad 仍 fp32（master 未动）。
- **收敛**：dim768 短训 150 步，bf16-b16 loss 轨迹对住 fp32-b16（step50 `4.40` vs `4.40`、step149 `3.984` vs `3.988`），单调降、无发散。
**闸门 ③ 显存 + 吞吐（dim768/18L/24h×32 ffn2048 seq256, steady-state）**：
| config | per-rank batch | 峰值显存 | tok/s | fits 32GB? |
|---|---|---|---|---|
| fp32 | 16 (v4 fallback) | 27.2 GB | 31.5K | ✅ |
| **bf16** | 16 | **19.3 GB（−29%）** | **35.5K（+13%）** | ✅ |
| fp32 | 32 | — | — | ❌ **OOM** |
| **bf16** | **32（甜点区）** | **31.1 GB** | **40.8K** | ✅ **解 OOM** |
→ 同 batch：bf16 显存 −29% / tok/s +13%；**bf16 解 fp32-batch32 OOM**，batch32 达 40.8K tok/s（+29% vs fp32-b16）。KI-2 标 **FIXED**。
</content>
</invoke>

90
docs/12-activation-recompute.md Normal file
View File

@@ -0,0 +1,90 @@
# Phase T13: 激活重计算（gradient checkpointing）— Design Document
> KI-3 的具体落地。autograd tape 为反向保存了所有中间激活；dim768/bf16 在单卡 32GB 能跑 batch32（T12 解 OOM），但**容量轴放大到 dim1024 会再次 OOM**——激活显存随 dim 线性增长。激活重计算用「多一次前向」换显存：段内激活不在前向保存，反向时**重算该段前向**重建局部 tape 再回传。峰值激活从「所有 block 同时在显存」降到「~一个 block + 每 block 的输入」→ dim1024 batch32 装得下。
## Goal
在**不动非重计算路径任何数值**的前提下，新增一个 **opt-in 的 per-block 激活重计算**：
1. **正确性硬闸门（exact）**：重计算是**数学精确**的——同一段前向、同一输入、同一（未变的）参数值、确定性 kernel ⟹ 重算出的激活与原激活逐位相同，回传的梯度与非重计算版一致。直接的 **on-vs-off 梯度对拍**（紧容差）+ 全套回归（T4 grad-check、T5 overfit+PyTorch 对拍、T6 AdamW、T8 DDP loss-match+跨 rank、xserv 闭环）开 `--recompute` 全绿。**绝不提交一个改变梯度的重计算。**
2. **显存（payoff）**：测 dim768 batch32 峰值显存 on vs off（应降）；确认 **dim1024 batch32 现在装得下**（不开重计算时 OOM，开了 fit）。
3. **吞吐**：测 tok/s on vs off（多一次前向，预计慢 ~20–35%）——报告 compute/memory 权衡。
## 什么是激活重计算
反向传播需要前向的中间激活（如 SwiGLU 的 `gate`、attention 的 `probs`）来算梯度。define-by-run 的 tape 默认把它们全部留在显存，直到对应 op 的 backward 跑完。模型越深、激活越多，峰值显存越高。
**梯度检查点**把模型切成若干**段（segment）**。前向时，段内 op **不记到 tape**（detached / no-grad），只保留**段的输入**（参数作为 leaf 本就常驻，不算激活）。反向时，当段的 output-grad 到达，**从保存的输入重跑该段前向**（输入作为 require-grad 的 leaf），用上游 grad seed 重算出的 output，在局部 tape 上回传，得到输入梯度（并累加参数梯度），然后释放局部 tape。
代价：每段多一次前向（约 +1/3 的总 FLOPs，因为反向本就 ~2× 前向）。收益：峰值激活从「所有段」降到「~一段 + 每段输入」。
**切粒度 = 每个 transformer block**：一个 block（attention 子块 + MLP 子块 + 两个残差）是天然的段边界，输入/输出都是 `[batch*seq, dim]` 的残差流张量，接口最干净。
## Module Layout（surgical：非重计算路径逐字节不动）
### 1. `xtrain-autodiff::checkpoint` — `checkpoint` 高阶原语
新增 `checkpoint(segment_fn, input, params) -> Var`，类比 `torch.utils.checkpoint`：
- `segment_fn: Fn(&Var, &[Var]) -> Var`——从单个输入 `x` 和参数 slice `p` 构建段前向、返回段输出。必须**确定性**、只依赖 `x` 和 `p`（这是重算精确的前提）。
- **前向（不 tape 内部）**：把 `input`/`params` detach 成新 leaf（`Var::leaf(v.value())`），跑 `segment_fn` 得到 `out_local`，**只取 `out_local.value()`**。局部 `Var` 出作用域即 drop → 段内激活立即释放。checkpoint 节点的 parents = `[input, ..params]`（参数梯度落进优化器读的同一批 leaf）。
- **反向（重算）**：闭包捕获 `Rc<segment_fn>`。给定 `dout`，从 `parents` 当前值重建 detached leaf，重跑 `segment_fn` 重建局部 tape，调 `out_local.backward_seeded(dout)`，再把 `x_det.grad()` push 给 `parents[0]`、各 `param_det.grad()` push 给对应参数 parent。闭包结束 → 局部 tape drop → 重算激活释放。
### 2. `xtrain-autodiff::tape` — `backward_seeded`
引擎原 `backward()` 只能从标量 root 出发（seed `ones_like` + 断言 numel==1）。段输出一般**非标量**，故新增 `Var::backward_seeded(seed)`：同样的 topo + 反向遍历，但用显式上游 grad seed（不断言标量）。`backward()` 退化为「seed ones」的薄包装——标量 loss 路径逐字节不变。
### 3. `xtrain-model::TinyTransformer` — per-block 包裹 + `recompute` 开关
- 把 block 前向体抽成**自由函数** `block_forward(cfg, compute_dtype, batch, seq, input, params)`（不借 `&self`，才能在反向闭包里重跑）；同时把 `linear / norm_gamma / attention / swiglu_mlp` 改成参数化 `(cfg, compute_dtype)` 的自由函数。`Block::block_params()` 给出 11 个 leaf 的固定序（与 `params()` 每 block 段一致）。
- `forward_batched` 的 block 循环：`recompute` 开 → `checkpoint(seg, &h, &b.block_params())`；关 → 直接 `block_forward(...)`（**与之前完全同图**）。
- `with_recompute(bool)` builder（opt-in，默认关 = 原 tape，数值逐字节同）。
### 4. `xtrain-train` / `xtrain-distributed` — `--recompute` flag
- `bin/train` / `bin/train_ddp` 加 `--recompute`，调 `model.with_recompute(true)`。AdamW / clip / checkpoint / DDP all-reduce **不改**（梯度语义与非重计算一致）。
## Key Design Decisions
- **切粒度 = 每个 block**：接口最干净（输入输出都是残差流 `[B*S, dim]`），且峰值激活降到 ~1 个 block。比「整模型一段」省得多，比「逐 op」简单得多。
- **参数作为 checkpoint 节点的 parents**：参数 leaf 跨前向/反向不变（只 grad 槽变），重算用**当前**参数值即原值。把它们列为 parents，重算恢复的参数梯度直接 push 到优化器读的同一批 leaf → **DDP/AdamW 零改动**。
- **detached leaf 隔离局部 tape**：前向/反向都把输入和参数 detach 成新 leaf，使 `segment_fn` 构建的图与外层 tape 断开。前向丢弃局部图（释放激活）；反向局部图回传完即 drop（释放重算激活）。
- **`backward_seeded` 而非改 `backward`**：段输出非标量，需要用上游 output-grad 作 seed 回传局部 tape。新增方法、原标量 `backward()` 不动。
- **重算精确 → 梯度逐位一致（硬闸门）**：同一 `segment_fn`、同一输入值、同一参数值、确定性前向 kernel ⟹ 重算 output 与原 output 相同，局部反向就是该段的普通解析反向。故输入/参数梯度与非重计算版一致——这是绝不能违反的闸门。
## 与 bf16 / DDP / batched 的组合
- **bf16（T12）**：`segment_fn` 就是不变的 block 前向，重算跑**同一条 bf16 路径**；`cast` 算子的 grad 升精度桥（bf16→fp32）照常。重计算节点的参数 parents 是 fp32 master leaf，恢复的是 fp32 梯度。on-vs-off 对拍同时跑 fp32 和 bf16 两路。
- **DDP（T8）**：每个 rank 独立 checkpoint 自己的前向/反向；恢复的参数梯度落进各 rank 的 `.grad()` 槽，再被 all-reduce 取均值——分布式路径不感知重计算。
- **batched（T10）**：段输入/输出透明带 `[batch*seq, …]` batch 维；`checkpoint` 与形状无关。
## 验证方法
### 1. 正确性（exact，硬闸门）
- **on-vs-off 梯度对拍**（`crates/xtrain-model/tests/recompute.rs`）：同 init 建两个模型（recompute on/off），跑同一 batched loss+backward，断言**前向 logits、loss、每个参数梯度**在紧容差内一致——参数化跑 **fp32 和 bf16** 两路。fp32 期望近逐位（容差 1e-4），bf16 仅放松到 bf16 舍入级（非重计算误差）。
- **全套回归开 `--recompute`**：T4 15 算子 grad-check、T5 overfit 27/27 + PyTorch 对拍、T6 AdamW、T8 DDP loss-match + 跨 rank、**xserv 闭环 md5**——全绿。
### 2. 显存（payoff）
- dash5 1× RTX 5090 32GB，dim768/18L batch32 seq256，bf16：测峰值显存 recompute on vs off（应降）。
- **dim1024 batch32**：先验证不开重计算 **OOM**，再验证开了 **fit**——capture 实际 `nvidia-smi` 峰值。
### 3. 吞吐
- 同 config 测 steady-state tok/s recompute on vs off，报告慢多少（预计 ~20–35%，多一次前向）。
## 实测结果（dash5 1× RTX 5090 32GB, bf16, batch32 seq256, steady-state）
**正确性（exact，硬闸门）**：on-vs-off 梯度对拍 —— **fp32 与 bf16 双路都逐位一致**：loss rel `0.00e0`、logits max rel `0.00e0`、**每个参数梯度 max rel `0.00e0`**（不是「在容差内」，是逐位相同——证实重算确实精确）。全套回归开/关重计算全绿：T4 15 算子 grad-check、5 结构、batched、bf16、overfit、AdamW（GPU+host）、GEMM、checkpoint roundtrip、**T8 DDP loss 对单卡 5.67e-7 + 跨 rank 0.0**；DDP+recompute 2 卡短训 loss 单调降（11.079→6.010）。
**显存 + 吞吐**（dim768 = 18L/24h×32/ffn2048 core 127M；dim1024 = 18L/32h×32/ffn2730 core 226M）：
| config | per-rank batch | 峰值显存 | tok/s | fits 32GB? |
|---|---|---|---|---|
| dim768 recompute **off** | 32 | 31144 MiB | 39.7K | ✅ |
| **dim768 recompute on** | 32 | **14562 MiB（−53%）** | **31.5K（−20%）** | ✅ |
| **dim1024** recompute **off** | 32 | 32100 → **OOM** | — | ❌ **OOM** |
| **dim1024 recompute on** | 32 | **16596 MiB** | 23.1K | ✅ **解 OOM** |
→ dim768：重计算把峰值显存 **31.1→14.6GB（−53%，~砍半激活）**，代价 tok/s **−20%**（多一次前向，落在预测 20–35% 区间）。dim1024 batch32：**不开 OOM（撞 32100/32607MiB 上限）→ 开了 16.6GB 稳稳装下**，~23K tok/s 训练正常收敛 —— **KI-3 的目标达成，dim1024 解锁**。

183
docs/13-flash-attention.md Normal file
View File

@@ -0,0 +1,183 @@
# Phase T14: 融合 Flash-Attention Kernel — Design Document
## Goal
T10 把 attention 批量化了，但它的 SDPA 走的是 **「物化 N×N scores」** 的组合路径：
`cublasSgemmStridedBatched`（Q·Kᵀ）→ 一个 causal-softmax kernel（写出整张 probs）→
`cublasSgemmStridedBatched`（P·V），**3 次 launch + 一张 `[bh, S, S]` 的 scores/probs 张量**
常驻显存（反向还要缓存这张 probs）。S 一大，这张 N×N 就成了激活显存与带宽的主导项。
T14 的目标：手写一个**单 kernel 的 fused flash-attention**——streaming / online softmax、**tiled
over KV**、**绝不物化 N×N**。前向一发 kernel 直接吐出 `out[bh,S,hd]`（外加 `O(N)` 的 logsumexp）；
反向一发 kernel（flash 式：重算 scores + dQ/dK/dV，同样不物化 N×N）。接进 model + autograd 作
**opt-in `--flash`**，默认保留 T10 的 composed 路径以便 A/B。
**硬闸门是诚实正确性**：新 kernel 的 dQ/dK/dV finite-diff grad-check 过；fwd/bwd 对现有 composed-SDPA
路径数值贴合（进 bf16 容差）；PyTorch SDPA 对拍 B>1；峰值显存↓（不物化 scores）+ tok/s before/after 实测；
全回归套（含 xserv 闭环 md5）开/关 flag 都绿——默认（flag off）图不变 → 不回归。
## 什么是 flash-attention
标准 attention 是 `O = softmax(causal(Q·Kᵀ/√d)) · V`，朴素实现把 `S[i,j] = Qᵢ·Kⱼ/√d` 整张
`[S,S]` 算出来、softmax、再乘 V——显存 `O(S²)`、HBM 读写 `O(S²)`。
**flash-attention** 的洞察：softmax 可以 **online（streaming）** 地算。把 K/V 切成若干 **tile**，对一个
query 行 `i`，依次扫过 KV tile，用 **running max `m` + running sum `l`** 维护 softmax 的归一化，并把
部分加权的 `V` 累加进一个 `[hd]` 的 accumulator `acc`，每来一个新 tile 就用「新旧 max 的差」对旧 `acc`/`l`
做 rescale。扫完所有 tile，`out = acc / l`。**整张 `[S,S]` 从不落地**——只有 `[hd]` 的 acc 和两个标量
在寄存器/共享内存里流动。峰值激活从 `O(S²)` 降到 `O(S·hd)`（就是 O 本身）。
online softmax 的核心递推（block `j` 的部分 logits 行 `s_j`，旧状态 `m, l, acc`）：
```text
m_new = max(m, max_k s_j[k])
p = exp(s_j - m_new) # 本 tile 的未归一化权重
l = l * exp(m - m_new) + sum(p) # 旧 sum 先 rescale，再加本 tile
acc = acc * exp(m - m_new) + p · V_tile # 旧 acc 同样 rescale，再加本 tile 贡献
m = m_new
# 扫完所有 tile：
out = acc / l
L = m + log(l) # logsumefp，O(N) 存给反向
```
**因果 mask 内联**：query 全局位置 = `i % S`（沿用 T10 的 per-seq 复位约定），KV 位置 `j` 满足
`j > i%S` 的列直接当 `-inf`（`p=0`）。tile 整块在对角线之上可**直接 skip**（causal 的天然稀疏，省一半算力）。
**反向（flash 式，[Dao 2022] 的标准做法）**：不缓存 probs，从 Q/K/V + 前向存的 `L[bh,S]` **重算** scores。
关键预计算 `D[i] = Σ_d dOᵢ[d]·Oᵢ[d]`（每 query 一个标量，`O(N)`），则对每个 `(i,j)`：
```text
s_ij = Qᵢ·Kⱼ * scale # 重算 logit
p_ij = exp(s_ij - L[i]) # 重算 softmax 权重（L 是前向存的 logsumexp）
dp_ij = dOᵢ · Vⱼ # 对 P 的梯度
ds_ij = p_ij * (dp_ij - D[i]) * scale # softmax 雅可比，化简掉了显式 N×N
dQᵢ += ds_ij * Kⱼ ; dKⱼ += ds_ij * Qᵢ ; dVⱼ += p_ij * dOᵢ
```
`ds = P ∘ (dP - D)` 是 softmax 反向用 `Σⱼ Pⱼ·dPⱼ = D`（因为 `D[i]=Σ dOᵢ·Oᵢ = Σⱼ Pᵢⱼ dPᵢⱼ`）化简的结果，
**不需要 N×N 的 softmax 雅可比矩阵**。同样 tiled、同样不物化 N×N。
## Module Layout（surgical：composed 路径逐字节不动，flash 全程新增并行路径）
```
csrc/ops/flash_attention.cu # 新：fwd kernel（online softmax，tiled KV）+ bwd kernel（重算 + dQ/dK/dV）
crates/xtrain-cuda/
├── src/ffi.rs # +launch_flash_attention_fwd_f32 / _bwd_f32 声明
└── build.rs # +flash_attention.cu
crates/xtrain-tensor/src/tensor.rs # +Tensor::flash_attention / flash_attention_backward（fwd 存 logsumexp L；bf16 upcast→f32 kernel→downcast）
crates/xtrain-autodiff/
├── src/ops.rs # +ops::flash_attention 节点（前向调 fwd，缓存 L，反向调 bwd）
└── tests/autograd.rs # +flash_attention(batched) dQ/dK/dV grad-check
crates/xtrain-model/
├── src/model.rs # attention() 按 use_flash 选 ops::attention | ops::flash_attention；+with_flash(bool) builder；flash 标志透传 block_forward（recompute 段内也走 flash）
└── tests/flash.rs # 新：flash == composed（fwd logits + 每参数梯度），参数化 fp32/bf16
crates/xtrain-train/src/bin/train.rs # +--flash flag → model.with_flash(true)
crates/xtrain-distributed/src/bin/train_ddp.rs # +--flash flag（DDP 路径）
crates/xtrain-model/tests/parity_dump.rs # PyTorch B>1 对拍跑两遍：composed 与 flash（共用 PyTorch oracle）
```
## Key Design Decisions
### ① 一个 block 负责一行 query（先做对，再谈快）
最直接、最易验证正确的并行划分：**`grid = bh * S`，每个 block 算一整行 query 的 `out[bh, i, :]`**。
block 内 `hd` 个线程（hd ≤ 128，正好一个 warp 多一点），共享 `m/l` 标量 + `acc[hd]`。block 顺序扫
KV tile（tile 宽 `BK`，沿 `j` 维），每个 tile：线程并行算 `BK` 个 logit（点积 over hd 用 block-reduce）、
求 tile max、online-rescale `m/l/acc`、累加 `p·V`。扫完写 `out = acc/l` 与 `L[i] = m + log(l)`。
**为什么先这样而不是 FA2 的 query-tile 划分**：本项目的硬闸门是**正确性 + 不物化 N×N + 显存↓**，不是
打榜峰值 FLOPs。一行一 block 的版本：(a) online softmax 与 N×N skip 已经完全落地（显存与带宽收益拿到），
(b) 代码直白、逐 query 行可对拍，正确性风险最低。它**不会**比 cuBLAS 两发 GEMM 更快（cuBLAS tensor-core
吃满），所以 tok/s 上 flash 在我们这种 `hd=32` 小头维下大概率**持平或略慢**——这正是 flash 的已知权衡
（flash 的胜场是**显存**，不是小模型的 wall-clock）。把这点诚实写进 perf 表，不掩饰。
### ② 前向只存 `L[bh,S]`（logsumefp），不存 probs
composed 路径反向要缓存整张 `probs[bh,S,S]`（`O(N²)`）。flash 反向**只需要前向的 logsumexp
`L[i]=m_i+log(l_i)`**（每 query 一个 fp32，`O(N)`）即可重算任意 `p_ij = exp(Qᵢ·Kⱼ·scale - L[i])`。
所以 fwd kernel 顺手把 `L` 写出来，autograd 节点缓存它（外加 Q/K/V/O parents 本就在）。**这就是显存闸门的来源**：
attention 的反向缓存从 `[bh,S,S]` 砍到 `[bh,S]`。
### ③ 反向用 `D[i]=Σ dOᵢ·Oᵢ` 化简 softmax 雅可比
softmax 反向通项 `ds_ij = p_ij·(dp_ij - Σ_k p_ik·dp_ik)`。注意 `Σ_k p_ik·dp_ik = Σ_k p_ik (dOᵢ·V_k)
= dOᵢ·(Σ_k p_ik V_k) = dOᵢ·Oᵢ = D[i]`。所以一趟先算 `D[bh,S]`（每行 `dO·O` 的点积，`O(N)`），反向
扫 KV tile 时直接 `ds = p·(dp - D)·scale`，**不需要再算或物化整行的 `Σ p·dp`**。
dQ/dK/dV 三者：dQ 由「该 query 行」累加（block 私有，无竞争）；dK/dV 跨 query 行累加同一个 `(j)`
→ 用 `atomicAdd` 到全局 dK/dV（fp32 原子加，确定 race-free）。
### ④ bf16：kernel 内 fp32，边界 cast（与 composed 路径一致的数值策略）
T10/T12 的 composed attention 对 bf16 也是 **softmax 用 fp32**（scores 升 f32 → kernel → probs 降回 bf16）。
flash 沿用同策略，最省心且数值最稳：bf16 模式下 `flash_attention` 把 Q/K/V `to_dtype(F32)` 喂给 fp32 kernel，
`out` 再 `to_dtype(BF16)`；反向同理。kernel 本身只有一份 fp32 实现。这样 flash 的 bf16 数值与 composed 的
bf16 数值是**同一套 fp32 softmax 算的**，只差 GEMM rounding（cuBLAS tensor-core vs kernel 内 fp32 FMA）→ 落在
既有 bf16 容差内。`L` 始终 fp32。
> 备选（不采纳）：bf16 全程 in-kernel half。收益是少两次 cast，但 (a) 引入与 composed 不同的 softmax 累加路径，
> 威胁 on-vs-off 贴合闸门；(b) 本规模 attention 非瓶颈。escape hatch：先 fp32-core 把正确性钉死，纯 half flash 留 follow-up。
### ⑤ opt-in 透传：`use_flash` 是运行时旗标，不是架构
`use_flash` 不进 `Config`（它不改模型尺寸、不改导出、不该污染 `num_params`），而是 `TinyTransformer` 的一个
`bool` 字段 + `with_flash(bool)` builder（对齐 `with_recompute` / `with_compute_dtype`）。`block_forward` 已经
是 `(cfg, cdt, …)` 的自由函数（T13 为 recompute 抽的），给它加一个 `flash: bool` 形参，model 的 `attention()`
据此选 `ops::attention`（composed）或 `ops::flash_attention`。recompute 闭包捕获 `flash`（`Copy`）→ **重算段内也走
flash**，flash×recompute 组合天然成立。默认 `false` = composed 路径**逐字节不变**（硬闸门：默认图不变 → 不回归）。
## 验证方法
**硬闸门全绿（dash5 实跑 capture）：**
### 1. 正确性
- **新 kernel dQ/dK/dV finite-diff grad-check**（`xtrain-autodiff/tests/autograd.rs::flash_attention_batched_bwd`）：
与既有 `attention_batched_bwd` 同构（`L = sum(W∘out)`，中心差分），断 dQ/dK/dV 在 `cfg_nonlinear`/`cfg_linear` 容差内。
- **flash == composed**（`xtrain-model/tests/flash.rs`）：同 init 两个模型（flash on/off），同一 batched
loss + backward，断**前向 logits / loss / 每参数梯度**在紧容差内一致；参数化 fp32（近逐位）与 bf16（bf16 舍入级）。
- **PyTorch SDPA 对拍 B>1**（`parity_dump.rs` + `parity.py`）：等价 PyTorch 模型（per-seq RoPE、per-seq causal、
QK-norm、SwiGLU）对拍 forward logits + 全部参数梯度——**composed 与 flash 两条都跑**，共用同一 PyTorch oracle。
- **全回归套开/关 `--flash`**：autograd 15、structural、batched==looped、bf16、recompute（逐位）、overfit 27/27、
AdamW（GPU bit-exact + host 对 torch）、DDP loss-match + 跨 rank、**xserv 闭环（导出 safetensors → md5 对 registry →
xserv 贪心逐 token 一致）**。flag off 默认图不变 → composed 数值不回归。
### 2. 显存（payoff）—— 不物化 N×N 的直接收益
dash5 1× RTX 5090，同 config，nvidia-smi 峰值，flash off vs on：attention 反向缓存 `[bh,S,S]→[bh,S]`，
峰值显存应↓（尤其 seq 大时）。capture 实际数字进表。
### 3. 吞吐
同 config steady-state tok/s flash off vs on。预期：本规模 `hd=32` 下 flash kernel **持平或略慢于** cuBLAS 双
GEMM（小头维喂不满 tensor-core 是 flash 的已知权衡，胜场在显存）——诚实报告，不为绿而调。
## 实测结果（dash5 1× RTX 5090）
**正确性（硬闸门全绿）：**
| 闸门 | 结果 |
|---|---|
| ① 新 kernel dQ/dK/dV finite-diff grad-check | **过** — dQ 9.3e-3 / dK 1.7e-2 / dV 5.6e-4（单 tile 干净区；多 tile 由②兜） |
| flash fwd 对 composed | max rel **6.7e-5** |
| flash bwd 对（已 grad-check 的）composed bwd | dQ **1.7e-5** / dK 1.2e-5 / dV 4.3e-5 |
| ② flash==composed（model 级，logits/loss/每参数梯度） | fp32: loss rel **0.0**、logits 1.7e-4、grad 4.4e-5；bf16: loss 1.5e-4、logits mean 1.6e-3/p99 5.9e-3、grad scaled-mean 1.2e-2 |
| ③ PyTorch SDPA 对拍 B>1（flash 路径，共用 composed oracle） | loss relerr **4.98e-8**、logits **7.92e-6**、25 参数 grad 全进 rtol 0.02 |
| ⑤ 回归套（flag off 默认 + flash 路径都测）：autograd 18 / structural 5 / batched / bf16 / **flash 3** / overfit 27/27 / recompute 2 / AdamW(GPU+host) / GEMM / DDP 2 / checkpoint-roundtrip | **全绿** |
| ⑤ xserv 闭环 md5（v3 ckpt 用 T14 代码重导 safetensors） | **逐位一致** `b04fc9f9a0c9af04c47d9ca649aea12e`（与 registry 同）→ 默认 export 零漂移 |
| ⑤ xserv 闭环（flash 训练 → 导出 → xserv 服务贪心） | flash-训出 coherent TinyStories；xserv(BF16) 对 xtrain(F32) 贪心：3 prompt 中 "One day" 逐 token 一致，其余在 ~0.5% BF16 漂移处晚分叉（与 v1/v2/v3 同款） |
> **finite-diff 的诚实记录**：长 softmax（seq>tile）会产生大量近零梯度元素，中心差分在那些元素上不可靠（出现伪 0.0 / 符号翻转——不是 backward bug）。故 ① 的 finite-diff 跑**单 tile 干净区**（seq=5，对齐既有 composed grad-check 的良态区），**多 tile 的 streaming/online 路径**用「flash bwd 对已 grad-check 的 composed bwd」（seq=40，dQ 1.7e-5）兜——比 finite-diff 更利。dQ/dK 用 eps=2e-3 压低 f32 舍入项（~4e-4 小梯度上舍入项压过截断项）。**没有为凑绿放宽容差**。
**④ 显存 + 吞吐（payoff vs tradeoff，dim768=8L/12h×64/ffn3072, bf16, steady-state）：**
| config | path | 峰值显存 | tok/s |
|---|---|---|---|
| batch8 seq1024 | composed (off) | 24670 MiB | **58.6K** |
| batch8 seq1024 | **flash (on)** | **20736 MiB（−16%）** | 25.0K（−57%, ~2.3× 慢） |
| batch2 seq2048 | composed (off) | 17264 MiB | 36.7K |
| batch2 seq2048 | **flash (on)** | **13246 MiB（−23%）** | 13.2K（−64%） |
→ **显存按预期降**（不物化 `[bh,S,S]`），且**收益随 seq 增长**（seq1024 −16% → seq2048 −23%，O(S²) 砍掉）。
**tok/s 如设计 ① 预测的「持平或略慢」实为 ~2.3–2.8× 慢**：hd=64 的小头维下，手写「一行一 block + 串行扫 KV」kernel 喂不满 SM，干不过 cuBLAS tensor-core 的两发批量 GEMM——这正是 flash 的已知权衡（**胜场在显存，不是小模型 wall-clock**），诚实报告不掩饰。两个落地的优化（softmax 权重缓存进 shared 省 hd× 的 expf；dK/dV 原子加摊到全 block 而非串行在列 owner 内）把 backward 从 6.8× 慢拉到 2.3× 慢——主瓶颈是 backward 的跨行原子累加（FA2 用 K-block 拥有 dK/dV 的独立 pass 解，本版未做，留 follow-up）。
> **escape hatch（follow-up，未做，记给后续）**：① FA2 式 query-tile 划分（一 block 多 query 行，K/V 进 shared 复用）提 SM 占用；② backward 的 dK/dV 改 K-block-owned 独立 pass 消跨行原子；③ 纯 bf16 in-kernel（省两次 cast）。本规模 attention 非训练瓶颈、且会动数值贴合闸门，按 escape hatch 推迟——T14 先把**正确性 + 不物化 N×N + 显存↓**钉死。

Some files were not shown because too many files have changed in this diff Show More