# 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)