fix: make checkpoint resume explicit and stop claiming optimizer state

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

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

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

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<PathBuf>,
+
+        /// 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,
     },
 }

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(())

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),
     }
 }

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"),

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;