feat(trainer): async checkpoint saves with CPU optimizer offload

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 <N>`. 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:

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."""
 

src/prime_rl/trainer/ckpt.py

@@ -3,24 +3,29 @@
 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,
     set_model_state_dict,
     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,15 +391,29 @@ 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).
         if self.world.is_master:
             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:

src/prime_rl/trainer/optim/base.py

@@ -1,17 +1,29 @@
 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: ...
 
     @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
 

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)

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

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:

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: