Gahow Wang
5353b38402
forward_batched(ids[B*S], batch)/loss_batched: run B equal-length sequences as ONE forward over flattened [B*S] ids, so every linear is one big [B*S,dim] GEMM. Attention reshapes to [B*nh,S,hd], runs the fused batched causal SDPA (per-seq mask + RoPE period=S, no cross-sequence attention), writes back [B*S,dim]. The old per-(batch,head) loop + host-round-tripping split/merge_heads + the additive causal_mask leaf are gone. forward(ids[seq]) is now forward_batched(ids,1), so the sampler / inference path (batch=1) is unchanged. +batched_ids_tensor helper. New batched.rs test: batched forward == looped single-sequence (logits identical 0.0, grads 6.4e-4, loss identical). PyTorch parity now exercises B>1 (B=2,S=4): loss 5e-8, logits 6.9e-6, all 25 param grads within rtol — verifying per-seq RoPE position + per-seq causal masking. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
28 lines
1.3 KiB
Rust
28 lines
1.3 KiB
Rust
//! 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};
|