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

Compare commits

base: gahow:92acf9f4137b442da0ed7bb01c5ab081dc730b83
Branches Tags
gahow:main
...
compare: gahow:main
Branches Tags
gahow:main

153 Commits
92acf9f413 ... 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
Gahow Wang
224f750ee4 autograd: tape engine + grad accumulation 
Var = Rc<RefCell<VarNode>> on a define-by-run tape: value + optional grad +
parents + backward closure. backward() seeds a scalar loss, walks reverse
topo order, and pushes grads to parents. push_grad always SUMs into the grad
slot — the fan-out accumulation path T3 lacked. Per-crate build.rs emits the
no_cuda cfg (does not propagate); engine gated, grad_check stays host-only.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:44:17 +08:00
Gahow Wang
5aef3742d6 ops: transformer op fwd/bwd CUDA kernels + Tensor wrappers 
add/mul/add_bias(+sum_rows)/rms_norm/silu/rope/softmax/cross_entropy,
each with its analytic backward, in csrc/ops/nn.cu (inlined warp/block
reductions). FFI declarations + nn.cu in build.rs (no_cuda gated). Tensor
gains the matching thin wrappers; DType grows I32 for cross-entropy targets.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:44:09 +08:00
Gahow Wang
88fbe0a85d gemm: realistic f32 tolerances in GEMM acceptance tests 
Forward: compare via matrix relative error (max abs error / max|ref|)
instead of a per-element ratio, so near-zero outputs where two correct
f32 GEMMs differ only in rounding order don't inflate the metric.
Backward: L = sum(W∘C) is bilinear, so central differences are
truncation-free — use eps=1e-2 (sharper f32 resolution of the
difference) and atol=1e-3 to floor near-zero-gradient subtraction noise.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:28:57 +08:00
Gahow Wang
dde2fde297 docs: Phase T3 — GEMM fwd/bwd + finite-diff 
Design doc covering the tiled forward, the dA/dB math + how transpose is
handled (materialize + reuse forward), the cuBLAS row-major reference, and
the finite-diff harness design + how T4 reuses it per-op.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:27:03 +08:00
Gahow Wang
1384044f27 gemm: GPU acceptance tests vs cuBLAS + finite-diff 
Forward: hand-written tiled GEMM vs cuBLAS sgemm on random matrices
(square / non-tile-aligned rect / 256³), max relative error < 1e-3, using
the row-major⟺col-major identity to drive cuBLAS without explicit
transposes. Backward: scalar loss L = sum(W∘C) (so dC = W), dA/dB from
matmul_backward checked against the finite-diff harness. Gated behind
not(no_cuda).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:26:58 +08:00
Gahow Wang
08c88bf360 gemm: tiled F32 forward + transpose + backward (dA/dB) 
Hand-written tiled GEMM (csrc/ops/gemm.cu, TILE_SIZE=32, FP32 accumulate,
boundary-masked) plus an out-of-place transpose kernel. Wire both through
xtrain-cuda FFI (no_cuda-gated) and expose at the tensor level:
Tensor::matmul, transpose_2d, and matmul_backward computing
dA = dC·Bᵀ and dB = Aᵀ·dC by materializing transposes and reusing the
forward. Also declare cuBLAS sgemm FFI + link cublas, used only as a
correctness reference in tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:26:51 +08:00
Gahow Wang
9ca98efd98 autodiff: finite-diff gradient-check harness 
New xtrain-autodiff crate with a reusable central finite-difference
gradient check: grad_check(x, shape, f, analytic_grad, cfg) compares an
analytic gradient against (f(x+ε)-f(x-ε))/2ε per element with a relative
tolerance. Host-only (no CUDA): the loss closure owns any GPU work, so
T4's per-op backward checks can reuse it directly. Includes host unit
tests (sum(x²) grad 2x passes; a wrong grad is rejected).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:26:42 +08:00
Gahow Wang
fbd07a578c tensor: minimal Tensor crate over xtrain-cuda 
New xtrain-tensor crate: DType (F32), shape/stride helpers, Arc-counted
host/device Storage with CPU↔CUDA copy, and a contiguous Tensor with
creation, host↔device transfer, and a scale() op driving the elementwise
kernel. GPU integration tests (host↔device roundtrip + scale correctness)
gated behind not(no_cuda); a thin build.rs emits the no_cuda cfg so the
kernel call sites compile out locally.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:13:06 +08:00
Gahow Wang
63dc05fd10 tensor: add scale elementwise CUDA kernel + FFI 
New csrc/ops/elementwise.cu (out[i]=in[i]*alpha), compiled by
xtrain-cuda/build.rs and exposed via launch_scale_f32 FFI, gated behind
not(no_cuda) like the existing vecadd smoke test.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:13:06 +08:00
Gahow Wang
8557a289a2 docs: Phase T2 — tensor abstraction 
Design doc for the minimal tensor layer: DType/shape/Storage/Tensor,
host↔device copy, and one elementwise kernel (scale) wired end-to-end.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:12:55 +08:00
Gahow Wang
c1b204296b docs: backfill T1 build-chain 
T1 shipped without a design doc; capture the Rust↔CUDA build chain
(build.rs+nvcc, no_cuda cfg pattern, RAII GpuBuffer, gitea↔dash5 flow).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:12:55 +08:00
135 changed files with 47501 additions and 45 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

264
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"
@@ -12,21 +21,276 @@ dependencies = [
"shlex", "shlex",
] ]
[[package]]
name = "cfg-if"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
[[package]]
name = "crunchy"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5"
[[package]] [[package]]
name = "find-msvc-tools" name = "find-msvc-tools"
version = "0.1.9" version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582"
[[package]]
name = "half"
version = "2.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b"
dependencies = [
"cfg-if",
"crunchy",
"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]]
name = "proc-macro2"
version = "1.0.106"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
dependencies = [
"unicode-ident",
]
[[package]]
name = "quote"
version = "1.0.45"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
dependencies = [
"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"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba"
[[package]]
name = "smallvec"
version = "1.15.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90"
[[package]]
name = "syn"
version = "2.0.117"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99"
dependencies = [
"proc-macro2",
"quote",
"unicode-ident",
]
[[package]]
name = "unicode-ident"
version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
[[package]]
name = "xserv-tokenizer"
version = "0.1.0"
dependencies = [
"regex",
"serde",
"serde_json",
]
[[package]]
name = "xtrain-autodiff"
version = "0.1.0"
dependencies = [
"xtrain-cuda",
"xtrain-tensor",
]
[[package]] [[package]]
name = "xtrain-cuda" name = "xtrain-cuda"
version = "0.1.0" version = "0.1.0"
dependencies = [ 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]]
name = "xtrain-tensor"
version = "0.1.0"
dependencies = [
"half",
"smallvec",
"xtrain-autodiff",
"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]]
name = "zerocopy"
version = "0.8.52"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f"
dependencies = [
"zerocopy-derive",
]
[[package]]
name = "zerocopy-derive"
version = "0.8.52"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "zmij"
version = "1.0.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"

10
Cargo.toml
View File

@@ -2,9 +2,19 @@
resolver = "2" resolver = "2"
members = [ members = [
"crates/xtrain-cuda", "crates/xtrain-cuda",
"crates/xtrain-tensor",
"crates/xtrain-autodiff",
"crates/xtrain-model",
"crates/xtrain-optim",
"crates/xtrain-train",
"crates/xtrain-distributed",
] ]
[workspace.package] [workspace.package]
version = "0.1.0" version = "0.1.0"
edition = "2024" edition = "2024"
license = "MIT" license = "MIT"
[workspace.dependencies]
half = "2"
smallvec = "1"

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).

11
crates/xtrain-autodiff/Cargo.toml Normal file
View File

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

27
crates/xtrain-autodiff/build.rs Normal file
View File

@@ -0,0 +1,27 @@
use std::env;
use std::path::Path;
use std::process::Command;
// xtrain-autodiff's `Var` tape calls GPU ops (via xtrain-tensor), so it gates
// those call sites behind `not(no_cuda)` — the same per-crate convention as
// xtrain-cuda / xtrain-tensor (cfg does not propagate across crates). The
// grad_check harness itself is host-only and always compiles. This script only
// detects nvcc and emits the cfg; it compiles no CUDA.
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()
}

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.
}),
)
}

126
crates/xtrain-autodiff/src/finite_diff.rs Normal file
View File

