chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,604 @@
|
||||
# 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 argparse
|
||||
import math
|
||||
import re
|
||||
import shutil
|
||||
from collections import OrderedDict
|
||||
|
||||
import paddle
|
||||
|
||||
|
||||
class ParallelConfig:
|
||||
def __init__(self, mp: int, pp: int, vpp: int = 1, sharding: int = 1):
|
||||
self.mp = mp
|
||||
self.pp = pp
|
||||
self.vpp = vpp
|
||||
self.sharding = sharding
|
||||
|
||||
def pipe_parallel_group(self, i: int, j: int):
|
||||
ans = []
|
||||
for k in range(self.pp):
|
||||
ans.append((i, j, k))
|
||||
return ans
|
||||
|
||||
|
||||
class LayerReNamingHelper:
|
||||
def __init__(self, template: str):
|
||||
self._template = template
|
||||
self._i = -1
|
||||
self._last_old_layer_name = None
|
||||
|
||||
def get_new_layer_name(self, old_layer_name: str):
|
||||
old_layer_name = old_layer_name.split(".")[0]
|
||||
if (
|
||||
self._last_old_layer_name is None
|
||||
or old_layer_name != self._last_old_layer_name
|
||||
):
|
||||
self._i = self._i + 1
|
||||
self._last_old_layer_name = old_layer_name
|
||||
return self._template.format(self._i)
|
||||
|
||||
|
||||
class LayerReNamingManager:
|
||||
def __init__(self):
|
||||
self._renaming_helpers = OrderedDict()
|
||||
self._renaming_helpers["linear"] = LayerReNamingHelper("linear_{}")
|
||||
self._renaming_helpers["layer_norm"] = LayerReNamingHelper(
|
||||
"layer_norm_{}"
|
||||
)
|
||||
self._renaming_helpers["embedding"] = LayerReNamingHelper(
|
||||
"embedding_{}"
|
||||
)
|
||||
|
||||
def get_new_layer_name(self, old_name: str):
|
||||
layer_name = ""
|
||||
for k, v in self._renaming_helpers.items():
|
||||
if old_name.startswith(k):
|
||||
layer_name = v.get_new_layer_name(old_name)
|
||||
break
|
||||
return layer_name
|
||||
|
||||
def get_new_param_name(self, old_name: str):
|
||||
names = old_name.split(".")
|
||||
layer_name = self.get_new_layer_name(names[0])
|
||||
assert layer_name, f"can not rename layer {names[0]}"
|
||||
names[0] = layer_name
|
||||
return ".".join(names)
|
||||
|
||||
|
||||
class PipeLineModelAdaptor:
|
||||
def __init__(
|
||||
self,
|
||||
src_parallel_config: ParallelConfig,
|
||||
dst_parallel_config: ParallelConfig,
|
||||
transformer_layer_num: int,
|
||||
segment_method: str = "layer",
|
||||
):
|
||||
self._src_parallel_config = src_parallel_config
|
||||
self._dst_parallel_config = dst_parallel_config
|
||||
self._transformer_layer_num = transformer_layer_num
|
||||
self._segment_method = segment_method
|
||||
|
||||
def apply(self, src_model_path: str, dst_model_path: str):
|
||||
for i in range(self._src_parallel_config.mp):
|
||||
for j in range(self._src_parallel_config.sharding):
|
||||
# TODO(liuzhenhai): use multiple process
|
||||
layers = []
|
||||
|
||||
# 1、extract layers in the same pp group
|
||||
group = self._src_parallel_config.pipe_parallel_group(i, j)
|
||||
src_dirs = [
|
||||
"{}/mp_{:0>2d}_sharding_{:0>2d}_pp_{:0>2d}".format(
|
||||
src_model_path, *e
|
||||
)
|
||||
for e in group
|
||||
]
|
||||
# first rank extract shared layer
|
||||
with_shared = True
|
||||
for dir in src_dirs:
|
||||
print(f"extract layer params in dir {dir}")
|
||||
layers.extend(self.extract_layers(dir, with_shared))
|
||||
with_shared = False
|
||||
# 2、sort and unique layers
|
||||
layers = self.sort_layers(layers)
|
||||
|
||||
# 3、resplit layers among pp group according new pp config
|
||||
layer_segments = self.segment_layers(
|
||||
layers, self._dst_parallel_config, self._segment_method
|
||||
)
|
||||
dst_group = self._dst_parallel_config.pipe_parallel_group(i, j)
|
||||
dst_dirs = [
|
||||
"{}/mp_{:0>2d}_sharding_{:0>2d}_pp_{:0>2d}".format(
|
||||
dst_model_path, *e
|
||||
)
|
||||
for e in dst_group
|
||||
]
|
||||
|
||||
# 4、merge layers belonging to the same node
|
||||
for layer_segment, dir_ in zip(layer_segments, dst_dirs):
|
||||
print(f"merge {len(layer_segment)} layers to {dir_}")
|
||||
self.merge_layers(layer_segment, dir_)
|
||||
|
||||
# 5、copy meta_state.pdopt
|
||||
for src_dir, dst_dir in zip(src_dirs, dst_dirs):
|
||||
shutil.copyfile(
|
||||
f"{src_dir}/meta_state.pdopt",
|
||||
f"{dst_dir}/meta_state.pdopt",
|
||||
)
|
||||
|
||||
def peek_model(self, model_dir: str):
|
||||
for i in range(self._src_parallel_config.mp):
|
||||
for j in range(self._src_parallel_config.sharding):
|
||||
group = self._src_parallel_config.pipe_parallel_group(i, j)
|
||||
dirs = [
|
||||
"{}/mp_{:0>2d}_sharding_{:0>2d}_pp_{:0>2d}".format(
|
||||
model_dir, *e
|
||||
)
|
||||
for e in group
|
||||
]
|
||||
for dir in dirs:
|
||||
print(f"peek partial model in {dir}:")
|
||||
self.peek_partial_model(dir)
|
||||
|
||||
def peek_partial_model(self, sub_dir: str):
|
||||
state_dict = paddle.load(f"{sub_dir}/model.pdparams")
|
||||
for k, v in state_dict.items():
|
||||
print(f"\t{k} -> {v.name}")
|
||||
|
||||
def extract_layers(self, dir: str, with_shared: bool):
|
||||
opt = paddle.load(dir + "/model_state.pdopt")
|
||||
params = paddle.load(dir + "/model.pdparams")
|
||||
shared_layer_parsed = False
|
||||
# tname -> (layer, param_name)
|
||||
tname_to_layer_and_pname = {}
|
||||
for k, v in params.items():
|
||||
layer = self._extract_layer_name(k)
|
||||
assert layer
|
||||
# special treatment for embedding layer, skip duplicated shared layer
|
||||
# shared layer may exist or not, if it exist it share weight with _layers.0
|
||||
# _layers.shared_layers.embed.word_embeddings.weight -> embedding_0.w_0
|
||||
# _layers.shared_layers.embed.position_embeddings.weight -> embedding_1.w_0
|
||||
# _layers.0.word_embeddings.weight -> embedding_0.w_0
|
||||
# _layers.0.position_embeddings.weight -> embedding_1.w_0
|
||||
shared_layer_parsed = shared_layer_parsed or (
|
||||
"_layers.shared_layers" in layer
|
||||
)
|
||||
if (
|
||||
"_layers.shared_layers" not in layer
|
||||
and ("word_embeddings" in k or "position_embeddings" in k)
|
||||
and shared_layer_parsed
|
||||
):
|
||||
continue
|
||||
tname_to_layer_and_pname[v.name] = (layer, k)
|
||||
|
||||
# get opt-> param mapping
|
||||
tensor_names = list(tname_to_layer_and_pname.keys())
|
||||
opt_names = [
|
||||
e for e in opt.keys() if e not in ["master_weights", "LR_Scheduler"]
|
||||
]
|
||||
opt_to_t = self._opt_name_to_tname(tensor_names, opt_names)
|
||||
# gather tensors belonging to one layer together
|
||||
layers = OrderedDict()
|
||||
for k, v in params.items():
|
||||
layer, p = tname_to_layer_and_pname[v.name]
|
||||
if layer not in layers:
|
||||
layers[layer] = {}
|
||||
layers[layer]["opt"] = OrderedDict()
|
||||
layers[layer]["params"] = OrderedDict()
|
||||
layers[layer]["master_weights"] = OrderedDict()
|
||||
layers[layer]["params"][p] = v
|
||||
|
||||
for k, v in opt.items():
|
||||
if k in ["master_weights", "LR_Scheduler"]:
|
||||
continue
|
||||
layer, _ = tname_to_layer_and_pname[opt_to_t[v.name]]
|
||||
layers[layer]["opt"][k] = v
|
||||
|
||||
if "master_weights" in opt:
|
||||
for k, v in opt["master_weights"].items():
|
||||
layer, _ = tname_to_layer_and_pname[k]
|
||||
layers[layer]["master_weights"][k] = v
|
||||
|
||||
if "LR_Scheduler" in opt:
|
||||
for layer in layers:
|
||||
layers[layer]["LR_Scheduler"] = opt["LR_Scheduler"]
|
||||
|
||||
ans = []
|
||||
|
||||
for layer_name, layer in layers.items():
|
||||
# special treatment for embedding layer
|
||||
if (not with_shared) and "shared_layers" in layer_name:
|
||||
continue
|
||||
file_name = f"./tmp_layer_files/{layer_name}.tmp"
|
||||
paddle.save(layer, file_name)
|
||||
ans.append((layer_name, file_name))
|
||||
print(f"save layer {layer_name} to {file_name}")
|
||||
return ans
|
||||
|
||||
def sort_layers(self, layers: list):
|
||||
def priority(elem):
|
||||
layer_name = elem[0]
|
||||
if "shared_layers" in layer_name:
|
||||
return -0.5
|
||||
match = re.search(
|
||||
r"^_layers((\.\d+)+|(\.shared_layers\.[^\.]+))", layer_name
|
||||
)
|
||||
assert match, f"{layer_name} not a valid layer name"
|
||||
return float(match.group(1).lstrip("."))
|
||||
|
||||
# strictly sort layers
|
||||
print("before sort {}".format("|".join([e[0] for e in layers])))
|
||||
layers.sort(key=priority)
|
||||
# unique
|
||||
unique_layers = []
|
||||
for e in layers:
|
||||
if unique_layers and e[0] == unique_layers[-1][0]:
|
||||
continue
|
||||
unique_layers.append(e)
|
||||
print("after sort {} ".format("|".join([e[0] for e in unique_layers])))
|
||||
return unique_layers
|
||||
|
||||
def segment_layers(
|
||||
self,
|
||||
layers: list,
|
||||
config: ParallelConfig,
|
||||
segment_method: str = "layer",
|
||||
):
|
||||
layer_num = len(layers)
|
||||
stage_num = config.pp * config.vpp
|
||||
|
||||
# segment by weights
|
||||
def segment_by_layer():
|
||||
# assume model is of the structure below
|
||||
# embedding -> n*(transformer layer) -> [optional output layer]
|
||||
# segment index
|
||||
weights = [0 for _ in range(layer_num)]
|
||||
non_zero_layers = range(1, layer_num - 1)
|
||||
# input layer is embedding
|
||||
if self._transformer_layer_num:
|
||||
assert self._transformer_layer_num < layer_num
|
||||
non_zero_layers = range(1, 1 + self._transformer_layer_num)
|
||||
for i in non_zero_layers:
|
||||
weights[i] = 1
|
||||
|
||||
part_size = sum(weights) // stage_num
|
||||
result = [0 for _ in range(stage_num + 1)]
|
||||
memory_counter = 0
|
||||
result_idx = 1
|
||||
for idx, weight in enumerate(weights):
|
||||
memory_counter += weight
|
||||
if memory_counter == part_size:
|
||||
result[result_idx] = idx + 1
|
||||
result_idx += 1
|
||||
memory_counter = 0
|
||||
result[stage_num] = layer_num
|
||||
return result
|
||||
|
||||
def segment_uniform():
|
||||
result = [0 for _ in range(stage_num + 1)]
|
||||
part_size = math.floor(layer_num / stage_num)
|
||||
extra_layers = layer_num % stage_num
|
||||
for i in range(1, stage_num):
|
||||
offset = 1 if i > (stage_num - extra_layers) else 0
|
||||
result[i] = int(
|
||||
min(result[i - 1] + part_size + offset, layer_num)
|
||||
)
|
||||
result[stage_num] = layer_num
|
||||
return result
|
||||
|
||||
result = (
|
||||
segment_uniform()
|
||||
if (segment_method == "uniform")
|
||||
else segment_by_layer()
|
||||
)
|
||||
index_segments = [[] for _ in range(config.pp)]
|
||||
for i in range(stage_num):
|
||||
index_segments[i % config.pp].append((result[i], result[i + 1]))
|
||||
|
||||
# name layers
|
||||
segments = [[] for i in range(config.pp)]
|
||||
for i in range(config.pp):
|
||||
for start, end in index_segments[i]:
|
||||
for j in range(start, end):
|
||||
if config.vpp > 1:
|
||||
segments[i].append(
|
||||
(
|
||||
[f"_layers.{start}.{j - start}"],
|
||||
layers[j][1],
|
||||
)
|
||||
)
|
||||
else:
|
||||
segments[i].append(([f"_layers.{j}"], layers[j][1]))
|
||||
|
||||
shared_layer_exist = any(
|
||||
"_layers.shared_layers" in e[0] for e in layers
|
||||
)
|
||||
if shared_layer_exist:
|
||||
# special treatment for shared layer
|
||||
if config.vpp > 1:
|
||||
segments[0] = [
|
||||
([layers[0][0], segments[0][0][0][0]], layers[0][1]),
|
||||
*segments[0][1:],
|
||||
]
|
||||
else:
|
||||
segments[0] = [([layers[0][0]], layers[0][1]), *segments[0][1:]]
|
||||
|
||||
for i in range(1, config.pp):
|
||||
segments[i] = [([layers[0][0]], layers[0][1])] + segments[i]
|
||||
|
||||
for pp_rank, segs in enumerate(segments):
|
||||
print(f"segment result for pp_rank {pp_rank}:")
|
||||
print(50 * "=")
|
||||
for seg in segs:
|
||||
print(f"{seg[0]} => {seg[1]}")
|
||||
return segments
|
||||
|
||||
def merge_layers(self, layers_segment: list, save_dir: str):
|
||||
params = OrderedDict()
|
||||
opt = OrderedDict()
|
||||
master_weights = OrderedDict()
|
||||
renaming_manager = LayerReNamingManager()
|
||||
|
||||
def merge(src, dst, map_k=None):
|
||||
for k, v in src.items():
|
||||
k = map_k(k) if map_k is not None else k
|
||||
dst[k] = v
|
||||
|
||||
lr_scheduler = None
|
||||
for layer_names, file_path in layers_segment:
|
||||
print(f"load {file_path}")
|
||||
layer = paddle.load(file_path)
|
||||
|
||||
def get_param_name_mapper(layer_name):
|
||||
# replace layer name
|
||||
def map_param_name(param_name):
|
||||
layer_pre = self._extract_layer_name(param_name)
|
||||
return layer_name + param_name[len(layer_pre) :]
|
||||
|
||||
return map_param_name
|
||||
|
||||
(
|
||||
layer_params,
|
||||
layer_opt,
|
||||
layer_master_weight,
|
||||
) = self._map_tensor_names(
|
||||
layer["params"],
|
||||
layer["opt"],
|
||||
layer["master_weights"],
|
||||
renaming_manager,
|
||||
)
|
||||
for layer_name in layer_names:
|
||||
merge(layer_params, params, get_param_name_mapper(layer_name))
|
||||
merge(layer_opt, opt)
|
||||
merge(layer_master_weight, master_weights)
|
||||
lr_scheduler = layer["LR_Scheduler"]
|
||||
|
||||
opt = self._pack_opt_state_dict(opt, master_weights, lr_scheduler)
|
||||
paddle.save(params, save_dir + "/model.pdparams")
|
||||
paddle.save(opt, save_dir + "/model_state.pdopt")
|
||||
|
||||
def _pack_opt_state_dict(self, opt, master_weights, lr_scheduler):
|
||||
opt["master_weights"] = master_weights
|
||||
opt["LR_Scheduler"] = lr_scheduler
|
||||
return opt
|
||||
|
||||
def _extract_layer_name(self, param_name: str):
|
||||
match = re.search(
|
||||
r"^_layers((\.\d+)+|(\.shared_layers\.[^\.]+))", param_name
|
||||
)
|
||||
layer_name = ""
|
||||
return "" if (not match) else match.group()
|
||||
|
||||
# map opt names to tensor name
|
||||
def _opt_name_to_tname(self, tensor_names, opt_names):
|
||||
tensor_names = set(tensor_names)
|
||||
all_names = []
|
||||
all_names.extend(list(tensor_names))
|
||||
all_names.extend(opt_names)
|
||||
all_names.sort()
|
||||
pre_t_name = ""
|
||||
opt_to_t = {}
|
||||
for n in all_names:
|
||||
if n in tensor_names:
|
||||
# we get a param
|
||||
pre_t_name = n
|
||||
else:
|
||||
assert pre_t_name
|
||||
opt_to_t[n] = pre_t_name
|
||||
return opt_to_t
|
||||
|
||||
def _map_tensor_names(self, params, opt, master_weights, renaming_manager):
|
||||
opt_renamed = OrderedDict()
|
||||
master_weights_renamed = OrderedDict()
|
||||
# old name to new name
|
||||
t_name_mapping = {}
|
||||
# map tensor names
|
||||
for k, v in params.items():
|
||||
t_name_mapping[v.name] = renaming_manager.get_new_param_name(v.name)
|
||||
v.name = t_name_mapping[v.name]
|
||||
# map opt names
|
||||
opt_to_tname = self._opt_name_to_tname(
|
||||
t_name_mapping.keys(), opt.keys()
|
||||
)
|
||||
for k, v in opt.items():
|
||||
old_t_name = opt_to_tname[k]
|
||||
t_name = t_name_mapping[old_t_name]
|
||||
opt_name = t_name + k[len(old_t_name) :]
|
||||
v.name = opt_name
|
||||
opt_renamed[opt_name] = v
|
||||
|
||||
# map master names
|
||||
for k, v in master_weights.items():
|
||||
t_name = t_name_mapping[k]
|
||||
v.name = t_name + v.name[len(k) :]
|
||||
master_weights_renamed[t_name] = v
|
||||
return (params, opt_renamed, master_weights_renamed)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
prog='model converter', description='converter a model'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--src_path',
|
||||
type=str,
|
||||
default="./output/epoch_0_step_30",
|
||||
help='path of the model to convert',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--dst_path',
|
||||
type=str,
|
||||
default="./test_adapt",
|
||||
help='path to saved the converted model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--src_mp',
|
||||
type=int,
|
||||
default=2,
|
||||
help='mp degree of the origin training task that dumped this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--src_pp',
|
||||
type=int,
|
||||
default=2,
|
||||
help='pp degree of the origin training task that dumped this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--src_vp',
|
||||
type=int,
|
||||
default=2,
|
||||
help='vp degree of the origin training task that dumped this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--dst_mp',
|
||||
type=int,
|
||||
default=None,
|
||||
help='mp degree of the origin training task that dumped this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--dst_pp',
|
||||
type=int,
|
||||
default=None,
|
||||
help='pp degree of the expected training task that would recover this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--dst_vp',
|
||||
type=int,
|
||||
default=2,
|
||||
help='vp degree of the expected training task that would recover this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--sharding',
|
||||
type=int,
|
||||
default=1,
|
||||
help=" sharding degree of both the origin training task that dumped this model and the expected training task that would recover this model",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--method',
|
||||
type=str,
|
||||
default="adapt_model",
|
||||
help='vp degree of the expected training task that would recover this model',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--segment_method',
|
||||
type=str,
|
||||
default="layer",
|
||||
help='method to segment layers to pp or vp stages',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--transformer_layer_num',
|
||||
type=int,
|
||||
default=0,
|
||||
help='transformer_layer_num of the model',
|
||||
)
|
||||
# assume model is of the structure below
|
||||
# embedding -> n*[transformer layer] -> optional output layer
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.dst_mp is None:
|
||||
args.dst_mp = args.src_mp
|
||||
if args.dst_pp is None:
|
||||
args.dst_pp = args.src_pp
|
||||
|
||||
assert args.src_mp == args.dst_mp, (
|
||||
f"src mp {args.src_mp} dst mp {args.dst_mp}"
|
||||
)
|
||||
|
||||
assert args.method in [
|
||||
'peek_model',
|
||||
'adapt_model',
|
||||
], "method should be in ['peek_model', 'adapt_model']"
|
||||
assert args.segment_method in [
|
||||
"uniform",
|
||||
"layer",
|
||||
], "segment_method should be 'uniform' or 'layer"
|
||||
|
||||
print(
|
||||
f"adapt model dumped by task with pp degree:{args.src_pp}, vp degree:{args.src_vp}, mp degree:{args.src_mp} to task with pp degree:{args.dst_pp}, vp degree:{args.dst_vp}, mp degree:{args.dst_mp}"
|
||||
)
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def adaptor_from_args(args):
|
||||
src_parallel_config = ParallelConfig(
|
||||
args.src_mp, args.src_pp, args.src_vp, args.sharding
|
||||
)
|
||||
|
||||
dst_parallel_config = ParallelConfig(
|
||||
args.dst_mp, args.dst_pp, args.dst_vp, args.sharding
|
||||
)
|
||||
|
||||
adaptor = PipeLineModelAdaptor(
|
||||
src_parallel_config,
|
||||
dst_parallel_config,
|
||||
args.transformer_layer_num,
|
||||
args.segment_method,
|
||||
)
|
||||
return adaptor
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
adaptor = adaptor_from_args(args)
|
||||
if args.method == "peek_model":
|
||||
adaptor.peek_model(args.dst_path)
|
||||
elif args.method == "adapt_model":
|
||||
adaptor.apply(args.src_path, args.dst_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
"""
|
||||
Usage:
|
||||
python pp_parallel_adaptor.py --src_mp xxx --src_path xxx --method \
|
||||
adapt_model/peek_model --dst_path xxx --sharding xxx --segment_method xxx --transformer_layer_num xxx
|
||||
|
||||
for the meaning of a specific arg, please use:
|
||||
python pp_parallel_adaptor.py -h
|
||||
"""
|
||||
main()
|
||||
Reference in New Issue
Block a user