model: tiny RoPE+RMSNorm+SwiGLU transformer + overfit test

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

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-15 16:05:20 +08:00
parent 0acfa5df11
commit e3912c2380
8 changed files with 466 additions and 0 deletions

View File

@@ -0,0 +1,26 @@
//! 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
//! (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()`.
//!
//! 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, ids_tensor, param_to_host};