[refactor] remove Registry and change the order of init (THUDM#398)

* [refactor] remove Registry and change the order of init

* bugfix

Author: zhuzilin

slime/backends/megatron_utils/actor.py

@@ -16,7 +16,6 @@
 from megatron.core import mpu
 from transformers import AutoConfig, AutoTokenizer
 
-from slime.ray.registry import get_actors
 from slime.ray.train_actor import TrainRayActor
 from slime.utils.data import process_rollout_data
 from slime.utils.distributed_utils import get_gloo_group, init_process_group
@@ -404,8 +403,7 @@ def update_weights(self):
 
         if not self.connected:
             self.connected = True
-            rollout_engines = get_actors("rollout")
-            rollout_engine_lock = get_actors("rollout_lock", 0)
+            rollout_engines, rollout_engine_lock = ray.get(self.rollout_manager.get_rollout_engines_and_lock.remote())
             self.weight_updator.connect_rollout_engines(rollout_engines, rollout_engine_lock)
             dist.barrier(group=get_gloo_group())
 

slime/ray/actor_group.py

@@ -122,19 +122,24 @@ def async_train(self, rollout_id, rollout_data_ref):
         """Do one rollout training"""
         return [actor.train.remote(rollout_id, rollout_data_ref) for actor in self._actor_handlers]
 
-    def async_save_model(self, step_id):
+    def save_model(self, step_id):
         """Save actor model on rank 0."""
-        return [actor.save_model.remote(step_id) for actor in self._actor_handlers]
+        return ray.get([actor.save_model.remote(step_id) for actor in self._actor_handlers])
 
-    def async_update_weights(self):
+    def update_weights(self):
         """Broadcast weights from rank 0 to all other ranks."""
-        return [actor.update_weights.remote() for actor in self._actor_handlers]
+        return ray.get([actor.update_weights.remote() for actor in self._actor_handlers])
 
-    def async_offload(self):
-        return [actor.sleep.remote(("model")) for actor in self._actor_handlers]
+    def offload(self):
+        return ray.get([actor.sleep.remote(("model")) for actor in self._actor_handlers])
 
-    def async_connect(self, critic_group):
-        return [
-            actor.connect_actor_critic.remote((critic))
-            for actor, critic in zip(self._actor_handlers, critic_group._actor_handlers)
-        ]
+    def connect(self, critic_group):
+        return ray.get(
+            [
+                actor.connect_actor_critic.remote((critic))
+                for actor, critic in zip(self._actor_handlers, critic_group._actor_handlers)
+            ]
+        )
+
+    def set_rollout_manager(self, rollout_manager):
+        return ray.get([actor.set_rollout_manager.remote(rollout_manager) for actor in self._actor_handlers])

slime/ray/placement_group.py

@@ -131,8 +131,58 @@ def create_training_group(args, pg, wandb_run_id):
     return actor_model
 
 
+def create_training_models(args, pgs, wandb_run_id):
+    actor_model = allocate_train_group(
+        args=args,
+        num_nodes=args.actor_num_nodes,
+        num_gpus_per_node=args.actor_num_gpus_per_node,
+        pg=pgs["actor"],
+        wandb_run_id=wandb_run_id,
+    )
+    if args.use_critic:
+        critic_model = allocate_train_group(
+            args=args,
+            num_nodes=args.critic_num_nodes,
+            num_gpus_per_node=args.critic_num_gpus_per_node,
+            pg=pgs["critic"],
+            wandb_run_id=wandb_run_id,
+        )
+        critic_init_handle = critic_model.async_init(args, role="critic", with_ref=False)
+    else:
+        critic_model = None
+
+    start_rollout_ids = ray.get(
+        actor_model.async_init(args, role="actor", with_ref=args.kl_coef != 0 or args.use_kl_loss)
+    )
+
+    assert len(set(start_rollout_ids)) == 1
+    if args.start_rollout_id is None:
+        args.start_rollout_id = start_rollout_ids[0]
+
+    if args.use_critic:
+        ray.get(critic_init_handle)
+        actor_model.connect(critic_model)
+
+    return actor_model, critic_model
+
+
 def create_rollout_manager(args, pg, wandb_run_id):
-    return RolloutManager.options(
+    rollout_manager = RolloutManager.options(
         num_cpus=1,
         num_gpus=0,
     ).remote(args, pg, wandb_run_id=wandb_run_id)
+
+    if args.rollout_global_dataset:
+        ray.get(rollout_manager.load.remote(args.start_rollout_id - 1))
+
+    # TODO: extract this to single function
+    rollout_engines, rollout_engine_lock = ray.get(rollout_manager.get_rollout_engines_and_lock.remote())
+
+    # calculate num_rollout from num_epoch
+    num_rollout_per_epoch = None
+    if args.num_rollout is None:
+        num_rollout_per_epoch = ray.get(rollout_manager.get_num_rollout_per_epoch.remote())
+        args.num_rollout = num_rollout_per_epoch * args.num_epoch
+    assert args.num_rollout > 0
+
+    return rollout_manager, num_rollout_per_epoch

slime/ray/registry.py

@@ -93,3 +93,10 @@ def save_model(self, iteration):
     @abc.abstractmethod
     def update_weights(self):
         raise NotImplementedError
+
+    @abc.abstractmethod
+    def connect_actor_critic(self, critic_group):
+        raise NotImplementedError
+
+    def set_rollout_manager(self, rollout_manager):
+        self.rollout_manager = rollout_manager

slime/ray/train_actor.py

@@ -1,8 +1,7 @@
 import ray
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS
 
-from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_group
-from slime.ray.registry import register_actor
+from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
 from slime.utils.arguments import parse_args
 from slime.utils.wandb_utils import init_wandb_primary
 
@@ -12,54 +11,19 @@ def train(args):
     pgs = create_placement_groups(args)
     wandb_run_id = init_wandb_primary(args)
 
-    actor_model = create_training_group(args, pgs["actor"], wandb_run_id=wandb_run_id)
-    if args.use_critic:
-        critic_model = create_training_group(args, pgs["critic"], wandb_run_id=wandb_run_id)
+    # create the actor and critic models
+    actor_model, critic_model = create_training_models(args, pgs, wandb_run_id=wandb_run_id)
 
     # create the rollout manager, with sglang engines inside.
-    rollout_manager = create_rollout_manager(args, pgs["rollout"], wandb_run_id=wandb_run_id)
-
-    # TODO: extract this to single function
-    rollout_engines, rollout_engine_lock = ray.get(rollout_manager.get_rollout_engines_and_lock.remote())
-    for i, rollout_engine in enumerate(rollout_engines):
-        register_actor("rollout", i, rollout_engine)
-    register_actor("rollout_lock", 0, rollout_engine_lock)
-    for i, actor in enumerate(actor_model._actor_handlers):
-        register_actor("actor", i, actor)
-    if args.use_critic:
-        for i, critic in enumerate(critic_model._actor_handlers):
-            register_actor("critic", i, critic)
-
-    # calculate num_rollout from num_epoch
-    num_rollout_per_epoch = None
-    if args.num_rollout is None:
-        num_rollout_per_epoch = ray.get(rollout_manager.get_num_rollout_per_epoch.remote())
-        args.num_rollout = num_rollout_per_epoch * args.num_epoch
-    assert args.num_rollout > 0
-
-    # sync the initialization (model initalization, load checkpoint, etc.)
-    if args.use_critic:
-        critic_init_handle = critic_model.async_init(args, role="critic", with_ref=False)
-
-    start_rollout_ids = ray.get(
-        actor_model.async_init(args, role="actor", with_ref=args.kl_coef != 0 or args.use_kl_loss)
-    )
-    assert len(set(start_rollout_ids)) == 1
-    if args.start_rollout_id is None:
-        args.start_rollout_id = start_rollout_ids[0]
-
-    if args.rollout_global_dataset:
-        ray.get(rollout_manager.load.remote(args.start_rollout_id - 1))
-
-    if args.use_critic:
-        ray.get(critic_init_handle)
-        ray.get(actor_model.async_connect(critic_model))
+    rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"], wandb_run_id=wandb_run_id)
+
+    actor_model.set_rollout_manager(rollout_manager)
 
     if args.offload:
         ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]))
 
     # always update weight first so that sglang has the loaded weights from training.
