chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,855 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user