@@ -0,0 +1,126 @@
//! Central finite-difference gradient check.
/// A scalar loss as a function of a flat parameter vector with a given shape.
/// `(data, shape) -> loss`. The closure owns how `data` becomes a `Tensor`
/// (e.g. `Tensor::from_slice(data, shape).to_device(...)`) and runs forward.
pub type ParamFn<'a> = dyn Fn(&[f32], &[usize]) -> f32 + 'a;
#[derive(Debug, Clone, Copy)]
pub struct GradCheckConfig {
/// Perturbation magnitude per element.
pub eps: f32,
/// Max allowed relative error: `|num - ana| / (|num| + |ana| + atol)`.
pub rel_tol: f32,
/// Absolute floor in the denominator, so near-zero grads don't blow up the
/// relative error.
pub atol: f32,
}
impl Default for GradCheckConfig {
fn default() -> Self {
// eps=1e-3 balances truncation error (∝ eps²) against the ~1e-7 f32
// rounding noise on f(x±eps); rel_tol=2e-2 is the usual slack for an
// f32 GPU GEMM checked against an f32 central difference.
Self {
eps: 1e-3,
rel_tol: 2e-2,
atol: 1e-4,
}
}
}
#[derive(Debug, Clone)]
pub struct GradCheckResult {
pub passed: bool,
pub max_rel_err: f32,
/// Index of the worst element (largest relative error).
pub worst_index: usize,
pub worst_numeric: f32,
pub worst_analytic: f32,
}
/// Check `analytic_grad` against the central finite difference of `f` at `x`.
///
/// - `x`: flat parameter values (the point at which the gradient is taken).
/// - `shape`: logical shape passed through to `f`.
/// - `f`: scalar loss; called `2 * x.len()` times.
/// - `analytic_grad`: candidate gradient, same length as `x`.
pub fn grad_check(
x: &[f32],
shape: &[usize],
f: &ParamFn,
analytic_grad: &[f32],
cfg: GradCheckConfig,
) -> GradCheckResult {
assert_eq!(
x.len(),
analytic_grad.len(),
"param/grad length mismatch: {} vs {}",
x.len(),
analytic_grad.len()
);
let mut perturbed = x.to_vec();
let mut max_rel_err = 0.0f32;
let mut worst_index = 0;
let mut worst_numeric = 0.0f32;
let mut worst_analytic = 0.0f32;
for i in 0..x.len() {
let orig = x[i];
perturbed[i] = orig + cfg.eps;
let f_plus = f(&perturbed, shape);
perturbed[i] = orig - cfg.eps;
let f_minus = f(&perturbed, shape);
perturbed[i] = orig; // restore for the next element
let numeric = (f_plus - f_minus) / (2.0 * cfg.eps);
let analytic = analytic_grad[i];
let rel_err = (numeric - analytic).abs() / (numeric.abs() + analytic.abs() + cfg.atol);
if rel_err > max_rel_err {
max_rel_err = rel_err;
worst_index = i;
worst_numeric = numeric;
worst_analytic = analytic;
}
}
GradCheckResult {
passed: max_rel_err <= cfg.rel_tol,
max_rel_err,
worst_index,
worst_numeric,
worst_analytic,
}
}
#[cfg(test)]
mod tests {
use super::*;
// Host-only sanity check (no GPU): loss = sum(x²), grad = 2x.
#[test]
fn quadratic_grad_check() {
let x = vec![1.0f32, -2.0, 3.0, 0.5];
let f = |v: &[f32], _shape: &[usize]| v.iter().map(|t| t * t).sum::<f32>();
let grad: Vec<f32> = x.iter().map(|t| 2.0 * t).collect();
let res = grad_check(&x, &[4], &f, &grad, GradCheckConfig::default());
assert!(res.passed, "max_rel_err = {}", res.max_rel_err);
}
// A deliberately wrong gradient must be rejected.
#[test]
fn wrong_grad_is_rejected() {
let x = vec![1.0f32, 2.0, 3.0];
let f = |v: &[f32], _shape: &[usize]| v.iter().map(|t| t * t).sum::<f32>();
let bad_grad = vec![0.0f32, 0.0, 0.0];
let res = grad_check(&x, &[3], &f, &bad_grad, GradCheckConfig::default());
assert!(!res.passed);
}
}

28
crates/xtrain-autodiff/src/lib.rs Normal file
View File

@@ -0,0 +1,28 @@
//! Reusable numerical-gradient checking for xtrain (Phase T3+).
//!
//! Given a scalar loss `f(x)` and an analytic gradient `g`, verify that `g`
//! matches the central finite-difference estimate
//! `(f(x+ε·eᵢ) - f(x-ε·eᵢ)) / 2ε` for every element `i`, within a relative
//! tolerance. Later phases (T4 autograd) reuse this per-op: wrap each op's
//! forward as the loss, run its backward to get `g`, and `grad_check`.
//!
//! The harness is host-only and dtype-agnostic at this layer: it works on a
//! flat `&[f32]` parameter vector + shape and a closure. The closure is free to
//! push the data to the GPU and run kernels — that detail stays out of here.
pub mod finite_diff;
pub use finite_diff::{GradCheckConfig, GradCheckResult, ParamFn, grad_check};
// Tape-based autograd engine + differentiable ops (Phase T4). These call GPU
// kernels via xtrain-tensor, so they are gated behind `not(no_cuda)` (the
// per-crate convention); the grad_check harness above stays host-only.
#[cfg(not(no_cuda))]
pub mod checkpoint;
#[cfg(not(no_cuda))]
pub mod ops;
#[cfg(not(no_cuda))]
pub mod tape;
#[cfg(not(no_cuda))]
pub use tape::Var;

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

187
crates/xtrain-autodiff/src/tape.rs Normal file
View File

@@ -0,0 +1,187 @@
//! Tape-based reverse-mode autograd (Phase T4).
//!
//! A [`Var`] is a reference-counted node wrapping a value [`Tensor`], an optional
//! accumulated gradient, and — for non-leaf nodes — the parents it was computed
//! from plus a backward closure. Forward ops (see [`crate::ops`]) build these
//! nodes; [`Var::backward`] walks the graph in reverse topological order and
//! pushes gradients to parents, **summing** on fan-out (a tensor consumed by
//! several ops).
//!
//! The graph is dynamic (define-by-run) and the design favours clarity over
//! speed: each op synchronizes its own kernels (T3 has no streams yet), and
//! gradient accumulation is an explicit elementwise add.
#![cfg(not(no_cuda))]
use std::cell::RefCell;
use std::rc::Rc;
use xtrain_tensor::Tensor;
/// Backward closure: given this node's accumulated grad and its parents, compute
/// and accumulate each parent's gradient contribution.
pub type BackwardFn = Box<dyn Fn(&Tensor, &[Var])>;
pub struct VarNode {
pub value: Tensor,
pub grad: Option<Tensor>,
parents: Vec<Var>,
backward: Option<BackwardFn>,
}
/// A node in the autograd tape. Cheap to clone (bumps the `Rc`); clones share
/// the same underlying node, which is how fan-out is detected.
#[derive(Clone)]
pub struct Var(Rc<RefCell<VarNode>>);
impl Var {
/// A leaf node (parameter / input). After `backward`, its `.grad()` holds the
/// accumulated gradient of the loss w.r.t. this tensor.
pub fn leaf(value: Tensor) -> Self {
Var(Rc::new(RefCell::new(VarNode {
value,
grad: None,
parents: Vec::new(),
backward: None,
})))
}
/// Build a non-leaf node from a forward `value`, its `parents`, and a
/// `backward` closure. Used by the op constructors in [`crate::ops`] and by
/// callers that compose custom nodes (e.g. a loss reduction or a transpose
/// the built-in op set doesn't cover).
pub fn from_op(value: Tensor, parents: Vec<Var>, backward: BackwardFn) -> Self {
Var(Rc::new(RefCell::new(VarNode {
value,
grad: None,
parents,
backward: Some(backward),
})))
}
/// Clone of the value tensor (cheap: storage is `Arc`-shared).
pub fn value(&self) -> Tensor {
self.0.borrow().value.clone()
}
/// The accumulated gradient, if any (populated after `backward`).
pub fn grad(&self) -> Option<Tensor> {
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.
fn id(&self) -> *const RefCell<VarNode> {
Rc::as_ptr(&self.0)
}
/// Accumulate `g` into this node's grad slot (SUM — the fan-out rule).
fn accumulate(&self, g: Tensor) {
let mut node = self.0.borrow_mut();
match node.grad.take() {
None => node.grad = Some(g),
Some(existing) => node.grad = Some(existing.add(&g)),
}
}
/// Reverse-mode backward from this node (treated as a scalar loss: its
/// upstream grad is seeded to ones). Populates `.grad` on every node that
/// transitively feeds it.
///
/// `loss` must be a scalar (single element) so the seed `dL/dL = 1` is
/// unambiguous, matching the `L = sum(W∘out)` grad-check convention.
pub fn backward(&self) {
assert_eq!(
self.value().numel(),
1,
"backward() expects a scalar loss; got 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.
let mut topo: Vec<Var> = Vec::new();
let mut visited: Vec<*const RefCell<VarNode>> = Vec::new();
build_topo(self, &mut topo, &mut visited);
// 2. Seed this node's gradient with the supplied upstream grad.
self.accumulate(seed);
// 3. Walk in reverse: each node hands its grad to its parents' closures.
for var in topo.iter().rev() {
let (grad, parents, backward) = {
let node = var.0.borrow();
(
node.grad.clone(),
node.parents.clone(),
node.backward.is_some(),
)
};
let grad = match grad {
Some(g) => g,
None => continue, // node didn't contribute to the loss
};
if backward {
// Borrow the closure out, run it, then drop the borrow.
let node = var.0.borrow();
let bw = node.backward.as_ref().unwrap();
bw(&grad, &parents);
}
}
}
/// Drop the parents/closure so the graph can be freed, keeping value+grad.
/// (Not needed for tests; provided for completeness of the engine.)
pub fn detach_graph(&self) {
let mut node = self.0.borrow_mut();
node.parents.clear();
node.backward = None;
}
/// Distribute `g` to a parent (called from op backward closures). Every node
/// accumulates its grad: intermediates need it for the chain rule, leaves
/// expose it as the result. This SUM is what makes fan-out correct.
pub fn push_grad(parent: &Var, g: Tensor) {
parent.accumulate(g);
}
}
/// Post-order DFS: append a node only after all its parents, dedup by identity.
fn build_topo(var: &Var, topo: &mut Vec<Var>, visited: &mut Vec<*const RefCell<VarNode>>) {
if visited.contains(&var.id()) {
return;
}
visited.push(var.id());
let parents = var.0.borrow().parents.clone();
for p in &parents {
build_topo(p, topo, visited);
}
topo.push(var.clone());
}
/// A ones tensor matching `t`'s shape/device (the backward seed for a scalar).
fn ones_like(t: &Tensor) -> Tensor {
let host = vec![1.0f32; t.numel()];
Tensor::from_slice(&host, t.shape()).to_device(t.device())
}

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

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

@@ -22,12 +22,24 @@ fn main() {
println!("cargo:rustc-link-search=native={cuda_path}/lib64"); println!("cargo:rustc-link-search=native={cuda_path}/lib64");
println!("cargo:rustc-link-lib=dylib=cudart"); println!("cargo:rustc-link-lib=dylib=cudart");
println!("cargo:rustc-link-lib=dylib=cuda"); println!("cargo:rustc-link-lib=dylib=cuda");
// cuBLAS is used only as a correctness reference for the hand-written GEMM.
println!("cargo:rustc-link-lib=dylib=cublas");
cc::Build::new() cc::Build::new()
.cuda(true) .cuda(true)
.cudart("shared") .cudart("shared")
.flag("-gencode=arch=compute_120,code=sm_120") .flag("-gencode=arch=compute_120,code=sm_120")
.file("../../csrc/test/vecadd.cu") .file("../../csrc/test/vecadd.cu")
.file("../../csrc/ops/elementwise.cu")
.file("../../csrc/ops/gemm.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}");
});
}

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

