Gahow Wang
77a82bfeee
Training loop (train_loop.rs): sample batch_size sequences, forward loss + backward (tape SUMs grads), clip_grad_norm with ×1/batch averaging, AdamW step with scheduled lr, zero_grad; logs loss/lr/gnorm/tok-s and checkpoints periodically; returns the loss trace. Checkpoint (checkpoint.rs): flat little-endian dump of params() in order (magic/version/count + per-param ndim/dims/f32 data); load_into validates and overwrites a matching model's params via set_value (exact f32 round-trip). Sampler (sample.rs): autoregressive greedy / temperature generation — re-runs forward on the growing prefix (model is single-sequence, RoPE pos=row). bin/train.rs: end-to-end entry — load tokenizer+corpus, train a tiny 4-layer model for a bounded budget, checkpoint, print samples. no_cuda stub keeps it buildable on a GPU-less host. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
23 lines
721 B
Rust
23 lines
721 B
Rust
//! 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;
|
|
|
|
#[cfg(not(no_cuda))]
|
|
pub mod checkpoint;
|
|
#[cfg(not(no_cuda))]
|
|
pub mod sample;
|
|
#[cfg(not(no_cuda))]
|
|
mod train_loop;
|
|
|
|
#[cfg(not(no_cuda))]
|
|
pub use train_loop::{TrainConfig, train};
|