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:
26
crates/xtrain-model/src/lib.rs
Normal file
26
crates/xtrain-model/src/lib.rs
Normal 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};
|
||||
Reference in New Issue
Block a user