@@ -13,20 +13,611 @@ 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;
} }
// The vector-add smoke-test kernel, compiled from csrc/test/vecadd.cu by build.rs. // GPU kernels compiled from csrc/ by build.rs. Only linked when CUDA is
// Only linked when CUDA is actually compiled (i.e. nvcc was present). // actually compiled (i.e. nvcc was present).
#[cfg(not(no_cuda))] #[cfg(not(no_cuda))]
unsafe extern "C" { unsafe extern "C" {
// Vector-add smoke test (csrc/test/vecadd.cu).
pub fn launch_vecadd_f32(a: *const f32, b: *const f32, c: *mut f32, n: i32, stream: CudaStream); pub fn launch_vecadd_f32(a: *const f32, b: *const f32, c: *mut f32, n: i32, stream: CudaStream);
// Elementwise scale: out[i] = in[i] * alpha (csrc/ops/elementwise.cu).
pub fn launch_scale_f32(
input: *const f32,
out: *mut f32,
alpha: f32,
n: i32,
stream: CudaStream,
);
// Tiled GEMM: C = A @ B, row-major F32. A:[M,K] B:[K,N] C:[M,N]
// (csrc/ops/gemm.cu).
pub fn launch_gemm_tiled_f32(
a: *const f32,
b: *const f32,
c: *mut f32,
m: i32,
n: i32,
k: i32,
stream: CudaStream,
);
// Out-of-place 2D transpose: out[j,i] = in[i,j]. in:[rows,cols] row-major,
// out:[cols,rows] row-major (csrc/ops/gemm.cu).
pub fn launch_transpose_f32(
input: *const f32,
out: *mut f32,
rows: i32,
cols: i32,
stream: CudaStream,
);
}
// Transformer / autograd op kernels (csrc/ops/nn.cu). Forward + backward for the
// ops the Phase T4 tape engine needs. All F32, row-major, contiguous.
#[cfg(not(no_cuda))]
unsafe extern "C" {
// Elementwise: out = a + b ; out = a * b.
pub fn launch_add_f32(a: *const f32, b: *const f32, out: *mut f32, n: i32, s: CudaStream);
pub fn launch_mul_f32(a: *const f32, b: *const f32, out: *mut f32, n: i32, s: CudaStream);
// Broadcast bias add: out[r,c] = x[r,c] + bias[c]. x:[rows,cols], bias:[cols].
pub fn launch_add_bias_f32(
x: *const f32,
bias: *const f32,
out: *mut f32,
rows: i32,
cols: i32,
s: CudaStream,
);
// Column-sum (over rows): dbias[c] = sum_r dout[r,c]. Bias backward.
pub fn launch_sum_rows_f32(
dout: *const f32,
dbias: *mut f32,
rows: i32,
cols: i32,
s: CudaStream,
);
// RMSNorm forward: writes y[rows,cols] and inv_rms[rows] (cached for bwd).
pub fn launch_rms_norm_f32(
x: *const f32,
gamma: *const f32,
y: *mut f32,
inv_rms: *mut f32,
rows: i32,
cols: i32,
eps: f32,
s: CudaStream,
);
pub fn launch_rms_norm_dx_f32(
x: *const f32,
gamma: *const f32,
dy: *const f32,
inv_rms: *const f32,
dx: *mut f32,
rows: i32,
cols: i32,
s: CudaStream,
);
pub fn launch_rms_norm_dgamma_f32(
x: *const f32,
dy: *const f32,
inv_rms: *const f32,
dgamma: *mut f32,
rows: i32,
cols: i32,
s: CudaStream,
);
// SiLU: y = x*sigmoid(x); backward dx.
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);
// 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(
x: *const f32,
y: *mut f32,
tokens: i32,
heads: i32,
head_dim: i32,
theta: f32,
period: i32,
s: CudaStream,
);
// RoPE at an absolute position offset (KV-cache decode, forward only): row
// `tok`'s position is `pos0 + tok` (no modulo). For a single decode token
// (tokens == 1) the one row sits at absolute position `pos0`.
pub fn launch_rope_at_f32(
x: *const f32,
y: *mut f32,
tokens: i32,
heads: i32,
head_dim: i32,
theta: f32,
pos0: i32,
s: CudaStream,
);
// RoPE with a per-row absolute position (batched KV-cache decode, M2b): row
// `tok`'s position is `positions[tok]`. Forward only.
pub fn launch_rope_pos_f32(
x: *const f32,
positions: *const i32,
y: *mut f32,
tokens: i32,
heads: i32,
head_dim: i32,
theta: f32,
s: CudaStream,
);
// Concatenate along the sequence dim: a:[bh,ta,hd], b:[bh,tb,hd] →
// out:[bh,ta+tb,hd] (device-side KV-cache append, M2c).
pub fn launch_cat_seq_f32(
a: *const f32,
b: *const f32,
out: *mut f32,
bh: i32,
ta_hd: i32,
tb_hd: i32,
s: CudaStream,
);
// Per-row scale: y[r,c] = x[r,c] * s[r] (GRPO policy-gradient backward).
pub fn launch_scale_rows_f32(
x: *const f32,
s: *const f32,
y: *mut f32,
rows: i32,
cols: i32,
stream: CudaStream,
);
pub fn launch_rope_dx_f32(
dy: *const f32,
dx: *mut f32,
tokens: i32,
heads: i32,
head_dim: i32,
theta: f32,
period: i32,
s: CudaStream,
);
// Row-wise softmax + Jacobian backward.
pub fn launch_softmax_f32(x: *const f32, y: *mut f32, rows: i32, cols: i32, s: CudaStream);
pub fn launch_softmax_dx_f32(
y: *const f32,
dy: *const f32,
dx: *mut f32,
rows: i32,
cols: i32,
s: CudaStream,
);
// Cross-entropy: fwd writes probs[rows,cols] + per-row loss[rows];
// bwd dx = scale*(probs - onehot).
pub fn launch_cross_entropy_fwd_f32(
x: *const f32,
target: *const i32,
probs: *mut f32,
loss: *mut f32,
rows: i32,
cols: i32,
s: CudaStream,
);
pub fn launch_cross_entropy_dx_f32(
probs: *const f32,
target: *const i32,
dx: *mut f32,
rows: i32,
cols: i32,
scale: f32,
s: CudaStream,
);
}
// Structural ops for the tiny transformer (csrc/ops/model.cu): token embedding
// (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))]
pub type CublasHandle = *mut c_void;
#[cfg(not(no_cuda))]
unsafe extern "C" {
pub fn cublasCreate_v2(handle: *mut CublasHandle) -> i32;
pub fn cublasDestroy_v2(handle: CublasHandle) -> i32;
pub fn cublasSgemm_v2(
handle: CublasHandle,
transa: i32,
transb: i32,
m: i32,
n: i32,
k: i32,
alpha: *const f32,
a: *const f32,
lda: i32,
b: *const f32,
ldb: i32,
beta: *const f32,
c: *mut f32,
ldc: 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))]
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]
);
}

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

@@ -0,0 +1,12 @@
[package]
name = "xtrain-tensor"
version.workspace = true
edition.workspace = true
[dependencies]
xtrain-cuda = { path = "../xtrain-cuda" }
half.workspace = true
smallvec.workspace = true
[dev-dependencies]
xtrain-autodiff = { path = "../xtrain-autodiff" }

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

@@ -0,0 +1,26 @@
use std::env;
use std::path::Path;
use std::process::Command;
// xtrain-tensor calls GPU kernels (via xtrain-cuda's FFI), so it gates those
// call sites behind `not(no_cuda)` — the same convention xtrain-cuda uses. This
// build script only detects nvcc and emits that cfg; it compiles no CUDA itself
// (the kernels are built by xtrain-cuda'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: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()
}

77
crates/xtrain-tensor/src/dtype.rs Normal file
View File

