wehub-resource-sync
1133 lines
43 KiB
Python
1133 lines
43 KiB
Python
# Copyright (c) 2022 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 hashlib
|
|
import itertools
|
|
import math
|
|
import time
|
|
from collections import defaultdict
|
|
|
|
import numpy as np
|
|
|
|
from ...process_mesh import ProcessMesh
|
|
from ..completion import Completer
|
|
from ..cost import CostEstimator
|
|
from ..dist_context import _node_id
|
|
from ..dist_op import DistributedOperator
|
|
from ..operators.common import find_compatible_distributed_operator_impls
|
|
from ..parallelizer_v2 import Parallelizer
|
|
from .trial import Trial, TrialStatus
|
|
from .tunable_space import TunableSpace
|
|
from .tunable_variable import Boolean, IntRange
|
|
|
|
|
|
class ParallelTuner:
|
|
def __init__(
|
|
self,
|
|
dist_context,
|
|
mode="train",
|
|
max_trials=25,
|
|
tuner_id=None,
|
|
seed=None,
|
|
logger=None,
|
|
loop_count=10,
|
|
):
|
|
self._loop_count = loop_count
|
|
self._estimator = None
|
|
self._dist_context = dist_context
|
|
assert self._dist_context._is_initialized
|
|
self._mode = mode
|
|
self._cluster = self._dist_context.cluster
|
|
self._num_machines = self._cluster.get_num_machines()
|
|
self._num_devices_per_machine = (
|
|
self._cluster.get_num_devices_per_machine()
|
|
)
|
|
self._space = TunableSpace()
|
|
self._objective = "time"
|
|
self._direction = "min"
|
|
self._max_trials = max_trials
|
|
self._tuner_id = tuner_id
|
|
self._seed = seed if seed is not None else 9999
|
|
|
|
print(
|
|
"seed",
|
|
self._seed,
|
|
"mode",
|
|
self._mode,
|
|
"num_machines",
|
|
self._num_machines,
|
|
"num_devices_per_machine",
|
|
self._num_devices_per_machine,
|
|
flush=True,
|
|
)
|
|
self._seed_state = self._seed
|
|
self._logger = logger
|
|
self._max_collisions = 3
|
|
self._tried_values = set()
|
|
self._num_trials = 0
|
|
self._rng = np.random.default_rng(self._seed)
|
|
|
|
# Search the op types in the include_op_types,
|
|
# and will search all op types if it is empty.
|
|
# Exclude the op types in the exclude_op_types
|
|
# from the search list.
|
|
self._exclude_op_types = []
|
|
self._include_op_types = []
|
|
# The final dist ops will be searched after considering
|
|
# the include_op_types and exclude_op_types.
|
|
self._concerned_dist_ops = {}
|
|
|
|
self._op_id_to_dist_attr_candidates = defaultdict(list)
|
|
self._cached_dims_mapping_candidates = {}
|
|
self._cached_candidates_info = defaultdict(list)
|
|
|
|
self._special_ops = [
|
|
"create_py_reader",
|
|
"create_double_buffer_reader",
|
|
"read",
|
|
"while",
|
|
"read_from_array",
|
|
"write_to_array",
|
|
]
|
|
|
|
# Each parallel strategy has two elements. The First one is for distributed tensors,
|
|
# the second element is for distributed tensors, the third element is for process meshes.
|
|
self._init_parallel_strategy = [None, None, None]
|
|
self._best_parallel_strategy = [None, None, None]
|
|
|
|
self._completer = Completer(self._dist_context)
|
|
|
|
self._parallelizer = Parallelizer(
|
|
self._mode, self._completer, self._dist_context
|
|
)
|
|
|
|
def _generate_combination(
|
|
self,
|
|
elements,
|
|
target,
|
|
idx,
|
|
partial_candidate,
|
|
candidates,
|
|
num_candidates=None,
|
|
):
|
|
if target == 0:
|
|
candidates.append(copy.deepcopy(partial_candidate))
|
|
return
|
|
|
|
if (
|
|
target < 0
|
|
or idx == len(elements)
|
|
or len(candidates) > num_candidates
|
|
):
|
|
return
|
|
|
|
# Use
|
|
partial_candidate.append(elements[idx])
|
|
self._generate_combination(
|
|
elements,
|
|
target - elements[idx],
|
|
idx,
|
|
partial_candidate,
|
|
candidates,
|
|
num_candidates,
|
|
)
|
|
# Not use
|
|
partial_candidate.pop()
|
|
self._generate_combination(
|
|
elements,
|
|
target,
|
|
idx + 1,
|
|
partial_candidate,
|
|
candidates,
|
|
num_candidates,
|
|
)
|
|
|
|
def _permute_combination(
|
|
self,
|
|
combination,
|
|
target,
|
|
check,
|
|
partial_candidate,
|
|
candidates,
|
|
num_candidates=None,
|
|
skip_prob=None,
|
|
):
|
|
if num_candidates is not None and len(candidates) == num_candidates:
|
|
return
|
|
|
|
if len(partial_candidate) == len(combination):
|
|
candidates.append(partial_candidate)
|
|
return
|
|
|
|
for i in range(len(combination)):
|
|
if check[i] == 1:
|
|
continue
|
|
if self._rng.choice([True, False], p=[skip_prob, 1 - skip_prob]):
|
|
continue
|
|
if (
|
|
i > 0
|
|
and combination[i] == combination[i - 1]
|
|
and check[i - 1] == 0
|
|
):
|
|
continue
|
|
check[i] = 1
|
|
self._permute_combination(
|
|
combination,
|
|
target,
|
|
check,
|
|
[*partial_candidate, combination[i]],
|
|
candidates,
|
|
num_candidates,
|
|
skip_prob,
|
|
)
|
|
check[i] = 0
|
|
|
|
def _partition_number(self, target):
|
|
log2_target = int(math.log2(target))
|
|
elements = [pow(2, i) for i in range(log2_target)]
|
|
if pow(2, log2_target) == target:
|
|
elements.append(target)
|
|
seed_candidates = []
|
|
num_seed_candidates = 1000
|
|
partial_results = []
|
|
self._generate_combination(
|
|
elements,
|
|
target,
|
|
0,
|
|
partial_results,
|
|
seed_candidates,
|
|
num_seed_candidates,
|
|
)
|
|
|
|
candidates = []
|
|
for seed_candidate in seed_candidates:
|
|
cur_candidates = []
|
|
num_cur_candidates = 16
|
|
seed_candidate.sort()
|
|
check = [0 for i in range(len(seed_candidate))]
|
|
if target <= 8:
|
|
skip_prob = 0.0
|
|
else:
|
|
skip_prob = len(seed_candidate) / target
|
|
self._permute_combination(
|
|
seed_candidate,
|
|
target,
|
|
check,
|
|
[],
|
|
cur_candidates,
|
|
num_cur_candidates,
|
|
skip_prob,
|
|
)
|
|
candidates.extend(cur_candidates)
|
|
return candidates
|
|
|
|
def _partition_devices(self, num_machines, num_devices_per_machine):
|
|
inter_node_partitions = self._partition_number(num_machines)
|
|
intra_node_partitions = self._partition_number(num_devices_per_machine)
|
|
return inter_node_partitions, intra_node_partitions
|
|
|
|
def _generate_process_mesh_list(
|
|
self, inter_node_partition, intra_node_partition
|
|
):
|
|
process_mesh_list = []
|
|
start_row = 0
|
|
start_col = 0
|
|
for m in inter_node_partition:
|
|
start_col = 0
|
|
for n in intra_node_partition:
|
|
process_mesh = []
|
|
for p in range(m):
|
|
start = (
|
|
start_row + p
|
|
) * self._num_devices_per_machine + start_col
|
|
tmp = []
|
|
for q in range(n):
|
|
tmp.append(start + q)
|
|
process_mesh.append(tmp)
|
|
process_mesh_list.append(copy.deepcopy(process_mesh))
|
|
start_col += n
|
|
start_row += m
|
|
return process_mesh_list
|
|
|
|
def _generate_dims_mapping_candidates_helper(
|
|
self, dims_mapping, dims_list, start, visited, candidates
|
|
):
|
|
if start == len(dims_mapping) or all(visited):
|
|
candidates.append(copy.deepcopy(dims_mapping))
|
|
return
|
|
|
|
for idx, dim in enumerate(dims_list):
|
|
if not visited[idx]:
|
|
dims_mapping[start] = dim
|
|
visited[idx] = True
|
|
self._generate_dims_mapping_candidates_helper(
|
|
dims_mapping, dims_list, start + 1, visited, candidates
|
|
)
|
|
visited[idx] = False
|
|
dims_mapping[start] = -1
|
|
self._generate_dims_mapping_candidates_helper(
|
|
dims_mapping, dims_list, start + 1, visited, candidates
|
|
)
|
|
|
|
def _generate_dims_mapping_candidates(
|
|
self, dims_mapping_len, process_mesh_len
|
|
):
|
|
assert dims_mapping_len >= 1 and process_mesh_len >= 1
|
|
key = (dims_mapping_len, process_mesh_len)
|
|
if key in self._cached_dims_mapping_candidates:
|
|
return self._cached_dims_mapping_candidates[key]
|
|
candidates = []
|
|
dims_mapping = [-1 for i in range(dims_mapping_len)]
|
|
dims_list = list(range(process_mesh_len))
|
|
visited = [False for i in range(process_mesh_len)]
|
|
self._generate_dims_mapping_candidates_helper(
|
|
dims_mapping, dims_list, 0, visited, candidates
|
|
)
|
|
self._cached_dims_mapping_candidates[key] = candidates
|
|
return candidates
|
|
|
|
def _generate_dist_attr_candidates(self, op_id, dist_op):
|
|
# For now, only allow the process meshes have two dimensions
|
|
process_mesh_len = 2
|
|
serial_op = dist_op.serial_op
|
|
op_dist_attr = dist_op.dist_attr
|
|
if serial_op.type in self._special_ops:
|
|
return [copy.deepcopy(op_dist_attr)]
|
|
key = []
|
|
key.append(serial_op.type)
|
|
for input_name in serial_op.input_names:
|
|
key.append(input_name)
|
|
for input_arg_name in serial_op.input(input_name):
|
|
key.append(
|
|
len(op_dist_attr.get_input_dims_mapping(input_arg_name))
|
|
)
|
|
for output_name in serial_op.output_names:
|
|
key.append(output_name)
|
|
for output_arg_name in serial_op.output(output_name):
|
|
key.append(
|
|
len(op_dist_attr.get_output_dims_mapping(output_arg_name))
|
|
)
|
|
key = tuple(key)
|
|
|
|
if key in self._cached_candidates_info:
|
|
cached_dist_attr_candidates = []
|
|
cached_input_arg_names = self._cached_candidates_info[key][0]
|
|
cached_output_arg_names = self._cached_candidates_info[key][1]
|
|
for cached_dist_attr in self._cached_candidates_info[key][2]:
|
|
new_op_dist_attr = copy.deepcopy(dist_op.dist_attr)
|
|
i = 0
|
|
for input_name in serial_op.input_names:
|
|
for input_arg_name in serial_op.input(input_name):
|
|
cached_dims_mapping = (
|
|
cached_dist_attr.get_input_dims_mapping(
|
|
cached_input_arg_names[i]
|
|
)
|
|
)
|
|
new_op_dist_attr.set_input_dims_mapping(
|
|
input_arg_name, cached_dims_mapping
|
|
)
|
|
i += 1
|
|
i = 0
|
|
for output_name in serial_op.output_names:
|
|
for output_arg_name in serial_op.output(output_name):
|
|
cached_dims_mapping = (
|
|
cached_dist_attr.get_output_dims_mapping(
|
|
cached_output_arg_names[i]
|
|
)
|
|
)
|
|
new_op_dist_attr.set_output_dims_mapping(
|
|
output_arg_name, cached_dims_mapping
|
|
)
|
|
i += 1
|
|
cached_dist_attr_candidates.append(new_op_dist_attr)
|
|
return cached_dist_attr_candidates
|
|
|
|
# cached_candidates_info = []
|
|
input_arg_names = []
|
|
for input_name in serial_op.input_names:
|
|
for input_arg_name in serial_op.input(input_name):
|
|
input_arg_names.append(input_arg_name)
|
|
self._cached_candidates_info[key].append(input_arg_names)
|
|
# cached_candidates_info.append(input_arg_names)
|
|
output_arg_names = []
|
|
for output_name in serial_op.output_names:
|
|
for output_arg_name in serial_op.output(output_name):
|
|
output_arg_names.append(output_arg_name)
|
|
self._cached_candidates_info[key].append(output_arg_names)
|
|
# cached_candidates_info.append(output_arg_names)
|
|
|
|
new_op_dist_attr = copy.deepcopy(dist_op.dist_attr)
|
|
# Find valid dims_mapping candidates for inputs
|
|
input_names = []
|
|
dims_mapping_generated = []
|
|
inputs_dist_attrs = op_dist_attr.inputs_dist_attrs
|
|
for tensor_name, tensor_dist_attr in inputs_dist_attrs.items():
|
|
original_dims_mapping = tensor_dist_attr.dims_mapping
|
|
dims_mapping_len = len(original_dims_mapping)
|
|
input_names.append(tensor_name)
|
|
if dims_mapping_len < 1:
|
|
dims_mapping_generated.append(
|
|
[copy.deepcopy(original_dims_mapping)]
|
|
)
|
|
else:
|
|
dims_mapping_generated.append(
|
|
self._generate_dims_mapping_candidates(
|
|
dims_mapping_len, process_mesh_len
|
|
)
|
|
)
|
|
input_dims_mapping_candidates = []
|
|
for dims_mapping_list in itertools.product(*dims_mapping_generated):
|
|
dims_mapping_list = list(dims_mapping_list)
|
|
assert len(dims_mapping_list) == len(input_names)
|
|
for i, dims_mapping in enumerate(dims_mapping_list):
|
|
new_op_dist_attr.set_input_dims_mapping(
|
|
input_names[i], dims_mapping
|
|
)
|
|
new_dist_op = DistributedOperator(
|
|
dist_op.serial_op, new_op_dist_attr
|
|
)
|
|
dist_op_impls = find_compatible_distributed_operator_impls(
|
|
new_dist_op, fwd=True
|
|
)
|
|
if dist_op_impls is not None:
|
|
input_dims_mapping_candidates.append(dims_mapping_list)
|
|
|
|
# Find valid dims_mapping candidates for outputs
|
|
output_names = []
|
|
dims_mapping_generated = []
|
|
outputs_dist_attrs = op_dist_attr.outputs_dist_attrs
|
|
for tensor_name, tensor_dist_attr in outputs_dist_attrs.items():
|
|
original_dims_mapping = tensor_dist_attr.dims_mapping
|
|
dims_mapping_len = len(original_dims_mapping)
|
|
output_names.append(tensor_name)
|
|
if dims_mapping_len < 1:
|
|
dims_mapping_generated.append(
|
|
[copy.deepcopy(original_dims_mapping)]
|
|
)
|
|
else:
|
|
dims_mapping_generated.append(
|
|
self._generate_dims_mapping_candidates(
|
|
dims_mapping_len, process_mesh_len
|
|
)
|
|
)
|
|
output_dims_mapping_candidates = []
|
|
for dims_mapping_list in itertools.product(*dims_mapping_generated):
|
|
dims_mapping_list = list(dims_mapping_list)
|
|
assert len(dims_mapping_list) == len(output_names)
|
|
for i, dims_mapping in enumerate(dims_mapping_list):
|
|
new_op_dist_attr.set_output_dims_mapping(
|
|
output_names[i], dims_mapping
|
|
)
|
|
new_dist_op = DistributedOperator(
|
|
dist_op.serial_op, new_op_dist_attr
|
|
)
|
|
dist_op_impls = find_compatible_distributed_operator_impls(
|
|
new_dist_op, fwd=False
|
|
)
|
|
if dist_op_impls is not None:
|
|
output_dims_mapping_candidates.append(dims_mapping_list)
|
|
|
|
if not input_dims_mapping_candidates and output_dims_mapping_candidates:
|
|
inout_dims_mapping_generated = [
|
|
[[[-2]]],
|
|
output_dims_mapping_candidates,
|
|
]
|
|
elif (
|
|
input_dims_mapping_candidates and not output_dims_mapping_candidates
|
|
):
|
|
inout_dims_mapping_generated = [
|
|
input_dims_mapping_candidates,
|
|
[[[-2]]],
|
|
]
|
|
elif (
|
|
not input_dims_mapping_candidates
|
|
and not output_dims_mapping_candidates
|
|
):
|
|
inout_dims_mapping_generated = [[[[-2]]], [[[-2]]]]
|
|
else:
|
|
inout_dims_mapping_generated = [
|
|
input_dims_mapping_candidates,
|
|
output_dims_mapping_candidates,
|
|
]
|
|
# Find valid dims_mapping generated for both inputs and outputs
|
|
cached_dist_attr_candidates = []
|
|
for inout_dims_mapping_list in itertools.product(
|
|
*inout_dims_mapping_generated
|
|
):
|
|
assert len(inout_dims_mapping_list) == 2
|
|
if input_dims_mapping_candidates:
|
|
assert len(inout_dims_mapping_list[0]) == len(input_names)
|
|
if output_dims_mapping_candidates:
|
|
assert len(inout_dims_mapping_list[1]) == len(output_names)
|
|
# set the dims_mappings for inputs
|
|
for i, dims_mapping in enumerate(inout_dims_mapping_list[0]):
|
|
if dims_mapping != [-2]:
|
|
new_op_dist_attr.set_input_dims_mapping(
|
|
input_names[i], dims_mapping
|
|
)
|
|
# set the dims_mappings for outputs
|
|
for i, dims_mapping in enumerate(inout_dims_mapping_list[1]):
|
|
if dims_mapping != [-2]:
|
|
new_op_dist_attr.set_output_dims_mapping(
|
|
output_names[i], dims_mapping
|
|
)
|
|
new_dist_op = DistributedOperator(
|
|
dist_op.serial_op, new_op_dist_attr
|
|
)
|
|
dist_op_impls = find_compatible_distributed_operator_impls(
|
|
new_dist_op, partial=False
|
|
)
|
|
if dist_op_impls is None:
|
|
continue
|
|
for dist_op_impl in dist_op_impls:
|
|
new_op_dist_attr.impl_type = dist_op_impl.type
|
|
new_op_dist_attr.impl_idx = dist_op_impl.idx
|
|
cached_dist_attr_candidates.append(
|
|
copy.deepcopy(new_op_dist_attr)
|
|
)
|
|
self._cached_candidates_info[key].append(cached_dist_attr_candidates)
|
|
return self._cached_candidates_info[key][2]
|
|
|
|
def construct_space(self):
|
|
inter_node_partitions, intra_node_partitions = self._partition_devices(
|
|
self._num_machines, self._num_devices_per_machine
|
|
)
|
|
self._space.choice(
|
|
"inter_node_partitions",
|
|
inter_node_partitions,
|
|
default=inter_node_partitions[0],
|
|
)
|
|
self._space.choice(
|
|
"intra_node_partitions",
|
|
intra_node_partitions,
|
|
default=intra_node_partitions[0],
|
|
)
|
|
|
|
dist_ops = self._dist_context._dist_ops_for_program
|
|
for op_id, dist_op in dist_ops.items():
|
|
op_type = dist_op.serial_op.type
|
|
if self._include_op_types:
|
|
if op_type in self._include_op_types:
|
|
self._concerned_dist_ops[op_id] = dist_op
|
|
else:
|
|
self._concerned_dist_ops[op_id] = dist_op
|
|
|
|
for op_id, dist_op in self._concerned_dist_ops.items():
|
|
op_type = dist_op.serial_op.type
|
|
if op_type in self._exclude_op_types:
|
|
del self._concerned_dist_ops[op_id]
|
|
|
|
print(
|
|
"Number of the concerned dist ops",
|
|
len(self._concerned_dist_ops),
|
|
flush=True,
|
|
)
|
|
search_space = 1
|
|
for op_id, dist_op in self._concerned_dist_ops.items():
|
|
op_dist_attr_candidates = self._generate_dist_attr_candidates(
|
|
op_id, dist_op
|
|
)
|
|
search_space *= len(op_dist_attr_candidates)
|
|
self._space.choice(
|
|
str(op_id),
|
|
op_dist_attr_candidates,
|
|
default=op_dist_attr_candidates[0],
|
|
)
|
|
|
|
def _compute_values_hash(self, values):
|
|
keys = sorted(values.keys())
|
|
s = "".join(str(k) + "=" + str(values[k]) for k in keys)
|
|
return hashlib.sha256(s.encode("utf-8")).hexdigest()[:32]
|
|
|
|
def _random_values(self):
|
|
space = TunableSpace()
|
|
collisions = 0
|
|
while True:
|
|
for v in self._space.variables.values():
|
|
space._register(v)
|
|
space.values[v.name] = v.random(self._seed_state)
|
|
self._seed_state += 1
|
|
values = space.values
|
|
values_hash = self._compute_values_hash(values)
|
|
if values_hash in self._tried_values:
|
|
collisions += 1
|
|
if collisions > self._max_collisions:
|
|
return None
|
|
continue
|
|
self._tried_values.add(values_hash)
|
|
break
|
|
return values
|
|
|
|
def _populate_space(self):
|
|
values = self._random_values()
|
|
if values is None:
|
|
return {"status": TrialStatus.STOPPED, "values": None}
|
|
return {"status": TrialStatus.RUNNING, "values": values}
|
|
|
|
def _create_trial(self):
|
|
trial_id = f"{{:0{len(str(self._max_trials))}d}}"
|
|
trial_id = trial_id.format(self._num_trials)
|
|
|
|
if self._max_trials and self._num_trials >= self._max_trials:
|
|
status = TrialStatus.STOPPED
|
|
values = None
|
|
else:
|
|
results = self._populate_space()
|
|
status = results["status"]
|
|
values = results["values"]
|
|
|
|
space = TunableSpace()
|
|
space.variables = self._space.variables
|
|
space.values = values
|
|
trial = Trial(tunable_space=space, trial_id=trial_id, status=status)
|
|
self._num_trials += 1
|
|
return trial
|
|
|
|
def _generate_pipeline_starts(self, process_mesh_list):
|
|
total_ops = len(self._dist_context._dist_ops_for_program)
|
|
total_stages = len(process_mesh_list)
|
|
ops_per_stage = total_ops // total_stages
|
|
if ops_per_stage == 0:
|
|
return None
|
|
# Compute the initial pipeline starts
|
|
pipeline_starts = []
|
|
start = 0
|
|
pipeline_starts.append(0)
|
|
# The pipeline_starts have total_stages+1 items, and
|
|
# at least have 2 items.
|
|
for _ in process_mesh_list:
|
|
start += ops_per_stage
|
|
pipeline_starts.append(start)
|
|
pipeline_starts[-1] = total_ops
|
|
# Adjust the pipeline starts by random selection
|
|
directions = []
|
|
sizes = []
|
|
half_ops_per_stage = ops_per_stage // 2
|
|
if half_ops_per_stage > 0 and total_stages > 1:
|
|
new_pipeline_starts = []
|
|
# Don't change the first start
|
|
new_pipeline_starts.append(0)
|
|
# Consider the starts except the first and the last one
|
|
for _ in pipeline_starts[1:-1]:
|
|
directions.append(Boolean("direction"))
|
|
sizes.append(
|
|
IntRange(
|
|
"size", start=0, stop=half_ops_per_stage, endpoint=True
|
|
)
|
|
)
|
|
for i, start in enumerate(pipeline_starts[1:-1]):
|
|
direction = directions[i].random(self._seed)
|
|
size = sizes[i].random(self._seed)
|
|
if direction:
|
|
# Subtract 1 from size to avoid the overlapping of new starts
|
|
new_start = start - (size - 1)
|
|
else:
|
|
new_start = start + size
|
|
new_pipeline_starts.append(new_start)
|
|
# Don't change the last start
|
|
new_pipeline_starts.append(pipeline_starts[-1])
|
|
# Validate the new starts
|
|
print(
|
|
"Adjusted pipeline starts",
|
|
new_pipeline_starts,
|
|
half_ops_per_stage,
|
|
pipeline_starts,
|
|
flush=True,
|
|
)
|
|
for i, new_start in enumerate(new_pipeline_starts[1:]):
|
|
assert new_start > new_pipeline_starts[i]
|
|
return new_pipeline_starts
|
|
else:
|
|
print(
|
|
"Non-adjusted pipeline starts",
|
|
pipeline_starts,
|
|
half_ops_per_stage,
|
|
flush=True,
|
|
)
|
|
return pipeline_starts
|
|
|
|
def _apply_pipeline_partition(self, process_mesh_list):
|
|
op_id_to_process_mesh = {}
|
|
total_ops = len(self._dist_context._dist_ops_for_program)
|
|
total_stages = len(process_mesh_list)
|
|
ops_per_stage = total_ops // total_stages
|
|
if ops_per_stage == 0:
|
|
return None
|
|
pipeline_starts = self._generate_pipeline_starts(process_mesh_list)
|
|
start_idx = 1
|
|
sorted_op_ids = sorted(self._dist_context._dist_ops_for_program.keys())
|
|
for idx, op_id in enumerate(sorted_op_ids):
|
|
if idx < pipeline_starts[start_idx]:
|
|
op_id_to_process_mesh[op_id] = process_mesh_list[start_idx - 1]
|
|
else:
|
|
start_idx += 1
|
|
op_id_to_process_mesh[op_id] = process_mesh_list[start_idx - 1]
|
|
return op_id_to_process_mesh
|
|
|
|
def _amend_dist_attr(self):
|
|
# 1) Reshape the process mesh of [1, x] to [x] or [x, 1] to [x],
|
|
# and amend the corresponding dims_mapping.
|
|
# 2) Set the dim_mapping to -1 when the shape cannot be divided
|
|
# by the corresponding processes.
|
|
for dist_op in self._dist_context._dist_ops_for_program.values():
|
|
dist_attr = dist_op.dist_attr
|
|
process_mesh = dist_attr.process_mesh
|
|
if process_mesh is None:
|
|
continue
|
|
assert process_mesh.ndim == 2
|
|
dim_of_one = None
|
|
dim_of_other = None
|
|
if process_mesh.shape[0] == 1:
|
|
dim_of_one = 0
|
|
dim_of_other = 1
|
|
elif process_mesh.shape[1] == 1:
|
|
dim_of_one = 1
|
|
dim_of_other = 0
|
|
|
|
if dim_of_one is not None:
|
|
dist_attr.process_mesh = ProcessMesh(process_mesh.process_ids)
|
|
self._dist_context.add_process_mesh(dist_attr.process_mesh)
|
|
|
|
for arg_name in dist_attr.inputs_dist_attrs.keys():
|
|
new_dims_mapping = []
|
|
dims_mapping = dist_attr.get_input_dims_mapping(arg_name)
|
|
for dim_mapping in dims_mapping:
|
|
if dim_mapping == dim_of_one:
|
|
new_dims_mapping.append(-1)
|
|
elif dim_mapping == dim_of_other:
|
|
new_dims_mapping.append(0)
|
|
else:
|
|
new_dims_mapping.append(dim_mapping)
|
|
dist_attr.set_input_dims_mapping(arg_name, new_dims_mapping)
|
|
|
|
dims_mapping = dist_attr.get_input_dims_mapping(arg_name)
|
|
# dynamic_dims = dist_attr.get_input_dynamic_dims(arg_name)
|
|
process_mesh = dist_attr.process_mesh
|
|
process_shape = process_mesh.shape
|
|
tensor = dist_op.get_serial_input(arg_name)
|
|
if dims_mapping:
|
|
tensor_shape = tensor.shape
|
|
else:
|
|
continue
|
|
for i, dim_mapping in enumerate(dims_mapping):
|
|
# if dim_mapping != -1 \
|
|
# and (tensor_shape[i] % process_shape[dim_mapping] != 0 \
|
|
# or dynamic_dims[i] == 1):
|
|
if dim_mapping != -1 and (
|
|
tensor_shape[i] % process_shape[dim_mapping] != 0
|
|
):
|
|
dims_mapping[i] = -1
|
|
# it is a fix-bug
|
|
if dim_mapping != -1 and process_shape[dim_mapping] == 1:
|
|
dims_mapping[i] = -1
|
|
|
|
for arg_name in dist_attr.outputs_dist_attrs.keys():
|
|
new_dims_mapping = []
|
|
dims_mapping = dist_attr.get_output_dims_mapping(arg_name)
|
|
for dim_mapping in dims_mapping:
|
|
if dim_mapping == dim_of_one:
|
|
new_dims_mapping.append(-1)
|
|
elif dim_mapping == dim_of_other:
|
|
new_dims_mapping.append(0)
|
|
else:
|
|
new_dims_mapping.append(dim_mapping)
|
|
dist_attr.set_output_dims_mapping(arg_name, new_dims_mapping)
|
|
|
|
dims_mapping = dist_attr.get_output_dims_mapping(arg_name)
|
|
# dynamic_dims = dist_attr.get_output_dynamic_dims(arg_name)
|
|
process_mesh = dist_attr.process_mesh
|
|
process_shape = process_mesh.shape
|
|
|
|
tensor = dist_op.get_serial_output(arg_name)
|
|
if dims_mapping:
|
|
tensor_shape = tensor.shape
|
|
else:
|
|
continue
|
|
for i, dim_mapping in enumerate(dims_mapping):
|
|
if dim_mapping != -1 and (
|
|
tensor_shape[i] % process_shape[dim_mapping] != 0
|
|
):
|
|
dims_mapping[i] = -1
|
|
# it is a fix-bug
|
|
if dim_mapping != -1 and process_shape[dim_mapping] == 1:
|
|
dims_mapping[i] = -1
|
|
dist_op_impls = find_compatible_distributed_operator_impls(
|
|
dist_op, partial=False
|
|
)
|
|
serial_op_type = dist_op.serial_op.type
|
|
|
|
if dist_op_impls is not None and (
|
|
serial_op_type != "fused_softmax_mask_upper_triangle"
|
|
or self._check_fused_softmax_mask_upper_triangle(dist_op)
|
|
):
|
|
dist_op.dist_attr.impl_type = dist_op_impls[0].type
|
|
dist_op.dist_attr.impl_idx = dist_op_impls[0].idx
|
|
else:
|
|
# Use the default dist op impl
|
|
for arg_name in dist_attr.inputs_dist_attrs.keys():
|
|
dims_mapping = dist_attr.get_input_dims_mapping(arg_name)
|
|
for i, _ in enumerate(dims_mapping):
|
|
dims_mapping[i] = -1
|
|
for arg_name in dist_attr.outputs_dist_attrs.keys():
|
|
dims_mapping = dist_attr.get_output_dims_mapping(arg_name)
|
|
for i, _ in enumerate(dims_mapping):
|
|
dims_mapping[i] = -1
|
|
dist_op.dist_attr.impl_type = "default"
|
|
dist_op.dist_attr.impl_idx = 0
|
|
|
|
def _check_fused_softmax_mask_upper_triangle(self, dist_op):
|
|
"""The last_but_one dim should be equal to last dim."""
|
|
input_name = dist_op.serial_op.input_arg_names[0]
|
|
input_dims_mapping = dist_op.dist_attr.get_input_dims_mapping(
|
|
input_name
|
|
)
|
|
topology = dist_op.dist_attr.process_mesh.shape
|
|
input_tensor = dist_op.get_serial_input(input_name)
|
|
last_but_one_dim = (
|
|
input_tensor.shape[-2] // topology[input_dims_mapping[-2]]
|
|
if input_dims_mapping[-2] != -1
|
|
else input_tensor.shape[-2]
|
|
)
|
|
last_dim = (
|
|
input_tensor.shape[-1] // topology[input_dims_mapping[-1]]
|
|
if input_dims_mapping[-1] != -1
|
|
else input_tensor.shape[-1]
|
|
)
|
|
if last_but_one_dim == last_dim:
|
|
return True
|
|
return False
|
|
|
|
def _eval_trial(self, trial):
|
|
if self._num_trials == 0:
|
|
num_prev_trials = 0
|
|
else:
|
|
num_prev_trials = self._num_trials - 1
|
|
|
|
results = None
|
|
|
|
start_time = time.time()
|
|
|
|
inter_node_partition = trial.space.values["inter_node_partitions"]
|
|
intra_node_partition = trial.space.values["intra_node_partitions"]
|
|
process_mesh_list = self._generate_process_mesh_list(
|
|
inter_node_partition, intra_node_partition
|
|
)
|
|
print("\tprocess_mesh list", process_mesh_list, flush=True)
|
|
op_id_to_process_mesh = self._apply_pipeline_partition(
|
|
process_mesh_list
|
|
)
|
|
if op_id_to_process_mesh is None:
|
|
print("Operators are less than pipeline stages", flush=True)
|
|
return results
|
|
|
|
op_id_to_dist_attr = {}
|
|
for name, value in trial.space.values.items():
|
|
if (
|
|
name != "inter_node_partitions"
|
|
and name != "intra_node_partitions"
|
|
):
|
|
op_id_to_dist_attr[int(name)] = value
|
|
|
|
end_time = time.time()
|
|
cur_sample_time = end_time - start_time
|
|
self._sample_time = (
|
|
num_prev_trials * self._sample_time + cur_sample_time
|
|
) / self._num_trials
|
|
print(
|
|
"\tsample_time",
|
|
num_prev_trials,
|
|
self._num_trials,
|
|
self._sample_time,
|
|
cur_sample_time,
|
|
flush=True,
|
|
)
|
|
|
|
assert len(op_id_to_process_mesh) == len(op_id_to_dist_attr)
|
|
|
|
start_time = time.time()
|
|
for op_id, process_mesh in op_id_to_process_mesh.items():
|
|
dist_op = self._dist_context._dist_ops_for_program[op_id]
|
|
dist_op.dist_attr = copy.deepcopy(op_id_to_dist_attr[op_id])
|
|
assert (
|
|
dist_op.dist_attr.impl_type
|
|
== op_id_to_dist_attr[op_id].impl_type
|
|
)
|
|
assert (
|
|
dist_op.dist_attr.impl_idx == op_id_to_dist_attr[op_id].impl_idx
|
|
)
|
|
dist_op.dist_attr.process_mesh = ProcessMesh(process_mesh)
|
|
self._amend_dist_attr()
|
|
|
|
self._completer._complete_tensor_dist_attr_by_op()
|
|
|
|
self._dist_context.block_state.parse_forward_blocks(
|
|
self._dist_context.serial_main_program
|
|
)
|
|
|
|
end_time = time.time()
|
|
cur_complete_time = end_time - start_time
|
|
self._complete_time = (
|
|
num_prev_trials * self._complete_time + cur_complete_time
|
|
) / self._num_trials
|
|
print(
|
|
"\tcomplete_time",
|
|
num_prev_trials,
|
|
self._num_trials,
|
|
self._complete_time,
|
|
cur_complete_time,
|
|
flush=True,
|
|
)
|
|
|
|
start_time = time.time()
|
|
estimate_time = self._estimate_trial()
|
|
end_time = time.time()
|
|
cur_estimate_time = end_time - start_time
|
|
self._estimate_time = (
|
|
num_prev_trials * self._estimate_time + cur_estimate_time
|
|
) / self._num_trials
|
|
print(
|
|
"\testimate_time",
|
|
num_prev_trials,
|
|
self._num_trials,
|
|
self._estimate_time,
|
|
cur_estimate_time,
|
|
estimate_time,
|
|
flush=True,
|
|
)
|
|
|
|
results = {"estimate_time": estimate_time}
|
|
return results
|
|
|
|
def _update_trail(self, trial, metrics, step=0):
|
|
for metric_name, metric_value in metrics.items():
|
|
trial.recorder.update(metric_name, metric_value, step=step)
|
|
return trial.status
|
|
|
|
def _estimate_trial(self):
|
|
assert self._cluster is not None
|
|
if self._mode == "eval":
|
|
self._estimator = CostEstimator(
|
|
self._dist_context.serial_main_program,
|
|
self._cluster,
|
|
loop_count=self._loop_count,
|
|
)
|
|
elif self._mode == "predict":
|
|
self._estimator = CostEstimator(
|
|
self._dist_context.serial_main_program,
|
|
self._cluster,
|
|
loop_count=self._loop_count,
|
|
)
|
|
elif self._mode == "train":
|
|
# get serial main program with backward
|
|
serial_main_program = self._dist_context.serial_main_program
|
|
serial_startup_program = self._dist_context.serial_startup_program
|
|
serial_optimizer = self._dist_context.serial_optimizer
|
|
|
|
# Generate backward
|
|
serial_loss = self._dist_context.serial_fetch_vars["loss"][0]
|
|
params_grads = self._parallelizer._generate_backward(
|
|
serial_main_program, serial_startup_program, serial_loss
|
|
)
|
|
|
|
# Generate optimizer
|
|
optimizer_ops = self._parallelizer._generate_optimizer(
|
|
serial_main_program,
|
|
serial_startup_program,
|
|
serial_optimizer,
|
|
params_grads,
|
|
)
|
|
self._estimator = CostEstimator(
|
|
serial_main_program, self._cluster, loop_count=self._loop_count
|
|
)
|
|
|
|
max_memory = self._estimator._estimate_max_memory_by_dist_op(
|
|
self._dist_context
|
|
)
|
|
print("\tmax_memory", f"{max_memory:,}", flush=True)
|
|
# The max memory must be less than 80% 32GB (hard code)
|
|
if max_memory > 32 * 0.8 * 1024 * 1024 * 1024:
|
|
return math.inf
|
|
else:
|
|
global_cost = self._estimator.estimate(self._dist_context)
|
|
return global_cost.time
|
|
|
|
def _store_init_parallel_strategy(self):
|
|
# If there is no annotation information, use the dp as the initial parallel strategy.
|
|
# TODO: we should need a better way to set up the initial parallel strategy.
|
|
if (
|
|
not self._dist_context.has_annotation
|
|
or not self._dist_context.process_meshes
|
|
):
|
|
ranks = self._num_machines * self._num_devices_per_machine
|
|
tensor_node = self._dist_context._serial_ordered_tensor_nodes[0]
|
|
tensor_node_id = _node_id(tensor_node)
|
|
tensor = self._dist_context._dist_tensors_for_graph[
|
|
tensor_node_id
|
|
].serial_tensor
|
|
tensor_dist_attr = self._dist_context._dist_tensors_for_graph[
|
|
tensor_node_id
|
|
].dist_attr
|
|
tensor_dist_attr.process_mesh = ProcessMesh(list(range(ranks)))
|
|
self._dist_context._process_meshes.append(
|
|
tensor_dist_attr.process_mesh
|
|
)
|
|
tensor_dist_attr.dims_mapping = [0] + [
|
|
-1 for _ in range(len(tensor.shape) - 1)
|
|
]
|
|
tensor_dist_attr.mark_annotated("process_mesh")
|
|
tensor_dist_attr.mark_annotated("dims_mapping")
|
|
print("Use dp as the init parallel strategy!", flush=True)
|
|
|
|
# Do the sharding propagation
|
|
self._completer.complete_forward_annotation()
|
|
self._dist_context.block_state.parse_forward_blocks(
|
|
self._dist_context.serial_main_program
|
|
)
|
|
|
|
# Backup the initial parallel strategy
|
|
self._init_parallel_strategy[0] = copy.deepcopy(
|
|
self._dist_context._dist_tensors_for_program
|
|
)
|
|
self._init_parallel_strategy[1] = copy.deepcopy(
|
|
self._dist_context._dist_ops_for_program
|
|
)
|
|
self._init_parallel_strategy[2] = copy.deepcopy(
|
|
self._dist_context.process_meshes
|
|
)
|
|
|
|
# Initialize the best parallel strategy to the initial one
|
|
self._best_parallel_strategy[0] = copy.deepcopy(
|
|
self._dist_context._dist_tensors_for_program
|
|
)
|
|
self._best_parallel_strategy[1] = copy.deepcopy(
|
|
self._dist_context._dist_ops_for_program
|
|
)
|
|
self._best_parallel_strategy[2] = copy.deepcopy(
|
|
self._dist_context._process_meshes
|
|
)
|
|
|
|
def _store_best_parallel_strategy(self):
|
|
# Swap the best and the current parallel strategy
|
|
tmp = [None, None, None]
|
|
tmp[0] = self._best_parallel_strategy[0]
|
|
tmp[1] = self._best_parallel_strategy[1]
|
|
tmp[2] = self._best_parallel_strategy[2]
|
|
self._best_parallel_strategy[0] = (
|
|
self._dist_context._dist_tensors_for_program
|
|
)
|
|
self._best_parallel_strategy[1] = (
|
|
self._dist_context._dist_ops_for_program
|
|
)
|
|
self._best_parallel_strategy[2] = self._dist_context._process_meshes
|
|
self._dist_context._dist_tensors_for_program = tmp[0]
|
|
self._dist_context._dist_ops_for_program = tmp[1]
|
|
self._dist_context._process_meshes = tmp[2]
|
|
|
|
def tune(self):
|
|
global_start_time = time.time()
|
|
self._dist_context._backup(serial=True, dist=True)
|
|
# This store statement must follow the above backup statement
|
|
self._store_init_parallel_strategy()
|
|
init_time = self._estimate_trial() # estimate_trial when init
|
|
# We have to restore the distributed context, because the estimation of one trail need to
|
|
# generate the backward and update parts. Since we will do the tuning process,
|
|
# here we only need to reset all distributed information to the default one.
|
|
self._dist_context._restore(
|
|
serial=True,
|
|
serial_mode="to_backup",
|
|
dist=True,
|
|
dist_mode="to_default",
|
|
)
|
|
|
|
best_time = init_time
|
|
start_time = time.time()
|
|
self.construct_space()
|
|
end_time = time.time()
|
|
print(
|
|
"construct_space time",
|
|
self._num_trials,
|
|
end_time - start_time,
|
|
flush=True,
|
|
)
|
|
create_trial_time = 0.0
|
|
eval_trial_time = 0.0
|
|
self._sample_time = 0.0
|
|
self._complete_time = 0.0
|
|
self._estimate_time = 0.0
|
|
while True:
|
|
start_time = time.time()
|
|
trial = self._create_trial()
|
|
if self._num_trials == 0:
|
|
num_prev_trials = 0
|
|
else:
|
|
num_prev_trials = self._num_trials - 1
|
|
end_time = time.time()
|
|
cur_create_trial_time = end_time - start_time
|
|
create_trial_time = (
|
|
num_prev_trials * create_trial_time + cur_create_trial_time
|
|
) / self._num_trials
|
|
print(
|
|
"create_trial time",
|
|
num_prev_trials,
|
|
self._num_trials,
|
|
create_trial_time,
|
|
cur_create_trial_time,
|
|
flush=True,
|
|
)
|
|
if trial.status == TrialStatus.STOPPED:
|
|
break
|
|
# We need to backup the distributed context, because the evaluation of one trail will
|
|
# generate the backward and update parts which may change the context.
|
|
# However, the distributed information of the context aren't backup since a new one is used.
|
|
self._dist_context._backup(serial=True, dist=False)
|
|
|
|
start_time = time.time()
|
|
results = self._eval_trial(trial)
|
|
end_time = time.time()
|
|
cur_eval_trial_time = end_time - start_time
|
|
eval_trial_time = (
|
|
num_prev_trials * eval_trial_time + cur_eval_trial_time
|
|
) / self._num_trials
|
|
print(
|
|
"eval_trial time",
|
|
num_prev_trials,
|
|
self._num_trials,
|
|
eval_trial_time,
|
|
cur_eval_trial_time,
|
|
"\n",
|
|
flush=True,
|
|
)
|
|
|
|
cur_time = results["estimate_time"]
|
|
if cur_time < best_time:
|
|
self._update_trail(trial, results)
|
|
self._store_best_parallel_strategy()
|
|
best_time = cur_time
|
|
# We need to restore the distributed context and reset the distributed information to the default.
|
|
self._dist_context._restore(
|
|
serial=True,
|
|
serial_mode="to_backup",
|
|
dist=True,
|
|
dist_mode="to_default",
|
|
)
|
|
# Select the best parallel strategy
|
|
self._dist_context._dist_tensors_for_program = (
|
|
self._best_parallel_strategy[0]
|
|
)
|
|
self._dist_context._dist_ops_for_program = self._best_parallel_strategy[
|
|
1
|
|
]
|
|
self._dist_context._process_meshes = self._best_parallel_strategy[2]
|