gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(rl): wire checkpoint save/resume into training loop

Browse Source 
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>
This commit is contained in:
jgrusewski
2026-05-27 20:30:50 +02:00
parent e7dc272a14
commit e9bddba6c1
1 changed files with 30 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
crates/ml-alpha/examples/alpha_rl_train.rs
View File

@@ -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<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> {
@@ -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();