Fix(checkpoint): add resume/pause in save_model() for offload_train (fixes #1886)

[Procrastinatorrrr]

🐛 Bug Summary

Checkpoint save crashes with --colocate and offload_train=True (regression from #1856, fixes #1886)

Error

Traceback (most recent call last):
  File "/root/slime/train.py", line 107, in <module>
    train(args)
  File "/root/slime/train.py", line 88, in train
    save(rollout_id)
  File "/root/slime/train.py", line 54, in save
    actor_model.save_model(
  File "/root/slime/slime/ray/actor_group.py", line 133, in save_model
    return ray.get([actor.save_model.remote(rollout_id, force_sync=force_sync) for actor in self._actor_handlers])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/ray/_private/client_mode_hook.py", line 107, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/ray/_private/worker.py", line 2980, in get
    values, debugger_breakpoint = worker.get_objects(
                                  ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/ray/_private/worker.py", line 1023, in get_objects
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(AcceleratorError): �[36mray::MegatronTrainRayActor.save_model()�[39m (pid=12545, ip=172.17.0.2, actor_id=7b15b4911ec1f78e875ffaec02000000, repr=<slime.backends.megatron_utils.actor.MegatronTrainRayActor object at 0x7123dfe3d640>)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/slime/slime/utils/timer.py", line 78, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/root/slime/slime/backends/megatron_utils/actor.py", line 525, in save_model
    save(rollout_id, self.model, self.optimizer, self.opt_param_scheduler)
  File "/root/slime/slime/backends/megatron_utils/model.py", line 765, in save
    save_checkpoint(
  File "/root/Megatron-LM/megatron/training/checkpointing.py", line 635, in save_checkpoint
    async_save_request = dist_checkpointing.save(state_dict, checkpoint_name, save_strategy,
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 407, in save
    sharded_state_dict, state_dict = save_preprocess(
                                     ^^^^^^^^^^^^^^^^
  File "/root/Megatron-LM/megatron/core/dist_checkpointing/state_dict_utils.py", line 56, in save_preprocess
    determine_global_metadata(sharded_part)[1],
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/Megatron-LM/megatron/core/dist_checkpointing/validation.py", line 518, in determine_global_metadata
    torch.distributed.all_gather_object(global_metadata, local_metadata)
  File "/root/slime/slime/utils/reloadable_process_group.py", line 67, in new_function
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/c10d_logger.py", line 81, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/distributed_c10d.py", line 3167, in all_gather_object
    input_tensor, local_size = _object_to_tensor(obj, current_device, group)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/distributed_c10d.py", line 3076, in _object_to_tensor
    byte_tensor = torch.ByteTensor(byte_storage).to(device)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.AcceleratorError: CUDA error: invalid argument
Search for `cudaErrorInvalidValue' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.

Root Cause

PR #1856 introduced unconditional self.sleep() at the end of train(), which calls
torch_memory_saver.pause() to release GPU memory. However, save_model() only calls
reload_process_groups() but missing torch_memory_saver.resume(), causing it to
operate on a paused model state.

Call stack when crash occurs:

train.py:85   → actor_model.async_train()
                 ↓
actor.py:371  → self.sleep()              # Model paused here (pause #1)
                 ↓
train.py:87-88 → should_run_periodic_action() → True (save interval reached)
                 ↓
train.py:89   → save(rollout_id)         # Tries to save paused model
                 ↓
actor.py:518  → reload_process_groups()   # ✅ Process groups rebuilt
                 ↓ ❌ MISSING: torch_memory_saver.resume()
actor.py:525  → save(paused_model)        # 💥 CRASH: CUDA error or double-pause

Fix

Complete the save_model() offload lifecycle by adding proper resume/pause:

  def save_model(self, rollout_id, force_sync=False):
      if self.args.debug_rollout_only:
          return
  
-     # torch dist may trigger nccl communication during saving.
      if self.args.offload_train:
+         # [FIX-#1856] Resume model before saving when using offload_train
          reload_process_groups()
+         torch_memory_saver.resume()       # Wake up paused model
+         clear_memory()                    # Ensure sufficient memory available
  
      if self.args.async_save:
          from megatron.training.async_utils import maybe_finalize_async_save
          maybe_finalize_async_save(blocking=True)
  
      save(rollout_id, self.model, self.optimizer, self.opt_param_scheduler)
  
      if force_sync and self.args.async_save:
          maybe_finalize_async_save(blocking=True)
  
      if self.args.save_hf is not None and self.role == "actor":
          from slime.backends.megatron_utils.model import save_hf_model
          save_hf_model(self.args, rollout_id, self.model)
  
-     if self.args.offload_train:
-         destroy_process_groups()
+     # [FIX-#1856] Pause model again after saving to free GPU memory
+     if self.args.offload_train:
+         clear_memory(clear_host_memory=True)
+         destroy_process_groups()
+         torch_memory_saver.pause()        # Release GPU memory again

Verification

Test environment:

• Model: Qwen3.5-4B (32 layers, 2560 hidden, TP=2)

• Hardware: 2× NVIDIA H200 (140GB each)

• Config: --colocate --advantage-estimator grpo --save-interval 1 --offload

• Result: ✅ Checkpoint saved successfully at every rollout without errors

• Log excerpt:

  [36m(MegatronTrainRayActor pid=12281)[0m   [2026-05-04 13:58:44.399834] successfully saved checkpoint from iteration       1 to /data/output/checkpoints/ [ t 1/2, p 1/1 ]
  

Files Changed

• slime/backends/megatron_utils/actor.py: Complete save_model() offload lifecycle (+6 lines)

Impact Analysis

• ✅ Minimal change: Only modifies save_model() method, no changes to train() or offload_train()
• ✅ Preserves all optimizations: Auto-sleep in train() remains active for memory efficiency
• ✅ No side effects: Only affects offload_train=True scenario with periodic saves
• ✅ Backwards compatible: All existing configurations work unchanged

Testing

• Required: Please add run-ci-ckpt label (validates checkpoint save path)
• Optional: run-ci-short for broader coverage

📚 Related Issues

• Fixes: [Question] Checkpoint save fails with --colocate + --save-interval after #1856 — am I missing anything? #1886 (Checkpoint save fails with --colocate + --save-interval after refactor/ppo #1856)
• Regression from: PR refactor/ppo #1856 (75af5297) - "refactor/ppo"
• Affects: All users of --colocate with offload_train=True who need periodic checkpoint saves
• Impact severity: High (training progress lost, cannot resume from checkpoints)
• Workaround: Disable --colocate or remove --save-interval (not practical for long training runs)

[feji3769]

Thanks for the fix! This solution lets me save successfully.

[lilei199908]

Thanks for the fix!

[Procrastinatorrrr]

CI Failure Summary

Both CI failures are pre-existing issues unrelated to this PR's changes.

---

1. run-ci-ckpt (Checkpoint Save/Load Test)

Error: ValueError: Cannot merge two lists with different lengths (81 and 52)
Location: actor.py:84 → init() → load_checkpoint() (checkpoint loading, not saving)

Why Unrelated:
This PR only modifies:

• save_model() — adds torch_memory_saver.resume/pause around save operation
• train() — moves sleep() into critic branch

Neither change affects init(), load_checkpoint(), or optimizer state structure.

The error occurs during Phase 2 (checkpoint loading) of test_qwen3_4B_ckpt.py.

---

2. test_plugin_runtime_hook_contracts::test_runtime_hook_callsite_is_stable[rollout_data_postprocess]

Error: assert 'self.rollout_data_postprocess(self.args)' in <actor.py source> — assertion failed

Root Cause: Signature mismatch between test expectation and actual code:

	Test Expects	Actual Code
Call	self.rollout_data_postprocess(self.args)	self.rollout_data_postprocess(self.args, rollout_id, rollout_data)
Params	(args,)	(args, rollout_id, rollout_data)

Why Unrelated: This PR does not modify rollout_data_postprocess or its call site. The actual code at actor.py:467 already passes 3 parameters, suggesting this hook's signature was changed in a prior commit (possibly PR #1856 or later) without updating the contract test.

---

Request

@maintainer Could you please verify if both tests pass on the current main branch? If they also fail on main, this confirms these are pre-existing regressions.

[lilei199908]

CI Failure Summary

Both CI failures are pre-existing issues unrelated to this PR's changes.

1. run-ci-ckpt (Checkpoint Save/Load Test)

Error: ValueError: Cannot merge two lists with different lengths (81 and 52) Location: actor.py:84 → init() → load_checkpoint() (checkpoint loading, not saving)

Why Unrelated: This PR only modifies:

• save_model() — adds torch_memory_saver.resume/pause around save operation
• train() — moves sleep() into critic branch

Neither change affects init(), load_checkpoint(), or optimizer state structure.

The error occurs during Phase 2 (checkpoint loading) of test_qwen3_4B_ckpt.py.

2. test_plugin_runtime_hook_contracts::test_runtime_hook_callsite_is_stable[rollout_data_postprocess]

Error: assert 'self.rollout_data_postprocess(self.args)' in <actor.py source> — assertion failed

Root Cause: Signature mismatch between test expectation and actual code:

Test Expects Actual Code
Call self.rollout_data_postprocess(self.args) self.rollout_data_postprocess(self.args, rollout_id, rollout_data)
Params (args,) (args, rollout_id, rollout_data)
Why Unrelated: This PR does not modify rollout_data_postprocess or its call site. The actual code at actor.py:467 already passes 3 parameters, suggesting this hook's signature was changed in a prior commit (possibly PR #1856 or later) without updating the contract test.

Request

@maintainer Could you please verify if both tests pass on the current main branch? If they also fail on main, this confirms these are pre-existing regressions.

CI Failure Summary

Both CI failures are pre-existing issues unrelated to this PR's changes.

1. run-ci-ckpt (Checkpoint Save/Load Test)

Error: ValueError: Cannot merge two lists with different lengths (81 and 52) Location: actor.py:84 → init() → load_checkpoint() (checkpoint loading, not saving)

Why Unrelated: This PR only modifies:

• save_model() — adds torch_memory_saver.resume/pause around save operation
• train() — moves sleep() into critic branch

Neither change affects init(), load_checkpoint(), or optimizer state structure.

The error occurs during Phase 2 (checkpoint loading) of test_qwen3_4B_ckpt.py.

2. test_plugin_runtime_hook_contracts::test_runtime_hook_callsite_is_stable[rollout_data_postprocess]

Error: assert 'self.rollout_data_postprocess(self.args)' in <actor.py source> — assertion failed

Root Cause: Signature mismatch between test expectation and actual code:

Test Expects Actual Code
Call self.rollout_data_postprocess(self.args) self.rollout_data_postprocess(self.args, rollout_id, rollout_data)
Params (args,) (args, rollout_id, rollout_data)
Why Unrelated: This PR does not modify rollout_data_postprocess or its call site. The actual code at actor.py:467 already passes 3 parameters, suggesting this hook's signature was changed in a prior commit (possibly PR #1856 or later) without updating the contract test.

Request

@maintainer Could you please verify if both tests pass on the current main branch? If they also fail on main, this confirms these are pre-existing regressions.

yes，these are pre-existing regressions and i'am trying to slove it. thanks