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>
91 lines
3.0 KiB
Rust
91 lines
3.0 KiB
Rust
//! Checkpoint save/load. Dumps the model's `params()` (in their stable order) to
|
|
//! a flat binary file and reloads them into a model with matching architecture.
|
|
//!
|
|
//! Format (little-endian):
|
|
//! ```text
|
|
//! magic : u32 = 0x58545254 ("XTRT")
|
|
//! version : u32 = 1
|
|
//! n_params: u32
|
|
//! repeat n_params times:
|
|
//! ndim : u32
|
|
//! dims : [u32; ndim]
|
|
//! data : [f32; prod(dims)]
|
|
//! ```
|
|
//! Architecture/config is NOT stored here — the caller rebuilds the model from
|
|
//! the same `Config` and `load_into`s the params (the round-trip and resume both
|
|
//! know their config). Gated behind `not(no_cuda)` (it round-trips GPU tensors).
|
|
|
|
#![cfg(not(no_cuda))]
|
|
|
|
use std::fs::File;
|
|
use std::io::{BufReader, BufWriter, Read, Write};
|
|
use std::path::Path;
|
|
use xtrain_autodiff::tape::Var;
|
|
use xtrain_tensor::{Device, Tensor};
|
|
|
|
const MAGIC: u32 = 0x5854_5254;
|
|
const VERSION: u32 = 1;
|
|
|
|
/// Write every parameter (value) to `path` in `params()` order.
|
|
pub fn save(path: &Path, params: &[Var]) -> std::io::Result<()> {
|
|
let mut w = BufWriter::new(File::create(path)?);
|
|
w.write_all(&MAGIC.to_le_bytes())?;
|
|
w.write_all(&VERSION.to_le_bytes())?;
|
|
w.write_all(&(params.len() as u32).to_le_bytes())?;
|
|
for p in params {
|
|
let v = p.value().to_device(Device::Cpu);
|
|
let shape = v.shape();
|
|
w.write_all(&(shape.len() as u32).to_le_bytes())?;
|
|
for &d in shape {
|
|
w.write_all(&(d as u32).to_le_bytes())?;
|
|
}
|
|
for &x in v.as_slice::<f32>() {
|
|
w.write_all(&x.to_le_bytes())?;
|
|
}
|
|
}
|
|
w.flush()
|
|
}
|
|
|
|
/// Read a checkpoint and overwrite each parameter's value in `params` (in order).
|
|
/// Shapes must match the saved ones. Tensors are placed on each param's device.
|
|
pub fn load_into(path: &Path, params: &[Var]) -> std::io::Result<()> {
|
|
let mut r = BufReader::new(File::open(path)?);
|
|
assert_eq!(read_u32(&mut r)?, MAGIC, "bad checkpoint magic");
|
|
assert_eq!(read_u32(&mut r)?, VERSION, "unsupported checkpoint version");
|
|
let n = read_u32(&mut r)? as usize;
|
|
assert_eq!(n, params.len(), "checkpoint param count != model");
|
|
|
|
for p in params {
|
|
let ndim = read_u32(&mut r)? as usize;
|
|
let mut dims = Vec::with_capacity(ndim);
|
|
for _ in 0..ndim {
|
|
dims.push(read_u32(&mut r)? as usize);
|
|
}
|
|
let numel: usize = dims.iter().product();
|
|
let mut data = vec![0.0f32; numel];
|
|
for slot in data.iter_mut() {
|
|
*slot = read_f32(&mut r)?;
|
|
}
|
|
let device = p.value().device();
|
|
assert_eq!(
|
|
p.value().shape(),
|
|
dims.as_slice(),
|
|
"checkpoint shape mismatch"
|
|
);
|
|
p.set_value(Tensor::from_slice(&data, &dims).to_device(device));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn read_u32<R: Read>(r: &mut R) -> std::io::Result<u32> {
|
|
let mut b = [0u8; 4];
|
|
r.read_exact(&mut b)?;
|
|
Ok(u32::from_le_bytes(b))
|
|
}
|
|
|
|
fn read_f32<R: Read>(r: &mut R) -> std::io::Result<f32> {
|
|
let mut b = [0u8; 4];
|
|
r.read_exact(&mut b)?;
|
|
Ok(f32::from_le_bytes(b))
|
|
}
|