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>
This commit is contained in:
@@ -246,6 +246,14 @@ struct Cli {
|
|||||||
/// it produces a non-finite loss (G8 NaN abort).
|
/// it produces a non-finite loss (G8 NaN abort).
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
diag_jsonl: Option<PathBuf>,
|
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> {
|
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 frd_gate_total: u64 = 0;
|
||||||
let mut heat_cap_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();
|
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 +
|
// GPU-resident data loading: sample_and_gather + gather_next +
|
||||||
// gather_current + gather_frd_labels all run as GPU kernels
|
// gather_current + gather_frd_labels all run as GPU kernels
|
||||||
// inside step_with_lobsim_gpu. Zero CPU data work per step.
|
// 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")?;
|
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 {
|
if step % cli.log_every == 0 || step == cli.n_steps - 1 {
|
||||||
diag.flush().context("diag: flush")?;
|
diag.flush().context("diag: flush")?;
|
||||||
let elapsed = t_start.elapsed().as_secs_f32();
|
let elapsed = t_start.elapsed().as_secs_f32();
|
||||||
|
|||||||
Reference in New Issue
Block a user