@@ -0,0 +1,77 @@
//! Tensor data types.
//!
//! T2 only needs `F32`; `BF16` was added in T12 for mixed-precision training
//! (bf16 linears / activations, fp32 master weights — see
//! `docs/11-bf16-mixed-precision.md`). The enum + `TensorDType` trait keep call
//! sites dtype-polymorphic.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DType {
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).
I32,
}
impl DType {
pub fn size_bytes(self) -> usize {
match self {
DType::F32 => 4,
DType::BF16 => 2,
DType::I32 => 4,
}
}
pub fn name(self) -> &'static str {
match self {
DType::F32 => "f32",
DType::BF16 => "bf16",
DType::I32 => "i32",
}
}
}
impl std::fmt::Display for DType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.name())
}
}
/// Rust types that can back a tensor. Gives `from_slice`/`as_slice` type safety.
pub trait TensorDType: Copy + Send + Sync + 'static {
const DTYPE: DType;
fn to_f64(self) -> f64;
fn from_f64(v: f64) -> Self;
}
impl TensorDType for f32 {
const DTYPE: DType = DType::F32;
fn to_f64(self) -> f64 {
self as f64
}
fn from_f64(v: f64) -> Self {
v as 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 {
const DTYPE: DType = DType::I32;
fn to_f64(self) -> f64 {
self as f64
}
fn from_f64(v: f64) -> Self {
v as i32
}
}

15
crates/xtrain-tensor/src/lib.rs Normal file
View File

@@ -0,0 +1,15 @@
//! Minimal tensor abstraction for xtrain (Phase T2).
//!
//! Provides a `DType`, shape/stride helpers, reference-counted host/device
//! `Storage`, and a `Tensor` with creation, host↔device transfer, and one
//! elementwise CUDA op (`scale`) wired end-to-end.
pub mod dtype;
pub mod shape;
pub mod storage;
pub mod tensor;
pub use dtype::{DType, TensorDType};
pub use shape::Dims;
pub use storage::{Device, Storage};
pub use tensor::Tensor;

57
crates/xtrain-tensor/src/shape.rs Normal file
View File

@@ -0,0 +1,57 @@
//! Shape / stride helpers. Strides are in **elements** (not bytes), row-major.
use smallvec::SmallVec;
/// Inline storage for the common ≤4D case; spills to the heap beyond that.
pub type Dims = SmallVec<[usize; 4]>;
/// Row-major (C order) contiguous strides for a shape.
/// Example: `[2, 3, 4]` => `[12, 4, 1]`.
pub fn contiguous_strides(shape: &[usize]) -> Dims {
let mut strides: Dims = SmallVec::with_capacity(shape.len());
strides.resize(shape.len(), 0);
if shape.is_empty() {
return strides;
}
strides[shape.len() - 1] = 1;
for i in (0..shape.len() - 1).rev() {
strides[i] = strides[i + 1] * shape[i + 1];
}
strides
}
/// True if `strides` describe a row-major contiguous layout for `shape`.
/// A mismatched stride on a size-1 dimension is fine (it is never stepped).
pub fn is_contiguous(shape: &[usize], strides: &[usize]) -> bool {
let ndim = shape.len();
let mut expected = 1usize;
for d in (0..ndim).rev() {
if shape[d] != 1 && strides[d] != expected {
return false;
}
expected *= shape[d];
}
true
}
/// Total element count.
pub fn num_elements(shape: &[usize]) -> usize {
shape.iter().product()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn contiguous_strides_basic() {
assert_eq!(contiguous_strides(&[2, 3, 4]).as_slice(), &[12, 4, 1]);
assert_eq!(contiguous_strides(&[5]).as_slice(), &[1]);
}
#[test]
fn is_contiguous_detects_transpose() {
assert!(is_contiguous(&[2, 3], &[3, 1]));
assert!(!is_contiguous(&[3, 2], &[1, 3]));
}
}

111
crates/xtrain-tensor/src/storage.rs Normal file
View File

@@ -0,0 +1,111 @@
//! Tensor storage: host (CPU) bytes or a GPU buffer, reference-counted so
//! views (clones with different shape/strides) can share the backing data.
use std::sync::Arc;
use xtrain_cuda::{GpuBuffer, Result as CudaResult};
enum StorageInner {
Cpu { data: Vec<u8> },
Cuda { buffer: GpuBuffer, device: u32 },
}
/// Reference-counted tensor storage. Cloning is cheap (bumps the `Arc`).
#[derive(Clone)]
pub struct Storage(Arc<StorageInner>);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Device {
Cpu,
Cuda(u32),
}
impl std::fmt::Display for Device {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Device::Cpu => write!(f, "cpu"),
Device::Cuda(i) => write!(f, "cuda:{i}"),
}
}
}
impl Storage {
pub fn cpu(data: Vec<u8>) -> Self {
Self(Arc::new(StorageInner::Cpu { data }))
}
pub fn cuda(buffer: GpuBuffer, device: u32) -> Self {
Self(Arc::new(StorageInner::Cuda { buffer, device }))
}
pub fn device(&self) -> Device {
match self.0.as_ref() {
StorageInner::Cpu { .. } => Device::Cpu,
StorageInner::Cuda { device, .. } => Device::Cuda(*device),
}
}
pub fn len_bytes(&self) -> usize {
match self.0.as_ref() {
StorageInner::Cpu { data } => data.len(),
StorageInner::Cuda { buffer, .. } => buffer.len(),
}
}
/// Read-only view of CPU bytes. Panics if the storage lives on the GPU.
pub fn as_cpu_bytes(&self) -> &[u8] {
match self.0.as_ref() {
StorageInner::Cpu { data } => data,
StorageInner::Cuda { .. } => panic!("cannot read GPU storage as CPU bytes"),
}
}
/// Borrow the GPU buffer. Panics if the storage lives on the CPU.
pub fn gpu_buffer(&self) -> &GpuBuffer {
match self.0.as_ref() {
StorageInner::Cuda { buffer, .. } => buffer,
StorageInner::Cpu { .. } => panic!("cannot read CPU storage as GPU buffer"),
}
}
/// Copy to another device. Returns a clone of the `Arc` when already there.
/// T2 supports CPU↔CUDA(0); device-to-device copy across GPUs is out of scope.
pub fn to_device(&self, target: Device) -> CudaResult<Self> {
let current = self.device();
if current == target {
return Ok(self.clone());
}
match (current, target) {
(Device::Cpu, Device::Cuda(dev)) => {
let host = self.as_cpu_bytes();
let mut buf = GpuBuffer::alloc(host.len())?;
buf.copy_from_host(host)?;
Ok(Storage::cuda(buf, dev))
}
(Device::Cuda(_), Device::Cpu) => {
let src = self.gpu_buffer();
let mut host = vec![0u8; src.len()];
src.copy_to_host(&mut host)?;
Ok(Storage::cpu(host))
}
(Device::Cuda(_), Device::Cuda(_)) => {
panic!("cross-GPU storage transfer is not supported in T2")
}
_ => unreachable!(),
}
}
/// Zeroed storage on the given device.
pub fn zeros(len_bytes: usize, device: Device) -> CudaResult<Self> {
match device {
Device::Cpu => Ok(Storage::cpu(vec![0u8; len_bytes])),
Device::Cuda(dev) => {
// 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)?;
buf.memset(0)?;
Ok(Storage::cuda(buf, dev))
}
}
}
}

1719
crates/xtrain-tensor/src/tensor.rs Normal file
View File

File diff suppressed because it is too large Load Diff

202
crates/xtrain-tensor/tests/gemm.rs Normal file
View File