-    ray.get(actor_model.async_update_weights())
+    actor_model.update_weights()
 
     if args.offload:
         ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE]))
@@ -88,21 +52,23 @@ def train(args):
             (rollout_id + 1) % args.save_interval == 0
             or (num_rollout_per_epoch is not None and (rollout_id + 1) % num_rollout_per_epoch == 0)
         ):
-            ray.get(actor_model.async_save_model(rollout_id))
+            actor_model.save_model(rollout_id)
+            if args.use_critic:
+                critic_model.save_model(rollout_id)
             if args.rollout_global_dataset:
                 ray.get(rollout_manager.save.remote(rollout_id))
 
         if args.offload:
             if args.use_critic:
-                ray.get(critic_model.async_offload())
+                critic_model.offload()
                 if rollout_id >= args.num_critic_only_steps:
-                    ray.get(actor_model.async_offload())
+                    actor_model.offload()
             else:
-                ray.get(actor_model.async_offload())
+                actor_model.offload()
 
             ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]))
 
-        ray.get(actor_model.async_update_weights())
+        actor_model.update_weights()
 
         if args.offload:
             ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE]))

train.py

@@ -1,7 +1,6 @@
 import ray
 
-from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_group
-from slime.ray.registry import register_actor
+from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_train_models
 from slime.utils.arguments import parse_args
 from slime.utils.wandb_utils import init_wandb_primary
 
