wehub-resource-sync
343 lines
12 KiB
Python
343 lines
12 KiB
Python
# 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 functools
|
|
import operator
|
|
import os
|
|
from collections import deque
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
|
|
from .cluster import DeviceType
|
|
from .graph import Graph
|
|
from .process_group import get_process_group
|
|
|
|
|
|
def is_collective_comm_op(op):
|
|
comm_list = [
|
|
"all_gather",
|
|
"all_reduce",
|
|
"broadcast",
|
|
]
|
|
reduce_type = [
|
|
dist.ReduceOp.SUM,
|
|
dist.ReduceOp.MIN,
|
|
dist.ReduceOp.MAX,
|
|
dist.ReduceOp.PROD,
|
|
]
|
|
if op.type == "reduce" and op.attr("reduce_type") in reduce_type:
|
|
return True
|
|
if op.type in comm_list:
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
|
|
def is_p2p_comm_op(op):
|
|
comm_list = ["send_v2", "recv_v2"]
|
|
if op.type in comm_list:
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
|
|
def get_dtype_bytes(dtype):
|
|
num_bytes = 0
|
|
if dtype == paddle.float64:
|
|
num_bytes = 8
|
|
elif dtype == paddle.float32:
|
|
num_bytes = 4
|
|
elif dtype == paddle.float16:
|
|
num_bytes = 2
|
|
elif dtype == paddle.bfloat16:
|
|
num_bytes = 2
|
|
elif dtype == paddle.int64:
|
|
num_bytes = 8
|
|
elif dtype == paddle.int32:
|
|
num_bytes = 4
|
|
elif dtype == paddle.int16:
|
|
num_bytes = 2
|
|
elif dtype == paddle.int8:
|
|
num_bytes = 1
|
|
elif dtype == paddle.uint8:
|
|
num_bytes = 1
|
|
else:
|
|
raise ValueError(f"Unrecognized dtype {dtype}.")
|
|
return num_bytes
|
|
|
|
|
|
def get_comm_volume(comm_op, src_rank, tgt_rank):
|
|
comm_volume = None
|
|
if src_rank == tgt_rank:
|
|
return comm_volume
|
|
comm_op_type = comm_op.type
|
|
if comm_op_type != "recv_v2":
|
|
tensor_name = comm_op.input_arg_names[0]
|
|
else:
|
|
tensor_name = comm_op.output_arg_names[0]
|
|
tensor = comm_op.block._find_var_recursive(tensor_name)
|
|
assert tensor is not None
|
|
tensor_shape = tensor.shape
|
|
# Skip the batch dim
|
|
new_tensor_shape = []
|
|
for val in tensor_shape:
|
|
if val == -1:
|
|
print("Warning: -1 in the tensor shape.")
|
|
new_tensor_shape.append(1)
|
|
else:
|
|
new_tensor_shape.append(val)
|
|
tensor_size = functools.reduce(operator.mul, new_tensor_shape, 1)
|
|
tensor_bytes = tensor_size * get_dtype_bytes(tensor.dtype)
|
|
if "c_allreduce" in comm_op_type or "all_reduce" in comm_op_type:
|
|
comm_volume = 2 * tensor_bytes
|
|
elif "all_gather" in comm_op_type:
|
|
comm_volume = tensor_bytes
|
|
elif "broadcast" in comm_op_type:
|
|
if comm_op.attr("root") == src_rank:
|
|
comm_volume = tensor_bytes
|
|
else:
|
|
comm_volume = None
|
|
elif "c_reduce" in comm_op_type:
|
|
if comm_op.attr("root_id") == src_rank:
|
|
comm_volume = None
|
|
else:
|
|
comm_volume = tensor_bytes
|
|
elif "reduce" == comm_op_type:
|
|
if comm_op.attr("root_id") == src_rank:
|
|
comm_volume = None
|
|
else:
|
|
comm_volume = tensor_bytes
|
|
elif "send_v2" in comm_op_type:
|
|
if comm_op.attr("peer") == tgt_rank:
|
|
comm_volume = tensor_bytes
|
|
else:
|
|
comm_volume = None
|
|
elif "recv_v2" in comm_op_type:
|
|
comm_volume = None
|
|
else:
|
|
raise ValueError("Unrecognized communication operator.")
|
|
return comm_volume
|
|
|
|
|
|
def analyze_comm_requirements_from_op(op, rank, g_process_group_map):
|
|
comm_requirements_to_ranks = {}
|
|
if is_collective_comm_op(op):
|
|
process_group_id = op.attr("ring_id")
|
|
process_group = get_process_group(process_group_id, g_process_group_map)
|
|
if rank not in process_group.ranks:
|
|
return comm_requirements_to_ranks
|
|
for tgt_rank in process_group.ranks:
|
|
comm_volume = get_comm_volume(op, rank, tgt_rank)
|
|
if comm_volume is not None:
|
|
comm_requirements_to_ranks[tgt_rank] = {}
|
|
comm_requirements_to_ranks[tgt_rank]["comm_volume"] = (
|
|
comm_volume
|
|
)
|
|
elif is_p2p_comm_op(op):
|
|
tgt_rank = op.attr("peer")
|
|
comm_volume = get_comm_volume(op, rank, tgt_rank)
|
|
if comm_volume is not None:
|
|
comm_requirements_to_ranks[tgt_rank] = {}
|
|
comm_requirements_to_ranks[tgt_rank]["comm_volume"] = comm_volume
|
|
else:
|
|
comm_requirements_to_ranks = {}
|
|
return comm_requirements_to_ranks
|
|
|
|
|
|
def analyze_requirements_for_program(src_info, rank):
|
|
program = src_info[0]
|
|
g_process_group_map = src_info[1]
|
|
resource_requirements = {}
|
|
comm_requirements_to_ranks = {}
|
|
# only support device_type and only support GPU for now
|
|
resource_requirements["device_type"] = DeviceType.GPU
|
|
for block in program.blocks:
|
|
for op in block.ops:
|
|
cur_comm_requirements_to_ranks = analyze_comm_requirements_from_op(
|
|
op, rank, g_process_group_map
|
|
)
|
|
for tgt_rank, link_info in cur_comm_requirements_to_ranks.items():
|
|
if tgt_rank in comm_requirements_to_ranks:
|
|
comm_requirements_to_ranks[tgt_rank]["comm_volume"] += (
|
|
link_info["comm_volume"]
|
|
)
|
|
else:
|
|
comm_requirements_to_ranks[tgt_rank] = {}
|
|
comm_requirements_to_ranks[tgt_rank]["comm_volume"] = (
|
|
link_info["comm_volume"]
|
|
)
|
|
return resource_requirements, comm_requirements_to_ranks
|
|
|
|
|
|
def build_process_graph(distributed_program):
|
|
graph = Graph()
|
|
for src_rank, src_info in distributed_program.items():
|
|
(
|
|
resource_requirements,
|
|
comm_requirements_to_ranks,
|
|
) = analyze_requirements_for_program(src_info, src_rank)
|
|
graph.add_node(src_rank, resource_requirements=resource_requirements)
|
|
for tgt_rank, comm_requirements in comm_requirements_to_ranks.items():
|
|
graph.add_edge(
|
|
src_rank, tgt_rank, comm_requirements=comm_requirements
|
|
)
|
|
return graph
|
|
|
|
|
|
def build_cluster_graph(cluster):
|
|
graph = Graph()
|
|
cuda_visible_devices_env = os.getenv("CUDA_VISIBLE_DEVICES")
|
|
cuda_visible_devices = []
|
|
if cuda_visible_devices_env is not None and cuda_visible_devices_env != "":
|
|
cuda_visible_devices = [
|
|
int(d.strip()) for d in cuda_visible_devices_env.split(",")
|
|
]
|
|
for machine in cluster.machines.values():
|
|
for device in machine.devices.values():
|
|
graph.add_node(device.global_id, device=device)
|
|
if (
|
|
cuda_visible_devices
|
|
and device.local_id not in cuda_visible_devices
|
|
):
|
|
graph.nodes[device.global_id]["occupied"] = True
|
|
else:
|
|
graph.nodes[device.global_id]["occupied"] = False
|
|
for link in machine.links.values():
|
|
graph.add_edge(
|
|
link.source.global_id, link.target.global_id, link=link
|
|
)
|
|
return graph
|
|
|
|
|
|
def mapping(distributed_program, cluster):
|
|
# A very simple mapping algorithm only for GPUs.
|
|
# Here we assume one process will be mapped to one GPU.
|
|
# In the future, more mapping configurations and algorithms will be supported.
|
|
process_graph = build_process_graph(distributed_program)
|
|
|
|
cluster_graph = build_cluster_graph(cluster)
|
|
|
|
for cur_rank_node in process_graph:
|
|
cur_rank_node["visited"] = False
|
|
|
|
def sort_by_comm_volume(rank_edge):
|
|
return rank_edge["comm_requirements"]["comm_volume"]
|
|
|
|
def sort_by_comm_bandwidth(device_edge):
|
|
return device_edge["link"].bandwidth
|
|
|
|
def select_unvisited_rank_node(rank_node_list):
|
|
selected_rank_node = None
|
|
for rank_node in rank_node_list:
|
|
if rank_node["visited"] is False:
|
|
selected_rank_node = rank_node
|
|
return selected_rank_node
|
|
|
|
queue = deque()
|
|
root_rank_node = select_unvisited_rank_node(
|
|
list(process_graph.nodes.values())
|
|
)
|
|
while root_rank_node is not None:
|
|
queue.append(root_rank_node)
|
|
while queue:
|
|
cur_rank_node = queue.popleft()
|
|
if cur_rank_node["visited"]:
|
|
continue
|
|
device_type = cur_rank_node["resource_requirements"]["device_type"]
|
|
cur_device_node = None
|
|
for device_node in cluster_graph.nodes.values():
|
|
if (device_node["device"].type == device_type) and (
|
|
not device_node["occupied"]
|
|
):
|
|
device_node["occupied"] = True
|
|
cur_rank_node["visited"] = True
|
|
cur_rank_node["device"] = device_node["device"]
|
|
cur_device_node = device_node
|
|
break
|
|
assert cur_device_node, (
|
|
"Cannot find a device to satisfy the requirement."
|
|
)
|
|
|
|
nbr_rank_edges = []
|
|
for nbr_rank_node_id, nbr_rank_edge in process_graph.adjs[
|
|
cur_rank_node.id
|
|
].items():
|
|
assert (
|
|
nbr_rank_edge.src_id == cur_rank_node.id
|
|
and nbr_rank_edge.tgt_id == nbr_rank_node_id
|
|
)
|
|
queue.append(process_graph.nodes[nbr_rank_node_id])
|
|
nbr_rank_edges.append(nbr_rank_edge)
|
|
nbr_rank_edges.sort(key=sort_by_comm_volume)
|
|
|
|
nbr_device_edges = []
|
|
for nbr_device_edge in cluster_graph.adjs[
|
|
cur_device_node.id
|
|
].values():
|
|
nbr_device_edges.append(nbr_device_edge)
|
|
nbr_device_edges.sort(key=sort_by_comm_bandwidth)
|
|
|
|
for nbr_rank_edge in nbr_rank_edges:
|
|
src_rank_node = process_graph.nodes[nbr_rank_edge.src_id][
|
|
"visited"
|
|
]
|
|
if src_rank_node:
|
|
continue
|
|
device_type = src_rank_node["resource_requirements"][
|
|
"device_type"
|
|
]
|
|
nbr_rank_node = process_graph.nodes[nbr_rank_edge.tgt_id]
|
|
for nbr_device_edge in nbr_device_edges:
|
|
nbr_device_node = cluster_graph.nodes[
|
|
nbr_device_edge.tgt_id
|
|
]
|
|
if (nbr_device_node["device"].type == device_type) and (
|
|
not nbr_device_node["occupied"]
|
|
):
|
|
nbr_device_node["occupied"] = True
|
|
nbr_rank_node["visited"] = True
|
|
nbr_rank_node["device"] = nbr_device_node["device"]
|
|
break
|
|
root_rank_node = select_unvisited_rank_node(
|
|
list(process_graph.nodes.values())
|
|
)
|
|
|
|
rank_mapping = {}
|
|
for rank, rank_node in process_graph.nodes.items():
|
|
device = rank_node["device"]
|
|
machine = device.machine
|
|
if machine.id in rank_mapping:
|
|
rank_mapping[machine.id]["hostname"] = machine.hostname
|
|
rank_mapping[machine.id]["addr"] = machine.addr
|
|
rank_mapping[machine.id]["port"] = machine.port
|
|
if rank not in rank_mapping[machine.id]["ranks"]:
|
|
rank_mapping[machine.id]["ranks"][rank] = []
|
|
rank_mapping[machine.id]["ranks"][rank].append(device.local_id)
|
|
else:
|
|
rank_mapping[machine.id]["ranks"][rank].append(device.local_id)
|
|
else:
|
|
rank_mapping[machine.id] = {}
|
|
rank_mapping[machine.id]["hostname"] = machine.hostname
|
|
rank_mapping[machine.id]["addr"] = machine.addr
|
|
rank_mapping[machine.id]["port"] = machine.port
|
|
rank_mapping[machine.id]["ranks"] = {}
|
|
rank_mapping[machine.id]["ranks"][rank] = []
|
|
rank_mapping[machine.id]["ranks"][rank].append(device.local_id)
|
|
for machine_mapping in rank_mapping.values():
|
|
for rank_devices in machine_mapping["ranks"].values():
|
|
rank_devices.sort()
|
|
|
|
return rank_mapping
|