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

Compare commits

base: gahow:main
Branches Tags
gahow:main
..
compare: gahow:59e33f60e3967e8da00bd39606243fe426f47a42
Branches Tags
gahow:main

2 Commits
main ... 59e33f60e3

Author SHA1 Message Date
Gahow Wang
59e33f60e3 docs: Phase T4 — autograd engine 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
 2026-06-15 15:47:43 +08:00
Gahow Wang
f3a1188f0e 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:47:43 +08:00
121 changed files with 180 additions and 44192 deletions
Expand all files Collapse all files

8
.gitignore vendored
View File

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

163
Cargo.lock generated
View File

@@ -2,15 +2,6 @@
# It is not intended for manual editing.
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]]
name = "cc"
version = "1.2.64"
@@ -50,18 +41,6 @@ dependencies = [
"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"
@@ -80,88 +59,6 @@ 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]]
name = "shlex"
version = "2.0.1"
@@ -191,20 +88,10 @@ 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",
]
@@ -215,36 +102,6 @@ dependencies = [
"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"
@@ -255,20 +112,6 @@ dependencies = [
"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"
@@ -288,9 +131,3 @@ dependencies = [
"quote",
"syn",
]
[[package]]
name = "zmij"
version = "1.0.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"

4
Cargo.toml
View File

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

232
README.md
View File

@@ -1,210 +1,50 @@
# xtrain
A from-scratch **Rust + CUDA** LLM **training** engine — the sibling of **xserv** (the
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.
A from-scratch **Rust + CUDA** LLM **training** engine — the sibling of
[xserv](https://github.com/) (the inference side). GPU-first.
> **Status: complete — three phases.**
> **Phase 1** = the from-scratch full stack (T1–T13) + an 8-version scaling study (v0–v8):
> hand-write the whole training-systems stack, then map the data-vs-capacity frontier.
> **Phase 2** = systems-stack depth (T14–T18): hand-write the five deferred training-stack
> features — fused flash-attention, real GQA, gradient accumulation, process-per-GPU DDP,
> dropout. **Phase 3** = one Chinchilla-style double-axis run (v9): dim1280 true-GQA +
> 6.01B FineWeb tokens, validating the v8 conclusion that data and capacity must scale
> together. Trains Qwen3-compatible LMs whose weights load into **xserv**; deterministic
> gates stay byte-identical, while large BF16 checkpoints are served and checked for
> prompt-level drift. This README is the capstone; per-topic detail lives in [`docs/`](docs/).
The goal is to learn the full training-systems stack by hand: autograd / backward
passes / optimizers (AdamW) / the training loop / distributed logic. Heavy lifting
is borrowed where it makes sense (GEMM → cuBLAS after a hand-written version,
multi-GPU comms → NCCL, tokenizer → reused from xserv), but the core is written
from scratch. The target architecture is a tiny modern transformer
(RoPE + RMSNorm + SwiGLU, ~1–30M params) whose forward aligns with xserv's Qwen3,
so the backward passes map one-to-one onto xserv's existing forward kernels and
trained weights can flow back into xserv.
---
## Status
## What got built (from scratch, by hand)
Bootstrapping (P0). This repo currently contains only the project skeleton and a
working Rust↔CUDA build chain, verified by a trivial vector-add CUDA kernel.
7 crates, no ML framework — only cuBLAS / NCCL / safetensors as deliberate "heavy-lifting"
borrows, the rest hand-written CUDA + Rust:
## Layout
| crate | what's hand-written |
|---|---|
| `xtrain-cuda` | CUDA Runtime FFI, RAII `GpuBuffer`, **caching/pool allocator**, cuBLAS (sgemm + bf16 GemmEx) bindings |
| `xtrain-tensor` | tensor (dtype/shape/strides/storage), elementwise + transpose + embedding kernels |
| `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) |
| `xtrain-model` | tiny **Qwen3-style** transformer (RoPE + RMSNorm + QK-norm + SwiGLU), batched forward, **GQA** (`num_kv_heads<num_heads`), residual/MLP **dropout** |
| `xtrain-optim` | hand-written **AdamW** (host + GPU kernels) |
| `xtrain-train` | training loop, LR schedule, grad clip, **gradient accumulation**, checkpoint, BPE corpus + cache, samplers, safetensors export |
| `xtrain-distributed` | **NCCL DDP** (thread-per-GPU + torchrun-style process-per-GPU launcher / cross-process `ncclUniqueId`, all-reduce) |
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,
```
xtrain/
├── Cargo.toml # workspace
├── csrc/ # CUDA sources (.cu)
│ └── test/vecadd.cu # trivial element-wise vector-add (smoke test)
└── crates/
└── xtrain-cuda/ # CUDA Runtime FFI + build.rs (nvcc → sm_120)
├── build.rs # compiles csrc/*.cu via the `cc` crate, links cudart
├── src/ # ffi / error / device / memory
└── tests/ # vecadd smoke test
```
Build/test the engine itself (CUDA compiles + runs on the GPU box; host-side `cargo check`
works anywhere via the `no_cuda` cfg):
The build mirrors xserv's approach: `build.rs` invokes `nvcc` (via the `cc` crate)
to compile `csrc/*.cu` targeting `sm_120` (RTX 5090) and links them into the Rust
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
export PATH=/usr/local/cuda/bin:$HOME/.cargo/bin:$PATH
cargo test --workspace # autograd grad-checks, PyTorch parity, DDP, etc.
cargo build
cargo test -p xtrain-cuda -- --nocapture # runs the vecadd smoke test
```
## Doc index
- [`docs/evolution.md`](docs/evolution.md) — per-milestone changes across algorithm / architecture / infra / dataset.
- [`docs/runs/README.md`](docs/runs/README.md) — the v0–v10 comparison; [`docs/runs/0N-*.md`](docs/runs/) — per-run detail.
- [`docs/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).
On a machine without `nvcc`/GPU, `build.rs` detects the missing toolchain, skips
CUDA compilation, and sets a `no_cuda` cfg — so host-side `cargo check` still
works (the GPU smoke test is compiled out).

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

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

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

@@ -18,8 +18,6 @@ pub use finite_diff::{GradCheckConfig, GradCheckResult, ParamFn, grad_check};
// 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;

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

@@ -13,27 +13,7 @@
#![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));
}),
)
}
use xtrain_tensor::Tensor;
/// `C = A @ B` (2D). Backward: `dA = dC @ Bᵀ`, `dB = Aᵀ @ dC`.
pub fn matmul(a: &Var, b: &Var) -> Var {
@@ -140,42 +120,15 @@ 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);
/// RoPE (rotate_half) over `x:[tokens,heads,head_dim]`. Orthogonal map, so the
/// backward is the inverse rotation of `dy` — no cached forward values needed.
pub fn rope(x: &Var, theta: f32) -> Var {
let out = x.value().rope(theta);
Var::from_op(
out,
vec![x.clone()],
Box::new(move |dy, parents| {
Var::push_grad(&parents[0], Tensor::rope_backward(dy, theta, period));
Var::push_grad(&parents[0], Tensor::rope_backward(dy, theta));
}),
)
}
@@ -193,237 +146,15 @@ pub fn softmax(x: &Var) -> Var {
)
}
/// 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`,
/// row. Returns a scalar [`Var`]. Backward: `dx = (probs - onehot)/rows`,
/// scaled by the upstream scalar grad.
pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
// CE math is fp32 (cross_entropy upcasts bf16 logits internally + caches fp32
// 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);
let rows = x.value().shape()[0];
// Mean loss as a host scalar wrapped back into a [1] tensor.
let mean = per_row.to_device(xtrain_tensor::Device::Cpu);
let mean_val: f32 = mean.as_slice::<f32>().iter().sum::<f32>() / valid_rows as f32;
let mean_val: f32 = mean.as_slice::<f32>().iter().sum::<f32>() / rows as f32;
let loss = Tensor::from_slice(&[mean_val], &[1]).to_device(x.value().device());
let target = target.clone();
@@ -433,251 +164,9 @@ pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
Box::new(move |d, parents| {
// `d` is the scalar upstream grad (1.0 when this is the loss root).
let upstream = d.to_device(xtrain_tensor::Device::Cpu).as_slice::<f32>()[0];
let scale = upstream / valid_rows as f32;
let scale = upstream / rows as f32;
let dx = Tensor::cross_entropy_backward(&probs, &target, scale);
Var::push_grad(&parents[0], dx.to_dtype(logit_dtype));
}),
)
}
/// Per-sequence log-probability: `Σ log πθ(target)` over the non-ignored
/// (`target ≥ 0`) positions — the quantity DPO (M3) compares between policy and
/// reference. `target` is `[rows]` I32 carrying `-100` (ignore) at masked positions
/// (e.g. the prompt) and the gold token id elsewhere; ignored positions contribute
/// 0, exactly like the SFT cross-entropy masking. Returns a scalar `[1]` Var.
///
/// Reuses the CE forward (per-row `−log p(target)`) and backward, so no new kernel:
/// `seq_logprob = −Σ per_row`, and `d(seq_logprob)/d(logits) = −(probs − onehot)`
/// = `cross_entropy_backward(probs, target, −upstream)` (a SUM, so no mean
/// division — contrast [`cross_entropy`], which divides by `valid_rows`).
pub fn seq_logprob(x: &Var, target: &Tensor) -> Var {
let logit_dtype = x.value().dtype();
let (probs, per_row) = x.value().cross_entropy(target);
// per_row[r] = −log p(target_r), and is 0 for ignored rows (target < 0), so the
// sum already counts only the supervised (completion) positions.
let sum_neg_lp: f32 = per_row
.to_device(xtrain_tensor::Device::Cpu)
.as_slice::<f32>()
.iter()
.sum();
let out = Tensor::from_slice(&[-sum_neg_lp], &[1]).to_device(x.value().device());
let target = target.clone();
Var::from_op(
out,
vec![x.clone()],
Box::new(move |d, parents| {
let upstream = d.to_device(xtrain_tensor::Device::Cpu).as_slice::<f32>()[0];
// d(Σ log p)/d(logits) = −(probs − onehot); SUM, so no /valid_rows.
let dx = Tensor::cross_entropy_backward(&probs, &target, -upstream);
Var::push_grad(&parents[0], dx.to_dtype(logit_dtype));
}),
)
}
/// DPO loss (Rafailov et al., M3) for one preference pair, as a scalar `[1]` Var
/// whose two parents are the POLICY sequence-logprobs of the chosen and rejected
/// completions (from [`seq_logprob`]); the REFERENCE logprobs are constants
/// (precomputed once from the frozen SFT model). With
/// `Δ = β·[(lpθ_chosen − lpref_chosen) − (lpθ_rejected − lpref_rejected)]`
/// the loss is `L = −log σ(Δ) = softplus(−Δ)`. Only the policy terms carry gradient:
/// `∂L/∂lpθ_chosen = −β·(1−σ(Δ))`, `∂L/∂lpθ_rejected = +β·(1−σ(Δ))`.
/// Degenerate points the M3 gate pins: `πθ == πref` ⇒ `Δ = 0`, `L = log 2`, implicit
/// reward 0; `β → 0` ⇒ gradient → 0. Same formula as TRL
/// (`-logsigmoid(β·(pol_c − pol_r − (ref_c − ref_r)))`).
pub fn dpo_loss(
lp_pol_chosen: &Var,
lp_pol_rejected: &Var,
lp_ref_chosen: f32,
lp_ref_rejected: f32,
beta: f32,
) -> Var {
use xtrain_tensor::Device;
let scalar = |v: &Var| v.value().to_device(Device::Cpu).as_slice::<f32>()[0];
let pc = scalar(lp_pol_chosen);
let pr = scalar(lp_pol_rejected);
let delta = beta * ((pc - lp_ref_chosen) - (pr - lp_ref_rejected));
// L = softplus(−Δ) = log(1 + e^{−Δ}) (numerically stable).
let nd = -delta;
let l = nd.max(0.0) + (-(nd.abs())).exp().ln_1p();
let dev = lp_pol_chosen.value().device();
let out = Tensor::from_slice(&[l], &[1]).to_device(dev);
Var::from_op(
out,
vec![lp_pol_chosen.clone(), lp_pol_rejected.clone()],
Box::new(move |d, parents| {
let up = d.to_device(Device::Cpu).as_slice::<f32>()[0];
// s = σ(−Δ) = 1 − σ(Δ); ∂L/∂Δ = −s, and ∂Δ/∂pc = β, ∂Δ/∂pr = −β.
let s = 1.0 / (1.0 + delta.exp());
let g = up * beta * s;
let dev = parents[0].value().device();
Var::push_grad(&parents[0], Tensor::from_slice(&[-g], &[1]).to_device(dev));
Var::push_grad(&parents[1], Tensor::from_slice(&[g], &[1]).to_device(dev));
}),
)
}
/// GRPO clipped policy-gradient loss (M4) for ONE completion, a scalar `[1]` Var
/// with the policy logits as the single parent. Per non-ignored (completion) token
/// `t` (`target[t] ≥ 0`):
/// `logπθ_t = log softmax(logits[t])[target_t]` (`= −per_row[t]` of cross_entropy)
/// `ρ_t = exp(logπθ_t − logp_old[t])`
/// `pg_t = min(ρ_t·A, clip(ρ_t, 1−ε, 1+ε)·A)`
/// `kl_t = exp(logp_ref[t] − logπθ_t) − (logp_ref[t] − logπθ_t) − 1` (k3 estimator)
/// `L = −mean_t pg_t + β·mean_t kl_t` over the `N` completion tokens.
///
/// `advantage` `A` is the group-relative advantage (constant per completion in
/// GRPO); `logp_old`/`logp_ref` are per-position constants (old policy at rollout
/// time / frozen reference). Backward reuses the CE machinery + the per-row
/// `scale_rows`: `dL/dlogits[t,:] = g_t·(onehot − probs)[t,:]` with
/// `g_t = −(1/N)A·ρ_t·[unclipped active] + (β/N)(1 − exp(logp_ref_t − logπθ_t))`.
/// Degenerate points the gate pins: `A=0` ⇒ only the KL term; `ε→∞` ⇒ vanilla PG
/// (no clip); `β=0` ⇒ no KL term.
#[allow(clippy::too_many_arguments)]
pub fn clipped_pg_loss(
logits: &Var,
target: &Tensor,
logp_old: &[f32],
logp_ref: &[f32],
advantage: f32,
eps: f32,
beta: f32,
) -> Var {
use xtrain_tensor::Device;
let logit_dtype = logits.value().dtype();
let (probs, per_row) = logits.value().cross_entropy(target);
let rows = per_row.shape()[0];
let per_row_h = per_row.to_device(Device::Cpu).as_slice::<f32>().to_vec();
let target_h = target.to_device(Device::Cpu).as_slice::<i32>().to_vec();
assert_eq!(logp_old.len(), rows, "logp_old must have one entry per position");
assert_eq!(logp_ref.len(), rows, "logp_ref must have one entry per position");
let mut s = vec![0f32; rows]; // per-row scale for cross_entropy_backward(·,·,1.0)
let (mut pg_sum, mut kl_sum, mut n) = (0f32, 0f32, 0f32);
for t in 0..rows {
if target_h[t] < 0 {
continue; // masked (prompt) position — no contribution, no gradient
}
n += 1.0;
let lp = -per_row_h[t]; // logπθ_t
let ratio = (lp - logp_old[t]).exp();
let clipped = ratio.clamp(1.0 - eps, 1.0 + eps);
let (unclipped_term, clipped_term) = (ratio * advantage, clipped * advantage);
pg_sum += unclipped_term.min(clipped_term);
let active = unclipped_term <= clipped_term; // min picks unclipped ⇒ grad flows
let d = logp_ref[t] - lp;
kl_sum += d.exp() - d - 1.0;
let pg_grad = if active { -advantage * ratio } else { 0.0 };
let kl_grad = beta * (1.0 - d.exp());
s[t] = -(pg_grad + kl_grad); // dL/dlogits = g·(onehot−probs) = −g·(probs−onehot)
}
let inv_n = if n > 0.0 { 1.0 / n } else { 1.0 };
for v in &mut s {
*v *= inv_n;
}
let loss_val = -pg_sum * inv_n + beta * kl_sum * inv_n;
let dev = logits.value().device();
let out = Tensor::from_slice(&[loss_val], &[1]).to_device(dev);
let s_dev = Tensor::from_slice(&s, &[rows]).to_device(dev);
let target = target.clone();
Var::from_op(
out,
vec![logits.clone()],
Box::new(move |d, parents| {
let up = d.to_device(Device::Cpu).as_slice::<f32>()[0];
// (probs − onehot), masked rows already 0; per-row scale by s; × upstream.
let ce = Tensor::cross_entropy_backward(&probs, &target, 1.0);
let mut dx = ce.scale_rows(&s_dev);
if up != 1.0 {
dx = dx.scale(up);
}
Var::push_grad(&parents[0], dx.to_dtype(logit_dtype));
}),
)
}
/// Batched GRPO clipped-PG loss over `N` ragged completions packed into ONE
/// `forward_batched` (M2d): `logits` is `[R, vocab]` with `R = N·Lmax` rows in
/// sequence-major order (sample 0's `Lmax` rows, then sample 1's, …), each ragged
/// completion right-padded to the batch's `Lmax`. Prompt AND pad rows are masked
/// (`target < 0`), so they contribute nothing and carry no gradient — the
/// **right-pad-is-free-under-causal-attention** property (a real completion row
/// never attends to the trailing pad rows, so its logits equal the unpadded
/// single-sequence forward's).
///
/// Unlike the per-sample [`clipped_pg_loss`] (which folds a single scalar
/// `advantage` and a global `1/N_tokens` normaliser), this op takes **per-row**
/// `advantage[t]` (the owning sample's group-relative `A`) and **per-row**
/// `weight[t]` (the full normaliser, e.g. `1/(N_samples · n_s)` for sample `s`'s
/// completion rows, `0` at masked rows). It does NOT compute its own `inv_n`. With
/// `weight[t] = 1/(N_samples·n_s)` and `advantage[t] = A_s` this is **bit-equivalent
/// to the looped path** `Σ_s scale·(1/n_s)·clipped_pg_loss_s` (`scale = 1/N_samples`):
/// the per-row backward is local (`cross_entropy_backward` is row-wise), so the
/// batched row-`t` gradient equals the looped sample-`s` row-`t` gradient, and the
/// scalar loss equals the looped weighted sum. (`tests/autograd.rs`:
/// `clipped_pg_loss_batched_matches_looped`.) Degenerate points match
/// [`clipped_pg_loss`] (`A=0` ⇒ KL only; `ε→∞` ⇒ vanilla PG; `β=0` ⇒ no KL).
#[allow(clippy::too_many_arguments)]
pub fn clipped_pg_loss_batched(
logits: &Var,
target: &Tensor,
logp_old: &[f32],
logp_ref: &[f32],
advantage: &[f32],
weight: &[f32],
eps: f32,
beta: f32,
) -> Var {
use xtrain_tensor::Device;
let logit_dtype = logits.value().dtype();
let (probs, per_row) = logits.value().cross_entropy(target);
let rows = per_row.shape()[0];
let per_row_h = per_row.to_device(Device::Cpu).as_slice::<f32>().to_vec();
let target_h = target.to_device(Device::Cpu).as_slice::<i32>().to_vec();
assert_eq!(logp_old.len(), rows, "logp_old must have one entry per row");
assert_eq!(logp_ref.len(), rows, "logp_ref must have one entry per row");
assert_eq!(advantage.len(), rows, "advantage must have one entry per row");
assert_eq!(weight.len(), rows, "weight must have one entry per row");
let mut s = vec![0f32; rows]; // per-row scale for cross_entropy_backward(·,·,1.0)
let mut loss_val = 0f32;
for t in 0..rows {
if target_h[t] < 0 {
continue; // masked (prompt or pad) row — no contribution, no gradient
}
let (a, w) = (advantage[t], weight[t]);
let lp = -per_row_h[t]; // logπθ_t
let ratio = (lp - logp_old[t]).exp();
let clipped = ratio.clamp(1.0 - eps, 1.0 + eps);
let (unclipped_term, clipped_term) = (ratio * a, clipped * a);
let pg_t = unclipped_term.min(clipped_term);
let active = unclipped_term <= clipped_term; // min picks unclipped ⇒ grad flows
let d = logp_ref[t] - lp;
let kl_t = d.exp() - d - 1.0;
let pg_grad = if active { -a * ratio } else { 0.0 };
let kl_grad = beta * (1.0 - d.exp());
// The full per-row normaliser is folded into s (no global inv_n here).
s[t] = -(pg_grad + kl_grad) * w;
loss_val += (-pg_t + beta * kl_t) * w;
}
let dev = logits.value().device();
let out = Tensor::from_slice(&[loss_val], &[1]).to_device(dev);
let s_dev = Tensor::from_slice(&s, &[rows]).to_device(dev);
let target = target.clone();
Var::from_op(
out,
vec![logits.clone()],
Box::new(move |d, parents| {
let up = d.to_device(Device::Cpu).as_slice::<f32>()[0];
let ce = Tensor::cross_entropy_backward(&probs, &target, 1.0);
let mut dx = ce.scale_rows(&s_dev);
if up != 1.0 {
dx = dx.scale(up);
}
Var::push_grad(&parents[0], dx.to_dtype(logit_dtype));
Var::push_grad(&parents[0], dx);
}),
)
}

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

@@ -68,19 +68,6 @@ impl Var {
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)
@@ -108,24 +95,14 @@ impl Var {
"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.
// 2. Seed the loss gradient with ones.
let seed = ones_like(&self.value());
self.accumulate(seed);
// 3. Walk in reverse: each node hands its grad to its parents' closures.

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

@@ -327,12 +327,12 @@ fn rope_bwd() {
let w = fill(n, 82);
let x = Var::leaf(cuda(&x_h, &[tokens, heads, head_dim]));
let out = ops::rope(&x, theta, tokens);
let out = ops::rope(&x, theta);
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).rope(theta, tokens), &wf);
let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).rope(theta), &wf);
report(
"rope dX",
&grad_check(
@@ -345,38 +345,6 @@ fn rope_bwd() {
);
}
// ---- rope batched (per-sequence position = row % period) ----
// tokens = B*S laid end to end; period = S. Sequences 2 and 3 re-use positions
// 0..S, so the kernel's `tok % period` must reset RoPE per sequence.
#[test]
fn rope_batched_bwd() {
require_gpu();
let (b, s, heads, head_dim) = (3, 4, 2, 8);
let tokens = b * s;
let n = tokens * heads * head_dim;
let theta = 10000.0;
let x_h = fill(n, 83);
let w = fill(n, 84);
let x = Var::leaf(cuda(&x_h, &[tokens, heads, head_dim]));
let out = ops::rope(&x, theta, s);
scalar_loss(&out, &w).backward();
let dx = x.grad().unwrap().to_device(Device::Cpu);
let wf = w.clone();
let lx = move |v: &[f32], sh: &[usize]| weighted_sum(&cuda(v, sh).rope(theta, s), &wf);
report(
"rope batched dX",
&grad_check(
&x_h,
&[tokens, heads, head_dim],
&lx,
dx.as_slice::<f32>(),
cfg_linear(),
),
);
}
// ---- softmax ----
#[test]
fn softmax_bwd() {
@@ -533,435 +501,6 @@ fn attention_composed_bwd() {
);
}
// ---- transpose_4d12 ([a,b,c,d] -> [a,c,b,d]) ----
#[test]
fn transpose_4d12_bwd() {
require_gpu();
let (a, b, c, d) = (2, 3, 4, 5);
let n = a * b * c * d;
let x_h = fill(n, 131);
let w = fill(n, 132);
let x = Var::leaf(cuda(&x_h, &[a, b, c, d]));
let out = ops::transpose_4d12(&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_4d12(), &wf);
report(
"transpose_4d12 dX",
&grad_check(&x_h, &[a, b, c, d], &lx, dx.as_slice::<f32>(), cfg_linear()),
);
}
// ---- fused batched causal attention (the T10 op) ----
// q,k,v: [bh, seq, hd]. Grad-check dq/dk/dv against finite-diff of L = sum(W∘out).
// bh = 2 (e.g. batch 1 × 2 heads, or 2 sequences × 1 head) exercises the batched
// GEMM stride; the causal mask is applied inside the op.
#[test]
fn attention_batched_bwd() {
require_gpu();
let (bh, seq, hd) = (2, 5, 6);
let n = bh * seq * hd;
let scale = 1.0 / (hd as f32).sqrt();
let q_h = fill(n, 141);
let k_h = fill(n, 142);
let v_h = fill(n, 143);
let w = fill(n, 144);
let q = Var::leaf(cuda(&q_h, &[bh, seq, hd]));
let k = Var::leaf(cuda(&k_h, &[bh, seq, hd]));
let v = Var::leaf(cuda(&v_h, &[bh, seq, hd]));
let out = ops::attention(&q, &k, &v, scale);
scalar_loss(&out, &w).backward();
let dq = q.grad().unwrap().to_device(Device::Cpu);
let dk = k.grad().unwrap().to_device(Device::Cpu);
let dv = v.grad().unwrap().to_device(Device::Cpu);
let fwd = move |qh: &[f32], kh: &[f32], vh: &[f32]| -> f32 {
let qv = cuda(qh, &[bh, seq, hd]);
let kv = cuda(kh, &[bh, seq, hd]);
let vv = cuda(vh, &[bh, seq, hd]);
let (o, _) = qv.attention(&kv, &vv, scale);
weighted_sum(&o, &w)
};
let (kf, vf, ff) = (k_h.clone(), v_h.clone(), fwd.clone());
let lq = move |x: &[f32], _s: &[usize]| ff(x, &kf, &vf);
report(
"attn(batched) dQ",
&grad_check(
&q_h,
&[bh, seq, hd],
&lq,
dq.as_slice::<f32>(),
cfg_nonlinear(),
),
);
let (qf, vf, ff) = (q_h.clone(), v_h.clone(), fwd.clone());
let lk = move |x: &[f32], _s: &[usize]| ff(&qf, x, &vf);
report(
"attn(batched) dK",
&grad_check(
&k_h,
&[bh, seq, hd],
&lk,
dk.as_slice::<f32>(),
cfg_nonlinear(),
),
);
let (qf, kf, ff) = (q_h.clone(), k_h.clone(), fwd.clone());
let lv = move |x: &[f32], _s: &[usize]| ff(&qf, &kf, x);
report(
"attn(batched) dV",
&grad_check(
&v_h,
&[bh, seq, hd],
&lv,
dv.as_slice::<f32>(),
cfg_linear(),
),
);
}
// ---- fused FLASH causal attention (the T14 op) ----
// Same structure + dimensions as attention_batched_bwd (bh=2,seq=5,hd=6), but
// exercises ops::flash_attention. Grad-check dq/dk/dv against finite-diff of
// L=sum(W∘out). This is the SINGLE-tile regime (seq<FA_TILE=32), matching the
// trusted composed grad-check's clean near-zero behavior; the MULTI-tile online-
// softmax path (seq>FA_TILE) is validated against the already-grad-checked
// composed backward by `flash_bwd_matches_composed_bwd` (seq=40) — sharper than
// finite-diff, which is unreliable on the near-zero grad elements a long softmax
// produces.
#[test]
fn flash_attention_batched_bwd() {
require_gpu();
let (bh, seq, hd) = (2, 5, 6);
let n = bh * seq * hd;
let scale = 1.0 / (hd as f32).sqrt();
// Scale Q/K up so the softmax is non-uniform (sharper attention) → the dQ/dK
// gradients are well-conditioned, not the near-zero saddle values a uniform
// softmax produces (those make central finite-diff give spurious 0.0 / sign
// flips that aren't backward bugs — cf. flash_bwd_matches_composed_bwd).
let q_h: Vec<f32> = fill(n, 241).iter().map(|v| v * 2.5).collect();
let k_h: Vec<f32> = fill(n, 242).iter().map(|v| v * 2.5).collect();
let v_h = fill(n, 243);
let w = fill(n, 244);
let q = Var::leaf(cuda(&q_h, &[bh, seq, hd]));
let k = Var::leaf(cuda(&k_h, &[bh, seq, hd]));
let v = Var::leaf(cuda(&v_h, &[bh, seq, hd]));
let out = ops::flash_attention(&q, &k, &v, scale);
scalar_loss(&out, &w).backward();
let dq = q.grad().unwrap().to_device(Device::Cpu);
let dk = k.grad().unwrap().to_device(Device::Cpu);
let dv = v.grad().unwrap().to_device(Device::Cpu);
let fwd = move |qh: &[f32], kh: &[f32], vh: &[f32]| -> f32 {
let qv = cuda(qh, &[bh, seq, hd]);
let kv = cuda(kh, &[bh, seq, hd]);
let vv = cuda(vh, &[bh, seq, hd]);
let (o, _) = qv.flash_attention(&kv, &vv, scale);
weighted_sum(&o, &w)
};
// Attention dQ/dK carry softmax curvature; for the small grad magnitudes here
// a larger eps (2e-3) cuts the f32 rounding term (∝|L|/eps) that dominates the
// O(eps²) truncation on a ~4e-4 grad. (dV is exactly linear → cfg_linear.)
let cfg_attn = GradCheckConfig {
eps: 2e-3,
rel_tol: 3e-2,
atol: 1e-3,
};
let (kf, vf, ff) = (k_h.clone(), v_h.clone(), fwd.clone());
let lq = move |x: &[f32], _s: &[usize]| ff(x, &kf, &vf);
report(
"flash dQ",
&grad_check(&q_h, &[bh, seq, hd], &lq, dq.as_slice::<f32>(), cfg_attn),
);
let (qf, vf, ff) = (q_h.clone(), v_h.clone(), fwd.clone());
let lk = move |x: &[f32], _s: &[usize]| ff(&qf, x, &vf);
report(
"flash dK",
&grad_check(&k_h, &[bh, seq, hd], &lk, dk.as_slice::<f32>(), cfg_attn),
);
let (qf, kf, ff) = (q_h.clone(), k_h.clone(), fwd.clone());
let lv = move |x: &[f32], _s: &[usize]| ff(&qf, &kf, x);
report(
"flash dV",
&grad_check(
&v_h,
&[bh, seq, hd],
&lv,
dv.as_slice::<f32>(),
cfg_linear(),
),
);
}
// flash forward must equal the composed attention forward (same SDPA math).
#[test]
fn flash_matches_composed_fwd() {
require_gpu();
let (bh, seq, hd) = (2, 40, 16);
let n = bh * seq * hd;
let scale = 1.0 / (hd as f32).sqrt();
let q = cuda(&fill(n, 341), &[bh, seq, hd]);
let k = cuda(&fill(n, 342), &[bh, seq, hd]);
let v = cuda(&fill(n, 343), &[bh, seq, hd]);
let (oc, _) = q.attention(&k, &v, scale);
let (of, _) = q.flash_attention(&k, &v, scale);
let oc = oc.to_device(Device::Cpu);
let of = of.to_device(Device::Cpu);
let max_rel = oc
.as_slice::<f32>()
.iter()
.zip(of.as_slice::<f32>())
.map(|(c, f)| (c - f).abs() / (c.abs() + 1e-6))
.fold(0.0f32, f32::max);
println!("flash-vs-composed fwd max rel: {max_rel:.3e}");
assert!(
max_rel < 1e-4,
"flash fwd diverges from composed: {max_rel:.3e}"
);
}
// flash backward must equal the (already grad-checked) composed backward. This is
// a sharper test than finite-diff: both share the trusted composed forward as the
// reference, so it isolates the flash bwd dQ/dK/dV math from finite-diff noise on
// near-zero gradient elements.
#[test]
fn flash_bwd_matches_composed_bwd() {
require_gpu();
let (bh, seq, hd) = (2, 40, 16);
let n = bh * seq * hd;
let scale = 1.0 / (hd as f32).sqrt();
let q_h = fill(n, 441);
let k_h = fill(n, 442);
let v_h = fill(n, 443);
let w = fill(n, 444);
let run = |flash: bool| -> (Vec<f32>, Vec<f32>, Vec<f32>) {
let q = Var::leaf(cuda(&q_h, &[bh, seq, hd]));
let k = Var::leaf(cuda(&k_h, &[bh, seq, hd]));
let v = Var::leaf(cuda(&v_h, &[bh, seq, hd]));
let out = if flash {
ops::flash_attention(&q, &k, &v, scale)
} else {
ops::attention(&q, &k, &v, scale)
};
scalar_loss(&out, &w).backward();
let g = |x: &Var| {
x.grad()
.unwrap()
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec()
};
(g(&q), g(&k), g(&v))
};
let (cq, ck, cv) = run(false);
let (fq, fk, fv) = run(true);
let maxrel = |a: &[f32], b: &[f32]| -> f32 {
a.iter()
.zip(b)
.map(|(x, y)| (x - y).abs() / (x.abs() + y.abs() + 1e-4))
.fold(0.0f32, f32::max)
};
let (rq, rk, rv) = (maxrel(&cq, &fq), maxrel(&ck, &fk), maxrel(&cv, &fv));
println!("flash-vs-composed bwd max rel: dQ {rq:.3e} dK {rk:.3e} dV {rv:.3e}");
assert!(rq < 2e-2, "dQ diverges: {rq:.3e}");
assert!(rk < 2e-2, "dK diverges: {rk:.3e}");
assert!(rv < 2e-2, "dV diverges: {rv:.3e}");
}
// ---- GQA repeat_kv head broadcast (Phase T15) ----
//
// repeat_kv expands K/V from [batch·num_kv, seq, hd] to [batch·nh, seq, hd]; each
// kv head is broadcast to its `group = nh/num_kv` query heads. The forward is a
// gather (a linear map), so finite-diff is clean. The CRITICAL gate is the
// BACKWARD: a kv head receives the SUM of the `group` query heads sharing it —
// the multi-group-to-one grad accumulation GQA correctness hinges on. We grad-check
// din against finite-diff of L = sum(W∘out) with group>1, plus assert the forward
// actually broadcasts and that group==1 is exact identity.
#[test]
fn repeat_kv_grad() {
require_gpu();
// batch 2, num_kv 2 → bh_kv 4 input rows; nh 6 → group 3, bh_q 12 output rows.
let (batch, num_kv, nh, seq, hd) = (2usize, 2usize, 6usize, 4usize, 5usize);
let n_in = batch * num_kv * seq * hd;
let n_out = batch * nh * seq * hd;
let x_h = fill(n_in, 711);
let w = fill(n_out, 712);
let kv = Var::leaf(cuda(&x_h, &[batch * num_kv, seq, hd]));
let out = ops::repeat_kv(&kv, nh, batch);
assert_eq!(out.value().shape(), &[batch * nh, seq, hd]);
// Forward sanity: out head (b·nh + qh) must equal in head (b·num_kv + qh/group).
let group = nh / num_kv;
let out_h = out
.value()
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec();
let row = seq * hd;
for b in 0..batch {
for qh in 0..nh {
let kvh = qh / group;
let o0 = (b * nh + qh) * row;
let i0 = (b * num_kv + kvh) * row;
for e in 0..row {
assert_eq!(out_h[o0 + e], x_h[i0 + e], "repeat_kv fwd mismatch");
}
}
}
scalar_loss(&out, &w).backward();
let din = kv.grad().unwrap().to_device(Device::Cpu);
let fwd = move |xh: &[f32], _s: &[usize]| -> f32 {
let kv = cuda(xh, &[batch * num_kv, seq, hd]);
let o = kv.repeat_kv(nh, batch);
weighted_sum(&o, &w)
};
// repeat_kv is exactly linear (gather/sum), so the linear-op tolerances apply.
report(
"repeat_kv din",
&grad_check(
&x_h,
&[batch * num_kv, seq, hd],
&fwd,
din.as_slice::<f32>(),
cfg_linear(),
),
);
}
// group==1 (num_kv == nh) must be a bit-exact identity in BOTH directions — this is
// the regression guard that makes the MHA path (kv_heads == n_heads) unchanged.
#[test]
fn repeat_kv_identity_group1() {
require_gpu();
let (batch, nh, seq, hd) = (2usize, 3usize, 4usize, 5usize);
let n = batch * nh * seq * hd;
let x_h = fill(n, 721);
let w = fill(n, 722);
let kv = Var::leaf(cuda(&x_h, &[batch * nh, seq, hd]));
let out = ops::repeat_kv(&kv, nh, batch); // group 1
let out_h = out
.value()
.to_device(Device::Cpu)
.as_slice::<f32>()
.to_vec();
assert_eq!(out_h, x_h, "group-1 repeat_kv fwd must be identity");
scalar_loss(&out, &w).backward();
let din = kv.grad().unwrap().to_device(Device::Cpu);
// dL/din = w exactly (identity forward → grad passes through unchanged).
for (g, expect) in din.as_slice::<f32>().iter().zip(&w) {
assert_eq!(*g, *expect, "group-1 repeat_kv bwd must be identity");
}
}
// ---- dropout (Phase T18) ----
//
// Fixed-seed finite-diff grad-check. Under a fixed `seed` the mask is constant
// (it depends only on (seed, index), NOT on x), so dropout is a fixed elementwise
// linear map `out_i = c_i·x_i` and the central difference of L is differentiable:
// the ± perturbation of each x_i sees the SAME mask. The forward function in the
// closure calls `ops::dropout(x, p, SEED)` with the same SEED, so it reproduces
// the same mask both times.
#[test]
fn dropout_bwd() {
require_gpu();
const SEED: u64 = 0xD120_FE5E;
let p = 0.3f32;
let (m, n) = (16, 12);
let x_h = fill(m * n, 71);
let w = fill(m * n, 72);
let x = Var::leaf(cuda(&x_h, &[m, n]));
let out = ops::dropout(&x, p, SEED);
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]| {
let o = ops::dropout(&Var::leaf(cuda(v, s)), p, SEED);
weighted_sum(&o.value(), &wf)
};
report(
"dropout dX",
&grad_check(&x_h, &[m, n], &lx, dx.as_slice::<f32>(), cfg_linear()),
);
}
// Inverted-dropout expectation + keep-rate check. Over a large tensor and a sweep
// of seeds, the mean of dropout(x) tracks the mean of x (E[out] ≈ x, the inverted
// 1/(1-p) scaling), and the kept fraction tracks 1-p (the RNG is ~Bernoulli).
#[test]
fn dropout_expectation_and_keep_rate() {
require_gpu();
let p = 0.25f32;
let n = 200_000usize;
let x_h = vec![1.0f32; n]; // mean(x) = 1 → mean(out) should ≈ 1
let x = cuda(&x_h, &[n]);
let trials = 8;
let mut mean_out_acc = 0.0f64;
let mut keep_acc = 0.0f64;
for t in 0..trials {
let (out, mask) = x.dropout(p, 0x5EED_0000 + t as u64);
let out_h = out.to_device(Device::Cpu);
let mask_h = mask.to_device(Device::Cpu);
let mean_out: f64 = out_h
.as_slice::<f32>()
.iter()
.map(|&v| v as f64)
.sum::<f64>()
/ n as f64;
let kept = mask_h
.as_slice::<f32>()
.iter()
.filter(|&&m| m != 0.0)
.count();
mean_out_acc += mean_out;
keep_acc += kept as f64 / n as f64;
}
let mean_out = mean_out_acc / trials as f64;
let keep_rate = keep_acc / trials as f64;
println!(
"dropout p={p}: E[out]={mean_out:.5} (input mean 1.0), keep_rate={keep_rate:.5} (1-p={:.3})",
1.0 - p
);
assert!(
(mean_out - 1.0).abs() < 0.01,
"E[out] {mean_out} not ≈ input mean 1.0 (inverted scaling broken)"
);
assert!(
(keep_rate - (1.0 - p) as f64).abs() < 0.01,
"keep_rate {keep_rate} not ≈ 1-p {}",
1.0 - p
);
}
// p=0 is a no-op (the op returns x.clone(), no node) → output is bit-identical to
// x and its grad flows straight through (the default-graph regression guard at the
// op level; the model-level bit-identity is in xtrain-model/tests/dropout.rs).
#[test]
fn dropout_p0_is_identity() {
require_gpu();
let (m, n) = (8, 5);
let x_h = fill(m * n, 91);
let x = cuda(&x_h, &[m, n]);
let (out, _mask) = x.dropout(0.0, 12345);
let out_h = out.to_device(Device::Cpu);
for (a, b) in x_h.iter().zip(out_h.as_slice::<f32>()) {
assert_eq!(*a, *b, "p=0 dropout must be identity");
}
}
// --- test helpers ---
// Scalar loss node L = sum(W ∘ out): wraps a fixed-weight Var and reduces. We
@@ -1005,266 +544,3 @@ fn transpose_var(x: &Var) -> Var {
}),
)
}
// seq_logprob (M3 DPO): Σ log p(target) over non-ignored rows. Grad-check with a
// completion mask — rows 0,1 are -100 (prompt, contribute 0), rows 2..6 supervised.
#[test]
fn seq_logprob_bwd() {
require_gpu();
let (rows, cols) = (6usize, 9usize);
let x_h = fill(rows * cols, 202);
let targets: Vec<i32> = (0..rows)
.map(|r| if r < 2 { -100 } else { (r * 2 % cols) as i32 })
.collect();
let target = Tensor::from_slice(&targets, &[rows]).to_device(Device::Cuda(0));
let x = Var::leaf(cuda(&x_h, &[rows, cols]));
let lp = ops::seq_logprob(&x, &target);
lp.backward();
let dx = x.grad().unwrap().to_device(Device::Cpu);
// Numeric scalar = seq_logprob = −Σ per_row (per_row is 0 for ignored rows).
let tgt = targets.clone();
let lx = move |v: &[f32], s: &[usize]| {
let t = Tensor::from_slice(&tgt, &[rows]).to_device(Device::Cuda(0));
let (_, per_row) = cuda(v, s).cross_entropy(&t);
-per_row
.to_device(Device::Cpu)
.as_slice::<f32>()
.iter()
.sum::<f32>()
};
report(
"seq_logprob dX",
&grad_check(&x_h, &[rows, cols], &lx, dx.as_slice::<f32>(), cfg_nonlinear()),
);
}
// dpo_loss (M3): scalar DPO loss with the two policy logprobs as parents. Grad-check
// each parent (finite diff of softplus(−Δ)) + the degenerate points the gate pins:
// policy==reference ⇒ Δ=0, L=log2, grads ∓β/2; β=0 ⇒ grads 0.
#[test]
fn dpo_loss_bwd_and_degenerate() {
require_gpu();
let (ref_c, ref_r, beta) = (0.5f32, 0.9f32, 0.1f32);
let (pc0, pr0) = (1.2f32, 0.7f32);
let softplus = |z: f32| z.max(0.0) + (-(z.abs())).exp().ln_1p();
let pc = Var::leaf(cuda(&[pc0], &[1]));
let pr = Var::leaf(cuda(&[pr0], &[1]));
let l = ops::dpo_loss(&pc, &pr, ref_c, ref_r, beta);
l.backward();
let dpc = pc.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
let dpr = pr.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
let l_of_pc = move |v: &[f32], _s: &[usize]| softplus(-(beta * ((v[0] - ref_c) - (pr0 - ref_r))));
report("dpo_loss dpc", &grad_check(&[pc0], &[1], &l_of_pc, &[dpc], cfg_nonlinear()));
let l_of_pr = move |v: &[f32], _s: &[usize]| softplus(-(beta * ((pc0 - ref_c) - (v[0] - ref_r))));
report("dpo_loss dpr", &grad_check(&[pr0], &[1], &l_of_pr, &[dpr], cfg_nonlinear()));
// Degenerate 1: policy == reference ⇒ Δ=0 ⇒ L=log2, grads = (∓β/2).
let pc2 = Var::leaf(cuda(&[ref_c], &[1]));
let pr2 = Var::leaf(cuda(&[ref_r], &[1]));
let l2 = ops::dpo_loss(&pc2, &pr2, ref_c, ref_r, beta);
let lval = l2.value().to_device(Device::Cpu).as_slice::<f32>()[0];
l2.backward();
let d2c = pc2.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
let d2r = pr2.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
assert!((lval - 2f32.ln()).abs() < 1e-5, "L at Δ=0 must be log2, got {lval}");
assert!(
(d2c + beta * 0.5).abs() < 1e-5 && (d2r - beta * 0.5).abs() < 1e-5,
"grads at Δ=0 must be ∓β/2, got ({d2c},{d2r})"
);
// Degenerate 2: β=0 ⇒ grads 0.
let pc3 = Var::leaf(cuda(&[pc0], &[1]));
let pr3 = Var::leaf(cuda(&[pr0], &[1]));
let l3 = ops::dpo_loss(&pc3, &pr3, ref_c, ref_r, 0.0);
l3.backward();
let d3c = pc3.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>()[0];
assert!(d3c.abs() < 1e-9, "β=0 ⇒ grad 0, got {d3c}");
println!("dpo_loss OK: grad-check (dpc,dpr) + degenerate (Δ=0→log2 & ∓β/2, β=0→0)");
}
// clipped_pg_loss (M4 GRPO): per-token clipped PG + k3 KL, one completion. Grad-check
// the active (in-trust-region) path + the A=0 (KL-only) path, plus value-level
// degenerate checks (ε→∞ ⇒ vanilla PG, β=0 ⇒ no KL).
#[test]
fn clipped_pg_loss_bwd_and_degenerate() {
require_gpu();
let (rows, cols) = (6usize, 10usize);
let x_h = fill(rows * cols, 303);
// rows 0,1 masked (prompt); 2..6 supervised (completion).
let targets: Vec<i32> = (0..rows)
.map(|r| if r < 2 { -100 } else { (r * 2 % cols) as i32 })
.collect();
let mk_target = || Tensor::from_slice(&targets, &[rows]).to_device(Device::Cuda(0));
// logp_old = logπθ at the base logits ⇒ ρ≈1 (in trust region → active path).
let (_, per_row0) = cuda(&x_h, &[rows, cols]).cross_entropy(&mk_target());
let logp_old: Vec<f32> = per_row0
.to_device(Device::Cpu)
.as_slice::<f32>()
.iter()
.map(|p| -p)
.collect();
let logp_ref: Vec<f32> = logp_old.iter().map(|l| l - 0.3).collect(); // exercise KL
let (eps, beta) = (0.2f32, 0.1f32);
// Host replica of the forward loss as a function of per-row CE values.
let host_loss = {
let (tg, lo, lr) = (targets.clone(), logp_old.clone(), logp_ref.clone());
move |per_row_h: &[f32], a: f32, e: f32, b: f32| -> f32 {
let (mut pg, mut kl, mut n) = (0f32, 0f32, 0f32);
for t in 0..per_row_h.len() {
if tg[t] < 0 {
continue;
}
n += 1.0;
let lp = -per_row_h[t];
let ratio = (lp - lo[t]).exp();
let clipped = ratio.clamp(1.0 - e, 1.0 + e);
pg += (ratio * a).min(clipped * a);
let d = lr[t] - lp;
kl += d.exp() - d - 1.0;
}
let inv = if n > 0.0 { 1.0 / n } else { 1.0 };
-pg * inv + b * kl * inv
}
};
let per_row_of = |v: &[f32], s: &[usize]| {
let (_, pr) = cuda(v, s).cross_entropy(&mk_target());
pr.to_device(Device::Cpu).as_slice::<f32>().to_vec()
};
// (1) grad-check the active PG path (A>0, ρ≈1).
let adv = 0.7f32;
let x = Var::leaf(cuda(&x_h, &[rows, cols]));
let loss = ops::clipped_pg_loss(&x, &mk_target(), &logp_old, &logp_ref, adv, eps, beta);
loss.backward();
let dx = x.grad().unwrap().to_device(Device::Cpu);
let hl = host_loss.clone();
let lx = move |v: &[f32], s: &[usize]| hl(&per_row_of(v, s), adv, eps, beta);
report(
"clipped_pg dX (active)",
&grad_check(&x_h, &[rows, cols], &lx, dx.as_slice::<f32>(), cfg_nonlinear()),
);
// (2) grad-check the A=0 path (loss = β·mean KL; PG gradient must vanish).
let x0 = Var::leaf(cuda(&x_h, &[rows, cols]));
let loss0 = ops::clipped_pg_loss(&x0, &mk_target(), &logp_old, &logp_ref, 0.0, eps, beta);
loss0.backward();
let dx0 = x0.grad().unwrap().to_device(Device::Cpu);
let hl0 = host_loss.clone();
let lx0 = move |v: &[f32], s: &[usize]| hl0(&per_row_of(v, s), 0.0, eps, beta);
report(
"clipped_pg dX (A=0, KL only)",
&grad_check(&x_h, &[rows, cols], &lx0, dx0.as_slice::<f32>(), cfg_nonlinear()),
);
// (3) ε→∞ ⇒ vanilla PG (no clip): loss value == −mean(ρA) + β·mean KL.
let big = 1e9f32;
let lv = ops::clipped_pg_loss(&Var::leaf(cuda(&x_h, &[rows, cols])), &mk_target(), &logp_old, &logp_ref, adv, big, beta);
let got = lv.value().to_device(Device::Cpu).as_slice::<f32>()[0];
let pr0 = per_row_of(&x_h, &[rows, cols]);
let want = host_loss(&pr0, adv, big, beta);
assert!((got - want).abs() < 1e-4, "ε→∞ vanilla loss mismatch: {got} vs {want}");
// (4) β=0 ⇒ no KL term (loss == −mean pg only).
let lvb = ops::clipped_pg_loss(&Var::leaf(cuda(&x_h, &[rows, cols])), &mk_target(), &logp_old, &logp_ref, adv, eps, 0.0);
let gotb = lvb.value().to_device(Device::Cpu).as_slice::<f32>()[0];
let wantb = host_loss(&pr0, adv, eps, 0.0);
assert!((gotb - wantb).abs() < 1e-5, "β=0 loss mismatch: {gotb} vs {wantb}");
println!("clipped_pg_loss OK: grad-check (active + A=0) + degenerate (ε→∞ vanilla, β=0 no KL)");
}
// clipped_pg_loss_batched (M2d): N ragged completions packed + right-padded into ONE
// forward must equal the looped per-sample path Σ_s (1/N)·clipped_pg_loss_s. The
// per-row CE backward is row-local, so folding weight = 1/(N·n_s) into the batched
// op reproduces the looped gradient and weighted-sum loss bit-for-bit (f32 path).
#[test]
fn clipped_pg_loss_batched_matches_looped() {
require_gpu();
let (n, lmax, cols) = (3usize, 5usize, 10usize);
let rows = n * lmax;
let x_h = fill(rows * cols, 909);
// Per sample: row 0 = prompt (-100); rows 1..real_len = completion; rest = pad
// (-100). Different real_len ⇒ n_s = {2, 3, 1} completion rows.
let real_len = [3usize, 4, 2];
let adv_s = [0.7f32, -0.5, 0.3];
let mut targets = vec![-100i32; rows];
for s in 0..n {
for r in 1..real_len[s] {
let t = s * lmax + r;
targets[t] = ((t * 3) % cols) as i32;
}
}
let mk_target = || Tensor::from_slice(&targets, &[rows]).to_device(Device::Cuda(0));
// logp_old ≈ logπθ at base logits (ρ≈1), logp_ref offset to exercise the KL term.
let (_, per_row0) = cuda(&x_h, &[rows, cols]).cross_entropy(&mk_target());
let logp_old: Vec<f32> = per_row0
.to_device(Device::Cpu)
.as_slice::<f32>()
.iter()
.map(|p| -p)
.collect();
let logp_ref: Vec<f32> = logp_old.iter().map(|l| l - 0.3).collect();
let (eps, beta) = (0.2f32, 0.1f32);
// Per-row advantage (sample's A) + per-row weight 1/(N·n_s) (full normaliser).
let n_of = |s: usize| (0..lmax).filter(|&r| targets[s * lmax + r] >= 0).count() as f32;
let mut advantage = vec![0f32; rows];
let mut weight = vec![0f32; rows];
for s in 0..n {
let w = (1.0 / n as f32) * (1.0 / n_of(s));
for r in 0..lmax {
advantage[s * lmax + r] = adv_s[s];
weight[s * lmax + r] = w;
}
}
// Batched: one packed [R, vocab] forward + one backward.
let xb = Var::leaf(cuda(&x_h, &[rows, cols]));
let lb = ops::clipped_pg_loss_batched(
&xb, &mk_target(), &logp_old, &logp_ref, &advantage, &weight, eps, beta,
);
lb.backward();
let gb = xb.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>().to_vec();
let lb_val = lb.value().to_device(Device::Cpu).as_slice::<f32>()[0];
// Looped reference: per-sample slice → clipped_pg_loss → scale(1/N) → backward.
let mut g_ref = vec![0f32; rows * cols];
let mut loss_ref = 0f32;
for s in 0..n {
let r0 = s * lmax;
let xs_h = x_h[r0 * cols..(r0 + lmax) * cols].to_vec();
let tgt_s: Vec<i32> = targets[r0..r0 + lmax].to_vec();
let lo_s = logp_old[r0..r0 + lmax].to_vec();
let lr_s = logp_ref[r0..r0 + lmax].to_vec();
let xs = Var::leaf(cuda(&xs_h, &[lmax, cols]));
let tgt = Tensor::from_slice(&tgt_s, &[lmax]).to_device(Device::Cuda(0));
let ls = ops::clipped_pg_loss(&xs, &tgt, &lo_s, &lr_s, adv_s[s], eps, beta);
let scaled = ops::scale(&ls, 1.0 / n as f32);
scaled.backward();
let gs = xs.grad().unwrap().to_device(Device::Cpu).as_slice::<f32>().to_vec();
g_ref[r0 * cols..(r0 + lmax) * cols].copy_from_slice(&gs);
loss_ref += scaled.value().to_device(Device::Cpu).as_slice::<f32>()[0];
}
let max_g = gb
.iter()
.zip(&g_ref)
.map(|(a, b)| (a - b).abs())
.fold(0.0f32, f32::max);
assert!(
(lb_val - loss_ref).abs() < 1e-5,
"batched loss {lb_val} vs looped {loss_ref}"
);
assert!(max_g < 1e-5, "batched grad vs looped: max|Δ| = {max_g}");
println!(
"clipped_pg_loss_batched OK: loss Δ={:.2e}, grad max|Δ|={:.2e} (== looped Σ_s 1/N·pg_s)",
(lb_val - loss_ref).abs(),
max_g
);
}

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

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

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

@@ -33,13 +33,6 @@ fn main() {
.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");
}

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

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

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

@@ -13,14 +13,12 @@ unsafe extern "C" {
// --- Device ---
pub fn cudaGetDeviceCount(count: *mut i32) -> i32;
pub fn cudaSetDevice(device: i32) -> i32;
pub fn cudaGetDevice(device: *mut i32) -> i32;
pub fn cudaDeviceSynchronize() -> i32;
// --- Memory ---
pub fn cudaMalloc(devptr: *mut *mut u8, size: usize) -> i32;
pub fn cudaFree(devptr: *mut u8) -> 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 ---
pub fn cudaGetErrorString(error: i32) -> *const c_char;
@@ -126,9 +124,7 @@ unsafe extern "C" {
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.
// RoPE (rotate_half), x:[tokens,heads,head_dim], position = token index.
pub fn launch_rope_f32(
x: *const f32,
y: *mut f32,
@@ -136,54 +132,8 @@ unsafe extern "C" {
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,
@@ -191,7 +141,6 @@ unsafe extern "C" {
heads: i32,
head_dim: i32,
theta: f32,
period: i32,
s: CudaStream,
);
@@ -228,181 +177,8 @@ unsafe extern "C" {
);
}
// 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.
// cuBLAS — used ONLY as a correctness reference for the hand-written GEMM in
// tests. Declared (and linked, see build.rs) only when CUDA is compiled in.
#[cfg(not(no_cuda))]
pub type CublasHandle = *mut c_void;
@@ -426,198 +202,7 @@ unsafe extern "C" {
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,10 +1,7 @@
#[cfg(not(no_cuda))]
pub mod cublas;
pub mod device;
pub mod error;
pub mod ffi;
pub mod memory;
mod pool;
pub use error::{CudaError, Result};
pub use memory::GpuBuffer;

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

@@ -1,37 +1,18 @@
use crate::error::{self, Result};
use crate::ffi;
use crate::pool;
/// 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.
/// RAII wrapper around a GPU memory allocation. Dropping frees the memory.
pub struct GpuBuffer {
ptr: *mut u8,
len: usize,
cap: usize,
device: i32,
}
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> {
assert!(len > 0, "cannot allocate 0 bytes on GPU");
let a = pool::acquire(len)?;
Ok(Self {
ptr: a.ptr,
len,
cap: a.cap,
device: a.device,
})
let mut ptr = std::ptr::null_mut();
error::check(unsafe { ffi::cudaMalloc(&mut ptr, len) })?;
Ok(Self { ptr, len })
}
pub fn len(&self) -> usize {
@@ -65,20 +46,13 @@ impl GpuBuffer {
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 {
fn drop(&mut self) {
// Return to the device pool for reuse (no cudaFree). The pool retains
// 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);
if !self.ptr.is_null() {
unsafe { ffi::cudaFree(self.ptr) };
}
}
}

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

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

@@ -1,13 +0,0 @@
[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
View File

@@ -1,33 +0,0 @@
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
View File

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

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

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

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

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

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

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

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

@@ -1,12 +0,0 @@
[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
View File

@@ -1,26 +0,0 @@
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
View File

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

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

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

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

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

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

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

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

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

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

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

@@ -1,198 +0,0 @@
#!/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
View File

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

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

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

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

@@ -1,10 +0,0 @@
[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
View File

@@ -1,26 +0,0 @@
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
View File

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

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

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

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

@@ -1,16 +1,12 @@
//! 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.
//! T2 only needs `F32`, but the enum + `TensorDType` trait are structured so
//! half-precision types (F16/BF16) can be added later (T7 mixed precision)
//! without touching call sites.
#[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,
}
@@ -19,7 +15,6 @@ impl DType {
pub fn size_bytes(self) -> usize {
match self {
DType::F32 => 4,
DType::BF16 => 2,
DType::I32 => 4,
}
}
@@ -27,7 +22,6 @@ impl DType {
pub fn name(self) -> &'static str {
match self {
DType::F32 => "f32",
DType::BF16 => "bf16",
DType::I32 => "i32",
}
}
@@ -56,16 +50,6 @@ impl TensorDType for f32 {
}
}
impl TensorDType for half::bf16 {
const DTYPE: DType = DType::BF16;
fn to_f64(self) -> f64 {
self.to_f64()
}
fn from_f64(v: f64) -> Self {
half::bf16::from_f64(v)
}
}
impl TensorDType for i32 {
const DTYPE: DType = DType::I32;
fn to_f64(self) -> f64 {

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

@@ -99,11 +99,9 @@ impl Storage {
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.
// No device memset in T2: stage zeros from the host.
let mut buf = GpuBuffer::alloc(len_bytes)?;
buf.memset(0)?;
buf.copy_from_host(&vec![0u8; len_bytes])?;
Ok(Storage::cuda(buf, dev))
}
}

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

File diff suppressed because it is too large Load Diff

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

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

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

@@ -1,31 +0,0 @@
[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
View File

@@ -1,26 +0,0 @@
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
View File

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -1,162 +0,0 @@
//! Batched GRPO training-side forwards (post-training M2d). After M2b/M2c made the
//! rollout cheap, the GRPO **step** is dominated by the per-sample full-sequence
//! forwards: the `per_token_logp` captures (policy + reference) and the inner
//! clipped-PG `forward`/`backward`s — each a single-sequence `forward` over a short
//! ragged completion. This module packs the `N = B·G` ragged samples of a step into
//! ONE `forward_batched`, amortising the per-launch overhead across N (the same win
//! M2b gave the rollout).
//!
//! The enabling property: **right-padding is free under causal attention.** Pad each
//! ragged completion on the RIGHT to the batch's `Lmax`; a real completion row is at
//! an earlier position than the trailing pad, and causal masking forbids attending
//! forward, so its logits are bit-identical to the unpadded single-sequence forward.
//! The pad rows' own outputs are garbage but are masked out (`target = -100`).
//!
//! Both the looped (baseline) and batched paths live here so they share one source of
//! truth — `bin/bench_grpo_batch` A/Bs them (timing + a closeness gate), and the
//! per-row equivalence of the loss op is pinned by `clipped_pg_loss_batched_matches_looped`
//! in `xtrain-autodiff/tests/autograd.rs`.
#![cfg(not(no_cuda))]
use xtrain_autodiff::ops;
use xtrain_model::{TinyTransformer, ids_tensor};
use xtrain_tensor::{Device, Tensor};
/// One framed completion of a GRPO step: the next-token `(input, target)` pair
/// (prompt positions masked to `-100` in `target`), its group-relative `adv`, and the
/// per-position rollout-time / reference logprobs the clipped-PG loss needs.
pub struct PgSample {
pub input: Vec<i32>,
pub target: Vec<i32>,
pub adv: f32,
pub logp_old: Vec<f32>,
pub logp_ref: Vec<f32>,
}
// ------------------------------- looped (baseline) -------------------------------
/// Per-position `logπ(target_t)` of one framed `(input, target)` pair (= `−per_row`
/// of cross_entropy; masked positions are 0). One single-sequence forward, no grad.
pub fn per_token_logp(model: &TinyTransformer, device: Device, input: &[i32], target: &[i32]) -> Vec<f32> {
let logits = model.forward(&ids_tensor(input, device)).value();
let (_, per_row) = logits.cross_entropy(&ids_tensor(target, device));
per_row
.to_device(Device::Cpu)
.as_slice::<f32>()
.iter()
.map(|p| -p)
.collect()
}
/// One inner clipped-PG epoch the looped way: per sample, a single-sequence forward +
/// [`ops::clipped_pg_loss`] scaled by `1/N` + backward (grads accumulate on `model`'s
/// params). Returns the summed scaled loss. Caller does clip + opt.step + zero_grad.
pub fn inner_pg_step_looped(
model: &TinyTransformer,
device: Device,
batch: &[PgSample],
eps: f32,
beta: f32,
) -> f32 {
let scale = 1.0 / batch.len() as f32;
let mut total = 0f32;
for s in batch {
let logits = model.forward(&ids_tensor(&s.input, device));
let loss = ops::clipped_pg_loss(&logits, &ids_tensor(&s.target, device), &s.logp_old, &s.logp_ref, s.adv, eps, beta);
let scaled = ops::scale(&loss, scale);
total += scaled.value().to_device(Device::Cpu).as_slice::<f32>()[0];
scaled.backward();
}
total
}
// ------------------------------- batched (M2d) -----------------------------------
/// Right-pad `m` ragged `i32` rows (each `< lmax` long) to `[m*lmax]` sequence-major,
/// filling with `pad`. Used for both the id stream (pad = 0, arbitrary) and the target
/// stream (pad = −100, ignored by cross_entropy).
fn pack_i32(rows: &[&[i32]], lmax: usize, pad: i32) -> Vec<i32> {
let mut flat = vec![pad; rows.len() * lmax];
for (i, r) in rows.iter().enumerate() {
flat[i * lmax..i * lmax + r.len()].copy_from_slice(r);
}
flat
}
/// Batched [`per_token_logp`]: pack `samples` (each `(input, target)`) right-padded to
/// `Lmax`, run ONE `forward_batched(batch = N)`, and slice each sample's `logπ` back to
/// its real length. Equal to looping [`per_token_logp`] (right-pad is free under causal
/// attention), to bf16 batch-reduction tolerance. `samples` are processed in chunks of
/// `micro` (≥1) to bound the `[chunk*Lmax, vocab]` logits memory.
pub fn per_token_logp_batched(
model: &TinyTransformer,
device: Device,
samples: &[(Vec<i32>, Vec<i32>)],
micro: usize,
) -> Vec<Vec<f32>> {
let mut out = Vec::with_capacity(samples.len());
for chunk in samples.chunks(micro.max(1)) {
let m = chunk.len();
let lmax = chunk.iter().map(|(i, _)| i.len()).max().unwrap();
let ins: Vec<&[i32]> = chunk.iter().map(|(i, _)| i.as_slice()).collect();
let tgs: Vec<&[i32]> = chunk.iter().map(|(_, t)| t.as_slice()).collect();
let ids = Tensor::from_slice(&pack_i32(&ins, lmax, 0), &[m * lmax]).to_device(device);
let tgt = Tensor::from_slice(&pack_i32(&tgs, lmax, -100), &[m * lmax]).to_device(device);
let logits = model.forward_batched(&ids, m).value();
let (_, per_row) = logits.cross_entropy(&tgt);
let pr = per_row.to_device(Device::Cpu).as_slice::<f32>().to_vec();
for (i, (inp, _)) in chunk.iter().enumerate() {
let b = i * lmax;
out.push((0..inp.len()).map(|r| -pr[b + r]).collect());
}
}
out
}
/// One inner clipped-PG epoch, batched: pack the batch (in `micro`-sized chunks) and run
/// ONE `forward_batched` + [`ops::clipped_pg_loss_batched`] + backward per chunk. The
/// per-row `weight = 1/(N·n_s)` uses the GLOBAL `N = batch.len()` (not the chunk size),
/// so chunked grad-accumulation reproduces the looped `Σ_s (1/N)(1/n_s)…` exactly.
/// Returns the summed loss. Caller does clip + opt.step + zero_grad.
pub fn inner_pg_step_batched(
model: &TinyTransformer,
device: Device,
batch: &[PgSample],
eps: f32,
beta: f32,
micro: usize,
) -> f32 {
let inv_n = 1.0 / batch.len() as f32;
let mut total = 0f32;
for chunk in batch.chunks(micro.max(1)) {
let m = chunk.len();
let lmax = chunk.iter().map(|s| s.input.len()).max().unwrap();
let ins: Vec<&[i32]> = chunk.iter().map(|s| s.input.as_slice()).collect();
let tgs: Vec<&[i32]> = chunk.iter().map(|s| s.target.as_slice()).collect();
let ids = Tensor::from_slice(&pack_i32(&ins, lmax, 0), &[m * lmax]).to_device(device);
let tgt = Tensor::from_slice(&pack_i32(&tgs, lmax, -100), &[m * lmax]).to_device(device);
let mut logp_old = vec![0f32; m * lmax];
let mut logp_ref = vec![0f32; m * lmax];
let mut advantage = vec![0f32; m * lmax];
let mut weight = vec![0f32; m * lmax];
for (i, s) in chunk.iter().enumerate() {
let b = i * lmax;
let li = s.input.len();
logp_old[b..b + li].copy_from_slice(&s.logp_old);
logp_ref[b..b + li].copy_from_slice(&s.logp_ref);
let n_s = s.target.iter().filter(|&&t| t >= 0).count().max(1) as f32;
let w = inv_n / n_s; // = 1/(N · n_s)
for r in 0..lmax {
advantage[b + r] = s.adv;
weight[b + r] = w;
}
}
let logits = model.forward_batched(&ids, m);
let loss = ops::clipped_pg_loss_batched(&logits, &tgt, &logp_old, &logp_ref, &advantage, &weight, eps, beta);
total += loss.value().to_device(Device::Cpu).as_slice::<f32>()[0];
loss.backward();
}
total
}

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

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

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

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

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

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

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

@@ -1,203 +0,0 @@
#!/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
View File

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

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

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

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

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

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

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

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

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

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

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

281
csrc/ops/flash_attention.cu
View File

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

88
csrc/ops/model.cu
View File

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

125
csrc/ops/nn.cu
View File

@@ -215,20 +215,14 @@ void launch_silu_dx_f32(const float* x, const float* dy, float* dx, int n, void*
// dx[i+h] = dy[i+h]*cos - dy[i]*sin
// =====================================================================
// `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) {
__global__ void rope_k(const float* x, 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 = tok % period;
float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
float angle = (float)pos * freq;
float angle = (float)tok * 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];
@@ -236,112 +230,20 @@ __global__ void rope_k(const float* x, float* y, int heads, int head_dim,
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) {
int head_dim, float theta, 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_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta);
}
// RoPE at an absolute position offset (KV-cache decode-time, forward only). Same
// rotate_half as rope_k, but row `tok`'s position is `pos0 + tok` (no modulo) —
// a single new decode token sits at absolute position pos0. The training rope_k
// (position = tok % period) is left untouched, so this adds no training-path risk.
__global__ void rope_at_k(const float* x, float* y, int heads, int head_dim,
float theta, int pos0) {
__global__ void rope_dx_k(const float* dy, float* dx, 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 = 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 angle = (float)tok * 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];
@@ -349,10 +251,10 @@ __global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim,
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) {
int head_dim, float theta, 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);
rope_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(dy, dx, heads, head_dim, theta);
}
// =====================================================================
@@ -428,7 +330,7 @@ __global__ void cross_entropy_fwd_k(const float* x, const int* target,
for (int c = threadIdx.x; c < cols; c += blockDim.x) pr[c] *= inv;
if (threadIdx.x == 0) {
int t = target[r];
loss[r] = t < 0 ? 0.0f : -logf(pr[t]);
loss[r] = -logf(pr[t]);
}
}
void launch_cross_entropy_fwd_f32(const float* x, const int* target,
@@ -444,13 +346,8 @@ __global__ void cross_entropy_dx_k(const float* probs, const int* target,
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= rows * cols) return;
int r = i / cols, c = i % cols;
int t = target[r];
if (t < 0) {
dx[i] = 0.0f;
} else {
float g = probs[i] - (c == t ? 1.0f : 0.0f);
dx[i] = g * scale;
}
float g = probs[i] - (c == target[r] ? 1.0f : 0.0f);
dx[i] = g * scale;
}
void launch_cross_entropy_dx_f32(const float* probs, const int* target,
float* dx, int rows, int cols, float scale, void* s) {

86
csrc/ops/optim.cu
View File

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

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

File diff suppressed because it is too large Load Diff

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

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

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

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

139
docs/06-performance.md
View File

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

157
docs/07-distributed.md
View File

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

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

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

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

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

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

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

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

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

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

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

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

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

180
docs/14-gqa.md
View File

@@ -1,180 +0,0 @@
# Phase T15: Grouped-Query Attention (GQA) — Design Document
## Goal
到 T14 为止，xtrain 的 attention 都是 **MHA**（`num_kv_heads = num_heads`）——每个
query 头有自己独立的 K/V 头。导出 xserv 时 `num_key_value_heads = num_attention_heads`
（退化 GQA，docs/08）。
T15 做**真正的 grouped-query attention**：`num_kv_heads < num_heads`，K/V 只投影到
`num_kv_heads · head_dim`，每个 KV 头被一组 `group = num_heads / num_kv_heads` 个 query
头**共享**（repeat_kv / broadcast）。GQA 是现代 LLM（Llama-2-70B、Qwen2/3、Mistral）的标配
——它把 KV cache 显存（推理）与 K/V 投影参数（训练）按 `group` 倍压缩，几乎不掉质量。
**硬闸门是诚实正确性**，重点在 **repeat_kv 的反向梯度累加**：一个 KV 头被 `group` 个 query 头
共享，反向时这 `group` 个 query 头各自对该 KV 头的梯度必须**正确求和**回到那一个共享 KV 头上。
这条「多组 q 头梯度汇到一个 kv 头」的累加路径是本任务最易错处，单列为首要 grad-check 闸门。
GQA 必须**同时**接进 T14 的 fused flash kernel（优先）与旧 composed/batched SDPA 路径，且
`num_kv_heads == num_heads`（`group = 1`）时与现有 MHA 路径**逐位一致**（回归保护）。
## 什么是 GQA
MHA：`num_heads` 个 query 头，每个配一个独立 K/V 头。
MQA（multi-query）：所有 query 头共享**一个** K/V 头（极端）。
GQA：折中——`num_kv_heads` 个 K/V 头，每个被 `group = num_heads/num_kv_heads` 个相邻 query
头共享。`num_kv_heads = num_heads` 退化为 MHA，`num_kv_heads = 1` 退化为 MQA。
```
num_heads = 8, num_kv_heads = 2 → group = 4
q heads: 0 1 2 3 4 5 6 7
kv heads: 0 0 0 0 1 1 1 1 # q head qh 用 kv head qh/group（相邻分组，连续）
```
**分组约定必须与 xserv repeat_kv 一致**（闭环命门）：xserv 的 `repeat_kv`
（`crates/xserv-model/src/qwen3.rs`）把 kv 头 `kvh` 复制到目标头 `dst = kvh*group + r`
（`r∈[0,group)`），即**query 头 `qh` 读 kv 头 `qh/group`，组内 query 头连续**。xtrain 的
repeat_kv 用同一映射，否则导出的 `q_proj` 行块与 kv 头对不上 → 闭环必崩。
## Module Layout（surgical：复用已验证的两条 SDPA，GQA = 头维 broadcast op）
```
csrc/ops/repeat_kv.cu # 新：repeat_kv fwd（头块 gather）+ bwd（组内 group 行求和，无 atomic，确定性）
crates/xtrain-cuda/
├── src/ffi.rs # +launch_repeat_kv_fwd_f32 / _bwd_f32 声明（no_cuda 门控）
└── build.rs # +repeat_kv.cu
crates/xtrain-tensor/src/tensor.rs # +Tensor::repeat_kv / repeat_kv_backward（[B*kvh,S,hd]→[B*nh,S,hd]；bf16 upcast→f32→downcast）
crates/xtrain-autodiff/
├── src/ops.rs # +ops::repeat_kv 节点（fwd broadcast，bwd 组内求和）
└── tests/autograd.rs # +repeat_kv grad-check（含 group>1 的多组梯度累加）
crates/xtrain-model/
├── src/config.rs # +num_kv_heads 字段（默认 = n_heads → MHA）；from_arch 加形参；num_params 计 K/V 投影按 kv_dim
├── src/model.rs # wk/wv 投影到 kv_dim；attention() 在 SDPA 前对 K/V 做 ops::repeat_kv；两条路径都吃到 GQA
└── tests/gqa.rs # 新：GQA(group>1) flash==composed + group=1 与 MHA 逐位一致
crates/xtrain-train/src/bin/train.rs # +--kv-heads flag
crates/xtrain-distributed/src/bin/train_ddp.rs # +--kv-heads flag（DDP 路径）
crates/xtrain-train/src/bin/export_safetensors.rs # +--kv-heads；config.json 写真 num_key_value_heads
crates/xtrain-model/tests/parity{.py,_dump.rs} # PyTorch 对拍加 GQA（kv 投影 + repeat_kv）
```
## Key Design Decisions
### ① GQA = K/V 头维 broadcast op，喂给**未改动**的两条 SDPA（不写第三套 attention）
T14 已经有两条**逐位/数值都验证过**的 SDPA：composed（`ops::attention`）与 fused flash
（`ops::flash_attention`），二者都吃 `[bh, S, hd]`（`bh = batch·heads`）。GQA 的本质只是「K/V
比 Q 少 `group` 倍头，用前把每个 kv 头复制 `group` 份」。所以**最外科**的做法：
- wk/wv 投影到 `kv_dim = num_kv_heads · head_dim`，按 `[B, num_kv, S, hd] → [B·num_kv, S, hd]`
排好（和 Q 的 `[B·nh, S, hd]` 同流水线，只是头数不同）。
- 在调 SDPA 之前，对 K、V 各做一个新 autograd op `ops::repeat_kv`，把 `[B·num_kv, S, hd]`
**broadcast** 成 `[B·nh, S, hd]`（输出行 `b·nh + qh` = 输入行 `b·num_kv + qh/group` 的拷贝）。
- 之后 `ops::attention` / `ops::flash_attention` **一字不改**——它们看到的就是满头的
`[B·nh, S, hd]`，GQA 对两条路径**同时、免费**生效。flash kernel / composed kernel 都不用碰。
**为什么不在 kernel 内做 GQA**：那要给 flash fwd/bwd 两个 kernel 各加 kv-head 索引、给 composed
的两次 strided GEMM 各算 GQA stride，且两套都要重测——是「第三套 attention 改动」。而 broadcast-op
方案：(a) 两条 SDPA 路径零改动、其 T14 闸门不回归；(b) repeat_kv 的 fwd/bwd 是独立可 grad-check
的小 op，正确性风险隔离在一处；(c) 关键的「多组 q 头梯度汇到一个 kv 头」就是 repeat_kv 的**反向**，
干净地落在一个 op 上单测。代价是 K/V 在显存里被物化成满头 `[B·nh,S,hd]`（多 `group` 倍）——本规模
（训练、seq 不极端）可接受；真要省这份显存是 follow-up（kernel 内 GQA 读取），记进逃生舱不在 T15 做。
> 备选（不采纳）：flash/composed kernel 内直接按 `kv_head = q_head/group` 索引 K/V。省 broadcast
> 物化，但动两套已验证 kernel + 重写两套 backward 的 kv 累加，违反「不写第三套 attention」与回归保护。
> escape hatch：先 broadcast-op 把正确性 + 闭环钉死，kernel-内 GQA（省显存）留 follow-up。
### ② repeat_kv 的反向 = 组内求和（确定性，无 atomic）
`repeat_kv` 前向：`out[b·nh + qh] = in[b·num_kv + qh/group]`（按 `S·hd` 整行拷贝）。
反向是它的**转置**：一个 kv 头收到它那 `group` 个 query 头的梯度之**和**：
```
din[b·num_kv + kvh] = Σ_{r=0}^{group-1} dout[b·nh + kvh·group + r]
```
这正是闸门要求的「多组 q 头梯度累加到一个 kv 头」。实现上**不用 atomicAdd**：每个输入
（kv-head, 元素）由唯一一个 block 负责，它**串行累加自己那 group 个连续源行**——天然 race-free
且**run-to-run 确定**（不像 flash bwd 的跨行 atomic 反向有归约序不确定问题）。`group=1` 时
反向退化为单行拷贝（identity）。
autograd 层面其实也可以靠引擎的扇出 SUM（把一个 kv Var 喂给 group 个下游），但那样图里要
显式建 group 份 view、且 flash/composed 的 batched 布局不是按头切的——做成一个专门的
broadcast op，fwd/bwd 各一发 kernel，最简且能单独 grad-check。
### ③ `num_kv_heads` 进 Config（它改模型尺寸/导出），默认 = n_heads → 退化 MHA
不同于 T14 的 `use_flash`（运行时旗标，不进 Config），`num_kv_heads` **改 K/V 投影的形状、改参数量、
改导出的 `num_key_value_heads`**——它是**架构**的一部分，必须进 `Config` 并落进 checkpoint/导出。
- `Config` 加 `num_kv_heads: usize`；`from_arch` 加该形参；`Config::tiny()` 默认 `num_kv_heads =
n_heads`（MHA）。约束：`num_heads % num_kv_heads == 0`（断言）。
- `num_params()`：K/V 投影从 `2·dim·dim` 改成 `2·dim·(num_kv_heads·head_dim)`；QK-norm 的
`k_norm` 仍是 `[head_dim]`（per-head，作用在单个 head 向量上，与头数无关）→ 不变。
- **`num_kv_heads == n_heads` 时 `group=1`**：`ops::repeat_kv` 是 identity（fwd 单行拷贝、bwd 单行
拷贝），wk/wv 形状回到 `[dim,dim]` → 整条图与 T14 的 MHA 路径**逐位一致**（回归保护闸门）。
> wk/wv 形状从 `[dim,dim]` 变成 `[dim, kv_dim]`：`Block` 里 wk/wv 的 `mk(&[dim, kv_dim])`，
> `params()`/`block_params()` 顺序不变（还是 attn_norm,wq,wk,wv,q_norm,k_norm,wo,...），只是
> wk/wv 的 shape 跟着 Config。导出转置照旧按各自 shape 走（`transpose` 读 `v.value().shape()`）。
### ④ bf16 / recompute / dropout / DDP 全部自动兼容
- **bf16**：`Tensor::repeat_kv` 沿用全 repo 一致的 cast 策略——bf16 入则 upcast f32 → kernel →
downcast；kernel 只一份 f32。`ops::repeat_kv` 的 fwd/bwd 都在 SDPA 之前/之后，dtype 与 K/V 流一致。
- **recompute（T13）**：repeat_kv 在 `block_forward` 内、`attention()` 里，重算段重跑 `attention()`
自然重跑 repeat_kv（无额外状态，确定性）→ 梯度仍逐位一致。
- **dropout（T18）**：dropout 接在 attn/mlp 子块**输出**，与 attention 内部的 repeat_kv 正交，不交互。
- **DDP**：repeat_kv 不引入新参数；wk/wv 变小（kv_dim）只是参数张量小一圈，`params()` 泛化迭代
+ all-reduce 照旧；跨 rank 一致性不受影响。
### ⑤ 导出 xserv：写真 `num_key_value_heads`，分组约定对齐 repeat_kv
`export_safetensors.rs` 的 `config.json` 把 `num_key_value_heads` 从「= num_attention_heads」改成
**真 `cfg.num_kv_heads`**；`--kv-heads` flag 传入（须与训练 ckpt 一致）。q/k/v_proj 各自按其 shape
转置导出（k/v_proj 现在是 `[kv_dim, dim]`，xserv loader 期望的 GQA 形状）。xserv 的 `repeat_kv`
用 `dst = kvh·group + r` 分组，与 ① 的 xtrain 约定**逐头对齐** → 同一份权重在两侧前向数学一致，
闭环（贪心逐 token 一致）成立。
## 验证方法
全部 `#![cfg(not(no_cuda))]` 门控，本地 `cargo check`/`fmt`，构建+实跑在 dash5（8× RTX 5090）。
### 1. 正确性（硬闸门全绿，dash5 实跑 capture）
- **repeat_kv finite-diff grad-check**（`autograd.rs::repeat_kv_grad`）：**核心闸门**——`group>1`
（如 bh: 2 kv 头 → 6 q 头）下 grad-check `din`，验证「多组 q 头梯度求和到一个 kv 头」的反向。
外加 `group=1` identity 自检。
- **GQA flash==composed**（`gqa.rs`）：真 GQA 配置（`num_kv_heads < n_heads`，如 8 头/2 kv 头）下，
flash on/off 两个同 init 模型，断 forward logits / loss / **每参数梯度**一致（fp32 紧容差 + bf16
舍入带）——尤其 wk/wv 的梯度（它们经过 repeat_kv 反向的组内求和）。
- **group=1 与 MHA 逐位一致**（`gqa.rs`）：`num_kv_heads = n_heads` 的模型对 T14 的 MHA 模型，
forward + 每参数梯度 `|Δ|=0`（回归保护）。
- **PyTorch GQA 对拍 B>1**（`parity_dump.rs` + `parity.py`）：等价 PyTorch 模型加 GQA（k/v 投影到
kv_dim + `repeat_interleave(group)` 分组，与 xserv/xtrain 约定一致），对拍 forward logits + 全部
参数梯度（composed 与 flash 两条都跑，共用同一 oracle）。
- **小 GQA 配置短训收敛**：一个真 GQA 小模型短训，loss 单调降、无 NaN、采样连贯。
- **全回归套开/关**：autograd / structural / batched==looped / bf16 / recompute（逐位）/ overfit 27/27 /
AdamW（GPU bit-exact + host 对 torch）/ DDP loss-match + 跨 rank（`--test-threads=1`）/ flash /
grad_accum / dropout / **xserv 闭环 md5**。MHA 默认（kv=heads）图不变 → 不回归。
### 2. 闭环（payoff）—— 真 GQA 端到端
导出一个 `num_key_value_heads < num_attention_heads` 的 GQA checkpoint → xserv 加载 → 贪心生成
**对 xtrain 自身逐 token 一致**（BF16 推理 vs f32 训练，与 v1–v8 同款判据）。这是 GQA 真正落地的证明：
训练侧的分组、导出的分组、推理侧 xserv 的 repeat_kv 分组三方对齐。
## 实测结果（dash5 1× / 2× RTX 5090）
**硬闸门全绿：**
| 闸门 | 结果 |
|---|---|
| ① repeat_kv grad-check（**多组 q 头梯度求和到一个 kv 头**，group=3） | **过** — din max_rel **2.05e-4**；group=1 identity 双向**逐位**（fwd/bwd |Δ|=0） |
| GQA flash==composed（model 级 8h/2kv，logits/loss/每参数梯度） | fp32: loss rel **0.0**、logits 3.0e-4、grad **4.1e-5**；bf16: loss 9.0e-5、logits mean 2.9e-3/p99 1.0e-2、grad scaled-mean 8.9e-3 |
| group=1 对 MHA**逐位一致**（回归保护） | **过** — logits + loss + 全部梯度 |Δ|=0 |
| ② PyTorch GQA 对拍 B>1（composed & flash，repeat_interleave 分组对齐） | composed: loss **1.74e-8**/logits 2.04e-5/25 grad 进 rtol；flash: loss 1.74e-8/logits 2.28e-5/25 grad 进 rtol |
| ③ 小 GQA 配置短训收敛（8h/2kv/hd32/4L/ffn1024，600 步） | train **10.90→3.15** 无 NaN、gnorm 稳 ~1.2、采样连贯英文（~200K tok/s） |
| ④ **xserv 闭环真 GQA**（导出 `num_key_value_heads=2 < num_attention_heads=8`，xserv 加载 `heads=8/2 kv`，贪心） | "One day"/"The little" 两 prompt **逐 token 一致**；"Once upon a time" 在 `...Lily's mommy ` 处 BF16 漂移晚分叉（said vs came）——与 v1/v2/v3/T14 同款判据 |
| ⑤ 回归套：autograd 23（含 repeat_kv 2）/ structural 5 / batched / bf16 / flash 2 / **gqa 4** / overfit 27/27 / recompute 2 / dropout 6 / grad_accum 3 / checkpoint-roundtrip / AdamW(host 对 torch 4.8e-6) / DDP 3(`--test-threads=1`, loss 5.67e-7+跨 rank 一致) / GEMM / tensor | **全绿** |
| ⑤ MHA 默认 export md5（v3 ckpt 用 T15 代码重导 safetensors） | **逐位一致** `b04fc9f9a0c9af04c47d9ca649aea12e`（与 registry/T14 同）→ 默认（kv=heads）export 零漂移 |
> **诚实记录**：闭环 2/3 prompt 完全 token-identical、1/3 在 BF16 漂移点晚分叉——这恰证明 GQA 分组**正确**：若 kv→q 头映射错，attention 会从第一个生成 token 起就崩（不会是深处近-tie 的晚分叉）。GQA 把 K/V 在显存里物化成满头 `[B·nh,S,hd]`（broadcast-op 方案的代价）——本规模可接受，kernel-内 GQA（省这份显存）留 follow-up。未为凑绿放宽任何容差。

165
docs/15-grad-accum.md
View File

@@ -1,165 +0,0 @@
# Phase T16: Gradient Accumulation — Design Document
## Goal
在已有的训练 loop（T6/T10）与 DDP（T8）之上，加 **micro-batch 梯度累积**：把 `accum_steps=N`
个 **micro-step** 的梯度在 tape 里累加起来，再做**一次** `AdamW.step` + `zero_grad`——得到
**有效 batch = N × micro_batch** 的更新，而显存只占**一个 micro-batch** 的激活峰值（不随 N 增长）。
两条硬约束：
1. **数值等效**：`accum_steps=N`（N 个 micro-step 后一次 step）必须对住「一个 N× 大 batch
的单 step」——梯度/loss 在仓内既有容差内**逐位贴合**。这是核心等效性证明。
2. **DDP 只在累积边界通信**：`world>1` 下，N 个 micro-step 里**只在最后一个**做 all-reduce
（中间 micro-step **跳过跨卡通信**），最终喂给优化器的仍是 global 有效 batch 的均值梯度，
loss 对单卡。
并暴露 train 入口的 `--accum-steps` flag。`accum_steps=1` 必须对当前无累积路径**逐位一致**
（回归保护）。
**不做**：micro-batch 间变 LR / 变 batch（恒定 micro_batch）；累积里换 dropout RNG（T18 才有
dropout）；ZeRO（T17）。本 Phase 只动**优化器 step 的节奏**与 **DDP 通信门控**，复用 tape 既有
的 SUM 累加。
## Module Layout
```
crates/xtrain-train/src/
├── train_loop.rs # TrainConfig += accum_steps；inner micro-loop（缩放 loss + tape SUM）
└── bin/train.rs # 新 --accum-steps flag；打印有效 batch
crates/xtrain-distributed/src/
└── ddp.rs # DdpConfig += accum_steps；all-reduce 门控到累积边界
crates/xtrain-train/tests/
└── grad_accum.rs # 等效性硬闸门 + accum_steps=1 逐位回归（单卡）
crates/xtrain-distributed/tests/
└── ddp_correctness.rs # += DDP+accum 对单卡（复用既有 ddp_matches… 框架）
docs/15-grad-accum.md # 本文
```
无新 crate、无新 kernel、无新 autograd op——梯度累积是**纯调度**：tape 早已 SUM 累加，
缩放用既有 `ops::scale`，DDP 通信用既有 `all_reduce_average_grads`，只是改**调用节奏与门控**。
## Key Design Decisions
### ① 等效性的数学：缩放每个 micro-loss 为 `1/N`
模型的 `loss_batched` 是 **CE-mean over `batch*seq` 行**（见 `model.rs`）。设一个 micro-batch 有
`B` 序列、seq 长 `S`，记某 micro-step 那批 `B*S` 行的 per-row 梯度之和为 `Σ_micro`：
- **大 batch 基线**（有效 batch `N·B`）：一次 `loss_batched(N·B 序列)` = CE-mean over `N·B·S` 行
→ backward 给 `G_big = Σ_all / (N·B·S)`，其中 `Σ_all = Σ_n Σ_micro_n`。
- **累积**（N 个 micro-step，每个 `B`）：micro-step n 的 `loss_batched(B)` = CE-mean over `B·S` 行
→ 若直接 backward 得 `Σ_micro_n / (B·S)`；**N 个 backward 之间不 `zero_grad`**，tape SUM 累加 →
`Σ_n Σ_micro_n / (B·S) = Σ_all / (B·S) = N · G_big`。
差一个因子 N。修正：**每个 micro-loss 先 `ops::scale(loss, 1/N)` 再 backward**——`scale` 的
backward 把上游梯度乘 `1/N`（见 `ops.rs`），于是每个 micro 贡献 `Σ_micro_n / (N·B·S)`，
累积后 `Σ_all / (N·B·S) = G_big`，**与大 batch 逐位等效**（仅 fp 求和顺序不同 → 进容差，和
T8 DDP-vs-单卡同性质）。
> 为什么不在 clip 里用 `pre_scale=1/N`？clip 的 `pre_scale` 已被 batch-mean 占用（=1.0）。
> 在 loss 上 `scale(1/N)` 更内聚：缩放穿过既有 autograd，不碰 clip/optimizer，且 `N=1` 时
> `scale(1.0)` 的 backward 是恒等乘 1 —— 这正是 `accum_steps=1` 逐位回归的保证（见 ④）。
报告的 step-loss = N 个 micro 的**原始** loss（未缩放值）之和 / N = 有效 batch 的 mean loss，
和大 batch 的单一 mean loss 一致（同样仅求和顺序差）。
### ② 单卡 train loop：inner micro-loop
每个 optimizer step：
```text
for micro in 0..N:
抽 B 序列 → loss = loss_batched(B)
step_loss_acc += raw_loss(loss) # 累报告用的原始 loss
scale(loss, 1/N).backward() # tape SUM 累加缩放后的梯度
# —— 累积边界 ——
clip_grad_norm_gpu(params, max_norm, 1.0) # 梯度已是有效 batch 均值
opt.step(lr); zero_grad()
losses.push(step_loss_acc / N)
tokens_seen += N * B * S # 有效 batch tok
```
`accum_steps` 默认 1 → micro-loop 跑一次、`scale(loss,1.0)`、不在 micro 间 zero_grad（本就如此）
→ 与现路径完全等价。**每个 micro-step 的计算图在它自己的 backward 后即可释放**（Rust `Rc` 在
循环变量出作用域时 drop），所以**显存峰值 = 单个 micro-batch 的激活**，不随 N 增长（③ 实测）。
抽样次序保持：单卡仍是连续从 RNG 抽 `N·B` 序列；与「大 batch 抽 `N·B`」逐序列对齐，只是分 N 组
forward——并集同序，所以 `Σ_all` 的项一致。
### ③ 显存平 + 有效 batch 实测
「显存不随 N 增长」是 grad-accum 的卖点，要**实测**而非断言：固定有效 batch `E = N·B`，跑
`(N=1,B=E)`（大 batch）vs `(N=E,B=1)`（极端累积），用 `nvidia-smi`/`cudaMemGetInfo` 量峰值显存——
后者应**显著低**（少 N× 激活）。train 入口打印 `effective batch = accum_steps × batch`。
### ④ `accum_steps=1` 逐位回归
`N=1` 时 inner loop 跑一次、`scale(loss, 1.0)`。`ops::scale(_, 1.0)` 的 fwd 是
`value.scale(1.0)`、bwd 是 `grad.scale(1.0)`——数学恒等。为**绝对**逐位（连一次 `×1.0` kernel
都不引入），实现里 `N==1` 直接 `loss.backward()`（跳过 scale），与现路径**字节一致**。测试
`accum1_bit_identical_to_no_accum` 锁这条。
### ⑤ DDP：all-reduce 门控到累积边界
T8 的 `all_reduce_average_grads(params)` 每 step 调一次。grad-accum 下**只在最后一个
micro-step 之后调一次**——中间 micro-step 的 backward 只在本卡 tape 里 SUM，**不发 NCCL**。
均值的账（沿用 T8 的「通信里 /world，clip 里 /b_local」拆分，再叠加 ① 的 /N）：
```text
每卡每 micro: scale(loss, 1/N).backward() → 本卡 tape SUM 该 micro 的 (Σ_micro / N)/...
N 个 micro 后, 本卡 grad = Σ_{micro∈本卡所有micro} ... = 本卡 N·B_local 行的 (1/N) 缩放和
all-reduce(sum)+/world (累积边界一次): 跨卡求和后 /world
→ 每卡持有 Σ_global,(N·B) / (N · world · ?) # 见下：用 1/N·scale 替代单卡的 1/b
clip pre_scale = 1.0
```
精确推导：每卡每 micro 的 `loss_batched(B_local)` 是 **本卡 mean over `B_local·S` 行**。
`scale(1/N)` 后 backward = `Σ_local_micro / (N · B_local · S)`。N 个 micro tape SUM →
`Σ_local_all / (N · B_local · S)`，其中 `Σ_local_all` = 本卡 `N·B_local` 行之和。
`all_reduce(sum)` 跨 world 卡 → `Σ_global_all / (N · B_local · S)`（`Σ_global_all` = 全
`world·N·B_local = N·B_global` 行之和）；`/world` → `Σ_global_all / (N · B_local · S · world)`
`= Σ_global_all / (N · B_global · S)`（因 `B_global = world·B_local`）。这正是**有效 batch
`N·B_global` 的 mean 梯度**——与单卡「有效 batch `N·B_global` 的大 batch 单 step」逐位等效
（求和顺序差进容差）。
> 关键正确性点：`all_reduce_average_grads` 里的 `/world` 是按 **world** 缩放（与 N 无关）；N 的
> 那个 `1/N` 已由 ① 的 `scale` 在每个 micro 的 backward 里完成。两者正交，不会互相污染。
> 单卡（`world=1`）退化：all-reduce 是 no-op，`/world=1`，只剩 ① 的 `1/N` → 与 ② 一致。
DDP 报告 loss = N 个 micro 的本卡原始 loss·B_local 之和、跨卡 all-reduce(sum)、/(N·B_global)。
### ⑥ 不变量小结
| | 单卡基线（大 batch E） | 单卡 accum（N×B=E） | DDP accum（world, N×B_local·world=E） |
|---|---|---|---|
| loss 缩放 | 无（CE-mean） | 每 micro `×1/N` | 每 micro `×1/N` |
| grad 累加 | tape SUM 一批 | tape SUM N 批 | tape SUM N 批/卡 |
| 跨卡通信 | — | — | **仅累积边界 1 次** all-reduce + /world |
| clip pre_scale | 1.0 | 1.0 | 1.0 |
| 显存峰值 | E 的激活 | **B 的激活** | **B_local 的激活** |
## 验证方法（验收，全部 dash5 实跑 capture）
GPU 测试 `#[cfg(not(no_cuda))]` 门控。
1. **等效性（核心硬闸门）** `grad_accum.rs::accum_equiv_big_batch`：同 init、同数据同序，
跑「`accum_steps=N`, micro_batch=B」与「`accum_steps=1`, batch=N·B」各一 step，断言
①loss、②**每个参数的 grad** rel-err 进 fp 容差（求和顺序差，~1e-4 量级，对齐 recompute/DDP
闸门约定）。多步版（跑 K 个 optimizer step）再断言**终参**贴合（误差不发散）。
2. **`accum_steps=1` 逐位回归** `grad_accum.rs::accum1_bit_identical`：`accum_steps=1` 与现
no-accum 路径同 init/同数据 → 每参数 grad `max|Δ| == 0.0`（④ 跳过 scale，字节一致）。
3. **DDP+accum 对单卡** `ddp_correctness.rs`（扩既有 `ddp_matches_single_gpu…`）：单卡
有效 batch `E` 的大 batch baseline vs `world=2 + accum_steps=N`（每卡每 micro `B_local`，
`world·N·B_local=E`）→ loss 轨迹 `max_rel<1e-3`、跨 rank 参数一致、且 only-at-boundary 通信
（micro 间不发 NCCL，由实现保证 + 不变量推导）。
4. **显存平 + 有效 batch** ：固定有效 batch，量 `(N=1,大batch)` vs `(N=大,micro=1)` 峰值显存
（后者显著低），train 入口打印 effective batch。capture nvidia-smi。
5. **全回归套**：autograd grad-check / structural / batched==looped / bf16 / recompute（逐位）/
overfit 27/27 / AdamW（GPU bit-exact + host vs torch）/ DDP loss-match + 跨 rank / **xserv
闭环 md5**——`accum_steps=1` 默认值保证全部不回归。

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