[Fix] ppo rollout engins for distribute

slime/backends/megatron_utils/actor.py

@@ -261,31 +261,19 @@ def train(self, rollout_id, rollout_data_ref):
     def train_critic(self, rollout_id, rollout_data):
         # Create data iterator for log_probs and train.
         data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
-        values = forward_only(
-            get_values,
-            self.args,
-            self.model,
-            data_iterator,
-            num_microbatches,
-        )["values"]
-
-        if rollout_id < self.args.num_critic_only_steps:
-            # we will only use the shape of log_probs in this situation
-            log_probs = values
-            ref_log_probs = values
-        else:
-            values, log_probs, ref_log_probs = sync_actor_critic_data(
-                self.args, values, None, None, self._actor_critic_groups
-            )
-
         rollout_data.update(
-            {
-                "values": values,
-                "log_probs": log_probs,
-                "ref_log_probs": ref_log_probs,
-            }
+            forward_only(
+                get_values,
+                self.args,
+                self.model,
+                data_iterator,
+                num_microbatches,
+            )
         )
 
+        if rollout_id >= self.args.num_critic_only_steps:
+            sync_actor_critic_data(self.args, rollout_data, self._actor_critic_groups)
+
         compute_advantages_and_returns(self.args, rollout_data)
 
         self.args.loss_type = "value_loss"
@@ -308,33 +296,32 @@ def train_actor(self, rollout_id, rollout_data):
                 if "ref" in self.weights:
                     if self.args.use_routing_replay:
                         os.environ["ROUTING_REPLAY_STAGE"] = "fallthrough"