@@ -0,0 +1,202 @@
// GPU acceptance tests for the hand-written GEMM forward + backward (Phase T3).
// Gated behind `not(no_cuda)`: on a GPU-less machine these compile out so host
// `cargo check` stays green; they run on dash5.
#![cfg(not(no_cuda))]
use xtrain_autodiff::{GradCheckConfig, grad_check};
use xtrain_cuda::device;
use xtrain_tensor::{Device, Tensor};
// Deterministic pseudo-random fill in [-0.5, 0.5), seeded by a linear
// congruential generator so tests are reproducible without an RNG dep.
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();
}
// --- cuBLAS reference (correctness oracle for the hand-written kernel) ---
/// Row-major `C = A @ B` via cuBLAS `Sgemm` (which is column-major).
/// Identity: row-major C = A@B ⟺ column-major Cᵀ = Bᵀ @ Aᵀ. We hand cuBLAS
/// our row-major B and A as-is (it reads them as the col-major transposes) with
/// OP_N, swapped order, and m=N, n=M, k=K. Output lands row-major in `c`.
fn cublas_matmul(a: &Tensor, b: &Tensor) -> Tensor {
let m = a.shape()[0];
let k = a.shape()[1];
let n = b.shape()[1];
let c = Tensor::zeros(&[m, n], xtrain_tensor::DType::F32, a.device());
let alpha = 1.0f32;
let beta = 0.0f32;
unsafe {
let mut handle = std::ptr::null_mut();
assert_eq!(xtrain_cuda::ffi::cublasCreate_v2(&mut handle), 0);
let status = xtrain_cuda::ffi::cublasSgemm_v2(
handle,
xtrain_cuda::ffi::CUBLAS_OP_N,
xtrain_cuda::ffi::CUBLAS_OP_N,
n as i32,
m as i32,
k as i32,
&alpha,
b.data_ptr() as *const f32,
n as i32,
a.data_ptr() as *const f32,
k as i32,
&beta,
c.data_ptr() as *mut f32,
n as i32,
);
assert_eq!(status, 0, "cublasSgemm failed: {status}");
device::synchronize().unwrap();
xtrain_cuda::ffi::cublasDestroy_v2(handle);
}
c
}
// Matrix relative error: max element-wise abs error normalized by the
// magnitude scale of the reference. Using a single global denominator avoids
// individual near-zero outputs (where two correct f32 GEMMs differ only in
// rounding order) blowing up a per-element ratio.
fn max_rel_err(got: &[f32], reference: &[f32]) -> f32 {
let scale = reference
.iter()
.fold(0.0f32, |m, r| m.max(r.abs()))
.max(1e-6);
let max_abs = got
.iter()
.zip(reference)
.map(|(g, r)| (g - r).abs())
.fold(0.0f32, f32::max);
max_abs / scale
}
// --- Forward: hand-written tiled GEMM vs cuBLAS sgemm ---
fn run_fwd(m: usize, k: usize, n: usize) {
require_gpu();
let a = Tensor::from_slice(&fill(m * k, 1), &[m, k]).to_device(Device::Cuda(0));
let b = Tensor::from_slice(&fill(k * n, 2), &[k, n]).to_device(Device::Cuda(0));
let mine = a.matmul(&b).to_device(Device::Cpu);
let reference = cublas_matmul(&a, &b).to_device(Device::Cpu);
let rel = max_rel_err(mine.as_slice::<f32>(), reference.as_slice::<f32>());
println!("fwd GEMM [{m}x{k}]@[{k}x{n}] vs cuBLAS: max_rel_err = {rel:.3e}");
assert!(rel < 1e-3, "fwd rel-err {rel} too high for {m}x{k}x{n}");
}
#[test]
fn fwd_square() {
run_fwd(64, 64, 64);
}
#[test]
fn fwd_rect() {
run_fwd(65, 97, 33); // non-tile-aligned dims exercise the boundary masking
}
#[test]
fn fwd_large() {
run_fwd(256, 256, 256);
}
// --- Backward: dA, dB vs the finite-difference harness ---
//
// Scalar loss L = sum(W ∘ C) with C = A @ B and W fixed random weights.
// Then dC = W, dA = dC @ Bᵀ, dB = Aᵀ @ dC (matmul_backward). We check each of
// dA and dB against central differences of L w.r.t. that input.
fn run_bwd(m: usize, k: usize, n: usize) {
require_gpu();
let a_host = fill(m * k, 11);
let b_host = fill(k * n, 22);
let w_host = fill(m * n, 33); // loss weights, == dC
let a = Tensor::from_slice(&a_host, &[m, k]).to_device(Device::Cuda(0));
let b = Tensor::from_slice(&b_host, &[k, n]).to_device(Device::Cuda(0));
let dc = Tensor::from_slice(&w_host, &[m, n]).to_device(Device::Cuda(0));
let (da, db) = Tensor::matmul_backward(&a, &b, &dc);
let da_host = da.to_device(Device::Cpu);
let db_host = db.to_device(Device::Cpu);
// L = sum(W∘C) is bilinear, so it is *exactly linear* in A (B fixed) and in
// B (A fixed): central differences carry no truncation error, and a larger
// eps only sharpens the f32 resolution of f(x+eps)-f(x-eps). atol floors the
// denominator at the ~1e-3 gradient scale so near-zero grads (pure f32
// subtraction noise) don't dominate the relative error.
let cfg = GradCheckConfig {
eps: 1e-2,
rel_tol: 2e-2,
atol: 1e-3,
};
// Check dA: vary A, hold B fixed.
let b_fixed = b_host.clone();
let w_fixed = w_host.clone();
let loss_a = move |a_vals: &[f32], a_shape: &[usize]| -> f32 {
let av = Tensor::from_slice(a_vals, a_shape).to_device(Device::Cuda(0));
let bv = Tensor::from_slice(&b_fixed, &[k, n]).to_device(Device::Cuda(0));
let c = av.matmul(&bv).to_device(Device::Cpu);
c.as_slice::<f32>()
.iter()
.zip(&w_fixed)
.map(|(c, w)| c * w)
.sum()
};
let res_a = grad_check(&a_host, &[m, k], &loss_a, da_host.as_slice::<f32>(), cfg);
println!(
"bwd dA [{m}x{k}]: max_rel_err = {:.3e} (worst num={:.5} ana={:.5} @ {})",
res_a.max_rel_err, res_a.worst_numeric, res_a.worst_analytic, res_a.worst_index
);
assert!(res_a.passed, "dA grad-check failed: {:?}", res_a);
// Check dB: vary B, hold A fixed.
let a_fixed = a_host.clone();
let w_fixed2 = w_host.clone();
let loss_b = move |b_vals: &[f32], b_shape: &[usize]| -> f32 {
let av = Tensor::from_slice(&a_fixed, &[m, k]).to_device(Device::Cuda(0));
let bv = Tensor::from_slice(b_vals, b_shape).to_device(Device::Cuda(0));
let c = av.matmul(&bv).to_device(Device::Cpu);
c.as_slice::<f32>()
.iter()
.zip(&w_fixed2)
.map(|(c, w)| c * w)
.sum()
};
let res_b = grad_check(&b_host, &[k, n], &loss_b, db_host.as_slice::<f32>(), cfg);
println!(
"bwd dB [{k}x{n}]: max_rel_err = {:.3e} (worst num={:.5} ana={:.5} @ {})",
res_b.max_rel_err, res_b.worst_numeric, res_b.worst_analytic, res_b.worst_index
);
assert!(res_b.passed, "dB grad-check failed: {:?}", res_b);
}
#[test]
fn bwd_square() {
run_bwd(16, 16, 16);
}
#[test]
fn bwd_rect() {
run_bwd(12, 20, 8);
}

225
crates/xtrain-tensor/tests/integration.rs Normal file
View File

