Buckets:
| """Run directories, deterministic state, and atomic SFT checkpoints.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import os | |
| import uuid | |
| from collections.abc import Callable | |
| from pathlib import Path | |
| import torch | |
| import yaml | |
| from shared.artifacts import ( | |
| acquire_run_writer, | |
| cache_fingerprint, | |
| claim_run_directory, | |
| environment_record, | |
| file_sha256, | |
| model_fingerprint, | |
| resolve_run_directory, | |
| restore_rng_state, | |
| rng_state, | |
| save_adapter, | |
| save_adapter_atomic, | |
| seed_everything, | |
| source_fingerprint, | |
| write_json, | |
| ) | |
| from .config import SFTConfig, load_config | |
| def _ddp_resume_contract( | |
| config: SFTConfig, *, epoch_extension: bool = False | |
| ) -> dict: | |
| """Return the mathematical run contract, excluding only DDP execution.""" | |
| value = config.as_dict() | |
| value["run"]["output_dir"] = None | |
| value["run"]["output_root"] = None | |
| value["run"].pop("version", None) | |
| value["run"].pop("dataset") | |
| value["run"].pop("category") | |
| value["checkpoint"]["resume_from"] = None | |
| value["checkpoint"].pop("resume_mode") | |
| value["trainer"].pop("ddp_find_unused_parameters") | |
| value["trainer"].pop("ddp_static_graph") | |
| value["trainer"].pop("gradient_sync") | |
| if epoch_extension: | |
| value["run"].pop("name") | |
| value["trainer"].pop("epochs") | |
| return value | |
| def _validate_ddp_execution_resume( | |
| checkpoint: Path, | |
| payload: dict, | |
| *, | |
| config: SFTConfig, | |
| provenance_fingerprints: dict[str, str], | |
| world_size: int, | |
| run_dir: Path | None, | |
| ) -> None: | |
| if run_dir is None: | |
| raise ValueError("DDP execution-change resume requires the child run directory") | |
| submission_path = run_dir / "_submission" / "submission.json" | |
| if not submission_path.is_file(): | |
| raise ValueError("DDP execution-change resume requires submission metadata") | |
| submission = json.loads(submission_path.read_text(encoding="utf-8")) | |
| parent_value = submission.get("parent_run") | |
| if not parent_value: | |
| raise ValueError("DDP execution-change resume requires a parent-linked run") | |
| if isinstance(parent_value, dict): | |
| parent = Path(parent_value["run_dir"]).resolve() | |
| parent_manifest = json.loads( | |
| (parent / "_submission" / "manifest.json").read_text(encoding="utf-8") | |
| ) | |
| if parent_value.get("source_hash") != parent_manifest.get("copied_hash"): | |
| raise ValueError("declared parent source hash does not match its snapshot") | |
| else: | |
| parent = Path(parent_value).resolve() | |
| if checkpoint.resolve().parent.parent != parent: | |
| raise ValueError("resume checkpoint does not belong to the declared parent run") | |
| checkpoint_config = load_config(checkpoint / "config.yaml") | |
| # Verify the stored configuration in its original schema, before newly | |
| # introduced execution defaults are materialized by load_config. | |
| saved_config = yaml.safe_load((checkpoint / "config.yaml").read_text()) | |
| saved_config["run"]["output_dir"] = None | |
| saved_config["run"]["output_root"] = None | |
| saved_config["run"].pop("dataset") | |
| saved_config["run"].pop("category") | |
| saved_config["run"].pop("version", None) | |
| saved_config["checkpoint"]["resume_from"] = None | |
| saved_fingerprint = hashlib.sha256( | |
| json.dumps(saved_config, sort_keys=True, separators=(",", ":")).encode() | |
| ).hexdigest() | |
| if payload["config_fingerprint"] != saved_fingerprint: | |
| raise ValueError("checkpoint config fingerprint disagrees with checkpoint config") | |
| epoch_extension = config.checkpoint.resume_mode == "epoch_extension" | |
| if _ddp_resume_contract( | |
| checkpoint_config, epoch_extension=epoch_extension | |
| ) != _ddp_resume_contract(config, epoch_extension=epoch_extension): | |
| raise ValueError("DDP execution-change resume has a non-DDP configuration change") | |
| old_ddp = ( | |
| checkpoint_config.trainer.ddp_find_unused_parameters, | |
| checkpoint_config.trainer.ddp_static_graph, | |
| ) | |
| new_ddp = ( | |
| config.trainer.ddp_find_unused_parameters, | |
| config.trainer.ddp_static_graph, | |
| ) | |
| if epoch_extension: | |
| if config.trainer.epochs <= checkpoint_config.trainer.epochs: | |
| raise ValueError("epoch extension must increase trainer.epochs") | |
| if checkpoint_config.scheduler.kind != "constant": | |
| raise ValueError("epoch extension currently requires a constant scheduler") | |
| summary_path = parent / "summary.json" | |
| if not summary_path.is_file(): | |
| raise ValueError("epoch extension requires a completed parent summary") | |
| summary = json.loads(summary_path.read_text(encoding="utf-8")) | |
| successful = int(payload["trainer"]["successful_updates"]) | |
| if successful != int(summary["successful_updates"]): | |
| raise ValueError("epoch extension checkpoint is not the parent final step") | |
| if successful != int(summary["target_updates"]): | |
| raise ValueError("epoch extension parent did not complete its target") | |
| if config.checkpoint.resume_mode == "ddp_execution_change" and old_ddp == new_ddp: | |
| raise ValueError("DDP execution-change resume did not change the DDP policy") | |
| if ( | |
| config.checkpoint.resume_mode == "ddp_execution_change" | |
| and saved_config["trainer"].get("gradient_sync", "ddp") | |
| != config.trainer.gradient_sync | |
| ): | |
| raise ValueError("gradient_sync migration requires distributed_execution_change") | |
| if payload["world_size"] != world_size: | |
| raise ValueError("resume world_size mismatch") | |
| parent_record = json.loads((parent / "config.json").read_text(encoding="utf-8")) | |
| if ( | |
| parent_record.get("provenance_fingerprints") | |
| != payload["provenance_fingerprints"] | |
| ): | |
| raise ValueError("parent provenance does not match the checkpoint") | |
| old_provenance = dict(payload["provenance_fingerprints"]) | |
| new_provenance = dict(provenance_fingerprints) | |
| old_provenance.pop("source", None) | |
| new_provenance.pop("source", None) | |
| if old_provenance != new_provenance: | |
| raise ValueError("DDP execution-change resume has non-source provenance drift") | |
| def resolve_cuda_device_index(*, local_rank: int, world_size: int) -> int: | |
| """Map a queue's physical GPU assignment to the process-visible ordinal.""" | |
| if world_size > 1: | |
| return local_rank | |
| assigned = os.environ.get("GPUQ_ASSIGNED_ID") | |
| if assigned is None: | |
| return local_rank | |
| visible = [ | |
| value.strip() | |
| for value in os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",") | |
| if value.strip() | |
| ] | |
| if assigned in visible: | |
| return visible.index(assigned) | |
| # Some schedulers expose UUIDs instead of physical indices. A one-device | |
| # process can still unambiguously use logical cuda:0. | |
| if visible and torch.cuda.device_count() == 1: | |
| return 0 | |
| return int(assigned) | |
| def save_checkpoint( | |
| *, | |
| checkpoints_dir: Path, | |
| model, | |
| processor, | |
| optimizer, | |
| scheduler, | |
| trainer_state: dict, | |
| config: SFTConfig, | |
| provenance_fingerprints: dict[str, str], | |
| rank: int, | |
| world_size: int, | |
| ema_state: dict | None = None, | |
| control_group=None, | |
| ) -> Path: | |
| """Collect per-rank RNG and atomically publish a complete checkpoint.""" | |
| import torch.distributed as dist | |
| step = trainer_state["successful_updates"] | |
| final = checkpoints_dir / f"step-{step:06d}" | |
| if rank == 0: | |
| if final.exists(): | |
| raise FileExistsError(f"checkpoint already exists: {final}") | |
| temporary = checkpoints_dir / f".tmp-step-{step:06d}-{uuid.uuid4().hex}" | |
| temporary.mkdir(parents=True, exist_ok=False) | |
| name = str(temporary) | |
| else: | |
| name = "" | |
| names = [name] | |
| if world_size > 1: | |
| dist.broadcast_object_list(names, src=0, group=control_group) | |
| temporary = Path(names[0]) | |
| torch.save(rng_state(), temporary / f"rng-rank-{rank:02d}.pt") | |
| if world_size > 1: | |
| dist.barrier(group=control_group) | |
| if rank == 0: | |
| save_adapter(model, processor, temporary / "adapter") | |
| torch.save( | |
| { | |
| "trainer": trainer_state, | |
| "optimizer": optimizer.state_dict(), | |
| "scheduler": scheduler.state_dict(), | |
| "config_fingerprint": config.fingerprint(), | |
| "provenance_fingerprints": provenance_fingerprints, | |
| "world_size": world_size, | |
| "ema_teacher": ema_state, | |
| }, | |
| temporary / "training_state.pt", | |
| ) | |
| checkpoint_config = config.as_dict() | |
| checkpoint_config["checkpoint"]["resume_from"] = None | |
| (temporary / "config.yaml").write_text( | |
| yaml.safe_dump(checkpoint_config, sort_keys=False), encoding="utf-8" | |
| ) | |
| (temporary / "COMPLETE").write_text("complete\n", encoding="utf-8") | |
| temporary.rename(final) | |
| if world_size > 1: | |
| dist.barrier(group=control_group) | |
| return final | |
| def load_training_state( | |
| checkpoint: str | Path, | |
| *, | |
| optimizer, | |
| scheduler, | |
| config: SFTConfig, | |
| provenance_fingerprints: dict[str, str], | |
| rank: int, | |
| world_size: int, | |
| ema_state_loader: Callable[[dict], None] | None = None, | |
| run_dir: str | Path | None = None, | |
| ) -> dict: | |
| checkpoint = Path(checkpoint) | |
| if not (checkpoint / "COMPLETE").exists(): | |
| raise ValueError(f"checkpoint is incomplete: {checkpoint}") | |
| payload = torch.load( | |
| checkpoint / "training_state.pt", map_location="cpu", weights_only=False | |
| ) | |
| expected = { | |
| "config_fingerprint": config.fingerprint(), | |
| "provenance_fingerprints": provenance_fingerprints, | |
| "world_size": world_size, | |
| } | |
| mismatches = {key: value for key, value in expected.items() if payload[key] != value} | |
| if mismatches and config.checkpoint.resume_mode in { | |
| "ddp_execution_change", | |
| "distributed_execution_change", | |
| "epoch_extension", | |
| }: | |
| _validate_ddp_execution_resume( | |
| checkpoint, | |
| payload, | |
| config=config, | |
| provenance_fingerprints=provenance_fingerprints, | |
| world_size=world_size, | |
| run_dir=Path(run_dir) if run_dir is not None else None, | |
| ) | |
| elif mismatches: | |
| key, value = next(iter(mismatches.items())) | |
| raise ValueError( | |
| f"resume {key} mismatch: checkpoint={payload[key]!r}, current={value!r}" | |
| ) | |
| optimizer.load_state_dict(payload["optimizer"]) | |
| scheduler.load_state_dict(payload["scheduler"]) | |
| if ema_state_loader is not None: | |
| if payload.get("ema_teacher") is None: | |
| raise ValueError("resume checkpoint is missing EMA teacher state") | |
| ema_state_loader(payload["ema_teacher"]) | |
| rank_path = checkpoint / f"rng-rank-{rank:02d}.pt" | |
| if not rank_path.exists(): | |
| raise FileNotFoundError(f"rank RNG state is missing: {rank_path}") | |
| restore_rng_state(torch.load(rank_path, map_location="cpu", weights_only=False)) | |
| return payload["trainer"] | |
| __all__ = [ | |
| "acquire_run_writer", | |
| "cache_fingerprint", | |
| "claim_run_directory", | |
| "environment_record", | |
| "file_sha256", | |
| "load_training_state", | |
| "model_fingerprint", | |
| "resolve_cuda_device_index", | |
| "resolve_run_directory", | |
| "save_adapter", | |
| "save_adapter_atomic", | |
| "save_checkpoint", | |
| "seed_everything", | |
| "source_fingerprint", | |
| "write_json", | |
| ] | |
Xet Storage Details
- Size:
- 11.5 kB
- Xet hash:
- a2d2f58d1e42d27cd90ce8edb5a8f3f16cf5196a3ff9a0ae1e09acc7f4caf836
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.