[samsja]

Opt-in async trainer checkpoint saves when CPU optimizer offload is on (model.offload = optimizer or full). The optimizer state already lives in CPU RAM, so the background thread writes it from there. It is not copied again.

What changes

• New ckpt.async_save on the trainer CheckpointConfig (default false). It only applies with CPU optimizer offload. Without offload, or with async_save = false, the save is synchronous as before.

• At save time, on the main thread:

  • GPU tensors (model weights, and in full mode the compute params) are copied to CPU. This is a real snapshot.
  • The CPU-resident optimizer state (and full mode's fp32 masters) is passed to DCP by reference. In optimizer mode the async path skips prepare_checkpoint_save, so the state no longer goes CPU→GPU just so DCP can copy it back.
  • The state dict containers are rebuilt, so the writer does not share the optimizer's live per-parameter state dicts.
  • Packed (fused) parameters whose packing dim is sharded (e.g. qkv with the default shard_fused_on_dim1 = false) get all-gathered when they are split into their logical entries. For those entries only, the CPU state is copied to GPU first, so the gather runs over NCCL. The gathered result is a fresh copy either way.
  • Then dcp.async_save runs with a stager that makes no copy, on a gloo group used only by background writes.

• Consistency: OffloadOptimizer.hold_state_updates(future) makes the optimizer wait for the write right before it next updates its CPU state in place:

  • CPUOffloadOptimizer.step(): after the GPU step, before the states are copied back into the pinned buffers.
  • FullCPUOffloadOptimizer: on the main thread at begin_backward(final_backward=True), before any CPU optimizer chunk is scheduled, and at the start of the synchronous step() path.

  The rest of the next step's forward/backward overlaps the write. A failed write raises at that point.

• CheckpointManager: waits for the previous write before it starts a new save. A step is added to ckpt_steps (and so becomes visible to keep_last) only after its write has finished. wait() runs after the final save in RL and SFT, before mark_saved on pause. The background write time is logged at info level. time/save_ckpt measures the blocking part.

• load is unchanged.

CPU RAM cost

• Optimizer state: nothing extra.
• One CPU copy of each rank's GPU weight shard, held until the write finishes.
• Logical entries of packed params that need a gather (see above) are full replicated copies on every rank until the write finishes. The synchronous path gathers the same tensors, but on GPU. With model.fusions.shard_fused_on_dim1 = true no gather is needed.

Measurements

Qwen3-0.6B SFT (reverse-text), 2x RTX PRO 6000 (FSDP 2), batch 32 × 1024, ckpt.interval = 2, keep_last = 2. Each checkpoint is about 7.7 GB on local disk. time/save_ckpt at steps 2 / 4 / 6:

model.offload	sync	async	background write
optimizer	7.5 / 6.6 / 7.0 s	0.95 / 0.63 / 0.20 s	5.1–6.5 s
full	9.1 / 7.9 / 8.2 s	0.84 / 0.30 / 0.25 s	6.4–8.0 s

On this small model the write (about 6 s) takes longer than a step (3.4 s / 4.2 s). So the step after a save waits about 2.5 s (optimizer) or 3.5 s (full) for the write to finish. Each save then costs about 3 s of stall in optimizer mode instead of about 7 s. With longer steps the write fully overlaps.

Related

• feat(trainer): async DCP checkpoint saves #3874 (closed): the earlier attempt. It staged the full state dict in CPU memory. This PR makes no staging copy of the CPU optimizer state.
• perf(optim): overlap optimizer state offload copies with the AdamW step #4008 touches state_offload.py. This PR adds one line there, before the copy back to CPU in step().
• fix(ckpt): restore optimizer param groups when resuming with CPU offload #4088 touches the offload branch of AppState.load_state_dict. This PR does not change load.

Test plan

• Correctness, both offload modes, 2 GPUs: a temporary test patch did a synchronous reference save of the same state right before each async save. The async checkpoint was then compared with the reference tensor by tensor at steps 2 and 4.
  • With the guard: bitwise equal (0 / 1241 tensors mismatched for optimizer, 0 / 1551 for full). This also holds with shard_fused_on_dim1 = true.
  • With the guard disabled: 170 (optimizer) and 243 (full) optimizer tensors differ (exp_avg, exp_avg_sq, masters). So the test catches a missing guard.
• Resume from the async step-4 checkpoint, both modes: the step-5 loss equals the original run's exactly (4.3025 / 4.3888). Step 6 is within run-to-run noise.
• .metadata exists for every saved step. keep_last = 2 cleanup kept the right steps.
• uv run pytest tests/unit/train/test_state_offload.py (GPU) and tests/unit/test_configs.py tests/unit/utils/test_pathing.py tests/unit/orchestrator/test_config_update.py. ruff check and ruff format.