[refactor] Add actor registry (THUDM#359)

Author: zhuzilin

slime/backends/fsdp_utils/actor.py

@@ -9,6 +9,7 @@
 from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor, AutoTokenizer
 
 import wandb
+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
@@ -95,6 +96,7 @@ def init(self, args, role, wandb_run_id, with_ref: bool = False):  # type: ignor
         self.update_cpu_params_dict(self.weights["actor"])
 
         self.weight_updator = UpdateWeightFromTensor(self.args, self.model)
+        self.connected = False
 
         if self.args.offload:
             self.sleep(("model"))
@@ -122,15 +124,6 @@ def save_model(self, iteration):
 
         raise NotImplementedError()
 
-    def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
-        self.rollout_engines = rollout_engines
-
-        if self.args.debug_train_only or self.args.debug_rollout_only:
-            return
-
-        self.weight_updator.connect_rollout_engines(rollout_engines, rollout_engine_lock)
-        dist.barrier(group=get_gloo_group())
-
     def compute_log_prob(
         self,
         model_tag,
@@ -392,6 +385,13 @@ def update_weights(self):  # type: ignore[override]
         if self.args.debug_train_only or self.args.debug_rollout_only:
             return
 
+        if not self.connected:
+            self.connected = True
+            rollout_engines = get_actors("rollout")
+            rollout_engine_lock = get_actors("rollout_lock", 0)
+            self.weight_updator.connect_rollout_engines(rollout_engines, rollout_engine_lock)
+            dist.barrier(group=get_gloo_group())
+
         if self.args.offload:
             # TODO: don't wake up here
             self.wake_up(("model"))

slime/backends/megatron_utils/actor.py

@@ -14,6 +14,7 @@
 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
@@ -88,6 +89,7 @@ def init(self, args, role, wandb_run_id, with_ref=False):
             quantization_config=getattr(self.hf_config, "quantization_config", None),
             vocab_size=self.tokenizer.vocab_size if self.args.vocab_size is None else self.args.vocab_size,
         )
+        self.connected = False
 
         # empty cache after initialization
         clear_memory()
@@ -384,15 +386,6 @@ def save_model(self, iteration):
 
         save(iteration, self.model, self.optimizer, self.opt_param_scheduler)
 
-    def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
-        self.rollout_engines = rollout_engines
-
-        if self.args.debug_train_only or self.args.debug_rollout_only:
-            return
-
-        self.weight_updator.connect_rollout_engines(rollout_engines, rollout_engine_lock)
-        dist.barrier(group=get_gloo_group())
-
     @timer
     def update_weights(self):
         if self.args.debug_train_only or self.args.debug_rollout_only:
@@ -401,6 +394,13 @@ def update_weights(self):
         if self.args.offload and hasattr(mpu, "reload_process_groups"):
             mpu.reload_process_groups()
 
+        if not self.connected:
+            self.connected = True
+            rollout_engines = get_actors("rollout")
+            rollout_engine_lock = get_actors("rollout_lock", 0)
+            self.weight_updator.connect_rollout_engines(rollout_engines, rollout_engine_lock)
+            dist.barrier(group=get_gloo_group())
+
         with torch_memory_saver.disable() if self.args.offload and not torch.version.hip else nullcontext():
             print_memory("before update_weights")
             self.weight_updator.update_weights()

slime/backends/xtuner_utils/actor.py

@@ -10,6 +10,7 @@
 from xtuner.v1.model import get_model_config_from_hf
 
 import wandb
+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
@@ -67,6 +68,7 @@ def init(self, args, role, wandb_run_id, with_ref: bool = False):
         self.sp_mesh = self.data_mesh["sp"]
 
         self.weight_updator = UpdateWeightFromDistributed(args, self.model)
+        self.connected = False
 
     def sleep(self, tags):
         if not getattr(self.args, "offload", False):
@@ -87,15 +89,6 @@ def save_model(self, iteration):
         path = f"{self.args.save}/iter_{iteration:07}/hf"
         self.model.save_hf(path, save_dtype=torch.bfloat16)
 
-    def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
-        self.rollout_engines = rollout_engines
-
-        if self.args.debug_train_only or self.args.debug_rollout_only:
-            return
-
-        self.weight_updator.connect_rollout_engines(rollout_engines, rollout_engine_lock)
-        dist.barrier(group=get_gloo_group())
-
     def get_rollout_data(self, rollout_data_ref):
         dp_rank = dist.get_rank() // self.args.sp_size
         dp_size = dist.get_world_size() // self.args.sp_size
@@ -267,6 +260,13 @@ def update_weights(self):  # type: ignore[override]
         if self.args.debug_train_only or self.args.debug_rollout_only:
             return
 
+        if not self.connected:
+            self.connected = True
+            rollout_engines = get_actors("rollout")
+            rollout_engine_lock = get_actors("rollout_lock", 0)
+            self.weight_updator.connect_rollout_engines(rollout_engines, rollout_engine_lock)
+            dist.barrier(group=get_gloo_group())
+
         if self.args.offload:
             # TODO: don't wake up here
             self.wake_up(("model"))

slime/ray/actor_group.py

@@ -37,11 +37,13 @@ def __init__(
         pg: tuple[PlacementGroup, list[int]],
         wandb_run_id: Optional[str] = None,
         num_gpus_per_actor: float = 1,
+        role: str = "actor",
     ) -> None:
         self.args = args
         self._num_nodes = num_nodes
         self._num_gpus_per_node = num_gpus_per_node
         self._wandb_run_id = wandb_run_id
+        self.role = role
 
         # Allocate the GPUs for actors w/o instantiating them
         self._allocate_gpus_for_actor(pg, num_gpus_per_actor, wandb_run_id=wandb_run_id)
@@ -113,18 +115,6 @@ def async_init(self, args, role, with_ref=False):
         self.args = args
         return [actor.init.remote(args, role, self._wandb_run_id, with_ref=with_ref) for actor in self._actor_handlers]
 
-    def async_init_weight_update_connections(self, rollout):
-        """
-        Connect rollout engines and actors, e.g. initialize the process group between them
-        to update weights after each training stage.
-        """
-        self.rollout = rollout
-        rollout_engines, rollout_engine_lock = ray.get(rollout.get_rollout_engines_and_lock.remote())
-        return [
-            actor.connect_rollout_engines.remote(rollout_engines, rollout_engine_lock)
-            for actor in self._actor_handlers
-        ]
-
     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]

