[Bugfix] stabilize actor offload wake/sleep flow (follow-up to #48)

[CalvinXKY]

Summary

Follow-up to #48. This PR keeps the actor-side stability fix on top of #48's IPC contract/coordinator/gather changes.

Scope in this PR:

• slime/backends/megatron_utils/actor.py
  • clear routing replay before offload sleep when enabled
  • only destroy process groups in sleep path for non-colocate critic reconnect case
  • restore actor tag after wake-up (self._switch_model("actor"))

Related IPC correctness is handled in #48 (coordinator = slot-start rank, slot-wide gloo gather, single-RPC version-with-data).

Test case: tests/test_qwen3_30B_A3B.py on A800, 8× colocate, Megatron TP=4, rollout-num-gpus-per-engine=8.

Failure mode

We hit a chain of failures while bringing up colocated IPC weight sync for Qwen3-30B-A3B. Below are the observed symptoms, NFS log paths, and key excerpts.

0) Checkpoint load appeared stuck (NFS mount)

Symptom: training looked hung at Megatron load:

loading release distributed checkpoint from /root/Qwen3-30B-A3B_torch_dist

Cause: /root/Qwen3-30B-A3B_torch_dist was a symlink to NFS (/data/nfs_87/xky/models/...), so ~57G torch_dist was read over the mount.

Mitigation (operational, not in this PR): rsync checkpoint to container local disk (/root/local_models/Qwen3-30B-A3B_torch_dist) before training.

Log: /data/nfs_87/xky/logs/test_qwen3_30B_A3B_20260527_074737.log

---

1) Duplicate start_weight_update in colocate IPC path

Symptom:

RuntimeError: start_weight_update called while a weight update is already active. Call finish_weight_update first.

Root cause: for one 8-GPU vLLM engine with Megatron TP=4, coordinator gating used tp_rank == 0. Both global rank 0 and rank 4 could enter the coordinator path and call start_weight_update twice.

Fix: #48 - gate coordinator on rank == slot_start (lowest global rank in the engine GPU slot).

Log: /data/nfs_87/xky/logs/test_qwen3_30B_A3B_20260527_025345.log (around line 1823)

(Worker_TP1) RuntimeError: start_weight_update called while a weight update is already active. Call finish_weight_update first.

---

2) Missing IPC handles during tensor transfer

Symptom:

ValueError: IPC handle not found for GPU UUID d2c381a3-522a-0b4f-8e39-20a919fc253f. Available UUIDs: [...]

Root cause: IPC payloads were gathered only within the Megatron TP subgroup (4 ranks), but the 8-GPU vLLM engine needs handles from the full engine slot (8 ranks).

Fix: #48 - build per-slot gloo group and all_gather_object across the full slot, not TP subgroup.

Log: /data/nfs_87/xky/logs/test_qwen3_30B_A3B_20260527_031623.log (line 1809)

---

3) torch_memory_saver crash after first successful weight sync

Symptom: Ray actor workers died right after update_weights reached 100%:

[torch_memory_saver.cpp] CUresult error: 1 (invalid argument)  file=csrc/core.cpp func=free line=81
ray.exceptions.ActorDiedError: The actor died unexpectedly before finishing this task.

Context: first IPC sync completed (Update weights: 100%|...| 115/115, finish_weight_update 200 OK), then crash during offload/sleep transition after rollout 0 train.

Fix: this PR (actor.py) - stabilize sleep/wake: clear routing replay, avoid unconditional PG destroy on colocate sleep, restore actor model tag after wake-up.

Log: /data/nfs_87/xky/logs/test_qwen3_30B_A3B_20260527_080824.log (lines 2823, 2853)

---

4) Ray host-memory kill during second update_weights

Symptom:

ray.exceptions.OutOfMemoryError: 1 worker(s) were killed due to the node running low on memory.
Memory ... was 952.93GB / 1003.08GB (0.950004), which exceeds the memory usage threshold of 0.950000
Top memory users: ray::MegatronTrainRayActor.update_weights (~99GB x 8 ranks)

Context: not GPU OOM - Ray killed workers when host RAM crossed the default 95% threshold during the second post-train weight sync.

Mitigation (operational): raise RAY_memory_usage_threshold=0.99 when starting isolated Ray in the run script.

Log: /data/nfs_87/xky/logs/test_qwen3_30B_A3B_20260527_092114.log (line 2700)

