diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index b625d5ef7..1b3b79251 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -255,6 +255,14 @@ struct Cli { /// diagnostic construction on most steps and let the GPU run free. #[arg(long, default_value = "1")] diag_every: usize, + + /// Resume from checkpoint file path. + #[arg(long)] + resume_from: Option, + + /// 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 { @@ -987,8 +995,17 @@ fn main() -> Result<()> { // and ~140ms of host-side Rust overhead per step. trainer.enable_mega_graph(); + 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. @@ -1111,6 +1128,18 @@ fn main() -> Result<()> { // blocks on diag I/O. let _ = diag_tx.try_send(frame); + 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 { let isv = trainer.isv_host_slice(); let elapsed = t_start.elapsed().as_secs_f32();