diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index 1439fa6216..07d26c08b8 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/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: diff --git a/slime/backends/megatron_utils/cp_utils.py b/slime/backends/megatron_utils/cp_utils.py index 6ee9193963..3ba6adfdf2 100644 --- a/slime/backends/megatron_utils/cp_utils.py +++ b/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) diff --git a/slime/backends/megatron_utils/data.py b/slime/backends/megatron_utils/data.py index ffcfb70429..d19d902d35 100644 --- a/slime/backends/megatron_utils/data.py +++ b/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}) diff --git a/slime/backends/megatron_utils/loss.py b/slime/backends/megatron_utils/loss.py index 50b1ffad3e..fc492be7a7 100644 --- a/slime/backends/megatron_utils/loss.py +++ b/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) diff --git a/slime/backends/megatron_utils/update_weight_utils.py b/slime/backends/megatron_utils/update_weight_utils.py index dedbea7512..654a8364bb 100644 --- a/slime/backends/megatron_utils/update_weight_utils.py +++ b/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 diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index aeafedd3cc..1f739dd3e4 100644 --- a/slime/utils/arguments.py +++ b/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 diff --git a/slime/utils/ppo_utils.py b/slime/utils/ppo_utils.py index 1236f5c4ae..60472c794a 100644 --- a/slime/utils/ppo_utils.py +++ b/slime/utils/ppo_utils.py @@ -164,44 +164,13 @@ def get_reinforce_plus_plus_returns( final_returns_chunks = [] for i in range(len(rewards)): local_kl_chunk = kl[i] - device, dtype = local_kl_chunk.device, local_kl_chunk.dtype total_len, response_len = total_lengths[i], response_lengths[i] - prompt_len = total_len - response_len if cp_size > 1: - from slime.backends.megatron_utils.cp_utils import get_logits_and_tokens_offset_with_cp - - # Step 1: Gather all KL chunks and token_offsets from all ranks - _, _, _, token_offsets = get_logits_and_tokens_offset_with_cp(total_len, response_len) - - object_to_gather = {"kl_chunk": local_kl_chunk.cpu(), "offsets": token_offsets} - gathered_objects = [None] * cp_size - dist.all_gather_object(gathered_objects, object_to_gather, group=mpu.get_context_parallel_group()) - - # Step 2: Reconstruct the full response tensor by splitting and placing each part. - full_kl_response = torch.zeros(response_len, device=device, dtype=dtype) - for obj in gathered_objects: - kl_chunk = obj["kl_chunk"].to(device) - global_offsets = obj["offsets"] - - # Calculate the lengths of part_0 and part_1 for this specific chunk. - s0, e0 = global_offsets[0] - s1, e1 = global_offsets[1] - res_s0, res_e0 = max(0, s0 - prompt_len), max(0, e0 - prompt_len) - res_s1, res_e1 = max(0, s1 - prompt_len), max(0, e1 - prompt_len) - len0 = res_e0 - res_s0 - len1 = res_e1 - res_s1 - - if kl_chunk.numel() > 0: - # Split the received contiguous chunk back into its zigzag parts. - kl_part_0, kl_part_1 = torch.split(kl_chunk, [len0, len1]) - - # Place each part in its own correct location. - if kl_part_0.numel() > 0: - full_kl_response[res_s0:res_e0] = kl_part_0 - if kl_part_1.numel() > 0: - full_kl_response[res_s1:res_e1] = kl_part_1 + # Step 1,2:Gather all chunks and token_offsets from all ranks and reconstruct the full response tensor by splitting and placing each part + from slime.backends.megatron_utils.cp_utils import all_gather_with_cp + full_kl_response = all_gather_with_cp(local_kl_chunk, total_len, response_len) else: full_kl_response = local_kl_chunk @@ -221,23 +190,9 @@ def get_reinforce_plus_plus_returns( # Step 4: Pick up the results corresponding to our local chunk's parts. if cp_size > 1: - local_returns_chunk_parts = [] - local_s0, local_e0 = token_offsets[0] - local_s1, local_e1 = token_offsets[1] - local_res_s0, local_res_e0 = max(0, local_s0 - prompt_len), max(0, local_e0 - prompt_len) - local_res_s1, local_res_e1 = max(0, local_s1 - prompt_len), max(0, local_e1 - prompt_len) - - if local_res_e0 > local_res_s0: - local_returns_chunk_parts.append(returns_for_seq[local_res_s0:local_res_e0]) - if local_res_e1 > local_res_s1: - local_returns_chunk_parts.append(returns_for_seq[local_res_s1:local_res_e1]) - - local_returns_chunk = ( - torch.cat(local_returns_chunk_parts) - if local_returns_chunk_parts - else torch.tensor([], device=device, dtype=dtype) - ) + from slime.backends.megatron_utils.cp_utils import slice_log_prob_with_cp + local_returns_chunk = slice_log_prob_with_cp(returns_for_seq, total_len, response_len) else: local_returns_chunk = returns_for_seq @@ -276,6 +231,8 @@ def get_reinforce_plus_plus_baseline_advantages( def get_advantages_and_returns( + total_len: int, + response_len: int, values: torch.Tensor, rewards: torch.Tensor, gamma: float, @@ -301,17 +258,38 @@ def get_advantages_and_returns( - advantages: Tensor of shape (response_size,) - returns: Tensor of shape (response_size,) """ + from megatron.core import mpu + + cp_size = mpu.get_context_parallel_world_size() + if cp_size > 1: + from slime.backends.megatron_utils.cp_utils import all_gather_with_cp + + full_rewards = all_gather_with_cp(rewards, total_len, response_len) + full_values = all_gather_with_cp(values, total_len, response_len) + else: + full_rewards = rewards + full_values = values + lastgaelam = 0 advantages_reversed = [] - response_length = rewards.size(0) - for t in reversed(range(response_length)): - nextvalues = values[t + 1] if t < response_length - 1 else 0.0 - delta = rewards[t] + gamma * nextvalues - values[t] + for t in reversed(range(response_len)): + nextvalues = full_values[t + 1] if t < response_len - 1 else 0.0 + delta = full_rewards[t] + gamma * nextvalues - full_values[t] lastgaelam = delta + gamma * lambd * lastgaelam advantages_reversed.append(lastgaelam) - advantages = torch.tensor(advantages_reversed[::-1], dtype=values.dtype, device=values.device) - returns = advantages + values + full_advantages = torch.tensor(advantages_reversed[::-1], dtype=full_values.dtype, device=full_values.device) + full_returns = full_advantages + full_values + + if cp_size > 0: + from slime.backends.megatron_utils.cp_utils import slice_log_prob_with_cp + + advantages = slice_log_prob_with_cp(full_advantages, total_len, response_len) + returns = slice_log_prob_with_cp(full_returns, total_len, response_len) + else: + advantages = full_advantages + returns = full_returns + return advantages.detach(), returns