[feat] support fault tolerant for rollout engines (THUDM#405)

* [feat] support fault tolerant for rollout engines

* support fault tolerant for UpdateWeightFromDistributed

* bugfix

* bugfix

Author: zhuzilin

slime/backends/fsdp_utils/actor.py

@@ -99,7 +99,6 @@ def init(self, args, role, wandb_run_id, with_ref: bool = False):  # type: ignor
 
         self.update_cpu_params_dict(self.weights["actor"])
 
-        self.connected = False
         self.weight_updator = (
             UpdateWeightFromTensor(self.args, self.model)
             if self.args.colocate
@@ -405,9 +404,10 @@ 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, rollout_engine_lock = ray.get(self.rollout_manager.get_rollout_engines_and_lock.remote())
+        rollout_engines, rollout_engine_lock, num_new_engines = ray.get(
+            self.rollout_manager.get_rollout_engines_and_lock.remote()
+        )
+        if num_new_engines > 0:
             self.weight_updator.connect_rollout_engines(rollout_engines, rollout_engine_lock)
             dist.barrier(group=get_gloo_group())
 

slime/backends/megatron_utils/actor.py

@@ -401,9 +401,10 @@ 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, rollout_engine_lock = ray.get(self.rollout_manager.get_rollout_engines_and_lock.remote())
+        rollout_engines, rollout_engine_lock, num_new_engines = ray.get(
+            self.rollout_manager.get_rollout_engines_and_lock.remote()
+        )
+        if num_new_engines > 0:
             self.weight_updator.connect_rollout_engines(rollout_engines, rollout_engine_lock)
             dist.barrier(group=get_gloo_group())
 

slime/backends/megatron_utils/update_weight_utils.py

@@ -305,6 +305,17 @@ def __init__(self, args, model, weights, *, model_name, quantization_config, voc
         self.param_info_buckets = get_param_info_buckets(self.args, self.model)
         self.weight_version = 0
 
+        # create the group within megatron.
+        for start_rank in range(0, dist.get_world_size(), self.args.rollout_num_gpus_per_engine):
+            end_rank = start_rank + self.args.rollout_num_gpus_per_engine
+            group_ranks = list(range(start_rank, end_rank))
+            new_group = dist.new_group(ranks=group_ranks, backend="gloo")
+            if dist.get_rank() in group_ranks:
+                self._ipc_gather_group = new_group
+                self._ipc_gather_src = start_rank
+
+        self._model_update_groups = None
+
     def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
         self.rollout_engines = rollout_engines
         colocate_engine_nums = (
@@ -322,6 +333,11 @@ def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
             )
             self._group_name = "slime"
             if self._is_distributed_src_rank:
+                if self._model_update_groups is not None:
+                    disconnect_rollout_engines_from_distributed(
+                        self.args, self._group_name, self._model_update_groups, self.distributed_rollout_engines
+                    )
+
                 self._model_update_groups = connect_rollout_engines_from_distributed(
                     self.args, self._group_name, self.distributed_rollout_engines
                 )
@@ -331,13 +347,7 @@ def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
             start_rank = i * self.args.rollout_num_gpus_per_engine
             end_rank = (i + 1) * self.args.rollout_num_gpus_per_engine
             group_ranks = list(range(start_rank, end_rank))
-            new_group = dist.new_group(
-                ranks=group_ranks,
-                backend="gloo",
-            )
             if dist.get_rank() in group_ranks:
-                self._ipc_gather_src = start_rank
-                self._ipc_gather_group = new_group
                 self._ipc_engine = engine
 
     @torch.no_grad()
@@ -496,6 +506,7 @@ def __init__(self, args, model, weights, *, model_name, quantization_config, voc
         self.vocab_size = vocab_size
         self.quantization_config = quantization_config
         self.weight_version = 0
+        self._model_update_groups = None
 
     def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
         self.rollout_engines = rollout_engines
@@ -512,6 +523,10 @@ def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
             self._group_name = f"slime-pp_{pp_rank}"
 
         if self._is_pp_src_rank:
+            if self._model_update_groups is not None:
+                disconnect_rollout_engines_from_distributed(
+                    self.args, self._group_name, self._model_update_groups, self.rollout_engines
+                )
             self._model_update_groups = connect_rollout_engines_from_distributed(
                 self.args, self._group_name, rollout_engines
             )
@@ -670,6 +685,12 @@ def connect_rollout_engines_from_distributed(args, group_name, rollout_engines):
     return model_update_groups
 
 
+def disconnect_rollout_engines_from_distributed(args, group_name, model_update_groups, rollout_engines):
+    refs = [engine.destroy_weights_update_group.remote(group_name) for engine in rollout_engines]
+    dist.destroy_process_group(model_update_groups)
+    ray.get(refs)
+
+
 def update_weights_from_distributed(args, group_name, group, weight_version, rollout_engines, converted_named_tensors):
     refs = [
         engine.update_weights_from_distributed.remote(

slime/backends/sglang_utils/sglang_engine.py

@@ -149,6 +149,28 @@ def _make_request(self, endpoint: str, payload: Optional[dict] = None):
         response.raise_for_status()
         return response.json()
 
+    def health_generate(self, timeout: float = 5.0) -> bool:
+        """Run /health_generate on the underlying SGLang HTTP server.
+
+        Args:
+            timeout: Timeout for the health request in seconds.
+
+        Returns:
+            True if the server responds with HTTP 200.
+
+        Raises:
+            requests.RequestException: If the request fails for any reason, including timeout.
+        """
+        if self.node_rank != 0:
+            return True
+
+        response = requests.get(
+            f"http://{self.server_args.host}:{self.server_args.port}/health_generate",
+            timeout=timeout,
+        )
+        response.raise_for_status()
+        return True
+
     def update_weights_from_tensor(
         self,
         serialized_named_tensors: List[str],
@@ -179,7 +201,7 @@ def flush_cache(self):
         if self.node_rank != 0:
             return
         # flush cache will not return status_code 200 when there are pending requests
-        while True:
+        for _ in range(60):
             try:
                 response = requests.get(f"http://{self.server_args.host}:{self.server_args.port}/flush_cache")
                 if response.status_code == 200:
@@ -188,7 +210,10 @@ def flush_cache(self):
                 raise e
             except Exception as e:
                 print(f"Error flushing cache: {e}")
+                time.sleep(1)
                 continue
+        else:
+            raise TimeoutError("Timeout while flushing cache.")
 
     def shutdown(self):
         requests.post(
@@ -230,6 +255,18 @@ def init_weights_update_group(self, master_address, master_port, rank_offset, wo
             },
         )
 
+    def destroy_weights_update_group(self, group_name):
+        try:
+            return self._make_request(
+                "destroy_weights_update_group",
+                {
+                    "group_name": group_name,
+                },
+            )
+        except:
+            # catch the case there the engine is just created and does not have the group.
+            pass
+
     def update_weights_from_distributed(
         self, names, dtypes, shapes, group_name, flush_cache=False, weight_version: Optional[str] = None
     ):

slime/backends/xtuner_utils/actor.py

@@ -260,9 +260,10 @@ 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, rollout_engine_lock = ray.get(self.rollout_manager.get_rollout_engines_and_lock.remote())
+        rollout_engines, rollout_engine_lock, num_new_engines = ray.get(
+            self.rollout_manager.get_rollout_engines_and_lock.remote()
+        )
+        if num_new_engines > 0:
             self.weight_updator.connect_rollout_engines(rollout_engines, rollout_engine_lock)
             dist.barrier(group=get_gloo_group())
 

slime/ray/placement_group.py

@@ -175,9 +175,6 @@ def create_rollout_manager(args, pg, 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: