gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions 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 66ec7f75f4
commit 7f4cd86421
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

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