# Copyright (c) 2021 PaddlePaddle Authors. 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. import copy import queue from enum import Enum import numpy as np import paddle from paddle.distributed.fleet.meta_optimizers.common import OpRole from paddle.framework import core SUCC = 0 # successor PRED = 1 # predecessor class CostNodeType(Enum): DEFAULT = 0 COMPUTATION = 1 COMMUNICATION = 2 VARIABLE = 3 MERGED = 4 NOP = 5 class Cost: def __init__(self): self.runtime = None self.static_mem = None self.peak_mem = None class CostModelMode(Enum): DEFAULT = 0 BENCHMARKING = 1 # costs based on trial runs ANALYSIS = 2 # costs based on analysis MIXED = 3 class CostNode: def __init__(self, node, node_type, id=None): self.id = id self.node = node self.type = node_type self._cost = 0 self.is_optim = False self.is_bwd = False @property def cost(self): return self._cost @cost.setter def cost(self, cost): if cost < 0: raise ValueError('Cost must be above 0.') self._cost = cost class MergedOpsCostNode(CostNode): def __init__(self, node_type, id=None, base_node_list=None, is_bwd=False): super().__init__(None, node_type, id) self.node_list = base_node_list self.is_bwd = is_bwd class CommOpCostNode(CostNode): def __init__( self, node, node_type, id=None, comm_node_list=None, is_bwd=False ): super().__init__(node, node_type, id) self.node_list = comm_node_list self.ranks = [] self.comm_type = node.type self.is_bwd = is_bwd def set_ranks(self, ranks): self.ranks = ranks def set_shapes(self, input_shape, output_shape): self.input_shape = input_shape self.output_shape = output_shape def init_comm_cost(self, cluster=None): # ref: https://github.com/NVIDIA/nccl-tests/blob/master/doc/PERFORMANCE.md # should get from `cluster` BANDWIDTH = 32 * 1024 / 1000 # MB/ms, V100 PCIe num_ranks = len(self.ranks) comm_volume = np.prod(self.input_shape) * 4 if 'allreduce' in self.comm_type: self._cost = comm_volume / ( BANDWIDTH * num_ranks / (2 * (num_ranks - 1)) ) elif 'gather' in self.comm_type: self._cost = comm_volume / (BANDWIDTH * num_ranks / (num_ranks - 1)) elif 'broadcast' in self.comm_type: self._cost = comm_volume / BANDWIDTH elif 'send' in self.comm_type or 'recv' in self.comm_type: self._cost = comm_volume / BANDWIDTH else: self._cost = 0 class TensorCostNode(CostNode): def __init__( self, node, node_type, id=None, base_node_list=None, batch_size=None, shared_node_id=None, ): super().__init__(node, node_type, id) if node.name == "create_py_reader_0" or node.name == "double_buffer_0": self.shape = [2, 2] self.dtype = paddle.float32 else: self.shape = node.shape self.dtype = node.dtype self.dtype_factor = 1 self.persistable = None self.shared_node_id = shared_node_id if self.dtype == paddle.float32 or node.dtype == paddle.int32: self.dtype_factor *= 4 elif node.dtype == paddle.int64: self.dtype_factor *= 8 elif node.dtype == paddle.uint8: self.dtype_factor = 1 else: self.dtype_factor = 2 # raise NotImplementedError("{} not counted".format(node.dtype)) self.batch_size = None if batch_size is not None: self.batch_size = batch_size def get_size(self): p = 1 for i in self.node.shape: if i == -1: # deal with placeholder assert self.batch_size is not None, "Batch size not decided." i = self.batch_size p *= i return p class CompOpCostNode(CostNode): def __init__(self, node, node_type, id=None, is_bwd=False, is_optim=False): super().__init__(node, node_type, id) self.is_bwd = is_bwd self.is_optim = is_optim def init_comp_cost(self, cost_data): # TODO: improve base.CostModel for more specific cost_data op_id = self.node.desc.id() if op_id in cost_data.keys(): self.cost = cost_data[op_id] else: self.cost = 0.0 class PipeEvent: def __init__(self, stage_id, event_name, duration, start_time=-1): self.stage_id = stage_id self.name = event_name self.duration = duration self.s_time = start_time self.e_time = -1 class CostModel: def __init__( self, mode=CostModelMode.BENCHMARKING, cluster=None, batch_size=1, microbatch_num=1, opcall_overhead=0, standalone_cost_data=None, pipeline_config=None, ): self.mode = mode # parameters self.opcall_overhead = opcall_overhead self.batch_size = batch_size self.microbatch_num = microbatch_num self.nodes = {} # name -> node self.origin_graph = {} # original graph self.op_graph = {} # op graph (no variables nodes) self.runtime_graph = {} # runtime graph, for simulation self.cluster = cluster self.cost_data = standalone_cost_data self.pp2rank = pipeline_config if self.pp2rank is not None: self.rank2pp = {} for stage_idx, ranks in enumerate(self.pp2rank): for rank in ranks: self.rank2pp[rank] = stage_idx else: self.rank2pp = None self.ring2rank = {} self.fwd_time = [] self.bwd_time = [] self.optim_time = [] def _parse_sub_program(self, program, nodes, graph, cost_data, sub_idx): assert len(program.blocks) == 1, ( "Program more than 1 block not supported." ) block = program.blocks[0] var_id = "lod_tensor_blocking_queue_0" new_var = program.global_block().create_var( name=var_id, dtype=paddle.float32, type=core.VarDesc.VarType.DENSE_TENSOR, ) nodes[var_id] = TensorCostNode( new_var, CostNodeType.VARIABLE, "lod_tensor_blocking_queue_0" ) for var in block.vars.values(): var_id = var.name # if var.name == "create_py_reader_0" or var.name == "double_buffer_0": # continue nodes[var_id] = TensorCostNode(var, CostNodeType.VARIABLE, var_id) graph[var_id] = [[], []] for op in block.ops: op_id = op.type + "_" + str(op.idx) if ( op.type.startswith('c_') or op.type.startswith('send') or op.type.startswith('recv') ): is_bwd = False if ( op.type.startswith('c_') and op.type != "c_sync_calc_stream" and not op.type.startswith('c_embedding') ): ring_id = op.attr('ring_id') if ring_id not in self.ring2rank: self.ring2rank[ring_id] = set() self.ring2rank[ring_id].add(sub_idx) is_bwd = '@GRAD' in op.output('Out')[0] elif op.type.startswith('recv'): is_bwd = '@GRAD' in op.output('Out')[0] elif op.type.startswith('send'): is_bwd = '@GRAD' in op.input('X')[0] op_node = CommOpCostNode( op, CostNodeType.COMMUNICATION, op_id, is_bwd ) else: is_bwd = ( int(op.attr('op_role')) == int(OpRole.Backward) ) or "@GRAD" in op.input_arg_names is_optim = 'LearningRate' in op.input_names op_node = CompOpCostNode( op, CostNodeType.COMPUTATION, op_id, is_bwd, is_optim ) op_node.init_comp_cost(cost_data) nodes[op_id] = op_node graph[op_id] = [[], []] comm_input_shape = [0] comm_output_shape = [0] for i in range(len(op.input_names)): try: var_id = op.input(op.input_names[i])[0] var_node = nodes[var_id] graph[op_id][PRED].append(var_node.id) graph[var_id][SUCC].append(op_node.id) comm_input_shape = var_node.shape except: continue for i in range(len(op.output_names)): try: var_id = op.output(op.output_names[i])[0] var_node = nodes[var_id] graph[op_id][SUCC].append(var_node.id) graph[var_id][PRED].append(op_node.id) comm_output_shape = var_node.shape except: continue if op_node.type == CostNodeType.COMMUNICATION: op_node.set_shapes(comm_input_shape, comm_output_shape) # resolve hazard: rename the r/w hazard variable nodes to ensure self.origin_graph is a DAG new_var_dict = {} for node_id, node in nodes.items(): if node.type == CostNodeType.VARIABLE and node.node.persistable: write_op_cnt = 0 for pred_id in graph[node_id][PRED]: pred = nodes[pred_id] if pred.type == CostNodeType.COMPUTATION and ( pred_id in graph[node_id][SUCC] ): graph[pred_id][SUCC].remove(node_id) graph[node_id][PRED].remove(pred_id) write_op_cnt += 1 new_var_id = node_id + f'_write_{write_op_cnt}' new_var = TensorCostNode( node.node, CostNodeType.VARIABLE, new_var_id, shared_node_id=node_id, ) graph[new_var_id] = [[], []] graph[pred_id][SUCC].append(new_var_id) graph[new_var_id][PRED].append(pred_id) new_var_dict[new_var_id] = new_var for k, v in new_var_dict.items(): nodes[k] = v return nodes def parse_program(self, distributed_program): self.distributed_program = distributed_program self.total_rank = len(self.distributed_program) sub_prog_cnt = len(distributed_program) self.nodes = [] * sub_prog_cnt self.origin_graph = [] * sub_prog_cnt # original graph self.op_graph = [] * sub_prog_cnt # op graph (no variables nodes) self.runtime_graph = [] * sub_prog_cnt # runtime graph, for simulation for sub_idx, sub_prog in enumerate(distributed_program): self.nodes.append({}) self.origin_graph.append({}) self.op_graph.append({}) self.runtime_graph.append({}) self._parse_sub_program( sub_prog, self.nodes[sub_idx], self.origin_graph[sub_idx], self.cost_data[ 0 if self.rank2pp is None else self.rank2pp[sub_idx] ], sub_idx, ) return self.nodes def _find_succ_op(self, node_id, sub_idx=0): succ_ops_id = [] for succ_id in self.origin_graph[sub_idx][node_id][SUCC]: succ = self.nodes[sub_idx][succ_id] if ( succ.type == CostNodeType.COMMUNICATION or succ.type == CostNodeType.COMPUTATION ): succ_ops_id.append(succ_id) elif succ.type == CostNodeType.VARIABLE: succ_ops_id = succ_ops_id + self._find_succ_op(succ_id, sub_idx) else: raise NotImplementedError( f'This type of node not supported yet:{succ.type}' ) return succ_ops_id def build_op_graph(self): for sub_idx in range(self.total_rank): op_nodes_id = [] for node_id, node in self.nodes[sub_idx].items(): if node.type == CostNodeType.VARIABLE: continue self.op_graph[sub_idx][node_id] = [[], []] op_nodes_id.append(node_id) for op_id in op_nodes_id: succ_nodes_id = self._find_succ_op(op_id, sub_idx) self.op_graph[sub_idx][op_id][SUCC] = succ_nodes_id for succ_id in succ_nodes_id: self.op_graph[sub_idx][succ_id][PRED].append(op_id) def build_runtime_graph(self): self.runtime_graph = copy.deepcopy(self.op_graph) def eliminate_multi_edges(self, graph=None): for node_id, edges in graph.items(): graph[node_id][PRED] = list(set(edges[PRED])) graph[node_id][SUCC] = list(set(edges[SUCC])) def merge_comm(self): for sub_idx in range(self.total_rank): for node_id, edges in self.op_graph[sub_idx].items(): node = self.nodes[sub_idx][node_id] if ( node_id.startswith('c_') and not node.id.startswith("c_sync_calc_stream") and not node.id.startswith('c_embedding') ): ring_id = node.node.attr('ring_id') node.set_ranks(list(self.ring2rank[ring_id])) node.init_comm_cost(self.cluster) elif node_id.startswith('send') or node_id.startswith('recv'): peer_rank = node.node.attr('peer') node.set_ranks([sub_idx, peer_rank]) node.init_comm_cost(self.cluster) else: pass # Not communication op def _merge_node(self, to_merge_node_list, merge_type='linear', nodes=None): nodes_list = [] node_cost = 0 for node in to_merge_node_list: if isinstance(node, MergedOpsCostNode): nodes_list += node.node_list else: nodes_list.append(node.id) if merge_type == 'linear': node_cost += node.cost elif merge_type == 'branch': node_cost = max(node_cost, node.cost) else: raise NotImplementedError( f'This type of merging is not supported:{merge_type}' ) merged_node_id = 'merged_' + str(len(nodes)) is_bwd = to_merge_node_list[0].is_bwd merged_node = MergedOpsCostNode( CostNodeType.MERGED, id=merged_node_id, base_node_list=nodes_list, is_bwd=is_bwd, ) merged_node.cost = node_cost return merged_node_id, merged_node def merge_linear(self): r''' This method does the following: If X depends on Y only, they must be run sequentially. [ e.g. A ->- C ->- D D and E depends on C only.] [ B ->-/ \->- E C depends on A and B. ] We merge X and Y into a new node and sum up their cost time. ''' cnt = 0 for sub_idx in range(self.total_rank): cnt += self._merge_linear( self.nodes[sub_idx], self.runtime_graph[sub_idx], is_bwd=False ) cnt += self._merge_linear( self.nodes[sub_idx], self.runtime_graph[sub_idx], is_bwd=True ) return cnt def merge_branch(self): r''' This method does the following: If a node has more than one successor, there is *branch*. [ e.g. A ->- B ->- D ] [ \->- C ->- / , B and C can be run at the same time ] case 1: if B or C is null (or D is directly dependent on A), it's equivalent to A->C->D or A->B->D, fall back to self.merge_linear case 2: if both B and C are some op, merged_cost = max(cost(B), cost(C)) ''' cnt = 0 for sub_idx in range(self.total_rank): cnt += self._merge_branch( self.nodes[sub_idx], self.runtime_graph[sub_idx], is_bwd=False ) cnt += self._merge_branch( self.nodes[sub_idx], self.runtime_graph[sub_idx], is_bwd=True ) return cnt def _merge_linear(self, nodes, runtime_graph, is_bwd=False): reduct_cnt = 0 rt_nodes_id = list(runtime_graph.keys()) for node_id in rt_nodes_id: if node_id not in runtime_graph.keys(): continue node = nodes[node_id] if not is_bwd == node.is_bwd or node.is_optim: continue edges = runtime_graph[node_id] ind = len(edges[PRED]) # in_degree if ind == 1: # only depend on one node pred_id = edges[PRED][0] pred = nodes[pred_id] merged_node_id, merged_node = self._merge_node( [node, pred], merge_type='linear', nodes=nodes ) nodes[merged_node_id] = merged_node runtime_graph[merged_node_id] = [[], []] # delete edges and add new edges succ = None try: runtime_graph[merged_node_id][SUCC] = copy.deepcopy( edges[SUCC] ) if len(runtime_graph[pred_id][SUCC]) > 1: # predecessor has more than 1 successor # the merged_node is to inherit the rest of its successors succ = runtime_graph[pred_id][SUCC] succ.remove(node_id) runtime_graph[merged_node_id][SUCC] += succ runtime_graph[merged_node_id][PRED] = runtime_graph[ pred_id ][PRED] except: pass try: for i in runtime_graph[pred_id][PRED]: try: runtime_graph[i][SUCC].remove(pred_id) except: continue runtime_graph[i][SUCC].append(merged_node_id) except: pass try: for i in edges[SUCC]: runtime_graph[i][PRED].remove(node_id) runtime_graph[i][PRED].append(merged_node_id) except: pass if succ is not None: for i in succ: try: runtime_graph[i][PRED].remove(pred_id) except: continue runtime_graph[i][PRED].append(merged_node_id) runtime_graph.pop(node_id) try: runtime_graph.pop(pred_id) except: continue reduct_cnt += 1 self.eliminate_multi_edges(runtime_graph) break return reduct_cnt # the number of nodes that have been reduced def _merge_branch(self, nodes, runtime_graph, is_bwd=False): reduct_cnt = 0 rt_nodes_id = list(runtime_graph.keys()) for node_id in rt_nodes_id: node = nodes[node_id] if not is_bwd == node.is_bwd or node.is_optim: continue edges = runtime_graph[node_id] outd = len(edges[SUCC]) # out_degree if outd > 1: # branch out succ_nodes_id = edges[SUCC] succ_to_elim = [] for succ_id in succ_nodes_id: for succ_2_id in succ_nodes_id: try: tmp = runtime_graph[succ_2_id][SUCC] except: continue if succ_id in tmp: succ_to_elim.append(succ_id) break for id in succ_to_elim: edges[SUCC].remove(id) runtime_graph[id][PRED].remove(node_id) reduct_cnt += 1 to_merge = True try: if ( len(edges[SUCC]) < 1 or len(runtime_graph[edges[SUCC][0]][SUCC]) < 1 ): continue except: continue end_node_id = runtime_graph[edges[SUCC][0]][SUCC][0] for i in succ_nodes_id: try: if ( len(runtime_graph[i][SUCC]) != 1 or runtime_graph[i][SUCC][0] != end_node_id ): to_merge = False # if branches has different end node, we don't merge them break except: continue if to_merge and len(succ_nodes_id) > 1: to_merge_node_list = [nodes[i] for i in succ_nodes_id] merged_node_id, merged_node = self._merge_node( to_merge_node_list, merge_type='branch', nodes=nodes ) nodes[merged_node_id] = merged_node runtime_graph[merged_node_id] = [[], []] # delete edges and add new edges runtime_graph[merged_node_id][SUCC] = [end_node_id] runtime_graph[merged_node_id][PRED] = edges[PRED] runtime_graph[end_node_id][PRED] = [merged_node_id] runtime_graph[node_id][SUCC] = [merged_node_id] try: for i in succ_nodes_id: runtime_graph.pop(i) reduct_cnt += len(to_merge_node_list) - 1 break except: pass return reduct_cnt def get_runtime_cost(self): def get_node_cost(node): node_cost = node.cost + self.opcall_overhead if isinstance(node, MergedOpsCostNode): for it in node.node_list: node_cost += self.opcall_overhead return node_cost for sub_idx in range(self.total_rank): fwd_cost = 0 bwd_cost = 0 optim_cost = 0 for node_id in self.runtime_graph[sub_idx].keys(): node = self.nodes[sub_idx][node_id] if node.is_optim: optim_cost += get_node_cost(node) elif node.is_bwd: bwd_cost += get_node_cost(node) else: fwd_cost += get_node_cost(node) self.fwd_time.append(fwd_cost) self.bwd_time.append(bwd_cost) self.optim_time.append(optim_cost) return self.fwd_time, self.bwd_time, self.optim_time def get_mem(self): static_list = [] top_list = [] for sub_idx in range(self.total_rank): static_mem, cur_mem, top_mem = self._simulate_mem( self.nodes[sub_idx], self.origin_graph[sub_idx] ) static_list.append(static_mem) top_list.append(top_mem) return static_list, top_list def _simulate_mem(self, nodes, origin_graph): q = queue.Queue(1024) sim_graph = copy.deepcopy(origin_graph) for node_id, node in nodes.items(): if len(sim_graph[node_id][PRED]) == 0: q.put(node_id) q.put('nop') cur_mem = 0 top_mem = -1 static_mem = 0 while not q.empty(): node_id = q.get() node = None size = 0 if node_id == 'nop': top_mem = max(cur_mem, top_mem) if q.empty(): break else: q.put(node_id) continue else: node = nodes[node_id] if node.type == CostNodeType.VARIABLE: size = node.get_size() if node.node.persistable: static_mem += size cur_mem += size edges = sim_graph[node_id] if not ( node.type == CostNodeType.VARIABLE and node.node.persistable ): for succ_id in edges[SUCC]: sim_graph[succ_id][PRED].remove(node_id) if len(sim_graph[succ_id][PRED]) == 0: q.put(succ_id) for pred_id in edges[PRED]: pred = nodes if pred.type == CostNodeType.VARIABLE: sim_graph[pred_id][SUCC].remove(node_id) if ( len(sim_graph[pred_id][SUCC]) == 0 and not pred.node.persistable ): cur_mem -= pred.get_size() return static_mem, cur_mem, top_mem def get_pipeline_time(self): if self.pp2rank is None: return self.fwd_time[0] + self.bwd_time[0] + self.optim_time[0] else: return self._simulate_pipeline() def _simulate_pipeline(self): stage_num = len(self.pp2rank) event_list = [] global_time = [0] * stage_num total_time = 0 fwd_cnt = list(range(stage_num, 0, -1)) bwd_cnt = [self.microbatch_num] * stage_num q = queue.Queue(1024) for i in range(self.microbatch_num): q.put(PipeEvent(0, 'fwd', self.fwd_time[0])) while not q.empty(): e = q.get() stid = e.stage_id if e.name == 'fwd': if fwd_cnt[stid] > 0: e.s_time = max(global_time[stid], e.s_time) e.e_time = e.s_time + e.duration event_list.append(e) if stid != stage_num - 1: q.put( PipeEvent( stid + 1, 'fwd', self.fwd_time[stid + 1], start_time=e.e_time, ) ) else: q.put( PipeEvent( stid, 'bwd', self.bwd_time[stid], start_time=e.e_time, ) ) fwd_cnt[stid] -= 1 global_time[stid] = e.e_time else: q.put(e) elif e.name == 'bwd': e.s_time = max(global_time[stid], e.s_time) e.e_time = e.s_time + e.duration event_list.append(e) if stid != 0: q.put( PipeEvent( stid - 1, 'bwd', self.bwd_time[stid - 1], start_time=e.e_time, ) ) fwd_cnt[stid] += 1 bwd_cnt[stid] -= 1 if bwd_cnt[stid] == 0: q.put( PipeEvent( stid, 'optim', self.optim_time[stid], start_time=e.e_time, ) ) global_time[stid] = e.e_time elif e.name == 'optim': e.s_time = max(global_time[stid], e.s_time) e.e_time = e.s_time + e.duration event_list.append(e) global_time[stid] = e.e_time else: raise NotImplementedError( f'This type of pipe event is not supported yet.{e.name}' ) for t in global_time: total_time = max(total_time, t) return total_time def get_cost(self): cost = Cost() static_mem, peak_mem = self.get_mem() cost.static_mem = static_mem cost.peak_mem = peak_mem self.merge_comm() while True: cnt = 0 cnt += self.merge_linear() cnt += self.merge_branch() if cnt == 0: # can't be further merged break self.get_runtime_cost() cost.runtime = self.get_pipeline_time() return cost def init(self, distributed_program): self.parse_program(distributed_program) self.build_op_graph() for sub_idx in range(self.total_rank): self.eliminate_multi_edges(self.op_graph[sub_idx]) self.build_runtime_graph() def estimate_cost( distributed_program, cluster, pipeline_config, standalone_cost_data, batch_size, ): """ Estimated cost from distributed program, cluster model and distributed settings. Args: distributed_program(list): list of paddle programs cluster(Cluster): cluster model standalone_cost_data(CostData): cost data given by paddle.core batch_size(int): batch size of the training workload pipeline_config(list): configuration of pipeline stage allocation """ # the following line is left for now, cluster model will be involved in the future assert cluster is None, "For now, cluster remains None" cm_ctx = CostModel( cluster=cluster, batch_size=batch_size, standalone_cost_data=standalone_cost_data, pipeline_config=pipeline_config, ) cm_ctx.init(distributed_program) cost = cm_ctx.get_cost() return cost