@@ -0,0 +1,225 @@
// GPU integration tests for the tensor abstraction. Both require nvcc + a GPU,
// so they are gated behind `not(no_cuda)`. On a GPU-less machine build.rs sets
// the `no_cuda` cfg and these compile out, keeping host `cargo check` green.
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_tensor::{Device, Tensor};
/// (a) Host → device → host roundtrip preserves the data exactly.
#[test]
fn host_device_roundtrip() {
assert!(
device::device_count().expect("device count") > 0,
"no CUDA device"
);
device::set_device(0).unwrap();
let host: Vec<f32> = (0..1024).map(|i| i as f32 * 0.5).collect();
let cpu = Tensor::from_slice(&host, &[1024]);
let gpu = cpu.to_device(Device::Cuda(0));
assert_eq!(gpu.device(), Device::Cuda(0));
assert_eq!(gpu.shape(), &[1024]);
let back = gpu.to_device(Device::Cpu);
assert_eq!(back.device(), Device::Cpu);
assert_eq!(back.as_slice::<f32>(), host.as_slice());
println!("roundtrip OK: {} elems preserved", host.len());
}
/// (b) The elementwise `scale` kernel produces correct results.
#[test]
fn elementwise_scale_kernel() {
assert!(
device::device_count().expect("device count") > 0,
"no CUDA device"
);
device::set_device(0).unwrap();
let host: Vec<f32> = (0..2048).map(|i| i as f32).collect();
let alpha = 3.0f32;
let expected: Vec<f32> = host.iter().map(|x| x * alpha).collect();
let gpu = Tensor::from_slice(&host, &[2048]).to_device(Device::Cuda(0));
let scaled = gpu.scale(alpha);
let result = scaled.to_device(Device::Cpu);
assert_eq!(result.shape(), &[2048]);
assert_eq!(result.as_slice::<f32>(), expected.as_slice());
let r = result.as_slice::<f32>();
println!(
"scale OK (alpha={alpha}): first={} mid={} last={} ({} elems)",
r[0],
r[r.len() / 2],
r[r.len() - 1],
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"

17
csrc/ops/elementwise.cu Normal file
View File

@@ -0,0 +1,17 @@
extern "C" {
// out[i] = in[i] * alpha (in-place safe: out may alias in)
__global__ void scale_f32(const float* in, float* out, float alpha, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
out[idx] = in[idx] * alpha;
}
}
void launch_scale_f32(const float* in, float* out, float alpha, int n, void* stream) {
int block = 256;
int grid = (n + block - 1) / block;
scale_f32<<<grid, block, 0, (cudaStream_t)stream>>>(in, out, alpha, n);
}
}

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"

75
csrc/ops/gemm.cu Normal file
View File

@@ -0,0 +1,75 @@
extern "C" {
// Tiled GEMM (shared memory). C = A @ B, all row-major F32.
// A: [M, K], B: [K, N], C: [M, N].
// Each block computes a TILE_SIZE x TILE_SIZE tile of C, cooperatively loading
// tiles of A and B into shared memory. FP32 accumulation.
#define TILE_SIZE 32
__global__ void gemm_tiled_f32(
const float* A, const float* B, float* C,
int M, int N, int K
) {
__shared__ float As[TILE_SIZE][TILE_SIZE];
__shared__ float Bs[TILE_SIZE][TILE_SIZE];
int row = blockIdx.y * TILE_SIZE + threadIdx.y;
int col = blockIdx.x * TILE_SIZE + threadIdx.x;
float sum = 0.0f;
for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
int a_col = t * TILE_SIZE + threadIdx.x;
if (row < M && a_col < K) {
As[threadIdx.y][threadIdx.x] = A[row * K + a_col];
} else {
As[threadIdx.y][threadIdx.x] = 0.0f;
}
int b_row = t * TILE_SIZE + threadIdx.y;
if (b_row < K && col < N) {
Bs[threadIdx.y][threadIdx.x] = B[b_row * N + col];
} else {
Bs[threadIdx.y][threadIdx.x] = 0.0f;
}
__syncthreads();
for (int k = 0; k < TILE_SIZE; k++) {
sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
}
__syncthreads();
}
if (row < M && col < N) {
C[row * N + col] = sum;
}
}
void launch_gemm_tiled_f32(
const float* A, const float* B, float* C,
int M, int N, int K, void* stream
) {
dim3 block(TILE_SIZE, TILE_SIZE);
dim3 grid((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE);
gemm_tiled_f32<<<grid, block, 0, (cudaStream_t)stream>>>(A, B, C, M, N, K);
}
// Out-of-place transpose: out[j, i] = in[i, j].
// in: [rows, cols] row-major, out: [cols, rows] row-major.
__global__ void transpose_f32(const float* in, float* out, int rows, int cols) {
int col = blockIdx.x * blockDim.x + threadIdx.x; // over cols of `in`
int row = blockIdx.y * blockDim.y + threadIdx.y; // over rows of `in`
if (row < rows && col < cols) {
out[col * rows + row] = in[row * cols + col];
}
}
void launch_transpose_f32(const float* in, float* out, int rows, int cols, void* stream) {
dim3 block(16, 16);
dim3 grid((cols + block.x - 1) / block.x, (rows + block.y - 1) / block.y);
transpose_f32<<<grid, block, 0, (cudaStream_t)stream>>>(in, out, rows, cols);
}
}

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"

461
csrc/ops/nn.cu Normal file
View File

@@ -0,0 +1,461 @@
// Forward + backward CUDA kernels for the transformer ops the autograd engine
// (Phase T4) needs: elementwise add/mul, broadcast bias add + its row-sum
// backward, RMSNorm, SiLU, RoPE, row-wise softmax, and cross-entropy.
//
// All F32, row-major, contiguous. Forward kernels mirror xserv
// (docs/04-transformer-kernels.md, docs/05-attention.md); the backward kernels
// are new (xserv is inference-only). Reduction helpers are inlined here so this
// file is self-contained (no shared header), matching the existing csrc/ layout.
#include <math.h>
extern "C" {
// --- Warp / block reductions (sum + max), block handles one row ---
__device__ __forceinline__ float warp_reduce_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 warp_reduce_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 block_reduce_sum(float v) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int warp = threadIdx.x >> 5;
int nwarps = (blockDim.x + 31) >> 5;
v = warp_reduce_sum(v);
if (lane == 0) shared[warp] = v;
__syncthreads();
v = (threadIdx.x < nwarps) ? shared[threadIdx.x] : 0.0f;
if (warp == 0) v = warp_reduce_sum(v);
// broadcast warp-0 lane-0 result to whole block
__shared__ float bcast;
if (threadIdx.x == 0) bcast = v;
__syncthreads();
return bcast;
}
__device__ __forceinline__ float block_reduce_max(float v) {
__shared__ float shared[32];
int lane = threadIdx.x & 31;
int warp = threadIdx.x >> 5;
int nwarps = (blockDim.x + 31) >> 5;
v = warp_reduce_max(v);
if (lane == 0) shared[warp] = v;
__syncthreads();
v = (threadIdx.x < nwarps) ? shared[threadIdx.x] : -INFINITY;
if (warp == 0) v = warp_reduce_max(v);
__shared__ float bcast;
if (threadIdx.x == 0) bcast = v;
__syncthreads();
return bcast;
}
// =====================================================================
// Elementwise add / mul (same-shape)
// =====================================================================
__global__ void add_k(const float* a, const float* b, float* out, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) out[i] = a[i] + b[i];
}
void launch_add_f32(const float* a, const float* b, float* out, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
add_k<<<grid, blk, 0, (cudaStream_t)s>>>(a, b, out, n);
}
__global__ void mul_k(const float* a, const float* b, float* out, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) out[i] = a[i] * b[i];
}
void launch_mul_f32(const float* a, const float* b, float* out, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
mul_k<<<grid, blk, 0, (cudaStream_t)s>>>(a, b, out, n);
}
// =====================================================================
// Broadcast bias add: out[r,c] = x[r,c] + bias[c] (x:[rows,cols])
// Backward for bias is a column-sum (sum over rows): dbias[c] = sum_r dout[r,c].
// =====================================================================
__global__ void add_bias_k(const float* x, const float* bias, float* out,
int rows, int cols) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < rows * cols) out[i] = x[i] + bias[i % cols];
}
void launch_add_bias_f32(const float* x, const float* bias, float* out,
int rows, int cols, void* s) {
int n = rows * cols, blk = 256, grid = (n + blk - 1) / blk;
add_bias_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, bias, out, rows, cols);
}
// dbias[c] = sum_r dout[r,c]. One block per column, threads stride over rows.
__global__ void sum_rows_k(const float* dout, float* dbias, int rows, int cols) {
int col = blockIdx.x;
float acc = 0.0f;
for (int r = threadIdx.x; r < rows; r += blockDim.x)
acc += dout[r * cols + col];
acc = block_reduce_sum(acc);
if (threadIdx.x == 0) dbias[col] = acc;
}
void launch_sum_rows_f32(const float* dout, float* dbias, int rows, int cols, void* s) {
int blk = 256;
sum_rows_k<<<cols, blk, 0, (cudaStream_t)s>>>(dout, dbias, rows, cols);
}
// =====================================================================
// RMSNorm: y[r,c] = x[r,c] * inv_rms[r] * gamma[c], inv_rms = rsqrt(mean(x²)+eps)
// x:[rows,cols], gamma:[cols]. Forward also writes inv_rms[rows] for backward.
// =====================================================================
__global__ void rms_norm_k(const float* x, const float* gamma, float* y,
float* inv_rms, int rows, int cols, float eps) {
int r = blockIdx.x;
const float* xr = x + r * cols;
float* yr = y + r * cols;
float ss = 0.0f;
for (int c = threadIdx.x; c < cols; c += blockDim.x) ss += xr[c] * xr[c];
ss = block_reduce_sum(ss);
float ir = rsqrtf(ss / cols + eps);
if (threadIdx.x == 0) inv_rms[r] = ir;
for (int c = threadIdx.x; c < cols; c += blockDim.x)
yr[c] = xr[c] * ir * gamma[c];
}
void launch_rms_norm_f32(const float* x, const float* gamma, float* y,
float* inv_rms, int rows, int cols, float eps, void* s) {
int blk = cols < 1024 ? cols : 1024;
if (blk < 32) blk = 32;
rms_norm_k<<<rows, blk, 0, (cudaStream_t)s>>>(x, gamma, y, inv_rms, rows, cols, eps);
}
// RMSNorm backward.
// Let g[c] = dy[r,c]*gamma[c], ir = inv_rms[r], n = cols.
// dx[r,c] = ir*g[c] - x[r,c]*ir³/n * sum_c(g[c]*x[r,c])
// dgamma[c] = sum_r dy[r,c] * x[r,c] * ir (accumulated across rows)
__global__ void rms_norm_dx_k(const float* x, const float* gamma, const float* dy,
const float* inv_rms, float* dx, int rows, int cols) {
int r = blockIdx.x;
const float* xr = x + r * cols;
const float* dyr = dy + r * cols;
float* dxr = dx + r * cols;
float ir = inv_rms[r];
float dot = 0.0f; // sum_c g[c]*x[c]
for (int c = threadIdx.x; c < cols; c += blockDim.x)
dot += dyr[c] * gamma[c] * xr[c];
dot = block_reduce_sum(dot);
float coeff = ir * ir * ir / (float)cols * dot;
for (int c = threadIdx.x; c < cols; c += blockDim.x)
dxr[c] = ir * dyr[c] * gamma[c] - xr[c] * coeff;
}
void launch_rms_norm_dx_f32(const float* x, const float* gamma, const float* dy,
const float* inv_rms, float* dx, int rows, int cols, void* s) {
int blk = cols < 1024 ? cols : 1024;
if (blk < 32) blk = 32;
rms_norm_dx_k<<<rows, blk, 0, (cudaStream_t)s>>>(x, gamma, dy, inv_rms, dx, rows, cols);
}
// dgamma[c] = sum_r dy[r,c] * x[r,c] * inv_rms[r]. One block per column.
__global__ void rms_norm_dgamma_k(const float* x, const float* dy, const float* inv_rms,
float* dgamma, int rows, int cols) {
int col = blockIdx.x;
float acc = 0.0f;
for (int r = threadIdx.x; r < rows; r += blockDim.x)
acc += dy[r * cols + col] * x[r * cols + col] * inv_rms[r];
acc = block_reduce_sum(acc);
if (threadIdx.x == 0) dgamma[col] = acc;
}
void launch_rms_norm_dgamma_f32(const float* x, const float* dy, const float* inv_rms,
float* dgamma, int rows, int cols, void* s) {
int blk = 256;
rms_norm_dgamma_k<<<cols, blk, 0, (cudaStream_t)s>>>(x, dy, inv_rms, dgamma, rows, cols);
}
// =====================================================================
// SiLU: y = x * sigmoid(x). Backward: dx = dy * (sig + x*sig*(1-sig)).
// =====================================================================
__global__ void silu_k(const float* x, float* y, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) { float xv = x[i]; y[i] = xv / (1.0f + expf(-xv)); }
}
void launch_silu_f32(const float* x, float* y, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
silu_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, n);
}
__global__ void silu_dx_k(const float* x, const float* dy, float* dx, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float xv = x[i];
float sig = 1.0f / (1.0f + expf(-xv));
dx[i] = dy[i] * (sig + xv * sig * (1.0f - sig));
}
}
void launch_silu_dx_f32(const float* x, const float* dy, float* dx, int n, void* s) {
int blk = 256, grid = (n + blk - 1) / blk;
silu_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, dy, dx, n);
}
// =====================================================================
// RoPE (rotate_half layout). x:[tokens, heads, head_dim]; position = token index.
// y[i] = x[i]*cos - x[i+h]*sin
// y[i+h] = x[i+h]*cos + x[i]*sin (i in [0,half), h=half_dim)
// freq[i] = theta^(-2i/head_dim); angle = pos*freq[i].
// Backward is the inverse (transpose) rotation: apply +angle's transpose ≡ -angle.
// dx[i] = dy[i]*cos + dy[i+h]*sin
// dx[i+h] = dy[i+h]*cos - dy[i]*sin
// =====================================================================
// `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 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);
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_f32(const float* x, float* y, int tokens, int heads,
int head_dim, float theta, int period, void* s) {
dim3 grid(tokens, heads);
int blk = head_dim / 2;
rope_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta, period);
}
// RoPE at an absolute position offset (KV-cache decode-time, forward only). Same
// rotate_half as rope_k, but row `tok`'s position is `pos0 + tok` (no modulo) —
// a single new decode token sits at absolute position pos0. The training rope_k
// (position = tok % period) is left untouched, so this adds no training-path risk.
__global__ void rope_at_k(const float* x, float* y, int heads, int head_dim,
float theta, int pos0) {
int tok = blockIdx.x;
int head = blockIdx.y;
int half = head_dim / 2;
int i = threadIdx.x;
if (i >= half) return;
int pos = pos0 + tok;
float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
float angle = (float)pos * freq;
float c = cosf(angle), sn = sinf(angle);
int base = (tok * heads + head) * head_dim;
float x0 = x[base + i], x1 = x[base + i + half];
y[base + i] = x0 * c - x1 * sn;
y[base + i + half] = x1 * c + x0 * sn;
}
void launch_rope_at_f32(const float* x, float* y, int tokens, int heads,
int head_dim, float theta, int pos0, void* s) {
dim3 grid(tokens, heads);
int blk = head_dim / 2;
rope_at_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta, pos0);
}
// RoPE with a PER-ROW absolute position (batched KV-cache decode, M2b): row `tok`'s
// position is `positions[tok]` (an i32 per token). For G-way batched decode all G
// rows share one decode position; for ragged batches each row carries its own.
// Forward only; the training rope_k is untouched.
__global__ void rope_pos_k(const float* x, const int* positions, float* y,
int heads, int head_dim, float theta) {
int tok = blockIdx.x;
int head = blockIdx.y;
int half = head_dim / 2;
int i = threadIdx.x;
if (i >= half) return;
int pos = positions[tok];
float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
float angle = (float)pos * freq;
float c = cosf(angle), sn = sinf(angle);
int base = (tok * heads + head) * head_dim;
float x0 = x[base + i], x1 = x[base + i + half];
y[base + i] = x0 * c - x1 * sn;
y[base + i + half] = x1 * c + x0 * sn;
}
void launch_rope_pos_f32(const float* x, const int* positions, float* y,
int tokens, int heads, int head_dim, float theta, void* s) {
dim3 grid(tokens, heads);
int blk = head_dim / 2;
rope_pos_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, positions, y, heads, head_dim, theta);
}
// Concatenate along the sequence (middle) dim: a:[bh,ta,hd], b:[bh,tb,hd] →
// out:[bh,ta+tb,hd] with out[:, :ta]=a, out[:, ta:]=b. The device-side KV-cache
// append (M2c): keeps K/V on the GPU and grows by one token per step, removing the
// host round-trip the M2a/M2b host cache paid. One block per bh row.
__global__ void cat_seq_k(const float* a, const float* b, float* out,
int ta_hd, int tb_hd) {
int i = blockIdx.x; // bh row
int o_hd = ta_hd + tb_hd;
const float* ar = a + (long)i * ta_hd;
const float* br = b + (long)i * tb_hd;
float* outr = out + (long)i * o_hd;
for (int j = threadIdx.x; j < ta_hd; j += blockDim.x) outr[j] = ar[j];
for (int j = threadIdx.x; j < tb_hd; j += blockDim.x) outr[ta_hd + j] = br[j];
}
void launch_cat_seq_f32(const float* a, const float* b, float* out,
int bh, int ta_hd, int tb_hd, void* s) {
cat_seq_k<<<bh, 256, 0, (cudaStream_t)s>>>(a, b, out, ta_hd, tb_hd);
}
// Per-row scale: y[r,c] = x[r,c] * s[r]. One block per row. Used by the GRPO
// (M4) policy-gradient backward, where each completion token's row of
// (probs − onehot) is scaled by its own per-token coefficient.
__global__ void scale_rows_k(const float* x, const float* s, float* y,
int rows, int cols) {
int r = blockIdx.x;
float sr = s[r];
for (int c = threadIdx.x; c < cols; c += blockDim.x)
y[r * cols + c] = x[r * cols + c] * sr;
}
void launch_scale_rows_f32(const float* x, const float* s, float* y,
int rows, int cols, void* st) {
int blk = cols < 1024 ? cols : 1024;
if (blk < 32) blk = 32;
scale_rows_k<<<rows, blk, 0, (cudaStream_t)st>>>(x, s, y, rows, cols);
}
__global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim,
float theta, int period) {
int tok = blockIdx.x;
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);
int base = (tok * heads + head) * head_dim;
float d0 = dy[base + i], d1 = dy[base + i + half];
dx[base + i] = d0 * c + d1 * sn;
dx[base + i + half] = d1 * c - d0 * sn;
}
void launch_rope_dx_f32(const float* dy, float* dx, int tokens, int heads,
int head_dim, float theta, int period, void* s) {
dim3 grid(tokens, heads);
int blk = head_dim / 2;
rope_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(dy, dx, heads, head_dim, theta, period);
}
// =====================================================================
// Row-wise safe softmax. x:[rows,cols] → y. Backward (Jacobian):
// dx[r,c] = y[r,c] * (dy[r,c] - sum_c'(dy[r,c']*y[r,c']))
// =====================================================================
__global__ void softmax_k(const float* x, float* y, int rows, int cols) {
int r = blockIdx.x;
const float* xr = x + r * cols;
float* yr = y + r * cols;
float m = -INFINITY;
for (int c = threadIdx.x; c < cols; c += blockDim.x) m = fmaxf(m, xr[c]);
m = block_reduce_max(m);
float sum = 0.0f;
for (int c = threadIdx.x; c < cols; c += blockDim.x) {
float e = expf(xr[c] - m);
yr[c] = e;
sum += e;
}
sum = block_reduce_sum(sum);
float inv = 1.0f / sum;
for (int c = threadIdx.x; c < cols; c += blockDim.x) yr[c] *= inv;
}
void launch_softmax_f32(const float* x, float* y, int rows, int cols, void* s) {
int blk = cols < 1024 ? cols : 1024;
if (blk < 32) blk = 32;
softmax_k<<<rows, blk, 0, (cudaStream_t)s>>>(x, y, rows, cols);
}
__global__ void softmax_dx_k(const float* y, const float* dy, float* dx,
int rows, int cols) {
int r = blockIdx.x;
const float* yr = y + r * cols;
const float* dyr = dy + r * cols;
float* dxr = dx + r * cols;
float dot = 0.0f; // sum_c dy*y
for (int c = threadIdx.x; c < cols; c += blockDim.x) dot += dyr[c] * yr[c];
dot = block_reduce_sum(dot);
for (int c = threadIdx.x; c < cols; c += blockDim.x)
dxr[c] = yr[c] * (dyr[c] - dot);
}
void launch_softmax_dx_f32(const float* y, const float* dy, float* dx,
int rows, int cols, void* s) {
int blk = cols < 1024 ? cols : 1024;
if (blk < 32) blk = 32;
softmax_dx_k<<<rows, blk, 0, (cudaStream_t)s>>>(y, dy, dx, rows, cols);
}
// =====================================================================
// Cross-entropy over logits x:[rows,cols] with int target per row.
// Forward writes per-row loss[r] = -log(softmax(x)[target]) and the softmax
// probs[r,:] (cached for backward). Backward: dx[r,c] = (probs[r,c]-onehot)/rows
// (mean reduction; the *rows scale folds the 1/rows of mean loss into dx).
// =====================================================================
__global__ void cross_entropy_fwd_k(const float* x, const int* target,
float* probs, float* loss, int rows, int cols) {
int r = blockIdx.x;
const float* xr = x + r * cols;
float* pr = probs + r * cols;
float m = -INFINITY;
for (int c = threadIdx.x; c < cols; c += blockDim.x) m = fmaxf(m, xr[c]);
m = block_reduce_max(m);
float sum = 0.0f;
for (int c = threadIdx.x; c < cols; c += blockDim.x) {
float e = expf(xr[c] - m);
pr[c] = e;
sum += e;
}
sum = block_reduce_sum(sum);
float inv = 1.0f / sum;
for (int c = threadIdx.x; c < cols; c += blockDim.x) pr[c] *= inv;
if (threadIdx.x == 0) {
int t = target[r];
loss[r] = t < 0 ? 0.0f : -logf(pr[t]);
}
}
void launch_cross_entropy_fwd_f32(const float* x, const int* target,
float* probs, float* loss, int rows, int cols, void* s) {
int blk = cols < 1024 ? cols : 1024;
if (blk < 32) blk = 32;
cross_entropy_fwd_k<<<rows, blk, 0, (cudaStream_t)s>>>(x, target, probs, loss, rows, cols);
}
// dx[r,c] = scale * (probs[r,c] - [c==target]). scale = upstream/rows.
__global__ void cross_entropy_dx_k(const float* probs, const int* target,
float* dx, int rows, int cols, float scale) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= rows * cols) return;
int r = i / cols, c = i % cols;
int t = target[r];
if (t < 0) {
dx[i] = 0.0f;
} else {
float g = probs[i] - (c == t ? 1.0f : 0.0f);
dx[i] = g * scale;
}
}
void launch_cross_entropy_dx_f32(const float* probs, const int* target,
float* dx, int rows, int cols, float scale, void* s) {
int n = rows * cols, blk = 256, grid = (n + blk - 1) / blk;
cross_entropy_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(probs, target, dx, rows, cols, scale);
}
} // extern "C"

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

