wehub-resource-sync
157 lines
6.7 KiB
Python
157 lines
6.7 KiB
Python
# 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 argparse
|
|
import json
|
|
import os
|
|
|
|
import numpy as np
|
|
import paddle
|
|
|
|
from paddlenlp.generation import GenerationConfig
|
|
from paddlenlp.transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
|
|
from paddlenlp.transformers.model_utils import load_tp_checkpoint
|
|
from paddlenlp.trl import llm_utils
|
|
|
|
|
|
def parse_arguments():
|
|
"""
|
|
parse_arguments
|
|
"""
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model_path", default=None, type=str, required=True, help="The directory of model.")
|
|
parser.add_argument("--output_path", default=None, type=str, help="The directory of split model")
|
|
parser.add_argument("--model_rank_id", default=None, type=int, help="Input model mp degree.")
|
|
parser.add_argument("--dtype", default="float16", type=str, help="The dtype of model weights.")
|
|
return parser.parse_args()
|
|
|
|
|
|
def split(args):
|
|
"""
|
|
Split model weight
|
|
"""
|
|
rank, nranks = llm_utils.init_dist_env()
|
|
|
|
if args.output_path is None:
|
|
args.output_path = os.path.join(args.model_path, f"{nranks}_ranks")
|
|
|
|
paddle.set_default_dtype(args.dtype)
|
|
|
|
config = AutoConfig.from_pretrained(args.model_path)
|
|
config.tensor_parallel_degree = nranks
|
|
config.tensor_parallel_rank = rank
|
|
|
|
generation_config = GenerationConfig.from_pretrained(args.model_path)
|
|
model = AutoModelForCausalLM.from_pretrained(args.model_path)
|
|
tokenizer = AutoTokenizer.from_pretrained(args.model_path)
|
|
|
|
if args.model_rank_id is not None:
|
|
model_path = os.path.join(args.model_path, f"model_state.tp0{args.model_rank_id - 1}.pdparams")
|
|
assert os.path.isfile(model_path), f"{model_path} not exist"
|
|
state_dict = load_tp_checkpoint(args.model_path, model, config)
|
|
model_rank = args.model_rank_id
|
|
save_base_rank = model_rank * nranks
|
|
else:
|
|
state_dict = load_tp_checkpoint(args.model_path, model, config)
|
|
model_rank = 0
|
|
save_base_rank = 0
|
|
|
|
weight_file = os.path.join(args.output_path, f"model_state.tp0{rank + save_base_rank}.pdparams")
|
|
paddle.save(state_dict, weight_file)
|
|
|
|
# process weight scales
|
|
possible_weight_scales_path = os.path.join(args.model_path, f"weight_scales_{model_rank}.json")
|
|
if os.path.exists(possible_weight_scales_path) and rank == 0:
|
|
with open(possible_weight_scales_path, "r") as f:
|
|
weight_scales_dict = json.load(f)
|
|
|
|
processed_weight_scales = [{} for i in range(nranks)]
|
|
for k, v in weight_scales_dict.items():
|
|
if "self_attn.q_proj" in k:
|
|
splited_value = np.split(np.array(v), nranks, axis=-1)
|
|
for tp_rank in range(nranks):
|
|
processed_weight_scales[tp_rank][k] = splited_value[tp_rank].tolist()
|
|
elif "self_attn.k_proj" in k:
|
|
splited_value = np.split(np.array(v), nranks, axis=-1)
|
|
for tp_rank in range(nranks):
|
|
processed_weight_scales[tp_rank][k] = splited_value[tp_rank].tolist()
|
|
elif "self_attn.v_proj" in k:
|
|
splited_value = np.split(np.array(v), nranks, axis=-1)
|
|
for tp_rank in range(nranks):
|
|
processed_weight_scales[tp_rank][k] = splited_value[tp_rank].tolist()
|
|
elif "self_attn.o_proj" in k:
|
|
for tp_rank in range(nranks):
|
|
processed_weight_scales[tp_rank][k] = v
|
|
elif "mlp.gate_proj" in k:
|
|
splited_value = np.split(np.array(v), nranks, axis=-1)
|
|
for tp_rank in range(nranks):
|
|
processed_weight_scales[tp_rank][k] = splited_value[tp_rank].tolist()
|
|
elif "mlp.up_proj" in k:
|
|
splited_value = np.split(np.array(v), nranks, axis=-1)
|
|
for tp_rank in range(nranks):
|
|
processed_weight_scales[tp_rank][k] = splited_value[tp_rank].tolist()
|
|
elif "mlp.down_proj" in k:
|
|
for tp_rank in range(nranks):
|
|
processed_weight_scales[tp_rank][k] = v
|
|
else:
|
|
raise ValueError(f"key {k} is not supported!")
|
|
|
|
for tp_rank in range(nranks):
|
|
save_path = os.path.join(args.output_path, f"weight_scales_{tp_rank + save_base_rank}.json")
|
|
with open(save_path, "w") as f:
|
|
print("weight scale save_path:", save_path)
|
|
json.dump(processed_weight_scales[tp_rank], f)
|
|
|
|
# process cachekv scales
|
|
possible_cache_path = os.path.join(args.model_path, f"cachekv_scales_{model_rank}.json")
|
|
if os.path.exists(possible_cache_path) and rank == 0:
|
|
with open(possible_cache_path, "r") as f:
|
|
cache_dict = json.load(f)
|
|
|
|
processed_cachekv_scales = [{} for i in range(nranks)]
|
|
for k, v in cache_dict.items():
|
|
v = np.array(v).flatten()
|
|
splited_value = np.split(np.array(v), nranks, axis=-1)
|
|
for tp_rank in range(nranks):
|
|
processed_cachekv_scales[tp_rank][k] = splited_value[tp_rank].tolist()
|
|
for tp_rank in range(nranks):
|
|
save_path = os.path.join(args.output_path, f"cachekv_scales_{tp_rank + save_base_rank}.json")
|
|
print("cachekv scale save_path:", save_path)
|
|
with open(save_path, "w") as f:
|
|
json.dump(processed_cachekv_scales[tp_rank], f)
|
|
|
|
# process act scales
|
|
possible_act_scales_path = os.path.join(args.model_path, f"act_scales_{model_rank}.json")
|
|
if os.path.exists(possible_act_scales_path) and rank == 0:
|
|
with open(possible_act_scales_path, "r") as f:
|
|
act_scale = json.load(f)
|
|
for tp_rank in range(nranks):
|
|
save_path = os.path.join(args.output_path, f"act_scales_{tp_rank + save_base_rank}.json")
|
|
with open(save_path, "w") as outf:
|
|
print("act scale save_path:", save_path)
|
|
json.dump(act_scale, outf)
|
|
|
|
if rank == 0:
|
|
tokenizer.save_pretrained(args.output_path)
|
|
config.save_pretrained(args.output_path)
|
|
generation_config.save_pretrained(args.output_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
"""
|
|
Script to split model weight.
|
|
"""
|
|
args = parse_arguments()
|
|
split(args)
|