-                    ref_log_probs = self.compute_log_prob(
-                        "ref",
-                        data_iterator,
-                        num_microbatches,
-                        store_prefix="ref_",
+                    rollout_data.update(
+                        self.compute_log_prob(
+                            "ref",
+                            data_iterator,
+                            num_microbatches,
+                            store_prefix="ref_",
+                        )
                     )
-                    rollout_data.update(ref_log_probs)
 
                 if self.args.use_routing_replay:
                     os.environ["ROUTING_REPLAY_STAGE"] = "record"
-                log_probs = self.compute_log_prob(
-                    "old_actor" if self.args.keep_old_actor else "actor",
-                    data_iterator,
-                    num_microbatches,
-                    store_prefix="",
+                rollout_data.update(
+                    self.compute_log_prob(
+                        "old_actor" if self.args.keep_old_actor else "actor",
+                        data_iterator,
+                        num_microbatches,
+                        store_prefix="",
+                    )
                 )
-                rollout_data.update(log_probs)
 
                 if self.args.use_critic:
-                    values, log_probs, ref_log_probs = sync_actor_critic_data(
+                    sync_actor_critic_data(
                         self.args,
-                        None,
-                        log_probs["log_probs"],
-                        ref_log_probs["ref_log_probs"] if (self.args.kl_coef != 0 or self.args.use_kl_loss) else None,
+                        rollout_data,
                         self._actor_critic_groups,
                     )
-                    rollout_data.update({"values": values})
 
                 # when there is old actor, we need to update the model params to actor manually
                 if "old_actor" in self.weights:

slime/backends/megatron_utils/cp_utils.py

@@ -1,3 +1,5 @@
+from typing import Union
+
 import torch
 import torch.distributed as dist
 import torch.nn.functional as F
@@ -170,7 +172,7 @@ def slice_with_cp(tokens: torch.Tensor, pad_value):
     return torch.cat([tokens[start_1:end_1], tokens[start_2:end_2]])
 
 
-def slice_log_prob_with_cp(log_prob: list[float], total_length: int, response_length: int):
+def slice_log_prob_with_cp(log_prob: Union[list[float], torch.Tensor], total_length: int, response_length: int):
     assert len(log_prob) == response_length
 
     cp_size = mpu.get_context_parallel_world_size()
@@ -183,4 +185,8 @@ def slice_log_prob_with_cp(log_prob: list[float], total_length: int, response_le
 
     chunk_1 = log_prob[logits_offset[0][0] - (prompt_length - 1) : logits_offset[0][1] - (prompt_length - 1)]
     chunk_2 = log_prob[logits_offset[1][0] - (prompt_length - 1) : logits_offset[1][1] - (prompt_length - 1)]
-    return chunk_1 + chunk_2
+
+    if isinstance(log_prob, list):
+        return chunk_1 + chunk_2
+    else:
+        return torch.cat([chunk_1, chunk_2], dim=0)

slime/backends/megatron_utils/data.py

@@ -390,20 +390,24 @@ def log_perf_data(rollout_id, args):
 
 def sync_actor_critic_data(
     args,
-    values: Optional[list[torch.Tensor]] = None,
-    log_probs: Optional[list[torch.Tensor]] = None,
-    ref_log_probs: Optional[list[torch.Tensor]] = None,
+    rollout_data: Optional[dict[str, list[torch.Tensor]]] = None,
     group: Optional[dist.ProcessGroup] = None,
 ):
+    values, log_probs, ref_log_probs = map(rollout_data.get, ("values", "log_probs", "ref_log_probs"))
+
+    # return when not the pp last stage
+    if not values and not log_probs:
+        return
+
     handles = []
 
-    if values is None:
+    if not values:
         values = [torch.empty_like(log_prob) for log_prob in log_probs]
     for value in values:
         handles.append(dist.broadcast(value, src=1, group=group, async_op=True))
 
     if args.kl_coef != 0 or args.use_kl_loss:
-        if log_probs is None:
+        if not log_probs:
             ref_log_probs = [torch.empty_like(value) for value in values]
             log_probs = [torch.empty_like(value) for value in values]
         for ref_log_prob, log_prob in zip(ref_log_probs, log_probs):
@@ -412,4 +416,5 @@ def sync_actor_critic_data(
 
     for handle in handles:
         handle.wait()
-    return values, log_probs, ref_log_probs
+
+    rollout_data.update({"values": values, "log_probs": log_probs, "ref_log_probs": ref_log_probs})

slime/backends/megatron_utils/loss.py

@@ -137,7 +137,7 @@ def compute_advantages_and_returns(args, rollout_data):
     if log_probs is None and values is None:
         return
 
-    if args.kl_coef == 0:
+    if args.kl_coef == 0 or not log_probs:
         # when kl_coef is 0, we won't compute ref_log_prob
         xs = log_probs if log_probs is not None else values
         kl = [torch.zeros_like(x, dtype=torch.float32, device=x.device) for x in xs]
@@ -163,13 +163,17 @@ def compute_advantages_and_returns(args, rollout_data):
         rewards = []
         for reward, k in zip(old_rewards, kl):
             k *= -args.kl_coef
-            k[-1] += reward
+            cp_rank = mpu.get_context_parallel_rank()
+            if cp_rank == 0:
+                k[-1] += reward
             rewards.append(k)
         advantages, returns = list(
             zip(
                 *[
-                    get_advantages_and_returns(value, reward, args.gamma, args.lambd)
-                    for value, reward in zip(values, rewards)
+                    get_advantages_and_returns(total_length, response_length, value, reward, args.gamma, args.lambd)
+                    for total_length, response_length, value, reward in zip(
+                        total_lengths, response_lengths, values, rewards
+                    )
                 ]
             )
         )
@@ -363,7 +367,7 @@ def value_loss_function(args, batch, logits, sum_of_sample_mean):
         total_lengths=batch["total_lengths"],
         response_lengths=batch["response_lengths"],
     )
-    values = torch.cat([value.squeeze(-1) for value in values["values"]], dim=0)
+    values = torch.cat([value.flatten() for value in values["values"]], dim=0)
 
     returns = torch.cat(batch["returns"], dim=0)
 

slime/backends/megatron_utils/update_weight_utils.py

@@ -307,6 +307,24 @@ def __init__(self, args, model, weights, *, model_name, quantization_config, voc
 
     def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
         self.rollout_engines = rollout_engines
+        colocate_engine_nums = (
+            self.args.actor_num_nodes * self.args.actor_num_gpus_per_node // self.args.rollout_num_gpus_per_engine
+        )
+        self.use_distribute = len(rollout_engines) > colocate_engine_nums
+
+        if self.use_distribute:
+            self.rollout_engines = rollout_engines[:colocate_engine_nums]
+            self.distributed_rollout_engines = rollout_engines[colocate_engine_nums:]
+            self._is_distributed_src_rank = (
+                mpu.get_data_parallel_rank(with_context_parallel=True) == 0
+                and mpu.get_tensor_model_parallel_rank() == 0
+                and mpu.get_pipeline_model_parallel_rank() == 0
+            )
+            self._group_name = "slime"
+            if self._is_distributed_src_rank:
+                self._model_update_groups = connect_rollout_engines_from_distributed(
+                    self.args, self._group_name, self.distributed_rollout_engines
+                )
 
         # Here we assume the gpu id of rollout engines and train actors are the same.
         for i, engine in enumerate(self.rollout_engines):
@@ -400,7 +418,22 @@ def _update_bucket_weights_from_tensor(self, param_infos):
             converted_named_tensors.extend(
                 convert_to_hf(self.args, self.model_name, info.name, param, self.quantization_config)
             )
-        self._update_converted_params_from_tensor(converted_named_tensors)
+
+        refs = self._update_converted_params_from_tensor(converted_named_tensors)
+
+        if self.use_distribute and self._is_distributed_src_rank:
+            refs.extend(
+                update_weights_from_distributed(
+                    self.args,
+                    self._group_name,
+                    self._model_update_groups,
+                    self.weight_version,
+                    self.distributed_rollout_engines,
+                    converted_named_tensors,
+                )
+            )
+
+        ray.get(refs)
 
     def _update_converted_params_from_tensor(self, converted_named_tensors):
         if use_flattened_tensor_bucket:
@@ -451,7 +484,8 @@ def _update_converted_params_from_tensor(self, converted_named_tensors):
                     "weight_version": str(self.weight_version),
                 }
                 refs.append(self._ipc_engine.update_weights_from_tensor.remote(**kwargs))
-            ray.get(refs)
+            return refs
+        return []
 
 
 class UpdateWeightFromDistributed:
@@ -478,31 +512,9 @@ def connect_rollout_engines(self, rollout_engines, rollout_engine_lock):
             self._group_name = f"slime-pp_{pp_rank}"
 
         if self._is_pp_src_rank:
-            master_address = ray._private.services.get_node_ip_address()
-            with socket.socket() as sock:
-                sock.bind(("", 0))
-                master_port = sock.getsockname()[1]
-            world_size = self.args.rollout_num_gpus + 1
-
-            refs = [
-                engine.init_weights_update_group.remote(
-                    master_address,
-                    master_port,
-                    i * self.args.rollout_num_gpus_per_engine + 1,
-                    world_size,
-                    self._group_name,
-                    backend="nccl",
-                )
-                for i, engine in enumerate(self.rollout_engines)
-            ]
-            self._model_update_groups = init_process_group(
-                backend="nccl",
-                init_method=f"tcp://{master_address}:{master_port}",
-                world_size=world_size,
-                rank=0,
-                group_name=self._group_name,
+            self._model_update_groups = connect_rollout_engines_from_distributed(
+                self.args, self._group_name, rollout_engines
             )
-            ray.get(refs)
 
     @torch.no_grad()
     def update_weights(self):
@@ -606,31 +618,74 @@ def _update_expert_bucket_weights_from_distributed(self, named_tensors, pbar=Non
         converted_hf_tensors = []
         for name, param in all_gathered_params:
             converted_hf_tensors += convert_to_hf(self.args, self.model_name, name, param, self.quantization_config)
-        self._update_bucket_weights_from_distributed(converted_hf_tensors, pbar=pbar)
+
+        self._update_bucket_weights_from_distributed(converted_hf_tensors, pbar)
 
     def _update_bucket_weights_from_distributed(self, converted_named_tensors, pbar=None):
         # lock the rollout engines to prevent dead lock on broadcast.
         while not ray.get(self.rollout_engine_lock.acquire.remote()):
             time.sleep(0.1)
 
-        refs = [
-            engine.update_weights_from_distributed.remote(
-                names=[name for name, _ in converted_named_tensors],
-                dtypes=[param.dtype for _, param in converted_named_tensors],
-                shapes=[param.shape for _, param in converted_named_tensors],
-                group_name=self._group_name,
-                weight_version=str(self.weight_version),
-            )
-            for engine in self.rollout_engines
-        ]
-
-        handles = []
-        for _, param in converted_named_tensors:
-            handles.append(dist.broadcast(param.data, 0, group=self._model_update_groups, async_op=True))
-        for handle in handles:
-            handle.wait()
+        refs = update_weights_from_distributed(
+            self.args,
+            self._group_name,
+            self._model_update_groups,
+            self.weight_version,
+            self.rollout_engines,
+            converted_named_tensors,
+        )
 
         ray.get(refs)
         converted_named_tensors.clear()
         ray.get(self.rollout_engine_lock.release.remote())
         pbar.update(1)
+
+
+def connect_rollout_engines_from_distributed(args, group_name, rollout_engines):
+    master_address = ray._private.services.get_node_ip_address()
+    with socket.socket() as sock:
+        sock.bind(("", 0))
+        master_port = sock.getsockname()[1]
+    world_size = len(rollout_engines) * args.rollout_num_gpus_per_engine + 1
+
+    refs = [
+        engine.init_weights_update_group.remote(
+            master_address,
+            master_port,
+            i * args.rollout_num_gpus_per_engine + 1,
+            world_size,
+            group_name,
+            backend="nccl",
+        )
+        for i, engine in enumerate(rollout_engines)
+    ]
+    model_update_groups = init_process_group(
+        backend="nccl",
+        init_method=f"tcp://{master_address}:{master_port}",
+        world_size=world_size,
+        rank=0,
+        group_name=group_name,
+    )
+    ray.get(refs)
+    return model_update_groups
+
+
+def update_weights_from_distributed(args, group_name, group, weight_version, rollout_engines, converted_named_tensors):
+    refs = [
+        engine.update_weights_from_distributed.remote(
+            names=[name for name, _ in converted_named_tensors],
+            dtypes=[param.dtype for _, param in converted_named_tensors],
+            shapes=[param.shape for _, param in converted_named_tensors],
+            group_name=group_name,
+            weight_version=str(weight_version),
+        )
+        for engine in rollout_engines
+    ]
+
+    handles = []
+    for _, param in converted_named_tensors:
+        handles.append(dist.broadcast(param.data, 0, group=group, async_op=True))
+    for handle in handles:
+        handle.wait()
+
+    return refs

slime/utils/arguments.py

@@ -1082,6 +1082,8 @@ def slime_validate_args(args):
                 f"* actor_num_nodes {args.actor_num_nodes}, overriding rollout_num_gpus to match actor_num_gpus_per_node * actor_num_nodes."
             )
             args.rollout_num_gpus = args.actor_num_gpus_per_node * args.actor_num_nodes
+            if args.use_critic:
+                args.rollout_num_gpus += args.critic_num_gpus_per_node * args.critic_num_nodes
 
     if args.eval_function_path is None:
         args.eval_function_path = args.rollout_function_path