@@ -12,56 +11,16 @@ def train(args):
     pgs = create_placement_groups(args)
     wandb_run_id = init_wandb_primary(args)
 
-    actor_model = create_training_group(args, pgs["actor"], wandb_run_id=wandb_run_id)
-    if args.use_critic:
-        critic_model = create_training_group(args, pgs["critic"], wandb_run_id=wandb_run_id)
+    # create the actor and critic models
+    actor_model, critic_model = create_train_models(args, pgs, wandb_run_id=wandb_run_id)
 
     # create the rollout manager, with sglang engines inside.
-    rollout_manager = create_rollout_manager(args, pgs["rollout"], wandb_run_id=wandb_run_id)
+    rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"], wandb_run_id=wandb_run_id)
 
-    # TODO: extract this to single function
-    rollout_engines, rollout_engine_lock = ray.get(rollout_manager.get_rollout_engines_and_lock.remote())
-    for i, rollout_engine in enumerate(rollout_engines):
-        register_actor("rollout", i, rollout_engine)
-    register_actor("rollout_lock", 0, rollout_engine_lock)
-    for i, actor in enumerate(actor_model._actor_handlers):
-        register_actor("actor", i, actor)
-    if args.use_critic:
-        for i, critic in enumerate(critic_model._actor_handlers):
-            register_actor("critic", i, critic)
-
-    # calculate num_rollout from num_epoch
-    num_rollout_per_epoch = None
-    if args.num_rollout is None:
-        num_rollout_per_epoch = ray.get(rollout_manager.get_num_rollout_per_epoch.remote())
-        args.num_rollout = num_rollout_per_epoch * args.num_epoch
-    assert args.num_rollout > 0
-
-    # sync the initialization (model initalization, load checkpoint, etc.)
-    if args.use_critic:
-        critic_init_handle = critic_model.async_init(args, role="critic", with_ref=False)
-
-    start_rollout_ids = ray.get(
-        actor_model.async_init(args, role="actor", with_ref=args.kl_coef != 0 or args.use_kl_loss)
-    )
-
-    assert len(set(start_rollout_ids)) == 1
-    if args.start_rollout_id is None:
-        args.start_rollout_id = start_rollout_ids[0]
-
-    if args.rollout_global_dataset:
-        ray.get(rollout_manager.load.remote(args.start_rollout_id - 1))
-
-    if args.use_critic:
-        ray.get(critic_init_handle)
-        ray.get(actor_model.async_connect(critic_model))
-
-    if args.use_critic:
-        ray.get(critic_init_handle)
-        ray.get(actor_model.async_connect(critic_model))
+    actor_model.set_rollout_manager(rollout_manager)
 
     # always update weight first so that sglang has the loaded weights from training.
-    ray.get(actor_model.async_update_weights())
+    actor_model.update_weights()
 
     # async train loop.
     rollout_data_next_future = rollout_manager.generate.remote(args.start_rollout_id)
@@ -86,15 +45,17 @@ def train(args):
             (rollout_id + 1) % args.save_interval == 0
             or (num_rollout_per_epoch is not None and (rollout_id + 1) % num_rollout_per_epoch == 0)
         ):
-            ray.get(actor_model.async_save_model(rollout_id))
+            actor_model.save_model(rollout_id)
+            if args.use_critic:
+                critic_model.save_model(rollout_id)
             if args.rollout_global_dataset:
                 ray.get(rollout_manager.save.remote(rollout_id))
 
         if (rollout_id + 1) % args.update_weights_interval == 0:
             # sync generate before update weights to prevent update weight in the middle of generation
             rollout_data_curr_ref = ray.get(rollout_data_next_future)
             rollout_data_next_future = None
-            ray.get(actor_model.async_update_weights())
+            actor_model.update_weights()
 
         if args.eval_interval is not None and (
             (rollout_id + 1) % args.eval_interval == 0