Buckets:
| """Atomic OPD checkpoints and exact next-update recovery.""" | |
| from __future__ import annotations | |
| import uuid | |
| from pathlib import Path | |
| from typing import Callable | |
| import torch | |
| import yaml | |
| from shared.artifacts import restore_rng_state, rng_state, save_adapter | |
| def latest_complete_checkpoint(checkpoints_dir: Path) -> Path: | |
| candidates = sorted( | |
| ( | |
| path | |
| for path in checkpoints_dir.glob("step-*") | |
| if path.is_dir() and (path / "COMPLETE").is_file() | |
| ), | |
| key=lambda path: int(path.name.removeprefix("step-")), | |
| ) | |
| if not candidates: | |
| raise FileNotFoundError( | |
| f"in-place OPD resume has no complete checkpoint: {checkpoints_dir}" | |
| ) | |
| return candidates[-1] | |
| def save_opd_checkpoint( | |
| *, | |
| checkpoints_dir: Path, | |
| model, | |
| processor, | |
| optimizer, | |
| scheduler, | |
| trainer_state: dict, | |
| config, | |
| provenance_fingerprints: dict[str, str], | |
| teacher_state: dict, | |
| world_size: int = 1, | |
| rank_rng_states: list[dict] | None = None, | |
| ) -> Path: | |
| """Publish a checkpoint only after a complete update and cleared queue.""" | |
| required = {"attempted_updates", "successful_updates", "query_update"} | |
| if not required.issubset(trainer_state): | |
| raise ValueError( | |
| "OPD trainer state is missing " | |
| f"{sorted(required - trainer_state.keys())}" | |
| ) | |
| if any(key in trainer_state for key in ("rollouts", "failure_queue")): | |
| raise ValueError( | |
| "sampled rollouts and update-local failure queue are not checkpoint state" | |
| ) | |
| if teacher_state.get("update_count") not in (None, trainer_state["successful_updates"]): | |
| raise ValueError("teacher update count differs from successful OPD updates") | |
| checkpoints_dir.mkdir(parents=True, exist_ok=True) | |
| step = int(trainer_state["successful_updates"]) | |
| final = checkpoints_dir / f"step-{step:06d}" | |
| if final.exists(): | |
| raise FileExistsError(f"checkpoint already exists: {final}") | |
| temporary = checkpoints_dir / f".tmp-step-{step:06d}-{uuid.uuid4().hex}" | |
| temporary.mkdir(parents=False, exist_ok=False) | |
| try: | |
| save_adapter(model, processor, temporary / "adapter") | |
| if rank_rng_states is None: | |
| rank_rng_states = [rng_state()] | |
| if len(rank_rng_states) != world_size: | |
| raise ValueError("rank RNG states must match OPD world size") | |
| if world_size == 1: | |
| torch.save(rank_rng_states[0], temporary / "rng.pt") | |
| else: | |
| for rank, state in enumerate(rank_rng_states): | |
| torch.save(state, temporary / f"rng-rank-{rank:02d}.pt") | |
| torch.save( | |
| { | |
| "schema": ( | |
| "opd_checkpoint_v1" if world_size == 1 else "opd_checkpoint_v2" | |
| ), | |
| "trainer": trainer_state, | |
| "optimizer": optimizer.state_dict(), | |
| "scheduler": scheduler.state_dict(), | |
| "config_fingerprint": config.fingerprint(), | |
| "provenance_fingerprints": provenance_fingerprints, | |
| "world_size": world_size, | |
| "teacher": teacher_state, | |
| }, | |
| temporary / "training_state.pt", | |
| ) | |
| (temporary / "config.yaml").write_text( | |
| yaml.safe_dump(config.as_dict(), sort_keys=False), encoding="utf-8" | |
| ) | |
| (temporary / "COMPLETE").write_text("complete\n", encoding="utf-8") | |
| temporary.rename(final) | |
| except Exception: | |
| # Keep a failed staging directory for forensic inspection; it is never | |
| # considered resumable because COMPLETE is absent. | |
| raise | |
| return final | |
| def load_opd_checkpoint( | |
| checkpoint: str | Path, | |
| *, | |
| optimizer, | |
| scheduler, | |
| config, | |
| provenance_fingerprints: dict[str, str], | |
| teacher_state_loader: Callable[[dict], None], | |
| rank: int = 0, | |
| world_size: int = 1, | |
| ) -> dict: | |
| checkpoint = Path(checkpoint) | |
| if not (checkpoint / "COMPLETE").is_file(): | |
| raise ValueError(f"OPD checkpoint is incomplete: {checkpoint}") | |
| payload = torch.load( | |
| checkpoint / "training_state.pt", map_location="cpu", weights_only=False | |
| ) | |
| if payload.get("schema") not in {"opd_checkpoint_v1", "opd_checkpoint_v2"}: | |
| raise ValueError("unsupported OPD checkpoint schema") | |
| expected = { | |
| "config_fingerprint": config.fingerprint(), | |
| "provenance_fingerprints": provenance_fingerprints, | |
| "world_size": world_size, | |
| } | |
| cap_changed = False | |
| for key, value in expected.items(): | |
| if payload.get(key) != value: | |
| if key != "world_size" and getattr(getattr(config, "checkpoint", None), "resume_mode", "exact") in {"execution_change", "rollout_limit_change"}: | |
| from .resume import validate_execution_change | |
| cap_changed = validate_execution_change(checkpoint, payload, config, provenance_fingerprints) | |
| continue | |
| raise ValueError(f"OPD resume {key} mismatch") | |
| optimizer.load_state_dict(payload["optimizer"]) | |
| scheduler.load_state_dict(payload["scheduler"]) | |
| teacher_state_loader(payload["teacher"]) | |
| rng_path = ( | |
| checkpoint / "rng.pt" | |
| if payload["schema"] == "opd_checkpoint_v1" | |
| else checkpoint / f"rng-rank-{rank:02d}.pt" | |
| ) | |
| restore_rng_state(torch.load(rng_path, map_location="cpu", weights_only=False)) | |
| state = payload["trainer"] | |
| if cap_changed: | |
| # A new stopping rule needs its own validation baseline; optimizer, | |
| # teacher, query and training counters remain those of the parent. | |
| state.pop("validation_baseline", None) | |
| state.pop("last_validation_artifact", None) | |
| state.pop("validation_baseline_step", None) | |
| state["rollout_limit_changed_at_update"] = state["successful_updates"] | |
| if payload["teacher"].get("update_count") not in ( | |
| None, | |
| state["successful_updates"], | |
| ): | |
| raise ValueError("restored teacher update count differs from OPD state") | |
| return state | |
| __all__ = [ | |
| "latest_complete_checkpoint", | |
| "load_opd_checkpoint", | |
| "save_opd_checkpoint", | |
| ] | |
Xet Storage Details
- Size:
- 6.17 kB
- Xet hash:
- 31bd49f12393941cf2c593356d1398fc7fcc2180ac659ab0950d259762d0e56e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.