Buckets:
| """Atomic GRPO checkpoints with immutable run-start reference recovery.""" | |
| from __future__ import annotations | |
| import json | |
| import re | |
| import uuid | |
| from pathlib import Path | |
| import torch | |
| import yaml | |
| from shared.artifacts import ( | |
| file_sha256, | |
| restore_rng_state, | |
| rng_state, | |
| save_adapter, | |
| write_json, | |
| ) | |
| CHECKPOINT_MANIFEST_SCHEMA = "spatial_mmcot_grpo_checkpoint_manifest_v1" | |
| CHECKPOINT_COMPLETE_SCHEMA = "spatial_mmcot_grpo_checkpoint_complete_v1" | |
| def _checkpoint_inventory(root: Path) -> dict[str, dict[str, int | str]]: | |
| files = {} | |
| for path in sorted(root.rglob("*")): | |
| if path.is_symlink(): | |
| raise ValueError(f"RL checkpoint cannot contain a symlink: {path}") | |
| if not path.is_file() or path in { | |
| root / "manifest.json", | |
| root / "COMPLETE", | |
| }: | |
| continue | |
| relative = str(path.relative_to(root)) | |
| files[relative] = {"size": path.stat().st_size, "sha256": file_sha256(path)} | |
| return files | |
| def verify_rl_checkpoint(checkpoint: str | Path) -> dict: | |
| """Verify the atomic marker, complete inventory, and every file hash.""" | |
| root = Path(checkpoint) | |
| try: | |
| complete = json.loads((root / "COMPLETE").read_text(encoding="utf-8")) | |
| manifest = json.loads((root / "manifest.json").read_text(encoding="utf-8")) | |
| except (OSError, json.JSONDecodeError) as error: | |
| raise ValueError(f"RL checkpoint completion metadata is invalid: {root}") from error | |
| if not isinstance(complete, dict) or not isinstance(manifest, dict): | |
| raise ValueError(f"RL checkpoint completion metadata is invalid: {root}") | |
| if ( | |
| complete.get("schema") != CHECKPOINT_COMPLETE_SCHEMA | |
| or manifest.get("schema") != CHECKPOINT_MANIFEST_SCHEMA | |
| or complete.get("manifest_sha256") != file_sha256(root / "manifest.json") | |
| ): | |
| raise ValueError(f"RL checkpoint completion manifest differs: {root}") | |
| expected = manifest.get("files") | |
| if not isinstance(expected, dict) or expected != _checkpoint_inventory(root): | |
| raise ValueError(f"RL checkpoint file hash/inventory differs: {root}") | |
| required = {"training_state.pt", "rng.pt", "config.yaml"} | |
| if not required.issubset(expected) or not any( | |
| name.startswith("adapter/") for name in expected | |
| ): | |
| raise ValueError(f"RL checkpoint manifest omits required files: {root}") | |
| return manifest | |
| def save_reference_state(run_dir: Path, reference) -> Path: | |
| final = run_dir / "reference" | |
| if final.exists(): | |
| if not (final / "COMPLETE").is_file(): | |
| raise FileExistsError(f"incomplete frozen reference exists: {final}") | |
| recorded = json.loads((final / "descriptor.json").read_text()) | |
| if recorded != reference.descriptor(): | |
| raise ValueError("existing frozen-reference descriptor differs") | |
| return final | |
| temporary = run_dir / f".reference.tmp-{uuid.uuid4().hex}" | |
| temporary.mkdir(parents=False, exist_ok=False) | |
| torch.save(reference.state_dict(), temporary / "state.pt") | |
| write_json(temporary / "descriptor.json", reference.descriptor()) | |
| (temporary / "COMPLETE").write_text("complete\n", encoding="utf-8") | |
| temporary.rename(final) | |
| return final | |
| def load_reference_state(run_dir: Path, reference) -> None: | |
| root = run_dir / "reference" | |
| if not (root / "COMPLETE").is_file(): | |
| raise ValueError(f"frozen reference state is incomplete: {root}") | |
| descriptor = json.loads((root / "descriptor.json").read_text()) | |
| state = torch.load(root / "state.pt", map_location="cpu", weights_only=False) | |
| if descriptor != { | |
| key: value for key, value in state.items() if key != "shadow" | |
| }: | |
| raise ValueError("frozen-reference state and descriptor differ") | |
| reference.load_state_dict(state) | |
| def latest_complete_checkpoint(checkpoints_dir: Path) -> Path: | |
| pattern = re.compile( | |
| r"^step-(\d+)-attempt-(\d+)(?:-retry-(\d+))?$" | |
| ) | |
| def key(path: Path) -> tuple[int, int, int]: | |
| match = pattern.fullmatch(path.name) | |
| if match is None: | |
| return (-1, -1, -1) | |
| step, attempt, retry = match.groups() | |
| return int(step), int(attempt), int(retry or 0) | |
| named = sorted( | |
| ( | |
| path | |
| for path in checkpoints_dir.glob("step-*") | |
| if path.is_dir() and pattern.fullmatch(path.name) | |
| ), | |
| key=key, | |
| reverse=True, | |
| ) | |
| if not named: | |
| raise FileNotFoundError( | |
| f"in-place RL resume has no checkpoint candidate: {checkpoints_dir}" | |
| ) | |
| newest = named[0] | |
| try: | |
| verify_rl_checkpoint(newest) | |
| except ValueError as error: | |
| raise ValueError( | |
| f"newest GRPO checkpoint is corrupt; refusing implicit rollback: {newest}" | |
| ) from error | |
| return newest | |
| def save_rl_checkpoint( | |
| *, | |
| checkpoints_dir: Path, | |
| actor, | |
| processor, | |
| optimizer, | |
| scheduler, | |
| trainer_state: dict, | |
| config, | |
| provenance_fingerprints: dict[str, str], | |
| reference, | |
| artifact_roots: tuple[Path, ...] = (), | |
| world_size: int = 1, | |
| rank_rng_states: list | None = None, | |
| ) -> Path: | |
| required = {"attempted_updates", "successful_updates", "query_update"} | |
| if not required.issubset(trainer_state): | |
| raise ValueError("RL trainer state is missing update cursors") | |
| if any(key in trainer_state for key in ("trajectories", "rollout_groups")): | |
| raise ValueError("fresh sampled trajectories cannot enter RL checkpoints") | |
| run_dir = checkpoints_dir.parent | |
| resolved_run = run_dir.resolve() | |
| relative_artifacts = [] | |
| for artifact in artifact_roots: | |
| resolved = Path(artifact).resolve(strict=True) | |
| if not resolved.is_relative_to(resolved_run) or not all( | |
| (resolved / name).is_file() for name in ("update.json", "metrics.json") | |
| ): | |
| raise ValueError(f"checkpoint artifact is outside/incomplete: {artifact}") | |
| relative_artifacts.append(str(resolved.relative_to(resolved_run))) | |
| save_reference_state(run_dir, reference) | |
| checkpoints_dir.mkdir(parents=True, exist_ok=True) | |
| step = int(trainer_state["successful_updates"]) | |
| attempt = int(trainer_state["attempted_updates"]) | |
| base = checkpoints_dir / f"step-{step:06d}-attempt-{attempt:06d}" | |
| final = base | |
| retry = 0 | |
| while final.exists(): | |
| try: | |
| verify_rl_checkpoint(final) | |
| except ValueError: | |
| retry += 1 | |
| final = checkpoints_dir / f"{base.name}-retry-{retry:04d}" | |
| continue | |
| raise FileExistsError(f"RL checkpoint already exists: {final}") | |
| temporary = checkpoints_dir / ( | |
| f".tmp-{final.name}-{uuid.uuid4().hex}" | |
| ) | |
| temporary.mkdir(parents=False, exist_ok=False) | |
| checkpoint_config = config.as_dict() | |
| resume_parent = None | |
| if isinstance(checkpoint_config.get("checkpoint"), dict): | |
| resume_parent = checkpoint_config["checkpoint"].get("resume_from") | |
| checkpoint_config["checkpoint"]["resume_from"] = None | |
| save_adapter(actor, processor, temporary / "adapter") | |
| torch.save(rank_rng_states if rank_rng_states is not None else rng_state(), temporary / "rng.pt") | |
| torch.save( | |
| { | |
| "schema": "spatial_mmcot_grpo_checkpoint_v1", | |
| "trainer": trainer_state, | |
| "optimizer": optimizer.state_dict(), | |
| "scheduler": scheduler.state_dict(), | |
| "config_fingerprint": config.fingerprint(), | |
| "provenance_fingerprints": provenance_fingerprints, | |
| "world_size": world_size, | |
| "reference_descriptor": reference.descriptor(), | |
| "committed_artifacts": relative_artifacts, | |
| "resume_parent": resume_parent, | |
| }, | |
| temporary / "training_state.pt", | |
| ) | |
| (temporary / "config.yaml").write_text( | |
| yaml.safe_dump(checkpoint_config, sort_keys=False), encoding="utf-8" | |
| ) | |
| write_json( | |
| temporary / "manifest.json", | |
| { | |
| "schema": CHECKPOINT_MANIFEST_SCHEMA, | |
| "files": _checkpoint_inventory(temporary), | |
| }, | |
| ) | |
| write_json( | |
| temporary / "COMPLETE", | |
| { | |
| "schema": CHECKPOINT_COMPLETE_SCHEMA, | |
| "manifest_sha256": file_sha256(temporary / "manifest.json"), | |
| }, | |
| ) | |
| temporary.rename(final) | |
| _commit_artifacts(run_dir, final, relative_artifacts, step, attempt) | |
| return final | |
| def _commit_artifacts( | |
| run_dir: Path, | |
| checkpoint: Path, | |
| relative_artifacts: list[str], | |
| successful_updates: int, | |
| attempted_updates: int, | |
| ) -> None: | |
| resolved_run = run_dir.resolve() | |
| resolved_checkpoint = checkpoint.resolve(strict=True) | |
| if not resolved_checkpoint.is_relative_to(resolved_run): | |
| raise ValueError("RL checkpoint is outside its run directory") | |
| descriptor = { | |
| "schema": "spatial_mmcot_grpo_artifact_commit_v1", | |
| "checkpoint": str(resolved_checkpoint.relative_to(resolved_run)), | |
| "successful_updates": successful_updates, | |
| "attempted_updates": attempted_updates, | |
| } | |
| for relative in relative_artifacts: | |
| artifact = (resolved_run / relative).resolve(strict=True) | |
| if not artifact.is_relative_to(resolved_run) or not all( | |
| (artifact / name).is_file() | |
| for name in ("update.json", "metrics.json") | |
| ): | |
| raise ValueError(f"committed RL artifact is outside/incomplete: {relative}") | |
| marker = artifact / "COMMITTED.json" | |
| if marker.is_file(): | |
| if json.loads(marker.read_text(encoding="utf-8")) != descriptor: | |
| raise ValueError(f"RL artifact commit marker differs: {marker}") | |
| else: | |
| write_json(marker, descriptor) | |
| def load_rl_checkpoint( | |
| checkpoint: str | Path, | |
| *, | |
| optimizer, | |
| scheduler, | |
| config, | |
| provenance_fingerprints: dict[str, str], | |
| reference, | |
| world_size: int = 1, | |
| commit_artifacts: bool = True, | |
| rank: int = 0, | |
| ) -> dict: | |
| checkpoint = Path(checkpoint) | |
| verify_rl_checkpoint(checkpoint) | |
| payload = torch.load( | |
| checkpoint / "training_state.pt", map_location="cpu", weights_only=False | |
| ) | |
| if payload.get("schema") != "spatial_mmcot_grpo_checkpoint_v1": | |
| raise ValueError("unsupported RL checkpoint schema") | |
| expected = { | |
| "config_fingerprint": config.fingerprint(), | |
| "provenance_fingerprints": provenance_fingerprints, | |
| "world_size": world_size, | |
| } | |
| differences = { | |
| key: {"expected": payload["provenance_fingerprints"].get(key), "actual": provenance_fingerprints.get(key)} | |
| for key in sorted(payload["provenance_fingerprints"].keys() | provenance_fingerprints.keys()) | |
| if payload["provenance_fingerprints"].get(key) != provenance_fingerprints.get(key) | |
| } | |
| report = {"checkpoint": str(checkpoint), "rank": rank, "differences": differences} | |
| print("RL_RESUME_PROVENANCE " + json.dumps(report, sort_keys=True), flush=True) | |
| output_dir = getattr(getattr(config, "run", None), "output_dir", None) | |
| if output_dir: | |
| diagnostics = Path(output_dir) / "diagnostics" | |
| diagnostics.mkdir(parents=True, exist_ok=True) | |
| write_json(diagnostics / f"resume-rank-{rank}-{uuid.uuid4().hex}.json", report) | |
| for key, value in expected.items(): | |
| if payload.get(key) != value: | |
| if getattr(getattr(config, "checkpoint", None), "resume_mode", "exact") != "execution_change": | |
| raise ValueError(f"RL resume {key} mismatch: expected={payload.get(key)!r}, actual={value!r}") | |
| from .resume import validate_execution_change | |
| validate_execution_change(checkpoint, payload, config, provenance_fingerprints) | |
| break | |
| if commit_artifacts: | |
| _commit_artifacts( | |
| checkpoint.parent.parent, | |
| checkpoint, | |
| list(payload.get("committed_artifacts", ())), | |
| int(payload["trainer"]["successful_updates"]), | |
| int(payload["trainer"]["attempted_updates"]), | |
| ) | |
| optimizer.load_state_dict(payload["optimizer"]) | |
| scheduler.load_state_dict(payload["scheduler"]) | |
| load_reference_state(checkpoint.parent.parent, reference) | |
| if reference.descriptor() != payload["reference_descriptor"]: | |
| raise ValueError("restored RL reference descriptor differs") | |
| saved_rng = torch.load(checkpoint / "rng.pt", map_location="cpu", weights_only=False) | |
| if isinstance(saved_rng, list): | |
| if len(saved_rng) != world_size: | |
| raise ValueError("checkpoint RNG rank count differs") | |
| saved_rng = saved_rng[rank] | |
| restore_rng_state(saved_rng) | |
| return payload["trainer"] | |
| __all__ = [ | |
| "latest_complete_checkpoint", | |
| "load_reference_state", | |
| "load_rl_checkpoint", | |
| "save_reference_state", | |
| "save_rl_checkpoint", | |
| "verify_rl_checkpoint", | |
| ] | |
Xet Storage Details
- Size:
- 12.9 kB
- Xet hash:
- f7139e969f80c97dc242edbbbe66ecf34f412fa381a3477bbeec278b5ac47bdd
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.