# Copyright 2024 Alibaba Group Holding Limited. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Sync parameters""" import concurrent.futures import traceback import time from collections import defaultdict from concurrent.futures import ThreadPoolExecutor from itertools import cycle from typing import List, Dict from queue import PriorityQueue import torch from tqdm import tqdm from chatlearn.launcher.initialize import patch_ray from chatlearn.utils import future from chatlearn.utils import utils from chatlearn.utils.constant import LORA_WEIGHT_PREFIX from chatlearn.utils.constant import PARAM_SYNC_COMM_TYPE from chatlearn.utils.constant import ROUTED_EXPERT_REGROUPING_COMM_TYPE from chatlearn.utils.global_vars import get_args from chatlearn.utils.logger import logger from chatlearn.utils.utils import execute_in_parallel from chatlearn.utils.timer import Timers from chatlearn.synchronizer.scheduler import CollectiveTask, parallel_execute_collective_tasks from . import get_synchronizer patch_ray() class ParameterSyncGroup: """ParameterSyncGroup""" def __init__(self, src_model, dst_model, group_name, frequency, error_signal): self.src_model = src_model self.dst_model = dst_model self.synchronizer = get_synchronizer(src_model, dst_model) self.group_name = group_name self.error_signal = error_signal self.send_recv_actor_mappings = defaultdict(list) self.recv_send_actor_mappings = defaultdict(list) self.send_recv_actor_mappings_stage2 = defaultdict(list) self.recv_send_actor_mappings_stage2 = defaultdict(list) self.actor2rank = {} self.actor2model = {} self._debug = get_args().runtime_args.debug self._num_src_pipeline_stage = None self._num_dst_pipeline_stage = None self._num_src_expert_parallel = None self._num_dst_expert_parallel = None self._num_src_tensor_parallel = None self._num_dst_tensor_parallel = None self._send_recv_param_names = {} self._actor2pipe = {} self._actor2tp = {} self._actor2ep = {} self._actor2dp = {} self._comm_type = get_args().runtime_args.param_sync_comm_type if src_model.colocate_with(dst_model) and self._comm_type == PARAM_SYNC_COMM_TYPE.BROADCAST: if self.num_src_tensor_parallel % 2 == 1 and self.num_dst_tensor_parallel % 2 == 1: logger.warning("Only support PARAM_SYNC_COMM_TYPE.BROADCAST when TP SIZE is even number, use P2P instead") self._comm_type = PARAM_SYNC_COMM_TYPE.P2P self.concurrent_comm = get_args().runtime_args.concurrent_comm self._enable_lora = self.src_model.module_args.lora.enable_lora # sync every n episodes, n = 0 for no param sync self._frequency = frequency self._free_sync_collective_group = get_args().runtime_args.free_sync_collective_group self._is_collective_group_created = True self.collective_groups = [] self.groups2actors = {} # group_name -> []actors self.src_dp_size = future.get(self.src_model.replicas[0].all_actors[0].get_data_parallel_size.remote()) self.send_actors_to_regroup_routed_experts = None self._comm_type_to_regroup_routed_experts = get_args().runtime_args.routed_expert_regrouping_comm_type assert self._comm_type_to_regroup_routed_experts in \ [ROUTED_EXPERT_REGROUPING_COMM_TYPE.ALLGATHER, ROUTED_EXPERT_REGROUPING_COMM_TYPE.ALLTOALL], \ f"Only support 'allgather' or 'alltoall' for routed expert regrouping, while {self._comm_type_to_regroup_routed_experts}" if self._comm_type_to_regroup_routed_experts == ROUTED_EXPERT_REGROUPING_COMM_TYPE.ALLTOALL: if self.num_dst_tensor_parallel * self.num_dst_expert_parallel != self.num_src_tensor_parallel * self.num_src_expert_parallel: logger.info("Only support ROUTED_EXPERT_REGROUPING_COMM_TYPE.ALLTOALL when src tp eqs dst tp, use 'allgather' instead.") self._comm_type_to_regroup_routed_experts = ROUTED_EXPERT_REGROUPING_COMM_TYPE.ALLGATHER logger.info(f"Set ROUTED_EXPERT_REGROUPING_COMM_TYPE = {self._comm_type_to_regroup_routed_experts}.") self.sorted_send_actors = None self.sorted_send_actors_stage2 = None self.actor2synchronizer = {} self.setup_collective_group() self.setup_rank_mapping() self.timers = Timers() def get_group_name(self, actors): return f"{self.group_name}_" + "_".join(str(self.actor2rank[actor]) for actor in actors) @property def frequency(self): return self._frequency def get_or_cache(self, actor, func_name, *args, **kwargs): def inner_func(*args, **kwargs): return future.get(getattr(getattr(actor, func_name), 'remote')(*args, **kwargs)) cached_name = str(actor) + "_" + func_name if hasattr(self, cached_name): cached = getattr(self, cached_name) else: cached = {} setattr(self, cached_name, cached) return utils.get_or_cache(cached, actor, inner_func, *args, **kwargs) def is_same_gpu(self, src_actor, dst_actor): src_gpu = self.get_or_cache(src_actor, "get_visible_gpus") dst_gpu = self.get_or_cache(dst_actor, "get_visible_gpus") src_address = self.get_or_cache(src_actor, "get_address") dst_address = self.get_or_cache(dst_actor, "get_address") return src_gpu == dst_gpu and src_address == dst_address @property def num_src_pipeline_stage(self): if self._num_src_pipeline_stage is None: self._num_src_pipeline_stage = future.get(self.src_model.replicas[0].all_actors[0].pipeline_model_parallel_size.remote()) return self._num_src_pipeline_stage @property def num_dst_pipeline_stage(self): if self._num_dst_pipeline_stage is None: self._num_dst_pipeline_stage = future.get(self.dst_model.replicas[0].all_actors[0].pipeline_model_parallel_size.remote()) return self._num_dst_pipeline_stage @property def num_src_tensor_parallel(self): if self._num_src_tensor_parallel is None: self._num_src_tensor_parallel = future.get(self.src_model.replicas[0].all_actors[0].tensor_model_parallel_size.remote()) return self._num_src_tensor_parallel @property def num_dst_tensor_parallel(self): if self._num_dst_tensor_parallel is None: self._num_dst_tensor_parallel = future.get(self.dst_model.replicas[0].all_actors[0].tensor_model_parallel_size.remote()) return self._num_dst_tensor_parallel @property def num_src_expert_parallel(self): if self._num_src_expert_parallel is None: self._num_src_expert_parallel = future.get(self.src_model.replicas[0].all_actors[0].expert_model_parallel_size.remote()) return self._num_src_expert_parallel @property def num_dst_expert_parallel(self): if self._num_dst_expert_parallel is None: self._num_dst_expert_parallel = future.get(self.dst_model.replicas[0].all_actors[0].expert_model_parallel_size.remote()) return self._num_dst_expert_parallel def setup_collective_group(self): """ Set up collective group for parameter sync. The group consists of all actors from both model's (src & dst) replicas. For P2P Sync, create the group directly. Assigning each actor an uniq rank id. For Broadcast, assigning the rank id only, and defer the group creation during the first parameter sync happens TODO-1: Maybe can remove the P2P sync method in the future for code clarity TODO-2: the _setup_ranks is a public method of DistActor class but named with a leading _ which makes it looks like a private method """ refs = [] # we put src_model first, so we don't need to change the rank of training model models = [self.src_model, self.dst_model] world_size = sum(model.actor_num for model in models) rank_offset = 0 for model in models: for replica in model.replicas: if self._comm_type == PARAM_SYNC_COMM_TYPE.P2P: refs += replica._setup_collective_group(rank_offset, world_size, self.group_name) else: replica._setup_ranks(rank_offset) rank_offset += replica.actor_num if refs: future.get(refs) logger.info(f"init collective group done for {self.group_name}") def destroy_collective_group(self): refs = [] try: refs.extend(self.src_model.destroy_collective_group()) refs.extend(self.dst_model.destroy_collective_group()) future.wait(refs) logger.info(f"destroy_collective_group success for {self.group_name}") except Exception as e: logger.exception(f"destroy_collective_group fail for {self.group_name} {e}") def setup_rank_mapping(self): self.tp_num_mapping = self.num_dst_tensor_parallel // self.num_src_tensor_parallel if self.tp_num_mapping == 1: self.build_rank_mapping() else: self.build_rank_mapping_two_stage() def insert_actor2rank(self, actor, rank: int): if actor not in self.actor2rank: self.actor2rank[actor] = rank def insert_actor2model(self, actor, model): if actor not in self.actor2model: self.actor2model[actor] = model def add_routed_experts_regrouping_actor(self, model, ranks_group: List): for replica_ranks_group in ranks_group: if isinstance(replica_ranks_group[0], list): for tp_ranks in replica_ranks_group: for rank in tp_ranks: actor = model.get_actor(rank) self.insert_actor2rank(actor, rank) self.insert_actor2model(actor, model) else: for rank in replica_ranks_group: actor = model.get_actor(rank) self.insert_actor2rank(actor, rank) self.insert_actor2model(actor, model) # pylint: disable=unused-argument def empty_add_recv_actor(self, src_rank, dst_rank): return def warmup_groups(self): return def add_recv_actor(self, src_rank, dst_rank): src_actor = self.src_model.get_actor(src_rank) self.insert_actor2rank(src_actor, src_rank) self.insert_actor2model(src_actor, self.src_model) dst_actor = self.dst_model.get_actor(dst_rank) self.insert_actor2rank(dst_actor, dst_rank) self.insert_actor2model(dst_actor, self.dst_model) src_gpu = self.get_or_cache(src_actor, "get_visible_gpus") dst_gpu = self.get_or_cache(dst_actor, "get_visible_gpus") src_tp_rank = self.get_actor_tp_rank(src_actor) dst_tp_rank = self.get_actor_tp_rank(dst_actor) src_pp_rank = self.get_actor_pipe_rank(src_actor) dst_pp_rank = self.get_actor_pipe_rank(dst_actor) src_ep_rank = self.get_actor_ep_rank(src_actor) dst_ep_rank = self.get_actor_ep_rank(dst_actor) logger.debug(f"build rank mapping from {src_rank} to {dst_rank}, from gpu {src_gpu} to {dst_gpu}, " + f"from pipe_stage {src_pp_rank} to {dst_pp_rank}, " + f"from tp rank {src_tp_rank} to {dst_tp_rank}, " + f"from ep rank {src_ep_rank} to {dst_ep_rank}.") self.send_recv_actor_mappings[src_actor].append(dst_actor) self.recv_send_actor_mappings[dst_actor].append(src_actor) def add_recv_actor_stage2(self, src_rank, dst_rank): src_actor = self.dst_model.get_actor(src_rank) self.insert_actor2rank(src_actor, src_rank) self.insert_actor2model(src_actor, self.dst_model) # stage 2 sends from dst_model to dst_model dst_actor = self.dst_model.get_actor(dst_rank) self.insert_actor2rank(dst_actor, dst_rank) self.insert_actor2model(dst_actor, self.dst_model) src_gpu = future.get(src_actor.get_visible_gpus.remote()) dst_gpu = future.get(dst_actor.get_visible_gpus.remote()) # TODO(jiangle.jl): support ep/cp. src_tp_rank = self.get_actor_tp_rank(src_actor) dst_tp_rank = self.get_actor_tp_rank(dst_actor) src_pp_rank = self.get_actor_pipe_rank(src_actor) dst_pp_rank = self.get_actor_pipe_rank(dst_actor) logger.debug(f"build rank mapping from {src_rank} to {dst_rank}, from gpu {src_gpu} to {dst_gpu}, " + \ f"from pipe_stage {src_pp_rank} to {dst_pp_rank}, " + \ f"from tp rank {src_tp_rank} to {dst_tp_rank}") self.send_recv_actor_mappings_stage2[src_actor].append(dst_actor) self.recv_send_actor_mappings_stage2[dst_actor].append(src_actor) def set_send_actors_to_regroup_routed_experts(self, src_replica_ranks_group): if self.send_actors_to_regroup_routed_experts is None: self.send_actors_to_regroup_routed_experts = [] for src_replica_ranks in src_replica_ranks_group: self.send_actors_to_regroup_routed_experts.append([]) if isinstance(src_replica_ranks[0], list): for src_tp_ranks in src_replica_ranks: self.send_actors_to_regroup_routed_experts[-1].extend( [self.src_model.get_actor(src_rank) for src_rank in src_tp_ranks]) else: self.send_actors_to_regroup_routed_experts[-1].extend( [self.src_model.get_actor(src_rank) for src_rank in src_replica_ranks]) def get_src_and_dst_dp_ranks(self, is_except_routed_experts=False): """ Return: The DP Group List for src & dst model [[DP-0], [DP-1] ... [DP-N]] """ dst_dp_ranks = self.dst_model.all_ranks local_src_ranks = future.get(self.src_model.replicas[0].get_local_param_ranks()) if local_src_ranks[0] is None or dst_dp_ranks is None: if self._debug: logger.warning( f"DEBUG MODE! src_dp_ranks {local_src_ranks} or dst_dp_ranks: {dst_dp_ranks} is None, " "make sure they have values in real application.") return local_src_ranks, dst_dp_ranks else: raise Exception(f"src_dp_ranks {local_src_ranks} or dst_dp_ranks {dst_dp_ranks} should not be None") dp_rank_to_ranks = defaultdict(list) for local_ranks, dp_rank in local_src_ranks: dp_rank_to_ranks[dp_rank].append(local_ranks[dp_rank]) if is_except_routed_experts: # for weight except routed expert, ep_size using for data parallel. # TODO-1 The logic here is a little bit complicate, it would be better to move to a seperate function # TODO-2 The logic here is about HEP, would be better called from class ParameterSyncGroupwithHEP src_hep_size = self.num_src_expert_parallel * self.num_src_tensor_parallel new_dict = defaultdict(list) idx = 0 for dp_rank, values in dp_rank_to_ranks.items(): assert len(values) % src_hep_size == 0, ( f"len of values({len(values)}) for dp_rank {dp_rank} must be divisible by hep size({src_hep_size})" f" when call get_src_and_dst_dp_ranks_for_except_routed_experts." ) pp_blocks = [values[i:i + src_hep_size] for i in range(0, len(values), src_hep_size)] sub_blocks_per_pp = [] for block in pp_blocks: sub_block_size = src_hep_size // self.num_src_expert_parallel sub_blocks = [block[i:i + sub_block_size] for i in range(0, src_hep_size, sub_block_size)] sub_blocks_per_pp.append(sub_blocks) for i in range(self.num_src_expert_parallel): merged_group = [] for sub_blocks in sub_blocks_per_pp: merged_group.extend(sub_blocks[i]) new_dict[idx].extend(merged_group) idx += 1 src_dp_ranks = [i[1] for i in sorted(new_dict.items())] else: src_dp_ranks = [i[1] for i in sorted(dp_rank_to_ranks.items())] return src_dp_ranks, dst_dp_ranks def get_load_balance_dst_rank( self, lb_dst_offset_pq_dict, s_idx, start, src_rank, dst_replica_ranks_group, d_idx, pre_allocate=False ): """Get the dst_rank for load balance when gpu collides. """ dst_tp_indices = sorted([ s_idx * self.tp_num_mapping + (start + i) % self.tp_num_mapping for i in range(self.tp_num_mapping) ]) indexed_dst_tp_group = tuple(dst_replica_ranks_group[d_idx][dst_tp_index] for dst_tp_index in dst_tp_indices) # Construct a priority queue (PQ) to retrieve `dst_rank` for load balancing when gpu collides. # The key of the PQ is (hit_time, max_seq_num), meaning that the rank is used for `hit_time` times, # while `max_seq_num` further sorts `dst_rank` when `hit_time` remains the same. if indexed_dst_tp_group not in lb_dst_offset_pq_dict: pq = PriorityQueue() max_seq_num = 0 hit_time = 0 while max_seq_num < self.tp_num_mapping: pq.put(( hit_time, max_seq_num, s_idx * self.tp_num_mapping + (start + max_seq_num) % self.tp_num_mapping )) max_seq_num += 1 lb_dst_offset_pq_dict[indexed_dst_tp_group] = [pq, max_seq_num] else: max_seq_num = lb_dst_offset_pq_dict[indexed_dst_tp_group][1] # Each time, we retrieve the first value of the PQ. # 1. If the first `dst_rank` will encounter gpu collision with `src_rank`, we retrieve it, set `seq_num` to # `max_seq_num`, increase `max_seq_num` by 1, and finally insert <(hit_time, seq_num), offset> back # to the PQ. # 2. If the first `dst_rank` won't encounter gpu collision with `src_rank`, we insert it to another PQ (called # legal_lb_recv_offset_pq). After looping through all legal solutions, we will retrieve the first load-balance one. # 3. If we cannot find a legal solution after `self.tp_num_mapping` times, all `src_rank` will encounter gpu collision # with `dst_rank`, we throw a runtime exception. lb_recv_offset_pq = lb_dst_offset_pq_dict[indexed_dst_tp_group][0] legal_lb_recv_offset_pq = PriorityQueue() assert len(lb_recv_offset_pq.queue) == self.tp_num_mapping, ( "length of the load-balance recv_offset priority queue must be equal to tp_num_mapping, " f"got {len(lb_recv_offset_pq.queue)} and {self.tp_num_mapping}." ) is_collide = False for _ in range(self.tp_num_mapping): hit_time, seq_num, offset = lb_recv_offset_pq.get() dst_rank = dst_replica_ranks_group[d_idx][offset] logger.debug(f"Trying to match {src_rank} and {dst_rank} (hit={hit_time}), remaining queue={lb_recv_offset_pq.queue})") src_actor = self.src_model.get_actor(src_rank) dst_actor = self.dst_model.get_actor(dst_rank) if self.is_same_gpu(src_actor, dst_actor): logger.info( f"src_rank ({src_rank}) will share the same gpu with dst_rank ({dst_rank}). " "This is not allowed in NCCL send-recv. ChatLearn will skip dst_rank to the next legal one." ) is_collide = True lb_recv_offset_pq.put((hit_time, max_seq_num, offset)) max_seq_num += 1 lb_dst_offset_pq_dict[indexed_dst_tp_group][1] = max_seq_num else: legal_lb_recv_offset_pq.put((hit_time, seq_num, offset)) logger.debug(f"legal_lb_recv_offset_pq={legal_lb_recv_offset_pq.queue}") # if pre_allocate is True and no collide, we directly return if pre_allocate is True and is_collide is False: while len(legal_lb_recv_offset_pq.queue) > 0: lb_recv_offset_pq.put(legal_lb_recv_offset_pq.get()) return None, False # there must be at least one legal recv offset if len(legal_lb_recv_offset_pq.queue) == 0: raise RuntimeError( f"Rank mapping solution is infeasible because src_rank ({src_rank}) will collide with all candidates." ) # extract the first legal one to keep load balance hit_time, seq_num, offset = legal_lb_recv_offset_pq.get() lb_recv_offset_pq.put((hit_time + 1, seq_num, offset)) # put other solutions back to lb_recv_offset_pq while len(legal_lb_recv_offset_pq.queue) > 0: lb_recv_offset_pq.put(legal_lb_recv_offset_pq.get()) logger.debug(f"after retrieving, lb_recv_offset_pq = {lb_recv_offset_pq.queue}") # return dst_rank dst_rank = dst_replica_ranks_group[d_idx][offset] return dst_rank, is_collide def build_rank_mapping(self, add_recv_actor_fn=None): """ setup rank mapping for src parameter and dst parameter get rank for one src_model, without model replicas for each DP Group: for each TP & EP Group: for each PP Group: mapping[src] = dst """ if add_recv_actor_fn is None: add_recv_actor_fn = self.add_recv_actor src_dp_ranks, dst_dp_ranks = self.get_src_and_dst_dp_ranks() if self._debug and (src_dp_ranks[0] is None or dst_dp_ranks is None): return if self.src_model.colocate_with(self.dst_model) and self.num_src_tensor_parallel % 2 == 1: replica_rank_iter = cycle(reversed(src_dp_ranks)) else: replica_rank_iter = cycle(iter(src_dp_ranks)) logger.debug(f"src_dp_ranks: {src_dp_ranks}") logger.debug(f"dst_dp_ranks: {dst_dp_ranks}") assert self.num_src_pipeline_stage % self.num_dst_pipeline_stage == 0 def split_ranks_by_tp_and_ep_size(ranks, tp_size : int = 1, ep_size : int = 1): tp_and_ep_size = tp_size * ep_size return [ranks[i:i + tp_and_ep_size] for i in range(0, len(ranks), tp_and_ep_size)] for dst_replica_ranks in dst_dp_ranks: src_replica_ranks = next(replica_rank_iter) src_replica_ranks_group = split_ranks_by_tp_and_ep_size(src_replica_ranks, self.num_src_tensor_parallel, self.num_src_expert_parallel) dst_replica_ranks_group = split_ranks_by_tp_and_ep_size(dst_replica_ranks, self.num_dst_tensor_parallel, self.num_dst_expert_parallel) self.set_send_actors_to_regroup_routed_experts(src_replica_ranks_group) pipe_map_interval = self.num_src_pipeline_stage // self.num_dst_pipeline_stage for i, src_tp_group in enumerate(src_replica_ranks_group): j = i // pipe_map_interval for src_rank, dst_rank in zip(src_tp_group, dst_replica_ranks_group[j]): add_recv_actor_fn(src_rank, dst_rank) # pylint: disable=unused-argument def build_rank_mapping_for_ep(self, add_recv_actor_fn=None): # Currently, we do not support build rank mapping for expert parallelism raise NotImplementedError("ChatLearn does not support build rank mapping from Megatron-LM for expert parallelism") def build_rank_mapping_two_stage(self, add_recv_actor_fn=None): # setup rank mapping for src parameter and dst parameter # get rank for one src_model, without model replicas if add_recv_actor_fn is None: add_recv_actor_stage1_fn = self.add_recv_actor add_recv_actor_stage2_fn = self.add_recv_actor_stage2 else: assert len(add_recv_actor_fn) == 2, ( "The length of add_recv_actor_fn should be 2. The first one is a function handler for communication stage 1, " "while the second one is a function handler for communication stage 2." ) add_recv_actor_stage1_fn = add_recv_actor_fn[0] add_recv_actor_stage2_fn = add_recv_actor_fn[1] src_ranks, dst_ranks = self.get_src_and_dst_dp_ranks(is_except_routed_experts=True) if self._debug and (src_ranks[0] is None or dst_ranks is None): return replica_rank_iter = cycle(iter(src_ranks)) logger.debug(f"src_ranks: {src_ranks}") logger.debug(f"dst_ranks: {dst_ranks}") assert self.num_dst_tensor_parallel % self.num_src_tensor_parallel == 0, \ "currently we require mod value equals to zero for tensor_model_parallel_size of dst_model and that of src_model while " + \ f"src model {self.src_model.name}(TP={self.num_src_tensor_parallel}) and " + \ f"dst model {self.dst_model.name}(TP={self.num_dst_tensor_parallel})" assert self.num_src_pipeline_stage % self.num_dst_pipeline_stage == 0 def split_ranks_by_tp_and_ep_size(ranks, tp_size, ep_size): if ep_size > 1: sort_ranks_on_grouped_tp = [] index = 0 tp_index = 0 for _ in range(len(ranks)): sort_ranks_on_grouped_tp.append(index) if tp_index < tp_size - 1: index += 1 tp_index += 1 else: start_index = index + 1 - tp_size index = start_index + (ep_size * tp_size) tp_index = 0 if index >= len(ranks): index = (index + tp_size) % len(ranks) else: sort_ranks_on_grouped_tp = ranks return [sort_ranks_on_grouped_tp[i:i + tp_size] for i in range(0, len(sort_ranks_on_grouped_tp), tp_size)] pair_list = [] p2p_list = [] src_replica_offset = 0 lb_dst_offset_pq_dict = {} for dst_replica_ranks in dst_ranks: src_replica_ranks = next(replica_rank_iter) # for weight except routed expert, ep_size using for data parallel. src_replica_ranks_group = split_ranks_by_tp_and_ep_size(src_replica_ranks, self.num_src_tensor_parallel, 1) dst_replica_ranks_group = split_ranks_by_tp_and_ep_size(dst_replica_ranks, self.num_dst_tensor_parallel, self.num_dst_expert_parallel) logger.debug(f"src_replica_ranks_group: {src_replica_ranks_group}") logger.debug(f"dst_replica_ranks_group: {dst_replica_ranks_group}") pipe_map_interval = self.num_src_pipeline_stage // self.num_dst_pipeline_stage assert pipe_map_interval >= 1, \ f"dst_pp expected to divide src_pp, while src_pp {self.num_src_pipeline_stage} and dst_pp {self.num_dst_pipeline_stage}" # stage 1: comm pairs that broadcast params from trainer to inference model # Each rank in trainer holds weights for tp_num_mapping ranks in inference model. # For example: trainer_tp = 2, inference_tp = 4 => tp_num_mapping = inference_tp // trainer_tp = 2 # Weight mapping from training to inference: # [0] -> [0', 1'] # [1] -> [2', 3'] # To avoid p2p communication on the same gpu, we only broadcast params to first rank in weight_mapping_group. # Comm mapping from training to inference: # [0] -> [0'] # [1] -> [2'] # Firstly, pre-allocate for those gpu collisions uncollided_index_to_start_j = {} for i, src_tp_group in enumerate(src_replica_ranks_group): if i < src_replica_offset: continue j = (i - src_replica_offset) // pipe_map_interval if j == self.num_dst_pipeline_stage: src_replica_offset = i break if self.tp_num_mapping == 1: start = 0 else: mod_i = (i - src_replica_offset) % self.tp_num_mapping start = mod_i if (i - src_replica_offset) < self.tp_num_mapping else (self.tp_num_mapping - mod_i - 1) % self.tp_num_mapping for s_idx, src_rank in enumerate(src_tp_group): dst_rank, is_collide = self.get_load_balance_dst_rank( lb_dst_offset_pq_dict, s_idx, start, src_rank, dst_replica_ranks_group, j, pre_allocate=True ) if is_collide: add_recv_actor_stage1_fn(src_rank, dst_rank) pair_list.append((src_rank, dst_rank)) else: assert dst_rank is None uncollided_index_to_start_j.update({(i, s_idx) : (start, j)}) # Then, allocate src_ranks without gpu collisions for i, src_tp_group in enumerate(src_replica_ranks_group): for s_idx, src_rank in enumerate(src_tp_group): if (i, s_idx) not in uncollided_index_to_start_j: continue start, j = uncollided_index_to_start_j.get((i, s_idx)) dst_rank, _ = self.get_load_balance_dst_rank( lb_dst_offset_pq_dict, s_idx, start, src_rank, dst_replica_ranks_group, j, pre_allocate=False ) add_recv_actor_stage1_fn(src_rank, dst_rank) pair_list.append((src_rank, dst_rank)) # stage 2: comm pairs that broadcast params from first rank to the other ranks for each weight_mapping_group # Comm mapping in each weight_mapping_group of inference: # [0'] -> [1'] # [2'] -> [3'] recv_ranks = [pair[1] for pair in pair_list] def p2p_pair_grouping(tuples): for s_idx, src_rank in enumerate(tuples): for d_idx, dst_rank in enumerate(tuples): if s_idx == d_idx or src_rank not in recv_ranks: # pylint: disable=cell-var-from-loop continue add_recv_actor_stage2_fn(src_rank, dst_rank) p2p_list.append((src_rank, dst_rank)) for dst_tp_group in dst_replica_ranks_group: dst_tp_group = split_ranks_by_tp_and_ep_size(dst_tp_group, self.tp_num_mapping, 1) for tuples in dst_tp_group: p2p_pair_grouping(tuples) logger.info(f"comm pair_list : {pair_list}") logger.info(f"comm p2p_list : {p2p_list}") def _clear_sync_send_recv_parameters(self, rank_mappings:List): if len(rank_mappings) == 0: return refs = [] flagged_actors = set() for rank_mapping in rank_mappings: if len(rank_mapping) == 0: continue for send_actor, recv_actors in rank_mapping.items(): if send_actor not in flagged_actors: refs.append(send_actor.clear_sync_send_recv_parameters.remote()) flagged_actors.add(send_actor) for recv_actor in recv_actors: if recv_actor not in flagged_actors: refs.append(recv_actor.clear_sync_send_recv_parameters.remote()) flagged_actors.add(recv_actor) future.get(refs) def _clear_send_recv_param_names(self): self._send_recv_param_names = {} def _clear_sorted_send_actors(self, sorted_send_actors_list:List): if len(sorted_send_actors_list) == 0: return for sorted_send_actors in sorted_send_actors_list: if sorted_send_actors is not None: sorted_send_actors = None def clear_cache(self, sorted_send_actors_list=None, rank_mapping_list=None): if sorted_send_actors_list is None: sorted_send_actors_list = [ self.send_actors_to_regroup_routed_experts, self.sorted_send_actors, self.sorted_send_actors_stage2 ] if rank_mapping_list is None: rank_mapping_list = [self.send_recv_actor_mappings, self.send_recv_actor_mappings_stage2] self._clear_sync_send_recv_parameters(rank_mapping_list) self._clear_send_recv_param_names() self._clear_sorted_send_actors(sorted_send_actors_list) def validate_sync_results(self, send_actor, recv_actors, requires_grad, filter_fn=None, param_group="default"): assert param_group in ("default", "routed", "except_routed"), ( f"param_group must be one of 'default', 'routed', or 'except_routed', got {param_group}." ) def validate(): src_names, dst_names = self.set_sync_param_names(send_actor, recv_actors[0], requires_grad, filter_fn, param_group) # check the value of src model and tgt model pipe_stage = self.get_actor_pipe_rank(send_actor) res = [send_actor.reset_sync_parameters.remote(src_names, pipe_stage)] for recv_actor in recv_actors: res.append(recv_actor.reset_sync_parameters.remote(dst_names, pipe_stage)) future.wait(res) src_names, dst_names = future.get([send_actor.get_parameter_to_sync_names.remote(pipe_stage), recv_actors[0].get_parameter_to_sync_names.remote(pipe_stage)]) assert len(src_names) == len(dst_names), ( f"expect the length of src_names and dst_names being the same, got {len(src_names)} and {len(dst_names)}" ) # check the value of src model and tgt model names = list(zip(src_names, dst_names)) for src_name, dst_name in tqdm(names): if param_group in ("default", "except_routed"): src_tensor = future.get(send_actor.get_parameter_to_sync.remote(src_name, pipe_stage, True, self.tp_num_mapping > 1)) elif param_group == "routed": src_tensor = future.get(send_actor.get_parameter_to_sync.remote(src_name, pipe_stage, True)) if src_tensor.isnan().any(): raise RuntimeError(f"weight {src_name} from send actor is nan, please check checkpoint or training process.") src_tensor_shape = src_tensor.shape for recv_actor in recv_actors: dst_tensor = future.get(recv_actor.get_parameter_to_sync.remote(dst_name, pipe_stage, True)) if dst_tensor.isnan().any(): raise RuntimeError(f"weight {dst_name} in recv actor is nan, please check param sync.") if param_group in ("default", "except_routed"): if self.tp_num_mapping == 1: # for trainer_tp == inference_tp assert src_tensor.shape == dst_tensor.shape, ( f"after weight sync {src_name}: {src_tensor.shape} and {dst_name}: {dst_tensor.shape} do not match." ) assert torch.allclose(src_tensor, dst_tensor, atol=1e-06), ( f"after weight sync {src_name}: {src_tensor} and {dst_name}: {dst_tensor} do not match." ) else: # for inference_tp % trainer_tp == 0 and inference_tp > trainer_tp dst_tensor_shape = dst_tensor.shape src_tensor = src_tensor.reshape(-1) dst_tensor = dst_tensor.reshape(-1) tp_slice = self.actor2rank[recv_actor] % self.tp_num_mapping if src_tensor.shape == dst_tensor.shape: src_tensor_slice = src_tensor else: assert ( src_tensor.shape[0] % dst_tensor.shape[0] == 0 and src_tensor.shape[0] // dst_tensor.shape[0] == self.tp_num_mapping ), ( f"num of elements in src_tensor must be divided by that of dst_tensor. " f"while src {src_name}: {src_tensor_shape} and dst {dst_name}: {dst_tensor_shape}." ) start = dst_tensor.shape[0] * tp_slice end = start + dst_tensor.shape[0] src_tensor_slice = src_tensor[start:end] assert torch.allclose(src_tensor_slice, dst_tensor, atol=1e-06), ( f"after weight sync {src_name}_{tp_slice}: " f"{src_tensor_slice.view(dst_tensor_shape)} and {dst_name}: {dst_tensor.view(dst_tensor_shape)} do not match." ) elif param_group == "routed": assert self.hep_num_mapping == 1 assert src_tensor.shape == dst_tensor.shape, ( f"after weight sync {src_name}: {src_tensor.shape} and {dst_name}: {dst_tensor.shape} do not match." ) assert torch.allclose(src_tensor, dst_tensor, atol=1e-06), ( f"after weight sync {src_name}: {src_tensor} and {dst_name}: {dst_tensor} do not match." ) return True logger.info("Going to validate transmitted tensors...") validate() logger.info("Validation passed!") def set_sync_param_names_stage2(self, send_actor, recv_actor, to_rank, requires_grad, filter_fn=None, param_group="default"): send_names, _ = self.set_sync_param_names(send_actor, send_actor, requires_grad, filter_fn, param_group) refs = [] refs.append(send_actor.set_send_parameters.remote(send_names, self.get_actor_pipe_rank(send_actor))) refs.append(recv_actor.set_recv_parameters.remote(to_rank, send_names, self.get_actor_pipe_rank(recv_actor))) future.get(refs) return send_names, send_names def sync_broadcast_two_stage(self, actors, group_name, requires_grad=None, stage2=False, filter_fn=None, param_group="default"): send_actor = actors[0] start_time = time.time() stage_str = "STAGE1" if stage2 is False else "STAGE2" for rank, recv_actor in enumerate(actors[1:]): if stage2: self.set_sync_param_names_stage2(send_actor, recv_actor, self.actor2rank[recv_actor], requires_grad, filter_fn, param_group) else: self.set_sync_param_names(send_actor, recv_actor, requires_grad, filter_fn, param_group) pipe_stage = self.get_actor_pipe_rank(send_actor) shape_refs = [] shape_refs.append(send_actor.get_parameter_shape.remote(pipe_stage)) shape_refs.append(recv_actor.get_parameter_shape.remote(pipe_stage)) send_shape_list, recv_shape_list = future.get(shape_refs) buffer_num = {} tp_division = {} for send_name_and_shape, recv_name_and_shape in zip(send_shape_list, recv_shape_list): send_param_num = send_name_and_shape[1].numel() recv_param_num = recv_name_and_shape[1].numel() # for group query attention, tensor might consist of tp part and dp part. ele_buffer_num = 1 if send_param_num == recv_param_num else self.tp_num_mapping buffer_num[recv_name_and_shape[0]] = ele_buffer_num tp_division[send_name_and_shape[0]] = ele_buffer_num refs = [] refs.append(recv_actor.set_tp_num_mapping.remote(self.tp_num_mapping)) refs.append(recv_actor.set_buffer_num.remote(buffer_num)) refs.append(send_actor.set_tp_num_mapping.remote(self.tp_num_mapping)) refs.append(send_actor.set_tp_division.remote(tp_division)) future.get(refs) refs = [] pipe_stage = self.get_actor_pipe_rank(send_actor) send_rank = 0 if stage2: assert len(actors) == 2, f"expect only 2 actors for stage2. \ sync params of relative rank to other slices of inference model. while {len(actors)}" for rank, actor in enumerate(actors): sync_buffer_rank = self.actor2rank[actors[1]] if rank == 0 and stage2 else 0 ref = actor.broadcast_parameter_two_stage.remote( self.actor2rank[actor], sync_buffer_rank, rank, send_rank, group_name, pipe_stage, stage2) refs.append(ref) rets = future.wait(refs, return_output=True) logger.info(f"sync_broadcast_two_stage done {stage_str} {group_name} using {time.time()-start_time} seconds") return rets def sync_broadcast(self, actors, group_name, requires_grad=None, filter_fn=None, param_group="default"): send_actor = actors[0] for recv_actor in actors[1:]: self.set_sync_param_names(send_actor, recv_actor, requires_grad, filter_fn, param_group) pipe_stage = self.get_actor_pipe_rank(send_actor) refs = [] for rank, actor in enumerate(actors): ref = actor.broadcast_parameter.remote(rank, 0, group_name, pipe_stage) refs.append(ref) future.wait(refs, return_output=True) def sync_allgather(self, actors, group_name, requires_grad=None, filter_fn=None): # Currently, only routed experts are to be all-gathered. for actor in actors: self.set_sync_param_names(actor, actor, requires_grad, filter_fn, param_group="routed", should_map_name=False) pipe_stage = self.get_actor_pipe_rank(actors[0]) refs = [] for actor in actors: ref = actor.allgather_routed_expert_parameter.remote(group_name, pipe_stage) refs.append(ref) future.wait(refs, return_output=True) def sync_alltoall(self, actors, requires_grad=None, filter_fn=None): # Currently, only routed experts are to be synced with all-to-all. for actor in actors: self.set_sync_param_names(actor, actor, requires_grad, filter_fn, param_group="routed", should_map_name=False) pipe_stage = self.get_actor_pipe_rank(actors[0]) refs = [] logger.info(f"apply alltoall among {[self.actor2rank[actor] for actor in actors]}") for actor in actors: ref = actor.alltoall_routed_expert_parameter.remote(pipe_stage) refs.append(ref) future.wait(refs, return_output=True) def _sync_send_recv(self, send_actor, recv_actor, requires_grad=None, filter_fn=None, param_group="default"): self.set_sync_param_names(send_actor, recv_actor, requires_grad, filter_fn, param_group) pipe_stage = self.get_actor_pipe_rank(send_actor) is_the_same_gpu = self.is_same_gpu(send_actor, recv_actor) if is_the_same_gpu: name2ref = send_actor.ray_put_parameter.remote(self.group_name, pipe_stage) recv_ref = recv_actor.ray_get_parameter.remote(self.group_name, name2ref, pipe_stage) future.get(recv_ref) else: send_ref = send_actor.send_parameter.remote(self.actor2rank[recv_actor], self.group_name, pipe_stage) recv_ref = recv_actor.recv_parameter.remote(self.actor2rank[send_actor], self.group_name, pipe_stage) future.get([send_ref, recv_ref]) logger.debug(f"sync all parameters from {send_actor} to {recv_actor}") def sync_send_recv(self, send_actor, recv_actor, requires_grad=None, filter_fn=None, param_group="default"): try: self._sync_send_recv(send_actor, recv_actor, requires_grad, filter_fn, param_group) except Exception: future.get(self.error_signal.set.remote(traceback.format_exc())) def check_param_names(self, send_actor, recv_actor, src_names, dst_names): ref0 = send_actor.check_param_exists.remote(src_names) ref1 = recv_actor.check_param_exists.remote(dst_names) states = future.get([ref0, ref1]) if not states[0]: raise RuntimeError(f"Check src parameters to sync fail {src_names}") if not states[1]: raise RuntimeError(f"Check dst parameters to sync fail {dst_names}") def get_actor_pipe_rank(self, actor): def inner_func(): return future.get(actor.pipeline_parallel_rank.remote()) return utils.get_or_cache(self._actor2pipe, actor, inner_func) def get_actor_tp_rank(self, actor): def inner_func(): return future.get(actor.tensor_parallel_rank.remote()) return utils.get_or_cache(self._actor2tp, actor, inner_func) def get_actor_ep_rank(self, actor): def inner_func(): return future.get(actor.expert_parallel_rank.remote()) return utils.get_or_cache(self._actor2ep, actor, inner_func) def get_actor_dp_rank(self, actor): def inner_func(): return future.get(actor.get_data_parallel_rank.remote()) return utils.get_or_cache(self._actor2dp, actor, inner_func) def _set_sync_param_names(self, send_actor, recv_actor, requires_grad=None, filter_fn=None, param_group="default", should_map_name=True): if requires_grad is None: requires_grad = True if self._enable_lora: # TODO(jiangle.jl): support freeze layer. requires_grad = False assert param_group in ("default", "routed", "except_routed"), ( f"param_group must be one of 'default', 'routed', or 'except_routed', got {param_group}." ) if self.num_src_pipeline_stage > 1: dst_pipe_rank = self.get_actor_pipe_rank(recv_actor) dst_layer_offset = self.get_or_cache(recv_actor, "get_pipeline_stage_layer_offset") dst_src_mappings = future.get(send_actor.build_pipeline_layer_name_mapping.remote( self.num_dst_pipeline_stage, dst_pipe_rank, dst_layer_offset, requires_grad=requires_grad)) dst_names = list(dst_src_mappings.keys()) src_names = list(dst_src_mappings.values()) else: src_names = dst_names = future.get(send_actor.get_parameter_names.remote(requires_grad=requires_grad)) if self._enable_lora: src_names = [ele for ele in src_names if LORA_WEIGHT_PREFIX not in ele] dst_names = [ele for ele in dst_names if LORA_WEIGHT_PREFIX not in ele] if filter_fn is not None: src_names = filter_fn(src_names) dst_names = filter_fn(dst_names) synchronizer = get_synchronizer(self.src_model, self.dst_model) if should_map_name: src_names, dst_names = synchronizer.map_name_from_src_to_dst(send_actor, recv_actor, src_names, dst_names) else: # For routed experts which need to regroup expert first in trainer actors. synchronizer.map_name_from_src_to_dst(send_actor, recv_actor, src_names, dst_names) self.actor2synchronizer[send_actor] = synchronizer future.wait(send_actor.set_synchronizer.remote(synchronizer)) self.check_param_names(send_actor, recv_actor, src_names, dst_names) dst_model = self.actor2model[recv_actor] if self.tp_num_mapping > 1 and ((not dst_model.use_vllm_backend and param_group != "routed") or dst_model.use_vllm_backend): key = (recv_actor, recv_actor, param_group) if key not in self._send_recv_param_names: self._send_recv_param_names[key] = (dst_names, dst_names) else: dst_names0 = self._send_recv_param_names[key][0] dst_names0 += dst_names self._send_recv_param_names[key] = (dst_names0, dst_names0) if not self.synchronizer.is_parameter_changed: pipe_stage = self.get_actor_pipe_rank(send_actor) refs = [] refs.append(send_actor.set_sync_parameters.remote(src_names, pipe_stage)) refs.append(recv_actor.set_sync_parameters.remote(dst_names, pipe_stage)) future.get(refs) return src_names, dst_names def set_sync_param_names(self, send_actor, recv_actor, requires_grad=None, filter_fn=None, param_group="default", should_map_name=True): src_names, dst_names = utils.get_or_cache(self._send_recv_param_names, (send_actor, recv_actor, param_group), \ lambda: self._set_sync_param_names(send_actor, recv_actor, requires_grad, filter_fn, param_group, should_map_name)) logger.debug(f"{self.actor2rank[send_actor]} -> {self.actor2rank[recv_actor]}: {src_names[:5]} -> {dst_names[:5]}") pipe_stage = self.get_actor_pipe_rank(send_actor) refs = [] refs.append(send_actor.reset_sync_parameters.remote(src_names, pipe_stage)) refs.append(recv_actor.reset_sync_parameters.remote(dst_names, pipe_stage)) future.get(refs) return src_names, dst_names def create_broadcast_group(self, send_actor, recv_actors, group_name=None, param_group="default"): actor_groups = [send_actor] actor_groups.extend(recv_actors) # Use self.actor2rank to ensure a globally unique number within a param_group. send_actor_rank = self.actor2rank[send_actor] recv_actor_ranks = '_'.join([str(self.actor2rank[actor]) for actor in recv_actors]) # Always include self.group_name to ensure the name of a param_group is unique. if group_name is None: group_name = self.group_name elif not group_name.startswith(self.group_name + "_"): group_name = self.group_name + "_" + group_name finalized_group_name = f"{group_name}_{param_group}_from_{send_actor_rank}_to_{recv_actor_ranks}" logger.debug(f"finalized_group_name is {finalized_group_name}") logger.debug(f"current collevtive_groups is {self.collective_groups}") logger.debug(f"send_actor: {send_actor}, recv_actors: {recv_actors}") if finalized_group_name not in self.collective_groups: refs = [] for rank, actor in enumerate(actor_groups): ref = actor.setup_collective_group.remote(rank, len(actor_groups), "nccl", finalized_group_name) refs.append(ref) future.wait(refs) self.collective_groups.append(finalized_group_name) self.groups2actors[finalized_group_name] = tuple(actor_groups) return actor_groups, finalized_group_name def create_allgather_group(self, actor_groups, group_name=None): # Use self.actor2rank to ensure a globally unique number within a param_group. actor_ranks = '_'.join([str(self.actor2rank[actor]) for actor in actor_groups]) # Always include self.group_name to ensure the name of a param_group is unique. if group_name is None: group_name = self.group_name elif not group_name.startswith(self.group_name + "_"): group_name = self.group_name + "_" + group_name finalized_group_name = f"{group_name}_routed_among_{actor_ranks}" logger.debug(f"finalized_group_name is {finalized_group_name}") logger.debug(f"current collevtive_groups is {self.collective_groups}") if finalized_group_name not in self.collective_groups: refs = [] for rank, actor in enumerate(actor_groups): ref = actor.setup_collective_group.remote(rank, len(actor_groups), "nccl", finalized_group_name) refs.append(ref) future.wait(refs) self.collective_groups.append(finalized_group_name) return actor_groups, finalized_group_name def sort_send_actors(self, send_recv_actor_mappings, sorted_send_actors): if sorted_send_actors is not None: return sorted_send_actors dp2send_actors = defaultdict(list) for send_actor in send_recv_actor_mappings: dp2send_actors[self.get_actor_dp_rank(send_actor)].append(send_actor) for dp_rank in dp2send_actors: send_actors = dp2send_actors[dp_rank] dp2send_actors[dp_rank] = sorted(send_actors, key=lambda x: self.actor2rank[x]) sorted_send_actors = [] dp_rank = 0 while len(sorted_send_actors) < len(send_recv_actor_mappings): sorted_send_actors.append(dp2send_actors[dp_rank].pop(0)) dp_rank += 1 # dp_rank not in dp2send_actors happens when inference replica number less than training replica number if dp_rank == self.src_dp_size or dp_rank not in dp2send_actors: dp_rank = 0 assert len(send_recv_actor_mappings) == len(sorted_send_actors) return sorted_send_actors def sync_broadcast_second_stage_internal(self, group_name, thread_group, requires_grad=None, filter_fn=None, param_group="default", dryrun=False): max_workers = len(thread_group) logger.info(f"Use {max_workers} workers for second_stage_internal broadcasting.") with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [] for idx, actor_group in enumerate(thread_group): send_actor, recv_actor = actor_group group_name_with_idx = f"{group_name}_{idx}" actor_groups, finalized_group_name = self.create_broadcast_group( send_actor, [recv_actor], group_name=group_name_with_idx, param_group=param_group ) if dryrun: continue futures.append(executor.submit( self.sync_broadcast_two_stage, actor_groups, finalized_group_name, requires_grad, True, filter_fn, param_group)) for _future in concurrent.futures.as_completed(futures): try: _future.result() except Exception as e: traceback.print_exc() raise RuntimeError(f"Parameter sync thread generated an exception: {e}") # pylint: disable=raise-missing-from concurrent.futures.wait(futures) def sync_broadcast_second_stage(self, group_name, thread_groups, requires_grad=None, filter_fn=None, param_group="default", dryrun=False): tp_size = self.num_dst_tensor_parallel num_thread_groups = len(thread_groups) // tp_size new_thread_groups = [thread_groups[tp_size*i:tp_size*(i+1)] for i in range(num_thread_groups)] if not new_thread_groups: new_thread_groups = [thread_groups] max_workers = 1 logger.info(f"Use {max_workers} workers for second_stage broadcasting.") with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [] for idx, thread_group in enumerate(new_thread_groups): group_name_with_idx = f"{group_name}_{idx}" futures.append(executor.submit( self.sync_broadcast_second_stage_internal, group_name_with_idx, thread_group, requires_grad, filter_fn, param_group, dryrun)) for _future in concurrent.futures.as_completed(futures): try: _future.result() except Exception as e: traceback.print_exc() raise RuntimeError(f"Parameter sync thread generated an exception: {e}") # pylint: disable=raise-missing-from concurrent.futures.wait(futures) def sync_broadcast_multi_threads( self, sorted_send_actors, send_recv_actor_mappings, max_workers=1, requires_grad=None, group_name=None, stage2=False, filter_fn=None, param_group="default", dryrun=False): if stage2: thread_group = [] for send_actor in sorted_send_actors: recv_actors = send_recv_actor_mappings[send_actor] for recv_actor in recv_actors: thread_group.append((send_actor, recv_actor)) actor_groups_to_sync = [] for group in thread_group: new_actor_group_flag = True for idx, actor_groups in enumerate(actor_groups_to_sync): in_actor_group = False for actor_group in actor_groups: if group[0] in actor_group or group[1] in actor_group: in_actor_group = True if not in_actor_group: new_actor_group_flag = False actor_groups_to_sync[idx].append(group) #pylint: disable=unnecessary-list-index-lookup break if new_actor_group_flag or not actor_groups_to_sync: actor_groups_to_sync.append([group]) for group_idx, actor_groups in enumerate(actor_groups_to_sync): if self._comm_type == PARAM_SYNC_COMM_TYPE.BROADCAST: self.sync_broadcast_second_stage( f"{group_name}_{group_idx}", actor_groups, requires_grad, filter_fn, param_group, dryrun=dryrun ) else: raise RuntimeError("support p2p only for scenes that trainer_tp not equal to inference_tp.") else: max_workers = len(sorted_send_actors) logger.info(f"Use {max_workers} workers for first_stage broadcasting.") with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [] for send_actor in sorted_send_actors: recv_actors = send_recv_actor_mappings[send_actor] if self._comm_type == PARAM_SYNC_COMM_TYPE.BROADCAST: actor_groups, finalized_group_name = self.create_broadcast_group( send_actor, recv_actors, group_name=group_name, param_group=param_group ) if not dryrun: futures.append(executor.submit( self.sync_broadcast_two_stage, actor_groups, finalized_group_name, requires_grad, stage2, filter_fn, param_group )) else: raise RuntimeError("support p2p only for scenes that trainer_tp not equal to inference_tp.") for _future in concurrent.futures.as_completed(futures): try: _future.result() except Exception as e: traceback.print_exc() raise RuntimeError(f"Parameter sync thread generated an exception: {e}") # pylint: disable=raise-missing-from concurrent.futures.wait(futures) def sync_allgather_multi_threads( self, send_actors, max_workers=1, requires_grad=None, group_name=None, filter_fn=None ): send_actors_to_allgather_routed_experts = send_actors[0] logger.info(f"Use {max_workers} workers for allgather multiprocessing.") with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [] for allgather_actors in send_actors_to_allgather_routed_experts: actor_groups, finalized_group_name = self.create_allgather_group(allgather_actors, group_name=group_name) futures.append(executor.submit( self.sync_allgather, actor_groups, finalized_group_name, requires_grad, filter_fn=filter_fn )) for _future in concurrent.futures.as_completed(futures): try: _future.result() except Exception as e: raise RuntimeError(f"Parameter sync thread generated an exception: {e}") # pylint: disable=raise-missing-from concurrent.futures.wait(futures) def sync_alltoall_multi_threads( self, send_actors, max_workers=1, requires_grad=None, filter_fn=None ): send_actors_to_alltoall_routed_experts = send_actors[0] max_workers = len(send_actors_to_alltoall_routed_experts) logger.info(f"Use {max_workers} workers for alltoall multiprocessing.") with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [] for actor_groups in send_actors_to_alltoall_routed_experts: futures.append(executor.submit( self.sync_alltoall, actor_groups, requires_grad, filter_fn=filter_fn )) for _future in concurrent.futures.as_completed(futures): try: _future.result() except Exception as e: raise RuntimeError(f"Parameter sync thread generated an exception: {e}") # pylint: disable=raise-missing-from concurrent.futures.wait(futures) def check_and_setup_collective_group(self): if not self._is_collective_group_created: # Re-create collective group only when it is destroyed before. assert self._free_sync_collective_group self.setup_collective_group() def check_and_destroy_collective_group(self): if self._free_sync_collective_group: self.destroy_collective_group() self._is_collective_group_created = False self.collective_groups = [] self.groups2actors = {} def check_and_fuse_lora(self, enable_lora, actor_mapping): send_actors_set = set() def check_and_fuse_lora_internal(actor_mapping_item): for send_actor in actor_mapping_item: if enable_lora and send_actor not in send_actors_set: ref = send_actor.fuse_lora_layer.remote() state = future.get([ref]) assert state, "Check fuse lora layer fail." send_actors_set.add(send_actor) if isinstance(actor_mapping, List): for actor_mapping_item in actor_mapping: if actor_mapping_item is None: continue check_and_fuse_lora_internal(actor_mapping_item) elif isinstance(actor_mapping, Dict): if actor_mapping is None: return check_and_fuse_lora_internal(actor_mapping) else: raise ValueError("unrecognized type for actor_mapping, expect: List or Dict") def check_and_unfuse_lora(self, enable_lora, actor_mapping): send_actors_set = set() def check_and_unfuse_lora_internal(actor_mapping_item): for send_actor in actor_mapping_item: if self._enable_lora and send_actor not in send_actors_set: ref = send_actor.unfuse_lora_layer.remote() state = future.get([ref]) assert state, "Check unfuse lora layer fail." send_actors_set.add(send_actor) if isinstance(actor_mapping, List): for actor_mapping_item in actor_mapping: if actor_mapping_item is None: continue check_and_unfuse_lora_internal(actor_mapping_item) elif isinstance(actor_mapping, Dict): if actor_mapping is None: return check_and_unfuse_lora_internal(actor_mapping) else: raise ValueError("unrecognized type for actor_mapping, expect: List or Dict") def validate_sync_results_parallel(self, actor_mappings_list:List, requires_grad=None, validate=False, filter_fn=None, param_group="default"): if self._debug or validate: assert len(actor_mappings_list) in (1, 2), f"The length of actor mapping list should be 1 or 2, but got {len(actor_mappings_list)}." args = [] for send_actor, recv_actors in actor_mappings_list[0].items(): for recv_actor in recv_actors: if len(actor_mappings_list) == 1: args.append((send_actor, [recv_actor], requires_grad, filter_fn, param_group)) elif len(actor_mappings_list) == 2: recv_actors_stage2 = actor_mappings_list[1].get(recv_actor, []) args.append((send_actor, [recv_actor] + recv_actors_stage2, requires_grad, filter_fn, param_group)) if self._debug: for arg in args: self.validate_sync_results(arg[0], arg[1], arg[2], arg[3], arg[4]) else: execute_in_parallel(self.validate_sync_results, args) def _calculate_max_workers(self, sorted_send_actors, actor_mappings=None): max_workers = get_args().runtime_args.param_sync_max_workers if max_workers is None: max_workers = max(self.src_model.total_gpu // self.num_src_pipeline_stage, 1) if max_workers == -1: if self._comm_type == PARAM_SYNC_COMM_TYPE.BROADCAST: max_workers = len(sorted_send_actors) else: assert actor_mappings is not None, ( "actor_mappings should not be None when max_workers is -1 and " "communication type for parameter synchronization is not broadcast." ) max_workers = len(sorted_send_actors) * len(actor_mappings[sorted_send_actors[0]]) return max_workers def _multi_thread_sync_for_tp_num_mapping_gt_1( self, send_actors:List, actor_mappings:List, requires_grad=None, filter_fn=None, param_group="default", dryrun=False ): assert len(send_actors) == 2, ( f"Expect the length of send_actors being 2 for TP num mapping greater than 1, but got {len(send_actors)}." ) send_actors_stage1 = send_actors[0] # pylint: disable=unused-variable send_actors_stage2 = send_actors[1] # pylint: disable=unused-variable assert len(actor_mappings) == 2, ( f"Expect the length of actor_mappings being 2 for TP num mapping greater than 1, but got {len(actor_mappings)}." ) actor_mappings_stage1 = actor_mappings[0] actor_mappings_stage2 = actor_mappings[1] # stage 1 self.timers("stage1").start() sorted_send_actors_stage1 = list(actor_mappings_stage1.keys()) max_workers = self._calculate_max_workers(sorted_send_actors_stage1, actor_mappings_stage1) group_name = self.group_name + "_stage1_comm" self.sync_broadcast_multi_threads( sorted_send_actors_stage1, actor_mappings_stage1, max_workers, requires_grad, group_name=group_name, stage2=False, filter_fn=filter_fn, param_group=param_group, dryrun=dryrun ) self.timers("stage1").stop() logger.info(f"finish stage1| {self.timers.log(names=['stage1'])}") # stage 2 self.timers("stage2").start() sorted_send_actors_stage2 = list(actor_mappings_stage2.keys()) max_workers = self._calculate_max_workers(sorted_send_actors_stage2, actor_mappings_stage2) group_name = self.group_name + "_stage2_comm" self.sync_broadcast_multi_threads( sorted_send_actors_stage2, actor_mappings_stage2, max_workers, requires_grad, group_name=group_name, stage2=True, filter_fn=filter_fn, param_group=param_group, dryrun=dryrun) self.timers("stage2").stop() logger.info(f"finish stage2| {self.timers.log(names=['stage2'])}") def split_sync_groups(self, send_actors, actor_mappings): groups = [] for send_actor in send_actors: recv_actors = actor_mappings[send_actor] rank_dict = [self.actor2rank[actor] for actor in [send_actor] + recv_actors] # gen groups placed = False for group in groups: if set(group["values"]).isdisjoint(rank_dict): group["keys"].append(send_actor) group["values"] = group["values"] + rank_dict placed = True break if not placed: groups.append({ "keys": [send_actor], "values": rank_dict.copy() }) total_elements = sum(len(group["keys"]) for group in groups) assert total_elements == len(send_actors), \ (f"needed total elements of groups {total_elements} == len of send_actors \ {len(send_actors)} in param sync.") for group in groups: assert len(group["values"]) == len(set(group["values"])), \ (f"the elements must be all different in group: {group['values']}") logger.info(f"split_sync_groups: {groups}") return [g["keys"] for g in groups] def _multi_thread_sync_for_tp_num_mapping_eq_1( self, send_actors_list:List, actor_mappings_list:List, requires_grad=None, filter_fn=None, param_group="default", dryrun=False ): assert len(send_actors_list) == 1 and len(actor_mappings_list) == 1 send_actors = send_actors_list[0] actor_mappings = actor_mappings_list[0] sorted_send_actors = self.sort_send_actors(actor_mappings, send_actors) max_workers = self._calculate_max_workers(sorted_send_actors, actor_mappings) src_pp_size = self.num_src_pipeline_stage groups = self.split_sync_groups(sorted_send_actors, actor_mappings) logger.info(f"Use {max_workers} workers for tp_num_mapping_eq_1 synchoronization, \ src_pp_size: {src_pp_size}, groups: {len(groups)}.") with ThreadPoolExecutor(max_workers=max_workers) as executor: for group in groups: t1 = time.time() futures = [] for send_actor in group: recv_actors = actor_mappings[send_actor] logger.info(f"Sending from {[self.actor2rank[send_actor]]} to {[self.actor2rank[actor] for actor in recv_actors]}.") if self._comm_type == PARAM_SYNC_COMM_TYPE.BROADCAST: actor_groups, finalized_group_name = self.create_broadcast_group(send_actor, recv_actors, param_group=param_group) if dryrun: continue futures.append(executor.submit( self.sync_broadcast, actor_groups, finalized_group_name, requires_grad, filter_fn=filter_fn, param_group=param_group )) else: for recv_actor in recv_actors: if dryrun: continue futures.append(executor.submit( self.sync_send_recv, send_actor, recv_actor, requires_grad, filter_fn=filter_fn, param_group=param_group )) t2 = time.time() for _future in concurrent.futures.as_completed(futures): try: _future.result() except Exception as e: traceback.print_exc() raise RuntimeError(f"Parameter sync thread generated an exception: {e}") from e concurrent.futures.wait(futures) t3 = time.time() logger.info(f"sync for tp_num_mapping_eq_1, submit time(s):{(t2-t1)}, sync time(s):{(t3-t2)}") def _single_thread_sync(self, actor_mappings_list:List, requires_grad=None, filter_fn=None, param_group="default"): assert len(actor_mappings_list) == 1 actor_mappings = actor_mappings_list[0] for send_actor, recv_actors in actor_mappings.items(): if self._comm_type == PARAM_SYNC_COMM_TYPE.BROADCAST: actor_groups, finalized_group_name = self.create_broadcast_group(send_actor, recv_actors, param_group=param_group) self.sync_broadcast(actor_groups, finalized_group_name, requires_grad, filter_fn=filter_fn, param_group=param_group) else: for recv_actor in recv_actors: self.sync_send_recv(send_actor, recv_actor, requires_grad, filter_fn=filter_fn, param_group=param_group) def recover_synchronizer(self): refs = [] for actor, synchronizer in self.actor2synchronizer.items(): refs.append(actor.set_synchronizer.remote(synchronizer)) future.wait(refs) def reset_synchronizer(self): refs = [] for actor, _ in self.actor2synchronizer.items(): refs.append(actor.set_synchronizer.remote(None)) future.wait(refs) def sync(self, requires_grad=None, validate=False, dryrun=False): self.recover_synchronizer() self.check_and_setup_collective_group() self.check_and_fuse_lora(self._enable_lora, self.send_recv_actor_mappings) send_actors_list : List = [] actor_mappings_list : List = [] if self.concurrent_comm: if self.tp_num_mapping > 1: send_actors_list = [self.sorted_send_actors, self.sorted_send_actors_stage2] actor_mappings_list = [self.send_recv_actor_mappings, self.send_recv_actor_mappings_stage2] self._multi_thread_sync_for_tp_num_mapping_gt_1( send_actors_list, actor_mappings_list, requires_grad=requires_grad ) else: send_actors_list = [self.sorted_send_actors] actor_mappings_list = [self.send_recv_actor_mappings] self._multi_thread_sync_for_tp_num_mapping_eq_1( send_actors_list, actor_mappings_list, requires_grad=requires_grad ) else: actor_mappings_list = [self.send_recv_actor_mappings] self._single_thread_sync( actor_mappings_list, requires_grad=requires_grad ) assert len(actor_mappings_list) >= 1 self.check_and_unfuse_lora(self._enable_lora, self.send_recv_actor_mappings) self.validate_sync_results_parallel(actor_mappings_list, requires_grad, validate) self.check_and_destroy_collective_group() self.reset_synchronizer() logger.info(f"Group {self.group_name} sync all parameters done, comm_type {self._comm_type}") class ParameterSyncGroupwithHEP(ParameterSyncGroup): """ParameterSyncGroup for Hyper Expert Parallel (HEP). Note that in HEP, EP size for routed experts is different from that for Megatorn-LM. For routed experts, the new EP size (we call it HEP size for clarification) = mpu.ep_size x mpu.tp_size, while Megatron-LM set the EP size as mpu.ep_size. However, the EP size of shared experts in HEP is equal to that in Megatron -LM (which is 1). In this case, routed experts treat TP and EP altogether as EP, and shared experts ignore EP just like other non-expert weights. Therefore, we manage seperate parameter sync groups for routed expert weigts and weights except routed experts in this class. """ def __init__(self, src_model, dst_model, group_name, frequency, error_signal): self.send_recv_actor_mappings_for_routed_experts = defaultdict(list) self.recv_send_actor_mappings_for_routed_experts = defaultdict(list) self._num_src_hyper_expert_parallel = None self._num_dst_hyper_expert_parallel = None self._actor2hep = {} self.sorted_send_actors_for_routed_experts = None super().__init__(src_model, dst_model, group_name, frequency, error_signal) def setup_rank_mapping(self): """ For now, we only allow parameter sync between two models that dst model tp size >= src model tp size && dst model ep size <= src model ep size There is no more restiction regarding PP and DP """ self.tp_num_mapping = self.num_dst_tensor_parallel // self.num_src_tensor_parallel self.ep_num_mapping = self.num_dst_expert_parallel / self.num_src_expert_parallel self.hep_num_mapping = self.num_dst_hyper_expert_parallel / self.num_src_hyper_expert_parallel assert self.tp_num_mapping >= 1, ( f"Currently, tensor parallel world size for training ({self.num_src_tensor_parallel}) should be" f"less or equal to tensor parallel world size for inference ({self.num_dst_tensor_parallel}) with HEP enabled." ) assert self.ep_num_mapping <= 1, ( f"Currently, expert parallel world size for training ({self.num_src_expert_parallel}) should be" f"greater or equal to expert parallel world size for inference ({self.num_dst_expert_parallel}) with HEP enabled." ) if self.dst_model.use_vllm_backend: if self.tp_num_mapping == 1: if self.ep_num_mapping == 1: self.build_rank_mapping() else: self.build_rank_mapping_for_ep() elif self.tp_num_mapping > 1: if self.hep_num_mapping == 1: self.build_rank_mapping_for_ep(add_recv_actor_fn=self.add_recv_actor_for_routed_experts) # only add all-gather actors self.build_rank_mapping_for_params_except_routed_expert() else: self.build_rank_mapping_for_ep(add_recv_actor_fn=self.empty_add_recv_actor) # only add all-gather actors self.build_rank_mapping_two_stage() else: raise NotImplementedError( f"ChatLearn does not support synchronizing from larger tp size ({self.num_src_tensor_parallel})" f"to smaller tp size ({self.num_dst_tensor_parallel}) currently." ) else: if self.ep_num_mapping == 1 and self.tp_num_mapping == 1: self.build_rank_mapping() elif self.hep_num_mapping == 1: # In this case, routed experts are mapped one by one, while params except routed experts are split by TP. self.build_rank_mapping_for_routed_experts() self.build_rank_mapping_for_params_except_routed_expert() else: # We do not support other cases for HEP. Please note that tp_num_mapping > 1 with ep_num_mapping = 1 is also unsupported. raise NotImplementedError( "ChatLearn does not support inequivalent EP x TP between training and inference with Hyper Expert Parallel (HEP) enabled and " f"inference model is an instance of `MegatronModule`. Your current setting is " f"EP{self.num_src_expert_parallel} TP{self.num_src_tensor_parallel} for training model `{self.src_model.name}` " f"and EP{self.num_dst_expert_parallel} TP{self.num_dst_tensor_parallel} for inference model `{self.dst_model.name}`." ) def build_rank_mapping_for_ep(self, add_recv_actor_fn=None): # setup rank mapping for src parameter and dst parameter # get rank for one src_model, without model replicas if add_recv_actor_fn is None: add_recv_actor_fn = self.add_recv_actor src_dp_ranks, dst_dp_ranks = self.get_src_and_dst_dp_ranks() if self._debug and (src_dp_ranks[0] is None or dst_dp_ranks is None): return assert len(src_dp_ranks[0]) % len(dst_dp_ranks[0]) == 0, \ f"src training model ranks should be times of dst ranks, but got {len(src_dp_ranks[0])} and {len(dst_dp_ranks[0])}" if self.src_model.colocate_with(self.dst_model) and self.num_src_tensor_parallel % 2 == 1: replica_rank_iter = cycle(reversed(src_dp_ranks)) else: replica_rank_iter = cycle(iter(src_dp_ranks)) logger.debug(f"src_dp_ranks: {src_dp_ranks}") logger.debug(f"dst_dp_ranks: {dst_dp_ranks}") assert self.num_src_pipeline_stage % self.num_dst_pipeline_stage == 0 def split_ranks_by_ep_and_tp_size(ranks, tp_size : int = 1, ep_size : int = 1): tp_and_ep_size = tp_size * ep_size return [[ranks[i:i + tp_size] for i in range(j, j + tp_and_ep_size, tp_size)] for j in range(0, len(ranks), tp_and_ep_size)] src_replica_ranks2offset = {} is_first_time_set_send_actors = True for dst_replica_ranks in dst_dp_ranks: src_replica_ranks = next(replica_rank_iter) if tuple(src_replica_ranks) not in src_replica_ranks2offset: src_replica_ranks2offset[tuple(src_replica_ranks)] = 0 is_first_time_set_send_actors = True else: is_first_time_set_send_actors = False src_replica_ranks_group = split_ranks_by_ep_and_tp_size(src_replica_ranks, self.num_src_tensor_parallel, self.num_src_expert_parallel) # Since dst replica is vllm and it doesn't have ep, the function will organize dst_replica_ranks_group as [pp[tp]] naturally. dst_replica_ranks_group = split_ranks_by_ep_and_tp_size(dst_replica_ranks, self.num_dst_tensor_parallel, self.num_dst_expert_parallel) if is_first_time_set_send_actors: self.set_send_actors_to_regroup_routed_experts(src_replica_ranks_group) self.add_routed_experts_regrouping_actor(self.src_model, src_replica_ranks_group) if add_recv_actor_fn is self.empty_add_recv_actor: continue pipe_map_interval = self.num_src_pipeline_stage // self.num_dst_pipeline_stage for i, src_ep_and_tp_group in enumerate(src_replica_ranks_group): j = i // pipe_map_interval first_src_tp_group = src_ep_and_tp_group[0] assert len(dst_replica_ranks_group[j][0]) % len(first_src_tp_group) == 0, ( "TP size of dst model should be times of src model, " f"but got {len(dst_replica_ranks_group[j][0])} and {len(first_src_tp_group)}" ) len_dst_div_src = len(dst_replica_ranks_group[j][0]) // len(first_src_tp_group) concated_src_tp_group = [] offset = src_replica_ranks2offset[tuple(src_replica_ranks)] # cycled concatenate src tp group to ensure len(concat_src_tp_group) == len(dst_replica_ranks_group[j][0]) for k in range(len_dst_div_src): concated_src_tp_group.extend(src_ep_and_tp_group[int((offset + k) % len(src_ep_and_tp_group))]) for src_rank, dst_rank in zip(concated_src_tp_group, dst_replica_ranks_group[j][0]): add_recv_actor_fn(src_rank, dst_rank) src_replica_ranks2offset[tuple(src_replica_ranks)] = int( (src_replica_ranks2offset[tuple(src_replica_ranks)] + len_dst_div_src) % len(src_ep_and_tp_group) ) if self._debug: def debug_msg_for_actor_mappings(actor_mapping): if actor_mapping is None: return for k, v_list in actor_mapping.items(): for v in v_list: logger.debug(f"actor_mappings: {self.actor2rank[k]} -> {self.actor2rank[v]}") debug_msg_for_actor_mappings(self.send_recv_actor_mappings) debug_msg_for_actor_mappings(self.send_recv_actor_mappings_for_routed_experts) for regroup_actors in self.send_actors_to_regroup_routed_experts: count += 1 cat_str = "_".join(str(self.actor2rank[actor]) for actor in regroup_actors) logger.info(f"{self._comm_type_to_regroup_routed_experts} actors: {cat_str}") for k, v_list in self.send_recv_actor_mappings.items(): for v in v_list: logger.info(f"send_recv_actor_mappings: {self.actor2rank[k]} -> {self.actor2rank[v]}") def add_recv_actor_for_routed_experts(self, src_rank, dst_rank): src_actor = self.src_model.get_actor(src_rank) self.insert_actor2rank(src_actor, src_rank) self.insert_actor2model(src_actor, self.src_model) dst_actor = self.dst_model.get_actor(dst_rank) self.insert_actor2rank(dst_actor, dst_rank) self.insert_actor2model(dst_actor, self.dst_model) src_gpu = self.get_or_cache(src_actor, "get_visible_gpus") dst_gpu = self.get_or_cache(dst_actor, "get_visible_gpus") src_tp_rank = self.get_actor_tp_rank(src_actor) dst_tp_rank = self.get_actor_tp_rank(dst_actor) src_pp_rank = self.get_actor_pipe_rank(src_actor) dst_pp_rank = self.get_actor_pipe_rank(dst_actor) src_ep_rank = self.get_actor_ep_rank(src_actor) dst_ep_rank = self.get_actor_ep_rank(dst_actor) src_hep_rank = self.get_actor_hep_rank(src_actor) dst_hep_rank = self.get_actor_hep_rank(dst_actor) logger.debug(f"build rank mapping from {src_rank} to {dst_rank}, from gpu {src_gpu} to {dst_gpu}, " + f"from pipe_stage {src_pp_rank} to {dst_pp_rank}, " + f"from tp rank {src_tp_rank} to {dst_tp_rank}, " + f"from ep rank {src_ep_rank} to {dst_ep_rank}, " + f"from hep rank {src_hep_rank} to {dst_hep_rank}.") self.send_recv_actor_mappings_for_routed_experts[src_actor].append(dst_actor) self.recv_send_actor_mappings_for_routed_experts[dst_actor].append(src_actor) @property def num_src_hyper_expert_parallel(self): if self._num_src_hyper_expert_parallel is None: self._num_src_hyper_expert_parallel = future.get(self.src_model.replicas[0].all_actors[0].tensor_and_expert_model_parallel_size.remote()) return self._num_src_hyper_expert_parallel @property def num_dst_hyper_expert_parallel(self): if self._num_dst_hyper_expert_parallel is None: self._num_dst_hyper_expert_parallel = future.get(self.dst_model.replicas[0].all_actors[0].tensor_and_expert_model_parallel_size.remote()) return self._num_dst_hyper_expert_parallel def get_actor_hep_rank(self, actor): def inner_func(): return future.get(actor.tensor_and_expert_model_parallel_size.remote()) return utils.get_or_cache(self._actor2hep, actor, inner_func) def build_rank_mapping_for_routed_experts(self): self.build_rank_mapping(add_recv_actor_fn=self.add_recv_actor_for_routed_experts) def build_rank_mapping_for_params_except_routed_expert(self): self.build_rank_mapping_two_stage(add_recv_actor_fn=None) def routed_experts_filter(self, name_list: List[str]): filted_names = [name for name in name_list if 'mlp.experts' in name] return filted_names def params_except_routed_expert_filter(self, name_list: List[str]): filted_names = [name for name in name_list if 'mlp.experts' not in name] return filted_names def clear_cache(self, sorted_send_actors_list=None, rank_mapping_list=None): if sorted_send_actors_list is None: sorted_send_actors_list = [ self.sorted_send_actors, self.sorted_send_actors_stage2, self.send_actors_to_regroup_routed_experts, self.sorted_send_actors_for_routed_experts ] if rank_mapping_list is None: rank_mapping_list = [ self.send_recv_actor_mappings, self.send_recv_actor_mappings_stage2, self.send_recv_actor_mappings_for_routed_experts ] self._clear_sync_send_recv_parameters(rank_mapping_list) self._clear_send_recv_param_names() self._clear_sorted_send_actors(sorted_send_actors_list) def warmup_groups(self): def warmup_tasks_func(task): actors = task.actors group = task.group refs = [] refs.append(actors[0].broadcast_dummy_tensor_send.remote(0, group)) for actor in actors[1:]: refs.append(actor.broadcast_dummy_tensor_recv.remote(0, group)) future.wait(refs) tasks = [] actors_set = set() for group_name, actors in self.groups2actors.items(): # filter actors if the same collective ring actor_ids = [self.actor2rank[actor] for actor in actors] key = tuple(sorted(actor_ids)) if key not in actors_set: tasks.append(CollectiveTask(actors, group_name)) parallel_execute_collective_tasks(tasks, warmup_tasks_func) def _synchronize_all_moe_parameters(self, requires_grad=None, validate=False, dryrun=False): self.check_and_setup_collective_group() send_actors_list : List = [ self.sorted_send_actors, self.sorted_send_actors_stage2 ] actor_mappings_list : List = [ self.send_recv_actor_mappings, self.send_recv_actor_mappings_stage2, self.send_actors_to_regroup_routed_experts, ] self.check_and_fuse_lora(self._enable_lora, actor_mappings_list) if self.concurrent_comm: assert self.dst_model.use_vllm_backend max_workers = self._calculate_max_workers(self.send_actors_to_regroup_routed_experts) if self._comm_type_to_regroup_routed_experts == ROUTED_EXPERT_REGROUPING_COMM_TYPE.ALLGATHER: # allgather routed experts only self.sync_allgather_multi_threads( [self.send_actors_to_regroup_routed_experts], max_workers=max_workers, requires_grad=requires_grad, group_name=self.group_name + "_allgather", filter_fn=self.routed_experts_filter) elif self._comm_type_to_regroup_routed_experts == ROUTED_EXPERT_REGROUPING_COMM_TYPE.ALLTOALL: if not dryrun: logger.info("start to alltoall router experts. ") start_time = time.time() # alltoall routed experts only self.sync_alltoall_multi_threads( [self.send_actors_to_regroup_routed_experts], max_workers=max_workers, requires_grad=requires_grad, filter_fn=self.routed_experts_filter) logger.info("complete to alltoall router experts using {time.time()-start_time:.2f} seconds ") # sync everything to inference model if self.tp_num_mapping == 1: logger.info("start to sync all moe experts") send_actors_list = [self.sorted_send_actors] actor_mappings_list = [self.send_recv_actor_mappings] self._multi_thread_sync_for_tp_num_mapping_eq_1( send_actors_list, actor_mappings_list, requires_grad=requires_grad, filter_fn=None, param_group="default", dryrun=dryrun ) logger.info("complete to sync all moe experts") elif self.tp_num_mapping > 1: # First, synchronize routed experts. logger.info("start to sync routed expert weights.") start_time = time.time() self._synchronize_routed_experts(requires_grad=requires_grad, validate=validate, dryrun=dryrun) logger.info(f"complete to sync routed expert weights. [stage1-1] using {time.time()-start_time:.2f} seconds") self.clear_cache( sorted_send_actors_list = [ self.send_actors_to_regroup_routed_experts, self.sorted_send_actors_for_routed_experts ], rank_mapping_list=[ self.send_recv_actor_mappings_for_routed_experts ] ) # Then, synchronize parameters except routed experts logger.info("start to sync parameters except routed eperts.") self._synchronize_params_except_routed_experts(requires_grad=requires_grad, validate=validate, dryrun=dryrun) logger.info("complete to sync parameters except routed experts.") self.reset_synchronizer() self.clear_cache( sorted_send_actors_list = [ self.sorted_send_actors, self.sorted_send_actors_stage2, ], rank_mapping_list = [ self.send_recv_actor_mappings, self.send_recv_actor_mappings_stage2 ] ) else: raise NotImplementedError( f"ChatLearn does not support synchronizing from larger tp size ({self.num_src_tensor_parallel})" f"to smaller tp size ({self.num_dst_tensor_parallel}) currently." ) else: raise NotImplementedError( "ChatLearn supports only concurrent_comm for training models with HEP enabled and inference with vLLM" ) self.check_and_unfuse_lora(self._enable_lora, actor_mappings_list) self.validate_sync_results_parallel(actor_mappings_list, requires_grad, validate) self.check_and_destroy_collective_group() self.reset_synchronizer() logger.info(f"Group {self.group_name} sync all parameters done, comm_type {self._comm_type}") def _synchronize_routed_experts(self, requires_grad=None, validate=False, dryrun=False): self.check_and_setup_collective_group() self.check_and_fuse_lora(self._enable_lora, self.send_recv_actor_mappings_for_routed_experts) send_actors_list : List = [] actor_mappings_list : List = [] if self.concurrent_comm: send_actors_list = [self.sorted_send_actors_for_routed_experts] actor_mappings_list = [self.send_recv_actor_mappings_for_routed_experts] self._multi_thread_sync_for_tp_num_mapping_eq_1( send_actors_list, actor_mappings_list, requires_grad=requires_grad, filter_fn=self.routed_experts_filter, param_group="routed", dryrun=dryrun, ) else: actor_mappings_list = [self.send_recv_actor_mappings_for_routed_experts] self._single_thread_sync( self.send_recv_actor_mappings_for_routed_experts, requires_grad=requires_grad, filter_fn=self.routed_experts_filter, param_group="routed", ) assert len(actor_mappings_list) >= 1 self.check_and_unfuse_lora(self._enable_lora, self.send_recv_actor_mappings_for_routed_experts) self.validate_sync_results_parallel( actor_mappings_list, requires_grad, validate, filter_fn=self.routed_experts_filter, param_group="routed" ) self.check_and_destroy_collective_group() logger.info(f"Group {self.group_name} sync all parameters done, comm_type {self._comm_type}") def _synchronize_params_except_routed_experts(self, requires_grad=None, validate=False, dryrun=False): self.check_and_setup_collective_group() self.check_and_fuse_lora(self._enable_lora, self.send_recv_actor_mappings) send_actors_list : List = [] actor_mappings_list : List = [] if self.concurrent_comm: if self.tp_num_mapping > 1: send_actors_list = [self.sorted_send_actors, self.sorted_send_actors_stage2] actor_mappings_list = [self.send_recv_actor_mappings, self.send_recv_actor_mappings_stage2] self._multi_thread_sync_for_tp_num_mapping_gt_1( send_actors_list, actor_mappings_list, requires_grad=requires_grad, filter_fn=self.params_except_routed_expert_filter, param_group="except_routed", dryrun=dryrun ) else: send_actors_list = [self.sorted_send_actors] actor_mappings_list = [self.send_recv_actor_mappings] self._multi_thread_sync_for_tp_num_mapping_eq_1( send_actors_list, actor_mappings_list, requires_grad=requires_grad, filter_fn=self.params_except_routed_expert_filter, param_group="except_routed" ) else: actor_mappings_list = [self.send_recv_actor_mappings] self._single_thread_sync( actor_mappings_list, requires_grad=requires_grad, filter_fn=self.params_except_routed_expert_filter, param_group="except_routed" ) self.check_and_unfuse_lora(self._enable_lora, self.send_recv_actor_mappings) self.validate_sync_results_parallel( actor_mappings_list, requires_grad, validate, filter_fn=self.params_except_routed_expert_filter, param_group="except_routed" ) self.check_and_destroy_collective_group() logger.info(f"Group {self.group_name} sync all parameters done, comm_type {self._comm_type}") def sync(self, requires_grad=None, validate=False, dryrun=False): if self.dst_model.use_vllm_backend: self.recover_synchronizer() self._synchronize_all_moe_parameters(requires_grad=requires_grad, validate=validate, dryrun=dryrun) else: if self.ep_num_mapping == 1 and self.tp_num_mapping == 1: # synchronization is the same as base class when applying Qwen + Qwen super().sync(requires_grad, validate) return self.recover_synchronizer() # First, synchronize routed experts. self._synchronize_routed_experts(requires_grad=requires_grad, validate=validate) self.clear_cache( sorted_send_actors_list = [ self.send_actors_to_regroup_routed_experts, self.sorted_send_actors_for_routed_experts ], rank_mapping_list=[ self.send_recv_actor_mappings_for_routed_experts ] ) # Then, synchronize parameters except routed experts self._synchronize_params_except_routed_experts(requires_grad=requires_grad, validate=validate) self.reset_synchronizer() self.clear_cache( sorted_send_actors_list = [ self.sorted_send_actors, self.sorted_send_actors_stage2, ], rank_mapping_list = [ self.send_recv_actor_mappings, self.send_recv_actor_mappings_stage2 ] )