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