diff --git a/CHANGELOG.md b/CHANGELOG.md index 10c451f..3aa72d3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed — checkpointing +- **Resume was implicit and could silently do nothing.** `train()` loaded any checkpoint it + found with no flag and no way to opt out, restoring the step counter with it — so + re-running `glm-train --steps 80` after an 80-step run restored `step = 80`, fell straight + through the loop, and reported a successful run having trained nothing. Resume is now + opt-in via `--resume`, an ignored checkpoint directory is called out on stdout, and a + resumed run already at `max_steps` fails with both numbers named instead of exiting quietly. +- **`save_checkpoint` fell back to the literal name `"unknown"`** when the parameter and + name lists disagreed in length, which collides in the safetensors map and silently drops + every parameter after the first. It now errors. + +### Removed — checkpointing +- `save_optimizer_state`, which claimed to save optimizer state and wrote `{"step": N}`. + candle's `AdamW` keeps its moments in a private struct with no accessor, so honouring the + claim would mean writing and maintaining our own optimizer — not worth it for a CPU + playground. The config field, the `optimizer_step_*.json` file, the empty `if let` that + read it back, and the claim in the README are all gone; what resume does restore is now + documented. + ### Changed — dependencies - Took every open Dependabot bump in one pass: candle-core and candle-nn 0.8.4 to 0.11.0, tokenizers 0.21.4 to 0.23.1, safetensors 0.4.5 to 0.8.0, rand 0.8 to 0.9, rand_distr 0.4 diff --git a/README.md b/README.md index 771bd6f..c9bfeed 100644 --- a/README.md +++ b/README.md @@ -420,7 +420,7 @@ training: - **YAML config** — Declarative training configuration via `--config` - **DataLoader** — Train/eval split, shuffling, random windowing - **LR Scheduler** — Cosine decay with linear warmup -- **Safetensors Checkpoints** — Save/load model + optimizer state +- **Safetensors Checkpoints** — Save/load weights, step counter and LR schedule - **Gradient Accumulation** — Configurable accumulation steps - **Gradient Clipping** — By global norm - **Evaluation Loop** — Periodic validation with fixed seed @@ -433,8 +433,16 @@ cargo run --release -- glm-train --data-path data --steps 500 # or drive it from a YAML config cargo run --release -- glm-train --config configs/train.yaml --steps 500 + +# continue an earlier run +cargo run --release -- glm-train --config configs/train.yaml --steps 1000 --resume ``` +Resume is opt-in. Without `--resume` an existing checkpoint directory is left alone and +training starts from step 0. `--resume` restores the weights, the step counter and the +learning-rate schedule position — but not AdamW's moments, which candle keeps private, so +the loss rises briefly on the first steps after a resume. + --- ## 🧩 Model Configurations diff --git a/configs/train.yaml b/configs/train.yaml index 26a33fb..96a8ca0 100644 --- a/configs/train.yaml +++ b/configs/train.yaml @@ -50,9 +50,9 @@ training: log_every: 10 eval_steps: 100 - # Checkpointing + # Checkpointing. Resume with `glm-train --resume`; it restores weights, the step + # counter and the LR schedule, but not AdamW's moments. checkpoint_dir: "glm_checkpoint" - save_optimizer_state: true keep_last_n_checkpoints: 3 # Precision diff --git a/src/cli.rs b/src/cli.rs index 4fc3cdd..0a4d108 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -97,5 +97,12 @@ pub enum Commands { /// YAML training config (see configs/train.yaml). Defaults are used if omitted. #[arg(short, long)] config: Option, + + /// Continue from the newest checkpoint instead of starting from step 0 + /// + /// Restores weights, the step counter and the learning-rate schedule. + /// AdamW's moments are not restored, so the loss briefly rises after a resume. + #[arg(long)] + resume: bool, }, } diff --git a/src/commands/glm_train.rs b/src/commands/glm_train.rs index 1b56ce0..7a62158 100644 --- a/src/commands/glm_train.rs +++ b/src/commands/glm_train.rs @@ -6,7 +6,13 @@ use candle_core::Device; use crate::cli::Cli; use crate::training::{GLMTrainer, TrainConfig}; -pub fn run(cli: &Cli, data_path: &str, steps: usize, config_path: Option<&Path>) -> Result<()> { +pub fn run( + cli: &Cli, + data_path: &str, + steps: usize, + config_path: Option<&Path>, + resume: bool, +) -> Result<()> { let device = Device::Cpu; let dtype_str = if cli.f16 { "FP16" } else { "FP32" }; @@ -52,7 +58,7 @@ pub fn run(cli: &Cli, data_path: &str, steps: usize, config_path: Option<&Path>) let mut trainer = GLMTrainer::from_config(&config, &device)?; println!("Starting training for {steps} steps...\n"); - trainer.train(data_dir, &device)?; + trainer.train(data_dir, &device, resume)?; println!("\nTraining complete!"); Ok(()) diff --git a/src/main.rs b/src/main.rs index 369ace3..88448ce 100644 --- a/src/main.rs +++ b/src/main.rs @@ -33,6 +33,7 @@ fn main() -> anyhow::Result<()> { data_path, steps, config, - } => commands::glm_train::run(&cli, data_path, *steps, config.as_deref()), + resume, + } => commands::glm_train::run(&cli, data_path, *steps, config.as_deref(), *resume), } } diff --git a/src/training/config.rs b/src/training/config.rs index c804cba..5368a56 100644 --- a/src/training/config.rs +++ b/src/training/config.rs @@ -65,7 +65,6 @@ pub struct TrainingConfig { pub log_every: usize, pub eval_steps: usize, pub checkpoint_dir: PathBuf, - pub save_optimizer_state: bool, pub keep_last_n_checkpoints: usize, pub dtype: String, pub tokenizer_path: PathBuf, @@ -94,7 +93,6 @@ impl Default for TrainingConfig { log_every: 10, eval_steps: 100, checkpoint_dir: PathBuf::from("glm_checkpoint"), - save_optimizer_state: true, keep_last_n_checkpoints: 3, dtype: "f32".to_string(), tokenizer_path: PathBuf::from("codegen_weights/tokenizer.json"), diff --git a/src/training/mod.rs b/src/training/mod.rs index 692428c..4f18ee2 100644 --- a/src/training/mod.rs +++ b/src/training/mod.rs @@ -5,7 +5,7 @@ //! - Data loading with train/eval split ([`data::DataLoader`]) //! - Learning rate scheduling with warmup ([`lr_scheduler::LrScheduler`]) //! - Gradient accumulation and clipping ([`train::clip_grad_norm`]) -//! - Safetensors checkpoint save/load +//! - Safetensors checkpoint save/load, and opt-in resume ([`train::GLMTrainer::load_checkpoint`]) pub mod config; pub mod data; diff --git a/src/training/train.rs b/src/training/train.rs index de819c9..66c4270 100644 --- a/src/training/train.rs +++ b/src/training/train.rs @@ -169,7 +169,13 @@ impl GLMTrainer { Ok(loss_scalar) } - pub fn train(&mut self, data_dir: &Path, device: &Device) -> Result<()> { + /// Run the training loop. + /// + /// With `resume`, weights and the step counter are restored from the newest + /// checkpoint in `checkpoint_dir`. Without it, any existing checkpoints are left + /// alone and training starts from step 0 — resuming used to happen implicitly, + /// which meant a stale directory could silently turn a run into a no-op. + pub fn train(&mut self, data_dir: &Path, device: &Device, resume: bool) -> Result<()> { // Load data let mut examples = load_data(data_dir, &self.tokenizer)?; @@ -207,8 +213,30 @@ impl GLMTrainer { self.config.seed + 1, )?; - // Try to load checkpoint - self.load_checkpoint()?; + if resume { + if !self.load_checkpoint()? { + println!( + "--resume: no checkpoint in {:?}, starting from step 0", + self.config.checkpoint_dir + ); + } + } else if self.latest_checkpoint().is_some() { + println!( + "\x1b[33mNote: checkpoints exist in {:?} but --resume was not passed; \ + starting from step 0 and overwriting.\x1b[0m", + self.config.checkpoint_dir + ); + } + + // A resumed run that is already at the limit would otherwise fall straight + // through the loop and report success without training anything. + if self.step >= self.config.max_steps { + return Err(candle_core::Error::Msg(format!( + "nothing to do: resumed at step {} but max_steps is {}. \ + Pass a larger --steps to continue training.", + self.step, self.config.max_steps + ))); + } println!("Starting training from step {}", self.step); println!(" Max steps: {}", self.config.max_steps); @@ -293,33 +321,28 @@ impl GLMTrainer { let params = self.model.param_vars(); let names = param_names(self.glm_config.num_layers); - - // Save as safetensors - let mut tensor_data = Vec::new(); - for (i, var) in params.iter().enumerate() { - let name = names.get(i).map(|s| s.as_str()).unwrap_or("unknown"); - let tensor = var.as_tensor(); - tensor_data.push((name.to_string(), tensor)); + // Saving and loading pair these two positionally. A mismatch used to fall back + // to the literal name "unknown", which collides in the safetensors map and + // silently drops every parameter after the first. + if params.len() != names.len() { + return Err(candle_core::Error::Msg(format!( + "checkpoint layout mismatch: {} parameters but {} names", + params.len(), + names.len() + ))); } + let tensors: Vec<(String, &Tensor)> = names + .iter() + .cloned() + .zip(params.iter().map(|var| var.as_tensor())) + .collect(); + let path = dir.join(format!("model_step_{:06}.safetensors", self.step)); - save_safetensors(&path, &tensor_data)?; - - // Save optimizer state (simplified - just step count) - if self.config.save_optimizer_state { - let opt_path = dir.join(format!("optimizer_step_{:06}.json", self.step)); - let state = serde_json::json!({ - "step": self.step, - }); - let opt_json = serde_json::to_string_pretty(&state).map_err(|e| { - candle_core::Error::Msg(format!("Failed to serialize optimizer state: {e}")) - })?; - std::fs::write(&opt_path, opt_json).map_err(|e| { - candle_core::Error::Msg(format!("Failed to write optimizer state: {e}")) - })?; - } + save_safetensors(&path, &tensors)?; - // Save training state + // Only the step and learning rate are restorable: candle's AdamW keeps its + // moments in a private struct with no accessor, so they cannot be saved. let state_path = dir.join("training_state.json"); let state = serde_json::json!({ "step": self.step, @@ -334,14 +357,11 @@ impl GLMTrainer { Ok(()) } - pub fn load_checkpoint(&mut self) -> Result<()> { + /// Path of the highest-numbered `model_step_*.safetensors` in the checkpoint dir. + fn latest_checkpoint(&self) -> Option { let dir = Path::new(&self.config.checkpoint_dir); - if !dir.exists() { - return Ok(()); - } - - // Find latest checkpoint - let mut checkpoints: Vec = std::fs::read_dir(dir)? + let mut checkpoints: Vec = std::fs::read_dir(dir) + .ok()? .filter_map(|e| e.ok()) .map(|e| e.path()) .filter(|p| p.extension().and_then(|s| s.to_str()) == Some("safetensors")) @@ -352,47 +372,38 @@ impl GLMTrainer { .starts_with("model_step_") }) .collect(); - checkpoints.sort(); + checkpoints.pop() + } - if let Some(latest) = checkpoints.last() { - println!("Loading checkpoint from {:?}", latest); - load_safetensors(latest, &mut self.model)?; - - // Load optimizer state - let opt_path = dir.join( - latest - .file_name() - .unwrap() - .to_string_lossy() - .replace("model_", "optimizer_") - .replace(".safetensors", ".json"), - ); - if opt_path.exists() { - if let Ok(content) = std::fs::read_to_string(&opt_path) { - if let Ok(state) = serde_json::from_str::(&content) { - if let Some(_step) = state.get("step").and_then(|v| v.as_u64()) { - // Note: AdamW doesn't expose step_count setter, but we track our own step - } - } - } - } - - // Load training state - let state_path = dir.join("training_state.json"); - if state_path.exists() { - let content = std::fs::read_to_string(&state_path) - .map_err(|e| candle_core::Error::Msg(format!("Failed to read state: {e}")))?; - let state: serde_json::Value = serde_json::from_str(&content) - .map_err(|e| candle_core::Error::Msg(format!("Failed to parse state: {e}")))?; - self.step = state.get("step").and_then(|v| v.as_u64()).unwrap_or(0) as usize; - self.lr_scheduler.set_step(self.step); - } + /// Restore weights, step counter and learning-rate schedule position from the + /// newest checkpoint. Returns `false` if there was nothing to restore. + /// + /// AdamW's moments are **not** restored — candle keeps them private — so the first + /// steps after a resume run with a cold optimizer and the loss briefly rises. + pub fn load_checkpoint(&mut self) -> Result { + let Some(latest) = self.latest_checkpoint() else { + return Ok(false); + }; - println!("Resumed from step {}", self.step); + println!("Loading checkpoint from {:?}", latest); + load_safetensors(&latest, &mut self.model)?; + + let state_path = Path::new(&self.config.checkpoint_dir).join("training_state.json"); + if state_path.exists() { + let content = std::fs::read_to_string(&state_path) + .map_err(|e| candle_core::Error::Msg(format!("Failed to read state: {e}")))?; + let state: serde_json::Value = serde_json::from_str(&content) + .map_err(|e| candle_core::Error::Msg(format!("Failed to parse state: {e}")))?; + self.step = state.get("step").and_then(|v| v.as_u64()).unwrap_or(0) as usize; + self.lr_scheduler.set_step(self.step); } - Ok(()) + println!( + "Resumed from step {} (optimizer moments restart cold)", + self.step + ); + Ok(true) } fn cleanup_old_checkpoints(&self) -> Result<()> { @@ -413,15 +424,6 @@ impl GLMTrainer { while model_checkpoints.len() > self.config.keep_last_n_checkpoints { let oldest = model_checkpoints.remove(0); let _ = std::fs::remove_file(&oldest); - let opt_file = dir.join( - oldest - .file_name() - .unwrap() - .to_string_lossy() - .replace("model_", "optimizer_") - .replace(".safetensors", ".json"), - ); - let _ = std::fs::remove_file(&opt_file); } Ok(()) @@ -682,4 +684,82 @@ mod tests { assert!((after - before).abs() < 1e-6, "{before} changed to {after}"); Ok(()) } + + fn tiny_model_config() -> GLMConfig { + GLMConfig { + vocab_size: 64, + hidden_dim: 32, + num_layers: 2, + num_heads: 4, + ffn_dim: 64, + max_seq_len: 16, + ..Default::default() + } + } + + /// Saving and loading pair `param_vars()` with `param_names()` positionally, so + /// the two must agree in length and the names must be unique. If they ever drift, + /// every parameter after the first mismatch loads into the wrong slot. + #[test] + fn parameter_names_match_parameter_vars() -> Result<()> { + let device = Device::Cpu; + let config = tiny_model_config(); + let model = TrainableGLMModel::new(config.clone(), &device)?; + + let names = param_names(config.num_layers); + assert_eq!( + names.len(), + model.param_vars().len(), + "name list and parameter list have drifted apart" + ); + + let unique: std::collections::HashSet<&String> = names.iter().collect(); + assert_eq!( + unique.len(), + names.len(), + "duplicate parameter names: {names:?}" + ); + Ok(()) + } + + /// A restored model must equal the saved one, parameter for parameter. + #[test] + fn checkpoint_round_trip_restores_every_parameter() -> Result<()> { + let device = Device::Cpu; + let config = tiny_model_config(); + let saved = TrainableGLMModel::new(config.clone(), &device)?; + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("model_step_000001.safetensors"); + let names = param_names(config.num_layers); + let vars = saved.param_vars(); + let tensors: Vec<(String, &Tensor)> = names + .iter() + .cloned() + .zip(vars.iter().map(|v| v.as_tensor())) + .collect(); + save_safetensors(&path, &tensors)?; + + // A second model starts from different random weights. + let mut restored = TrainableGLMModel::new(config.clone(), &device)?; + let before = (restored.param_vars()[0].as_tensor() - vars[0].as_tensor())? + .abs()? + .max_all()? + .to_scalar::()?; + assert!(before > 0.0, "the two models started out identical"); + + load_safetensors(&path, &mut restored)?; + + for (name, (expected, actual)) in names + .iter() + .zip(vars.iter().zip(restored.param_vars().iter())) + { + let diff = (actual.as_tensor() - expected.as_tensor())? + .abs()? + .max_all()? + .to_scalar::()?; + assert!(diff < 1e-6, "{name} differs by {diff} after a round trip"); + } + Ok(()) + } }