# 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. # ============================================================================== """Executor""" import threading from collections import defaultdict from itertools import cycle from ray.util.queue import Queue import torch from chatlearn.models.vllm_module_v2 import VLLMModuleV2 from chatlearn.runtime.model_flow import ModelFlow from chatlearn.utils import future from chatlearn.utils.constant import CHATLEARN_REGROUP_TAG, INDEX_TAG from chatlearn.utils.constant import LOG_START from chatlearn.utils.global_vars import get_args from chatlearn.utils.logger import logger from .utils import encode_data, decode_data from .utils import FlowParser def split_list(lst, n): assert len(lst) % n == 0, f"{len(lst)} % {n} != 0" k = len(lst) // n return [lst[i*k:(i+1)*k] for i in range(n)] def split_along_batch(tensors, num_splits): res = [{} for _ in range(num_splits)] if tensors is None: return res for key in tensors.keys(): to_batch = tensors[key] if isinstance(to_batch, torch.Tensor): batched = to_batch.chunk(num_splits) elif isinstance(to_batch, list): batched = split_list(to_batch, num_splits) else: raise Exception(f"unknown types key: {key} and {type(to_batch)} to split: {key} {tensors.keys()} {to_batch}") for idx, ele in enumerate(batched): res[idx][key] = ele return res # pylint: disable=not-callable class Executor: """Executor""" def __init__(self, model_flow): """ Executor Args ---- models : List[BaseModule] a list of modules """ self._set_flow(model_flow) self.args = get_args().runtime_args self.model_flow = None self.local_models = self.models self._batch_per_episode = -1 self.is_eval = False self._timers = None self.model2iter = {} self.merged_buffer = defaultdict(dict) self._metric_list = [] def set_timers(self, _timers): self._timers = _timers @property def timers(self): return self._timers def _set_flow(self, flow): """ Set compution flow Args ---- flow : callable a function that defines model computation flow Returns ------- Executor return self """ self._flow = flow self.model_to_call_funcs = FlowParser().parse(flow) for model, func_names in self.model_to_call_funcs.items(): model.call_funcs += func_names self.models = list(self.model_to_call_funcs.keys()) return self @property def first_node(self): return self.model_flow.model_nodes[0] @property def first_model(self): return self.first_node.model def update_models(self, models): # update local model with remote models new_models = [] name_to_new_models = {model.name: model for model in models} for model in self.local_models: dist_model = name_to_new_models[model.name] dist_model.group_dist_actors_by_tp_rank() new_models.append(dist_model) self.models = new_models if self.args is None: self.args = get_args().runtime_args def setup(self): self._models_and_results_to_wait = [] self.model_flow = ModelFlow(self) self.model_flow.trace(self.models, self._flow) self.models = [model_node.model for model_node in self.model_flow.model_nodes] self.model_locks = {model_node: threading.Lock() for model_node in self.model_flow.model_nodes} def _next_model(self, model): if len(model.replicas) == 1: return model.replicas[0] if model not in self.model2iter: self.model2iter[model] = cycle(iter(model.replicas)) return next(self.model2iter[model]) def get_merged_data(self, queues, encode=True, micro_batch_index=None, model_node=None, trainable=False): mb0 = None if micro_batch_index is not None: mb0 = micro_batch_index data_list = [None] * len(queues) merged_buffer = self.merged_buffer[model_node] for index, queue in enumerate(queues): if index not in merged_buffer: merged_buffer[index] = {} if mb0 in merged_buffer[index]: data_list[index] = merged_buffer[index].pop(mb0) continue while True: flag = False while queue.qsize() == 0: if mb0 in merged_buffer[index]: data_list[index] = merged_buffer[index].pop(mb0) flag = True break if flag: break encoded_data = queue.get() mb, data = decode_data(encoded_data) if mb0 is None: mb0 = mb if isinstance(data, list) and not trainable: data = data[-1] if mb == mb0: data_list[index] = data break merged_buffer[index][mb] = data if encode: return encode_data(mb0, data_list) return data_list def get_merged_data_locked(self, queues, encode=True, micro_batch_index=None, model_node=None, trainable=False): with self.model_locks[model_node]: return self.get_merged_data(queues, encode, micro_batch_index, model_node, trainable) @staticmethod def align_out_queues(queues, encode=False): # TODO: deal with one2many scene out_queues = [] min_qsize = min([ele.qsize() for ele in queues]) # pylint: disable=consider-using-generator for queue in queues: num_producers = queue.qsize() if num_producers == min_qsize: out_queues.append(queue) continue assert num_producers % min_qsize == 0 out_queue = Queue() res_list = [] while queue.qsize() > 0: res = queue.get() res = decode_data(res)[1] if encode else res res_list.append(res) division = num_producers // min_qsize in_qsize = len(res_list) out_qsize = in_qsize // division for q_idx in range(out_qsize): start = q_idx * division end = start + division out_queue.put(encode_data(q_idx, {CHATLEARN_REGROUP_TAG:res_list[start:end]})) out_queues.append(out_queue) return out_queues def get_all_merged_data(self, queues, out_queue, encode=True): logger.info(f"{LOG_START} start to align output queues with sizes {[ele.qsize() for ele in queues]}.") queues = self.align_out_queues(queues, True) logger.info(f"{LOG_START} complete to align output queues, sizes of output_queues are {[ele.qsize() for ele in queues]}.") queue0 = queues[0] while queue0.qsize() > 0: res = self.get_merged_data(queues, encode) out_queue.put(res) def rebatch_all_merged_data(self, model_node, in_queues, is_eval=False):# pylint: disable=unused-argument if not model_node.input_nodes: return in_queues out_queues = [None] * len(in_queues) num_consumers = self.batch_per_episode(model_node.model) for index, (input_node, in_queue) in enumerate(zip(model_node.input_nodes, in_queues)): num_producers = self.batch_per_episode(input_node.model) if num_producers == num_consumers: out_queues[index] = in_queue else: out_queues[index] = Queue() res_list = [] while in_queue.qsize() > 0: res = in_queue.get() res = decode_data(res)[1] res_list.append(res) if num_producers > num_consumers: # Deal with the case where num_producers > num_consumers assert num_producers % num_consumers == 0, \ f"many2one: num_producers: {num_producers}, num_consumers: {num_consumers}, len inqueue: {len(in_queues)}" division = num_producers // num_consumers in_qsize = len(res_list) out_qsize = in_qsize // division for q_idx in range(out_qsize): start = q_idx * division end = start + division out_queues[index].put(encode_data(q_idx, {CHATLEARN_REGROUP_TAG:res_list[start:end]})) else: # Deal with the case where num_producers < num_consumers # TODO: add index for one2many case assert num_consumers % num_producers == 0, \ f"one2many: num_producers: {num_producers}, num_consumers: {num_consumers}, len inqueue: {len(in_queues)}" division = num_consumers // num_producers in_qsize = len(res_list) out_qsize = in_qsize * division for q_idx in range(out_qsize): start = q_idx // division end = start + 1 out_queues[index].put(encode_data(q_idx, {CHATLEARN_REGROUP_TAG: res_list[start:end], INDEX_TAG: (q_idx % division, division)})) return out_queues def get_next_data(self, in_queue, model_node, micro_batch_index): if isinstance(in_queue, list): if len(in_queue) > 0: # this should happen for inference models, will trigger bug for training models # since training models accept a list of remote object, which has the same # behavior for models accept multiple inputs # we need to deal with it later assert not model_node.trainable data = self.get_merged_data_locked(in_queue, micro_batch_index=micro_batch_index, model_node=model_node, trainable=model_node.trainable) mb, query = decode_data(data) else: mb, query = micro_batch_index, [] else: data = self.get_merged_data_locked([in_queue], micro_batch_index=micro_batch_index, model_node=model_node, trainable=model_node.trainable) assert len(data['data']) == 1 data['data'] = data['data'][0] mb, query = decode_data(data) query = [query] return mb, query def generate_step_one_model_internal(self, model_node, in_queue, step_num, replica, func_name="forward_step", to_empty_cache=None, is_eval=False, to_onload=None, to_offload=None, micro_batch_index=None): """ Args: model: DistModel in_queue: Queue step_num: int replica: current model replica of DistModel func_name: str to_empty_cache: None or boolean """ model = model_node.model kwargs = {} replica_num = len(model.replicas) output = [] if isinstance(replica.model, VLLMModuleV2): last_step_start = max(self.num_iteration(model) - replica_num, 0) is_last_batch = step_num >= last_step_start kwargs["is_last_batch"] = is_last_batch if is_eval is not None: kwargs["is_eval"] = is_eval if to_empty_cache is not None: kwargs["to_empty_cache"] = to_empty_cache if to_onload is not None: kwargs["to_onload"] = to_onload if to_offload is not None: kwargs["to_offload"] = to_offload mb, query = self.get_next_data(in_queue, model_node, micro_batch_index) assert isinstance(query, list) ret = replica.call_actor_remote_func(replica.vllm_engine, func_name, *query, **kwargs) output.append((ret, mb)) else: last_step_start = max(self.num_iteration(model) - replica_num, 0) is_last_batch = step_num >= last_step_start kwargs["is_last_batch"] = is_last_batch if to_empty_cache is not None: kwargs["to_empty_cache"] = to_empty_cache if to_onload is not None: kwargs["to_onload"] = to_onload if to_offload is not None: kwargs["to_offload"] = to_offload if is_eval is not None: kwargs["is_eval"] = is_eval for _, actors in replica.dp_rank_to_actors.items(): mb, query = self.get_next_data(in_queue, model_node, micro_batch_index) assert isinstance(query, list) for actor in actors: ret = replica.call_actor_remote_func(actor, func_name, *query, **kwargs) output.append((ret, mb)) return output def generate_step_one_model(self, model_node, replica, in_queue, out_queue, step_num, func_name="forward_step", to_empty_cache=None, is_eval=False, to_onload=None, to_offload=None, micro_batch_index=None): """ Args: model: DistModel in_queue: Queue out_queue: Queue step_num: int func_name: str to_empty_cache: None or boolean """ model = model_node.model # output is a list of tuple, each tuple is (remote_refs, mb) output = self.generate_step_one_model_internal(model_node, in_queue, step_num, replica, func_name, to_empty_cache, is_eval, to_onload, to_offload, micro_batch_index) num_dp_rank = len(replica.dp_rank_to_actors) if model.module_args.zero_size == 1: # If (tp > 1 or pp > 1) and ep = 1 for current model, its `output` will be a list whose # length is the number of Actors. In this case, all members in the list # are the same, and we choose output[-1] to put into out_queue. # If (tp > 1 or pp > 1) and ep > 1, we choose last output for each dp rank to put into # out_queue. if model.module_args.expert_model_parallel_size == 1 and num_dp_rank == 1: result = [output[-1]] else: num_output = len(output) assert num_output % num_dp_rank == 0, ( f"The number of outputs ({num_output}) must be divisible by " f"the number of dp_ranks ({num_dp_rank}) in a replica." ) interval = num_output // num_dp_rank result = [output[i] for i in range(interval - 1, num_output, interval)] else: result = output if isinstance(out_queue, list): for oq in out_queue: for res, mb in result: oq.put(encode_data(mb, res)) else: for res, mb in result: out_queue.put(encode_data(mb, res)) # To ensure all Actors are finished synchronously, all remote refs should be returned # note that ray wait does not support tuple type, return a list of list remote_refs = [item[0] for item in output] return out_queue, remote_refs def regroup_inqueue(self, model_node, queues, is_eval=False): if self.args.policy_to_regroup_queue == "global_barrier": # barrier to regroup all queues of producer node if not isinstance(queues, list): queues = [queues] logger.info(f"{LOG_START} regroup_inqueue in_queue {model_node}: {[ele.qsize() for ele in queues]}") out_queues = self.rebatch_all_merged_data(model_node, queues, is_eval=is_eval) logger.info(f"{LOG_START} regroup_inqueue out_queues {model_node}: {[ele.qsize() for ele in out_queues]}") return out_queues else: raise RuntimeError(f"Unsupported policy_to_regroup_queue {self.args.policy_to_regroup_queue}.") def compute_loop_one_model(self, model_node, num_batch=None): logger.info(f"{LOG_START} start compute_loop for {model_node}, is_eval={self.is_eval}") model = model_node.model is_eval = self.is_eval if num_batch is None: num_batch = self.num_iteration(model) func_name = model_node.func_name if model_node.remote_objects_to_wait: logger.info(f"{LOG_START} start to wait colocate models to finish for {model_node}") model_node.wait_colocate_models_to_finish(self.timers, func_name) logger.info(f"{LOG_START} complete to wait colocate models to finish for {model_node}") replica_num = len(model.replicas) last_step_start = max(num_batch - replica_num, 0) in_queue = model_node.get_input_queues() logger.info(f"{LOG_START} start to regroup in_queue for {model_node}") in_queue = self.regroup_inqueue(model_node, in_queue, is_eval=is_eval) logger.info(f"{LOG_START} complete to regroup in_queue for {model_node}") if isinstance(in_queue, list) and len(in_queue) == 1: in_queue = in_queue[0] results = [] logger.info(f"{LOG_START} start to generate_step_one_model for {model_node}") self.timers(f"{model.name}").start() for step in range(num_batch): to_empty_cache = step >= last_step_start and (model.is_colocate or model.module_args.force_free_memory) to_onload = step < replica_num and ((model.is_colocate and model.enable_offload) or model.module_args.force_free_memory) to_offload = step >= last_step_start and ((model.is_colocate and model.enable_offload) or model.module_args.force_free_memory) replica = self._next_model(model) _, data = self.generate_step_one_model(model_node, replica, in_queue, model_node.out_queues, step, func_name, to_empty_cache, is_eval=is_eval, to_onload=to_onload, to_offload=to_offload) results.append(data) self.timers(f"{model.name}").stop() if model_node.next_colocate_node: # before the execution of next colocate model, perform the wait, since we want to empty the cache. logger.info( f"{LOG_START} Model {model_node.next_colocate_node} will wait model {model} to finish since they are colocated") self._models_and_results_to_wait = model_node.next_colocate_node.add_dependent_colocate_model_results( model_node, results, self._models_and_results_to_wait) elif model.colocate_models or model.trainable: # 1. the model may colocate with training/inference, so we should wait until the end of compute_loop # 2. the model is trainable and it does not have next_colocate_model, we should make sure it is finished before parameter_sync # so we add them to a temp list logger.info(f"{LOG_START} Sync {model} in the end of {self.__class__.__name__}") self._models_and_results_to_wait.append((model_node, results)) def compute_loop(self, out_queue, num_batch=None): for model_group in self.model_flow.flow_topology: for model_node in model_group: self.compute_loop_one_model(model_node, num_batch) data = [None] * len(self.model_flow.return_model_nodes) for model_node in self.model_flow.model_nodes: self.timers(f"{model_node.model.name}").start() if model_node in self.model_flow.return_model_nodes: # let the results order follow model_node order data[self.model_flow.return_model_nodes.index(model_node)] = model_node.out_queues[-1] self.timers(f"{model_node.model.name}").stop() model_names = [] results = [] for model, result in self._models_and_results_to_wait: model_names.append(model.name) results.extend(result) if results: for model_name in model_names: self.timers(f"{model_name}").start() func_name = self.model_flow.model_nodes[0].func_name future.wait(results, f"{model_names} {func_name}") for model_name in model_names: self.timers(f"{model_name}").stop() self._models_and_results_to_wait = [] if data: self.get_all_merged_data(data, out_queue, encode=False) def setup_queues(self): data_queues = [] out_queue = Queue() for model_node in self.model_flow.input_consumers: data_queue = Queue() data_queues.append(data_queue) model_node.set_input_queue(data_queue) for model_node in self.model_flow.model_nodes: num_out_queue = len(model_node.output_nodes) if model_node in self.model_flow.return_model_nodes: num_out_queue += 1 model_node.set_out_queues([Queue() for _ in range(num_out_queue)]) return data_queues, out_queue # pylint: disable=not-callable