From 7194fd1cd7db049bd2f35afa43779bdeb1a0bd57 Mon Sep 17 00:00:00 2001 From: sami jaghouar Date: Thu, 29 Oct 2026 17:17:48 +0000 Subject: [PATCH] feat(trainer): async checkpoint saves with CPU optimizer offload --- docs/training.md | 2 + .../src/prime_rl/configs/trainer.py | 3 + src/prime_rl/trainer/ckpt.py | 124 ++++++++++++++++-- src/prime_rl/trainer/optim/base.py | 12 ++ src/prime_rl/trainer/optim/offload.py | 6 + src/prime_rl/trainer/optim/state_offload.py | 3 +- src/prime_rl/trainer/rl/train.py | 1 + src/prime_rl/trainer/sft/train.py | 1 + 8 files changed, 139 insertions(+), 13 deletions(-) diff --git a/docs/training.md b/docs/training.md index eb800a1d96e..d0d817166dd 100644 --- a/docs/training.md +++ b/docs/training.md @@ -332,6 +332,8 @@ uv run rl @ rl.toml --ckpt.interval 25 --ckpt.keep-last 3 # rolling window of 3 uv run rl @ rl.toml --ckpt.interval 25 --ckpt.keep-interval 100 # …plus permanent every 100 ``` +With CPU optimizer offload (`model.offload` = `optimizer` or `full`), `ckpt.async_save = true` on the trainer (`--trainer.ckpt.async-save` for `rl`, `--ckpt.async-save` for `sft`) writes the trainer checkpoint in a background thread. The save copies the GPU weights to CPU and passes the CPU optimizer state to the writer without copying it. Training continues, and the next optimizer update of the CPU state waits until the write is done. `time/save_ckpt` then measures only the blocking part. A step counts toward `keep_last` once its write has finished, and the trainer waits for the final write before it exits. + ### Resuming a Run Re-run the same launch command and pass `--resume` (latest checkpoint) or `--resume.step `. Resuming reuses the run directory, so the run needs a name you can point back at — launch with `--run.name` (or pass the first run's auto-generated name). Make sure `--max-steps` is at least the target final step, not the remaining delta: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 4be7a8daaeb..5b9b7035109 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -495,6 +495,9 @@ class CheckpointConfig(BaseConfig): keep_interval: int | None = Field(None, ge=1) """Keep checkpoints at every N steps permanently (e.g. ``keep_interval=100`` keeps step 100, 200, ...). If None, no interval-based keeping.""" + async_save: bool = False + """Write checkpoints in a background thread while training continues. Only applies with CPU optimizer offload (``model.offload`` = ``optimizer`` or ``full``); otherwise saves stay synchronous.""" + skip_scheduler: bool = False """Skip loading the scheduler from checkpoint.""" diff --git a/src/prime_rl/trainer/ckpt.py b/src/prime_rl/trainer/ckpt.py index 8514be4eeac..7698d0119b8 100644 --- a/src/prime_rl/trainer/ckpt.py +++ b/src/prime_rl/trainer/ckpt.py @@ -3,12 +3,14 @@ import json import shutil import time +from concurrent.futures import Future from dataclasses import asdict, dataclass from pathlib import Path from typing import Any import torch from torch import nn +from torch.distributed.checkpoint.staging import AsyncStager from torch.distributed.checkpoint.state_dict import ( get_model_state_dict, get_state_dict, @@ -16,11 +18,14 @@ set_state_dict, ) from torch.distributed.checkpoint.state_dict_loader import load as dcp_load +from torch.distributed.checkpoint.state_dict_saver import async_save as dcp_async_save from torch.distributed.checkpoint.state_dict_saver import save as dcp_save from torch.distributed.checkpoint.stateful import Stateful +from torch.distributed.tensor import DTensor, Shard from torch.nn import Module from torch.optim.lr_scheduler import LRScheduler from torch.optim.optimizer import Optimizer +from torch.utils._pytree import tree_map_only from torchdata.stateful_dataloader import StatefulDataLoader from prime_rl.configs.shared import ResumeConfig @@ -29,6 +34,7 @@ from prime_rl.trainer.models.base import PreTrainedModelPrimeRL from prime_rl.trainer.models.conversion_ops import view_mapping from prime_rl.trainer.models.fusions import ( + get_model_packed_parameters, join_loaded_optimizer_state_for_runtime, split_packed_optimizer_state_for_checkpoint, write_back_loaded_packed_optimizer_state, @@ -55,6 +61,48 @@ def _try_rmtree(path: Path, logger) -> None: logger.warning(f"Failed to remove {path}: {e}, skipping cleanup") +class _NoCopyStager(AsyncStager): + """Hands the state dict to the background writer as is: it is already a CPU snapshot.""" + + _synchronize_after_execute = False + + def stage(self, state_dict: dict[str, Any]) -> dict[str, Any]: + return state_dict + + +def _copy_gpu_tensors_to_cpu(state_dict: dict[str, Any]) -> dict[str, Any]: + """Copy GPU tensors to CPU and keep CPU tensors by reference. The containers are rebuilt, so + the result shares no dict or list with the live optimizer state.""" + + def to_cpu(tensor: torch.Tensor) -> torch.Tensor: + if isinstance(tensor, DTensor): + if tensor._local_tensor.is_cpu: + return tensor + return DTensor(tensor._local_tensor.cpu(), tensor._spec, requires_grad=False) + return tensor if tensor.is_cpu else tensor.cpu() + + return tree_map_only(torch.Tensor, to_cpu, state_dict) + + +def _copy_cpu_state_split_on_sharded_dim_to_gpu( + model: nn.Module, optimizer_state_dict: dict[str, Any] +) -> dict[str, Any]: + """Splitting a packed parameter's state along its sharded dim all-gathers it, so CPU-resident + state is copied to GPU for that gather to run over NCCL.""" + state = dict(optimizer_state_dict["state"]) + for packed_info in get_model_packed_parameters(model): + if packed_info.fqn not in state: + continue + sharded_on_packing_dim = Shard(packed_info.spec.dim) + state[packed_info.fqn] = { + key: DTensor(value._local_tensor.cuda(), value._spec, requires_grad=False) + if isinstance(value, DTensor) and value._local_tensor.is_cpu and sharded_on_packing_dim in value.placements + else value + for key, value in state[packed_info.fqn].items() + } + return {**optimizer_state_dict, "state": state} + + class AppState(Stateful): """ A wrapper for checkpointing the trainer with sharded weights and optimizer @@ -86,14 +134,12 @@ def _get_checkpoint_optimizers(self) -> list[Optimizer]: def _has_cpu_offload(self) -> bool: return any(isinstance(optimizer, OffloadOptimizer) for optimizer in self.optimizers) - def state_dict(self) -> dict[str, Any]: - for optimizer in self.optimizers: - if isinstance(optimizer, OffloadOptimizer): - optimizer.prepare_checkpoint_save() - + def _build_state_dict(self, background_save: bool = False) -> dict[str, Any]: # Automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT checkpoint_optimizers = self._get_checkpoint_optimizers() model_state_dict, optimizer_state_dict = get_state_dict(self.model, checkpoint_optimizers) + if background_save: + optimizer_state_dict = _copy_cpu_state_split_on_sharded_dim_to_gpu(self.model, optimizer_state_dict) # Runtime fusions own one physical state tensor per fused parameter; the checkpoint # sees it under the canonical names. This dict is also the template dcp_load writes # into, and load_state_dict packs the loaded entries back into the runtime state. @@ -107,6 +153,14 @@ def state_dict(self) -> dict[str, Any]: if self.progress is not None: progress_state_dict = asdict(self.progress) state_dict["progress"] = progress_state_dict + return state_dict + + def state_dict(self) -> dict[str, Any]: + for optimizer in self.optimizers: + if isinstance(optimizer, OffloadOptimizer): + optimizer.prepare_checkpoint_save() + + state_dict = self._build_state_dict() for optimizer in self.optimizers: if isinstance(optimizer, OffloadOptimizer): @@ -119,6 +173,11 @@ def state_dict(self) -> dict[str, Any]: return state_dict + def cpu_state_dict(self) -> dict[str, Any]: + """The state dict for a background save: GPU tensors are copied to CPU, while the CPU-resident + optimizer state is passed by reference, so its owner must not update it until the write is done.""" + return _copy_gpu_tensors_to_cpu(self._build_state_dict(background_save=True)) + def load_state_dict(self, state_dict: dict[str, Any]): checkpoint_optimizers = self._get_checkpoint_optimizers() has_cpu_offload = self._has_cpu_offload() @@ -184,6 +243,10 @@ def __init__(self, output_dir: Path, config: CheckpointConfig, resume: ResumeCon else: self.ckpt_steps = all_steps + # Background writes run their DCP collectives on their own gloo group, created on the first one + self._async_save_group: torch.distributed.ProcessGroup | None = None + self._pending_save: tuple[int, Future] | None = None + def get_ckpt_path(self, step: int) -> Path: """Get the path to write the trainer checkpoint for a given step.""" return get_step_path(self.ckpt_dir, step) / "trainer" @@ -196,13 +259,12 @@ def save_to_path( scheduler: LRScheduler, progress: Progress, dataloader: StatefulDataLoader | None = None, - ): - """Save the trainer checkpoint to a given path.""" + ) -> Future | None: + """Save the trainer checkpoint to a given path. A background save returns the future of its write.""" self.logger.debug(f"Saving training checkpoint to {path}") start_time = time.perf_counter() - # Create checkpoint state - state_dict = {"app": AppState(model, optimizers, scheduler, progress)} + app_state = AppState(model, optimizers, scheduler, progress) # Checkpoint the local dataloader if dataloader is not None: @@ -217,9 +279,33 @@ def save_to_path( # Before the sharded state, so the mapping is in place once .metadata appears self._save_mapping(path, model) - dcp_save(state_dict, checkpoint_id=path, process_group=get_dcp_group()) + offload_optimizers = [optimizer for optimizer in optimizers if isinstance(optimizer, OffloadOptimizer)] + if self.config.async_save and offload_optimizers: + if self._async_save_group is None: + self._async_save_group = get_dcp_group() + state_dict = {"app": app_state.cpu_state_dict()} + write_start_time = time.perf_counter() + future = dcp_async_save( + state_dict, checkpoint_id=path, process_group=self._async_save_group, async_stager=_NoCopyStager() + ) + for optimizer in offload_optimizers: + optimizer.hold_state_updates(future) + + def log_write_time(future: Future) -> None: + if future.exception() is None: + write_time = format_time(time.perf_counter() - write_start_time) + self.logger.info(f"Wrote training checkpoint to {path} in the background in {write_time}") + + future.add_done_callback(log_write_time) + self.logger.debug( + f"Started background save of training checkpoint in {format_time(write_start_time - start_time)}" + ) + return future + + dcp_save({"app": app_state}, checkpoint_id=path, process_group=get_dcp_group()) self.logger.debug(f"Saved training checkpoint in {format_time(time.perf_counter() - start_time)}") + return None def _save_mapping(self, path: Path, model: nn.Module) -> None: """Write the prime -> HF conversion of the saved model weights as views, so a reader can @@ -305,7 +391,9 @@ def save( progress: Progress, dataloader: StatefulDataLoader | None = None, ) -> None: - """Save the full checkpoint state for a specified step.""" + """Save the full checkpoint state for a specified step. A background save is only added to + ``ckpt_steps`` once ``wait`` has seen its write finish.""" + self.wait() ckpt_path = self.get_ckpt_path(step) # Master-only mkdir + barrier: concurrent mkdir from every rank can # re-raise FileExistsError on a parallel FS (see save_to_path). @@ -313,7 +401,19 @@ def save( ckpt_path.parent.mkdir(parents=True, exist_ok=True) torch.distributed.barrier(group=get_dcp_group()) - self.save_to_path(ckpt_path, model, optimizers, scheduler, progress, dataloader) + future = self.save_to_path(ckpt_path, model, optimizers, scheduler, progress, dataloader) + if future is None: + bisect.insort(self.ckpt_steps, step) + else: + self._pending_save = (step, future) + + def wait(self) -> None: + """Block until the pending background save is written, raising its error if it failed.""" + if self._pending_save is None: + return + step, future = self._pending_save + self._pending_save = None + future.result() bisect.insort(self.ckpt_steps, step) def maybe_clean(self) -> None: diff --git a/src/prime_rl/trainer/optim/base.py b/src/prime_rl/trainer/optim/base.py index bbdf364f089..7a4f7105295 100644 --- a/src/prime_rl/trainer/optim/base.py +++ b/src/prime_rl/trainer/optim/base.py @@ -1,10 +1,13 @@ from abc import ABC, abstractmethod +from concurrent.futures import Future from typing import TypeAlias from torch.optim import Optimizer class OffloadOptimizer(ABC): + _pending_save: Future | None = None + @property @abstractmethod def base_optimizer(self) -> Optimizer: ... @@ -12,6 +15,15 @@ def base_optimizer(self) -> Optimizer: ... @abstractmethod def checkpoint_optimizer(self) -> Optimizer: ... + def hold_state_updates(self, save: Future) -> None: + """Delay the next in-place update of the CPU optimizer state until ``save`` has written it.""" + self._pending_save = save + + def _wait_for_pending_save(self) -> None: + if self._pending_save is not None: + self._pending_save.result() + self._pending_save = None + def prepare_checkpoint_save(self) -> None: pass diff --git a/src/prime_rl/trainer/optim/offload.py b/src/prime_rl/trainer/optim/offload.py index f0b698c9939..540fe4201dd 100644 --- a/src/prime_rl/trainer/optim/offload.py +++ b/src/prime_rl/trainer/optim/offload.py @@ -86,12 +86,14 @@ def __init__( timeout_seconds: float, target_chunk_numel: int, chunk_ready_callback: Callable[[int], None], + before_optimizer_chunks: Callable[[], None], gradient_dtypes: dict[int, torch.dtype], compute_dtypes: dict[int, torch.dtype], ): self._chunks = chunks self._dp_replicate = dp_replicate self._chunk_ready_callback = chunk_ready_callback + self._before_optimizer_chunks = before_optimizer_chunks self._timeout_seconds = timeout_seconds # Per-step wall-time accumulators; each key is written by exactly one pipeline thread. self._timings: dict[str, float] = {} @@ -493,6 +495,8 @@ def begin_step(self, gradient_scale: float, *, overlap_optimizer: bool) -> None: self._completed_chunks = 0 def begin_backward(self, *, final_backward: bool) -> None: + if final_backward and self._overlap_optimizer: + self._before_optimizer_chunks() with self._condition: self._raise_worker_error() if self._backward_open: @@ -740,6 +744,7 @@ def __init__( timeout_seconds=self._TIMEOUT_SECONDS, target_chunk_numel=self._TARGET_CPU_CHUNK_NUMEL, chunk_ready_callback=self._step_cpu_chunk, + before_optimizer_chunks=self._wait_for_pending_save, gradient_dtypes={id(master.model_param): master.gradient_dtype for master in master_weights.values()}, compute_dtypes={id(master.model_param): master.compute_dtype for master in master_weights.values()}, ) @@ -911,6 +916,7 @@ def step(self, closure=None): self._initialized = True return # Synchronous path (validation steps and checkpoint boundaries). + self._wait_for_pending_save() gradients_by_chunk = [self._gradient_manager.load_cpu_chunk(i) for i in range(len(self._chunks))] for i in range(len(self._chunks)): output_slots, compute_params = self._gradient_manager.acquire_output_chunk(i) diff --git a/src/prime_rl/trainer/optim/state_offload.py b/src/prime_rl/trainer/optim/state_offload.py index 15a7f1eca2c..9d6daaca4bb 100644 --- a/src/prime_rl/trainer/optim/state_offload.py +++ b/src/prime_rl/trainer/optim/state_offload.py @@ -68,7 +68,8 @@ def step(self, closure=None): # Run optimizer step result = self.optimizer.step(closure) - # Move states back to CPU + # Move states back to CPU, into the pinned buffers a background checkpoint save may still read + self._wait_for_pending_save() self._move_states("cpu") return result diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index ce15234661e..581e4a2dada 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -824,6 +824,7 @@ def train(config: TrainerConfig): if config.ckpt is not None or paused: logger.info(f"Saving final checkpoint at step {progress.step}") ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress) + ckpt_manager.wait() if paused: controller.mark_saved(progress.step) else: diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 33774a943ff..62631158d3a 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -676,6 +676,7 @@ def is_online_eval_step(step: int) -> bool: if config.ckpt is not None or paused: logger.info(f"Saving final checkpoint at step {progress.step}") ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) + ckpt_manager.wait() if paused: controller.mark_saved(progress.step) else: