From 7f4cd86421a279955b949c950c7842e77a5bf584 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Wed, 27 May 2026 20:30:50 +0200
Subject: [PATCH] feat(rl): wire checkpoint save/resume into training loop

CLI flags: --resume-from <path> and --checkpoint-every <n> (default 5000).
Resume loads checkpoint and continues from saved step. Save writes
rolling checkpoints (keeps last 2, deletes older).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 crates/ml-alpha/examples/alpha_rl_train.rs | 31 +++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs
index b81590201..a7feeb545 100644
--- a/crates/ml-alpha/examples/alpha_rl_train.rs
+++ b/crates/ml-alpha/examples/alpha_rl_train.rs
@@ -246,6 +246,14 @@ struct Cli {
     /// it produces a non-finite loss (G8 NaN abort).
     #[arg(long)]
     diag_jsonl: Option<PathBuf>,
+
+    /// Resume from checkpoint file path.
+    #[arg(long)]
+    resume_from: Option<PathBuf>,
+
+    /// Checkpoint interval in steps (default: 5000). Set to 0 to disable.
+    #[arg(long, default_value = "5000")]
+    checkpoint_every: usize,
 }
 
 fn parse_instrument_mode(s: &str) -> Result<InstrumentFilter, String> {
@@ -504,8 +512,17 @@ fn main() -> Result<()> {
     let mut frd_gate_total: u64 = 0;
     let mut heat_cap_total: u64 = 0;
 
+    let start_step = if let Some(ref ckpt_path) = cli.resume_from {
+        let resumed_step = trainer.load_checkpoint(&dev, ckpt_path)
+            .with_context(|| format!("resume from {}", ckpt_path.display()))?;
+        eprintln!("resumed from checkpoint at step {resumed_step}");
+        resumed_step as usize
+    } else {
+        0
+    };
+
     let t_start = std::time::Instant::now();
-    for step in 0..cli.n_steps {
+    for step in start_step..cli.n_steps {
         // GPU-resident data loading: sample_and_gather + gather_next +
         // gather_current + gather_frd_labels all run as GPU kernels
         // inside step_with_lobsim_gpu. Zero CPU data work per step.
@@ -1099,6 +1116,18 @@ fn main() -> Result<()> {
         });
         writeln!(diag, "{}", record).context("diag: writeln jsonl")?;
 
+        if cli.checkpoint_every > 0 && step > 0 && step % cli.checkpoint_every == 0 {
+            let ckpt_path = cli.out.join(format!("checkpoint-{step}.bin"));
+            trainer.save_checkpoint(&ckpt_path, step as u64)
+                .with_context(|| format!("checkpoint at step {step}"))?;
+            eprintln!("checkpoint saved: {}", ckpt_path.display());
+            let old_step = step.saturating_sub(cli.checkpoint_every * 2);
+            if old_step > 0 {
+                let old_path = cli.out.join(format!("checkpoint-{old_step}.bin"));
+                let _ = std::fs::remove_file(&old_path);
+            }
+        }
+
         if step % cli.log_every == 0 || step == cli.n_steps - 1 {
             diag.flush().context("diag: flush")?;
             let elapsed = t_start.elapsed().as_secs_f32();