diff --git a/.github/workflows/package-install.yml b/.github/workflows/package-install.yml index 4c07877ee..63150c0e7 100644 --- a/.github/workflows/package-install.yml +++ b/.github/workflows/package-install.yml @@ -42,7 +42,7 @@ jobs: cd "$project_dir" uv init --name art-base-install-smoke --python 3.12 --bare uv add "openpipe-art @ file://${wheel_path}" - uv run python -c "import importlib.util; assert importlib.util.find_spec('numpy') is None; assert importlib.util.find_spec('torch') is None; import art; from art.pipeline_trainer import PipelineTrainer; print(art.__name__, PipelineTrainer.__name__)" + uv run python -c "import importlib.util; assert importlib.util.find_spec('numpy') is not None; assert importlib.util.find_spec('torch') is None; import art; from art.pipeline_trainer import PipelineTrainer; print(art.__name__, PipelineTrainer.__name__)" uv add "weave==0.52.41" uv run python -c "from weave.trace.settings import override_settings; import art" @@ -58,7 +58,9 @@ jobs: project_dir="$(mktemp -d)" cd "$project_dir" uv init --name art-install-smoke --python 3.12 --bare - uv add "openpipe-art[backend] @ file://${wheel_path}" + uv add --index https://download.pytorch.org/whl/cu128 \ + --index-strategy unsafe-best-match \ + "openpipe-art[backend] @ file://${wheel_path}" uv sync - name: Smoke test distributed wheel surface @@ -80,7 +82,9 @@ jobs: --index-url https://download.pytorch.org/whl/cpu \ "torch==2.11.0" uv pip install --python .venv/bin/python \ - "openpipe-art[distributed] @ file://${wheel_path}" + "openpipe-art @ file://${wheel_path}" \ + 'aiohttp>=3.13.0' 'msgspec>=0.21.0' 'torchmonarch==0.6.0' \ + 'transformers>=5.2.0,<=5.12.1' uv pip install --python .venv/bin/python --no-deps \ "openpipe-art[distributed,megatron] @ file://${wheel_path}" diff --git a/.github/workflows/trainer-rank-gpu.yml b/.github/workflows/trainer-rank-gpu.yml index bac8bab20..47b268da0 100644 --- a/.github/workflows/trainer-rank-gpu.yml +++ b/.github/workflows/trainer-rank-gpu.yml @@ -45,7 +45,7 @@ jobs: git cat-file -e "${BASE_SHA}^{commit}" 2>/dev/null \ || git fetch --no-tags origin "${BASE_SHA}" - pattern='^src/art/megatron/(_hybrid_ep/|hybrid_ep_setup\.py|setup\.sh|prefix_tree(_packing|_state)?\.py|context_parallel/|flex_attn/|gdn/|lora\.py|megatron_patches\.py|training/(finalize_grads|microbatches)\.py)' + pattern='^src/art/(trainer_rank/|megatron/(_hybrid_ep/|hybrid_ep_setup\.py|setup\.sh|prefix_tree(_packing|_state)?\.py|context_parallel/|flex_attn/|gdn/|lora\.py|megatron_patches\.py|training/(finalize_grads|microbatches)\.py))' changed="$(git diff --name-only "${BASE_SHA}...${HEAD_SHA}")" critical="$(printf '%s\n' "${changed}" | grep -E "${pattern}" || true)" if [ -n "${critical}" ]; then diff --git a/dev/trainer_rank_checkpoint_acceptance.py b/dev/trainer_rank_checkpoint_acceptance.py new file mode 100644 index 000000000..41257fed0 --- /dev/null +++ b/dev/trainer_rank_checkpoint_acceptance.py @@ -0,0 +1,107 @@ +"""Exercise canonical TrainerRank checkpoints under ``torchrun``. + +The driver runs this module repeatedly with different rank counts. ``step-save`` +loads a LoRA/checkpoint, applies one deterministic optimizer step, and saves a +canonical checkpoint. ``step-export`` applies the same step and exports LoRA +weights for cross-topology comparison. +""" + +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import json +import os +from pathlib import Path +import time + +import torch +import torch.distributed as dist + +from art.trainer_rank import AdamParams, TrainerRank + + +def _args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("operation", choices=("step-save", "step-export", "load")) + parser.add_argument("--source", required=True) + parser.add_argument("--output") + parser.add_argument("--model", default="Qwen/Qwen3.5-4B") + parser.add_argument("--layers", type=int, default=1) + parser.add_argument("--grad", type=float, default=1e-4) + parser.add_argument("--output-json") + return parser.parse_args() + + +def _digest(path: Path) -> str: + digest = hashlib.sha256() + for item in sorted(path.rglob("*")): + if item.is_file(): + digest.update(item.relative_to(path).as_posix().encode()) + digest.update(item.read_bytes()) + return digest.hexdigest() + + +async def _load(trainer: TrainerRank, source: str) -> None: + await trainer.load_checkpoint(source) + + +def main() -> None: + args = _args() + os.environ.setdefault("ART_MEGATRON_TENSOR_MODEL_PARALLEL_SIZE", "1") + os.environ.setdefault("ART_MEGATRON_CONTEXT_PARALLEL_SIZE", "1") + os.environ.setdefault("ART_MEGATRON_PIPELINE_MODEL_PARALLEL_SIZE", "1") + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group("nccl") + rank = dist.get_rank() + started = time.perf_counter() + try: + from art.megatron import train as megatron_train + + runtime = megatron_train.build_training_runtime( + model_identifier=args.model, + provider_configure=lambda provider: setattr( + provider, "num_layers", args.layers + ), + print_env=rank == 0, + ) + trainer = TrainerRank(runtime) + asyncio.run(_load(trainer, args.source)) + loaded = time.perf_counter() + if args.operation != "load": + slot = trainer._checkpoint_slots[args.source] + for parameter in slot.params: + parameter.grad = torch.full_like(parameter, args.grad) + metrics = trainer.optim_step( + params=AdamParams(learning_rate=3e-4, grad_clip_norm=0), + scale_grads=1 / dist.get_world_size(), + ) + if metrics["update_successful"] != 1: + raise RuntimeError(f"optimizer step failed: {metrics}") + if args.output is None: + raise ValueError("--output is required") + if args.operation == "step-save": + trainer.save_checkpoint(args.output) + else: + trainer.export_lora(args.output) + dist.barrier() + if rank == 0: + output = None if args.output is None else Path(args.output) + payload = { + "world_size": dist.get_world_size(), + "load_seconds": loaded - started, + "total_seconds": time.perf_counter() - started, + "output_digest": None if output is None else _digest(output), + "peak_gpu_bytes": torch.cuda.max_memory_allocated(), + } + encoded = json.dumps(payload, sort_keys=True) + print(encoded, flush=True) + if args.output_json: + Path(args.output_json).write_text(encoded + "\n") + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/examples/multinode/program.py b/examples/multinode/program.py index 070ce226b..06f2dd7e6 100644 --- a/examples/multinode/program.py +++ b/examples/multinode/program.py @@ -65,7 +65,10 @@ async def main(hosts: Any) -> None: ) executor.set_workers(workers) model = art.TrainableModel( - name="multinode-smoke", project="art", base_model="not-loaded" + name="multinode-smoke", + project="art", + base_model="not-loaded", + run_name="multinode-smoke", ) trajectories: list[art.Trajectory] = [] for answer in REWARDS: diff --git a/pyproject.toml b/pyproject.toml index 6c71cfa8f..f84e1b6a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "polars>=1.26.0", "tblib>=3.0.0", "nest-asyncio>=1.6.0", + "numpy<2; python_version < '3.13'", "setproctitle>=1.3.6", ] diff --git a/scripts/ci/trainer-rank-checkpoint-acceptance.sky.yaml b/scripts/ci/trainer-rank-checkpoint-acceptance.sky.yaml new file mode 100644 index 000000000..247cd6b11 --- /dev/null +++ b/scripts/ci/trainer-rank-checkpoint-acceptance.sky.yaml @@ -0,0 +1,54 @@ +name: trainer-rank-checkpoint-acceptance + +workdir: . + +resources: + accelerators: H200:2 + image_id: docker:docker.io/bradhiltonnw/art-gpu:latest + +setup: | + uv sync --frozen --extra megatron --group dev + uv run --no-sync python -m art.megatron.hybrid_ep_setup + +run: | + set -euo pipefail + result=/tmp/art-checkpoint-acceptance + fixture="$(find . -maxdepth 1 -type d -name '.live-checkpoint-fixture.*' -print -quit)" + test -n "${fixture}" + rm -rf "${result}" + mkdir -p "${result}" + + exercise() { + ranks="$1" + operation="$2" + source="$3" + output="$4" + uv run --no-sync torchrun --standalone --nproc-per-node="${ranks}" \ + dev/trainer_rank_checkpoint_acceptance.py "${operation}" \ + --source "${source}" --output "${output}" \ + --output-json "${output}.json" + } + + exercise 1 step-save "${fixture}" "${result}/source-1" + exercise 1 step-export "${result}/source-1" "${result}/restore-1-to-1" + exercise 2 step-export "${result}/source-1" "${result}/restore-1-to-2" + diff -qr "${result}/restore-1-to-1" "${result}/restore-1-to-2" + + exercise 2 step-save "${fixture}" "${result}/source-2" + exercise 2 step-export "${result}/source-2" "${result}/restore-2-to-2" + exercise 1 step-export "${result}/source-2" "${result}/restore-2-to-1" + diff -qr "${result}/restore-2-to-2" "${result}/restore-2-to-1" + cat "${result}"/*.json + +config: + kubernetes: + pod_config: + spec: + schedulerName: binpack-scheduler + activeDeadlineSeconds: 3600 + containers: + - name: ray-node + imagePullPolicy: Always + env: + - name: UV_LINK_MODE + value: copy diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index 15c413f56..6037506a3 100644 --- a/src/art/trainer_rank/__init__.py +++ b/src/art/trainer_rank/__init__.py @@ -1,28 +1,14 @@ from __future__ import annotations -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Literal, TypedDict, cast, overload +import asyncio +from collections.abc import Callable, Iterable, Iterator, Sequence +from typing import TYPE_CHECKING, Literal, cast, overload import torch import torch.distributed as dist - -class TrainerRankOptimizerLayout(TypedDict): - parallel: tuple[int, int, int, int, int, int, int, int] - parameters: tuple[ - tuple[tuple[int, ...], str, str, bool, int | None, str, tuple[int, ...]], - ..., - ] - - -class TrainerRankOptimizerState(TypedDict): - format_version: Literal[1] - layout: TrainerRankOptimizerLayout - master_params: tuple[torch.Tensor, ...] - optimizer: dict[str, object] - - -from . import _impl # noqa: E402 +from . import _impl +from ._checkpoint import CheckpointManifest, materialize_lora, validate_checkpoint AdapterSelection = _impl.AdapterSelection AdamParams = _impl.AdamParams @@ -42,7 +28,8 @@ class TrainerRankOptimizerState(TypedDict): TrainerRankMemoryError = _impl.TrainerRankMemoryError TrainerRankSlotStateError = _impl.TrainerRankSlotStateError Unset = _impl.Unset -_PushedSlot = _impl._PushedSlot +MaterializedCheckpoint = _impl.MaterializedCheckpoint +PushedCheckpoint = _impl.PushedCheckpoint if TYPE_CHECKING: from art.megatron.train import TrainingRuntime @@ -56,9 +43,9 @@ class TrainerRankOptimizerState(TypedDict): MicroBatchStats, TopK, TrainerRankMemoryError, - TrainerRankOptimizerLayout, - TrainerRankOptimizerState, TrainerRankSlotStateError, + MaterializedCheckpoint, + PushedCheckpoint, ): _public_type.__module__ = __name__ del _public_type @@ -85,55 +72,51 @@ def __init__( def zero_grad(self) -> None: super().zero_grad() - def set_checkpoint(self, name: str | None) -> None: - super().set_checkpoint(name) + def prefetch_checkpoints( + self, + *checkpoints: str | MaterializedCheckpoint, + ) -> asyncio.Task[None]: + return super().prefetch_checkpoints(*checkpoints) - def set_lora(self, name: str | None) -> None: - super().set_lora(name) + def load_checkpoint( + self, checkpoint: str | MaterializedCheckpoint | None + ) -> asyncio.Task[None]: + return super().load_checkpoint(checkpoint) - def push_checkpoint(self, name: str | None) -> _PushedSlot: - return super().push_checkpoint(name) + def push_checkpoint( + self, checkpoint: str | MaterializedCheckpoint | None + ) -> PushedCheckpoint: + return super().push_checkpoint(checkpoint) - def push_lora(self, name: str | None) -> _PushedSlot: - return super().push_lora(name) + def pop_checkpoint(self) -> None: + super().pop_checkpoint() - def pop_pushed_lora_or_checkpoint(self) -> None: - super().pop_pushed_lora_or_checkpoint() + def save_checkpoint( + self, + output_dir: str, + checkpoint_path: str | Literal["active"] = "active", + ) -> None: + super().save_checkpoint(output_dir, checkpoint_path) - def load_checkpoint_slot( + def prepare_checkpoint_save( self, - name: str, - adapter_model: Mapping[str, torch.Tensor], - *, - optimizer_state: TrainerRankOptimizerState | None = None, - alpha: float | None = None, - adapter_config: Mapping[str, object] | None = None, - ) -> int: - return super().load_checkpoint_slot( - name, - adapter_model, - optimizer_state=optimizer_state, - alpha=alpha, - adapter_config=adapter_config, - ) + output_dir: str, + checkpoint_path: str | Literal["active"] = "active", + ) -> None: + super().prepare_checkpoint_save(output_dir, checkpoint_path) - def checkpoint_slot_optimizer_state( - self, name: str - ) -> TrainerRankOptimizerState | None: - return super().checkpoint_slot_optimizer_state(name) + def finish_checkpoint_save(self, output_dir: str) -> None: + super().finish_checkpoint_save(output_dir) - def save_checkpoint_slot_lora(self, name: str, output_dir: str) -> None: - """Collectively publish a trained checkpoint slot as a vLLM LoRA.""" - super().save_checkpoint_slot_lora(name, output_dir) + def abort_checkpoint_save(self, output_dir: str) -> None: + super().abort_checkpoint_save(output_dir) - def load_lora_slot( + def export_lora( self, - name: str, - adapter_model: Mapping[str, torch.Tensor], - *, - alpha: float | None = None, + output_dir: str, + checkpoint_path: str | Literal["active"] = "active", ) -> int: - return super().load_lora_slot(name, adapter_model, alpha=alpha) + return super().export_lora(output_dir, checkpoint_path) @overload def forward_micro_batches( @@ -283,15 +266,18 @@ def optim_step( __all__ = [ "AdapterSelection", "AdamParams", + "CheckpointManifest", "ForwardInput", "ForwardOutput", "MicroBatch", "MicroBatchStats", + "MaterializedCheckpoint", + "materialize_lora", "TopK", "TrainerRank", "TrainerRankMemoryError", - "TrainerRankOptimizerLayout", - "TrainerRankOptimizerState", + "PushedCheckpoint", "TrainerRankSlotStateError", "Unset", + "validate_checkpoint", ] diff --git a/src/art/trainer_rank/_checkpoint.py b/src/art/trainer_rank/_checkpoint.py new file mode 100644 index 000000000..4ab7a01cd --- /dev/null +++ b/src/art/trainer_rank/_checkpoint.py @@ -0,0 +1,1387 @@ +"""Topology-portable persistence for TrainerRank checkpoints.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +import hashlib +import importlib +import json +import os +from pathlib import Path, PurePosixPath, PureWindowsPath +import shutil +import struct +import threading +from typing import TYPE_CHECKING, Literal, TypedDict, cast +import uuid + +import torch +import torch.distributed as dist + +if TYPE_CHECKING: + from art.megatron.lora import LoRA, LoraShardMeta, LoRASlotRef + from art.trainer_rank._impl import ( + TrainerRank, + _AdapterConfig, + _DynamicOptimizer, + ) + +FORMAT = 1 +MANIFEST_FILE = "checkpoint.json" +_ART_FORMAT_KEY = "art_lora_format" +_ART_FORMAT = "art-trainer-rank-v1" + + +class OptimizerConfig(TypedDict): + learning_rate: float + beta1: float + beta2: float + eps: float + weight_decay: float + + +class CheckpointManifest(TypedDict): + format_version: Literal[1] + base_model_name_or_path: str + optimizer: OptimizerConfig | None + parameters: dict[str, list[str]] + steps: dict[str, float] + files: dict[str, str] + digest: str + + +@dataclass(frozen=True) +class PreparedCheckpoint: + path: Path + config: dict[str, object] + keys: tuple[str, ...] + manifest: CheckpointManifest | None + digest: str + + +@dataclass(frozen=True) +class LocalOptimizerState: + masters: tuple[torch.Tensor, ...] + exp_avgs: tuple[torch.Tensor, ...] + exp_avg_sqs: tuple[torch.Tensor, ...] + steps: tuple[float, ...] + config: OptimizerConfig + + +@dataclass(frozen=True) +class _LocalShard: + metadata: LoraShardMeta + file: str + + +@dataclass(frozen=True) +class _PreparedSave: + sequence: int + snapshot: Path + reservation: Path + destination: Path + config: dict[str, object] + shards: tuple[_LocalShard, ...] + optimizer: OptimizerConfig | None + + +@dataclass(frozen=True) +class _FinalizedSave: + sequence: int + outcome: Literal["finish", "abort"] + + +type _SlotSnapshot = tuple[ + tuple[ + "LoRA", + dict["LoRASlotRef", str], + dict[str, torch.nn.Module], + dict[str, "LoRASlotRef"], + ], + ..., +] + + +def _distributed() -> bool: + return dist.is_available() and dist.is_initialized() + + +def _rank() -> int: + return dist.get_rank() if _distributed() else 0 + + +def _gather[T](value: T, group: dist.ProcessGroup | None = None) -> tuple[T, ...]: + if not _distributed(): + return (value,) + values: list[T | None] = [None] * dist.get_world_size(group) + dist.all_gather_object(values, value, group=group) + return tuple(cast(T, item) for item in values) + + +def raise_distributed( + error: BaseException | None, + phase: str, + group: dist.ProcessGroup | None = None, +) -> None: + errors = _gather(None if error is None else repr(error), group) + if not any(errors): + return + if error is not None: + raise error + raise RuntimeError( + f"Another rank failed to {phase}: {next(item for item in errors if item)}" + ) + + +def _safe_relative(value: str) -> PurePosixPath: + path = PurePosixPath(value) + if ( + not value + or path.is_absolute() + or ".." in path.parts + or PureWindowsPath(value).drive + or "\\" in value + ): + raise RuntimeError(f"Unsafe checkpoint path: {value!r}") + return path + + +def _hash_files(root: Path, files: Iterable[str], *, seed: bytes = b"") -> str: + digest = hashlib.blake2b(digest_size=32) + digest.update(seed) + for relative in sorted(files): + digest.update(relative.encode()) + with (root / _safe_relative(relative)).open("rb") as handle: + while chunk := handle.read(8 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _manifest_seed(manifest: Mapping[str, object]) -> bytes: + value = {**manifest, "digest": ""} + return json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + + +def _file_digest(path: Path) -> str: + digest = hashlib.blake2b(digest_size=32) + with path.open("rb") as handle: + while chunk := handle.read(8 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _manifest_digest(manifest: Mapping[str, object]) -> str: + return hashlib.blake2b(_manifest_seed(manifest), digest_size=32).hexdigest() + + +def _validate_manifest( + manifest: CheckpointManifest, + *, + adapter_keys: set[str], + config: Mapping[str, object], +) -> set[str]: + if manifest.get("format_version") != FORMAT: + raise RuntimeError("Unsupported ART checkpoint format") + digest = manifest.get("digest") + file_digests = manifest.get("files") + if ( + not isinstance(digest, str) + or len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + or not isinstance(file_digests, dict) + or any( + not isinstance(path, str) + or not isinstance(value, str) + or len(value) != 64 + or any(character not in "0123456789abcdef" for character in value) + for path, value in file_digests.items() + ) + ): + raise RuntimeError("Checkpoint digest is invalid") + if manifest.get("base_model_name_or_path") != config.get("base_model_name_or_path"): + raise RuntimeError( + "Checkpoint manifest and adapter config name different models" + ) + optimizer = manifest.get("optimizer") + parameters = manifest.get("parameters") + steps = manifest.get("steps") + if not isinstance(parameters, dict) or not isinstance(steps, dict): + raise RuntimeError("Checkpoint optimizer mapping is invalid") + files: set[str] = set() + if optimizer is None: + if parameters or steps: + raise RuntimeError("LoRA-only checkpoint contains optimizer metadata") + else: + required = {"learning_rate", "beta1", "beta2", "eps", "weight_decay"} + optimizer_values = cast(dict[str, object], optimizer) + if ( + not isinstance(optimizer, dict) + or set(optimizer_values) != required + or any( + not isinstance(optimizer_values[key], int | float) + or isinstance(optimizer_values[key], bool) + for key in required + ) + ): + raise RuntimeError("Checkpoint optimizer config is invalid") + if set(parameters) != adapter_keys or set(steps) != adapter_keys: + raise RuntimeError( + "Checkpoint optimizer mapping differs from adapter tensors: " + f"parameters={sorted(set(parameters) ^ adapter_keys)[:8]} " + f"steps={sorted(set(steps) ^ adapter_keys)[:8]}" + ) + for key, record in parameters.items(): + if ( + not isinstance(key, str) + or not isinstance(record, list | tuple) + or len(record) != 3 + or not all(isinstance(item, str) for item in record) + ): + raise RuntimeError( + f"Checkpoint optimizer mapping is invalid for {key!r}" + ) + normalized = [_safe_relative(item).as_posix() for item in record] + parameters[key] = normalized + files.update(normalized) + if any( + not isinstance(value, int | float) or isinstance(value, bool) + for value in steps.values() + ): + raise RuntimeError("Checkpoint optimizer steps are invalid") + expected_files = { + "adapter_config.json", + "adapter_model.safetensors", + *files, + } + if set(file_digests) != expected_files: + raise RuntimeError("Checkpoint file digest mapping is invalid") + return files + + +def prepare_checkpoint( + path: str, *, artifact_entries: Iterable[str] | None = None +) -> PreparedCheckpoint: + root = Path(path).resolve(strict=True) + if not root.is_dir(): + raise FileNotFoundError(f"Checkpoint is not a directory: {path}") + from art.megatron.model_support.lora_disk import load_adapter_config, safe_open + + config = cast(dict[str, object], load_adapter_config(root)) + adapter = root / "adapter_model.safetensors" + with safe_open(adapter, framework="pt") as handle: + keys = tuple(sorted(handle.keys())) + manifest_path = root / MANIFEST_FILE + manifest: CheckpointManifest | None = None + if manifest_path.is_file(): + value = json.loads(manifest_path.read_text()) + if not isinstance(value, dict) or value.get("format_version") != FORMAT: + raise RuntimeError("Unsupported ART checkpoint format") + manifest = cast(CheckpointManifest, value) + if config.get(_ART_FORMAT_KEY) != _ART_FORMAT: + raise RuntimeError("Canonical checkpoint adapter format is invalid") + optimizer_files = _validate_manifest( + manifest, adapter_keys=set(keys), config=config + ) + files = { + "adapter_config.json", + "adapter_model.safetensors", + MANIFEST_FILE, + *optimizer_files, + } + expected = manifest["digest"] + actual = _manifest_digest(manifest) + if actual != expected: + raise RuntimeError(f"Checkpoint digest mismatch: {actual} != {expected}") + if artifact_entries is None: + downloaded = files - {MANIFEST_FILE} + else: + available = {_safe_relative(entry).as_posix() for entry in artifact_entries} + if missing := sorted(files - available): + raise RuntimeError( + f"Checkpoint artifact is missing entries: {missing[:8]}" + ) + downloaded = {"adapter_config.json", "adapter_model.safetensors"} + for relative in downloaded: + file_actual = _file_digest(root / relative) + if file_actual != manifest["files"][relative]: + raise RuntimeError( + f"Checkpoint file digest mismatch for {relative}: " + f"{file_actual} != {manifest['files'][relative]}" + ) + else: + if artifact_entries is not None: + raise RuntimeError("Checkpoint artifact lacks a canonical manifest") + actual = _hash_files(root, ("adapter_config.json", "adapter_model.safetensors")) + return PreparedCheckpoint(root, config, keys, manifest, actual) + + +def validate_checkpoint( + path: str | Path, *, require_optimizer: bool = False +) -> CheckpointManifest | None: + prepared = prepare_checkpoint(str(path)) + if require_optimizer and ( + prepared.manifest is None or prepared.manifest["optimizer"] is None + ): + raise RuntimeError("Checkpoint does not contain optimizer state") + return prepared.manifest + + +def materialize_lora( + path: str | Path, + output_dir: str | Path, + *, + require_optimizer: bool = False, + artifact_entries: Iterable[str] | None = None, + expected_digest: str | None = None, +) -> None: + source = prepare_checkpoint(str(path), artifact_entries=artifact_entries) + if expected_digest is not None and source.digest != expected_digest: + raise RuntimeError( + f"Checkpoint digest mismatch: {source.digest} != {expected_digest}" + ) + if require_optimizer and ( + source.manifest is None or source.manifest["optimizer"] is None + ): + raise RuntimeError("Checkpoint does not contain optimizer state") + destination = Path(output_dir) + if destination.exists() and any(destination.iterdir()): + raise FileExistsError(f"LoRA output directory is not empty: {destination}") + destination.mkdir(parents=True, exist_ok=True) + for name in ("adapter_config.json", "adapter_model.safetensors"): + shutil.copy2(source.path / name, destination / name) + from art.megatron.model_support.lora_disk import normalize_lora_checkpoint_to_vllm + + normalize_lora_checkpoint_to_vllm(destination) + + +def _optimizer_config(dynamic: _DynamicOptimizer) -> OptimizerConfig: + group = dynamic.optimizer.param_groups[0] + beta1, beta2 = group["betas"] + return { + "learning_rate": float(group["lr"]), + "beta1": float(beta1), + "beta2": float(beta2), + "eps": float(group["eps"]), + "weight_decay": float(group["weight_decay"]), + } + + +def _validate_save_state(trainer: TrainerRank, name: str) -> _AdapterConfig: + slot = trainer._checkpoint_slots.get(name) + if slot is None or slot.config is None: + raise trainer._slot_state_error(f"Unknown checkpoint: {name!r}") + if trainer._checkpoint_grad_flags((name,))[0]: + raise trainer._slot_state_error( + f"Checkpoint {name!r} has accumulated gradients" + ) + return slot.config + + +def _local_state( + trainer: TrainerRank, name: str, snapshot: Path +) -> tuple[tuple[_LocalShard, ...], OptimizerConfig | None]: + from art.megatron.lora import LoRA + from art.megatron.weights.lora_publish import collect_local_lora_entries + + ref = trainer._slot_ref(name) + tensors, metadata = collect_local_lora_entries( + trainer.runtime.model, {}, owner_rank=_rank(), slot_ref=ref + ) + dynamic = trainer._checkpoint_slots[name].optimizer + optimizer = None if dynamic is None else _optimizer_config(dynamic) + masters = ( + {} + if dynamic is None + else { + id(param): master + for param, master in zip( + trainer._checkpoint_slots[name].params, + dynamic.master_params, + strict=True, + ) + } + ) + by_key = {item.key: item for item in metadata} + payloads: dict[str, dict[str, torch.Tensor]] = {} + metadata_by_block: dict[str, list[LoraShardMeta]] = {} + for item in metadata: + payloads.setdefault(item.block, {})[f"lora/{item.key}"] = ( + tensors[item.key].cpu().contiguous() + ) + metadata_by_block.setdefault(item.block, []).append(item) + if dynamic is not None: + for chunk in trainer.runtime.model: + for module in chunk.modules(): + if not isinstance(module, LoRA): + continue + for key, param, expert in module._export_items(ref): + item = by_key.get(key) + if item is None: + continue + master = masters[id(param)] + state = dynamic.optimizer.state.get(master, {}) + values = ( + master, + cast(torch.Tensor | None, state.get("exp_avg")), + cast(torch.Tensor | None, state.get("exp_avg_sq")), + ) + for component, value in zip( + ("master", "exp_avg", "exp_avg_sq"), values, strict=True + ): + value = torch.zeros_like(master) if value is None else value + local = value if expert is None else value[expert] + payloads[item.block][f"{component}/{key}"] = ( + local.T.float().cpu().contiguous() + ) + step = state.get("step", 0.0) + payloads[item.block][f"step/{key}"] = torch.tensor(float(step)) + records: list[_LocalShard] = [] + for index, block in enumerate(sorted(payloads)): + relative = f"block-{index:06d}.safetensors" + importlib.import_module("safetensors.torch").save_file( + payloads[block], snapshot / relative + ) + records.extend(_LocalShard(item, relative) for item in metadata_by_block[block]) + return tuple(records), optimizer + + +def prepare_checkpoint_save( + trainer: TrainerRank, output_dir: str, checkpoint_name: str +) -> None: + with trainer._checkpoint_prepare_lock: + group = _ensure_group(trainer) + identity = (output_dir, checkpoint_name) + if any(value != identity for value in _gather(identity, group)): + raise RuntimeError("Checkpoint save identity differs across ranks") + with trainer._checkpoint_save_condition: + pending = ( + output_dir in trainer._checkpoint_preparing_saves + or output_dir in trainer._prepared_checkpoint_saves + ) + if any(value != pending for value in _gather(pending, group)): + raise RuntimeError( + f"Checkpoint save state differs across ranks: {output_dir}" + ) + if pending: + raise RuntimeError(f"Checkpoint save is already pending: {output_dir}") + with trainer._checkpoint_save_condition: + trainer._checkpoint_preparing_saves.add(output_dir) + try: + known = ( + checkpoint_name in trainer._checkpoint_slots + and trainer._checkpoint_slots[checkpoint_name].config is not None + ) + if not all(_gather(known, group)): + raise trainer._slot_state_error( + f"Unknown checkpoint on at least one rank: {checkpoint_name!r}" + ) + config = deepcopy(_validate_save_state(trainer, checkpoint_name)) + if any(value != config for value in _gather(config, group)): + raise trainer._slot_state_error( + f"Checkpoint {checkpoint_name!r} configuration differs across ranks" + ) + except BaseException: + with trainer._checkpoint_save_condition: + trainer._checkpoint_preparing_saves.discard(output_dir) + raise + destination = Path(output_dir) + reservation = destination.with_name(f".{destination.name}.reserved") + snapshot = destination.with_name( + f".{destination.name}.snapshot-r{_rank()}-{uuid.uuid4().hex}" + ) + error: BaseException | None = None + prepared: _PreparedSave | None = None + shards: tuple[_LocalShard, ...] | None = None + optimizer: OptimizerConfig | None = None + reservation_created = False + with trainer._checkpoint_save_condition: + sequence = trainer._checkpoint_save_sequence + trainer._checkpoint_save_sequence += 1 + if any(value != sequence for value in _gather(sequence, group)): + with trainer._checkpoint_save_condition: + trainer._checkpoint_preparing_saves.discard(output_dir) + with trainer._checkpoint_save_condition: + trainer._checkpoint_save_skipped.add(sequence) + _advance_save_queue(trainer, sequence) + raise RuntimeError("Checkpoint save order differs across ranks") + try: + if _rank() == 0: + reservation.mkdir(parents=True) + reservation_created = True + snapshot.mkdir(parents=True) + shards, optimizer = _local_state(trainer, checkpoint_name, snapshot) + except BaseException as exc: + error = exc + try: + raise_distributed(error, "prepare checkpoint", group) + if any(value != optimizer for value in _gather(optimizer, group)): + raise trainer._slot_state_error( + f"Checkpoint {checkpoint_name!r} optimizer differs across ranks" + ) + assert shards is not None + prepared = _PreparedSave( + sequence, + snapshot, + reservation, + destination, + dict(config), + shards, + optimizer, + ) + except BaseException as failure: + cleanup = _cleanup_paths( + [snapshot, *([reservation] if reservation_created else [])] + ) + with trainer._checkpoint_save_condition: + trainer._checkpoint_preparing_saves.discard(output_dir) + trainer._checkpoint_save_skipped.add(sequence) + _advance_save_queue(trainer, sequence) + cleanup_failure: BaseException | None = None + try: + raise_distributed(cleanup, "clean up checkpoint preparation", group) + except BaseException as exc: + cleanup_failure = exc + if cleanup_failure is not None: + raise BaseExceptionGroup( + "checkpoint preparation and cleanup both failed", + [failure, cleanup_failure], + ) from None + raise failure + assert prepared is not None + with trainer._checkpoint_save_condition: + trainer._prepared_checkpoint_saves[output_dir] = prepared + trainer._finalized_checkpoint_saves.pop(output_dir, None) + trainer._checkpoint_preparing_saves.discard(output_dir) + trainer._checkpoint_save_condition.notify_all() + + +def _read_snapshot( + prepared: _PreparedSave, relative: str, prefix: str, keys: Iterable[str] +) -> dict[str, torch.Tensor]: + load = importlib.import_module("safetensors.torch").load_file + payload = load(prepared.snapshot / relative) + return {key: payload[f"{prefix}/{key}"] for key in keys} + + +def _merge_component( + prepared: _PreparedSave, + metadata: Sequence[LoraShardMeta], + component: str, + group: dist.ProcessGroup | None, +) -> dict[str, torch.Tensor]: + from art.megatron.weights.lora_publish import merge_sharded_adapter_entries + + owned = [item for item in metadata if item.owner_rank == _rank()] + local: dict[str, torch.Tensor] = {} + error: BaseException | None = None + try: + if owned: + files = { + record.file for record in prepared.shards if record.metadata in owned + } + for relative in files: + keys = [ + item.key + for item in owned + if next( + record.file + for record in prepared.shards + if record.metadata == item + ) + == relative + ] + local.update(_read_snapshot(prepared, relative, component, keys)) + except BaseException as exc: + error = exc + raise_distributed(error, f"read checkpoint {component} block", group) + exchanged: dict[tuple[int, str], torch.Tensor] = {} + for item in sorted(metadata, key=lambda value: (value.owner_rank, value.key)): + identity = (item.owner_rank, item.key) + if _rank() == item.owner_rank: + tensor = local[item.key].contiguous() + if _rank() == 0: + exchanged[identity] = tensor + else: + dist.send(tensor, dst=0, group=group) + elif _rank() == 0: + dtype = ( + getattr(torch, item.dtype_name) + if component == "lora" + else torch.float32 + ) + tensor = torch.empty(item.shape, dtype=dtype) + dist.recv(tensor, src=item.owner_rank, group=group) + exchanged[identity] = tensor + entries: dict[str, list[tuple[dict[str, object], torch.Tensor]]] = {} + merged: dict[str, torch.Tensor] = {} + error = None + if _rank() == 0: + try: + for item in metadata: + entries.setdefault(item.key, []).append( + (item.manifest, exchanged[(item.owner_rank, item.key)]) + ) + merged = merge_sharded_adapter_entries(entries) # type: ignore[arg-type] + except BaseException as exc: + error = exc + raise_distributed(error, f"merge checkpoint {component} block", group) + return merged + + +def _consolidate(shards: Sequence[Path], output: Path) -> None: + sources: dict[str, tuple[Path, int, int, int, dict[str, object]]] = {} + for shard in shards: + with shard.open("rb") as handle: + header_size = struct.unpack(" None: + error: BaseException | None = None + if _rank() == 0: + try: + action() + except BaseException as exc: + error = exc + raise_distributed(error, phase, group) + + +def _finish(trainer: TrainerRank, prepared: _PreparedSave) -> None: + from art.megatron.model_support.lora_disk import save_adapter_config + + group = _ensure_finalize_group(trainer) + metadata = [item for values in _gather(prepared.shards, group) for item in values] + identities: set[tuple[str, int]] = set() + selected: list[LoraShardMeta] = [] + for item in sorted(metadata, key=lambda value: value.metadata.owner_rank): + identity = (item.metadata.key, int(item.metadata.manifest.get("shard_rank", 0))) + if identity not in identities: + identities.add(identity) + selected.append(item.metadata) + blocks = sorted({item.block for item in selected}) + temporary = prepared.destination.with_name( + f".{prepared.destination.name}.tmp-{uuid.uuid4().hex}" + ) + _rank_zero_phase( + lambda: temporary.mkdir(parents=True), "create checkpoint output", group + ) + parameters: dict[str, list[str]] = {} + steps: dict[str, float] = {} + lora_shards: list[Path] = [] + try: + for index, block in enumerate(blocks): + block_metadata = [item for item in selected if item.block == block] + lora = _merge_component(prepared, block_metadata, "lora", group) + relative = f".adapter-{index:06d}.safetensors" + _rank_zero_phase( + lambda: importlib.import_module("safetensors.torch").save_file( + lora, temporary / relative + ), + "write checkpoint adapter block", + group, + ) + if _rank() == 0: + lora_shards.append(temporary / relative) + if prepared.optimizer is None: + continue + files: list[str] = [] + for component in ("master", "exp_avg", "exp_avg_sq"): + tensors = _merge_component(prepared, block_metadata, component, group) + relative = f"optimizer/{component}-{index:06d}.safetensors" + + def write_optimizer_block() -> None: + (temporary / "optimizer").mkdir(exist_ok=True) + importlib.import_module("safetensors.torch").save_file( + tensors, temporary / relative + ) + + _rank_zero_phase( + write_optimizer_block, "write checkpoint optimizer block", group + ) + files.append(relative) + if _rank() == 0: + for key in (item.key for item in block_metadata): + parameters[key] = list(files) + owned = [item for item in block_metadata if item.owner_rank == _rank()] + local_steps: dict[str, float] = {} + error: BaseException | None = None + try: + for relative in { + record.file + for record in prepared.shards + if record.metadata in owned + }: + load = importlib.import_module("safetensors.torch").load_file + payload = load(prepared.snapshot / relative) + local_steps.update( + (key.removeprefix("step/"), float(value.item())) + for key, value in payload.items() + if key.startswith("step/") + ) + except BaseException as exc: + error = exc + raise_distributed(error, "read checkpoint optimizer steps", group) + step_values: dict[str, set[float]] = {} + for values in _gather(local_steps, group): + for key, value in values.items(): + step_values.setdefault(key, set()).add(value) + if mismatched := { + key: values for key, values in step_values.items() if len(values) != 1 + }: + raise trainer._slot_state_error( + f"Optimizer shard steps differ: {mismatched}" + ) + steps.update((key, values.pop()) for key, values in step_values.items()) + + def commit() -> None: + _consolidate(lora_shards, temporary / "adapter_model.safetensors") + for shard in lora_shards: + shard.unlink() + save_adapter_config( + temporary, {**prepared.config, _ART_FORMAT_KEY: _ART_FORMAT} + ) + manifest: CheckpointManifest = { + "format_version": FORMAT, + "base_model_name_or_path": str( + prepared.config["base_model_name_or_path"] + ), + "optimizer": prepared.optimizer, + "parameters": parameters, + "steps": steps, + "files": {}, + "digest": "", + } + artifact_files = { + "adapter_config.json", + "adapter_model.safetensors", + *(file for record in parameters.values() for file in record), + } + manifest["files"] = { + relative: _file_digest(temporary / relative) + for relative in artifact_files + } + manifest["digest"] = _manifest_digest(manifest) + (temporary / MANIFEST_FILE).write_text( + json.dumps(manifest, indent=2) + "\n" + ) + if prepared.destination.exists(): + if ( + prepare_checkpoint(str(prepared.destination)).digest + != manifest["digest"] + ): + raise FileExistsError( + "Checkpoint path already contains different state: " + f"{prepared.destination}" + ) + shutil.rmtree(temporary) + else: + os.replace(temporary, prepared.destination) + + _rank_zero_phase(commit, "commit checkpoint", group) + except BaseException: + if _rank() == 0: + shutil.rmtree(temporary, ignore_errors=True) + raise + + +def _advance_save_queue(trainer: TrainerRank, sequence: int) -> None: + with trainer._checkpoint_save_condition: + if sequence == trainer._checkpoint_save_next: + trainer._checkpoint_save_next += 1 + while trainer._checkpoint_save_next in trainer._checkpoint_save_skipped: + trainer._checkpoint_save_skipped.remove(trainer._checkpoint_save_next) + trainer._checkpoint_save_next += 1 + trainer._checkpoint_save_condition.notify_all() + + +def _cleanup_paths(paths: Iterable[Path]) -> BaseException | None: + errors: list[BaseException] = [] + for path in paths: + try: + shutil.rmtree(path) + except FileNotFoundError: + pass + except BaseException as exc: + errors.append(exc) + return BaseExceptionGroup("checkpoint cleanup failed", errors) if errors else None + + +def _claim_finalization( + trainer: TrainerRank, + output_dir: str, + action: Literal["finish", "abort"], +) -> _PreparedSave | None: + with trainer._checkpoint_save_condition: + while True: + prepared = trainer._prepared_checkpoint_saves.get(output_dir) + if prepared is None: + if output_dir in trainer._finalized_checkpoint_saves: + return None + if action == "abort": + return None + raise RuntimeError(f"Checkpoint save was not prepared: {output_dir}") + outcome = trainer._checkpoint_save_outcomes.get(output_dir) + if outcome is not None and outcome != action: + raise RuntimeError( + f"Checkpoint save was already {outcome}ed: {output_dir}" + ) + if output_dir in trainer._checkpoint_finalizing_saves: + trainer._checkpoint_save_condition.wait() + continue + if outcome is None and prepared.sequence != trainer._checkpoint_save_next: + raise RuntimeError( + "Checkpoint saves must be finalized in preparation order: " + f"expected sequence {trainer._checkpoint_save_next}, got " + f"{prepared.sequence}" + ) + trainer._checkpoint_finalizing_saves[output_dir] = action + return prepared + + +def _finalize_checkpoint_save( + trainer: TrainerRank, + output_dir: str, + action: Literal["finish", "abort"], +) -> None: + group = _ensure_finalize_group(trainer) + with trainer._checkpoint_finalize_lock: + with trainer._checkpoint_save_condition: + local = trainer._prepared_checkpoint_saves.get(output_dir) + finalized = trainer._finalized_checkpoint_saves.get(output_dir) + sequence = ( + local.sequence + if local is not None + else None + if finalized is None + else finalized.sequence + ) + outcome = ( + trainer._checkpoint_save_outcomes.get(output_dir) + if finalized is None + else finalized.outcome + ) + states = _gather((action, output_dir, sequence, outcome), group) + if any(not isinstance(value, tuple) or len(value) != 4 for value in states): + raise RuntimeError("Checkpoint finalization protocol is out of sync") + if any(value[:3] != states[0][:3] for value in states): + raise RuntimeError("Checkpoint save actions differ across ranks") + outcomes = {value[3] for value in states} + if len(outcomes) != 1: + raise RuntimeError("Checkpoint save outcomes differ across ranks") + if sequence is None: + if action == "abort": + return + raise RuntimeError(f"Checkpoint save was not prepared: {output_dir}") + finalized_ranks = _gather(finalized is not None, group) + if all(finalized_ranks): + if outcome == "finish" or action == "abort": + return + raise RuntimeError(f"Checkpoint save was already {outcome}ed: {output_dir}") + if outcome is not None and outcome != action: + raise RuntimeError(f"Checkpoint save was already {outcome}ed: {output_dir}") + prepared = ( + _claim_finalization(trainer, output_dir, action) + if finalized is None + else None + ) + assert prepared is not None or finalized is not None + error: BaseException | None = None + cleanup_failed = True + try: + if finalized is None and outcome is None and action == "finish": + try: + assert prepared is not None + _finish(trainer, prepared) + except BaseException as exc: + error = exc + if finalized is None and outcome is None: + assert prepared is not None + outcome = action if error is None else "abort" + with trainer._checkpoint_save_condition: + trainer._checkpoint_save_outcomes[output_dir] = outcome + if outcome == "abort": + trainer._checkpoint_save_skipped.add(prepared.sequence) + _advance_save_queue(trainer, prepared.sequence) + cleanup = None + if prepared is not None: + paths = [prepared.snapshot] + if _rank() == 0: + paths.append(prepared.reservation) + cleanup = _cleanup_paths(paths) + try: + failures = _gather( + ( + None if error is None else repr(error), + None if cleanup is None else repr(cleanup), + ), + group, + ) + except BaseException as exc: + local_failures = [ + *([error] if error is not None else []), + *([cleanup] if cleanup is not None else []), + exc, + ] + if len(local_failures) == 1: + raise local_failures[0] + raise BaseExceptionGroup( + "checkpoint finalization and coordination failed", local_failures + ) from None + if any( + not isinstance(failure, tuple) + or len(failure) != 2 + or any( + value is not None and not isinstance(value, str) + for value in failure + ) + for failure in failures + ): + raise RuntimeError("Checkpoint finalization protocol is out of sync") + cleanup_failed = any( + cleanup_error is not None for _, cleanup_error in failures + ) + local_failures = [ + *([error] if error is not None else []), + *([cleanup] if cleanup is not None else []), + ] + if local_failures: + if len(local_failures) == 1: + raise local_failures[0] + raise BaseExceptionGroup( + "checkpoint finalization failed", local_failures + ) + if remote := next((failure for failure in failures if any(failure)), None): + raise RuntimeError( + f"Another rank failed to {action} checkpoint: {remote}" + ) + finally: + with trainer._checkpoint_save_condition: + trainer._checkpoint_finalizing_saves.pop(output_dir, None) + if not cleanup_failed: + trainer._prepared_checkpoint_saves.pop(output_dir, None) + trainer._checkpoint_save_outcomes.pop(output_dir, None) + assert outcome is not None + trainer._finalized_checkpoint_saves[output_dir] = _FinalizedSave( + sequence, outcome + ) + trainer._checkpoint_save_condition.notify_all() + + +def finish_checkpoint_save(trainer: TrainerRank, output_dir: str) -> None: + _finalize_checkpoint_save(trainer, output_dir, "finish") + + +def abort_checkpoint_save(trainer: TrainerRank, output_dir: str) -> None: + _finalize_checkpoint_save(trainer, output_dir, "abort") + + +def _load_adapter( + trainer: TrainerRank, source: PreparedCheckpoint, keys: Iterable[str] +) -> dict[str, torch.Tensor]: + if source.manifest is None: + from art.megatron.model_support.lora_disk import ( + load_lora_tensors_for_megatron, + ) + + loaded = load_lora_tensors_for_megatron( + source.path, handler=trainer.runtime.model_support_handler + ) + return {key: value for key, value in loaded.items() if key in set(keys)} + safe_open = importlib.import_module("safetensors").safe_open + with safe_open(source.path / "adapter_model.safetensors", framework="pt") as handle: + available = set(handle.keys()) + return {key: handle.get_tensor(key) for key in keys if key in available} + + +def _localized( + module: LoRA, tensor: torch.Tensor, parameter: torch.nn.Parameter +) -> torch.Tensor: + return module._localized_weight(tensor, into=parameter).contiguous() + + +def _slot_snapshot(trainer: TrainerRank) -> _SlotSnapshot: + from art.megatron.lora import LoRA + + return tuple( + ( + module, + dict(module._slot_keys), + dict(module._slot_modules.items()), + { + key: cast(LoRASlotRef, getattr(slot, "ref")) + for key, slot in module._slot_modules.items() + }, + ) + for chunk in trainer.runtime.model + for module in chunk.modules() + if isinstance(module, LoRA) + ) + + +def _restore_slots(snapshot: _SlotSnapshot) -> None: + for module, keys, slots, refs in snapshot: + for key, slot in slots.items(): + setattr(slot, "ref", refs[key]) + module._slot_keys = keys + module._slot_modules = torch.nn.ModuleDict(slots) + + +def _commit_slot(trainer: TrainerRank, source: str, destination: str) -> None: + from art.megatron.lora import LoRA + + source_ref = trainer._slot_ref(source) + destination_ref = trainer._slot_ref(destination) + for chunk in trainer.runtime.model: + for module in chunk.modules(): + if not isinstance(module, LoRA): + continue + source_key = module._slot_keys.pop(source_ref, None) + destination_key = module._slot_keys.pop(destination_ref, None) + if source_key is None: + if destination_key is not None: + del module._slot_modules[destination_key] + continue + slot = module._slot_modules[source_key] + setattr(slot, "ref", destination_ref) + target_key = destination_key or source_key + module._slot_keys[destination_ref] = target_key + if target_key != source_key: + module._slot_modules[target_key] = slot + del module._slot_modules[source_key] + + +def _optimizer_state( + trainer: TrainerRank, source: PreparedCheckpoint, name: str +) -> LocalOptimizerState: + assert source.manifest is not None and source.manifest["optimizer"] is not None + from art.megatron.lora import LoRA + + ref = trainer._slot_ref(name) + components: dict[str, list[torch.Tensor]] = { + "master": [], + "exp_avg": [], + "exp_avg_sq": [], + } + steps: list[float] = [] + sites: list[tuple[LoRA, str, torch.nn.Parameter, list[str], list[list[str]]]] = [] + file_keys: dict[str, set[str]] = {} + for chunk in trainer.runtime.model: + for module in chunk.modules(): + if not isinstance(module, LoRA) or module._slot(ref) is None: + continue + for suffix, parameter in module._lora_params(ref): + suffix = suffix.removesuffix(".weight") + keys = [ + key + for key in module._expected_weight_keys(suffix) + if isinstance(key, str) + ] + records = [source.manifest["parameters"][key] for key in keys] + sites.append((module, suffix, parameter, keys, records)) + for record in records: + for filename in record: + file_keys.setdefault(filename, set()).update(keys) + safe_open = importlib.import_module("safetensors").safe_open + loaded: dict[str, dict[str, torch.Tensor]] = {} + for filename, keys in file_keys.items(): + with safe_open(source.path / filename, framework="pt") as handle: + loaded[filename] = {key: handle.get_tensor(key) for key in keys} + for module, suffix, parameter, keys, records in sites: + if not keys: + for component in components.values(): + component.append(torch.zeros_like(parameter)) + steps.append(0.0) + continue + for index, component in enumerate(components): + tensors = { + key: loaded[record[index]][key] + for key, record in zip(keys, records, strict=True) + } + full = module._adapter_weight(tensors, suffix=suffix) + components[component].append(_localized(module, full, parameter)) + key_steps = {source.manifest["steps"][key] for key in keys} + if len(key_steps) != 1: + raise RuntimeError(f"Optimizer steps differ for {keys}") + steps.append(key_steps.pop()) + return LocalOptimizerState( + tuple(components["master"]), + tuple(components["exp_avg"]), + tuple(components["exp_avg_sq"]), + tuple(steps), + source.manifest["optimizer"], + ) + + +def _phase[T]( + action: Callable[[], T], phase: str, group: dist.ProcessGroup | None +) -> T: + result: T | None = None + error: BaseException | None = None + try: + result = action() + except BaseException as exc: + error = exc + raise_distributed(error, phase, group) + return cast(T, result) + + +def _validate_base_model( + trainer: TrainerRank, + source: PreparedCheckpoint, + config: Mapping[str, object], +) -> None: + configured = str(config["base_model_name_or_path"]) + if ( + source.manifest is not None + and source.manifest["base_model_name_or_path"] != configured + ): + raise trainer._slot_state_error( + "Checkpoint manifest and adapter config name different base models" + ) + runtime_model = getattr(trainer.runtime, "model_identifier", None) + if runtime_model is not None and runtime_model != configured: + raise trainer._slot_state_error( + f"Checkpoint base model {configured!r} differs from runtime model " + f"{runtime_model!r}" + ) + supported = tuple( + getattr(getattr(trainer.runtime, "model_support_spec", None), "model_names", ()) + ) + if supported and configured not in supported: + raise trainer._slot_state_error( + f"Checkpoint base model {configured!r} is incompatible with this runtime" + ) + + +def _rollback_load( + trainer: TrainerRank, + snapshot: _SlotSnapshot, + temporary: str, + name: str, + previous: object, + group: dist.ProcessGroup | None, +) -> None: + def rollback() -> None: + _restore_slots(snapshot) + trainer._checkpoint_slots.pop(temporary, None) + if previous is None: + trainer._checkpoint_slots.pop(name, None) + else: + from art.trainer_rank._impl import _CheckpointSlot + + trainer._checkpoint_slots[name] = cast(_CheckpointSlot, previous) + + _phase(rollback, "roll back checkpoint load", group) + + +def load_checkpoint( + trainer: TrainerRank, source: PreparedCheckpoint, name: str +) -> None: + group = _ensure_group(trainer) + if any(value != source.digest for value in _gather(source.digest, group)): + raise trainer._slot_state_error( + f"Checkpoint {name!r} content differs across ranks" + ) + config = _phase( + lambda: trainer._validate_checkpoint_adapter_config( + name, source.config, alpha=None + ), + "validate checkpoint config", + group, + ) + assert config is not None + if any(value != config for value in _gather(config, group)): + raise trainer._slot_state_error( + f"Checkpoint {name!r} configuration differs across ranks" + ) + _phase( + lambda: _validate_base_model(trainer, source, config), + "validate checkpoint base model", + group, + ) + _phase( + lambda: trainer._guard_slot_can_load(trainer._slot_ref(name)), + "validate checkpoint target", + group, + ) + local_keys = trainer._local_lora_adapter_templates() + adapter = _phase( + lambda: _load_adapter(trainer, source, local_keys), + "read checkpoint adapter", + group, + ) + prepared_adapter = _phase( + lambda: trainer._prepare_adapter_model( + name, adapter, canonicalized=source.manifest is not None + ), + "localize checkpoint adapter", + group, + ) + expected = {key for keys in _gather(tuple(prepared_adapter), group) for key in keys} + if source.manifest is not None and expected != set(source.keys): + raise trainer._slot_state_error( + "Checkpoint tensor coverage differs from runtime" + ) + temporary = f"__art_loading_{uuid.uuid4().hex}" + snapshot = _slot_snapshot(trainer) + previous = trainer._checkpoint_slots.get(name) + try: + loaded = _phase( + lambda: trainer._load_checkpoint_slot( + temporary, + prepared_adapter, + alpha=float(config["lora_alpha"]), + _prepared=True, + ), + "stage checkpoint adapter", + group, + ) + params = _phase( + lambda: trainer._validate_checkpoint_consistency( + temporary, loaded, expected + ), + "validate staged checkpoint", + group, + ) + from art.trainer_rank._impl import _CheckpointSlot + + trainer._checkpoint_slots[temporary] = _CheckpointSlot(params, config) + _phase( + lambda: trainer._validate_loaded_checkpoint_config(temporary, config), + "validate loaded checkpoint config", + group, + ) + if source.manifest is not None and source.manifest["optimizer"] is not None: + optimizer_state = _phase( + lambda: _optimizer_state(trainer, source, temporary), + "read checkpoint optimizer", + group, + ) + trainer._checkpoint_slots[temporary].optimizer = _phase( + lambda: trainer._restore_canonical_optimizer( + temporary, optimizer_state + ), + "restore checkpoint optimizer", + group, + ) + + def commit() -> None: + _commit_slot(trainer, temporary, name) + staged = trainer._checkpoint_slots.pop(temporary) + staged.revision = 0 if previous is None else previous.revision + 1 + trainer._checkpoint_slots[name] = staged + + _phase(commit, "commit checkpoint", group) + except BaseException: + _rollback_load(trainer, snapshot, temporary, name, previous, group) + raise + + +def export_lora(trainer: TrainerRank, output_dir: str, checkpoint_name: str) -> int: + group = _ensure_group(trainer) + slot = None + error: BaseException | None = None + try: + slot = trainer._checkpoint_slots.get(checkpoint_name) + if slot is None: + raise ValueError(f"Unknown checkpoint: {checkpoint_name!r}") + if slot.config is None: + raise trainer._slot_state_error( + f"Checkpoint {checkpoint_name!r} has no adapter_config" + ) + except BaseException as exc: + error = exc + raise_distributed(error, "validate LoRA export", group) + assert slot is not None and slot.config is not None + identity = (dict(slot.config), slot.revision) + if any(value != identity for value in _gather(identity, group)): + raise trainer._slot_state_error( + f"Checkpoint {checkpoint_name!r} differs across ranks" + ) + from art.megatron.weights.lora_publish import save_vllm_lora_from_model + + error = None + try: + save_vllm_lora_from_model( + model=trainer.runtime.model, + adapter_dtypes={}, + handler=trainer.runtime.model_support_handler, + adapter_config=dict(slot.config), + output_dir=output_dir, + rank=trainer.runtime.rank, + world_size=trainer.runtime.world_size, + slot_ref=trainer._slot_ref(checkpoint_name), + ) + except BaseException as exc: + error = exc + raise_distributed(error, "export LoRA", group) + return slot.revision + + +def _ensure_groups( + trainer: TrainerRank, +) -> tuple[dist.ProcessGroup | None, dist.ProcessGroup | None]: + if not hasattr(trainer, "_checkpoint_process_group"): + trainer._checkpoint_process_group = None + if not hasattr(trainer, "_checkpoint_finalize_process_group"): + trainer._checkpoint_finalize_process_group = None + if not hasattr(trainer, "_checkpoint_group_lock"): + trainer._checkpoint_group_lock = threading.Lock() + if _distributed(): + with trainer._checkpoint_group_lock: + if trainer._checkpoint_process_group is None: + trainer._checkpoint_process_group = dist.new_group(backend="gloo") + if trainer._checkpoint_finalize_process_group is None: + trainer._checkpoint_finalize_process_group = dist.new_group( + backend="gloo" + ) + return ( + trainer._checkpoint_process_group, + trainer._checkpoint_finalize_process_group, + ) + + +def _ensure_group(trainer: TrainerRank) -> dist.ProcessGroup | None: + return _ensure_groups(trainer)[0] + + +def _ensure_finalize_group(trainer: TrainerRank) -> dist.ProcessGroup | None: + return _ensure_groups(trainer)[1] diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index c0007f598..bc11710b4 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -2,7 +2,11 @@ from __future__ import annotations +import asyncio from collections.abc import ( + Awaitable, + Callable, + Generator, Iterable, Iterator, Mapping, @@ -11,11 +15,14 @@ from copy import deepcopy from dataclasses import dataclass import os +from pathlib import Path +import threading from types import TracebackType from typing import ( TYPE_CHECKING, Generic, Literal, + NotRequired, Self, TypedDict, TypeVar, @@ -47,7 +54,12 @@ from art.megatron.lora import LoRASlotRef from art.megatron.prefix_tree_state import PrefixTreeAttentionState from art.megatron.train import TrainingRuntime - from art.trainer_rank import TrainerRankOptimizerLayout, TrainerRankOptimizerState + from art.trainer_rank._checkpoint import ( + LocalOptimizerState, + PreparedCheckpoint, + _FinalizedSave, + _PreparedSave, + ) @dataclass(frozen=True) @@ -76,9 +88,14 @@ class TopK: class _AdapterConfig(TypedDict): base_model_name_or_path: str + revision: NotRequired[str | None] r: int lora_alpha: float target_modules: str | list[str] + num_attention_heads: NotRequired[int] + num_key_value_heads: NotRequired[int] + head_dim: NotRequired[int] + hidden_size: NotRequired[int] class _Unset: @@ -91,7 +108,6 @@ class _Unset: @dataclass(frozen=True) class _LocalLoRASlotRef: - kind: Literal["checkpoint", "lora"] name: str | None @@ -111,7 +127,6 @@ class ForwardInput(Generic[LogprobsT, TopKT, LogitsT, HiddenStatesT]): logits: bool = False hidden_states: bool = False checkpoint: AdapterSelection = Unset - lora: AdapterSelection = Unset @overload def __new__( @@ -123,7 +138,6 @@ def __new__( logits: Literal[False] = False, hidden_states: Literal[False] = False, checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[None, None, None, None]": ... @overload @@ -136,7 +150,6 @@ def __new__( logits: Literal[False] = False, hidden_states: Literal[False] = False, checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[torch.Tensor, None, None, None]": ... @overload @@ -149,7 +162,6 @@ def __new__( logits: Literal[False] = False, hidden_states: Literal[False] = False, checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[None, TopK, None, None]": ... @overload @@ -162,7 +174,6 @@ def __new__( logits: Literal[True], hidden_states: Literal[False] = False, checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[None, None, torch.Tensor, None]": ... @overload @@ -175,7 +186,6 @@ def __new__( logits: Literal[False] = False, hidden_states: Literal[True], checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[None, None, None, torch.Tensor]": ... @overload @@ -188,7 +198,6 @@ def __new__( logits: Literal[False] = False, hidden_states: Literal[False] = False, checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[torch.Tensor, TopK, None, None]": ... @overload @@ -201,7 +210,6 @@ def __new__( logits: Literal[True], hidden_states: Literal[False] = False, checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[torch.Tensor, None, torch.Tensor, None]": ... @overload @@ -214,7 +222,6 @@ def __new__( logits: Literal[False] = False, hidden_states: Literal[True], checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[torch.Tensor, None, None, torch.Tensor]": ... @overload @@ -227,7 +234,6 @@ def __new__( logits: Literal[True], hidden_states: Literal[False] = False, checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[None, TopK, torch.Tensor, None]": ... @overload @@ -240,7 +246,6 @@ def __new__( logits: Literal[False] = False, hidden_states: Literal[True], checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[None, TopK, None, torch.Tensor]": ... @overload @@ -253,7 +258,6 @@ def __new__( logits: Literal[True], hidden_states: Literal[True], checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[None, None, torch.Tensor, torch.Tensor]": ... @overload @@ -266,7 +270,6 @@ def __new__( logits: Literal[True], hidden_states: Literal[False] = False, checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[torch.Tensor, TopK, torch.Tensor, None]": ... @overload @@ -279,7 +282,6 @@ def __new__( logits: Literal[False] = False, hidden_states: Literal[True], checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[torch.Tensor, TopK, None, torch.Tensor]": ... @overload @@ -292,7 +294,6 @@ def __new__( logits: Literal[True], hidden_states: Literal[True], checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[torch.Tensor, None, torch.Tensor, torch.Tensor]": ... @overload @@ -305,7 +306,6 @@ def __new__( logits: Literal[True], hidden_states: Literal[True], checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[None, TopK, torch.Tensor, torch.Tensor]": ... @overload @@ -318,7 +318,6 @@ def __new__( logits: Literal[True], hidden_states: Literal[True], checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[torch.Tensor, TopK, torch.Tensor, torch.Tensor]": ... @overload @@ -331,7 +330,6 @@ def __new__( logits: bool = False, hidden_states: bool = False, checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> "ForwardInput[torch.Tensor | None, TopK | None, torch.Tensor | None, torch.Tensor | None]": ... def __new__( @@ -343,15 +341,12 @@ def __new__( logits: bool = False, hidden_states: bool = False, checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> Self: return object.__new__(cls) def __post_init__(self) -> None: if self.top_k is not None and self.top_k < 1: raise ValueError("top_k must be >= 1") - if self.checkpoint is not Unset and self.lora is not Unset: - raise ValueError("ForwardInput cannot set both checkpoint and lora") type AnyForwardInput = ForwardInput[ @@ -454,27 +449,106 @@ class _DynamicOptimizer: master_params: tuple[torch.nn.Parameter, ...] -@dataclass(frozen=True) -class _PushedSlot: - trainer: "TrainerRank" - ref: "LoRASlotRef" +@dataclass +class _CheckpointSlot: + params: tuple[torch.nn.Parameter, ...] = () + config: _AdapterConfig | None = None + optimizer: _DynamicOptimizer | None = None + revision: int = 0 + - def __enter__(self) -> "_PushedSlot": +@dataclass(frozen=True) +class MaterializedCheckpoint: + """A logical checkpoint and its rank-local materialized directory.""" + + path: str + directory: str + + +@dataclass +class PushedCheckpoint: + _trainer: "TrainerRank" + _path: str | None + _directory: str | None + _task: asyncio.Task[None] | None = None + _entered: bool = False + _closed: bool = False + + def __await__(self) -> Generator[object, None, None]: + return self._ensure_task().__await__() + + def __enter__(self) -> "PushedCheckpoint": + if self._entered or self._closed: + raise RuntimeError("Pushed checkpoint context cannot be entered twice") + if self._task is not None: + if not self._task.done(): + raise RuntimeError( + "Checkpoint push is running asynchronously; use 'async with'" + ) + self._task.result() + else: + self._trainer._push_checkpoint_sync(self._path, self._directory) + self._entered = True return self def __exit__( self, - exc_type: type[BaseException] | None, - exc_value: BaseException | None, + exception_type: type[BaseException] | None, + exception: BaseException | None, traceback: TracebackType | None, ) -> bool: - if not self.trainer._slot_stack or self.trainer._slot_stack[-1] != self.ref: - raise RuntimeError( - "Pushed LoRA/checkpoint stack changed before context exit" - ) - self.trainer.pop_pushed_lora_or_checkpoint() + self._exit(exception) return False + async def __aenter__(self) -> "PushedCheckpoint": + if self._entered or self._closed: + raise RuntimeError("Pushed checkpoint context cannot be entered twice") + task = self._ensure_task() + try: + await task + except asyncio.CancelledError: + if task.done() and not task.cancelled() and task.exception() is None: + self._entered = True + self._pop() + raise + self._entered = True + return self + + async def __aexit__( + self, + exception_type: type[BaseException] | None, + exception: BaseException | None, + traceback: TracebackType | None, + ) -> bool: + self._exit(exception) + return False + + def _ensure_task(self) -> asyncio.Task[None]: + if self._task is None: + self._task = self._trainer._activate_checkpoint(self._path, self._directory) + return self._task + + def _exit(self, body_error: BaseException | None) -> None: + try: + self._pop() + except BaseException as pop_error: + if body_error is not None: + raise BaseExceptionGroup( + "checkpoint context body and cleanup both failed", + [body_error, pop_error], + ) from None + raise + + def _pop(self) -> None: + if not self._entered: + return + ref = self._trainer._slot_ref(self._path) + if not self._trainer._slot_stack or self._trainer._slot_stack[-1] != ref: + raise RuntimeError("Pushed checkpoint stack changed before context exit") + self._trainer.pop_checkpoint() + self._entered = False + self._closed = True + @dataclass(frozen=True) class _ForwardItem: @@ -569,11 +643,23 @@ def __init__( ) self._default_slot_ref: LoRASlotRef | None = None self._slot_stack: list[LoRASlotRef] = [] - self._dynamic_optimizers: dict[str, _DynamicOptimizer] = {} - self._checkpoint_slot_params_by_name: dict[ - str, tuple[torch.nn.Parameter, ...] - ] = {} - self._checkpoint_slot_adapter_configs: dict[str, _AdapterConfig] = {} + self._checkpoint_slots: dict[str, _CheckpointSlot] = {} + self._checkpoint_prefetches: dict[str, asyncio.Task[PreparedCheckpoint]] = {} + self._checkpoint_mutation_tail: asyncio.Task[None] | None = None + self._checkpoint_process_group: dist.ProcessGroup | None = None + self._checkpoint_finalize_process_group: dist.ProcessGroup | None = None + self._checkpoint_group_lock = threading.Lock() + self._checkpoint_prepare_lock = threading.Lock() + self._checkpoint_finalize_lock = threading.Lock() + self._checkpoint_save_condition = threading.Condition() + self._checkpoint_save_sequence = 0 + self._checkpoint_save_next = 0 + self._checkpoint_save_skipped: set[int] = set() + self._checkpoint_preparing_saves: set[str] = set() + self._checkpoint_finalizing_saves: dict[str, Literal["finish", "abort"]] = {} + self._checkpoint_save_outcomes: dict[str, Literal["finish", "abort"]] = {} + self._prepared_checkpoint_saves: dict[str, _PreparedSave] = {} + self._finalized_checkpoint_saves: dict[str, _FinalizedSave] = {} self._pending_slot_graphs: dict[ LoRASlotRef, list[weakref.ReferenceType[torch.Tensor]] ] = {} @@ -593,126 +679,237 @@ def zero_grad(self) -> None: optimizer = self.runtime.optimizer if optimizer is not None: optimizer.zero_grad() - for params in self._checkpoint_slot_params_by_name.values(): - for param in params: + for slot in self._checkpoint_slots.values(): + for param in slot.params: param.grad = None self._prune_slot_graphs() - def set_checkpoint(self, name: str | None) -> None: - self._set_default_slot(self._slot_ref("checkpoint", name)) + def prefetch_checkpoints( + self, *checkpoints: str | MaterializedCheckpoint + ) -> asyncio.Task[None]: + sources = tuple( + self._checkpoint_source(checkpoint)[1] for checkpoint in checkpoints + ) + assert all(source is not None for source in sources) + + async def prefetch() -> None: + await asyncio.gather( + *( + self._prefetch_checkpoint(source) + for source in sources + if source is not None + ) + ) + + return asyncio.create_task(prefetch()) + + def load_checkpoint( + self, checkpoint: str | MaterializedCheckpoint | None + ) -> asyncio.Task[None]: + logical, source = self._checkpoint_source(checkpoint) + return self._load_checkpoint(logical, source) - def set_lora(self, name: str | None) -> None: - self._set_default_slot(self._slot_ref("lora", name)) + def _load_checkpoint( + self, logical_path: str | None, source_path: str | None + ) -> asyncio.Task[None]: + prefetch = ( + None + if source_path is None + else asyncio.create_task(self._prefetch_checkpoint(source_path)) + ) + + async def load() -> None: + if self._slot_stack: + raise RuntimeError("Cannot load a checkpoint while one is pushed") + if logical_path is None: + self._set_default_slot(self._slot_ref(None)) + return + assert source_path is not None and prefetch is not None + await self._load_checkpoint_path( + logical_path, source_path=source_path, prefetch=prefetch + ) + self._set_default_slot(self._slot_ref(logical_path)) + + return self._checkpoint_mutation_task(load) + + def push_checkpoint( + self, checkpoint: str | MaterializedCheckpoint | None + ) -> PushedCheckpoint: + logical, directory = self._checkpoint_source(checkpoint) + return PushedCheckpoint(self, logical, directory) + + def _activate_checkpoint( + self, logical_path: str | None, source_path: str | None + ) -> asyncio.Task[None]: + prefetch = ( + asyncio.create_task(self._prefetch_checkpoint(source_path)) + if source_path is not None and logical_path not in self._checkpoint_slots + else None + ) - def push_checkpoint(self, name: str | None) -> _PushedSlot: - ref = self._slot_ref("checkpoint", name) - self._slot_stack.append(ref) - return _PushedSlot(self, ref) + async def push() -> None: + if prefetch is not None and logical_path not in self._checkpoint_slots: + assert logical_path is not None and source_path is not None + await self._load_checkpoint_path( + logical_path, source_path=source_path, prefetch=prefetch + ) + self._slot_stack.append(self._slot_ref(logical_path)) - def push_lora(self, name: str | None) -> _PushedSlot: - ref = self._slot_ref("lora", name) - self._slot_stack.append(ref) - return _PushedSlot(self, ref) + return self._checkpoint_mutation_task(push) - def pop_pushed_lora_or_checkpoint(self) -> None: + def _push_checkpoint_sync( + self, logical_path: str | None, source_path: str | None + ) -> None: + predecessor = self._checkpoint_mutation_tail + if predecessor is not None: + if not predecessor.done(): + raise RuntimeError( + "A checkpoint mutation is running asynchronously; use 'async with'" + ) + if not predecessor.cancelled(): + predecessor.exception() + if source_path is not None and logical_path not in self._checkpoint_slots: + assert logical_path is not None + from . import _checkpoint + + source = _checkpoint.prepare_checkpoint(source_path) + _checkpoint.load_checkpoint(self, source, logical_path) + self._slot_stack.append(self._slot_ref(logical_path)) + + def _checkpoint_mutation_task( + self, operation: Callable[[], Awaitable[None]] + ) -> asyncio.Task[None]: + predecessor = self._checkpoint_mutation_tail + + async def ordered() -> None: + if predecessor is not None: + try: + await asyncio.shield(predecessor) + except asyncio.CancelledError: + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + except Exception: + pass + await operation() + + task = asyncio.create_task(ordered()) + self._checkpoint_mutation_tail = task + return task + + def pop_checkpoint(self) -> None: if not self._slot_stack: - raise RuntimeError("No pushed LoRA or checkpoint to pop") + raise RuntimeError("No pushed checkpoint to pop") self._slot_stack.pop() - def load_checkpoint_slot( + def save_checkpoint( self, - name: str, - adapter_model: Mapping[str, torch.Tensor], - *, - optimizer_state: TrainerRankOptimizerState | None = None, - alpha: float | None = None, - adapter_config: Mapping[str, object] | None = None, - ) -> int: - config = self._validate_checkpoint_slot_adapter_config( - name, adapter_config, alpha=alpha - ) - loaded = self._load_slot( - "checkpoint", - name, - adapter_model, - trainable=True, - alpha=alpha if config is None else float(config["lora_alpha"]), + output_dir: str, + checkpoint_path: str | Literal["active"] = "active", + ) -> None: + self.prepare_checkpoint_save(output_dir, checkpoint_path) + self.finish_checkpoint_save(output_dir) + + def prepare_checkpoint_save( + self, + output_dir: str, + checkpoint_path: str | Literal["active"] = "active", + ) -> None: + from . import _checkpoint + + _checkpoint.prepare_checkpoint_save( + self, output_dir, self._resolve_checkpoint_name(checkpoint_path) ) - slot_params = self._validate_dynamic_slot_consistency( - "checkpoint", name, loaded + + def finish_checkpoint_save(self, output_dir: str) -> None: + from . import _checkpoint + + _checkpoint.finish_checkpoint_save(self, output_dir) + + def abort_checkpoint_save(self, output_dir: str) -> None: + from . import _checkpoint + + _checkpoint.abort_checkpoint_save(self, output_dir) + + def export_lora( + self, + output_dir: str, + checkpoint_path: str | Literal["active"] = "active", + ) -> int: + from . import _checkpoint + + return _checkpoint.export_lora( + self, output_dir, self._resolve_checkpoint_name(checkpoint_path) ) - if config is not None: - self._validate_loaded_checkpoint_slot_config(name, config) - self._checkpoint_slot_params_by_name[name] = slot_params - if optimizer_state is None: - self._dynamic_optimizers.pop(name, None) - else: - self._dynamic_optimizers[name] = self._restore_dynamic_optimizer( - name, optimizer_state - ) - configs = getattr(self, "_checkpoint_slot_adapter_configs", None) - if configs is None: - configs = self._checkpoint_slot_adapter_configs = {} - if config is None: - configs.pop(name, None) - else: - configs[name] = config - return loaded - - def checkpoint_slot_optimizer_state( - self, name: str - ) -> TrainerRankOptimizerState | None: - if name not in self._checkpoint_slot_params_by_name: - raise ValueError(f"Unknown checkpoint slot: {name!r}") - dynamic = self._dynamic_optimizers.get(name) - if dynamic is None: - return None - state: TrainerRankOptimizerState = { - "format_version": 1, - "layout": self._dynamic_optimizer_layout(name), - "master_params": tuple( - param.detach().cpu().clone() for param in dynamic.master_params - ), - "optimizer": cast( - dict[str, object], - _state_to_cpu(dynamic.optimizer.state_dict()), - ), - } - return state - def save_checkpoint_slot_lora(self, name: str, output_dir: str) -> None: - """Collectively publish a trained checkpoint slot as a vLLM LoRA.""" - known = name in self._checkpoint_slot_params_by_name - if dist.is_available() and dist.is_initialized(): - gathered: list[tuple[str, bool] | None] = [None] * dist.get_world_size() - dist.all_gather_object(gathered, (name, known)) - if any(state != (name, True) for state in gathered): - raise ValueError( - "Checkpoint slot publish requires the same loaded name on all " - f"ranks; got {gathered}" - ) - if not known: - raise ValueError(f"Unknown checkpoint slot: {name!r}") - config = getattr(self, "_checkpoint_slot_adapter_configs", {}).get(name) - if config is None: - raise TrainerRankSlotStateError( - f"Checkpoint slot {name!r} was loaded without adapter_config; " - "reload it with adapter_config=... before publishing." + @staticmethod + def _checkpoint_source_key(path: str) -> str: + return str(Path(path).resolve()) + + @staticmethod + def _checkpoint_source( + checkpoint: str | MaterializedCheckpoint | None, + ) -> tuple[str | None, str | None]: + if isinstance(checkpoint, MaterializedCheckpoint): + return checkpoint.path, checkpoint.directory + return checkpoint, checkpoint + + async def _prefetch_checkpoint(self, source_path: str) -> PreparedCheckpoint: + key = self._checkpoint_source_key(source_path) + task = self._checkpoint_prefetches.get(key) + if task is None: + from ._checkpoint import prepare_checkpoint + + task = self._checkpoint_prefetches[key] = asyncio.create_task( + asyncio.to_thread(prepare_checkpoint, key) ) - from art.megatron.weights.lora_publish import save_vllm_lora_from_model + try: + return await asyncio.shield(task) + except BaseException: + if task.done(): + self._checkpoint_prefetches.pop(key, None) + raise - save_vllm_lora_from_model( - model=self.runtime.model, - adapter_dtypes={}, - handler=self.runtime.model_support_handler, - adapter_config=config, - output_dir=output_dir, - rank=self.runtime.rank, - world_size=self.runtime.world_size, - slot_ref=self._slot_ref("checkpoint", name), - ) + async def _load_checkpoint_path( + self, + logical_path: str, + *, + source_path: str, + prefetch: asyncio.Task[PreparedCheckpoint], + ) -> None: + from . import _checkpoint - def _validate_checkpoint_slot_adapter_config( + key = self._checkpoint_source_key(source_path) + source: PreparedCheckpoint | None = None + error: BaseException | None = None + try: + source = await asyncio.shield(prefetch) + except BaseException as exc: + error = exc + group = _checkpoint._ensure_group(self) + _checkpoint.raise_distributed(error, "prepare checkpoint", group) + assert source is not None + _checkpoint.load_checkpoint(self, source, logical_path) + self._checkpoint_prefetches.pop(key, None) + + def _resolve_checkpoint_name(self, checkpoint_path: str | Literal["active"]) -> str: + if checkpoint_path != "active": + return checkpoint_path + ref = self._slot_stack[-1] if self._slot_stack else self._default_slot_ref + if ref is None or ref.name is None: + raise TrainerRankSlotStateError("No active trainable checkpoint") + return ref.name + + @staticmethod + def _slot_state_error(message: str) -> TrainerRankSlotStateError: + return TrainerRankSlotStateError(message) + + def _checkpoint_group(self) -> dist.ProcessGroup | None: + from ._checkpoint import _ensure_group + + return _ensure_group(self) + + def _validate_checkpoint_adapter_config( self, name: str, adapter_config: Mapping[str, object] | None, @@ -722,7 +919,7 @@ def _validate_checkpoint_slot_adapter_config( config = None if adapter_config is None else deepcopy(dict(adapter_config)) if dist.is_available() and dist.is_initialized(): gathered: list[dict[str, object] | None] = [None] * dist.get_world_size() - dist.all_gather_object(gathered, config) + dist.all_gather_object(gathered, config, group=self._checkpoint_group()) if any(value != config for value in gathered): raise ValueError( f"Adapter config for checkpoint slot {name!r} differs across ranks" @@ -742,6 +939,20 @@ def _validate_checkpoint_slot_adapter_config( raise TypeError( "adapter_config['base_model_name_or_path'] must be a string" ) + if base_model.startswith(("Qwen/Qwen3.5-", "Qwen/Qwen3.6-")): + dimensions = { + "num_attention_heads": getattr( + self.runtime.provider, "num_attention_heads", None + ), + "num_key_value_heads": getattr( + self.runtime.provider, "num_query_groups", None + ), + "head_dim": getattr(self.runtime.provider, "kv_channels", None), + "hidden_size": getattr(self.runtime.provider, "hidden_size", None), + } + for key, value in dimensions.items(): + if value is not None: + config[key] = int(value) if not isinstance(rank, int) or isinstance(rank, bool): raise TypeError("adapter_config['r'] must be an integer") if not isinstance(config_alpha_value, int | float) or isinstance( @@ -764,12 +975,12 @@ def _validate_checkpoint_slot_adapter_config( ) return cast(_AdapterConfig, config) - def _validate_loaded_checkpoint_slot_config( + def _validate_loaded_checkpoint_config( self, name: str, config: _AdapterConfig ) -> None: from art.megatron.lora import LoRA - ref = self._slot_ref("checkpoint", name) + ref = self._slot_ref(name) slots = [ slot for chunk in self.runtime.model @@ -785,19 +996,6 @@ def _validate_loaded_checkpoint_slot_config( f"rank/alpha={expected}, loaded weights use {sorted(actual)}" ) - def load_lora_slot( - self, - name: str, - adapter_model: Mapping[str, torch.Tensor], - *, - alpha: float | None = None, - ) -> int: - loaded = self._load_slot( - "lora", name, adapter_model, trainable=False, alpha=alpha - ) - self._validate_dynamic_slot_consistency("lora", name, loaded) - return loaded - @overload def forward_micro_batches( self, @@ -994,63 +1192,86 @@ def optim_step( scale_grads=scale_grads, ) - def _load_slot( + def _load_checkpoint_slot( self, - kind: Literal["checkpoint", "lora"], name: str, adapter_model: Mapping[str, torch.Tensor], *, - trainable: bool, - alpha: float | None, + alpha: float, + _prepared: bool = False, ) -> int: if self._slot_stack: - raise RuntimeError("Cannot load a LoRA/checkpoint while a slot is pushed") - adapter_model = self._prepare_adapter_model(kind, name, adapter_model) - from art.megatron.lora import LORA_ALPHA, load_lora_slot_into_model + raise RuntimeError("Cannot load a checkpoint while one is pushed") + adapter_model = self._prepare_adapter_model( + name, adapter_model, canonicalized=_prepared + ) + from art.megatron.lora import load_lora_slot_into_model - ref = self._slot_ref(kind, name) + ref = self._slot_ref(name) self._guard_slot_can_load(ref) + self._compact_lora_slot_keys() return load_lora_slot_into_model( self.runtime.model, ref, adapter_model, - alpha=LORA_ALPHA if alpha is None else alpha, - requires_grad=trainable, + alpha=alpha, + requires_grad=True, ) + def _compact_lora_slot_keys(self) -> None: + from art.megatron.lora import LoRA + + for chunk in self.runtime.model: + for module in chunk.modules(): + if not isinstance(module, LoRA): + continue + slots = [ + (ref, module._slot_modules[key]) + for ref, key in module._slot_keys.items() + ] + module._slot_keys = { + ref: f"slot_{index}" for index, (ref, _slot) in enumerate(slots) + } + module._slot_modules = torch.nn.ModuleDict( + {f"slot_{index}": slot for index, (_ref, slot) in enumerate(slots)} + ) + def _prepare_adapter_model( self, - kind: Literal["checkpoint", "lora"], name: str, adapter_model: Mapping[str, torch.Tensor], + *, + canonicalized: bool = False, ) -> dict[str, torch.Tensor]: templates = self._local_lora_adapter_templates() keys = set(adapter_model) expected = set(templates) if dist.is_available() and dist.is_initialized(): gathered: list[set[str] | None] = [None] * dist.get_world_size() - dist.all_gather_object(gathered, expected) + dist.all_gather_object(gathered, expected, group=self._checkpoint_group()) expected = set().union(*(value for value in gathered if value is not None)) if unknown := sorted(keys - expected): preview = ", ".join(repr(key) for key in unknown[:8]) more = "" if len(unknown) <= 8 else f", ... +{len(unknown) - 8} more" raise ValueError( - f"Adapter for {kind} slot {name!r} contains keys that do not match " - f"installed LoRA wrapper sites: {preview}{more}. Configure the " - "Megatron runtime with matching LoRA target modules before loading." + f"Checkpoint {name!r} contains keys that do not match installed " + f"LoRA wrapper sites: {preview}{more}. Configure the Megatron " + "runtime with matching LoRA target modules before loading." ) local_state = { key: tensor for key, tensor in adapter_model.items() if key in templates } adapter_model = ( - self.runtime.model_support_handler.canonicalize_loaded_lora_state( + local_state + if canonicalized + else self.runtime.model_support_handler.canonicalize_loaded_lora_state( local_state, self.runtime.model ) ) if set(adapter_model) != set(local_state): raise TrainerRankSlotStateError( "Model-specific LoRA canonicalization changed the adapter key set " - f"for {kind} slot {name!r}." + f"for checkpoint {name!r}." ) return { key: tensor.to( @@ -1080,82 +1301,83 @@ def _local_lora_adapter_templates(self) -> dict[str, torch.Tensor]: ) return templates + def _iter_slot_parameters(self, ref: "LoRASlotRef") -> Iterator[torch.nn.Parameter]: + from art.megatron.lora import iter_lora_slot_parameters + + return iter_lora_slot_parameters(self.runtime.model, ref) + + def _local_parameter_key_groups(self, name: str) -> tuple[tuple[str, ...], ...]: + ref = self._slot_ref(name) + return tuple( + tuple(str(key) for key in expected(str(suffix).removesuffix(".weight"))) + for chunk in self.runtime.model + for module in chunk.modules() + if (lora_params := getattr(module, "_lora_params", None)) is not None + if (expected := getattr(module, "_expected_weight_keys", None)) is not None + for suffix, _param in lora_params(ref) + ) + + def _validate_checkpoint_consistency( + self, name: str, loaded_sites: int, expected_keys: set[str] + ) -> tuple[torch.nn.Parameter, ...]: + params = tuple(self._iter_slot_parameters(self._slot_ref(name))) + local_keys = { + key for group in self._local_parameter_key_groups(name) for key in group + } + gathered = ( + [local_keys] + if not (dist.is_available() and dist.is_initialized()) + else [None] * dist.get_world_size() + ) + if dist.is_available() and dist.is_initialized(): + dist.all_gather_object(gathered, local_keys, group=self._checkpoint_group()) + covered = set().union(*(keys for keys in gathered if keys is not None)) + if loaded_sites < 1 or covered != expected_keys: + raise TrainerRankSlotStateError( + f"Checkpoint {name!r} has inconsistent distributed coverage" + ) + return params + def _set_default_slot(self, ref: "LoRASlotRef") -> None: if self._slot_stack: - raise RuntimeError("Cannot set a LoRA/checkpoint while a slot is pushed") + raise RuntimeError("Cannot select a checkpoint while one is pushed") self._default_slot_ref = ref @staticmethod - def _slot_ref( - kind: Literal["checkpoint", "lora"], name: str | None - ) -> "LoRASlotRef": + def _slot_ref(name: str | None) -> "LoRASlotRef": try: from art.megatron.lora import LoRASlotRef except ModuleNotFoundError as exc: if exc.name is None or not exc.name.startswith("megatron"): raise - return cast("LoRASlotRef", _LocalLoRASlotRef(kind=kind, name=name)) - - return LoRASlotRef(kind=kind, name=name) + return cast("LoRASlotRef", _LocalLoRASlotRef(name=name)) - def _validate_dynamic_slot_consistency( - self, - kind: Literal["checkpoint", "lora"], - name: str, - loaded_sites: int, - ) -> tuple[torch.nn.Parameter, ...]: - from art.megatron.lora import iter_lora_slot_parameters - - ref = self._slot_ref(kind, name) - params = tuple(iter_lora_slot_parameters(self.runtime.model, ref)) - if not (dist.is_available() and dist.is_initialized()): - return params - - signature = tuple( - ( - tuple(param.shape), - str(param.dtype), - bool(getattr(param, "allreduce", True)), - str(getattr(param, "grad_sync_domain", "tp_default")), - str(getattr(param, "grad_sync_op", "none")), - ) - for param in params - ) - local = (int(loaded_sites), signature) - gathered: list[tuple[int, object] | None] = [None] * dist.get_world_size() - dist.all_gather_object(gathered, local) - ranks = [state for state in gathered if state is not None] - if all(state == ranks[0] for state in ranks[1:]): - return params - raise RuntimeError( - f"Dynamic LoRA slot {kind}:{name} is not loaded consistently across " - "distributed ranks. This usually means a sharded/exported LoRA state " - "dict was passed directly to TrainerRank; gather or materialize the " - "full adapter state before loading a dynamic slot. " - f"Loaded-site counts by rank: {[state[0] for state in ranks]}." - ) + return LoRASlotRef(kind="checkpoint", name=name) def _resolve_slot_ref(self, request: AnyForwardInput) -> "LoRASlotRef | None": if request.checkpoint is not Unset: - return self._slot_ref("checkpoint", cast(str | None, request.checkpoint)) - if request.lora is not Unset: - return self._slot_ref("lora", cast(str | None, request.lora)) + name = cast(str | None, request.checkpoint) + if name is not None and name not in self._checkpoint_slots: + raise TrainerRankSlotStateError( + f"Forward input selects unloaded checkpoint {name!r}" + ) + return self._slot_ref(name) if self._slot_stack: return self._slot_stack[-1] if self._default_slot_ref is not None: return self._default_slot_ref - return self._slot_ref("checkpoint", None) + return self._slot_ref(None) def _selected_dynamic_checkpoints( self, checkpoints: Sequence[str] | None, ) -> tuple[str, ...]: - loaded = set(self._checkpoint_slot_params_by_name) + loaded = set(self._checkpoint_slots) if not loaded: raise TrainerRankSlotStateError( "TrainerRank.optim_step requires a loaded checkpoint slot. Call " - "load_checkpoint_slot(...) and run backward on outputs produced by " + "load_checkpoint(...) and run backward on outputs produced by " "that slot before stepping." ) requested = ( @@ -1198,7 +1420,7 @@ def _checkpoint_grad_flags(self, names: Sequence[str]) -> tuple[bool, ...]: [ any( param.grad is not None - for param in self._checkpoint_slot_params_by_name[name] + for param in self._checkpoint_slots[name].params ) for name in names ], @@ -1222,7 +1444,7 @@ def _dynamic_optim_step( selected = [] for name in checkpoint_names: self._guard_checkpoint_can_step(name) - slot_params = self._checkpoint_slot_params_by_name[name] + slot_params = self._checkpoint_slots[name].params slot_grads = self._reduce_dynamic_grads( slot_params, scale_grads=scale_grads ) @@ -1258,7 +1480,8 @@ def _dynamic_optim_step( ): model.copy_(master) model.grad = None - self._prune_slot_graphs(self._slot_ref("checkpoint", name)) + self._prune_slot_graphs(self._slot_ref(name)) + self._checkpoint_slots[name].revision += 1 return { "learning_rate": float(params.learning_rate), "grad_norm": float(grad_norm), @@ -1271,10 +1494,11 @@ def _dynamic_optimizer( name: str, params: AdamParams, ) -> _DynamicOptimizer: - dynamic = self._dynamic_optimizers.get(name) + slot = self._checkpoint_slots[name] + dynamic = slot.optimizer if dynamic is None: dynamic = self._new_dynamic_optimizer(name, params) - self._dynamic_optimizers[name] = dynamic + slot.optimizer = dynamic return dynamic for group in dynamic.optimizer.param_groups: group["lr"] = params.learning_rate @@ -1289,7 +1513,7 @@ def _new_dynamic_optimizer( *, master_params: Sequence[torch.Tensor] | None = None, ) -> _DynamicOptimizer: - model_params = self._checkpoint_slot_params_by_name[name] + model_params = self._checkpoint_slots[name].params sources = model_params if master_params is None else tuple(master_params) if len(sources) != len(model_params) or any( not isinstance(source, torch.Tensor) for source in sources @@ -1298,6 +1522,13 @@ def _new_dynamic_optimizer( f"Optimizer state for checkpoint slot {name!r} has " f"{len(sources)} master parameters; expected {len(model_params)}." ) + if any( + tuple(source.shape) != tuple(model.shape) + for source, model in zip(sources, model_params, strict=True) + ): + raise TrainerRankSlotStateError( + f"Optimizer master parameter shape does not match checkpoint {name!r}" + ) masters = tuple( torch.nn.Parameter( source.detach().to(device=model.device, dtype=torch.float32).clone() @@ -1316,55 +1547,40 @@ def _new_dynamic_optimizer( ) return _DynamicOptimizer(optimizer, masters) - def _restore_dynamic_optimizer( + def _restore_canonical_optimizer( self, name: str, - state: TrainerRankOptimizerState, + state: "LocalOptimizerState", ) -> _DynamicOptimizer: - if state.get("format_version") != 1: - raise TrainerRankSlotStateError( - f"Unsupported optimizer state format for checkpoint slot {name!r}." - ) - if state.get("layout") != self._dynamic_optimizer_layout(name): - raise TrainerRankSlotStateError( - f"Optimizer state for checkpoint slot {name!r} was saved for a " - "different topology or parameter layout. Save and restore one " - "optimizer shard per TrainerRank with matching TP/EP/ETP ranks." - ) - master_params = state.get("master_params") - optimizer_state = state.get("optimizer") - if not isinstance(master_params, Sequence) or not isinstance( - optimizer_state, Mapping - ): - raise TrainerRankSlotStateError( - f"Optimizer state for checkpoint slot {name!r} is incomplete." - ) dynamic = self._new_dynamic_optimizer( name, - AdamParams(learning_rate=0.0), - master_params=cast(Sequence[torch.Tensor], master_params), - ) - try: - dynamic.optimizer.load_state_dict( - {str(key): value for key, value in optimizer_state.items()} - ) - except ValueError as exc: - raise TrainerRankSlotStateError( - f"Optimizer state for checkpoint slot {name!r} does not match the " - "loaded slot parameter groups." - ) from exc - for param in dynamic.master_params: - for state_name, value in dynamic.optimizer.state.get(param, {}).items(): - if ( - isinstance(value, torch.Tensor) - and int(value.ndim) > 0 - and tuple(value.shape) != tuple(param.shape) - ): - raise TrainerRankSlotStateError( - f"Optimizer state {state_name!r} for checkpoint slot " - f"{name!r} has shape {tuple(value.shape)}, but the loaded " - f"slot parameter has shape {tuple(param.shape)}." - ) + AdamParams( + learning_rate=state.config["learning_rate"], + beta1=state.config["beta1"], + beta2=state.config["beta2"], + weight_decay=state.config["weight_decay"], + ), + master_params=state.masters, + ) + dynamic.optimizer.param_groups[0]["eps"] = state.config["eps"] + for master, exp_avg, exp_avg_sq, step in zip( + dynamic.master_params, + state.exp_avgs, + state.exp_avg_sqs, + state.steps, + strict=True, + ): + if tuple(exp_avg.shape) != tuple(master.shape) or tuple( + exp_avg_sq.shape + ) != tuple(master.shape): + raise TrainerRankSlotStateError( + f"Canonical optimizer moment shape does not match {name!r}" + ) + dynamic.optimizer.state[master] = { + "step": torch.tensor(step, dtype=torch.float32), + "exp_avg": exp_avg.to(master.device, torch.float32).clone(), + "exp_avg_sq": exp_avg_sq.to(master.device, torch.float32).clone(), + } self._zero_dynamic_optimizer_padding(name, dynamic) return dynamic @@ -1382,13 +1598,13 @@ def _zero_dynamic_optimizer_padding( value.masked_fill_(mask, 0) def _dynamic_optimizer_padding_masks(self, name: str) -> tuple[torch.Tensor, ...]: - params = self._checkpoint_slot_params_by_name[name] + params = self._checkpoint_slots[name].params masks = tuple(torch.zeros_like(param, dtype=torch.bool) for param in params) param_indices = {id(param): index for index, param in enumerate(params)} exported: dict[str, torch.Tensor] = {} owners: dict[str, tuple[int, int | None]] = {} mapped_indices: set[int] = set() - ref = self._slot_ref("checkpoint", name) + ref = self._slot_ref(name) for chunk in self.runtime.model: for module in chunk.modules(): @@ -1405,29 +1621,29 @@ def _dynamic_optimizer_padding_masks(self, name: str) -> tuple[torch.Tensor, ... if int(param.ndim) == 3: if len(keys) != int(param.shape[0]): raise TrainerRankSlotStateError( - f"Cannot map optimizer padding for checkpoint slot " + f"Cannot map optimizer padding for checkpoint " f"{name!r}: {len(keys)} adapter keys describe " f"{int(param.shape[0])} local experts." ) for expert, key in enumerate(keys): exported[str(key)] = torch.ones_like(param[expert].T) owners[str(key)] = (index, expert) - else: - if len(keys) != 1: - raise TrainerRankSlotStateError( - f"Cannot map optimizer padding for checkpoint slot " - f"{name!r}: expected one adapter key, got {len(keys)}." - ) + elif len(keys) == 1: key = str(keys[0]) exported[key] = torch.ones_like(param.T) owners[key] = (index, None) + else: + raise TrainerRankSlotStateError( + f"Cannot map optimizer padding for checkpoint {name!r}: " + f"expected one adapter key, got {len(keys)}." + ) if mapped_indices and ( missing := sorted(set(range(len(params))) - mapped_indices) ): raise TrainerRankSlotStateError( - f"Cannot map optimizer padding for checkpoint slot {name!r}: " - f"parameter indices {missing} do not belong to installed LoRA sites." + f"Cannot map optimizer padding for checkpoint {name!r}: parameter " + f"indices {missing} do not belong to installed LoRA sites." ) canonical = self.runtime.model_support_handler.canonicalize_loaded_lora_state( @@ -1496,39 +1712,6 @@ def add( coalesced_all_reduce(bucket_grads, group=group, op=op) return grads - def _dynamic_optimizer_layout(self, name: str) -> TrainerRankOptimizerLayout: - parameters = cast( - tuple[ - tuple[ - tuple[int, ...], - str, - str, - bool, - int | None, - str, - tuple[int, ...], - ], - ..., - ], - tuple( - ( - tuple(param.shape), - str(param.dtype), - str(getattr(param, "lora_shard_domain", "tp")), - bool(getattr(param, "lora_tp_sharded", False)), - getattr(param, "lora_tp_shard_dim", None), - str(getattr(param, "lora_tp_shard_strategy", "uniform")), - tuple(getattr(param, "lora_tp_component_sizes", ())), - ) - for param in self._checkpoint_slot_params_by_name[name] - ), - ) - layout: TrainerRankOptimizerLayout = { - "parallel": _parallel_optimizer_coordinates(), - "parameters": parameters, - } - return layout - def _select_next_micro_batch( self, items: Sequence[ForwardInputsT], @@ -1935,7 +2118,7 @@ def _guard_slot_can_load(self, ref: "LoRASlotRef") -> None: if not self._has_live_slot_graph(ref): return raise TrainerRankSlotStateError( - f"Cannot load {ref.kind} slot {ref.name!r} while outputs from an " + f"Cannot load checkpoint {ref.name!r} while outputs from an " "earlier forward using that slot still have a live backward graph. " "Activation checkpoint recompute resolves slots by name, so replacing " "the slot before backward can compute gradients with different LoRA " @@ -1945,7 +2128,7 @@ def _guard_slot_can_load(self, ref: "LoRASlotRef") -> None: ) def _guard_checkpoint_can_step(self, name: str) -> None: - ref = self._slot_ref("checkpoint", name) + ref = self._slot_ref(name) if not self._has_live_slot_graph(ref): return raise TrainerRankSlotStateError( @@ -2920,36 +3103,6 @@ def _include_in_distributed_grad_norm(param: torch.nn.Parameter) -> bool: return shard_group is None or shard_group.size() <= 1 or shard_group.rank() == 0 -def _parallel_optimizer_coordinates() -> tuple[int, int, int, int, int, int, int, int]: - if not (dist.is_available() and dist.is_initialized()): - return (1, 0, 1, 0, 1, 0, 1, 0) - from megatron.core import parallel_state as ps - - expert_tp_group = ps.get_expert_tensor_parallel_group(check_initialized=False) - return ( - int(ps.get_tensor_model_parallel_world_size()), - int(ps.get_tensor_model_parallel_rank()), - int(ps.get_expert_model_parallel_world_size()), - int(ps.get_expert_model_parallel_rank()), - 1 if expert_tp_group is None else int(expert_tp_group.size()), - 0 if expert_tp_group is None else int(expert_tp_group.rank()), - int(ps.get_pipeline_model_parallel_world_size()), - int(ps.get_pipeline_model_parallel_rank()), - ) - - -def _state_to_cpu(value: object) -> object: - if isinstance(value, torch.Tensor): - return value.detach().cpu().clone() - if isinstance(value, Mapping): - return {key: _state_to_cpu(item) for key, item in value.items()} - if isinstance(value, tuple): - return tuple(_state_to_cpu(item) for item in value) - if isinstance(value, list): - return [_state_to_cpu(item) for item in value] - return value - - def _vocab_parallel_target_logprobs( local_logits: torch.Tensor, labels: torch.Tensor, diff --git a/tests/integration/megatron/lora/test_dynamic_lora_slots.py b/tests/integration/megatron/lora/test_dynamic_lora_slots.py index e6029f35d..d1ba1688c 100644 --- a/tests/integration/megatron/lora/test_dynamic_lora_slots.py +++ b/tests/integration/megatron/lora/test_dynamic_lora_slots.py @@ -6,6 +6,7 @@ from pathlib import Path import socket from types import SimpleNamespace +from typing import cast import pytest @@ -17,9 +18,15 @@ import torch.multiprocessing as mp # noqa: E402 from art.megatron.lora import LoRA, LoRASlotRef, use_lora_slot # noqa: E402 +from art.trainer_rank._checkpoint import ( # noqa: E402 + LocalOptimizerState, + OptimizerConfig, + _commit_slot, +) from art.trainer_rank._impl import ( # noqa: E402 AdamParams, TrainerRank, + _CheckpointSlot, _distributed_grad_norm, _vocab_parallel_log_z, _vocab_parallel_target_logprobs, @@ -72,7 +79,7 @@ def test_dynamic_lora_slots_capture_recompute_context_and_step_independently() - key: value.cpu().double() for key, value in _adapter("dense", rank=3, seed=7).items() } - trainer.load_checkpoint_slot("CPU", cpu_adapter) + _install_checkpoint(trainer, "CPU", cpu_adapter) cpu_slot = lora._slot(LoRASlotRef("checkpoint", "CPU")) assert cpu_slot is not None assert cpu_slot.A_T.device == lora.A_T.device @@ -82,7 +89,7 @@ def test_dynamic_lora_slots_capture_recompute_context_and_step_independently() - with trainer.push_checkpoint("A"): assert trainer._slot_stack[-1] == ref_a - with trainer.push_lora(None): + with trainer.push_checkpoint(None): assert trainer._slot_stack[-1].name is None assert trainer._slot_stack[-1] == ref_a assert trainer._slot_stack == [] @@ -100,6 +107,41 @@ def test_dynamic_lora_slots_capture_recompute_context_and_step_independently() - _assert_reload_replaces_slot_optimizer(ref_a, lora, trainer) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required.") +def test_checkpoint_reload_does_not_alias_the_next_slot() -> None: + with _single_rank_model_parallel(): + device = torch.device("cuda") + first = LoRA("first", 4, 5, 2, 32, torch.float32, device) + second = LoRA("second", 4, 5, 2, 32, torch.float32, device) + trainer = _trainer_for(first, device) + trainer.runtime.model = [torch.nn.Sequential(first, second)] + + def adapter( + seed: int, *, include_second: bool = True + ) -> dict[str, torch.Tensor]: + return _adapter("first", rank=2, seed=seed) | ( + _adapter("second", rank=2, seed=seed + 10) if include_second else {} + ) + + def stage(destination: str, state: dict[str, torch.Tensor]) -> None: + temporary = f"temporary-{destination}" + trainer._load_checkpoint_slot(temporary, state, alpha=32.0) + _commit_slot(trainer, temporary, destination) + + stage("A", adapter(1)) + stage("B", adapter(2)) + slot_b = second._slot(trainer._slot_ref("B")) + assert slot_b is not None + expected_b = slot_b.A_T.detach().clone() + stage("A", adapter(3, include_second=False)) + stage("C", adapter(4)) + + slot_b = second._slot(trainer._slot_ref("B")) + slot_c = second._slot(trainer._slot_ref("C")) + assert slot_b is not None and slot_c is not None and slot_b is not slot_c + torch.testing.assert_close(slot_b.A_T, expected_b) + + @pytest.mark.parametrize("tp_size", (2, 4)) def test_trainer_rank_tp_head_backward_matches_unsharded_oracle( tp_size: int, @@ -261,8 +303,7 @@ def _assert_distributed_optimizer_restore(device: torch.device) -> None: with use_lora_slot(ref): lora(x).sum().backward() trainer.optim_step(params=params, checkpoints=["A"]) - state = trainer.checkpoint_slot_optimizer_state("A") - assert state is not None + state = _optimizer_state(trainer, "A") slot = lora._slot(ref) assert slot is not None adapter = { @@ -275,7 +316,10 @@ def _assert_distributed_optimizer_restore(device: torch.device) -> None: restored_lora = LoRA("dense", 4, 5, 2, 32, torch.float32, device) restored = _trainer_for(restored_lora, device) - restored.load_checkpoint_slot("A", adapter, optimizer_state=state) + _install_checkpoint(restored, "A", adapter) + restored._checkpoint_slots["A"].optimizer = restored._restore_canonical_optimizer( + "A", state + ) with use_lora_slot(ref): restored_lora(x).sum().backward() restored.optim_step(params=params, checkpoints=["A"]) @@ -373,13 +417,13 @@ def _assert_reload_replaces_slot_optimizer( trainer: TrainerRank, ) -> None: assert ref.name is not None - old_params = trainer._checkpoint_slot_params_by_name[ref.name] - assert ref.name in trainer._dynamic_optimizers + old_params = trainer._checkpoint_slots[ref.name].params + assert trainer._checkpoint_slots[ref.name].optimizer is not None - trainer.load_checkpoint_slot(ref.name, _adapter("dense", rank=3, seed=9)) + _install_checkpoint(trainer, ref.name, _adapter("dense", rank=3, seed=9)) - new_params = trainer._checkpoint_slot_params_by_name[ref.name] - assert ref.name not in trainer._dynamic_optimizers + new_params = trainer._checkpoint_slots[ref.name].params + assert trainer._checkpoint_slots[ref.name].optimizer is None assert [tuple(param.shape) for param in new_params] == [(4, 3), (3, 5)] assert all(old is not new for old, new in zip(old_params, new_params, strict=True)) slot = lora._slot(ref) @@ -387,6 +431,46 @@ def _assert_reload_replaces_slot_optimizer( assert slot.rank == 3 +def _install_checkpoint( + trainer: TrainerRank, name: str, adapter: dict[str, torch.Tensor] +) -> int: + loaded = trainer._load_checkpoint_slot(name, adapter, alpha=32.0) + previous = trainer._checkpoint_slots.get(name) + trainer._checkpoint_slots[name] = _CheckpointSlot( + tuple(trainer._iter_slot_parameters(trainer._slot_ref(name))), + revision=0 if previous is None else previous.revision + 1, + ) + return loaded + + +def _optimizer_state(trainer: TrainerRank, name: str) -> LocalOptimizerState: + dynamic = trainer._checkpoint_slots[name].optimizer + assert dynamic is not None + states = [ + cast(dict[str, torch.Tensor], dynamic.optimizer.state[master]) + for master in dynamic.master_params + ] + group = dynamic.optimizer.param_groups[0] + beta1, beta2 = cast(tuple[float, float], group["betas"]) + return LocalOptimizerState( + masters=tuple( + master.detach().cpu().clone() for master in dynamic.master_params + ), + exp_avgs=tuple(state["exp_avg"].detach().cpu().clone() for state in states), + exp_avg_sqs=tuple( + state["exp_avg_sq"].detach().cpu().clone() for state in states + ), + steps=tuple(float(state["step"].item()) for state in states), + config=OptimizerConfig( + learning_rate=float(group["lr"]), + beta1=beta1, + beta2=beta2, + eps=float(group["eps"]), + weight_decay=float(group["weight_decay"]), + ), + ) + + def _trainer_for(lora: LoRA, device: torch.device) -> TrainerRank: trainer = TrainerRank.__new__(TrainerRank) trainer.runtime = SimpleNamespace( @@ -397,11 +481,14 @@ def _trainer_for(lora: LoRA, device: torch.device) -> TrainerRank: trainer.device = device trainer._slot_stack = [] trainer._default_slot_ref = None - trainer._dynamic_optimizers = {} - trainer._checkpoint_slot_params_by_name = { - "A": tuple(lora.lora_slot_params(LoRASlotRef("checkpoint", "A"))), - "B": tuple(lora.lora_slot_params(LoRASlotRef("checkpoint", "B"))), + trainer._checkpoint_slots = { + name: _CheckpointSlot( + tuple(lora.lora_slot_params(LoRASlotRef("checkpoint", name))) + ) + for name in ("A", "B") } + trainer._checkpoint_prefetches = {} + trainer._checkpoint_mutation_tail = None return trainer diff --git a/tests/integration/megatron/lora/test_lora_disk_codecs.py b/tests/integration/megatron/lora/test_lora_disk_codecs.py index d6aedd437..219473381 100644 --- a/tests/integration/megatron/lora/test_lora_disk_codecs.py +++ b/tests/integration/megatron/lora/test_lora_disk_codecs.py @@ -38,6 +38,7 @@ save_vllm_lora_from_model, ) from art.trainer_rank import TrainerRank +from art.trainer_rank._impl import _AdapterConfig, _CheckpointSlot from art.utils.convert_moe_lora import convert_checkpoint_if_needed REPO_ROOT = Path(__file__).parents[4] @@ -1584,14 +1585,16 @@ def test_trainer_rank_publishes_named_checkpoint_slot_without_mutating_base( ) trainer._slot_stack = [] trainer._pending_slot_graphs = {} - trainer._dynamic_optimizers = {} - trainer._checkpoint_slot_params_by_name = {} - trainer._checkpoint_slot_adapter_configs = {} + trainer._checkpoint_slots = {} config = _config("Qwen/Qwen3-8B", rank=2, alpha=2) - assert trainer.load_checkpoint_slot("student", adapter, adapter_config=config) == 1 + assert trainer._load_checkpoint_slot("student", adapter, alpha=2) == 1 + trainer._checkpoint_slots["student"] = _CheckpointSlot( + tuple(trainer._iter_slot_parameters(trainer._slot_ref("student"))), + cast(_AdapterConfig, config), + ) output_dir = tmp_path / "checkpoint" - trainer.save_checkpoint_slot_lora("student", str(output_dir)) + assert trainer.export_lora(str(output_dir), "student") == 0 _assert_tensors_equal(load_file(output_dir / "adapter_model.safetensors"), adapter) assert json.loads((output_dir / "adapter_config.json").read_text()) == { diff --git a/tests/unit/test_preprocessing_tokenize.py b/tests/unit/test_preprocessing_tokenize.py index 380feb543..35a15aaac 100644 --- a/tests/unit/test_preprocessing_tokenize.py +++ b/tests/unit/test_preprocessing_tokenize.py @@ -5,10 +5,15 @@ import pytest from transformers.tokenization_utils_base import BatchEncoding -from art.megatron.model_support.handlers.gemma4 import ( - GEMMA4_DENSE_HANDLER, - GEMMA4_MOE_HANDLER, -) +try: + from art.megatron.model_support.handlers.gemma4 import ( + GEMMA4_DENSE_HANDLER, + GEMMA4_MOE_HANDLER, + ) +except ModuleNotFoundError as error: + if error.name is None or not error.name.startswith("megatron"): + raise + pytest.skip("Megatron is not installed", allow_module_level=True) from art.preprocessing.tokenize import ( _normalize_tool_call_arguments_for_chat_template, tokenize_sft_batch, diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index 042dad3c1..7c5720c06 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -1,15 +1,23 @@ from __future__ import annotations +import asyncio from collections.abc import Iterable from dataclasses import dataclass +from datetime import timedelta import gc from importlib.util import find_spec import inspect +import json +from pathlib import Path +import threading +import time from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast import pytest import torch +import torch.distributed as dist +import torch.multiprocessing as mp from art.megatron.prefix_tree_packing import prefix_tree_pack from art.trainer_rank import ( @@ -17,16 +25,33 @@ AdapterSelection, ForwardInput, ForwardOutput, + MaterializedCheckpoint, TopK, TrainerRank, TrainerRankMemoryError, - TrainerRankOptimizerLayout, - TrainerRankOptimizerState, TrainerRankSlotStateError, Unset, ) +from art.trainer_rank._checkpoint import ( + CheckpointManifest, + LocalOptimizerState, + OptimizerConfig, + PreparedCheckpoint, + _file_digest, + _FinalizedSave, + _manifest_digest, + _merge_component, + _PreparedSave, + _validate_save_state, + abort_checkpoint_save, + finish_checkpoint_save, + materialize_lora, + prepare_checkpoint, + prepare_checkpoint_save, +) from art.trainer_rank._impl import ( _anchor_disconnected_outputs, + _CheckpointSlot, _MemoryCheck, _MemoryProfile, _validate_top_k, @@ -46,8 +71,6 @@ def test_public_types_have_canonical_module_paths() -> None: assert { "AdapterSelection", - "TrainerRankOptimizerLayout", - "TrainerRankOptimizerState", "Unset", } <= set(art.trainer_rank.__all__) for public_type in ( @@ -57,8 +80,6 @@ def test_public_types_have_canonical_module_paths() -> None: TopK, TrainerRank, TrainerRankMemoryError, - TrainerRankOptimizerLayout, - TrainerRankOptimizerState, TrainerRankSlotStateError, ): assert public_type.__module__ == "art.trainer_rank" @@ -99,7 +120,6 @@ def zero_grad(self) -> None: @dataclass(frozen=True) class _SlotRef: - kind: str name: str | None @@ -122,14 +142,21 @@ def _runtime( model_support_handler=SimpleNamespace( build_gdn_execution_spec=True, canonicalize_loaded_lora_state=lambda state, _model: state, + from_vllm_lora_tensors=lambda state, **_kwargs: state, + to_vllm_lora_tensors=lambda state, **kwargs: ( + state, + kwargs["adapter_config"], + ), zero_internal_padding_grads=lambda _model: None, zero_internal_padding_params=lambda _model: None, ), + rank=0, + world_size=1, ) # type: ignore -def _slot_ref(kind: str, name: str | None) -> "LoRASlotRef": - return _SlotRef(kind, name) # type: ignore +def _slot_ref(name: str | None) -> "LoRASlotRef": + return _SlotRef(name) # type: ignore def _target_request(token: int) -> ForwardInput[torch.Tensor, None, None, None]: @@ -180,7 +207,7 @@ def _trainer_with_checkpoint( ) -> tuple[TrainerRank, torch.nn.Parameter]: trainer = TrainerRank(_runtime()) param = torch.nn.Parameter(value.clone()) - trainer._checkpoint_slot_params_by_name["student"] = (param,) + trainer._checkpoint_slots.setdefault("student", _CheckpointSlot()).params = (param,) monkeypatch.setattr( trainer, "_reduce_dynamic_grads", @@ -210,8 +237,7 @@ def _tracked_targets( def test_forward_input_validation() -> None: with pytest.raises(ValueError, match="top_k must be >= 1"): ForwardInput(input_tokens=torch.tensor([1]), top_k=0) - with pytest.raises(ValueError, match="cannot set both checkpoint and lora"): - ForwardInput(input_tokens=torch.tensor([1]), checkpoint="a", lora="b") + assert "lora" not in ForwardInput.__dataclass_fields__ with pytest.raises(ValueError, match="top_k=9 exceeds vocabulary size 8"): _validate_top_k(9, _Model()) @@ -223,7 +249,40 @@ def test_forward_input_distinguishes_unset_and_base_checkpoint( request = ForwardInput(input_tokens=torch.tensor([1]), checkpoint=checkpoint) assert request.checkpoint is expected - assert request.lora is Unset + + +def test_dp_rank_forward_rejects_unloaded_explicit_checkpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) + _stub_forward(monkeypatch, trainer) + request = ForwardInput( + input_tokens=torch.tensor([1]), + target_tokens=torch.tensor([1]), + checkpoint="typo", + ) + + with pytest.raises(TrainerRankSlotStateError, match="unloaded.*'typo'"): + trainer.dp_rank_forward([request]) + + +@pytest.mark.parametrize("checkpoint", (None, "student")) +def test_dp_rank_forward_accepts_base_or_loaded_explicit_checkpoint( + monkeypatch: pytest.MonkeyPatch, + checkpoint: str | None, +) -> None: + trainer = TrainerRank(_runtime()) + trainer._checkpoint_slots.setdefault("student", _CheckpointSlot()).params = () + _stub_forward(monkeypatch, trainer) + request = ForwardInput( + input_tokens=torch.tensor([1]), + target_tokens=torch.tensor([1]), + checkpoint=checkpoint, + ) + + output = trainer.dp_rank_forward([request]) + + assert isinstance(output[0], ForwardOutput) def test_forward_input_preserves_public_runtime_shape() -> None: @@ -388,15 +447,359 @@ def test_hybridep_rejects_buffer_growth_with_live_graph( ) -def test_trainer_rank_adapter_stack_errors() -> None: +async def test_trainer_rank_checkpoint_stack_errors() -> None: trainer = TrainerRank(_runtime()) - with pytest.raises(RuntimeError, match="No pushed LoRA or checkpoint"): - trainer.pop_pushed_lora_or_checkpoint() + with pytest.raises(RuntimeError, match="No pushed checkpoint"): + trainer.pop_checkpoint() trainer._slot_stack.append(object()) # type: ignore - for load in (trainer.load_checkpoint_slot, trainer.load_lora_slot): - with pytest.raises(RuntimeError, match="Cannot load a LoRA/checkpoint"): - load("teacher", {}) + with pytest.raises(RuntimeError, match="Cannot load a checkpoint"): + await trainer.load_checkpoint("teacher") + + +async def test_checkpoint_tasks_and_async_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) + trainer._checkpoint_prefetches = {} + fetched: list[str] = [] + + async def prefetch(path: str) -> object: + fetched.append(path) + return object() + + def install(trainer: TrainerRank, _source: object, path: str) -> None: + trainer._checkpoint_slots.setdefault(path, _CheckpointSlot()).params = () + trainer._checkpoint_slots[path].revision = 0 + + monkeypatch.setattr(trainer, "_prefetch_checkpoint", prefetch) + monkeypatch.setattr("art.trainer_rank._checkpoint.load_checkpoint", install) + + task = trainer.load_checkpoint("student") + assert isinstance(task, asyncio.Task) + await task + assert fetched == ["student"] + assert trainer._default_slot_ref == trainer._slot_ref("student") + + task = trainer.prefetch_checkpoints("teacher", "reference") + assert isinstance(task, asyncio.Task) + await task + assert fetched[-2:] == ["teacher", "reference"] + + pushed = trainer.push_checkpoint("student") + await pushed + assert trainer._slot_stack == [trainer._slot_ref("student")] + trainer.pop_checkpoint() + async with trainer.push_checkpoint("student"): + assert trainer._slot_stack == [trainer._slot_ref("student")] + async with trainer.push_checkpoint("missing"): + assert trainer._slot_stack == [ + trainer._slot_ref("student"), + trainer._slot_ref("missing"), + ] + assert trainer._slot_stack == [trainer._slot_ref("student")] + assert trainer._slot_stack == [] + assert fetched[-1] == "missing" + + +def test_checkpoint_sync_context_and_body_error_preservation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) + trainer._checkpoint_slots["student"] = _CheckpointSlot() + with trainer.push_checkpoint("student"): + assert trainer._slot_stack == [trainer._slot_ref("student")] + assert trainer._slot_stack == [] + + pushed = trainer.push_checkpoint("student") + monkeypatch.setattr( + trainer, + "pop_checkpoint", + lambda: (_ for _ in ()).throw(RuntimeError("cleanup failed")), + ) + with pytest.raises(ExceptionGroup) as captured: + with pushed: + raise ValueError("body failed") + assert {type(error) for error in captured.value.exceptions} == { + ValueError, + RuntimeError, + } + + +async def test_checkpoint_context_cancellation_after_successful_push_cleans_stack( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) + trainer._checkpoint_slots.setdefault("student", _CheckpointSlot()).params = () + parent = asyncio.current_task() + assert parent is not None + original_slot_ref = trainer._slot_ref + cancellation_scheduled = False + + def cancel_parent_after_resolving(path: str | None): + nonlocal cancellation_scheduled + ref = original_slot_ref(path) + if not cancellation_scheduled: + cancellation_scheduled = True + asyncio.get_running_loop().call_soon(parent.cancel) + return ref + + monkeypatch.setattr(trainer, "_slot_ref", cancel_parent_after_resolving) + pushed = trainer.push_checkpoint("student") + entered = False + + with pytest.raises(asyncio.CancelledError): + async with pushed: + entered = True + + assert cancellation_scheduled + assert pushed._task is not None + assert pushed._task.done() and not pushed._task.cancelled() + assert not entered + assert trainer._slot_stack == [] + + +async def test_checkpoint_context_cancellation_while_push_is_pending_does_not_leak( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) + started = asyncio.Event() + release = asyncio.Event() + + async def prefetch(_path: str) -> object: + started.set() + await release.wait() + return object() + + monkeypatch.setattr(trainer, "_prefetch_checkpoint", prefetch) + pushed = trainer.push_checkpoint("student") + entering = asyncio.create_task(pushed.__aenter__()) + await started.wait() + entering.cancel() + with pytest.raises(asyncio.CancelledError): + await entering + release.set() + await asyncio.sleep(0) + + assert pushed._task is not None and pushed._task.cancelled() + assert trainer._slot_stack == [] + + +def test_pushed_checkpoint_cannot_be_reused() -> None: + trainer = TrainerRank(_runtime()) + trainer._checkpoint_slots["student"] = _CheckpointSlot() + pushed = trainer.push_checkpoint("student") + + with pushed: + pass + with pytest.raises(RuntimeError, match="cannot be entered twice"): + with pushed: + pass + + +async def test_shared_checkpoint_prefetch_survives_waiter_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) + started = asyncio.Event() + release = asyncio.Event() + source = object() + + async def delayed_to_thread(_function: object, *_args: object) -> object: + started.set() + await release.wait() + return source + + monkeypatch.setattr(asyncio, "to_thread", delayed_to_thread) + first = asyncio.create_task(trainer._prefetch_checkpoint("student")) + second = asyncio.create_task(trainer._prefetch_checkpoint("student")) + await started.wait() + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + release.set() + + assert await second is source + [cached] = trainer._checkpoint_prefetches.values() + assert cached.result() is source + + +async def test_shared_checkpoint_prefetch_serves_successful_waiters( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) + started = asyncio.Event() + release = asyncio.Event() + source = object() + calls = 0 + + async def delayed_to_thread(_function: object, *_args: object) -> object: + nonlocal calls + calls += 1 + started.set() + await release.wait() + return source + + monkeypatch.setattr(asyncio, "to_thread", delayed_to_thread) + first = asyncio.create_task(trainer._prefetch_checkpoint("student")) + second = asyncio.create_task(trainer._prefetch_checkpoint("student")) + await started.wait() + await asyncio.sleep(0) + release.set() + + assert await asyncio.gather(first, second) == [source, source] + assert calls == 1 + [cached] = trainer._checkpoint_prefetches.values() + assert cached.result() is source + + +async def test_materialized_sources_keep_logical_checkpoint_identities( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + trainer = TrainerRank(_runtime()) + prepared: list[str] = [] + installed: list[tuple[str, object]] = [] + + def prepare(source_path: str) -> object: + prepared.append(source_path) + return object() + + def install(trainer: TrainerRank, source: object, logical_path: str) -> None: + installed.append((logical_path, source)) + trainer._checkpoint_slots.setdefault( + logical_path, _CheckpointSlot() + ).params = () + trainer._checkpoint_slots[logical_path].revision = 0 + + monkeypatch.setattr("art.trainer_rank._checkpoint.prepare_checkpoint", prepare) + monkeypatch.setattr("art.trainer_rank._checkpoint.load_checkpoint", install) + root_a = str(tmp_path / "immutable-a") + root_b = str(tmp_path / "immutable-b") + logical_a = "wandb-artifact:///entity/project/run:step1" + logical_b = "wandb-artifact:///entity/project/run-teacher:step1" + logical_c = "wandb-artifact:///entity/project/run-reference:step1" + + await asyncio.gather( + trainer.load_checkpoint(MaterializedCheckpoint(logical_a, root_a)), + trainer.load_checkpoint(MaterializedCheckpoint(logical_b, root_a)), + ) + assert prepared == [trainer._checkpoint_source_key(root_a)] + + await trainer.prefetch_checkpoints(MaterializedCheckpoint(logical_c, root_b)) + await trainer.load_checkpoint(MaterializedCheckpoint(logical_c, root_b)) + assert sorted(prepared) == sorted( + (trainer._checkpoint_source_key(root_a), trainer._checkpoint_source_key(root_b)) + ) + assert [logical_path for logical_path, _source in installed] == [ + logical_a, + logical_b, + logical_c, + ] + assert set(trainer._checkpoint_slots) == { + logical_a, + logical_b, + logical_c, + } + assert trainer._default_slot_ref == trainer._slot_ref(logical_c) + for logical_path in (logical_a, logical_b, logical_c): + request = ForwardInput(input_tokens=torch.tensor([1]), checkpoint=logical_path) + assert trainer._resolve_slot_ref(request) == trainer._slot_ref(logical_path) + + refreshed_root = str(tmp_path / "immutable-new") + await trainer.load_checkpoint(MaterializedCheckpoint(logical_a, refreshed_root)) + assert installed[-1][0] == logical_a + assert prepared[-1] == trainer._checkpoint_source_key(refreshed_root) + + prepared_before_push = tuple(prepared) + pushed = trainer.push_checkpoint( + MaterializedCheckpoint(logical_a, str(tmp_path / "unused-while-loaded")) + ) + await pushed + assert tuple(prepared) == prepared_before_push + assert trainer._slot_stack == [trainer._slot_ref(logical_a)] + trainer.pop_checkpoint() + + +async def test_prefetch_does_not_silently_ignore_empty_materialized_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) + seen: list[str] = [] + + async def prefetch(path: str) -> object: + seen.append(path) + return object() + + monkeypatch.setattr(trainer, "_prefetch_checkpoint", prefetch) + await trainer.prefetch_checkpoints(MaterializedCheckpoint("logical", "")) + assert seen == [""] + + +async def test_checkpoint_mutations_follow_call_order_and_recover_from_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) + trainer._checkpoint_prefetches = {} + started = {name: asyncio.Event() for name in ("first", "second")} + ready = {name: asyncio.Event() for name in ("first", "second")} + installed: list[str] = [] + + async def prefetch(path: str) -> object: + event = started.get(path) + if event is not None: + event.set() + await ready[path].wait() + return object() + + def install(trainer: TrainerRank, _source: object, path: str) -> None: + installed.append(path) + if path == "bad": + raise RuntimeError("injected load failure") + trainer._checkpoint_slots.setdefault(path, _CheckpointSlot()).params = () + trainer._checkpoint_slots[path].revision = 0 + + monkeypatch.setattr(trainer, "_prefetch_checkpoint", prefetch) + monkeypatch.setattr("art.trainer_rank._checkpoint.load_checkpoint", install) + first = trainer.load_checkpoint("first") + second = trainer.load_checkpoint("second") + await asyncio.gather(*(event.wait() for event in started.values())) + ready["second"].set() + await asyncio.sleep(0) + assert installed == [] + ready["first"].set() + await asyncio.gather(first, second) + assert installed == ["first", "second"] + + bad = trainer.load_checkpoint("bad") + after = trainer.load_checkpoint("after") + with pytest.raises(RuntimeError, match="injected load failure"): + await bad + await after + assert installed[-2:] == ["bad", "after"] + + +@pytest.mark.parametrize("local_failure", [False, True]) +async def test_checkpoint_prefetch_failures_are_coordinated( + monkeypatch: pytest.MonkeyPatch, local_failure: bool +) -> None: + trainer = TrainerRank(_runtime()) + + async def prefetch(_path: str) -> object: + if local_failure: + raise OSError("rank-local prefetch failed") + return object() + + def coordinated( + error: BaseException | None, phase: str, _group: object | None = None + ) -> None: + assert phase == "prepare checkpoint" + assert isinstance(error, OSError) is local_failure + raise RuntimeError("a rank failed to prepare checkpoint") + + monkeypatch.setattr(trainer, "_prefetch_checkpoint", prefetch) + monkeypatch.setattr("art.trainer_rank._checkpoint.raise_distributed", coordinated) + with pytest.raises(RuntimeError, match="a rank failed"): + await trainer.load_checkpoint("student") def test_trainer_rank_rejects_adapter_keys_without_installed_lora_site() -> None: @@ -405,11 +808,10 @@ def test_trainer_rank_rejects_adapter_keys_without_installed_lora_site() -> None "base.layer.lora_A.weight": torch.empty(1), "base.layer.lora_B.weight": torch.empty(1), } - trainer._prepare_adapter_model("checkpoint", "student", valid) + trainer._prepare_adapter_model("student", valid) with pytest.raises(ValueError, match="matching LoRA target modules"): trainer._prepare_adapter_model( - "checkpoint", "student", {**valid, "base.other.lora_A.weight": torch.empty(1)}, ) @@ -423,7 +825,7 @@ def test_trainer_rank_normalizes_adapter_tensors_to_installed_site() -> None: "base.layer.lora_B.weight": torch.ones(5, 3, dtype=torch.float32), } - normalized = trainer._prepare_adapter_model("checkpoint", "student", adapter) + normalized = trainer._prepare_adapter_model("student", adapter) assert all(tensor.device == site.A_T.device for tensor in normalized.values()) assert all(tensor.dtype == torch.bfloat16 for tensor in normalized.values()) @@ -438,20 +840,41 @@ def test_checkpoint_slot_adapter_config_is_validated_and_copied() -> None: "target_modules": ["q_proj"], } - retained = trainer._validate_checkpoint_slot_adapter_config( - "student", config, alpha=16 - ) + retained = trainer._validate_checkpoint_adapter_config("student", config, alpha=16) assert retained == config config["target_modules"].append("v_proj") # type: ignore[union-attr] assert retained is not None assert retained["target_modules"] == ["q_proj"] with pytest.raises(ValueError, match="conflicts"): - trainer._validate_checkpoint_slot_adapter_config("student", config, alpha=32) + trainer._validate_checkpoint_adapter_config("student", config, alpha=32) with pytest.raises(ValueError, match="missing"): - trainer._validate_checkpoint_slot_adapter_config( - "student", {"r": 8}, alpha=None - ) + trainer._validate_checkpoint_adapter_config("student", {"r": 8}, alpha=None) + + +def test_qwen35_checkpoint_adapter_config_captures_attention_dimensions() -> None: + runtime = _runtime() + runtime.provider.num_attention_heads = 16 + runtime.provider.num_query_groups = 4 + runtime.provider.kv_channels = 128 + trainer = TrainerRank(runtime) + + retained = trainer._validate_checkpoint_adapter_config( + "student", + { + "base_model_name_or_path": "Qwen/Qwen3.5-4B", + "r": 8, + "lora_alpha": 16, + "target_modules": ["q_proj"], + }, + alpha=16, + ) + + assert retained is not None + assert retained["num_attention_heads"] == 16 + assert retained["num_key_value_heads"] == 4 + assert retained["head_dim"] == 128 + assert retained["hidden_size"] == 4 @pytest.mark.parametrize( @@ -477,7 +900,7 @@ def test_checkpoint_slot_adapter_config_rejects_invalid_field_types( config[field] = value with pytest.raises(TypeError, match=field): - trainer._validate_checkpoint_slot_adapter_config("student", config, alpha=None) + trainer._validate_checkpoint_adapter_config("student", config, alpha=None) def test_checkpoint_slot_adapter_config_rejects_cross_rank_mismatch( @@ -487,43 +910,21 @@ def test_checkpoint_slot_adapter_config_rejects_cross_rank_mismatch( monkeypatch.setattr("art.trainer_rank.dist.is_initialized", lambda: True) monkeypatch.setattr("art.trainer_rank.dist.get_world_size", lambda: 2) - def gather(output: list[object], value: object) -> None: - output[:] = [value, {"different": True}] + checkpoint_group = cast(dist.ProcessGroup, object()) + trainer._checkpoint_process_group = checkpoint_group + trainer._checkpoint_finalize_process_group = cast(dist.ProcessGroup, object()) + + def gather( + output: list[object], value: object, *, group: object | None = None + ) -> None: + assert group is checkpoint_group + revision = value[1] if isinstance(value, tuple) and len(value) == 2 else None + output[:] = [value, ({"different": True}, revision)] monkeypatch.setattr("art.trainer_rank.dist.all_gather_object", gather) with pytest.raises(ValueError, match="differs across ranks"): - trainer._validate_checkpoint_slot_adapter_config("student", None, alpha=None) - - -def test_load_checkpoint_slot_retains_config_and_uses_its_alpha( - monkeypatch: pytest.MonkeyPatch, -) -> None: - trainer = TrainerRank(_runtime()) - seen: dict[str, object] = {} - monkeypatch.setattr( - trainer, - "_load_slot", - lambda *_args, **kwargs: seen.update(kwargs) or 1, - ) - monkeypatch.setattr(trainer, "_validate_dynamic_slot_consistency", lambda *_: ()) - monkeypatch.setattr( - trainer, "_validate_loaded_checkpoint_slot_config", lambda *_: None - ) - config = { - "base_model_name_or_path": "Qwen/Qwen3-8B", - "r": 8, - "lora_alpha": 16, - "target_modules": ["q_proj"], - } - - trainer.load_checkpoint_slot("student", {}, adapter_config=config) - - assert seen["alpha"] == 16 - assert trainer._checkpoint_slot_adapter_configs["student"] == config - trainer.load_checkpoint_slot("student", {}, alpha=7) - assert seen["alpha"] == 7 - assert "student" not in trainer._checkpoint_slot_adapter_configs + trainer._validate_checkpoint_adapter_config("student", None, alpha=None) @pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") @@ -533,12 +934,20 @@ def test_slot_load_canonicalizes_only_local_incoming_adapter( calls: list[tuple[dict[str, torch.Tensor], object]] = [] loaded_state: dict[str, torch.Tensor] = {} runtime = _runtime() - runtime.model_support_handler.canonicalize_loaded_lora_state = lambda state, model: ( - calls.append((state, model)) - or {key: torch.zeros_like(value) for key, value in state.items()} + monkeypatch.setattr( + runtime.model_support_handler, + "canonicalize_loaded_lora_state", + lambda state, model: ( + calls.append((state, model)) + or {key: torch.zeros_like(value) for key, value in state.items()} + ), ) - runtime.model_support_handler.zero_internal_padding_params = lambda _model: ( - pytest.fail("slot load must not mutate unrelated slot parameters") + monkeypatch.setattr( + runtime.model_support_handler, + "zero_internal_padding_params", + lambda _model: pytest.fail( + "slot load must not mutate unrelated slot parameters" + ), ) trainer = TrainerRank(runtime) monkeypatch.setattr( @@ -548,7 +957,17 @@ def test_slot_load_canonicalizes_only_local_incoming_adapter( monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2) - def gather_expected(values: list[set[str] | None], local: set[str]) -> None: + checkpoint_group = cast(dist.ProcessGroup, object()) + trainer._checkpoint_process_group = checkpoint_group + trainer._checkpoint_finalize_process_group = cast(dist.ProcessGroup, object()) + + def gather_expected( + values: list[set[str] | None], + local: set[str], + *, + group: object | None = None, + ) -> None: + assert group is checkpoint_group values[:] = [local, {"remote_weight"}] monkeypatch.setattr(torch.distributed, "all_gather_object", gather_expected) @@ -569,7 +988,7 @@ def load_slot( ) adapter = {"weight": torch.ones(1), "remote_weight": torch.ones(1)} - trainer._load_slot("checkpoint", "student", adapter, trainable=True, alpha=None) + trainer._load_checkpoint_slot("student", adapter, alpha=1.0) assert calls == [({"weight": adapter["weight"]}, runtime.model)] torch.testing.assert_close(loaded_state["weight"], torch.zeros(1)) @@ -577,15 +996,788 @@ def load_slot( torch.testing.assert_close(adapter["weight"], torch.ones(1)) -def test_checkpoint_slot_publish_requires_retained_adapter_config() -> None: +def test_checkpoint_export_requires_retained_adapter_config() -> None: trainer = TrainerRank(_runtime()) - with pytest.raises(ValueError, match="Unknown checkpoint slot"): - trainer.save_checkpoint_slot_lora("missing", "/unused") - - trainer._checkpoint_slot_params_by_name["student"] = () + with pytest.raises(ValueError, match="Unknown checkpoint"): + trainer.export_lora("/unused", "missing") + trainer._checkpoint_slots.setdefault("student", _CheckpointSlot()).params = () with pytest.raises(TrainerRankSlotStateError, match="adapter_config"): - trainer.save_checkpoint_slot_lora("student", "/unused") + trainer.export_lora("/unused", "student") + + +def test_checkpoint_save_rejects_accumulated_gradients() -> None: + trainer = TrainerRank(_runtime()) + parameter = torch.nn.Parameter(torch.ones(2)) + parameter.grad = torch.ones_like(parameter) + trainer._checkpoint_slots.setdefault("student", _CheckpointSlot()).params = ( + parameter, + ) + trainer._checkpoint_slots["student"].config = { + "base_model_name_or_path": "test", + "r": 1, + "lora_alpha": 1, + "target_modules": [], + } + + with pytest.raises(TrainerRankSlotStateError, match="accumulated gradients"): + _validate_save_state(trainer, "student") + + +def _canonical_checkpoint(root: Path) -> CheckpointManifest: + from safetensors.torch import save_file + + root.mkdir() + config = { + "base_model_name_or_path": "test/model", + "r": 1, + "lora_alpha": 1, + "target_modules": ["q_proj"], + "art_lora_format": "art-trainer-rank-v1", + } + (root / "adapter_config.json").write_text(json.dumps(config)) + key = "layer.q_proj.lora_A.weight" + save_file({key: torch.ones(1, 2)}, root / "adapter_model.safetensors") + (root / "optimizer").mkdir() + files = [] + for component in ("master", "exp_avg", "exp_avg_sq"): + relative = f"optimizer/{component}.safetensors" + save_file({key: torch.ones(1, 2)}, root / relative) + files.append(relative) + manifest: CheckpointManifest = { + "format_version": 1, + "base_model_name_or_path": "test/model", + "optimizer": OptimizerConfig( + learning_rate=1e-3, + beta1=0.9, + beta2=0.99, + eps=1e-8, + weight_decay=0.1, + ), + "parameters": {key: files}, + "steps": {key: 3.0}, + "files": {}, + "digest": "", + } + payloads = {"adapter_config.json", "adapter_model.safetensors", *files} + manifest["files"] = { + relative: _file_digest(root / relative) for relative in payloads + } + manifest["digest"] = _manifest_digest(manifest) + (root / "checkpoint.json").write_text(json.dumps(manifest)) + return manifest + + +@pytest.mark.parametrize( + "mutate", + ( + lambda manifest: manifest["steps"].update({next(iter(manifest["steps"])): 4.0}), + lambda manifest: manifest["optimizer"].update({"eps": 1e-6}), # type: ignore[union-attr] + lambda manifest: manifest["parameters"].update( + {next(iter(manifest["parameters"])): ("../bad", "x", "y")} + ), + ), +) +def test_checkpoint_manifest_semantics_are_authenticated( + tmp_path: Path, mutate: object +) -> None: + root = tmp_path / "checkpoint" + manifest = _canonical_checkpoint(root) + cast(Any, mutate)(manifest) + (root / "checkpoint.json").write_text(json.dumps(manifest)) + + with pytest.raises(RuntimeError, match="digest mismatch|Unsafe checkpoint"): + prepare_checkpoint(str(root)) + + +@pytest.mark.parametrize("extra", (True, False)) +def test_checkpoint_optimizer_mapping_must_match_adapter( + tmp_path: Path, extra: bool +) -> None: + root = tmp_path / "checkpoint" + manifest = _canonical_checkpoint(root) + if extra: + manifest["parameters"]["unexpected"] = next( + iter(manifest["parameters"].values()) + ) + manifest["steps"]["unexpected"] = 0 + else: + manifest["parameters"].pop(next(iter(manifest["parameters"]))) + (root / "checkpoint.json").write_text(json.dumps(manifest)) + + with pytest.raises(RuntimeError, match="mapping differs"): + prepare_checkpoint(str(root)) + + +def test_materialize_lora_validates_exact_artifact_without_optimizer_downloads( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + source = tmp_path / "source" + manifest = _canonical_checkpoint(source) + local = tmp_path / "local" + local.mkdir() + for name in ("adapter_config.json", "adapter_model.safetensors", "checkpoint.json"): + (local / name).write_bytes((source / name).read_bytes()) + entries = { + "adapter_config.json", + "adapter_model.safetensors", + "checkpoint.json", + *(file for files in manifest["parameters"].values() for file in files), + } + monkeypatch.setattr( + "art.megatron.model_support.lora_disk.normalize_lora_checkpoint_to_vllm", + lambda _path: None, + ) + + output = tmp_path / "output" + materialize_lora( + local, + output, + require_optimizer=True, + artifact_entries=entries, + expected_digest=manifest["digest"], + ) + assert {path.name for path in output.iterdir()} == { + "adapter_config.json", + "adapter_model.safetensors", + } + + from safetensors.torch import save_file + + save_file( + {next(iter(manifest["parameters"])): torch.zeros(1, 2)}, + local / "adapter_model.safetensors", + ) + with pytest.raises(RuntimeError, match="file digest mismatch"): + materialize_lora( + local, + tmp_path / "corrupt-output", + require_optimizer=True, + artifact_entries=entries, + expected_digest=manifest["digest"], + ) + + manifest["files"]["adapter_model.safetensors"] = _file_digest( + local / "adapter_model.safetensors" + ) + (local / "checkpoint.json").write_text(json.dumps(manifest)) + with pytest.raises(RuntimeError, match="Checkpoint digest mismatch"): + materialize_lora( + local, + tmp_path / "tampered-manifest-output", + require_optimizer=True, + artifact_entries=entries, + expected_digest=manifest["digest"], + ) + (local / "checkpoint.json").write_bytes((source / "checkpoint.json").read_bytes()) + + with pytest.raises(RuntimeError, match="digest mismatch"): + materialize_lora( + local, + tmp_path / "bad-digest", + artifact_entries=entries, + expected_digest="bad", + ) + with pytest.raises(RuntimeError, match="missing entries"): + materialize_lora( + local, + tmp_path / "missing-entry", + require_optimizer=True, + artifact_entries={"adapter_config.json", "adapter_model.safetensors"}, + ) + + +def _save_state_trainer() -> TrainerRank: + trainer = TrainerRank(_runtime()) + trainer._checkpoint_process_group = None + return trainer + + +def _prepared_save(root: Path, sequence: int) -> _PreparedSave: + snapshot = root / f"snapshot-{sequence}" + reservation = root / f"reserved-{sequence}" + snapshot.mkdir() + reservation.mkdir() + return _PreparedSave( + sequence=sequence, + snapshot=snapshot, + reservation=reservation, + destination=root / f"output-{sequence}", + config={}, + shards=(), + optimizer=None, + ) + + +def test_optimizer_shards_are_received_as_float32( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from art.megatron.lora import LoraShardMeta + + prepared = _prepared_save(tmp_path, 0) + metadata = [ + LoraShardMeta( + "weight", + 1, + (2, 3), + "bfloat16", + {"kind": "replicated"}, + "block", + ) + ] + received: list[torch.dtype] = [] + monkeypatch.setattr("art.trainer_rank._checkpoint._rank", lambda: 0) + monkeypatch.setattr( + "art.trainer_rank._checkpoint.raise_distributed", lambda *_args: None + ) + + def recv(tensor: torch.Tensor, **_kwargs: object) -> None: + received.append(tensor.dtype) + + monkeypatch.setattr(dist, "recv", recv) + monkeypatch.setattr( + "art.megatron.weights.lora_publish.merge_sharded_adapter_entries", + lambda entries: { + key: values[0][1] for key, values in cast(dict, entries).items() + }, + ) + + merged = _merge_component(prepared, metadata, "master", None) + + assert received == [torch.float32] + assert merged["weight"].dtype == torch.float32 + + +def test_checkpoint_fifo_abort_and_failure_do_not_block_later_save( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + trainer = _save_state_trainer() + first = _prepared_save(tmp_path, 0) + second = _prepared_save(tmp_path, 1) + third = _prepared_save(tmp_path, 2) + trainer._prepared_checkpoint_saves = { + "first": first, + "second": second, + "third": third, + } + calls: list[int] = [] + + def finalize(_trainer: TrainerRank, prepared: _PreparedSave) -> None: + calls.append(prepared.sequence) + if prepared.sequence == 1: + raise RuntimeError("injected finalization failure") + + monkeypatch.setattr("art.trainer_rank._checkpoint._finish", finalize) + abort_checkpoint_save(trainer, "first") + with pytest.raises(RuntimeError, match="injected"): + finish_checkpoint_save(trainer, "second") + finish_checkpoint_save(trainer, "third") + + assert calls == [1, 2] + assert trainer._checkpoint_save_next == 3 + + +def test_checkpoint_out_of_order_finalization_fails_without_blocking( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + trainer = _save_state_trainer() + first = _prepared_save(tmp_path, 0) + second = _prepared_save(tmp_path, 1) + trainer._prepared_checkpoint_saves = {"first": first, "second": second} + monkeypatch.setattr("art.trainer_rank._checkpoint._finish", lambda *_args: None) + + with pytest.raises(RuntimeError, match="finalized in preparation order"): + finish_checkpoint_save(trainer, "second") + finish_checkpoint_save(trainer, "first") + finish_checkpoint_save(trainer, "second") + + assert trainer._checkpoint_save_next == 2 + + +def test_concurrent_checkpoint_finish_runs_once( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + trainer = _save_state_trainer() + prepared = _prepared_save(tmp_path, 0) + trainer._prepared_checkpoint_saves = {"save": prepared} + entered = threading.Event() + release = threading.Event() + calls = 0 + + def finalize(_trainer: TrainerRank, _prepared: _PreparedSave) -> None: + nonlocal calls + calls += 1 + entered.set() + assert release.wait(timeout=2) + + monkeypatch.setattr("art.trainer_rank._checkpoint._finish", finalize) + errors: list[BaseException] = [] + + def finish() -> None: + try: + finish_checkpoint_save(trainer, "save") + except BaseException as exc: + errors.append(exc) + + threads = [threading.Thread(target=finish) for _ in range(2)] + threads[0].start() + assert entered.wait(timeout=2) + threads[1].start() + time.sleep(0.05) + release.set() + for thread in threads: + thread.join(timeout=2) + assert not thread.is_alive() + + assert not errors + assert calls == 1 + + +@pytest.mark.parametrize("action", ("finish", "abort")) +def test_checkpoint_cleanup_failure_can_be_retried( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + action: str, +) -> None: + from art.trainer_rank import _checkpoint + + trainer = _save_state_trainer() + prepared = _prepared_save(tmp_path, 0) + trainer._prepared_checkpoint_saves = {"save": prepared} + finalizations = 0 + + def finalize(_trainer: TrainerRank, _prepared: _PreparedSave) -> None: + nonlocal finalizations + finalizations += 1 + + original = _checkpoint.shutil.rmtree + failed = False + + def fail_once(path: Path, ignore_errors: bool = False, **_: object) -> None: + nonlocal failed + if Path(path) == prepared.snapshot and not failed: + failed = True + raise OSError("injected cleanup failure") + original(path, ignore_errors=ignore_errors) + + monkeypatch.setattr(_checkpoint, "_finish", finalize) + monkeypatch.setattr(_checkpoint.shutil, "rmtree", fail_once) + operation = finish_checkpoint_save if action == "finish" else abort_checkpoint_save + with pytest.raises(BaseExceptionGroup, match="cleanup failed"): + operation(trainer, "save") + operation(trainer, "save") + + assert finalizations == (1 if action == "finish" else 0) + assert "save" not in trainer._prepared_checkpoint_saves + assert not prepared.snapshot.exists() + assert not prepared.reservation.exists() + + +def test_checkpoint_cleanup_gather_failure_releases_finalizer( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from art.trainer_rank import _checkpoint + + trainer = _save_state_trainer() + prepared = _prepared_save(tmp_path, 0) + trainer._prepared_checkpoint_saves = {"save": prepared} + original = _checkpoint._gather + failed = False + + def fail_once( + value: object, group: dist.ProcessGroup | None = None + ) -> tuple[object, ...]: + nonlocal failed + if isinstance(value, tuple) and len(value) == 2 and not failed: + failed = True + raise RuntimeError("injected cleanup gather failure") + return original(value, group) + + monkeypatch.setattr(_checkpoint, "_finish", lambda *_: None) + monkeypatch.setattr(_checkpoint, "_gather", fail_once) + with pytest.raises(RuntimeError, match="cleanup gather"): + finish_checkpoint_save(trainer, "save") + assert "save" not in trainer._checkpoint_finalizing_saves + finish_checkpoint_save(trainer, "save") + assert "save" not in trainer._prepared_checkpoint_saves + + +def test_checkpoint_asymmetric_cleanup_gather_can_converge( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from art.trainer_rank import _checkpoint + + completed = _save_state_trainer() + retained = _save_state_trainer() + retained_root = tmp_path / "retained" + retained_root.mkdir() + retained_save = _prepared_save(retained_root, 0) + completed._finalized_checkpoint_saves["save"] = _FinalizedSave(0, "finish") + retained._prepared_checkpoint_saves["save"] = retained_save + retained._checkpoint_save_outcomes["save"] = "finish" + + def mixed( + value: object, _group: dist.ProcessGroup | None = None + ) -> tuple[object, ...]: + if isinstance(value, bool): + return (True, False) + return (value, value) + + monkeypatch.setattr(_checkpoint, "_gather", mixed) + finish_checkpoint_save(completed, "save") + finish_checkpoint_save(retained, "save") + assert "save" in completed._finalized_checkpoint_saves + assert "save" in retained._finalized_checkpoint_saves + assert "save" not in retained._prepared_checkpoint_saves + + +def test_checkpoint_cleanup_gather_preserves_finish_error( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from art.trainer_rank import _checkpoint + + trainer = _save_state_trainer() + prepared = _prepared_save(tmp_path, 0) + trainer._prepared_checkpoint_saves = {"save": prepared} + original = _checkpoint._gather + + def fail_cleanup( + value: object, group: dist.ProcessGroup | None = None + ) -> tuple[object, ...]: + if isinstance(value, tuple) and len(value) == 2: + raise RuntimeError("cleanup collective failed") + return original(value, group) + + def fail_finish(*_: object) -> None: + raise ValueError("snapshot failed") + + monkeypatch.setattr(_checkpoint, "_finish", fail_finish) + monkeypatch.setattr( + _checkpoint, "_cleanup_paths", lambda *_: OSError("unlink failed") + ) + monkeypatch.setattr(_checkpoint, "_gather", fail_cleanup) + with pytest.raises(BaseExceptionGroup) as raised: + finish_checkpoint_save(trainer, "save") + assert any(isinstance(error, ValueError) for error in raised.value.exceptions) + assert any(isinstance(error, OSError) for error in raised.value.exceptions) + assert any(isinstance(error, RuntimeError) for error in raised.value.exceptions) + + +def test_checkpoint_prepare_preserves_foreign_reservation(tmp_path: Path) -> None: + trainer = _save_state_trainer() + trainer._checkpoint_slots["student"] = _CheckpointSlot( + config={ + "base_model_name_or_path": "test", + "r": 1, + "lora_alpha": 1, + "target_modules": [], + } + ) + output = tmp_path / "save" + reservation = tmp_path / ".save.reserved" + reservation.mkdir() + marker = reservation / "owner" + marker.write_text("foreign") + + with pytest.raises(FileExistsError): + prepare_checkpoint_save(trainer, str(output), "student") + + assert marker.read_text() == "foreign" + + +def test_checkpoint_prepare_reports_snapshot_cleanup_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from art.trainer_rank import _checkpoint + + trainer = _save_state_trainer() + trainer._checkpoint_slots["student"] = _CheckpointSlot( + config={ + "base_model_name_or_path": "test", + "r": 1, + "lora_alpha": 1, + "target_modules": [], + } + ) + original = _checkpoint.shutil.rmtree + + def fail_snapshot(path: Path, ignore_errors: bool = False, **_: object) -> None: + if ".snapshot-" in Path(path).name: + raise OSError("injected cleanup failure") + original(path, ignore_errors=ignore_errors) + + monkeypatch.setattr( + _checkpoint, + "_local_state", + lambda *_args: (_ for _ in ()).throw(RuntimeError("snapshot failed")), + ) + monkeypatch.setattr(_checkpoint.shutil, "rmtree", fail_snapshot) + + with pytest.raises(BaseExceptionGroup) as captured: + prepare_checkpoint_save(trainer, str(tmp_path / "save"), "student") + messages = " ".join(str(error) for error in captured.value.exceptions) + assert "snapshot failed" in messages + assert "cleanup" in messages + + +def _checkpoint_load_failure_worker( + rank: int, world_size: int, init_method: str, phase: str +) -> None: + dist.init_process_group( + "gloo", + rank=rank, + world_size=world_size, + init_method=init_method, + timeout=timedelta(seconds=15), + ) + from art.trainer_rank import _checkpoint as checkpoint_module + + originals = ( + checkpoint_module._load_adapter, + checkpoint_module._optimizer_state, + checkpoint_module._commit_slot, + ) + try: + trainer = TrainerRank.__new__(TrainerRank) + trainer.runtime = SimpleNamespace( + model=[], + model_identifier=None, + model_support_spec=None, + provider=SimpleNamespace(), + ) + trainer._checkpoint_process_group = None + trainer._checkpoint_slots = {} + trainer._slot_stack = [] + trainer._local_lora_adapter_templates = lambda: {} # type: ignore[method-assign] + trainer._guard_slot_can_load = lambda _ref: None # type: ignore[method-assign] + trainer._load_checkpoint_slot = lambda *_args, **_kwargs: 1 # type: ignore[method-assign] + trainer._validate_checkpoint_consistency = lambda *_args: () # type: ignore[method-assign] + trainer._validate_loaded_checkpoint_config = lambda *_args: None # type: ignore[method-assign] + trainer._restore_canonical_optimizer = lambda *_args: cast(Any, object()) # type: ignore[method-assign] + if phase == "export": + if rank == 1: + trainer._checkpoint_slots["student"] = _CheckpointSlot( + config={ + "base_model_name_or_path": "test/model", + "r": 1, + "lora_alpha": 1, + "target_modules": [], + } + ) + with pytest.raises((ValueError, RuntimeError), match="Unknown|Another"): + checkpoint_module.export_lora(trainer, "/unused", "student") + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == world_size + return + + optimizer = ( + OptimizerConfig( + learning_rate=1e-3, + beta1=0.9, + beta2=0.99, + eps=1e-8, + weight_decay=0.1, + ) + if phase == "optimizer" + else None + ) + manifest: CheckpointManifest | None = ( + { + "format_version": 1, + "base_model_name_or_path": "test/model", + "optimizer": optimizer, + "parameters": {}, + "steps": {}, + "files": {}, + "digest": "digest", + } + if phase != "read" + else None + ) + source = PreparedCheckpoint( + Path("/unused"), + { + "base_model_name_or_path": "test/model", + "r": 1, + "lora_alpha": 1, + "target_modules": [], + }, + (), + manifest, + "digest", + ) + + setattr( + checkpoint_module, + "_load_adapter", + ( + lambda *_args: ( + (_ for _ in ()).throw(RuntimeError("injected snapshot read")) + if phase == "read" and rank == 1 + else {} + ) + ), + ) + setattr( + checkpoint_module, + "_optimizer_state", + ( + lambda *_args: ( + (_ for _ in ()).throw(RuntimeError("injected optimizer read")) + if phase == "optimizer" and rank == 1 + else LocalOptimizerState( + (), (), (), (), cast(OptimizerConfig, optimizer) + ) + ) + ), + ) + setattr( + checkpoint_module, + "_commit_slot", + ( + lambda *_args: ( + (_ for _ in ()).throw(RuntimeError("injected rank-zero commit")) + if phase == "commit" and rank == 0 + else None + ) + ), + ) + + with pytest.raises(RuntimeError, match="injected|Another rank failed"): + checkpoint_module.load_checkpoint(trainer, source, "student") + assert "student" not in trainer._checkpoint_slots + assert not any( + name.startswith("__art_loading_") for name in trainer._checkpoint_slots + ) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == world_size + finally: + for name, value in zip( + ("_load_adapter", "_optimizer_state", "_commit_slot"), + originals, + strict=True, + ): + setattr(checkpoint_module, name, value) + dist.destroy_process_group() + + +@pytest.mark.parametrize("phase", ("read", "optimizer", "commit", "export")) +def test_checkpoint_load_failure_is_collective_and_transactional( + tmp_path: Path, phase: str +) -> None: + context = mp.spawn( + _checkpoint_load_failure_worker, + args=(2, f"file://{tmp_path / f'load-{phase}'}", phase), + nprocs=2, + join=False, + ) + deadline = time.monotonic() + 90 + while time.monotonic() < deadline: + if context.join(timeout=1): + return + else: + for process in context.processes: + process.terminate() + pytest.fail(f"collective checkpoint {phase} failure test hung") + + +@pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") +def test_real_checkpoint_codec_restores_exact_next_optimizer_step( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from art.megatron import lora as lora_module + from art.megatron.lora import LoRA + from art.trainer_rank import _checkpoint as checkpoint_module + + monkeypatch.setattr(lora_module.ps, "get_expert_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + lora_module.ps, + "get_data_parallel_rank", + lambda **_kwargs: 0, + ) + config = cast( + Any, + { + "base_model_name_or_path": "test/model", + "r": 2, + "lora_alpha": 2, + "target_modules": ["q_proj"], + }, + ) + adapter = { + "layer.q_proj.lora_A.weight": torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]), + "layer.q_proj.lora_B.weight": torch.tensor( + [[0.2, 0.1], [0.3, 0.4], [0.5, 0.6], [0.7, 0.8]] + ), + } + adam = AdamParams( + learning_rate=3e-4, + beta1=0.8, + beta2=0.95, + weight_decay=0.1, + grad_clip_norm=10, + ) + + def make_trainer() -> TrainerRank: + lora = LoRA("layer.q_proj", 3, 4, 2, 2, torch.float32, torch.device("cpu")) + trainer = TrainerRank(_runtime(lora)) + loaded = trainer._load_checkpoint_slot("student", adapter, alpha=2) + params = trainer._validate_checkpoint_consistency( + "student", loaded, set(adapter) + ) + trainer._checkpoint_slots["student"] = _CheckpointSlot(params, config) + monkeypatch.setattr( + trainer, + "_reduce_dynamic_grads", + lambda params, **_kwargs: tuple(item.grad.float() for item in params), + ) + return trainer + + original = make_trainer() + for parameter in original._checkpoint_slots["student"].params: + parameter.grad = torch.full_like(parameter, 0.25) + original.optim_step(params=adam) + output = tmp_path / "exact" + original.save_checkpoint(str(output), "student") + original.save_checkpoint(str(output), "student") + assert not list(tmp_path.glob(".exact.snapshot-*")) + assert not (tmp_path / ".exact.reserved").exists() + prepared = prepare_checkpoint(str(output)) + assert prepared.manifest is not None + assert prepared.manifest["optimizer"] is not None + + restored_lora = LoRA("layer.q_proj", 3, 4, 2, 2, torch.float32, torch.device("cpu")) + restored = TrainerRank(_runtime(restored_lora)) + monkeypatch.setattr( + restored, + "_reduce_dynamic_grads", + lambda params, **_kwargs: tuple(item.grad.float() for item in params), + ) + checkpoint_module.load_checkpoint(restored, prepared, "student") + + for trainer in (original, restored): + for parameter in trainer._checkpoint_slots["student"].params: + parameter.grad = torch.full_like(parameter, -0.125) + trainer.optim_step(params=adam) + for actual, expected in zip( + restored._checkpoint_slots["student"].params, + original._checkpoint_slots["student"].params, + strict=True, + ): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + restored_optimizer = restored._checkpoint_slots["student"].optimizer + original_optimizer = original._checkpoint_slots["student"].optimizer + assert restored_optimizer is not None and original_optimizer is not None + _assert_nested_tensors_equal( + restored_optimizer.optimizer.state_dict(), + original_optimizer.optimizer.state_dict(), + ) + with pytest.raises(FileExistsError, match="different state"): + original.save_checkpoint(str(output), "student") + assert not list(tmp_path.glob(".exact.snapshot-*")) + assert not (tmp_path / ".exact.reserved").exists() def test_trainer_rank_default_forward_uses_explicit_base_slot() -> None: @@ -596,7 +1788,6 @@ def test_trainer_rank_default_forward_uses_explicit_base_slot() -> None: assert len(plan.groups) == 1 slot = plan.groups[0].slot_ref assert slot is not None - assert getattr(slot, "kind") == "checkpoint" assert getattr(slot, "name") is None @@ -612,7 +1803,7 @@ def test_optim_step_requires_loaded_checkpoint_slot() -> None: def test_optim_step_rejects_loaded_slots_without_grads() -> None: trainer = TrainerRank(_runtime()) - trainer._checkpoint_slot_params_by_name["student"] = ( + trainer._checkpoint_slots.setdefault("student", _CheckpointSlot()).params = ( torch.nn.Parameter(torch.ones(2)), ) @@ -632,8 +1823,10 @@ def test_optim_step_rejects_explicit_slot_subset_with_missing_grads( ready = torch.nn.Parameter(torch.ones(2)) missing = torch.nn.Parameter(torch.ones(2)) ready.grad = torch.ones_like(ready) - trainer._checkpoint_slot_params_by_name["ready"] = (ready,) - trainer._checkpoint_slot_params_by_name["missing"] = (missing,) + trainer._checkpoint_slots.setdefault("ready", _CheckpointSlot()).params = (ready,) + trainer._checkpoint_slots.setdefault("missing", _CheckpointSlot()).params = ( + missing, + ) monkeypatch.setattr( trainer, "_reduce_dynamic_grads", @@ -654,8 +1847,10 @@ def test_optim_step_implicitly_steps_only_slots_with_grads( ready = torch.nn.Parameter(torch.ones(2)) untouched = torch.nn.Parameter(torch.ones(2)) ready.grad = torch.ones_like(ready) - trainer._checkpoint_slot_params_by_name["ready"] = (ready,) - trainer._checkpoint_slot_params_by_name["untouched"] = (untouched,) + trainer._checkpoint_slots.setdefault("ready", _CheckpointSlot()).params = (ready,) + trainer._checkpoint_slots.setdefault("untouched", _CheckpointSlot()).params = ( + untouched, + ) monkeypatch.setattr( trainer, "_reduce_dynamic_grads", @@ -668,8 +1863,8 @@ def test_optim_step_implicitly_steps_only_slots_with_grads( params=AdamParams(learning_rate=1e-2, weight_decay=0.0, grad_clip_norm=10.0) ) - assert "ready" in trainer._dynamic_optimizers - assert "untouched" not in trainer._dynamic_optimizers + assert trainer._checkpoint_slots["ready"].optimizer is not None + assert trainer._checkpoint_slots["untouched"].optimizer is None assert not torch.equal(before_ready, ready) torch.testing.assert_close(untouched, before_untouched) @@ -687,12 +1882,20 @@ def zero_padding_grads(_model: object) -> None: assert param.grad is not None param.grad[-1] = 0.0 - runtime.model_support_handler.zero_internal_padding_grads = zero_padding_grads - runtime.model_support_handler.zero_internal_padding_params = lambda _model: ( - pytest.fail("slot step must not mutate unrelated slot parameters") + monkeypatch.setattr( + runtime.model_support_handler, + "zero_internal_padding_grads", + zero_padding_grads, + ) + monkeypatch.setattr( + runtime.model_support_handler, + "zero_internal_padding_params", + lambda _model: pytest.fail( + "slot step must not mutate unrelated slot parameters" + ), ) trainer = TrainerRank(runtime) - trainer._checkpoint_slot_params_by_name["student"] = (param,) + trainer._checkpoint_slots.setdefault("student", _CheckpointSlot()).params = (param,) monkeypatch.setattr( trainer, "_reduce_dynamic_grads", @@ -711,7 +1914,7 @@ def zero_padding_grads(_model: object) -> None: assert param[-1].item() == 0.0 -def test_checkpoint_slot_optimizer_state_reproduces_exact_next_step( +def test_canonical_optimizer_state_reproduces_exact_next_step( monkeypatch: pytest.MonkeyPatch, ) -> None: adam = AdamParams( @@ -721,31 +1924,51 @@ def test_checkpoint_slot_optimizer_state_reproduces_exact_next_step( weight_decay=0.1, grad_clip_norm=10.0, ) - original, original_param = _trainer_with_checkpoint( monkeypatch, torch.tensor([0.5, -0.25], dtype=torch.bfloat16) ) + original._checkpoint_slots["student"].revision = 0 original_param.grad = torch.tensor([0.2, -0.4], dtype=torch.bfloat16) original.optim_step(params=adam) - state = original.checkpoint_slot_optimizer_state("student") - assert state is not None + dynamic = original._checkpoint_slots["student"].optimizer + assert dynamic is not None + optimizer_state = dynamic.optimizer.state[dynamic.master_params[0]] + group = dynamic.optimizer.param_groups[0] + beta1, beta2 = cast(tuple[float, float], group["betas"]) + state = LocalOptimizerState( + masters=tuple(param.detach().clone() for param in dynamic.master_params), + exp_avgs=(cast(torch.Tensor, optimizer_state["exp_avg"]).clone(),), + exp_avg_sqs=(cast(torch.Tensor, optimizer_state["exp_avg_sq"]).clone(),), + steps=(float(cast(torch.Tensor, optimizer_state["step"]).item()),), + config=OptimizerConfig( + learning_rate=float(group["lr"]), + beta1=beta1, + beta2=beta2, + eps=float(group["eps"]), + weight_decay=float(group["weight_decay"]), + ), + ) restored, restored_param = _trainer_with_checkpoint( monkeypatch, original_param.detach() ) - restored._dynamic_optimizers["student"] = restored._restore_dynamic_optimizer( - "student", state - ) + restored._checkpoint_slots[ + "student" + ].optimizer = restored._restore_canonical_optimizer("student", state) for param in (original_param, restored_param): param.grad = torch.tensor([-0.3, 0.1], dtype=torch.bfloat16) original.optim_step(params=adam) restored.optim_step(params=adam) torch.testing.assert_close(restored_param, original_param, atol=0, rtol=0) - original_state = original.checkpoint_slot_optimizer_state("student") - restored_state = restored.checkpoint_slot_optimizer_state("student") - assert original_state is not None and restored_state is not None - _assert_nested_tensors_equal(restored_state, original_state) + restored_optimizer = restored._checkpoint_slots["student"].optimizer + original_optimizer = original._checkpoint_slots["student"].optimizer + assert restored_optimizer is not None and original_optimizer is not None + _assert_nested_tensors_equal( + restored_optimizer.optimizer.state_dict(), + original_optimizer.optimizer.state_dict(), + ) + assert original._checkpoint_slots["student"].revision == 2 def test_dynamic_optimizer_keeps_fp32_master_weight_and_moments( @@ -765,7 +1988,8 @@ def test_dynamic_optimizer_keeps_fp32_master_weight_and_moments( ) ) - dynamic = trainer._dynamic_optimizers["student"] + dynamic = trainer._checkpoint_slots["student"].optimizer + assert dynamic is not None assert dynamic.master_params[0].dtype == torch.float32 assert param.item() < torch.tensor(0.1, dtype=torch.bfloat16).item() state = dynamic.optimizer.state[dynamic.master_params[0]] @@ -773,35 +1997,26 @@ def test_dynamic_optimizer_keeps_fp32_master_weight_and_moments( assert state["exp_avg_sq"].dtype == torch.float32 -@pytest.mark.parametrize( - ("corruption", "error"), - ( - ("layout", "topology or parameter layout"), - ("missing_master", "master parameters"), - ("shape", "topology or parameter layout"), - ), -) -def test_checkpoint_slot_optimizer_state_rejects_incompatible_state( - corruption: str, - error: str, +def test_canonical_optimizer_rejects_incompatible_local_shape( monkeypatch: pytest.MonkeyPatch, ) -> None: - trainer, param = _trainer_with_checkpoint(monkeypatch, torch.ones(2)) - param.grad = torch.ones_like(param) - trainer.optim_step( - params=AdamParams(learning_rate=1e-2, weight_decay=0.0, grad_clip_norm=10.0) - ) - state = trainer.checkpoint_slot_optimizer_state("student") - assert state is not None - if corruption == "layout": - cast(dict[str, object], state)["layout"] = {"different": True} - elif corruption == "missing_master": - state["master_params"] = () - restored, _ = _trainer_with_checkpoint( - monkeypatch, torch.ones(3 if corruption == "shape" else 2) + trainer, _ = _trainer_with_checkpoint(monkeypatch, torch.ones(2)) + state = LocalOptimizerState( + masters=(torch.ones(3),), + exp_avgs=(torch.zeros(3),), + exp_avg_sqs=(torch.zeros(3),), + steps=(1.0,), + config=OptimizerConfig( + learning_rate=1e-3, + beta1=0.9, + beta2=0.99, + eps=1e-8, + weight_decay=0.0, + ), ) - with pytest.raises(TrainerRankSlotStateError, match=error): - restored._restore_dynamic_optimizer("student", state) + + with pytest.raises(TrainerRankSlotStateError, match="master parameter shape"): + trainer._restore_canonical_optimizer("student", state) @pytest.mark.parametrize("operation", ("load", "step")) @@ -810,7 +2025,7 @@ def test_trainer_rank_rejects_mutating_slot_with_pending_graph( monkeypatch: pytest.MonkeyPatch, ) -> None: trainer = TrainerRank(_runtime()) - ref = _slot_ref("checkpoint", "teacher") + ref = _slot_ref("teacher") monkeypatch.setattr(trainer, "_slot_ref", _slot_ref) target = _tracked_targets(trainer, ref, 2)[0] guard = ( @@ -837,7 +2052,7 @@ def test_trainer_rank_step_allows_missing_slot_graph_bookkeeping( def test_trainer_rank_zero_grad_does_not_clear_live_slot_graphs() -> None: trainer = TrainerRank(_runtime()) - ref = _slot_ref("lora", "teacher") + ref = _slot_ref("teacher") output = ForwardOutput( None, TopK( @@ -858,7 +2073,7 @@ def test_trainer_rank_zero_grad_does_not_clear_live_slot_graphs() -> None: def test_trainer_rank_retained_backward_keeps_slot_graph_guard() -> None: trainer = TrainerRank(_runtime()) - ref = _slot_ref("checkpoint", "teacher") + ref = _slot_ref("teacher") target = _tracked_targets(trainer, ref, 2)[0] target.sum().backward(retain_graph=True) @@ -871,7 +2086,7 @@ def test_trainer_rank_retained_backward_keeps_slot_graph_guard() -> None: def test_trainer_rank_tracks_each_independent_output_graph() -> None: trainer = TrainerRank(_runtime()) - ref = _slot_ref("checkpoint", "teacher") + ref = _slot_ref("teacher") first, second = _tracked_targets(trainer, ref, 2, 3) first.sum().backward() @@ -884,7 +2099,7 @@ def test_trainer_rank_tracks_each_independent_output_graph() -> None: def test_trainer_rank_tracks_graph_after_output_is_replaced_by_loss() -> None: trainer = TrainerRank(_runtime()) - ref = _slot_ref("checkpoint", "teacher") + ref = _slot_ref("teacher") target = _tracked_targets(trainer, ref, 2)[0] loss = target.sum() del target @@ -899,7 +2114,7 @@ def test_trainer_rank_tracks_graph_after_output_is_replaced_by_loss() -> None: def test_trainer_rank_releases_abandoned_output_graph() -> None: trainer = TrainerRank(_runtime()) - ref = _slot_ref("checkpoint", "teacher") + ref = _slot_ref("teacher") target = _tracked_targets(trainer, ref, 2)[0] del target gc.collect() @@ -1102,6 +2317,7 @@ def test_forward_micro_batches_rejects_mismatched_replicated_counts( monkeypatch.setattr(trainer_rank.dist, "is_available", lambda: True) monkeypatch.setattr(trainer_rank.dist, "is_initialized", lambda: True) monkeypatch.setattr(trainer_rank.dist, "get_world_size", lambda: 2) + monkeypatch.setattr(trainer_rank.dist, "all_reduce", lambda *_args, **_kwargs: None) def gather(output, value): output[:] = [value, value + 1] diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py index 509e102d5..ff1803366 100644 --- a/tests/unit/test_trainer_rank_weird_shapes.py +++ b/tests/unit/test_trainer_rank_weird_shapes.py @@ -21,6 +21,7 @@ Unset, ) from art.trainer_rank._impl import ( + _CheckpointSlot, _flatten, _MemoryCheck, _MemoryProfile, @@ -68,7 +69,6 @@ def _target_request( logits: bool = False, hidden_states: bool = False, checkpoint: AdapterSelection = Unset, - lora: AdapterSelection = Unset, ) -> ForwardInput: labels = ( tokens @@ -85,7 +85,6 @@ def _target_request( logits=logits, hidden_states=hidden_states, checkpoint=checkpoint, - lora=lora, ) @@ -442,13 +441,15 @@ def test_heterogeneous_slots_split_packing_without_losing_output_estimates( monkeypatch.setattr( TrainerRank, "_slot_ref", - staticmethod(lambda kind, name: (kind, name)), + staticmethod(lambda name: name), ) - rank.set_checkpoint("student") + rank._default_slot_ref = rank._slot_ref("student") + for name in ("student", "teacher", "critic"): + rank._checkpoint_slots.setdefault(name, _CheckpointSlot()).params = () requests = [ _target_request(_tokens(1, 2, 3), top_k=3), _target_request(_tokens(1, 2, 4), checkpoint=None, logits=True), - _target_request(_tokens(1, 2, 5), lora="teacher", hidden_states=True), + _target_request(_tokens(1, 2, 5), checkpoint="teacher", hidden_states=True), _target_request(_tokens(1, 2, 6), checkpoint="critic", target_count=4), ] @@ -462,10 +463,10 @@ def test_heterogeneous_slots_split_packing_without_losing_output_estimates( assert signature == plan.signature assert plan.signature.slot_group_count == 4 assert {group.slot_ref for group in plan.groups} == { - ("checkpoint", "student"), - ("checkpoint", None), - ("lora", "teacher"), - ("checkpoint", "critic"), + "student", + None, + "teacher", + "critic", } diff --git a/uv.lock b/uv.lock index b93378a40..e79abda3e 100644 --- a/uv.lock +++ b/uv.lock @@ -5427,6 +5427,7 @@ dependencies = [ { name = "anthropic" }, { name = "litellm" }, { name = "nest-asyncio" }, + { name = "numpy" }, { name = "openai" }, { name = "polars" }, { name = "pydantic" }, @@ -5664,6 +5665,7 @@ requires-dist = [ { name = "nest-asyncio", specifier = ">=1.6.0" }, { name = "ninja", marker = "extra == 'megatron'", specifier = ">=1.11.1" }, { name = "ninja", marker = "extra == 'megatron-cu130'", specifier = ">=1.11.1" }, + { name = "numpy", marker = "python_full_version < '3.13'", specifier = "<2" }, { name = "numpy", marker = "extra == 'megatron'", specifier = "<2" }, { name = "numpy", marker = "extra == 'megatron-cu130'", specifier = "<2" }, { name = "numpy", marker = "extra == 'tinker'", specifier = "<2" },