---

Fix (this PR)

slime/backends/megatron_utils/actor.py

• clear RoutingReplay before offload sleep when --use-routing-replay
• only destroy_process_groups() in sleep when non-colocate critic reconnect path applies
• after wake_up(), call self._switch_model("actor") so the active model tag is restored

Why actor.py change is needed (with #48)

#48 fixes IPC coordinator selection and slot-wide handle gather. Even after IPC sync succeeds, colocate + offload_train still needs a stable actor sleep/wake path. Without these actor adjustments we saw intermittent torch_memory_saver / ActorDiedError crashes after successful weight updates.

Test plan

• Reproduced failures (1)-(4) above before fixes
• pre-commit run --all-files
• python -m pytest tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py -q - 7 passed
• End-to-end on A800 with local checkpoint copy

Successful run after fixes:

Log: /data/nfs_87/xky/logs/test_qwen3_30B_A3B_20260527_094318.log

[2026-05-27 10:01:27] ... perf 2: {...}
[2026-05-27 10:01:53] ... rollout 2: {...}
[2026-05-27 10:02:52] ... step 2: {...}
Update weights: 100%|██████████| 115/115
2026-05-27 10:03:48 ... Job 'raysubmit_RDDgpg3WPpdD7nCX' succeeded

• Run full colocate CI matrix on clean image

[gemini-code-assist]

Code Review

This pull request updates the weight synchronization logic to use a dedicated Gloo process group (_ipc_slot_group) covering the entire engine GPU slot, rather than relying on the Megatron tensor parallel group. It also introduces corresponding unit tests and minor lifecycle fixes in the actor backend. However, a critical issue was identified where dist.new_group is called within a rank-conditional block; because it is a collective operation, this will lead to a distributed deadlock. All ranks must invoke dist.new_group collectively.

[gemini-code-assist]

In PyTorch distributed, dist.new_group is a collective call that must be invoked by all ranks in the default process group (the entire world), even if they are not going to be members of the new group.

Calling dist.new_group inside the rank-conditional block if start <= rank < end: means only a subset of ranks will execute it, which will cause a distributed hang/deadlock during initialization when there are multiple colocated engines or when some ranks do not belong to the current engine's GPU slot.

To fix this, you should create the Gloo group collectively for all ranks outside of the conditional block:

        for i, engine in enumerate(self._colocated_engines):
            start = colocate_gpu_offsets[i]
            end = start + colocate_gpu_counts[i]
            slot_ranks = list(range(start, end))
            # Every rank in the default group must call new_group collectively
            slot_group = dist.new_group(ranks=slot_ranks, backend="gloo")
            
            rank = dist.get_rank()
            if start <= rank < end:
                self._ipc_engine = engine
                self._ipc_engine_slot_start = start
                self._ipc_engine_slot_end = end
                self._ipc_slot_group = slot_group
                # First global rank in the engine GPU slot (not Megatron TP rank 0).
                if rank == start:
                    self._ipc_engine_coordinator = True

[CalvinXKY]

test_qwen3_30B_A3B result:

[knlnguyen1802]

LGTM, thanks for the fix

[aoshen02]

I don't see slime& miles do these modification, wonder if the error is just because of experiment setting problem.

[CalvinXKY]

I don't see slime& miles do these modification, wonder if the error is just because of experiment setting problem.

Sounds reasonable. I'll reproduce it.

[CalvinXKY]

Closure Note
Tested PR #53 fixes on latest main with Qwen3-4B colocate PPO training. Both unpatched and patched runs completed successfully with nearly identical train_rollout_logprob_abs_diff (~0.013–0.018), confirming the fixes do not affect basic colocate PPO correctness.
However, the three fixes are defensive improvements for specific configurations:

1. RoutingReplay.clear_all() in sleep: Only triggers with use_routing_replay=True (not tested here). Prevents stale routing cache.
2. Move destroy_process_groups() inside if block: In colocate mode, the if condition is always False, so this was never called. The fix matters for non-colocate offload mode.
3. _switch_model("actor") in wake_up: Safety guard to ensure actor model is active after wake_up. May be implicitly handled in colocate via weight sync.
   Verdict: The fixes are valid defensive improvements but do not fix observable bugs in the standard colocate PPO workflow. Closing as low-priority / nice-to-have.