slime/ray/registry.py

@@ -0,0 +1,36 @@
+import ray
+
+
+@ray.remote
+class Registry:
+    def __init__(self):
+        self.actors = {}
+
+    def set(self, role, key, actor):
+        if role not in self.actors:
+            self.actors[role] = {}
+        self.actors[role][key] = actor
+
+    def get(self, role: str, key=None):
+        actors = self.actors[role]
+        if key is None:
+            return list(actors.values())
+        return actors[key]
+
+
+REGISTRY = None
+
+
+def register_actor(role, key, actor):
+    try:
+        registry = ray.get_actor("slime_actor_registry")
+    except ValueError:
+        global REGISTRY
+        REGISTRY = Registry.options(name="slime_actor_registry").remote()
+        registry = REGISTRY
+    registry.set.remote(role, key, actor)
+
+
+def get_actors(role, key=None):
+    registry = ray.get_actor("slime_actor_registry")
+    return ray.get(registry.get.remote(role, key))

slime/ray/rollout.py

@@ -230,6 +230,7 @@ def _create_rollout_engines(args, pg):
                     "env_vars": {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST}
                     | {
                         "SGL_JIT_DEEPGEMM_PRECOMPILE": "false",
+                        "SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true",
                     }
                 },
             ).remote(args, rank=i)

slime/ray/train_actor.py

@@ -82,10 +82,6 @@ def sleep(self, tags):
     def wake_up(self, tags):
         raise NotImplementedError
 
-    @abc.abstractmethod
-    def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
-        raise NotImplementedError
-
     @abc.abstractmethod
     def train(self, rollout_id, rollout_data_ref):
         raise NotImplementedError

slime/ray/utils.py

@@ -3,6 +3,8 @@
 
 import ray
 import torch
+from slime.ray.ray_actor import RayActor
+
 
 # Refer to
 # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/nvidia_gpu.py#L95-L96
@@ -34,7 +36,7 @@ def get_physical_gpu_id():
 
 
 @ray.remote
-class Lock:
+class Lock(RayActor):
     def __init__(self):
         self._locked = False  # False: unlocked, True: locked
 

slime/utils/data.py

@@ -73,7 +73,7 @@ def __init__(
             # TODO: this is slow.
             if max_length is not None:
                 if not multimodal_keys:
-                    if len(tokenizer(data[prompt_key])["input_ids"]) > max_length:
+                    if len(prompt) > max_length:
                         continue
 
             self.origin_samples.append(

train.py

@@ -2,6 +2,7 @@
 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.utils.arguments import parse_args
 from slime.utils.wandb_utils import init_wandb_primary
 
@@ -18,6 +19,17 @@ def train(args):
     # 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:
@@ -32,17 +44,13 @@ def train(args):
     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))
 
-    # initialize the connection for weight update during training
-    ray.get(actor_model.async_init_weight_update_connections(rollout_manager))
-
     if args.use_critic:
         ray.get(critic_init_handle)
         ray.get(actor_model.async_connect(critic_model))