fix: release the allocator cache after every checkpoint load

src/prime_rl/trainer/ckpt.py

@@ -154,11 +154,13 @@ def load_state_dict(self, state_dict: dict[str, Any]):
 
         # state_dict is the same dict object that dcp_load held internally; clearing
         # it drops the last references to the loaded tensor wrappers so the cuda
-        # allocator can release whatever blocks it cached during the read.
-        if has_cpu_offload:
-            state_dict.clear()
-            gc.collect()
-            torch.cuda.empty_cache()
+        # allocator can release whatever blocks it cached during the read. The
+        # non-offload path needs this most: its optimizer state is resident on the
+        # GPU, and under expandable_segments a cached block stays mapped and so is
+        # unavailable to NCCL and cuBLAS workspaces until empty_cache returns it.
+        state_dict.clear()
+        gc.collect()
+        torch.cuda.empty_cache()
 
 
 class CheckpointManager:

src/prime_rl/trainer/sft/train.py

@@ -389,6 +389,12 @@ def is_online_eval_step(step: int) -> bool:
         if world.is_master:
             prune_broadcasts_beyond(config.run_dir, startup_version)
         logger.info(f"Broadcasting startup policy weights (v{startup_version}) for online evals")
+        # The per-layer gather peaks well above the resident weights, and on a resume
+        # it lands on the pool the checkpoint load just left behind. Release cached
+        # blocks first, draining pending work so empty_cache cannot free a block out
+        # from under a running kernel. Mirrors the RL trainer's broadcast.
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
         weight_sender.broadcast(model, startup_version)
 
     logger.info(f"Starting training loop (max_steps={config.max_steps or 'infinite'})")