chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
@@ -0,0 +1,532 @@
|
||||
# Copyright (c) 2023 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 collections
|
||||
import copy
|
||||
import os
|
||||
import pickle
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.base import core
|
||||
from paddle.base.framework import Program
|
||||
from paddle.distributed.auto_parallel.static.converter import Converter
|
||||
from paddle.distributed.auto_parallel.static.dist_context import (
|
||||
get_default_distributed_context,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.utils import (
|
||||
is_backward_op,
|
||||
is_forward_op,
|
||||
is_loss_op,
|
||||
)
|
||||
from paddle.static.io import deserialize_program
|
||||
|
||||
_valid_types = [
|
||||
core.VarDesc.VarType.DENSE_TENSOR,
|
||||
core.VarDesc.VarType.SELECTED_ROWS,
|
||||
core.VarDesc.VarType.DENSE_TENSOR_ARRAY,
|
||||
]
|
||||
|
||||
paddle.enable_static()
|
||||
|
||||
|
||||
class AutoAlignTool:
|
||||
"""
|
||||
This is an automatic parallel precision alignment tool。
|
||||
"""
|
||||
|
||||
def __init__(self, program: Program, step=1, fetch_list=None):
|
||||
"""Set some initialization information of the tool.
|
||||
step: Step when returning a specific variable name。
|
||||
fetch_list: initialization fetch_list.When a specific step is not reached, return this.
|
||||
It can combine with Engine class。
|
||||
example:in Engine.fit function,like this
|
||||
try:
|
||||
fetch_list = []
|
||||
align_tool = AutoAlignTool(self.main_program, 0, fetch_names)
|
||||
level = 0
|
||||
fetch_list = align_tool.get_var(level, step)
|
||||
outs = self._executor.run(
|
||||
self.main_program,
|
||||
fetch_list=fetch_list,
|
||||
use_program_cache=self._strategy.use_cache,
|
||||
return_numpy=self._strategy.return_numpy,
|
||||
)
|
||||
if fetch_list != fetch_names:
|
||||
align_tool.save(dir_path, outs, fetch_list, self._dist_contexts["train"], self.serial)
|
||||
exit(0)
|
||||
except core.EOFException:
|
||||
break
|
||||
"""
|
||||
assert isinstance(program, Program)
|
||||
self._program = program
|
||||
self._blocks = program.blocks
|
||||
self._step = step
|
||||
self._fetch_list = fetch_list
|
||||
assert self._blocks is not None
|
||||
|
||||
def set_step(self, step):
|
||||
self._step = step
|
||||
|
||||
def get_var(self, level, step):
|
||||
"""
|
||||
level must be in [0,1,2,3,4,5].
|
||||
"""
|
||||
if step != self._step or step == -1:
|
||||
return self._fetch_list
|
||||
if level == 0:
|
||||
return self.get_loss_lr_var()
|
||||
elif level == 1:
|
||||
return self.get_data_var()
|
||||
elif level == 2:
|
||||
return self.get_param_var()
|
||||
elif level == 3:
|
||||
return self.get_param_grad_var()
|
||||
elif level == 4:
|
||||
return self.get_forward_tmp_var()
|
||||
elif level == 5:
|
||||
return self.get_backward_tmp_var()
|
||||
else:
|
||||
raise ValueError
|
||||
|
||||
def set_program(self, program: Program):
|
||||
assert isinstance(program, Program)
|
||||
self._program = program
|
||||
self._blocks = program.blocks
|
||||
assert self._blocks is not None
|
||||
|
||||
def get_loss_lr_var(self):
|
||||
"""
|
||||
Returns the variable name of learning rate and loss
|
||||
"""
|
||||
fetch_set = set()
|
||||
loss_ops = []
|
||||
for block in self._blocks:
|
||||
for op in block.ops:
|
||||
if is_loss_op(op):
|
||||
assert len(op.desc.output_arg_names()) == 1, (
|
||||
"loss op should only output loss var"
|
||||
)
|
||||
loss_ops.append(op)
|
||||
|
||||
for block in self._blocks:
|
||||
for varname in block.vars:
|
||||
var = block._find_var_recursive(varname)
|
||||
|
||||
if var is None or var.type not in _valid_types:
|
||||
continue
|
||||
|
||||
if "learning_rate" in var.name:
|
||||
fetch_set.add(var.name)
|
||||
|
||||
for loss_op in loss_ops:
|
||||
fetch_set.add(loss_op.output_arg_names[0])
|
||||
|
||||
return list(fetch_set)
|
||||
|
||||
def get_data_var(self):
|
||||
"""
|
||||
Returns the variable name of data.
|
||||
"""
|
||||
fetch_set = set()
|
||||
for block in self._blocks:
|
||||
for varname in block.vars:
|
||||
var = block._find_var_recursive(varname)
|
||||
|
||||
if var is None or var.type not in _valid_types:
|
||||
continue
|
||||
|
||||
if var.is_data:
|
||||
fetch_set.add(var.name)
|
||||
return list(fetch_set)
|
||||
|
||||
def get_param_var(self):
|
||||
"""
|
||||
Returns the variable name of parameters.
|
||||
"""
|
||||
fetch_set = set()
|
||||
for block in self._blocks:
|
||||
for op in block.ops:
|
||||
if is_backward_op(op):
|
||||
break
|
||||
for varname in op.input_arg_names + op.output_arg_names:
|
||||
var = block._find_var_recursive(varname)
|
||||
if var is None or var.type not in _valid_types:
|
||||
continue
|
||||
if var.is_parameter:
|
||||
fetch_set.add(varname)
|
||||
|
||||
return list(fetch_set)
|
||||
|
||||
def get_param_grad_var(self):
|
||||
"""
|
||||
Returns the variable name of parameters' gradient.
|
||||
"""
|
||||
fetch_set = set()
|
||||
for block in self._blocks:
|
||||
for op in block.ops:
|
||||
if is_forward_op(op):
|
||||
continue
|
||||
for varname in op.input_arg_names + op.output_arg_names:
|
||||
if "@GRAD" not in varname:
|
||||
continue
|
||||
fwd_varname = varname.split("@GRAD")[0]
|
||||
fwd_var = block._find_var_recursive(fwd_varname)
|
||||
if fwd_var is None or fwd_var.type not in _valid_types:
|
||||
continue
|
||||
if fwd_var.is_parameter is False:
|
||||
continue
|
||||
var = block._find_var_recursive(varname)
|
||||
if var is None or var.type not in _valid_types:
|
||||
continue
|
||||
fetch_set.add(varname)
|
||||
|
||||
return list(fetch_set)
|
||||
|
||||
def get_forward_tmp_var(self):
|
||||
"""
|
||||
Returns the name of the temporary variable in the forward propagation
|
||||
"""
|
||||
fetch_set = set()
|
||||
loss_lr_list = self.get_loss_lr_var()
|
||||
for block in self._blocks:
|
||||
for op in block.ops:
|
||||
if is_backward_op(op):
|
||||
break
|
||||
for varname in op.input_arg_names + op.output_arg_names:
|
||||
if varname in loss_lr_list:
|
||||
continue
|
||||
var = block._find_var_recursive(varname)
|
||||
if var is None or var.type not in _valid_types:
|
||||
continue
|
||||
if var.is_data or var.is_parameter:
|
||||
continue
|
||||
fetch_set.add(varname)
|
||||
|
||||
return list(fetch_set)
|
||||
|
||||
def get_backward_tmp_var(self):
|
||||
"""
|
||||
Returns the name of a temporary variable in back-propagation
|
||||
"""
|
||||
fetch_set = set()
|
||||
loss_lr_list = self.get_loss_lr_var()
|
||||
forward_tmp_list = self.get_forward_tmp_var()
|
||||
for block in self._blocks:
|
||||
for op in block.ops:
|
||||
if is_backward_op(op):
|
||||
for varname in op.input_arg_names + op.output_arg_names:
|
||||
if (
|
||||
varname in loss_lr_list
|
||||
or varname in forward_tmp_list
|
||||
):
|
||||
continue
|
||||
if "@GRAD" in varname:
|
||||
fwd_varname = varname.split("@GRAD")[0]
|
||||
fwd_var = block._find_var_recursive(fwd_varname)
|
||||
if (
|
||||
fwd_var is not None
|
||||
and fwd_var.type in _valid_types
|
||||
):
|
||||
if fwd_var.is_parameter:
|
||||
continue
|
||||
var = block._find_var_recursive(varname)
|
||||
if var is None or var.type not in _valid_types:
|
||||
continue
|
||||
if var.is_data or var.is_parameter:
|
||||
continue
|
||||
fetch_set.add(varname)
|
||||
|
||||
return list(fetch_set)
|
||||
|
||||
def save(self, save_dir, vars, fetch_list, dist_context=None):
|
||||
"""
|
||||
save fetch variables, distributed properties of variables and program.
|
||||
"""
|
||||
if os.path.exists(save_dir) is False:
|
||||
os.mkdir(save_dir)
|
||||
if dist_context is None:
|
||||
dist_context = get_default_distributed_context()
|
||||
assert os.path.exists(save_dir)
|
||||
if dist.get_world_size() == 1:
|
||||
vars_path = os.path.join(save_dir, "vars.pkl")
|
||||
program_path = os.path.join(save_dir, "program.pdmodel")
|
||||
dist_attr_path = os.path.join(save_dir, "dist_attr.pkl")
|
||||
else:
|
||||
vars_path = os.path.join(
|
||||
save_dir, f"vars_rank{dist.get_rank()}.pkl"
|
||||
)
|
||||
program_path = os.path.join(
|
||||
save_dir, f"program_rank{dist.get_rank()}.pdmodel"
|
||||
)
|
||||
dist_attr_path = os.path.join(
|
||||
save_dir, f"dist_attr_rank{dist.get_rank()}.pkl"
|
||||
)
|
||||
if vars is not None:
|
||||
vars_dict = {}
|
||||
assert len(fetch_list) == len(vars)
|
||||
for i in range(len(fetch_list)):
|
||||
if vars[i] is None:
|
||||
continue
|
||||
vars_dict[fetch_list[i]] = vars[i]
|
||||
with open(vars_path, "wb") as f:
|
||||
pickle.dump(vars_dict, f)
|
||||
dist_attr = {}
|
||||
for var in self._program.list_vars():
|
||||
if var.name not in fetch_list:
|
||||
continue
|
||||
tensor_dist_attr = (
|
||||
dist_context.get_tensor_dist_attr_for_program(var)
|
||||
)
|
||||
if tensor_dist_attr is None:
|
||||
continue
|
||||
process_mesh = tensor_dist_attr.process_mesh
|
||||
dims_mapping = tensor_dist_attr.dims_mapping
|
||||
dist_attr[var.name] = {
|
||||
"process_shape": process_mesh.shape,
|
||||
"process_group": process_mesh.process_ids,
|
||||
"dims_mapping": dims_mapping,
|
||||
}
|
||||
if len(dist_attr) > 0:
|
||||
with open(dist_attr_path, "wb") as f:
|
||||
pickle.dump(dist_attr, f)
|
||||
if self._program is not None:
|
||||
with open(program_path, "wb") as f:
|
||||
f.write(self._program.desc.serialize_to_string())
|
||||
|
||||
@staticmethod
|
||||
def load(save_dir):
|
||||
assert os.path.exists(save_dir)
|
||||
filename_list = sorted(os.listdir(save_dir))
|
||||
vars_list = []
|
||||
program_list = []
|
||||
dist_attr_list = []
|
||||
for filename in filename_list:
|
||||
filepath = os.path.join(save_dir, filename)
|
||||
assert os.path.isfile(filepath)
|
||||
if "vars" in filename:
|
||||
assert filename.endswith("pkl")
|
||||
with open(filepath, "rb") as f:
|
||||
from paddle.framework.restricted_unpickler import (
|
||||
safe_load_pickle,
|
||||
)
|
||||
|
||||
vars_list.append(safe_load_pickle(f))
|
||||
elif "program" in filename:
|
||||
assert filename.endswith("pdmodel")
|
||||
with open(filepath, "rb") as f:
|
||||
program_string = f.read()
|
||||
program_list.append(deserialize_program(program_string))
|
||||
elif "dist_attr" in filename:
|
||||
assert filename.endswith("pkl")
|
||||
with open(filepath, "rb") as f:
|
||||
from paddle.framework.restricted_unpickler import (
|
||||
safe_load_pickle,
|
||||
)
|
||||
|
||||
dist_attr_list.append(safe_load_pickle(f))
|
||||
|
||||
dist_attr_map = {}
|
||||
for dist_attrs in dist_attr_list:
|
||||
for dist_attr_name in dist_attrs.keys():
|
||||
if dist_attr_name not in dist_attr_map:
|
||||
dist_attr_map[dist_attr_name] = dist_attrs[dist_attr_name]
|
||||
assert len(vars_list) == len(program_list)
|
||||
return vars_list, program_list, dist_attr_map
|
||||
|
||||
@staticmethod
|
||||
def convert_src_tensor_2_dst_tensor(vars_list, src_attr_map, dst_attr_map):
|
||||
"""
|
||||
Converter is a class object for auto parallel to convert tensors from
|
||||
one parallel strategy to another one. Tensors will merge and slice value
|
||||
with their strategy when strategies are different.
|
||||
But like dp to pp or dp to serial is not supported.
|
||||
"""
|
||||
assert len(vars_list) >= 1
|
||||
# if dist_attr_map is None or len(dist_attr_map) == 0 or len(vars_list) == 1:
|
||||
if src_attr_map is None or len(src_attr_map) == 0:
|
||||
return vars_list[0]
|
||||
|
||||
dst_strategies = {}
|
||||
src_strategies = {}
|
||||
tensors_dict = {}
|
||||
|
||||
convert_tensor_dict = None
|
||||
for var_name in src_attr_map.keys():
|
||||
assert var_name not in dst_strategies
|
||||
dist_vars = []
|
||||
for vars in vars_list:
|
||||
if var_name in vars.keys():
|
||||
dist_vars.append(vars[var_name])
|
||||
if len(dist_vars) == 0:
|
||||
continue
|
||||
|
||||
if var_name in dst_attr_map and var_name in src_attr_map:
|
||||
dst_strategies[var_name] = copy.deepcopy(dst_attr_map[var_name])
|
||||
src_strategies[var_name] = copy.deepcopy(src_attr_map[var_name])
|
||||
tensors_dict[var_name] = dist_vars
|
||||
|
||||
if src_attr_map == dst_attr_map:
|
||||
return tensors_dict
|
||||
converter = Converter(tensors_dict, src_strategies, dst_strategies)
|
||||
convert_tensor_dict = converter.convert()
|
||||
|
||||
return convert_tensor_dict
|
||||
|
||||
@staticmethod
|
||||
def find_diff_vars(fixed_vars_map, query_vars_map):
|
||||
"""
|
||||
Found two variable names with different variable lists
|
||||
"""
|
||||
diff_var_name_list = set()
|
||||
for var_name in fixed_vars_map.keys():
|
||||
if var_name in query_vars_map:
|
||||
fixed_vars = fixed_vars_map[var_name]
|
||||
query_vars = query_vars_map[var_name]
|
||||
if isinstance(fixed_vars, np.ndarray):
|
||||
fixed_vars = [fixed_vars]
|
||||
if isinstance(query_vars, np.ndarray):
|
||||
query_vars = [query_vars]
|
||||
|
||||
length = min(len(fixed_vars), len(query_vars))
|
||||
if len(fixed_vars) != len(query_vars):
|
||||
print()
|
||||
for i in range(length):
|
||||
if not np.allclose(fixed_vars[i], query_vars[i]):
|
||||
diff_var_name_list.add(var_name)
|
||||
return diff_var_name_list
|
||||
|
||||
@staticmethod
|
||||
def diff_information(right_dir, wrong_dir):
|
||||
"""
|
||||
Find the corresponding operator according to the variable name.
|
||||
"""
|
||||
(
|
||||
right_vars_list,
|
||||
right_program_list,
|
||||
right_dist_attr_map,
|
||||
) = AutoAlignTool.load(right_dir)
|
||||
(
|
||||
wrong_vars_list,
|
||||
wrong_program_list,
|
||||
wrong_dist_attr_map,
|
||||
) = AutoAlignTool.load(wrong_dir)
|
||||
right_tensors_dict = AutoAlignTool.convert_src_tensor_2_dst_tensor(
|
||||
right_vars_list, right_dist_attr_map, right_dist_attr_map
|
||||
)
|
||||
wrong_tensors_dict = AutoAlignTool.convert_src_tensor_2_dst_tensor(
|
||||
wrong_vars_list, wrong_dist_attr_map, right_dist_attr_map
|
||||
)
|
||||
|
||||
diff_var_name_list = AutoAlignTool.find_diff_vars(
|
||||
right_tensors_dict, wrong_tensors_dict
|
||||
)
|
||||
|
||||
diff_ops_varname_dict = collections.OrderedDict()
|
||||
|
||||
for program in wrong_program_list:
|
||||
for block in program.blocks:
|
||||
for op in block.ops:
|
||||
for varname in op.input_arg_names + op.output_arg_names:
|
||||
if varname in diff_var_name_list:
|
||||
if len(diff_ops_varname_dict) == 0:
|
||||
print(
|
||||
"first different op:\n",
|
||||
op,
|
||||
f"\ndifferent varname is:{varname}",
|
||||
)
|
||||
if op not in diff_ops_varname_dict:
|
||||
diff_ops_varname_dict[op] = [varname]
|
||||
else:
|
||||
diff_ops_varname_dict[op].append(varname)
|
||||
|
||||
return diff_ops_varname_dict
|
||||
|
||||
@staticmethod
|
||||
def diff_information_from_dirs(right_dirs, wrong_dirs):
|
||||
right_vars_list = []
|
||||
right_program_list = []
|
||||
right_dist_attr_map = {}
|
||||
for right_dir in right_dirs:
|
||||
(
|
||||
tmp_vars_list,
|
||||
right_program_list,
|
||||
tmp_dist_attr_map,
|
||||
) = AutoAlignTool.load(right_dir)
|
||||
if len(right_vars_list) == 0:
|
||||
right_vars_list = tmp_vars_list
|
||||
else:
|
||||
for i in range(len(tmp_vars_list)):
|
||||
vars_list = tmp_vars_list[i]
|
||||
for key in vars_list.keys():
|
||||
if key not in right_vars_list[i].keys():
|
||||
right_vars_list[i][key] = vars_list[key]
|
||||
|
||||
for key in tmp_dist_attr_map.keys():
|
||||
if key not in right_dist_attr_map:
|
||||
right_dist_attr_map[key] = tmp_dist_attr_map[key]
|
||||
|
||||
wrong_vars_list = []
|
||||
wrong_program_list = []
|
||||
wrong_dist_attr_map = {}
|
||||
for wrong_dir in wrong_dirs:
|
||||
(
|
||||
tmp_vars_list,
|
||||
wrong_program_list,
|
||||
tmp_dist_attr_map,
|
||||
) = AutoAlignTool.load(wrong_dir)
|
||||
if len(wrong_vars_list) == 0:
|
||||
wrong_vars_list = tmp_vars_list
|
||||
else:
|
||||
for i in range(len(tmp_vars_list)):
|
||||
vars_list = tmp_vars_list[i]
|
||||
for key in vars_list.keys():
|
||||
if key not in wrong_vars_list[i].keys():
|
||||
wrong_vars_list[i][key] = vars_list[key]
|
||||
|
||||
for key in tmp_dist_attr_map.keys():
|
||||
if key not in wrong_dist_attr_map:
|
||||
wrong_dist_attr_map[key] = tmp_dist_attr_map[key]
|
||||
|
||||
right_tensors_dict = AutoAlignTool.convert_src_tensor_2_dst_tensor(
|
||||
right_vars_list, right_dist_attr_map, right_dist_attr_map
|
||||
)
|
||||
wrong_tensors_dict = AutoAlignTool.convert_src_tensor_2_dst_tensor(
|
||||
wrong_vars_list, wrong_dist_attr_map, right_dist_attr_map
|
||||
)
|
||||
diff_var_name_list = AutoAlignTool.find_diff_vars(
|
||||
right_tensors_dict, wrong_tensors_dict
|
||||
)
|
||||
|
||||
diff_ops_varname_dict = collections.OrderedDict()
|
||||
|
||||
for program in wrong_program_list:
|
||||
for block in program.blocks:
|
||||
for op in block.ops:
|
||||
for varname in op.input_arg_names + op.output_arg_names:
|
||||
if varname in diff_var_name_list:
|
||||
if len(diff_ops_varname_dict) == 0:
|
||||
print(
|
||||
"first different op:\n",
|
||||
op,
|
||||
f"\ndifferent varname is:{varname}",
|
||||
)
|
||||
if op not in diff_ops_varname_dict:
|
||||
diff_ops_varname_dict[op] = [varname]
|
||||
else:
|
||||
diff_ops_varname_dict[op].append(varname)
|
||||
|
||||
return diff_ops_varname_dict
|
||||
@@ -0,0 +1,244 @@
|
||||
# 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 os
|
||||
import time
|
||||
|
||||
import paddle
|
||||
from paddle.hapi.callbacks import (
|
||||
Callback,
|
||||
CallbackList,
|
||||
LRScheduler,
|
||||
ModelCheckpoint,
|
||||
ProgBarLogger,
|
||||
)
|
||||
|
||||
from ..interface import CollectionNames, get_collection
|
||||
|
||||
|
||||
def config_callbacks(
|
||||
callbacks=None,
|
||||
engine=None,
|
||||
batch_size=None,
|
||||
epochs=None,
|
||||
steps=None,
|
||||
log_freq=2,
|
||||
verbose=2,
|
||||
save_freq=1,
|
||||
save_dir=None,
|
||||
metrics=None,
|
||||
acc_step=1,
|
||||
mode='train',
|
||||
):
|
||||
cbks = callbacks or []
|
||||
cbks = cbks if isinstance(cbks, (list, tuple)) else [cbks]
|
||||
|
||||
if not any(isinstance(k, ProgBarLogger) for k in cbks) and verbose:
|
||||
cbks = [ProgBarLoggerAuto(log_freq, verbose=verbose), *cbks]
|
||||
|
||||
if not any(isinstance(k, LRScheduler) for k in cbks):
|
||||
cbks = [LRSchedulerAuto(), *cbks]
|
||||
|
||||
if not any(isinstance(k, ModelCheckpoint) for k in cbks):
|
||||
cbks = [*cbks, ModelCheckpointAuto(save_freq, save_dir)]
|
||||
|
||||
if not any(isinstance(k, Profiler) for k in cbks) and verbose == 3:
|
||||
cbks = [*cbks, Profiler(timer_only=True)]
|
||||
|
||||
if not any(isinstance(k, History) for k in cbks):
|
||||
cbks = [*cbks, History()]
|
||||
|
||||
for i, k in enumerate(cbks):
|
||||
if isinstance(k, ProgBarLogger):
|
||||
cbks[i] = ProgBarLoggerAuto(k.log_freq, k.verbose)
|
||||
if isinstance(k, LRScheduler):
|
||||
cbks[i] = LRSchedulerAuto(k.by_step, k.by_epoch)
|
||||
if isinstance(k, ModelCheckpoint):
|
||||
cbks[i] = ModelCheckpointAuto(k.save_freq, k.save_dir)
|
||||
|
||||
cbk_list = CallbackList(cbks)
|
||||
cbk_list.set_model(engine)
|
||||
metrics = metrics or [] if mode != 'test' else []
|
||||
params = {
|
||||
'batch_size': batch_size,
|
||||
'epochs': epochs,
|
||||
'steps': steps,
|
||||
'verbose': verbose,
|
||||
'metrics': metrics,
|
||||
'acc_step': acc_step,
|
||||
}
|
||||
cbk_list.set_params(params)
|
||||
return cbk_list
|
||||
|
||||
|
||||
class ProgBarLoggerAuto(ProgBarLogger):
|
||||
def __init__(self, log_freq=1, verbose=2):
|
||||
super().__init__(log_freq, verbose)
|
||||
|
||||
def _is_print(self):
|
||||
return True
|
||||
|
||||
def _updates(self, logs, mode):
|
||||
values = []
|
||||
metrics = getattr(self, f'{mode}_metrics')
|
||||
progbar = getattr(self, f'{mode}_progbar')
|
||||
steps = getattr(self, f'{mode}_step')
|
||||
|
||||
for k in metrics:
|
||||
if k in logs:
|
||||
values.append((k, logs[k]))
|
||||
|
||||
if 'lr' in logs:
|
||||
values.append(('lr', logs['lr']))
|
||||
|
||||
fetches_logs = logs.get('fetches', {})
|
||||
collect_logging = get_collection(CollectionNames.LOGGING)
|
||||
for name, var in collect_logging:
|
||||
k = name or var.name
|
||||
if k in fetches_logs:
|
||||
values.append((k, fetches_logs[k]))
|
||||
|
||||
out_logs = logs.get('outputs', {})
|
||||
for k in out_logs:
|
||||
values.append((k, out_logs[k]))
|
||||
|
||||
if self.verbose == 3 and hasattr(self, f'_{mode}_timer'):
|
||||
timer = getattr(self, f'_{mode}_timer')
|
||||
cnt = timer['count'] if timer['count'] > 0 else 1.0
|
||||
samples = timer['samples'] if timer['samples'] > 0 else 1.0
|
||||
values.append(
|
||||
('avg_reader_cost', "%.5f sec" % (timer['data_time'] / cnt))
|
||||
)
|
||||
values.append(
|
||||
('avg_batch_cost', "%.5f sec" % (timer['batch_time'] / cnt))
|
||||
)
|
||||
values.append(
|
||||
(
|
||||
'ips',
|
||||
"%.5f samples/sec"
|
||||
% (samples / (timer['data_time'] + timer['batch_time'])),
|
||||
)
|
||||
)
|
||||
timer['count'] = 0
|
||||
timer['samples'] = 0
|
||||
timer['data_time'] = 0.0
|
||||
timer['batch_time'] = 0.0
|
||||
|
||||
progbar.update(steps, values)
|
||||
|
||||
def on_eval_batch_end(self, step, logs=None):
|
||||
logs = logs or {}
|
||||
self.eval_step += 1
|
||||
samples = self.params['batch_size']
|
||||
self.evaled_samples += samples
|
||||
|
||||
self._eval_timer['batch_time'] += (
|
||||
time.time() - self._eval_timer['batch_data_end_time']
|
||||
)
|
||||
self._eval_timer['count'] += 1
|
||||
samples = self.params['batch_size']
|
||||
self._eval_timer['samples'] += samples
|
||||
|
||||
if self._is_print() and self.eval_step % self.log_freq == 0:
|
||||
if self.eval_steps is None or self.eval_step < self.eval_steps:
|
||||
self._updates(logs, 'eval')
|
||||
|
||||
self._eval_timer['batch_start_time'] = time.time()
|
||||
|
||||
|
||||
class LRSchedulerAuto(LRScheduler):
|
||||
def __init__(self, by_step=True, by_epoch=False):
|
||||
super().__init__(by_step, by_epoch)
|
||||
|
||||
def on_epoch_begin(self, epoch=None, logs=None):
|
||||
self.acc_step = self.params["acc_step"]
|
||||
self.epoch = epoch
|
||||
self.train_step = 0
|
||||
|
||||
def on_train_batch_end(self, step, logs=None):
|
||||
self.train_step += 1
|
||||
|
||||
if self.by_step and self.train_step % self.acc_step == 0:
|
||||
if (
|
||||
self.model.optimizer
|
||||
and hasattr(self.model.optimizer, '_learning_rate')
|
||||
and isinstance(
|
||||
self.model.optimizer._learning_rate,
|
||||
paddle.optimizer.lr.LRScheduler,
|
||||
)
|
||||
):
|
||||
self.model.optimizer._learning_rate.step()
|
||||
|
||||
|
||||
class History(Callback):
|
||||
def __init__(self):
|
||||
self.history = {}
|
||||
|
||||
def on_train_begin(self, logs=None):
|
||||
self.epoch = []
|
||||
|
||||
def on_epoch_end(self, epoch, logs=None):
|
||||
logs = logs or {}
|
||||
self.epoch.append(epoch)
|
||||
for k, v in logs.items():
|
||||
self.history.setdefault(k, []).append(v)
|
||||
|
||||
self.model.history = self
|
||||
|
||||
|
||||
class Profiler(Callback):
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.prof = paddle.profiler.Profiler(*args, **kwargs)
|
||||
|
||||
def on_epoch_begin(self, epoch=None, logs=None):
|
||||
self.epoch = epoch
|
||||
self.train_step = 0
|
||||
self.batch_size = self.params["batch_size"]
|
||||
self.steps = self.params['steps']
|
||||
|
||||
def on_train_begin(self, logs=None):
|
||||
self.prof.start()
|
||||
|
||||
def on_train_batch_end(self, step, logs=None):
|
||||
self.train_step += 1
|
||||
self.prof.step(num_samples=self.batch_size)
|
||||
print(
|
||||
"step {}:{}".format(
|
||||
self.train_step, self.prof.step_info(unit='samples')
|
||||
)
|
||||
)
|
||||
|
||||
def on_train_end(self, logs=None):
|
||||
self.prof.stop()
|
||||
self.prof.summary()
|
||||
|
||||
|
||||
class ModelCheckpointAuto(ModelCheckpoint):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def _is_save(self):
|
||||
return self.model and self.save_dir
|
||||
|
||||
def on_epoch_end(self, epoch, logs=None):
|
||||
if self._is_save() and (self.epoch + 1) % self.save_freq == 0:
|
||||
path = f'{self.save_dir}/epoch{epoch}'
|
||||
print(f'save checkpoint at {os.path.abspath(path)}')
|
||||
self.model.save(path)
|
||||
|
||||
def on_train_end(self, logs=None):
|
||||
if self._is_save():
|
||||
path = f'{self.save_dir}/final'
|
||||
print(f'save checkpoint at {os.path.abspath(path)}')
|
||||
self.model.save(path)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,129 @@
|
||||
# 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.
|
||||
|
||||
from enum import IntEnum, unique
|
||||
|
||||
import numpy as np
|
||||
|
||||
from paddle.framework import core
|
||||
|
||||
|
||||
@unique
|
||||
class DeviceType(IntEnum):
|
||||
UNKNOWN = 0
|
||||
CPU = 1
|
||||
GPU = 2
|
||||
XPU = 3
|
||||
DCU = 5
|
||||
NIC = 6
|
||||
|
||||
|
||||
@unique
|
||||
class LinkType(IntEnum):
|
||||
UNKNOWN = 0
|
||||
LOC = 1
|
||||
SYS = 2
|
||||
PHB = 3
|
||||
PIX = 4
|
||||
PIB = 5
|
||||
NVL = 6
|
||||
NVB = 7
|
||||
NET = 8
|
||||
|
||||
|
||||
class DeviceMesh(core.DeviceMesh):
|
||||
r"""
|
||||
The class `DeviceMesh` describes the topology of physical devices.
|
||||
|
||||
Args:
|
||||
mesh (list|numpy.array): an N-dimensional array describes the topology
|
||||
of logical processes.
|
||||
dim_names (list, optional): the i-th element of this list gives the name of the
|
||||
i-th dimension.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
|
||||
>>> import paddle
|
||||
>>> import paddle.distributed as dist
|
||||
|
||||
>>> paddle.enable_static()
|
||||
|
||||
>>> mesh = dist.DeviceMesh([[2, 4, 5], [0, 1, 3]])
|
||||
>>> assert mesh.shape == [2, 3]
|
||||
>>> assert mesh.device_ids == [2, 4, 5, 0, 1, 3]
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, name, mesh, dim_names=None):
|
||||
self._name = name
|
||||
|
||||
if not isinstance(mesh, list) and not isinstance(mesh, np.ndarray):
|
||||
raise ValueError(
|
||||
'The mesh must be an instance of list or np.ndarray.'
|
||||
)
|
||||
if isinstance(mesh, list):
|
||||
mesh = np.array(mesh)
|
||||
|
||||
self._mesh = mesh
|
||||
|
||||
self._shape = list(self._mesh.shape)
|
||||
|
||||
self._device_ids = self._mesh.flatten().tolist()
|
||||
assert all(isinstance(p, int) for p in self._device_ids), (
|
||||
"All elements of the mesh be integer"
|
||||
)
|
||||
assert min(self._device_ids) >= 0, (
|
||||
'All elements of the mesh must be >= 0.'
|
||||
)
|
||||
unique_device_ids = set(self._device_ids)
|
||||
assert len(unique_device_ids) == len(self._device_ids), (
|
||||
'All elements of the mesh must be unique.'
|
||||
)
|
||||
|
||||
if dim_names is not None:
|
||||
assert len(dim_names) == len(self._shape), (
|
||||
"The length of dims_names must be same as the shape of the mesh."
|
||||
)
|
||||
self._dim_names = dim_names
|
||||
else:
|
||||
self._dim_names = ["d" + str(i) for i in range(len(self._shape))]
|
||||
|
||||
# Follow the requirement for using pybind11
|
||||
core.DeviceMesh.__init__(
|
||||
self, self._name, self._shape, self._device_ids, self._dim_names
|
||||
)
|
||||
|
||||
@property
|
||||
def mesh(self):
|
||||
return self._mesh
|
||||
|
||||
|
||||
# class Cluster:
|
||||
# """
|
||||
# The cluster represents the hardware resource.
|
||||
# """
|
||||
|
||||
# def __init__(self):
|
||||
# self._device_meshes = {}
|
||||
|
||||
# def device_mesh(self, device_mesh_name):
|
||||
# return self._device_meshes[device_mesh_name]
|
||||
|
||||
# def add_device_mesh(self, device_mesh):
|
||||
# self._device_meshes[device_mesh.name] = device_mesh
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,543 @@
|
||||
# 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 logging
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
|
||||
from ...utils.log_utils import get_logger
|
||||
|
||||
|
||||
class Converter:
|
||||
"""
|
||||
Converter is a class object for auto parallel to convert tensors from
|
||||
one parallel strategy to another one. Tensors will merge and slice value
|
||||
with their strategy when strategies are different.
|
||||
"""
|
||||
|
||||
def __init__(self, tensors_dict, pre_strategy, cur_strategy):
|
||||
"""
|
||||
Args:
|
||||
tensors_dict(dict): tensors' value of all ranks that to be converted.
|
||||
key is tensor's name(str), value is all ranks' data(list(numpy.ndarray))
|
||||
pre_strategy(dict): tensors' distributed attribute of last training process.
|
||||
key is tensor's name(str), value is tensor's distributed attribute in last
|
||||
training process.
|
||||
cur_strategy(dict): tensors' distributed attribute of current rank.
|
||||
key is tensor's name(str), value is tensor's distributed attribute in current
|
||||
rank.
|
||||
"""
|
||||
self._tensors_dict = self._check_tensor_dict(tensors_dict)
|
||||
self._pre_strategy = self._check_pre_strategy(pre_strategy)
|
||||
self._cur_strategy = self._check_cur_strategy(cur_strategy)
|
||||
self._logger = get_logger(logging.INFO)
|
||||
|
||||
def _check_tensor_dict(self, tensors_dict):
|
||||
if not tensors_dict:
|
||||
raise ValueError(
|
||||
"'tensors_dict' is None, "
|
||||
"the tensors to be converted cannot be None."
|
||||
)
|
||||
if not isinstance(tensors_dict, dict):
|
||||
raise TypeError(
|
||||
f"The type of 'tensors_dict' should be 'dict', but got '{type(tensors_dict)}'."
|
||||
)
|
||||
return tensors_dict
|
||||
|
||||
def _check_pre_strategy(self, pre_strategy):
|
||||
if not pre_strategy:
|
||||
raise ValueError(
|
||||
"'pre_strategy' is None, there are not tensors in pre process."
|
||||
)
|
||||
if not isinstance(pre_strategy, dict):
|
||||
raise TypeError(
|
||||
"The type of 'pre_strategy' should be 'dict', "
|
||||
f"but got '{type(pre_strategy)}'."
|
||||
)
|
||||
return pre_strategy
|
||||
|
||||
def _check_cur_strategy(self, cur_strategy):
|
||||
if not cur_strategy:
|
||||
warnings.warn(
|
||||
"'cur_strategy' is None, there are not tensors in cur process"
|
||||
)
|
||||
if not isinstance(cur_strategy, dict):
|
||||
raise TypeError(
|
||||
"The type of 'cur_strategy' should be 'dict', "
|
||||
f"but got '{type(cur_strategy)}'."
|
||||
)
|
||||
return cur_strategy
|
||||
|
||||
def convert(self, strict=True):
|
||||
"""
|
||||
Convert tensors
|
||||
|
||||
Args:
|
||||
strict(bool): whether to strict convert tensor with tensor's name. If False, it will
|
||||
convert tensors by prefix matching. Otherwise, tensors will be converted with
|
||||
their name strictly.
|
||||
|
||||
Returns:
|
||||
converted tensors(dict)
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
|
||||
>>> import numpy as np
|
||||
>>> from paddle.distributed.auto_parallel.static.converter import Converter
|
||||
>>> complete_tensors = np.arange(4).reshape([2, 2])
|
||||
>>> partial_tensors = np.split(complete_tensors, 2, axis=0)
|
||||
>>> name = "tmp_0"
|
||||
>>> tensors_dict = {name: partial_tensors}
|
||||
>>> strategy_1 = {
|
||||
... name: {
|
||||
... "process_shape": [2],
|
||||
... "process_group": [0, 1],
|
||||
... "dims_mapping": [0, -1],
|
||||
... },
|
||||
... }
|
||||
>>> strategy_2 = {
|
||||
... name: {
|
||||
... "process_shape": [2],
|
||||
... "process_group": [0, 1],
|
||||
... "dims_mapping": [-1, -1],
|
||||
... },
|
||||
... }
|
||||
>>> converter = Converter(tensors_dict, strategy_1, strategy_2)
|
||||
>>> result = converter.convert()
|
||||
>>> # the result's value is equal to `complete_tensors`
|
||||
"""
|
||||
tensors_dict = {}
|
||||
# the name which is in cur_process but not in pre_process
|
||||
tensor_not_in_pre = []
|
||||
# the name which is in pre_process but not in cur_process
|
||||
tensor_not_in_cur = []
|
||||
# the name which is in strategy but not in ckpt files
|
||||
tensor_not_in_ckpt = []
|
||||
self._logger.info("Start to convert tensors.")
|
||||
for tensor_name in self._cur_strategy:
|
||||
if tensor_name not in self._pre_strategy:
|
||||
tensor_not_in_pre.append(tensor_name)
|
||||
continue
|
||||
if tensor_name not in self._tensors_dict:
|
||||
tensor_not_in_ckpt.append(tensor_name)
|
||||
continue
|
||||
self._pre_name = tensor_name
|
||||
self._cur_name = tensor_name
|
||||
tensor_list = self._tensors_dict[tensor_name]
|
||||
pre_dist_attr = self._pre_strategy[tensor_name]
|
||||
cur_dist_attr = self._cur_strategy[tensor_name]
|
||||
try:
|
||||
tensors_dict[tensor_name] = Converter.merge_and_slice(
|
||||
tensor_list, pre_dist_attr, cur_dist_attr
|
||||
)
|
||||
except ValueError as err:
|
||||
raise ValueError(
|
||||
f"Fail to convert tensor '{tensor_name}'. {err}"
|
||||
)
|
||||
|
||||
for tensor_name in self._pre_strategy:
|
||||
if tensor_name not in self._cur_strategy:
|
||||
tensor_not_in_cur.append(tensor_name)
|
||||
|
||||
if not strict:
|
||||
(
|
||||
tensors_dict,
|
||||
tensor_match_with_pre,
|
||||
tensor_match_with_cur,
|
||||
) = self.convert_with_prefix_match(
|
||||
tensors_dict, tensor_not_in_pre, tensor_not_in_cur
|
||||
)
|
||||
else:
|
||||
tensors_dict, tensor_match_with_pre, tensor_match_with_cur = (
|
||||
tensors_dict,
|
||||
[],
|
||||
[],
|
||||
)
|
||||
|
||||
tensor_not_in_pre = set(tensor_not_in_pre) - set(tensor_match_with_pre)
|
||||
tensor_not_in_cur = set(tensor_not_in_cur) - set(tensor_match_with_cur)
|
||||
if tensor_not_in_pre:
|
||||
warnings.warn(
|
||||
f"tensors [{tensor_not_in_pre}] are not found in last training strategy."
|
||||
)
|
||||
if tensor_not_in_cur:
|
||||
warnings.warn(
|
||||
f"tensors [{tensor_not_in_cur}] are not found in current training strategy."
|
||||
)
|
||||
if tensor_not_in_ckpt:
|
||||
warnings.warn(
|
||||
f"tensors [{tensor_not_in_ckpt}] are found in pre_strategy, but are not found"
|
||||
"in checkpoint files, please check your checkpoint files."
|
||||
)
|
||||
|
||||
return tensors_dict
|
||||
|
||||
def convert_with_prefix_match(
|
||||
self, tensors_dict, tensor_not_in_pre, tensor_not_in_cur
|
||||
):
|
||||
# the name which in cur_process and can match with pre_process
|
||||
tensor_match_with_pre = []
|
||||
# the name which in pre_process and can match with cur_process
|
||||
tensor_match_with_cur = []
|
||||
for cur_name in tensor_not_in_pre:
|
||||
prefix_name = cur_name
|
||||
while prefix_name.find("_") != -1:
|
||||
prefix_name = prefix_name[: prefix_name.rfind("_")]
|
||||
for pre_name in tensor_not_in_cur:
|
||||
if prefix_name in pre_name:
|
||||
# 'cur_name' of cur_process can match with 'pre_name' of pre_process
|
||||
self._pre_name = pre_name
|
||||
self._cur_name = cur_name
|
||||
pre_tensor_list = self._tensors_dict[pre_name]
|
||||
pre_dist_attr = self._pre_strategy[pre_name]
|
||||
cur_dist_attr = self._cur_strategy[cur_name]
|
||||
try:
|
||||
tensors_dict[cur_name] = Converter.merge_and_slice(
|
||||
pre_tensor_list, pre_dist_attr, cur_dist_attr
|
||||
)
|
||||
except ValueError as err:
|
||||
raise ValueError(
|
||||
f"Fail to convert tensor '{cur_name}' by '{pre_name}'. {err}"
|
||||
)
|
||||
self._logger.info(
|
||||
f"tensor [{cur_name}] is matched with tensor [{pre_name}]"
|
||||
)
|
||||
tensor_match_with_pre.append(cur_name)
|
||||
tensor_match_with_cur.append(pre_name)
|
||||
break
|
||||
break
|
||||
|
||||
return tensors_dict, tensor_match_with_pre, tensor_match_with_cur
|
||||
|
||||
@staticmethod
|
||||
def merge_and_slice(tensor_list, pre_dist_attr, cur_dist_attr):
|
||||
"""
|
||||
Merge tensors with previous dist_attr and slice tensors with current dist_attr
|
||||
|
||||
Returns:
|
||||
tensor(numpy.narray): a tensor's value of current rank.
|
||||
"""
|
||||
assert isinstance(tensor_list, list)
|
||||
assert all(isinstance(p, np.ndarray) for p in tensor_list)
|
||||
|
||||
if pre_dist_attr == cur_dist_attr:
|
||||
# skip merge and slice tensor
|
||||
rank_id = paddle.distributed.get_rank()
|
||||
index = cur_dist_attr["process_group"].index(rank_id)
|
||||
tensor = tensor_list[index]
|
||||
else:
|
||||
pre_dims_mapping = pre_dist_attr["dims_mapping"]
|
||||
cur_dims_mapping = cur_dist_attr["dims_mapping"]
|
||||
|
||||
if len(pre_dims_mapping) and (
|
||||
len(set(pre_dims_mapping)) > 1 or -1 not in pre_dims_mapping
|
||||
):
|
||||
# merge tensor
|
||||
tensor = Converter.merge_with_dist_attr(
|
||||
tensor_list, pre_dist_attr
|
||||
)
|
||||
else:
|
||||
# skip merge tensor
|
||||
tensor = tensor_list[0]
|
||||
|
||||
if len(cur_dims_mapping) and (
|
||||
len(set(cur_dims_mapping)) > 1 or -1 not in cur_dims_mapping
|
||||
):
|
||||
# slice tensor
|
||||
tensor = Converter.slice_with_dist_attr(tensor, cur_dist_attr)
|
||||
|
||||
return tensor
|
||||
|
||||
@staticmethod
|
||||
def merge_with_dist_attr(tensor_list, dist_attr):
|
||||
"""Merge tensor with distributed attribute"""
|
||||
from .reshard import Resharder
|
||||
|
||||
dims_mapping = dist_attr["dims_mapping"]
|
||||
process_shape = dist_attr["process_shape"]
|
||||
process_group = dist_attr["process_group"]
|
||||
# get the complete shape of the tensor
|
||||
complete_shape = Resharder.compute_complete_shape(
|
||||
tensor_list[0].shape, process_shape, dims_mapping
|
||||
)
|
||||
# merge the tensor with dist_attr
|
||||
partition_tensor_list = []
|
||||
merged_partition = []
|
||||
for process in process_group:
|
||||
partition_index = Resharder.compute_partition_index(
|
||||
process,
|
||||
complete_shape,
|
||||
dims_mapping,
|
||||
process_shape,
|
||||
process_group,
|
||||
)
|
||||
index = process_group.index(process)
|
||||
if partition_index not in merged_partition:
|
||||
merged_partition.append(partition_index)
|
||||
Converter.merge(
|
||||
partition_tensor_list,
|
||||
tensor_list[index],
|
||||
partition_index,
|
||||
complete_shape,
|
||||
)
|
||||
|
||||
if len(partition_tensor_list) != 1:
|
||||
raise ValueError(
|
||||
f"Fail to merge tensor with dist_attr '{dist_attr}'."
|
||||
)
|
||||
complete_tensor = partition_tensor_list[0][0]
|
||||
return complete_tensor
|
||||
|
||||
@staticmethod
|
||||
def slice_with_dist_attr(tensor, dist_attr):
|
||||
"""Slice tensor with distributed attribute"""
|
||||
dims_mapping = dist_attr["dims_mapping"]
|
||||
if len(dims_mapping) == 0:
|
||||
# NOTE: scalar tensor no need to split
|
||||
return tensor
|
||||
process_shape = dist_attr["process_shape"]
|
||||
process_group = dist_attr["process_group"]
|
||||
# slice the tensor with dist_attr
|
||||
partition_index_list = Converter._get_split_indices(
|
||||
tensor.shape, dims_mapping, process_shape, process_group
|
||||
)
|
||||
sliced_tensor_list = Converter.split(
|
||||
tensor, partition_index_list, len(partition_index_list)
|
||||
)
|
||||
# get the current tensor's index in sliced_tensor_list
|
||||
rank_id = paddle.distributed.get_rank()
|
||||
sliced_tensor_index = Converter._get_sliced_index(
|
||||
rank_id, tensor.shape, dims_mapping, process_shape, process_group
|
||||
)
|
||||
if sliced_tensor_index not in range(len(sliced_tensor_list)):
|
||||
raise ValueError(
|
||||
f"Fail to slice tensor with dist_attr '{dist_attr}'."
|
||||
)
|
||||
sliced_tensor = sliced_tensor_list[sliced_tensor_index]
|
||||
return sliced_tensor
|
||||
|
||||
@staticmethod
|
||||
def merge(partition_tensor_list, tensor, partition_index, complete_shape):
|
||||
"""
|
||||
Merge partial tensors to a complete.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
|
||||
>>> import numpy as np
|
||||
>>> import paddle
|
||||
>>> from paddle.distributed.auto_parallel.static.converter import Converter
|
||||
>>> partition_tensor_list = [(np.array([[[1.11, 1.12]]]), [[0, 1], [0, 1], [0, 2]])]
|
||||
>>> tensor = np.array([[[1.13, 1.14]]])
|
||||
>>> partition_index = [[0, 1], [0, 1], [2, 4]]
|
||||
>>> complete_shape = [3, 2]
|
||||
|
||||
>>> Converter.merge(partition_tensor_list, tensor, partition_index, complete_shape)
|
||||
>>> print(partition_tensor_list)
|
||||
[(array([[[1.11, 1.12, 1.13, 1.14]]]), [[0, 1], [0, 1], [0, 4]])]
|
||||
"""
|
||||
from .reshard import Resharder
|
||||
|
||||
if len(partition_tensor_list) == 1:
|
||||
is_complete_data = True
|
||||
for idx, item in enumerate(partition_tensor_list[0][1]):
|
||||
if item[0] != 0 or item[1] != complete_shape[idx]:
|
||||
is_complete_data = False
|
||||
break
|
||||
if is_complete_data:
|
||||
return
|
||||
|
||||
if not partition_tensor_list:
|
||||
partition_tensor_list.append((tensor, partition_index))
|
||||
else:
|
||||
i = 0
|
||||
while i < len(partition_tensor_list):
|
||||
(
|
||||
concat_axis,
|
||||
first_order,
|
||||
new_partition,
|
||||
) = Resharder.compute_concat_info(
|
||||
partition_tensor_list[i][1], partition_index
|
||||
)
|
||||
if concat_axis != -1:
|
||||
if first_order == 0:
|
||||
new_tensor = np.concatenate(
|
||||
(partition_tensor_list[i][0], tensor),
|
||||
axis=concat_axis,
|
||||
)
|
||||
else:
|
||||
new_tensor = np.concatenate(
|
||||
(tensor, partition_tensor_list[i][0]),
|
||||
axis=concat_axis,
|
||||
)
|
||||
|
||||
partition_tensor_list.pop(i)
|
||||
Converter.merge(
|
||||
partition_tensor_list,
|
||||
new_tensor,
|
||||
new_partition,
|
||||
complete_shape,
|
||||
)
|
||||
break
|
||||
i += 1
|
||||
|
||||
@staticmethod
|
||||
def split(complete_tensor, partition_index_list, length):
|
||||
"""
|
||||
Slice a complete tensor.
|
||||
|
||||
Returns:
|
||||
sliced_tensor_list(list): sliced tensors with 'partition_index_list'
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
|
||||
>>> import numpy as np
|
||||
>>> from paddle.distributed.auto_parallel.static.converter import Converter
|
||||
>>> complete_tensor = np.array([[[1.11, 1.12, 1.13, 1.14, 1.15, 1.16]]])
|
||||
>>> rank = 2
|
||||
>>> complete_shape = [1, 1, 6]
|
||||
>>> dims_mapping = [-1, -1, 0]
|
||||
>>> process_shape = [3]
|
||||
>>> process_group = [0, 1, 2]
|
||||
|
||||
>>> sliced_tensor_list = Converter.split(complete_tensor, [[], [], [2, 4]], 3)
|
||||
>>> print(sliced_tensor_list)
|
||||
[array([[[1.11, 1.12]]]), array([[[1.13, 1.14]]]), array([[[1.15, 1.16]]])]
|
||||
"""
|
||||
sliced_tensor_list = []
|
||||
axis = len(complete_tensor.shape) - length
|
||||
sliced_tensor = np.split(
|
||||
complete_tensor, partition_index_list[axis], axis=axis
|
||||
)
|
||||
if length == 1:
|
||||
return sliced_tensor
|
||||
for tensor in sliced_tensor:
|
||||
sliced_tensor_list.extend(
|
||||
Converter.split(tensor, partition_index_list, length - 1)
|
||||
)
|
||||
return sliced_tensor_list
|
||||
|
||||
@staticmethod
|
||||
def _get_split_indices(
|
||||
complete_shape, dims_mapping, process_shape, process_group
|
||||
):
|
||||
"""
|
||||
Get split indices of every dimension.
|
||||
|
||||
Returns:
|
||||
split_indices_list(list): the split indices of every dimension of the tensor
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
|
||||
>>> import numpy as np
|
||||
>>> from paddle.distributed.auto_parallel.static.utils import _get_split_indices
|
||||
>>> complete_tensor = np.array([[[1.11, 1.12, 1.13, 1.14, 1.15, 1.16]]])
|
||||
>>> complete_shape = [1, 1, 6]
|
||||
>>> dims_mapping = [-1, -1, 0]
|
||||
>>> process_shape = [3]
|
||||
>>> process_group = [0, 1, 2]
|
||||
|
||||
>>> index = _get_split_indices(complete_shape, dims_mapping, process_shape, process_group)
|
||||
>>> print(index)
|
||||
[[], [], [2, 4]]
|
||||
"""
|
||||
from .reshard import Resharder
|
||||
|
||||
split_indices_list = []
|
||||
for process in process_group:
|
||||
partition_index = Resharder.compute_partition_index(
|
||||
process,
|
||||
complete_shape,
|
||||
dims_mapping,
|
||||
process_shape,
|
||||
process_group,
|
||||
)
|
||||
if split_indices_list:
|
||||
for dim in range(len(partition_index)):
|
||||
split_indices_list[dim].extend(partition_index[dim])
|
||||
else:
|
||||
split_indices_list = partition_index
|
||||
split_indices_list = list(
|
||||
map(
|
||||
lambda x, y: list(set(x) - {y} - {0}),
|
||||
split_indices_list,
|
||||
complete_shape,
|
||||
)
|
||||
)
|
||||
split_indices_list = [sorted(x) for x in split_indices_list]
|
||||
return split_indices_list
|
||||
|
||||
@staticmethod
|
||||
def _get_sliced_index(
|
||||
rank_id, complete_shape, dims_mapping, process_shape, process_group
|
||||
):
|
||||
"""
|
||||
Get sliced_tensor's index of current rank in all sliced tensors list.
|
||||
|
||||
Returns:
|
||||
sliced_tensor_index(int): the index of sliced tensor in sliced_tensor_list
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
|
||||
>>> import numpy as np
|
||||
>>> from paddle.distributed.auto_parallel.static.converter import Converter
|
||||
>>> complete_tensor = np.array([[[1.11, 1.12, 1.13, 1.14, 1.15, 1.16]]])
|
||||
>>> rank = 2
|
||||
>>> complete_shape = [1, 1, 6]
|
||||
>>> dims_mapping = [-1, -1, 0]
|
||||
>>> process_shape = [3]
|
||||
>>> process_group = [0, 1, 2]
|
||||
|
||||
>>> index = Converter._get_sliced_index(
|
||||
... rank,
|
||||
... complete_shape,
|
||||
... dims_mapping,
|
||||
... process_shape,
|
||||
... process_group,
|
||||
... )
|
||||
>>> print(index)
|
||||
2
|
||||
"""
|
||||
from .reshard import Resharder
|
||||
|
||||
partition_index = Resharder.compute_partition_index(
|
||||
rank_id, complete_shape, dims_mapping, process_shape, process_group
|
||||
)
|
||||
sliced_index = 0
|
||||
for i, shape in enumerate(complete_shape):
|
||||
if dims_mapping[i] == -1:
|
||||
slice_shape = shape
|
||||
else:
|
||||
slice_shape = shape // process_shape[dims_mapping[i]]
|
||||
if slice_shape == 1:
|
||||
index = partition_index[i][0]
|
||||
else:
|
||||
index = (partition_index[i][0] + 1) // slice_shape
|
||||
sliced_index = sliced_index * (shape // slice_shape) + index
|
||||
return sliced_index
|
||||
@@ -0,0 +1,62 @@
|
||||
# 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
|
||||
|
||||
from .base_cost import ( # noqa: F401
|
||||
CommContext,
|
||||
Cost,
|
||||
_g_op_cost_factory,
|
||||
build_comm_costs_from_descs,
|
||||
build_comm_desc,
|
||||
build_comm_desc_from_dist_op,
|
||||
build_comp_costs_from_descs,
|
||||
build_comp_desc_from_dist_op,
|
||||
build_comp_desc_str_for_predict,
|
||||
build_dp_costs,
|
||||
calc_time_by_cost_model,
|
||||
)
|
||||
from .comm_op_cost import ( # noqa: F401
|
||||
AllgatherOpCost,
|
||||
AllReduceOpCost,
|
||||
AllreduceSumOpCost,
|
||||
BroadcastOpCost,
|
||||
IdentityOpCost,
|
||||
RecvOpCost,
|
||||
SendOpCost,
|
||||
)
|
||||
from .comp_op_cost import ( # noqa: F401
|
||||
ConcatOpCost,
|
||||
EmbeddingGradOpCost,
|
||||
EmbeddingOpCost,
|
||||
FillConstantBatchSizeLikeOpCost,
|
||||
MatmulGradOpCost,
|
||||
MatmulOpCost,
|
||||
MatmulV2GradOpCost,
|
||||
MatmulV2OpCost,
|
||||
MulGradOpCost,
|
||||
MulOpCost,
|
||||
Reshape2GradOpCost,
|
||||
Reshape2OpCost,
|
||||
SliceOpCost,
|
||||
SoftmaxGradOpCost,
|
||||
SoftmaxOpCost,
|
||||
SplitOpCost,
|
||||
Transpose2GradOpCost,
|
||||
Transpose2OpCost,
|
||||
)
|
||||
from .estimate_cost import CostEstimator # noqa: F401
|
||||
from .op_runtime_cost import ( # noqa: F401
|
||||
check_if_op_supports_runtime_profiling,
|
||||
measure_program_real_op_cost,
|
||||
)
|
||||
from .tensor_cost import TensorCost # noqa: F401
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,316 @@
|
||||
# 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 math
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
|
||||
from .base_cost import CommOpCost, register_op_cost
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AllreduceSumOpCost(CommOpCost):
|
||||
OP_TYPE = "c_allreduce_sum"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
# use tree if cross machine and use ring if in a single machine
|
||||
time = None
|
||||
cluster = self.comm_context.cluster
|
||||
if not cluster.cross_machine(self.group_ranks):
|
||||
time = self.calc_time_ring()
|
||||
else:
|
||||
time = self.calc_time_tree()
|
||||
|
||||
return time
|
||||
|
||||
def calc_time_ring(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
alpha += (
|
||||
2
|
||||
* (self.rank_count - self.machine_count)
|
||||
* self.comm_context.intra_ring
|
||||
)
|
||||
alpha += (
|
||||
2
|
||||
* (self.machine_count - 1)
|
||||
* (
|
||||
self.comm_context.inter_ring
|
||||
+ self.hops * self.comm_context.switch
|
||||
)
|
||||
)
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = (
|
||||
alpha
|
||||
+ 2
|
||||
* (self.rank_count - 1)
|
||||
/ self.rank_count
|
||||
* self.comm_count
|
||||
* beta
|
||||
)
|
||||
|
||||
return time
|
||||
|
||||
def calc_time_tree(self):
|
||||
alpha = self.comm_context.base_tree
|
||||
alpha += (
|
||||
2
|
||||
* (self.rank_count / self.machine_count - 1)
|
||||
* self.comm_context.intra_tree
|
||||
)
|
||||
alpha += math.log2(self.machine_count) * (
|
||||
self.comm_context.inter_tree + self.hops * self.comm_context.switch
|
||||
)
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
|
||||
time = alpha + 2 * self.comm_count * beta
|
||||
|
||||
return time
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AllReduceOpCost(CommOpCost):
|
||||
OP_TYPE = "all_reduce"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
# use tree if cross machine and use ring if in a single machine
|
||||
time = None
|
||||
cluster = self.comm_context.cluster
|
||||
if not cluster.cross_machine(self.group_ranks):
|
||||
time = self.calc_time_ring()
|
||||
else:
|
||||
time = self.calc_time_tree()
|
||||
|
||||
return time
|
||||
|
||||
def calc_time_ring(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
alpha += (
|
||||
2
|
||||
* (self.rank_count - self.machine_count)
|
||||
* self.comm_context.intra_ring
|
||||
)
|
||||
alpha += (
|
||||
2
|
||||
* (self.machine_count - 1)
|
||||
* (
|
||||
self.comm_context.inter_ring
|
||||
+ self.hops * self.comm_context.switch
|
||||
)
|
||||
)
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = (
|
||||
alpha
|
||||
+ 2
|
||||
* (self.rank_count - 1)
|
||||
/ self.rank_count
|
||||
* self.comm_count
|
||||
* beta
|
||||
)
|
||||
|
||||
return time
|
||||
|
||||
def calc_time_tree(self):
|
||||
alpha = self.comm_context.base_tree
|
||||
alpha += (
|
||||
2
|
||||
* (self.rank_count / self.machine_count - 1)
|
||||
* self.comm_context.intra_tree
|
||||
)
|
||||
alpha += math.log2(self.machine_count) * (
|
||||
self.comm_context.inter_tree + self.hops * self.comm_context.switch
|
||||
)
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
|
||||
time = alpha + 2 * self.comm_count * beta
|
||||
|
||||
return time
|
||||
|
||||
@property
|
||||
def comm_count(self):
|
||||
from ..reshard import get_var_with_recursion
|
||||
|
||||
if self._comm_count is None:
|
||||
dtype = None
|
||||
shape = None
|
||||
if self.op is not None:
|
||||
vars = self.op.block.vars
|
||||
try:
|
||||
var_name = self.op.input("x")[0]
|
||||
except:
|
||||
var_name = self.op.output("out")[0]
|
||||
var = get_var_with_recursion(
|
||||
var_name, self.op.block, self.op.block.program
|
||||
)
|
||||
dtype = var.dtype
|
||||
shape = var.shape
|
||||
elif self.op_desc is not None:
|
||||
dtype = self.op_desc["inputs"]["x"][0][0]
|
||||
shape = self.op_desc["inputs"]["x"][0][1]
|
||||
|
||||
factor = None
|
||||
if dtype == paddle.float32 or dtype == paddle.int32:
|
||||
factor = 4
|
||||
else:
|
||||
raise ValueError(f"Unsupported comm dtype {dtype}")
|
||||
comm_count = int(np.prod(shape)) * factor
|
||||
self._comm_count = comm_count
|
||||
|
||||
return self._comm_count
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AllgatherOpCost(CommOpCost):
|
||||
OP_TYPE = "all_gather"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
time = self.calc_time_ring()
|
||||
return time
|
||||
|
||||
def calc_time_ring(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
alpha += (
|
||||
self.rank_count - self.machine_count
|
||||
) * self.comm_context.intra_ring
|
||||
alpha += (self.machine_count - 1) * (
|
||||
self.comm_context.inter_ring + self.hops * self.comm_context.switch
|
||||
)
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = (
|
||||
alpha
|
||||
+ (self.rank_count - 1) / self.rank_count * self.comm_count * beta
|
||||
)
|
||||
return time
|
||||
|
||||
@property
|
||||
def comm_count(self):
|
||||
from ..reshard import get_var_with_recursion
|
||||
|
||||
if self._comm_count is None:
|
||||
dtype = None
|
||||
shape = None
|
||||
if self.op is not None:
|
||||
vars = self.op.block.vars
|
||||
try:
|
||||
var_name = self.op.input("x")[0]
|
||||
except:
|
||||
var_name = self.op.output("out")[0]
|
||||
var = get_var_with_recursion(
|
||||
var_name, self.op.block, self.op.block.program
|
||||
)
|
||||
dtype = var.dtype
|
||||
shape = var.shape
|
||||
elif self.op_desc is not None:
|
||||
dtype = self.op_desc["inputs"]["X"][0][0]
|
||||
shape = self.op_desc["inputs"]["X"][0][1]
|
||||
|
||||
factor = None
|
||||
if dtype == paddle.float32 or dtype == paddle.int32:
|
||||
factor = 4
|
||||
else:
|
||||
raise ValueError(f"Unsupported comm dtype {dtype}")
|
||||
comm_count = int(np.prod(shape)) * factor
|
||||
self._comm_count = comm_count
|
||||
|
||||
return self._comm_count
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class BroadcastOpCost(CommOpCost):
|
||||
OP_TYPE = "broadcast"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
time = self.calc_time_ring()
|
||||
return time
|
||||
|
||||
def calc_time_ring(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
if self.machine_count > 1:
|
||||
alpha += (
|
||||
self.comm_context.inter_ring
|
||||
+ self.hops * self.comm_context.switch
|
||||
)
|
||||
else:
|
||||
alpha += self.comm_context.intra_ring
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = alpha + self.comm_count * beta
|
||||
|
||||
return time
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class IdentityOpCost(CommOpCost):
|
||||
OP_TYPE = "c_identity"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
return self.comm_count * 1 / (144 * 1e3)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class RecvOpCost(CommOpCost):
|
||||
OP_TYPE = "recv_v2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
if self.machine_count > 1:
|
||||
alpha += (
|
||||
self.comm_context.inter_ring
|
||||
+ self.hops * self.comm_context.switch
|
||||
)
|
||||
else:
|
||||
alpha += self.comm_context.intra_ring
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = alpha + self.comm_count * beta
|
||||
return time
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SendOpCost(CommOpCost):
|
||||
OP_TYPE = "send_v2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
if self.machine_count > 1:
|
||||
alpha += (
|
||||
self.comm_context.inter_ring
|
||||
+ self.hops * self.comm_context.switch
|
||||
)
|
||||
else:
|
||||
alpha += self.comm_context.intra_ring
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = alpha + self.comm_count * beta
|
||||
|
||||
return time
|
||||
@@ -0,0 +1,591 @@
|
||||
# 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
|
||||
|
||||
from .base_cost import CompOpCost, register_op_cost
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AdamOpCost(CompOpCost):
|
||||
OP_TYPE = "adam"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ArgsortOpCost(CompOpCost):
|
||||
OP_TYPE = "argsort"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AssignOpCost(CompOpCost):
|
||||
OP_TYPE = "assign"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AssignValueOpCost(CompOpCost):
|
||||
OP_TYPE = "assign_value"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class BeamSearchOpCost(CompOpCost):
|
||||
OP_TYPE = "beam_search"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class BeamSearchDecodeOpCost(CompOpCost):
|
||||
OP_TYPE = "beam_search_decode"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class CastOpCost(CompOpCost):
|
||||
OP_TYPE = "cast"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ConcatOpCost(CompOpCost):
|
||||
OP_TYPE = "concat"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class DropoutOpCost(CompOpCost):
|
||||
OP_TYPE = "dropout"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class DropoutGradOpCost(CompOpCost):
|
||||
OP_TYPE = "dropout_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseAddOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_add"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseAddGradOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_add_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseDivOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_div"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseDivGradOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_div_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseMulOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_mul"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseMulGradOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_mul_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseSubOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_sub"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseSubGradOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_sub_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class EqualOpCost(CompOpCost):
|
||||
OP_TYPE = "equal"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class EmbeddingOpCost(CompOpCost):
|
||||
OP_TYPE = "c_embedding"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class EmbeddingGradOpCost(CompOpCost):
|
||||
OP_TYPE = "c_embedding_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class FillConstantOpCost(CompOpCost):
|
||||
OP_TYPE = "fill_constant"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class FillConstantBatchSizeLikeOpCost(CompOpCost):
|
||||
OP_TYPE = "fill_constant_batch_size_like"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class FusedSoftmaxMaskUpperTriangleOpCost(CompOpCost):
|
||||
OP_TYPE = "fused_softmax_mask_upper_triangle"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class FusedSoftmaxMaskUpperTriangleGradOpCost(CompOpCost):
|
||||
OP_TYPE = "fused_softmax_mask_upper_triangle_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class GatherOpCost(CompOpCost):
|
||||
OP_TYPE = "gather"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class GeluOpCost(CompOpCost):
|
||||
OP_TYPE = "gelu"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class GeluGradOpCost(CompOpCost):
|
||||
OP_TYPE = "gelu_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class GreaterEqualOpCost(CompOpCost):
|
||||
OP_TYPE = "greater_equal"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class IncrementOpCost(CompOpCost):
|
||||
OP_TYPE = "increment"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class IsEmptyOpCost(CompOpCost):
|
||||
OP_TYPE = "is_empty"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LayerNormOpCost(CompOpCost):
|
||||
OP_TYPE = "layer_norm"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LayerNormGradOpCost(CompOpCost):
|
||||
OP_TYPE = "layer_norm_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LessThanOpCost(CompOpCost):
|
||||
OP_TYPE = "less_than"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LogicalNotOpCost(CompOpCost):
|
||||
OP_TYPE = "logical_not"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LogicalAndOpCost(CompOpCost):
|
||||
OP_TYPE = "logical_and"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LodResetOpCost(CompOpCost):
|
||||
OP_TYPE = "lod_reset"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LogOpCost(CompOpCost):
|
||||
OP_TYPE = "log"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LookupTableV2OpCost(CompOpCost):
|
||||
OP_TYPE = "lookup_table_v2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LookupTableV2GradOpCost(CompOpCost):
|
||||
OP_TYPE = "lookup_table_v2_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MatmulOpCost(CompOpCost):
|
||||
OP_TYPE = "matmul"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MatmulGradOpCost(CompOpCost):
|
||||
OP_TYPE = "matmul_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MatmulV2OpCost(CompOpCost):
|
||||
OP_TYPE = "matmul_v2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MatmulV2GradOpCost(CompOpCost):
|
||||
OP_TYPE = "matmul_v2_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MemcpyOpCost(CompOpCost):
|
||||
OP_TYPE = "memcpy"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MulOpCost(CompOpCost):
|
||||
OP_TYPE = "mul"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MulGradOpCost(CompOpCost):
|
||||
OP_TYPE = "mul_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class OneHotOpCost(CompOpCost):
|
||||
OP_TYPE = "one_hot"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ReadFromArrayOpCost(CompOpCost):
|
||||
OP_TYPE = "read_from_array"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ReduceSumOpCost(CompOpCost):
|
||||
OP_TYPE = "reduce_sum"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ReduceSumGradOpCost(CompOpCost):
|
||||
OP_TYPE = "reduce_sum_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Reshape2OpCost(CompOpCost):
|
||||
OP_TYPE = "reshape2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Reshape2GradOpCost(CompOpCost):
|
||||
OP_TYPE = "reshape2_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ReduceMeanOpCost(CompOpCost):
|
||||
OP_TYPE = "reduce_mean"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ReduceMeanGradOpCost(CompOpCost):
|
||||
OP_TYPE = "reduce_mean_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ScaleOpCost(CompOpCost):
|
||||
OP_TYPE = "scale"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ShapeOpCost(CompOpCost):
|
||||
OP_TYPE = "shape"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SliceOpCost(CompOpCost):
|
||||
OP_TYPE = "slice"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SoftmaxOpCost(CompOpCost):
|
||||
OP_TYPE = "softmax"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SoftmaxGradOpCost(CompOpCost):
|
||||
OP_TYPE = "softmax_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SoftmaxWithCrossEntropyOpCost(CompOpCost):
|
||||
OP_TYPE = "softmax_with_cross_entropy"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SoftmaxWithCrossEntropyGradOpCost(CompOpCost):
|
||||
OP_TYPE = "softmax_with_cross_entropy_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SplitOpCost(CompOpCost):
|
||||
OP_TYPE = "split"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Squeeze2OpCost(CompOpCost):
|
||||
OP_TYPE = "squeeze2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SquareOpCost(CompOpCost):
|
||||
OP_TYPE = "square"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SquareGradOpCost(CompOpCost):
|
||||
OP_TYPE = "square_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SumOpCost(CompOpCost):
|
||||
OP_TYPE = "sum"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class TopKOpCost(CompOpCost):
|
||||
OP_TYPE = "top_k"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Transpose2OpCost(CompOpCost):
|
||||
OP_TYPE = "transpose2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Transpose2GradOpCost(CompOpCost):
|
||||
OP_TYPE = "transpose2_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Unsqueeze2OpCost(CompOpCost):
|
||||
OP_TYPE = "unsqueeze2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class WriteToArrayOpCost(CompOpCost):
|
||||
OP_TYPE = "write_to_array"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
@@ -0,0 +1,671 @@
|
||||
# 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
|
||||
|
||||
from collections import OrderedDict
|
||||
from functools import reduce
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
||||
|
||||
from ..dist_tensor import DistributedTensor
|
||||
from ..operators.common import get_distributed_operator_impl_container
|
||||
from .base_cost import Cost
|
||||
|
||||
|
||||
class CostEstimator:
|
||||
_special_op_type = ["fused_attention", "fused_feedforward"]
|
||||
|
||||
def __init__(
|
||||
self, program, cluster, mode="modeling", rank=None, loop_count=10
|
||||
):
|
||||
self._program = program
|
||||
self._cluster = cluster
|
||||
self._check_mode(mode)
|
||||
self._mode = mode
|
||||
self._rank = rank if rank is not None else paddle.distributed.get_rank()
|
||||
self._loop_count = loop_count
|
||||
self._global_cost = Cost()
|
||||
self._local_cost_mapping = {}
|
||||
self._detailed_cost = OrderedDict() # {`op_id`: {"reshard": [], "dist_op": [], "local_cost": local_cost}}}
|
||||
self._bubble_time_mapping = {}
|
||||
self._ordered_ops = []
|
||||
self.max_memories = {}
|
||||
self.max_memory = None
|
||||
|
||||
@property
|
||||
def loop_count(self):
|
||||
return self._loop_count
|
||||
|
||||
@property
|
||||
def detailed_cost(self):
|
||||
return self._detailed_cost
|
||||
|
||||
@property
|
||||
def program(self):
|
||||
return self._program
|
||||
|
||||
@property
|
||||
def rank(self):
|
||||
return self._rank
|
||||
|
||||
@property
|
||||
def dist_context(self):
|
||||
return self._dist_context
|
||||
|
||||
@property
|
||||
def cluster(self):
|
||||
return self._cluster
|
||||
|
||||
@property
|
||||
def mode(self):
|
||||
return self._mode
|
||||
|
||||
@property
|
||||
def global_cost(self):
|
||||
max_time = 0
|
||||
memory = 0
|
||||
flops = 0
|
||||
for rank in self._local_cost_mapping:
|
||||
cost = self._local_cost_mapping[rank]
|
||||
if cost.time > max_time:
|
||||
max_time = cost.time
|
||||
memory += cost.memory
|
||||
flops += cost.flops
|
||||
self._global_cost.time = max_time
|
||||
self._global_cost.memory = memory
|
||||
self._global_cost.flops = flops
|
||||
return self._global_cost
|
||||
|
||||
def local_cost(self, rank=None):
|
||||
rank = self.rank if rank is None else rank
|
||||
if rank not in self._local_cost_mapping:
|
||||
self._local_cost_mapping[rank] = Cost()
|
||||
|
||||
return self._local_cost_mapping[rank]
|
||||
|
||||
def local_bubble_time(self, rank=None):
|
||||
rank = self.rank if rank is None else rank
|
||||
return self._bubble_time_mapping[rank]
|
||||
|
||||
def _check_mode(self, mode):
|
||||
if mode not in ["modeling", "profiling"]:
|
||||
raise ValueError(
|
||||
f"Just support modeling and profiling, but got {mode}"
|
||||
)
|
||||
|
||||
def _is_special_var_name(self, var_name):
|
||||
special_var_name = ["lod_tensor_blocking_queue_0"]
|
||||
if var_name in special_var_name:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _estimate_core(self, dist_context, resharder, block):
|
||||
from ..reshard import get_var_with_recursion
|
||||
|
||||
ops = block.ops
|
||||
loop_count = None
|
||||
if block.desc.id != self.program.global_block().desc.id:
|
||||
loop_count = self.loop_count
|
||||
else:
|
||||
loop_count = 1
|
||||
for i in range(loop_count):
|
||||
for op in ops:
|
||||
self._detailed_cost[op.desc.id()] = OrderedDict()
|
||||
# If in the while sub block, the detail of cost is the last cost
|
||||
detail = self._detailed_cost[op.desc.id()]
|
||||
detail["reshard_cost"] = OrderedDict() #
|
||||
detail["dist_op_cost"] = []
|
||||
if int(op.attr('op_role')) == int(OpRole.Optimize):
|
||||
continue
|
||||
if op.type in [
|
||||
"create_py_reader",
|
||||
"create_double_buffer_reader",
|
||||
"read",
|
||||
]:
|
||||
continue
|
||||
|
||||
# NOTE: It does not support nested loop and just supports while op when op has sub block now.
|
||||
if op.type == "while":
|
||||
while_block = self.program.blocks[op.attr("sub_block").id]
|
||||
self._estimate_core(dist_context, resharder, while_block)
|
||||
continue
|
||||
|
||||
for var_name in op.input_arg_names:
|
||||
if self._is_special_var_name(var_name):
|
||||
continue
|
||||
var = get_var_with_recursion(var_name, block, self.program)
|
||||
reshard_cost = resharder.get_cost(op, var, self.cluster)
|
||||
|
||||
# Calc reshard cost
|
||||
if reshard_cost is not None:
|
||||
detail["reshard_cost"][var_name] = reshard_cost
|
||||
|
||||
comm_costs = reshard_cost[0]
|
||||
local_comp_cost = reshard_cost[1]
|
||||
for comm_cost in comm_costs:
|
||||
# Time is cumulative in global cost and local cost, but memory and flops just are cumulative in global cost.
|
||||
# Comm sync
|
||||
for item in comm_cost:
|
||||
group_ranks, cost = item
|
||||
max_time = None
|
||||
cost_time = {}
|
||||
for rank in group_ranks:
|
||||
rank_cost = self.local_cost(rank)
|
||||
cost_time[rank] = rank_cost.time
|
||||
if max_time is None:
|
||||
max_time = rank_cost.time
|
||||
else:
|
||||
if max_time < rank_cost.time:
|
||||
max_time = rank_cost.time
|
||||
|
||||
for rank in group_ranks:
|
||||
self.local_cost(rank).time = (
|
||||
max_time + cost.time
|
||||
)
|
||||
|
||||
if rank not in self._bubble_time_mapping:
|
||||
self._bubble_time_mapping[rank] = 0
|
||||
|
||||
self._bubble_time_mapping[rank] += (
|
||||
max_time - cost_time[rank]
|
||||
)
|
||||
|
||||
for rank in local_comp_cost:
|
||||
for comp_cost in local_comp_cost[rank]:
|
||||
self.local_cost(rank).time += comp_cost.time
|
||||
|
||||
# Calc dist op cost
|
||||
dist_op = dist_context.get_dist_op_for_program(op)
|
||||
if not dist_op:
|
||||
continue
|
||||
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
processes = op_dist_attr.process_mesh.process_ids
|
||||
|
||||
container = get_distributed_operator_impl_container(
|
||||
op_dist_attr.impl_type
|
||||
)
|
||||
dist_impl = container.impls[op_dist_attr.impl_idx]
|
||||
|
||||
dist_op_cost = dist_impl.calc_cost(
|
||||
op.attr('op_role'), dist_op, dist_context, self.cluster
|
||||
)
|
||||
detail["dist_op_cost"] = dist_op_cost
|
||||
|
||||
if dist_op_cost is None:
|
||||
assert (
|
||||
dist_op.serial_op.type in CostEstimator._special_op_type
|
||||
)
|
||||
continue
|
||||
for item in dist_op_cost:
|
||||
if isinstance(item, list):
|
||||
# Comm sync
|
||||
for comm_op_cost in item:
|
||||
max_time = None
|
||||
cost_time = {}
|
||||
group_ranks = comm_op_cost.group_ranks
|
||||
for rank in comm_op_cost.group_ranks:
|
||||
rank_cost = self.local_cost(rank)
|
||||
cost_time[rank] = rank_cost.time
|
||||
if max_time is None:
|
||||
max_time = rank_cost.time
|
||||
else:
|
||||
if max_time < rank_cost.time:
|
||||
max_time = rank_cost.time
|
||||
for rank in group_ranks:
|
||||
self.local_cost(rank).time = (
|
||||
max_time + comm_op_cost.time
|
||||
if op.attr('op_role') != OpRole.Backward
|
||||
else max_time + 0.9 * comm_op_cost.time
|
||||
)
|
||||
if rank not in self._bubble_time_mapping:
|
||||
self._bubble_time_mapping[rank] = 0
|
||||
self._bubble_time_mapping[rank] += (
|
||||
max_time - cost_time[rank]
|
||||
)
|
||||
elif isinstance(item, dict):
|
||||
# Op just one
|
||||
for rank in processes:
|
||||
# DP+PP+MP
|
||||
if rank not in item:
|
||||
continue
|
||||
self.local_cost(rank).time += item[rank].time
|
||||
|
||||
def prepare(self):
|
||||
self._global_cost = Cost()
|
||||
self._local_cost_mapping = {}
|
||||
self._detailed_cost = OrderedDict()
|
||||
self._bubble_time_mapping = {}
|
||||
|
||||
def _calculate_bytes(self, sizes, dtype):
|
||||
if sizes:
|
||||
total_count = reduce(lambda x, y: x * y, sizes, 1)
|
||||
else:
|
||||
total_count = 0
|
||||
|
||||
if dtype == paddle.float64 or dtype == paddle.int64:
|
||||
dtype_factor = 8
|
||||
elif dtype == paddle.float32 or dtype == paddle.int32:
|
||||
dtype_factor = 4
|
||||
elif (
|
||||
dtype == paddle.float16
|
||||
or dtype == paddle.bfloat16
|
||||
or dtype == paddle.int16
|
||||
):
|
||||
dtype_factor = 2
|
||||
elif dtype == paddle.int8 or dtype == paddle.uint8:
|
||||
dtype_factor = 1
|
||||
else:
|
||||
dtype_factor = 8
|
||||
|
||||
memory = total_count * dtype_factor
|
||||
return memory
|
||||
|
||||
def _estimate_max_memory_by_dist_op(self, dist_context):
|
||||
# This estimation will be improved, now reshard and inplace are not considered.
|
||||
# Persist var is not free.
|
||||
def _convert_pm_and_dm_to_str(process_mesh, dims_mapping):
|
||||
processes = ",".join([str(x) for x in process_mesh.process_ids])
|
||||
topology = ",".join([str(x) for x in process_mesh.shape])
|
||||
dims_mapping = ",".join([str(x) for x in dims_mapping])
|
||||
result = processes + topology + dims_mapping
|
||||
return result
|
||||
|
||||
memories = {}
|
||||
self.max_memories = {}
|
||||
var_info = {} # var_name: [[process_mesh, dims_mapping], [id]], [[process_mesh, dims_mapping], [id]]}
|
||||
|
||||
for block in self.program.blocks:
|
||||
for op in block.ops:
|
||||
self._ordered_ops.append([op.desc.id(), op])
|
||||
self._ordered_ops.sort(key=lambda x: x[0])
|
||||
|
||||
parameters = set()
|
||||
for op_id, op in self._ordered_ops:
|
||||
if op.type in [
|
||||
"create_py_reader",
|
||||
"create_double_buffer_reader",
|
||||
"read",
|
||||
]:
|
||||
continue
|
||||
dist_op = dist_context.get_dist_op_for_program(op)
|
||||
if not dist_op:
|
||||
continue
|
||||
process_mesh = dist_op.dist_attr.process_mesh
|
||||
for var_name in op.input_arg_names:
|
||||
input_dims_mapping = dist_op.dist_attr.get_input_dims_mapping(
|
||||
var_name
|
||||
)
|
||||
|
||||
if var_name not in var_info:
|
||||
var_info[var_name] = {}
|
||||
key = _convert_pm_and_dm_to_str(
|
||||
process_mesh, input_dims_mapping
|
||||
)
|
||||
if key not in var_info[var_name]:
|
||||
var_info[var_name][key] = {}
|
||||
# It is even partition now
|
||||
if "position" not in var_info[var_name][key]:
|
||||
var_info[var_name][key]["position"] = []
|
||||
var_info[var_name][key]["position"].append(op_id)
|
||||
|
||||
if "memory" not in var_info[var_name][key]:
|
||||
var = dist_op.get_serial_input(var_name)
|
||||
global_sizes = var.shape
|
||||
dtype = var.dtype
|
||||
sizes = DistributedTensor.get_local_sizes(
|
||||
global_sizes,
|
||||
input_dims_mapping,
|
||||
process_mesh.shape,
|
||||
process_mesh.process_ids,
|
||||
)
|
||||
var_info[var_name][key]["memory"] = self._calculate_bytes(
|
||||
sizes, dtype
|
||||
)
|
||||
if var.persistable:
|
||||
name = var_name + key
|
||||
if name not in parameters:
|
||||
parameters.add(name)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in memories:
|
||||
memories[process] = 0
|
||||
memories[process] += var_info[var_name][key][
|
||||
"memory"
|
||||
]
|
||||
|
||||
for var_name in op.output_arg_names:
|
||||
output_dims_mapping = dist_op.dist_attr.get_output_dims_mapping(
|
||||
var_name
|
||||
)
|
||||
if var_name not in var_info:
|
||||
var_info[var_name] = {}
|
||||
key = _convert_pm_and_dm_to_str(
|
||||
process_mesh, output_dims_mapping
|
||||
)
|
||||
if key not in var_info[var_name]:
|
||||
var_info[var_name][key] = {}
|
||||
if "position" not in var_info[var_name][key]:
|
||||
var_info[var_name][key]["position"] = []
|
||||
var_info[var_name][key]["position"].append(op_id)
|
||||
|
||||
if "memory" not in var_info[var_name][key]:
|
||||
var = dist_op.get_serial_output(var_name)
|
||||
global_sizes = var.shape
|
||||
dtype = var.dtype
|
||||
sizes = DistributedTensor.get_local_sizes(
|
||||
global_sizes,
|
||||
output_dims_mapping,
|
||||
process_mesh.shape,
|
||||
process_mesh.process_ids,
|
||||
)
|
||||
var_info[var_name][key]["memory"] = self._calculate_bytes(
|
||||
sizes, dtype
|
||||
)
|
||||
if var.persistable:
|
||||
name = var_name + key
|
||||
if name not in parameters:
|
||||
parameters.add(name)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in memories:
|
||||
memories[process] = 0
|
||||
memories[process] += var_info[var_name][key][
|
||||
"memory"
|
||||
]
|
||||
|
||||
has_used_vars = set()
|
||||
not_calc_vars = set()
|
||||
for op_id, op in self._ordered_ops:
|
||||
if op.type in [
|
||||
"create_py_reader",
|
||||
"create_double_buffer_reader",
|
||||
"read",
|
||||
]:
|
||||
continue
|
||||
can_free_memories = {}
|
||||
can_free_vars = set()
|
||||
dist_op = dist_context.get_dist_op_for_program(op)
|
||||
if not dist_op:
|
||||
continue
|
||||
process_mesh = dist_op.dist_attr.process_mesh
|
||||
for var_name in op.input_arg_names:
|
||||
input_dims_mapping = dist_op.dist_attr.get_input_dims_mapping(
|
||||
var_name
|
||||
)
|
||||
key = _convert_pm_and_dm_to_str(
|
||||
process_mesh, input_dims_mapping
|
||||
)
|
||||
has_used_var = var_name + key
|
||||
var = dist_op.get_serial_input(var_name)
|
||||
# Not used
|
||||
if (
|
||||
has_used_var not in has_used_vars
|
||||
and has_used_var not in parameters
|
||||
):
|
||||
if has_used_var in not_calc_vars:
|
||||
continue
|
||||
has_used_vars.add(has_used_var)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in memories:
|
||||
memories[process] = 0
|
||||
memories[process] += var_info[var_name][key]["memory"]
|
||||
# Used
|
||||
if op_id == var_info[var_name][key]["position"][-1]:
|
||||
if (
|
||||
has_used_var not in can_free_vars
|
||||
and not var.persistable
|
||||
):
|
||||
can_free_vars.add(has_used_var)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in can_free_memories:
|
||||
can_free_memories[process] = 0
|
||||
can_free_memories[process] += var_info[var_name][
|
||||
key
|
||||
]["memory"]
|
||||
|
||||
for var_name in op.output_arg_names:
|
||||
output_dims_mapping = dist_op.dist_attr.get_output_dims_mapping(
|
||||
var_name
|
||||
)
|
||||
key = _convert_pm_and_dm_to_str(
|
||||
process_mesh, output_dims_mapping
|
||||
)
|
||||
has_used_var = var_name + key
|
||||
var = dist_op.get_serial_output(var_name)
|
||||
if (
|
||||
op.type == "reshape2"
|
||||
or op.type == "transpose2"
|
||||
or op.type == "elementwise_add"
|
||||
):
|
||||
not_calc_vars.add(has_used_var)
|
||||
continue
|
||||
# Not used
|
||||
if (
|
||||
has_used_var not in has_used_vars
|
||||
and has_used_var not in parameters
|
||||
):
|
||||
has_used_vars.add(has_used_var)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in memories:
|
||||
memories[process] = 0
|
||||
memories[process] += var_info[var_name][key]["memory"]
|
||||
# Used
|
||||
if op_id == var_info[var_name][key]["position"][-1]:
|
||||
if (
|
||||
has_used_var not in can_free_vars
|
||||
and not var.persistable
|
||||
):
|
||||
can_free_vars.add(has_used_var)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in can_free_memories:
|
||||
can_free_memories[process] = 0
|
||||
can_free_memories[process] += var_info[var_name][
|
||||
key
|
||||
]["memory"]
|
||||
|
||||
# Calc peak memory
|
||||
for process in memories:
|
||||
if process not in self.max_memories:
|
||||
self.max_memories[process] = memories[process]
|
||||
else:
|
||||
if memories[process] > self.max_memories[process]:
|
||||
self.max_memories[process] = memories[process]
|
||||
# Free memory
|
||||
for process in can_free_memories:
|
||||
if process in memories:
|
||||
memories[process] -= can_free_memories[process]
|
||||
|
||||
# Calculate the max memory in all ranks
|
||||
max_memory = max(self.max_memories.values())
|
||||
self.max_memory = max_memory
|
||||
|
||||
return max_memory
|
||||
|
||||
def estimate(self, dist_context, resharder=None):
|
||||
self.prepare()
|
||||
from ..reshard import Resharder
|
||||
|
||||
resharder = (
|
||||
Resharder(self.program, None, self.rank, dist_context, [])
|
||||
if resharder is None
|
||||
else resharder
|
||||
)
|
||||
|
||||
block = self.program.global_block()
|
||||
self._estimate_core(dist_context, resharder, block)
|
||||
|
||||
return self.global_cost
|
||||
|
||||
def _print_tag(self, max_len, length):
|
||||
tag = "+" + "-" * max_len
|
||||
for i in range(length):
|
||||
print(tag, end="")
|
||||
if i == length - 1:
|
||||
print("+")
|
||||
|
||||
def _print_vals(self, vals, max_len):
|
||||
for idx, val in enumerate(vals):
|
||||
s = "|" + str(val).center(max_len)
|
||||
print(s, end="")
|
||||
if idx == len(vals) - 1:
|
||||
print("|")
|
||||
|
||||
def _pretty_print_memory_cost(self):
|
||||
"""Print memory of every rank prettily."""
|
||||
if not self.max_memories or not self.max_memory:
|
||||
raise ValueError("Please calculate memory cost before print.")
|
||||
|
||||
# Padding automatically
|
||||
max_len = 0
|
||||
header = ["Rank", "Memory(MiB)"]
|
||||
memories = [
|
||||
int(item // 1e6) for item in list(self.max_memories.values())
|
||||
]
|
||||
for memory in memories + header:
|
||||
if len(str(memory)) > max_len:
|
||||
max_len = len(str(memory))
|
||||
max_len += 4 # for pretty print of center
|
||||
|
||||
# Print tag
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
# Print header
|
||||
self._print_vals(header, max_len)
|
||||
|
||||
# Print tag
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
# Print rank and its memory
|
||||
for i in range(len(self.max_memories)):
|
||||
memory = memories[i]
|
||||
vals = [i, memory]
|
||||
self._print_vals(vals, max_len)
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
def _pretty_print_global(self):
|
||||
"""Print global execution time and max memory prettily."""
|
||||
if not self.max_memories or not self.max_memory:
|
||||
raise ValueError("Please calculate cost before print.")
|
||||
|
||||
# Padding automatically
|
||||
max_len = 0
|
||||
header = ["Execution Time(us)", "Max Memory(MiB)"]
|
||||
vals = [round(self.global_cost.time, 3), int(self.max_memory // 1e6)]
|
||||
for memory in vals + header:
|
||||
if len(str(memory)) > max_len:
|
||||
max_len = len(str(memory))
|
||||
max_len += 4 # for pretty print of center
|
||||
|
||||
# Print tag
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
# Print header
|
||||
self._print_vals(header, max_len)
|
||||
|
||||
# Print tag
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
# Print exec time and max memory
|
||||
self._print_vals(vals, max_len)
|
||||
|
||||
# Print tag
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
def pretty_print_cost(self):
|
||||
"""Print cost prettily."""
|
||||
print("The global execution time and max memory are as follows:")
|
||||
self._pretty_print_global()
|
||||
print("The memory of every rank is as follows:")
|
||||
self._pretty_print_memory_cost()
|
||||
|
||||
|
||||
def get_cost_from_engine(engine, mode):
|
||||
import copy
|
||||
|
||||
from ..utils import to_list
|
||||
|
||||
# Construct cost estimator by original main program
|
||||
serial_main_prog = (
|
||||
engine._fwd_main_progs[mode].clone()
|
||||
if mode in engine._fwd_main_progs
|
||||
else engine._orig_main_prog.clone()
|
||||
)
|
||||
|
||||
serial_startup_prog = (
|
||||
engine._fwd_dist_contexts[mode]._original_serial_main_program.clone()
|
||||
if mode in engine._fwd_dist_contexts
|
||||
else engine._orig_startup_prog.clone()
|
||||
)
|
||||
losses = (
|
||||
to_list(engine._loss)
|
||||
if (
|
||||
not isinstance(engine._loss, paddle.nn.Layer)
|
||||
and not callable(engine._loss)
|
||||
)
|
||||
else engine._losses
|
||||
)
|
||||
serial_optimizer = copy.deepcopy(engine._orig_optimizer)
|
||||
if mode in engine._fwd_dist_contexts:
|
||||
dist_context = copy.deepcopy(engine._fwd_dist_contexts[mode])
|
||||
else:
|
||||
from ..dist_context import DistributedContext
|
||||
|
||||
dist_context = DistributedContext(
|
||||
serial_main_prog,
|
||||
serial_startup_prog,
|
||||
serial_optimizer,
|
||||
losses,
|
||||
{},
|
||||
{"loss": losses},
|
||||
engine._cluster,
|
||||
engine._strategy,
|
||||
)
|
||||
from ..completion import Completer
|
||||
|
||||
completer = Completer(dist_context)
|
||||
completer.complete_forward_annotation()
|
||||
dist_context.block_state.parse_forward_blocks(
|
||||
dist_context.serial_main_program
|
||||
)
|
||||
|
||||
if mode == "eval" or mode == "predict":
|
||||
cost_estimator = CostEstimator(serial_main_prog, engine._cluster)
|
||||
elif mode == "train":
|
||||
from ..parallelizer_v2 import Parallelizer
|
||||
|
||||
# Get serial main program with backward
|
||||
parallelizer = Parallelizer(mode, completer, dist_context)
|
||||
# Generate backward
|
||||
loss_name = dist_context.serial_loss.name
|
||||
serial_loss = serial_main_prog.global_block()._var_recursive(loss_name)
|
||||
params_grads = parallelizer._generate_backward(
|
||||
serial_main_prog, serial_startup_prog, serial_loss
|
||||
)
|
||||
|
||||
# Generate optimizer
|
||||
optimizer_ops = parallelizer._generate_optimizer(
|
||||
serial_main_prog,
|
||||
serial_startup_prog,
|
||||
serial_optimizer,
|
||||
params_grads,
|
||||
)
|
||||
cost_estimator = CostEstimator(serial_main_prog, engine._cluster)
|
||||
|
||||
# Estimate global_cost and max memory
|
||||
global_cost = cost_estimator.estimate(dist_context)
|
||||
max_memory = cost_estimator._estimate_max_memory_by_dist_op(dist_context)
|
||||
|
||||
# Print the cost
|
||||
cost_estimator.pretty_print_cost()
|
||||
|
||||
return global_cost, max_memory
|
||||
@@ -0,0 +1,320 @@
|
||||
# Copyright (c) 2023 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 logging
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.base import core
|
||||
from paddle.base.data_feeder import convert_dtype
|
||||
from paddle.base.executor import (
|
||||
_as_lodtensor,
|
||||
_StandaloneExecutor,
|
||||
check_feed_shape_type,
|
||||
)
|
||||
from paddle.base.framework import Operator, Program
|
||||
from paddle.distributed.auto_parallel.static.utils import get_logger, is_comm_op
|
||||
|
||||
|
||||
def check_if_op_supports_runtime_profiling(op):
|
||||
return not is_comm_op(op)
|
||||
|
||||
|
||||
def _measure_program_real_op_cost_multipass(program, place, run_iters, verbose):
|
||||
'''
|
||||
Run op profiling for a single pass. Internal function, do not call this directly.
|
||||
'''
|
||||
|
||||
# clone the program to avoid accidental change made to the vanilla program.
|
||||
cloned_program = program.clone()
|
||||
cloned_main_block = cloned_program.global_block()
|
||||
|
||||
# We will run the executor in a newly created scope, so that our
|
||||
# executor will not pollute the global scope when running. Since
|
||||
# we created a brand new scope, we need to manually create input
|
||||
# tensors and network parameters and feed fake data into them.
|
||||
scope = core.Scope()
|
||||
|
||||
logger = get_logger(log_level=logging.INFO)
|
||||
|
||||
def _analyze_graph_and_collect_all_vars_with_zero_in_degree():
|
||||
var_in_degree = {}
|
||||
|
||||
def _collect_op_input_var_names(op: Operator):
|
||||
input_var_names = []
|
||||
for input_name in op.input_names:
|
||||
input_var_names += op.input(input_name)
|
||||
return input_var_names
|
||||
|
||||
def _collect_op_output_var_names(op: Operator):
|
||||
output_var_names = []
|
||||
for output_name in op.output_names:
|
||||
output_var_names += op.output(output_name)
|
||||
return output_var_names
|
||||
|
||||
def _record_op_output_vars_in_degree(in_var_names, out_var_names):
|
||||
for out_var_name in out_var_names:
|
||||
if out_var_name in in_var_names:
|
||||
# NOTE (liuchenghao): if an op's input var is its output var,
|
||||
# this means this var forms an in-place connection to itself,
|
||||
# in this situation we need to ignore this variable, this way
|
||||
# we can ensure that vars with zero in-degree are dangling vars
|
||||
# and they should be created manually before program executes.
|
||||
continue
|
||||
var_in_degree[out_var_name] += 1
|
||||
|
||||
def _filter_vars_with_zero_in_degree_and_ignore_feed_fetch_vars():
|
||||
filtered_vars = []
|
||||
for var_name in var_in_degree:
|
||||
if var_name in ['feed', 'fetch']:
|
||||
continue
|
||||
if var_in_degree[var_name] == 0:
|
||||
filtered_vars.append(var_name)
|
||||
return filtered_vars
|
||||
|
||||
for op in cloned_main_block.ops:
|
||||
op: Operator
|
||||
if is_comm_op(op):
|
||||
# ignore communication op from graph, because sometimes we want to profile a sub-graph
|
||||
# and these dangling operators will not work (no graph to communicate to/from)
|
||||
continue
|
||||
input_var_names, output_var_names = (
|
||||
_collect_op_input_var_names(op),
|
||||
_collect_op_output_var_names(op),
|
||||
)
|
||||
for var_name in input_var_names + output_var_names:
|
||||
if var_name not in var_in_degree:
|
||||
var_in_degree[var_name] = 0
|
||||
_record_op_output_vars_in_degree(input_var_names, output_var_names)
|
||||
return _filter_vars_with_zero_in_degree_and_ignore_feed_fetch_vars()
|
||||
|
||||
def _alloc_and_fill_var(var_name):
|
||||
supported_var_dtypes = [
|
||||
"paddle.float16",
|
||||
"paddle.float32",
|
||||
"paddle.float64",
|
||||
"paddle.int8",
|
||||
"paddle.int16",
|
||||
"paddle.int32",
|
||||
"paddle.int64",
|
||||
"paddle.bool",
|
||||
]
|
||||
var = cloned_main_block.var(var_name)
|
||||
var_shape = var.shape
|
||||
var_dtype = var.dtype
|
||||
assert str(var_dtype) in supported_var_dtypes, (
|
||||
'Found unsupported variable dtype: "{}", current supported '
|
||||
'dtype(s) is/are: [{}]. '.format(
|
||||
str(var_dtype), ", ".join(supported_var_dtypes)
|
||||
)
|
||||
)
|
||||
(
|
||||
logger.info(
|
||||
f'[+] var: "{var_name}", shape={var_shape}, dtype="{var_dtype}".\n'
|
||||
)
|
||||
if verbose
|
||||
else None
|
||||
)
|
||||
np_dtype = (
|
||||
convert_dtype(var_dtype)
|
||||
if isinstance(var_dtype, core.VarDesc.VarType)
|
||||
else var_dtype
|
||||
)
|
||||
if str(var_dtype).find('int') != -1:
|
||||
# target variable's type is int* (uint*, int*), it is highly possible that
|
||||
# the target variable contains indices (such as lookup_table op's input var)
|
||||
# for safety we need to fill it with all one instead of random numbers
|
||||
# NOTE (liuchenghao): filling with zero will generate "division by zero" error
|
||||
# in mod ops, so filling with one seems to be the simplest way to make it work,
|
||||
# although it is possible that for array with only one element, index "1" is
|
||||
# invalid, that situation is very rare and we don't need to care about it now.
|
||||
new_tensor = np.array(np.ones(var_shape)).astype(np_dtype)
|
||||
else:
|
||||
# target variable's type is float*, we treat it as an ordinary tensor, fill it
|
||||
# with random gaussian numbers
|
||||
new_tensor = np.array(np.random.randn(*var_shape)).astype(np_dtype)
|
||||
new_tensor = _as_lodtensor(new_tensor, place, var_dtype)
|
||||
check_feed_shape_type(var, new_tensor)
|
||||
core.set_variable(scope, new_tensor, var_name)
|
||||
return new_tensor
|
||||
|
||||
def _configure_feed_ops_and_return_feed_names():
|
||||
"""
|
||||
configure feed op,
|
||||
1. alloc feed op output var storage
|
||||
2. fill feed op's input var
|
||||
return feed var names
|
||||
"""
|
||||
|
||||
feed_names = []
|
||||
has_feed_op = False
|
||||
for op in cloned_main_block.ops:
|
||||
if op.type == "feed":
|
||||
has_feed_op = True
|
||||
out_var_name = op.desc.output('Out')[0]
|
||||
in_var_name = op.desc.input('X')[0] # this is usually "feed"
|
||||
input_index = op.desc.attr('col')
|
||||
new_tensor = _alloc_and_fill_var(out_var_name)
|
||||
core.set_feed_variable(
|
||||
scope, new_tensor, in_var_name, input_index
|
||||
)
|
||||
feed_names.append(out_var_name)
|
||||
if not has_feed_op:
|
||||
(
|
||||
logger.info("WARNING: program does not have any feed op.\n")
|
||||
if verbose
|
||||
else None
|
||||
)
|
||||
return feed_names
|
||||
|
||||
for var_name in _analyze_graph_and_collect_all_vars_with_zero_in_degree():
|
||||
_alloc_and_fill_var(var_name)
|
||||
feed_names = _configure_feed_ops_and_return_feed_names()
|
||||
|
||||
# build a simple plan from program and run profiling
|
||||
plan = core.Plan([core.Job("default")], {"default": cloned_program.desc})
|
||||
exe = _StandaloneExecutor(place, plan, scope)
|
||||
|
||||
num_ops = len(cloned_main_block.ops)
|
||||
prof_results = [[None for _ in range(run_iters)] for _ in range(num_ops)]
|
||||
|
||||
for iter_id in range(run_iters):
|
||||
# for each iteration, run profiling and retrieve modified version of program desc
|
||||
program_desc = exe.run_profile(feed_names)
|
||||
|
||||
# rebuild program object from the new program desc
|
||||
temp_program = cloned_program.clone()
|
||||
temp_program._rebuild_from_desc(program_desc)
|
||||
temp_main_block = temp_program.global_block()
|
||||
|
||||
# collect profiling result
|
||||
for op_id, temp_op in zip(
|
||||
range(len(temp_main_block.ops)), temp_main_block.ops
|
||||
):
|
||||
run_time_us = temp_op.dist_attr.run_time_us
|
||||
prof_results[op_id][iter_id] = (
|
||||
run_time_us
|
||||
if check_if_op_supports_runtime_profiling(temp_op)
|
||||
and run_time_us >= 0.0
|
||||
else None
|
||||
)
|
||||
return prof_results
|
||||
|
||||
|
||||
def measure_program_real_op_cost(
|
||||
program: paddle.static.Program,
|
||||
run_iters: int = 8,
|
||||
place=paddle.base.framework._current_expected_place(),
|
||||
verbose_level: int = 0,
|
||||
):
|
||||
'''
|
||||
Description
|
||||
-----------
|
||||
Measuring real op run time (us) with respect to the given "program" and "place".
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
@param program: paddle.static.Program
|
||||
The program object waiting to be executed.
|
||||
@param run_iters: int
|
||||
Specify how many iterations will be run during profiling. Larger value tends
|
||||
to give more accurate profiling result but requires more time.
|
||||
@param place: paddle.CPUPlace | paddle.CUDAPlace
|
||||
Where the program is going to be executed.
|
||||
@param verbose_level: int
|
||||
Set up verbose level during profiling. Can be set to one of the following:
|
||||
0 = turn off, don't output anything,
|
||||
1 = output profiling messages only,
|
||||
2 = output profiling and debug messages.
|
||||
|
||||
Returns
|
||||
-----------
|
||||
Nothing to return. This API will write op run time directly into program object.
|
||||
For example, to retrieve the run time for the first op in program, use:
|
||||
>>> program.global_block().ops[0].dist_attr.run_time_us
|
||||
|
||||
Note
|
||||
-----------
|
||||
Not all ops support runtime profiling. Currently communication ops do not support
|
||||
runtime profiling feature since their execution times rely on other ops. To check
|
||||
if an op supports runtime profiling, use:
|
||||
>>> check_if_op_supports_runtime_profiling(op)
|
||||
where "op" is an instance of "paddle.base.framework.Operator".
|
||||
|
||||
Example
|
||||
-----------
|
||||
* Profiling a simple program from scratch:
|
||||
>>> from paddle.distributed.auto_parallel.static.utils import (
|
||||
... measure_program_real_op_cost,
|
||||
... )
|
||||
>>> program = ... # build your own program object here.
|
||||
>>> measure_program_real_op_cost(
|
||||
>>> program, verbose_level=1
|
||||
>>> )
|
||||
* Profiling a program which is already embedded into an Executor or some other class instance:
|
||||
>>> import paddle
|
||||
>>> from paddle.distributed.auto_parallel.static.utils import (
|
||||
... measure_program_real_op_cost,
|
||||
... )
|
||||
>>> place: str = paddle.device.get_device() # here we assume place = "cuda:x"
|
||||
>>> place = paddle.CUDAPlace(int(place.split(':')[1]))
|
||||
>>> # here "program" is an inner object that has already been built before
|
||||
>>> measure_program_real_op_cost(program, verbose_level=1)
|
||||
'''
|
||||
|
||||
assert isinstance(program, Program), (
|
||||
f'"program" should be a instance of "paddle.base.framework.Program" but got type "{type(program).__name__}".'
|
||||
)
|
||||
supported_places = [
|
||||
paddle.CUDAPlace,
|
||||
]
|
||||
assert any(
|
||||
isinstance(place, supported_place)
|
||||
for supported_place in supported_places
|
||||
), (
|
||||
f'Current place ({place}) does not support runtime profiling. "place" should be one of the following: {supported_places}.'
|
||||
)
|
||||
assert isinstance(run_iters, int) and run_iters >= 1, (
|
||||
'Invalid parameter run_iters set. run_iters should be an integer >= 1.'
|
||||
)
|
||||
if run_iters == 1:
|
||||
warnings.warn(
|
||||
'run_iters was set to 1, profiling results might be inaccurate due to outliers.'
|
||||
)
|
||||
|
||||
logger = get_logger(log_level=logging.INFO)
|
||||
|
||||
# run profiling multiple times and record op run time of each run
|
||||
prof_results = _measure_program_real_op_cost_multipass(
|
||||
program, place, run_iters, verbose=(verbose_level >= 2)
|
||||
)
|
||||
op_num = len(prof_results)
|
||||
for op_id, op in zip(range(op_num), program.global_block().ops):
|
||||
op_runtime_us_final = None
|
||||
if prof_results[op_id][0] is not None:
|
||||
op_runtime_us_final = np.median(prof_results[op_id])
|
||||
if (
|
||||
op_runtime_us_final is not None
|
||||
and check_if_op_supports_runtime_profiling(op)
|
||||
):
|
||||
op.dist_attr.run_time_us = op_runtime_us_final
|
||||
(
|
||||
logger.info(
|
||||
f"{op_id!s:>4} {op.type!s:>32} {op_runtime_us_final:.1f} us"
|
||||
)
|
||||
if verbose_level >= 1
|
||||
else None
|
||||
)
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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
|
||||
|
||||
from functools import reduce
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.auto_parallel.static.dist_tensor import (
|
||||
DistributedTensor,
|
||||
)
|
||||
from paddle.static import Variable
|
||||
|
||||
from .base_cost import Cost
|
||||
|
||||
|
||||
class TensorCost:
|
||||
def __init__(self, tensor=None, dist_tensor=None, shape=None, dtype=None):
|
||||
self._check_args(tensor, dist_tensor, shape, dtype)
|
||||
self._tensor = tensor
|
||||
self._dist_tensor = dist_tensor
|
||||
self._shape = shape
|
||||
self._dtype = dtype
|
||||
self._cost = self.calc_cost()
|
||||
|
||||
@property
|
||||
def tensor(self):
|
||||
return self._tensor
|
||||
|
||||
@property
|
||||
def dist_tensor(self):
|
||||
return self._dist_tensor
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self._shape
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self._dtype
|
||||
|
||||
def _check_args(self, tensor, dist_tensor, shape, dtype):
|
||||
if tensor is not None:
|
||||
assert shape is None and dist_tensor is None and dtype is None
|
||||
|
||||
if not isinstance(tensor, Variable):
|
||||
raise TypeError(
|
||||
f"Please check tensor type is Variable, but got {type(tensor)}"
|
||||
)
|
||||
|
||||
elif dist_tensor is not None:
|
||||
assert tensor is None and shape is None
|
||||
if not isinstance(dist_tensor, DistributedTensor):
|
||||
raise TypeError(
|
||||
f"Please check dist_tensor type is DistributedTensor, but got {type(dist_tensor)}"
|
||||
)
|
||||
|
||||
elif shape is not None:
|
||||
assert tensor is None and dist_tensor is None and dtype is not None
|
||||
if not isinstance(shape, (list, set)):
|
||||
raise TypeError(
|
||||
f"Please check shape type is list or set, but got {type(shape)}"
|
||||
)
|
||||
|
||||
elif dtype is not None:
|
||||
assert tensor is None and dist_tensor is None and shape is not None
|
||||
|
||||
@property
|
||||
def cost(self):
|
||||
return self._cost
|
||||
|
||||
def calc_cost(self):
|
||||
dtype = None
|
||||
shape = None
|
||||
|
||||
if self.dist_tensor:
|
||||
shape = self.dist_tensor.local_sizes()
|
||||
dtype = self.dist_tensor.serial_tensor.dtype
|
||||
elif self.tensor:
|
||||
shape = self.tensor.shape
|
||||
dtype = self.tensor.dtype
|
||||
elif self.shape and self.dtype:
|
||||
shape = self.shape
|
||||
dtype = self.dtype
|
||||
|
||||
total_count = reduce(lambda x, y: x * y, shape, 1)
|
||||
|
||||
if dtype == paddle.float32 or dtype == paddle.int32:
|
||||
dtype_factor = 4
|
||||
elif dtype == paddle.int64:
|
||||
dtype_factor = 8
|
||||
elif dtype == paddle.uint8:
|
||||
dtype_factor = 1
|
||||
else:
|
||||
dtype_factor = 2
|
||||
|
||||
memory = total_count * dtype_factor
|
||||
assert memory >= 0
|
||||
cost = Cost(memory=memory)
|
||||
|
||||
return cost
|
||||
@@ -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
|
||||
@@ -0,0 +1,19 @@
|
||||
# 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
|
||||
|
||||
from paddle.base.core import ( # noqa: F401
|
||||
DistTensorSpec,
|
||||
OperatorDistAttr,
|
||||
TensorDistAttr,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,62 @@
|
||||
# Copyright (c) 2023 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
|
||||
|
||||
from paddle.static import InputSpec
|
||||
|
||||
from ..placement_type import get_shard_spec
|
||||
from .utils import convert_to_dims_mapping
|
||||
|
||||
|
||||
class DistributedInputSpec(InputSpec):
|
||||
def __init__(
|
||||
self,
|
||||
shape,
|
||||
dtype='float32',
|
||||
name=None,
|
||||
stop_gradient=False,
|
||||
mesh=None,
|
||||
placements=None,
|
||||
local_shape=None,
|
||||
):
|
||||
super().__init__(shape, dtype, name, stop_gradient)
|
||||
self.mesh = copy.deepcopy(mesh)
|
||||
sharding_specs = get_shard_spec(mesh, placements, len(self.shape))
|
||||
self.dims_mapping = convert_to_dims_mapping(sharding_specs, mesh)
|
||||
self.local_shape = local_shape
|
||||
|
||||
@classmethod
|
||||
def from_dtensor(cls, dtensor, name=None, shape=None):
|
||||
"""
|
||||
Generates a DistributedInputSpec based on dist tensor.
|
||||
|
||||
Args:
|
||||
dtensor: the dist tensor.
|
||||
|
||||
Returns:
|
||||
A DistributedInputSpec instance generated from dtensor.
|
||||
"""
|
||||
return cls(
|
||||
shape=dtensor.shape if shape is None else shape,
|
||||
dtype=dtensor.dtype,
|
||||
name=name,
|
||||
stop_gradient=dtensor.stop_gradient,
|
||||
mesh=dtensor.process_mesh,
|
||||
placements=dtensor.placements,
|
||||
local_shape=dtensor._local_value().shape,
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return f"{super().__repr__()}, mesh:{self.mesh}, placements:{self.dims_mapping}"
|
||||
@@ -0,0 +1,290 @@
|
||||
# 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 abc
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.io import BatchSampler, IterableDataset
|
||||
from paddle.io.dataloader.batch_sampler import (
|
||||
DistributedBatchSampler,
|
||||
_InfiniteIterableSampler,
|
||||
)
|
||||
from paddle.io.dataloader.dataloader_iter import (
|
||||
_DatasetKind,
|
||||
default_collate_fn,
|
||||
default_convert_fn,
|
||||
)
|
||||
|
||||
|
||||
class DistributedDataLoaderBase(metaclass=abc.ABCMeta):
|
||||
@abc.abstractmethod
|
||||
def __iter__(self):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class DistributedDataLoaderFromGenerator(DistributedDataLoaderBase):
|
||||
def __init__(
|
||||
self,
|
||||
dataset,
|
||||
feed_list=None,
|
||||
capacity=None,
|
||||
use_double_buffer=True,
|
||||
iterable=True,
|
||||
return_list=False,
|
||||
use_multiprocess=False,
|
||||
drop_last=True,
|
||||
places=None,
|
||||
batch_size=1,
|
||||
epochs=1,
|
||||
steps_per_epoch=None,
|
||||
collate_fn=None,
|
||||
split_data=True,
|
||||
data_parallel_world_size=[],
|
||||
data_parallel_rank=[],
|
||||
acc_steps=1,
|
||||
):
|
||||
self.dataset = dataset
|
||||
self.feed_list = feed_list
|
||||
self.capacity = capacity
|
||||
self.use_double_buffer = use_double_buffer
|
||||
self.iterable = iterable
|
||||
self.return_list = return_list
|
||||
self.use_multiprocess = use_multiprocess
|
||||
self.drop_last = drop_last
|
||||
self.places = places
|
||||
self.batch_size = batch_size
|
||||
self.epochs = epochs
|
||||
self.steps_per_epoch = steps_per_epoch
|
||||
self.collate_fn = collate_fn
|
||||
self.split_data = split_data
|
||||
assert len(data_parallel_world_size) == len(feed_list)
|
||||
assert len(data_parallel_rank) == len(feed_list)
|
||||
self.dp_world_sizes = data_parallel_world_size
|
||||
self.dp_ranks = data_parallel_rank
|
||||
self.acc_steps = acc_steps
|
||||
|
||||
if isinstance(dataset, IterableDataset):
|
||||
self.dataset_kind = _DatasetKind.ITER
|
||||
else:
|
||||
self.dataset_kind = _DatasetKind.MAP
|
||||
|
||||
if self.batch_size is None:
|
||||
self.batch_sampler = None
|
||||
else:
|
||||
if isinstance(dataset, IterableDataset):
|
||||
self.batch_sampler = _InfiniteIterableSampler(
|
||||
dataset, batch_size
|
||||
)
|
||||
else:
|
||||
self.batch_sampler = BatchSampler(
|
||||
dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
drop_last=drop_last,
|
||||
)
|
||||
|
||||
self.auto_collate_batch = self.batch_sampler is not None
|
||||
self.sampler_iter = iter(self.index_sampler)
|
||||
|
||||
if self.auto_collate_batch:
|
||||
self.collate_fn = collate_fn or default_collate_fn
|
||||
else:
|
||||
self.collate_fn = collate_fn or default_convert_fn
|
||||
|
||||
self.dataset_fetcher = _DatasetKind.create_fetcher(
|
||||
self.dataset_kind,
|
||||
self.dataset,
|
||||
self.auto_collate_batch,
|
||||
self.collate_fn,
|
||||
self.drop_last,
|
||||
)
|
||||
|
||||
self._steps = self._infer_steps()
|
||||
self._inner_dataloader = self._create_inner_dataloader()
|
||||
|
||||
def __iter__(self):
|
||||
self._cur_step = 0
|
||||
self._inner_dataloader.start()
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if not self._steps:
|
||||
self._cur_step += 1
|
||||
return None
|
||||
elif self._cur_step < self._steps:
|
||||
self._cur_step += 1
|
||||
return None
|
||||
else:
|
||||
self._inner_dataloader.reset()
|
||||
self.sampler_iter = iter(self.index_sampler)
|
||||
raise StopIteration
|
||||
|
||||
def _infer_steps(self):
|
||||
if isinstance(self.steps_per_epoch, int) and self.steps_per_epoch > 0:
|
||||
return self.steps_per_epoch
|
||||
try:
|
||||
if isinstance(self.dataset, IterableDataset):
|
||||
steps_per_epoch = None
|
||||
elif self.batch_size is None:
|
||||
steps_per_epoch = len(self.dataset) // self.acc_steps
|
||||
else:
|
||||
steps_per_epoch = (
|
||||
len(self.dataset) // self.batch_size // self.acc_steps
|
||||
)
|
||||
except:
|
||||
raise ValueError(
|
||||
"Please set `steps_per_epoch` or implement `__len__` method in dataset class."
|
||||
)
|
||||
return steps_per_epoch
|
||||
|
||||
@property
|
||||
def index_sampler(self):
|
||||
if self.auto_collate_batch:
|
||||
return self.batch_sampler
|
||||
else:
|
||||
if self.dataset_kind == _DatasetKind.MAP:
|
||||
return list(range(len(self.dataset)))
|
||||
else:
|
||||
return _InfiniteIterableSampler(self.dataset, 1)
|
||||
|
||||
def _create_inner_dataloader(self):
|
||||
def data_generator():
|
||||
while True:
|
||||
try:
|
||||
indices = next(self.sampler_iter)
|
||||
batch = self.dataset_fetcher.fetch(indices)
|
||||
if batch is None:
|
||||
break
|
||||
except StopIteration:
|
||||
self.dataset_fetcher = _DatasetKind.create_fetcher(
|
||||
self.dataset_kind,
|
||||
self.dataset,
|
||||
self.auto_collate_batch,
|
||||
self.collate_fn,
|
||||
self.drop_last,
|
||||
)
|
||||
break
|
||||
|
||||
partial_data = []
|
||||
for i, d in enumerate(batch):
|
||||
array = np.array(d)
|
||||
if not self.split_data:
|
||||
partial_data.append(array)
|
||||
continue
|
||||
|
||||
batch_size = array.shape[0]
|
||||
assert batch_size % self.dp_world_sizes[i] == 0, (
|
||||
f"batch_size [{batch_size}] is not divisible by dp_world_size [{self.dp_world_sizes[i]}]"
|
||||
)
|
||||
partial_data.append(
|
||||
np.split(array, self.dp_world_sizes[i])[
|
||||
self.dp_ranks[i]
|
||||
]
|
||||
)
|
||||
|
||||
yield partial_data
|
||||
|
||||
dataloader = paddle.base.io.DataLoader.from_generator(
|
||||
feed_list=self.feed_list,
|
||||
capacity=self.capacity,
|
||||
use_double_buffer=self.use_double_buffer,
|
||||
# iterable=self.iterable,
|
||||
iterable=False,
|
||||
return_list=self.return_list,
|
||||
use_multiprocess=self.use_multiprocess,
|
||||
drop_last=self.drop_last,
|
||||
)
|
||||
dataloader.set_batch_generator(data_generator, self.places)
|
||||
|
||||
return dataloader
|
||||
|
||||
|
||||
class DistributedDataLoader(DistributedDataLoaderBase):
|
||||
def __init__(
|
||||
self,
|
||||
dataset,
|
||||
feed_list=None,
|
||||
places=None,
|
||||
return_list=True,
|
||||
batch_size=1,
|
||||
shuffle=False,
|
||||
drop_last=False,
|
||||
collate_fn=None,
|
||||
num_workers=0,
|
||||
use_buffer_reader=True,
|
||||
use_shared_memory=True,
|
||||
timeout=0,
|
||||
worker_init_fn=None,
|
||||
epochs=1,
|
||||
steps_per_epoch=None,
|
||||
split_data=True,
|
||||
data_parallel_world_size=[],
|
||||
data_parallel_rank=[],
|
||||
):
|
||||
self.dataset = dataset
|
||||
self.feed_list = feed_list
|
||||
self.return_list = return_list
|
||||
self.places = places
|
||||
self.batch_size = batch_size
|
||||
self.shuffle = shuffle
|
||||
self.drop_last = drop_last
|
||||
self.collate_fn = collate_fn
|
||||
self.num_workers = num_workers
|
||||
self.use_buffer_reader = use_buffer_reader
|
||||
self.use_shared_memory = use_shared_memory
|
||||
self.timeout = timeout
|
||||
self.worker_init_fn = worker_init_fn
|
||||
self.epochs = epochs
|
||||
self.steps_per_epoch = steps_per_epoch
|
||||
self.dp_world_sizes = data_parallel_world_size
|
||||
self.dp_ranks = data_parallel_rank
|
||||
self.split_data = split_data
|
||||
|
||||
if self.batch_size is None:
|
||||
self.batch_sampler = None
|
||||
else:
|
||||
self.batch_sampler = DistributedBatchSampler(
|
||||
dataset=self.dataset,
|
||||
batch_size=self.batch_size,
|
||||
num_replicas=self.dp_world_sizes[0],
|
||||
rank=self.dp_ranks[0],
|
||||
shuffle=self.shuffle,
|
||||
drop_last=self.drop_last,
|
||||
)
|
||||
|
||||
self._dataloader = paddle.io.DataLoader(
|
||||
self.dataset,
|
||||
feed_list=self.feed_list,
|
||||
places=self.places,
|
||||
return_list=self.return_list,
|
||||
batch_sampler=self.batch_sampler,
|
||||
batch_size=1 if self.batch_sampler else self.batch_size,
|
||||
collate_fn=self.collate_fn,
|
||||
num_workers=self.num_workers,
|
||||
use_buffer_reader=self.use_buffer_reader,
|
||||
use_shared_memory=self.use_shared_memory,
|
||||
timeout=self.timeout,
|
||||
worker_init_fn=self.worker_init_fn,
|
||||
)
|
||||
|
||||
def __len__(self):
|
||||
return len(self._dataloader)
|
||||
|
||||
def __iter__(self):
|
||||
return self._dataloader.__iter__()
|
||||
|
||||
def __call__(self):
|
||||
return self._dataloader.__iter__()
|
||||
@@ -0,0 +1,323 @@
|
||||
# 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 paddle
|
||||
from paddle.static import Variable
|
||||
|
||||
from .dist_attribute import OperatorDistAttr
|
||||
from .utils import (
|
||||
__no_shape_var_type__,
|
||||
convert_to_shard_spec,
|
||||
verify_shard_spec,
|
||||
)
|
||||
|
||||
|
||||
class DistributedOperator:
|
||||
def __init__(self, serial_op, dist_attr=None):
|
||||
self._serial_op = serial_op
|
||||
if dist_attr is not None and isinstance(dist_attr, OperatorDistAttr):
|
||||
# TODO: remove this deepcopy after we fix the issue
|
||||
self._dist_attr = copy.deepcopy(dist_attr)
|
||||
# self._dist_attr = dist_attr
|
||||
# TODO: Do we really need to write back to serial op？
|
||||
self._serial_op.dist_attr = dist_attr
|
||||
else:
|
||||
assert dist_attr is None, f"{dist_attr}"
|
||||
# Use the dist attr of serial_op to do the initialization
|
||||
self._dist_attr = self._serial_op.dist_attr
|
||||
self._serial_inputs = {}
|
||||
self._serial_outputs = {}
|
||||
|
||||
@property
|
||||
def serial_op(self):
|
||||
return self._serial_op
|
||||
|
||||
@property
|
||||
def dist_attr(self):
|
||||
return self._dist_attr
|
||||
|
||||
@dist_attr.setter
|
||||
def dist_attr(self, dist_attr):
|
||||
self._dist_attr = dist_attr
|
||||
# TODO: Do we really need to write back to serial op？
|
||||
self._serial_op.dist_attr = dist_attr
|
||||
|
||||
def get_serial_input(self, name):
|
||||
if self._serial_op.type == "create_py_reader":
|
||||
tensor = None
|
||||
elif self._serial_op.block._find_var_recursive(name) is not None:
|
||||
tensor = self._serial_op.block._var_recursive(name)
|
||||
else:
|
||||
tensor = None
|
||||
return tensor
|
||||
|
||||
def get_serial_output(self, name):
|
||||
tensor = self._serial_op.block._var_recursive(name)
|
||||
return tensor
|
||||
|
||||
def validate_dist_attr(self):
|
||||
if "read" in self.serial_op.type or "while" == self.serial_op.type:
|
||||
return True
|
||||
for name in self.serial_op.input_arg_names:
|
||||
input_dist_attr = self.dist_attr.get_input_dist_attr(name)
|
||||
dims_mapping = input_dist_attr.dims_mapping
|
||||
if self.get_serial_input(name).type in __no_shape_var_type__:
|
||||
shape = []
|
||||
else:
|
||||
shape = self.get_serial_input(name).shape
|
||||
if len(shape) != len(dims_mapping):
|
||||
return False
|
||||
for i in range(len(dims_mapping)):
|
||||
if dims_mapping[i] < -1 or dims_mapping[i] >= len(
|
||||
self.dist_attr.process_mesh.shape
|
||||
):
|
||||
return False
|
||||
for i in range(len(self.dist_attr.process_mesh.shape)):
|
||||
if dims_mapping.count(i) > 1:
|
||||
return False
|
||||
if self.dist_attr.process_mesh != input_dist_attr.process_mesh:
|
||||
return False
|
||||
|
||||
for name in self.serial_op.output_arg_names:
|
||||
output_dist_attr = self.dist_attr.get_output_dist_attr(name)
|
||||
dims_mapping = output_dist_attr.dims_mapping
|
||||
if self.get_serial_output(name).type in __no_shape_var_type__:
|
||||
shape = []
|
||||
else:
|
||||
shape = self.get_serial_output(name).shape
|
||||
if len(shape) != len(dims_mapping):
|
||||
return False
|
||||
for i in range(len(dims_mapping)):
|
||||
if dims_mapping[i] < -1 or dims_mapping[i] >= len(
|
||||
self.dist_attr.process_mesh.shape
|
||||
):
|
||||
return False
|
||||
for i in range(len(self.dist_attr.process_mesh.shape)):
|
||||
if dims_mapping.count(i) > 1:
|
||||
return False
|
||||
if self.dist_attr.process_mesh != output_dist_attr.process_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def __str__(self):
|
||||
str = f"{{op type: {self.serial_op.desc.type()}, op id: {self.serial_op.desc.id()}, op original_id: {self.serial_op.desc.original_id()}"
|
||||
|
||||
# str += ", {}".format(self.dist_attr)
|
||||
# return str
|
||||
|
||||
if self.dist_attr.is_annotated("process_mesh"):
|
||||
annotated_str = "annotated"
|
||||
else:
|
||||
annotated_str = "non-annotated"
|
||||
str += (
|
||||
f", process_mesh ({annotated_str}): {self.dist_attr.process_mesh}"
|
||||
)
|
||||
|
||||
str += f" , execution_stream: {self.dist_attr.execution_stream}"
|
||||
|
||||
for arg_name in self.serial_op.desc.input_arg_names():
|
||||
try:
|
||||
dims_mapping = self.dist_attr.get_input_dims_mapping(arg_name)
|
||||
except IndexError:
|
||||
raise IndexError(
|
||||
f"There is not input var '{arg_name}''s dist_attr in current op '{self.serial_op.desc.type()}'"
|
||||
)
|
||||
if self.dist_attr.is_annotated_input_dims_mapping(arg_name):
|
||||
annotated_str = "annotated"
|
||||
else:
|
||||
annotated_str = "non-annotated"
|
||||
if self.get_serial_input(arg_name) is not None:
|
||||
if self.get_serial_input(arg_name).is_parameter:
|
||||
is_parameter_str = "parameter"
|
||||
else:
|
||||
is_parameter_str = "non-parameter"
|
||||
else:
|
||||
is_parameter_str = "non-parameter"
|
||||
|
||||
# partial
|
||||
input_dist_attr = self.dist_attr.get_input_dist_attr(arg_name)
|
||||
partial_dims = sorted(input_dist_attr._partial_dims())
|
||||
|
||||
str += f"; {arg_name}'s dims_mapping (input, {annotated_str}, {is_parameter_str}): {dims_mapping}, partial on dims: {partial_dims}"
|
||||
|
||||
for arg_name in self.serial_op.desc.output_arg_names():
|
||||
try:
|
||||
dims_mapping = self.dist_attr.get_output_dims_mapping(arg_name)
|
||||
except IndexError:
|
||||
raise IndexError(
|
||||
f"There is not output var '{arg_name}''s dist_attr in current op '{self.serial_op.desc.type()}'"
|
||||
)
|
||||
if self.dist_attr.is_annotated_output_dims_mapping(arg_name):
|
||||
annotated_str = "annotated"
|
||||
else:
|
||||
annotated_str = "non-annotated"
|
||||
if self.get_serial_output(arg_name) is not None:
|
||||
if self.get_serial_output(arg_name).is_parameter:
|
||||
is_parameter_str = "parameter"
|
||||
else:
|
||||
is_parameter_str = "non-parameter"
|
||||
else:
|
||||
is_parameter_str = "non-parameter"
|
||||
|
||||
# partial
|
||||
output_dist_attr = self.dist_attr.get_output_dist_attr(arg_name)
|
||||
partial_dims = sorted(output_dist_attr._partial_dims())
|
||||
|
||||
str += f"; {arg_name}'s dims_mapping (output, {annotated_str}, {is_parameter_str}): {dims_mapping}, partial on dims: {partial_dims}"
|
||||
|
||||
str += f", dist_impl idx: {self.dist_attr.impl_idx} , dist_impl type: {self.dist_attr.impl_type}, chunk_id: {self.dist_attr.chunk_id} }}"
|
||||
|
||||
return str
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
cls = self.__class__
|
||||
result = cls.__new__(cls)
|
||||
memo[id(self)] = result
|
||||
for k, v in self.__dict__.items():
|
||||
if (
|
||||
k == "_serial_op"
|
||||
or k == "_serial_inputs"
|
||||
or k == "_serial_outputs"
|
||||
):
|
||||
setattr(result, k, v)
|
||||
else:
|
||||
setattr(result, k, copy.deepcopy(v, memo))
|
||||
return result
|
||||
|
||||
|
||||
class DistributedOperatorHelper:
|
||||
def __init__(
|
||||
self,
|
||||
serial_op,
|
||||
process_mesh,
|
||||
in_dims_mappings,
|
||||
out_dims_mappings,
|
||||
kwargs,
|
||||
):
|
||||
self._serial_op = serial_op
|
||||
self._process_mesh = process_mesh
|
||||
self._in_dims_mappings = in_dims_mappings
|
||||
self._out_dims_mappings = out_dims_mappings
|
||||
self._chunk_id = kwargs["chunk_id"] if "chunk_id" in kwargs else 0
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
tensor_to_dims_mapping = {}
|
||||
index = 0
|
||||
if self._in_dims_mappings:
|
||||
assert len(args) + len(kwargs) == len(self._in_dims_mappings), (
|
||||
f"The length of dims_mapping {len(self._in_dims_mappings)} does not matching the length output {len(args) + len(kwargs)}."
|
||||
)
|
||||
for arg in args:
|
||||
if isinstance(arg, Variable) and self._in_dims_mappings:
|
||||
tensor_to_dims_mapping[arg.name] = self._in_dims_mappings[index]
|
||||
index += 1
|
||||
for arg in kwargs.values() and self._in_dims_mappings:
|
||||
if isinstance(arg, Variable):
|
||||
tensor_to_dims_mapping[arg.name] = self._in_dims_mappings[index]
|
||||
index += 1
|
||||
|
||||
default_prog = paddle.static.default_main_program()
|
||||
cur_block = default_prog.current_block()
|
||||
op_size = len(cur_block.ops)
|
||||
if paddle.base.dygraph.base.in_to_static_mode():
|
||||
output = paddle.jit.dy2static.convert_call_func.convert_call(
|
||||
self._serial_op
|
||||
)(*args, **kwargs)
|
||||
else:
|
||||
output = self._serial_op(*args, **kwargs)
|
||||
new_op_size = len(cur_block.ops)
|
||||
|
||||
if isinstance(output, (tuple, list)):
|
||||
new_output = list(output)
|
||||
elif isinstance(output, Variable):
|
||||
new_output = [output]
|
||||
else:
|
||||
raise ValueError("Unrecognized output.")
|
||||
|
||||
if self._out_dims_mappings:
|
||||
assert len(new_output) == len(self._out_dims_mappings), (
|
||||
f"The length of dims_mapping {len(self._out_dims_mappings)} does not matching the length output {len(new_output)}."
|
||||
)
|
||||
for i, item in enumerate(new_output):
|
||||
if isinstance(item, Variable) and self._out_dims_mappings:
|
||||
tensor_to_dims_mapping[item.name] = self._out_dims_mappings[i]
|
||||
|
||||
from .dist_context import get_default_distributed_context
|
||||
|
||||
default_dist_ctx = get_default_distributed_context()
|
||||
for idx in range(op_size, new_op_size):
|
||||
op = cur_block.ops[idx]
|
||||
dist_op = DistributedOperator(op)
|
||||
for name in dist_op.serial_op.input_arg_names:
|
||||
if name in tensor_to_dims_mapping.keys():
|
||||
tensor = dist_op.get_serial_input(name)
|
||||
tensor_dist_attr = dist_op.dist_attr.get_input_dist_attr(
|
||||
name
|
||||
)
|
||||
dims_mapping = tensor_to_dims_mapping[name]
|
||||
if tensor is None:
|
||||
tensor_shape = []
|
||||
else:
|
||||
if tensor.type in __no_shape_var_type__:
|
||||
tensor_shape = []
|
||||
else:
|
||||
tensor_shape = tensor.shape
|
||||
if dims_mapping is not None:
|
||||
dims_mapping = tensor_to_dims_mapping[name]
|
||||
shard_spec = convert_to_shard_spec(
|
||||
dims_mapping, self._process_mesh
|
||||
)
|
||||
assert verify_shard_spec(
|
||||
shard_spec, tensor_shape, self._process_mesh
|
||||
), (
|
||||
f"For tensor {name}, shard_spec {shard_spec} is invalid with tensor_shape {tensor_shape} and process_mesh {self._process_mesh}."
|
||||
)
|
||||
tensor_dist_attr.dims_mapping = dims_mapping
|
||||
tensor_dist_attr.mark_annotated("dims_mapping")
|
||||
for name in dist_op.serial_op.output_arg_names:
|
||||
if name in tensor_to_dims_mapping.keys():
|
||||
tensor = dist_op.get_serial_output(name)
|
||||
tensor_dist_attr = dist_op.dist_attr.get_output_dist_attr(
|
||||
name
|
||||
)
|
||||
dims_mapping = tensor_to_dims_mapping[name]
|
||||
if tensor is None:
|
||||
tensor_shape = []
|
||||
else:
|
||||
if tensor.type in __no_shape_var_type__:
|
||||
tensor_shape = []
|
||||
else:
|
||||
tensor_shape = tensor.shape
|
||||
if dims_mapping is not None:
|
||||
dims_mapping = tensor_to_dims_mapping[name]
|
||||
shard_spec = convert_to_shard_spec(
|
||||
dims_mapping, self._process_mesh
|
||||
)
|
||||
assert verify_shard_spec(
|
||||
shard_spec, tensor_shape, self._process_mesh
|
||||
), (
|
||||
f"For tensor {name}, shard_spec {shard_spec} is invalid with tensor_shape {tensor_shape} and process_mesh {self._process_mesh}."
|
||||
)
|
||||
tensor_dist_attr.dims_mapping = dims_mapping
|
||||
tensor_dist_attr.mark_annotated("dims_mapping")
|
||||
dist_op.dist_attr.process_mesh = self._process_mesh
|
||||
dist_op.dist_attr.chunk_id = self._chunk_id
|
||||
if self._process_mesh is not None:
|
||||
dist_op.dist_attr.mark_annotated("process_mesh")
|
||||
default_dist_ctx.add_dist_op_for_program(dist_op)
|
||||
default_dist_ctx.add_process_mesh(self._process_mesh)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,261 @@
|
||||
# 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 errno
|
||||
import logging
|
||||
import os
|
||||
import pickle
|
||||
import re
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.framework import core
|
||||
|
||||
from ...utils.log_utils import get_logger
|
||||
from .process_group import _g_process_group_map
|
||||
from .utils import get_dist_attr
|
||||
|
||||
|
||||
def check_filename(re_exp, filename):
|
||||
if re.search(re_exp, filename):
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def _process_path(path):
|
||||
filename = os.path.basename(path)
|
||||
if filename == "":
|
||||
raise ValueError(
|
||||
"path should be of 'dirname/filename' format, but received filename is empty string"
|
||||
)
|
||||
try:
|
||||
dirname = os.path.dirname(path)
|
||||
os.makedirs(dirname)
|
||||
except OSError as e:
|
||||
if e.errno != errno.EEXIST:
|
||||
raise
|
||||
return dirname, filename
|
||||
|
||||
|
||||
class DistributedSaver:
|
||||
def __init__(self):
|
||||
self._logger = get_logger(logging.INFO)
|
||||
|
||||
def save(self, path, serial_program, dist_main_program, dist_context):
|
||||
def _save_state(program, path, mode="param"):
|
||||
state = {
|
||||
k: np.array(v) for k, v in program.state_dict(mode).items()
|
||||
}
|
||||
with open(path, "wb") as f:
|
||||
pickle.dump(state, f)
|
||||
|
||||
dirname, filename = _process_path(path)
|
||||
|
||||
rank_id = paddle.distributed.get_rank()
|
||||
# save serial program when rank id is 0
|
||||
if rank_id == 0:
|
||||
self._save_rank_mapping(dirname)
|
||||
serial_model_filename = filename + "_serial.pdmodel"
|
||||
serial_model_path = os.path.join(dirname, serial_model_filename)
|
||||
with open(serial_model_path, "wb") as f:
|
||||
f.write(serial_program.desc.serialize_to_string())
|
||||
|
||||
# save distributed main program
|
||||
dist_model_filename = filename + "_dist" + str(rank_id) + ".pdmodel"
|
||||
dist_model_path = os.path.join(dirname, dist_model_filename)
|
||||
with open(dist_model_path, "wb") as f:
|
||||
f.write(dist_main_program.desc.serialize_to_string())
|
||||
|
||||
# save distributed attribute
|
||||
dist_attr_filename = filename + "_dist" + str(rank_id) + ".pdattr"
|
||||
dist_attr_path = os.path.join(dirname, dist_attr_filename)
|
||||
dist_attrs = get_dist_attr(dist_main_program, dist_context)
|
||||
with open(dist_attr_path, "wb") as f:
|
||||
pickle.dump(dist_attrs, f)
|
||||
|
||||
# save distributed params
|
||||
dist_param_filename = filename + "_dist" + str(rank_id) + ".pdparams"
|
||||
dist_param_path = os.path.join(dirname, dist_param_filename)
|
||||
_save_state(dist_main_program, dist_param_path)
|
||||
|
||||
# save distributed opt states
|
||||
dist_opt_filename = filename + "_dist" + str(rank_id) + ".pdopt"
|
||||
dist_opt_path = os.path.join(dirname, dist_opt_filename)
|
||||
_save_state(dist_main_program, dist_opt_path, "opt")
|
||||
|
||||
# TODO:save cluster.json
|
||||
|
||||
def load(self, path, load_optimizer=True):
|
||||
# TODO: if `program` is None, load `path.pdmodel`.
|
||||
def _load_file(filename, dirname, suffix="pdparams"):
|
||||
file_list = []
|
||||
for file in os.listdir(dirname):
|
||||
if check_filename(f'{filename}(.*)_dist(.*).{suffix}', file):
|
||||
file_list.append(os.path.join(dirname, file))
|
||||
file_list.sort()
|
||||
return file_list
|
||||
|
||||
def _load_state(filename, dirname, suffix="pdparams"):
|
||||
file_list = _load_file(filename, dirname, suffix)
|
||||
state_dict = {}
|
||||
for file in file_list:
|
||||
with open(file, 'rb') as f:
|
||||
from paddle.framework.restricted_unpickler import (
|
||||
safe_load_pickle,
|
||||
)
|
||||
|
||||
state_dict_info = safe_load_pickle(f, encoding='latin1')
|
||||
for name, value in state_dict_info.items():
|
||||
if name in state_dict:
|
||||
state_dict[name].append(np.array(value))
|
||||
else:
|
||||
state_dict[name] = [np.array(value)]
|
||||
self._logger.info(f"Load param file: {file_list}")
|
||||
return state_dict
|
||||
|
||||
filename = os.path.basename(path)
|
||||
if filename == "":
|
||||
raise ValueError(
|
||||
"path should be of 'dirname/filename' format, but received filename is empty string"
|
||||
)
|
||||
dirname = os.path.dirname(path)
|
||||
|
||||
# load path.pdparam and path.pdopt
|
||||
param_state_dict = _load_state(filename, dirname)
|
||||
opt_state_dict = (
|
||||
_load_state(filename, dirname, "pdopt") if load_optimizer else {}
|
||||
)
|
||||
state_dict = dict(param_state_dict, **opt_state_dict)
|
||||
|
||||
# load path.pdattr
|
||||
dist_attr_file_list = _load_file(filename, dirname, "pdattr")
|
||||
self._logger.info(
|
||||
f"Load distributed attribute file: {dist_attr_file_list}"
|
||||
)
|
||||
dist_attr = {}
|
||||
for dist_attr_file in dist_attr_file_list:
|
||||
with open(dist_attr_file, 'rb') as f:
|
||||
from paddle.framework.restricted_unpickler import (
|
||||
safe_load_pickle,
|
||||
)
|
||||
|
||||
dist_attr_info = safe_load_pickle(f, encoding='latin1')
|
||||
for name, attr in dist_attr_info.items():
|
||||
if name not in dist_attr:
|
||||
dist_attr[name] = attr
|
||||
|
||||
return state_dict, dist_attr
|
||||
|
||||
def save_inference_model(self, path, feed_vars, fetch_vars, exe, **kwargs):
|
||||
dirname, filename = _process_path(path)
|
||||
|
||||
# save distributed inference program
|
||||
rank_id = paddle.distributed.get_rank()
|
||||
if rank_id == 0:
|
||||
self._save_rank_mapping(dirname)
|
||||
op_role_key = core.op_proto_and_checker_maker.kOpRoleAttrName()
|
||||
op_role_forward = int(core.op_proto_and_checker_maker.OpRole.Forward)
|
||||
|
||||
dist_main_prog = kwargs.get('program', None)
|
||||
if not dist_main_prog:
|
||||
dist_main_prog = paddle.static.default_main_program()
|
||||
global_block = dist_main_prog.global_block()
|
||||
|
||||
ops = global_block.ops
|
||||
feed_vars_names = [x.name for x in feed_vars]
|
||||
fetch_vars_names = [x.name for x in fetch_vars]
|
||||
|
||||
last_idx = -1
|
||||
for idx, op in enumerate(ops):
|
||||
if op.attr(op_role_key) != op_role_forward:
|
||||
continue
|
||||
if op.type == "read" or op.type == "feed" or op.type == 'recv_v2':
|
||||
feed_vars_names += op.output("Out")
|
||||
if op.type == "send_v2":
|
||||
fetch_vars_names += op.input("X")
|
||||
last_idx = max(idx, last_idx)
|
||||
for out_name in op.output_arg_names:
|
||||
if out_name in fetch_vars_names:
|
||||
last_idx = max(idx, last_idx)
|
||||
|
||||
used_inputs = []
|
||||
used_outputs = []
|
||||
for idx, op in enumerate(ops):
|
||||
if idx > last_idx:
|
||||
break
|
||||
used_inputs += op.input_arg_names
|
||||
used_outputs += op.output_arg_names
|
||||
|
||||
# delete duplicated elements and keep order
|
||||
feed_vars_names = list({}.fromkeys(feed_vars_names).keys())
|
||||
used_inputs = list({}.fromkeys(used_inputs).keys())
|
||||
fetch_vars_names = list({}.fromkeys(fetch_vars_names).keys())
|
||||
used_outputs = list({}.fromkeys(used_outputs).keys())
|
||||
|
||||
dist_feed_vars_names = [
|
||||
var_name for var_name in feed_vars_names if var_name in used_inputs
|
||||
]
|
||||
dist_fetch_vars_names = [
|
||||
var_name
|
||||
for var_name in fetch_vars_names
|
||||
if var_name in used_outputs
|
||||
]
|
||||
|
||||
dist_feed_vars = list(
|
||||
reversed([global_block.vars[name] for name in dist_feed_vars_names])
|
||||
)
|
||||
dist_fetch_vars = [
|
||||
global_block.vars[name] for name in dist_fetch_vars_names
|
||||
]
|
||||
|
||||
dist_filename = filename + "_dist" + str(rank_id)
|
||||
dist_path = os.path.join(dirname, dist_filename)
|
||||
legacy_format = kwargs.get("legacy_format", False)
|
||||
paddle.static.save_inference_model(
|
||||
dist_path,
|
||||
dist_feed_vars,
|
||||
dist_fetch_vars,
|
||||
exe,
|
||||
program=dist_main_prog,
|
||||
legacy_format=legacy_format,
|
||||
)
|
||||
|
||||
def _save_rank_mapping(self, dirname):
|
||||
path = os.path.join(dirname, 'rank_mapping.csv')
|
||||
f = open(path, 'w')
|
||||
f.write('[ring_id -> ranks]\n')
|
||||
for process_group in _g_process_group_map.values():
|
||||
ring_id = process_group._group_id
|
||||
ranks = [str(rank) for rank in process_group._ranks]
|
||||
id_to_rank = str(ring_id) + "," + ",".join(ranks) + '\n'
|
||||
f.write(id_to_rank)
|
||||
id_to_rank = ""
|
||||
f.write('[rank -> ring_ids]\n')
|
||||
rank_to_id_dict = {}
|
||||
for process_group in _g_process_group_map.values():
|
||||
ring_id = process_group._group_id
|
||||
for rank in process_group._ranks:
|
||||
if rank in rank_to_id_dict:
|
||||
rank_to_id_dict[rank].append(str(ring_id))
|
||||
else:
|
||||
rank_to_id_dict[rank] = [str(ring_id)]
|
||||
rank_to_id = ""
|
||||
for item, val in rank_to_id_dict.items():
|
||||
rank_to_id += str(item) + ","
|
||||
rank_to_id += ",".join(val) + "\n"
|
||||
f.write(rank_to_id)
|
||||
rank_to_id = ""
|
||||
f.close()
|
||||
@@ -0,0 +1,412 @@
|
||||
# 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 inspect
|
||||
|
||||
import paddle
|
||||
from paddle.framework import Block
|
||||
from paddle.static import Parameter, Variable
|
||||
|
||||
from .dist_attribute import TensorDistAttr
|
||||
from .utils import __no_shape_var_type__, _linear_idx2coordinate
|
||||
|
||||
|
||||
class DistributedTensor:
|
||||
"""
|
||||
DistributedTensor represents the distribution of tensor on the process group and
|
||||
local tensors can be created by DistributedTensor.
|
||||
Only support even sharding now and uneven sharding will be supported in the future.
|
||||
Local tensor information can be obtained from the DistributedTensor instance object,
|
||||
or obtained by the static methods provided by DistributedTensor,
|
||||
including shard (i.e. the index in the serial tensor), offsets, and sizes.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _validate_sizes_and_dist_attr(
|
||||
sizes, dims_mapping, topology, processes, rank=None, shard_sizes=None
|
||||
):
|
||||
if not (
|
||||
isinstance(sizes, (list, tuple))
|
||||
and all(isinstance(x, int) and x >= 0 for x in sizes)
|
||||
):
|
||||
raise ValueError(
|
||||
f"The sizes must be list or tuple and item in sizes must be non-negative integer, but got {sizes}"
|
||||
)
|
||||
if not (
|
||||
isinstance(dims_mapping, (list, tuple))
|
||||
and all(isinstance(x, int) and x >= -1 for x in dims_mapping)
|
||||
):
|
||||
raise ValueError(
|
||||
f"The dims_mapping must be list or tuple and item in dims_mapping must >= -1, but got {dims_mapping}"
|
||||
)
|
||||
if not (
|
||||
isinstance(processes, (list, tuple))
|
||||
and all(isinstance(x, int) and x >= 0 for x in processes)
|
||||
):
|
||||
raise ValueError(
|
||||
f"The processes must be list or tuple and item in processes must be integer, but got {processes}"
|
||||
)
|
||||
if not (
|
||||
isinstance(topology, (list, tuple))
|
||||
and all(isinstance(x, int) and x > 0 for x in topology)
|
||||
):
|
||||
raise ValueError(
|
||||
f"The topology must be list or tuple and item in topology must be non-negative integer, but got {topology}"
|
||||
)
|
||||
if rank is not None and not (isinstance(rank, int) and rank >= 0):
|
||||
raise ValueError(f"The rank must >= 0, but got {rank}")
|
||||
|
||||
# # NOTE: Only support even sharding now
|
||||
# if shard_sizes is not None:
|
||||
# raise ValueError("Only support even sharding now.")
|
||||
|
||||
@staticmethod
|
||||
def get_local_sizes(
|
||||
global_sizes,
|
||||
dims_mapping,
|
||||
topology,
|
||||
processes,
|
||||
rank=None,
|
||||
shard_sizes=None,
|
||||
):
|
||||
DistributedTensor._validate_sizes_and_dist_attr(
|
||||
global_sizes, dims_mapping, topology, processes, rank, shard_sizes
|
||||
)
|
||||
|
||||
local_sizes = []
|
||||
# for even sharding, the local sizes of every rank are equal
|
||||
|
||||
for idx, item in enumerate(global_sizes):
|
||||
# This is a trick to avoid dims_mapping is []
|
||||
val = dims_mapping[idx] if idx < len(dims_mapping) else -1
|
||||
if val == -1:
|
||||
local_sizes.append(item)
|
||||
else:
|
||||
local_sizes.append(item // topology[dims_mapping[idx]])
|
||||
|
||||
return local_sizes
|
||||
|
||||
@staticmethod
|
||||
def get_local_offsets(
|
||||
global_sizes, dims_mapping, topology, processes, rank, shard_sizes=None
|
||||
):
|
||||
local_sizes = DistributedTensor.get_local_sizes(
|
||||
global_sizes, dims_mapping, topology, processes, rank, shard_sizes
|
||||
)
|
||||
local_offsets = []
|
||||
rank_relative = processes.index(rank)
|
||||
coordinate = _linear_idx2coordinate(topology, rank_relative)
|
||||
|
||||
for i in range(len(global_sizes)):
|
||||
if dims_mapping[i] == -1:
|
||||
local_offsets.append(0)
|
||||
else:
|
||||
local_offsets.append(
|
||||
coordinate[dims_mapping[i]] * local_sizes[i]
|
||||
)
|
||||
return local_offsets
|
||||
|
||||
@staticmethod
|
||||
def get_global_sizes(
|
||||
local_sizes,
|
||||
dims_mapping,
|
||||
topology,
|
||||
processes,
|
||||
rank=None,
|
||||
shard_sizes=None,
|
||||
):
|
||||
DistributedTensor._validate_sizes_and_dist_attr(
|
||||
local_sizes, dims_mapping, topology, processes, rank, shard_sizes
|
||||
)
|
||||
global_sizes = []
|
||||
for idx, item in enumerate(local_sizes):
|
||||
if dims_mapping[idx] == -1:
|
||||
global_sizes.append(item)
|
||||
else:
|
||||
global_sizes.append(item * topology[dims_mapping[idx]])
|
||||
return global_sizes
|
||||
|
||||
@staticmethod
|
||||
def get_local_shard(
|
||||
global_sizes, dims_mapping, topology, processes, rank, shard_sizes=None
|
||||
):
|
||||
local_offsets = DistributedTensor.get_local_offsets(
|
||||
global_sizes, dims_mapping, topology, processes, rank, shard_sizes
|
||||
)
|
||||
local_sizes = DistributedTensor.get_local_sizes(
|
||||
global_sizes, dims_mapping, topology, processes, rank, shard_sizes
|
||||
)
|
||||
assert len(local_sizes) == len(local_offsets), (
|
||||
f"The length of local_sizes must be equal to local_offsets, but got {len(local_sizes)} and {len(local_offsets)}."
|
||||
)
|
||||
|
||||
local_end_offsets = [
|
||||
x[0] + x[1] for x in zip(local_offsets, local_sizes)
|
||||
]
|
||||
local_shard = list(zip(local_offsets, local_end_offsets))
|
||||
return local_shard
|
||||
|
||||
def __init__(self, serial_tensor, dist_attr=None, dist_context=None):
|
||||
self._serial_tensor = serial_tensor
|
||||
if dist_attr is not None and isinstance(dist_attr, TensorDistAttr):
|
||||
# TODO: remove this deepcopy after we fix the issue
|
||||
self._dist_attr = copy.deepcopy(dist_attr)
|
||||
# self._dist_attr = dist_attr
|
||||
# TODO: Do we really need to write dist_attr back to serial_tensor？
|
||||
self._serial_tensor.dist_attr = dist_attr
|
||||
else:
|
||||
assert dist_attr is None, f"{dist_attr}"
|
||||
# Use the dist attr of serial_tensor to do the initialization
|
||||
self._dist_attr = self._serial_tensor.dist_attr
|
||||
|
||||
self._batch_dim = 0
|
||||
self._local_offsets_map = {}
|
||||
self._local_shard_map = {}
|
||||
self._local_tensor_map = {}
|
||||
|
||||
from .dist_context import get_default_distributed_context
|
||||
|
||||
self._dist_context = (
|
||||
dist_context
|
||||
if dist_context is not None
|
||||
else get_default_distributed_context()
|
||||
)
|
||||
# TODO: Add Automatically to dist_context after initialized and it will be adapted in the future.
|
||||
# self._dist_context.add_dist_tensor_for_program(self)
|
||||
|
||||
@property
|
||||
def serial_tensor(self):
|
||||
return self._serial_tensor
|
||||
|
||||
@property
|
||||
def dist_attr(self):
|
||||
return self._dist_attr
|
||||
|
||||
@dist_attr.setter
|
||||
def dist_attr(self, dist_attr):
|
||||
self._dist_attr = dist_attr
|
||||
# TODO: Do we really need to write back dist_attr to serial_tensor？
|
||||
self._serial_tensor.dist_attr = dist_attr
|
||||
|
||||
@property
|
||||
def dist_context(self):
|
||||
return self._dist_context
|
||||
|
||||
# def _init_default_dist_attr(self):
|
||||
# if self._dist_attr.dims_mapping is None:
|
||||
# if self.serial_tensor.type in __no_shape_var_type__:
|
||||
# tensor_shape = []
|
||||
# else:
|
||||
# tensor_shape = self._serial_tensor.shape
|
||||
# tensor_dims_mapping = [-1 for _ in range(len(tensor_shape))]
|
||||
# self._dist_attr.dims_mapping = tensor_dims_mapping
|
||||
|
||||
def validate_dist_attr(self):
|
||||
if self.serial_tensor.type in __no_shape_var_type__:
|
||||
return True
|
||||
tensor_shape = self.serial_tensor.shape
|
||||
if len(tensor_shape) != len(self.dist_attr.dims_mapping):
|
||||
return False
|
||||
for i in range(len(self.dist_attr.dims_mapping)):
|
||||
if self.dist_attr.dims_mapping[
|
||||
i
|
||||
] < -1 or self.dist_attr.dims_mapping[i] >= len(
|
||||
self.dist_attr.process_mesh.shape
|
||||
):
|
||||
return False
|
||||
for i in range(len(self.dist_attr.process_mesh.shape)):
|
||||
if self.dist_attr.dims_mapping.count(i) > 1:
|
||||
return False
|
||||
return True
|
||||
|
||||
def local_sizes(self, rank=None):
|
||||
"""Get local sizes of the given rank."""
|
||||
rank = paddle.distributed.get_rank() if rank is None else rank
|
||||
global_sizes = self.serial_tensor.shape
|
||||
dims_mapping = self.dist_attr.dims_mapping
|
||||
# shard_sizes = self.dist_attr.shard_sizes
|
||||
processes = self.dist_attr.process_mesh.process_ids
|
||||
topology = self.dist_attr.process_mesh.shape
|
||||
local_sizes = DistributedTensor.get_local_sizes(
|
||||
global_sizes, dims_mapping, topology, processes, rank
|
||||
)
|
||||
|
||||
return local_sizes
|
||||
|
||||
def local_offsets(self, rank=None):
|
||||
rank = paddle.distributed.get_rank() if rank is None else rank
|
||||
local_offsets = None
|
||||
if rank in self._local_offsets_map.keys():
|
||||
local_offsets = self._local_offsets_map[rank]
|
||||
else:
|
||||
global_sizes = self.serial_tensor.shape
|
||||
dims_mapping = self.dist_attr.dims_mapping
|
||||
# shard_sizes = self.dist_attr.shard_sizes
|
||||
processes = self.dist_attr.process_mesh.process_ids
|
||||
topology = self.dist_attr.process_mesh.shape
|
||||
local_offsets = DistributedTensor.get_local_offsets(
|
||||
global_sizes, dims_mapping, topology, processes, rank
|
||||
)
|
||||
self._local_offsets_map[rank] = local_offsets
|
||||
|
||||
return local_offsets
|
||||
|
||||
def global_sizes(self):
|
||||
return self.serial_tensor.shape
|
||||
|
||||
def local_shard(self, rank=None):
|
||||
rank = paddle.distributed.get_rank() if rank is None else rank
|
||||
local_shard = None
|
||||
if rank in self._local_shard_map.keys():
|
||||
local_shard = self._local_shard_map[rank]
|
||||
else:
|
||||
global_sizes = self.serial_tensor.shape
|
||||
dims_mapping = self.dist_attr.dims_mapping
|
||||
# shard_sizes = self.dist_attr.shard_sizes
|
||||
processes = self.dist_attr.process_mesh.process_ids
|
||||
topology = self.dist_attr.process_mesh.shape
|
||||
local_shard = DistributedTensor.get_local_shard(
|
||||
global_sizes, dims_mapping, topology, processes, rank
|
||||
)
|
||||
self._local_shard_map[rank] = local_shard
|
||||
|
||||
return local_shard
|
||||
|
||||
def new_local_tensor(self, block=None, rank=None, name=None):
|
||||
"""
|
||||
Create a new local tensor of serial tensor corresponding to rank.
|
||||
Args:
|
||||
block (Block): The block contains the new tensor. Default value is recommend and it will be created in the block of dist main program corresponding to the serial tensor block id. Default: None.
|
||||
rank (int): The rank id. Default value is recommend and it will be the current rank. Default: None.
|
||||
"""
|
||||
|
||||
def _copy_kwargs(serial_tensor):
|
||||
kwargs = {}
|
||||
no_need_copy_args = ["self", "block", "shape", "name"]
|
||||
arg_spec = inspect.getfullargspec(Variable.__init__)
|
||||
|
||||
for key in arg_spec.args:
|
||||
# TODO: Check the copied attribute from serial tensor whether valid
|
||||
if key in no_need_copy_args:
|
||||
continue
|
||||
elif key not in kwargs:
|
||||
if key == "type":
|
||||
kwargs[key] = serial_tensor.desc.type()
|
||||
elif key == "dtype":
|
||||
kwargs[key] = serial_tensor.desc.dtype()
|
||||
elif key == "lod_level":
|
||||
kwargs[key] = serial_tensor.desc.lod_level()
|
||||
elif key == "persistable":
|
||||
kwargs[key] = serial_tensor.desc.persistable()
|
||||
elif key == "stop_gradient":
|
||||
kwargs[key] = serial_tensor.desc.stop_gradient()
|
||||
elif key == "need_check_feed":
|
||||
kwargs[key] = serial_tensor.desc.need_check_feed()
|
||||
# TODO: Get capacity by framework
|
||||
elif key == "capacity":
|
||||
continue
|
||||
else:
|
||||
kwargs[key] = self.serial_tensor.__dict__[key]
|
||||
|
||||
if isinstance(serial_tensor, Parameter):
|
||||
kwargs["trainable"] = serial_tensor.trainable
|
||||
kwargs["optimize_attr"] = serial_tensor.trainable
|
||||
kwargs["regularizer"] = serial_tensor.regularizer
|
||||
kwargs["do_model_average"] = serial_tensor.do_model_average
|
||||
kwargs["need_clip"] = serial_tensor.need_clip
|
||||
kwargs["is_distributed"] = serial_tensor.is_distributed
|
||||
kwargs["is_parameter"] = serial_tensor.is_parameter
|
||||
|
||||
return kwargs
|
||||
|
||||
if rank is not None and not (isinstance(rank, int) and rank >= 0):
|
||||
raise ValueError(f"The rank must >= 0, but got {rank}")
|
||||
if block is not None and not isinstance(block, Block):
|
||||
raise TypeError(f"The block must be Block, but got {type(block)}.")
|
||||
rank = paddle.distributed.get_rank() if rank is None else rank
|
||||
|
||||
if block is None:
|
||||
block_id = self.serial_tensor.block.idx
|
||||
block = self.dist_context.dist_main_programs[rank].block(block_id)
|
||||
|
||||
# copy serial tensor attribute
|
||||
kwargs = _copy_kwargs(self.serial_tensor)
|
||||
kwargs["name"] = name
|
||||
kwargs["shape"] = self.local_sizes(rank)
|
||||
|
||||
if isinstance(self.serial_tensor, Parameter):
|
||||
kwargs.pop("persistable")
|
||||
local_tensor = Parameter(block=block, **kwargs)
|
||||
else:
|
||||
local_tensor = block.create_var(**kwargs)
|
||||
|
||||
# TODO: Set original id when set original_id is approved
|
||||
local_tensor.desc.set_original_id(self.serial_tensor.desc.id())
|
||||
self._local_tensor_map[rank] = local_tensor
|
||||
return local_tensor
|
||||
|
||||
def local_tensor(self, rank=None):
|
||||
rank = paddle.distributed.get_rank() if rank is None else rank
|
||||
assert rank in self._local_tensor_map, (
|
||||
f"The rank {rank} local tensor has not been created."
|
||||
)
|
||||
return self._local_tensor_map[rank]
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
cls = self.__class__
|
||||
result = cls.__new__(cls)
|
||||
memo[id(self)] = result
|
||||
for k, v in self.__dict__.items():
|
||||
if k == "_serial_tensor" or k == "_local_tensor_map":
|
||||
setattr(result, k, v)
|
||||
else:
|
||||
setattr(result, k, copy.deepcopy(v, memo))
|
||||
return result
|
||||
|
||||
def __str__(self):
|
||||
str = f"{{tensor name: {self.serial_tensor.desc.name()}, tensor id: {self.serial_tensor.desc.id()}, tensor original_id {self.serial_tensor.desc.original_id()}"
|
||||
|
||||
# str += ", {}".format(self.dist_attr)
|
||||
# return str
|
||||
|
||||
if self.dist_attr.is_annotated("process_mesh"):
|
||||
annotated_str = "annotated"
|
||||
else:
|
||||
annotated_str = "non-annotated"
|
||||
str += (
|
||||
f", process_mesh ({annotated_str}): {self.dist_attr.process_mesh}"
|
||||
)
|
||||
|
||||
str += f", is_parameter: {self.serial_tensor.is_parameter}"
|
||||
str += f", chunk_id: {self.dist_attr.chunk_id}"
|
||||
|
||||
if self.dist_attr.is_annotated("dims_mapping"):
|
||||
annotated_str = "annotated"
|
||||
else:
|
||||
annotated_str = "non-annotated"
|
||||
str += f", dims_mapping ({annotated_str}): {self.dist_attr.dims_mapping} }}"
|
||||
|
||||
# if self.dist_attr.is_annotated("shard_mask"):
|
||||
# annotated_str = "annotated"
|
||||
# else:
|
||||
# annotated_str = "non-annotated"
|
||||
# str += ", shard_mask ({}): {}".format(annotated_str, None)
|
||||
|
||||
# if self.dist_attr.is_annotated("offload_device"):
|
||||
# annotated_str = "annotated"
|
||||
# else:
|
||||
# annotated_str = "non-annotated"
|
||||
# str += ", offload_device ({}): {} }}".format(annotated_str, None)
|
||||
return str
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,186 @@
|
||||
# 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
|
||||
|
||||
from collections import OrderedDict
|
||||
|
||||
|
||||
class Node:
|
||||
def __init__(self, id, **attrs):
|
||||
# Each node must has a unique id
|
||||
self._id = id
|
||||
# Attributes for Node
|
||||
self._attrs = {}
|
||||
self._attrs.update(attrs)
|
||||
|
||||
@property
|
||||
def id(self):
|
||||
return self._id
|
||||
|
||||
@property
|
||||
def attrs(self):
|
||||
return self._attrs
|
||||
|
||||
def __getitem__(self, attr_name):
|
||||
return self._attrs[attr_name]
|
||||
|
||||
def __setitem__(self, attr_name, attr_value):
|
||||
self._attrs[attr_name] = attr_value
|
||||
|
||||
def __contains__(self, attr_name):
|
||||
try:
|
||||
return attr_name in self._attrs
|
||||
except TypeError:
|
||||
return False
|
||||
|
||||
def __str__(self):
|
||||
str = f"(id: {self.id}, attrs: {self.attrs})"
|
||||
return str
|
||||
|
||||
|
||||
class Edge:
|
||||
def __init__(self, src_id, tgt_id, **attrs):
|
||||
# The id of source node in an Edge
|
||||
self._src_id = src_id
|
||||
# The id of target node in an Edge
|
||||
self._tgt_id = tgt_id
|
||||
# Attributes for Edge
|
||||
self._attrs = {}
|
||||
self._attrs.update(attrs)
|
||||
|
||||
@property
|
||||
def src_id(self):
|
||||
return self._src_id
|
||||
|
||||
@property
|
||||
def tgt_id(self):
|
||||
return self._tgt_id
|
||||
|
||||
@property
|
||||
def attrs(self):
|
||||
return self._attrs
|
||||
|
||||
def __getitem__(self, attr_name):
|
||||
return self._attrs[attr_name]
|
||||
|
||||
def __setitem__(self, attr_name, attr_value):
|
||||
self._attrs[attr_name] = attr_value
|
||||
|
||||
def __contains__(self, attr_name):
|
||||
try:
|
||||
return attr_name in self._attrs
|
||||
except TypeError:
|
||||
return False
|
||||
|
||||
def __str__(self):
|
||||
str = ""
|
||||
str += f"(src_id: {self.src_id}, tgt_id: {self.tgt_id}, attrs: {self._attrs})"
|
||||
return str
|
||||
|
||||
|
||||
class Graph:
|
||||
def __init__(self, **attrs):
|
||||
# _nodes is dict for storing the nodes of the graph.
|
||||
# The key of this dict is the node id.
|
||||
self._nodes = {}
|
||||
# _adjs is a dict of dict for storing the adjacency of the graph.
|
||||
# The key of the outer dict is the node id of the source node and
|
||||
# the key of the inner dict is the node id of the target node.
|
||||
self._adjs = {}
|
||||
# Attributes for Graph
|
||||
self._attrs = {}
|
||||
self._attrs.update(attrs)
|
||||
self._reverse_adjs = {}
|
||||
self._attr_to_nodes = {}
|
||||
|
||||
@property
|
||||
def nodes(self):
|
||||
return self._nodes
|
||||
|
||||
@property
|
||||
def attrs(self):
|
||||
return self._attrs
|
||||
|
||||
@property
|
||||
def adjs(self):
|
||||
return self._adjs
|
||||
|
||||
def add_node(self, node_id, **attrs):
|
||||
if node_id is None:
|
||||
raise ValueError("None cannot be a node")
|
||||
if node_id not in self._nodes:
|
||||
node = Node(node_id, **attrs)
|
||||
self._nodes[node_id] = node
|
||||
self._adjs[node_id] = {}
|
||||
self._reverse_adjs[node_id] = []
|
||||
else:
|
||||
self._nodes[node_id].attrs.update(attrs)
|
||||
|
||||
return self._nodes[node_id]
|
||||
|
||||
def add_edge(self, src_id, tgt_id, **attrs):
|
||||
# add nodes
|
||||
if src_id is None:
|
||||
raise ValueError("None cannot be a node")
|
||||
if tgt_id is None:
|
||||
raise ValueError("None cannot be a node")
|
||||
if src_id not in self._nodes:
|
||||
src_node = Node(src_id)
|
||||
self._nodes[src_id] = src_node
|
||||
# for one tensor to multiple ops
|
||||
self._adjs[src_id] = OrderedDict()
|
||||
self._reverse_adjs[src_id] = []
|
||||
if tgt_id not in self._nodes:
|
||||
tgt_node = Node(tgt_id)
|
||||
self._nodes[tgt_id] = tgt_node
|
||||
# for one tensor to multiple ops
|
||||
self._adjs[tgt_id] = OrderedDict()
|
||||
self._reverse_adjs[tgt_id] = []
|
||||
# add the edge
|
||||
edge = Edge(src_id, tgt_id, **attrs)
|
||||
self._adjs[src_id][tgt_id] = edge
|
||||
|
||||
# add the reverse adj
|
||||
self._reverse_adjs[tgt_id].append(self.nodes[src_id])
|
||||
return edge
|
||||
|
||||
def __len__(self):
|
||||
return len(self._nodes)
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._nodes.values())
|
||||
|
||||
def __getitem__(self, node_id):
|
||||
# Return the adjacency of a node
|
||||
return self._adjs[node_id]
|
||||
|
||||
def __contains__(self, node_id):
|
||||
# Check whether a node in the graph
|
||||
try:
|
||||
return node_id in self._nodes
|
||||
except TypeError:
|
||||
return False
|
||||
|
||||
def __str__(self):
|
||||
str = ""
|
||||
str += "**************Nodes**************\n"
|
||||
for node_id in self.nodes:
|
||||
str += f"{self.nodes[node_id]}\n"
|
||||
|
||||
str += "**************Edges**************\n"
|
||||
for src_id in self.adjs:
|
||||
str += f"--------------{src_id}--------------\n"
|
||||
for idx, tgt_id in enumerate(self.adjs[src_id]):
|
||||
str += f"{self.adjs[src_id][tgt_id]}\n"
|
||||
|
||||
return str
|
||||
@@ -0,0 +1,673 @@
|
||||
# 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 inspect
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
|
||||
import paddle
|
||||
from paddle import core
|
||||
from paddle.jit import not_to_static, to_static
|
||||
from paddle.jit.dy2static.program_translator import (
|
||||
ProgramTranslator,
|
||||
StaticFunction,
|
||||
)
|
||||
from paddle.jit.dy2static.utils import as_not_paddle_func
|
||||
from paddle.nn import Layer
|
||||
from paddle.static import Parameter, global_scope, program_guard
|
||||
from paddle.static.amp.fp16_utils import (
|
||||
DEFAULT_AMP_OPTIONS,
|
||||
prepare_op_amp_options,
|
||||
)
|
||||
|
||||
from .converter import Converter
|
||||
from .dist_attribute import TensorDistAttr
|
||||
from .process_group import get_world_process_group
|
||||
from .utils import get_logger, to_list
|
||||
|
||||
|
||||
class ProxyLayer(Layer):
|
||||
"""
|
||||
ProxyLayer implements all logic for converting dygraph model into
|
||||
static Program IR. Meanwhile, it provides conventional interfaces for
|
||||
auto parallel to visit feed/fetch/loss/metric variables.
|
||||
"""
|
||||
|
||||
def __init__(self, layer, loss_func, metrics):
|
||||
super().__init__()
|
||||
# NOTE: All verify logics are finished in Engine.Prepare
|
||||
self.inner_layer = layer
|
||||
self.loss_func = loss_func
|
||||
self.metrics = metrics
|
||||
# train / eval / predict
|
||||
self.mode = None
|
||||
|
||||
# generated program vars
|
||||
self._input_vars = defaultdict(list)
|
||||
self._label_vars = defaultdict(list)
|
||||
self._output_vars = defaultdict(list)
|
||||
self._loss_vars = defaultdict(list)
|
||||
self._loss_names = defaultdict(list)
|
||||
self._metric_vars = defaultdict(list)
|
||||
|
||||
# Consider ProxyLayer as not Paddle inner function because it contains
|
||||
# user-defined layer.
|
||||
for fn_name in [
|
||||
"_train",
|
||||
"_eval",
|
||||
"_predict",
|
||||
"call_loss",
|
||||
"call_metrics",
|
||||
]:
|
||||
as_not_paddle_func(
|
||||
f"{inspect.getmodule(ProxyLayer).__name__}.ProxyLayer.{fn_name}"
|
||||
)
|
||||
|
||||
@paddle.jit.not_to_static
|
||||
def append_loss_to_shadow_output(self, mode):
|
||||
name = paddle.utils.unique_name.generate('loss')
|
||||
paddle._C_ops.set_persistable_value(self._loss_vars[mode], name)
|
||||
self._loss_names[mode] = name
|
||||
|
||||
def _train(self, inputs, labels):
|
||||
"""
|
||||
Train process of inner_layer with forward/loss/metric logic.
|
||||
"""
|
||||
# step 1. save feed variables of Program
|
||||
mode = 'train'
|
||||
self._input_vars[mode] = inputs
|
||||
self._label_vars[mode] = labels
|
||||
|
||||
# step 2. call inner_layer.forward
|
||||
self._output_vars[mode] = self.inner_layer(*inputs)
|
||||
|
||||
# step 3. calculate loss if needed
|
||||
new_inputs = self._prepare(self.output_vars, labels)
|
||||
self._loss_vars[mode] = self.call_loss(new_inputs)
|
||||
if paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
|
||||
"FLAGS_enable_pir_api"
|
||||
]:
|
||||
self.append_loss_to_shadow_output(mode)
|
||||
|
||||
# step 4. calculate metrics if needed
|
||||
self._metric_vars[mode] = self.call_metrics(new_inputs)
|
||||
|
||||
def _eval(self, inputs, labels):
|
||||
"""
|
||||
Evaluate process of inner_layer with forward/loss/metric logic.
|
||||
"""
|
||||
# TODO(dev): we can reuse codes with self._train after making
|
||||
# sure if they can.
|
||||
|
||||
# step 1. save feed variables of Program
|
||||
mode = 'eval'
|
||||
self._input_vars[mode] = inputs
|
||||
self._label_vars[mode] = labels
|
||||
|
||||
# step 2. call inner_layer.forward
|
||||
self._output_vars[mode] = self.inner_layer(*inputs)
|
||||
|
||||
# step 3. calculate loss if needed
|
||||
new_inputs = self._prepare(self.output_vars, labels)
|
||||
self._loss_vars[mode] = self.call_loss(new_inputs)
|
||||
if paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
|
||||
"FLAGS_enable_pir_api"
|
||||
]:
|
||||
self.append_loss_to_shadow_output(mode)
|
||||
|
||||
# step 4. calculate metrics if needed
|
||||
self._metric_vars[mode] = self.call_metrics(new_inputs)
|
||||
|
||||
def _predict(self, inputs, labels):
|
||||
"""
|
||||
Predict process of inner_layer with forward logic.
|
||||
"""
|
||||
# step 1. save feed variables of Program
|
||||
mode = 'predict'
|
||||
self._input_vars[mode] = inputs
|
||||
self._label_vars[mode] = labels
|
||||
|
||||
# step 2. call inner_layer.forward
|
||||
self._output_vars[mode] = self.inner_layer(*inputs)
|
||||
|
||||
@not_to_static
|
||||
def _prepare(self, outputs, labels):
|
||||
"""
|
||||
Concat outputs and labels as a single list
|
||||
|
||||
NOTE(dev): We use @not_to_static to avoid AST Analysis.
|
||||
"""
|
||||
return to_list(outputs) + to_list(labels)
|
||||
|
||||
def call_loss(self, inputs):
|
||||
"""
|
||||
Apply Loss Function on outputs and labels.
|
||||
|
||||
Args:
|
||||
inputs: List[Variable]
|
||||
|
||||
Returns: List[Variable]
|
||||
"""
|
||||
res = []
|
||||
if self.loss_func is not None:
|
||||
res = self.loss_func(*inputs)
|
||||
return res
|
||||
|
||||
def call_metrics(self, inputs):
|
||||
"""
|
||||
Apply Metrics Function on outputs and labels.
|
||||
|
||||
Args:
|
||||
inputs: List[Variable]
|
||||
|
||||
Returns: List[Variable]
|
||||
"""
|
||||
outs = []
|
||||
for metric in self.metrics:
|
||||
outs.append(to_list(metric.compute(*inputs)))
|
||||
|
||||
return outs
|
||||
|
||||
def set_mode(self, mode):
|
||||
self.mode = mode
|
||||
self.training = mode == 'train'
|
||||
|
||||
def clone(self):
|
||||
return ProxyLayer(self.inner_layer, self.loss_func, self.metrics)
|
||||
|
||||
@property
|
||||
def input_vars(self):
|
||||
return self._input_vars[self.mode]
|
||||
|
||||
@property
|
||||
def label_vars(self):
|
||||
return self._label_vars[self.mode]
|
||||
|
||||
@property
|
||||
def output_vars(self):
|
||||
return self._output_vars[self.mode]
|
||||
|
||||
@property
|
||||
def loss_vars(self):
|
||||
return self._loss_vars[self.mode]
|
||||
|
||||
@property
|
||||
def loss_names(self):
|
||||
return self._loss_names[self.mode]
|
||||
|
||||
@property
|
||||
def metric_vars(self):
|
||||
return self._metric_vars[self.mode]
|
||||
|
||||
@property
|
||||
def startup_program(self):
|
||||
return self.inner_layer._startup_program()
|
||||
|
||||
|
||||
class BuildInfo:
|
||||
def __init__(self):
|
||||
self.clear()
|
||||
|
||||
def has_cache(self, mode, update=False):
|
||||
is_cache = self.states[mode]
|
||||
if update:
|
||||
self.cache(mode)
|
||||
return is_cache
|
||||
|
||||
def cache(self, mode):
|
||||
self.states[mode] = True
|
||||
|
||||
def clear(self):
|
||||
self.states = defaultdict(bool)
|
||||
|
||||
|
||||
class ProgramHelper:
|
||||
"""
|
||||
A Helper class for Engine to provides different Program IR according specified 'mode'.
|
||||
"""
|
||||
|
||||
def __init__(self, layer, loss_func, metrics, inputs_spec, labels_spec):
|
||||
# original model config information
|
||||
# TODO(Aurelius84): Implement append_backward and optimizer in ProxyLayer
|
||||
# after distribute engine satisfy basic condition.
|
||||
self.proxy_layer = ProxyLayer(layer, loss_func, metrics)
|
||||
self.inputs_spec = inputs_spec
|
||||
self.labels_spec = labels_spec
|
||||
|
||||
self.build_info = BuildInfo()
|
||||
self._logger = get_logger(logging.INFO)
|
||||
self.lazy_init = False
|
||||
self._all_params_dist_attr = {}
|
||||
|
||||
def reset(self):
|
||||
"""
|
||||
Reset all state of current Object.
|
||||
"""
|
||||
self.build_info.clear()
|
||||
self.proxy_layer = self.proxy_layer.clone()
|
||||
|
||||
def build_program(self, mode):
|
||||
"""
|
||||
Convert dygraph model into static Program IR.
|
||||
"""
|
||||
assert mode in ['train', 'eval', 'predict']
|
||||
self.proxy_layer.set_mode(mode)
|
||||
# skip if we has already built program.
|
||||
if self.build_info.has_cache(mode, True):
|
||||
self._logger.info(
|
||||
f"Already build program with mode = {mode}, use cached program."
|
||||
)
|
||||
return
|
||||
|
||||
self._logger.info(f"start to build program for mode = {mode}.")
|
||||
input_spec = [self.inputs_spec, self.labels_spec]
|
||||
static_func = to_static(
|
||||
self.static_func(), input_spec=input_spec, full_graph=True
|
||||
)
|
||||
|
||||
func_name = '_' + mode
|
||||
setattr(self.proxy_layer, func_name, static_func)
|
||||
|
||||
# NOTE(dev): Because @to_static is a Lazy mechanism, so we explicitly call this to trigger
|
||||
# generating Program IR immediately.
|
||||
concrete_program = getattr(self.proxy_layer, func_name).concrete_program
|
||||
|
||||
# TODO(zhiqiu): prepare_op_amp_options is not supported for PIR program
|
||||
# It will to use dynamic-static unified amp in pir program, and there is
|
||||
# no need to fit for prepare_op_amp_options
|
||||
if not paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
|
||||
"FLAGS_enable_pir_api"
|
||||
]:
|
||||
prepare_op_amp_options(
|
||||
concrete_program.main_program,
|
||||
ProgramTranslator.get_instance()._amp_records,
|
||||
DEFAULT_AMP_OPTIONS,
|
||||
)
|
||||
self._build_startup_program()
|
||||
|
||||
def _build_startup_program(self):
|
||||
"""
|
||||
Create and Sync parameters into startup program.
|
||||
"""
|
||||
startup_program = self.startup_program
|
||||
if len(startup_program.global_block().ops) > 1:
|
||||
self.lazy_init = True
|
||||
return
|
||||
|
||||
for param in self.concrete_program.parameters:
|
||||
Parameter(
|
||||
name=param.name,
|
||||
desc=param,
|
||||
type=param.type,
|
||||
shape=param.shape,
|
||||
dtype=param.dtype,
|
||||
stop_gradient=param.stop_gradient,
|
||||
block=startup_program.global_block(),
|
||||
)
|
||||
|
||||
def apply_optimizer(self, optimizer):
|
||||
"""
|
||||
Append backward and generate optimizer operations.
|
||||
"""
|
||||
self._verify_optimizer(optimizer)
|
||||
self._logger.info(
|
||||
"start to apply optimizer: %s ", type(optimizer).__name__
|
||||
)
|
||||
# clear optimizer parameters
|
||||
original_params = optimizer._parameter_list
|
||||
optimizer._parameter_list = None
|
||||
with program_guard(self.main_program, self.startup_program):
|
||||
res = optimizer.minimize(self.loss_vars[0])
|
||||
|
||||
# restore optimizer parameters
|
||||
optimizer._parameter_list = original_params
|
||||
return res
|
||||
|
||||
def _verify_optimizer(self, optimizer):
|
||||
assert optimizer is not None
|
||||
assert hasattr(optimizer, "minimize"), (
|
||||
"Optimizer must have minimize() method."
|
||||
)
|
||||
assert self.proxy_layer.mode == 'train', (
|
||||
f"Required mode == 'train', but received '{self.proxy_layer.mode}'"
|
||||
)
|
||||
assert len(self.loss_vars) == 1, (
|
||||
f"Required len(loss_vars) == 1, but received len(loss_vars) = {len(self.loss_vars)}"
|
||||
)
|
||||
|
||||
def to(self, mode):
|
||||
"""
|
||||
Switch underly proxy layer mode into target mode.
|
||||
"""
|
||||
assert mode in ['train', 'eval', 'predict']
|
||||
func = getattr(self.proxy_layer, '_' + mode)
|
||||
assert isinstance(func, StaticFunction), (
|
||||
"Please call build_program(mode) firstly."
|
||||
)
|
||||
self.proxy_layer.set_mode(mode)
|
||||
|
||||
def static_func(self):
|
||||
"""
|
||||
Return StaticFunction instance with underly target mode.
|
||||
"""
|
||||
assert self.proxy_layer.mode in [
|
||||
'train',
|
||||
'eval',
|
||||
'predict',
|
||||
], "Please call build_program(mode) firstly."
|
||||
func_name = '_' + self.proxy_layer.mode
|
||||
return getattr(self.proxy_layer, func_name)
|
||||
|
||||
def init_pir(self, main_program, place):
|
||||
# collect all params in current dist program
|
||||
param_values = main_program.global_block().all_parameters()
|
||||
value_name_to_value = {}
|
||||
dy_param_name_to_pir_param_name = {}
|
||||
for value in param_values:
|
||||
value_name_to_value[value.name] = value
|
||||
|
||||
dy_params = self.concrete_program.parameters[0]
|
||||
pir_param = self.concrete_program.parameters[1]
|
||||
|
||||
for i in range(len(pir_param)):
|
||||
if pir_param[i].name in value_name_to_value:
|
||||
dy_param_name_to_pir_param_name[dy_params[i].name] = pir_param[
|
||||
i
|
||||
].name
|
||||
|
||||
is_comm = False
|
||||
for param in dy_params:
|
||||
if param.is_dist():
|
||||
process_mesh, dims_mapping = self._all_params_dist_attr[
|
||||
param.name
|
||||
]
|
||||
var_dist_attr = TensorDistAttr()
|
||||
var_dist_attr.process_mesh = process_mesh
|
||||
var_dist_attr.dims_mapping = dims_mapping
|
||||
is_comm = True
|
||||
with paddle.no_grad():
|
||||
tmp = paddle.base.core.reshard(param, var_dist_attr)
|
||||
if tmp._is_initialized():
|
||||
param.get_tensor()._share_data_with(tmp.get_tensor())
|
||||
else:
|
||||
# Only setting the "param" to "None" can't release the memory
|
||||
param.get_tensor()._clear()
|
||||
param = None
|
||||
|
||||
# create var in scope and share parameters to scope
|
||||
if param is None:
|
||||
continue
|
||||
if param.name not in dy_param_name_to_pir_param_name:
|
||||
# Release the redundant params
|
||||
param.get_tensor()._clear()
|
||||
continue
|
||||
if not param._is_initialized():
|
||||
continue
|
||||
if param.is_dense():
|
||||
value_name = dy_param_name_to_pir_param_name[param.name]
|
||||
value = value_name_to_value[value_name]
|
||||
# get param_var's dist_attr
|
||||
assert value.is_dist_dense_tensor_type(), (
|
||||
f"param [{value.name}] is not dist tensor type"
|
||||
)
|
||||
dist_attr = {
|
||||
"dims_mapping": value.dist_attr().dims_mapping,
|
||||
"process_shape": value.dist_attr().process_mesh.shape,
|
||||
"process_group": value.dist_attr().process_mesh.process_ids,
|
||||
}
|
||||
# slice param_value with dist_attr
|
||||
# share sliced_param_value with param_tensor in global_scope
|
||||
pir_scope_param = global_scope().var(value_name).get_tensor()
|
||||
sliced_param = Converter.slice_with_dist_attr(
|
||||
param.numpy(), dist_attr
|
||||
)
|
||||
pir_scope_param.set(sliced_param, place)
|
||||
param.get_tensor()._clear()
|
||||
|
||||
elif param.is_dist():
|
||||
value_name = dy_param_name_to_pir_param_name[param.name]
|
||||
value = value_name_to_value[value_name]
|
||||
# assert value.is_dist_dense_tensor_type(), "param [{}] is not dist tensor type".format(value.name)
|
||||
pir_scope_param = global_scope().var(value_name).get_tensor()
|
||||
pir_scope_param._share_data_with(
|
||||
param.get_tensor().get_tensor()
|
||||
)
|
||||
param.get_tensor()._clear()
|
||||
|
||||
world_group = get_world_process_group()
|
||||
if (
|
||||
is_comm
|
||||
and world_group.nranks > 1
|
||||
and paddle.distributed.get_world_size() > 1
|
||||
):
|
||||
paddle.disable_static()
|
||||
barrier_tensor = paddle.full([1], 1, dtype="int32")
|
||||
# barrier is not available in xpu for now
|
||||
if not paddle.framework.core.is_compiled_with_xpu():
|
||||
paddle._legacy_C_ops.barrier(
|
||||
barrier_tensor, barrier_tensor, 'ring_id', 0
|
||||
)
|
||||
paddle.enable_static()
|
||||
|
||||
def init(self, main_program, place, dist_context):
|
||||
if self.lazy_init:
|
||||
return
|
||||
|
||||
amp_strategy = dist_context.strategy.amp
|
||||
amp_config = copy.deepcopy(amp_strategy.to_dict())
|
||||
need_cast_parameter = amp_strategy.enable and amp_config["level"] in [
|
||||
"o2",
|
||||
"o3",
|
||||
]
|
||||
is_comm = False
|
||||
for param in self.concrete_program.parameters:
|
||||
if param.is_dist():
|
||||
serial_main_program = self.concrete_program.main_program
|
||||
var = serial_main_program.global_block().vars[param.name]
|
||||
var_dist_attr = dist_context.get_tensor_dist_attr_for_program(
|
||||
var
|
||||
)
|
||||
is_comm = True
|
||||
# No need to construct backward.
|
||||
with paddle.no_grad():
|
||||
tmp = paddle.base.core.reshard(param, var_dist_attr)
|
||||
if tmp._is_initialized():
|
||||
param.get_tensor()._share_data_with(tmp.get_tensor())
|
||||
else:
|
||||
# Only setting the "param" to "None" can't release the memory
|
||||
param.get_tensor()._clear()
|
||||
param = None
|
||||
paddle.device.synchronize()
|
||||
|
||||
# create var in scope and share parameters to scope
|
||||
if param is None:
|
||||
continue
|
||||
if param.name not in main_program.global_block().vars:
|
||||
# Release the redundant params
|
||||
param.get_tensor()._clear()
|
||||
continue
|
||||
if not param._is_initialized():
|
||||
continue
|
||||
if param.is_dense():
|
||||
# get param_var's dist_attr
|
||||
var = main_program.global_block().vars[param.name]
|
||||
var_dist_attr = dist_context.get_tensor_dist_attr_for_program(
|
||||
var
|
||||
)
|
||||
dist_attr = {
|
||||
"dims_mapping": var_dist_attr.dims_mapping,
|
||||
"process_shape": var_dist_attr.process_mesh.shape,
|
||||
"process_group": var_dist_attr.process_mesh.process_ids,
|
||||
}
|
||||
# slice param_value with dist_attr
|
||||
# share sliced_param_value with param_tensor in global_scope
|
||||
param_tensor = global_scope().var(param.name).get_tensor()
|
||||
sliced_param = Converter.slice_with_dist_attr(
|
||||
param.numpy(), dist_attr
|
||||
)
|
||||
param_tensor.set(sliced_param, place)
|
||||
if not need_cast_parameter:
|
||||
param.get_tensor()._clear()
|
||||
elif param.is_dist():
|
||||
dense_tensor = global_scope().var(param.name).get_tensor()
|
||||
dense_tensor._share_data_with(param.get_tensor().get_tensor())
|
||||
|
||||
# transform the parameter in eager mode for amp.
|
||||
if need_cast_parameter:
|
||||
for param in self.concrete_program.parameters:
|
||||
amp_dtype = amp_config["dtype"]
|
||||
scope_var = global_scope().find_var(param.name)
|
||||
# The parameter is not in this rank.
|
||||
if not scope_var:
|
||||
continue
|
||||
# The parameter do not need to transform
|
||||
if param.dtype in [paddle.float16, paddle.bfloat16]:
|
||||
continue
|
||||
scope_tensor = global_scope().var(param.name).get_tensor()
|
||||
assert scope_var and scope_tensor._is_initialized(), (
|
||||
f"Parameter: {param.name} is not put into global_scope or not initialized."
|
||||
)
|
||||
param_used = param
|
||||
# For the params without dist_attr.
|
||||
# NOTE(lizhiyu): In principle, each param should have dist_attr.
|
||||
if param.is_dense():
|
||||
# get param_var's dist_attr
|
||||
var = main_program.global_block().vars[param.name]
|
||||
var_dist_attr = (
|
||||
dist_context.get_tensor_dist_attr_for_program(var)
|
||||
)
|
||||
dist_attr = {
|
||||
"dims_mapping": var_dist_attr.dims_mapping,
|
||||
"process_shape": var_dist_attr.process_mesh.shape,
|
||||
"process_group": var_dist_attr.process_mesh.process_ids,
|
||||
}
|
||||
# slice param_value with dist_attr
|
||||
sliced_param = Converter.slice_with_dist_attr(
|
||||
param.numpy(), dist_attr
|
||||
)
|
||||
with paddle.base.dygraph.guard():
|
||||
param_used = paddle.to_tensor(
|
||||
sliced_param, place=param.place
|
||||
)
|
||||
param.get_tensor()._clear()
|
||||
with paddle.base.dygraph.guard():
|
||||
if amp_dtype == "float16":
|
||||
with (
|
||||
paddle.no_grad(),
|
||||
paddle.base.framework._dygraph_place_guard(
|
||||
place=place
|
||||
),
|
||||
):
|
||||
t_casted = param_used.cast(
|
||||
dtype=core.VarDesc.VarType.FP16
|
||||
)
|
||||
elif amp_dtype == "bfloat16":
|
||||
with (
|
||||
paddle.no_grad(),
|
||||
paddle.base.framework._dygraph_place_guard(
|
||||
place=place
|
||||
),
|
||||
):
|
||||
t_casted = param_used.cast(
|
||||
dtype=core.VarDesc.VarType.BF16
|
||||
)
|
||||
# NOTE(lizhiyu): Clear the origin param. Don't use `param_used.get_tensor().get_tensor()._clear()` to
|
||||
# clear the `DistTensor`, because it can't clear the `_holder`,
|
||||
# which `param_used.get_tensor().get_tensor()` will copy one `DenseTensor`.
|
||||
param_used.get_tensor()._clear()
|
||||
if t_casted.is_dist():
|
||||
scope_tensor._share_data_with(
|
||||
t_casted.get_tensor().get_tensor()
|
||||
)
|
||||
else:
|
||||
scope_tensor._share_data_with(t_casted.get_tensor())
|
||||
|
||||
world_group = get_world_process_group()
|
||||
if (
|
||||
is_comm
|
||||
and world_group.nranks > 1
|
||||
and paddle.distributed.get_world_size() > 1
|
||||
):
|
||||
paddle.disable_static()
|
||||
barrier_tensor = paddle.full([1], 1, dtype="int32")
|
||||
# barrier is not available in xpu for now
|
||||
if not paddle.framework.core.is_compiled_with_xpu():
|
||||
paddle._legacy_C_ops.barrier(
|
||||
barrier_tensor, barrier_tensor, 'ring_id', 0
|
||||
)
|
||||
paddle.enable_static()
|
||||
|
||||
def cache_whole_graph_dist_attr(self, all_params):
|
||||
for param_value in all_params:
|
||||
dist_attr = param_value.dist_attr()
|
||||
if dist_attr:
|
||||
process_mesh = dist_attr.process_mesh
|
||||
dims_mapping = dist_attr.dims_mapping
|
||||
self._all_params_dist_attr[param_value.name] = [
|
||||
process_mesh,
|
||||
dims_mapping,
|
||||
]
|
||||
|
||||
@property
|
||||
def concrete_program(self):
|
||||
return self.static_func().concrete_program
|
||||
|
||||
@property
|
||||
def main_program(self):
|
||||
return self.concrete_program.main_program
|
||||
|
||||
@property
|
||||
def startup_program(self):
|
||||
try:
|
||||
return self.proxy_layer.startup_program
|
||||
except Exception as err:
|
||||
self._logger.warning(
|
||||
"The startup_program is not built by `lazy init`."
|
||||
)
|
||||
if isinstance(err, AssertionError):
|
||||
return self.concrete_program.startup_program
|
||||
raise err
|
||||
|
||||
@property
|
||||
def input_vars(self):
|
||||
return to_list(self.proxy_layer.input_vars)
|
||||
|
||||
@property
|
||||
def output_vars(self):
|
||||
return to_list(self.proxy_layer.output_vars)
|
||||
|
||||
@property
|
||||
def label_vars(self):
|
||||
return to_list(self.proxy_layer.label_vars)
|
||||
|
||||
@property
|
||||
def loss_vars(self):
|
||||
return to_list(self.proxy_layer.loss_vars)
|
||||
|
||||
@property
|
||||
def loss_names(self):
|
||||
return to_list(self.proxy_layer.loss_names)
|
||||
|
||||
@property
|
||||
def metric_vars(self):
|
||||
return to_list(self.proxy_layer.metric_vars)
|
||||
|
||||
def named_parameters(self):
|
||||
static_func = self.static_func()
|
||||
partial_program = static_func.get_concrete_program(
|
||||
self.inputs_spec, self.labels_spec
|
||||
)[-1]
|
||||
# TODO(xiongkun): support pir in the feature.
|
||||
return {param.name: param for param in partial_program._params}
|
||||
@@ -0,0 +1,342 @@
|
||||
# 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
|
||||
@@ -0,0 +1,138 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
|
||||
from .reshard_funcs.base_reshard_func import is_replicated
|
||||
from .utils import _complete_op_dist_attr
|
||||
|
||||
dist_skip_op_list = [
|
||||
"builtin.combine",
|
||||
"builtin.split",
|
||||
"cf.yield",
|
||||
"cf.tuple_push",
|
||||
"cf.tuple_pop",
|
||||
"cf.stack_create",
|
||||
"pd_op.pylayer",
|
||||
]
|
||||
|
||||
|
||||
def verify_dist_block(block):
|
||||
for op in block.ops:
|
||||
if op.name() in dist_skip_op_list:
|
||||
continue
|
||||
if op.name() == "dist_op.shard_tensor":
|
||||
raise RuntimeError("Block still contain shard_tensor_op.")
|
||||
# Note (luchang): Temp fix, remove unused parameter 'op'.
|
||||
# Will be removed in the future.
|
||||
if op.name() == "builtin.parameter":
|
||||
if op.result(0).use_empty():
|
||||
op.erase()
|
||||
continue
|
||||
|
||||
|
||||
def apply_mix2dist_pass(program, block=None):
|
||||
if block is None:
|
||||
block = program.global_block()
|
||||
deleted_ops = []
|
||||
for op in block.ops:
|
||||
for inner_block in op.blocks():
|
||||
apply_mix2dist_pass(program, block=inner_block)
|
||||
if op.name() != "dist_op.shard_tensor":
|
||||
continue
|
||||
shard_operand_value = op.operand_source(0)
|
||||
if not shard_operand_value.has_one_use():
|
||||
raise RuntimeError(
|
||||
f"shard_tensor is supposed to be called right after tensor is created, the use_count of tensor to be sharded is {shard_operand_value.use_count}, which is "
|
||||
"not Supported in right now."
|
||||
)
|
||||
shard_result_value = op.result(0)
|
||||
shard_result_value.replace_all_uses_with(shard_operand_value)
|
||||
deleted_ops.append(op)
|
||||
prev_op = shard_operand_value.get_defining_op()
|
||||
if (
|
||||
prev_op.name() == "builtin.parameter"
|
||||
or prev_op.name() == "pd_op.data"
|
||||
):
|
||||
prev_op.dist_attr = op.dist_attr
|
||||
shard_operand_value.set_type(shard_result_value.type())
|
||||
shard_operand_value.stop_gradient = shard_result_value.stop_gradient
|
||||
shard_operand_value.persistable = shard_result_value.persistable
|
||||
elif (
|
||||
prev_op.name() == "pd_op.randint"
|
||||
or prev_op.name() == "pd_op.gaussian"
|
||||
):
|
||||
mesh = shard_result_value.dist_attr().process_mesh
|
||||
# input
|
||||
shape_value = prev_op.operand_source(0)
|
||||
dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
mesh, [-1 for _ in range(len(shape_value.shape))], {}
|
||||
)
|
||||
shape_value.update_dist_attr(dist_attr)
|
||||
# op
|
||||
prev_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
mesh, [dist_attr], [shard_result_value.dist_attr()]
|
||||
)
|
||||
)
|
||||
# deal with full_int_array op
|
||||
prev_prev_op = shape_value.get_defining_op()
|
||||
prev_prev_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
mesh, [], [dist_attr]
|
||||
)
|
||||
)
|
||||
# output
|
||||
shard_operand_value.set_type(shard_result_value.type())
|
||||
shard_operand_value.stop_gradient = shard_result_value.stop_gradient
|
||||
shard_operand_value.persistable = shard_result_value.persistable
|
||||
else:
|
||||
dist_attr = shard_result_value.dist_attr()
|
||||
if not is_replicated(dist_attr):
|
||||
raise RuntimeError(
|
||||
f"{prev_op} is not support sharded by shard_tensor op in pir mode."
|
||||
)
|
||||
mesh = dist_attr.process_mesh
|
||||
ops_list = [prev_op]
|
||||
while len(ops_list) != 0:
|
||||
cur_op = ops_list.pop()
|
||||
if cur_op.dist_attr is not None:
|
||||
continue
|
||||
operand_attrs = []
|
||||
result_attrs = []
|
||||
for input in cur_op.operands_source():
|
||||
dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
mesh, [-1 for _ in range(len(input.shape))], {}
|
||||
)
|
||||
)
|
||||
operand_attrs.append(dist_attr)
|
||||
ops_list.append(input.get_defining_op())
|
||||
for result in cur_op.results():
|
||||
dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
mesh, [-1 for _ in range(len(result.shape))], {}
|
||||
)
|
||||
)
|
||||
result.update_dist_attr(dist_attr)
|
||||
result_attrs.append(dist_attr)
|
||||
cur_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
mesh, operand_attrs, result_attrs
|
||||
)
|
||||
)
|
||||
for op in deleted_ops:
|
||||
op.erase()
|
||||
_complete_op_dist_attr(program, block=block)
|
||||
verify_dist_block(block)
|
||||
@@ -0,0 +1,61 @@
|
||||
# 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 os
|
||||
|
||||
from . import ( # noqa: F401
|
||||
dist_assign,
|
||||
dist_check_finite_and_unscale,
|
||||
dist_concat,
|
||||
dist_default,
|
||||
dist_dropout,
|
||||
dist_eltwise,
|
||||
dist_embedding,
|
||||
dist_expand_as,
|
||||
dist_fill_constant_batch_size_like,
|
||||
dist_flash_attn,
|
||||
dist_fused_attention,
|
||||
dist_fused_dropout_add,
|
||||
dist_fused_feedforward,
|
||||
dist_fused_rms_norm,
|
||||
dist_fused_rope,
|
||||
dist_gather_nd,
|
||||
dist_layer_norm,
|
||||
dist_matmul,
|
||||
dist_pnorm,
|
||||
dist_reduce_sum_p,
|
||||
dist_reshape,
|
||||
dist_scale,
|
||||
dist_shape,
|
||||
dist_slice,
|
||||
dist_softmax,
|
||||
dist_split,
|
||||
dist_stack,
|
||||
dist_strided_slice,
|
||||
dist_tile,
|
||||
dist_transpose,
|
||||
dist_unsqueeze2,
|
||||
dist_update_loss_scaling,
|
||||
)
|
||||
from .common import ( # noqa: F401
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
find_compatible_distributed_operator_impls,
|
||||
find_distributed_operator_impl_container,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
|
||||
parallel_ce = os.getenv("PARALLEL_CROSS_ENTROPY")
|
||||
if parallel_ce == "true":
|
||||
from . import dist_cross_entropy # noqa: F401
|
||||
@@ -0,0 +1,838 @@
|
||||
# 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 abc
|
||||
import logging
|
||||
import warnings
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.base.log_helper import get_logger
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
|
||||
|
||||
from ..dist_attribute import OperatorDistAttr
|
||||
from ..process_group import new_process_group
|
||||
from ..utils import (
|
||||
_get_comm_group,
|
||||
_get_corresponding_rank,
|
||||
compute_compatible_dims_mapping,
|
||||
is_optimize_op,
|
||||
set_dist_op_desc_original_id,
|
||||
)
|
||||
|
||||
_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
_g_distributed_operator_impl_containers = {}
|
||||
|
||||
_g_elementwise_ops = [
|
||||
"assign",
|
||||
"elementwise",
|
||||
"gelu",
|
||||
# "dropout",
|
||||
"scale",
|
||||
"relu",
|
||||
"cast",
|
||||
# "gather",
|
||||
# "concat",
|
||||
"silu",
|
||||
"fused_softmax_mask_upper_triangle",
|
||||
]
|
||||
BACKWARD_ONLY_DIST_OPS = {'check_finite_and_unscale', 'update_loss_scaling'}
|
||||
|
||||
_gradient_sync_by_partial_ops = [
|
||||
"matmul_v2_grad",
|
||||
"elementwise_add_grad",
|
||||
"layer_norm_grad",
|
||||
"lookup_table_v2_grad",
|
||||
# "conv",
|
||||
]
|
||||
|
||||
|
||||
class ParallelMode:
|
||||
"""
|
||||
the parallel mode for communication or auxiliary operator
|
||||
"""
|
||||
|
||||
DataParallel = "auto_parallel/data_parallel"
|
||||
TensorParallel = "auto_parallel/tensor_parallel"
|
||||
PipelineParallel = "auto_parallel/pipeline_parallel"
|
||||
MoEParallel = "auto_parallel/moe_parallel"
|
||||
|
||||
|
||||
class SyncMode:
|
||||
"""
|
||||
the synchronization mode for communication or auxiliary operator
|
||||
"""
|
||||
|
||||
AmpFlagSync = "auto_parallel/amp_flag_synchronization"
|
||||
GlobalNormSync = "auto_parallel/global_norm_synchronization"
|
||||
|
||||
|
||||
def is_elementwise_op(op_type):
|
||||
if op_type in _g_elementwise_ops:
|
||||
return True
|
||||
if "elementwise" in op_type:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class DistributedOperatorImplContainer(abc.ABC):
|
||||
def __init__(self, op_type):
|
||||
self._type = op_type
|
||||
self._impls = []
|
||||
|
||||
@property
|
||||
def type(self):
|
||||
return self._type
|
||||
|
||||
@type.setter
|
||||
def type(self, op_type):
|
||||
self._type = op_type
|
||||
|
||||
@property
|
||||
def impls(self):
|
||||
return self._impls
|
||||
|
||||
def register_impl(self, dist_impl):
|
||||
assert self.type == dist_impl.type, (
|
||||
"Op type of container must be same as that of the implementation."
|
||||
)
|
||||
impl_idx = len(self.impls)
|
||||
dist_impl.idx = impl_idx
|
||||
self._impls.append(dist_impl)
|
||||
|
||||
def get_impl(self, impl_idx):
|
||||
return self._impls[impl_idx]
|
||||
|
||||
def get_input_compatible_impls(self, dist_op):
|
||||
compatible_impls = []
|
||||
for impl in self.impls:
|
||||
if impl.is_input_compatible(dist_op):
|
||||
compatible_impls.append(impl)
|
||||
return compatible_impls
|
||||
|
||||
def get_output_compatible_impls(self, dist_op):
|
||||
compatible_impls = []
|
||||
for impl in self.impls:
|
||||
if impl.is_output_compatible(dist_op):
|
||||
compatible_impls.append(impl)
|
||||
return compatible_impls
|
||||
|
||||
def get_compatible_impls(self, dist_op):
|
||||
compatible_impls = []
|
||||
for impl in self.impls:
|
||||
if impl.is_auto_compatible(dist_op):
|
||||
compatible_impls.append(impl)
|
||||
return compatible_impls
|
||||
|
||||
# (NOTE) Currently, both DistributedOperatorImplContainer and DistributedOperatorImpl have update_dims_mapping method.
|
||||
# But this method is supposed to be maintained by DistributedOperatorImplContainer, and we are ongoing adding method
|
||||
# to DistributedOperatorImplContainer and removing those in DistributedOperatorImpl.
|
||||
# @abc.abstractmethod
|
||||
def update_dims_mapping(self, dist_op):
|
||||
raise NotImplementedError("Please Implement this method in Subclass.")
|
||||
|
||||
# (NOTE) Currently we has limited DistributedOperatorImpls for an op to deal with different parallel patterns of this op.
|
||||
# This function help to choose the correct DistributedOperatorImpl based on the result from InferSPMD.
|
||||
# @abc.abstractmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
raise NotImplementedError("Please Implement this method in Subclass.")
|
||||
|
||||
|
||||
class DistributedOperatorImpl(abc.ABC):
|
||||
def __init__(self, name):
|
||||
self._name = name
|
||||
self._type = None
|
||||
self._idx = None
|
||||
self._forward_implemented = False
|
||||
self._backward_implemented = False
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
return self._name
|
||||
|
||||
@name.setter
|
||||
def name(self, name):
|
||||
self._name = name
|
||||
|
||||
@property
|
||||
def type(self):
|
||||
return self._type
|
||||
|
||||
@type.setter
|
||||
def type(self, op_type):
|
||||
self._type = op_type
|
||||
|
||||
@property
|
||||
def idx(self):
|
||||
return self._idx
|
||||
|
||||
@idx.setter
|
||||
def idx(self, impl_idx):
|
||||
self._idx = impl_idx
|
||||
|
||||
# to be deprecated
|
||||
@abc.abstractmethod
|
||||
def is_input_compatible(self, dist_op):
|
||||
raise NotImplementedError("Please Implement this method in Subclass.")
|
||||
|
||||
# to be deprecated
|
||||
@abc.abstractmethod
|
||||
def is_output_compatible(self, dist_op):
|
||||
raise NotImplementedError("Please Implement this method in Subclass.")
|
||||
|
||||
# to be deprecated
|
||||
@abc.abstractmethod
|
||||
def is_auto_compatible(self, dist_op):
|
||||
raise NotImplementedError("Please Implement this method in Subclass.")
|
||||
|
||||
@staticmethod
|
||||
@abc.abstractmethod
|
||||
def forward(dist_ctx, *args, **kwargs):
|
||||
raise NotImplementedError("Please Implement this method in Subclass.")
|
||||
|
||||
@staticmethod
|
||||
@abc.abstractmethod
|
||||
def backward(dist_ctx, *grad_outputs, **kwargs):
|
||||
raise NotImplementedError("Please Implement this method in Subclass.")
|
||||
|
||||
# to be deprecated
|
||||
def update_dims_mapping(self, dist_op):
|
||||
raise NotImplementedError("Please Implement this method in Subclass.")
|
||||
|
||||
|
||||
def register_distributed_operator_impl_container(container):
|
||||
global _g_distributed_operator_impl_containers
|
||||
_g_distributed_operator_impl_containers[container.type] = container
|
||||
|
||||
|
||||
def get_distributed_operator_impl_container(op_type):
|
||||
global _g_distributed_operator_impl_containers
|
||||
return _g_distributed_operator_impl_containers.get(op_type, None)
|
||||
|
||||
|
||||
def register_distributed_operator_impl(op_type, dist_impl):
|
||||
dist_op_impl_container = get_distributed_operator_impl_container(op_type)
|
||||
if dist_op_impl_container is not None:
|
||||
dist_impl.type = op_type
|
||||
dist_op_impl_container.register_impl(dist_impl)
|
||||
else:
|
||||
raise AssertionError(
|
||||
"Must register distributed operator registry first."
|
||||
)
|
||||
|
||||
|
||||
def find_compatible_distributed_operator_impls(dist_op, fwd=True, partial=True):
|
||||
"""
|
||||
Here just return the first compatible implementation.
|
||||
This will be improved by cost model in the future.
|
||||
"""
|
||||
op_type = dist_op.serial_op.type
|
||||
dist_op_impl_container = get_distributed_operator_impl_container(op_type)
|
||||
dist_op_eltwise_impl_container = get_distributed_operator_impl_container(
|
||||
"elementwise"
|
||||
)
|
||||
dist_op_default_impl_container = get_distributed_operator_impl_container(
|
||||
"default"
|
||||
)
|
||||
compatible_impls = []
|
||||
if partial:
|
||||
if fwd:
|
||||
# First, find impls in the corresponding container
|
||||
if dist_op_impl_container:
|
||||
compatible_impls.extend(
|
||||
dist_op_impl_container.get_input_compatible_impls(dist_op)
|
||||
)
|
||||
# Second, find impls in the elementwise container
|
||||
if dist_op_eltwise_impl_container and is_elementwise_op(op_type):
|
||||
compatible_impls.extend(
|
||||
dist_op_eltwise_impl_container.get_input_compatible_impls(
|
||||
dist_op
|
||||
)
|
||||
)
|
||||
# Third, find impls in the default container
|
||||
if dist_op_default_impl_container:
|
||||
compatible_impls.extend(
|
||||
dist_op_default_impl_container.get_input_compatible_impls(
|
||||
dist_op
|
||||
)
|
||||
)
|
||||
else:
|
||||
# First, find impls in the corresponding container
|
||||
if dist_op_impl_container:
|
||||
compatible_impls.extend(
|
||||
dist_op_impl_container.get_output_compatible_impls(dist_op)
|
||||
)
|
||||
# Second, find impls in the elementwise container
|
||||
if dist_op_eltwise_impl_container and is_elementwise_op(op_type):
|
||||
compatible_impls.extend(
|
||||
dist_op_eltwise_impl_container.get_output_compatible_impls(
|
||||
dist_op
|
||||
)
|
||||
)
|
||||
# Third, find impls in the default container
|
||||
if dist_op_default_impl_container:
|
||||
compatible_impls.extend(
|
||||
dist_op_default_impl_container.get_output_compatible_impls(
|
||||
dist_op
|
||||
)
|
||||
)
|
||||
else:
|
||||
# First, find impls in the corresponding container
|
||||
if dist_op_impl_container:
|
||||
compatible_impls.extend(
|
||||
dist_op_impl_container.get_compatible_impls(dist_op)
|
||||
)
|
||||
# Second, find impls in the elementwise container
|
||||
if dist_op_eltwise_impl_container and is_elementwise_op(op_type):
|
||||
compatible_impls.extend(
|
||||
dist_op_eltwise_impl_container.get_compatible_impls(dist_op)
|
||||
)
|
||||
# Third, find impls in the default container
|
||||
if dist_op_default_impl_container:
|
||||
compatible_impls.extend(
|
||||
dist_op_default_impl_container.get_compatible_impls(dist_op)
|
||||
)
|
||||
|
||||
if compatible_impls:
|
||||
# For now, just return the first compatible impl
|
||||
# best_compatible_impl = compatible_impls[0]
|
||||
best_compatible_impl = compatible_impls
|
||||
else:
|
||||
best_compatible_impl = None
|
||||
return best_compatible_impl
|
||||
|
||||
|
||||
def find_distributed_operator_impl_container(dist_op):
|
||||
"""
|
||||
Return a unique container for dist op.
|
||||
If not specific container found, default container will be return.
|
||||
"""
|
||||
op_type = dist_op.serial_op.type
|
||||
|
||||
# Op has a match container
|
||||
dist_op_impl_container = get_distributed_operator_impl_container(op_type)
|
||||
if dist_op_impl_container is None:
|
||||
# if op is register to elemwise spmd rule and has NO specific container implemented
|
||||
if is_elementwise_op(op_type):
|
||||
dist_op_impl_container = get_distributed_operator_impl_container(
|
||||
"elementwise"
|
||||
)
|
||||
# default container for all bottom line cases
|
||||
else:
|
||||
dist_op_impl_container = get_distributed_operator_impl_container(
|
||||
"default"
|
||||
)
|
||||
|
||||
_logger.debug(
|
||||
f"Op [{op_type}] Complete DistAttr using {type(dist_op_impl_container).__name__}"
|
||||
)
|
||||
return dist_op_impl_container
|
||||
|
||||
|
||||
def is_parameter_related(varname, block, dist_context=None):
|
||||
# TODO(zhaoyingli): maintain a dict in dist_context to record all variables which are be renamed
|
||||
if ".subprog_" in varname:
|
||||
varname = varname[: varname.index(".subprog_")]
|
||||
if ".cast_fp" in varname:
|
||||
varname = varname[: varname.index(".cast_fp")]
|
||||
if ".cast_bf" in varname:
|
||||
varname = varname[: varname.index(".cast_bf")]
|
||||
if ".quantized" in varname:
|
||||
varname = varname[: varname.index(".quantized")]
|
||||
assert block._find_var_recursive(varname), (
|
||||
f"cannot find var {varname} in cur block"
|
||||
)
|
||||
var = block._var_recursive(varname)
|
||||
# NOTE(hack method): to find the param which is resharded
|
||||
if dist_context and "@RESHARD" in varname:
|
||||
varname = varname[: varname.index("@RESHARD")]
|
||||
serial_program = dist_context.serial_main_program
|
||||
var = serial_program.global_block()._find_var_recursive(varname)
|
||||
if var is None:
|
||||
return False
|
||||
# NOTE(liym27): when Y_var is not a parameter, but Y_var is resharded by a parameter.
|
||||
elif "reshard_api" in varname:
|
||||
for op in block.ops:
|
||||
if op.type == "assign" and varname in op.output("Out"):
|
||||
in_varname = op.input("X")[0]
|
||||
var = block._find_var_recursive(in_varname)
|
||||
if var is not None and var.is_parameter:
|
||||
return True
|
||||
return var.is_parameter
|
||||
|
||||
|
||||
def infer_shape(block, src_var, src_var_dist_attr, op_input_dist_attr):
|
||||
var_shape = block._var_recursive(src_var.name).shape
|
||||
var_topology = src_var_dist_attr.process_mesh.shape
|
||||
var_dims_mapping = src_var_dist_attr.dims_mapping
|
||||
|
||||
complete_shape = []
|
||||
for idx, shape in enumerate(var_shape):
|
||||
if var_dims_mapping[idx] == -1:
|
||||
complete_shape.append(shape)
|
||||
else:
|
||||
new_shape = shape * var_topology[var_dims_mapping[idx]]
|
||||
complete_shape.append(new_shape)
|
||||
|
||||
exact_shape = []
|
||||
input_topology = op_input_dist_attr.process_mesh.shape
|
||||
input_dims_mapping = op_input_dist_attr.dims_mapping
|
||||
for idx, shape in enumerate(complete_shape):
|
||||
if input_dims_mapping[idx] == -1:
|
||||
exact_shape.append(shape)
|
||||
else:
|
||||
new_shape = shape // input_topology[input_dims_mapping[idx]]
|
||||
exact_shape.append(new_shape)
|
||||
|
||||
return exact_shape
|
||||
|
||||
|
||||
def set_comm_op_dist_attr_for_program(
|
||||
new_op, process_mesh, tensor_dist_attr, ctx, **kwargs
|
||||
):
|
||||
assert process_mesh is not None
|
||||
assert tensor_dist_attr is not None
|
||||
|
||||
new_op_dist_attr = OperatorDistAttr()
|
||||
new_op_dist_attr.process_mesh = process_mesh
|
||||
if "chunk_id" in kwargs:
|
||||
new_op_dist_attr.chunk_id = kwargs["chunk_id"]
|
||||
for input_varname in new_op.desc.input_arg_names():
|
||||
new_op_dist_attr.set_input_dist_attr(input_varname, tensor_dist_attr)
|
||||
for output_varname in new_op.desc.output_arg_names():
|
||||
new_op_dist_attr.set_output_dist_attr(output_varname, tensor_dist_attr)
|
||||
ctx.set_op_dist_attr_for_program(new_op, new_op_dist_attr)
|
||||
|
||||
|
||||
def naive_copy_op_dist_attr_for_program(new_op, ref_op, ctx):
|
||||
ref_dist_attr = ctx.get_op_dist_attr_for_program(ref_op)
|
||||
new_op_dist_attr = OperatorDistAttr()
|
||||
new_op_dist_attr.process_mesh = ref_dist_attr.process_mesh
|
||||
new_op_dist_attr.impl_type = ref_dist_attr.impl_type
|
||||
new_op_dist_attr.impl_idx = ref_dist_attr.impl_idx
|
||||
new_op_dist_attr.chunk_id = ref_dist_attr.chunk_id
|
||||
|
||||
for input_name in ref_op.input_names:
|
||||
assert input_name in new_op.input_names
|
||||
assert len(ref_op.input(input_name)) == 1
|
||||
assert len(new_op.input(input_name)) == 1
|
||||
|
||||
ref_tensor_dist_attr = ref_dist_attr.get_input_dist_attr(
|
||||
ref_op.input(input_name)[0]
|
||||
)
|
||||
new_op_dist_attr.set_input_dist_attr(
|
||||
new_op.input(input_name)[0], ref_tensor_dist_attr
|
||||
)
|
||||
|
||||
for output_name in ref_op.output_names:
|
||||
assert output_name in new_op.output_names
|
||||
assert len(ref_op.output(output_name)) == 1
|
||||
assert len(new_op.output(output_name)) == 1
|
||||
|
||||
ref_tensor_dist_attr = ref_dist_attr.get_output_dist_attr(
|
||||
ref_op.output(output_name)[0]
|
||||
)
|
||||
new_op_dist_attr.set_output_dist_attr(
|
||||
new_op.output(output_name)[0], ref_tensor_dist_attr
|
||||
)
|
||||
|
||||
ctx.set_op_dist_attr_for_program(new_op, new_op_dist_attr)
|
||||
|
||||
|
||||
def get_data_parallel_group(dist_ctx, op, act_grad_names, rank):
|
||||
"""
|
||||
deduce the data parallel communication group for current operator.
|
||||
|
||||
Args:
|
||||
dist_ctx (DistributedContext): dist context.
|
||||
op (Operator): the current (backward) operator which might need.
|
||||
act_grad_names (list): list of input activation grads variable name to the current operator.
|
||||
rank (int): global ranks index for current process.
|
||||
"""
|
||||
dp_group = None
|
||||
|
||||
op_dist_attr = dist_ctx.get_op_dist_attr_for_program(op)
|
||||
process_mesh = op_dist_attr.process_mesh
|
||||
mesh_shape = process_mesh.shape
|
||||
# FIXME Hack for Pipeline Parallelism where the current operator
|
||||
# not belong to the mesh the current rank belong to.
|
||||
if rank not in process_mesh.process_ids:
|
||||
rank = _get_corresponding_rank(dist_ctx, process_mesh, rank)
|
||||
|
||||
for var_name in act_grad_names:
|
||||
var_dim_mapping = op_dist_attr.get_input_dims_mapping(var_name)
|
||||
# consider that the variable's shape is [], which is 0-D
|
||||
# TODO utilize the batch_dim attr instead of "0" in future
|
||||
batch_size_axis = var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
|
||||
|
||||
if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh.process_ids,
|
||||
process_mesh.shape,
|
||||
batch_size_axis,
|
||||
rank,
|
||||
)
|
||||
dp_group = new_process_group(group_ranks)
|
||||
break
|
||||
if dp_group is not None:
|
||||
return [dp_group]
|
||||
else:
|
||||
return []
|
||||
|
||||
|
||||
def sync_and_scale_gradients(dist_ctx, op, groups, allreduce_var_names):
|
||||
"""
|
||||
insert the allreduce and scale ops for gradients of model
|
||||
parameters for operator in data parallelism.
|
||||
|
||||
Args:
|
||||
dist_ctx (DistributedContext): dist context.
|
||||
op (Operator): the current (backward) operator which might need.
|
||||
allreduce_var_names (list): list of the parameter's grads variable name in the current operator output.
|
||||
"""
|
||||
|
||||
op_dist_attr = dist_ctx.get_op_dist_attr_for_program(op)
|
||||
process_mesh = op_dist_attr.process_mesh
|
||||
chunk_id = op_dist_attr.chunk_id
|
||||
dist_op_context = dist_ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
|
||||
reduce_type = dist.ReduceOp.SUM
|
||||
need_scale = dist_ctx.gradient_scale
|
||||
|
||||
for group in groups:
|
||||
group_size = len(group.ranks)
|
||||
|
||||
for var_name in allreduce_var_names:
|
||||
added_ops = []
|
||||
grad_var = main_block.var(var_name)
|
||||
allreduce_op = main_block.append_op(
|
||||
type='all_reduce',
|
||||
inputs={'x': [grad_var]},
|
||||
outputs={'out': [grad_var]},
|
||||
attrs={
|
||||
'ring_id': group.id,
|
||||
'reduce_type': reduce_type,
|
||||
OP_ROLE_KEY: OpRole.Backward,
|
||||
},
|
||||
)
|
||||
allreduce_op._set_attr(
|
||||
'op_namescope', '/' + ParallelMode.DataParallel
|
||||
)
|
||||
added_ops.append(allreduce_op)
|
||||
|
||||
if need_scale:
|
||||
scale_op = main_block.append_op(
|
||||
type='scale',
|
||||
inputs={'X': grad_var},
|
||||
outputs={'Out': grad_var},
|
||||
attrs={
|
||||
'scale': 1.0 / group_size,
|
||||
OP_ROLE_KEY: OpRole.Backward,
|
||||
},
|
||||
)
|
||||
scale_op._set_attr(
|
||||
'op_namescope', '/' + ParallelMode.DataParallel
|
||||
)
|
||||
added_ops.append(scale_op)
|
||||
|
||||
dims_mapping = op_dist_attr.get_output_dims_mapping(grad_var.name)
|
||||
assert dims_mapping is not None, (
|
||||
f"Unexpected: dims_mapping of output [{grad_var.name}] of op [{op_dist_attr.op_type}] is None"
|
||||
)
|
||||
# NOTE auxiliary op's dist attr should follow dist_op not dist_tensor
|
||||
for new_op in added_ops:
|
||||
new_op_attr = OperatorDistAttr()
|
||||
new_op_attr.process_mesh = process_mesh
|
||||
new_op_attr.chunk_id = chunk_id
|
||||
new_op_attr.set_output_dims_mapping(grad_var.name, dims_mapping)
|
||||
new_op_attr.set_input_dims_mapping(grad_var.name, dims_mapping)
|
||||
dist_ctx.set_op_dist_attr_for_program(new_op, new_op_attr)
|
||||
|
||||
|
||||
def get_partial_groups(dist_ctx, op, out_grad_names, rank):
|
||||
"""
|
||||
deduce the partial communication group for current operator output vars.
|
||||
|
||||
Args:
|
||||
dist_ctx (DistributedContext): dist context.
|
||||
op (Operator): the current (backward) operator which might need.
|
||||
out_grad_names (list): list of the output parameter's grads variable name of the current operator.
|
||||
rank (int): global ranks index for current process.
|
||||
"""
|
||||
op_dist_attr = dist_ctx.get_op_dist_attr_for_program(op)
|
||||
process_mesh = op_dist_attr.process_mesh
|
||||
mesh_shape = process_mesh.shape
|
||||
|
||||
groups = []
|
||||
|
||||
partial_dims = None
|
||||
for var_name in out_grad_names:
|
||||
var_dist_attr = op_dist_attr.get_output_dist_attr(var_name)
|
||||
if partial_dims is None:
|
||||
partial_dims = var_dist_attr._partial_dims()
|
||||
else:
|
||||
assert partial_dims == var_dist_attr._partial_dims(), (
|
||||
f"Partial dims of outputs {out_grad_names} of op [{op.type}] is not consistent"
|
||||
)
|
||||
|
||||
partial_dims = list(partial_dims)
|
||||
partial_dims.sort()
|
||||
|
||||
# FIXME Hack for Pipeline Parallelism where the current operator
|
||||
# not belong to the mesh the current rank belong to.
|
||||
if rank not in process_mesh.process_ids:
|
||||
rank = _get_corresponding_rank(dist_ctx, process_mesh, rank)
|
||||
|
||||
for dim in partial_dims:
|
||||
if mesh_shape[dim] > 1:
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh.process_ids,
|
||||
process_mesh.shape,
|
||||
dim,
|
||||
rank,
|
||||
)
|
||||
groups.append(new_process_group(group_ranks))
|
||||
|
||||
return groups
|
||||
|
||||
|
||||
def gradient_synchronization(
|
||||
dist_ctx, op, act_grad_names, out_grad_names, rank
|
||||
):
|
||||
"""
|
||||
conduct the allreduce and scaling for gradients of model
|
||||
parameters for operator in parallelism train.
|
||||
|
||||
Args:
|
||||
dist_ctx (DistributedContext): dist context.
|
||||
op (Operator): the current (backward) operator which might need.
|
||||
act_grad_names (list): list of input activation grads variable name to the current operator.
|
||||
out_grad_names (list): list of the output parameter's grads variable name of the current operator.
|
||||
rank (int): global ranks index for current process.
|
||||
"""
|
||||
|
||||
if not is_in_backward_phase(dist_ctx):
|
||||
return
|
||||
|
||||
if (
|
||||
is_optimize_op(op)
|
||||
or len(act_grad_names) == 0
|
||||
or len(out_grad_names) == 0
|
||||
):
|
||||
return
|
||||
|
||||
if op.type in _gradient_sync_by_partial_ops:
|
||||
sync_groups = get_partial_groups(dist_ctx, op, out_grad_names, rank)
|
||||
# NOTE we reverse the following old branch to support operators (e.g. fuse operators) that haven't been adopted for partial inferspmd,
|
||||
# and remove this branch after all operators are adopted for partial inferspmd.
|
||||
else:
|
||||
sync_groups = get_data_parallel_group(
|
||||
dist_ctx, op, act_grad_names, rank
|
||||
)
|
||||
|
||||
if len(sync_groups) < 1:
|
||||
return
|
||||
|
||||
sync_and_scale_gradients(dist_ctx, op, sync_groups, out_grad_names)
|
||||
|
||||
|
||||
def is_data_parallel_scale_op(op):
|
||||
return (
|
||||
op.type == "scale"
|
||||
and op.desc.has_attr("op_namescope")
|
||||
and ParallelMode.DataParallel in op.desc.attr("op_namescope")
|
||||
)
|
||||
|
||||
|
||||
def is_data_parallel_reduce_op(op):
|
||||
is_allreduce_op = op.type in [
|
||||
"c_allreduce_sum",
|
||||
"c_allreduce_avg",
|
||||
]
|
||||
is_all_reduce_op = op.type == "all_reduce" and op.desc.attr(
|
||||
"reduce_type"
|
||||
) in [
|
||||
dist.ReduceOp.SUM,
|
||||
dist.ReduceOp.AVG,
|
||||
]
|
||||
is_reduce_op = op.type == "reduce" and op.desc.attr("reduce_type") in [
|
||||
dist.ReduceOp.SUM,
|
||||
dist.ReduceOp.AVG,
|
||||
]
|
||||
return (
|
||||
(is_allreduce_op or is_all_reduce_op or is_reduce_op)
|
||||
and op.desc.has_attr("op_namescope")
|
||||
and ParallelMode.DataParallel in op.desc.attr("op_namescope")
|
||||
)
|
||||
|
||||
|
||||
def is_amp_flag_sync_op(op):
|
||||
return (
|
||||
op.type == "all_reduce"
|
||||
and op.desc.attr("op_type") == paddle.distributed.ReduceOp.MAX
|
||||
and op.desc.has_attr("op_namescope")
|
||||
and SyncMode.AmpFlagSync in op.desc.attr("op_namescope")
|
||||
)
|
||||
|
||||
|
||||
def is_global_norm_sync_op(op):
|
||||
return (
|
||||
op.type == "all_reduce"
|
||||
and op.desc.attr("reduce_type") == dist.ReduceOp.SUM
|
||||
and op.desc.has_attr("op_namescope")
|
||||
and SyncMode.GlobalNormSync in op.desc.attr("op_namescope")
|
||||
)
|
||||
|
||||
|
||||
def is_in_backward_phase(dist_ctx):
|
||||
# NOTE currently high-order differential in Paddle dose NOT distinguish gradient computation operators
|
||||
# in Forward phase and operators in Backward phase (both with op_role=1), which will mislead
|
||||
# auto parallel to add gradient synchronization for gradient computation operators in Forward phase.
|
||||
# we use this FLAG to distinguish these two phases temporarily.
|
||||
|
||||
return dist_ctx.dist_op_context.in_backward_phase()
|
||||
|
||||
|
||||
def merge_forward_backward_dims_mapping(fw_results, bw_results):
|
||||
flatten_fw_inputs = paddle.utils.flatten(fw_results[0])
|
||||
flatten_fw_outputs = paddle.utils.flatten(fw_results[1])
|
||||
flatten_bw_inputs = paddle.utils.flatten(bw_results[0])
|
||||
flatten_bw_outputs = paddle.utils.flatten(bw_results[1])
|
||||
ninputs = len(flatten_fw_inputs)
|
||||
noutputs = len(flatten_fw_outputs)
|
||||
inferred_input_dims_mappings = []
|
||||
inferred_output_dims_mappings = []
|
||||
|
||||
for i in range(ninputs):
|
||||
compatible_dims_mapping = compute_compatible_dims_mapping(
|
||||
[
|
||||
flatten_fw_inputs[i].dims_mapping,
|
||||
flatten_bw_inputs[i].dims_mapping,
|
||||
]
|
||||
)
|
||||
inferred_input_dims_mappings.append(compatible_dims_mapping)
|
||||
|
||||
for i in range(noutputs):
|
||||
compatible_dims_mapping = compute_compatible_dims_mapping(
|
||||
[
|
||||
flatten_fw_outputs[i].dims_mapping,
|
||||
flatten_bw_outputs[i].dims_mapping,
|
||||
]
|
||||
)
|
||||
inferred_output_dims_mappings.append(compatible_dims_mapping)
|
||||
return inferred_input_dims_mappings, inferred_output_dims_mappings
|
||||
|
||||
|
||||
def update_op_dims_mapping(
|
||||
dist_op, input_arg_names, output_arg_names, fw_results, bw_results
|
||||
):
|
||||
(
|
||||
inferred_input_dims_mappings,
|
||||
inferred_output_dims_mappings,
|
||||
) = merge_forward_backward_dims_mapping(fw_results, bw_results)
|
||||
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
changed = False
|
||||
if len(input_arg_names) != len(inferred_input_dims_mappings):
|
||||
warnings.warn(
|
||||
f"dims mapping is NOT Match, inferred [{len(inferred_input_dims_mappings)}], original: [{len(input_arg_names)}]; dist op: [{dist_op}]"
|
||||
)
|
||||
if len(output_arg_names) != len(inferred_output_dims_mappings):
|
||||
warnings.warn(
|
||||
f"dims mapping is NOT Match, inferred [{len(inferred_output_dims_mappings)}], original: [{len(output_arg_names)}]; dist op: [{dist_op}]"
|
||||
)
|
||||
|
||||
for i in range(len(input_arg_names)):
|
||||
original_dims_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
input_arg_names[i]
|
||||
)
|
||||
inferred_dims_mapping = inferred_input_dims_mappings[i]
|
||||
if (inferred_dims_mapping is not None) and (
|
||||
original_dims_mapping != inferred_dims_mapping
|
||||
):
|
||||
_logger.debug(
|
||||
f"Changed: Op [{dist_op.serial_op.type}], name [{input_arg_names[i]}], Original [{original_dims_mapping}], Inferred [{inferred_dims_mapping}]"
|
||||
)
|
||||
changed = True
|
||||
op_dist_attr.set_input_dims_mapping(
|
||||
input_arg_names[i], inferred_dims_mapping
|
||||
)
|
||||
# TODO support partial for inputs
|
||||
|
||||
for i in range(len(output_arg_names)):
|
||||
original_dims_mapping = op_dist_attr.get_output_dims_mapping(
|
||||
output_arg_names[i]
|
||||
)
|
||||
inferred_dims_mapping = inferred_output_dims_mappings[i]
|
||||
if (inferred_dims_mapping is not None) and (
|
||||
original_dims_mapping != inferred_dims_mapping
|
||||
):
|
||||
_logger.debug(
|
||||
f"Changed: Op [{dist_op.serial_op.type}], name [{output_arg_names[i]}], Original [{original_dims_mapping}], Inferred [{inferred_dims_mapping}]"
|
||||
)
|
||||
changed = True
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
output_arg_names[i], inferred_dims_mapping
|
||||
)
|
||||
|
||||
# NOTE in partial stage-I, we infer partial for output in infer_forward only
|
||||
output_dist_attr = op_dist_attr.get_output_dist_attr(
|
||||
output_arg_names[i]
|
||||
)
|
||||
output_idx = output_arg_names.index(output_arg_names[i])
|
||||
if (
|
||||
fw_results[1][output_idx]._partial_dims()
|
||||
!= output_dist_attr._partial_dims()
|
||||
):
|
||||
# _logger.info(
|
||||
# "Changed: Op [{}], tensor name [{}], Original partial on [{}], Inferred partial on [{}]".format(
|
||||
# dist_op.serial_op.type,
|
||||
# output_arg_names[i],
|
||||
# output_dist_attr._partial_dims(),
|
||||
# fw_results[1][output_idx]._partial_dims(),
|
||||
# )
|
||||
# )
|
||||
output_dist_attr._clean_partial_status()
|
||||
output_dist_attr._set_partial_dims(
|
||||
list(fw_results[1][0]._partial_dims())
|
||||
)
|
||||
changed = True
|
||||
|
||||
return changed
|
||||
|
||||
|
||||
def get_default_distributed_operator_impl():
|
||||
dist_op_default_impl_container = get_distributed_operator_impl_container(
|
||||
"default"
|
||||
)
|
||||
num_impls = len(dist_op_default_impl_container.impls)
|
||||
assert num_impls == 1, f"Default dist op has [{num_impls}] impls"
|
||||
return dist_op_default_impl_container.get_impl(0)
|
||||
|
||||
|
||||
def copy_op_without_infer_shape(src_op, block, ctx, varname_kwargs):
|
||||
new_op = block.append_op(type='nop')
|
||||
new_op_desc = new_op.desc
|
||||
new_op_desc.copy_from(src_op.desc)
|
||||
set_dist_op_desc_original_id(new_op_desc, src_op.desc, ctx)
|
||||
for input_name in src_op.desc.input_names():
|
||||
new_op_desc.set_input(input_name, varname_kwargs[input_name])
|
||||
for output_name in src_op.desc.output_names():
|
||||
new_op_desc.set_output(output_name, varname_kwargs[output_name])
|
||||
# TODO: should we add a new dist attr for the new op here?
|
||||
return new_op
|
||||
@@ -0,0 +1,90 @@
|
||||
# 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.
|
||||
|
||||
from ..utils import compute_compatible_and_update_dim_mapping
|
||||
from .common import DistributedOperatorImpl, DistributedOperatorImplContainer
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedAssign(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
# TODO remove assign dist op
|
||||
# register_distributed_operator_impl_container(DistributedAssign("assign"))
|
||||
|
||||
|
||||
class DistributedAssignImpl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
if x_dims_mapping != out_dims_mapping:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
for i in range(len(x_dims_mapping)):
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[x_dims_mapping, out_dims_mapping], [i, i]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
# register_distributed_operator_impl("assign", DistributedAssignImpl("assign"))
|
||||
+206
@@ -0,0 +1,206 @@
|
||||
# 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 paddle
|
||||
from paddle.distributed.auto_parallel.static.process_group import (
|
||||
get_world_process_group,
|
||||
)
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
|
||||
from paddle.framework import core
|
||||
|
||||
from ..dist_attribute import OperatorDistAttr
|
||||
from ..process_group import new_process_group
|
||||
from ..utils import set_dist_op_desc_original_id, set_var_dist_attr
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
SyncMode,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
|
||||
world_process_group = get_world_process_group()
|
||||
|
||||
|
||||
class DistributedCheckFiniteAndUnscale(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedCheckFiniteAndUnscale("check_finite_and_unscale")
|
||||
)
|
||||
|
||||
|
||||
class DistributedCheckFiniteAndUnscaleImpl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._name = name
|
||||
self._forward_implemented = False
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
raise RuntimeError(
|
||||
"DistributedCheckFiniteAndUnscaleImpl's is_input_compatible should not be called !"
|
||||
)
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
raise RuntimeError(
|
||||
"DistributedCheckFiniteAndUnscaleImpl's is_output_compatible should not be called !"
|
||||
)
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
raise RuntimeError(
|
||||
"DistributedCheckFiniteAndUnscaleImpl's is_auto_compatible should not be called !"
|
||||
)
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
raise RuntimeError(
|
||||
"DistributedCheckFiniteAndUnscaleImpl's update_dims_mapping should not be called !"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
raise RuntimeError(
|
||||
"DistributedCheckFiniteAndUnscaleImpl's forward should not be called !"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
# by now the backward function only insert the gradient allreduce for dist op itself
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.main_block
|
||||
backward_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
|
||||
assert dist_attr is not None, (
|
||||
f"backward op [{backward_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
assert rank_id in dist_attr.process_mesh.process_ids
|
||||
|
||||
assert 'X' in kwargs, "input [{}] is not given".format('X')
|
||||
assert 'Scale' in kwargs, "input [{}] is not given".format('Scale')
|
||||
assert 'Out' in kwargs, "input [{}] is not given".format('Out')
|
||||
assert 'FoundInfinite' in kwargs, "output [{}] is not given".format(
|
||||
'FoundInfinite'
|
||||
)
|
||||
|
||||
assert len(kwargs['Scale']) == 1, (
|
||||
"check_finite_and_unscale input Scale take 1 variable but got {}".format(
|
||||
kwargs['Scale']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['FoundInfinite']) == 1, (
|
||||
"check_finite_and_unscale input FoundInfinite take 1 variable but got {}".format(
|
||||
kwargs['FoundInfinite']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['X']) == len(kwargs['Out']), (
|
||||
"check_finite_and_unscale got [{}] X and [{}] Out, which are supposed to be equal".format(
|
||||
len(kwargs['X']), len(kwargs['Out'])
|
||||
)
|
||||
)
|
||||
|
||||
filter_vars = []
|
||||
for varname in kwargs['X']:
|
||||
if (
|
||||
rank_id
|
||||
in ctx.get_tensor_dist_attr_for_program(
|
||||
main_block._var_recursive(varname)
|
||||
).process_mesh.process_ids
|
||||
):
|
||||
filter_vars.append(varname)
|
||||
|
||||
# replicate op in dist program
|
||||
dist_op = main_block.append_op(type='nop')
|
||||
dist_op_desc = dist_op.desc
|
||||
dist_op_desc.copy_from(backward_op.desc)
|
||||
set_dist_op_desc_original_id(dist_op_desc, backward_op.desc, ctx)
|
||||
dist_op_desc.set_input('X', filter_vars)
|
||||
dist_op_desc.set_output('Out', filter_vars)
|
||||
# TODO: should we add a new dist attr for the new op here?
|
||||
|
||||
# sync result
|
||||
group = new_process_group(world_process_group.ranks)
|
||||
|
||||
inf_var = main_block._var_recursive(kwargs['FoundInfinite'][0])
|
||||
inf_var_int32 = main_block.create_var(
|
||||
name=inf_var.name + "@cast_int32",
|
||||
shape=inf_var.shape,
|
||||
dtype=core.VarDesc.VarType.INT32,
|
||||
)
|
||||
set_var_dist_attr(
|
||||
ctx,
|
||||
inf_var_int32,
|
||||
ctx.get_tensor_dist_attr_for_program(inf_var).dims_mapping,
|
||||
ctx.get_tensor_dist_attr_for_program(inf_var).process_mesh,
|
||||
)
|
||||
cast_op1 = main_block.append_op(
|
||||
type='cast',
|
||||
inputs={'X': inf_var},
|
||||
outputs={'Out': inf_var_int32},
|
||||
attrs={
|
||||
"in_dtype": inf_var.dtype,
|
||||
"out_dtype": inf_var_int32.dtype,
|
||||
OP_ROLE_KEY: OpRole.Optimize,
|
||||
},
|
||||
)
|
||||
allreduce_op = main_block.append_op(
|
||||
type='all_reduce',
|
||||
inputs={'x': inf_var_int32},
|
||||
outputs={'out': inf_var_int32},
|
||||
attrs={
|
||||
'ring_id': group.id,
|
||||
'op_type': paddle.distributed.ReduceOp.MAX,
|
||||
OP_ROLE_KEY: OpRole.Optimize,
|
||||
},
|
||||
)
|
||||
allreduce_op._set_attr('op_namescope', '/' + SyncMode.AmpFlagSync)
|
||||
cast_op2 = main_block.append_op(
|
||||
type='cast',
|
||||
inputs={'X': inf_var_int32},
|
||||
outputs={'Out': inf_var},
|
||||
attrs={
|
||||
"in_dtype": inf_var_int32.dtype,
|
||||
"out_dtype": inf_var.dtype,
|
||||
OP_ROLE_KEY: OpRole.Optimize,
|
||||
},
|
||||
)
|
||||
|
||||
for op in [cast_op1, allreduce_op, cast_op2]:
|
||||
new_op_dist_attr = OperatorDistAttr()
|
||||
for varname in op.input_arg_names:
|
||||
var_dist_attr = ctx.get_tensor_dist_attr_for_program(
|
||||
main_block._var_recursive(varname)
|
||||
)
|
||||
assert var_dist_attr is not None
|
||||
new_op_dist_attr.set_input_dims_mapping(
|
||||
varname, var_dist_attr.dims_mapping
|
||||
)
|
||||
for varname in op.output_arg_names:
|
||||
var_dist_attr = ctx.get_tensor_dist_attr_for_program(
|
||||
main_block._var_recursive(varname)
|
||||
)
|
||||
new_op_dist_attr.set_output_dims_mapping(
|
||||
varname, var_dist_attr.dims_mapping
|
||||
)
|
||||
new_op_dist_attr.process_mesh = var_dist_attr.process_mesh
|
||||
ctx.set_op_dist_attr_for_program(op, new_op_dist_attr)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"check_finite_and_unscale",
|
||||
DistributedCheckFiniteAndUnscaleImpl("check_finite_and_unscale"),
|
||||
)
|
||||
@@ -0,0 +1,76 @@
|
||||
# Copyright (c) 2023 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
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..utils import get_dist_tensor_spec
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
|
||||
class DistributedConcat(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
|
||||
axis_tensor = op_desc.input('AxisTensor')
|
||||
assert len(axis_tensor) == 0, (
|
||||
"Please use axis attr instead of AxisTensor"
|
||||
)
|
||||
|
||||
input_arg_names = op_desc.input_arg_names()
|
||||
output_arg_names = op_desc.output_arg_names()
|
||||
axis = op_desc.attr('axis')
|
||||
|
||||
input_specs = []
|
||||
for name in input_arg_names:
|
||||
input_specs.append(get_dist_tensor_spec(dist_op, name))
|
||||
output_spec = get_dist_tensor_spec(dist_op, output_arg_names[0], False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("concat")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(input_specs, axis)
|
||||
bw_results = rule.infer_backward(input_specs, output_spec, axis)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op,
|
||||
input_arg_names,
|
||||
output_arg_names,
|
||||
fw_results,
|
||||
bw_results,
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedConcat("concat"))
|
||||
@@ -0,0 +1,528 @@
|
||||
# Copyright (c) 2023 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
|
||||
|
||||
from paddle.common_ops_import import check_variable_and_dtype
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..dist_attribute import OperatorDistAttr
|
||||
from ..process_group import new_process_group
|
||||
from ..utils import (
|
||||
_get_comm_group,
|
||||
_get_corresponding_rank,
|
||||
get_dist_tensor_spec,
|
||||
is_dim_shard,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
ParallelMode,
|
||||
copy_op_without_infer_shape,
|
||||
naive_copy_op_dist_attr_for_program,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
|
||||
class DistributedCrossEntropy(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
|
||||
logits_name = op_desc.input('Logits')[0]
|
||||
label_name = op_desc.input('Label')[0]
|
||||
loss_name = op_desc.output('Loss')[0]
|
||||
softmax_name = op_desc.output('Softmax')[0]
|
||||
|
||||
soft_label = op_desc.attr('soft_label')
|
||||
ignore_index = op_desc.attr('ignore_index')
|
||||
numeric_stable_mode = op_desc.attr('numeric_stable_mode')
|
||||
axis = op_desc.attr('axis')
|
||||
|
||||
logits_spec = get_dist_tensor_spec(dist_op, logits_name)
|
||||
label_spec = get_dist_tensor_spec(dist_op, label_name)
|
||||
loss_spec = get_dist_tensor_spec(dist_op, loss_name, False)
|
||||
softmax_spec = get_dist_tensor_spec(dist_op, softmax_name, False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("softmax_with_cross_entropy")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(
|
||||
logits_spec,
|
||||
label_spec,
|
||||
soft_label,
|
||||
True,
|
||||
numeric_stable_mode,
|
||||
ignore_index,
|
||||
axis,
|
||||
)
|
||||
bw_results = rule.infer_backward(
|
||||
logits_spec,
|
||||
label_spec,
|
||||
softmax_spec,
|
||||
loss_spec,
|
||||
soft_label,
|
||||
True,
|
||||
numeric_stable_mode,
|
||||
ignore_index,
|
||||
axis,
|
||||
)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op,
|
||||
[logits_name, label_name],
|
||||
[softmax_name, loss_name],
|
||||
fw_results,
|
||||
bw_results,
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
op_dist_attr.impl_type = op_desc.type()
|
||||
|
||||
logits_name = op_desc.input('Logits')[0]
|
||||
|
||||
soft_label = op_desc.attr('soft_label')
|
||||
axis = op_desc.attr('axis')
|
||||
|
||||
logits_dims_mapping = copy.deepcopy(
|
||||
op_dist_attr.get_input_dims_mapping(logits_name)
|
||||
)
|
||||
logits_ndim = len(logits_dims_mapping)
|
||||
axis = axis + logits_ndim if axis < 0 else axis
|
||||
|
||||
if is_dim_shard(logits_dims_mapping[axis]):
|
||||
assert soft_label is False, (
|
||||
"parallel_cross_entropy does not support soft_label now."
|
||||
)
|
||||
assert axis == logits_ndim - 1, (
|
||||
"parallel_cross_entropy can only support shard on the last dim now."
|
||||
)
|
||||
op_dist_attr.impl_idx = 1
|
||||
else:
|
||||
op_dist_attr.impl_idx = 0
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedCrossEntropy("softmax_with_cross_entropy")
|
||||
)
|
||||
|
||||
|
||||
# The softmax_norm axis is not sharded
|
||||
class DistributedCrossEntropyImpl0(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
"""
|
||||
kwargs: inputname_mapping & outputname_mapping
|
||||
"""
|
||||
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
assert op_dist_attr is not None, (
|
||||
f"forward op [{src_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
# check validation of inputs / outputs
|
||||
assert 'Logits' in kwargs, "input [Logits] is not given"
|
||||
assert 'Label' in kwargs, "input [Label] is not given"
|
||||
assert 'Loss' in kwargs, "output [Loss] is not given"
|
||||
assert 'Softmax' in kwargs, "output [Softmax] is not given"
|
||||
|
||||
assert len(kwargs['Logits']) == 1, (
|
||||
"input [Logits] take 1 variable but got {}".format(kwargs['Logits'])
|
||||
)
|
||||
assert len(kwargs['Label']) == 1, (
|
||||
"input [Label] take 1 variable but got {}".format(kwargs['Label'])
|
||||
)
|
||||
|
||||
logits_var = main_block._var_recursive(kwargs['Logits'][0])
|
||||
label_var = main_block._var_recursive(kwargs['Label'][0])
|
||||
loss_var = main_block._var_recursive(kwargs['Loss'][0])
|
||||
softmax_var = main_block._var_recursive(kwargs['Softmax'][0])
|
||||
|
||||
# got dist attribute info
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
process_mesh_group = op_dist_attr.process_mesh.process_ids
|
||||
|
||||
# FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
|
||||
if rank_id not in process_mesh_group:
|
||||
rank_id = _get_corresponding_rank(
|
||||
ctx, op_dist_attr.process_mesh, rank_id
|
||||
)
|
||||
|
||||
check_variable_and_dtype(
|
||||
logits_var,
|
||||
'input',
|
||||
['bfloat16', 'float16', 'float32', 'float64'],
|
||||
'cross_entropy_with_softmax',
|
||||
)
|
||||
check_variable_and_dtype(
|
||||
label_var,
|
||||
'input',
|
||||
['int32', 'int64', 'float32', 'float64'],
|
||||
'cross_entropy_with_softmax',
|
||||
)
|
||||
check_variable_and_dtype(
|
||||
loss_var,
|
||||
'output',
|
||||
['bfloat16', 'float16', 'float32', 'float64'],
|
||||
'cross_entropy_with_softmax',
|
||||
)
|
||||
check_variable_and_dtype(
|
||||
softmax_var,
|
||||
'output',
|
||||
['bfloat16', 'float16', 'float32', 'float64'],
|
||||
'cross_entropy_with_softmax',
|
||||
)
|
||||
copy_op_without_infer_shape(src_op, main_block, ctx, kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
backward_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
|
||||
|
||||
assert op_dist_attr is not None, (
|
||||
f"backward op [{backward_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
# check validation of inputs / outputs
|
||||
assert 'Softmax' in kwargs, "input [Logits] is not given"
|
||||
assert 'Label' in kwargs, "input [Label] is not given"
|
||||
assert 'Loss@GRAD' in kwargs, "input [Loss@GRAD] is not given"
|
||||
assert 'Logits@GRAD' in kwargs, "output [Logits@GRAD] is not given"
|
||||
|
||||
assert len(kwargs['Softmax']) == 1, (
|
||||
"input [Softmax] take 1 variable but got {}".format(
|
||||
kwargs['Softmax']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['Label']) == 1, (
|
||||
"input [Label] take 1 variable but got {}".format(kwargs['Label'])
|
||||
)
|
||||
assert len(kwargs['Loss@GRAD']) == 1, (
|
||||
"input [Loss@GRAD] take 1 variable but got {}".format(kwargs['Out'])
|
||||
)
|
||||
assert len(kwargs['Logits@GRAD']) == 1, (
|
||||
"output [Logits@GRAD] take 1 variable but got {}".format(
|
||||
kwargs['Logits@GRAD']
|
||||
)
|
||||
)
|
||||
# replicate op in dist program
|
||||
copy_op_without_infer_shape(backward_op, main_block, ctx, kwargs)
|
||||
|
||||
|
||||
class DistributedCrossEntropyImpl1(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
"""
|
||||
kwargs: inputname_mapping & outputname_mapping
|
||||
"""
|
||||
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
assert op_dist_attr is not None, (
|
||||
f"forward op [{src_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
# check validation of inputs / outputs
|
||||
assert 'Logits' in kwargs, "input [Logits] is not given"
|
||||
assert 'Label' in kwargs, "input [Label] is not given"
|
||||
assert 'Loss' in kwargs, "output [Loss] is not given"
|
||||
assert 'Softmax' in kwargs, "output [Softmax] is not given"
|
||||
|
||||
assert len(kwargs['Logits']) == 1, (
|
||||
"input [Logits] take 1 variable but got {}".format(kwargs['Logits'])
|
||||
)
|
||||
assert len(kwargs['Label']) == 1, (
|
||||
"input [Label] take 1 variable but got {}".format(kwargs['Label'])
|
||||
)
|
||||
|
||||
logits_var = main_block._var_recursive(kwargs['Logits'][0])
|
||||
label_var = main_block._var_recursive(kwargs['Label'][0])
|
||||
loss_var = main_block._var_recursive(kwargs['Loss'][0])
|
||||
softmax_var = main_block._var_recursive(kwargs['Softmax'][0])
|
||||
|
||||
# got dist attribute info
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
process_mesh_group = op_dist_attr.process_mesh.process_ids
|
||||
|
||||
# FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
|
||||
if rank_id not in process_mesh_group:
|
||||
rank_id = _get_corresponding_rank(
|
||||
ctx, op_dist_attr.process_mesh, rank_id
|
||||
)
|
||||
|
||||
check_variable_and_dtype(
|
||||
logits_var,
|
||||
'input',
|
||||
['bfloat16', 'float16', 'float32', 'float64'],
|
||||
'c_softmax_with_cross_entropy',
|
||||
)
|
||||
check_variable_and_dtype(
|
||||
label_var,
|
||||
'input',
|
||||
['int32', 'int64', 'float32', 'float64'],
|
||||
'c_softmax_with_cross_entropy',
|
||||
)
|
||||
check_variable_and_dtype(
|
||||
loss_var,
|
||||
'output',
|
||||
['bfloat16', 'float16', 'float32', 'float64'],
|
||||
'c_softmax_with_cross_entropy',
|
||||
)
|
||||
check_variable_and_dtype(
|
||||
softmax_var,
|
||||
'output',
|
||||
['bfloat16', 'float16', 'float32', 'float64'],
|
||||
'c_softmax_with_cross_entropy',
|
||||
)
|
||||
|
||||
# infer new var shape with op dist attr
|
||||
# the dims mapping in dist_op may be different from that in tensor
|
||||
# so we should
|
||||
loss_dist_attr = ctx.get_tensor_dist_attr_for_program(loss_var)
|
||||
assert loss_dist_attr is not None
|
||||
softmax_dist_attr = ctx.get_tensor_dist_attr_for_program(softmax_var)
|
||||
assert softmax_dist_attr is not None
|
||||
op_dist_attr_loss = op_dist_attr.get_output_dist_attr(loss_var.name)
|
||||
assert op_dist_attr_loss is not None
|
||||
op_dist_attr_softmax = op_dist_attr.get_output_dist_attr(
|
||||
softmax_var.name
|
||||
)
|
||||
assert op_dist_attr_softmax is not None
|
||||
|
||||
# TODO calculate ring id
|
||||
softmax_axis = src_op.desc.attr('axis')
|
||||
logits_dims_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
logits_var.name
|
||||
)
|
||||
parallel_axis = logits_dims_mapping[softmax_axis]
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh_group, process_mesh_shape, parallel_axis, rank_id
|
||||
)
|
||||
group = new_process_group(group_ranks)
|
||||
|
||||
c_cross_entropy_op = main_block.append_op(
|
||||
type='c_softmax_with_cross_entropy',
|
||||
inputs={
|
||||
'Logits': logits_var,
|
||||
'Label': label_var,
|
||||
},
|
||||
outputs={
|
||||
'Loss': loss_var,
|
||||
'Softmax': softmax_var,
|
||||
},
|
||||
attrs={
|
||||
'ring_id': group.id,
|
||||
'rank': group.local_rank(rank_id),
|
||||
'nranks': group.nranks,
|
||||
'ignore_index': src_op.desc.attr('ignore_index'),
|
||||
OP_ROLE_KEY: src_op.attr('op_role'),
|
||||
},
|
||||
)
|
||||
naive_copy_op_dist_attr_for_program(c_cross_entropy_op, src_op, ctx)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
backward_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
|
||||
|
||||
assert op_dist_attr is not None, (
|
||||
f"backward op [{backward_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
# check validation of inputs / outputs
|
||||
assert 'Softmax' in kwargs, "input [Softmax] is not given"
|
||||
assert 'Label' in kwargs, "input [Label] is not given"
|
||||
assert 'Loss@GRAD' in kwargs, "input [Loss@GRAD] is not given"
|
||||
assert 'Logits@GRAD' in kwargs, "output [Logits@GRAD] is not given"
|
||||
|
||||
assert len(kwargs['Softmax']) == 1, (
|
||||
"input [Softmax] take 1 variable but got {}".format(
|
||||
kwargs['Softmax']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['Label']) == 1, (
|
||||
"input [Label] take 1 variable but got {}".format(kwargs['Label'])
|
||||
)
|
||||
assert len(kwargs['Loss@GRAD']) == 1, (
|
||||
"input [Loss@GRAD] take 1 variable but got {}".format(
|
||||
kwargs['Loss@GRAD']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['Logits@GRAD']) == 1, (
|
||||
"output [Logits@GRAD] take 1 variable but got {}".format(
|
||||
kwargs['Logits@GRAD']
|
||||
)
|
||||
)
|
||||
|
||||
# got dist attribute info
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
process_mesh_group = op_dist_attr.process_mesh.process_ids
|
||||
|
||||
# FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
|
||||
if rank_id not in process_mesh_group:
|
||||
rank_id = _get_corresponding_rank(
|
||||
ctx, op_dist_attr.process_mesh, rank_id
|
||||
)
|
||||
|
||||
for op in main_block.ops:
|
||||
# the output value of reduce_mean_grad is 1/numel, so when the
|
||||
# tensor is sharded, we should insert a scale op to make the
|
||||
# grad correct.
|
||||
if (
|
||||
op.type == "reduce_mean_grad"
|
||||
and kwargs['Loss@GRAD'][0] in op.output_arg_names
|
||||
):
|
||||
loss_grad_var = main_block._var_recursive(
|
||||
kwargs['Loss@GRAD'][0]
|
||||
)
|
||||
loss_grad_dims_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
loss_grad_var.name
|
||||
)
|
||||
degree = 1.0
|
||||
for i in range(len(loss_grad_dims_mapping) - 1):
|
||||
if loss_grad_dims_mapping[i] != -1:
|
||||
degree *= process_mesh_shape[loss_grad_dims_mapping[i]]
|
||||
if degree > 1:
|
||||
scale_op = main_block.append_op(
|
||||
type='scale',
|
||||
inputs={'X': loss_grad_var},
|
||||
outputs={'Out': loss_grad_var},
|
||||
attrs={
|
||||
'scale': 1.0 / degree,
|
||||
OP_ROLE_KEY: OpRole.Backward,
|
||||
},
|
||||
)
|
||||
scale_op._set_attr(
|
||||
'op_namescope', '/' + ParallelMode.DataParallel
|
||||
)
|
||||
dims_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
loss_grad_var.name
|
||||
)
|
||||
scale_op_attr = OperatorDistAttr()
|
||||
scale_op_attr.process_mesh = op_dist_attr.process_mesh
|
||||
scale_op_attr.chunk_id = op_dist_attr.chunk_id
|
||||
scale_op_attr.set_output_dims_mapping(
|
||||
loss_grad_var.name, dims_mapping
|
||||
)
|
||||
scale_op_attr.set_input_dims_mapping(
|
||||
loss_grad_var.name, dims_mapping
|
||||
)
|
||||
ctx.set_op_dist_attr_for_program(scale_op, scale_op_attr)
|
||||
|
||||
# TODO calculate ring id
|
||||
softmax_axis = backward_op.desc.attr('axis')
|
||||
# softmax_dims_mapping is the same as logits_dims_mapping
|
||||
softmax_dims_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
kwargs['Softmax'][0]
|
||||
)
|
||||
parallel_axis = softmax_dims_mapping[softmax_axis]
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh_group, process_mesh_shape, parallel_axis, rank_id
|
||||
)
|
||||
group = new_process_group(group_ranks)
|
||||
|
||||
cross_entropy_grad_op_desc = main_block.append_op(type='nop').desc
|
||||
cross_entropy_grad_op_desc.set_type("c_softmax_with_cross_entropy_grad")
|
||||
cross_entropy_grad_op_desc.set_input('Softmax', [kwargs['Softmax'][0]])
|
||||
cross_entropy_grad_op_desc.set_input('Label', [kwargs['Label'][0]])
|
||||
cross_entropy_grad_op_desc.set_input(
|
||||
'Loss@GRAD', [kwargs['Loss@GRAD'][0]]
|
||||
)
|
||||
cross_entropy_grad_op_desc.set_output(
|
||||
'Logits@GRAD', [kwargs['Logits@GRAD'][0]]
|
||||
)
|
||||
|
||||
ignore_index = backward_op.desc.attr('ignore_index')
|
||||
# the ignore_index attribute in c_cross_entropy_grad kernel
|
||||
# is int64_t type, so we should set this attribute with
|
||||
# _set_int64_attr. Otherwise ignore_index will be int32 type,
|
||||
# causing an error.
|
||||
cross_entropy_grad_op_desc._set_int64_attr('ignore_index', ignore_index)
|
||||
cross_entropy_grad_op_desc._set_attr('ring_id', group.id)
|
||||
cross_entropy_grad_op_desc._set_attr('rank', group.local_rank(rank_id))
|
||||
cross_entropy_grad_op_desc._set_attr('nranks', group.nranks)
|
||||
cross_entropy_grad_op_desc._set_attr(OP_ROLE_KEY, OpRole.Backward)
|
||||
|
||||
cross_entropy_grad_op = main_block.ops[-1]
|
||||
naive_copy_op_dist_attr_for_program(
|
||||
cross_entropy_grad_op, backward_op, ctx
|
||||
)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"softmax_with_cross_entropy", DistributedCrossEntropyImpl0("cross_entropy")
|
||||
)
|
||||
register_distributed_operator_impl(
|
||||
"softmax_with_cross_entropy",
|
||||
DistributedCrossEntropyImpl1("c_cross_entropy"),
|
||||
)
|
||||
@@ -0,0 +1,681 @@
|
||||
# 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 paddle
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
|
||||
|
||||
from ..completion import contains_spmd_rule, get_phi_spmd_rule
|
||||
from ..cost import (
|
||||
_g_op_cost_factory,
|
||||
build_comp_costs_from_descs,
|
||||
build_comp_desc_from_dist_op,
|
||||
build_dp_costs,
|
||||
)
|
||||
from ..dist_attribute import DistTensorSpec, OperatorDistAttr
|
||||
from ..process_group import new_process_group
|
||||
from ..utils import (
|
||||
_get_comm_group,
|
||||
_get_corresponding_rank,
|
||||
compute_compatible_dim_mapping,
|
||||
get_dist_tensor_spec,
|
||||
is_prim_op,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
copy_op_without_infer_shape,
|
||||
get_default_distributed_operator_impl,
|
||||
gradient_synchronization,
|
||||
is_parameter_related,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
set_comm_op_dist_attr_for_program,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
__op_not_need_param_init__ = ["while", "cond"]
|
||||
__op_has_shape_attr__ = [
|
||||
"fill_constant_batch_size_like",
|
||||
"fill_constant",
|
||||
"expand_v2",
|
||||
"expand_as_v2",
|
||||
]
|
||||
|
||||
|
||||
def prim_operator_data_parallel_functor(ctx, src_op):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
|
||||
var_name = src_op.output_arg_names[0]
|
||||
if var_name in ctx.grads_params:
|
||||
assert var_name not in ctx.synced_gradient, (
|
||||
f"in primitive mode, grad is already {var_name} synced"
|
||||
)
|
||||
ctx.synced_gradient.add(var_name)
|
||||
sync_group = new_process_group(ctx.data_parallel_group)
|
||||
|
||||
allreduce_op = main_block.append_op(
|
||||
type='all_reduce',
|
||||
inputs={'x': [var_name]},
|
||||
outputs={'out': [var_name]},
|
||||
attrs={
|
||||
'ring_id': sync_group.id,
|
||||
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
||||
OP_ROLE_KEY: OpRole.Backward,
|
||||
},
|
||||
)
|
||||
|
||||
param = ctx.grads_params[var_name]
|
||||
startup_block = dist_op_context.startup_block
|
||||
new_op = startup_block.append_op(
|
||||
type='broadcast',
|
||||
inputs={'x': [param]},
|
||||
outputs={'out': [param]},
|
||||
attrs={
|
||||
'ring_id': sync_group.id,
|
||||
'root': 0,
|
||||
OP_ROLE_KEY: OpRole.Forward,
|
||||
},
|
||||
)
|
||||
|
||||
grad_var = main_block._var_recursive(var_name)
|
||||
dims_mapping = ctx.get_tensor_dist_attr_for_program(
|
||||
grad_var
|
||||
).dims_mapping
|
||||
dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
process_mesh = dist_attr.process_mesh
|
||||
op_attr = OperatorDistAttr()
|
||||
op_attr.process_mesh = process_mesh
|
||||
op_attr.set_output_dims_mapping(grad_var.name, dims_mapping)
|
||||
op_attr.set_input_dims_mapping(grad_var.name, dims_mapping)
|
||||
ctx.set_op_dist_attr_for_program(allreduce_op, op_attr)
|
||||
|
||||
|
||||
class DistributedDefault(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
input_arg_names = op_desc.input_arg_names()
|
||||
output_arg_names = op_desc.output_arg_names()
|
||||
main_block = dist_op.serial_op.block
|
||||
|
||||
num_inputs = len(input_arg_names)
|
||||
input_specs = []
|
||||
for i in range(num_inputs):
|
||||
assert not is_parameter_related(input_arg_names[i], main_block), (
|
||||
f"input {input_arg_names[i]} of op {dist_op.serial_op} is parameter, op should not use default rule."
|
||||
)
|
||||
input_specs.append(
|
||||
get_dist_tensor_spec(dist_op, input_arg_names[i])
|
||||
)
|
||||
num_outputs = len(output_arg_names)
|
||||
output_specs = []
|
||||
for i in range(num_outputs):
|
||||
assert not is_parameter_related(output_arg_names[i], main_block), (
|
||||
f"output {output_arg_names[i]} of op {dist_op.serial_op} is parameter, op should not use default rule."
|
||||
)
|
||||
output_specs.append(
|
||||
get_dist_tensor_spec(dist_op, output_arg_names[i], False)
|
||||
)
|
||||
|
||||
# step2: infer spmd
|
||||
if contains_spmd_rule(dist_op.serial_op.type):
|
||||
# when some inputs are optional, the input_arg_names will be less than input_names
|
||||
# and we can pass empty DistTensorSpec() as argument
|
||||
if len(op_desc.input_names()) > len(op_desc.input_arg_names()):
|
||||
for i in range(
|
||||
len(op_desc.input_names()) - len(op_desc.input_arg_names())
|
||||
):
|
||||
input_specs.append(DistTensorSpec())
|
||||
rule = get_phi_spmd_rule(dist_op.serial_op.type)
|
||||
fw_results = rule.infer_forward(*input_specs)
|
||||
bw_results = rule.infer_backward(*input_specs, output_specs)
|
||||
else:
|
||||
rule = get_phi_spmd_rule('default_')
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(input_specs, output_specs)
|
||||
bw_results = rule.infer_backward(input_specs, output_specs)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op, input_arg_names, output_arg_names, fw_results, bw_results
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
# all op use default dist operator impl.
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedDefault("default"))
|
||||
|
||||
|
||||
# Replicated Default
|
||||
class DistributedDefaultImpl0(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def calc_cost(self, op_role, dist_op, ctx, cluster):
|
||||
"""Calculate the cost by the op role."""
|
||||
cost = None
|
||||
if int(op_role) == int(OpRole.Backward):
|
||||
cost = self.calc_bwd_cost(dist_op, ctx, cluster)
|
||||
else:
|
||||
cost = self.calc_fwd_cost(dist_op, ctx, cluster)
|
||||
assert cost is not None
|
||||
return cost
|
||||
|
||||
def calc_fwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
processes = dist_op.dist_attr.process_mesh.process_ids
|
||||
op_type = dist_op.serial_op.type
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
_g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res_cost = [cost_mapping]
|
||||
|
||||
return res_cost
|
||||
|
||||
def calc_bwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
res = []
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
dist_attr = dist_op.dist_attr
|
||||
process_mesh = dist_attr.process_mesh
|
||||
processes = process_mesh.process_ids
|
||||
backward_op = dist_op.serial_op
|
||||
op_type = backward_op.type
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
_g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
main_block = backward_op.block
|
||||
need_gradient_allreduce = False
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and not is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
|
||||
mesh_shape = process_mesh.shape
|
||||
batch_size_axis = (
|
||||
var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
|
||||
)
|
||||
if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
|
||||
need_gradient_allreduce = True
|
||||
break
|
||||
|
||||
if need_gradient_allreduce:
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(
|
||||
varname
|
||||
)
|
||||
mesh_shape = process_mesh.shape
|
||||
parallel_axis = batch_size_axis
|
||||
attrs = {"use_calc_stream": True}
|
||||
var_names = [varname + "@GRAD"]
|
||||
build_dp_costs(
|
||||
res,
|
||||
dist_op,
|
||||
ctx,
|
||||
var_names,
|
||||
attrs,
|
||||
parallel_axis,
|
||||
cluster,
|
||||
)
|
||||
return res
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
batch_dim_mappings = []
|
||||
input_names = op_desc.input_names()
|
||||
xshape_arg_names = []
|
||||
if "XShape" in input_names:
|
||||
xshape_arg_names = op_desc.input("XShape")
|
||||
for arg_name in op_desc.input_arg_names():
|
||||
serial_tensor = dist_op.get_serial_input(arg_name)
|
||||
dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
|
||||
if serial_tensor.is_parameter:
|
||||
for mapping in dims_mapping:
|
||||
if mapping != -1:
|
||||
return False
|
||||
continue
|
||||
if arg_name not in xshape_arg_names:
|
||||
if len(dims_mapping) > 1:
|
||||
for mapping in dims_mapping[1:]:
|
||||
if mapping != -1:
|
||||
return False
|
||||
if len(dims_mapping) >= 1:
|
||||
batch_dim_mappings.append(dims_mapping[0])
|
||||
else:
|
||||
if dims_mapping[0] != -1:
|
||||
return False
|
||||
if len(dims_mapping) > 2:
|
||||
for mapping in dims_mapping[2:]:
|
||||
if mapping != -1:
|
||||
return False
|
||||
if len(dims_mapping) >= 2:
|
||||
batch_dim_mappings.append(dims_mapping[1])
|
||||
|
||||
if compute_compatible_dim_mapping(batch_dim_mappings) is None:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
output_names = op_desc.output_names()
|
||||
batch_dim_mappings = []
|
||||
xshape_arg_names = []
|
||||
if "XShape" in output_names:
|
||||
xshape_arg_names = op_desc.output("XShape")
|
||||
for arg_name in op_desc.output_arg_names():
|
||||
serial_tensor = dist_op.get_serial_output(arg_name)
|
||||
dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
|
||||
if serial_tensor.is_parameter:
|
||||
for mapping in dims_mapping:
|
||||
if mapping != -1:
|
||||
return False
|
||||
continue
|
||||
if arg_name not in xshape_arg_names:
|
||||
if len(dims_mapping) > 1:
|
||||
for mapping in dims_mapping[1:]:
|
||||
if mapping != -1:
|
||||
return False
|
||||
if len(dims_mapping) >= 1:
|
||||
batch_dim_mappings.append(dims_mapping[0])
|
||||
else:
|
||||
if dims_mapping[0] != -1:
|
||||
return False
|
||||
if len(dims_mapping) > 2:
|
||||
for mapping in dims_mapping[2:]:
|
||||
if mapping != -1:
|
||||
return False
|
||||
if len(dims_mapping) >= 2:
|
||||
batch_dim_mappings.append(dims_mapping[1])
|
||||
|
||||
if compute_compatible_dim_mapping(batch_dim_mappings) is None:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
batch_dim_mappings = []
|
||||
# Check input compatibility
|
||||
input_names = op_desc.input_names()
|
||||
xshape_arg_names = []
|
||||
if "XShape" in input_names:
|
||||
xshape_arg_names = op_desc.input("XShape")
|
||||
for arg_name in op_desc.input_arg_names():
|
||||
serial_tensor = dist_op.get_serial_input(arg_name)
|
||||
dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
|
||||
if serial_tensor is not None and serial_tensor.is_parameter:
|
||||
for mapping in dims_mapping:
|
||||
if mapping != -1:
|
||||
return False
|
||||
continue
|
||||
if arg_name not in xshape_arg_names:
|
||||
if len(dims_mapping) > 1:
|
||||
for mapping in dims_mapping[1:]:
|
||||
if mapping != -1:
|
||||
return False
|
||||
if len(dims_mapping) >= 1:
|
||||
batch_dim_mappings.append(dims_mapping[0])
|
||||
else:
|
||||
if dims_mapping[0] != -1:
|
||||
return False
|
||||
if len(dims_mapping) > 2:
|
||||
for mapping in dims_mapping[2:]:
|
||||
if mapping != -1:
|
||||
return False
|
||||
if len(dims_mapping) >= 2:
|
||||
batch_dim_mappings.append(dims_mapping[1])
|
||||
|
||||
# Check output compatibility
|
||||
output_names = op_desc.output_names()
|
||||
xshape_arg_names = []
|
||||
if "XShape" in output_names:
|
||||
xshape_arg_names = op_desc.output("XShape")
|
||||
for arg_name in op_desc.output_arg_names():
|
||||
serial_tensor = dist_op.get_serial_output(arg_name)
|
||||
dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
|
||||
if serial_tensor is not None and serial_tensor.is_parameter:
|
||||
for mapping in dims_mapping:
|
||||
if mapping != -1:
|
||||
return False
|
||||
continue
|
||||
if arg_name not in xshape_arg_names:
|
||||
if len(dims_mapping) > 1:
|
||||
for mapping in dims_mapping[1:]:
|
||||
if mapping != -1:
|
||||
return False
|
||||
if len(dims_mapping) >= 1:
|
||||
batch_dim_mappings.append(dims_mapping[0])
|
||||
else:
|
||||
if dims_mapping[0] != -1:
|
||||
return False
|
||||
if len(dims_mapping) > 2:
|
||||
for mapping in dims_mapping[2:]:
|
||||
if mapping != -1:
|
||||
return False
|
||||
if len(dims_mapping) >= 2:
|
||||
batch_dim_mappings.append(dims_mapping[1])
|
||||
|
||||
# Check batch dim mapping compatibility
|
||||
if not all(
|
||||
batch_dim_mappings[0] == dim_mapping
|
||||
for dim_mapping in batch_dim_mappings
|
||||
):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
|
||||
if op_desc.type() == "while":
|
||||
return False
|
||||
|
||||
input_names = op_desc.input_names()
|
||||
input_xshape_arg_names = []
|
||||
if "XShape" in input_names:
|
||||
input_xshape_arg_names = op_desc.input("XShape")
|
||||
|
||||
output_names = op_desc.output_names()
|
||||
output_xshape_arg_names = []
|
||||
if "XShape" in output_names:
|
||||
output_xshape_arg_names = op_desc.output("XShape")
|
||||
|
||||
batch_dim_mappings = []
|
||||
for arg_name in op_desc.input_arg_names():
|
||||
serial_tensor = dist_op.get_serial_input(arg_name)
|
||||
if serial_tensor.is_parameter:
|
||||
continue
|
||||
dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
|
||||
if arg_name not in input_xshape_arg_names:
|
||||
if len(dims_mapping) >= 1:
|
||||
batch_dim_mappings.append(dims_mapping[0])
|
||||
else:
|
||||
batch_dim_mappings.append(dims_mapping[1])
|
||||
for arg_name in op_desc.output_arg_names():
|
||||
if op_desc.type() == 'fill_any_like':
|
||||
input_tensor = dist_op.get_serial_input(
|
||||
op_desc.input_arg_names()[0]
|
||||
)
|
||||
if input_tensor.is_parameter:
|
||||
continue
|
||||
serial_tensor = dist_op.get_serial_output(arg_name)
|
||||
if serial_tensor.is_parameter:
|
||||
continue
|
||||
dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
|
||||
if arg_name not in output_xshape_arg_names:
|
||||
if len(dims_mapping) >= 1:
|
||||
batch_dim_mappings.append(dims_mapping[0])
|
||||
else:
|
||||
batch_dim_mappings.append(dims_mapping[1])
|
||||
|
||||
if not batch_dim_mappings:
|
||||
return changed
|
||||
|
||||
compatible_dim_mapping = compute_compatible_dim_mapping(
|
||||
batch_dim_mappings
|
||||
)
|
||||
if compatible_dim_mapping is None:
|
||||
return False
|
||||
|
||||
for arg_name in op_desc.input_arg_names():
|
||||
serial_tensor = dist_op.get_serial_input(arg_name)
|
||||
if serial_tensor.is_parameter:
|
||||
continue
|
||||
dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
|
||||
if arg_name not in input_xshape_arg_names:
|
||||
if (
|
||||
len(dims_mapping) >= 1
|
||||
and compatible_dim_mapping != dims_mapping[0]
|
||||
):
|
||||
dims_mapping[0] = compatible_dim_mapping
|
||||
op_dist_attr.set_input_dims_mapping(arg_name, dims_mapping)
|
||||
changed = True
|
||||
else:
|
||||
if (
|
||||
len(dims_mapping) >= 2
|
||||
and compatible_dim_mapping != dims_mapping[1]
|
||||
):
|
||||
dims_mapping[1] = compatible_dim_mapping
|
||||
op_dist_attr.set_input_dims_mapping(arg_name, dims_mapping)
|
||||
changed = True
|
||||
for arg_name in op_desc.output_arg_names():
|
||||
if op_desc.type() == 'fill_any_like':
|
||||
input_tensor = dist_op.get_serial_input(
|
||||
op_desc.input_arg_names()[0]
|
||||
)
|
||||
if input_tensor.is_parameter:
|
||||
continue
|
||||
if op_desc.type() in ["shape", "slice"]:
|
||||
continue
|
||||
serial_tensor = dist_op.get_serial_output(arg_name)
|
||||
if serial_tensor.is_parameter:
|
||||
continue
|
||||
dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
|
||||
if arg_name not in output_xshape_arg_names:
|
||||
if (
|
||||
len(dims_mapping) >= 1
|
||||
and compatible_dim_mapping != dims_mapping[0]
|
||||
):
|
||||
dims_mapping[0] = compatible_dim_mapping
|
||||
op_dist_attr.set_output_dims_mapping(arg_name, dims_mapping)
|
||||
changed = True
|
||||
else:
|
||||
if (
|
||||
len(dims_mapping) >= 2
|
||||
and compatible_dim_mapping != dims_mapping[1]
|
||||
):
|
||||
dims_mapping[1] = compatible_dim_mapping
|
||||
op_dist_attr.set_output_dims_mapping(arg_name, dims_mapping)
|
||||
changed = True
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
|
||||
# check validation of inputs / outputs
|
||||
for input_name in src_op.desc.input_names():
|
||||
assert input_name in kwargs, f"input [{input_name}] is not given"
|
||||
assert len(kwargs[input_name]) == len(
|
||||
src_op.desc.input(input_name)
|
||||
), f"number of tensor for input [{input_name}] is not match"
|
||||
for output_name in src_op.desc.output_names():
|
||||
assert output_name in kwargs, f"input [{output_name}] is not given"
|
||||
assert len(kwargs[output_name]) == len(
|
||||
src_op.desc.output(output_name)
|
||||
), f"number of tensor for input [{output_name}] is not match"
|
||||
|
||||
# replicate op in dist program
|
||||
dst_op = copy_op_without_infer_shape(src_op, main_block, ctx, kwargs)
|
||||
|
||||
def get_shape_attr_name():
|
||||
for name in ["shape", "target_shape"]:
|
||||
if src_op.has_attr(name) and src_op.attr(name):
|
||||
return name
|
||||
return None
|
||||
|
||||
shape_attr_name = get_shape_attr_name()
|
||||
if shape_attr_name and src_op.type in __op_has_shape_attr__:
|
||||
shape_list = src_op.attr(shape_attr_name)
|
||||
Out_var = main_block._var_recursive(kwargs['Out'][0])
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
dim_mapping = op_dist_attr.get_output_dims_mapping(Out_var.name)
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
assert len(shape_list) == len(dim_mapping)
|
||||
# modify target shape
|
||||
for idx, axis in enumerate(dim_mapping):
|
||||
if axis >= 0:
|
||||
if len(shape_list) > idx:
|
||||
shape_list[idx] = (
|
||||
shape_list[idx] // process_mesh_shape[axis]
|
||||
)
|
||||
dst_op.desc._set_attr(shape_attr_name, shape_list)
|
||||
|
||||
# data parallel synchronization for primitive operators
|
||||
from paddle.incubate.autograd import prim_enabled
|
||||
|
||||
if prim_enabled():
|
||||
assert is_prim_op(src_op)
|
||||
prim_operator_data_parallel_functor(ctx, src_op)
|
||||
return
|
||||
|
||||
# param initialization sync
|
||||
if src_op.type in __op_not_need_param_init__:
|
||||
return
|
||||
|
||||
for varname in dst_op.desc.input_arg_names():
|
||||
if (
|
||||
startup_block.has_var(varname)
|
||||
and startup_block.var(varname).is_parameter
|
||||
and varname not in dist_op_context.already_init_sync_vars
|
||||
):
|
||||
dist_op_context.already_init_sync_vars.add(varname)
|
||||
param = startup_block.var(varname)
|
||||
param_dist_attr = ctx.get_tensor_dist_attr_for_program(param)
|
||||
process_mesh = param_dist_attr.process_mesh
|
||||
dims_mapping = param_dist_attr.dims_mapping
|
||||
|
||||
# FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
|
||||
if rank_id not in process_mesh.process_ids:
|
||||
rank_id = _get_corresponding_rank(
|
||||
ctx, process_mesh, rank_id
|
||||
)
|
||||
|
||||
# NOTE all not splited axis should be presented in mesh
|
||||
for axis, size in enumerate(process_mesh.shape):
|
||||
if size <= 1 or axis in dims_mapping:
|
||||
pass
|
||||
else:
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh.process_ids,
|
||||
process_mesh.shape,
|
||||
axis,
|
||||
rank_id,
|
||||
)
|
||||
sync_group = new_process_group(group_ranks)
|
||||
|
||||
new_op = startup_block.append_op(
|
||||
type='broadcast',
|
||||
inputs={'x': param},
|
||||
outputs={'out': param},
|
||||
attrs={
|
||||
'ring_id': sync_group.id,
|
||||
'root': 0,
|
||||
OP_ROLE_KEY: OpRole.Forward,
|
||||
},
|
||||
)
|
||||
set_comm_op_dist_attr_for_program(
|
||||
new_op,
|
||||
process_mesh,
|
||||
param_dist_attr,
|
||||
ctx,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
# by now the backward function only insert the gradient allreduce for dist op itself
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
backward_op = dist_op_context.cur_src_op
|
||||
dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
|
||||
assert dist_attr is not None, (
|
||||
f"backward op [{backward_op}] don't have dist attribute !"
|
||||
)
|
||||
rank_id = dist_op_context.rank_id
|
||||
|
||||
# check validation of inputs / outputs
|
||||
for input_name in backward_op.desc.input_names():
|
||||
assert input_name in kwargs, f"input [{input_name}] is not given"
|
||||
assert len(kwargs[input_name]) == len(
|
||||
backward_op.desc.input(input_name)
|
||||
), f"number of tensor for input [{input_name}] is not match"
|
||||
for output_name in backward_op.desc.output_names():
|
||||
assert output_name in kwargs, f"input [{output_name}] is not given"
|
||||
assert len(kwargs[output_name]) == len(
|
||||
backward_op.desc.output(output_name)
|
||||
), f"number of tensor for input [{output_name}] is not match"
|
||||
|
||||
# replicate op in dist program
|
||||
copy_op_without_infer_shape(backward_op, main_block, ctx, kwargs)
|
||||
|
||||
# data parallel gradient synchronization
|
||||
act_grad_names = []
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and not is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
act_grad_names.append(varname)
|
||||
|
||||
out_grad_names = []
|
||||
for output_name in backward_op.desc.output_names():
|
||||
for varname in backward_op.desc.output(output_name):
|
||||
if varname in kwargs["grad_var_to_var"]:
|
||||
fwd_name = kwargs["grad_var_to_var"][varname]
|
||||
if not main_block._find_var_recursive(fwd_name):
|
||||
continue
|
||||
if is_parameter_related(fwd_name, main_block):
|
||||
out_grad_names.append(varname)
|
||||
|
||||
gradient_synchronization(
|
||||
ctx, backward_op, act_grad_names, out_grad_names, rank_id
|
||||
)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"default", DistributedDefaultImpl0("replicate_parallel")
|
||||
)
|
||||
@@ -0,0 +1,238 @@
|
||||
# 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 logging
|
||||
|
||||
import paddle
|
||||
from paddle.base.log_helper import get_logger
|
||||
from paddle.framework import core
|
||||
from paddle.utils import unique_name
|
||||
|
||||
from ...random import determinate_rng, is_enable_auto_rand_ctrl
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..utils import (
|
||||
get_dist_tensor_spec,
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
|
||||
set_var_dist_attr,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
merge_forward_backward_dims_mapping,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
from .dist_eltwise import DistributedDefaultImpl0, DistributedElementwiseImpl0
|
||||
|
||||
_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
|
||||
class DistributedDropout(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
mask_name = op_desc.output('Mask')[0]
|
||||
# seed_name = op_desc.input('Seed')[0] // seed is a scalar and leave it to be unsharded
|
||||
|
||||
x_spec = get_dist_tensor_spec(dist_op, x_name)
|
||||
output_spec = get_dist_tensor_spec(dist_op, out_name, False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("dropout")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(x_spec)
|
||||
bw_results = rule.infer_backward(x_spec, output_spec)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op, [x_name], [out_name], fw_results, bw_results
|
||||
)
|
||||
|
||||
# step5: update mask and seed dropout special
|
||||
if changed:
|
||||
(
|
||||
_,
|
||||
inferred_output_dims_mappings,
|
||||
) = merge_forward_backward_dims_mapping(fw_results, bw_results)
|
||||
dist_op.dist_attr.set_output_dims_mapping(
|
||||
mask_name, inferred_output_dims_mappings[0]
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
# all dropout op use Dropout with Random Control dist operator impl.
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
op_dist_attr.impl_type = "dropout"
|
||||
op_dist_attr.impl_idx = 0
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedDropout("dropout"))
|
||||
|
||||
|
||||
# Dist Dropout with Random Control
|
||||
# Dropout re-use the compatible and cost function of elementwise
|
||||
class DistributedDropoutImpl0(DistributedElementwiseImpl0):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
assert op_dist_attr is not None, (
|
||||
f"forward op [{src_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
if is_enable_auto_rand_ctrl() and not op_dist_attr.is_recompute:
|
||||
# check validation of inputs / outputs
|
||||
assert 'X' in kwargs, "input [{}] is not given".format('X')
|
||||
assert len(kwargs['X']) == 1, (
|
||||
"input X should be only one tensor but got {}".format(
|
||||
kwargs['X']
|
||||
)
|
||||
)
|
||||
assert 'Seed' in kwargs, "input [{}] is not given".format('Seed')
|
||||
|
||||
if (
|
||||
src_op.has_attr("fix_seed")
|
||||
and src_op.attr("fix_seed")
|
||||
and src_op.has_attr("seed")
|
||||
and src_op.attr("seed")
|
||||
):
|
||||
_logger.info(
|
||||
f"Auto Parallel Random Control Skipped Since manual seed is set by user: {src_op}"
|
||||
)
|
||||
elif rank_id not in op_dist_attr.process_mesh.process_ids:
|
||||
pass
|
||||
# NOTE Adopt for recompute
|
||||
# If user already set seed, We should not modify it. But if the seed is added by recompute pass, it should be under control.
|
||||
# TODO in future recompute pass should happen after parallel partition. and remove this at that time.
|
||||
elif len(kwargs['Seed']) > 0 or len(src_op.input("Seed")) > 0:
|
||||
seed_var_name = kwargs['Seed'][0]
|
||||
if seed_var_name.startswith('rc_seed'):
|
||||
pre_op = main_block.ops[-1]
|
||||
assert (
|
||||
pre_op.type == "seed"
|
||||
and len(pre_op.attr("rng_name")) == 0
|
||||
), f"found exception op {pre_op}"
|
||||
|
||||
# determinate rng
|
||||
X_var = main_block._var_recursive(kwargs['X'][0])
|
||||
X_dims_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
X_var.name
|
||||
)
|
||||
process_mesh = op_dist_attr.process_mesh
|
||||
rng_name = determinate_rng(
|
||||
rank_id, X_dims_mapping, process_mesh
|
||||
)
|
||||
# make recompute seed under control
|
||||
pre_op._set_attr("rng_name", rng_name)
|
||||
pre_op._set_attr("deterministic", True)
|
||||
pre_op._set_attr("force_cpu", True)
|
||||
else:
|
||||
_logger.info(
|
||||
f"Auto Parallel Random Control Skipped Since manual seed is set by user: {src_op}"
|
||||
)
|
||||
else:
|
||||
# determinate rng
|
||||
X_var = main_block._var_recursive(kwargs['X'][0])
|
||||
X_dims_mapping = op_dist_attr.get_input_dims_mapping(X_var.name)
|
||||
process_mesh = op_dist_attr.process_mesh
|
||||
|
||||
rng_name = determinate_rng(
|
||||
rank_id, X_dims_mapping, process_mesh
|
||||
)
|
||||
assert rng_name is not None and rng_name != ""
|
||||
|
||||
# insert seed op
|
||||
seed_var = main_block.create_var(
|
||||
name=unique_name.generate_with_ignorable_key(
|
||||
".".join(["tensor_parallel_seed", 'tmp'])
|
||||
),
|
||||
dtype=paddle.int32,
|
||||
type=core.VarDesc.VarType.DENSE_TENSOR,
|
||||
persistable=False,
|
||||
stop_gradient=False,
|
||||
)
|
||||
|
||||
# set new seed_var's dist_attr
|
||||
seed_var_dims_mapping = [-1]
|
||||
seed_var_dist_attr = set_var_dist_attr(
|
||||
ctx,
|
||||
seed_var,
|
||||
seed_var_dims_mapping,
|
||||
process_mesh,
|
||||
chunk_id=op_dist_attr.chunk_id,
|
||||
)
|
||||
|
||||
# adopt for recompute
|
||||
# force_cpu to reduce sync copy from CPU->GPU->CPU, and reduce pipeline hang
|
||||
seed_op = main_block.append_op(
|
||||
type='seed',
|
||||
outputs={'Out': seed_var},
|
||||
attrs={
|
||||
'deterministic': True,
|
||||
'rng_name': rng_name,
|
||||
'force_cpu': True,
|
||||
},
|
||||
)
|
||||
seed_op._set_attr('op_namescope', 'auto_tensor_parallel_seed')
|
||||
# set new seed op's dist_attr
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
||||
seed_op,
|
||||
process_mesh,
|
||||
seed_var_dims_mapping,
|
||||
ctx,
|
||||
chunk_id=op_dist_attr.chunk_id,
|
||||
)
|
||||
|
||||
# modify dropout op
|
||||
src_op.desc.set_input("Seed", [seed_var.name])
|
||||
src_op.desc._set_attr("fix_seed", False)
|
||||
src_op.desc._set_attr("seed", 0)
|
||||
op_dist_attr.set_input_dist_attr(
|
||||
seed_var.name, seed_var_dist_attr
|
||||
)
|
||||
kwargs['Seed'] = [seed_var.name]
|
||||
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
# dropout backward is deterministic by mask, and not need for random state control
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"dropout", DistributedDropoutImpl0("random_control")
|
||||
)
|
||||
@@ -0,0 +1,400 @@
|
||||
# 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
|
||||
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..cost import (
|
||||
_g_op_cost_factory,
|
||||
build_comp_costs_from_descs,
|
||||
build_comp_desc_from_dist_op,
|
||||
build_dp_costs,
|
||||
)
|
||||
from ..utils import (
|
||||
compute_compatible_dim_mapping,
|
||||
compute_compatible_dims_mapping,
|
||||
get_dist_tensor_spec,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
is_elementwise_op,
|
||||
is_parameter_related,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedElementwise(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
assert len(op_desc.input_arg_names()) >= 1, (
|
||||
f"elementwise op [{op_desc.type}] has [{len(op_desc.input_arg_names())}] inputs"
|
||||
)
|
||||
input_arg_names = op_desc.input_arg_names()
|
||||
assert len(op_desc.output_arg_names()) == 1, (
|
||||
f"elementwise op [{dist_op.serial_op}] has [{len(op_desc.output_arg_names())}] outputs"
|
||||
)
|
||||
output_arg_name = op_desc.output_arg_names()[0]
|
||||
num_inputs = len(input_arg_names)
|
||||
|
||||
# TODO (zhangyichen) replace dist tensor specs by dist tensor in future.
|
||||
input_specs = []
|
||||
for i in range(num_inputs):
|
||||
input_specs.append(
|
||||
get_dist_tensor_spec(dist_op, input_arg_names[i])
|
||||
)
|
||||
output_spec = get_dist_tensor_spec(dist_op, output_arg_name, False)
|
||||
|
||||
# step2: infer spmd
|
||||
# TODO revise me
|
||||
op_type = op_desc.type()
|
||||
rule = get_phi_spmd_rule(op_type)
|
||||
fw_results = rule.infer_forward(*input_specs)
|
||||
bw_results = rule.infer_backward(*input_specs, output_spec)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op, input_arg_names, [output_arg_name], fw_results, bw_results
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
# NOTE this function will be remove once we use local reshard to replace distopimpls
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
# all elementwise op use default dist operator impl.
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedElementwise("elementwise")
|
||||
)
|
||||
|
||||
|
||||
# Replicated Elementwise
|
||||
class DistributedElementwiseImpl0(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = False
|
||||
self._backward_implemented = False
|
||||
|
||||
def calc_cost(self, op_role, dist_op, ctx, cluster):
|
||||
"""Calculate the cost by the op role."""
|
||||
cost = None
|
||||
if int(op_role) == int(OpRole.Backward):
|
||||
cost = self.calc_bwd_cost(dist_op, ctx, cluster)
|
||||
else:
|
||||
cost = self.calc_fwd_cost(dist_op, ctx, cluster)
|
||||
assert cost is not None
|
||||
return cost
|
||||
|
||||
def calc_fwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
processes = dist_op.dist_attr.process_mesh.process_ids
|
||||
op_type = dist_op.serial_op.type
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
_g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res_cost = [cost_mapping]
|
||||
|
||||
return res_cost
|
||||
|
||||
def calc_bwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
res = []
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
dist_attr = dist_op.dist_attr
|
||||
process_mesh = dist_attr.process_mesh
|
||||
processes = process_mesh.process_ids
|
||||
backward_op = dist_op.serial_op
|
||||
op_type = backward_op.type
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
_g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
main_block = backward_op.block
|
||||
need_gradient_allreduce = False
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and not is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
|
||||
mesh_shape = process_mesh.shape
|
||||
batch_size_axis = (
|
||||
var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
|
||||
)
|
||||
if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
|
||||
need_gradient_allreduce = True
|
||||
break
|
||||
|
||||
if need_gradient_allreduce:
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(
|
||||
varname
|
||||
)
|
||||
mesh_shape = process_mesh.shape
|
||||
batch_size_axis = var_dim_mapping[0]
|
||||
parallel_axis = batch_size_axis
|
||||
attrs = {"use_calc_stream": True}
|
||||
var_names = [varname + "@GRAD"]
|
||||
build_dp_costs(
|
||||
res,
|
||||
dist_op,
|
||||
ctx,
|
||||
var_names,
|
||||
attrs,
|
||||
parallel_axis,
|
||||
cluster,
|
||||
)
|
||||
return res
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
if not is_elementwise_op(op_desc.type()):
|
||||
return False
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
dims_mapping_list = []
|
||||
input_arg_names = op_desc.input_arg_names()
|
||||
max_dims_mapping_len = -1
|
||||
for arg_name in input_arg_names:
|
||||
dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
|
||||
if max_dims_mapping_len < len(dims_mapping):
|
||||
max_dims_mapping_len = len(dims_mapping)
|
||||
dims_mapping_list.append(dims_mapping)
|
||||
|
||||
for idx in range(max_dims_mapping_len):
|
||||
dim_mappings = []
|
||||
for dims_mapping in dims_mapping_list:
|
||||
if idx < len(dims_mapping):
|
||||
dim_mappings.append(dims_mapping[-(idx + 1)])
|
||||
if compute_compatible_dim_mapping(dim_mappings) is None:
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
if not is_elementwise_op(op_desc.type()):
|
||||
return False
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
dims_mapping_list = []
|
||||
output_arg_names = op_desc.output_arg_names()
|
||||
max_dims_mapping_len = -1
|
||||
for arg_name in output_arg_names:
|
||||
dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
|
||||
if max_dims_mapping_len < len(dims_mapping):
|
||||
max_dims_mapping_len = len(dims_mapping)
|
||||
dims_mapping_list.append(dims_mapping)
|
||||
|
||||
for idx in range(max_dims_mapping_len):
|
||||
dim_mappings = []
|
||||
for dims_mapping in dims_mapping_list:
|
||||
if idx < len(dims_mapping):
|
||||
dim_mappings.append(dims_mapping[-(idx + 1)])
|
||||
if compute_compatible_dim_mapping(dim_mappings) is None:
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
if not is_elementwise_op(op_desc.type()):
|
||||
return False
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
dims_mapping_list = []
|
||||
|
||||
input_arg_names = op_desc.input_arg_names()
|
||||
input_max_dims_mapping_len = -1
|
||||
for arg_name in input_arg_names:
|
||||
dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
|
||||
if input_max_dims_mapping_len < len(dims_mapping):
|
||||
input_max_dims_mapping_len = len(dims_mapping)
|
||||
dims_mapping_list.append(dims_mapping)
|
||||
|
||||
output_arg_names = op_desc.output_arg_names()
|
||||
output_max_dims_mapping_len = -1
|
||||
for arg_name in output_arg_names:
|
||||
dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
|
||||
if output_max_dims_mapping_len < len(dims_mapping):
|
||||
output_max_dims_mapping_len = len(dims_mapping)
|
||||
dims_mapping_list.append(dims_mapping)
|
||||
|
||||
assert input_max_dims_mapping_len == output_max_dims_mapping_len
|
||||
max_dims_mapping_len = input_max_dims_mapping_len
|
||||
|
||||
for idx in range(max_dims_mapping_len):
|
||||
dim_mappings = []
|
||||
for dims_mapping in dims_mapping_list:
|
||||
if idx < len(dims_mapping):
|
||||
dim_mappings.append(dims_mapping[-(idx + 1)])
|
||||
if not all(
|
||||
dim_mappings[0] == dim_mapping for dim_mapping in dim_mappings
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
dims_mapping_list = []
|
||||
|
||||
input_arg_names = op_desc.input_arg_names()
|
||||
input_dims_mapping_dict = {}
|
||||
input_dims_mapping_lens = {}
|
||||
input_max_dims_mapping_len = -1
|
||||
for arg_name in input_arg_names:
|
||||
dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
|
||||
if input_max_dims_mapping_len < len(dims_mapping):
|
||||
input_max_dims_mapping_len = len(dims_mapping)
|
||||
input_dims_mapping_dict[arg_name] = dims_mapping
|
||||
input_dims_mapping_lens[arg_name] = len(dims_mapping)
|
||||
for arg_name in input_arg_names:
|
||||
if input_dims_mapping_lens[arg_name] < input_max_dims_mapping_len:
|
||||
new_dims_mapping = [
|
||||
-1 for _ in range(input_max_dims_mapping_len)
|
||||
]
|
||||
for i in range(input_dims_mapping_lens[arg_name]):
|
||||
new_idx = (
|
||||
input_max_dims_mapping_len
|
||||
- input_dims_mapping_lens[arg_name]
|
||||
) + i
|
||||
new_dims_mapping[new_idx] = input_dims_mapping_dict[
|
||||
arg_name
|
||||
][i]
|
||||
dims_mapping_list.append(new_dims_mapping)
|
||||
else:
|
||||
dims_mapping_list.append(input_dims_mapping_dict[arg_name])
|
||||
|
||||
output_arg_names = op_desc.output_arg_names()
|
||||
output_dims_mapping_dict = {}
|
||||
output_dims_mapping_lens = {}
|
||||
output_max_dims_mapping_len = -1
|
||||
for arg_name in output_arg_names:
|
||||
dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
|
||||
if output_max_dims_mapping_len < len(dims_mapping):
|
||||
output_max_dims_mapping_len = len(dims_mapping)
|
||||
output_dims_mapping_dict[arg_name] = dims_mapping
|
||||
output_dims_mapping_lens[arg_name] = len(dims_mapping)
|
||||
for arg_name in output_arg_names:
|
||||
if output_dims_mapping_lens[arg_name] < output_max_dims_mapping_len:
|
||||
new_dims_mapping = [
|
||||
-1 for _ in range(output_max_dims_mapping_len)
|
||||
]
|
||||
for i in range(output_dims_mapping_lens[arg_name]):
|
||||
new_idx = (
|
||||
output_max_dims_mapping_len
|
||||
- output_dims_mapping_lens[arg_name]
|
||||
) + i
|
||||
new_dims_mapping[new_idx] = output_dims_mapping_dict[
|
||||
arg_name
|
||||
][i]
|
||||
dims_mapping_list.append(new_dims_mapping)
|
||||
else:
|
||||
dims_mapping_list.append(output_dims_mapping_dict[arg_name])
|
||||
|
||||
assert input_max_dims_mapping_len == output_max_dims_mapping_len
|
||||
max_dims_mapping_len = input_max_dims_mapping_len
|
||||
compatible_dims_mapping = compute_compatible_dims_mapping(
|
||||
dims_mapping_list
|
||||
)
|
||||
if compatible_dims_mapping is None:
|
||||
return False
|
||||
|
||||
for arg_name in input_arg_names:
|
||||
if input_dims_mapping_lens[arg_name] < max_dims_mapping_len:
|
||||
new_dims_mapping = [
|
||||
-1 for _ in range(input_dims_mapping_lens[arg_name])
|
||||
]
|
||||
for i in range(input_dims_mapping_lens[arg_name]):
|
||||
new_idx = (
|
||||
max_dims_mapping_len - input_dims_mapping_lens[arg_name]
|
||||
) + i
|
||||
new_dims_mapping[i] = compatible_dims_mapping[new_idx]
|
||||
if new_dims_mapping != input_dims_mapping_dict[arg_name]:
|
||||
op_dist_attr.set_input_dims_mapping(
|
||||
arg_name, new_dims_mapping
|
||||
)
|
||||
changed = True
|
||||
else:
|
||||
if compatible_dims_mapping != input_dims_mapping_dict[arg_name]:
|
||||
op_dist_attr.set_input_dims_mapping(
|
||||
arg_name, compatible_dims_mapping
|
||||
)
|
||||
changed = True
|
||||
|
||||
for arg_name in output_arg_names:
|
||||
if output_dims_mapping_lens[arg_name] < max_dims_mapping_len:
|
||||
new_dims_mapping = [
|
||||
-1 for _ in range(output_dims_mapping_lens[arg_name])
|
||||
]
|
||||
for i in range(output_dims_mapping_lens[arg_name]):
|
||||
new_idx = (
|
||||
max_dims_mapping_len
|
||||
- output_dims_mapping_lens[arg_name]
|
||||
) + i
|
||||
new_dims_mapping[i] = compatible_dims_mapping[new_idx]
|
||||
if new_dims_mapping != output_dims_mapping_dict[arg_name]:
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
arg_name, new_dims_mapping
|
||||
)
|
||||
changed = True
|
||||
else:
|
||||
if (
|
||||
compatible_dims_mapping
|
||||
!= output_dims_mapping_dict[arg_name]
|
||||
):
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
arg_name, compatible_dims_mapping
|
||||
)
|
||||
changed = True
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"elementwise", DistributedElementwiseImpl0("replicate_parallel")
|
||||
)
|
||||
@@ -0,0 +1,671 @@
|
||||
# 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 paddle
|
||||
from paddle.common_ops_import import check_variable_and_dtype
|
||||
from paddle.distributed.auto_parallel.static.cost.comm_op_cost import (
|
||||
AllReduceOpCost,
|
||||
IdentityOpCost,
|
||||
)
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
|
||||
from paddle.framework import core
|
||||
from paddle.utils import unique_name
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..cost import (
|
||||
EmbeddingGradOpCost,
|
||||
EmbeddingOpCost,
|
||||
build_comm_costs_from_descs,
|
||||
build_comm_desc_from_dist_op,
|
||||
build_comp_costs_from_descs,
|
||||
build_comp_desc_from_dist_op,
|
||||
build_dp_costs,
|
||||
)
|
||||
from ..dist_attribute import OperatorDistAttr
|
||||
from ..process_group import new_process_group
|
||||
from ..utils import (
|
||||
_get_comm_group,
|
||||
_get_corresponding_rank,
|
||||
_get_idx_in_axis,
|
||||
compute_compatible_and_update_dim_mapping,
|
||||
get_dist_tensor_spec,
|
||||
is_dim_replicate,
|
||||
is_dim_shard,
|
||||
set_var_dist_attr,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
ParallelMode,
|
||||
get_default_distributed_operator_impl,
|
||||
gradient_synchronization,
|
||||
naive_copy_op_dist_attr_for_program,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
set_comm_op_dist_attr_for_program,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
|
||||
class DistributedEmbedding(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
assert dist_op.serial_op.type == "lookup_table_v2", (
|
||||
f"{dist_op.serial_op.type} is not supported by dist embedding yet."
|
||||
)
|
||||
|
||||
x_name = op_desc.input('Ids')[0]
|
||||
w_name = op_desc.input('W')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
padding_idx = op_desc.attr('padding_idx')
|
||||
is_sparse = op_desc.attr('is_sparse')
|
||||
|
||||
x_spec = get_dist_tensor_spec(dist_op, x_name)
|
||||
w_spec = get_dist_tensor_spec(dist_op, w_name)
|
||||
output_spec = get_dist_tensor_spec(dist_op, out_name, False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("embedding")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(x_spec, w_spec, padding_idx, is_sparse)
|
||||
bw_results = rule.infer_backward(
|
||||
x_spec, w_spec, output_spec, padding_idx, is_sparse
|
||||
)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op, [x_name, w_name], [out_name], fw_results, bw_results
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
reverted = False
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
op_desc = dist_op.serial_op.desc
|
||||
out_name = op_desc.output('Out')[0]
|
||||
out_dist_attr = op_dist_attr.get_output_dist_attr(out_name)
|
||||
|
||||
# vocab parallel embedding
|
||||
if out_dist_attr._is_partial():
|
||||
op_dist_attr.impl_type = op_desc.type()
|
||||
op_dist_attr.impl_idx = 0
|
||||
# data parallel or col parallel of weight
|
||||
else:
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return reverted
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedEmbedding("lookup_table_v2")
|
||||
)
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedEmbedding("c_embedding")
|
||||
)
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedEmbedding("lookup_table")
|
||||
)
|
||||
|
||||
|
||||
def adopt_lookup_table_v1(ctx, main_block, src_op, Ids_var):
|
||||
assert len(Ids_var.shape) == 3, (
|
||||
f"input Ids to lookup_table should have 3 dimensions but got [{Ids_var.name}] with shape [{Ids_var.shape}]"
|
||||
)
|
||||
if not Ids_var.stop_gradient:
|
||||
raise NotImplementedError(
|
||||
'Requiring the gradient of Ids of lookup_table(v1) dist op is not currently supported. Please open an issue with details on your use case so that we can prioritize adding this (for instance, adversarial training for language model).'
|
||||
)
|
||||
|
||||
target_shape = list(Ids_var.shape[:-1])
|
||||
intermediate_var_0 = main_block.create_var(
|
||||
name=unique_name.generate_with_ignorable_key(
|
||||
".".join(["dist_reshape", 'tmp'])
|
||||
),
|
||||
dtype=Ids_var.dtype,
|
||||
shape=target_shape,
|
||||
type=core.VarDesc.VarType.DENSE_TENSOR,
|
||||
persistable=False,
|
||||
stop_gradient=True,
|
||||
)
|
||||
|
||||
target_shape = [0, *list(Ids_var.shape[:-1])]
|
||||
xshape_var = main_block.create_var(
|
||||
name=unique_name.generate_with_ignorable_key(
|
||||
".".join(["dist_Xshape", 'tmp'])
|
||||
),
|
||||
dtype=Ids_var.dtype,
|
||||
shape=target_shape,
|
||||
type=core.VarDesc.VarType.DENSE_TENSOR,
|
||||
persistable=False,
|
||||
stop_gradient=True,
|
||||
)
|
||||
|
||||
# TODO use inplace reshape for memory saving
|
||||
reshape_op = main_block.append_op(
|
||||
type='reshape2',
|
||||
inputs={'X': [Ids_var]},
|
||||
outputs={'Out': [intermediate_var_0], 'XShape': [xshape_var]},
|
||||
attrs={
|
||||
"shape": [0, -1],
|
||||
},
|
||||
)
|
||||
|
||||
# set dist attr
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
Ids_var_dist_attr = op_dist_attr.get_input_dist_attr(Ids_var.name)
|
||||
assert Ids_var_dist_attr is not None
|
||||
intermediate_var_0_dist_attr = set_var_dist_attr(
|
||||
ctx,
|
||||
intermediate_var_0,
|
||||
Ids_var_dist_attr.dims_mapping,
|
||||
Ids_var_dist_attr.process_mesh,
|
||||
chunk_id=Ids_var_dist_attr.chunk_id,
|
||||
)
|
||||
set_var_dist_attr(
|
||||
ctx,
|
||||
xshape_var,
|
||||
[-1, *list(Ids_var_dist_attr.dims_mapping)],
|
||||
Ids_var_dist_attr.process_mesh,
|
||||
chunk_id=Ids_var_dist_attr.chunk_id,
|
||||
)
|
||||
# rename src_op's input
|
||||
src_op._rename_input(Ids_var.name, intermediate_var_0.name)
|
||||
op_dist_attr.del_input_dist_attr(Ids_var.name)
|
||||
op_dist_attr.set_input_dist_attr(
|
||||
intermediate_var_0.name, intermediate_var_0_dist_attr
|
||||
)
|
||||
|
||||
new_op_dist_attr = OperatorDistAttr()
|
||||
new_op_dist_attr.process_mesh = Ids_var_dist_attr.process_mesh
|
||||
new_op_dist_attr.impl_type = "default"
|
||||
new_op_dist_attr.impl_idx = 0
|
||||
new_op_dist_attr.chunk_id = Ids_var_dist_attr.chunk_id
|
||||
new_op_dist_attr.set_input_dims_mapping(
|
||||
Ids_var.name, Ids_var_dist_attr.dims_mapping
|
||||
)
|
||||
new_op_dist_attr.set_output_dims_mapping(
|
||||
intermediate_var_0.name, Ids_var_dist_attr.dims_mapping
|
||||
)
|
||||
new_op_dist_attr.set_output_dims_mapping(
|
||||
xshape_var.name, [-1, *list(Ids_var_dist_attr.dims_mapping)]
|
||||
)
|
||||
ctx.set_op_dist_attr_for_program(reshape_op, new_op_dist_attr)
|
||||
|
||||
return intermediate_var_0
|
||||
|
||||
|
||||
# RowParallel
|
||||
class DistributedEmbeddingImpl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def calc_cost(self, op_role, dist_op, ctx, cluster):
|
||||
"""Calculate the cost by the op role."""
|
||||
cost = None
|
||||
if int(op_role) == int(OpRole.Forward):
|
||||
cost = self.calc_fwd_cost(dist_op, ctx, cluster)
|
||||
elif int(op_role) == int(OpRole.Backward):
|
||||
cost = self.calc_bwd_cost(dist_op, ctx, cluster)
|
||||
assert cost is not None
|
||||
return cost
|
||||
|
||||
def calc_fwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
processes = dist_op.dist_attr.process_mesh.process_ids
|
||||
# embedding need start_index
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
EmbeddingOpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
|
||||
serial_op = dist_op.serial_op
|
||||
parallel_axis = dist_op.dist_attr.get_input_dims_mapping(
|
||||
serial_op.input("W")[0]
|
||||
)[0]
|
||||
attrs = {"use_calc_stream": True, "use_model_parallel": True}
|
||||
var_names = serial_op.output("Out")
|
||||
all_reduce_sum_desc_mapping = build_comm_desc_from_dist_op(
|
||||
"all_reduce",
|
||||
dist_op,
|
||||
ctx,
|
||||
var_names,
|
||||
attrs=attrs,
|
||||
parallel_axis=parallel_axis,
|
||||
)
|
||||
|
||||
comm_op_cost_list = build_comm_costs_from_descs(
|
||||
AllReduceOpCost,
|
||||
ctx,
|
||||
processes,
|
||||
all_reduce_sum_desc_mapping,
|
||||
cluster,
|
||||
)
|
||||
|
||||
res_cost = [cost_mapping, comm_op_cost_list]
|
||||
|
||||
return res_cost
|
||||
|
||||
def calc_bwd_cost(self, dist_op, ctx, cluster):
|
||||
# by now the backward function only insert the gradient allreduce for dist op itself
|
||||
res = []
|
||||
backward_op = dist_op.serial_op
|
||||
main_block = backward_op.block
|
||||
dist_attr = dist_op.dist_attr
|
||||
|
||||
embedding_row_dim_mapping = dist_attr.get_input_dims_mapping(
|
||||
backward_op.input("W")[0]
|
||||
)[0]
|
||||
parallel_axis = embedding_row_dim_mapping
|
||||
attrs = {"use_calc_stream": True, "use_model_parallel": True}
|
||||
var_names = [backward_op.input("Out@GRAD")[0]]
|
||||
c_identity_desc_mapping = build_comm_desc_from_dist_op(
|
||||
"c_identity",
|
||||
dist_op,
|
||||
ctx,
|
||||
var_names,
|
||||
attrs=attrs,
|
||||
parallel_axis=parallel_axis,
|
||||
)
|
||||
|
||||
process_mesh = dist_attr.process_mesh
|
||||
processes = process_mesh.process_ids
|
||||
comm_op_cost_list = build_comm_costs_from_descs(
|
||||
IdentityOpCost, ctx, processes, c_identity_desc_mapping, cluster
|
||||
)
|
||||
res.append(comm_op_cost_list)
|
||||
|
||||
# calc comp op cost
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
EmbeddingGradOpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
# need gradient allreduce
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(
|
||||
backward_op.input("Ids")[0]
|
||||
)
|
||||
mesh_shape = process_mesh.shape
|
||||
batch_size_axis = var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
|
||||
if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
|
||||
parallel_axis = batch_size_axis
|
||||
attrs = {"use_calc_stream": True}
|
||||
var_names = [backward_op.output('W@GRAD')[0]]
|
||||
build_dp_costs(
|
||||
res, dist_op, ctx, var_names, attrs, parallel_axis, cluster
|
||||
)
|
||||
|
||||
return res
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
ids_name = op_desc.input('Ids')[0]
|
||||
w_name = op_desc.input('W')[0]
|
||||
ids_dims_mapping = op_dist_attr.get_input_dims_mapping(ids_name)
|
||||
w_dims_mapping = op_dist_attr.get_input_dims_mapping(w_name)
|
||||
if is_dim_replicate(w_dims_mapping[-2]) or is_dim_shard(
|
||||
w_dims_mapping[-1]
|
||||
):
|
||||
return False
|
||||
# Other dimensions must be replicate except the batch dimension
|
||||
for mapping in ids_dims_mapping[1:]:
|
||||
if is_dim_shard(mapping):
|
||||
return False
|
||||
|
||||
if is_dim_shard(ids_dims_mapping[0]) and is_dim_shard(
|
||||
w_dims_mapping[-2]
|
||||
):
|
||||
if ids_dims_mapping[0] == w_dims_mapping[-2]:
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
out_name = op_desc.output('Out')[0]
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
# Other dimensions must be replicate except the batch dimension
|
||||
for mapping in out_dims_mapping[1:]:
|
||||
if is_dim_shard(mapping):
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
ids_name = op_desc.input('Ids')[0]
|
||||
w_name = op_desc.input('W')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
ids_dims_mapping = op_dist_attr.get_input_dims_mapping(ids_name)
|
||||
w_dims_mapping = op_dist_attr.get_input_dims_mapping(w_name)
|
||||
|
||||
if ids_dims_mapping != out_dims_mapping[: len(ids_dims_mapping)]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
ids_name = op_desc.input('Ids')[0]
|
||||
w_name = op_desc.input('W')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
ids_dims_mapping = op_dist_attr.get_input_dims_mapping(ids_name)
|
||||
w_dims_mapping = op_dist_attr.get_input_dims_mapping(w_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
for i in range(len(ids_dims_mapping)):
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[ids_dims_mapping, out_dims_mapping], [i, i]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[w_dims_mapping, out_dims_mapping], [-1, -1]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(ids_name, ids_dims_mapping)
|
||||
op_dist_attr.set_input_dims_mapping(w_name, w_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
"""
|
||||
kwargs: inputname_mapping & outputname_mapping
|
||||
"""
|
||||
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
assert op_dist_attr is not None, (
|
||||
f"forward op [{src_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
# check validation of inputs / outputs
|
||||
assert 'Ids' in kwargs, "input [{}] is not given".format('Ids')
|
||||
assert 'W' in kwargs, "input [{}] is not given".format('W')
|
||||
assert 'Out' in kwargs, "output [{}] is not given".format('Out')
|
||||
|
||||
assert len(kwargs['Ids']) == 1, (
|
||||
"row_parallel_embedding input Ids take 1 variable but got {}".format(
|
||||
kwargs['Ids']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['W']) == 1, (
|
||||
"row_parallel_embedding input W take 1 variable but got {}".format(
|
||||
kwargs['W']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['Out']) == 1, (
|
||||
"row_parallel_embedding output Out take 1 variable but got {}".format(
|
||||
kwargs['Out']
|
||||
)
|
||||
)
|
||||
|
||||
Ids_var = main_block._var_recursive(kwargs['Ids'][0])
|
||||
Weight_var = main_block._var_recursive(kwargs['W'][0])
|
||||
Out_var = main_block._var_recursive(kwargs['Out'][0])
|
||||
|
||||
# support lookup_table_v1
|
||||
if src_op.type == 'lookup_table':
|
||||
Ids_var = adopt_lookup_table_v1(ctx, main_block, src_op, Ids_var)
|
||||
|
||||
# got dist attribute info
|
||||
embedding_row_dim_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
Weight_var.name
|
||||
)[0]
|
||||
assert embedding_row_dim_mapping >= 0, (
|
||||
f"row_parallel_embedding's row should be divided by a specific mesh axis, but got [{embedding_row_dim_mapping}]"
|
||||
)
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
process_mesh_group = op_dist_attr.process_mesh.process_ids
|
||||
|
||||
# FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
|
||||
if rank_id not in process_mesh_group:
|
||||
rank_id = _get_corresponding_rank(
|
||||
ctx, op_dist_attr.process_mesh, rank_id
|
||||
)
|
||||
|
||||
# A generalized method to calculate embedding offset using cartesian product
|
||||
relative_idx = _get_idx_in_axis(
|
||||
process_mesh_group,
|
||||
process_mesh_shape,
|
||||
embedding_row_dim_mapping,
|
||||
rank_id,
|
||||
)
|
||||
|
||||
per_part_size = Weight_var.shape[0]
|
||||
relative_idx = relative_idx * per_part_size
|
||||
|
||||
# TODO calculate ring id
|
||||
parallel_axis = embedding_row_dim_mapping
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh_group, process_mesh_shape, parallel_axis, rank_id
|
||||
)
|
||||
group = new_process_group(group_ranks)
|
||||
|
||||
# append op
|
||||
check_variable_and_dtype(
|
||||
Ids_var, 'input', ['int32', 'int64'], 'c_embedding'
|
||||
)
|
||||
|
||||
# infer new var shape with op dist attr
|
||||
out_tensor_dist_attr = ctx.get_tensor_dist_attr_for_program(Out_var)
|
||||
assert out_tensor_dist_attr is not None
|
||||
out_var_dist_attr = op_dist_attr.get_output_dist_attr(Out_var.name)
|
||||
assert out_var_dist_attr is not None
|
||||
|
||||
c_embedding_op_desc = main_block.append_op(type='nop').desc
|
||||
c_embedding_op_desc.set_type("c_embedding")
|
||||
c_embedding_op_desc.set_input('Ids', [Ids_var.name])
|
||||
c_embedding_op_desc.set_input('W', [Weight_var.name])
|
||||
c_embedding_op_desc.set_output('Out', [Out_var.name])
|
||||
c_embedding_op_desc._set_attr('start_index', relative_idx)
|
||||
c_embedding_op_desc._set_attr(OP_ROLE_KEY, src_op.attr('op_role'))
|
||||
c_embedding_op = main_block.ops[-1]
|
||||
assert c_embedding_op.type == "c_embedding"
|
||||
naive_copy_op_dist_attr_for_program(c_embedding_op, src_op, ctx)
|
||||
|
||||
# use_model_parallel
|
||||
all_reduce_sum_op = main_block.append_op(
|
||||
type='all_reduce',
|
||||
inputs={'x': [Out_var]},
|
||||
outputs={'out': [Out_var]},
|
||||
attrs={
|
||||
'ring_id': group.id,
|
||||
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
||||
'use_model_parallel': True,
|
||||
OP_ROLE_KEY: src_op.attr('op_role'),
|
||||
},
|
||||
)
|
||||
all_reduce_sum_op._set_attr(
|
||||
'op_namescope', '/' + ParallelMode.TensorParallel
|
||||
)
|
||||
# allreduce
|
||||
set_comm_op_dist_attr_for_program(
|
||||
all_reduce_sum_op,
|
||||
op_dist_attr.process_mesh,
|
||||
out_var_dist_attr,
|
||||
ctx,
|
||||
chunk_id=op_dist_attr.chunk_id,
|
||||
)
|
||||
|
||||
# param initialization sync
|
||||
if Weight_var.is_parameter and not op_dist_attr.is_recompute:
|
||||
if Weight_var.name in dist_op_context.already_init_sync_vars:
|
||||
return
|
||||
dist_op_context.already_init_sync_vars.add(Weight_var.name)
|
||||
param = startup_block.var(Weight_var.name)
|
||||
param_dist_attr = ctx.get_tensor_dist_attr_for_program(param)
|
||||
process_mesh = param_dist_attr.process_mesh
|
||||
dim_mapping = param_dist_attr.dims_mapping
|
||||
|
||||
# NOTE all not split axis should be presented in mesh
|
||||
for axis, size in enumerate(process_mesh.shape):
|
||||
if size <= 1 or axis in dim_mapping:
|
||||
pass
|
||||
else:
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh.process_ids,
|
||||
process_mesh.shape,
|
||||
axis,
|
||||
rank_id,
|
||||
)
|
||||
sync_group = new_process_group(group_ranks)
|
||||
|
||||
broadcast_op = startup_block.append_op(
|
||||
type='broadcast',
|
||||
inputs={'x': param},
|
||||
outputs={'out': param},
|
||||
attrs={
|
||||
'ring_id': sync_group.id,
|
||||
'root': 0,
|
||||
OP_ROLE_KEY: OpRole.Forward,
|
||||
},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
# by now the backward function only insert the gradient allreduce for dist op itself
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
backward_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
|
||||
assert dist_attr is not None, (
|
||||
f"backward op [{backward_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
# FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
|
||||
if rank_id not in dist_attr.process_mesh.process_ids:
|
||||
rank_id = _get_corresponding_rank(
|
||||
ctx, dist_attr.process_mesh, rank_id
|
||||
)
|
||||
|
||||
assert 'Ids' in kwargs, "input [{}] is not given".format('Ids')
|
||||
assert 'W' in kwargs, "input [{}] is not given".format('W')
|
||||
assert 'Out@GRAD' in kwargs, "input [{}] is not given".format('Out')
|
||||
assert 'W@GRAD' in kwargs, "output [{}] is not given".format('W@GRAD')
|
||||
|
||||
assert len(kwargs['Ids']) == 1, (
|
||||
"row_parallel_embedding input Ids take 1 variable but got {}".format(
|
||||
kwargs['Ids']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['W']) == 1, (
|
||||
"row_parallel_embedding input Ids take 1 variable but got {}".format(
|
||||
kwargs['W']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['Out@GRAD']) == 1, (
|
||||
"row_parallel_embedding input Ids take 1 variable but got {}".format(
|
||||
kwargs['Out']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['W@GRAD']) == 1, (
|
||||
"row_parallel_embedding output Ids take 1 variable but got {}".format(
|
||||
kwargs['W@GRAD']
|
||||
)
|
||||
)
|
||||
|
||||
Ids_var = main_block._var_recursive(kwargs['Ids'][0])
|
||||
Weight_var = main_block._var_recursive(kwargs['W'][0])
|
||||
Out_grad = main_block._var_recursive(kwargs['Out@GRAD'][0])
|
||||
Weight_grad = main_block._var_recursive(kwargs['W@GRAD'][0])
|
||||
|
||||
embedding_row_dim_mapping = dist_attr.get_input_dims_mapping(
|
||||
Weight_var.name
|
||||
)[0]
|
||||
assert embedding_row_dim_mapping >= 0, (
|
||||
f"row_parallel_embedding's row should be divided by a specific mesh axis, but got [{embedding_row_dim_mapping}]"
|
||||
)
|
||||
process_mesh_shape = dist_attr.process_mesh.shape
|
||||
process_mesh_group = dist_attr.process_mesh.process_ids
|
||||
|
||||
# A generalized method to calculate embedding offset using cartesian product
|
||||
relative_idx = _get_idx_in_axis(
|
||||
process_mesh_group,
|
||||
process_mesh_shape,
|
||||
embedding_row_dim_mapping,
|
||||
rank_id,
|
||||
)
|
||||
per_part_size = Weight_var.shape[0]
|
||||
relative_idx = relative_idx * per_part_size
|
||||
|
||||
c_embedding_grad_op_desc = main_block.append_op(type='nop').desc
|
||||
c_embedding_grad_op_desc.set_type("c_embedding_grad")
|
||||
c_embedding_grad_op_desc.set_input('Ids', [Ids_var.name])
|
||||
c_embedding_grad_op_desc.set_input('W', [Weight_var.name])
|
||||
c_embedding_grad_op_desc.set_input('Out@GRAD', [Out_grad.name])
|
||||
c_embedding_grad_op_desc.set_output('W@GRAD', [Weight_grad.name])
|
||||
c_embedding_grad_op_desc._set_attr('start_index', relative_idx)
|
||||
c_embedding_grad_op_desc._set_attr(OP_ROLE_KEY, OpRole.Backward)
|
||||
|
||||
c_embedding_grad_op = main_block.ops[-1]
|
||||
assert c_embedding_grad_op.type == "c_embedding_grad"
|
||||
naive_copy_op_dist_attr_for_program(
|
||||
c_embedding_grad_op, backward_op, ctx
|
||||
)
|
||||
|
||||
# data parallel gradient synchronization
|
||||
act_grad_names = [Ids_var.name]
|
||||
out_grad_names = [kwargs['W@GRAD'][0]]
|
||||
|
||||
gradient_synchronization(
|
||||
ctx, backward_op, act_grad_names, out_grad_names, rank_id
|
||||
)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"lookup_table_v2", DistributedEmbeddingImpl("row_parallel")
|
||||
)
|
||||
register_distributed_operator_impl(
|
||||
"c_embedding", DistributedEmbeddingImpl("row_parallel")
|
||||
)
|
||||
register_distributed_operator_impl(
|
||||
"lookup_table", DistributedEmbeddingImpl("row_parallel")
|
||||
)
|
||||
@@ -0,0 +1,80 @@
|
||||
# Copyright (c) 2023 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
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..utils import get_dist_tensor_spec
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
|
||||
class DistributedExpandAs(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
|
||||
input_arg_names = op_desc.input_arg_names()
|
||||
output_arg_names = op_desc.output_arg_names()
|
||||
target_shape = op_desc.attr('target_shape')
|
||||
|
||||
input_specs = []
|
||||
for name in input_arg_names:
|
||||
input_specs.append(get_dist_tensor_spec(dist_op, name))
|
||||
|
||||
assert len(input_specs) == 2
|
||||
|
||||
output_spec = get_dist_tensor_spec(dist_op, output_arg_names[0], False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("expand_as")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(
|
||||
input_specs[0], input_specs[1], target_shape
|
||||
)
|
||||
bw_results = rule.infer_backward(
|
||||
input_specs[0], input_specs[1], output_spec, target_shape
|
||||
)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op,
|
||||
input_arg_names,
|
||||
output_arg_names,
|
||||
fw_results,
|
||||
bw_results,
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedExpandAs("expand_as_v2")
|
||||
)
|
||||
+144
@@ -0,0 +1,144 @@
|
||||
# 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
|
||||
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
||||
|
||||
from ..cost import (
|
||||
FillConstantBatchSizeLikeOpCost,
|
||||
build_comp_costs_from_descs,
|
||||
build_comp_desc_from_dist_op,
|
||||
)
|
||||
from ..utils import compute_compatible_and_update_dim_mapping
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedFillConstantBatchSizeLike(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedFillConstantBatchSizeLike("fill_constant_batch_size_like")
|
||||
)
|
||||
|
||||
|
||||
class DistributedFillConstantBatchSizeLikeImpl0(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def calc_cost(self, op_role, dist_op, ctx, cluster):
|
||||
cost = None
|
||||
if int(op_role) == int(OpRole.Backward):
|
||||
raise ValueError(
|
||||
"The fill_constant_batch_size_like has no grad op."
|
||||
)
|
||||
else:
|
||||
cost = self.calc_fwd_cost(dist_op, ctx, cluster)
|
||||
assert cost is not None
|
||||
return cost
|
||||
|
||||
def calc_fwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
processes = dist_op.dist_attr.process_mesh.process_ids
|
||||
op_type = dist_op.serial_op.type
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
FillConstantBatchSizeLikeOpCost,
|
||||
ctx,
|
||||
processes,
|
||||
desc_mapping,
|
||||
cluster,
|
||||
)
|
||||
|
||||
res_cost = [cost_mapping]
|
||||
return res_cost
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
out_name = op_desc.output('Out')[0]
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
shape_list = op_desc.attr("shape")
|
||||
|
||||
if len(shape_list) != len(out_dims_mapping):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
out_name = op_desc.output('Out')[0]
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
in_name = op_desc.input('Input')[0]
|
||||
in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
|
||||
|
||||
# the dim_mapping of batch dimension should be the same
|
||||
return out_dims_mapping[0] == in_dims_mapping[0]
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('Input')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
# only the batch size dimension of input and output are relative.
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[x_dims_mapping, out_dims_mapping], [0, 0]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
"""
|
||||
kwargs: inputname_mapping & outputname_mapping
|
||||
"""
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"fill_constant_batch_size_like",
|
||||
DistributedFillConstantBatchSizeLikeImpl0("fill_by_shape"),
|
||||
)
|
||||
@@ -0,0 +1,97 @@
|
||||
# Copyright (c) 2023 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
|
||||
|
||||
from ...random import determinate_rng, is_enable_auto_rand_ctrl
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
from .dist_eltwise import DistributedElementwiseImpl0
|
||||
|
||||
|
||||
class DistributedFlashAttn(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedFlashAttn("flash_attn"))
|
||||
|
||||
|
||||
# Dist FlashAttn with Random Control
|
||||
# NOTE(zhiqiu): trick implementation, copy dist_attr of q,k,v to out
|
||||
class DistributedFlashAttnImpl0(DistributedElementwiseImpl0):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
|
||||
if (
|
||||
is_enable_auto_rand_ctrl()
|
||||
and not op_dist_attr.is_recompute
|
||||
and rank_id in op_dist_attr.process_mesh.process_ids
|
||||
):
|
||||
assert op_dist_attr is not None, (
|
||||
f"forward op [{src_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
if (
|
||||
len(kwargs.get('fixed_seed_offset', [])) > 0
|
||||
or len(src_op.input("fixed_seed_offset")) > 0
|
||||
):
|
||||
# TODO(kuizhiqing) recompute should go here
|
||||
pass
|
||||
else:
|
||||
# determinate rng
|
||||
q_var = main_block._var_recursive(kwargs['q'][0])
|
||||
k_var = main_block._var_recursive(kwargs['k'][0])
|
||||
q_dims_mapping = op_dist_attr.get_input_dims_mapping(q_var.name)
|
||||
k_dims_mapping = op_dist_attr.get_input_dims_mapping(k_var.name)
|
||||
process_mesh = op_dist_attr.process_mesh
|
||||
dims_mapping = [*q_dims_mapping[:3], q_dims_mapping[2]]
|
||||
|
||||
rng_name = determinate_rng(rank_id, dims_mapping, process_mesh)
|
||||
assert rng_name is not None and rng_name != ""
|
||||
|
||||
src_op._set_attr('rng_name', rng_name)
|
||||
|
||||
DistributedElementwiseImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
# dropout backward is deterministic by mask, and not need for random state control
|
||||
DistributedElementwiseImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"flash_attn", DistributedFlashAttnImpl0("random_control")
|
||||
)
|
||||
@@ -0,0 +1,235 @@
|
||||
# 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.
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from ..utils import (
|
||||
_get_comm_group,
|
||||
_get_corresponding_rank,
|
||||
compute_compatible_and_update_dim_mapping,
|
||||
is_dim_replicate,
|
||||
is_dim_shard,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedFusedAttention(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedFusedAttention("fused_attention")
|
||||
)
|
||||
|
||||
|
||||
class DistributedFusedAttentionImpl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
qkv_w = op_desc.input('QKVW')[0]
|
||||
qkv_bias = op_desc.input('QKVBias')[0]
|
||||
out_w = op_desc.input('OutLinearW')[0]
|
||||
out_bias = op_desc.input('OutLinearBias')[0]
|
||||
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
qkv_w_dims_mapping = op_dist_attr.get_input_dims_mapping(qkv_w)
|
||||
qkv_bias_dims_mapping = op_dist_attr.get_input_dims_mapping(qkv_bias)
|
||||
out_w_dims_mapping = op_dist_attr.get_input_dims_mapping(out_w)
|
||||
out_bias_dims_mapping = op_dist_attr.get_input_dims_mapping(out_bias)
|
||||
|
||||
head_axis = 1
|
||||
for mapping in x_dims_mapping[1:-1]:
|
||||
if is_dim_shard(mapping):
|
||||
return False
|
||||
if len(qkv_w_dims_mapping) != 4 or is_dim_replicate(
|
||||
qkv_w_dims_mapping[head_axis]
|
||||
):
|
||||
return False
|
||||
if len(qkv_bias_dims_mapping) != 3 or is_dim_replicate(
|
||||
qkv_bias_dims_mapping[head_axis]
|
||||
):
|
||||
return False
|
||||
if is_dim_replicate(out_w_dims_mapping[0]):
|
||||
return False
|
||||
if is_dim_shard(out_bias_dims_mapping[-1]):
|
||||
return False
|
||||
|
||||
replicated_dims = [
|
||||
qkv_w_dims_mapping[0],
|
||||
qkv_w_dims_mapping[-2],
|
||||
qkv_w_dims_mapping[-1],
|
||||
qkv_bias_dims_mapping[0],
|
||||
qkv_bias_dims_mapping[-1],
|
||||
out_w_dims_mapping[-1],
|
||||
out_bias_dims_mapping[-1],
|
||||
]
|
||||
for mapping in replicated_dims:
|
||||
if is_dim_shard(mapping):
|
||||
return False
|
||||
if qkv_bias_dims_mapping[head_axis] != qkv_w_dims_mapping[head_axis]:
|
||||
return False
|
||||
if qkv_bias_dims_mapping[head_axis] != out_w_dims_mapping[0]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
|
||||
# none of output should be sharded
|
||||
for out_name in op_desc.output_names():
|
||||
out = op_desc.output(out_name)[0]
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out)
|
||||
for mapping in out_dims_mapping[1:-1]:
|
||||
if is_dim_shard(mapping):
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_names = op_desc.output('Y')
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
for out_name in out_names:
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
if x_dims_mapping != out_dims_mapping:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_names = op_desc.output('Y')
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
|
||||
for out_name in out_names:
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
for i in range(len(x_dims_mapping)):
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[x_dims_mapping, out_dims_mapping], [i, i]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
out_name, out_dims_mapping
|
||||
)
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
|
||||
if rank_id not in op_dist_attr.process_mesh.process_ids:
|
||||
rank_id = _get_corresponding_rank(
|
||||
ctx, op_dist_attr.process_mesh, rank_id
|
||||
)
|
||||
|
||||
# infer logic comm presentation
|
||||
head_axis = 1
|
||||
qkv_w = src_op.input('QKVW')[0]
|
||||
qkv_w_col_dim_mapping = op_dist_attr.get_input_dims_mapping(qkv_w)[
|
||||
head_axis
|
||||
]
|
||||
assert qkv_w_col_dim_mapping >= 0, (
|
||||
f"col_parallel_matmul's row should be divided by a specific mesh axis, but got [{qkv_w_col_dim_mapping}]"
|
||||
)
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
process_mesh_group = op_dist_attr.process_mesh.process_ids
|
||||
|
||||
parallel_axis = qkv_w_col_dim_mapping
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh_group, process_mesh_shape, parallel_axis, rank_id
|
||||
)
|
||||
group = new_process_group(group_ranks)
|
||||
|
||||
# insert op
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
# setting comm id
|
||||
new_op = main_block.ops[-1]
|
||||
assert new_op.type == "fused_attention"
|
||||
new_op._set_attr("ring_id", int(group.id))
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
|
||||
if rank_id not in op_dist_attr.process_mesh.process_ids:
|
||||
rank_id = _get_corresponding_rank(
|
||||
ctx, op_dist_attr.process_mesh, rank_id
|
||||
)
|
||||
|
||||
# infer logic comm presentation
|
||||
out_w = src_op.input('OutLinearW')[0]
|
||||
out_w_col_dim_mapping = op_dist_attr.get_input_dims_mapping(out_w)[-1]
|
||||
assert out_w_col_dim_mapping >= 0, (
|
||||
f"col_parallel_matmul's row should be divided by a specific mesh axis, but got [{out_w_col_dim_mapping}]"
|
||||
)
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
process_mesh_group = op_dist_attr.process_mesh.process_ids
|
||||
|
||||
parallel_axis = out_w_col_dim_mapping
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh_group, process_mesh_shape, parallel_axis, rank_id
|
||||
)
|
||||
group = new_process_group(group_ranks)
|
||||
|
||||
# insert op
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
# setting comm id
|
||||
new_op = main_block.ops[-1]
|
||||
assert new_op.type == "fused_attention_grad"
|
||||
new_op._set_attr("ring_id", int(group.id))
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"fused_attention", DistributedFusedAttentionImpl("tensor_parallel")
|
||||
)
|
||||
@@ -0,0 +1,195 @@
|
||||
# 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 logging
|
||||
|
||||
import paddle
|
||||
from paddle.base.log_helper import get_logger
|
||||
from paddle.framework import core
|
||||
from paddle.utils import unique_name
|
||||
|
||||
from ...random import determinate_rng, is_enable_auto_rand_ctrl
|
||||
from ..utils import (
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
|
||||
set_var_dist_attr,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
from .dist_eltwise import DistributedDefaultImpl0, DistributedElementwiseImpl0
|
||||
|
||||
_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
|
||||
class DistributedDropout(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedDropout("fused_dropout_add")
|
||||
)
|
||||
|
||||
|
||||
# Dist Dropout with Random Control
|
||||
# Dropout re-use the compatible and cost function of elementwise
|
||||
class DistributedDropoutImpl0(DistributedElementwiseImpl0):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
|
||||
if is_enable_auto_rand_ctrl() and not op_dist_attr.is_recompute:
|
||||
assert op_dist_attr is not None, (
|
||||
f"forward op [{src_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
assert 'seed_tensor' in kwargs, "input [{}] is not given".format(
|
||||
'seed_tensor'
|
||||
)
|
||||
|
||||
if (
|
||||
src_op.has_attr("fix_seed")
|
||||
and src_op.attr("fix_seed")
|
||||
and src_op.has_attr("seed")
|
||||
and src_op.attr("seed")
|
||||
):
|
||||
_logger.info(
|
||||
f"Auto Parallel Random Control Skipped Since manual seed is set by user: {src_op}"
|
||||
)
|
||||
elif rank_id not in op_dist_attr.process_mesh.process_ids:
|
||||
pass
|
||||
elif (
|
||||
len(kwargs['seed_tensor']) > 0
|
||||
or len(src_op.input("seed_tensor")) > 0
|
||||
):
|
||||
seed_var_name = kwargs['seed_tensor'][0]
|
||||
if seed_var_name.startswith('rc_seed'):
|
||||
pre_op = main_block.ops[-1]
|
||||
assert (
|
||||
pre_op.type == "seed"
|
||||
and len(pre_op.attr("rng_name")) == 0
|
||||
), f"found exception op {pre_op}"
|
||||
|
||||
# determinate rng
|
||||
X_var = main_block._var_recursive(kwargs['x'][0])
|
||||
X_dims_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
X_var.name
|
||||
)
|
||||
process_mesh = op_dist_attr.process_mesh
|
||||
rng_name = determinate_rng(
|
||||
rank_id, X_dims_mapping, process_mesh
|
||||
)
|
||||
# make recompute seed under control
|
||||
pre_op._set_attr("rng_name", rng_name)
|
||||
pre_op._set_attr("deterministic", True)
|
||||
pre_op._set_attr("force_cpu", True)
|
||||
else:
|
||||
_logger.info(
|
||||
f"Auto Parallel Random Control Skipped Since manual seed is set by user: {src_op}"
|
||||
)
|
||||
else:
|
||||
# determinate rng
|
||||
X_var = main_block._var_recursive(kwargs['x'][0])
|
||||
X_dims_mapping = op_dist_attr.get_input_dims_mapping(X_var.name)
|
||||
process_mesh = op_dist_attr.process_mesh
|
||||
|
||||
rng_name = determinate_rng(
|
||||
rank_id, X_dims_mapping, process_mesh
|
||||
)
|
||||
assert rng_name is not None and rng_name != ""
|
||||
|
||||
# insert seed op
|
||||
seed_var = main_block.create_var(
|
||||
name=unique_name.generate_with_ignorable_key(
|
||||
".".join(["tensor_parallel_seed", 'tmp'])
|
||||
),
|
||||
dtype=paddle.int32,
|
||||
type=core.VarDesc.VarType.DENSE_TENSOR,
|
||||
persistable=False,
|
||||
stop_gradient=False,
|
||||
)
|
||||
|
||||
# set new seed_var's dist_attr
|
||||
seed_var_dims_mapping = [-1]
|
||||
seed_var_dist_attr = set_var_dist_attr(
|
||||
ctx,
|
||||
seed_var,
|
||||
seed_var_dims_mapping,
|
||||
process_mesh,
|
||||
chunk_id=op_dist_attr.chunk_id,
|
||||
)
|
||||
|
||||
# adopt for recompute
|
||||
# force_cpu to reduce sync copy from CPU->GPU->CPU, and reduce pipeline hang
|
||||
seed_op = main_block.append_op(
|
||||
type='seed',
|
||||
outputs={'Out': seed_var},
|
||||
attrs={
|
||||
'deterministic': True,
|
||||
'rng_name': rng_name,
|
||||
'force_cpu': True,
|
||||
},
|
||||
)
|
||||
seed_op._set_attr('op_namescope', 'auto_tensor_parallel_seed')
|
||||
# set new seed op's dist_attr
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
||||
seed_op,
|
||||
process_mesh,
|
||||
seed_var_dims_mapping,
|
||||
ctx,
|
||||
chunk_id=op_dist_attr.chunk_id,
|
||||
)
|
||||
|
||||
# modify dropout op
|
||||
src_op.desc.set_input("seed_tensor", [seed_var.name])
|
||||
src_op._remove_attr("fix_seed")
|
||||
src_op._remove_attr("seed")
|
||||
op_dist_attr.set_input_dist_attr(
|
||||
seed_var.name, seed_var_dist_attr
|
||||
)
|
||||
kwargs['seed_tensor'] = [seed_var.name]
|
||||
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
# dropout backward is deterministic by mask, and not need for random state control
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"fused_dropout_add", DistributedDropoutImpl0("random_control")
|
||||
)
|
||||
@@ -0,0 +1,228 @@
|
||||
# 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.
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from ..utils import (
|
||||
_get_comm_group,
|
||||
_get_corresponding_rank,
|
||||
compute_compatible_and_update_dim_mapping,
|
||||
is_dim_replicate,
|
||||
is_dim_shard,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedFusedFeedForward(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedFusedFeedForward("fused_feedforward")
|
||||
)
|
||||
|
||||
|
||||
class DistributedFusedFeedForwardImpl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
linear1_weight = op_desc.input('Linear1Weight')[0]
|
||||
linear1_bias = op_desc.input('Linear1Bias')[0]
|
||||
linear2_weight = op_desc.input('Linear2Weight')[0]
|
||||
linear2_bias = op_desc.input('Linear2Bias')[0]
|
||||
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
linear1_weight_dims_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
linear1_weight
|
||||
)
|
||||
linear1_bias_dims_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
linear1_bias
|
||||
)
|
||||
linear2_weight_dims_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
linear2_weight
|
||||
)
|
||||
linear2_bias_dims_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
linear2_bias
|
||||
)
|
||||
|
||||
for mapping in x_dims_mapping[1:-1]:
|
||||
if is_dim_shard(mapping):
|
||||
return False
|
||||
if is_dim_shard(linear1_weight_dims_mapping[-2]) or is_dim_replicate(
|
||||
linear1_weight_dims_mapping[-1]
|
||||
):
|
||||
return False
|
||||
if is_dim_replicate(linear1_bias_dims_mapping[-1]):
|
||||
return False
|
||||
if is_dim_replicate(linear2_weight_dims_mapping[-2]) or is_dim_shard(
|
||||
linear2_weight_dims_mapping[-1]
|
||||
):
|
||||
return False
|
||||
if is_dim_shard(linear2_bias_dims_mapping[-1]):
|
||||
return False
|
||||
if linear1_weight_dims_mapping[-1] != linear2_weight_dims_mapping[-2]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
|
||||
# none of output should be sharded
|
||||
for out_name in op_desc.output_names():
|
||||
out = op_desc.output(out_name)[0]
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out)
|
||||
for mapping in out_dims_mapping[1:-1]:
|
||||
if is_dim_shard(mapping):
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_names = op_desc.output('Out')
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
for out_name in out_names:
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
if x_dims_mapping != out_dims_mapping:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_names = op_desc.output('Out')
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
|
||||
for out_name in out_names:
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
for i in range(len(x_dims_mapping)):
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[x_dims_mapping, out_dims_mapping], [i, i]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
out_name, out_dims_mapping
|
||||
)
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
|
||||
if rank_id not in op_dist_attr.process_mesh.process_ids:
|
||||
rank_id = _get_corresponding_rank(
|
||||
ctx, op_dist_attr.process_mesh, rank_id
|
||||
)
|
||||
|
||||
# infer logic comm presentation
|
||||
linear1_weight = src_op.input('Linear1Weight')[0]
|
||||
linear1_weight_col_dim_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
linear1_weight
|
||||
)[-1]
|
||||
assert linear1_weight_col_dim_mapping >= 0, (
|
||||
f"col_parallel_matmul's row should be divided by a specific mesh axis, but got [{linear1_weight_col_dim_mapping}]"
|
||||
)
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
process_mesh_group = op_dist_attr.process_mesh.process_ids
|
||||
|
||||
parallel_axis = linear1_weight_col_dim_mapping
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh_group, process_mesh_shape, parallel_axis, rank_id
|
||||
)
|
||||
group = new_process_group(group_ranks)
|
||||
|
||||
# insert op
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
# setting comm id
|
||||
new_op = main_block.ops[-1]
|
||||
assert new_op.type == "fused_feedforward"
|
||||
new_op._set_attr("ring_id", int(group.id))
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
|
||||
if rank_id not in op_dist_attr.process_mesh.process_ids:
|
||||
rank_id = _get_corresponding_rank(
|
||||
ctx, op_dist_attr.process_mesh, rank_id
|
||||
)
|
||||
|
||||
# infer logic comm presentation
|
||||
linear2_weight = src_op.input('Linear2Weight')[0]
|
||||
linear2_weight_col_dim_mapping = op_dist_attr.get_input_dims_mapping(
|
||||
linear2_weight
|
||||
)[-1]
|
||||
assert linear2_weight_col_dim_mapping >= 0, (
|
||||
f"col_parallel_matmul's row should be divided by a specific mesh axis, but got [{linear2_weight_col_dim_mapping}]"
|
||||
)
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
process_mesh_group = op_dist_attr.process_mesh.process_ids
|
||||
|
||||
parallel_axis = linear2_weight_col_dim_mapping
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh_group, process_mesh_shape, parallel_axis, rank_id
|
||||
)
|
||||
group = new_process_group(group_ranks)
|
||||
|
||||
# insert op
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
# setting comm id
|
||||
new_op = main_block.ops[-1]
|
||||
assert new_op.type == "fused_feedforward_grad"
|
||||
new_op._set_attr("ring_id", int(group.id))
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"fused_feedforward", DistributedFusedFeedForwardImpl("tensor_parallel")
|
||||
)
|
||||
@@ -0,0 +1,94 @@
|
||||
# Copyright (c) 2024 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 logging
|
||||
|
||||
from paddle.base.log_helper import get_logger
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..utils import get_dist_tensor_spec
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
|
||||
class DistributedLayerNorm(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
|
||||
x_name = op_desc.input('x')[0]
|
||||
scale_name = op_desc.input('scale')[0]
|
||||
y_name = op_desc.output('y')[0]
|
||||
invvar_name = op_desc.output('invvar')[0]
|
||||
|
||||
x_spec = get_dist_tensor_spec(dist_op, x_name)
|
||||
scale_spec = get_dist_tensor_spec(dist_op, scale_name)
|
||||
|
||||
y_spec = get_dist_tensor_spec(dist_op, y_name, is_input=False)
|
||||
invvar_spec = get_dist_tensor_spec(dist_op, invvar_name, is_input=False)
|
||||
|
||||
epsilon = op_desc.attr('epsilon')
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("fused_rms_norm")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(x_spec, scale_spec, epsilon)
|
||||
bw_results = rule.infer_backward(
|
||||
x_spec,
|
||||
scale_spec,
|
||||
y_spec,
|
||||
invvar_spec,
|
||||
epsilon,
|
||||
)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op,
|
||||
[x_name, scale_name],
|
||||
[y_name, invvar_name],
|
||||
fw_results,
|
||||
bw_results,
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
|
||||
# default impl
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedLayerNorm("fused_rms_norm")
|
||||
)
|
||||
@@ -0,0 +1,189 @@
|
||||
# Copyright (c) 2023 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
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..dist_attribute import DistTensorSpec
|
||||
from ..utils import get_dist_tensor_spec
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
|
||||
class DistributedFusedRope(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args), build fake spec for optional args
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
input_parameters = op_desc.input_names()
|
||||
output_parameters = op_desc.output_names()
|
||||
is_input_arg_exist = lambda parameter: (
|
||||
parameter in input_parameters and op_desc.input(parameter)
|
||||
)
|
||||
is_output_arg_exist = lambda parameter: (
|
||||
parameter in output_parameters and op_desc.output(parameter)
|
||||
)
|
||||
|
||||
q = op_desc.input('q')[0]
|
||||
k = op_desc.input('k')[0] if is_input_arg_exist('k') else None
|
||||
v = op_desc.input('v')[0] if is_input_arg_exist('v') else None
|
||||
sin = op_desc.input('sin')[0] if is_input_arg_exist('sin') else None
|
||||
cos = op_desc.input('cos')[0] if is_input_arg_exist('cos') else None
|
||||
position_ids = (
|
||||
op_desc.input('position_ids')[0]
|
||||
if is_input_arg_exist('position_ids')
|
||||
else None
|
||||
)
|
||||
out_q = op_desc.output('out_q')[0]
|
||||
out_k = (
|
||||
op_desc.output('out_k')[0] if is_output_arg_exist('out_k') else None
|
||||
)
|
||||
out_v = (
|
||||
op_desc.output('out_v')[0] if is_output_arg_exist('out_v') else None
|
||||
)
|
||||
|
||||
q_spec = get_dist_tensor_spec(dist_op, q)
|
||||
k_spec = (
|
||||
get_dist_tensor_spec(dist_op, k)
|
||||
if k is not None
|
||||
else DistTensorSpec()
|
||||
)
|
||||
v_spec = (
|
||||
get_dist_tensor_spec(dist_op, v)
|
||||
if v is not None
|
||||
else DistTensorSpec()
|
||||
)
|
||||
sin_spec = (
|
||||
get_dist_tensor_spec(dist_op, sin)
|
||||
if sin is not None
|
||||
else DistTensorSpec()
|
||||
)
|
||||
cos_spec = (
|
||||
get_dist_tensor_spec(dist_op, cos)
|
||||
if cos is not None
|
||||
else DistTensorSpec()
|
||||
)
|
||||
position_ids_spec = (
|
||||
get_dist_tensor_spec(dist_op, position_ids)
|
||||
if position_ids is not None
|
||||
else DistTensorSpec()
|
||||
)
|
||||
out_q_spec = get_dist_tensor_spec(dist_op, out_q, is_input=False)
|
||||
out_k_spec = (
|
||||
get_dist_tensor_spec(dist_op, out_k, is_input=False)
|
||||
if out_k is not None
|
||||
else DistTensorSpec()
|
||||
)
|
||||
out_v_spec = (
|
||||
get_dist_tensor_spec(dist_op, out_v, is_input=False)
|
||||
if out_v is not None
|
||||
else DistTensorSpec()
|
||||
)
|
||||
|
||||
use_neox_rotary_style = op_desc.attr("use_neox_rotary_style")
|
||||
time_major = op_desc.attr("time_major")
|
||||
rotary_emb_base = op_desc.attr("rotary_emb_base")
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("fused_rotary_position_embedding")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(
|
||||
q_spec,
|
||||
k_spec,
|
||||
v_spec,
|
||||
sin_spec,
|
||||
cos_spec,
|
||||
position_ids_spec,
|
||||
use_neox_rotary_style,
|
||||
time_major,
|
||||
rotary_emb_base,
|
||||
)
|
||||
bw_results = rule.infer_backward(
|
||||
q_spec,
|
||||
k_spec,
|
||||
v_spec,
|
||||
sin_spec,
|
||||
cos_spec,
|
||||
position_ids_spec,
|
||||
out_q_spec,
|
||||
out_k_spec,
|
||||
out_v_spec,
|
||||
use_neox_rotary_style,
|
||||
time_major,
|
||||
rotary_emb_base,
|
||||
)
|
||||
|
||||
# remove optional args in spmd results
|
||||
input_args = [q, k, v, sin, cos, position_ids]
|
||||
output_args = [out_q, out_k, out_v]
|
||||
fw_and_bw_results_without_optional_arg = []
|
||||
for results in [fw_results, bw_results]:
|
||||
input_results = results[0]
|
||||
output_results = results[1]
|
||||
input_results_without_optional_arg = []
|
||||
output_results_without_optional_arg = []
|
||||
for idx, input_arg in enumerate(input_args):
|
||||
if input_arg is not None:
|
||||
input_results_without_optional_arg.append(
|
||||
input_results[idx]
|
||||
)
|
||||
for idx, output_arg in enumerate(output_args):
|
||||
if output_arg is not None:
|
||||
output_results_without_optional_arg.append(
|
||||
output_results[idx]
|
||||
)
|
||||
fw_and_bw_results_without_optional_arg.append(
|
||||
[
|
||||
input_results_without_optional_arg,
|
||||
output_results_without_optional_arg,
|
||||
]
|
||||
)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op,
|
||||
input_arg_names=[
|
||||
input_arg for input_arg in input_args if input_arg is not None
|
||||
],
|
||||
output_arg_names=[
|
||||
output_arg
|
||||
for output_arg in output_args
|
||||
if output_arg is not None
|
||||
],
|
||||
fw_results=fw_and_bw_results_without_optional_arg[0],
|
||||
bw_results=fw_and_bw_results_without_optional_arg[1],
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedFusedRope("fused_rotary_position_embedding")
|
||||
)
|
||||
@@ -0,0 +1,70 @@
|
||||
# Copyright (c) 2024 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
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..utils import get_dist_tensor_spec
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
|
||||
class DistributedGatherNd(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
|
||||
x_name = op_desc.input('X')[0]
|
||||
index_name = op_desc.input('Index')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
|
||||
x_specs = get_dist_tensor_spec(dist_op, x_name)
|
||||
index_specs = get_dist_tensor_spec(dist_op, index_name)
|
||||
output_spec = get_dist_tensor_spec(dist_op, out_name, False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("gather_nd")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(x_specs, index_specs)
|
||||
bw_results = rule.infer_backward(x_specs, index_specs, output_spec)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op,
|
||||
[x_name, index_name],
|
||||
[out_name],
|
||||
fw_results,
|
||||
bw_results,
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedGatherNd("gather_nd"))
|
||||
@@ -0,0 +1,151 @@
|
||||
# Copyright (c) 2023 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 logging
|
||||
|
||||
from paddle.base.log_helper import get_logger
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..dist_attribute import DistTensorSpec, TensorDistAttr
|
||||
from ..utils import get_dist_tensor_spec, is_dim_shard
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
|
||||
class DistributedLayerNorm(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
|
||||
x_name = op_desc.input('X')[0]
|
||||
scale_name = (
|
||||
op_desc.input('Scale')[0]
|
||||
if len(op_desc.input('Scale')) > 0
|
||||
else None
|
||||
)
|
||||
bias_name = (
|
||||
op_desc.input('Bias')[0] if len(op_desc.input('Bias')) > 0 else None
|
||||
)
|
||||
y_name = op_desc.output('Y')[0]
|
||||
var_name = op_desc.output('Variance')[0]
|
||||
mean_name = op_desc.output('Mean')[0]
|
||||
begin_norm_axis = op_desc.attr('begin_norm_axis')
|
||||
|
||||
x_spec = get_dist_tensor_spec(dist_op, x_name)
|
||||
scale_spec = (
|
||||
DistTensorSpec([0], TensorDistAttr())
|
||||
if scale_name is None
|
||||
else get_dist_tensor_spec(dist_op, scale_name)
|
||||
)
|
||||
bias_spec = (
|
||||
DistTensorSpec([0], TensorDistAttr())
|
||||
if bias_name is None
|
||||
else get_dist_tensor_spec(dist_op, bias_name)
|
||||
)
|
||||
y_spec = get_dist_tensor_spec(dist_op, y_name, False)
|
||||
var_spec = get_dist_tensor_spec(dist_op, var_name, False)
|
||||
mean_spec = get_dist_tensor_spec(dist_op, mean_name, False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("layer_norm")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(
|
||||
x_spec, scale_spec, bias_spec, 1.0, begin_norm_axis
|
||||
)
|
||||
bw_results = rule.infer_backward(
|
||||
x_spec,
|
||||
scale_spec,
|
||||
bias_spec,
|
||||
y_spec,
|
||||
var_spec,
|
||||
mean_spec,
|
||||
1.0,
|
||||
begin_norm_axis,
|
||||
)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
input_arg_names = [x_name]
|
||||
if scale_name is not None:
|
||||
input_arg_names.append(scale_name)
|
||||
if bias_name is not None:
|
||||
input_arg_names.append(bias_name)
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op,
|
||||
input_arg_names,
|
||||
[y_name, var_name, mean_name],
|
||||
fw_results,
|
||||
bw_results,
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
begin_norm_axis = op_desc.attr('begin_norm_axis')
|
||||
|
||||
# sharded on begin_norm_axis
|
||||
x_name = op_desc.input('X')[0]
|
||||
x_dims_mapping = copy.deepcopy(
|
||||
op_dist_attr.get_input_dims_mapping(x_name)
|
||||
)
|
||||
if (begin_norm_axis > 0) and is_dim_shard(
|
||||
x_dims_mapping[begin_norm_axis]
|
||||
):
|
||||
# TODO (ljz) support sharding on `begin_norm_axis`
|
||||
_logger.info(
|
||||
"sharding on `begin_norm_axis` is not supported yet, we resharded it as replicated"
|
||||
)
|
||||
x_dims_mapping[begin_norm_axis] = -1
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
|
||||
param_names = [op_desc.input('Scale')[0], op_desc.input('Bias')[0]]
|
||||
for p_name in param_names:
|
||||
p_dims_mapping = copy.deepcopy(
|
||||
op_dist_attr.get_input_dims_mapping(p_name)
|
||||
)
|
||||
p_dims_mapping[begin_norm_axis] = -1
|
||||
op_dist_attr.set_input_dims_mapping(p_name, p_dims_mapping)
|
||||
|
||||
y_name = op_desc.output('Y')[0]
|
||||
y_dims_mapping = copy.deepcopy(
|
||||
op_dist_attr.get_output_dims_mapping(y_name)
|
||||
)
|
||||
y_dims_mapping[begin_norm_axis] = -1
|
||||
op_dist_attr.set_input_dims_mapping(y_name, y_dims_mapping)
|
||||
|
||||
# default impl
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedLayerNorm("layer_norm"))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,387 @@
|
||||
# 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
|
||||
|
||||
from paddle.common_ops_import import check_dtype, check_variable_and_dtype
|
||||
from paddle.distributed.utils.stream_utils import ExecutionStreamType
|
||||
from paddle.framework import core
|
||||
from paddle.static import Operator
|
||||
|
||||
from ..dist_attribute import OperatorDistAttr, TensorDistAttr
|
||||
from ..process_group import new_process_group
|
||||
from ..utils import (
|
||||
_get_comm_group,
|
||||
_get_corresponding_rank,
|
||||
compute_compatible_dim_mapping,
|
||||
is_dim_replicate,
|
||||
is_dim_shard,
|
||||
set_dist_op_desc_original_id,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
|
||||
|
||||
class DistributedPNorm(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedPNorm("p_norm"))
|
||||
|
||||
|
||||
# Data Parallel
|
||||
class DistributedPNormImpl0(DistributedOperatorImpl):
|
||||
"""
|
||||
TODO: p_norm scene
|
||||
|
||||
1. axis == None, isinstance(p, (int, float)), asvector = True
|
||||
1.1 x_dims_mapping == [0, -1, -1]
|
||||
allgather input if it is split by dp group
|
||||
1.2 x_dims_mapping == [-1, 0, -1]
|
||||
allgather, split and concat input if it is split by mp group
|
||||
2. isinstance(axis, int), asvector = False
|
||||
1.1 axis == 0 and x_dims_mapping == [0, -1, -1]
|
||||
allgather input if it's input[0] is splited by dp group.
|
||||
1.2 axis == 1 and x_dims_mapping == [-1, 0, -1]
|
||||
allgather, split and concat input if it's input[1] is split by mp group
|
||||
"""
|
||||
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
axis = op_desc.attr('axis')
|
||||
asvector = op_desc.attr('asvector')
|
||||
x_name = op_desc.input('X')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
if is_dim_replicate(x_dims_mapping[0]):
|
||||
return False
|
||||
# Other dimensions must be replicate except the batch dimension
|
||||
for mapping in x_dims_mapping[1:]:
|
||||
if is_dim_shard(mapping):
|
||||
return False
|
||||
if not (axis == -1 and asvector) and not (axis == 0 and not asvector):
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (
|
||||
(not self.is_input_compatible(dist_op))
|
||||
or (not self.is_output_compatible(dist_op))
|
||||
or (not self.is_compatible(dist_op))
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
axis = op_desc.attr('axis')
|
||||
keepdim = op_desc.attr('keepdim')
|
||||
|
||||
batch_dim_mappings = []
|
||||
for arg_name in op_desc.input_arg_names():
|
||||
dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
|
||||
if len(dims_mapping) >= 1:
|
||||
batch_dim_mappings.append(dims_mapping[0])
|
||||
for arg_name in op_desc.output_arg_names():
|
||||
dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
|
||||
if len(dims_mapping) >= 1:
|
||||
batch_dim_mappings.append(dims_mapping[0])
|
||||
|
||||
compatible_dim_mapping = compute_compatible_dim_mapping(
|
||||
batch_dim_mappings
|
||||
)
|
||||
if compatible_dim_mapping is None:
|
||||
return False
|
||||
|
||||
for arg_name in op_desc.input_arg_names():
|
||||
dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
|
||||
if (
|
||||
len(dims_mapping) >= 1
|
||||
and compatible_dim_mapping != dims_mapping[0]
|
||||
):
|
||||
dims_mapping[0] = compatible_dim_mapping
|
||||
op_dist_attr.set_input_dims_mapping(arg_name, dims_mapping)
|
||||
changed = True
|
||||
|
||||
if axis == 0 and not keepdim:
|
||||
for arg_name in op_desc.output_arg_names():
|
||||
dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
|
||||
if len(dims_mapping) >= 1 and dims_mapping[0] != -1:
|
||||
dims_mapping[0] = -1
|
||||
op_dist_attr.set_output_dims_mapping(arg_name, dims_mapping)
|
||||
changed = True
|
||||
else:
|
||||
for arg_name in op_desc.output_arg_names():
|
||||
dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
|
||||
if (
|
||||
len(dims_mapping) >= 1
|
||||
and compatible_dim_mapping != dims_mapping[0]
|
||||
):
|
||||
dims_mapping[0] = compatible_dim_mapping
|
||||
op_dist_attr.set_output_dims_mapping(arg_name, dims_mapping)
|
||||
changed = True
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
assert op_dist_attr is not None
|
||||
|
||||
# check validation of inputs / outputs
|
||||
for input_name in src_op.desc.input_names():
|
||||
assert input_name in kwargs, f"input [{input_name}] is not given"
|
||||
assert len(kwargs[input_name]) == len(
|
||||
src_op.desc.input(input_name)
|
||||
), f"number of tensor for input [{input_name}] is not match"
|
||||
for output_name in src_op.desc.output_names():
|
||||
assert output_name in kwargs, f"input [{output_name}] is not given"
|
||||
assert len(kwargs[output_name]) == len(
|
||||
src_op.desc.output(output_name)
|
||||
), f"number of tensor for input [{output_name}] is not match"
|
||||
|
||||
if rank_id not in op_dist_attr.process_mesh.process_ids:
|
||||
rank_id = _get_corresponding_rank(
|
||||
ctx, op_dist_attr.process_mesh, rank_id
|
||||
)
|
||||
|
||||
X_var = main_block._var_recursive(kwargs['X'][0])
|
||||
in_dims_mapping = op_dist_attr.get_input_dims_mapping(X_var.name)
|
||||
for axis in range(len(in_dims_mapping)):
|
||||
if in_dims_mapping[axis] != -1:
|
||||
break
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
process_mesh_group = op_dist_attr.process_mesh.process_ids
|
||||
group_ranks = _get_comm_group(
|
||||
process_mesh_group, process_mesh_shape, axis, rank_id
|
||||
)
|
||||
group = new_process_group(group_ranks)
|
||||
|
||||
check_variable_and_dtype(
|
||||
X_var, 'x', ['float16', 'float32', 'float64'], 'norm'
|
||||
)
|
||||
check_dtype(
|
||||
X_var.dtype, 'dtype', ['float16', 'float32', 'float64'], 'norm'
|
||||
)
|
||||
|
||||
# 2. insert all_gather op
|
||||
# create all_gather output var
|
||||
allgather_out = main_block.create_var(
|
||||
name=".".join(["all_gather", X_var.name]),
|
||||
dtype=X_var.dtype,
|
||||
shape=X_var.shape,
|
||||
type=core.VarDesc.VarType.DENSE_TENSOR,
|
||||
persistable=False,
|
||||
stop_gradient=X_var.stop_gradient,
|
||||
)
|
||||
# set allgather_out tensor dist_attr
|
||||
allgather_out_dist_attr = TensorDistAttr()
|
||||
allgather_out_dist_attr.process_mesh = op_dist_attr.process_mesh
|
||||
allgather_out_dist_attr.chunk_id = op_dist_attr.chunk_id
|
||||
allgather_out_dist_attr.dims_mapping = [
|
||||
-1 for i in range(len(allgather_out.shape))
|
||||
]
|
||||
ctx.set_tensor_dist_attr_for_program(
|
||||
allgather_out, allgather_out_dist_attr
|
||||
)
|
||||
all_gather_op = main_block.append_op(
|
||||
type='all_gather',
|
||||
inputs={'x': [X_var]},
|
||||
outputs={'out': [allgather_out]},
|
||||
attrs={
|
||||
'ring_id': group.id,
|
||||
'use_calc_stream': True,
|
||||
'nranks': group.nranks,
|
||||
'op_role': src_op.attr('op_role'),
|
||||
},
|
||||
)
|
||||
# set all_gather op dist_attr
|
||||
allgather_op_dist_attr = OperatorDistAttr()
|
||||
allgather_op_dist_attr.process_mesh = op_dist_attr.process_mesh
|
||||
allgather_op_dist_attr.chunk_id = op_dist_attr.chunk_id
|
||||
allgather_op_dist_attr.set_input_dims_mapping(
|
||||
X_var.name, in_dims_mapping
|
||||
)
|
||||
allgather_op_dist_attr.set_output_dims_mapping(
|
||||
allgather_out.name, allgather_out_dist_attr.dims_mapping
|
||||
)
|
||||
allgather_op_dist_attr.execution_stream = (
|
||||
ExecutionStreamType.DefaultStream.value
|
||||
)
|
||||
ctx.set_op_dist_attr_for_program(all_gather_op, allgather_op_dist_attr)
|
||||
|
||||
# 3. copy p_norm op desc and reset input name
|
||||
# rename input
|
||||
kwargs['X'] = [allgather_out.name]
|
||||
# replicate op in dist program
|
||||
dist_op = main_block.append_op(type='nop')
|
||||
dist_op_desc = dist_op.desc
|
||||
dist_op_desc.copy_from(src_op.desc)
|
||||
set_dist_op_desc_original_id(dist_op_desc, src_op.desc, ctx)
|
||||
for input_name in src_op.desc.input_names():
|
||||
dist_op_desc.set_input(input_name, kwargs[input_name])
|
||||
for output_name in src_op.desc.output_names():
|
||||
dist_op_desc.set_output(output_name, kwargs[output_name])
|
||||
pnorm_op = Operator(main_block, dist_op_desc)
|
||||
op_dist_attr.set_input_dims_mapping(
|
||||
allgather_out.name, allgather_out_dist_attr.dims_mapping
|
||||
)
|
||||
# Remove the unrelated dist attr
|
||||
op_dist_attr.del_input_dist_attr(X_var.name)
|
||||
ctx.set_op_dist_attr_for_program(pnorm_op, op_dist_attr)
|
||||
# TODO: should we add a new dist attr for the new op here?
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
backward_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
|
||||
assert op_dist_attr is not None
|
||||
|
||||
# check validation of inputs / outputs
|
||||
for input_name in backward_op.desc.input_names():
|
||||
assert input_name in kwargs, f"input [{input_name}] is not given"
|
||||
assert len(kwargs[input_name]) == len(
|
||||
backward_op.desc.input(input_name)
|
||||
), f"number of tensor for input [{input_name}] is not match"
|
||||
for output_name in backward_op.desc.output_names():
|
||||
assert output_name in kwargs, f"input [{output_name}] is not given"
|
||||
assert len(kwargs[output_name]) == len(
|
||||
backward_op.desc.output(output_name)
|
||||
), f"number of tensor for input [{output_name}] is not match"
|
||||
|
||||
X_var = main_block._var_recursive(kwargs['X'][0])
|
||||
X_grad_var = main_block._var_recursive(kwargs['X@GRAD'][0])
|
||||
|
||||
# 1. copy p_norm_grad op and reset input name and output name
|
||||
new_kwargs = copy.deepcopy(kwargs)
|
||||
new_kwargs['X'] = [".".join(["all_gather", X_var.name])]
|
||||
new_X_var = main_block._var_recursive(new_kwargs['X'][0])
|
||||
new_X_grad = main_block.create_var(
|
||||
name=".".join(["all_gather", X_grad_var.name]),
|
||||
dtype=X_grad_var.dtype,
|
||||
shape=new_X_var.shape,
|
||||
type=core.VarDesc.VarType.DENSE_TENSOR,
|
||||
persistable=False,
|
||||
stop_gradient=X_grad_var.stop_gradient,
|
||||
)
|
||||
new_kwargs['X@GRAD'] = [new_X_grad.name]
|
||||
new_X_var_dist_attr = ctx.get_tensor_dist_attr_for_program(new_X_var)
|
||||
ctx.set_tensor_dist_attr_for_program(new_X_grad, new_X_var_dist_attr)
|
||||
# replicate op in dist program with new kwargs
|
||||
dist_op = main_block.append_op(type='nop')
|
||||
dist_op_desc = dist_op.desc
|
||||
dist_op_desc.copy_from(backward_op.desc)
|
||||
# Refer to the related dist op
|
||||
set_dist_op_desc_original_id(dist_op_desc, backward_op.desc, ctx)
|
||||
for input_name in backward_op.desc.input_names():
|
||||
dist_op_desc.set_input(input_name, new_kwargs[input_name])
|
||||
for output_name in backward_op.desc.output_names():
|
||||
dist_op_desc.set_output(output_name, new_kwargs[output_name])
|
||||
p_norm_grad_op = Operator(main_block, dist_op_desc)
|
||||
op_dist_attr.set_input_dims_mapping(
|
||||
new_X_var.name, new_X_var_dist_attr.dims_mapping
|
||||
)
|
||||
# Store X_grad_var dims_mapping for later use
|
||||
X_grad_var_dims_mapping = op_dist_attr.get_output_dims_mapping(
|
||||
X_grad_var.name
|
||||
)
|
||||
# Remove the unrelated dist attr
|
||||
op_dist_attr.del_input_dist_attr(X_var.name)
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
new_X_grad.name, new_X_var_dist_attr.dims_mapping
|
||||
)
|
||||
# Remove the unrelated dist attr
|
||||
op_dist_attr.del_output_dist_attr(X_grad_var.name)
|
||||
ctx.set_op_dist_attr_for_program(p_norm_grad_op, op_dist_attr)
|
||||
# TODO: should we add a new dist attr for the new op here?
|
||||
|
||||
# 2. insert slice op
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
process_mesh_group = op_dist_attr.process_mesh.process_ids
|
||||
dims_mapping = [0] + [-1 for _ in range(len(new_X_grad.shape) - 1)]
|
||||
from ..reshard import Resharder
|
||||
|
||||
partition_idx = Resharder.compute_partition_index(
|
||||
rank_id,
|
||||
new_X_grad.shape,
|
||||
dims_mapping,
|
||||
process_mesh_shape,
|
||||
process_mesh_group,
|
||||
)
|
||||
slice_starts = []
|
||||
slice_ends = []
|
||||
slices_axes = []
|
||||
for idx, item in enumerate(partition_idx):
|
||||
slice_starts.append(item[0])
|
||||
slice_ends.append(item[1])
|
||||
slices_axes.append(idx)
|
||||
|
||||
infer_flags = [1 for i in range(len(slices_axes))]
|
||||
attrs = {
|
||||
"axes": slices_axes,
|
||||
"starts": slice_starts,
|
||||
"ends": slice_ends,
|
||||
"infer_flags": infer_flags,
|
||||
"op_role": backward_op.attr('op_role'),
|
||||
}
|
||||
slice_op = main_block.append_op(
|
||||
type='slice',
|
||||
inputs={'Input': [new_X_grad]},
|
||||
outputs={'Out': [X_grad_var]},
|
||||
attrs=attrs,
|
||||
)
|
||||
slice_op_dist_attr = OperatorDistAttr()
|
||||
slice_op_dist_attr.process_mesh = op_dist_attr.process_mesh
|
||||
slice_op_dist_attr.chunk_id = op_dist_attr.chunk_id
|
||||
slice_op_dist_attr.set_input_dims_mapping(
|
||||
new_X_grad.name, new_X_var_dist_attr.dims_mapping
|
||||
)
|
||||
slice_op_dist_attr.set_output_dims_mapping(
|
||||
X_grad_var.name, X_grad_var_dims_mapping
|
||||
)
|
||||
ctx.set_op_dist_attr_for_program(slice_op, slice_op_dist_attr)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"p_norm", DistributedPNormImpl0("data_parallel")
|
||||
)
|
||||
@@ -0,0 +1,240 @@
|
||||
# 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 paddle
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..dist_attribute import OperatorDistAttr
|
||||
from ..process_group import new_process_group
|
||||
from ..utils import (
|
||||
get_dist_tensor_spec,
|
||||
is_dim_shard,
|
||||
set_dist_op_desc_original_id,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
|
||||
class DistributedReduceSum(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
assert len(op_desc.input_arg_names()) == 1, (
|
||||
f"reduce_sum op [{op_desc.type}] has [{len(op_desc.input_arg_names())}] inputs"
|
||||
)
|
||||
input_arg_name = op_desc.input_arg_names()[0]
|
||||
assert len(op_desc.output_arg_names()) == 1, (
|
||||
f"reduce_sum op [{op_desc.type}] has [{len(op_desc.output_arg_names())}] outputs"
|
||||
)
|
||||
output_arg_name = op_desc.output_arg_names()[0]
|
||||
keep_dim = op_desc.attr('keep_dim')
|
||||
dims = op_desc.attr('dim')
|
||||
|
||||
# TODO (zhangyichen) replace dist tensor spec by dist tensor in future.
|
||||
input_spec = get_dist_tensor_spec(dist_op, input_arg_name)
|
||||
output_spec = get_dist_tensor_spec(dist_op, output_arg_name, False)
|
||||
# len(dims) == 0 means reduce_all
|
||||
if len(dims) == 0:
|
||||
dims = list(range(len(input_spec.shape)))
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("reduce_sum")
|
||||
fw_results = rule.infer_forward(input_spec, dims, keep_dim)
|
||||
bw_results = rule.infer_backward(
|
||||
input_spec, output_spec, dims, keep_dim
|
||||
)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op, [input_arg_name], [output_arg_name], fw_results, bw_results
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
# NOTE this function will be remove once we use local reshard to replace distopimpls
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
op_desc = dist_op.serial_op.desc
|
||||
input_name = op_desc.input_arg_names()[0]
|
||||
input_dims_mapping = copy.deepcopy(
|
||||
op_dist_attr.get_input_dims_mapping(input_name)
|
||||
)
|
||||
axes = op_desc.attr('dim')
|
||||
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
reverted = False
|
||||
|
||||
def is_partial_reduce(axes, dims_mapping):
|
||||
# FIXME(ljz) Hack for performance:
|
||||
# if the reduce result is a scalar, it is the loss reduce in GPT case,
|
||||
# and if any axis of reduce input is sharded, the result loss would be partial.
|
||||
# BUT we keep the loss as partial instead of allreduce it for performance, since it would effect the backward.
|
||||
# we should use an optimization pass for the Hack in future.
|
||||
if len(axes) != 0 and (len(axes) < len(dims_mapping)):
|
||||
for axis in axes:
|
||||
if is_dim_shard(dims_mapping[axis]):
|
||||
return True # reverted
|
||||
return False
|
||||
|
||||
# if reduce_axis is sharded, the output is partial and need to be allreduce
|
||||
if is_partial_reduce(axes, input_dims_mapping):
|
||||
# TODO (ljz) support reduce where the reduce_axis is sharded
|
||||
dist_op.dist_attr = original_op_dist_attr
|
||||
reverted = True
|
||||
# if reduce_axis is unsharded, NO extra operator need.
|
||||
else:
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
return reverted
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedReduceSum("reduce_sum"))
|
||||
|
||||
|
||||
class DistributedReduceSumPrimitive(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedReduceSumPrimitive("reduce_sum_p")
|
||||
)
|
||||
|
||||
|
||||
# Batch Dimension ReduceSum Primitive
|
||||
class DistributedReduceSumPrimitiveImpl0(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
|
||||
return len(op_desc.input_arg_names()) == 1
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
outputs = op_desc.output_arg_names()
|
||||
|
||||
if len(outputs) != 1:
|
||||
return False
|
||||
|
||||
output_name = outputs[0]
|
||||
output_var = dist_op.serial_op.block._var_recursive(output_name)
|
||||
if output_var.shape != ():
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
|
||||
return self.is_input_compatible(dist_op) and self.is_output_compatible(
|
||||
dist_op
|
||||
)
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
startup_block = dist_op_context.startup_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
|
||||
# check validation of inputs / outputs
|
||||
for input_name in src_op.desc.input_names():
|
||||
assert input_name in kwargs, f"input [{input_name}] is not given"
|
||||
assert len(kwargs[input_name]) == len(
|
||||
src_op.desc.input(input_name)
|
||||
), f"number of tensor for input [{input_name}] is not match"
|
||||
for output_name in src_op.desc.output_names():
|
||||
assert output_name in kwargs, f"input [{output_name}] is not given"
|
||||
assert len(kwargs[output_name]) == len(
|
||||
src_op.desc.output(output_name)
|
||||
), f"number of tensor for input [{output_name}] is not match"
|
||||
|
||||
# replicate op in dist program
|
||||
dist_op = main_block.append_op(type='nop')
|
||||
dist_op_desc = dist_op.desc
|
||||
dist_op_desc.copy_from(src_op.desc)
|
||||
set_dist_op_desc_original_id(dist_op_desc, src_op.desc, ctx)
|
||||
for input_name in src_op.desc.input_names():
|
||||
dist_op_desc.set_input(input_name, kwargs[input_name])
|
||||
for output_name in src_op.desc.output_names():
|
||||
dist_op_desc.set_output(output_name, kwargs[output_name])
|
||||
# TODO: should we add a new dist attr for the new op here?
|
||||
|
||||
# batch dimension synchronization
|
||||
var_name = src_op.output_arg_names[0]
|
||||
sync_group = new_process_group(ctx.data_parallel_group)
|
||||
allreduce_op = main_block.append_op(
|
||||
type='all_reduce',
|
||||
inputs={'x': [var_name]},
|
||||
outputs={'out': [var_name]},
|
||||
attrs={
|
||||
'ring_id': sync_group.id,
|
||||
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
||||
OP_ROLE_KEY: OpRole.Forward,
|
||||
},
|
||||
)
|
||||
|
||||
# dist attr
|
||||
var = main_block._var_recursive(var_name)
|
||||
tensor_dist_attr = ctx.get_tensor_dist_attr_for_program(var)
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
new_op_attr = OperatorDistAttr()
|
||||
new_op_attr.process_mesh = op_dist_attr.process_mesh
|
||||
new_op_attr.set_output_dims_mapping(
|
||||
var.name, tensor_dist_attr.dims_mapping
|
||||
)
|
||||
new_op_attr.set_input_dims_mapping(
|
||||
var.name, tensor_dist_attr.dims_mapping
|
||||
)
|
||||
ctx.set_op_dist_attr_for_program(allreduce_op, new_op_attr)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
raise RuntimeError("primitive operator does NOT have backward function")
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"reduce_sum_p",
|
||||
DistributedReduceSumPrimitiveImpl0("batch_dimension_reduce_sum_p"),
|
||||
)
|
||||
@@ -0,0 +1,866 @@
|
||||
# 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
|
||||
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..cost import (
|
||||
Reshape2GradOpCost,
|
||||
Reshape2OpCost,
|
||||
build_comp_costs_from_descs,
|
||||
build_comp_desc_from_dist_op,
|
||||
build_dp_costs,
|
||||
)
|
||||
from ..utils import (
|
||||
compute_compatible_and_update_dim_mapping,
|
||||
get_dist_tensor_spec,
|
||||
is_dim_shard,
|
||||
set_dist_op_desc_original_id,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
is_parameter_related,
|
||||
merge_forward_backward_dims_mapping,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedReshape2(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
assert dist_op.serial_op.type == "reshape2", (
|
||||
f"{dist_op.serial_op.type} is not supported by dist reshape yet."
|
||||
)
|
||||
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
xshape_name = op_desc.output('XShape')[0]
|
||||
shape = op_desc.attr('shape')
|
||||
|
||||
x_spec = get_dist_tensor_spec(dist_op, x_name)
|
||||
output_spec = get_dist_tensor_spec(dist_op, out_name, False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("reshape")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(x_spec, shape)
|
||||
bw_results = rule.infer_backward(x_spec, output_spec, shape)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op, [x_name], [out_name], fw_results, bw_results
|
||||
)
|
||||
|
||||
# step4: update xshape
|
||||
inferred_input_dims_mappings, _ = merge_forward_backward_dims_mapping(
|
||||
fw_results, bw_results
|
||||
)
|
||||
dist_op.dist_attr.set_output_dims_mapping(
|
||||
xshape_name, [-1] + inferred_input_dims_mappings[0]
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
reverted = False
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
|
||||
# all reshape mapping to impl0
|
||||
op_dist_attr.impl_type = "reshape2"
|
||||
op_dist_attr.impl_idx = 0
|
||||
|
||||
return reverted
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedReshape2("reshape2"))
|
||||
|
||||
|
||||
class DistributedReshapeImpl0(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = False
|
||||
|
||||
def calc_cost(self, op_role, dist_op, ctx, cluster):
|
||||
cost = None
|
||||
if int(op_role) == int(OpRole.Backward):
|
||||
cost = self.calc_bwd_cost(dist_op, ctx, cluster)
|
||||
else:
|
||||
cost = self.calc_fwd_cost(dist_op, ctx, cluster)
|
||||
assert cost is not None
|
||||
return cost
|
||||
|
||||
def calc_fwd_cost(self, dist_op, ctx, cluster):
|
||||
res = []
|
||||
op = dist_op.serial_op
|
||||
dist_attr = dist_op.dist_attr
|
||||
|
||||
shape_list = op.desc.attr("shape")
|
||||
# got dist attribute info
|
||||
dim_mapping = dist_attr.get_output_dims_mapping(op.output("Out")[0])
|
||||
process_mesh_shape = dist_attr.process_mesh.shape
|
||||
|
||||
# modify target shape
|
||||
for idx, axis in enumerate(dim_mapping):
|
||||
if axis >= 0:
|
||||
if len(shape_list) > idx:
|
||||
shape_list[idx] = (
|
||||
shape_list[idx] // process_mesh_shape[axis]
|
||||
)
|
||||
|
||||
# calc comp op cost
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
processes = dist_attr.process_mesh.process_ids
|
||||
for key in desc_mapping:
|
||||
desc_mapping[key]["shape"] = shape_list
|
||||
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
Reshape2OpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
return res
|
||||
|
||||
def calc_bwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
res = []
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
dist_attr = dist_op.dist_attr
|
||||
process_mesh = dist_attr.process_mesh
|
||||
processes = process_mesh.process_ids
|
||||
op_type = dist_op.serial_op.type
|
||||
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
Reshape2GradOpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
backward_op = dist_op.serial_op
|
||||
main_block = backward_op.block
|
||||
need_gradient_allreduce = False
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
# NOTE input var's dim_mapping of backward op should be the same with input var instead of corresponding varname of forward op
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
|
||||
|
||||
mesh_shape = process_mesh.shape
|
||||
batch_size_axis = (
|
||||
var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
|
||||
)
|
||||
if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
|
||||
parallel_axis = batch_size_axis
|
||||
attrs = {"use_calc_stream": True}
|
||||
var_names = [varname + "@GRAD"]
|
||||
build_dp_costs(
|
||||
res,
|
||||
dist_op,
|
||||
ctx,
|
||||
var_names,
|
||||
attrs,
|
||||
parallel_axis,
|
||||
cluster,
|
||||
)
|
||||
|
||||
return res
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
if len(x_dims_mapping) != len(out_dims_mapping) - 1:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
if len(x_dims_mapping) != len(out_dims_mapping) - 1:
|
||||
return False
|
||||
|
||||
if is_dim_shard(out_dims_mapping[-1]):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_shape_name = op_desc.output('XShape')[0]
|
||||
x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
|
||||
x_shape_name
|
||||
)
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
for idx, dim_mapping in enumerate(out_dims_mapping[:-1]):
|
||||
if x_dims_mapping[idx] != dim_mapping:
|
||||
return False
|
||||
|
||||
if x_shape_dims_mapping[0] != -1:
|
||||
return False
|
||||
|
||||
if x_shape_dims_mapping[1:] != x_dims_mapping[:]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_shape_name = op_desc.output('XShape')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
|
||||
x_shape_name
|
||||
)
|
||||
|
||||
for i in range(len(x_dims_mapping)):
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[x_dims_mapping, out_dims_mapping], [i, i]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
|
||||
for i in range(len(x_dims_mapping)):
|
||||
x_shape_dims_mapping[i + 1] = x_dims_mapping[i]
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
x_shape_name, x_shape_dims_mapping
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
"""
|
||||
kwargs: inputname_mapping & outputname_mapping
|
||||
"""
|
||||
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
assert op_dist_attr is not None, (
|
||||
f"backward op [{src_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
# check validation of inputs / outputs
|
||||
for input_name in src_op.desc.input_names():
|
||||
assert input_name in kwargs, f"input [{input_name}] is not given"
|
||||
assert len(kwargs[input_name]) == len(
|
||||
src_op.desc.input(input_name)
|
||||
), f"number of tensor for input [{input_name}] is not match"
|
||||
for output_name in src_op.desc.output_names():
|
||||
assert output_name in kwargs, f"input [{output_name}] is not given"
|
||||
assert len(kwargs[output_name]) == len(
|
||||
src_op.desc.output(output_name)
|
||||
), f"number of tensor for input [{output_name}] is not match"
|
||||
|
||||
X_var = main_block._var_recursive(kwargs['X'][0])
|
||||
Out_var = main_block._var_recursive(kwargs['Out'][0])
|
||||
XShape_var = main_block._var_recursive(kwargs['XShape'][0])
|
||||
shape_list = src_op.desc.attr("shape")
|
||||
ShapeTensor_var_list = []
|
||||
for name in kwargs['ShapeTensor']:
|
||||
ShapeTensor_var_list.append(name)
|
||||
Shape_var_list = []
|
||||
for name in kwargs['Shape']:
|
||||
Shape_var_list.append(name)
|
||||
|
||||
# got dist attribute info
|
||||
dim_mapping = op_dist_attr.get_output_dims_mapping(Out_var.name)
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
|
||||
# modify target shape
|
||||
for idx, axis in enumerate(dim_mapping):
|
||||
if axis >= 0:
|
||||
if len(shape_list) > idx:
|
||||
shape_list[idx] = (
|
||||
shape_list[idx] // process_mesh_shape[axis]
|
||||
)
|
||||
|
||||
# create op
|
||||
new_op = main_block.append_op(type='nop')
|
||||
new_op_desc = new_op.desc
|
||||
new_op_desc.copy_from(src_op.desc)
|
||||
set_dist_op_desc_original_id(new_op_desc, src_op.desc, ctx)
|
||||
new_op_desc.set_input('ShapeTensor', ShapeTensor_var_list)
|
||||
new_op_desc.set_input('Shape', Shape_var_list)
|
||||
new_op_desc.set_input('X', [X_var.name])
|
||||
new_op_desc.set_output('XShape', [XShape_var.name])
|
||||
new_op_desc.set_output('Out', [Out_var.name])
|
||||
new_op_desc._set_attr('shape', shape_list)
|
||||
# TODO: should we add a new dist attr for the new op here?
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
class DistributedReshapeImpl1(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = False
|
||||
|
||||
def calc_cost(self, op_role, dist_op, ctx, cluster):
|
||||
cost = None
|
||||
if int(op_role) == int(OpRole.Backward):
|
||||
cost = self.calc_bwd_cost(dist_op, ctx, cluster)
|
||||
else:
|
||||
cost = self.calc_fwd_cost(dist_op, ctx, cluster)
|
||||
assert cost is not None
|
||||
return cost
|
||||
|
||||
def calc_fwd_cost(self, dist_op, ctx, cluster):
|
||||
res = []
|
||||
op = dist_op.serial_op
|
||||
dist_attr = dist_op.dist_attr
|
||||
|
||||
shape_list = op.desc.attr("shape")
|
||||
# got dist attribute info
|
||||
dim_mapping = dist_attr.get_output_dims_mapping(op.output("Out")[0])
|
||||
process_mesh_shape = dist_attr.process_mesh.shape
|
||||
|
||||
# modify target shape
|
||||
for idx, axis in enumerate(dim_mapping):
|
||||
if axis >= 0:
|
||||
if len(shape_list) > idx:
|
||||
shape_list[idx] = (
|
||||
shape_list[idx] // process_mesh_shape[axis]
|
||||
)
|
||||
|
||||
# calc comp op cost
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
processes = dist_attr.process_mesh.process_ids
|
||||
for key in desc_mapping:
|
||||
desc_mapping[key]["shape"] = shape_list
|
||||
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
Reshape2OpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
return res
|
||||
|
||||
def calc_bwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
res = []
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
dist_attr = dist_op.dist_attr
|
||||
process_mesh = dist_attr.process_mesh
|
||||
processes = process_mesh.process_ids
|
||||
op_type = dist_op.serial_op.type
|
||||
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
Reshape2GradOpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
backward_op = dist_op.serial_op
|
||||
main_block = backward_op.block
|
||||
need_gradient_allreduce = False
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and not is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
# NOTE input var's dim_mapping of backward op should be the same with input var instead of corresponding varname of forward op
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
|
||||
|
||||
mesh_shape = process_mesh.shape
|
||||
batch_size_axis = (
|
||||
var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
|
||||
)
|
||||
if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
|
||||
parallel_axis = batch_size_axis
|
||||
attrs = {"use_calc_stream": True}
|
||||
var_names = [varname + "@GRAD"]
|
||||
build_dp_costs(
|
||||
res,
|
||||
dist_op,
|
||||
ctx,
|
||||
var_names,
|
||||
attrs,
|
||||
parallel_axis,
|
||||
cluster,
|
||||
)
|
||||
|
||||
return res
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
if len(x_dims_mapping) != len(out_dims_mapping) + 1:
|
||||
return False
|
||||
|
||||
if is_dim_shard(x_dims_mapping[-1]):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
if len(x_dims_mapping) != len(out_dims_mapping) + 1:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_shape_name = op_desc.output('XShape')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
|
||||
x_shape_name
|
||||
)
|
||||
|
||||
if is_dim_shard(x_dims_mapping[-1]):
|
||||
return False
|
||||
|
||||
for idx, item in enumerate(x_dims_mapping[:-1]):
|
||||
if out_dims_mapping[idx] != item:
|
||||
return False
|
||||
|
||||
if x_shape_dims_mapping[0] != -1:
|
||||
return False
|
||||
|
||||
if x_shape_dims_mapping[1:] != x_dims_mapping[:]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_shape_name = op_desc.output('XShape')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
|
||||
x_shape_name
|
||||
)
|
||||
|
||||
for i in range(len(out_dims_mapping)):
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[x_dims_mapping, out_dims_mapping], [i, i]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
|
||||
for i in range(len(x_dims_mapping)):
|
||||
x_shape_dims_mapping[i + 1] = x_dims_mapping[i]
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
x_shape_name, x_shape_dims_mapping
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
"""
|
||||
kwargs: inputname_mapping & outputname_mapping
|
||||
"""
|
||||
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
assert op_dist_attr is not None, (
|
||||
f"backward op [{src_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
# check validation of inputs / outputs
|
||||
for input_name in src_op.desc.input_names():
|
||||
assert input_name in kwargs, f"input [{input_name}] is not given"
|
||||
assert len(kwargs[input_name]) == len(
|
||||
src_op.desc.input(input_name)
|
||||
), f"number of tensor for input [{input_name}] is not match"
|
||||
for output_name in src_op.desc.output_names():
|
||||
assert output_name in kwargs, f"input [{output_name}] is not given"
|
||||
assert len(kwargs[output_name]) == len(
|
||||
src_op.desc.output(output_name)
|
||||
), f"number of tensor for input [{output_name}] is not match"
|
||||
|
||||
X_var = main_block._var_recursive(kwargs['X'][0])
|
||||
Out_var = main_block._var_recursive(kwargs['Out'][0])
|
||||
XShape_var = main_block._var_recursive(kwargs['XShape'][0])
|
||||
shape_list = src_op.desc.attr("shape")
|
||||
ShapeTensor_var_list = []
|
||||
for name in kwargs['ShapeTensor']:
|
||||
ShapeTensor_var_list.append(name)
|
||||
Shape_var_list = []
|
||||
for name in kwargs['Shape']:
|
||||
Shape_var_list.append(name)
|
||||
|
||||
# got dist attribute info
|
||||
dim_mapping = op_dist_attr.get_output_dims_mapping(Out_var.name)
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
|
||||
# modify target shape
|
||||
for idx, axis in enumerate(dim_mapping):
|
||||
if axis >= 0:
|
||||
if len(shape_list) > idx:
|
||||
shape_list[idx] = (
|
||||
shape_list[idx] // process_mesh_shape[axis]
|
||||
)
|
||||
|
||||
# create op
|
||||
new_op = main_block.append_op(type='nop')
|
||||
new_op_desc = new_op.desc
|
||||
new_op_desc.copy_from(src_op.desc)
|
||||
set_dist_op_desc_original_id(new_op_desc, src_op.desc, ctx)
|
||||
new_op_desc.set_input('ShapeTensor', ShapeTensor_var_list)
|
||||
new_op_desc.set_input('Shape', Shape_var_list)
|
||||
new_op_desc.set_input('X', [X_var.name])
|
||||
new_op_desc.set_output('XShape', [XShape_var.name])
|
||||
new_op_desc.set_output('Out', [Out_var.name])
|
||||
new_op_desc._set_attr('shape', shape_list)
|
||||
# TODO: should we add a new dist attr for the new op here?
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
class DistributedReshapeImpl2(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = False
|
||||
|
||||
def calc_cost(self, op_role, dist_op, ctx, cluster):
|
||||
cost = None
|
||||
if int(op_role) == int(OpRole.Backward):
|
||||
cost = self.calc_bwd_cost(dist_op, ctx, cluster)
|
||||
else:
|
||||
cost = self.calc_fwd_cost(dist_op, ctx, cluster)
|
||||
assert cost is not None
|
||||
return cost
|
||||
|
||||
def calc_fwd_cost(self, dist_op, ctx, cluster):
|
||||
res = []
|
||||
op = dist_op.serial_op
|
||||
dist_attr = dist_op.dist_attr
|
||||
|
||||
shape_list = op.desc.attr("shape")
|
||||
# got dist attribute info
|
||||
dim_mapping = dist_attr.get_output_dims_mapping(op.output("Out")[0])
|
||||
process_mesh_shape = dist_attr.process_mesh.shape
|
||||
|
||||
# modify target shape
|
||||
for idx, axis in enumerate(dim_mapping):
|
||||
if axis >= 0:
|
||||
if len(shape_list) > idx:
|
||||
shape_list[idx] = (
|
||||
shape_list[idx] // process_mesh_shape[axis]
|
||||
)
|
||||
|
||||
# calc comp op cost
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
processes = dist_attr.process_mesh.process_ids
|
||||
for key in desc_mapping:
|
||||
desc_mapping[key]["shape"] = shape_list
|
||||
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
Reshape2OpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
return res
|
||||
|
||||
def calc_bwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
res = []
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
dist_attr = dist_op.dist_attr
|
||||
process_mesh = dist_attr.process_mesh
|
||||
processes = process_mesh.process_ids
|
||||
op_type = dist_op.serial_op.type
|
||||
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
Reshape2GradOpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
backward_op = dist_op.serial_op
|
||||
main_block = backward_op.block
|
||||
need_gradient_allreduce = False
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and not is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
# NOTE input var's dim_mapping of backward op should be the same with input var instead of corresponding varname of forward op
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
|
||||
|
||||
mesh_shape = process_mesh.shape
|
||||
batch_size_axis = (
|
||||
var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
|
||||
)
|
||||
if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
|
||||
parallel_axis = batch_size_axis
|
||||
attrs = {"use_calc_stream": True}
|
||||
var_names = [varname + "@GRAD"]
|
||||
build_dp_costs(
|
||||
res,
|
||||
dist_op,
|
||||
ctx,
|
||||
var_names,
|
||||
attrs,
|
||||
parallel_axis,
|
||||
cluster,
|
||||
)
|
||||
|
||||
return res
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
if len(x_dims_mapping) != len(out_dims_mapping):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_name = op_desc.input('X')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
if len(x_dims_mapping) != len(out_dims_mapping):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_shape_name = op_desc.output('XShape')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
|
||||
x_shape_name
|
||||
)
|
||||
|
||||
for idx, item in enumerate(x_dims_mapping[:-1]):
|
||||
if out_dims_mapping[idx] != item:
|
||||
return False
|
||||
|
||||
if x_shape_dims_mapping[0] != -1:
|
||||
return False
|
||||
|
||||
if x_shape_dims_mapping[1:] != out_dims_mapping[:]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_shape_name = op_desc.output('XShape')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
|
||||
x_shape_name
|
||||
)
|
||||
|
||||
for i in range(len(out_dims_mapping) - 1):
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[x_dims_mapping, out_dims_mapping], [i, i]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
|
||||
for i in range(len(out_dims_mapping)):
|
||||
x_shape_dims_mapping[i + 1] = out_dims_mapping[i]
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
x_shape_name, x_shape_dims_mapping
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
"""
|
||||
kwargs: inputname_mapping & outputname_mapping
|
||||
"""
|
||||
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.work_block
|
||||
src_op = dist_op_context.cur_src_op
|
||||
op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
|
||||
assert op_dist_attr is not None, (
|
||||
f"backward op [{src_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
# check validation of inputs / outputs
|
||||
for input_name in src_op.desc.input_names():
|
||||
assert input_name in kwargs, f"input [{input_name}] is not given"
|
||||
assert len(kwargs[input_name]) == len(
|
||||
src_op.desc.input(input_name)
|
||||
), f"number of tensor for input [{input_name}] is not match"
|
||||
for output_name in src_op.desc.output_names():
|
||||
assert output_name in kwargs, f"input [{output_name}] is not given"
|
||||
assert len(kwargs[output_name]) == len(
|
||||
src_op.desc.output(output_name)
|
||||
), f"number of tensor for input [{output_name}] is not match"
|
||||
|
||||
X_var = main_block._var_recursive(kwargs['X'][0])
|
||||
Out_var = main_block._var_recursive(kwargs['Out'][0])
|
||||
XShape_var = main_block._var_recursive(kwargs['XShape'][0])
|
||||
shape_list = src_op.desc.attr("shape")
|
||||
ShapeTensor_var_list = []
|
||||
for name in kwargs['ShapeTensor']:
|
||||
ShapeTensor_var_list.append(name)
|
||||
Shape_var_list = []
|
||||
for name in kwargs['Shape']:
|
||||
Shape_var_list.append(name)
|
||||
|
||||
# got dist attribute info
|
||||
out_dim_mapping = op_dist_attr.get_output_dims_mapping(Out_var.name)
|
||||
process_mesh_shape = op_dist_attr.process_mesh.shape
|
||||
|
||||
# modify target shape
|
||||
for idx, axis in enumerate(out_dim_mapping):
|
||||
if axis >= 0:
|
||||
if len(shape_list) > idx:
|
||||
shape_list[idx] = (
|
||||
shape_list[idx] // process_mesh_shape[axis]
|
||||
)
|
||||
|
||||
# create op
|
||||
new_op = main_block.append_op(type='nop')
|
||||
new_op_desc = new_op.desc
|
||||
new_op_desc.copy_from(src_op.desc)
|
||||
set_dist_op_desc_original_id(new_op_desc, src_op.desc, ctx)
|
||||
new_op_desc.set_input('ShapeTensor', ShapeTensor_var_list)
|
||||
new_op_desc.set_input('Shape', Shape_var_list)
|
||||
new_op_desc.set_input('X', [X_var.name])
|
||||
new_op_desc.set_output('XShape', [XShape_var.name])
|
||||
new_op_desc.set_output('Out', [Out_var.name])
|
||||
new_op_desc._set_attr('shape', shape_list)
|
||||
# TODO: should we add a new dist attr for the new op here?
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"reshape2", DistributedReshapeImpl0("add_one_dim_back")
|
||||
)
|
||||
register_distributed_operator_impl(
|
||||
"reshape2", DistributedReshapeImpl1("remove_one_dim_back")
|
||||
)
|
||||
register_distributed_operator_impl(
|
||||
"reshape2", DistributedReshapeImpl2("same_dim_shape")
|
||||
)
|
||||
@@ -0,0 +1,192 @@
|
||||
# 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.
|
||||
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
||||
|
||||
from ..cost import (
|
||||
_g_op_cost_factory,
|
||||
build_comp_costs_from_descs,
|
||||
build_comp_desc_from_dist_op,
|
||||
build_dp_costs,
|
||||
)
|
||||
from ..utils import compute_compatible_and_update_dim_mapping
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
is_parameter_related,
|
||||
)
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedScale(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
# TODO remove assign dist op
|
||||
# register_distributed_operator_impl_container(DistributedScale("scale"))
|
||||
# register_distributed_operator_impl_container(DistributedScale("fill_any_like"))
|
||||
# register_distributed_operator_impl_container(DistributedScale("where"))
|
||||
# register_distributed_operator_impl_container(DistributedScale("tanh"))
|
||||
|
||||
|
||||
class DistributedScaleImpl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def calc_cost(self, op_role, dist_op, ctx, cluster):
|
||||
"""Calculate the cost by the op role."""
|
||||
cost = None
|
||||
if int(op_role) == int(OpRole.Backward):
|
||||
cost = self.calc_bwd_cost(dist_op, ctx, cluster)
|
||||
else:
|
||||
cost = self.calc_fwd_cost(dist_op, ctx, cluster)
|
||||
assert cost is not None
|
||||
return cost
|
||||
|
||||
def calc_fwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
processes = dist_op.dist_attr.process_mesh.process_ids
|
||||
op_type = dist_op.serial_op.type
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
_g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res_cost = [cost_mapping]
|
||||
|
||||
return res_cost
|
||||
|
||||
def calc_bwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
res = []
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
dist_attr = dist_op.dist_attr
|
||||
process_mesh = dist_attr.process_mesh
|
||||
processes = process_mesh.process_ids
|
||||
backward_op = dist_op.serial_op
|
||||
op_type = backward_op.type
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
_g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
main_block = backward_op.block
|
||||
need_gradient_allreduce = False
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and not is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
|
||||
mesh_shape = process_mesh.shape
|
||||
batch_size_axis = (
|
||||
var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
|
||||
)
|
||||
if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
|
||||
need_gradient_allreduce = True
|
||||
break
|
||||
|
||||
if need_gradient_allreduce:
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(
|
||||
varname
|
||||
)
|
||||
mesh_shape = process_mesh.shape
|
||||
parallel_axis = batch_size_axis
|
||||
attrs = {"use_calc_stream": True}
|
||||
var_names = [varname + "@GRAD"]
|
||||
build_dp_costs(
|
||||
res,
|
||||
dist_op,
|
||||
ctx,
|
||||
var_names,
|
||||
attrs,
|
||||
parallel_axis,
|
||||
cluster,
|
||||
)
|
||||
return res
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
out_name = op_desc.output('Out')[0]
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
in_dims_mappings = []
|
||||
for in_name in op_desc.input_arg_names():
|
||||
in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
|
||||
in_dims_mappings.append(in_dims_mapping)
|
||||
|
||||
for x_dims_mapping in in_dims_mappings:
|
||||
if x_dims_mapping != out_dims_mapping:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
for i in range(len(x_dims_mapping)):
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[x_dims_mapping, out_dims_mapping], [i, i]
|
||||
)
|
||||
if dim_changed:
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
|
||||
changed = True
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
# register_distributed_operator_impl("scale", DistributedScaleImpl("scale"))
|
||||
# register_distributed_operator_impl(
|
||||
# "fill_any_like", DistributedScaleImpl("fill_any_like")
|
||||
# )
|
||||
# register_distributed_operator_impl("where", DistributedScaleImpl("where"))
|
||||
# register_distributed_operator_impl("tanh", DistributedScaleImpl("tanh"))
|
||||
@@ -0,0 +1,74 @@
|
||||
# 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.
|
||||
|
||||
from ..utils import is_dim_shard
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedShape(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedShape("shape"))
|
||||
|
||||
|
||||
class DistributedShapeImpl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
out_name = op_desc.output('Out')[0]
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
assert len(out_dims_mapping) == 1
|
||||
if is_dim_shard(out_dims_mapping[0]):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
register_distributed_operator_impl("shape", DistributedShapeImpl("shape"))
|
||||
@@ -0,0 +1,178 @@
|
||||
# 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.
|
||||
|
||||
|
||||
from ..utils import compute_compatible_dim_mapping, is_dim_shard
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedSlice(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedSlice("slice"))
|
||||
|
||||
|
||||
class DistributedSliceImpl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
in_name = op_desc.input('Input')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
in_var = dist_op.serial_op.block._var_recursive(in_name)
|
||||
out_var = dist_op.serial_op.block._var_recursive(out_name)
|
||||
axes = op_desc.attr('axes')
|
||||
in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
|
||||
for axis in axes:
|
||||
if (
|
||||
is_dim_shard(in_dims_mapping[axis])
|
||||
and in_var.shape[axis] != out_var.shape[axis]
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
in_name = op_desc.input('Input')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
in_var = dist_op.serial_op.block._var_recursive(in_name)
|
||||
out_var = dist_op.serial_op.block._var_recursive(out_name)
|
||||
axes = op_desc.attr('axes')
|
||||
decrease_axis = op_desc.attr('decrease_axis')
|
||||
in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
ref_indices = []
|
||||
for i in range(len(in_dims_mapping)):
|
||||
if i not in decrease_axis:
|
||||
ref_indices.append(i)
|
||||
if ref_indices == []:
|
||||
assert len(out_dims_mapping) == 0
|
||||
else:
|
||||
for i in range(len(out_dims_mapping)):
|
||||
ref_index = ref_indices[i]
|
||||
if (
|
||||
ref_index in axes
|
||||
and is_dim_shard(out_dims_mapping[i])
|
||||
and in_var.shape[ref_index] != out_var.shape[ref_index]
|
||||
):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
in_name = op_desc.input('Input')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
decrease_axis = op_desc.attr('decrease_axis')
|
||||
in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
if len(in_dims_mapping) - len(decrease_axis) != 0 and len(
|
||||
out_dims_mapping
|
||||
) != len(in_dims_mapping) - len(decrease_axis):
|
||||
return False
|
||||
|
||||
new_out_dims_mapping = []
|
||||
for i in range(len(in_dims_mapping)):
|
||||
if i not in decrease_axis:
|
||||
new_out_dims_mapping.append(in_dims_mapping[i])
|
||||
if new_out_dims_mapping == []:
|
||||
new_out_dims_mapping = [-1]
|
||||
if new_out_dims_mapping != out_dims_mapping:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (
|
||||
(not self.is_input_compatible(dist_op))
|
||||
or (not self.is_output_compatible(dist_op))
|
||||
or (not self.is_compatible(dist_op))
|
||||
):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
in_name = op_desc.input('Input')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
decrease_axis = op_desc.attr('decrease_axis')
|
||||
in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
ref_dims_mapping = []
|
||||
ref_indices = []
|
||||
for i in range(len(in_dims_mapping)):
|
||||
if i not in decrease_axis:
|
||||
ref_dims_mapping.append(in_dims_mapping[i])
|
||||
ref_indices.append(i)
|
||||
|
||||
if ref_dims_mapping == []:
|
||||
assert len(ref_dims_mapping) == len(out_dims_mapping)
|
||||
changed = False
|
||||
else:
|
||||
assert len(ref_dims_mapping) == len(out_dims_mapping)
|
||||
for i in range(len(out_dims_mapping)):
|
||||
compatible_dim_mapping = compute_compatible_dim_mapping(
|
||||
[out_dims_mapping[i], ref_dims_mapping[i]]
|
||||
)
|
||||
if compatible_dim_mapping is None:
|
||||
continue
|
||||
if ref_dims_mapping[i] != compatible_dim_mapping:
|
||||
in_dims_mapping[ref_indices[i]] = compatible_dim_mapping
|
||||
changed = True
|
||||
if out_dims_mapping[i] != compatible_dim_mapping:
|
||||
out_dims_mapping[i] = compatible_dim_mapping
|
||||
changed = True
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(in_name, in_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"slice", DistributedSliceImpl("decrease_in_axis")
|
||||
)
|
||||
@@ -0,0 +1,200 @@
|
||||
# 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
|
||||
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
||||
|
||||
from ..cost import (
|
||||
SoftmaxGradOpCost,
|
||||
SoftmaxOpCost,
|
||||
build_comp_costs_from_descs,
|
||||
build_comp_desc_from_dist_op,
|
||||
build_dp_costs,
|
||||
)
|
||||
from ..utils import compute_compatible_and_update_dim_mapping, is_dim_shard
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
is_parameter_related,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedSoftmax(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedSoftmax("softmax"))
|
||||
|
||||
|
||||
class DistributedSoftmaxImpl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = False
|
||||
self._backward_implemented = False
|
||||
|
||||
def calc_cost(self, op_role, dist_op, ctx, cluster):
|
||||
cost = None
|
||||
if int(op_role) == int(OpRole.Backward):
|
||||
cost = self.calc_bwd_cost(dist_op, ctx, cluster)
|
||||
else:
|
||||
cost = self.calc_fwd_cost(dist_op, ctx, cluster)
|
||||
assert cost is not None
|
||||
return cost
|
||||
|
||||
def calc_fwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
processes = dist_op.dist_attr.process_mesh.process_ids
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
SoftmaxOpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
|
||||
res_cost = [cost_mapping]
|
||||
return res_cost
|
||||
|
||||
def calc_bwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
res = []
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
dist_attr = dist_op.dist_attr
|
||||
process_mesh = dist_attr.process_mesh
|
||||
processes = process_mesh.process_ids
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
SoftmaxGradOpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
backward_op = dist_op.serial_op
|
||||
main_block = backward_op.block
|
||||
need_gradient_allreduce = False
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
# NOTE input var's dim_mapping of backward op should be the same with input var instead of corresponding varname of forward op
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
|
||||
|
||||
mesh_shape = process_mesh.shape
|
||||
batch_size_axis = (
|
||||
var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
|
||||
)
|
||||
if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
|
||||
parallel_axis = batch_size_axis
|
||||
attrs = {"use_calc_stream": True}
|
||||
var_names = [varname + "@GRAD"]
|
||||
build_dp_costs(
|
||||
res,
|
||||
dist_op,
|
||||
ctx,
|
||||
var_names,
|
||||
attrs,
|
||||
parallel_axis,
|
||||
cluster,
|
||||
)
|
||||
|
||||
return res
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
axis = op_desc.attr('axis')
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
|
||||
# if axis != -1 and axis != len(x_dims_mapping) - 1:
|
||||
# return False
|
||||
|
||||
if is_dim_shard(x_dims_mapping[axis]):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
out_name = op_desc.output('Out')[0]
|
||||
axis = op_desc.attr('axis')
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
# if axis != -1 and axis != len(out_dims_mapping) - 1:
|
||||
# return False
|
||||
|
||||
if is_dim_shard(out_dims_mapping[axis]):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
axis = op_desc.attr('axis')
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
# if axis != -1 and axis != len(x_dims_mapping) - 1:
|
||||
# return False
|
||||
|
||||
if x_dims_mapping != out_dims_mapping:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
|
||||
for i in range(len(x_dims_mapping)):
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[x_dims_mapping, out_dims_mapping], [i, i]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"softmax", DistributedSoftmaxImpl("replicate_last_axis")
|
||||
)
|
||||
@@ -0,0 +1,197 @@
|
||||
# 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.
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..utils import (
|
||||
compute_compatible_and_update_dim_mapping,
|
||||
get_dist_tensor_spec,
|
||||
is_dim_shard,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedSplit(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
|
||||
x_name = op_desc.input('X')[0]
|
||||
assert len(op_desc.input('AxisTensor')) == 0, (
|
||||
"Attribute AxisTensor is not supported by dist split."
|
||||
)
|
||||
assert len(op_desc.input('SectionsTensorList')) == 0, (
|
||||
"Attribute SectionsTensorList is not supported by dist split."
|
||||
)
|
||||
output_arg_names = op_desc.output('Out')
|
||||
|
||||
num = op_desc.attr('num')
|
||||
sections = op_desc.attr('sections')
|
||||
if num:
|
||||
assert (sections is None) or (len(sections) == 0), (
|
||||
f"Both Attributes of num: {num} and sections: {sections} are specified."
|
||||
)
|
||||
first_attr = num
|
||||
rule_type = "split_with_num"
|
||||
else:
|
||||
assert not num, (
|
||||
f"Both Attributes of num: {num} and sections: {sections} are specified."
|
||||
)
|
||||
first_attr = sections
|
||||
rule_type = "split"
|
||||
axis = op_desc.attr('axis')
|
||||
|
||||
x_spec = get_dist_tensor_spec(dist_op, x_name)
|
||||
num_outputs = len(output_arg_names)
|
||||
output_specs = []
|
||||
for i in range(num_outputs):
|
||||
output_specs.append(
|
||||
get_dist_tensor_spec(dist_op, output_arg_names[i], False)
|
||||
)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule(rule_type)
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(x_spec, first_attr, axis)
|
||||
bw_results = rule.infer_backward(x_spec, output_specs, first_attr, axis)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op, [x_name], output_arg_names, fw_results, bw_results
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
# all split op use default dist operator impl.
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedSplit("split"))
|
||||
register_distributed_operator_impl_container(DistributedSplit("split_with_num"))
|
||||
|
||||
|
||||
class DistributedSplitImpl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = True
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
axis = op_desc.attr('axis')
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
|
||||
if is_dim_shard(x_dims_mapping[axis]):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
out_names = op_desc.output('Out')
|
||||
axis = op_desc.attr('axis')
|
||||
for out_name in out_names:
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
if is_dim_shard(out_dims_mapping[axis]):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
axis = op_desc.attr('axis')
|
||||
out_names = op_desc.output('Out')
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
for out_name in out_names:
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
if x_dims_mapping != out_dims_mapping:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_names = op_desc.output('Out')
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
|
||||
for out_name in out_names:
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
for i in range(len(x_dims_mapping)):
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[x_dims_mapping, out_dims_mapping], [i, i]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
out_name, out_dims_mapping
|
||||
)
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
return changed
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (
|
||||
(not self.is_input_compatible(dist_op))
|
||||
or (not self.is_output_compatible(dist_op))
|
||||
or (not self.is_compatible(dist_op))
|
||||
):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"split", DistributedSplitImpl("replicate_in_axis")
|
||||
)
|
||||
@@ -0,0 +1,71 @@
|
||||
# Copyright (c) 2024 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
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..utils import get_dist_tensor_spec
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
|
||||
class DistributedStack(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
|
||||
input_arg_names = op_desc.input_arg_names()
|
||||
output_arg_names = op_desc.output_arg_names()
|
||||
axis = op_desc.attr('axis')
|
||||
|
||||
input_specs = []
|
||||
for name in input_arg_names:
|
||||
input_specs.append(get_dist_tensor_spec(dist_op, name))
|
||||
output_spec = get_dist_tensor_spec(dist_op, output_arg_names[0], False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("stack")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(input_specs, axis)
|
||||
bw_results = rule.infer_backward(input_specs, output_spec, axis)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op,
|
||||
input_arg_names,
|
||||
output_arg_names,
|
||||
fw_results,
|
||||
bw_results,
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedStack("stack"))
|
||||
@@ -0,0 +1,81 @@
|
||||
# Copyright (c) 2024 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.
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..utils import get_dist_tensor_spec
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
|
||||
class DistributedStridedSlice(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
|
||||
x_name = op_desc.input('Input')[0]
|
||||
y_name = op_desc.output('Out')[0]
|
||||
axes = op_desc.attr('axes')
|
||||
starts = op_desc.attr('starts')
|
||||
ends = op_desc.attr('ends')
|
||||
strides = op_desc.attr('strides')
|
||||
|
||||
x_spec = get_dist_tensor_spec(dist_op, x_name)
|
||||
y_spec = get_dist_tensor_spec(dist_op, y_name, False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("strided_slice")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(x_spec, axes, starts, ends, strides)
|
||||
bw_results = rule.infer_backward(
|
||||
x_spec,
|
||||
y_spec,
|
||||
axes,
|
||||
starts,
|
||||
ends,
|
||||
strides,
|
||||
)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op,
|
||||
[x_name],
|
||||
[y_name],
|
||||
fw_results,
|
||||
bw_results,
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedStridedSlice("strided_slice")
|
||||
)
|
||||
@@ -0,0 +1,72 @@
|
||||
# Copyright (c) 2024 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
|
||||
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..utils import (
|
||||
get_dist_tensor_spec,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
|
||||
class DistributedTile(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
assert dist_op.serial_op.type == "tile", (
|
||||
f"{dist_op.serial_op.type} is not supported by dist transpose yet."
|
||||
)
|
||||
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
repeat_times = op_desc.attr('repeat_times')
|
||||
|
||||
x_spec = get_dist_tensor_spec(dist_op, x_name)
|
||||
output_spec = get_dist_tensor_spec(dist_op, out_name, False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("tile")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(x_spec, repeat_times)
|
||||
bw_results = rule.infer_backward(x_spec, output_spec, repeat_times)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op, [x_name], [out_name], fw_results, bw_results
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
# all elementwise op use default dist operator impl.
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(DistributedTile("tile"))
|
||||
@@ -0,0 +1,270 @@
|
||||
# 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
|
||||
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..cost import (
|
||||
Transpose2GradOpCost,
|
||||
Transpose2OpCost,
|
||||
build_comp_costs_from_descs,
|
||||
build_comp_desc_from_dist_op,
|
||||
build_dp_costs,
|
||||
)
|
||||
from ..utils import (
|
||||
compute_compatible_and_update_dim_mapping,
|
||||
get_dist_tensor_spec,
|
||||
)
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
is_parameter_related,
|
||||
merge_forward_backward_dims_mapping,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
from .dist_default import DistributedDefaultImpl0
|
||||
|
||||
|
||||
class DistributedTranspose2(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
assert dist_op.serial_op.type == "transpose2", (
|
||||
f"{dist_op.serial_op.type} is not supported by dist transpose yet."
|
||||
)
|
||||
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
xshape_name = op_desc.output('XShape')[0]
|
||||
axes = op_desc.attr('axis')
|
||||
|
||||
x_spec = get_dist_tensor_spec(dist_op, x_name)
|
||||
output_spec = get_dist_tensor_spec(dist_op, out_name, False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("transpose")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(x_spec, axes)
|
||||
bw_results = rule.infer_backward(x_spec, output_spec, axes)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op, [x_name], [out_name], fw_results, bw_results
|
||||
)
|
||||
|
||||
# step4: update xshape
|
||||
inferred_input_dims_mappings, _ = merge_forward_backward_dims_mapping(
|
||||
fw_results, bw_results
|
||||
)
|
||||
dist_op.dist_attr.set_output_dims_mapping(
|
||||
xshape_name, [-1] + inferred_input_dims_mappings[0]
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
# all elementwise op use default dist operator impl.
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedTranspose2("transpose2")
|
||||
)
|
||||
|
||||
|
||||
class DistributedTranspose2Impl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._forward_implemented = False
|
||||
self._backward_implemented = False
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
return True
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
if (not self.is_input_compatible(dist_op)) or (
|
||||
not self.is_output_compatible(dist_op)
|
||||
):
|
||||
return False
|
||||
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
perm = op_desc.attr('axis')
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_shape_name = op_desc.output('XShape')[0]
|
||||
x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
|
||||
x_shape_name
|
||||
)
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
new_dims_mapping = [-1 for i in range(len(x_dims_mapping))]
|
||||
for i in range(len(x_dims_mapping)):
|
||||
new_dims_mapping[i] = x_dims_mapping[perm[i]]
|
||||
|
||||
if len(x_dims_mapping) != len(out_dims_mapping):
|
||||
return False
|
||||
|
||||
if new_dims_mapping != out_dims_mapping:
|
||||
return False
|
||||
|
||||
if x_shape_dims_mapping[0] != -1:
|
||||
return False
|
||||
|
||||
if x_shape_dims_mapping[1:] != x_dims_mapping[:]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
changed = False
|
||||
op_desc = dist_op.serial_op.desc
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
x_shape_name = op_desc.output('XShape')[0]
|
||||
x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
|
||||
out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
|
||||
x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
|
||||
x_shape_name
|
||||
)
|
||||
perm = op_desc.attr('axis')
|
||||
|
||||
assert len(x_dims_mapping) == len(perm)
|
||||
|
||||
new_dims_mapping = [-1 for i in range(len(x_dims_mapping))]
|
||||
for i in range(len(x_dims_mapping)):
|
||||
new_dims_mapping[i] = x_dims_mapping[perm[i]]
|
||||
|
||||
for i in range(len(out_dims_mapping)):
|
||||
dim_changed = compute_compatible_and_update_dim_mapping(
|
||||
[new_dims_mapping, out_dims_mapping], [i, i]
|
||||
)
|
||||
if dim_changed:
|
||||
changed = True
|
||||
|
||||
for i in range(len(x_dims_mapping)):
|
||||
if x_dims_mapping[perm[i]] != new_dims_mapping[i]:
|
||||
x_dims_mapping[perm[i]] = new_dims_mapping[i]
|
||||
changed = True
|
||||
|
||||
for i in range(len(x_dims_mapping)):
|
||||
x_shape_dims_mapping[i + 1] = x_dims_mapping[i]
|
||||
|
||||
if changed:
|
||||
op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
x_shape_name, x_shape_dims_mapping
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
def calc_cost(self, op_role, dist_op, ctx, cluster):
|
||||
cost = None
|
||||
if int(op_role) == int(OpRole.Backward):
|
||||
cost = self.calc_bwd_cost(dist_op, ctx, cluster)
|
||||
else:
|
||||
cost = self.calc_fwd_cost(dist_op, ctx, cluster)
|
||||
assert cost is not None
|
||||
return cost
|
||||
|
||||
def calc_fwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
processes = dist_op.dist_attr.process_mesh.process_ids
|
||||
op_type = dist_op.serial_op.type
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
Transpose2OpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
|
||||
res_cost = [cost_mapping]
|
||||
return res_cost
|
||||
|
||||
def calc_bwd_cost(self, dist_op, ctx, cluster):
|
||||
# calc comp op cost
|
||||
res = []
|
||||
desc_mapping = build_comp_desc_from_dist_op(
|
||||
dist_op=dist_op, dist_context=ctx
|
||||
)
|
||||
dist_attr = dist_op.dist_attr
|
||||
process_mesh = dist_attr.process_mesh
|
||||
processes = process_mesh.process_ids
|
||||
op_type = dist_op.serial_op.type
|
||||
cost_mapping = build_comp_costs_from_descs(
|
||||
Transpose2GradOpCost, ctx, processes, desc_mapping, cluster
|
||||
)
|
||||
res.append(cost_mapping)
|
||||
|
||||
backward_op = dist_op.serial_op
|
||||
main_block = backward_op.block
|
||||
need_gradient_allreduce = False
|
||||
for input_name in backward_op.desc.input_names():
|
||||
for varname in backward_op.desc.input(input_name):
|
||||
if "@GRAD" not in varname and is_parameter_related(
|
||||
varname, main_block
|
||||
):
|
||||
# NOTE input var's dim_mapping of backward op should be the same with input var instead of corresponding varname of forward op
|
||||
var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
|
||||
|
||||
mesh_shape = process_mesh.shape
|
||||
batch_size_axis = (
|
||||
var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
|
||||
)
|
||||
if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
|
||||
parallel_axis = batch_size_axis
|
||||
attrs = {"use_calc_stream": True}
|
||||
var_names = [varname + "@GRAD"]
|
||||
build_dp_costs(
|
||||
res,
|
||||
dist_op,
|
||||
ctx,
|
||||
var_names,
|
||||
attrs,
|
||||
parallel_axis,
|
||||
cluster,
|
||||
)
|
||||
return res
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"transpose2", DistributedTranspose2Impl("same_mapping_transpose")
|
||||
)
|
||||
@@ -0,0 +1,75 @@
|
||||
# Copyright (c) 2023 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
|
||||
|
||||
from ..completion import get_phi_spmd_rule
|
||||
from ..utils import get_dist_tensor_spec
|
||||
from .common import (
|
||||
DistributedOperatorImplContainer,
|
||||
get_default_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
update_op_dims_mapping,
|
||||
)
|
||||
|
||||
|
||||
class DistributedUnSqueeze2(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
@staticmethod
|
||||
def update_dims_mapping(dist_op):
|
||||
# step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
|
||||
op_desc = dist_op.serial_op.desc
|
||||
|
||||
axes_tensor = op_desc.input('AxesTensor')
|
||||
axes_tensor_list = op_desc.input('AxesTensorList')
|
||||
assert len(axes_tensor) == 0 and len(axes_tensor_list) == 0
|
||||
|
||||
x_name = op_desc.input('X')[0]
|
||||
out_name = op_desc.output('Out')[0]
|
||||
axes = op_desc.attr('axes')
|
||||
|
||||
input_spec = get_dist_tensor_spec(dist_op, x_name)
|
||||
output_spec = get_dist_tensor_spec(dist_op, out_name, False)
|
||||
|
||||
# step2: infer spmd
|
||||
rule = get_phi_spmd_rule("unsqueeze2")
|
||||
# tensor order following order in PHI definition
|
||||
fw_results = rule.infer_forward(input_spec, axes)
|
||||
bw_results = rule.infer_backward(input_spec, output_spec, axes)
|
||||
|
||||
# step3: update dist_attr
|
||||
# tensor order following order in PHI definition
|
||||
changed = update_op_dims_mapping(
|
||||
dist_op,
|
||||
[x_name],
|
||||
[out_name],
|
||||
fw_results,
|
||||
bw_results,
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
default_impl = get_default_distributed_operator_impl()
|
||||
op_dist_attr.impl_type = default_impl.type
|
||||
op_dist_attr.impl_idx = default_impl.idx
|
||||
|
||||
return False
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedUnSqueeze2("unsqueeze2")
|
||||
)
|
||||
@@ -0,0 +1,171 @@
|
||||
# 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
|
||||
|
||||
from ..utils import set_dist_op_desc_original_id
|
||||
from .common import (
|
||||
DistributedOperatorImpl,
|
||||
DistributedOperatorImplContainer,
|
||||
register_distributed_operator_impl,
|
||||
register_distributed_operator_impl_container,
|
||||
)
|
||||
|
||||
|
||||
class DistributedUpdateLossScaling(DistributedOperatorImplContainer):
|
||||
def __init__(self, op_type):
|
||||
super().__init__(op_type)
|
||||
|
||||
|
||||
register_distributed_operator_impl_container(
|
||||
DistributedUpdateLossScaling("update_loss_scaling")
|
||||
)
|
||||
|
||||
|
||||
class DistributedUpdateLossScalingImpl(DistributedOperatorImpl):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
self._name = name
|
||||
self._forward_implemented = False
|
||||
self._backward_implemented = True
|
||||
|
||||
def is_input_compatible(self, dist_op):
|
||||
raise RuntimeError(
|
||||
"DistributedUpdateLossScalingImpl's is_input_compatible should not be called !"
|
||||
)
|
||||
|
||||
def is_output_compatible(self, dist_op):
|
||||
raise RuntimeError(
|
||||
"DistributedUpdateLossScalingImpl's is_output_compatible should not be called !"
|
||||
)
|
||||
|
||||
def is_auto_compatible(self, dist_op):
|
||||
raise RuntimeError(
|
||||
"DistributedUpdateLossScalingImpl's is_auto_compatible should not be called !"
|
||||
)
|
||||
|
||||
def update_dims_mapping(self, dist_op):
|
||||
raise RuntimeError(
|
||||
"DistributedUpdateLossScalingImpl's update_dims_mapping should not be called !"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, *args, **kwargs):
|
||||
raise RuntimeError(
|
||||
"DistributedUpdateLossScalingImpl's forward should not be called !"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args, **kwargs):
|
||||
# the backward function only filter the gradient with current rank id
|
||||
dist_op_context = ctx.dist_op_context
|
||||
main_block = dist_op_context.main_block
|
||||
backward_op = dist_op_context.cur_src_op
|
||||
rank_id = dist_op_context.rank_id
|
||||
dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
|
||||
assert dist_attr is not None, (
|
||||
f"backward op [{backward_op}] don't have dist attribute !"
|
||||
)
|
||||
|
||||
assert rank_id in dist_attr.process_mesh.process_ids
|
||||
|
||||
assert 'X' in kwargs, "input [{}] is not given".format('X')
|
||||
assert 'FoundInfinite' in kwargs, "input [{}] is not given".format(
|
||||
'FoundInfinite'
|
||||
)
|
||||
assert 'PrevLossScaling' in kwargs, "input [{}] is not given".format(
|
||||
'PrevLossScaling'
|
||||
)
|
||||
assert 'InGoodSteps' in kwargs, "input [{}] is not given".format(
|
||||
'InGoodSteps'
|
||||
)
|
||||
assert 'InBadSteps' in kwargs, "input [{}] is not given".format(
|
||||
'InBadSteps'
|
||||
)
|
||||
|
||||
assert 'Out' in kwargs, "output [{}] is not given".format('Out')
|
||||
assert 'LossScaling' in kwargs, "output [{}] is not given".format(
|
||||
'LossScaling'
|
||||
)
|
||||
assert 'OutGoodSteps' in kwargs, "output [{}] is not given".format(
|
||||
'OutGoodSteps'
|
||||
)
|
||||
assert 'OutBadSteps' in kwargs, "output [{}] is not given".format(
|
||||
'OutBadSteps'
|
||||
)
|
||||
|
||||
assert len(kwargs['FoundInfinite']) == 1, (
|
||||
"update_loss_scaling input FoundInfinite take 1 variable but got {}".format(
|
||||
kwargs['FoundInfinite']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['PrevLossScaling']) == 1, (
|
||||
"update_loss_scaling input PrevLossScaling take 1 variable but got {}".format(
|
||||
kwargs['PrevLossScaling']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['InGoodSteps']) == 1, (
|
||||
"update_loss_scaling input InGoodSteps take 1 variable but got {}".format(
|
||||
kwargs['InGoodSteps']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['InBadSteps']) == 1, (
|
||||
"update_loss_scaling input InBadSteps take 1 variable but got {}".format(
|
||||
kwargs['InBadSteps']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['LossScaling']) == 1, (
|
||||
"update_loss_scaling output LossScaling take 1 variable but got {}".format(
|
||||
kwargs['LossScaling']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['OutGoodSteps']) == 1, (
|
||||
"update_loss_scaling output OutGoodSteps take 1 variable but got {}".format(
|
||||
kwargs['OutGoodSteps']
|
||||
)
|
||||
)
|
||||
assert len(kwargs['OutBadSteps']) == 1, (
|
||||
"update_loss_scaling output OutBadSteps take 1 variable but got {}".format(
|
||||
kwargs['OutBadSteps']
|
||||
)
|
||||
)
|
||||
|
||||
assert len(kwargs['X']) == len(kwargs['Out']), (
|
||||
"update_loss_scaling got [{}] X and [{}] Out, which are supposed to be equal".format(
|
||||
len(kwargs['X']), len(kwargs['Out'])
|
||||
)
|
||||
)
|
||||
|
||||
filter_vars = []
|
||||
for varname in kwargs['X']:
|
||||
if (
|
||||
rank_id
|
||||
in ctx.get_tensor_dist_attr_for_program(
|
||||
main_block._var_recursive(varname)
|
||||
).process_mesh.process_ids
|
||||
):
|
||||
filter_vars.append(varname)
|
||||
|
||||
# replicate op in dist program
|
||||
dist_op = main_block.append_op(type='nop')
|
||||
dist_op_desc = dist_op.desc
|
||||
dist_op_desc.copy_from(backward_op.desc)
|
||||
set_dist_op_desc_original_id(dist_op_desc, backward_op.desc, ctx)
|
||||
dist_op_desc.set_input('X', filter_vars)
|
||||
dist_op_desc.set_output('Out', filter_vars)
|
||||
# TODO: should we add a new dist attr for the new op here?
|
||||
|
||||
|
||||
register_distributed_operator_impl(
|
||||
"update_loss_scaling",
|
||||
DistributedUpdateLossScalingImpl("update_loss_scaling"),
|
||||
)
|
||||
@@ -0,0 +1,539 @@
|
||||
# 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 json
|
||||
import logging
|
||||
import os
|
||||
import pathlib
|
||||
import pickle
|
||||
import shlex
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.passes import PassContext, new_pass
|
||||
from paddle.distributed.utils.log_utils import get_logger
|
||||
from paddle.framework import core
|
||||
from paddle.static import append_backward, program_guard
|
||||
|
||||
from .cluster import Cluster
|
||||
from .completion import Completer
|
||||
from .dist_context import DistributedContext, set_default_distributed_context
|
||||
from .dist_op import DistributedOperator
|
||||
from .dist_tensor import DistributedTensor
|
||||
from .mapper import mapping
|
||||
from .partitioner import Partitioner
|
||||
from .planner import Planner
|
||||
from .process_group import (
|
||||
ProcessGroup,
|
||||
_g_process_group_map,
|
||||
get_all_process_groups,
|
||||
get_process_group,
|
||||
get_world_process_group,
|
||||
)
|
||||
from .reshard import Resharder
|
||||
from .utils import SerialProgramInfo, make_data_unshard
|
||||
|
||||
_logger = get_logger(logging.INFO)
|
||||
|
||||
|
||||
class AutoParallelizer:
|
||||
"""
|
||||
AutoParallelizer is the main controller class to do the auto parallel process.
|
||||
And the auto parallel process will be triggered in the wrapped parallelize function.
|
||||
To facilitate the auto parallelization, it will contain information about program, cluster and the
|
||||
related context. In this basic version, the program information will be retrieved from
|
||||
Fleet object, and the cluster information can be retrieved in the new created Cluster object,
|
||||
and the context information can be retrieved in the new created DistributedContext.
|
||||
"""
|
||||
|
||||
def __init__(self, fleet):
|
||||
self._fleet = fleet
|
||||
self._optimizer = self._fleet.user_defined_optimizer
|
||||
self._dist_strategy = self._fleet._user_defined_strategy
|
||||
self._dist_context = DistributedContext()
|
||||
self._cluster = None
|
||||
self._cluster_topo_path = os.getenv("PADDLE_CLUSTER_TOPO_PATH", None)
|
||||
if self._cluster_topo_path is not None:
|
||||
self._cluster = Cluster()
|
||||
self._cluster.build_from_file(self._cluster_topo_path)
|
||||
# Prepare information for auto mapping
|
||||
self._rank_mapping_path = os.getenv("PADDLE_RANK_MAPPING_PATH", None)
|
||||
enable_auto_mapping_env = os.getenv("PADDLE_ENABLE_AUTO_MAPPING", None)
|
||||
if enable_auto_mapping_env is None:
|
||||
self._enable_auto_mapping = False
|
||||
else:
|
||||
self._enable_auto_mapping = True
|
||||
self._pass_context = PassContext()
|
||||
|
||||
self._need_rank_mapping = os.getenv("PADDLE_NEED_RANK_MAPPING")
|
||||
self._need_rank_mapping = (
|
||||
True
|
||||
if self._need_rank_mapping
|
||||
and self._need_rank_mapping.lower() == 'true'
|
||||
else False
|
||||
)
|
||||
# self._pass_context = None
|
||||
|
||||
def _remove_distributed_attrs(self, main_program):
|
||||
suffix = core.kAutoParallelSuffix()
|
||||
# distributed attributes for variable have been removed
|
||||
# in previous process.
|
||||
for block in main_program.blocks:
|
||||
for op in block.ops:
|
||||
for attr_name in op.attr_names:
|
||||
if suffix in attr_name:
|
||||
op._remove_attr(attr_name)
|
||||
|
||||
def _apply_pre_optimization_passes(
|
||||
self, main_program, startup_program, loss, params_grads, no_grad_set
|
||||
):
|
||||
# apply amp pass
|
||||
if self._dist_strategy.amp:
|
||||
config = copy.deepcopy(self._dist_strategy.amp_configs)
|
||||
config["dist_context"] = self._dist_context
|
||||
config["params_grads"] = params_grads
|
||||
config["loss"] = loss
|
||||
if config["use_pure_fp16"]:
|
||||
config["base_opt"] = self._optimizer
|
||||
auto_parallel_fp16_pass = new_pass("auto_parallel_fp16", config)
|
||||
auto_parallel_fp16_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
loss = auto_parallel_fp16_pass.get_loss()
|
||||
else:
|
||||
auto_parallel_amp_pass = new_pass("auto_parallel_amp", config)
|
||||
auto_parallel_amp_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
loss = auto_parallel_amp_pass.get_loss()
|
||||
|
||||
# apply recompute pass
|
||||
if self._dist_strategy.recompute:
|
||||
config = copy.deepcopy(self._dist_strategy.recompute_configs)
|
||||
config["dist_context"] = self._dist_context
|
||||
config["no_grad_set"] = copy.deepcopy(no_grad_set)
|
||||
config["loss"] = loss
|
||||
auto_parallel_recompute_pass = new_pass(
|
||||
"auto_parallel_recompute", config
|
||||
)
|
||||
auto_parallel_recompute_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
|
||||
def _generate_backward(
|
||||
self,
|
||||
main_program,
|
||||
startup_program,
|
||||
loss,
|
||||
parameter_list,
|
||||
no_grad_set,
|
||||
callbacks,
|
||||
):
|
||||
with program_guard(main_program, startup_program):
|
||||
params_grads = append_backward(
|
||||
loss,
|
||||
parameter_list,
|
||||
no_grad_set,
|
||||
callbacks,
|
||||
distop_context=self._dist_context.dist_op_context,
|
||||
)
|
||||
self._completer = Completer(self._dist_context)
|
||||
self._completer.complete_backward_annotation(main_program)
|
||||
self._dist_context.block_state.parse_backward_blocks(main_program)
|
||||
return params_grads
|
||||
|
||||
def _apply_optimize(self, main_program, startup_program, params_grads):
|
||||
optimizer = copy.deepcopy(self._optimizer)
|
||||
with program_guard(main_program, startup_program):
|
||||
optimize_ops = optimizer.apply_gradients(params_grads)
|
||||
|
||||
self._dist_context._serial_optimizer = optimizer
|
||||
# update completion
|
||||
self._completer = Completer(self._dist_context)
|
||||
self._completer.complete_update_annotation(main_program)
|
||||
|
||||
return optimize_ops
|
||||
|
||||
def _apply_post_optimization_passes(
|
||||
self, main_program, startup_program, rank, params_grads
|
||||
):
|
||||
if self._dist_strategy.sharding:
|
||||
config = copy.deepcopy(self._dist_strategy.sharding_configs)
|
||||
config["dist_context"] = self._dist_context
|
||||
config["params_grads"] = params_grads
|
||||
config["global_rank"] = rank
|
||||
auto_parallel_sharding_pass = new_pass(
|
||||
"auto_parallel_sharding", config
|
||||
)
|
||||
auto_parallel_sharding_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
params_grads = self._pass_context.get_attr("params_grads")
|
||||
|
||||
config = copy.deepcopy(self._dist_strategy.sharding_configs)
|
||||
config["dist_context"] = self._dist_context
|
||||
config["params_grads"] = params_grads
|
||||
config["rank_id"] = rank
|
||||
auto_parallel_clip_pass = new_pass("auto_parallel_grad_clip", config)
|
||||
auto_parallel_clip_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
|
||||
if self._dist_strategy.gradient_merge:
|
||||
config = copy.deepcopy(self._dist_strategy.gradient_merge_configs)
|
||||
config["dist_context"] = self._dist_context
|
||||
config["params_grads"] = params_grads
|
||||
auto_parallel_gradient_merge_pass = new_pass(
|
||||
"auto_parallel_gradient_merge_pass", config
|
||||
)
|
||||
auto_parallel_gradient_merge_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
|
||||
def _get_dist_program(self, rank, dist_context=None, relaunch_phase=False):
|
||||
completed_main_program = None
|
||||
serial_main_program = self._main_program.clone()
|
||||
serial_startup_program = self._startup_program.clone()
|
||||
serial_loss = serial_main_program.global_block().var(self._loss.name)
|
||||
|
||||
# generating serial
|
||||
if dist_context is None:
|
||||
# Annotation completion
|
||||
self._dist_context = DistributedContext()
|
||||
_logger.info("Start annotation dist attr.")
|
||||
self._completer = Completer(self._dist_context)
|
||||
completed_main_program = (
|
||||
self._completer.complete_forward_annotation(serial_main_program)
|
||||
)
|
||||
else:
|
||||
completed_main_program = serial_main_program
|
||||
self._dist_context = copy.deepcopy(dist_context)
|
||||
|
||||
# parse forward sub block
|
||||
self._dist_context.block_state.parse_forward_blocks(serial_main_program)
|
||||
|
||||
# serial backward pass
|
||||
params_grads = self._generate_backward(
|
||||
completed_main_program,
|
||||
serial_startup_program,
|
||||
serial_loss,
|
||||
self._parameter_list,
|
||||
self._no_grad_set,
|
||||
self._callbacks,
|
||||
)
|
||||
|
||||
# serial forward pass
|
||||
self._apply_pre_optimization_passes(
|
||||
completed_main_program,
|
||||
serial_startup_program,
|
||||
serial_loss,
|
||||
params_grads,
|
||||
self._no_grad_set,
|
||||
)
|
||||
# Logical partition
|
||||
partitioner = Partitioner(self._dist_context, rank)
|
||||
(
|
||||
dist_main_prog,
|
||||
dist_startup_prog,
|
||||
dist_params_grads,
|
||||
) = partitioner.partition(
|
||||
completed_main_program, serial_startup_program, params_grads
|
||||
)
|
||||
|
||||
# TODO refactor the placement of optimizer
|
||||
# generate optimize program
|
||||
dist_optimize_ops = self._apply_optimize(
|
||||
dist_main_prog, dist_startup_prog, dist_params_grads
|
||||
)
|
||||
|
||||
make_data_unshard(dist_main_prog, dist_startup_prog, self._dist_context)
|
||||
|
||||
resharder = Resharder(
|
||||
dist_main_prog,
|
||||
dist_startup_prog,
|
||||
rank,
|
||||
self._dist_context,
|
||||
dist_params_grads,
|
||||
)
|
||||
resharder.reshard()
|
||||
|
||||
self._apply_post_optimization_passes(
|
||||
dist_main_prog, dist_startup_prog, rank, dist_params_grads
|
||||
)
|
||||
g_process_group_map = None
|
||||
if not relaunch_phase:
|
||||
g_process_group_map = copy.deepcopy(_g_process_group_map)
|
||||
_g_process_group_map.clear()
|
||||
_g_process_group_map[0] = ProcessGroup(0, [])
|
||||
for process_mesh in self._dist_context._process_meshes:
|
||||
_g_process_group_map[0].add_ranks(process_mesh.process_ids)
|
||||
return (
|
||||
dist_optimize_ops,
|
||||
dist_params_grads,
|
||||
dist_startup_prog,
|
||||
dist_main_prog,
|
||||
g_process_group_map,
|
||||
)
|
||||
|
||||
def parallelize(
|
||||
self,
|
||||
loss,
|
||||
startup_program,
|
||||
parameter_list=None,
|
||||
no_grad_set=None,
|
||||
callbacks=None,
|
||||
):
|
||||
assert startup_program is not None
|
||||
self._loss = loss
|
||||
self._startup_program = startup_program
|
||||
self._main_program = loss.block.program
|
||||
self._parameter_list = parameter_list
|
||||
self._no_grad_set = no_grad_set
|
||||
self._callbacks = callbacks
|
||||
|
||||
if self._enable_auto_mapping and self._need_rank_mapping:
|
||||
# Do the mapping pass before parallelization
|
||||
assert self._cluster is not None, (
|
||||
"The cluster must not be none when using auto mapping."
|
||||
)
|
||||
dist_programs = {}
|
||||
world_process_group = get_world_process_group()
|
||||
dist_context = None
|
||||
# auto search
|
||||
if self._dist_strategy.auto_search:
|
||||
logging.info("Start searching dist attr.")
|
||||
serial_program_info = SerialProgramInfo(
|
||||
self._main_program,
|
||||
self._startup_program,
|
||||
self._loss,
|
||||
self._optimizer,
|
||||
self._cluster,
|
||||
)
|
||||
planner = Planner(
|
||||
serial_program_info,
|
||||
self,
|
||||
algorithm_config={"name": "mcmc", "max_search_times": 5},
|
||||
)
|
||||
dist_context, _ = planner.search()
|
||||
logging.info("End searching dist attr.")
|
||||
|
||||
# serialize the dist context by planner
|
||||
if dist_context is not None:
|
||||
logging.info("Start serialize searched dist attr")
|
||||
cwd = pathlib.Path().cwd()
|
||||
searched_dist_context_path = os.path.join(
|
||||
cwd, f"searched_dist_context_{time.time()}.pkl"
|
||||
)
|
||||
saved_dist_context = {}
|
||||
ops_dist_attr = {}
|
||||
tensors_dist_attr = {}
|
||||
for key, dist_op in dist_context._dist_ops_for_program.items():
|
||||
ops_dist_attr[key] = dist_op.dist_attr
|
||||
for (
|
||||
key,
|
||||
dist_tensor,
|
||||
) in dist_context._dist_tensors_for_program.items():
|
||||
tensors_dist_attr[key] = dist_tensor.dist_attr
|
||||
saved_dist_context["ops_dist_attr"] = ops_dist_attr
|
||||
saved_dist_context["tensors_dist_attr"] = tensors_dist_attr
|
||||
saved_dist_context["process_meshes"] = (
|
||||
dist_context._process_meshes
|
||||
)
|
||||
with open(
|
||||
searched_dist_context_path, "wb"
|
||||
) as dist_context_file:
|
||||
pickle.dump(saved_dist_context, dist_context_file)
|
||||
os.environ['PADDLE_SEARCHED_DIST_CONTEXT_PATH'] = (
|
||||
searched_dist_context_path
|
||||
)
|
||||
logging.info(
|
||||
f"End serialize searched dist attr to {searched_dist_context_path}"
|
||||
)
|
||||
|
||||
for rank in world_process_group.ranks:
|
||||
(
|
||||
dist_optimize_ops,
|
||||
dist_params_grads,
|
||||
dist_startup_prog,
|
||||
dist_main_prog,
|
||||
g_process_group_map,
|
||||
) = self._get_dist_program(rank, dist_context)
|
||||
dist_programs[rank] = [dist_main_prog, g_process_group_map]
|
||||
|
||||
# Do the mapping between the distributed program graph and the cluster graph
|
||||
rank_mapping_dict = mapping(dist_programs, self._cluster)
|
||||
rank_mapping = list(rank_mapping_dict.values())
|
||||
|
||||
# Relaunch the training by using the rank mapping file
|
||||
with open(self._rank_mapping_path, "w") as rank_mapping_file:
|
||||
json.dump(rank_mapping, rank_mapping_file)
|
||||
|
||||
enable_elastic = os.getenv("PADDLE_ENABLE_ELASTIC")
|
||||
enable_elastic = (
|
||||
True
|
||||
if enable_elastic and enable_elastic.lower() == 'true'
|
||||
else False
|
||||
)
|
||||
if enable_elastic:
|
||||
print("Auto mapping finished, now do elastic re-launch")
|
||||
sys.exit(
|
||||
paddle.distributed.fleet.elastic.manager.ELASTIC_AUTO_PARALLEL_EXIT_CODE
|
||||
)
|
||||
|
||||
original_cmd_args = os.getenv("PADDLE_ORIGINAL_CMD_ARGS")
|
||||
rank_mapping_args = " ".join(
|
||||
["--rank_mapping_path", self._rank_mapping_path]
|
||||
)
|
||||
if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
|
||||
coverage_args = ["-m", "coverage", "run", "--branch", "-p"]
|
||||
else:
|
||||
coverage_args = []
|
||||
new_cmd_args = (
|
||||
"-m paddle.distributed.fleet.launch"
|
||||
+ " "
|
||||
+ rank_mapping_args
|
||||
+ " "
|
||||
+ original_cmd_args
|
||||
)
|
||||
new_cmd = [
|
||||
sys.executable,
|
||||
"-u",
|
||||
*coverage_args,
|
||||
*shlex.split(new_cmd_args),
|
||||
]
|
||||
new_process = subprocess.Popen(new_cmd)
|
||||
new_process.wait()
|
||||
assert new_process.returncode == 0, (
|
||||
"Launch failed with rank mapping"
|
||||
)
|
||||
print("Successfully do the second launch for auto mapping!")
|
||||
sys.exit(0)
|
||||
else:
|
||||
# Parallelization after the mapping pass
|
||||
rank = paddle.distributed.get_rank()
|
||||
dist_context = None
|
||||
searched_dist_context_path = os.getenv(
|
||||
"PADDLE_SEARCHED_DIST_CONTEXT_PATH", None
|
||||
)
|
||||
if searched_dist_context_path is not None:
|
||||
with open(
|
||||
searched_dist_context_path, "rb"
|
||||
) as dist_context_file:
|
||||
from paddle.framework.restricted_unpickler import (
|
||||
safe_load_pickle,
|
||||
)
|
||||
|
||||
saved_dist_context = safe_load_pickle(dist_context_file)
|
||||
dist_context = DistributedContext()
|
||||
for op in self._main_program.global_block().ops:
|
||||
dist_attr = saved_dist_context["ops_dist_attr"][
|
||||
op.desc.id()
|
||||
]
|
||||
dist_op = DistributedOperator(op, dist_attr)
|
||||
dist_context.add_dist_op_for_program(dist_op)
|
||||
|
||||
vars = self._main_program.global_block().vars
|
||||
for var in vars.values():
|
||||
dist_attr = saved_dist_context["tensors_dist_attr"][
|
||||
var.desc.id()
|
||||
]
|
||||
dist_tensor = DistributedTensor(var, dist_attr)
|
||||
dist_context.add_dist_tensor_for_program(dist_tensor)
|
||||
|
||||
dist_context._process_meshes = saved_dist_context[
|
||||
"process_meshes"
|
||||
]
|
||||
|
||||
else:
|
||||
if self._dist_strategy.auto_search:
|
||||
serial_program_info = SerialProgramInfo(
|
||||
self._main_program,
|
||||
self._startup_program,
|
||||
self._loss,
|
||||
self._optimizer,
|
||||
cluster=self._cluster,
|
||||
)
|
||||
planner = Planner(
|
||||
serial_program_info,
|
||||
self,
|
||||
algorithm_config={
|
||||
"name": "mcmc",
|
||||
"max_search_times": 5,
|
||||
},
|
||||
)
|
||||
dist_context, _ = planner.search()
|
||||
|
||||
# rebuild g_process_group
|
||||
if dist_context is not None:
|
||||
pg0 = get_process_group(0)
|
||||
for process_mesh in dist_context._process_meshes:
|
||||
pg0.add_ranks(process_mesh.process_ids)
|
||||
(
|
||||
dist_optimize_ops,
|
||||
dist_params_grads,
|
||||
dist_startup_prog,
|
||||
dist_main_prog,
|
||||
_,
|
||||
) = self._get_dist_program(rank, dist_context, relaunch_phase=True)
|
||||
|
||||
# NOTE: This is a trick to fix hang in pipeline mode when dist context is searched by planner
|
||||
if self._dist_strategy.auto_search:
|
||||
is_pipeline = False
|
||||
for op in dist_main_prog.global_block().ops:
|
||||
if op.type == "send_v2" or op.type == "recv_v2":
|
||||
is_pipeline = True
|
||||
break
|
||||
if is_pipeline:
|
||||
with paddle.static.program_guard(dist_main_prog):
|
||||
paddle.distributed.barrier()
|
||||
|
||||
# Traverse different rank programs and traverse each op of them,
|
||||
# instantiate communication by process_mapping.
|
||||
all_process_groups = get_all_process_groups()
|
||||
for process_group in all_process_groups:
|
||||
process_group.instantiate()
|
||||
|
||||
# Copy distributed info to the default context
|
||||
set_default_distributed_context(self._dist_context)
|
||||
|
||||
# The last step: remove all distributed attributes to be compatible
|
||||
# with inference.
|
||||
self._remove_distributed_attrs(dist_main_prog)
|
||||
|
||||
return (
|
||||
dist_optimize_ops,
|
||||
dist_params_grads,
|
||||
dist_startup_prog,
|
||||
dist_main_prog,
|
||||
)
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
cls = self.__class__
|
||||
result = cls.__new__(cls)
|
||||
memo[id(self)] = result
|
||||
for k, v in self.__dict__.items():
|
||||
if (
|
||||
k == "_main_program"
|
||||
or k == "_startup_program"
|
||||
or k == "_dist_context"
|
||||
or k == "_fleet"
|
||||
or k == "_loss"
|
||||
):
|
||||
setattr(result, k, v)
|
||||
else:
|
||||
setattr(result, k, copy.deepcopy(v, memo))
|
||||
return result
|
||||
@@ -0,0 +1,553 @@
|
||||
# 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 logging
|
||||
import os
|
||||
import time
|
||||
|
||||
from paddle.distributed.passes.pass_base import PassManager, new_pass
|
||||
from paddle.framework import get_flags
|
||||
from paddle.static import append_backward, program_guard
|
||||
|
||||
from ...utils.log_utils import get_logger
|
||||
from ..random import init_auto_parallel_rng
|
||||
from .partitioner import Partitioner
|
||||
from .process_group import get_world_process_group
|
||||
from .reshard import Resharder
|
||||
from .utils import (
|
||||
get_pp_stage,
|
||||
is_sequential_run,
|
||||
)
|
||||
|
||||
PIR_PASS = [
|
||||
'fused_gemm_epilogue_pass',
|
||||
'fused_linear_param_grad_add_pass',
|
||||
'fuse_allreduce_split_to_reducescatter_pass',
|
||||
'fused_dropout_add_pass',
|
||||
]
|
||||
|
||||
PIR_PYTHON_PASS = [
|
||||
'eliminate_transpose',
|
||||
]
|
||||
|
||||
|
||||
class Parallelizer:
|
||||
def __init__(self, mode, completer, dist_context):
|
||||
self._mode = mode
|
||||
self._completer = completer
|
||||
self._dist_context = dist_context
|
||||
assert self._dist_context._is_initialized
|
||||
self._pass_context = self._dist_context.pass_context
|
||||
self._strategy = self._dist_context.strategy
|
||||
self._logger = get_logger(logging.INFO)
|
||||
|
||||
@property
|
||||
def is_train(self):
|
||||
return self._mode == "train"
|
||||
|
||||
@property
|
||||
def is_test(self):
|
||||
return self._mode in ["eval", "predict"]
|
||||
|
||||
def parallel_all(self, parameter_list=None):
|
||||
world_process_group = get_world_process_group()
|
||||
all_ranks = world_process_group.ranks
|
||||
for rank in all_ranks:
|
||||
# self._dist_context._backup(serial=True, dist=True)
|
||||
self.parallel(rank, parameter_list)
|
||||
# self._dist_context._restore(serial=True, dist=True)
|
||||
|
||||
def parallel(self, rank, parameter_list=None):
|
||||
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
|
||||
if self.is_train and serial_optimizer:
|
||||
# Generate backward
|
||||
serial_loss = self._dist_context.serial_loss
|
||||
params_grads = self._generate_backward(
|
||||
serial_main_program,
|
||||
serial_startup_program,
|
||||
serial_loss,
|
||||
parameter_list,
|
||||
)
|
||||
# Apply pre optimization passes
|
||||
time0 = time.time()
|
||||
(
|
||||
serial_main_program,
|
||||
serial_startup_program,
|
||||
params_grads,
|
||||
) = self._apply_pre_optimization(
|
||||
serial_main_program,
|
||||
serial_startup_program,
|
||||
serial_loss,
|
||||
serial_optimizer,
|
||||
params_grads,
|
||||
)
|
||||
self._logger.debug(
|
||||
f"within parallel apply_pre_optimization time: {time.time() - time0}, mode {self._mode}"
|
||||
)
|
||||
# Do logical partition
|
||||
time0 = time.time()
|
||||
partitioner = Partitioner(self._dist_context, rank)
|
||||
(
|
||||
dist_main_prog,
|
||||
dist_startup_prog,
|
||||
dist_params_grads,
|
||||
) = partitioner.partition(
|
||||
serial_main_program, serial_startup_program, params_grads
|
||||
)
|
||||
|
||||
init_auto_parallel_rng()
|
||||
|
||||
self._logger.debug(
|
||||
f"within parallel partitioner time: {time.time() - time0}, mode {self._mode}"
|
||||
)
|
||||
# Generate optimizer
|
||||
time0 = time.time()
|
||||
self._generate_optimizer(
|
||||
dist_main_prog,
|
||||
dist_startup_prog,
|
||||
serial_optimizer,
|
||||
dist_params_grads,
|
||||
)
|
||||
self._logger.debug(
|
||||
f"within parallel optimizer time: {time.time() - time0}, mode {self._mode}"
|
||||
)
|
||||
|
||||
resharder = Resharder(
|
||||
dist_main_prog,
|
||||
dist_startup_prog,
|
||||
rank,
|
||||
self._dist_context,
|
||||
dist_params_grads,
|
||||
)
|
||||
resharder.reshard()
|
||||
self._logger.debug(
|
||||
f"within parallel reshard time: {time.time() - time0}, mode {self._mode}"
|
||||
)
|
||||
# Apply post optimization passes
|
||||
time0 = time.time()
|
||||
self._apply_post_optimization(
|
||||
dist_main_prog, dist_startup_prog, rank, dist_params_grads
|
||||
)
|
||||
self._logger.debug(
|
||||
f"within parallel apply_post_optimization time: {time.time() - time0}, mode {self._mode}"
|
||||
)
|
||||
else:
|
||||
# Apply pre optimization passes
|
||||
time0 = time.time()
|
||||
(
|
||||
serial_main_program,
|
||||
serial_startup_program,
|
||||
params_grads,
|
||||
) = self._apply_pre_optimization(
|
||||
serial_main_program, serial_startup_program, None, None, []
|
||||
)
|
||||
self._logger.debug(
|
||||
f"within parallel apply_pre_optimization time: {time.time() - time0}, mode {self._mode}"
|
||||
)
|
||||
# Do logical partition
|
||||
time0 = time.time()
|
||||
partitioner = Partitioner(self._dist_context, rank)
|
||||
(
|
||||
dist_main_prog,
|
||||
dist_startup_prog,
|
||||
dist_params_grads,
|
||||
) = partitioner.partition(
|
||||
serial_main_program, serial_startup_program, []
|
||||
)
|
||||
# Do reshard process
|
||||
self._logger.debug(
|
||||
f"within parallel partitioner time: {time.time() - time0}, mode {self._mode}"
|
||||
)
|
||||
time0 = time.time()
|
||||
# Do reshard process
|
||||
micro_bsz = (
|
||||
1
|
||||
if not self._strategy.pipeline.enable
|
||||
else self._strategy.pipeline.micro_batch_size
|
||||
)
|
||||
resharder = Resharder(
|
||||
dist_main_prog,
|
||||
dist_startup_prog,
|
||||
rank,
|
||||
self._dist_context,
|
||||
[],
|
||||
micro_bsz,
|
||||
)
|
||||
resharder.reshard()
|
||||
self._logger.debug(
|
||||
f"within parallel reshard time: {time.time() - time0}, mode {self._mode}"
|
||||
)
|
||||
# Apply post optimization passes
|
||||
time0 = time.time()
|
||||
self._apply_post_optimization(
|
||||
dist_main_prog, dist_startup_prog, rank, dist_params_grads
|
||||
)
|
||||
self._logger.debug(
|
||||
f"within parallel apply_post_optimization time: {time.time() - time0}, mode {self._mode}"
|
||||
)
|
||||
|
||||
# Clone program for test
|
||||
if self.is_test:
|
||||
pipeline_opt = dist_main_prog._pipeline_opt
|
||||
dist_main_prog = dist_main_prog.clone(for_test=True)
|
||||
dist_startup_prog = dist_startup_prog.clone(for_test=True)
|
||||
dist_main_prog._pipeline_opt = pipeline_opt
|
||||
|
||||
# Store the distributed programs for further usages
|
||||
self._dist_context.dist_main_programs[rank] = dist_main_prog
|
||||
self._dist_context.dist_startup_programs[rank] = dist_startup_prog
|
||||
|
||||
def _generate_backward(
|
||||
self, main_program, startup_program, loss, parameter_list=None
|
||||
):
|
||||
# NOTE(zhaoyinglia):
|
||||
# Guarantee the order of params_grads is same between dynamic mode and static mode
|
||||
# by making parameter_list equal to model.parameters(),
|
||||
# because the order affect the result of ClipGradByGLobalNorm.
|
||||
# If parameter_list is not None, the order of params_grads is same with parameter_list.
|
||||
# If parameter_list is None, params_grads will be as prog.global_block().all_parameters().
|
||||
with program_guard(main_program, startup_program):
|
||||
params_grads = append_backward(
|
||||
loss,
|
||||
parameter_list=parameter_list,
|
||||
distop_context=self._dist_context.dist_op_context,
|
||||
)
|
||||
self._completer.complete_backward_annotation(main_program)
|
||||
self._dist_context.block_state.parse_backward_blocks(main_program)
|
||||
return params_grads
|
||||
|
||||
def _generate_optimizer(
|
||||
self, main_program, startup_program, optimizer, params_grads
|
||||
):
|
||||
# NOTE:
|
||||
# 1. `apply_gradients` will add an Accumulator for a parameter only once,
|
||||
# but optimizer will be called repeatedly in re-launch, so optimizer need to be copied.
|
||||
# 2. lr_scheduler cannot be deepcopy, cause 'deepcopy' will lead to difference of learning_rate between executor and engine.
|
||||
learning_rate = optimizer._learning_rate
|
||||
new_optimizer = copy.deepcopy(optimizer)
|
||||
new_optimizer._learning_rate = learning_rate
|
||||
new_optimizer._sorted = False
|
||||
self._dist_context._serial_optimizer = optimizer
|
||||
self._dist_context._serial_optimizer._learning_rate = learning_rate
|
||||
|
||||
with (
|
||||
program_guard(main_program, startup_program),
|
||||
main_program.switch_name_generator_guard("opt_"),
|
||||
):
|
||||
optimizer_ops = new_optimizer.apply_gradients(params_grads)
|
||||
self._completer.complete_update_annotation(main_program)
|
||||
return optimizer_ops
|
||||
|
||||
def _apply_pre_optimization(
|
||||
self, main_program, startup_program, loss, optimizer, params_grads
|
||||
):
|
||||
if self._strategy is None:
|
||||
return
|
||||
|
||||
# apply amp pass on train/eval/predict
|
||||
if self._strategy.amp.enable:
|
||||
config = copy.deepcopy(self._strategy.amp.to_dict())
|
||||
config["dist_context"] = self._dist_context
|
||||
config["params_grads"] = params_grads
|
||||
config["loss"] = loss
|
||||
config["input_data"] = (
|
||||
self._dist_context.serial_feed_vars["inputs"]
|
||||
+ self._dist_context.serial_feed_vars["labels"]
|
||||
)
|
||||
self._logger.info(
|
||||
"Applying AMP-{}-{} ...".format(
|
||||
config["dtype"], config['level']
|
||||
),
|
||||
)
|
||||
if config['level'] == "o1":
|
||||
auto_parallel_amp_pass = new_pass("auto_parallel_amp", config)
|
||||
auto_parallel_amp_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
loss = auto_parallel_amp_pass.get_loss()
|
||||
elif config['level'] in ['o2', 'o3']:
|
||||
config["base_opt"] = optimizer
|
||||
auto_parallel_fp16_pass = new_pass("auto_parallel_fp16", config)
|
||||
auto_parallel_fp16_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
loss = auto_parallel_fp16_pass.get_loss()
|
||||
else:
|
||||
raise ValueError("AMP level should be one of o1, o2, o3")
|
||||
|
||||
# apply quantization pass
|
||||
# The pass can be applied when mode must be 'train'
|
||||
if self.is_train and self._strategy.qat.enable:
|
||||
config = copy.deepcopy(self._strategy.qat.to_dict())
|
||||
config["dist_context"] = self._dist_context
|
||||
config["params_grads"] = params_grads
|
||||
config["mode"] = self._mode
|
||||
config["loss"] = loss
|
||||
auto_parallel_quantization_pass = new_pass(
|
||||
"auto_parallel_quantization", config
|
||||
)
|
||||
auto_parallel_quantization_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
main_program = self._pass_context.get_attr("main_program")
|
||||
startup_program = self._pass_context.get_attr("startup_program")
|
||||
params_grads = self._pass_context.get_attr("params_grads")
|
||||
loss = self._pass_context.get_attr("loss")
|
||||
|
||||
# apply recompute pass
|
||||
# recompute is then train-only optimization
|
||||
if self.is_train and self._strategy.recompute.enable:
|
||||
config = copy.deepcopy(self._strategy.recompute.to_dict())
|
||||
config["dist_context"] = self._dist_context
|
||||
config["no_grad_set"] = None
|
||||
config["loss"] = loss
|
||||
auto_parallel_recompute_pass = new_pass(
|
||||
"auto_parallel_recompute", config
|
||||
)
|
||||
auto_parallel_recompute_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
|
||||
return main_program, startup_program, params_grads
|
||||
|
||||
def _check_dist_attr(self, program, num_model_chunks, dist_context):
|
||||
for _, block in enumerate(program.blocks):
|
||||
for _, op in enumerate(block.ops):
|
||||
op_dist_attr = dist_context.get_op_dist_attr_for_program(op)
|
||||
if op_dist_attr is None:
|
||||
raise ValueError(
|
||||
f"There is not dist_attr for op[{op.type}]."
|
||||
)
|
||||
|
||||
def _apply_post_optimization(
|
||||
self, main_program, startup_program, rank, params_grads
|
||||
):
|
||||
if self._strategy is None:
|
||||
return
|
||||
|
||||
# sequence parallel optimization
|
||||
if self._strategy.sp_optimization.enable:
|
||||
config = copy.deepcopy(self._strategy.sp_optimization.to_dict())
|
||||
config["dist_context"] = self._dist_context
|
||||
config["global_rank"] = rank
|
||||
sp_pass = new_pass(
|
||||
"auto_parallel_sequence_parallel_optimization", config
|
||||
)
|
||||
sp_pass.apply([main_program], [startup_program], self._pass_context)
|
||||
|
||||
# apply fused linear promotion pass
|
||||
if (
|
||||
self.is_train
|
||||
and self._strategy.fused_linear_promotion.enable
|
||||
and self._strategy.fused_passes.enable
|
||||
):
|
||||
if (
|
||||
len(self._strategy.fused_passes.fused_passes_list) > 0
|
||||
and "fuse_gemm_epilogue"
|
||||
in self._strategy.fused_passes.fused_passes_list
|
||||
):
|
||||
amp_config = None
|
||||
if self._strategy.amp.enable:
|
||||
amp_config = copy.deepcopy(self._strategy.amp.to_dict())
|
||||
config = {}
|
||||
config["dist_context"] = self._dist_context
|
||||
config["global_rank"] = rank
|
||||
config["enable_sp"] = self._strategy.sp_optimization.enable
|
||||
config["params_grads"] = params_grads
|
||||
config["amp_level"] = (
|
||||
amp_config['level'] if amp_config is not None else "o0"
|
||||
)
|
||||
fused_linear_promotion_pass = new_pass(
|
||||
"auto_parallel_fused_linear_promotion", config
|
||||
)
|
||||
fused_linear_promotion_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
|
||||
# apply master grad pass
|
||||
if self._strategy.amp.enable:
|
||||
amp_config = copy.deepcopy(self._strategy.amp.to_dict())
|
||||
config = {}
|
||||
config["dist_context"] = self._dist_context
|
||||
config["params_grads"] = params_grads
|
||||
config["completer"] = self._completer
|
||||
if amp_config['level'] == "o2" and amp_config["use_master_grad"]:
|
||||
master_grad_pass = new_pass(
|
||||
"auto_parallel_master_grad_pass", config
|
||||
)
|
||||
master_grad_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
|
||||
# data parallel optimization
|
||||
if self._strategy.dp_optimization.enable:
|
||||
config = copy.deepcopy(self._strategy.dp_optimization.to_dict())
|
||||
config["dist_context"] = self._dist_context
|
||||
config["global_rank"] = rank
|
||||
config["use_sharding"] = self._strategy.sharding.enable
|
||||
dp_pass = new_pass(
|
||||
"auto_parallel_data_parallel_optimization", config
|
||||
)
|
||||
dp_pass.apply([main_program], [startup_program], self._pass_context)
|
||||
|
||||
gradient_sync_after_accumulate = (
|
||||
self._strategy.dp_optimization.gradient_sync_after_accumulate
|
||||
)
|
||||
if gradient_sync_after_accumulate:
|
||||
global_params_grads = params_grads
|
||||
|
||||
if self._strategy.sharding.enable:
|
||||
config = copy.deepcopy(self._strategy.sharding.to_dict())
|
||||
config["dist_context"] = self._dist_context
|
||||
config["params_grads"] = params_grads
|
||||
config["global_rank"] = rank
|
||||
config["gradient_sync_after_accumulate"] = (
|
||||
gradient_sync_after_accumulate
|
||||
)
|
||||
if self._strategy.amp.enable:
|
||||
amp_config = copy.deepcopy(self._strategy.amp.to_dict())
|
||||
config["amp_dtype"] = amp_config['dtype']
|
||||
auto_parallel_sharding_pass = new_pass(
|
||||
"auto_parallel_sharding", config
|
||||
)
|
||||
auto_parallel_sharding_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
params_grads = self._pass_context.get_attr("params_grads")
|
||||
|
||||
if self._strategy.mp_optimization.allreduce_matmul_grad_overlapping:
|
||||
if int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1:
|
||||
self._logger.warning(
|
||||
"You set mp_optimization.allreduce_matmul_grad_overlapping=True, but you did not set environment "
|
||||
"variable CUDA_DEVICE_MAX_CONNECTIONS=1, which may leads to performance "
|
||||
"loss. Try to export CUDA_DEVICE_MAX_CONNECTIONS=1 for better performance."
|
||||
)
|
||||
|
||||
config = {
|
||||
"dist_context": self._dist_context,
|
||||
}
|
||||
allreduce_matmul_grad_overlapping_pass = new_pass(
|
||||
"allreduce_matmul_grad_overlapping", config
|
||||
)
|
||||
allreduce_matmul_grad_overlapping_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
|
||||
if self.is_train:
|
||||
# GradClip is train-only optimization
|
||||
config = copy.deepcopy(self._strategy.sharding.to_dict())
|
||||
config["dist_context"] = self._dist_context
|
||||
config["params_grads"] = params_grads
|
||||
config["rank_id"] = rank
|
||||
auto_parallel_clip_pass = new_pass(
|
||||
"auto_parallel_grad_clip", config
|
||||
)
|
||||
auto_parallel_clip_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
|
||||
if not is_sequential_run():
|
||||
# deps for newexe
|
||||
config = {}
|
||||
config["dist_context"] = self._dist_context
|
||||
APSED_pass = new_pass(
|
||||
"auto_parallel_supplement_explicit_dependencies", config
|
||||
)
|
||||
APSED_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
|
||||
if self.is_train and self._strategy.pipeline.enable:
|
||||
self._strategy.gradient_merge.enable = True
|
||||
self._strategy.gradient_merge.k_steps = (
|
||||
self._strategy.pipeline.accumulate_steps
|
||||
)
|
||||
self._strategy.gradient_merge.avg = True
|
||||
|
||||
# gradient_merge is then train-only optimization
|
||||
grad_to_global_grad = {}
|
||||
if self.is_train and self._strategy.gradient_merge.enable:
|
||||
config = copy.deepcopy(self._strategy.gradient_merge.to_dict())
|
||||
config["dist_context"] = self._dist_context
|
||||
config["grad_to_global_grad"] = grad_to_global_grad
|
||||
config["pipeline_mode"] = self._strategy.pipeline.schedule_mode
|
||||
if gradient_sync_after_accumulate:
|
||||
config["params_grads"] = global_params_grads
|
||||
config["gradient_sync_after_accumulate"] = (
|
||||
gradient_sync_after_accumulate
|
||||
)
|
||||
else:
|
||||
config["params_grads"] = params_grads
|
||||
auto_parallel_gradient_merge_pass = new_pass(
|
||||
"auto_parallel_gradient_merge_pass", config
|
||||
)
|
||||
auto_parallel_gradient_merge_pass.apply(
|
||||
[main_program], [startup_program], self._pass_context
|
||||
)
|
||||
|
||||
self._check_dist_attr(
|
||||
main_program,
|
||||
self._strategy.pipeline.vpp_degree,
|
||||
self._dist_context,
|
||||
)
|
||||
|
||||
enable_ir = get_flags("FLAGS_enable_pir_in_executor")[
|
||||
'FLAGS_enable_pir_in_executor'
|
||||
]
|
||||
ir_pass_list = []
|
||||
if self.is_train and self._strategy.fused_passes.enable:
|
||||
if len(self._strategy.fused_passes.fused_passes_list) > 0:
|
||||
program_pass_list = []
|
||||
for p in self._strategy.fused_passes.fused_passes_list:
|
||||
if enable_ir and p in (PIR_PASS + PIR_PYTHON_PASS):
|
||||
ir_pass_list.append(p)
|
||||
else:
|
||||
program_pass_list.append(new_pass(p))
|
||||
pass_manager = PassManager(program_pass_list)
|
||||
pass_manager.apply([main_program], [startup_program])
|
||||
|
||||
main_program._pass_opt = {}
|
||||
main_program._pass_opt['pass_list'] = ir_pass_list
|
||||
|
||||
if self.is_train and self._strategy.pipeline.enable:
|
||||
enable_send_recv_overlap = (
|
||||
self._strategy.pipeline.enable_send_recv_overlap
|
||||
)
|
||||
if (
|
||||
enable_send_recv_overlap
|
||||
and int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1
|
||||
):
|
||||
self._logger.warning(
|
||||
"You set pipeline.enable_send_recv_overlap=True, but you did not set environment "
|
||||
"variable CUDA_DEVICE_MAX_CONNECTIONS=1, which may leads to performance "
|
||||
"loss. Try to export CUDA_DEVICE_MAX_CONNECTIONS=1 for better performance."
|
||||
)
|
||||
|
||||
main_program._pipeline_opt = {}
|
||||
main_program._pipeline_opt["standalone_opt"] = {
|
||||
"enable_send_recv_overlap": enable_send_recv_overlap,
|
||||
"schedule_mode": self._strategy.pipeline.schedule_mode,
|
||||
"num_micro_batches": self._strategy.pipeline.accumulate_steps,
|
||||
"pp_degree": len(self._dist_context.process_meshes),
|
||||
"pp_stage": get_pp_stage(self._dist_context, rank),
|
||||
"vpp_degree": self._strategy.pipeline.vpp_degree,
|
||||
"dist_context": self._dist_context,
|
||||
"program_runtimes": self._strategy.pipeline.program_runtimes,
|
||||
"memory_limit_times": self._strategy.pipeline.memory_limit_times,
|
||||
"split_backward": self._strategy.pipeline.split_backward,
|
||||
"grad_to_global_grad": grad_to_global_grad,
|
||||
}
|
||||
@@ -0,0 +1,543 @@
|
||||
# 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
|
||||
from collections import defaultdict
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.auto_parallel.static.dist_context import (
|
||||
DistributedContext,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.operators.common import (
|
||||
get_distributed_operator_impl_container,
|
||||
)
|
||||
from paddle.framework import Program, core
|
||||
from paddle.static import Parameter
|
||||
|
||||
from .dist_attribute import OperatorDistAttr
|
||||
from .operators.common import BACKWARD_ONLY_DIST_OPS
|
||||
from .utils import (
|
||||
__no_shape_var_type__,
|
||||
is_backward_op,
|
||||
is_forward_op,
|
||||
is_loss_op,
|
||||
is_optimize_op,
|
||||
)
|
||||
|
||||
__varname_not_in_block__ = ["lod_tensor_blocking_queue"]
|
||||
|
||||
|
||||
class Partitioner:
|
||||
"""
|
||||
warning:: Partitioner is experimental and subject to change.
|
||||
|
||||
Partitioner convert a program into another program.
|
||||
Given a serial program which has been auto completed with shard annotation, the Partitioner
|
||||
convert the serial program into a "distributed" program. The Partitioner will modify the serial
|
||||
program in following two ways, which is also the major difference between serial and distributed program:
|
||||
1. partition op: replace a serial op into its corresponding dist op inferred from the shard annotation
|
||||
2. partition var: if a var is sharded, modify the shape of var according to its shard annotation
|
||||
|
||||
Partitioner is supposed to be call by the auto parallel framework, and not supposed to be directly called by user.
|
||||
"""
|
||||
|
||||
def __init__(self, dist_context, rank_id=0):
|
||||
"""
|
||||
Args:
|
||||
dist_context (DistributedContext): used to access the distributed_attr of var & op, every Partitioner object could maintain its own DistributedContext member, and partition program base on that shard scenario.
|
||||
rank_id (int): global rank id to which the partitioned distributed program belong.
|
||||
"""
|
||||
if not isinstance(dist_context, DistributedContext):
|
||||
raise TypeError(
|
||||
f"dist_context be DistributedContext, got {type(dist_context)} here"
|
||||
)
|
||||
|
||||
self._dist_context = dist_context
|
||||
self._rank_id = rank_id
|
||||
self._serial2dist_varname_mapping = defaultdict(
|
||||
dict
|
||||
) # block_id -> serial_varname -> dist_varname
|
||||
self._dist_varname_suffix = ""
|
||||
self._forward_op_id2forward_op = {}
|
||||
|
||||
def partition(
|
||||
self, serial_main_program, serial_startup_program, params_grads
|
||||
):
|
||||
if not isinstance(serial_main_program, (Program)):
|
||||
raise TypeError(
|
||||
f"main_program be paddle.framework.Program, got {type(serial_main_program)} here"
|
||||
)
|
||||
|
||||
# check if shard annotated serial program valid
|
||||
if not self._is_valid_annotated_program(serial_main_program):
|
||||
raise RuntimeError(
|
||||
"Not all vars or ops are annotated in main program !"
|
||||
)
|
||||
|
||||
# init distop helper
|
||||
dist_op_context = self._dist_context.dist_op_context
|
||||
dist_op_context.varname_mapping = self._serial2dist_varname_mapping
|
||||
dist_op_context.rank_id = self._rank_id
|
||||
|
||||
# partition startup program
|
||||
if serial_startup_program is None:
|
||||
partitioned_startup_prog = None
|
||||
else:
|
||||
partitioned_startup_prog = self.partition_startup_program(
|
||||
serial_main_program, serial_startup_program
|
||||
)
|
||||
dist_op_context.dst_startup_program = partitioned_startup_prog
|
||||
|
||||
# partition main program
|
||||
(
|
||||
partitioned_main_prog,
|
||||
partitioned_params_grads,
|
||||
) = self.partition_main_program(serial_main_program, params_grads)
|
||||
|
||||
return (
|
||||
partitioned_main_prog,
|
||||
partitioned_startup_prog,
|
||||
partitioned_params_grads,
|
||||
)
|
||||
|
||||
def partition_startup_program(
|
||||
self, serial_main_program, serial_startup_program
|
||||
):
|
||||
if not isinstance(serial_startup_program, (Program)):
|
||||
raise TypeError(
|
||||
f"dist_context be paddle.framework.Program, got {type(serial_startup_program)} here"
|
||||
)
|
||||
|
||||
partitioned_startup_prog = paddle.framework.Program()
|
||||
partitioned_startup_prog._name_generator = (
|
||||
serial_startup_program._name_generator.clone()
|
||||
)
|
||||
ref_block = serial_main_program.global_block()
|
||||
target_block = partitioned_startup_prog.global_block()
|
||||
var2shape = {}
|
||||
temp_varname_map = {}
|
||||
|
||||
# tensors
|
||||
for var in serial_startup_program.list_vars():
|
||||
assert var.persistable
|
||||
new_name = var.name + self._dist_varname_suffix
|
||||
temp_varname_map[var.name] = new_name
|
||||
target_shape = _partition_var(
|
||||
self._dist_context, ref_block, target_block, var.name, new_name
|
||||
)
|
||||
var2shape[new_name] = target_shape
|
||||
|
||||
# ops
|
||||
for op in serial_startup_program.global_block().ops:
|
||||
# TODO if var not belong to this rank, should be filtered
|
||||
output_vars = op.desc.output_arg_names()
|
||||
assert len(output_vars) == 1, (
|
||||
f"initializer should output only ONE variable, but got [{op.desc}]"
|
||||
)
|
||||
assert temp_varname_map[output_vars[0]] in var2shape, (
|
||||
f"try to initialize [{output_vars[0]}] which is not a persistable var"
|
||||
)
|
||||
new_op_desc = target_block.desc.append_op()
|
||||
new_op_desc.copy_from(op.desc)
|
||||
new_op_desc._rename_output(
|
||||
output_vars[0], temp_varname_map[output_vars[0]]
|
||||
)
|
||||
new_op_desc._set_attr(
|
||||
"shape", var2shape[temp_varname_map[output_vars[0]]]
|
||||
)
|
||||
target_block._sync_with_cpp()
|
||||
|
||||
# set distribute attribute
|
||||
new_op = target_block.ops[-1]
|
||||
assert new_op.type == new_op_desc.type()
|
||||
assert new_op.desc == new_op_desc
|
||||
output_var = target_block.var(output_vars[0])
|
||||
output_var_attr = (
|
||||
self._dist_context.get_tensor_dist_attr_for_program(output_var)
|
||||
)
|
||||
op_attr = OperatorDistAttr()
|
||||
op_attr.process_mesh = output_var_attr.process_mesh
|
||||
op_attr.set_output_dims_mapping(
|
||||
output_var.name, output_var_attr.dims_mapping
|
||||
)
|
||||
op_attr.set_input_dims_mapping(
|
||||
output_var.name, output_var_attr.dims_mapping
|
||||
)
|
||||
self._dist_context.set_op_dist_attr_for_program(new_op, op_attr)
|
||||
|
||||
return partitioned_startup_prog
|
||||
|
||||
def partition_main_program(self, serial_main_program, params_and_grads):
|
||||
"""
|
||||
1. partition variables
|
||||
2. replace local op with corresponding dist op
|
||||
"""
|
||||
|
||||
partitioned_main_prog = paddle.framework.Program()
|
||||
partitioned_main_prog._name_generator = (
|
||||
serial_main_program._name_generator.clone()
|
||||
)
|
||||
dist_op_context = self._dist_context.dist_op_context
|
||||
dist_op_context.dst_main_program = partitioned_main_prog
|
||||
|
||||
for idx in range(self._dist_context.block_state.nblock):
|
||||
ref_block = serial_main_program.blocks[idx]
|
||||
|
||||
if idx == 0:
|
||||
target_block = partitioned_main_prog.blocks[0]
|
||||
else:
|
||||
target_block = partitioned_main_prog._create_block(
|
||||
parent_idx=ref_block.parent_idx
|
||||
)
|
||||
assert ref_block.idx == target_block.idx
|
||||
target_block._set_forward_block_idx(ref_block.forward_block_idx)
|
||||
dist_op_context.work_block = target_block
|
||||
self.partition_block(ref_block, target_block)
|
||||
|
||||
partitioned_main_prog.current_block_idx = 0
|
||||
|
||||
# should reconnect the block_attr ptr to the correct block
|
||||
for block_id in range(self._dist_context.block_state.nblock):
|
||||
block = partitioned_main_prog.block(block_id)
|
||||
for op in block.ops:
|
||||
for attr_name in op.all_attrs():
|
||||
if op.attr_type(attr_name) == core.AttrType.BLOCK:
|
||||
relative_id = op._block_attr_id(attr_name)
|
||||
op._set_attr(
|
||||
attr_name, partitioned_main_prog.block(relative_id)
|
||||
)
|
||||
|
||||
partitioned_params_and_grads = []
|
||||
for p, g in params_and_grads:
|
||||
assert p.name in self._serial2dist_varname_mapping[0]
|
||||
dist_p = self._get_dist_var_by_serial_var(
|
||||
p, partitioned_main_prog, 0
|
||||
)
|
||||
if g is None:
|
||||
dist_g = None
|
||||
else:
|
||||
assert g.name in self._serial2dist_varname_mapping[0]
|
||||
dist_g = self._get_dist_var_by_serial_var(
|
||||
g, partitioned_main_prog, 0
|
||||
)
|
||||
partitioned_params_and_grads.append((dist_p, dist_g))
|
||||
|
||||
return partitioned_main_prog, partitioned_params_and_grads
|
||||
|
||||
def partition_block(self, ref_block, target_block):
|
||||
dist_op_context = self._dist_context.dist_op_context
|
||||
|
||||
last_fwd_op_idx = -1
|
||||
for idx, op in enumerate(ref_block.ops):
|
||||
if is_loss_op(op):
|
||||
last_fwd_op_idx = idx
|
||||
break
|
||||
|
||||
if last_fwd_op_idx == -1:
|
||||
last_fwd_op_idx = len(ref_block.ops)
|
||||
|
||||
for idx in range(len(ref_block.ops)):
|
||||
if idx <= last_fwd_op_idx:
|
||||
self._forward_op_id2forward_op[
|
||||
ref_block.ops[idx].desc.original_id()
|
||||
] = ref_block.ops[idx]
|
||||
|
||||
# partition
|
||||
appended_grad_times = 0
|
||||
for idx, op in enumerate(ref_block.ops):
|
||||
op_dist_attr = self._dist_context.get_op_dist_attr_for_program(op)
|
||||
if is_backward_op(op) and (
|
||||
is_forward_op(ref_block.ops[idx - 1])
|
||||
or is_loss_op(ref_block.ops[idx - 1])
|
||||
):
|
||||
if not op_dist_attr.is_recompute:
|
||||
appended_grad_times += 1
|
||||
|
||||
# partition input variables
|
||||
for serial_input_varname in op.desc.input_arg_names():
|
||||
if (
|
||||
serial_input_varname
|
||||
not in self._serial2dist_varname_mapping[
|
||||
ref_block.forward_block_idx
|
||||
]
|
||||
or serial_input_varname
|
||||
not in self._serial2dist_varname_mapping[ref_block.idx]
|
||||
):
|
||||
new_varname = (
|
||||
serial_input_varname + self._dist_varname_suffix
|
||||
)
|
||||
if ref_block.has_var(serial_input_varname):
|
||||
_partition_var(
|
||||
self._dist_context,
|
||||
ref_block,
|
||||
target_block,
|
||||
serial_input_varname,
|
||||
new_varname,
|
||||
)
|
||||
self._serial2dist_varname_mapping[ref_block.idx][
|
||||
serial_input_varname
|
||||
] = new_varname
|
||||
|
||||
# partition output vars
|
||||
for serial_output_varname in op.desc.output_arg_names():
|
||||
if (
|
||||
serial_output_varname
|
||||
not in self._serial2dist_varname_mapping[
|
||||
ref_block.forward_block_idx
|
||||
]
|
||||
or serial_output_varname
|
||||
not in self._serial2dist_varname_mapping[ref_block.idx]
|
||||
):
|
||||
new_varname = (
|
||||
serial_output_varname + self._dist_varname_suffix
|
||||
)
|
||||
if ref_block.has_var(serial_output_varname):
|
||||
_partition_var(
|
||||
self._dist_context,
|
||||
ref_block,
|
||||
target_block,
|
||||
serial_output_varname,
|
||||
new_varname,
|
||||
)
|
||||
self._serial2dist_varname_mapping[ref_block.idx][
|
||||
serial_output_varname
|
||||
] = new_varname
|
||||
|
||||
# partition op
|
||||
if is_forward_op(op) or op_dist_attr.is_recompute:
|
||||
kinputs, koutputs = dist_op_context.prepare_context(op)
|
||||
dist_op_forward_impl = _get_dist_op_forward_implement(
|
||||
op, self._dist_context
|
||||
)
|
||||
dist_op_forward_impl.forward(
|
||||
self._dist_context, **kinputs, **koutputs
|
||||
)
|
||||
|
||||
elif is_backward_op(op):
|
||||
kinputs, koutputs = dist_op_context.prepare_context(op)
|
||||
dist_op_backward_impl = _get_dist_op_backward_implement(
|
||||
op, self._dist_context, self._forward_op_id2forward_op
|
||||
)
|
||||
grad_var_to_var = (
|
||||
self._dist_context.dist_op_context.grad_var_to_var[
|
||||
appended_grad_times
|
||||
]
|
||||
)
|
||||
dist_op_backward_impl.backward(
|
||||
self._dist_context,
|
||||
**kinputs,
|
||||
**koutputs,
|
||||
**{"grad_var_to_var": grad_var_to_var},
|
||||
)
|
||||
elif is_optimize_op(op):
|
||||
# NOTE: BACKWARD_ONLY_DIST_OPS's op_role must be 2 because of 1F1B PASS
|
||||
kinputs, koutputs = dist_op_context.prepare_context(op)
|
||||
dist_op_opt_impl = _get_dist_op_backward_implement(
|
||||
op, self._dist_context, self._forward_op_id2forward_op
|
||||
)
|
||||
dist_op_opt_impl.backward(
|
||||
self._dist_context,
|
||||
**kinputs,
|
||||
**koutputs,
|
||||
**{"grad_var_to_var": {}},
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"partitioner only support forward and backward, optimize ops, but got {op}"
|
||||
)
|
||||
|
||||
def _is_valid_annotated_program(self, program):
|
||||
# TODO (ZJ-LIANG) should check all block
|
||||
ops = program.global_block().ops
|
||||
vars_ = program.list_vars()
|
||||
op_dist_attrs = [
|
||||
self._dist_context.get_op_dist_attr_for_program(op) for op in ops
|
||||
]
|
||||
var_dist_attrs = [
|
||||
self._dist_context.get_tensor_dist_attr_for_program(var)
|
||||
for var in vars_
|
||||
if (var.type not in __no_shape_var_type__)
|
||||
]
|
||||
|
||||
all_ops_annotated = all(
|
||||
dist_attr is not None for dist_attr in op_dist_attrs
|
||||
)
|
||||
all_vars_annotated = all(
|
||||
dist_attr is not None for dist_attr in var_dist_attrs
|
||||
)
|
||||
|
||||
return all_ops_annotated and all_vars_annotated
|
||||
|
||||
def _get_dist_var_by_serial_var(
|
||||
self, serial_var, partitioned_main_prog, block_id
|
||||
):
|
||||
block_idx = serial_var.block.idx
|
||||
target_block = partitioned_main_prog.blocks[block_idx]
|
||||
dist_var_name = self._serial2dist_varname_mapping[block_id][
|
||||
serial_var.name
|
||||
]
|
||||
assert target_block.has_var(dist_var_name)
|
||||
return target_block.var(dist_var_name)
|
||||
|
||||
|
||||
def _get_dist_shape(var, dist_attr):
|
||||
var_shape = var.shape
|
||||
mapping = dist_attr.dims_mapping
|
||||
mesh = dist_attr.process_mesh.shape
|
||||
if mapping == []:
|
||||
return var_shape
|
||||
|
||||
assert len(var_shape) == len(mapping), (
|
||||
f"variable shape [{var_shape}] and dim_mapping [{mapping}] is NOT match !"
|
||||
)
|
||||
new_shape = []
|
||||
for idx in range(len(var_shape)):
|
||||
if var_shape[idx] == -1 or mapping[idx] == -1:
|
||||
new_shape.append(var_shape[idx])
|
||||
else:
|
||||
assert var_shape[idx] % mesh[mapping[idx]] == 0, (
|
||||
f"un-event partition: var_shape[idx]=[{var_shape[idx]}], mesh[{mesh[mapping[idx]]}], {var.name}, {var_shape}, {mesh}, {mapping}"
|
||||
)
|
||||
new_shape.append(var_shape[idx] // mesh[mapping[idx]])
|
||||
|
||||
return new_shape
|
||||
|
||||
|
||||
def _partition_parameter(
|
||||
dist_context, src_var, dst_block, dst_varname, dst_shape
|
||||
):
|
||||
# NOTE hack to copied Parameter
|
||||
# not initialized parameter, need to initialize it
|
||||
copied_kwargs = {}
|
||||
copied_kwargs['trainable'] = src_var.trainable
|
||||
copied_kwargs['optimize_attr'] = src_var.optimize_attr
|
||||
copied_kwargs['regularizer'] = src_var.regularizer
|
||||
copied_kwargs['do_model_average'] = src_var.do_model_average
|
||||
copied_kwargs['need_clip'] = src_var.need_clip
|
||||
|
||||
param = Parameter(
|
||||
block=dst_block,
|
||||
type=src_var.type,
|
||||
name=dst_varname,
|
||||
shape=dst_shape,
|
||||
dtype=src_var.dtype,
|
||||
lod_level=src_var.lod_level,
|
||||
error_clip=src_var.error_clip,
|
||||
stop_gradient=src_var.stop_gradient,
|
||||
is_data=src_var.is_data,
|
||||
belong_to_optimizer=src_var.belong_to_optimizer,
|
||||
**copied_kwargs,
|
||||
)
|
||||
|
||||
return param
|
||||
|
||||
|
||||
def _partition_intermediate_var(
|
||||
dist_context, src_var, dst_block, dst_varname, dst_shape
|
||||
):
|
||||
var = dst_block.create_var(
|
||||
type=src_var.type,
|
||||
name=dst_varname,
|
||||
shape=dst_shape,
|
||||
dtype=src_var.dtype,
|
||||
lod_level=src_var.lod_level,
|
||||
persistable=src_var.persistable,
|
||||
error_clip=src_var.error_clip,
|
||||
stop_gradient=src_var.stop_gradient,
|
||||
is_data=src_var.is_data,
|
||||
belong_to_optimizer=src_var.belong_to_optimizer,
|
||||
)
|
||||
|
||||
return var
|
||||
|
||||
|
||||
def _partition_var(
|
||||
dist_context, src_block, dst_block, src_varname, dst_varname
|
||||
):
|
||||
"""
|
||||
partition include: split + replicate
|
||||
"""
|
||||
src_var = src_block.var(src_varname)
|
||||
|
||||
if src_var.type in __no_shape_var_type__:
|
||||
persist = getattr(src_var, 'persistable', False)
|
||||
new_var = dst_block.create_var(
|
||||
type=src_var.type,
|
||||
name=dst_varname,
|
||||
persistable=persist,
|
||||
stop_gradient=True,
|
||||
)
|
||||
target_shape = None
|
||||
else:
|
||||
dist_attr = dist_context.get_tensor_dist_attr_for_program(src_var)
|
||||
target_shape = _get_dist_shape(src_var, dist_attr)
|
||||
|
||||
if isinstance(src_var, Parameter):
|
||||
new_var = _partition_parameter(
|
||||
dist_context, src_var, dst_block, dst_varname, target_shape
|
||||
)
|
||||
else:
|
||||
new_var = _partition_intermediate_var(
|
||||
dist_context, src_var, dst_block, dst_varname, target_shape
|
||||
)
|
||||
|
||||
dist_attr = copy.deepcopy(
|
||||
dist_context.get_tensor_dist_attr_for_program(src_var)
|
||||
)
|
||||
assert dist_attr is not None
|
||||
dist_context.set_tensor_dist_attr_for_program(new_var, dist_attr)
|
||||
|
||||
return target_shape
|
||||
|
||||
|
||||
def _get_dist_op_backward_implement(
|
||||
backward_op, dist_context, forward_op_id2forward_op
|
||||
):
|
||||
dist_op_context = dist_context.dist_op_context
|
||||
if backward_op.desc.original_id() in dist_op_context.grad_op_id_to_op_id:
|
||||
forward_op_id = dist_op_context.grad_op_id_to_op_id[
|
||||
backward_op.desc.original_id()
|
||||
]
|
||||
forward_op = forward_op_id2forward_op[forward_op_id]
|
||||
forward_op_dist_attr = dist_context.get_op_dist_attr_for_program(
|
||||
forward_op
|
||||
)
|
||||
dist_op_impl_container = get_distributed_operator_impl_container(
|
||||
forward_op_dist_attr.impl_type
|
||||
)
|
||||
dist_op_impl = dist_op_impl_container.get_impl(
|
||||
forward_op_dist_attr.impl_idx
|
||||
)
|
||||
return dist_op_impl
|
||||
|
||||
# # NOTE trick for dist ops that only have backward implement
|
||||
if backward_op.type in BACKWARD_ONLY_DIST_OPS:
|
||||
op_dist_attr = dist_context.get_op_dist_attr_for_program(backward_op)
|
||||
assert op_dist_attr.impl_idx >= 0
|
||||
dist_op_impl = get_distributed_operator_impl_container(
|
||||
op_dist_attr.impl_type
|
||||
).get_impl(op_dist_attr.impl_idx)
|
||||
return dist_op_impl
|
||||
|
||||
dist_op = get_distributed_operator_impl_container("default")
|
||||
return dist_op.get_impl(0)
|
||||
|
||||
|
||||
def _get_dist_op_forward_implement(forward_op, dist_context):
|
||||
dist_attr = dist_context.get_op_dist_attr_for_program(forward_op)
|
||||
dist_op_impl_container = get_distributed_operator_impl_container(
|
||||
dist_attr.impl_type
|
||||
)
|
||||
dist_op_impl = dist_op_impl_container.get_impl(dist_attr.impl_idx)
|
||||
return dist_op_impl
|
||||
File diff suppressed because it is too large
Load Diff
+1112
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,180 @@
|
||||
# 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 logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
from paddle.distributed.auto_parallel.process_mesh import ProcessMesh
|
||||
from paddle.distributed.auto_parallel.static.dist_attribute import (
|
||||
OperatorDistAttr,
|
||||
TensorDistAttr,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.dist_op import DistributedOperator
|
||||
from paddle.distributed.auto_parallel.static.dist_tensor import (
|
||||
DistributedTensor,
|
||||
)
|
||||
|
||||
from ...utils.log_utils import get_logger
|
||||
from .completion import Completer
|
||||
from .dist_context import get_default_distributed_context
|
||||
from .tuner.parallel_tuner import ParallelTuner
|
||||
from .tuner.rule_based_tuner import RuleBasedTuner
|
||||
from .utils import is_naive_data_parallel
|
||||
|
||||
|
||||
class Planner:
|
||||
def __init__(self, mode, dist_context):
|
||||
self._mode = mode
|
||||
self._dist_context = dist_context
|
||||
self._load = False # load dist_attr from file
|
||||
|
||||
# NOTE: [HighOrderGrad]. There are grad ops in forward phase, and it need
|
||||
# dependency of backward-forward ops in forward completion.
|
||||
default_ctx = get_default_distributed_context()
|
||||
self._dist_context._dist_op_context = default_ctx.dist_op_context
|
||||
self._dist_context.data_parallel = default_ctx.data_parallel
|
||||
if not is_naive_data_parallel(self._dist_context):
|
||||
# Use SSA graph for complex parallelism
|
||||
self._dist_context.initialize(with_graph=True)
|
||||
else:
|
||||
# Use program for data parallel parallelism
|
||||
self._dist_context.initialize(with_graph=False)
|
||||
|
||||
self._completer = Completer(self._dist_context)
|
||||
|
||||
self._strategy = dist_context.strategy
|
||||
# set parallel tuner for auto search
|
||||
if self._strategy.auto_mode == "full_random":
|
||||
self._parallel_tuner = ParallelTuner(
|
||||
self._dist_context, mode=self._mode
|
||||
)
|
||||
elif self._strategy.auto_mode == "full_rule_based":
|
||||
self._parallel_tuner = RuleBasedTuner(
|
||||
self._dist_context, mode=self._mode
|
||||
)
|
||||
|
||||
@property
|
||||
def completer(self):
|
||||
return self._completer
|
||||
|
||||
def plan(self):
|
||||
logger = get_logger(logging.INFO)
|
||||
path = None
|
||||
if self._dist_context._json_config:
|
||||
try:
|
||||
path = self._dist_context._json_config["tuner_load_path"]
|
||||
except:
|
||||
path = None
|
||||
if path and os.path.exists(path):
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
from paddle.framework.restricted_unpickler import (
|
||||
safe_load_pickle,
|
||||
)
|
||||
|
||||
dist_attrs = safe_load_pickle(f)
|
||||
tensor_dist_attrs = dist_attrs["tensor"]
|
||||
op_dist_attrs = dist_attrs["op"]
|
||||
process_meshes = dist_attrs["process_meshes"]
|
||||
cluster = dist_attrs["cluster"]
|
||||
last_gpu_model = cluster.machines[0].devices[0].model
|
||||
last_gpu_memory = cluster.machines[0].devices[0].memory
|
||||
last_node_count = len(cluster.machines)
|
||||
last_device_count = len(cluster.get_all_devices("GPU"))
|
||||
|
||||
gpu_model = (
|
||||
self._dist_context.cluster.machines[0].devices[0].model
|
||||
)
|
||||
gpu_memory = (
|
||||
self._dist_context.cluster.machines[0].devices[0].memory
|
||||
)
|
||||
node_count = len(self._dist_context.cluster.machines)
|
||||
device_count = len(
|
||||
self._dist_context.cluster.get_all_devices("GPU")
|
||||
)
|
||||
if (
|
||||
gpu_model != last_gpu_model
|
||||
or gpu_memory != last_gpu_memory
|
||||
or last_node_count != node_count
|
||||
or device_count != last_device_count
|
||||
):
|
||||
logger.info(
|
||||
f"The cluster {node_count} nodes {device_count} {gpu_model} devices is different from the saved last cluster {last_node_count} nodes {last_device_count} {last_gpu_model} devices, so we run the planner again."
|
||||
)
|
||||
need_set_dist_attr = False
|
||||
else:
|
||||
need_set_dist_attr = True
|
||||
except:
|
||||
need_set_dist_attr = False
|
||||
|
||||
if need_set_dist_attr:
|
||||
for key in op_dist_attrs:
|
||||
serial_op = self._dist_context._dist_ops_for_program[
|
||||
key
|
||||
].serial_op
|
||||
# clear dist attr
|
||||
serial_op.dist_attr = OperatorDistAttr(serial_op.desc)
|
||||
serial_op.dist_attr.parse_from_string(op_dist_attrs[key])
|
||||
self._dist_context._dist_ops_for_program[key] = (
|
||||
DistributedOperator(serial_op)
|
||||
)
|
||||
|
||||
for key in tensor_dist_attrs:
|
||||
serial_tensor = (
|
||||
self._dist_context._dist_tensors_for_program[
|
||||
key
|
||||
].serial_tensor
|
||||
)
|
||||
# clear dist attr
|
||||
serial_tensor.dist_attr = TensorDistAttr(serial_tensor.desc)
|
||||
serial_tensor.dist_attr.parse_from_string(
|
||||
tensor_dist_attrs[key]
|
||||
)
|
||||
self._dist_context._dist_tensors_for_program[key] = (
|
||||
DistributedTensor(serial_tensor)
|
||||
)
|
||||
|
||||
process_meshes = []
|
||||
for item in dist_attrs["process_meshes"]:
|
||||
process_ids = item[0]
|
||||
shape = item[1]
|
||||
process_meshes.append(
|
||||
ProcessMesh(
|
||||
np.array(process_ids).reshape(shape).tolist()
|
||||
)
|
||||
)
|
||||
|
||||
self._dist_context.process_meshes = process_meshes
|
||||
self._load = True
|
||||
|
||||
logger.info(
|
||||
f"The parallel strategy has been loaded from {path}"
|
||||
)
|
||||
|
||||
if not self._load:
|
||||
if self._strategy.auto_mode != "semi":
|
||||
self._parallel_tuner.tune()
|
||||
else:
|
||||
self._completer.complete_forward_annotation()
|
||||
|
||||
if os.getenv("PADDLE_AUTO_PARALLEL_STAGE", "run") != "run":
|
||||
sys.exit()
|
||||
|
||||
# parse forward sub block
|
||||
self._dist_context.block_state.parse_forward_blocks(
|
||||
self._dist_context.serial_main_program
|
||||
)
|
||||
@@ -0,0 +1,275 @@
|
||||
# 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 hashlib
|
||||
from collections import OrderedDict
|
||||
|
||||
import paddle
|
||||
from paddle.framework import core
|
||||
|
||||
from ...collective import _get_global_env, _new_ring_id
|
||||
from ...utils.log_utils import get_logger
|
||||
from .utils import dygraph_guard
|
||||
|
||||
logger = get_logger("INFO", __name__)
|
||||
|
||||
|
||||
def get_all_process_groups():
|
||||
global _g_process_group_map
|
||||
return _g_process_group_map.values()
|
||||
|
||||
|
||||
def get_process_group(group_id, g_process_group_map=None):
|
||||
global _g_process_group_map
|
||||
return (
|
||||
_g_process_group_map.get(group_id, None)
|
||||
if g_process_group_map is None
|
||||
else g_process_group_map.get(group_id, None)
|
||||
)
|
||||
|
||||
|
||||
def get_world_process_group():
|
||||
global _g_process_group_map
|
||||
return _g_process_group_map[0]
|
||||
|
||||
|
||||
def clear_all_process_groups():
|
||||
global _g_process_group_map
|
||||
_g_process_group_map = {}
|
||||
_g_process_group_map[0] = ProcessGroup(0, [])
|
||||
|
||||
|
||||
def remove_process_group(ring_id):
|
||||
global _g_process_group_map
|
||||
if ring_id in _g_process_group_map:
|
||||
_g_process_group_map.pop(ring_id)
|
||||
|
||||
|
||||
def new_process_group(
|
||||
ranks, group_id=None, force_new_group=False, group_type=None
|
||||
):
|
||||
global _g_process_group_map
|
||||
if not force_new_group:
|
||||
# A key constructed from ranks is used for avoiding duplication
|
||||
new_key = '_'.join(map(str, ranks))
|
||||
for pg_id, pg in _g_process_group_map.items():
|
||||
cur_key = '_'.join(map(str, pg.ranks))
|
||||
if pg_id != 0 and new_key == cur_key:
|
||||
return pg
|
||||
# If not matching the existing one, construct a new process group
|
||||
num_groups = len(_g_process_group_map)
|
||||
# Note: our process group may interfere with the original implementation
|
||||
# so the created group id should start from the original _new_ring_id()
|
||||
if group_id is None:
|
||||
group_id = _new_ring_id() + num_groups + 1
|
||||
|
||||
new_pg = ProcessGroup(group_id, ranks, group_type)
|
||||
_g_process_group_map[group_id] = new_pg
|
||||
|
||||
return new_pg
|
||||
|
||||
|
||||
# This implementation refers to lots of Paddle/python/paddle/distributed/collective.py,
|
||||
# Fleet also has a collective helper which uses ops to initialize communication in
|
||||
# Paddle/python/paddle/distributed/fleet/meta_optimizers/common.py. We use the first one
|
||||
# because it seems simple. This should be enhanced to manage the process membership and
|
||||
# the instantiation process in a more general way. In the future, the process group may
|
||||
# handle the communication implementation choice.
|
||||
class ProcessGroup:
|
||||
def __init__(self, group_id, ranks, group_type=None):
|
||||
if group_id == 0 and get_process_group(0) is not None:
|
||||
assert group_id != 0, (
|
||||
"Process group id 0 is reserved for all ranks."
|
||||
)
|
||||
self._group_id = group_id
|
||||
self._ranks = ranks
|
||||
# Add the current ranks into group 0
|
||||
if group_id != 0:
|
||||
global _g_process_group_map
|
||||
_g_process_group_map[0].add_ranks(ranks)
|
||||
self._is_instantiate = False
|
||||
self._group_type = group_type
|
||||
|
||||
@property
|
||||
def id(self):
|
||||
return self._group_id
|
||||
|
||||
@property
|
||||
def ranks(self):
|
||||
return self._ranks
|
||||
|
||||
@property
|
||||
def nranks(self):
|
||||
return len(self._ranks)
|
||||
|
||||
@property
|
||||
def group_type(self):
|
||||
return self._group_type
|
||||
|
||||
def add_ranks(self, new_ranks):
|
||||
if set(new_ranks) <= set(self.ranks):
|
||||
return
|
||||
else:
|
||||
assert not self.is_instantiate(), (
|
||||
"Cannot add new ranks after instantiating the process group"
|
||||
)
|
||||
self._ranks.extend(new_ranks)
|
||||
self._ranks = list(set(self.ranks))
|
||||
|
||||
def local_rank(self, global_rank):
|
||||
if global_rank in self.ranks:
|
||||
return self.ranks.index(global_rank)
|
||||
else:
|
||||
raise AssertionError(
|
||||
f"Rank {global_rank} doesn't belong to this group"
|
||||
)
|
||||
|
||||
def is_instantiate(self):
|
||||
return self._is_instantiate
|
||||
|
||||
@dygraph_guard
|
||||
def instantiate(self):
|
||||
if self._is_instantiate:
|
||||
return
|
||||
ring_id = self.id
|
||||
genv = _get_global_env()
|
||||
global_rank = genv.rank
|
||||
|
||||
if self.nranks >= 2 and global_rank in self.ranks:
|
||||
logger.info(
|
||||
f"group_id: {self.id}, ranks: {self.ranks}, nranks: {self.nranks}, trainer_endpoints: {genv.current_endpoint}"
|
||||
)
|
||||
strategy = core.ParallelStrategy()
|
||||
strategy.nranks = self.nranks
|
||||
strategy.local_rank = self.local_rank(global_rank)
|
||||
strategy.trainer_endpoints = [
|
||||
genv.trainer_endpoints[i] for i in self.ranks
|
||||
]
|
||||
strategy.current_endpoint = genv.current_endpoint
|
||||
strategy.nrings = 1
|
||||
if core.is_compiled_with_cuda():
|
||||
place = core.CUDAPlace(genv.device_id)
|
||||
store = core.create_or_get_global_tcp_store()
|
||||
endpoints_str = ""
|
||||
for endpoint in strategy.trainer_endpoints:
|
||||
endpoints_str += endpoint
|
||||
endpoints_str += f"ring_id:{ring_id}"
|
||||
endpoints_str_hash = hashlib.md5(
|
||||
endpoints_str.encode(encoding='UTF-8')
|
||||
).hexdigest()
|
||||
|
||||
core.CommContextManager.set_device_id(genv.device_id)
|
||||
core.CommContextManager.create_nccl_comm_context(
|
||||
store,
|
||||
str(ring_id),
|
||||
strategy.local_rank,
|
||||
strategy.nranks,
|
||||
endpoints_str_hash,
|
||||
)
|
||||
elif core.is_compiled_with_xpu():
|
||||
place = core.XPUPlace(genv.device_id)
|
||||
store = core.create_or_get_global_tcp_store()
|
||||
endpoints_str = ""
|
||||
for endpoint in strategy.trainer_endpoints:
|
||||
endpoints_str += endpoint
|
||||
endpoints_str += f"ring_id:{ring_id}"
|
||||
endpoints_str_hash = hashlib.md5(
|
||||
endpoints_str.encode(encoding='UTF-8')
|
||||
).hexdigest()
|
||||
|
||||
core.CommContextManager.set_device_id(genv.device_id)
|
||||
core.CommContextManager.create_bkcl_comm_context(
|
||||
store,
|
||||
str(ring_id),
|
||||
strategy.local_rank,
|
||||
strategy.nranks,
|
||||
endpoints_str_hash,
|
||||
)
|
||||
elif genv.device_type in core.get_all_custom_device_type():
|
||||
place = core.CustomPlace(genv.device_type, genv.device_id)
|
||||
core.XCCLParallelContext(strategy, place).init_with_ring_id(
|
||||
ring_id
|
||||
)
|
||||
else:
|
||||
raise AssertionError('No CUDA device found')
|
||||
|
||||
if core.is_compiled_with_cuda():
|
||||
paddle.set_device(
|
||||
f'gpu:{paddle.distributed.ParallelEnv().dev_id}'
|
||||
)
|
||||
elif core.is_compiled_with_xpu():
|
||||
paddle.set_device(
|
||||
f'xpu:{paddle.distributed.ParallelEnv().dev_id}'
|
||||
)
|
||||
elif genv.device_type in core.get_all_custom_device_type():
|
||||
paddle.set_device(
|
||||
f'{paddle.distributed.ParallelEnv().device_type!s}:{paddle.distributed.ParallelEnv().dev_id}'
|
||||
)
|
||||
|
||||
# TODO(shenliang03): This is a temporary solution to solve the problem of
|
||||
# hang caused by cross-creation of new_group
|
||||
barrier_tensor = paddle.full([1], 1, dtype="int32")
|
||||
# barrier is not available in xpu for now
|
||||
if not paddle.framework.core.is_compiled_with_xpu():
|
||||
paddle._legacy_C_ops.barrier(
|
||||
barrier_tensor, barrier_tensor, 'ring_id', ring_id
|
||||
)
|
||||
|
||||
# NOTE(zhiqiu): to avoid send/recv hang in lazy init
|
||||
if self._group_type == 'p2p':
|
||||
alltoall_tmp = paddle.empty(
|
||||
shape=[self.nranks, self.nranks], dtype="int32"
|
||||
)
|
||||
paddle._legacy_C_ops.all_to_all(
|
||||
alltoall_tmp, 'use_calc_stream', True, 'ring_id', ring_id
|
||||
)
|
||||
paddle.device.cuda.synchronize()
|
||||
|
||||
if self.nranks > 1:
|
||||
barrier_tensor = paddle.full([1], 1, dtype="int32")
|
||||
# barrier is not available in xpu for now
|
||||
if not paddle.framework.core.is_compiled_with_xpu():
|
||||
paddle._legacy_C_ops.barrier(
|
||||
barrier_tensor, barrier_tensor, 'ring_id', 0
|
||||
)
|
||||
|
||||
self._is_instantiate = True
|
||||
|
||||
def is_member(self):
|
||||
return True
|
||||
|
||||
def __eq__(self, other):
|
||||
if not isinstance(other, ProcessGroup):
|
||||
return False
|
||||
if self.id != other.id:
|
||||
return False
|
||||
return True
|
||||
|
||||
def __ne__(self, other):
|
||||
return not self.__eq__(other)
|
||||
|
||||
def __str__(self):
|
||||
string = "id: {}, nranks: {}, ranks: {}.".format(
|
||||
self.id, self.nranks, ", ".join(map(str, self.ranks))
|
||||
)
|
||||
return string
|
||||
|
||||
def __hash__(self):
|
||||
return hash(self.__str__())
|
||||
|
||||
|
||||
# Note that Process group 0 is reserved for representing all ranks.
|
||||
# At the beginning, group 0 is empty and new ranks will be added automatically.
|
||||
_g_process_group_map = OrderedDict()
|
||||
_g_process_group_map[0] = ProcessGroup(0, [])
|
||||
@@ -0,0 +1,145 @@
|
||||
# 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 numpy as np
|
||||
|
||||
from paddle.framework import core
|
||||
|
||||
|
||||
class ProcessMesh(core.ProcessMesh):
|
||||
r"""
|
||||
The class `ProcessMesh` describes the topology of logical processes.
|
||||
|
||||
Args:
|
||||
mesh (list|numpy.array): an N-dimensional array describes the topology
|
||||
of logical processes.
|
||||
dim_names (list, optional): the i-th element of this list gives the name of the
|
||||
i-th dimension.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> import paddle.distributed as dist
|
||||
>>> paddle.enable_static()
|
||||
|
||||
>>> mesh = dist.ProcessMesh([[2, 4, 5], [0, 1, 3]])
|
||||
>>> assert mesh.shape == [2, 3]
|
||||
>>> assert mesh.process_ids == [2, 4, 5, 0, 1, 3]
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, mesh, dim_names=None):
|
||||
if not isinstance(mesh, list) and not isinstance(mesh, np.ndarray):
|
||||
raise ValueError(
|
||||
'The mesh must be an instance of list or np.ndarray.'
|
||||
)
|
||||
if isinstance(mesh, list):
|
||||
mesh = np.array(mesh)
|
||||
|
||||
self._mesh = mesh
|
||||
|
||||
self._shape = list(self._mesh.shape)
|
||||
|
||||
self._process_ids = self._mesh.flatten().tolist()
|
||||
if not all(isinstance(p, int) for p in self._process_ids):
|
||||
raise ValueError("All elements of the mesh must be integer")
|
||||
|
||||
if min(self._process_ids) < 0:
|
||||
raise ValueError('All elements of the mesh must be >= 0.')
|
||||
|
||||
unique_process_ids = set(self._process_ids)
|
||||
if len(unique_process_ids) != len(self._process_ids):
|
||||
raise ValueError('All elements of the mesh must be unique.')
|
||||
|
||||
if dim_names is not None:
|
||||
assert len(dim_names) == len(self._shape), (
|
||||
"The length of dims_names must be same as the shape of the mesh."
|
||||
)
|
||||
self._dim_names = dim_names
|
||||
else:
|
||||
self._dim_names = ["d" + str(i) for i in range(len(self._shape))]
|
||||
|
||||
# Follow the requirement for using pybind11
|
||||
core.ProcessMesh.__init__(
|
||||
self, self._shape, self._process_ids, self._dim_names
|
||||
)
|
||||
|
||||
@property
|
||||
def mesh(self):
|
||||
return self._mesh
|
||||
|
||||
|
||||
def compute_compatible_process_mesh(process_meshes):
|
||||
"""Compute the compatible process mesh given a list of process meshes."""
|
||||
if not process_meshes:
|
||||
return None
|
||||
|
||||
def _compute_compatible_of_two_process_meshes(pm1, pm2):
|
||||
if pm1 is None:
|
||||
return True, pm2
|
||||
if pm2 is None:
|
||||
return True, pm1
|
||||
if pm1 == pm2:
|
||||
return True, pm1
|
||||
if pm1.process_ids == pm2.process_ids:
|
||||
if len(pm1.shape) >= len(pm2.shape):
|
||||
return True, pm1
|
||||
else:
|
||||
return True, pm2
|
||||
process_set1 = set(pm1.process_ids)
|
||||
process_set2 = set(pm2.process_ids)
|
||||
if process_set1.issubset(process_set2):
|
||||
return True, pm2
|
||||
if process_set2.issubset(process_set1):
|
||||
return True, pm1
|
||||
return False, None
|
||||
|
||||
compatible_result = None
|
||||
for process_mesh in process_meshes:
|
||||
(
|
||||
compatible,
|
||||
compatible_result,
|
||||
) = _compute_compatible_of_two_process_meshes(
|
||||
compatible_result, process_mesh
|
||||
)
|
||||
if not compatible:
|
||||
return None
|
||||
if compatible_result.empty():
|
||||
return None
|
||||
if isinstance(compatible_result, core.ProcessMesh):
|
||||
mesh = np.array(compatible_result.process_ids).reshape(
|
||||
compatible_result.shape
|
||||
)
|
||||
return ProcessMesh(mesh, compatible_result.dim_names)
|
||||
elif isinstance(compatible_result, ProcessMesh):
|
||||
return ProcessMesh(compatible_result.mesh, compatible_result.dim_names)
|
||||
else:
|
||||
raise ValueError("Unrecognized ProcessMesh.")
|
||||
|
||||
|
||||
def merge_process_mesh(process_meshes):
|
||||
"""Merge a list of process meshes."""
|
||||
merged_process_mesh = None
|
||||
merged_process_ids = set()
|
||||
for process_mesh in process_meshes:
|
||||
if process_mesh is not None:
|
||||
process_ids = set(process_mesh.process_ids)
|
||||
merged_process_ids = merged_process_ids.union(process_ids)
|
||||
if len(merged_process_ids) != 0:
|
||||
merged_process_mesh = ProcessMesh(list(merged_process_ids))
|
||||
return merged_process_mesh
|
||||
@@ -0,0 +1,258 @@
|
||||
# Copyright (c) 2023 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 json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from argparse import ArgumentParser
|
||||
|
||||
import paddle
|
||||
from paddle.base.log_helper import get_logger
|
||||
|
||||
_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
|
||||
color_map = {
|
||||
"forward": "thread_state_running", # RGB: 126, 200, 148
|
||||
"backward": "rail_idle", # RGB: 238, 142, 0
|
||||
"optimizer": "rail_response", # RGB: 238, 142, 0
|
||||
"default": "thread_state_unknown", # RGB: 199, 155, 125
|
||||
}
|
||||
|
||||
ignore_job_type = ["recv_forward", "send_backward"]
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = ArgumentParser()
|
||||
device_count = paddle.device.cuda.device_count()
|
||||
all_devices = ",".join([str(i) for i in range(device_count)])
|
||||
parser.add_argument("--devices", type=str, default=all_devices)
|
||||
parser.add_argument("--log_dir", type=str, required=True)
|
||||
parser.add_argument("--multi_machine", action="store_true")
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def process_job_log(log_data, device_id, multi_machine_idx=-1):
|
||||
log_pattern = r'.*?Profiler Info: Job \((\d+)\), type = (\w+), micro_batch_id = (\d+), job_start_time = (\d+.\d+), job_end_time = (\d+.\d+)'
|
||||
matches = re.findall(log_pattern, log_data)
|
||||
events = []
|
||||
last_end_time = None
|
||||
|
||||
step_times = []
|
||||
step_start_time = 0
|
||||
step_end_time = 0
|
||||
|
||||
start_job_type = ""
|
||||
|
||||
for i, match in enumerate(matches):
|
||||
job_id, job_type, micro_batch_id, job_start_time, job_end_time = match
|
||||
|
||||
if job_type in ignore_job_type:
|
||||
continue
|
||||
|
||||
if job_type != "default" and start_job_type == "":
|
||||
start_job_type = job_type
|
||||
|
||||
start_time = float(job_start_time.strip()) * 1000
|
||||
end_time = float(job_end_time.strip()) * 1000
|
||||
|
||||
is_start_time_recorded = 0
|
||||
|
||||
if job_type == start_job_type and micro_batch_id == "0":
|
||||
if step_start_time != 0:
|
||||
step_times.append([step_start_time, step_end_time])
|
||||
step_start_time = start_time
|
||||
|
||||
step_end_time = end_time
|
||||
|
||||
tid_name = (
|
||||
"GPU" + str(device_id)
|
||||
if multi_machine_idx == -1
|
||||
else "GPU"
|
||||
+ str(device_id)
|
||||
+ "(machine:"
|
||||
+ str(multi_machine_idx)
|
||||
+ ")"
|
||||
)
|
||||
event_start = {
|
||||
"name": job_type + "_" + str(job_id),
|
||||
"cat": job_type,
|
||||
"ph": "B",
|
||||
"ts": start_time,
|
||||
"pid": 0,
|
||||
"tid": tid_name,
|
||||
}
|
||||
event_end = {
|
||||
"name": job_type + "_" + str(job_id),
|
||||
"cat": job_type,
|
||||
"ph": "E",
|
||||
"pid": 0,
|
||||
"ts": end_time,
|
||||
"tid": tid_name,
|
||||
}
|
||||
if job_type in color_map:
|
||||
event_start["cname"] = color_map[job_type]
|
||||
event_end["cname"] = color_map[job_type]
|
||||
|
||||
events.append(event_start)
|
||||
events.append(event_end)
|
||||
|
||||
last_end_time = end_time
|
||||
|
||||
step_times.append([step_start_time, step_end_time])
|
||||
return events, step_times
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
all_events = []
|
||||
step_infos = []
|
||||
start_step = 0
|
||||
machine_num = 1
|
||||
|
||||
def process_one_machine_log(log_dir, multi_machine_idx=-1):
|
||||
for device_id in args.devices.split(","):
|
||||
_logger.info(f"Process device {device_id}")
|
||||
device_id = int(device_id)
|
||||
log_file = os.path.join(log_dir, "workerlog." + str(device_id))
|
||||
with open(log_file, "r") as f:
|
||||
log_data = f.read()
|
||||
|
||||
start_step_pattern = (
|
||||
r'.*?Schedule Profiler start at step (\d+) and end at step.*'
|
||||
)
|
||||
start_step_match = re.findall(start_step_pattern, log_data)
|
||||
start_step = (
|
||||
int(start_step_match[0]) if len(start_step_match) > 0 else 0
|
||||
)
|
||||
|
||||
events, step_times = process_job_log(
|
||||
log_data, device_id, multi_machine_idx
|
||||
)
|
||||
all_events.extend(events)
|
||||
for i, info in enumerate(step_times):
|
||||
if len(step_infos) <= i:
|
||||
step_infos.append([float("inf"), float("-inf")])
|
||||
step_infos[i][0] = min(step_infos[i][0], info[0])
|
||||
step_infos[i][1] = max(step_infos[i][1], info[1])
|
||||
return start_step
|
||||
|
||||
if args.multi_machine:
|
||||
multi_machine_dirs = os.listdir(args.log_dir)
|
||||
multi_machine_dirs = [
|
||||
os.path.join(args.log_dir, d)
|
||||
for d in multi_machine_dirs
|
||||
if d.startswith("machine")
|
||||
and os.path.isdir(os.path.join(args.log_dir, d))
|
||||
]
|
||||
machine_num = len(multi_machine_dirs)
|
||||
for i, d in enumerate(multi_machine_dirs):
|
||||
_logger.info(f"Process machine {i}")
|
||||
start_step = max(process_one_machine_log(d, i), start_step)
|
||||
else:
|
||||
start_step = process_one_machine_log(args.log_dir)
|
||||
|
||||
for i, info in enumerate(step_infos):
|
||||
start_time = info[0]
|
||||
if i > 0:
|
||||
start_time = max(start_time, step_infos[i - 1][1])
|
||||
event_start = {
|
||||
"name": "step" + str(i + start_step),
|
||||
"cat": "step",
|
||||
"ph": "B",
|
||||
"ts": start_time,
|
||||
"pid": 0,
|
||||
"tid": "Step",
|
||||
"cname": color_map["default"],
|
||||
}
|
||||
event_end = {
|
||||
"name": "step" + str(i + start_step),
|
||||
"cat": "step",
|
||||
"ph": "E",
|
||||
"ts": info[1],
|
||||
"pid": 0,
|
||||
"tid": "Step",
|
||||
"cname": color_map["default"],
|
||||
}
|
||||
|
||||
all_events.append(event_start)
|
||||
all_events.append(event_end)
|
||||
|
||||
save_path = os.path.join(args.log_dir, "pipeline_profile.json")
|
||||
with open(save_path, "w") as f:
|
||||
f.write(json.dumps({"traceEvents": all_events}))
|
||||
_logger.info(f"Save pipeline profile to {save_path}")
|
||||
|
||||
# support Perfetto format
|
||||
save_path = os.path.join(args.log_dir, "pipeline_profile_perfetto.json")
|
||||
all_events.extend(
|
||||
[
|
||||
{
|
||||
"args": {"name": "STEP"},
|
||||
"cat": "__metadata",
|
||||
"name": "thread_name",
|
||||
"ph": "M",
|
||||
"pid": 0,
|
||||
"tid": 2333,
|
||||
"ts": 0,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
for i in range(machine_num):
|
||||
for j in range(len(args.devices.split(","))):
|
||||
if machine_num > 1:
|
||||
name = f"GPU:{j}(machine:{i})"
|
||||
tid = i * len(args.devices.split(",")) + j + 2334
|
||||
else:
|
||||
name = f"GPU:{j}"
|
||||
tid = j + 2334
|
||||
all_events.extend(
|
||||
[
|
||||
{
|
||||
"args": {"name": name},
|
||||
"cat": "__metadata",
|
||||
"name": "thread_name",
|
||||
"ph": "M",
|
||||
"pid": 0,
|
||||
"tid": tid,
|
||||
"ts": 0,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
json_str = json.dumps({"traceEvents": all_events})
|
||||
json_str = json_str.replace('"Step"', '2333')
|
||||
|
||||
for i in range(machine_num):
|
||||
for j in range(len(args.devices.split(","))):
|
||||
if machine_num > 1:
|
||||
json_str = json_str.replace(
|
||||
f'"GPU{j}(machine:{i})"',
|
||||
f'{i * len(args.devices.split(",")) + j + 2334}',
|
||||
)
|
||||
else:
|
||||
json_str = json_str.replace(f'"GPU{j}"', f'{j + 2334}')
|
||||
|
||||
with open(save_path, "w") as f:
|
||||
f.write(json_str)
|
||||
_logger.info(f"Save pipeline profile to {save_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,131 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
|
||||
# all registered reshard functions
|
||||
_g_reshard_func_list = []
|
||||
|
||||
|
||||
class ReshardFunction:
|
||||
def is_suitable(self, dist_tensor, dist_attr):
|
||||
raise NotImplementedError
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def choose_reshard_func(src_dist_attr, dst_dist_attr):
|
||||
global _g_reshard_func_list
|
||||
for reshard_func in _g_reshard_func_list:
|
||||
if reshard_func.is_suitable(src_dist_attr, dst_dist_attr):
|
||||
return reshard_func
|
||||
return None
|
||||
|
||||
|
||||
def register_reshard_func(reshard_func):
|
||||
global _g_reshard_func_list
|
||||
_g_reshard_func_list.append(reshard_func)
|
||||
|
||||
|
||||
def clean_reshard_funcs():
|
||||
global _g_reshard_func_list
|
||||
_g_reshard_func_list.clear()
|
||||
|
||||
|
||||
def is_shard(dist_attr):
|
||||
for v in dist_attr.dims_mapping:
|
||||
if v != -1:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_partial(dist_attr):
|
||||
if len(dist_attr.partial_status) > 0:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_replicated(dist_attr):
|
||||
dims_mapping_set = set(dist_attr.dims_mapping)
|
||||
if len(dist_attr.partial_status) == 0 and (
|
||||
len(dims_mapping_set) == 0
|
||||
or (len(dims_mapping_set) == 1 and -1 in dims_mapping_set)
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def copy_dist_attr_with_new_member(
|
||||
src_dist_attr,
|
||||
new_process_mesh=None,
|
||||
new_dims_mapping=None,
|
||||
new_partial_status=None,
|
||||
):
|
||||
if new_process_mesh is None:
|
||||
new_process_mesh = src_dist_attr.process_mesh
|
||||
if new_dims_mapping is None:
|
||||
new_dims_mapping = src_dist_attr.dims_mapping
|
||||
if new_partial_status is None:
|
||||
new_partial_status = src_dist_attr.partial_status
|
||||
|
||||
return paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
new_process_mesh,
|
||||
new_dims_mapping,
|
||||
new_partial_status,
|
||||
)
|
||||
|
||||
|
||||
def copy_op_attr_with_new_member(
|
||||
src_dist_attr,
|
||||
new_process_mesh=None,
|
||||
new_operands=None,
|
||||
new_results=None,
|
||||
new_chunk_id=None,
|
||||
):
|
||||
if new_process_mesh is None:
|
||||
new_process_mesh = src_dist_attr.process_mesh
|
||||
if new_operands is None:
|
||||
new_operands = src_dist_attr.operands()
|
||||
if new_results is None:
|
||||
new_results = src_dist_attr.results()
|
||||
if new_chunk_id is None:
|
||||
new_chunk_id = src_dist_attr.chunk_id
|
||||
|
||||
return paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
new_process_mesh,
|
||||
new_operands,
|
||||
new_results,
|
||||
new_chunk_id,
|
||||
)
|
||||
|
||||
|
||||
def copy_process_mesh_with_new_member(
|
||||
src_process_mesh,
|
||||
new_shape=None,
|
||||
new_process_ids=None,
|
||||
new_dim_names=None,
|
||||
):
|
||||
if new_shape is None:
|
||||
new_shape = src_process_mesh.shape
|
||||
if new_process_ids is None:
|
||||
new_process_ids = src_process_mesh.process_ids
|
||||
if new_dim_names is None:
|
||||
new_dim_names = src_process_mesh.dim_names
|
||||
|
||||
return paddle.base.libpaddle.pir.create_process_mesh(
|
||||
new_shape,
|
||||
new_process_ids,
|
||||
new_dim_names,
|
||||
)
|
||||
+94
@@ -0,0 +1,94 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
is_replicated,
|
||||
)
|
||||
from .nd_mesh_reshard_func import NdMeshReshardFunction
|
||||
|
||||
|
||||
class GlobalToSubMeshFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
# NOTE we could allow the src_dist_attr is not replicated and reshard it as replicated before go through the global_to_sub logic
|
||||
# but the dst_dist_attr should be replicated otherwise there will be un-defined result when change the mesh.
|
||||
if not is_replicated(dst_dist_attr):
|
||||
return False
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
if in_mesh.ndim > out_mesh.ndim + 1:
|
||||
return False
|
||||
if in_mesh.ndim == out_mesh.ndim:
|
||||
return set(out_mesh.process_ids) < set(in_mesh.process_ids)
|
||||
else:
|
||||
sub_meshes = paddle.base.libpaddle.pir.get_sub_meshes(in_mesh)
|
||||
return out_mesh in sub_meshes
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
# reshard operand as replicated before change the mesh.
|
||||
if not is_replicated(src_dist_attr):
|
||||
tmp_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[-1] * len(src_dist_attr.dims_mapping),
|
||||
{},
|
||||
)
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dist_attr
|
||||
)
|
||||
|
||||
pre_reshard_func = NdMeshReshardFunction()
|
||||
src_value = pre_reshard_func.reshard(
|
||||
src_dist_attr,
|
||||
tmp_dist_attr,
|
||||
src_value,
|
||||
tmp_dst_type,
|
||||
)
|
||||
src_dist_attr = tmp_dist_attr
|
||||
|
||||
if src_value.has_one_use():
|
||||
src_value.update_dist_attr(dst_dist_attr)
|
||||
prev_op = src_value.get_defining_op()
|
||||
op_dist_attr = prev_op.dist_attr
|
||||
op_mesh = op_dist_attr.process_mesh
|
||||
operands = op_dist_attr.operands()
|
||||
results = op_dist_attr.results()
|
||||
chunk_id = op_dist_attr.chunk_id
|
||||
results[src_value.index()] = dst_dist_attr
|
||||
prev_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
op_mesh, operands, results, chunk_id
|
||||
)
|
||||
)
|
||||
return src_value
|
||||
else:
|
||||
dst_value = paddle._C_ops.share_data_(src_value)
|
||||
share_data_op = dst_value.get_defining_op()
|
||||
# set dist type and dist attr
|
||||
dst_value.set_type(dst_type)
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
share_data_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[src_dist_attr],
|
||||
[dst_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
return dst_value
|
||||
@@ -0,0 +1,365 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.distributed.auto_parallel.static.utils import split_mesh
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
copy_dist_attr_with_new_member,
|
||||
is_partial,
|
||||
)
|
||||
from .p_to_r_reshard_func import PToRReshardFunction
|
||||
from .p_to_s_reshard_func import PToSReshardFunction
|
||||
from .r_to_p_reshard_func import RToPReshardFunction
|
||||
from .r_to_s_reshard_func import RToSReshardFunction
|
||||
from .s_to_r_reshard_func import SToRReshardFunction
|
||||
from .same_status_reshard_func import SameStatusReshardFunction
|
||||
|
||||
|
||||
def find_first_diff_shard_axis(src_dist_attr, dst_dist_attr):
|
||||
src_dims_mapping = src_dist_attr.dims_mapping
|
||||
dst_dims_mapping = dst_dist_attr.dims_mapping
|
||||
ndim = len(src_dims_mapping)
|
||||
for i in range(ndim - 1, -1, -1):
|
||||
if src_dims_mapping[i] != dst_dims_mapping[i]:
|
||||
return i
|
||||
return -1
|
||||
|
||||
|
||||
def get_1D_sub_process_mesh(process_mesh, mesh_dim):
|
||||
"""
|
||||
Get the 1-D sub process mesh on specific mesh_dim which:
|
||||
1) where the reshard should be performed.
|
||||
2) contains current process.
|
||||
|
||||
Args:
|
||||
process_mesh (ProcessMesh): the global process mesh.
|
||||
mesh_dim (int): the mesh dimension where the dist_tensor is
|
||||
sharded or partial.
|
||||
|
||||
e.g.
|
||||
1) process_mesh = [[0, 1, 2], [3, 4, 5]], axis = 0:
|
||||
process rank id returned sub mesh
|
||||
0 or 3 [0, 3]
|
||||
1 or 4 [1, 4]
|
||||
2 or 5 [2, 5]
|
||||
2) process_mesh = [[0, 1, 2], [3, 4, 5]], axis = 1:
|
||||
process rank id returned sub mesh
|
||||
0 or 1 or 2 [0, 1, 2]
|
||||
3 or 4 or 5 [3, 4, 5]
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
mesh_shape = process_mesh.shape
|
||||
dim_names = process_mesh.dim_names
|
||||
process_ids = np.array(process_mesh.process_ids).reshape(mesh_shape)
|
||||
|
||||
rank_id = dist.get_rank()
|
||||
# FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
|
||||
if rank_id not in process_mesh.process_ids:
|
||||
rank_id = process_mesh.process_ids[0]
|
||||
coord = list(np.where(process_ids == rank_id))
|
||||
coord[mesh_dim] = range(mesh_shape[mesh_dim])
|
||||
sub_process_ids = process_ids[tuple(coord)].flatten()
|
||||
sub_mesh_name = dim_names[mesh_dim]
|
||||
|
||||
return dist.ProcessMesh(sub_process_ids, [sub_mesh_name])
|
||||
|
||||
|
||||
class NdMeshReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
if out_mesh.ndim <= 1:
|
||||
return False
|
||||
|
||||
# check dims_mapping and partial_status
|
||||
if src_dist_attr == dst_dist_attr:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
"""
|
||||
Reshard on N-d mesh:
|
||||
1. Find the tensor dimensions where the dims_mapping values
|
||||
differ between src_dist_attr and dst_dist_attr.
|
||||
2. From higher to lower, convert the non-replicated dimensions
|
||||
in step1 to replicated using corresponding 1-D mesh functions.
|
||||
3. Convert the replicated dimensions in step2 to the status in
|
||||
dst_dist_attr with corresponding 1-D mesh functions.
|
||||
"""
|
||||
# Step1. find first dimension with different shard status in src_dist_attr
|
||||
# and dst_dist_attr.
|
||||
first_diff_axis = find_first_diff_shard_axis(
|
||||
src_dist_attr, dst_dist_attr
|
||||
)
|
||||
# out_value = src_value # intermediate result
|
||||
# src_type = src_value.type()
|
||||
tensor_ndim = len(src_value.shape)
|
||||
process_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
# Step2. Convert the non-replicated dimensions to replicated.
|
||||
|
||||
# Step2.1 convert shard status to replicated
|
||||
for i in range(first_diff_axis, -1, -1):
|
||||
in_mesh_axis = src_dist_attr.dims_mapping[i]
|
||||
out_mesh_axis = dst_dist_attr.dims_mapping[i]
|
||||
if in_mesh_axis == -1 or in_mesh_axis == out_mesh_axis:
|
||||
continue
|
||||
|
||||
# calculate the dist_attr after converting
|
||||
tmp_dims_mapping = src_dist_attr.dims_mapping
|
||||
tmp_dims_mapping[i] = -1
|
||||
tmp_dst_dist_attr = copy_dist_attr_with_new_member(
|
||||
src_dist_attr, new_dims_mapping=tmp_dims_mapping
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dst_dist_attr
|
||||
)
|
||||
sub_mesh_list = split_mesh(process_mesh, in_mesh_axis)
|
||||
for sub_mesh in sub_mesh_list:
|
||||
new_process_group(sorted(sub_mesh.process_ids))
|
||||
# get the process_mesh on specific axis
|
||||
sub_mesh = get_1D_sub_process_mesh(process_mesh, in_mesh_axis)
|
||||
|
||||
# calculate corresponding 1-D dist_attr of src_dst_attr
|
||||
in_one_dim_dims_mapping = [-1] * tensor_ndim
|
||||
in_one_dim_dims_mapping[i] = 0
|
||||
in_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh, in_one_dim_dims_mapping, {}
|
||||
)
|
||||
)
|
||||
|
||||
# calculate corresponding 1-D dist_attr of dst_dst_attr
|
||||
out_one_dim_dims_mapping = [-1] * tensor_ndim
|
||||
out_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh, out_one_dim_dims_mapping, {}
|
||||
)
|
||||
)
|
||||
|
||||
one_dim_func = SToRReshardFunction()
|
||||
src_value = one_dim_func.reshard(
|
||||
in_one_dim_dist_attr,
|
||||
out_one_dim_dist_attr,
|
||||
src_value,
|
||||
tmp_dst_type,
|
||||
)
|
||||
src_dist_attr = tmp_dst_dist_attr
|
||||
|
||||
# Step2.2. convert partial status to replicated
|
||||
if is_partial(src_dist_attr):
|
||||
in_partial_status = src_dist_attr.partial_status
|
||||
out_partial_status = dst_dist_attr.partial_status # read-only
|
||||
# convert each partial dim to replicated with corresponding
|
||||
# 1-D mesh function
|
||||
for partial_dim, partial_type in in_partial_status.items():
|
||||
if partial_dim in out_partial_status:
|
||||
if out_partial_status[partial_dim] != partial_type:
|
||||
raise NotImplementedError(
|
||||
f"Reshard tensor from one partial type {partial_type} to another partial type {out_partial_status[partial_dim]} is not supported yet."
|
||||
)
|
||||
continue
|
||||
|
||||
p_to_s = False
|
||||
if partial_dim in dst_dist_attr.dims_mapping:
|
||||
p_to_s = True
|
||||
shard_index = dst_dist_attr.dims_mapping.index(partial_dim)
|
||||
# get the partial status after converting
|
||||
tmp_partial_status = src_dist_attr.partial_status
|
||||
tmp_partial_status.pop(partial_dim)
|
||||
|
||||
tmp_dims_mapping = src_dist_attr.dims_mapping
|
||||
if p_to_s:
|
||||
tmp_dims_mapping[shard_index] = partial_dim
|
||||
|
||||
tmp_dst_dist_attr = copy_dist_attr_with_new_member(
|
||||
src_dist_attr,
|
||||
new_dims_mapping=tmp_dims_mapping,
|
||||
new_partial_status=tmp_partial_status,
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dst_dist_attr
|
||||
)
|
||||
sub_mesh_list = split_mesh(process_mesh, partial_dim)
|
||||
for sub_mesh in sub_mesh_list:
|
||||
new_process_group(sorted(sub_mesh.process_ids))
|
||||
# get the process_mesh on specific axis
|
||||
sub_mesh = get_1D_sub_process_mesh(process_mesh, partial_dim)
|
||||
|
||||
# calculate corresponding 1-D dist_attr of src_dst_attr
|
||||
in_one_dim_partial_status = {0: partial_type}
|
||||
in_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh,
|
||||
[-1] * tensor_ndim,
|
||||
in_one_dim_partial_status,
|
||||
)
|
||||
)
|
||||
out_one_dim_dims_mapping = [-1] * tensor_ndim
|
||||
one_dim_func = PToRReshardFunction()
|
||||
if p_to_s:
|
||||
out_one_dim_dims_mapping[shard_index] = 0
|
||||
one_dim_func = PToSReshardFunction()
|
||||
|
||||
# calculate corresponding 1-D dist_attr of dst_dst_attr
|
||||
out_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh,
|
||||
out_one_dim_dims_mapping,
|
||||
{},
|
||||
)
|
||||
)
|
||||
|
||||
src_value = one_dim_func.reshard(
|
||||
in_one_dim_dist_attr,
|
||||
out_one_dim_dist_attr,
|
||||
src_value,
|
||||
tmp_dst_type,
|
||||
)
|
||||
src_dist_attr = tmp_dst_dist_attr
|
||||
|
||||
# Step3. Convert the replicated status to the status in dst_dist_attr
|
||||
# Step3.1 convert replicated to partial
|
||||
if is_partial(dst_dist_attr):
|
||||
in_partial_status = src_dist_attr.partial_status
|
||||
out_partial_status = dst_dist_attr.partial_status
|
||||
for partial_dim, partial_type in out_partial_status.items():
|
||||
if partial_dim in in_partial_status:
|
||||
continue
|
||||
|
||||
sub_mesh = get_1D_sub_process_mesh(process_mesh, partial_dim)
|
||||
|
||||
in_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh,
|
||||
[-1] * tensor_ndim,
|
||||
{},
|
||||
)
|
||||
)
|
||||
out_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh, [-1] * tensor_ndim, {0: partial_type}
|
||||
)
|
||||
)
|
||||
|
||||
tmp_dst_dist_attr = copy_dist_attr_with_new_member(
|
||||
dst_dist_attr,
|
||||
new_partial_status={partial_dim: partial_type},
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dst_dist_attr
|
||||
)
|
||||
|
||||
src_value = RToPReshardFunction().reshard(
|
||||
in_one_dim_dist_attr,
|
||||
out_one_dim_dist_attr,
|
||||
src_value,
|
||||
tmp_dst_type,
|
||||
)
|
||||
src_dist_attr = tmp_dst_dist_attr
|
||||
|
||||
# Step3.2 convert replicated to shard
|
||||
for i in range(first_diff_axis, -1, -1):
|
||||
in_mesh_axis = src_dist_attr.dims_mapping[i]
|
||||
out_mesh_axis = dst_dist_attr.dims_mapping[i]
|
||||
if in_mesh_axis == out_mesh_axis:
|
||||
continue
|
||||
|
||||
# calculate the dist_attr after converting
|
||||
tmp_dims_mapping = src_dist_attr.dims_mapping
|
||||
tmp_dims_mapping[i] = out_mesh_axis
|
||||
tmp_dst_dist_attr = copy_dist_attr_with_new_member(
|
||||
src_dist_attr, new_dims_mapping=tmp_dims_mapping
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dst_dist_attr
|
||||
)
|
||||
|
||||
# get the process_mesh on specific axis
|
||||
sub_mesh = get_1D_sub_process_mesh(process_mesh, out_mesh_axis)
|
||||
|
||||
# calculate the corresponding 1-D input dist attr
|
||||
in_one_dim_dims_mapping = [-1] * tensor_ndim
|
||||
in_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh, in_one_dim_dims_mapping, {}
|
||||
)
|
||||
)
|
||||
|
||||
# calculate the corresponding 1-D output dist attr
|
||||
out_one_dim_dims_mapping = [-1] * tensor_ndim
|
||||
out_one_dim_dims_mapping[i] = 0
|
||||
out_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh, out_one_dim_dims_mapping, {}
|
||||
)
|
||||
)
|
||||
one_dim_func = RToSReshardFunction()
|
||||
src_value = one_dim_func.reshard(
|
||||
in_one_dim_dist_attr,
|
||||
out_one_dim_dist_attr,
|
||||
src_value,
|
||||
tmp_dst_type,
|
||||
)
|
||||
src_dist_attr = tmp_dst_dist_attr
|
||||
return src_value
|
||||
|
||||
|
||||
class NdMeshReshardFunctionCrossMesh(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh == out_mesh:
|
||||
return False
|
||||
if in_mesh.shape != out_mesh.shape:
|
||||
return False
|
||||
if out_mesh.ndim <= 1:
|
||||
return False
|
||||
if src_dist_attr == dst_dist_attr:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
same_status_func = SameStatusReshardFunction()
|
||||
tmp_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
src_dist_attr.dims_mapping,
|
||||
src_dist_attr.partial_status,
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dist_attr
|
||||
)
|
||||
src_value = same_status_func.reshard(
|
||||
src_dist_attr, tmp_dist_attr, src_value, tmp_dst_type
|
||||
)
|
||||
|
||||
nd_mesh_func = NdMeshReshardFunction()
|
||||
assert nd_mesh_func.is_suitable(tmp_dist_attr, dst_dist_attr), (
|
||||
f"Invoke the p to r reshard function is not valid from {tmp_dist_attr} to {dst_dist_attr}"
|
||||
)
|
||||
return nd_mesh_func.reshard(
|
||||
tmp_dist_attr, dst_dist_attr, src_value, dst_type
|
||||
)
|
||||
@@ -0,0 +1,113 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import ReshardFunction, is_partial, is_replicated
|
||||
from .same_status_reshard_func import SameStatusReshardFunction
|
||||
|
||||
|
||||
class PToRReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_partial(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_replicated(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
src_reduce_type = src_dist_attr.partial_status[0]
|
||||
# reduce_mean = False
|
||||
# if src_reduce_type == paddle.base.core.ReduceType.kRedAvg:
|
||||
# src_reduce_type = paddle.base.core.ReduceType.kRedSum
|
||||
# reduce_mean = True
|
||||
|
||||
group = new_process_group(sorted(src_mesh.process_ids))
|
||||
reduced_value = paddle._C_ops.all_reduce(
|
||||
src_value, group.id, int(src_reduce_type)
|
||||
)
|
||||
# set dist type and dist attr
|
||||
reduced_value.set_type(dst_type)
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
|
||||
reduced_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh,
|
||||
[src_dist_attr],
|
||||
[dst_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
return reduced_value
|
||||
|
||||
|
||||
class PToRReshardFunctionCrossMesh(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_partial(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_replicated(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if (
|
||||
in_mesh.ndim != 1
|
||||
or out_mesh.ndim != 1
|
||||
or in_mesh.shape != out_mesh.shape
|
||||
):
|
||||
return False
|
||||
|
||||
if in_mesh == out_mesh:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
same_status_func = SameStatusReshardFunction()
|
||||
tmp_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
src_dist_attr.dims_mapping,
|
||||
src_dist_attr.partial_status,
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dist_attr
|
||||
)
|
||||
src_value = same_status_func.reshard(
|
||||
src_dist_attr, tmp_dist_attr, src_value, tmp_dst_type
|
||||
)
|
||||
|
||||
p_to_r_func = PToRReshardFunction()
|
||||
assert p_to_r_func.is_suitable(tmp_dist_attr, dst_dist_attr), (
|
||||
f"Invoke the p to r reshard function is not valid from {tmp_dist_attr} to {dst_dist_attr}"
|
||||
)
|
||||
return p_to_r_func.reshard(
|
||||
tmp_dist_attr, dst_dist_attr, src_value, dst_type
|
||||
)
|
||||
@@ -0,0 +1,245 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.distributed.utils.stream_utils import ExecutionStreamType
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
copy_dist_attr_with_new_member,
|
||||
is_partial,
|
||||
is_shard,
|
||||
)
|
||||
|
||||
|
||||
class PToSReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_partial(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_shard(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
src_reduce_type = src_dist_attr.partial_status[0]
|
||||
assert src_reduce_type == paddle.base.core.ReduceType.kRedSum, (
|
||||
f"The p to s reshard func only support sum op, but received {src_reduce_type}"
|
||||
)
|
||||
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
|
||||
split_axis = dst_dist_attr.dims_mapping.index(0)
|
||||
num_of_process = len(src_dist_attr.process_mesh.process_ids)
|
||||
remainder_of_padding = src_value.shape[split_axis] % num_of_process
|
||||
is_balanced_split = remainder_of_padding == 0
|
||||
|
||||
permute = False
|
||||
if split_axis != 0:
|
||||
perm = list(range(0, len(src_value.shape)))
|
||||
perm[0] = split_axis
|
||||
perm[split_axis] = 0
|
||||
src_value = paddle._C_ops.transpose(src_value, perm)
|
||||
permute = True
|
||||
tmp_dims_mapping = dst_dist_attr.dims_mapping
|
||||
tmp_dims_mapping[split_axis] = -1
|
||||
tmp_dims_mapping[0] = 0
|
||||
dst_dist_attr = copy_dist_attr_with_new_member(
|
||||
dst_dist_attr, new_dims_mapping=tmp_dims_mapping
|
||||
)
|
||||
|
||||
if is_balanced_split:
|
||||
global_dst_attr = dst_type.as_dist_type().dist_attr()
|
||||
global_dims_mapping = global_dst_attr.dims_mapping
|
||||
axis = global_dims_mapping[0]
|
||||
global_dims_mapping[0] = global_dims_mapping[split_axis]
|
||||
global_dims_mapping[split_axis] = axis
|
||||
global_dist_attr = copy_dist_attr_with_new_member(
|
||||
global_dst_attr, new_dims_mapping=global_dims_mapping
|
||||
)
|
||||
dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), global_dist_attr
|
||||
)
|
||||
group = new_process_group(sorted(src_mesh.process_ids))
|
||||
dst_value = paddle._C_ops.reduce_scatter(
|
||||
src_value, group.id, num_of_process
|
||||
)
|
||||
dst_value.get_defining_op().set_execution_stream(
|
||||
ExecutionStreamType.DefaultStream.value
|
||||
)
|
||||
|
||||
# set dist type and dist attr
|
||||
dst_value.set_type(dst_type)
|
||||
dst_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh, [src_dist_attr], [dst_dist_attr], chunk_id
|
||||
)
|
||||
)
|
||||
|
||||
if split_axis != 0:
|
||||
dst_value = paddle._C_ops.transpose(dst_value, perm)
|
||||
return dst_value
|
||||
else:
|
||||
dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), dst_dist_attr
|
||||
)
|
||||
original_dims_mapping = dst_dist_attr.dims_mapping.copy()
|
||||
original_split_axis = split_axis
|
||||
split_axis = 0
|
||||
avg_size_on_split_axis = int(
|
||||
(src_value.shape[split_axis] + num_of_process - 1)
|
||||
/ num_of_process
|
||||
)
|
||||
padding_num = (
|
||||
avg_size_on_split_axis * num_of_process
|
||||
- src_value.shape[split_axis]
|
||||
)
|
||||
padding_shape = src_value._local_shape
|
||||
padding_shape[split_axis] = padding_num
|
||||
padding_tensor = paddle.full(
|
||||
padding_shape,
|
||||
0.0,
|
||||
src_value.dtype,
|
||||
)
|
||||
tmp_src_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
padding_tensor.type(), src_dist_attr
|
||||
)
|
||||
padding_tensor.set_type(tmp_src_type)
|
||||
padding_tensor.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh, [], [src_dist_attr], chunk_id
|
||||
)
|
||||
)
|
||||
concat_value = paddle._C_ops.concat(
|
||||
[src_value, padding_tensor], split_axis
|
||||
)
|
||||
axis_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[-1],
|
||||
{0: paddle.base.core.ReduceType.kRedSum},
|
||||
)
|
||||
)
|
||||
concat_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[
|
||||
paddle.base.libpaddle.pir.create_array_attribute(
|
||||
[src_dist_attr, src_dist_attr]
|
||||
),
|
||||
axis_dist_attr,
|
||||
],
|
||||
[src_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
|
||||
concat_global_shape = list(src_value.shape)
|
||||
concat_global_shape[split_axis] = (
|
||||
avg_size_on_split_axis * num_of_process
|
||||
)
|
||||
concat_type = paddle.pir.create_shaped_type(
|
||||
src_value.type(), concat_global_shape
|
||||
)
|
||||
concat_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
concat_type, src_dist_attr
|
||||
)
|
||||
concat_value.set_type(concat_type)
|
||||
|
||||
dst_value = self.reshard_p_to_s_with_padding(
|
||||
concat_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
padding_num,
|
||||
)
|
||||
if permute:
|
||||
dst_value = paddle._C_ops.transpose(dst_value, perm)
|
||||
split_axis = original_split_axis
|
||||
return dst_value
|
||||
|
||||
def reshard_p_to_s_with_padding(
|
||||
self,
|
||||
src_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
padding_num=0,
|
||||
):
|
||||
group = new_process_group(
|
||||
sorted(src_dist_attr.process_mesh.process_ids)
|
||||
)
|
||||
dst_value = paddle._C_ops.reduce_scatter(
|
||||
src_value, group.id, len(src_dist_attr.process_mesh.process_ids)
|
||||
)
|
||||
out_global_shape = dst_type.shape
|
||||
out_global_shape[split_axis] = (
|
||||
padding_num + out_global_shape[split_axis]
|
||||
)
|
||||
dst_tmp_type = paddle.pir.create_shaped_type(
|
||||
dst_value.type(), out_global_shape
|
||||
)
|
||||
dst_tmp_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
dst_tmp_type, dst_dist_attr
|
||||
)
|
||||
dst_value.set_type(dst_tmp_type)
|
||||
dst_value.get_defining_op().set_execution_stream(
|
||||
ExecutionStreamType.DefaultStream.value
|
||||
)
|
||||
dst_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[src_dist_attr],
|
||||
[dst_dist_attr],
|
||||
src_value.get_defining_op().dist_attr.chunk_id,
|
||||
)
|
||||
)
|
||||
if padding_num != 0:
|
||||
if dist.get_rank() == dst_dist_attr.process_mesh.process_ids[-1]:
|
||||
dst_value = paddle._C_ops.split(
|
||||
dst_value,
|
||||
[
|
||||
dst_value.shape[split_axis] - padding_num,
|
||||
padding_num,
|
||||
],
|
||||
0,
|
||||
)[0]
|
||||
dst_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
[dst_dist_attr],
|
||||
[dst_dist_attr],
|
||||
src_value.get_defining_op().dist_attr.chunk_id,
|
||||
)
|
||||
)
|
||||
else:
|
||||
dst_value.set_type(dst_type)
|
||||
return dst_value
|
||||
@@ -0,0 +1,73 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
is_partial,
|
||||
is_replicated,
|
||||
is_shard,
|
||||
)
|
||||
|
||||
|
||||
class RToPReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_replicated(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_partial(dst_dist_attr) or is_shard(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
dst_mesh = dst_dist_attr.process_mesh
|
||||
dst_reduce_type = dst_dist_attr.partial_status[0]
|
||||
local_rank = paddle.distributed.get_rank()
|
||||
|
||||
assert dst_reduce_type in [
|
||||
paddle.base.core.ReduceType.kRedSum,
|
||||
paddle.distributed.ReduceType.kRedAvg,
|
||||
paddle.distributed.ReduceType.kRedMax,
|
||||
], f"Unsupported reduce type {dst_reduce_type}"
|
||||
|
||||
if (
|
||||
dst_reduce_type == paddle.distributed.ReduceType.kRedSum
|
||||
and local_rank != 0
|
||||
):
|
||||
dst_value = paddle.full(src_value.shape, 0, dtype=src_value.dtype)
|
||||
else:
|
||||
dst_value = paddle.assign(src_value)
|
||||
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
dst_value.set_type(dst_type)
|
||||
dst_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
dst_mesh, [src_dist_attr], [dst_dist_attr], chunk_id
|
||||
)
|
||||
)
|
||||
|
||||
return dst_value
|
||||
@@ -0,0 +1,142 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
|
||||
from .base_reshard_func import ReshardFunction, is_replicated, is_shard
|
||||
from .same_status_reshard_func import SameStatusReshardFunction
|
||||
|
||||
|
||||
class RToSReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_replicated(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_shard(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
split_axis = -1
|
||||
mesh_axis = -1
|
||||
for idx, v in enumerate(dst_dist_attr.dims_mapping):
|
||||
if v != -1:
|
||||
split_axis = idx
|
||||
mesh_axis = v
|
||||
|
||||
mesh = src_dist_attr.process_mesh
|
||||
curr_global_rank = paddle.distributed.get_rank()
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
|
||||
if curr_global_rank in mesh.process_ids:
|
||||
total_nums = src_value.shape[split_axis]
|
||||
num_of_pieces = mesh.shape[mesh_axis]
|
||||
if num_of_pieces == 1:
|
||||
dst_value = paddle._C_ops.share_data_(src_value)
|
||||
share_data_op = dst_value.get_defining_op()
|
||||
# set dist type and dist attr
|
||||
dst_value.set_type(dst_type)
|
||||
share_data_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[src_dist_attr],
|
||||
[dst_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
return dst_value
|
||||
piece_len = (total_nums + num_of_pieces - 1) // num_of_pieces
|
||||
rank_relative = mesh.process_ids.index(curr_global_rank)
|
||||
start = rank_relative * piece_len
|
||||
end = start + piece_len
|
||||
if curr_global_rank == mesh.process_ids[-1]:
|
||||
end = total_nums
|
||||
|
||||
out_value = paddle.slice(src_value, [split_axis], [start], [end])
|
||||
|
||||
out_value.set_type(dst_type)
|
||||
out_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
mesh, [src_dist_attr], [dst_dist_attr], chunk_id
|
||||
)
|
||||
)
|
||||
return out_value
|
||||
# fake var will be removed in remove_other_rank_op_pass.
|
||||
fake_var = paddle._C_ops.reshard_v2(src_value, dst_dist_attr)
|
||||
fake_var.set_type(dst_type)
|
||||
return fake_var
|
||||
|
||||
|
||||
class RToSReshardFunctionCrossMesh(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_replicated(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_shard(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if (
|
||||
in_mesh.ndim != 1
|
||||
or out_mesh.ndim != 1
|
||||
or in_mesh.shape != out_mesh.shape
|
||||
):
|
||||
return False
|
||||
|
||||
if in_mesh == out_mesh:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
same_status_func = SameStatusReshardFunction()
|
||||
tmp_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
src_dist_attr.dims_mapping,
|
||||
src_dist_attr.partial_status,
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dist_attr
|
||||
)
|
||||
out_value = same_status_func.reshard(
|
||||
src_dist_attr, tmp_dist_attr, src_value, tmp_dst_type
|
||||
)
|
||||
|
||||
if out_value is None:
|
||||
return None
|
||||
|
||||
curr_global_rank = paddle.distributed.get_rank()
|
||||
if curr_global_rank in dst_dist_attr.process_mesh.process_ids:
|
||||
r_to_s_func = RToSReshardFunction()
|
||||
assert r_to_s_func.is_suitable(tmp_dist_attr, dst_dist_attr), (
|
||||
f"Invoke the r to s reshard function is not valid from {tmp_dist_attr} to {dst_dist_attr}"
|
||||
)
|
||||
return r_to_s_func.reshard(
|
||||
tmp_dist_attr, dst_dist_attr, out_value, dst_type
|
||||
)
|
||||
return None
|
||||
@@ -0,0 +1,59 @@
|
||||
# Copyright (c) 2024 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.
|
||||
|
||||
from .base_reshard_func import register_reshard_func
|
||||
from .global_to_sub_mesh_func import GlobalToSubMeshFunction
|
||||
from .nd_mesh_reshard_func import (
|
||||
NdMeshReshardFunction,
|
||||
NdMeshReshardFunctionCrossMesh,
|
||||
)
|
||||
from .p_to_r_reshard_func import (
|
||||
PToRReshardFunction,
|
||||
PToRReshardFunctionCrossMesh,
|
||||
)
|
||||
from .p_to_s_reshard_func import (
|
||||
PToSReshardFunction,
|
||||
)
|
||||
from .r_to_p_reshard_func import RToPReshardFunction
|
||||
from .r_to_s_reshard_func import (
|
||||
RToSReshardFunction,
|
||||
RToSReshardFunctionCrossMesh,
|
||||
)
|
||||
from .s_to_r_reshard_func import (
|
||||
SToRReshardFunction,
|
||||
SToRReshardFunctionCrossMesh,
|
||||
)
|
||||
from .s_to_s_reshard_func import SToSReshardFunction
|
||||
from .same_status_reshard_func import SameStatusReshardFunction
|
||||
from .sub_to_global_mesh_func import SubToGlobalMeshFunction
|
||||
|
||||
|
||||
def register_reshard_funcs():
|
||||
register_reshard_func(PToRReshardFunction())
|
||||
register_reshard_func(PToRReshardFunctionCrossMesh())
|
||||
register_reshard_func(PToSReshardFunction())
|
||||
register_reshard_func(RToSReshardFunction())
|
||||
register_reshard_func(RToSReshardFunctionCrossMesh())
|
||||
register_reshard_func(RToPReshardFunction())
|
||||
register_reshard_func(SameStatusReshardFunction())
|
||||
register_reshard_func(SToRReshardFunction())
|
||||
register_reshard_func(SToRReshardFunctionCrossMesh())
|
||||
register_reshard_func(NdMeshReshardFunction())
|
||||
register_reshard_func(NdMeshReshardFunctionCrossMesh())
|
||||
register_reshard_func(GlobalToSubMeshFunction())
|
||||
register_reshard_func(SubToGlobalMeshFunction())
|
||||
register_reshard_func(SToSReshardFunction())
|
||||
|
||||
|
||||
register_reshard_funcs()
|
||||
@@ -0,0 +1,363 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
copy_op_attr_with_new_member,
|
||||
is_replicated,
|
||||
is_shard,
|
||||
)
|
||||
from .same_status_reshard_func import SameStatusReshardFunction
|
||||
|
||||
|
||||
class SToRReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_shard(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_replicated(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def infer_allgather_dist_type(self, in_value, split_axis):
|
||||
tensor_ndim = len(in_value.shape)
|
||||
in_dist_attr = in_value.dist_attr()
|
||||
split_mesh_dim = in_dist_attr.dims_mapping[split_axis]
|
||||
mesh = in_dist_attr.process_mesh
|
||||
|
||||
# Calculate local shape. In nd_mesh_reshard, multiple tensor axis
|
||||
# may be shard and it will call this 1-D s_to_r function on each
|
||||
# axis. In this case, we should recompute the local and global shape.
|
||||
out_local_shape = list(in_value.shape)
|
||||
out_local_shape[split_axis] = int(
|
||||
(in_value.shape[split_axis] + mesh.shape[split_mesh_dim] - 1)
|
||||
/ mesh.shape[split_mesh_dim]
|
||||
)
|
||||
out_global_shape = list(out_local_shape)
|
||||
out_global_shape[0] *= mesh.shape[split_mesh_dim]
|
||||
out_type = paddle.pir.create_shaped_type(
|
||||
in_value.type(), out_global_shape
|
||||
)
|
||||
|
||||
out_dims_mapping = list(in_dist_attr.dims_mapping)
|
||||
out_dims_mapping[split_axis] = -1
|
||||
out_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
mesh, out_dims_mapping, in_dist_attr.partial_status
|
||||
)
|
||||
out_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
out_type, out_dist_attr
|
||||
)
|
||||
return out_type
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
if src_dist_attr.process_mesh.size == 1:
|
||||
dst_value = paddle._C_ops.share_data_(src_value)
|
||||
share_data_op = dst_value.get_defining_op()
|
||||
# set dist type and dist attr
|
||||
dst_value.set_type(dst_type)
|
||||
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
share_data_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[src_dist_attr],
|
||||
[dst_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
return dst_value
|
||||
|
||||
def get_split_axis_with_dims_mapping(dims_mapping):
|
||||
split_axis = {}
|
||||
for idx, v in enumerate(dims_mapping):
|
||||
if v != -1:
|
||||
split_axis[idx] = v
|
||||
return split_axis
|
||||
|
||||
split_axis_map = get_split_axis_with_dims_mapping(
|
||||
src_dist_attr.dims_mapping
|
||||
)
|
||||
|
||||
split_axis = -1
|
||||
for k, v in split_axis_map.items():
|
||||
split_axis = k
|
||||
break
|
||||
num_of_process = src_dist_attr.process_mesh.size
|
||||
num_of_padding = src_value.shape[split_axis] % num_of_process
|
||||
is_balanced_split = num_of_padding == 0
|
||||
|
||||
if is_balanced_split:
|
||||
new_value = self.reshard_s_to_r_with_padding(
|
||||
src_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
num_of_padding,
|
||||
)
|
||||
return new_value
|
||||
else:
|
||||
# find the last one
|
||||
need_padding = (
|
||||
paddle.distributed.get_rank()
|
||||
== src_dist_attr.process_mesh.process_ids[-1]
|
||||
)
|
||||
|
||||
# get padding_num
|
||||
avg_size_on_split_axis = int(
|
||||
(src_value.shape[split_axis] + num_of_process - 1)
|
||||
/ num_of_process
|
||||
)
|
||||
padding_num = (
|
||||
avg_size_on_split_axis * num_of_process
|
||||
- src_value.shape[split_axis]
|
||||
)
|
||||
if need_padding:
|
||||
# set right _local_shape
|
||||
local_shape_at_split_axis = src_value.shape[
|
||||
split_axis
|
||||
] - avg_size_on_split_axis * (num_of_process - 1)
|
||||
local_shape = src_value._local_shape
|
||||
local_shape[split_axis] = local_shape_at_split_axis
|
||||
tmp_src_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), src_dist_attr, list(local_shape)
|
||||
)
|
||||
src_value.set_type(tmp_src_type)
|
||||
padding_shape = src_value._local_shape
|
||||
padding_shape[split_axis] = padding_num
|
||||
padding_tensor = paddle.full(
|
||||
padding_shape,
|
||||
0.0,
|
||||
src_value.dtype,
|
||||
)
|
||||
tmp_src_type1 = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
padding_tensor.type(), dst_dist_attr
|
||||
)
|
||||
padding_tensor.set_type(tmp_src_type1)
|
||||
padding_tensor.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
dst_dist_attr.process_mesh, [], [dst_dist_attr]
|
||||
)
|
||||
)
|
||||
|
||||
concat_value = paddle._C_ops.concat(
|
||||
[src_value, padding_tensor], split_axis
|
||||
)
|
||||
# set concat dist_attr
|
||||
axis_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
src_dist_attr.process_mesh, [-1], {}
|
||||
)
|
||||
)
|
||||
concat_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[
|
||||
paddle.base.libpaddle.pir.create_array_attribute(
|
||||
[src_dist_attr, dst_dist_attr]
|
||||
),
|
||||
axis_dist_attr,
|
||||
],
|
||||
[src_dist_attr],
|
||||
)
|
||||
)
|
||||
# set concat_value type
|
||||
concat_global_shape = list(src_value.shape)
|
||||
concat_global_shape[split_axis] = (
|
||||
avg_size_on_split_axis * num_of_process
|
||||
)
|
||||
concat_type = paddle.pir.create_shaped_type(
|
||||
src_value.type(), concat_global_shape
|
||||
)
|
||||
concat_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
concat_type, src_dist_attr
|
||||
)
|
||||
concat_value.set_type(concat_type)
|
||||
|
||||
new_value = self.reshard_s_to_r_with_padding(
|
||||
concat_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
padding_num,
|
||||
)
|
||||
return new_value
|
||||
else:
|
||||
new_value = self.reshard_s_to_r_with_padding(
|
||||
src_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
padding_num,
|
||||
)
|
||||
return new_value
|
||||
|
||||
def reshard_s_to_r_with_padding(
|
||||
self,
|
||||
src_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
padding_num=0,
|
||||
):
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
num_of_process = len(src_mesh.process_ids)
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
|
||||
group = new_process_group(sorted(src_mesh.process_ids))
|
||||
allgather_value = paddle._C_ops.all_gather(
|
||||
src_value, group.id, num_of_process
|
||||
)
|
||||
allgather_type = self.infer_allgather_dist_type(src_value, split_axis)
|
||||
allgather_value.set_type(allgather_type)
|
||||
|
||||
# set op_dist_attr
|
||||
new_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
[-1] * len(dst_dist_attr.dims_mapping),
|
||||
dst_dist_attr.partial_status,
|
||||
)
|
||||
allgather_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh, [src_dist_attr], [new_dist_attr], chunk_id
|
||||
)
|
||||
)
|
||||
|
||||
if split_axis != 0 or padding_num != 0:
|
||||
allgather_op = allgather_value.get_defining_op()
|
||||
split_values = paddle._C_ops.split_with_num(
|
||||
allgather_op.result(0), num_of_process, 0
|
||||
)
|
||||
builtin_split_op = split_values[0].get_defining_op()
|
||||
pd_split_op = builtin_split_op.operand_source(0).get_defining_op()
|
||||
pd_split_op.dist_attr = copy_op_attr_with_new_member(
|
||||
pd_split_op.dist_attr, new_chunk_id=chunk_id
|
||||
)
|
||||
|
||||
# fix the split_with_num dist attribute.
|
||||
new_inner_types = []
|
||||
for sub_value in split_values:
|
||||
new_inner_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
sub_value.type(), allgather_value.dist_attr()
|
||||
)
|
||||
new_inner_types.append(new_inner_type)
|
||||
sub_value.set_type(new_inner_type)
|
||||
vec_type = paddle.base.libpaddle.pir.create_vec_type(
|
||||
new_inner_types
|
||||
)
|
||||
pd_split_op.result(0).set_type(vec_type)
|
||||
|
||||
if padding_num != 0:
|
||||
tmp_split_values = paddle._C_ops.split(
|
||||
split_values[-1],
|
||||
[
|
||||
split_values[-1].shape[split_axis] - padding_num,
|
||||
padding_num,
|
||||
],
|
||||
split_axis,
|
||||
)
|
||||
split_op = tmp_split_values.get_defining_op()
|
||||
split_op.dist_attr = copy_op_attr_with_new_member(
|
||||
split_op.dist_attr, new_chunk_id=chunk_id
|
||||
)
|
||||
split_values[-1] = tmp_split_values[0]
|
||||
concat_value = paddle._C_ops.concat(split_values, split_axis)
|
||||
concat_op = concat_value.get_defining_op()
|
||||
concat_op.dist_attr = copy_op_attr_with_new_member(
|
||||
concat_op.dist_attr, new_chunk_id=chunk_id
|
||||
)
|
||||
return concat_value
|
||||
else:
|
||||
concat_value = paddle._C_ops.concat(split_values, split_axis)
|
||||
# fold builtin.split op and builtin.combine op
|
||||
concat_op = concat_value.get_defining_op()
|
||||
concat_op.dist_attr = copy_op_attr_with_new_member(
|
||||
concat_op.dist_attr, new_chunk_id=chunk_id
|
||||
)
|
||||
builtin_combine_op = concat_op.operand_source(
|
||||
0
|
||||
).get_defining_op()
|
||||
concat_op.operand(0).set_source(pd_split_op.result(0))
|
||||
builtin_combine_op.erase()
|
||||
builtin_split_op.erase()
|
||||
return concat_value
|
||||
|
||||
return allgather_value
|
||||
|
||||
|
||||
class SToRReshardFunctionCrossMesh(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_shard(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_replicated(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if (
|
||||
in_mesh.ndim != 1
|
||||
or out_mesh.ndim != 1
|
||||
or in_mesh.shape != out_mesh.shape
|
||||
):
|
||||
return False
|
||||
|
||||
if in_mesh == out_mesh:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
same_status_func = SameStatusReshardFunction()
|
||||
tmp_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
src_dist_attr.dims_mapping,
|
||||
src_dist_attr.partial_status,
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dist_attr
|
||||
)
|
||||
out_value = same_status_func.reshard(
|
||||
src_dist_attr, tmp_dist_attr, src_value, tmp_dst_type
|
||||
)
|
||||
|
||||
s_to_r_func = SToRReshardFunction()
|
||||
assert s_to_r_func.is_suitable(tmp_dist_attr, dst_dist_attr), (
|
||||
f"Invoke the p to r reshard function is not valid from {tmp_dist_attr} to {dst_dist_attr}"
|
||||
)
|
||||
return s_to_r_func.reshard(
|
||||
tmp_dist_attr, dst_dist_attr, out_value, dst_type
|
||||
)
|
||||
@@ -0,0 +1,124 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
from paddle.distributed.utils.stream_utils import ExecutionStreamType
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
is_shard,
|
||||
)
|
||||
|
||||
|
||||
class SToSReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_shard(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_shard(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
"""
|
||||
Reshard from shard to shard status on 1D mesh.
|
||||
E.g. tensor shape: [B, S, H], mesh = [0, 1]
|
||||
1. [Shard(0)] --> [Shard(1)], N ranks:
|
||||
1). reshape from [B, S, H] -> [B, N, S/N, H]
|
||||
2). transpose from [B, N, S/N, H] -> [N, B, S/N, H]
|
||||
3). reshape from [N, B, S/N, H] -> [N*B, S/N, H]
|
||||
4). all to all communicate
|
||||
2. [Shard(1)] --> [Shard(0)], N ranks:
|
||||
1). all to all communicate
|
||||
2). reshape from [B, S, H] -> [N, B/N, S, H]
|
||||
3). transpose from [N, B/N, S, H] -> [B/N, N, S/N, H]
|
||||
4). reshape from [B/N, N, S/N, H] -> [B, S, H]
|
||||
"""
|
||||
in_split_axis = src_dist_attr.dims_mapping.index(0)
|
||||
out_split_axis = dst_dist_attr.dims_mapping.index(0)
|
||||
nranks = len(src_dist_attr.process_mesh.process_ids)
|
||||
|
||||
if out_split_axis != 0:
|
||||
pre_shape = copy.copy(src_value.shape)
|
||||
if pre_shape[out_split_axis] != -1:
|
||||
pre_shape[out_split_axis] = pre_shape[out_split_axis] // nranks
|
||||
pre_shape.insert(out_split_axis, nranks)
|
||||
out_reshape1 = paddle._C_ops.reshape(src_value, pre_shape)
|
||||
|
||||
axes = [out_split_axis]
|
||||
for i in range(len(pre_shape)):
|
||||
if i != out_split_axis:
|
||||
axes.append(i)
|
||||
out_transpose = paddle._C_ops.transpose(out_reshape1, axes)
|
||||
|
||||
pre_shape.pop(out_split_axis)
|
||||
if pre_shape[in_split_axis] != -1:
|
||||
pre_shape[in_split_axis] *= nranks
|
||||
in_all2all = paddle._C_ops.reshape(out_transpose, pre_shape)
|
||||
in_all2all_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), dst_dist_attr
|
||||
)
|
||||
in_all2all.set_type(in_all2all_type)
|
||||
else:
|
||||
in_all2all = paddle._C_ops.share_data_(src_value)
|
||||
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
group = new_process_group(sorted(src_mesh.process_ids))
|
||||
dst_value = paddle._C_ops.all_to_all(in_all2all, group.id)
|
||||
dst_value.get_defining_op().set_execution_stream(
|
||||
ExecutionStreamType.DefaultStream.value
|
||||
)
|
||||
out_all2all_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
in_all2all.type(), src_dist_attr
|
||||
)
|
||||
dst_value.set_type(out_all2all_type)
|
||||
dst_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh, [src_dist_attr], [dst_dist_attr], -1
|
||||
)
|
||||
)
|
||||
|
||||
if in_split_axis != 0:
|
||||
post_shape = copy.copy(src_value.shape)
|
||||
if post_shape[0] != -1:
|
||||
post_shape[0] = post_shape[0] // nranks
|
||||
post_shape.insert(0, nranks)
|
||||
dst_value = paddle.reshape(dst_value, post_shape)
|
||||
|
||||
axes = list(range(1, len(post_shape)))
|
||||
axes.insert(in_split_axis, 0)
|
||||
dst_value = paddle._C_ops.transpose(dst_value, axes)
|
||||
|
||||
post_shape.pop(0)
|
||||
if post_shape[in_split_axis] != -1:
|
||||
post_shape[in_split_axis] *= nranks
|
||||
dst_value = paddle._C_ops.reshape(dst_value, post_shape)
|
||||
|
||||
dst_value.set_type(dst_type)
|
||||
|
||||
return dst_value
|
||||
+158
@@ -0,0 +1,158 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
from paddle.distributed.passes.pass_utils import find_var_used_op_chunk_id
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import ReshardFunction
|
||||
|
||||
|
||||
class SameStatusReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if src_dist_attr.dims_mapping != dst_dist_attr.dims_mapping:
|
||||
return False
|
||||
if src_dist_attr.partial_dims != dst_dist_attr.partial_dims:
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh == out_mesh:
|
||||
return False
|
||||
if in_mesh.shape != out_mesh.shape:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
dst_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
all_process_ids = list(
|
||||
set(src_mesh.process_ids) | set(dst_mesh.process_ids)
|
||||
)
|
||||
all_process_ids = sorted(all_process_ids)
|
||||
|
||||
cur_global_rank = paddle.distributed.get_rank()
|
||||
|
||||
for src, dst in zip(src_mesh.process_ids, dst_mesh.process_ids):
|
||||
if src != dst:
|
||||
new_process_group([src, dst], group_type="p2p")
|
||||
new_process_group([dst, src], group_type="p2p")
|
||||
|
||||
is_send = True
|
||||
for src, dst in zip(src_mesh.process_ids, dst_mesh.process_ids):
|
||||
if src == cur_global_rank:
|
||||
chunk_id = -1
|
||||
if (
|
||||
src_value.get_defining_op().name() == "pd_op.add_n"
|
||||
and src_value.get_defining_op()
|
||||
.operand_source(0)
|
||||
.get_defining_op()
|
||||
.name()
|
||||
== "builtin.combine"
|
||||
):
|
||||
add_n_op = src_value.get_defining_op()
|
||||
combine_op = add_n_op.operand_source(0).get_defining_op()
|
||||
combine_op_chunk_id_list = []
|
||||
for input in combine_op.operands():
|
||||
if input.source().get_defining_op().dist_attr:
|
||||
combine_op_chunk_id_list.append(
|
||||
input.source()
|
||||
.get_defining_op()
|
||||
.dist_attr.chunk_id
|
||||
)
|
||||
else:
|
||||
combine_op_chunk_id_list.append(-1)
|
||||
# check combine_op operands chunk_id equal
|
||||
assert all(
|
||||
x == combine_op_chunk_id_list[0]
|
||||
for x in combine_op_chunk_id_list
|
||||
), "combine_op's operands has different chunk_id."
|
||||
chunk_id = combine_op_chunk_id_list[0]
|
||||
# reset add_n chunk_id
|
||||
add_n_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
add_n_op.dist_attr.process_mesh,
|
||||
add_n_op.dist_attr.operands(),
|
||||
add_n_op.dist_attr.results(),
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
else:
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = (
|
||||
src_value.get_defining_op().dist_attr.chunk_id
|
||||
)
|
||||
|
||||
comm_group = new_process_group([src, dst], group_type="p2p")
|
||||
paddle._C_ops.send_v2(
|
||||
src_value,
|
||||
comm_group.id,
|
||||
comm_group.ranks.index(dst),
|
||||
True,
|
||||
False,
|
||||
)
|
||||
point = paddle.base.libpaddle.pir.get_current_insertion_point()
|
||||
point.prev()
|
||||
new_op = point.get_operation()
|
||||
assert new_op.name() == "pd_op.send_v2"
|
||||
new_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh, [src_dist_attr], [], chunk_id
|
||||
)
|
||||
)
|
||||
break
|
||||
|
||||
elif dst == cur_global_rank:
|
||||
all_used_ops = src_value.all_used_ops()
|
||||
chunk_id = -1
|
||||
for used_op in all_used_ops:
|
||||
var = used_op.result(0)
|
||||
if var.dist_attr().process_mesh == dst_mesh:
|
||||
chunk_id = find_var_used_op_chunk_id(var)
|
||||
|
||||
assert -1 not in dst_type.shape, (
|
||||
"dynamic shape is not supported by pir-auto parallel yet."
|
||||
)
|
||||
|
||||
comm_group = new_process_group([src, dst], group_type="p2p")
|
||||
recv_value = paddle._C_ops.recv_v2(
|
||||
dst_type._local_shape,
|
||||
dst_type.dtype,
|
||||
comm_group.ranks.index(src),
|
||||
comm_group.id,
|
||||
True,
|
||||
False,
|
||||
)
|
||||
new_op = recv_value.get_defining_op()
|
||||
new_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
dst_mesh,
|
||||
[],
|
||||
[dst_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
recv_value.set_type(dst_type)
|
||||
is_send = False
|
||||
break
|
||||
|
||||
if is_send:
|
||||
# fake var will be removed in remove_other_rank_op_pass.
|
||||
fake_var = paddle._C_ops.reshard_v2(src_value, dst_dist_attr)
|
||||
fake_var.set_type(dst_type)
|
||||
return fake_var
|
||||
else:
|
||||
return recv_value
|
||||
+171
@@ -0,0 +1,171 @@
|
||||
# Copyright (c) 2024 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 paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.distributed.auto_parallel.placement_type import (
|
||||
check_placements_equal,
|
||||
)
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from ..utils import split_mesh
|
||||
from .base_reshard_func import ReshardFunction, copy_dist_attr_with_new_member
|
||||
|
||||
|
||||
def _mesh_equal_ignore_shape_one(mesh1, mesh2, dim: int):
|
||||
"""
|
||||
Check if two process meshes are equal, ignoring the shape value `1`
|
||||
in the specified dimension. This is used when mesh1 is a sub-mesh
|
||||
split from a global mesh, in this case, the shape of mesh1 is `1`
|
||||
in the split dim.
|
||||
E.g, the following two meshes are equal:
|
||||
mesh1: shape = [1,2,2], process_ids = [0,1,2,3]
|
||||
mesh2: shape = [2,2], process_ids = [0,1,2,3]
|
||||
"""
|
||||
assert dim >= 0 and dim < len(mesh1.shape), "invalid dim arg"
|
||||
if mesh1 == mesh2:
|
||||
return True
|
||||
|
||||
if mesh1.process_ids != mesh2.process_ids:
|
||||
return False
|
||||
|
||||
a_shape = copy.copy(mesh1.shape)
|
||||
b_shape = copy.copy(mesh2.shape)
|
||||
|
||||
if a_shape[dim] != 1:
|
||||
return False
|
||||
a_shape.pop(dim)
|
||||
|
||||
return a_shape == b_shape
|
||||
|
||||
|
||||
class SubToGlobalMeshFunction(ReshardFunction):
|
||||
"""
|
||||
Reshard from sub-mesh to global mesh, now only supports
|
||||
both input and output values are replicated, e.g.
|
||||
1. input: mesh:[0], placements:[Replicate()]
|
||||
output: mesh:[0,1], placements:[Replicate()]
|
||||
2. input: mesh:[0,1], placements:[Replicate()]
|
||||
output: mesh:[[0,1],[2,3]], placements:[Replicate(), Replicate()]
|
||||
"""
|
||||
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
sub_mesh_dim = paddle.base.core.sub_mesh_dim(out_mesh, in_mesh)
|
||||
if sub_mesh_dim == -1:
|
||||
return False
|
||||
sub_meshes, sub_placements = (
|
||||
dist.auto_parallel.api._get_sub_meshes_and_local_placements(
|
||||
out_mesh, dst_dist_attr.placements_attr, sub_mesh_dim
|
||||
)
|
||||
)
|
||||
if not check_placements_equal(
|
||||
src_dist_attr.placements_attr, sub_placements
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
dst_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
sub_mesh_dim = paddle.base.core.sub_mesh_dim(dst_mesh, src_mesh)
|
||||
sub_meshes = split_mesh(dst_mesh, sub_mesh_dim)
|
||||
dst_meshes = [
|
||||
mesh
|
||||
for mesh in sub_meshes
|
||||
if not _mesh_equal_ignore_shape_one(mesh, src_mesh, sub_mesh_dim)
|
||||
]
|
||||
|
||||
comm_group_ids = []
|
||||
root_ranks = []
|
||||
for p_id in src_mesh.process_ids:
|
||||
comm_group_ids.append([p_id])
|
||||
root_ranks.append(p_id)
|
||||
for i, group_ids in enumerate(comm_group_ids):
|
||||
for mesh in dst_meshes:
|
||||
group_ids.append(mesh.process_ids[i])
|
||||
|
||||
other_ranks = copy.copy(dst_mesh.process_ids)
|
||||
for rank in other_ranks:
|
||||
if rank in src_mesh.process_ids:
|
||||
other_ranks.remove(rank)
|
||||
|
||||
cur_rank = paddle.distributed.get_rank()
|
||||
|
||||
if cur_rank in src_mesh.process_ids:
|
||||
# the root rank will broadcast the src_value to other ranks
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
tmp_value = paddle._C_ops.share_data_(src_value)
|
||||
value_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), src_value.dist_attr()
|
||||
)
|
||||
tmp_value.set_type(value_type)
|
||||
op = tmp_value.get_defining_op()
|
||||
op.dist_attr = paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh, [src_dist_attr], [src_dist_attr], chunk_id
|
||||
)
|
||||
elif cur_rank in other_ranks:
|
||||
# create the buffer on other ranks for receiving the data
|
||||
tmp_value = paddle.zeros(dst_type.shape, dst_type.dtype)
|
||||
op = tmp_value.get_defining_op()
|
||||
mesh = paddle.distributed.ProcessMesh(other_ranks)
|
||||
value_dist_attr = copy_dist_attr_with_new_member(
|
||||
dst_dist_attr, new_process_mesh=mesh
|
||||
)
|
||||
value_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
dst_type, value_dist_attr
|
||||
)
|
||||
tmp_value.set_type(value_type)
|
||||
op.dist_attr = paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
mesh, [], [value_dist_attr]
|
||||
)
|
||||
else:
|
||||
# do nothing if the current rank is not in src_mesh and dst_mesh.
|
||||
# use reshard_op to create and return a fake value, and the fake
|
||||
# value will be removed 'remove_other_rank_op_pass'.
|
||||
fake_var = paddle._C_ops.reshard_v2(src_value, dst_dist_attr)
|
||||
return fake_var
|
||||
|
||||
# create communication groups
|
||||
for i, group_ids in enumerate(comm_group_ids):
|
||||
comm_group_ids[i] = sorted(group_ids)
|
||||
# the root arg in broadcast is the local index
|
||||
# of the rank in the communication group
|
||||
root_ranks[i] = comm_group_ids[i].index(root_ranks[i])
|
||||
|
||||
comm_groups = []
|
||||
for i, group_ids in enumerate(comm_group_ids):
|
||||
comm_groups.append(new_process_group(group_ids))
|
||||
if cur_rank in group_ids:
|
||||
cur_group_id = i
|
||||
|
||||
broadcast_value = paddle._C_ops.broadcast(
|
||||
tmp_value, comm_groups[cur_group_id].id, root_ranks[cur_group_id]
|
||||
)
|
||||
broadcast_value.set_type(dst_type)
|
||||
|
||||
broadcast_op = broadcast_value.get_defining_op()
|
||||
broadcast_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
dst_mesh, [src_dist_attr], [dst_dist_attr]
|
||||
)
|
||||
)
|
||||
|
||||
return broadcast_value
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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.
|
||||
|
||||
from .profiler import profiler # noqa: F401
|
||||
|
||||
__all__ = []
|
||||
@@ -0,0 +1,241 @@
|
||||
# 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 logging
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from ..utils import get_logger, is_recompute_op
|
||||
from .trial import (
|
||||
OptimizationTunerTrial as Trial,
|
||||
TrialStatus,
|
||||
)
|
||||
|
||||
|
||||
class AlgorithmBase(ABC):
|
||||
"""
|
||||
An Tuning algorithm is a class to find out an optimal configuration
|
||||
given the selected tuning optimization pass(es) and the arguments to be tuned.
|
||||
Different optimization pass(es) will correspond to a different algorithm,
|
||||
where different search space **pruning rules** will applied.
|
||||
|
||||
In another word, the key "algorithm" for this class is the
|
||||
search space pruning rules specific for the given optimization scenario.
|
||||
"""
|
||||
|
||||
_REGISTERED_ALGORITHMS = {}
|
||||
|
||||
name = None
|
||||
|
||||
@staticmethod
|
||||
def _register(algo_name, algo_class):
|
||||
assert issubclass(algo_class, AlgorithmBase)
|
||||
AlgorithmBase._REGISTERED_ALGORITHMS[algo_name] = algo_class
|
||||
|
||||
def __init__(self, config):
|
||||
self._config = config
|
||||
self._init_spaces()
|
||||
self._logger = get_logger(logging.INFO)
|
||||
self._changed_configs = []
|
||||
|
||||
@property
|
||||
def changed_configs(self):
|
||||
return self._changed_configs[:]
|
||||
|
||||
def collect_model_info(self, main_prog, startup_prog):
|
||||
"""
|
||||
Collect the model static info (from programs) that could be used to
|
||||
pruning candidate trials and saving tuning time. For instance,
|
||||
model info like number of model parameters and activation memory could be
|
||||
used to prune candidate trial and decide the next trial.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _init_spaces(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def next_trial(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def update(self, results):
|
||||
"""
|
||||
Update the algorithm with the results of last trial. Using this information is used to
|
||||
pruning the search space of the future trial.
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_config_from_trial(self, trial):
|
||||
"""
|
||||
Return a new fleet.DistributedStrategy with the configurations in trial.
|
||||
"""
|
||||
assert len(self._changed_configs) > 0
|
||||
new_strategy = copy.deepcopy(self._config.dist_strategy)
|
||||
for name in self._changed_configs:
|
||||
config = getattr(trial.space, name)
|
||||
setattr(new_strategy, name, config)
|
||||
return new_strategy
|
||||
|
||||
|
||||
def register_algor(name):
|
||||
def impl(cls):
|
||||
AlgorithmBase._register(name, cls)
|
||||
cls.name = name
|
||||
return cls
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
def new_algorithm(name, config):
|
||||
algor_class = AlgorithmBase._REGISTERED_ALGORITHMS.get(name)
|
||||
assert algor_class is not None, f"Algorithm {name} is not defined."
|
||||
algor_obj = algor_class(config)
|
||||
return algor_obj
|
||||
|
||||
|
||||
@register_algor("sharding")
|
||||
class ShardingStageAlgorithm(AlgorithmBase):
|
||||
# TODO import trial class & copy strategy
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self._changed_configs = ["sharding"]
|
||||
|
||||
def _init_spaces(self):
|
||||
self._max_stage = 3
|
||||
self._trial_idx = 0
|
||||
|
||||
stage_range = self._config.sharding.get("tuning_range", None)
|
||||
if stage_range:
|
||||
assert set(stage_range).issubset({0, 1, 2, 3}), (
|
||||
f"Sharding Stage should belong into range within 0 - 3 but got {stage_range}."
|
||||
)
|
||||
stage_range.sort(reverse=True)
|
||||
else:
|
||||
stage_range = list(range(self._max_stage + 1))
|
||||
stage_range.sort(reverse=True)
|
||||
self._stage_range = stage_range[:]
|
||||
self._total_num_trial = len(self._stage_range)
|
||||
|
||||
def next_trial(self):
|
||||
if self._trial_idx < self._total_num_trial:
|
||||
stage = self._stage_range[self._trial_idx]
|
||||
|
||||
new_strategy = copy.deepcopy(self._config.dist_strategy)
|
||||
sharding = new_strategy.sharding
|
||||
sharding.stage = stage
|
||||
|
||||
name = f"trial-sharding-stage{stage}"
|
||||
trial = Trial(new_strategy, name, self.changed_configs)
|
||||
|
||||
return trial
|
||||
else:
|
||||
return Trial(None, None, None, status=TrialStatus.STOPPED)
|
||||
|
||||
def update(self, results):
|
||||
et = results.get("ErrorType", None)
|
||||
if et and et == "ResourceExhaustedError":
|
||||
self._trial_idx = self._total_num_trial
|
||||
self._logger.info(
|
||||
"Last trial is failed with OOM, all remaining trials are pruned to save time !"
|
||||
)
|
||||
else:
|
||||
self._trial_idx += 1
|
||||
|
||||
|
||||
@register_algor("recompute")
|
||||
class RecomputeCheckpointAlgorithm(AlgorithmBase):
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self._changed_configs = ["recompute"]
|
||||
|
||||
def collect_model_info(self, main_prog, startup_prog):
|
||||
segments = []
|
||||
for op in main_prog.global_block().ops:
|
||||
if not is_recompute_op(op):
|
||||
continue
|
||||
|
||||
seg_name = op.attr('op_namescope')
|
||||
if seg_name not in segments:
|
||||
segments.append(seg_name)
|
||||
|
||||
self._total_num_trial = len(segments)
|
||||
self._tuning_segments = list(range(len(segments)))
|
||||
self._trial_left = 0
|
||||
self._trial_right = len(segments) - 1
|
||||
self._trial_idx = int(0 + (len(segments) - 1) / 2)
|
||||
|
||||
def _init_spaces(self):
|
||||
self._recompute_mode = "all"
|
||||
|
||||
def next_trial(self):
|
||||
if self._trial_idx < self._total_num_trial:
|
||||
if self._recompute_mode == "all":
|
||||
self._recompute_flag = False
|
||||
new_strategy = copy.deepcopy(self._config.dist_strategy)
|
||||
name = "trial-recompute-all-segments"
|
||||
return Trial(new_strategy, name, self.changed_configs)
|
||||
elif self._recompute_mode == "none":
|
||||
self._recompute_flag = False
|
||||
new_strategy = copy.deepcopy(self._config.dist_strategy)
|
||||
recompute = new_strategy.recompute
|
||||
recompute.enable = False
|
||||
name = "trial-recompute-none-segments"
|
||||
return Trial(new_strategy, name, self.changed_configs)
|
||||
elif self._recompute_mode == "part":
|
||||
new_no_recompute = self._tuning_segments[: self._trial_idx]
|
||||
new_strategy = copy.deepcopy(self._config.dist_strategy)
|
||||
recompute = new_strategy.recompute
|
||||
recompute.no_recompute_segments.extend(new_no_recompute)
|
||||
name = f"trial-recompute-part-segments-idx{self._trial_idx}"
|
||||
return Trial(new_strategy, name, self.changed_configs)
|
||||
else:
|
||||
return Trial(None, None, None, status=TrialStatus.STOPPED)
|
||||
|
||||
def update(self, results):
|
||||
et = results.get("ErrorType", None)
|
||||
if self._recompute_mode == "all":
|
||||
if et and et == "ResourceExhaustedError":
|
||||
self._trial_idx = self._total_num_trial
|
||||
self._logger.info(
|
||||
"Recompute all candidate segments is failed with OOM, please reduce model size or batch size."
|
||||
)
|
||||
else:
|
||||
self._recompute_mode = "none"
|
||||
elif self._recompute_mode == "none":
|
||||
if et and et == "ResourceExhaustedError":
|
||||
self._recompute_mode = "part"
|
||||
else:
|
||||
self._trial_idx = self._total_num_trial
|
||||
self._logger.info(
|
||||
"Recompute is unnecessary for this model size, which will reduce the Throughput."
|
||||
)
|
||||
else:
|
||||
if self._trail_left >= self._trail_right:
|
||||
self._trial_idx = self._total_num_trial
|
||||
elif et and et == "ResourceExhaustedError":
|
||||
self._trail_left = self._trail_left
|
||||
self._trail_right = self._trial_idx - 1
|
||||
self._trial_idx = int(
|
||||
self._trail_left
|
||||
+ (self._trail_right - self._trail_left) / 2
|
||||
)
|
||||
else:
|
||||
self._trail_left = self._trial_idx + 1
|
||||
self._trail_right = self._trail_right
|
||||
self._trial_idx = int(
|
||||
self._trail_left
|
||||
+ (self._trail_right - self._trail_left) / 2
|
||||
)
|
||||
@@ -0,0 +1,122 @@
|
||||
# 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 os
|
||||
|
||||
from ...strategy import Strategy
|
||||
|
||||
_tuning_supported_passes = ["sharding", "recompute"]
|
||||
|
||||
|
||||
def _get_pass_config(strategy, pass_name):
|
||||
config = getattr(strategy, pass_name)
|
||||
return config
|
||||
|
||||
|
||||
class TuningConfig:
|
||||
"""
|
||||
A uniform config wrap:
|
||||
distributed strategy: the user defined configuration for optimization pass
|
||||
tuning config: configuration for the tuning process: mode (profile or cost model), log dir, extra tuning config for optimization like search range for specific
|
||||
"""
|
||||
|
||||
def __init__(self, strategy):
|
||||
if not isinstance(strategy, Strategy):
|
||||
raise TypeError("'strategy' must be object of class `Strategy`.")
|
||||
|
||||
self._tuning_passes_name = set()
|
||||
self._dist_strategy = copy.deepcopy(strategy)
|
||||
self._mode = None
|
||||
self._profile_start_step = None
|
||||
self._profile_end_step = None
|
||||
self._project_dir = None
|
||||
self._max_num_trial = None
|
||||
self._early_stop = None
|
||||
self._debug = None
|
||||
|
||||
self._initialize()
|
||||
|
||||
@property
|
||||
def mode(self):
|
||||
return self._mode
|
||||
|
||||
@property
|
||||
def profile_start_step(self):
|
||||
return self._profile_start_step
|
||||
|
||||
@property
|
||||
def profile_end_step(self):
|
||||
return self._profile_end_step
|
||||
|
||||
@property
|
||||
def project_dir(self):
|
||||
return self._project_dir
|
||||
|
||||
@property
|
||||
def tuning_passes_name(self):
|
||||
return self._tuning_passes_name
|
||||
|
||||
@property
|
||||
def max_num_trial(self):
|
||||
return self._max_num_trial
|
||||
|
||||
@property
|
||||
def early_stop(self):
|
||||
return self._early_stop
|
||||
|
||||
@property
|
||||
def debug(self):
|
||||
return self._debug
|
||||
|
||||
@property
|
||||
def dist_strategy(self):
|
||||
return self._dist_strategy
|
||||
|
||||
# initialize config with user define value or default value
|
||||
def _initialize(self):
|
||||
tuning_strategy = self._dist_strategy.tuning
|
||||
|
||||
self._mode = tuning_strategy.get("mode", "PROFILE")
|
||||
self._profile_start_step = tuning_strategy.get("profile_start_step", 10)
|
||||
self._profile_end_step = tuning_strategy.get("profile_end_step", 30)
|
||||
self._max_num_trial = tuning_strategy.get("max_num_trial", 50)
|
||||
self._early_stop = tuning_strategy.get("early_stop", None)
|
||||
self._debug = tuning_strategy.get("debug", False)
|
||||
|
||||
project_dir = tuning_strategy.get("project_dir", None)
|
||||
if not project_dir:
|
||||
project_dir = os.path.join(os.getcwd(), "OptimizationTuning")
|
||||
self._project_dir = project_dir
|
||||
|
||||
for p in _tuning_supported_passes:
|
||||
if (
|
||||
getattr(self._dist_strategy, p)
|
||||
and _get_pass_config(self._dist_strategy, p).enable_tuning
|
||||
):
|
||||
# TODO distinguish different args of each passes
|
||||
self._tuning_passes_name.add(p)
|
||||
|
||||
p_strategy = getattr(self._dist_strategy, p)
|
||||
self.__dict__[p] = p_strategy
|
||||
|
||||
# # TODO verify the user defined configs
|
||||
# tuning_config_for_pass = tuning_strategy.get(p, None)
|
||||
# if tuning_config_for_pass:
|
||||
# for k, v in tuning_config_for_pass.items():
|
||||
# self.__dict__[p][k] = v
|
||||
|
||||
# (NOTE)tuning config ONLY wraps dist strategy for pass config which is to be tuned
|
||||
def __getattr__(self, item):
|
||||
return getattr(self._dist_strategy, item)
|
||||
@@ -0,0 +1,643 @@
|
||||
# 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 json
|
||||
import logging
|
||||
|
||||
# import yaml
|
||||
import os
|
||||
import pathlib
|
||||
import pickle
|
||||
import shlex
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.auto_parallel.static.completion import Completer
|
||||
from paddle.distributed.auto_parallel.static.dist_context import (
|
||||
DistributedContext,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.partitioner import Partitioner
|
||||
from paddle.distributed.auto_parallel.static.process_group import (
|
||||
clear_all_process_groups,
|
||||
get_all_process_groups,
|
||||
new_process_group,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.reshard import Resharder
|
||||
from paddle.distributed.auto_parallel.static.utils import debug_program
|
||||
from paddle.distributed.passes import PassContext, new_pass
|
||||
from paddle.static import append_backward, program_guard
|
||||
|
||||
from ..utils import get_logger
|
||||
from .algorithms import new_algorithm
|
||||
from .config import TuningConfig
|
||||
from .trial import TrialStatus
|
||||
|
||||
|
||||
def _get_new_params_grads(target_program, ref_program, ref_params_grads):
|
||||
ref_block = ref_program.global_block()
|
||||
target_block = target_program.global_block()
|
||||
target_params_grads = []
|
||||
|
||||
for p, g in ref_params_grads:
|
||||
# NOTE grad var might not be generated
|
||||
assert ref_block.has_var(p.name)
|
||||
assert target_block.has_var(p.name)
|
||||
new_p = target_block.var(p.name)
|
||||
if g:
|
||||
new_g = target_block.var(g.name)
|
||||
else:
|
||||
new_g = None
|
||||
|
||||
target_params_grads.append((new_p, new_g))
|
||||
|
||||
return target_params_grads
|
||||
|
||||
|
||||
def _get_new_loss(target_program, ref_program, loss):
|
||||
ref_block = ref_program.global_block()
|
||||
target_block = target_program.global_block()
|
||||
assert ref_block.has_var(loss.name)
|
||||
|
||||
return target_block.var(loss.name)
|
||||
|
||||
|
||||
def parse_process_groups():
|
||||
group_map = {}
|
||||
all_process_groups = get_all_process_groups()
|
||||
for process_group in all_process_groups:
|
||||
group_map[process_group.id] = process_group.ranks
|
||||
return group_map
|
||||
|
||||
|
||||
def get_metric(results):
|
||||
assert isinstance(results, dict), (
|
||||
f"results should be type of dictionary, but got {type(results)}."
|
||||
)
|
||||
if 'Throughput' in results and isinstance(results['Throughput'], float):
|
||||
return float(results['Throughput'])
|
||||
else:
|
||||
return -1.0
|
||||
|
||||
|
||||
def parse_results(results):
|
||||
if results['Throughput'] > 0:
|
||||
return "Throughput: {} step / s.".format(results['Throughput'])
|
||||
et = results.get("ErrorType", None)
|
||||
if et == "ResourceExhaustedError":
|
||||
return "Fail with OOM"
|
||||
else:
|
||||
return "Fail with UNKNOWN ERROR"
|
||||
|
||||
|
||||
# TODO only dependent on dist context
|
||||
# all env need to be start a new pass are member of dist context
|
||||
def _copy_context(ref_dist_context):
|
||||
# clear all process groups and recover the world process group
|
||||
clear_all_process_groups()
|
||||
ranks = []
|
||||
for process_mesh in ref_dist_context._process_meshes:
|
||||
ranks.extend(process_mesh.process_ids)
|
||||
new_process_group(list(set(ranks)))
|
||||
|
||||
new_dist_context = DistributedContext()
|
||||
new_dist_context._serial_main_program = (
|
||||
ref_dist_context.serial_main_program.clone(for_test=False)
|
||||
)
|
||||
new_dist_context._serial_startup_program = (
|
||||
ref_dist_context.serial_startup_program.clone(for_test=False)
|
||||
)
|
||||
|
||||
# mapping variable into new dist context
|
||||
if getattr(ref_dist_context, '_params_grads', None):
|
||||
new_dist_context._params_grads = _get_new_params_grads(
|
||||
new_dist_context.serial_main_program,
|
||||
ref_dist_context.serial_main_program,
|
||||
ref_dist_context._params_grads,
|
||||
)
|
||||
new_dist_context._serial_loss = _get_new_loss(
|
||||
new_dist_context.serial_main_program,
|
||||
ref_dist_context.serial_main_program,
|
||||
ref_dist_context.serial_loss,
|
||||
)
|
||||
|
||||
for key, var_list in ref_dist_context._serial_feed_vars.items():
|
||||
new_var_list = []
|
||||
for var in var_list:
|
||||
block_idx = var.block.idx
|
||||
var_name = var.name
|
||||
var = new_dist_context._serial_main_program.blocks[
|
||||
block_idx
|
||||
]._var_recursive(var_name)
|
||||
new_var_list.append(var)
|
||||
new_dist_context._serial_feed_vars[key] = new_var_list
|
||||
|
||||
for key, var_list in ref_dist_context._serial_fetch_vars.items():
|
||||
new_var_list = []
|
||||
# metrics is a list of list
|
||||
if key == "metrics":
|
||||
for inner_var_list in var_list:
|
||||
new_inner_var_list = []
|
||||
for var in inner_var_list:
|
||||
block_idx = var.block.idx
|
||||
var_name = var.name
|
||||
var = new_dist_context._serial_main_program.blocks[
|
||||
block_idx
|
||||
]._var_recursive(var_name)
|
||||
new_inner_var_list.append(var)
|
||||
new_var_list.append(new_inner_var_list)
|
||||
else:
|
||||
for var in var_list:
|
||||
block_idx = var.block.idx
|
||||
var_name = var.name
|
||||
var = new_dist_context._serial_main_program.blocks[
|
||||
block_idx
|
||||
]._var_recursive(var_name)
|
||||
new_var_list.append(var)
|
||||
new_dist_context._serial_fetch_vars[key] = new_var_list
|
||||
|
||||
# copy information in forward and backward
|
||||
new_dist_context._serial_optimizer = copy.deepcopy(
|
||||
ref_dist_context.serial_optimizer
|
||||
)
|
||||
new_dist_context._dist_tensors_for_program = copy.deepcopy(
|
||||
ref_dist_context._dist_tensors_for_program
|
||||
)
|
||||
new_dist_context._dist_ops_for_program = copy.deepcopy(
|
||||
ref_dist_context._dist_ops_for_program
|
||||
)
|
||||
for pm in ref_dist_context.process_meshes:
|
||||
new_dist_context.add_process_mesh(pm)
|
||||
new_dist_context._dist_op_context = copy.deepcopy(
|
||||
ref_dist_context._dist_op_context
|
||||
)
|
||||
new_dist_context._block_state = copy.deepcopy(ref_dist_context.block_state)
|
||||
|
||||
return new_dist_context
|
||||
|
||||
|
||||
class OptimizationTuner:
|
||||
"""
|
||||
OptimizationTuner is used to manage the tuning procedure of hyper-parameters (configs)
|
||||
of Optimization Pass in AutoParallel.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dist_context,
|
||||
dataset,
|
||||
inputs_spec,
|
||||
labels_spec,
|
||||
batch_size,
|
||||
rank,
|
||||
):
|
||||
self._config = TuningConfig(dist_context.strategy)
|
||||
# should not modify dist context from calling function
|
||||
self._baseline_dist_context = _copy_context(dist_context)
|
||||
self._baseline_completer = Completer(self._baseline_dist_context)
|
||||
|
||||
self._rank = rank
|
||||
self._inputs_spec = inputs_spec
|
||||
self._labels_spec = labels_spec
|
||||
self._dataset = dataset
|
||||
self._batch_size = batch_size
|
||||
|
||||
self._finished_trials = []
|
||||
self._best_metric = None
|
||||
self._best_iter = float("-inf")
|
||||
|
||||
self._logger = get_logger(logging.INFO)
|
||||
|
||||
self._build_programs_without_optimization()
|
||||
self._select_tuning_algorithm()
|
||||
|
||||
@property
|
||||
def project_dir(self):
|
||||
dirname = self._config.project_dir
|
||||
if not os.path.exists(dirname):
|
||||
if self.rank == 0:
|
||||
pathlib.Path(dirname).mkdir(parents=True, exist_ok=True)
|
||||
return dirname
|
||||
|
||||
@property
|
||||
def rank(self):
|
||||
return self._rank
|
||||
|
||||
@property
|
||||
def device_id(self):
|
||||
return paddle.distributed.ParallelEnv().device_id
|
||||
|
||||
# TODO Generate complete program with all parts like forward, backward, update
|
||||
# as well as parallelism transformation.
|
||||
def _build_programs_without_optimization(self):
|
||||
serial_main_program = self._baseline_dist_context.serial_main_program
|
||||
serial_startup_program = (
|
||||
self._baseline_dist_context.serial_startup_program
|
||||
)
|
||||
serial_loss = self._baseline_dist_context.serial_loss
|
||||
|
||||
with program_guard(serial_main_program, serial_startup_program):
|
||||
params_grads = append_backward(
|
||||
serial_loss,
|
||||
distop_context=self._baseline_dist_context.dist_op_context,
|
||||
)
|
||||
|
||||
self._baseline_completer.complete_backward_annotation(
|
||||
serial_main_program
|
||||
)
|
||||
self._baseline_dist_context.block_state.parse_backward_blocks(
|
||||
serial_main_program
|
||||
)
|
||||
self._baseline_dist_context._params_grads = params_grads
|
||||
|
||||
if self._config.debug:
|
||||
baseline_dir = os.path.join(self.project_dir, "baseline")
|
||||
if not os.path.exists(baseline_dir):
|
||||
pathlib.Path(baseline_dir).mkdir(parents=True, exist_ok=True)
|
||||
debug_program(
|
||||
self._baseline_dist_context._serial_main_program,
|
||||
baseline_dir,
|
||||
"main",
|
||||
)
|
||||
debug_program(
|
||||
self._baseline_dist_context._serial_startup_program,
|
||||
baseline_dir,
|
||||
"startup",
|
||||
)
|
||||
|
||||
def _select_tuning_algorithm(self):
|
||||
selected_passes_set = self._config.tuning_passes_name
|
||||
algorithm_name = "_".join(sorted(selected_passes_set))
|
||||
self._algorithm = new_algorithm(algorithm_name, self._config)
|
||||
|
||||
def _apply_optimization(self, trial):
|
||||
new_strategy = trial.space
|
||||
dist_context = _copy_context(self._baseline_dist_context)
|
||||
pass_context = PassContext()
|
||||
completer = Completer(dist_context)
|
||||
|
||||
main_program = dist_context.serial_main_program
|
||||
startup_program = dist_context.serial_startup_program
|
||||
|
||||
# applying optimization pass
|
||||
if new_strategy.amp.enable:
|
||||
config = copy.deepcopy(new_strategy.amp.to_dict())
|
||||
config["dist_context"] = dist_context
|
||||
config["params_grads"] = dist_context._params_grads
|
||||
# TODO AMP Pass should not use loss var
|
||||
config["loss"] = dist_context.serial_loss
|
||||
config["input_data"] = (
|
||||
self._baseline_dist_context.serial_feed_vars["inputs"]
|
||||
+ self._baseline_dist_context.serial_feed_vars["labels"]
|
||||
)
|
||||
if config["dtype"] == "float16" and config["level"] == "o2":
|
||||
config["base_opt"] = dist_context.serial_optimizer
|
||||
auto_parallel_fp16_pass = new_pass("auto_parallel_fp16", config)
|
||||
auto_parallel_fp16_pass.apply(
|
||||
[main_program], [startup_program], pass_context
|
||||
)
|
||||
dist_context._serial_loss = auto_parallel_fp16_pass.get_loss()
|
||||
else:
|
||||
auto_parallel_amp_pass = new_pass("auto_parallel_amp", config)
|
||||
auto_parallel_amp_pass.apply(
|
||||
[main_program], [startup_program], pass_context
|
||||
)
|
||||
dist_context._serial_loss = auto_parallel_amp_pass.get_loss()
|
||||
|
||||
if new_strategy.recompute.enable:
|
||||
config = copy.deepcopy(new_strategy.recompute.to_dict())
|
||||
config["dist_context"] = dist_context
|
||||
config["no_grad_set"] = None
|
||||
config["loss"] = dist_context.serial_loss
|
||||
auto_parallel_recompute_pass = new_pass(
|
||||
"auto_parallel_recompute", config
|
||||
)
|
||||
auto_parallel_recompute_pass.apply(
|
||||
[main_program], [startup_program], pass_context
|
||||
)
|
||||
|
||||
# Do logical partition
|
||||
partitioner = Partitioner(dist_context, self.rank)
|
||||
(
|
||||
dist_main_prog,
|
||||
dist_startup_prog,
|
||||
dist_params_grads,
|
||||
) = partitioner.partition(
|
||||
main_program, startup_program, dist_context._params_grads
|
||||
)
|
||||
|
||||
# Generate optimizer
|
||||
# FIXME should be remove from apply pass after pass support optimizers
|
||||
with (
|
||||
program_guard(dist_main_prog, dist_startup_prog),
|
||||
dist_main_prog.switch_name_generator_guard("opt_"),
|
||||
):
|
||||
optimizer_ops = dist_context.serial_optimizer.apply_gradients(
|
||||
dist_params_grads
|
||||
)
|
||||
completer.complete_update_annotation(dist_main_prog)
|
||||
|
||||
resharder = Resharder(
|
||||
dist_main_prog,
|
||||
dist_startup_prog,
|
||||
self.rank,
|
||||
dist_context,
|
||||
dist_params_grads,
|
||||
)
|
||||
resharder.reshard()
|
||||
|
||||
config = {}
|
||||
config["dist_context"] = dist_context
|
||||
config["global_rank"] = self.rank
|
||||
config["use_sharding"] = new_strategy.sharding.enable
|
||||
dp_pass = new_pass("auto_parallel_data_parallel_optimization", config)
|
||||
dp_pass.apply([dist_main_prog], [dist_startup_prog], pass_context)
|
||||
|
||||
if new_strategy.sharding.enable:
|
||||
config = copy.deepcopy(new_strategy.sharding.to_dict())
|
||||
config["dist_context"] = dist_context
|
||||
config["params_grads"] = dist_params_grads
|
||||
config["global_rank"] = self.rank
|
||||
auto_parallel_sharding_pass = new_pass(
|
||||
"auto_parallel_sharding", config
|
||||
)
|
||||
auto_parallel_sharding_pass.apply(
|
||||
[dist_main_prog], [dist_startup_prog], pass_context
|
||||
)
|
||||
dist_params_grads = pass_context.get_attr("params_grads")
|
||||
|
||||
# gradient clip
|
||||
config = copy.deepcopy(new_strategy.sharding.to_dict())
|
||||
config["dist_context"] = dist_context
|
||||
config["params_grads"] = dist_params_grads
|
||||
config["rank_id"] = self.rank
|
||||
auto_parallel_clip_pass = new_pass("auto_parallel_grad_clip", config)
|
||||
auto_parallel_clip_pass.apply(
|
||||
[dist_main_prog], [dist_startup_prog], pass_context
|
||||
)
|
||||
|
||||
if new_strategy.gradient_merge.enable:
|
||||
config = copy.deepcopy(new_strategy.gradient_merge.to_dict())
|
||||
config["dist_context"] = dist_context
|
||||
config["params_grads"] = dist_params_grads
|
||||
auto_parallel_gradient_merge_pass = new_pass(
|
||||
"auto_parallel_gradient_merge_pass", config
|
||||
)
|
||||
auto_parallel_gradient_merge_pass.apply(
|
||||
[dist_main_prog], [dist_startup_prog], pass_context
|
||||
)
|
||||
trial.main_program, trial.startup_program = (
|
||||
dist_main_prog,
|
||||
dist_startup_prog,
|
||||
)
|
||||
return trial
|
||||
|
||||
def _get_profile_context(self, trial, result_path):
|
||||
profile_ctx = {}
|
||||
|
||||
profile_ctx['distributed_env'] = copy.deepcopy(
|
||||
paddle.distributed.ParallelEnv()
|
||||
)
|
||||
profile_ctx['group_map'] = parse_process_groups()
|
||||
profile_ctx["loss_var_name"] = (
|
||||
self._baseline_dist_context.serial_loss.name
|
||||
)
|
||||
profile_ctx["main_program_decs"] = (
|
||||
trial.main_program.desc.serialize_to_string()
|
||||
)
|
||||
profile_ctx["startup_program_decs"] = (
|
||||
trial.startup_program.desc.serialize_to_string()
|
||||
)
|
||||
self._dataset.batch_size = self._batch_size
|
||||
self._dataset.input_names = self._get_input_names()
|
||||
|
||||
profile_ctx["dataset"] = self._dataset
|
||||
profile_ctx["result_filename"] = result_path
|
||||
|
||||
return profile_ctx
|
||||
|
||||
def _get_input_names(self):
|
||||
input_names = []
|
||||
for input_spec in self._inputs_spec[:] + self._labels_spec[:]:
|
||||
input_names.append(input_spec.name)
|
||||
return input_names
|
||||
|
||||
def _launch_profile(self, ctx_path, trial_dir):
|
||||
if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
|
||||
coverage_args = ["-m", "coverage", "run", "--branch", "-p"]
|
||||
else:
|
||||
coverage_args = []
|
||||
|
||||
profile_args = " ".join(
|
||||
[
|
||||
"--rank",
|
||||
str(self.rank),
|
||||
"--device_id",
|
||||
str(self.device_id),
|
||||
"--ctx_filename",
|
||||
ctx_path,
|
||||
"--profile_start_step",
|
||||
str(self._config.profile_start_step),
|
||||
"--profile_end_step",
|
||||
str(self._config.profile_end_step),
|
||||
]
|
||||
)
|
||||
cmd_args = (
|
||||
"-m paddle.distributed.auto_parallel.static.tuner.profiler"
|
||||
+ " "
|
||||
+ profile_args
|
||||
)
|
||||
cmd = [sys.executable, "-u", *coverage_args, *shlex.split(cmd_args)]
|
||||
|
||||
parent_env = copy.copy(os.environ.copy())
|
||||
# env flags need for profile
|
||||
new_env = {}
|
||||
new_env.update(parent_env)
|
||||
|
||||
# TODO if any rank hang or fail, kill all processes
|
||||
self._logger.debug("Executing cmd:\n{} .".format(" ".join(cmd)))
|
||||
# new_process = subprocess.Popen(cmd, env=new_env)
|
||||
with (
|
||||
open(
|
||||
os.path.join(trial_dir, "stdout.log" + str(self.rank)), "wb"
|
||||
) as out,
|
||||
open(
|
||||
os.path.join(trial_dir, "stderr.log" + str(self.rank)), "wb"
|
||||
) as err,
|
||||
):
|
||||
result = subprocess.Popen(cmd, stdout=out, stderr=err, env=new_env)
|
||||
result.wait()
|
||||
out.flush()
|
||||
err.flush()
|
||||
os.fsync(out)
|
||||
os.fsync(err)
|
||||
|
||||
def _profile_trial(self, trial):
|
||||
# Making working directory
|
||||
trial_dir = self._get_trial_dir(trial)
|
||||
if not os.path.exists(trial_dir):
|
||||
if self.rank == 0:
|
||||
pathlib.Path(trial_dir).mkdir(parents=True, exist_ok=True)
|
||||
else:
|
||||
while not os.path.exists(trial_dir):
|
||||
pass
|
||||
ctx_filename = "profile_ctx." + str(self.rank)
|
||||
ctx_path = os.path.join(trial_dir, ctx_filename)
|
||||
result_path = os.path.join(trial_dir, "result.json")
|
||||
|
||||
# Prepare Profile Context
|
||||
profile_ctx = self._get_profile_context(trial, result_path)
|
||||
with open(ctx_path, 'wb') as f:
|
||||
pickle.dump(profile_ctx, f, protocol=4)
|
||||
|
||||
if self._config.debug:
|
||||
debug_program(trial.main_program, trial_dir, "main_program")
|
||||
debug_program(trial.startup_program, trial_dir, "startup_program")
|
||||
|
||||
# Run
|
||||
self._launch_profile(ctx_path, trial_dir)
|
||||
|
||||
# Load results
|
||||
try:
|
||||
with open(result_path, 'r') as fp:
|
||||
results = json.load(fp)
|
||||
return results
|
||||
except FileNotFoundError:
|
||||
Error_results = {"Throughput": -1, "ErrorType": 'FatalError'}
|
||||
return Error_results
|
||||
|
||||
def _evaluate_trial(self, trial):
|
||||
self._logger.info(f"Trial {trial.name} evaluation start.")
|
||||
self._apply_optimization(trial)
|
||||
|
||||
if self._config.mode == "PROFILE":
|
||||
results = self._profile_trial(trial)
|
||||
|
||||
elif self._config.mode == "COSTMODEL":
|
||||
raise NotImplementedError(
|
||||
"COSTMODEL mode for optimization tuning is not supported yet!"
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"invalid evaluation mode: {self._config.mode}"
|
||||
)
|
||||
|
||||
self._logger.info(
|
||||
f"Trial {trial.name} evaluation finish with {parse_results(results)}."
|
||||
)
|
||||
return results
|
||||
|
||||
def _update(self, i, trial, results):
|
||||
self._finished_trials.append(trial)
|
||||
|
||||
cur_metric = get_metric(results)
|
||||
if self._best_metric is None or cur_metric > self._best_metric:
|
||||
self._best_metric = cur_metric
|
||||
self._best_iter = i
|
||||
|
||||
def _get_trial_dir(self, trial):
|
||||
return os.path.join(self.project_dir, trial.name)
|
||||
|
||||
def get_best_config(self):
|
||||
"""
|
||||
Return the best optimization configuration found in the tuning.
|
||||
|
||||
Returns:
|
||||
A object of fleet.DistributedStrategy with best configuration.
|
||||
"""
|
||||
assert self._best_iter >= 0, "The best configuration is not found yet !"
|
||||
best_trial = self._finished_trials[self._best_iter]
|
||||
return self._algorithm.get_config_from_trial(best_trial)
|
||||
|
||||
def summary(self):
|
||||
"""
|
||||
Display tuning result summary.
|
||||
"""
|
||||
# TODO summary with the trial_name with metric_of_trial
|
||||
best_trial = self._finished_trials[self._best_iter]
|
||||
summary_ = f"""
|
||||
Tuning Result Summary
|
||||
Run total {len(self._finished_trials)} trials with {(time.time() - self._tuning_start_time) / 60} min.
|
||||
The best trial is: [{best_trial.name}], whose configuration is following:
|
||||
"""
|
||||
summary_ += "\n" + best_trial.summary() + "\n"
|
||||
self._logger.info(summary_)
|
||||
with open(os.path.join(self.project_dir, "summary.txt"), "w+") as fw:
|
||||
fw.writelines(line + "\n" for line in summary_.split("\n"))
|
||||
|
||||
# full_strategy = self.get_best_config()
|
||||
# path = os.path.join(self.project_dir, "tuned_dist_strategy.yaml")
|
||||
# with open(path, 'w') as outfile:
|
||||
# yaml.dump(full_strategy, outfile, default_flow_style=False)
|
||||
|
||||
def clear(self):
|
||||
"""
|
||||
Clear the temporary file generated in tuning procedure.
|
||||
"""
|
||||
# TODO clear up zombie process created by tuning
|
||||
if not self._config.debug:
|
||||
for trial in self._finished_trials:
|
||||
trial_dir = self._get_trial_dir(trial)
|
||||
shutil.rmtree(trial_dir, ignore_errors=True)
|
||||
|
||||
def tune(self):
|
||||
"""
|
||||
Performs the search for best hyperparameter configurations
|
||||
for the selected optimization pass(es).
|
||||
"""
|
||||
|
||||
# step1: collect model info which might be used for
|
||||
# pruning the search space of the algorithm
|
||||
self._tuning_start_time = time.time()
|
||||
self._algorithm.collect_model_info(
|
||||
self._baseline_dist_context.serial_main_program,
|
||||
self._baseline_dist_context.serial_startup_program,
|
||||
)
|
||||
|
||||
# main search loop
|
||||
i = 0
|
||||
while i < self._config.max_num_trial:
|
||||
# step2: create a new trial
|
||||
trial = self._algorithm.next_trial()
|
||||
|
||||
if trial.status == TrialStatus.STOPPED:
|
||||
break
|
||||
|
||||
# step3: evaluate the trial
|
||||
results = self._evaluate_trial(trial)
|
||||
|
||||
# step4: update the algorithm with last result,
|
||||
# which could be used by algorithm to pruning the
|
||||
# remaining search space.
|
||||
self._algorithm.update(results)
|
||||
self._update(i, trial, results)
|
||||
|
||||
# early stop
|
||||
i += 1
|
||||
if (
|
||||
self._config.early_stop
|
||||
and self._config.early_stop <= i - self._best_iter
|
||||
):
|
||||
self._logger.info(
|
||||
f"Early stop the Tuning since there is no better trial found within [{self._config.early_stop}] trials"
|
||||
)
|
||||
break
|
||||
|
||||
# step5: summary the best config and return
|
||||
self.summary()
|
||||
|
||||
self.clear()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,296 @@
|
||||
# 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 argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.auto_parallel.static.dist_loader import (
|
||||
DistributedDataLoaderFromGenerator,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.process_group import (
|
||||
get_all_process_groups,
|
||||
new_process_group,
|
||||
)
|
||||
from paddle.distributed.collective import _get_global_env
|
||||
from paddle.framework import Program, _current_expected_place
|
||||
from paddle.static import Operator
|
||||
|
||||
paddle.enable_static()
|
||||
|
||||
|
||||
def str2bool(v):
|
||||
if v.lower() in ('yes', 'true', 't', 'y', '1'):
|
||||
return True
|
||||
elif v.lower() in ('no', 'false', 'f', 'n', '0'):
|
||||
return False
|
||||
else:
|
||||
raise argparse.ArgumentTypeError('Unsupported value encountered.')
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--profile_start_step",
|
||||
default=10,
|
||||
type=int,
|
||||
help="integer indicates the warmup step before starting profile.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile_end_step",
|
||||
default=30,
|
||||
type=int,
|
||||
help="integer indicates at the end step of profile.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rank",
|
||||
type=int,
|
||||
required=True,
|
||||
help="the rank id of the this process.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--device_id",
|
||||
type=int,
|
||||
required=True,
|
||||
help="the device id of the this process.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ctx_filename",
|
||||
type=str,
|
||||
required=True,
|
||||
help="the filename to the profile context file saved by optimization tuner",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def init_process_groups(group_map, rank):
|
||||
for group_id, ranks in group_map.items():
|
||||
if group_id == 0:
|
||||
continue
|
||||
new_process_group(ranks=ranks, group_id=group_id)
|
||||
|
||||
# TODO should instantiate global group first
|
||||
all_process_groups = get_all_process_groups()
|
||||
for process_group in all_process_groups:
|
||||
print(process_group)
|
||||
process_group.instantiate()
|
||||
|
||||
|
||||
def get_cpp_error_type(error):
|
||||
msg = str(error).splitlines()
|
||||
cpp_error_types = [
|
||||
'InvalidArgumentError',
|
||||
'NotFoundError',
|
||||
'OutOfRangeError',
|
||||
'AlreadyExistsError',
|
||||
'ResourceExhaustedError',
|
||||
'PreconditionNotMetError',
|
||||
'PermissionDeniedError',
|
||||
'ExecutionTimeoutError',
|
||||
'UnimplementedError',
|
||||
'UnavailableError',
|
||||
'FatalError',
|
||||
'ExternalError',
|
||||
]
|
||||
error_type = 'FatalError'
|
||||
for et in cpp_error_types:
|
||||
for line in msg:
|
||||
if et in line:
|
||||
return et
|
||||
return error_type
|
||||
|
||||
|
||||
def create_dataloader(
|
||||
main_program, startup_program, profile_ctx, epochs=1, steps_per_epoch=None
|
||||
):
|
||||
dataset = profile_ctx["dataset"]
|
||||
main_block = main_program.global_block()
|
||||
feed_list = []
|
||||
for name in dataset.input_names:
|
||||
if name in main_block.vars:
|
||||
feed_list.append(main_block.vars[name])
|
||||
|
||||
# remove the first three ops if multi run fit/evaluate/predict
|
||||
op_size = len(main_block.ops)
|
||||
if main_block.ops[0].type == 'create_py_reader':
|
||||
op_size -= 3
|
||||
for _ in range(3):
|
||||
main_block._remove_op(0, sync=False)
|
||||
|
||||
# insert read op at the end of program
|
||||
places = paddle.static.cuda_places()
|
||||
with paddle.static.program_guard(main_program, startup_program):
|
||||
dataloader = DistributedDataLoaderFromGenerator(
|
||||
dataset=dataset,
|
||||
feed_list=feed_list,
|
||||
capacity=70,
|
||||
places=places,
|
||||
batch_size=dataset.batch_size,
|
||||
epochs=epochs,
|
||||
steps_per_epoch=steps_per_epoch,
|
||||
data_parallel_world_size=dataset.dp_world_size,
|
||||
data_parallel_rank=dataset.dp_rank,
|
||||
)
|
||||
|
||||
# move read op from the end of program to the start of program
|
||||
new_op_size = len(main_block.ops)
|
||||
for _ in range(new_op_size - 1, op_size - 1, -1):
|
||||
op = main_block.ops[new_op_size - 1]
|
||||
new_op_desc = main_block.desc._prepend_op()
|
||||
new_op_desc.copy_from(op.desc)
|
||||
new_op = Operator(main_block, new_op_desc, type=new_op_desc.type())
|
||||
main_block.ops.insert(0, new_op)
|
||||
for _ in range(new_op_size - op_size):
|
||||
main_block._remove_op(new_op_size, sync=False)
|
||||
main_block._sync_with_cpp()
|
||||
return dataloader
|
||||
|
||||
|
||||
def init_comm(profile_ctx):
|
||||
# override the env for current process
|
||||
dist_env = profile_ctx['distributed_env']
|
||||
genv = _get_global_env()
|
||||
genv = dist_env
|
||||
print(
|
||||
f"current process rank: {genv.rank}, device_id: {genv.device_id}, ip: {genv.current_endpoint}."
|
||||
)
|
||||
|
||||
# init nccl comm
|
||||
group_map = profile_ctx['group_map']
|
||||
init_process_groups(group_map, args.rank)
|
||||
|
||||
|
||||
def load_programs(profile_ctx):
|
||||
main_program_desc_str = profile_ctx['main_program_decs']
|
||||
main_program = Program.parse_from_string(main_program_desc_str)
|
||||
|
||||
startup_program_decs_str = profile_ctx['startup_program_decs']
|
||||
startup_program = Program.parse_from_string(startup_program_decs_str)
|
||||
|
||||
loss_var_name = profile_ctx["loss_var_name"]
|
||||
assert main_program.global_block().has_var(loss_var_name)
|
||||
loss_var = main_program.global_block().var(loss_var_name)
|
||||
|
||||
return main_program, startup_program, loss_var
|
||||
|
||||
|
||||
def get_executor():
|
||||
place_type = _current_expected_place()
|
||||
if not isinstance(place_type, paddle.CUDAPlace):
|
||||
raise RuntimeError("OptimizationTuner only support CUDA GPU right now.")
|
||||
|
||||
genv = _get_global_env()
|
||||
place = paddle.CUDAPlace(genv.device_id)
|
||||
exe = paddle.static.Executor(place)
|
||||
return exe
|
||||
|
||||
|
||||
def profiler(args):
|
||||
"""
|
||||
main function to profile experiment for each pass hyper-parameter.
|
||||
"""
|
||||
# load ctx
|
||||
if not os.path.isfile(args.ctx_filename):
|
||||
raise ValueError(
|
||||
f"There is no profile context named {args.ctx_filename}."
|
||||
)
|
||||
with open(args.ctx_filename, 'rb') as f:
|
||||
from paddle.framework.restricted_unpickler import safe_load_pickle
|
||||
|
||||
profile_ctx = safe_load_pickle(f, encoding='latin1')
|
||||
|
||||
init_comm(profile_ctx)
|
||||
|
||||
main_program, startup_program, loss_var = load_programs(profile_ctx)
|
||||
|
||||
data_loader = create_dataloader(main_program, startup_program, profile_ctx)
|
||||
|
||||
result_path = profile_ctx["result_filename"]
|
||||
|
||||
exe = get_executor()
|
||||
|
||||
try:
|
||||
exe.run(startup_program)
|
||||
# profile main
|
||||
duration = 0
|
||||
eval_step = 0
|
||||
data_loader._inner_dataloader.start()
|
||||
while eval_step < args.profile_end_step:
|
||||
start_time = time.time()
|
||||
|
||||
loss = exe.run(
|
||||
main_program,
|
||||
fetch_list=[loss_var],
|
||||
use_program_cache=True,
|
||||
)
|
||||
|
||||
end_time = time.time()
|
||||
|
||||
if eval_step >= args.profile_start_step:
|
||||
duration += end_time - start_time
|
||||
|
||||
print(f"step: {eval_step}, loss_print: {loss[0]:f}")
|
||||
eval_step += 1
|
||||
|
||||
avg_tput = (
|
||||
1.0 * (args.profile_end_step - args.profile_start_step) / duration
|
||||
)
|
||||
|
||||
result_dict = {
|
||||
"Throughput": avg_tput,
|
||||
"ErrorType": None,
|
||||
}
|
||||
|
||||
if paddle.distributed.get_rank() == 0:
|
||||
with open(result_path, 'w') as fp:
|
||||
json.dump(result_dict, fp)
|
||||
|
||||
print(f"profile done! avg speed : {avg_tput} step / s.")
|
||||
|
||||
except paddle.framework.core.EOFException:
|
||||
data_loader._inner_dataloader.reset()
|
||||
|
||||
except Exception as e:
|
||||
error_type = get_cpp_error_type(e)
|
||||
result_dict = {
|
||||
"Throughput": -1,
|
||||
"ErrorType": error_type,
|
||||
}
|
||||
if not os.path.isfile(result_path):
|
||||
with open(result_path, 'w') as fp:
|
||||
json.dump(result_dict, fp)
|
||||
|
||||
print(f"profile failed with error: [{error_type}]")
|
||||
print(e)
|
||||
print(traceback.format_exc())
|
||||
|
||||
data_loader._inner_dataloader.reset()
|
||||
del data_loader._inner_dataloader
|
||||
sys.exit(1)
|
||||
|
||||
data_loader._inner_dataloader.reset()
|
||||
del data_loader._inner_dataloader
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
paddle.framework.set_flags({'FLAGS_new_executor_sequential_run': 1})
|
||||
args = parse_args()
|
||||
profiler(args)
|
||||
@@ -0,0 +1,216 @@
|
||||
# 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.
|
||||
|
||||
# Notice that the following codes are modified from KerasTuner for a different purpose.
|
||||
# Please refer to https://github.com/keras-team/keras-tuner/blob/master/keras_tuner/engine/metrics_tracking.py.
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class MetricRecord:
|
||||
"""
|
||||
One record for a single metric at a given execution step.
|
||||
"""
|
||||
|
||||
def __init__(self, value, step):
|
||||
self._value = value
|
||||
self._step = step
|
||||
|
||||
@property
|
||||
def value(self):
|
||||
return self._value
|
||||
|
||||
@value.setter
|
||||
def value(self, value):
|
||||
self._value = value
|
||||
|
||||
@property
|
||||
def step(self):
|
||||
return self._step
|
||||
|
||||
@step.setter
|
||||
def step(self, step):
|
||||
self._step = step
|
||||
|
||||
def mean(self):
|
||||
return np.mean(self.value)
|
||||
|
||||
def get_state(self):
|
||||
return {"value": self.value, "step": self.step}
|
||||
|
||||
@classmethod
|
||||
def from_state(cls, state):
|
||||
return cls(**state)
|
||||
|
||||
def __eq__(self, other):
|
||||
if not isinstance(other, MetricRecord):
|
||||
return False
|
||||
return other.value == self.value and other.step == self.step
|
||||
|
||||
def __repr__(self):
|
||||
return f"MetricRecord(value={self.value}, step={self.step})"
|
||||
|
||||
|
||||
class MetricRecords:
|
||||
"""
|
||||
Records of a single metric across different executions.
|
||||
"""
|
||||
|
||||
def __init__(self, direction="min"):
|
||||
if direction not in {"min", "max"}:
|
||||
raise ValueError(
|
||||
f"direction should be one of {{min, max}}, but got: {direction}."
|
||||
)
|
||||
self._direction = direction
|
||||
self._records = {}
|
||||
|
||||
@property
|
||||
def records(self):
|
||||
return sorted(self._records.values(), key=lambda r: r.step)
|
||||
|
||||
@records.setter
|
||||
def records(self, records):
|
||||
for r in records:
|
||||
self.update(r.value, step=r.step)
|
||||
|
||||
@property
|
||||
def direction(self):
|
||||
return self._direction
|
||||
|
||||
@direction.setter
|
||||
def direction(self, direction):
|
||||
self._direction = direction
|
||||
|
||||
def update(self, value, step=0):
|
||||
if step in self._records:
|
||||
self._records[step].set_value(value)
|
||||
else:
|
||||
self._records[step] = MetricRecord(value, step=step)
|
||||
|
||||
def get_best_value(self):
|
||||
values = [r.mean() for r in self._records.values()]
|
||||
if not values:
|
||||
return None
|
||||
if self._direction == "min":
|
||||
return np.nanmin(values)
|
||||
return np.nanmax(values)
|
||||
|
||||
def get_best_step(self):
|
||||
best_value = self.get_best_value()
|
||||
if best_value is None:
|
||||
return None
|
||||
for r in self._records.values():
|
||||
if r.mean() == best_value:
|
||||
return r.step
|
||||
|
||||
def get_statistics(self):
|
||||
records = self.records
|
||||
records_values = [r.mean() for r in records]
|
||||
if not len(records_values):
|
||||
return {}
|
||||
return {
|
||||
"min": float(np.nanmin(records_values)),
|
||||
"max": float(np.nanmax(records_values)),
|
||||
"mean": float(np.nanmean(records_values)),
|
||||
"median": float(np.nanmedian(records_values)),
|
||||
"var": float(np.nanvar(records_values)),
|
||||
"std": float(np.nanstd(records_values)),
|
||||
}
|
||||
|
||||
def get_state(self):
|
||||
state = {}
|
||||
state["direction"] = self._direction
|
||||
state["records"] = [r.get_state() for r in self.records]
|
||||
return state
|
||||
|
||||
@classmethod
|
||||
def from_state(cls, state):
|
||||
records = cls(state["direction"])
|
||||
records.records = [MetricRecord.from_state(r) for r in state["records"]]
|
||||
return records
|
||||
|
||||
|
||||
class MetricsRecorder:
|
||||
"""
|
||||
Record the values for all metrics.
|
||||
"""
|
||||
|
||||
def __init__(self, metrics=None):
|
||||
self._records = {}
|
||||
self.register_metrics(metrics)
|
||||
|
||||
@property
|
||||
def records(self):
|
||||
return self._records
|
||||
|
||||
def exists(self, name):
|
||||
return name in self._records
|
||||
|
||||
def register_metrics(self, metrics=None):
|
||||
metrics = metrics or []
|
||||
for metric in metrics:
|
||||
self.register(metric.name)
|
||||
|
||||
def register(self, name, direction=None):
|
||||
if self.exists(name):
|
||||
raise ValueError(f"Metric {name} have been registered.")
|
||||
if direction is None:
|
||||
direction = "min"
|
||||
self._records[name] = MetricRecords(direction)
|
||||
|
||||
def update(self, name, value, step=0):
|
||||
value = float(value)
|
||||
if not self.exists(name):
|
||||
self.register(name)
|
||||
|
||||
prev_best = self._records[name].get_best_value()
|
||||
self._records[name].update(value, step=step)
|
||||
new_best = self._records[name].get_best_value()
|
||||
|
||||
improved = new_best != prev_best
|
||||
return improved
|
||||
|
||||
def get_records(self, name):
|
||||
return self._records[name].records
|
||||
|
||||
def set_records(self, name, records):
|
||||
if not self.exists(name):
|
||||
self.register(name)
|
||||
self._records[name].records = records
|
||||
|
||||
def get_best_value(self, name):
|
||||
return self._records[name].get_best_value()
|
||||
|
||||
def get_best_step(self, name):
|
||||
return self._records[name].get_best_step()
|
||||
|
||||
def get_statistics(self, name):
|
||||
return self._records[name].get_statistics()
|
||||
|
||||
def get_state(self):
|
||||
return {
|
||||
"metrics": {
|
||||
name: metric_records.get_state()
|
||||
for name, metric_records in self._records.items()
|
||||
}
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_state(cls, state):
|
||||
recorder = cls()
|
||||
recorder._records = {
|
||||
name: MetricRecords.from_state(metric_records)
|
||||
for name, metric_records in state["metrics"].items()
|
||||
}
|
||||
return recorder
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,42 @@
|
||||
# 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.
|
||||
|
||||
# Notice that the following codes are modified from KerasTuner for a different purpose.
|
||||
# Please refer to https://github.com/keras-team/keras-tuner/blob/master/keras_tuner/engine/metrics_tracking.py.
|
||||
|
||||
import json
|
||||
|
||||
|
||||
class Storable:
|
||||
def get_state(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def set_state(self, state):
|
||||
raise NotImplementedError
|
||||
|
||||
def save(self, path):
|
||||
state = self.get_state()
|
||||
state_json = json.dumps(state)
|
||||
try:
|
||||
with open(path, "w") as f:
|
||||
f.write(state_json)
|
||||
return str(path)
|
||||
except OSError as e:
|
||||
raise OSError(f"Failed to save file at {path}: {e}") from e
|
||||
|
||||
def load(self, path):
|
||||
with open(path, "r") as f:
|
||||
state_data = f.read()
|
||||
state = json.loads(state_data)
|
||||
self.set_state(state)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,169 @@
|
||||
# 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.
|
||||
|
||||
# Notice that the following codes are modified from KerasTuner to implement our own tuner.
|
||||
# Please refer to https://github.com/keras-team/keras-tuner/blob/master/keras_tuner/engine/trial.py.
|
||||
|
||||
import hashlib
|
||||
import random
|
||||
import time
|
||||
|
||||
from .recorder import MetricsRecorder
|
||||
from .storable import Storable
|
||||
from .tunable_space import TunableSpace
|
||||
|
||||
|
||||
class TrialStatus:
|
||||
RUNNING = "RUNNING"
|
||||
COMPLETED = "COMPLETED"
|
||||
STOPPED = "STOPPED"
|
||||
INVALID = "INVALID"
|
||||
|
||||
|
||||
class Trial(Storable):
|
||||
def __init__(
|
||||
self, tunable_space, trial_id=None, status=TrialStatus.RUNNING
|
||||
):
|
||||
self._id = _generate_trial_id() if trial_id is None else trial_id
|
||||
self._space = tunable_space
|
||||
self._recorder = MetricsRecorder()
|
||||
self._score = None
|
||||
self._best_step = None
|
||||
self._status = status
|
||||
|
||||
@property
|
||||
def id(self):
|
||||
return self._id
|
||||
|
||||
@property
|
||||
def space(self):
|
||||
return self._space
|
||||
|
||||
@property
|
||||
def recorder(self):
|
||||
return self._recorder
|
||||
|
||||
@property
|
||||
def score(self):
|
||||
return self._score
|
||||
|
||||
@score.setter
|
||||
def score(self, score):
|
||||
self._score = score
|
||||
|
||||
@property
|
||||
def best_step(self):
|
||||
return self._best_step
|
||||
|
||||
@best_step.setter
|
||||
def best_step(self, best_step):
|
||||
self._best_step = best_step
|
||||
|
||||
@property
|
||||
def status(self):
|
||||
return self._status
|
||||
|
||||
@status.setter
|
||||
def status(self, status):
|
||||
self._status = status
|
||||
|
||||
def summary(self):
|
||||
print("Tunable space:")
|
||||
if self.space.values:
|
||||
for tv, value in self.space.values.items():
|
||||
print(tv + ":", value)
|
||||
|
||||
if self.score is not None:
|
||||
print(f"Score: {self.score}")
|
||||
|
||||
def get_state(self):
|
||||
return {
|
||||
"id": self.id,
|
||||
"space": self.space.get_state(),
|
||||
"recorder": self.recorder.get_state(),
|
||||
"score": self.score,
|
||||
"best_step": self.best_step,
|
||||
"status": self.status,
|
||||
}
|
||||
|
||||
def set_state(self, state):
|
||||
self._id = state["id"]
|
||||
self._space = TunableSpace.from_state(state["space"])
|
||||
self._recorder = MetricsRecorder.from_state(state["recorder"])
|
||||
self._score = state["score"]
|
||||
self._best_step = state["best_step"]
|
||||
self._status = state["status"]
|
||||
|
||||
@classmethod
|
||||
def from_state(cls, state):
|
||||
trial = cls(tunable_space=None)
|
||||
trial.set_state(state)
|
||||
return trial
|
||||
|
||||
|
||||
class OptimizationTunerTrial(Trial):
|
||||
def __init__(
|
||||
self,
|
||||
config,
|
||||
name,
|
||||
changed_configs,
|
||||
trial_id=None,
|
||||
status=TrialStatus.RUNNING,
|
||||
):
|
||||
super().__init__(config, trial_id, status)
|
||||
self._name = name
|
||||
self._changed_configs = changed_configs
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
return self._name
|
||||
|
||||
def summary(self):
|
||||
spacing = 2
|
||||
max_k = 38
|
||||
max_v = 38
|
||||
|
||||
length = max_k + max_v + spacing
|
||||
|
||||
h1_format = " " + f"|{{:^{length}s}}|\n"
|
||||
h2_format = " " + "|{{:>{}s}}{}{{:^{}s}}|\n".format(
|
||||
max_k, " " * spacing, max_v
|
||||
)
|
||||
|
||||
border = " +" + "".join(["="] * length) + "+"
|
||||
line = " +" + "".join(["-"] * length) + "+"
|
||||
|
||||
draws = border + "\n"
|
||||
draws += h1_format.format("")
|
||||
draws += h1_format.format("Tuned Configurations Overview")
|
||||
draws += h1_format.format("")
|
||||
|
||||
for name in self._changed_configs:
|
||||
draws += border + "\n"
|
||||
draws += h1_format.format(f"{name} auto=True <-> {name}")
|
||||
draws += line + "\n"
|
||||
my_configs = getattr(self.space, name)
|
||||
keys = my_configs.to_dict().keys()
|
||||
for key in keys:
|
||||
draws += h2_format.format(
|
||||
key, str(my_configs.to_dict().get(key, None))
|
||||
)
|
||||
|
||||
result_res = draws + border
|
||||
return result_res
|
||||
|
||||
|
||||
def _generate_trial_id():
|
||||
s = str(time.time()) + str(random.randint(1, int(1e7)))
|
||||
return hashlib.sha256(s.encode("utf-8")).hexdigest()[:32]
|
||||
@@ -0,0 +1,156 @@
|
||||
# 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.
|
||||
|
||||
# Notice that the following codes are modified from KerasTuner to implement our own tuner.
|
||||
# Please refer to https://github.com/keras-team/keras-tuner/blob/master/keras_tuner/engine/hyperparameters.py.
|
||||
|
||||
from .tunable_variable import Boolean, Choice, Fixed, FloatRange, IntRange
|
||||
|
||||
|
||||
class TunableSpace:
|
||||
"""
|
||||
A TunableSpace is constructed by the tunable variables.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# Tunable variables for this tunable variables
|
||||
self._variables = {}
|
||||
# Specific values corresponding to each tunable variable
|
||||
self._values = {}
|
||||
|
||||
@property
|
||||
def variables(self):
|
||||
return self._variables
|
||||
|
||||
@variables.setter
|
||||
def variables(self, variables):
|
||||
self._variables = variables
|
||||
|
||||
@property
|
||||
def values(self):
|
||||
return self._values
|
||||
|
||||
@values.setter
|
||||
def values(self, values):
|
||||
self._values = values
|
||||
|
||||
def get_value(self, name):
|
||||
if name in self.values:
|
||||
return self.values[name]
|
||||
else:
|
||||
raise KeyError(f"{name} does not exist.")
|
||||
|
||||
def set_value(self, name, value):
|
||||
if name in self.values:
|
||||
self.values[name] = value
|
||||
else:
|
||||
raise KeyError(f"{name} does not exist.")
|
||||
|
||||
def _exists(self, name):
|
||||
if name in self._variables:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _retrieve(self, tv):
|
||||
tv = tv.__class__.from_state(tv.get_state())
|
||||
if self._exists(tv.name):
|
||||
return self.get_value(tv.name)
|
||||
return self._register(tv)
|
||||
|
||||
def _register(self, tv):
|
||||
self._variables[tv.name] = tv
|
||||
if tv.name not in self.values:
|
||||
self.values[tv.name] = tv.default
|
||||
return self.values[tv.name]
|
||||
|
||||
def __getitem__(self, name):
|
||||
return self.get_value(name)
|
||||
|
||||
def __setitem__(self, name, value):
|
||||
self.set_value(name, value)
|
||||
|
||||
def __contains__(self, name):
|
||||
try:
|
||||
self.get_value(name)
|
||||
return True
|
||||
except (KeyError, ValueError):
|
||||
return False
|
||||
|
||||
def fixed(self, name, default):
|
||||
tv = Fixed(name=name, default=default)
|
||||
return self._retrieve(tv)
|
||||
|
||||
def boolean(self, name, default=False):
|
||||
tv = Boolean(name=name, default=default)
|
||||
return self._retrieve(tv)
|
||||
|
||||
def choice(self, name, values, default=None):
|
||||
tv = Choice(name=name, values=values, default=default)
|
||||
return self._retrieve(tv)
|
||||
|
||||
def int_range(self, name, start, stop, step=1, default=None):
|
||||
tv = IntRange(
|
||||
name=name, start=start, stop=stop, step=step, default=default
|
||||
)
|
||||
return self._retrieve(tv)
|
||||
|
||||
def float_range(self, name, start, stop, step=None, default=None):
|
||||
tv = FloatRange(
|
||||
name=name, start=start, stop=stop, step=step, default=default
|
||||
)
|
||||
return self._retrieve(tv)
|
||||
|
||||
def get_state(self):
|
||||
return {
|
||||
"variables": [
|
||||
{"class_name": v.__class__.__name__, "state": v.get_state()}
|
||||
for v in self._variables.values()
|
||||
],
|
||||
"values": dict(self.values.items()),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_state(cls, state):
|
||||
ts = cls()
|
||||
for v in state["variables"]:
|
||||
v = _deserialize_tunable_variable(v)
|
||||
ts._variables[v.name] = v
|
||||
ts._values = dict(state["values"].items())
|
||||
return ts
|
||||
|
||||
|
||||
def _deserialize_tunable_variable(state):
|
||||
classes = (Boolean, Fixed, Choice, IntRange, FloatRange)
|
||||
cls_name_to_cls = {cls.__name__: cls for cls in classes}
|
||||
|
||||
if isinstance(state, classes):
|
||||
return state
|
||||
|
||||
if (
|
||||
not isinstance(state, dict)
|
||||
or "class_name" not in state
|
||||
or "state" not in state
|
||||
):
|
||||
raise ValueError(
|
||||
f"Expect state to be a python dict containing class_name and state as keys, but found {state}"
|
||||
)
|
||||
|
||||
cls_name = state["class_name"]
|
||||
cls = cls_name_to_cls[cls_name]
|
||||
if cls is None:
|
||||
raise ValueError(f"Unknown class name {cls_name}")
|
||||
|
||||
cls_state = state["state"]
|
||||
deserialized_object = cls.from_state(cls_state)
|
||||
return deserialized_object
|
||||
@@ -0,0 +1,240 @@
|
||||
# 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.
|
||||
|
||||
# Notice that the following codes are modified from KerasTuner to implement our own tuner.
|
||||
# Please refer to https://github.com/keras-team/keras-tuner/blob/master/keras_tuner/engine/hyperparameters.py.
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class TunableVariable:
|
||||
"""
|
||||
TunableVariable base class.
|
||||
"""
|
||||
|
||||
def __init__(self, name, default=None):
|
||||
self.name = name
|
||||
self._default = default
|
||||
|
||||
@property
|
||||
def default(self):
|
||||
return self._default
|
||||
|
||||
def get_state(self):
|
||||
return {"name": self.name, "default": self.default}
|
||||
|
||||
@classmethod
|
||||
def from_state(cls, state):
|
||||
return cls(**state)
|
||||
|
||||
|
||||
class Fixed(TunableVariable):
|
||||
"""
|
||||
Fixed variable which cannot be changed.
|
||||
"""
|
||||
|
||||
def __init__(self, name, default):
|
||||
super().__init__(name=name, default=default)
|
||||
self.name = name
|
||||
if not isinstance(default, (str, int, float, bool)):
|
||||
raise ValueError(
|
||||
f"Fixed must be an str, int, float or bool, but found {default}"
|
||||
)
|
||||
self._default = default
|
||||
|
||||
def random(self, seed=None):
|
||||
return self._default
|
||||
|
||||
def __repr__(self):
|
||||
return f"Fixed(name: {self.name}, value: {self.default})"
|
||||
|
||||
|
||||
class Boolean(TunableVariable):
|
||||
"""
|
||||
Choice between True and False.
|
||||
"""
|
||||
|
||||
def __init__(self, name, default=False):
|
||||
super().__init__(name=name, default=default)
|
||||
if default not in {True, False}:
|
||||
raise ValueError(
|
||||
f"default must be a Python boolean, but got {default}"
|
||||
)
|
||||
|
||||
def random(self, seed=None):
|
||||
rng = np.random.default_rng(seed)
|
||||
return rng.choice((True, False))
|
||||
|
||||
def __repr__(self):
|
||||
return f'Boolean(name: "{self.name}", default: {self.default})'
|
||||
|
||||
|
||||
class Choice(TunableVariable):
|
||||
def __init__(self, name, values, default=None):
|
||||
super().__init__(name=name, default=default)
|
||||
|
||||
types = {type(v) for v in values}
|
||||
if len(types) > 1:
|
||||
raise TypeError(
|
||||
f"Choice can contain only one type of value, but found values: {values} with types: {types}."
|
||||
)
|
||||
self._is_unknown_type = False
|
||||
|
||||
if isinstance(values[0], str):
|
||||
values = [str(v) for v in values]
|
||||
if default is not None:
|
||||
default = str(default)
|
||||
elif isinstance(values[0], int):
|
||||
values = [int(v) for v in values]
|
||||
if default is not None:
|
||||
default = int(default)
|
||||
elif isinstance(values[0], float):
|
||||
values = [float(v) for v in values]
|
||||
if default is not None:
|
||||
default = float(default)
|
||||
elif isinstance(values[0], bool):
|
||||
values = [bool(v) for v in values]
|
||||
if default is not None:
|
||||
default = bool(default)
|
||||
else:
|
||||
self._is_unknown_type = True
|
||||
self._indices = list(range(len(values)))
|
||||
self.values = values
|
||||
|
||||
if default is not None and default not in values:
|
||||
raise ValueError(
|
||||
f"The default value should be one of the choices {values}, but found {default}"
|
||||
)
|
||||
self._default = default
|
||||
|
||||
@property
|
||||
def default(self):
|
||||
if self._default is None:
|
||||
if None in self.values:
|
||||
return None
|
||||
return self.values[0]
|
||||
return self._default
|
||||
|
||||
def random(self, seed=None):
|
||||
rng = np.random.default_rng(seed)
|
||||
if self._is_unknown_type:
|
||||
indice = rng.choice(self._indices)
|
||||
return self.values[indice]
|
||||
else:
|
||||
return rng.choice(self.values)
|
||||
|
||||
def get_state(self):
|
||||
state = super().get_state()
|
||||
state["values"] = self.values
|
||||
return state
|
||||
|
||||
def __repr__(self):
|
||||
return f'Choice(name: "{self.name}", values: {self.values}, default: {self.default})'
|
||||
|
||||
|
||||
class IntRange(TunableVariable):
|
||||
"""
|
||||
Integer range.
|
||||
"""
|
||||
|
||||
def __init__(self, name, start, stop, step=1, default=None, endpoint=False):
|
||||
super().__init__(name=name, default=default)
|
||||
self.start = self._check_int(start)
|
||||
self.stop = self._check_int(stop)
|
||||
self.step = self._check_int(step)
|
||||
self._default = default
|
||||
self.endpoint = endpoint
|
||||
|
||||
@property
|
||||
def default(self):
|
||||
if self._default is not None:
|
||||
return self._default
|
||||
return self.start
|
||||
|
||||
def random(self, seed=None):
|
||||
rng = np.random.default_rng(seed)
|
||||
value = (self.stop - self.start) * rng.random() + self.start
|
||||
if self.step is not None:
|
||||
if self.endpoint:
|
||||
values = np.arange(self.start, self.stop + 1e-7, step=self.step)
|
||||
else:
|
||||
values = np.arange(self.start, self.stop, step=self.step)
|
||||
closest_index = np.abs(values - value).argmin()
|
||||
value = values[closest_index]
|
||||
return int(value)
|
||||
|
||||
def get_state(self):
|
||||
state = super().get_state()
|
||||
state["start"] = self.start
|
||||
state["stop"] = self.stop
|
||||
state["step"] = self.step
|
||||
state["default"] = self._default
|
||||
return state
|
||||
|
||||
def _check_int(self, val):
|
||||
int_val = int(val)
|
||||
if int_val != val:
|
||||
raise ValueError(f"Expects val is an int, but found: {val}.")
|
||||
return int_val
|
||||
|
||||
def __repr__(self):
|
||||
return f"IntRange(name: {self.name}, start: {self.start}, stop: {self.stop}, step: {self.step}, default: {self.default})"
|
||||
|
||||
|
||||
class FloatRange(TunableVariable):
|
||||
"""
|
||||
Float range.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, name, start, stop, step=None, default=None, endpoint=False
|
||||
):
|
||||
super().__init__(name=name, default=default)
|
||||
self.stop = float(stop)
|
||||
self.start = float(start)
|
||||
if step is not None:
|
||||
self.step = float(step)
|
||||
else:
|
||||
self.step = None
|
||||
self._default = default
|
||||
self.endpoint = endpoint
|
||||
|
||||
@property
|
||||
def default(self):
|
||||
if self._default is not None:
|
||||
return self._default
|
||||
return self.start
|
||||
|
||||
def random(self, seed=None):
|
||||
rng = np.random.default_rng(seed)
|
||||
value = (self.stop - self.start) * rng.random() + self.start
|
||||
if self.step is not None:
|
||||
if self.endpoint:
|
||||
values = np.arange(self.start, self.stop + 1e-7, step=self.step)
|
||||
else:
|
||||
values = np.arange(self.start, self.stop, step=self.step)
|
||||
closest_index = np.abs(values - value).argmin()
|
||||
value = values[closest_index]
|
||||
return value
|
||||
|
||||
def get_state(self):
|
||||
state = super().get_state()
|
||||
state["start"] = self.start
|
||||
state["stop"] = self.stop
|
||||
state["step"] = self.step
|
||||
state["endpoint"] = self.endpoint
|
||||
return state
|
||||
|
||||
def __repr__(self):
|
||||
return f"FloatRange(name: {self.name}, start: {self.start}, stop: {self.stop}, step: {self.step}, default: {self.default}, endpoint: {self.endpoint})"
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user