From 4dc7543ba451db41edbde12cd07e9fe0bbec9de5 Mon Sep 17 00:00:00 2001 From: Procrastinatorrrr <18924614955@163.com> Date: Mon, 4 May 2026 22:04:43 +0800 Subject: [PATCH 1/2] Fix(checkpoint): add resume/pause in save_model() for offload_train (#1856 regression) When using --colocate with offload_train=True, train() calls self.sleep() which pauses the model via torch_memory_saver.pause(). The subsequent save_model() only rebuilds process groups but fails to resume the model, causing CUDA error during checkpoint save. This fix completes the save_model() offload lifecycle: - Add torch_memory_saver.resume() before save to wake up paused model - Add clear_memory() before save to ensure sufficient memory - Add clear_memory(clear_host_memory=True) after save for cleanup - Add torch_memory_saver.pause() after save to free GPU memory Fixes checkpoint save crash for all users of --colocate with --save-interval. Tested: GRPO training on Qwen3.5-4B with 2x H200 GPUs, checkpoints save successfully. --- slime/backends/megatron_utils/actor.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index adf9ad3a33..c886509931 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -516,6 +516,8 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: # torch dist may trigger nccl communication during saving. if self.args.offload_train: reload_process_groups() + torch_memory_saver.resume() + clear_memory() if self.args.async_save: from megatron.training.async_utils import maybe_finalize_async_save @@ -533,7 +535,9 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: save_hf_model(self.args, rollout_id, self.model) if self.args.offload_train: + clear_memory(clear_host_memory=True) destroy_process_groups() + torch_memory_saver.pause() @timer def update_weights(self) -> None: From d304ad2d5c1267c2c37133afc29e4c10e32efe47 Mon Sep 17 00:00:00 2001 From: Procrastinatorrrr <18924614955@163.com> Date: Tue, 5 May 2026 22:46:35 +0800 Subject: [PATCH 2/2] refactor: use wake_up()/sleep() for cleaner save_model --- slime/backends/megatron_utils/actor.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index c886509931..16f1748cc5 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -515,9 +515,7 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: # torch dist may trigger nccl communication during saving. if self.args.offload_train: - reload_process_groups() - torch_memory_saver.resume() - clear_memory() + self.wake_up() if self.args.async_save: from megatron.training.async_utils import maybe_finalize_async_save @@ -535,9 +533,7 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: save_hf_model(self.args, rollout_id, self.model) if self.args.offload_train: - clear_memory(clear_host_memory=True) - destroy_process_groups() - torch_memory_saver.pause() + self.sleep() @timer def update_weights(self) -> None: