[mikasenghaas]

Summary

Stacked on #4053. Ablations on that PR showed that one cleanup point is enough, so this PR replaces both of its changes:

• CheckpointManager.load_from_path: run gc.collect() + torch.cuda.empty_cache() right after dcp_load returns. This covers every load path: SFT and RL, with and without CPU offload.
• AppState.load_state_dict: remove the cleanup. The new cleanup after dcp_load makes it redundant on every path, including the offload-only cleanup that main already had.
• sft/train.py: remove the synchronize() + empty_cache() before the startup broadcast. It is not needed, and the sync added no safety.

Net diff against main: the cleanup moves from AppState.load_state_dict (offload paths only) to load_from_path (all paths).

Why

During a non-offload resume, peak allocated memory rises about 21 GiB above the steady state. That is about the size of the sharded fp32 parameters: DCP's _init_optim_state builds the optimizer state with a zero-grad step, which allocates zeros_like gradients for every parameter. 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 does not change, so the extra memory is cache fragmentation, not live tensors.

Verification

Same setup as the ablation in #4053: Qwen3-30B-A3B SFT on 1 node of H200s, with 6 trainer GPUs and 2 inference GPUs, online evals on, seq len 8192. Each resume loads a step-1 checkpoint and trains step 2. Numbers are rank 0, in GiB. Each configuration was tested with env-var toggles on top of #4053; those debug toggles are not part of this PR.

offload = "none":

Step 2	Reserved after load	Peak allocated	Peak reserved
Fresh run (no resume)	—	83.7	93.6
Resume on main	80.0	83.5	114.9
Resume on #4053	60.5	83.5	93.5
Resume, this PR	58.6	83.5	93.5

With this PR, reserved memory after the load equals allocated memory (58.6), and step 2 matches the fresh run. The ckpt.py change from #4053 on its own left 1.9 GiB (step-2 peak reserved 95.4).

CPU offload, to check that removing the AppState cleanup is safe:

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	— (job 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

RL was not tested. It uses the same CheckpointManager.load_from_path.

🤖 Generated with Claude Code

---

Note

Low Risk
Timing-only memory hygiene around checkpoint resume; no change to what is loaded or how training steps run.

Overview
Consolidates post-checkpoint-load CUDA memory cleanup into one place: CheckpointManager.load_from_path now runs gc.collect() and torch.cuda.empty_cache() immediately after dcp_load, so SFT/RL resume and all offload modes share the same fix for allocator-cache fragmentation from DCP’s optimizer materialization (zero-grad gradients).

Removes redundant cleanup elsewhere: AppState.load_state_dict no longer clears the load state_dict and empties the cache (that path was offload-specific and duplicated the new hook), and SFT startup no longer calls synchronize() + empty_cache() before the policy weight broadcast—resume loads are already cleaned up centrally.

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