81
docs/00-build-chain.md Normal file
View File

@@ -0,0 +1,81 @@
# Phase T1: Rust↔CUDA Build Chain — Design Document
> 回填文档：T1 随 commit `92acf9f` 落地（scaffold + 构建链路 + vecadd 冒烟测试），
> 当时未写文档，此处补记，供后续 Phase 回溯。
## Goal
打通「Rust 调 CUDA」的最小闭环：用 `build.rs` 调 `nvcc` 编 `csrc/*.cu`，
经手写 `extern "C"` FFI 链入 Rust；本地无 GPU/nvcc 时仍能 `cargo check`。
验收以一个 vector-add kernel 冒烟通过为准。
## Module Layout
```
xtrain/
├── Cargo.toml # workspace
├── csrc/
│ └── test/vecadd.cu # 冒烟用 c[i]=a[i]+b[i]
└── crates/xtrain-cuda/
├── build.rs # 检测 nvcc → 编 .cu / 链 cudart，否则发 no_cuda cfg
└── src/
├── ffi.rs # extern "C" 绑定（cudaMalloc/Memcpy/Free…+ kernel launch）
├── error.rs # CudaError + check(code)
├── device.rs # device_count / set_device / synchronize
└── memory.rs # RAII GpuBuffer（alloc/H2D/D2H/Drop）
```
## Key Design Decisions
### `build.rs` 检测 nvcc，缺失则发 `no_cuda` cfg
```rust
if !nvcc_available(&cuda_path) {
println!("cargo:warning=nvcc not found — skipping CUDA compilation (host-only build).");
println!("cargo:rustc-cfg=no_cuda");
return; // 不编 .cu、不链 cudart
}
// 否则：cc::Build::new().cuda(true)…-gencode=arch=compute_120,code=sm_120
```
- 本地机无 GPU/nvcc → 跳过 CUDA 编译，host 侧 Rust 照常 `cargo check`。
- dash5 有 nvcc 12.9 → 实编 `.cu`，目标 `sm_120`（RTX 5090）。
- 必须配 `println!("cargo:rustc-check-cfg=cfg(no_cuda)")` 声明自定义 cfg，
否则新版 rustc 报 `unexpected_cfgs` warning。
### `no_cuda` cfg 门控 GPU-only 代码
GPU 专属的 kernel-launch FFI 与集成测试用 `#[cfg(not(no_cuda))]` 包起来：
```rust
#[cfg(not(no_cuda))]
unsafe extern "C" {
pub fn launch_vecadd_f32(a: *const f32, b: *const f32, c: *mut f32, n: i32, stream: CudaStream);
}
```
测试文件顶用 `#![cfg(not(no_cuda))]` 整体门控。**注意**：`cudaMalloc` 等
runtime 符号在 `ffi.rs` 里**不**门控（恒声明），本地只 `cargo check`（不链接）所以无碍；
真正链接发生在 dash5（cudart 在）。约定：**本地只 check/fmt，链接+测试都在 dash5**。
### RAII `GpuBuffer`
`alloc`/`copy_from_host`(H2D)/`copy_to_host`(D2H)，`Drop` 调 `cudaFree`，
`unsafe impl Send`。后续张量层在其上搭 device 存储（见 `docs/01-tensor.md`）。
## 验证方法
```sh
ssh dash5
export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
cd ~/projects/xtrain && cargo test -p xtrain-cuda -- --nocapture
# vecadd OK: a[i]+b[i]，c[i]=3i 通过
```
本地（无 GPU）：`cargo check` 绿（build.rs 发 `no_cuda`，GPU 测试编译出局）。
## gitea ↔ dash5 同步流
- `origin = git@gitea:gahow/xtrain.git`，分支 `main`；git 身份用 `~/.gitconfig`（不 `-c` 覆盖）。
- 本地 `git push origin main` → dash5 `cd ~/projects/xtrain && git pull` → export PATH 后构建/测试。
- 连 gitea 会打印无害的 `Welcome to VyOS` banner，可忽略。

