mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2026-04-24 01:29:57 +08:00
jiangjiajun
1352 lines
48 KiB
Python
1352 lines
48 KiB
Python
"""
|
|
# 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 __future__ import annotations
|
|
|
|
import collections
|
|
import hashlib
|
|
import json
|
|
import multiprocessing as mp
|
|
import os
|
|
import random
|
|
import re
|
|
import struct
|
|
from functools import partial
|
|
from typing import Callable, Optional
|
|
|
|
import numpy as np
|
|
from paddlenlp.transformers import PretrainedTokenizer
|
|
from paddlenlp.transformers.model_utils import _add_variant
|
|
from paddlenlp.transformers.utils import paddlenlp_load
|
|
from paddlenlp.transformers.model_utils import load_tp_checkpoint
|
|
from safetensors import safe_open
|
|
|
|
from paddlenlp.utils.env import (
|
|
PADDLE_WEIGHTS_INDEX_NAME,
|
|
SAFE_MASTER_WEIGHTS_INDEX_NAME,
|
|
SAFE_PEFT_WEIGHTS_INDEX_NAME,
|
|
SAFE_WEIGHTS_INDEX_NAME,
|
|
)
|
|
from paddlenlp.utils.log import logger
|
|
from tqdm import tqdm
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
from paddle.common_ops_import import convert_dtype
|
|
from paddle.distributed import fleet
|
|
from paddlenlp.transformers import PretrainedTokenizer
|
|
from paddlenlp.transformers.model_utils import _add_variant, load_tp_checkpoint
|
|
from paddlenlp.transformers.utils import paddlenlp_load
|
|
from paddlenlp.utils.env import (PADDLE_WEIGHTS_INDEX_NAME,
|
|
SAFE_MASTER_WEIGHTS_INDEX_NAME,
|
|
SAFE_PEFT_WEIGHTS_INDEX_NAME,
|
|
SAFE_WEIGHTS_INDEX_NAME)
|
|
from paddlenlp.utils.log import logger
|
|
from safetensors import safe_open
|
|
from tqdm import tqdm
|
|
|
|
from fastdeploy.platforms import current_platform
|
|
|
|
from .tokenizer import ErnieBotTokenizer
|
|
import glob
|
|
|
|
MODEL_LIB_NAMES = [
|
|
"ernie_bot.modeling",
|
|
"ernie_bot.modeling_pp",
|
|
"ernie_bot.modeling_moe",
|
|
"ernie_bot.modeling_rm",
|
|
"ernie_bot.proxy_distill",
|
|
]
|
|
|
|
MAX_BSZ = 512
|
|
MAX_DRAFT_TOKENS = 6
|
|
|
|
|
|
class UniqueIDGenerator:
|
|
"""
|
|
The generator for the export model id
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
def generate_unique_id(self, state_dict):
|
|
"""
|
|
Generate the model id from the timestamp
|
|
"""
|
|
keys = state_dict.keys()
|
|
sorted_keys = sorted(keys)
|
|
first_key = sorted_keys[0]
|
|
first_parameter = state_dict[first_key].cast("float32")
|
|
# 假设模型参数是唯一的，通过第一个key来获取md5sum
|
|
model_md5 = hashlib.md5(str(
|
|
first_parameter.sum()).encode("utf-8")).hexdigest()
|
|
unique_id = f"{model_md5}-{random.randint(10000, 99999)}"
|
|
return unique_id
|
|
|
|
|
|
def load_sharded_checkpoint(folder, variant=None, return_numpy=False):
|
|
"""
|
|
|
|
This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being
|
|
loaded in the model.
|
|
|
|
Args:
|
|
folder (`str` or `os.PathLike`): A path to a folder containing the sharded checkpoint.
|
|
variant (`str`): The model variant.
|
|
|
|
"""
|
|
# Load the index
|
|
pdparams_file = os.path.join(folder,
|
|
_add_variant("model_state.pdparams", variant))
|
|
lora_pdparams_file = os.path.join(
|
|
folder, _add_variant("lora_model_state.pdparams", variant))
|
|
safetensors_file = os.path.join(folder,
|
|
_add_variant("model.safetensors", variant))
|
|
if os.path.isfile(pdparams_file):
|
|
return paddle.load(pdparams_file, return_numpy=return_numpy)
|
|
if os.path.isfile(lora_pdparams_file):
|
|
return paddle.load(lora_pdparams_file, return_numpy=return_numpy)
|
|
if os.path.isfile(safetensors_file):
|
|
try:
|
|
from paddlenlp.utils.safetensors import \
|
|
fast_load_file as safe_load_file
|
|
except ImportError:
|
|
from safetensors.numpy import load_file as safe_load_file
|
|
|
|
state_dict = safe_load_file(safetensors_file)
|
|
if not return_numpy:
|
|
for key in list(state_dict.keys()):
|
|
if isinstance(state_dict[key], np.ndarray):
|
|
state_dict[key] = paddle.Tensor(state_dict.pop(key),
|
|
zero_copy=True)
|
|
return state_dict
|
|
|
|
index_file = os.path.join(folder,
|
|
_add_variant(PADDLE_WEIGHTS_INDEX_NAME, variant))
|
|
safe_index_file = os.path.join(
|
|
folder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant))
|
|
safe_master_file = os.path.join(
|
|
folder, _add_variant(SAFE_MASTER_WEIGHTS_INDEX_NAME, variant))
|
|
safe_peft_file = os.path.join(
|
|
folder, _add_variant(SAFE_PEFT_WEIGHTS_INDEX_NAME, variant))
|
|
|
|
index_present = os.path.isfile(index_file)
|
|
safe_index_present = os.path.isfile(safe_index_file)
|
|
safe_master_present = os.path.isfile(safe_master_file)
|
|
safe_peft_present = os.path.isfile(safe_peft_file)
|
|
|
|
load_safe = False
|
|
load_index = None
|
|
if safe_index_present:
|
|
load_safe = True # load safe due to preference
|
|
load_index = safe_index_file
|
|
elif safe_master_present:
|
|
load_safe = True
|
|
load_index = safe_master_file
|
|
elif index_present:
|
|
load_index = index_file
|
|
elif safe_peft_present:
|
|
load_safe = True
|
|
load_index = safe_peft_file
|
|
else:
|
|
raise ValueError(
|
|
f"Could not find {index_file} or {safe_index_file} or {safe_peft_file}"
|
|
)
|
|
|
|
if load_safe:
|
|
try:
|
|
from paddlenlp.utils.safetensors import \
|
|
fast_load_file as safe_load_file
|
|
except ImportError:
|
|
from safetensors.numpy import load_file as safe_load_file
|
|
|
|
with open(load_index, "r", encoding="utf-8") as f:
|
|
index = json.load(f)
|
|
|
|
shard_files = list(set(index["weight_map"].values()))
|
|
loader = (safe_load_file if load_safe else partial(
|
|
paddlenlp_load, map_location="np" if return_numpy else "cpu"))
|
|
|
|
ret = {}
|
|
for shard_file in tqdm(shard_files):
|
|
state_dict = loader(os.path.join(folder, shard_file))
|
|
ret.update(state_dict)
|
|
|
|
if not return_numpy:
|
|
for key in list(ret.keys()):
|
|
if isinstance(ret[key], np.ndarray):
|
|
ret[key] = paddle.Tensor(ret.pop(key), zero_copy=True)
|
|
|
|
return ret
|
|
|
|
|
|
def convert_ndarray_dtype(np_array: np.ndarray,
|
|
target_dtype: str) -> np.ndarray:
|
|
"""convert ndarray
|
|
|
|
Args:
|
|
np_array (np.ndarray): numpy ndarray instance
|
|
target_dtype (str): the target dtype
|
|
|
|
Returns:
|
|
np.ndarray: converted numpy ndarray instance
|
|
"""
|
|
source_dtype = convert_dtype(np_array.dtype)
|
|
if source_dtype == "uint16" or target_dtype == "bfloat16":
|
|
if paddle.is_compiled_with_xpu():
|
|
# xpu not support bf16.
|
|
tensor = paddle.to_tensor(np_array, place=paddle.CPUPlace())
|
|
else:
|
|
tensor = paddle.to_tensor(np_array)
|
|
tensor = paddle.cast(tensor, target_dtype)
|
|
return tensor.numpy()
|
|
|
|
# TODO(wj-Mcat): device_guard will slow the converting
|
|
# with device_guard("cpu"):
|
|
# tensor = paddle.to_tensor(np_array)
|
|
# tensor = paddle.cast(tensor, target_dtype)
|
|
# return tensor.numpy()
|
|
|
|
if target_dtype == "bfloat16":
|
|
target_dtype = "uint16"
|
|
|
|
return np_array.astype(target_dtype)
|
|
|
|
|
|
def ernie_bot_postprocess_past_key_value(past_key_values):
|
|
"""
|
|
ernie_bot_postprocess_past_key_values
|
|
"""
|
|
Cache = collections.namedtuple("Cache", ["k", "v"])
|
|
# (layer_num, bs, prefixlen, head_num/tensor_parallel_degree, head_dim)*2
|
|
keys, values = paddle.transpose(past_key_values, perm=[2, 0, 1, 3,
|
|
4]).split(2)
|
|
|
|
past_key_values = []
|
|
for k, v in zip(keys, values):
|
|
past_key_values.append(Cache(k, v))
|
|
return past_key_values
|
|
|
|
|
|
def ernie_bot_pad_attention_mask(input_ids_shape, num_prefix_tokens,
|
|
attention_mask):
|
|
"""
|
|
ernie_bot_pad_attention_mask
|
|
"""
|
|
if attention_mask.dim() == 2:
|
|
attention_mask = attention_mask[:, None, None, :]
|
|
prefix_attention_mask = paddle.ones(
|
|
[input_ids_shape[0], 1, 1, num_prefix_tokens],
|
|
dtype=attention_mask.dtype,
|
|
)
|
|
else:
|
|
prefix_attention_mask = paddle.ones(
|
|
[input_ids_shape[0], 1, input_ids_shape[-1], num_prefix_tokens],
|
|
dtype=attention_mask.dtype,
|
|
)
|
|
return paddle.concat((prefix_attention_mask, attention_mask), axis=3)
|
|
|
|
|
|
def set_seed(seed: int):
|
|
"""
|
|
set random seed for all random modules
|
|
"""
|
|
paddle.seed(seed)
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
|
|
|
|
def get_infer_model_path(input_dir, model_prefix, is_export: bool = False):
|
|
"""when n_ranks = 1, infer_model_path is: `{input_dir}/{model_prefix}.pdiparams`
|
|
when n_ranks > 1, infer_model_path is: `{input_dir}/rank_{idx}/{model_prefix}.pdiparams`
|
|
|
|
Args:
|
|
input_dir (str): the base input_dir
|
|
model_prefix (str): the prefix name of model
|
|
|
|
Returns:
|
|
str: the path of infer model path
|
|
"""
|
|
n_ranks = dist.get_world_size()
|
|
try:
|
|
local_rank = dist.ParallelEnv().dev_id
|
|
except Exception:
|
|
logger.info(
|
|
"`dist.ParallelEnv().dev_id` is not supported on CPU devices,so set local_rank = 0."
|
|
)
|
|
local_rank = 0
|
|
if n_ranks > 1:
|
|
return os.path.join(input_dir, f"rank_{local_rank}", model_prefix)
|
|
|
|
# if n_ranks director exist, return N-rank directory
|
|
sub_rank_dir = os.path.join(input_dir, f"rank_{local_rank}")
|
|
|
|
if is_export:
|
|
return os.path.join(sub_rank_dir, model_prefix)
|
|
else:
|
|
# when inference, return sub_rank_dir when exists
|
|
if os.path.exists(sub_rank_dir):
|
|
return os.path.join(sub_rank_dir, model_prefix)
|
|
else:
|
|
return os.path.join(input_dir, model_prefix)
|
|
|
|
|
|
def pad_batch_data(insts, pad_id=0, return_seq_len=False, pad_style="right"):
|
|
"""Pad the instances to the max sequence length in batch."""
|
|
# pad to max input len i bsz
|
|
max_len = max(map(len, insts))
|
|
# pad to max input len
|
|
# max_len = args.max_len
|
|
if pad_style == "left":
|
|
inst_data = np.array([[pad_id] * (max_len - len(inst)) + list(inst)
|
|
for inst in insts])
|
|
else:
|
|
inst_data = np.array(
|
|
[list(inst) + [pad_id] * (max_len - len(inst)) for inst in insts])
|
|
if return_seq_len:
|
|
seq_len = np.array([len(inst) for inst in insts])
|
|
return inst_data.astype("int64").reshape([-1, max_len]), seq_len
|
|
else:
|
|
return inst_data.astype("int64").reshape([-1, max_len])
|
|
|
|
|
|
def load_prefix_weights(
|
|
prefix_path: str,
|
|
inference: bool = False,
|
|
batch_size: int = 1,
|
|
dtype: str = "bfloat16",
|
|
) -> np.ndarray | list[paddle.Tensor]:
|
|
"""load prefix weight by path
|
|
|
|
Args:
|
|
prefix_path (str): the path of prefix weight
|
|
"""
|
|
past_key_values = paddle.to_tensor(
|
|
np.load(f"{prefix_path}/pre_caches.npy")).unsqueeze(2)
|
|
|
|
if batch_size > 1:
|
|
past_key_values = paddle.concat([past_key_values] * batch_size, axis=2)
|
|
|
|
# .chatglm static model require one tensor, otherwise list of tensor
|
|
past_key_values = past_key_values.astype(dtype)
|
|
if inference:
|
|
return past_key_values.numpy()
|
|
return past_key_values
|
|
|
|
|
|
def build_for_generation(model, tokenizer: PretrainedTokenizer,
|
|
generation_kwargs: dict):
|
|
"""build `ErnieBotForGenerationFuse` to generate tokens
|
|
|
|
Args:
|
|
model (_type_): ErnieBotModel or ErnieBotFusedModel
|
|
tokenizer (PretrainedTokenizer): pretrained tokenizer
|
|
generation_kwargs (dict): generation_kwargs for model
|
|
|
|
Returns:
|
|
PretrainedModel: ErnieBotForGenerationFuse
|
|
"""
|
|
from ernie_bot.single_model_fused import ErnieBotForGenerationFuse
|
|
|
|
configs = {
|
|
"bos_token_id": tokenizer.bos_token_id,
|
|
"eos_token_id": tokenizer.eos_token_id,
|
|
"pad_token_id": tokenizer.pad_token_id,
|
|
"initializer_range": 0.02,
|
|
"fused_linear": False,
|
|
"min_dec_len": 1,
|
|
"max_dec_len": 1024,
|
|
"top_k": 0,
|
|
"top_p": 0.7,
|
|
"temperature": 0.95,
|
|
"use_topp_sampling": True,
|
|
"inference": True,
|
|
}
|
|
configs.update(generation_kwargs)
|
|
model = ErnieBotForGenerationFuse(model, configs=configs)
|
|
model.eval()
|
|
return model
|
|
|
|
|
|
def init_distributed_env() -> tuple[int, int]:
|
|
"""init distributed envs, and only support mp in ErnieBotModel
|
|
|
|
Returns:
|
|
tuple[int, int]: tensor_parallel_degree, tensor_parallel_rank
|
|
"""
|
|
tensor_parallel_degree = dist.get_world_size()
|
|
tensor_parallel_rank = 0
|
|
|
|
if tensor_parallel_degree > 1:
|
|
strategy = fleet.DistributedStrategy()
|
|
strategy.hybrid_configs = {
|
|
"dp_degree": 1,
|
|
"mp_degree": tensor_parallel_degree,
|
|
"pp_degree": 1,
|
|
"sharding_degree": 1,
|
|
}
|
|
|
|
fleet.init(is_collective=True, strategy=strategy)
|
|
hcg = fleet.get_hybrid_communicate_group()
|
|
tensor_parallel_rank = hcg.get_model_parallel_rank()
|
|
|
|
return tensor_parallel_degree, tensor_parallel_rank
|
|
|
|
|
|
def generate_rank_mapping(output_dir: str):
|
|
"""generate current distributed rank mapping file
|
|
|
|
Args:
|
|
output_dir (str): the directory of rank_mapping file
|
|
"""
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# must in distributed env
|
|
hcg = fleet.get_hybrid_communicate_group()
|
|
model_parallel_group = hcg.get_model_parallel_group()
|
|
ring_id = model_parallel_group.id
|
|
|
|
world_size = dist.get_world_size()
|
|
with open(os.path.join(output_dir, "rank_mapping.csv"), "w") as f:
|
|
f.write("[ring_id -> ranks]\n")
|
|
f.write(",".join(map(str, [0] + list(range(world_size)))) + "\n")
|
|
f.write(",".join(map(str, [ring_id] + list(range(world_size)))) + "\n")
|
|
|
|
f.write("[rank -> ring_ids]\n")
|
|
for i in range(world_size):
|
|
f.write(f"{i},0,{ring_id}\n")
|
|
|
|
|
|
def save_infer_result(trainer, dev_ds, k=100, src_length=256, tgt_length=512):
|
|
"""
|
|
save infer result into jsonl format
|
|
"""
|
|
from predict_generation import Predictor, batchfy_text
|
|
|
|
all_instructions = []
|
|
all_answers = []
|
|
all_output = []
|
|
|
|
# top k instruction from dev_ds
|
|
for i, ds in enumerate(dev_ds.data):
|
|
if i == k:
|
|
break
|
|
if "instruction" in ds:
|
|
all_instructions.append(ds["instruction"])
|
|
all_answers.append(ds["output"])
|
|
elif "src" in ds:
|
|
if isinstance(ds["src"], list):
|
|
all_instructions.append(ds["src"][0])
|
|
all_answers.append(ds["tgt"][0])
|
|
else:
|
|
all_instructions.append(ds["src"])
|
|
all_answers.append(ds["tgt"])
|
|
|
|
batch_texts = batchfy_text(all_instructions,
|
|
trainer.args.per_device_eval_batch_size)
|
|
predictor = Predictor(
|
|
tokenizer=trainer.tokenizer,
|
|
model=trainer.model,
|
|
src_length=src_length,
|
|
tgt_length=tgt_length,
|
|
)
|
|
|
|
# infer results
|
|
for bs, texts in enumerate(batch_texts):
|
|
outputs = predictor.predict(texts)
|
|
for i, (text, result) in enumerate(zip(texts, outputs["result"])):
|
|
out = {
|
|
"instruction":
|
|
text,
|
|
"answer":
|
|
all_answers[bs * trainer.args.per_device_eval_batch_size + i],
|
|
"output":
|
|
result,
|
|
}
|
|
all_output.append(out)
|
|
|
|
# save results
|
|
if trainer.args.tensor_parallel_rank == 0:
|
|
with open(os.path.join(trainer.args.output_dir, "infer_result.json"),
|
|
"w") as f:
|
|
for out in all_output:
|
|
f.write(json.dumps(out, ensure_ascii=False) + "\n")
|
|
|
|
|
|
def w4a8_weight_convert(state_dict):
|
|
"""W4A8 权重转换函数
|
|
Args:
|
|
state_dict (dict): state_dict of model
|
|
"""
|
|
|
|
def w4_weight_squash(value, name, w4a8_weight_bites_name_map):
|
|
weight_dq = value
|
|
# W8表象下的W4权重的absmax值为112，使用正负112进行权重类型判断
|
|
if weight_dq.max() == 112 or weight_dq.min() == -112:
|
|
weight_dq = weight_dq.cast("int8")
|
|
np_weight_dq = np.array(weight_dq, dtype="int8").view("uint8")
|
|
np_weight_dq_left_div_16 = (np_weight_dq / 16).astype("int8")
|
|
# weight_q = (weight_dq/16).cast('int8')
|
|
weight_q = paddle.to_tensor(np_weight_dq_left_div_16, dtype="int8")
|
|
logger.debug(f"int4 weight:{name}")
|
|
w4a8_weight_bites_name_map[name] = 4
|
|
return weight_q.cast("int8")
|
|
elif weight_dq.max() == 127 or weight_dq.min() == -128:
|
|
logger.debug(f"int8 weight:{name}")
|
|
w4a8_weight_bites_name_map[name] = 8
|
|
return weight_dq.cast("int8")
|
|
else:
|
|
logger.debug(f"fp16/bf16/float weight:{name}")
|
|
return weight_dq
|
|
|
|
w4a8_weight_bites_name_map = {}
|
|
for name, value in state_dict.items():
|
|
if value.dtype == "uint16":
|
|
weight_q = w4_weight_squash(
|
|
paddle.to_tensor(value).cast("float32"),
|
|
name,
|
|
w4a8_weight_bites_name_map,
|
|
)
|
|
state_dict[name] = weight_q.numpy(
|
|
) if weight_q is not None else value
|
|
del weight_q
|
|
w4a8_weight_bites_layers_map = {}
|
|
w4a8_weight_bites_layers_map["qkv_gemm_bits_map"] = []
|
|
w4a8_weight_bites_layers_map["out_gemm_bits_map"] = []
|
|
w4a8_weight_bites_layers_map["ffn1_gemm_bits_map"] = []
|
|
w4a8_weight_bites_layers_map["ffn2_gemm_bits_map"] = []
|
|
for name_keys, gemm_bits in w4a8_weight_bites_name_map.items():
|
|
if "qkv_proj" in name_keys:
|
|
w4a8_weight_bites_layers_map["qkv_gemm_bits_map"].append(gemm_bits)
|
|
elif "out_proj" in name_keys:
|
|
w4a8_weight_bites_layers_map["out_gemm_bits_map"].append(gemm_bits)
|
|
elif "linear1" in name_keys:
|
|
w4a8_weight_bites_layers_map["ffn1_gemm_bits_map"].append(
|
|
gemm_bits)
|
|
elif "linear2" in name_keys:
|
|
w4a8_weight_bites_layers_map["ffn2_gemm_bits_map"].append(
|
|
gemm_bits)
|
|
logger.debug(
|
|
f"w4a8_weight_bites_layers_map:{w4a8_weight_bites_layers_map}")
|
|
return state_dict, w4a8_weight_bites_layers_map
|
|
|
|
|
|
def _vocab_size_with_padding(vocab_size, div_unit, mp_degree):
|
|
padded_size = vocab_size
|
|
multiple = div_unit * mp_degree
|
|
while (padded_size % multiple) != 0:
|
|
padded_size += 1
|
|
# logger.warning(
|
|
# " > padded vocab (size: {}) with {} dummy tokens "
|
|
# "(new size: {})".format(vocab_size, padded_size - vocab_size, padded_size)
|
|
# )
|
|
return padded_size
|
|
|
|
|
|
def save_test_case(cases: list[list[dict]], file: str):
|
|
"""save test to result file
|
|
|
|
Args:
|
|
cases (list[list[dict]]): the content of case
|
|
file (str): the path of saved file
|
|
"""
|
|
with open(file, "w+", encoding="utf-8") as f:
|
|
for case in cases:
|
|
raw = json.dumps(case, ensure_ascii=False)
|
|
f.write(raw + "\n")
|
|
|
|
|
|
def infer_save_test_case(cases: list[list[dict]], file: str):
|
|
"""save test to result file
|
|
|
|
Args:
|
|
cases (list[list[dict]]): the content of case
|
|
file (str): the path of saved file
|
|
"""
|
|
with open(file, "a+", encoding="utf-8") as f:
|
|
for case in cases:
|
|
raw = json.dumps(case, ensure_ascii=False)
|
|
f.write(raw + "\n")
|
|
|
|
|
|
def deserialize_from_file(fp):
|
|
"""
|
|
deserialize a binary file into an array
|
|
"""
|
|
|
|
x_type = fp.read(1)
|
|
x_type_out = struct.unpack("c", x_type)[0]
|
|
# data
|
|
data_list = []
|
|
if x_type_out == b"0":
|
|
data = fp.read(4)
|
|
data_out = struct.unpack("f", data)[0]
|
|
while data:
|
|
data_out = struct.unpack("f", data)[0]
|
|
data_list.append(data_out)
|
|
data = fp.read(4)
|
|
elif x_type_out == b"1":
|
|
data = fp.read(8)
|
|
while data:
|
|
data_out = struct.unpack("l", data)[0]
|
|
data_list.append(data_out)
|
|
data = fp.read(8)
|
|
elif x_type_out == b"2":
|
|
data = fp.read(4)
|
|
while data:
|
|
data_out = struct.unpack("i", data)[0]
|
|
data_list.append(data_out)
|
|
data = fp.read(4)
|
|
else:
|
|
print("type error")
|
|
data_arr = np.array(data_list)
|
|
return data_arr
|
|
|
|
|
|
def read_res(
|
|
model_name_or_path,
|
|
output_tensor_max_shape,
|
|
result_queue: mp.Queue,
|
|
msg_queue_id=None,
|
|
use_ep=False,
|
|
ep_just_for_test=False,
|
|
tokenizer=None,
|
|
):
|
|
"""Read result from queue."""
|
|
if msg_queue_id is None:
|
|
if (current_platform.is_cuda() and
|
|
current_platform.available()) or paddle.is_compiled_with_xpu():
|
|
from fastdeploy.model_executor.ops.gpu import get_output
|
|
elif paddle.is_compiled_with_custom_device("npu"):
|
|
from paddle_custom_device.npu import get_output
|
|
else: # CPU
|
|
from fastdeploy.model_executor.ops.cpu import get_output
|
|
else:
|
|
if (current_platform.is_cuda() and
|
|
current_platform.available()) or paddle.is_compiled_with_xpu():
|
|
from fastdeploy.model_executor.ops.gpu import get_output_dynamic
|
|
elif paddle.is_compiled_with_custom_device("npu"):
|
|
from paddle_custom_device.npu import get_output_dynamic
|
|
else: # CPU
|
|
from fastdeploy.model_executor.ops.cpu import get_output_dynamic
|
|
|
|
if tokenizer is None:
|
|
tokenizer = ErnieBotTokenizer.from_pretrained(model_name_or_path)
|
|
|
|
paddle.device.set_device("cpu")
|
|
paddle.disable_static()
|
|
output_tensor = paddle.full(output_tensor_max_shape,
|
|
fill_value=2,
|
|
dtype="int64")
|
|
|
|
while True:
|
|
outputs = []
|
|
while True:
|
|
if msg_queue_id is None:
|
|
get_output(output_tensor, 0, True)
|
|
else:
|
|
get_output_dynamic(output_tensor, 0, True, msg_queue_id)
|
|
if int(output_tensor[0, 0]) == -2: # read none
|
|
continue
|
|
bsz = int(output_tensor[1, 0])
|
|
output_numpy = output_tensor[2:bsz + 2].numpy()
|
|
output_numpy[output_numpy == -1] = 2
|
|
outputs.append(output_numpy)
|
|
|
|
if int(output_tensor[0, 0]) < 0:
|
|
break
|
|
output = np.concatenate(outputs, axis=1)
|
|
seqs = tokenizer.batch_decode(
|
|
output.tolist(),
|
|
skip_special_tokens=True,
|
|
clean_up_tokenization_spaces=False,
|
|
)
|
|
if use_ep and (not ep_just_for_test):
|
|
print("seqs: ", seqs)
|
|
for i, seq in enumerate(seqs):
|
|
result_queue.put([i, len(output.tolist()[i]), seq])
|
|
|
|
|
|
def speculate_read_res(
|
|
model_name_or_path,
|
|
output_tensor_max_shape,
|
|
result_queue: mp.Queue,
|
|
msg_queue_id=None,
|
|
):
|
|
"""Read result from queue."""
|
|
if msg_queue_id is None:
|
|
from fastdeploy.model_executor.ops.gpu import speculate_get_output
|
|
else:
|
|
from fastdeploy.model_executor.ops.gpu import \
|
|
speculate_get_output_dynamic
|
|
|
|
tokenizer = ErnieBotTokenizer.from_pretrained(model_name_or_path)
|
|
paddle.device.set_device("cpu")
|
|
paddle.disable_static()
|
|
output_tensor = paddle.full(output_tensor_max_shape,
|
|
fill_value=2,
|
|
dtype="int64")
|
|
while True:
|
|
outputs = []
|
|
for _ in range(MAX_BSZ):
|
|
outputs.append([])
|
|
|
|
while True:
|
|
if msg_queue_id is None:
|
|
speculate_get_output(output_tensor, 0, True)
|
|
else:
|
|
speculate_get_output_dynamic(output_tensor, 0, True,
|
|
msg_queue_id)
|
|
if int(output_tensor[0]) == -2: # read none
|
|
continue
|
|
bsz = int(output_tensor[1])
|
|
accept_num = output_tensor[2:bsz + 2].numpy()
|
|
for bi in range(bsz):
|
|
outputs[bi].extend(
|
|
output_tensor.numpy()[2 + MAX_BSZ +
|
|
bi * MAX_DRAFT_TOKENS:2 + MAX_BSZ +
|
|
bi * MAX_DRAFT_TOKENS +
|
|
accept_num[bi]].tolist())
|
|
if int(output_tensor[0]) == -1:
|
|
break
|
|
|
|
seqs = tokenizer.batch_decode(
|
|
outputs,
|
|
skip_special_tokens=True,
|
|
clean_up_tokenization_spaces=False,
|
|
)
|
|
for i in range(bsz):
|
|
result_queue.put([i, len(outputs[i]), seqs[i]])
|
|
|
|
|
|
def calculate_effective_tokens(training_args, train_dataset, max_seq_len):
|
|
"""
|
|
Calculate the effective tokens during training.
|
|
"""
|
|
total_effective_tokens = 0
|
|
try:
|
|
data_parallel_degree = training_args.data_parallel_degree
|
|
except Exception:
|
|
data_parallel_degree = 1
|
|
if training_args.sharding_parallel_degree > 1:
|
|
sharding_parallel_degree = training_args.sharding_parallel_degree
|
|
else:
|
|
sharding_parallel_degree = 1
|
|
|
|
total_batch = (training_args.max_steps *
|
|
training_args.per_device_train_batch_size *
|
|
training_args.gradient_accumulation_steps *
|
|
sharding_parallel_degree * data_parallel_degree)
|
|
for i, data in enumerate(train_dataset):
|
|
if i == total_batch:
|
|
break
|
|
for dd in data:
|
|
total_effective_tokens += len(dd.token_ids)
|
|
total_tokens = total_batch * max_seq_len
|
|
|
|
return total_effective_tokens, total_tokens
|
|
|
|
|
|
def estimate_training(train_dataset, data_args, training_args, model_args):
|
|
"""
|
|
根据训练数据估算训练所需的步数。
|
|
|
|
Args:
|
|
- None
|
|
|
|
Returns:
|
|
- dict: 返回一个字典，包含了训练所需的步骤数信息。
|
|
|
|
"""
|
|
train_dataset.estimate = True
|
|
logger.info("Start to estimate max training steps...")
|
|
with open(data_args.train_task_config) as f:
|
|
train_task_group = json.load(f)
|
|
|
|
if len(train_task_group) > 1:
|
|
logger.warning(
|
|
"Suggest to use max_steps instead of num_train_epochs for multi source dataset."
|
|
)
|
|
logger.info(
|
|
"Multi source dataset detected, number of samples will be estimated by following rule. "
|
|
"num_samples = (source1_num_samples * prob1 + source2_num_samples * prob2 + ...) * epochs"
|
|
)
|
|
|
|
max_samples = train_dataset.max_estimate_samples
|
|
|
|
if training_args.max_estimate_samples != -1:
|
|
# Set estimate samples to max_estimate_samples
|
|
logger.warning(
|
|
"The results between sampling and non-sampling methods may differ."
|
|
)
|
|
train_dataset.max_estimate_samples = min(
|
|
training_args.max_estimate_samples,
|
|
train_dataset.max_estimate_samples)
|
|
|
|
if train_dataset.max_estimate_samples > 0:
|
|
train_batches = 0
|
|
train_tokens = 0
|
|
for sequences in train_dataset:
|
|
if not train_dataset.estimate:
|
|
break
|
|
train_batches += 1
|
|
for sequence in sequences:
|
|
train_tokens += len(sequence.token_ids)
|
|
|
|
train_tokens *= training_args.num_train_epochs
|
|
train_batches *= training_args.num_train_epochs
|
|
global_batch_size = (training_args.per_device_train_batch_size *
|
|
training_args.gradient_accumulation_steps *
|
|
max(training_args.data_parallel_degree, 1) *
|
|
max(training_args.sharding_parallel_degree, 1))
|
|
max_steps = int(np.ceil(train_batches / global_batch_size))
|
|
|
|
if max_samples != train_dataset.max_estimate_samples:
|
|
max_steps *= max_samples / train_dataset.max_estimate_samples
|
|
train_tokens *= max_samples / train_dataset.max_estimate_samples
|
|
train_dataset.used_samples *= (max_samples /
|
|
train_dataset.max_estimate_samples)
|
|
train_dataset.unused_samples *= (
|
|
max_samples / train_dataset.max_estimate_samples)
|
|
|
|
res = {
|
|
"num_train_epochs":
|
|
int(training_args.num_train_epochs),
|
|
"max_steps":
|
|
int(np.ceil(max_steps)),
|
|
"train_tokens":
|
|
int(train_tokens),
|
|
"global_batch_size":
|
|
int(global_batch_size),
|
|
"gradient_accumulation_steps":
|
|
training_args.gradient_accumulation_steps,
|
|
"warmup_steps":
|
|
int(np.ceil(0.1 * max_steps)),
|
|
"per_device_train_batch_size":
|
|
int(training_args.per_device_train_batch_size),
|
|
"tensor_parallel_degree":
|
|
int(training_args.tensor_parallel_degree),
|
|
"pipeline_parallel_degree":
|
|
int(training_args.pipeline_parallel_degree),
|
|
"sharding_parallel_degree":
|
|
int(training_args.sharding_parallel_degree),
|
|
"seed":
|
|
training_args.seed,
|
|
"num_samples_each_epoch":
|
|
data_args.num_samples_each_epoch,
|
|
"example_from_same_task_prob":
|
|
data_args.example_from_same_task_prob,
|
|
"pseudo_sampling_prob":
|
|
data_args.pseudo_sampling_prob,
|
|
"trigger_data_prob":
|
|
data_args.trigger_data_prob,
|
|
"max_seq_len":
|
|
int(data_args.max_seq_len),
|
|
"valid":
|
|
True,
|
|
"train_samples":
|
|
int(max_samples * training_args.num_train_epochs),
|
|
"estimate_samples":
|
|
int(train_dataset.max_estimate_samples),
|
|
"actual_train_samples":
|
|
int(train_dataset.used_samples * training_args.num_train_epochs),
|
|
"skip_samples":
|
|
int(train_dataset.unused_samples * training_args.num_train_epochs),
|
|
}
|
|
if hasattr(training_args, "num_of_gpus"):
|
|
res["num_of_gpus"] = training_args.num_of_gpus
|
|
|
|
if train_batches / training_args.num_train_epochs / global_batch_size < 1:
|
|
logger.warning(
|
|
"This dataset is too small, you'd better enlarge your dataset."
|
|
)
|
|
res["valid"] = False
|
|
|
|
if getattr(training_args, "estimation_output_file", None):
|
|
with open(training_args.estimation_output_file,
|
|
"w",
|
|
encoding="utf-8") as f:
|
|
json.dump(res, f)
|
|
|
|
return max_steps
|
|
else:
|
|
res = {
|
|
"num_train_epochs":
|
|
int(training_args.num_train_epochs),
|
|
"max_steps":
|
|
0,
|
|
"gradient_accumulation_steps":
|
|
training_args.gradient_accumulation_steps,
|
|
"train_tokens":
|
|
0,
|
|
"per_device_train_batch_size":
|
|
int(training_args.per_device_train_batch_size),
|
|
"tensor_parallel_degree":
|
|
int(training_args.tensor_parallel_degree),
|
|
"pipeline_parallel_degree":
|
|
int(training_args.pipeline_parallel_degree),
|
|
"sharding_parallel_degree":
|
|
int(training_args.sharding_parallel_degree),
|
|
"num_samples_each_epoch":
|
|
data_args.num_samples_each_epoch,
|
|
"example_from_same_task_prob":
|
|
data_args.example_from_same_task_prob,
|
|
"pseudo_sampling_prob":
|
|
data_args.pseudo_sampling_prob,
|
|
"trigger_data_prob":
|
|
data_args.trigger_data_prob,
|
|
"max_seq_len":
|
|
int(data_args.max_seq_len),
|
|
"seed":
|
|
data_args.seed,
|
|
"valid":
|
|
False,
|
|
"train_samples":
|
|
0,
|
|
}
|
|
if hasattr(training_args, "num_of_gpus"):
|
|
res["num_of_gpus"] = training_args.num_of_gpus
|
|
|
|
if getattr(training_args, "estimation_output_file", None):
|
|
with open(training_args.estimation_output_file,
|
|
"w",
|
|
encoding="utf-8") as f:
|
|
json.dump(res, f)
|
|
|
|
logger.error("No valid data found, please check your dataset format.")
|
|
return 0
|
|
|
|
|
|
def get_w4a8_gemm_config_tuple(file_root_path):
|
|
"""读取预配置的gemm 配置表
|
|
Args:
|
|
file_root_path (str): the directory of w4a8_gemm_config files
|
|
"""
|
|
|
|
def get_gemm_config_tuple_from_file(file):
|
|
gemm_tuple_list = []
|
|
for line in file:
|
|
line_split = line.split(" ")
|
|
gemm_tuple_list.append([
|
|
int(line_split[1]),
|
|
int(line_split[2]),
|
|
int(line_split[3]),
|
|
int(line_split[4]),
|
|
int(line_split[5]),
|
|
int(line_split[6]),
|
|
int(line_split[7]),
|
|
])
|
|
gemm_tuple_list.sort(key=lambda x: x[0])
|
|
gemm_tuple_numpy = np.array(gemm_tuple_list, dtype="int32")
|
|
gemm_tuple_numpy = gemm_tuple_numpy.flatten()
|
|
return gemm_tuple_numpy
|
|
|
|
qkv_gemm_config_tuple = []
|
|
out_linear_gemm_config_tuple = []
|
|
ffn1_gemm_config_tuple = []
|
|
ffn2_gemm_config_tuple = []
|
|
try:
|
|
qkv_tuned_gemm_config_log_path = os.path.join(
|
|
f"{file_root_path}", "qkv_tuned_gemm_config.log")
|
|
with open(qkv_tuned_gemm_config_log_path) as file:
|
|
qkv_gemm_config_tuple = get_gemm_config_tuple_from_file(file)
|
|
out_linear_tuned_gemm_config_log_path = os.path.join(
|
|
f"{file_root_path}", "out_linear_tuned_gemm_config.log")
|
|
with open(out_linear_tuned_gemm_config_log_path) as file:
|
|
out_linear_gemm_config_tuple = get_gemm_config_tuple_from_file(
|
|
file)
|
|
ffn1_tuned_gemm_config_log_path = os.path.join(
|
|
f"{file_root_path}", "ffn1_tuned_gemm_config.log")
|
|
with open(ffn1_tuned_gemm_config_log_path) as file:
|
|
ffn1_gemm_config_tuple = get_gemm_config_tuple_from_file(file)
|
|
ffn2_tuned_gemm_config_log_path = os.path.join(
|
|
f"{file_root_path}", "ffn2_tuned_gemm_config.log")
|
|
with open(ffn2_tuned_gemm_config_log_path) as file:
|
|
ffn2_gemm_config_tuple = get_gemm_config_tuple_from_file(file)
|
|
except Exception:
|
|
logger.warning(
|
|
"Found gemm config for W4A8 failed, using empty gemm tuple list for W4A8"
|
|
)
|
|
w4a8_gemm_config = {}
|
|
w4a8_gemm_config["qkv_gemm_config_tuple"] = qkv_gemm_config_tuple
|
|
w4a8_gemm_config[
|
|
"out_linear_gemm_config_tuple"] = out_linear_gemm_config_tuple
|
|
w4a8_gemm_config["ffn1_gemm_config_tuple"] = ffn1_gemm_config_tuple
|
|
w4a8_gemm_config["ffn2_gemm_config_tuple"] = ffn2_gemm_config_tuple
|
|
return w4a8_gemm_config
|
|
|
|
|
|
def update_refined_recompute(rr, sequence_parallel, lora=False):
|
|
"""update refined recompute dict."""
|
|
# if rr is a dict, return it directly
|
|
if isinstance(rr, dict):
|
|
return rr
|
|
if rr == "":
|
|
return {}
|
|
else:
|
|
|
|
rr_res = {
|
|
"mlp_row_ln": 0,
|
|
"attention_row_ln": 0,
|
|
"attention_column_ln": 0,
|
|
"mlp_column_ln": 0,
|
|
"flash_attn": 0,
|
|
}
|
|
ops = rr.split(",")
|
|
for op in ops:
|
|
if ":" not in op:
|
|
raise ValueError(
|
|
"Illegal refined_recompute input, please check.")
|
|
op_name, skip_num = op.split(":")[0], int(op.split(":")[1])
|
|
if op_name not in rr_res:
|
|
raise ValueError(
|
|
f"Refined recompute do not support {op_name}, please check."
|
|
)
|
|
|
|
if op_name in [
|
|
"mlp_row_ln",
|
|
"attention_row_ln",
|
|
"attention_column_ln",
|
|
"mlp_column_ln",
|
|
]:
|
|
if not sequence_parallel:
|
|
logger.warning(
|
|
f"Currently, the `{op_name}` op is only supported "
|
|
"when `sequence_parallel=True`. This refined recompute op will be ignored."
|
|
)
|
|
continue
|
|
if lora:
|
|
logger.warning(
|
|
"Currently, LoRA does not support refined recompute "
|
|
f"for the `{op_name}` op. This refined recompute op will be ignored."
|
|
)
|
|
continue
|
|
rr_res[op_name] = skip_num
|
|
|
|
return rr_res
|
|
|
|
|
|
def model_convert_fp8(model_path, device=None):
|
|
"""
|
|
Convert a model checkpoint from bf16/fp16 to fp8 format.
|
|
Args:
|
|
model_path (str): The path to the directory containing the model checkpoint files
|
|
(e.g., config.json and model_state.pdparams).
|
|
device (str, optional): The device to set for paddle, such as 'cpu' or 'gpu'.
|
|
If None, the default device is used.
|
|
|
|
Note:
|
|
This function requires non-smooth quantization 'act_scales' to be applied when using the converted model.
|
|
"""
|
|
if device is not None:
|
|
paddle.device.set_device(device)
|
|
|
|
config_path = os.path.join(model_path, "config.json")
|
|
with open(config_path, "r") as model_config_file:
|
|
model_config = json.load(model_config_file)
|
|
nums_layers = model_config["num_layers"]
|
|
|
|
weight_scales_path = os.path.join(model_path, "weight_scales_0.json")
|
|
with open(weight_scales_path, "r") as weight_scales_file:
|
|
weight_scales = json.load(weight_scales_file)
|
|
if "ernie.decoder.layers." + str(
|
|
0) + ".gate.weight_quanter" in weight_scales:
|
|
logger.info("FP8 model checkpoint already converted")
|
|
return
|
|
else:
|
|
logger.info("Converting model checkpoint to fp8...")
|
|
|
|
ffn1_weights_name = ".linear1.weight"
|
|
ffn1_bias_name = ".linear1.bias"
|
|
|
|
gate_weights_name = ".gate.weight"
|
|
up_weights_name = ".up.weight"
|
|
gate_bias_name = ".gate.bias"
|
|
up_bias_name = ".up.bias"
|
|
|
|
params_states = paddle.load(
|
|
os.path.join(model_path, "model_state.pdparams"))
|
|
new_path = os.path.join(model_path, "model_state.pdparams")
|
|
|
|
for i in range(0, nums_layers):
|
|
ffn1_weights = params_states["ernie.decoder.layers." + str(i) +
|
|
ffn1_weights_name]
|
|
ffn1_weights_0 = ffn1_weights[:, ::2]
|
|
ffn1_weights_1 = ffn1_weights[:, 1::2]
|
|
|
|
ffn1_weights_0_range = paddle.abs(ffn1_weights_0).max()
|
|
ffn1_weights_1_range = paddle.abs(ffn1_weights_1).max()
|
|
|
|
weight_scales["ernie.decoder.layers." + str(i) +
|
|
".gate.weight_quanter"] = (paddle.cast(
|
|
ffn1_weights_0_range, "float").numpy().tolist())
|
|
weight_scales["ernie.decoder.layers." + str(i) +
|
|
".up.weight_quanter"] = (paddle.cast(
|
|
ffn1_weights_1_range, "float").numpy().tolist())
|
|
params_states["ernie.decoder.layers." + str(i) +
|
|
gate_weights_name] = (ffn1_weights_0 * 448 /
|
|
ffn1_weights_0_range)
|
|
params_states["ernie.decoder.layers." + str(i) +
|
|
up_weights_name] = (ffn1_weights_1 * 448 /
|
|
ffn1_weights_1_range)
|
|
del params_states["ernie.decoder.layers." + str(i) + ffn1_weights_name]
|
|
|
|
ffn1_bias = params_states["ernie.decoder.layers." + str(i) +
|
|
ffn1_bias_name]
|
|
params_states["ernie.decoder.layers." + str(i) +
|
|
gate_bias_name] = ffn1_bias[::2]
|
|
params_states["ernie.decoder.layers." + str(i) +
|
|
up_bias_name] = ffn1_bias[1::2]
|
|
del params_states["ernie.decoder.layers." + str(i) + ffn1_bias_name]
|
|
|
|
with open(model_path + "/weight_scales_0.json", "w") as weight_scales_file:
|
|
json.dump(weight_scales, weight_scales_file)
|
|
|
|
paddle.save(params_states, new_path)
|
|
|
|
|
|
|
|
def load_ep_checkpoint(model_path, config, return_numpy=False, return_key_name=True):
|
|
"""
|
|
load ep checkpoint
|
|
"""
|
|
if return_key_name:
|
|
merge_path = os.path.join(model_path, "merged_tp1_state_split")
|
|
if os.path.isdir(merge_path):
|
|
# load keyname
|
|
|
|
state_dicts = []
|
|
files = glob.glob(model_path + "/merged_tp1_state_split/*")
|
|
for file_name in files:
|
|
try:
|
|
state_dicts += [
|
|
{file_name.split("/")[-1]: file_name}
|
|
] # save {layer_name: weight_file_name}
|
|
except Exception:
|
|
pass
|
|
new_state_dict = {}
|
|
for state_dict in state_dicts:
|
|
for key, value in state_dict.items():
|
|
new_state_dict[key] = value
|
|
state_dict = new_state_dict
|
|
else:
|
|
with open(
|
|
os.path.join(model_path, "model.safetensors.index.json"), "r"
|
|
) as f:
|
|
weight_map = json.load(f)["weight_map"]
|
|
state_dict = {
|
|
k: "[" + k + "]" + os.path.join(model_path, v)
|
|
for k, v in weight_map.items()
|
|
}
|
|
return state_dict
|
|
else:
|
|
# return_numpy=True cpu
|
|
# return_numpy=False gpu
|
|
with open(os.path.join(model_path, "model.safetensors.index.json"), "r") as f:
|
|
weight_list = json.load(f)["weight_map"]
|
|
filtered_map = {k: v for k, v in weight_list.items() if "experts" not in k}
|
|
num_local_ffn_keys = []
|
|
quant_suffix = (
|
|
"quant_weight"
|
|
if config.use_offline_quant and config.moe_quant_type != "default"
|
|
else ""
|
|
)
|
|
scale_suffix = (
|
|
"quant_scale"
|
|
if config.use_offline_quant and config.moe_quant_type != "default"
|
|
else ""
|
|
)
|
|
|
|
for i in range(config.moe_layer_start_index, config.num_layers):
|
|
for j in range(
|
|
config.num_experts_start_offset,
|
|
config.num_experts_start_offset + config.num_experts_per_rank,
|
|
):
|
|
ffn1_quant_key = f"ernie.layers.{i}.mlp.experts.{j}.up_gate_proj.weight.{quant_suffix}"
|
|
ffn2_quant_key = (
|
|
f"ernie.layers.{i}.mlp.experts.{j}.down_proj.weight.{quant_suffix}"
|
|
)
|
|
ffn1_scale_key = f"ernie.layers.{i}.mlp.experts.{j}.up_gate_proj.weight.{scale_suffix}"
|
|
ffn2_scale_key = (
|
|
f"ernie.layers.{i}.mlp.experts.{j}.down_proj.weight.{scale_suffix}"
|
|
)
|
|
num_local_ffn_keys.append(ffn1_quant_key)
|
|
num_local_ffn_keys.append(ffn2_quant_key)
|
|
num_local_ffn_keys.append(ffn1_scale_key)
|
|
num_local_ffn_keys.append(ffn2_scale_key)
|
|
|
|
for k in num_local_ffn_keys:
|
|
if k in weight_list:
|
|
filtered_map[k] = weight_list[k]
|
|
|
|
state_dict = {}
|
|
for k, safetensor_path in filtered_map.items():
|
|
with safe_open(
|
|
os.path.join(model_path, safetensor_path), framework="np", device="cpu"
|
|
) as f:
|
|
if k in f.keys():
|
|
weight = f.get_tensor(k)
|
|
if not return_numpy:
|
|
weight = paddle.Tensor(weight, zero_copy=True)
|
|
weight = weight._copy_to(
|
|
paddle.framework._current_expected_place(), False
|
|
)
|
|
state_dict[k] = weight
|
|
return state_dict
|
|
|
|
|
|
def get_safe_tensor_file(model_path):
|
|
"""
|
|
get_safe_tensor_file
|
|
"""
|
|
with open(os.path.join(model_path, "model.safetensors.index.json"),
|
|
"r") as f:
|
|
weight_map = json.load(f)["weight_map"]
|
|
safe_tensor_list = list(set(weight_map.values()))
|
|
key_name_list = list(set(weight_map.keys()))
|
|
safe_tensor_list = [os.path.join(model_path, v) for v in safe_tensor_list]
|
|
|
|
return key_name_list, safe_tensor_list
|
|
|
|
|
|
def safetensors_weights_iterator(safe_tensor_list: list[str], ):
|
|
"""
|
|
safetensors_weights_iterator
|
|
"""
|
|
for st_file in tqdm(
|
|
safe_tensor_list,
|
|
desc="Loading safetensors checkpoint shards",
|
|
):
|
|
with safe_open(st_file, framework="np") as f:
|
|
for name in f.keys():
|
|
param = f.get_tensor(name)
|
|
yield name, param
|
|
|
|
|
|
def get_state_dict(model_path, config):
|
|
"""
|
|
get_sate_dict
|
|
"""
|
|
state_dict = {}
|
|
_, safe_tensor_list = get_safe_tensor_file(
|
|
os.path.join(model_path, f"rank{config.tensor_parallel_rank}"))
|
|
weights_iterator = safetensors_weights_iterator(safe_tensor_list)
|
|
for name, weight in weights_iterator:
|
|
state_dict[name] = weight
|
|
return state_dict
|
|
|
|
|
|
def load_checkpoint(model_path, cls, config, return_numpy=True):
|
|
"""
|
|
load checkpoint
|
|
"""
|
|
if config.use_ep:
|
|
state_dict = load_ep_checkpoint(
|
|
model_path, config, return_numpy=True, return_key_name=True
|
|
)
|
|
else:
|
|
rank_dirs = [
|
|
f
|
|
for f in os.listdir(model_path)
|
|
if f.startswith("rank") and os.path.isdir(os.path.join(model_path, f))
|
|
]
|
|
if len(rank_dirs) > 1:
|
|
if config.tensor_parallel_degree != len(rank_dirs):
|
|
raise ValueError(
|
|
f"Your model only supports loading with tp{len(rank_dirs)}"
|
|
)
|
|
state_dict = get_state_dict(model_path, config)
|
|
else:
|
|
state_dict = load_tp_checkpoint(
|
|
model_path, cls, config, return_numpy=return_numpy
|
|
)
|
|
return state_dict
|
|
|
|
|
|
def parser_quant_type(quant_type):
|
|
"""
|
|
Parse the quantization type string and return the corresponding quantization types for weights,
|
|
activations, and custom.
|
|
|
|
Args:
|
|
quant_type (str): The quantization type string. It can be one of the following formats:
|
|
- "weight_only_int8" or "wint8": Only weights are quantized to int8.
|
|
- "weight_only_int4" or "wint4": Only weights are quantized to int4.
|
|
- A custom string in the format of "wxaybzcfp8", where 'x', 'y', 'z' are the quantization bitwidths
|
|
for weights, activations, and custom respectively,
|
|
and 'a', 'b', 'c' are the prefixes indicating the quantization types
|
|
(e.g., 'fp8' for floating-point 8-bit).
|
|
If a prefix is missing, the default quantization type will be used.
|
|
|
|
Returns:
|
|
tuple: A tuple of three strings representing the quantization types for weights, activations,
|
|
and custom respectively.
|
|
If the input is "weight_only_int8" or "wint8", returns ("int8", default_type, default_type).
|
|
If the input is "weight_only_int4" or "wint4", returns ("int4", default_type, default_type).
|
|
For custom strings, returns the parsed quantization types based on the input format.
|
|
|
|
Raises:
|
|
AssertionError: If the custom quantization type string format is incorrect.
|
|
"""
|
|
default_type = paddle.get_default_dtype()
|
|
conver_dict = {
|
|
"8": "int8",
|
|
"4": "int4",
|
|
"16": paddle.get_default_dtype,
|
|
"fp8": "float8_e4m3fn",
|
|
"fp16": "float16",
|
|
"bf16": "bfloat16",
|
|
"fp32": "float32"
|
|
}
|
|
cache_type = default_type
|
|
if "c8" in quant_type:
|
|
cache_type = "int8"
|
|
elif "cfp8" in quant_type:
|
|
cache_type = "fp8"
|
|
elif "c4" in quant_type:
|
|
cache_type = "int4"
|
|
if "weight_only_int8" in quant_type or "wint8" in quant_type:
|
|
return "int8", default_type, cache_type
|
|
elif "weight_only_int4" in quant_type or "wint4" in quant_type:
|
|
return "int4", default_type, cache_type
|
|
else:
|
|
# split quant type, eg. w4afp8c8 -> ['w', '4', 'a', 'fp8', 'c', '8']
|
|
pattern = f"({'|'.join(map(re.escape, ['w', 'a', 'c']))})"
|
|
splited_type = re.split(pattern, quant_type)
|
|
splited_type = [tmp_type for tmp_type in splited_type if tmp_type]
|
|
assert (len(splited_type) % 2 == 0 and len(splited_type)
|
|
<= 6), f"Quant type[{quant_type}] format error."
|
|
|
|
quant_type_list = []
|
|
if "w" in splited_type:
|
|
w_idx = splited_type.index("w")
|
|
quant_type_list.append(conver_dict[splited_type[w_idx + 1]])
|
|
else:
|
|
quant_type_list.append(default_type)
|
|
if "a" in splited_type:
|
|
a_idx = splited_type.index("a")
|
|
quant_type_list.append(conver_dict[splited_type[a_idx + 1]])
|
|
else:
|
|
quant_type_list.append(default_type)
|
|
if "c" in splited_type:
|
|
c_idx = splited_type.index("c")
|
|
quant_type_list.append(conver_dict[splited_type[c_idx + 1]])
|
|
else:
|
|
quant_type_list.append(default_type)
|
|
|
|
return quant_type_list[0], quant_type_list[1], quant_type_list[2]
|