[mikasenghaas]

Summary

• Release the CUDA allocator cache once after every checkpoint load: gc.collect() + torch.cuda.empty_cache() right after dcp_load returns in CheckpointManager.load_from_path.
• Remove the cleanup from AppState.load_state_dict. It ran only with CPU offload, and the new cleanup now covers that path too.

Net diff: one ckpt.py change. The cleanup moves from the offload-only path to every load path (SFT and RL, with and without CPU offload).

Problem

Without CPU offload (model.offload = "none"), a resume peaks at much more reserved memory than an uninterrupted run at the same step. A 24-node SFT run (192 ranks, 512K sequence length, 267.7 GiB GPUs) showed:

Step 901, same batch	Peak reserved
Fresh run, 900 steps in	216.8 GiB
Resumed from step 900	257.0 GiB (96%)

The resumed run logged expandable_segments: memory mapping failed with OOM during its first step.

Cause: DCP's _init_optim_state builds the optimizer state with a zero-grad optimizer.step(). This allocates a zeros_like gradient for every parameter, so allocated memory peaks about one fp32 parameter shard above the steady state during the load. When the load ends, those blocks stay in the allocator cache. The first step does not reuse them and reserves on top of them. Peak allocated memory does not change, so the extra memory is cache fragmentation, not live tensors. With CPU offload, the old cleanup in AppState.load_state_dict already released the cache, so only the non-offload path was affected.

Verification

Single-node runs: Qwen3-30B-A3B SFT on 1 node of H200s (6 trainer GPUs, 2 inference GPUs), online evals on, seq len 8192. A base run trains 1 step and saves a checkpoint. Each variant resumes from that checkpoint and trains step 2. Each change was toggled with an env var, so all variants share the same code and checkpoint. Numbers are rank 0, in GiB.

offload = "none". Ablation of each candidate change (the first two were in an earlier version of this PR):

Step 2	Reserved after load	Peak allocated	Peak reserved
Fresh run (no resume)	—	83.7	93.6
Resume, main	80.0	83.5	114.9 (+21.3)
Resume, unconditional cleanup in AppState.load_state_dict	60.5	83.5	95.4 (+1.9)
Resume, empty_cache() before the SFT startup broadcast	80.0	83.5	93.5
Resume, both of the above	60.5	83.5	93.5
Resume, this PR	58.6	83.5	93.5

• Peak allocated is the same in every row. Only the cached, unused memory changes.
• On main, the load leaves 21.4 GiB cached (reserved 80.0 vs allocated 58.6). Step 2 reserves about 31 GiB on top of it.
• With this PR, reserved memory after the load equals allocated memory, and step 2 matches the fresh run.
• The cleanup inside AppState.load_state_dict leaves 1.9 GiB cached: some load buffers are freed only after dcp_load returns.
• The empty_cache() before the startup broadcast also removes the gap, but only because it runs right after the load, and only when online evals are on. This PR does not need it.
• Memory outside the PyTorch cache stays at about 2.4 GiB in every run. The cached blocks do not take memory away from NCCL or cuBLAS.
• The offload = "none" variants ran twice, with identical numbers.

CPU offload. Removing the cleanup from AppState.load_state_dict changes nothing:

Resume	Reserved after load	Step-2 peak reserved
offload = "optimizer", main	19.5	82.0
offload = "optimizer", this PR	19.5	82.0
offload = "full", main	9.8	— (killed by an unrelated HF Hub 429 in inference)
offload = "full", both cleanups	9.8	35.1
offload = "full", this PR	9.8	34.9

Not tested: RL resumes (they use the same CheckpointManager.load_from_path), and the 24-node run above with this exact diff.

🤖 Generated with Claude Code

---

Note

Low Risk
Behavior change is limited to memory reclamation after resume; no changes to model, optimizer, or checkpoint I/O logic beyond when empty_cache runs.

Overview
Moves GPU cache cleanup to run on every checkpoint resume, not only when optimizer CPU offload is enabled.

Post-load gc.collect() and torch.cuda.empty_cache() are removed from AppState.load_state_dict (previously gated on has_cpu_offload, with optional state_dict.clear()). The same cleanup now runs in CheckpointManager.load_from_path immediately after dcp_load, with a note that DCP’s optimizer materialization can leave gradient allocations in the CUDA cache and inflate reserved memory on the first training step after resume.

This targets inflated peak reserved memory (and near-OOM on large SFT resumes with model.offload = "none") without changing checkpoint contents or load semantics.

^{Reviewed by Cursor Bugbot for commit 9193c73. Bugbot is set up for automated code reviews on this repo. Configure here.}

[garrett361]

Great if this change is helping with peak memory, but none of the added comment strings align with my understanding of what these. Also, claude is harsh on this point

The ckpt.py change is worth keeping, but almost every explanation the PR gives for it is wrong

lol

[garrett361]

@mikasenghaas did you test without this addition? Just with the ckpt.py edit? I don't believe this is necessary.

Also don't believe calling empty_cache should be able to free any block that's currently used by a kernel.

[garrett361]

I also think this is all more closely related to memory fragmentation. Unconditionally doing clear/collect/empty_cache is more aggressively dropping references and letting the allocator reset its cache before doing further allocations, which may let the newer allocations be less fragmented and reserve less memory.

Do you have the before/after peak allocated memory numbers, too? Would probably be informative.

[garrett361]

Approving this as I don't think any of this is harmful, but I don't think either comment string is accurate and only one of the two file changes is actually addressing what is described in the PR.