Gahow Wang 0150263055 perf: KI-3 fixed — dim1024 batch32 fits, mem 31.1→14.6GB, tok/s 39.7K→31.5K
Per-block activation recompute (T13) measured on dash5 (1× RTX 5090 32GB, bf16,
batch32 seq256, steady-state):

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

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

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-17 09:50:29 +08:00
2026-06-15 17:14:56 +08:00

xtrain

A from-scratch Rust + CUDA LLM training engine — the sibling of xserv (the inference side). GPU-first.

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, ~130M 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

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.

Layout

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

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

export PATH=/usr/local/cuda/bin:$HOME/.cargo/bin:$PATH
cargo build
cargo test -p xtrain-cuda -- --nocapture   # runs the vecadd smoke test

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

Description
No description provided
Readme 3.1 MiB
Languages
Rust 87.6%
Cuda 8.7%
Python 2.2%
Shell 1.5%