| """Atomic training checkpoints with full RNG state for exact resume.""" |
|
|
| from __future__ import annotations |
|
|
| import os |
| import random |
| import re |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| _STEP_RE = re.compile(r"step_(\d{7})\.pt$") |
|
|
|
|
| def checkpoint_path(directory: Path, step: int) -> Path: |
| return Path(directory) / f"step_{step:07d}.pt" |
|
|
|
|
| def save_checkpoint(path: Path, model, optimizer, scaler, step: int, extra: dict) -> None: |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| state = { |
| "model": model.state_dict(), |
| "optimizer": optimizer.state_dict() if optimizer is not None else None, |
| "scaler": scaler.state_dict() if scaler is not None else None, |
| "step": step, |
| "extra": extra, |
| "rng": { |
| "python": random.getstate(), |
| "numpy": np.random.get_state(), |
| "torch": torch.get_rng_state(), |
| "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, |
| }, |
| } |
| tmp = path.with_suffix(".tmp") |
| torch.save(state, tmp) |
| os.replace(tmp, path) |
|
|
|
|
| def load_checkpoint(path: Path, model, optimizer=None, scaler=None) -> dict: |
| state = torch.load(path, map_location="cpu", weights_only=False) |
| model.load_state_dict(state["model"]) |
| if optimizer is not None and state["optimizer"] is not None: |
| optimizer.load_state_dict(state["optimizer"]) |
| if scaler is not None and state["scaler"] is not None: |
| scaler.load_state_dict(state["scaler"]) |
| rng = state["rng"] |
| random.setstate(rng["python"]) |
| np.random.set_state(rng["numpy"]) |
| torch.set_rng_state(rng["torch"]) |
| if rng["cuda"] is not None and torch.cuda.is_available(): |
| torch.cuda.set_rng_state_all(rng["cuda"]) |
| return {"step": state["step"], "extra": state["extra"]} |
|
|
|
|
| def _step_files(directory: Path) -> list[tuple[int, Path]]: |
| out = [] |
| for p in Path(directory).glob("step_*.pt"): |
| m = _STEP_RE.search(p.name) |
| if m: |
| out.append((int(m.group(1)), p)) |
| return sorted(out) |
|
|
|
|
| def latest_checkpoint(directory: Path) -> Path | None: |
| files = _step_files(directory) if Path(directory).exists() else [] |
| return files[-1][1] if files else None |
|
|
|
|
| def prune_checkpoints(directory: Path, keep_last: int = 2) -> None: |
| for _, p in _step_files(directory)[:-keep_last]: |
| p.unlink() |
|
|
