Add clear_num_new_engines and some code cleanup (THUDM#1349)

Author: zhuzilin

slime/backends/fsdp_utils/actor.py

@@ -818,6 +818,8 @@ def update_weights(self) -> None:  # type: ignore[override]
         if num_new_engines > 0:
             self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock)
             dist.barrier(group=get_gloo_group())
+            if dist.get_rank() == 0:
+                ray.get(self.rollout_manager.clear_num_new_engines.remote())
 
         self.weight_updater.update_weights()
 

slime/backends/megatron_utils/actor.py

@@ -524,6 +524,8 @@ def update_weights(self) -> None:
         if num_new_engines > 0:
             self.weight_updater.connect_rollout_engines(rollout_engines, rollout_engine_lock)
             dist.barrier(group=get_gloo_group())
+            if dist.get_rank() == 0:
+                ray.get(self.rollout_manager.clear_num_new_engines.remote())
 
         with torch_memory_saver.disable() if self.args.offload_train else nullcontext():
             print_memory("before update_weights")

slime/ray/rollout.py

@@ -10,7 +10,7 @@
 import ray
 import torch
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
-from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS
+from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS
 
 from slime.backends.sglang_utils.sglang_engine import SGLangEngine
 from slime.rollout.base_types import call_rollout_fn
@@ -170,6 +170,12 @@ def onload(self, tags: list[str] | None = None):
             ]
         )
 
+    def onload_weights(self):
+        self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS])
+
+    def onload_kv(self):
+        self.onload(tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH])
+
     def recover_rollout_engines(self):
         """Restart any dead rollout engines and update num_new_engines for update_weights detection."""
         self.health_monitoring_pause()
@@ -187,6 +193,10 @@ def recover_rollout_engines(self):
 
         return self.rollout_engines, self.rollout_engine_lock, self.num_new_engines
 
+    def clear_num_new_engines(self):
+        # when fault tolerance is not enabled, we need to manually clear num_new_engines after update_weights
+        self.num_new_engines = 0
+
     def health_monitoring_pause(self) -> None:
         if self._health_monitor is not None:
             self._health_monitor.pause()

train.py

@@ -1,10 +1,4 @@
 import ray
-from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS
-
-try:
-    from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
-except ImportError:
-    GPU_MEMORY_TYPE_CUDA_GRAPH = None
 
 from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
 from slime.utils.arguments import parse_args
@@ -27,7 +21,7 @@ def train(args):
     actor_model, critic_model = create_training_models(args, pgs, rollout_manager)
 
     if args.offload_rollout:
-        ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]))
+        ray.get(rollout_manager.onload_weights.remote())
 
     # always update weight first so that sglang has the loaded weights from training.
     actor_model.update_weights()
@@ -36,9 +30,7 @@ def train(args):
         ray.get(rollout_manager.check_weights.remote(action="compare"))
 
     if args.offload_rollout:
-        if GPU_MEMORY_TYPE_CUDA_GRAPH is not None:
-            ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH]))
-        ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE]))
+        ray.get(rollout_manager.onload_kv.remote())
 
     # special case for eval-only
     if args.num_rollout == 0 and args.eval_interval is not None:
@@ -55,9 +47,19 @@ def offload_train():
         else:
             actor_model.clear_memory()
 
-    def onload_rollout():
-        if args.offload_rollout:
-            ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]))
+    def save(rollout_id):
+        if (not args.use_critic) or (rollout_id >= args.num_critic_only_steps):
+            actor_model.save_model(
+                rollout_id,
+                force_sync=rollout_id == args.num_rollout - 1,
+            )
+        if args.use_critic:
+            critic_model.save_model(
+                rollout_id,
+                force_sync=rollout_id == args.num_rollout - 1,
+            )
+        if args.rollout_global_dataset:
+            ray.get(rollout_manager.save.remote(rollout_id))
 
     # train loop.
     # note that for async training, one can change the position of the sync operation(ray.get).
@@ -79,27 +81,14 @@ def onload_rollout():
             ray.get(actor_model.async_train(rollout_id, rollout_data_ref))
 
         if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout):
-            if (not args.use_critic) or (rollout_id >= args.num_critic_only_steps):
-                actor_model.save_model(
-                    rollout_id,
-                    force_sync=rollout_id == args.num_rollout - 1,
-                )
-            if args.use_critic:
-                critic_model.save_model(
-                    rollout_id,
-                    force_sync=rollout_id == args.num_rollout - 1,
-                )
-            if args.rollout_global_dataset:
-                ray.get(rollout_manager.save.remote(rollout_id))
+            save(rollout_id)
 
         offload_train()
-        onload_rollout()
+        if args.offload_rollout:
+            ray.get(rollout_manager.onload_weights.remote())
         actor_model.update_weights()
-
         if args.offload_rollout:
-            if GPU_MEMORY_TYPE_CUDA_GRAPH is not None:
-                ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH]))
-            ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE]))
+            ray.get(rollout_manager.onload_kv.remote())
 
         if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch):
             ray.get(rollout_manager.eval.remote(rollout_id))