134
docs/01-tensor.md Normal file
View File

@@ -0,0 +1,134 @@
# Phase: Tensor & Device Buffer — Design Document
## Goal
在 T1 的 `xtrain-cuda`（`GpuBuffer`/`device`/`error`）之上搭最小张量抽象，
作为后续 GEMM / autograd / transformer 的数据基础。本 Phase 只做四件事：
1. `DType`（先 F32，可扩）+ shape/strides；
2. 引用计数的 host/device `Storage`；
3. `Tensor`：创建 + host↔device 拷贝；
4. **一个** elementwise CUDA kernel（`scale`，`out=in*alpha`）端到端打通张量 API。
**明确不做**（留给 T3+）：GEMM、autograd、broadcast、view/transpose、半精度。
## Module Layout
```
crates/xtrain-tensor/
├── Cargo.toml # 依赖 xtrain-cuda + half + smallvec
├── build.rs # 检测 nvcc，缺失则发 no_cuda cfg（与 xtrain-cuda 一致）
└── src/
├── lib.rs # re-exports
├── dtype.rs # DType{F32} + TensorDType trait
├── shape.rs # contiguous_strides / is_contiguous / num_elements
├── storage.rs # Storage(Arc) + Device，CPU↔CUDA 拷贝
└── tensor.rs # Tensor：创建 / 设备迁移 / as_slice / scale kernel
csrc/ops/elementwise.cu # scale_f32 + launch_scale_f32（由 xtrain-cuda/build.rs 编）
```
## Key Design Decisions
### DType + TensorDType trait（先 F32）
```rust
pub enum DType { F32 } // 后续 T7 混合精度再加 F16/BF16
pub trait TensorDType: Copy + Send + Sync + 'static {
const DTYPE: DType;
fn to_f64(self) -> f64;
fn from_f64(v: f64) -> Self;
}
```
trait 让 `from_slice<T>` / `as_slice<T>` 有类型安全。镜像 xserv 的结构，
但只实现 F32 一种——不提前引入用不到的类型。
### Storage 引用计数
```rust
#[derive(Clone)]
pub struct Storage(Arc<StorageInner>);
enum StorageInner {
Cpu { data: Vec<u8> },
Cuda { buffer: GpuBuffer, device: u32 },
}
```
- `Arc` 让未来的 view（transpose/slice）能共享底层数据；T2 暂不产生 view，但类型已就位。
- `to_device(target)`：同设备返回 `Arc` clone（零拷贝）；
CPU→CUDA 走 `GpuBuffer::alloc + copy_from_host`(H2D)，CUDA→CPU 走 `copy_to_host`(D2H)。
- 跨 GPU（CUDA→CUDA 不同卡）T2 不支持（`xtrain-cuda` 暂无 D2D），显式 panic 说明边界。
- `zeros` 在 GPU 上靠 host 端零缓冲 stage 上去（T2 无 device memset，简单优先；后续可加 kernel）。
### Strided Tensor（结构就位，T2 只产生 contiguous）
```rust
pub struct Tensor {
storage: Storage,
shape: Dims, // SmallVec<[usize;4]>，≤4D 免堆分配
strides: Dims, // 以元素为单位，row-major
offset: usize, // 给未来 slice 留的口子
dtype: DType,
}
```
- `strides`/`offset` 字段先放着，T2 创建的张量恒 contiguous、offset=0；
这样 T3+ 加 view 不必改结构体形状。
- `is_contiguous()` 校验 strides 是否匹配 shape（size-1 维度的 stride 不计）。
- `as_slice::<T>()` / `data_ptr()` 要求 contiguous；`data_ptr` 按 dtype 字节算偏移，
供 kernel launch 用。
### Elementwise kernel 端到端（scale）
CUDA 侧（`csrc/ops/elementwise.cu`）：
```cuda
__global__ void scale_f32(const float* in, float* out, float alpha, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) out[i] = in[i] * alpha;
}
```
FFI 声明放在 `xtrain-cuda/src/ffi.rs`（与既有 build 链路/`no_cuda` 门控同处），
张量层 `Tensor::scale(alpha)` 调它：
```rust
#[cfg(not(no_cuda))]
pub fn scale(&self, alpha: f32) -> Self { // out-of-place，要求 contiguous F32 CUDA 张量
let out = Tensor::zeros(&self.shape, self.dtype, self.device());
unsafe { xtrain_cuda::ffi::launch_scale_f32(self.data_ptr() as *const f32,
out.data_ptr() as *mut f32, alpha, self.numel() as i32, null_mut()); }
xtrain_cuda::device::synchronize().unwrap();
out
}
```
- kernel FFI 留在 `xtrain-cuda`（构建链路与 `no_cuda` cfg 都在那），张量层只调用——
避免在张量 crate 里再开一套 nvcc 编译。
- `scale` 用 `#[cfg(not(no_cuda))]` 门控；为此 `xtrain-tensor` 加了一个**只检测 nvcc、
不编译任何 .cu** 的 `build.rs`，发同名 `no_cuda` cfg（cfg 不跨 crate 传播，必须各自发）。
## 验证方法
GPU 测试用 `#![cfg(not(no_cuda))]` 门控，在 dash5 实跑：
```sh
ssh dash5
export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
cd ~/projects/xtrain && cargo test -p xtrain-tensor -- --nocapture
```
- **(a) host↔device 往返拷贝**：CPU 张量 → CUDA → 拷回 CPU，逐元素 `assert_eq` 原样。
- **(b) elementwise 正确性**：`scale(3.0)` 后拷回，对 `host[i]*3.0` 逐元素相等。
本地（无 GPU）：`cargo check --workspace --all-targets` + `cargo fmt --all -- --check` 绿；
GPU 测试编译出局（约定：本地只 check/fmt，链接+测试都在 dash5）。
## Takeaways
1. **cfg 不跨 crate**：`no_cuda` 由各 crate 自己的 `build.rs` 发；张量 crate 要门控 kernel 调用，
就得加一个轻量 build.rs（只检测、不编译）。
2. **结构先于功能**：`strides`/`offset` 先放进结构体，T3 加 view 时不动 shape，降低后续改动面。
3. **边界显式 panic**：跨 GPU 拷贝、非 contiguous as_slice 等 T2 不支持的路径直接 panic 写清原因，
而不是悄悄给错结果。
4. **kernel 收口在 xtrain-cuda**：构建链路单点，张量层保持纯 Rust 调用，符合 T1 立的约定。

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