diff --git a/crates/xtrain-train/src/bin/train.rs b/crates/xtrain-train/src/bin/train.rs index dcbe059..22ff49f 100644 --- a/crates/xtrain-train/src/bin/train.rs +++ b/crates/xtrain-train/src/bin/train.rs @@ -167,6 +167,20 @@ fn main() { } }); + // Eval-only mode: load a checkpoint and score it on the held-out val set, then + // exit. Used to put an EXISTING model (e.g. v0) and a new one on the same + // metric — the v0-vs-v1 val-loss comparison. The arch flags must match the ckpt. + if let Some(p) = args.iter().position(|a| a == "--eval-ckpt") { + let ckpt_path = PathBuf::from(args.get(p + 1).expect("--eval-ckpt ")); + xtrain_train::checkpoint::load_into(&ckpt_path, &model.params()) + .expect("load eval checkpoint"); + let v = valid.expect("--eval-ckpt needs --val-tokens > 0"); + let vl = xtrain_train::eval_loss(&model, device, &v, seq_len, eval_batches); + println!("eval-only: {} → val loss {vl:.4}", ckpt_path.display()); + sample_some(&model, device, &tok_path); + return; + } + let tcfg = TrainConfig { seq_len, batch_size, diff --git a/crates/xtrain-train/src/lib.rs b/crates/xtrain-train/src/lib.rs index 6a9c1d1..2c00d1d 100644 --- a/crates/xtrain-train/src/lib.rs +++ b/crates/xtrain-train/src/lib.rs @@ -19,4 +19,4 @@ pub mod sample; mod train_loop; #[cfg(not(no_cuda))] -pub use train_loop::{TrainConfig, TrainResult, train}; +pub use train_loop::{TrainConfig, TrainResult, eval_loss, train}; diff --git a/crates/xtrain-train/src/train_loop.rs b/crates/xtrain-train/src/train_loop.rs index 0acbcf3..0185e5a 100644 --- a/crates/xtrain-train/src/train_loop.rs +++ b/crates/xtrain-train/src/train_loop.rs @@ -153,8 +153,8 @@ pub fn train( /// Mean cross-entropy over `batches` deterministic, non-overlapping windows of /// the validation corpus (no backward — eval only). Deterministic so val loss is -/// comparable across steps and runs. -fn eval_loss( +/// comparable across steps and runs (and across models — the v0-vs-v1 metric). +pub fn eval_loss( model: &TinyTransformer, device: Device, valid: &Corpus,