From 687e884ad446896c6d68997552408d2069e02c29 Mon Sep 17 00:00:00 2001 From: Muhammad Ayyan Khan Date: Mon, 24 Aug 2026 04:59:26 +0500 Subject: [PATCH] fix: make checkpoint resume explicit and stop claiming optimizer state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resume was implicit. train() loaded any checkpoint it found in checkpoint_dir with no flag and no way to opt out, restoring the step counter along with the weights — so re-running `glm-train --steps 80` after an 80-step run restored step = 80, fell straight through the loop, and printed what looked like a successful run having trained nothing. Resume is now opt-in via --resume. Without it, an existing checkpoint directory is left alone and the run starts from step 0, with a note on stdout so a stale directory is visible rather than surprising. With it and a restored step already at max_steps, the run fails naming both numbers instead of exiting quietly. save_optimizer_state claimed to save optimizer state and wrote {"step": N} — no AdamW moments. candle keeps first_moment and second_moment in a private VarAdamW with no accessor, so honouring that claim means writing and maintaining our own optimizer, which is not worth it for a CPU playground with short runs. The config field, the optimizer_step_*.json file, the empty if-let that read it back, and the README claim are gone. What resume does restore — weights, step counter, LR schedule position — is now documented, along with the fact that Adam restarts cold and the loss briefly rises. save_checkpoint also fell back to the literal name "unknown" when the parameter and name lists disagreed in length. That collides in the safetensors map and silently drops every parameter after the first; it now errors. Verified end to end. Last loss before resuming was 5.0751 at step 40; the first after was 5.1103 at step 45, and the learning rate picked up at the right point on the cosine schedule. A cold start would have been back near 10.8. 104 tests pass, including a checkpoint round trip that asserts every parameter matches after a save and load, and a check that the parameter and name lists agree in length and carry no duplicates. --- CHANGELOG.md | 19 +++ README.md | 10 +- configs/train.yaml | 4 +- src/cli.rs | 7 ++ src/commands/glm_train.rs | 10 +- src/main.rs | 3 +- src/training/config.rs | 2 - src/training/mod.rs | 2 +- src/training/train.rs | 236 +++++++++++++++++++++++++------------- 9 files changed, 206 insertions(+), 87 deletions(-) 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(()) + } }