[Mergekit]update & add LoRA merge (#9811)
* add * fix bug * fix * add * add lora merge * add * add * add * add * add * add
This commit is contained in:
lugimzzz
@@ -17,10 +17,7 @@ import os
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from typing import List, Optional
|
||||
|
||||
import paddle
|
||||
|
||||
from paddlenlp.utils.env import MERGE_CONFIG_NAME
|
||||
from paddlenlp.utils.log import logger
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -30,7 +27,6 @@ class MergeConfig:
|
||||
"""
|
||||
|
||||
# Common parameters
|
||||
device: str = field(default="cpu", metadata={"help": "Device to use for the merge.ex cpu、 gpu、low_gpu_mem"})
|
||||
tensor_type: str = field(
|
||||
default="np", metadata={"help": "Tensor type to use for the merge. Choose np(CPU Only) or pd (CPU/GPU)"}
|
||||
)
|
||||
@@ -39,6 +35,8 @@ class MergeConfig:
|
||||
merge_method: str = field(default="linear", metadata={"help": "The merge strategy."})
|
||||
merge_type: str = field(default="linear", metadata={"help": "The type of merge process."})
|
||||
sparsify_type: str = field(default=None, metadata={"help": "The type of sparsify process."})
|
||||
split_pieces: int = field(default=8, metadata={"help": "Split large tensor to multi-piece"})
|
||||
max_tensor_mem: float = field(default=0.5, metadata={"help": "Split tensor if exceed setting max_tensor_mem."})
|
||||
|
||||
# Model parameters
|
||||
model_path_list: Optional[List[str]] = field(default=None, metadata={"help": "Merge model name or path list"})
|
||||
@@ -46,7 +44,11 @@ class MergeConfig:
|
||||
default=None, metadata={"help": "Merge model name or path string.(split by ',')"}
|
||||
)
|
||||
base_model_path: str = field(default=None, metadata={"help": "Base model name or path."})
|
||||
output_path: str = field(default=None, metadata={"help": "Base model name or path."})
|
||||
output_path: str = field(default=None, metadata={"help": "Output model name or path."})
|
||||
lora_model_path: str = field(default=None, metadata={"help": "LoRA model name or path."})
|
||||
copy_file_list: Optional[List[str]] = field(
|
||||
default=None, metadata={"help": "Copy file list from base model path or first model path."}
|
||||
)
|
||||
# merge parameters
|
||||
weight_list: Optional[List[float]] = field(
|
||||
default=None, metadata={"help": "Relative (or absolute if normalize=False) weighting of a given tensor"}
|
||||
@@ -75,32 +77,43 @@ class MergeConfig:
|
||||
os.makedirs(self.output_path, exist_ok=True)
|
||||
if self.tensor_type not in ["np", "pd"]:
|
||||
raise ValueError(f"Unsupported tensor type: {self.tensor_type}. Support 'np' and 'pd' only.")
|
||||
if self.device == "gpu" and self.tensor_type == "np":
|
||||
logger.warning("np only support cpu device, but got gpu. Setting `device` to `cpu`.")
|
||||
self.device = "cpu"
|
||||
if self.lora_model_path is not None:
|
||||
if self.base_model_path is None:
|
||||
raise ValueError("Please specify the base_model_path when using LoRA merge.")
|
||||
self.tensor_type = "pd"
|
||||
|
||||
elif self.merge_method not in ["linear", "ties", "slerp", "della_linear", "della", "dare_linear", "dare_ties"]:
|
||||
raise ValueError(
|
||||
f"Unsupported merge strategy: {self.merge_method}. Please choose one from ['linear', 'slerp']."
|
||||
)
|
||||
if self.model_path_str is not None:
|
||||
self.model_path_list = self.model_path_str.split(",")
|
||||
if self.model_path_list is not None:
|
||||
if not isinstance(self.model_path_list, list) or len(self.model_path_list) < 2:
|
||||
raise ValueError(f"Please specify the model_path_list at least two. But got {self.model_path_list}")
|
||||
if self.weight_list is None:
|
||||
self.weight_list = [1.0] * len(self.model_path_list)
|
||||
self.normalize = True
|
||||
if len(self.model_path_list) != len(self.weight_list):
|
||||
raise ValueError("The length of model_path_list and weight_list must be the same.")
|
||||
if self.reserve_p < 0 or self.reserve_p > 1:
|
||||
raise ValueError("reserve_p must be between 0 and 1.")
|
||||
if "della" in self.merge_method or self.sparsify_type == "magprune":
|
||||
if self.reserve_p <= self.epsilon / 2 or self.reserve_p >= (1 - self.epsilon):
|
||||
if self.lora_model_path is None:
|
||||
if self.merge_method not in [
|
||||
"linear",
|
||||
"ties",
|
||||
"slerp",
|
||||
"della_linear",
|
||||
"della",
|
||||
"dare_linear",
|
||||
"dare_ties",
|
||||
]:
|
||||
raise ValueError(
|
||||
f"Error: reserve_p +- epsilon/2 must be in the range (0, 1). reserve_p + epsilon/2 = {self.reserve_p + self.epsilon / 2 }, reserve_p - epsilon/2 = {self.reserve_p - self.epsilon / 2 }"
|
||||
f"Unsupported merge strategy: {self.merge_method}. Please choose one from ['linear', 'slerp', 'ties', 'della_linear', 'della', ']."
|
||||
)
|
||||
paddle.set_device(self.device)
|
||||
if self.model_path_str is not None:
|
||||
self.model_path_list = self.model_path_str.split(",")
|
||||
if self.model_path_list is not None:
|
||||
if not isinstance(self.model_path_list, list) or len(self.model_path_list) < 2:
|
||||
raise ValueError(
|
||||
f"Please specify the model_path_list at least two. But got {self.model_path_list}"
|
||||
)
|
||||
if self.weight_list is None:
|
||||
self.weight_list = [1.0] * len(self.model_path_list)
|
||||
self.normalize = True
|
||||
if len(self.model_path_list) != len(self.weight_list):
|
||||
raise ValueError("The length of model_path_list and weight_list must be the same.")
|
||||
if self.reserve_p < 0 or self.reserve_p > 1:
|
||||
raise ValueError("reserve_p must be between 0 and 1.")
|
||||
if "della" in self.merge_method or self.sparsify_type == "magprune":
|
||||
if self.reserve_p <= self.epsilon / 2 or self.reserve_p >= (1 - self.epsilon):
|
||||
raise ValueError(
|
||||
f"Error: reserve_p +- epsilon/2 must be in the range (0, 1). reserve_p + epsilon/2 = {self.reserve_p + self.epsilon / 2 }, reserve_p - epsilon/2 = {self.reserve_p - self.epsilon / 2 }"
|
||||
)
|
||||
|
||||
@property
|
||||
def __dict__(self):
|
||||
|
||||
@@ -48,11 +48,10 @@ class MergeMethod:
|
||||
tensor_output = sum(weight * tensor for weight, tensor in zip(weight_list, tensor_list))
|
||||
return tensor_output
|
||||
elif self.merge_config.tensor_type == "pd":
|
||||
stacked_tensors = paddle.stack(tensor_list, axis=0)
|
||||
weights = paddle.to_tensor(weight_list, dtype=stacked_tensors.dtype)
|
||||
weights = weights.reshape([-1] + [1] * (len(stacked_tensors.shape) - 1))
|
||||
weighted_sum = paddle.sum(stacked_tensors * weights, axis=0)
|
||||
return weighted_sum
|
||||
tensor_output = paddle.zeros_like(tensor_list[0])
|
||||
for i, tensor in enumerate(tensor_list):
|
||||
tensor_output += tensor * weight_list[i]
|
||||
return tensor_output
|
||||
else:
|
||||
raise ValueError(f"Unkonwn tensor type {self.merge_config.tensor_type}")
|
||||
|
||||
@@ -155,28 +154,34 @@ class MergeMethod:
|
||||
|
||||
elif self.merge_config.tensor_type == "pd":
|
||||
mask_dtype = tensor_list[0].dtype
|
||||
weight_list = self.merge_config.weight_list
|
||||
stacked_tensors = paddle.stack(tensor_list, axis=0)
|
||||
weights = paddle.to_tensor(weight_list, dtype=stacked_tensors.dtype)
|
||||
weights = weights.reshape([-1] + [1] * (len(stacked_tensors.shape) - 1))
|
||||
weighted_tensors = stacked_tensors * weights
|
||||
|
||||
# Elect majority sign
|
||||
if self.merge_config.ties_elect_type == "sum":
|
||||
majority_sign = (paddle.sum(weighted_tensors, axis=0) >= 0).astype(mask_dtype) * 2 - 1
|
||||
elif self.merge_config.ties_elect_type == "count":
|
||||
stacked_signs = paddle.sign(stacked_tensors).astype(mask_dtype)
|
||||
majority_sign = (paddle.sum(stacked_signs, axis=0) >= 0).astype(mask_dtype) * 2 - 1
|
||||
else:
|
||||
raise NotImplementedError(f"ties_elect_type: {self.merge_config.ties_elect_type} is unknown.")
|
||||
majority_sign = paddle.zeros_like(tensor_list[0])
|
||||
for i, tensor in enumerate(tensor_list):
|
||||
if self.merge_config.ties_elect_type == "sum":
|
||||
majority_sign += tensor * self.merge_config.weight_list[i]
|
||||
elif self.merge_config.ties_elect_type == "count":
|
||||
majority_sign += tensor.sign()
|
||||
else:
|
||||
raise NotImplementedError(f"ties_elect_type: {self.merge_config.ties_elect_type} is unknown.")
|
||||
majority_sign = (majority_sign >= 0).astype(mask_dtype) * 2 - 1
|
||||
|
||||
# Merge
|
||||
stacked_masks = (paddle.sign(weighted_tensors) == majority_sign).astype(mask_dtype)
|
||||
masked_tensors = stacked_masks * weighted_tensors
|
||||
merge_tensor = paddle.sum(masked_tensors, axis=0)
|
||||
merge_tensor = paddle.zeros_like(tensor_list[0])
|
||||
if self.merge_config.normalize:
|
||||
divisor = paddle.zeros_like(tensor_list[0])
|
||||
for i, tensor in enumerate(tensor_list):
|
||||
if self.merge_config.normalize:
|
||||
mask = (tensor.sign() == majority_sign).astype(mask_dtype) * self.merge_config.weight_list[i]
|
||||
divisor += mask
|
||||
merge_tensor += mask * tensor
|
||||
else:
|
||||
merge_tensor += (
|
||||
(tensor.sign() == majority_sign).astype(mask_dtype) * tensor * self.merge_config.weight_list[i]
|
||||
)
|
||||
|
||||
# Normalize
|
||||
if self.merge_config.normalize:
|
||||
weight_masks = stacked_masks * weights
|
||||
divisor = paddle.sum(weight_masks, axis=0)
|
||||
divisor = paddle.where(paddle.abs(divisor) < 1e-8, paddle.ones_like(divisor), divisor)
|
||||
merge_tensor /= divisor
|
||||
|
||||
|
||||
@@ -11,9 +11,10 @@
|
||||
# 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 gc
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
from multiprocessing import Process
|
||||
|
||||
import numpy as np
|
||||
@@ -22,6 +23,8 @@ import paddle.distributed as dist
|
||||
from safetensors import safe_open
|
||||
from safetensors.numpy import save_file
|
||||
|
||||
from paddlenlp.peft import LoRAConfig
|
||||
from paddlenlp.utils import device_guard
|
||||
from paddlenlp.utils.env import (
|
||||
LORA_WEIGHTS_NAME,
|
||||
PADDLE_MASTER_WEIGHTS_NAME,
|
||||
@@ -29,14 +32,18 @@ from paddlenlp.utils.env import (
|
||||
SAFE_MASTER_WEIGHTS_INDEX_NAME,
|
||||
SAFE_MASTER_WEIGHTS_NAME,
|
||||
SAFE_PEFT_WEIGHTS_INDEX_NAME,
|
||||
SAFE_PEFT_WEIGHTS_NAME,
|
||||
SAFE_WEIGHTS_INDEX_NAME,
|
||||
SAFE_WEIGHTS_NAME,
|
||||
)
|
||||
from paddlenlp.utils.log import logger
|
||||
from paddlenlp.utils.safetensors import fast_safe_open
|
||||
|
||||
from .merge_method import MergeMethod
|
||||
from .merge_utils import divide_positions
|
||||
from .merge_utils import (
|
||||
divide_lora_key_list,
|
||||
divide_positions,
|
||||
divide_safetensor_key_list,
|
||||
)
|
||||
from .sparsify_method import SparsifyMethod
|
||||
|
||||
SPARSIFY_MERGE_MAPPING = {
|
||||
@@ -57,6 +64,10 @@ class MergeModel:
|
||||
self.is_peft = False
|
||||
|
||||
def reset_merge_model(self, merge_config=None, merge_param_dict=None):
|
||||
self.is_cpu = "cpu" in paddle.device.get_device()
|
||||
if not self.is_cpu:
|
||||
if dist.get_world_size() > 1 and not paddle.distributed.is_initialized():
|
||||
dist.init_parallel_env()
|
||||
if merge_config is not None:
|
||||
self.merge_config = merge_config
|
||||
elif merge_param_dict is not None:
|
||||
@@ -76,207 +87,248 @@ class MergeModel:
|
||||
self.merge_method = MergeMethod(merge_config, sparsify_method)
|
||||
|
||||
def merge_model(self):
|
||||
if self.merge_config.lora_model_path is not None:
|
||||
self.merge_lora_model()
|
||||
else:
|
||||
if self.merge_config.tensor_type == "np" and not self.is_cpu:
|
||||
# Avoid memory allocated on GPU
|
||||
with device_guard():
|
||||
self.mergekit()
|
||||
else:
|
||||
self.mergekit()
|
||||
self.copy_file()
|
||||
|
||||
def copy_file(self):
|
||||
if self.merge_config.copy_file_list is not None:
|
||||
if self.merge_config.base_model_path is not None:
|
||||
src_path = self.merge_config.base_model_path
|
||||
else:
|
||||
src_path = self.merge_config.model_path_list[0]
|
||||
for file in self.merge_config.copy_file_list:
|
||||
src_file = os.path.join(src_path, file)
|
||||
dst_file = os.path.join(self.merge_config.output_path, file)
|
||||
if os.path.isfile(src_file):
|
||||
shutil.copy2(src_file, dst_file)
|
||||
else:
|
||||
logger.warning(f"Copy failed: {file} not found in {src_path}")
|
||||
|
||||
def mergekit(self):
|
||||
# Check model file type
|
||||
file_type_list = []
|
||||
for model_path in self.merge_config.model_path_list:
|
||||
file_type_list.append(self.check_model_path(model_path))
|
||||
if self.merge_config.base_model_path is not None:
|
||||
file_type_list.append(self.check_model_path(self.merge_config.base_model_path))
|
||||
if not all(file_type[1] is True for file_type in file_type_list) and not all(
|
||||
file_type[1] is False for file_type in file_type_list
|
||||
):
|
||||
raise ValueError("Please ensure that all models should be same type.")
|
||||
if all(
|
||||
file_type[0] == "safetensors" or file_type[0] == "safetensors_without_index"
|
||||
for file_type in file_type_list
|
||||
):
|
||||
|
||||
# Merge model (distinguish between safetensors and pdparams)
|
||||
if all(file_type == "safetensors" or file_type == "safetensors_without_index" for file_type in file_type_list):
|
||||
self.merge_safetensor_model(file_type_list)
|
||||
else:
|
||||
self.merge_mix_model(file_type_list)
|
||||
|
||||
def merge_mix_model(self, file_type_list):
|
||||
# Load model state dict
|
||||
state_dict_list = []
|
||||
for i, model_path in enumerate(self.merge_config.model_path_list):
|
||||
state_dict_list.append(self.get_model_state_dict(model_path, file_type_list[i][0]))
|
||||
state_dict_list.append(self.get_model_state_dict(model_path, file_type_list[i]))
|
||||
if self.merge_config.base_model_path is not None:
|
||||
state_dict_list.append(self.get_model_state_dict(self.merge_config.base_model_path, file_type_list[-1][0]))
|
||||
state_dict_list.append(self.get_model_state_dict(self.merge_config.base_model_path, file_type_list[-1]))
|
||||
|
||||
if not all(state_dict_list[0].keys() == state_dict.keys() for state_dict in state_dict_list):
|
||||
raise ValueError("State dict keys mismatch. Please make sure you load the correct weight file")
|
||||
if self.merge_config.base_model_path is not None:
|
||||
base_state_dict = state_dict_list.pop()
|
||||
base_file_type = file_type_list.pop()
|
||||
merge_state_dict = {}
|
||||
total_size = 0
|
||||
weight_map = {}
|
||||
for key in state_dict_list[0].keys():
|
||||
is_bf16 = False
|
||||
tensor_list = []
|
||||
for state_dict, file_type in zip(state_dict_list, file_type_list):
|
||||
if file_type[0] == "pdparams":
|
||||
if str(state_dict[key].dtype) == "paddle.bfloat16":
|
||||
is_bf16 = True
|
||||
state_dict[key] = state_dict[key].astype("float32").numpy()
|
||||
else:
|
||||
state_dict[key] = state_dict[key].numpy()
|
||||
elif str(state_dict[key].dtype) == "uint16":
|
||||
is_bf16 = True
|
||||
state_dict[key] = paddle.to_tensor(state_dict[key], dtype="bfloat16").astype("float32").numpy()
|
||||
tensor_list.append(state_dict[key])
|
||||
if self.merge_config.base_model_path is not None:
|
||||
if base_file_type[0] == "pdparams":
|
||||
if str(base_state_dict[key].dtype) == "paddle.bfloat16":
|
||||
base_state_dict[key] = base_state_dict[key].astype("float32").numpy()
|
||||
else:
|
||||
base_state_dict[key] = base_state_dict[key].numpy()
|
||||
elif str(base_state_dict[key].dtype) == "uint16":
|
||||
base_state_dict[key] = (
|
||||
paddle.to_tensor(base_state_dict[key], dtype="bfloat16").astype("float32").numpy()
|
||||
)
|
||||
tensor_list = [tensor - base_state_dict[key] for tensor in tensor_list]
|
||||
merge_state_dict[key] = self.merge_method.merge(tensor_list)
|
||||
if self.merge_config.base_model_path is not None:
|
||||
merge_state_dict[key] += base_state_dict[key]
|
||||
# dtype==bfloat16: numpy(float32) -> paddle(float32) -> paddle(bfloat16) -> numpy(uint16)
|
||||
if is_bf16:
|
||||
merge_state_dict[key] = (
|
||||
paddle.to_tensor(merge_state_dict[key], dtype="float32").astype("bfloat16").numpy()
|
||||
)
|
||||
total_size += np.prod(merge_state_dict[key].shape) * self.numpy_dtype_map[str(merge_state_dict[key].dtype)]
|
||||
if self.merge_config.merge_prefix == "model" and file_type_list[0][1] is True:
|
||||
weight_map[key] = "peft_model-00001-of-00001.safetensors"
|
||||
else:
|
||||
weight_map[key] = f"{self.merge_config.merge_prefix}-00001-of-00001.safetensors"
|
||||
# save safetensor file
|
||||
if self.merge_config.merge_prefix == "model" and file_type_list[0][1] is True:
|
||||
save_file(
|
||||
merge_state_dict,
|
||||
os.path.join(self.merge_config.output_path, "peft_model-00001-of-00001.safetensors"),
|
||||
metadata={"format": "np"},
|
||||
)
|
||||
else:
|
||||
save_file(
|
||||
merge_state_dict,
|
||||
os.path.join(
|
||||
self.merge_config.output_path, f"{self.merge_config.merge_prefix}-00001-of-00001.safetensors"
|
||||
),
|
||||
metadata={"format": "np"},
|
||||
)
|
||||
# save safe index file
|
||||
index = {"metadata": {"total_size": int(total_size)}, "weight_map": weight_map}
|
||||
if self.merge_config.merge_prefix == "model" and file_type_list[0][1] is True:
|
||||
save_index_file = os.path.join(self.merge_config.output_path, SAFE_PEFT_WEIGHTS_INDEX_NAME)
|
||||
else:
|
||||
save_index_file = os.path.join(self.merge_config.output_path, self.safe_index_name())
|
||||
with open(save_index_file, "w", encoding="utf-8") as f:
|
||||
content = json.dumps(index, indent=2) + "\n"
|
||||
f.write(content)
|
||||
# save merge config file
|
||||
self.merge_config.save_pretrained(self.merge_config.output_path)
|
||||
del state_dict_list
|
||||
del merge_state_dict
|
||||
if self.merge_config.base_model_path is not None:
|
||||
del base_state_dict
|
||||
gc.collect()
|
||||
|
||||
def get_model_state_dict(self, model_path, file_type):
|
||||
# Merge state dict
|
||||
merge_state_dict = {}
|
||||
index = {"metadata": {"total_size": 0}, "weight_map": {}}
|
||||
|
||||
key_list = list(state_dict_list[file_type_list.index("pdparams")].keys())
|
||||
model_num = len(state_dict_list)
|
||||
rank = dist.get_rank()
|
||||
positions = divide_positions(len(key_list), dist.get_world_size())
|
||||
local_keys = key_list[positions[rank] : positions[rank + 1]]
|
||||
for ii in range(len(positions) - 1):
|
||||
shard_file = f"{self.merge_config.merge_prefix}-{ii+1:05d}-of-{dist.get_world_size():05d}.safetensors"
|
||||
for key in key_list[positions[ii] : positions[ii + 1]]:
|
||||
index["weight_map"][key] = shard_file
|
||||
index["metadata"]["total_size"] += int(
|
||||
np.prod(state_dict_list[0][key].shape) * self.numpy_dtype_map[str(state_dict_list[0][key].dtype)]
|
||||
)
|
||||
for key in local_keys:
|
||||
# Tensor preprocess
|
||||
is_bf16 = str(state_dict_list[0][key].dtype) == "uint16"
|
||||
tensor_list = [state_dict_list[i].pop(key) for i in range(model_num)]
|
||||
tensor_mem = int(np.prod(tensor_list[0].shape) * self.numpy_dtype_map[str(tensor_list[0].dtype)]) / (
|
||||
1024**3
|
||||
)
|
||||
if self.merge_config.tensor_type == "pd" and tensor_mem > self.merge_config.max_tensor_mem:
|
||||
tensor_split_list = [
|
||||
np.array_split(tensor, self.merge_config.split_pieces, axis=0) for tensor in tensor_list
|
||||
]
|
||||
merge_split = []
|
||||
for sp in range(self.merge_config.split_pieces):
|
||||
tensor_list = [tensor_split[sp] for tensor_split in tensor_split_list]
|
||||
if is_bf16:
|
||||
tensor_list = [
|
||||
paddle.Tensor(tensor, zero_copy=True).astype("float32") for tensor in tensor_list
|
||||
]
|
||||
else:
|
||||
tensor_list = [paddle.Tensor(tensor, zero_copy=True) for tensor in tensor_list]
|
||||
if self.merge_config.base_model_path is not None:
|
||||
base_tensor = tensor_list.pop()
|
||||
tensor_list = [tensor - base_tensor for tensor in tensor_list]
|
||||
merge_tensor = self.merge_method.merge(tensor_list)
|
||||
if self.merge_config.base_model_path is not None:
|
||||
merge_tensor += base_tensor
|
||||
if is_bf16:
|
||||
merge_split.append(merge_tensor.astype("bfloat16").numpy())
|
||||
else:
|
||||
merge_split.append(merge_tensor.numpy())
|
||||
merge_state_dict[key] = np.concatenate(merge_split, axis=0)
|
||||
else:
|
||||
if self.merge_config.tensor_type == "pd":
|
||||
if is_bf16:
|
||||
tensor_list = [
|
||||
paddle.Tensor(tensor, zero_copy=True).astype("float32") for tensor in tensor_list
|
||||
]
|
||||
else:
|
||||
tensor_list = [paddle.Tensor(tensor, zero_copy=True) for tensor in tensor_list]
|
||||
elif self.merge_config.tensor_type == "np" and is_bf16:
|
||||
tensor_list = [
|
||||
paddle.Tensor(tensor, zero_copy=True).astype("float32").numpy() for tensor in tensor_list
|
||||
]
|
||||
|
||||
if self.merge_config.base_model_path is not None:
|
||||
base_tensor = tensor_list.pop()
|
||||
tensor_list = [tensor - base_tensor for tensor in tensor_list]
|
||||
merge_tensor = self.merge_method.merge(tensor_list)
|
||||
if self.merge_config.base_model_path is not None:
|
||||
merge_tensor += base_tensor
|
||||
if self.merge_config.tensor_type == "pd":
|
||||
if is_bf16:
|
||||
merge_state_dict[key] = merge_tensor.astype("bfloat16").numpy()
|
||||
else:
|
||||
merge_state_dict[key] = merge_tensor.numpy()
|
||||
elif self.merge_config.tensor_type == "np" and is_bf16:
|
||||
# dtype==bfloat16: numpy(float32) -> paddle(float32) -> paddle(bfloat16) -> numpy(uint16)
|
||||
merge_state_dict[key] = paddle.Tensor(merge_tensor, zero_copy=True).astype("bfloat16").numpy()
|
||||
|
||||
# Save safetensor file
|
||||
save_file(
|
||||
merge_state_dict,
|
||||
os.path.join(
|
||||
self.merge_config.output_path,
|
||||
f"{self.merge_config.merge_prefix}-{rank+1:05d}-of-{dist.get_world_size():05d}.safetensors",
|
||||
),
|
||||
metadata={"format": "np"},
|
||||
)
|
||||
# Save index file & merge config file
|
||||
if paddle.distributed.get_rank() == 0:
|
||||
save_index_file = os.path.join(self.merge_config.output_path, self.safe_index_name())
|
||||
with open(save_index_file, "w", encoding="utf-8") as f:
|
||||
f.write(json.dumps(index, indent=2) + "\n")
|
||||
self.merge_config.save_pretrained(self.merge_config.output_path)
|
||||
|
||||
def get_model_state_dict(self, model_path, file_type, key_list=None):
|
||||
if file_type == "safetensors":
|
||||
state_dict = {}
|
||||
with open(os.path.join(model_path, self.safe_index_name()), "r", encoding="utf-8") as f:
|
||||
index = json.load(f)
|
||||
for key in index["weight_map"].keys():
|
||||
with fast_safe_open(
|
||||
os.path.join(model_path, index["weight_map"][key]),
|
||||
framework="np",
|
||||
) as f:
|
||||
state_dict[key] = f.get_tensor(key)
|
||||
if key_list is None:
|
||||
files = set(index["weight_map"].values())
|
||||
for file in files:
|
||||
with fast_safe_open(os.path.join(model_path, file), framework="np") as f:
|
||||
for k in f.keys():
|
||||
state_dict[k] = f.get_tensor(k)
|
||||
else:
|
||||
file_map = {}
|
||||
for key in key_list:
|
||||
if index["weight_map"][key] not in file_map:
|
||||
file_map[index["weight_map"][key]] = [key]
|
||||
else:
|
||||
file_map[index["weight_map"][key]].append(key)
|
||||
for file in file_map.keys():
|
||||
with fast_safe_open(os.path.join(model_path, file), framework="np") as f:
|
||||
for k in file_map[file]:
|
||||
state_dict[k] = f.get_tensor(k)
|
||||
elif file_type == "safetensors_without_index":
|
||||
state_dict = {}
|
||||
with fast_safe_open(os.path.join(model_path, self.safe_weight_name()), framework="numpy") as f:
|
||||
for k in f.keys():
|
||||
tgt_key_list = f.keys() if key_list is None else key_list
|
||||
for k in tgt_key_list:
|
||||
state_dict[k] = f.get_tensor(k)
|
||||
elif file_type == "pdparams":
|
||||
state_dict = paddle.load(os.path.join(model_path, self.weight_name()))
|
||||
state_dict = np.load(os.path.join(model_path, self.weight_name()), allow_pickle=True)
|
||||
if "StructuredToParameterName@@" in state_dict.keys():
|
||||
state_dict.pop("StructuredToParameterName@@")
|
||||
elif file_type == "lora_pdparams":
|
||||
state_dict = np.load(os.path.join(model_path, LORA_WEIGHTS_NAME), allow_pickle=True)
|
||||
elif file_type == "lora_safetensors":
|
||||
state_dict = {}
|
||||
with open(os.path.join(model_path, SAFE_PEFT_WEIGHTS_INDEX_NAME), "r", encoding="utf-8") as f:
|
||||
index = json.load(f)
|
||||
files = set(index["weight_map"].values())
|
||||
for file in files:
|
||||
with fast_safe_open(os.path.join(model_path, file), framework="np") as f:
|
||||
for k in f.keys():
|
||||
state_dict[k] = f.get_tensor(k)
|
||||
else:
|
||||
raise ValueError(f"Unsupported file_type: {file_type}")
|
||||
return state_dict
|
||||
|
||||
def create_safetensor_index(self, model_path):
|
||||
weight_map = {}
|
||||
total_size = 0
|
||||
|
||||
with safe_open(os.path.join(model_path, self.safe_weight_name()), framework="numpy") as f:
|
||||
for key in f.keys():
|
||||
tensor = f.get_tensor(key)
|
||||
total_size += np.prod(tensor.shape) * self.numpy_dtype_map[str(tensor.dtype)]
|
||||
weight_map[key] = self.safe_weight_name()
|
||||
index = {"metadata": {"total_size": total_size}, "weight_map": weight_map}
|
||||
def get_safetensor_index(self, model_path, file_type):
|
||||
if file_type == "safetensors":
|
||||
with open(os.path.join(model_path, self.safe_index_name()), "r", encoding="utf-8") as f:
|
||||
index = json.load(f)
|
||||
elif file_type == "safetensors_without_index":
|
||||
weight_map = {}
|
||||
total_size = 0
|
||||
with safe_open(os.path.join(model_path, self.safe_weight_name()), framework="numpy") as f:
|
||||
for key in f.keys():
|
||||
tensor = f.get_tensor(key)
|
||||
total_size += int(np.prod(tensor.shape) * self.numpy_dtype_map[str(tensor.dtype)])
|
||||
weight_map[key] = self.safe_weight_name()
|
||||
index = {"metadata": {"total_size": total_size}, "weight_map": weight_map}
|
||||
return index
|
||||
|
||||
def merge_safetensor_model(self, file_type_list):
|
||||
use_gpu = self.merge_config.device == "gpu"
|
||||
|
||||
if use_gpu:
|
||||
rank = dist.get_rank()
|
||||
if dist.get_world_size() > 1:
|
||||
dist.init_parallel_env()
|
||||
|
||||
# Load index
|
||||
index_list = []
|
||||
model_path_list = self.merge_config.model_path_list
|
||||
model_path_list = self.merge_config.model_path_list.copy()
|
||||
if self.merge_config.base_model_path is not None:
|
||||
model_path_list += [self.merge_config.base_model_path]
|
||||
|
||||
for model_path, file_type in zip(model_path_list, file_type_list):
|
||||
if file_type[0] == "safetensors":
|
||||
if file_type[1] is False:
|
||||
with open(os.path.join(model_path, self.safe_index_name()), "r", encoding="utf-8") as f:
|
||||
index_list.append(json.load(f))
|
||||
else:
|
||||
with open(os.path.join(model_path, SAFE_PEFT_WEIGHTS_INDEX_NAME), "r", encoding="utf-8") as f:
|
||||
index_list.append(json.load(f))
|
||||
else:
|
||||
index = self.create_safetensor_index(model_path)
|
||||
index_list.append(index)
|
||||
index_list.append(self.get_safetensor_index(model_path, file_type))
|
||||
|
||||
# Check index
|
||||
if not all(index_list[0]["metadata"]["total_size"] == index["metadata"]["total_size"] for index in index_list):
|
||||
raise ValueError("Weights total_size mismatch. Please make sure you load the correct weight file")
|
||||
if not all(index_list[0]["weight_map"].keys() == index["weight_map"].keys() for index in index_list):
|
||||
raise ValueError("Weights weight_map mismatch. Please make sure you load the correct weight file")
|
||||
|
||||
# Initialize new index
|
||||
index = {}
|
||||
index["metadata"] = index_list[0]["metadata"]
|
||||
index["metadata"]["total_size"] = int(index["metadata"]["total_size"])
|
||||
index["weight_map"] = {}
|
||||
|
||||
key_list = list(index_list[0]["weight_map"].keys())
|
||||
if use_gpu:
|
||||
positions = divide_positions(len(key_list), dist.get_world_size())
|
||||
num = self.merge_config.n_process if self.is_cpu else dist.get_world_size()
|
||||
if file_type_list[0] == "safetensors" and len(set(index_list[0]["weight_map"].values())) >= num:
|
||||
key_list, positions = divide_safetensor_key_list(index_list[0]["weight_map"], num)
|
||||
else:
|
||||
positions = divide_positions(len(key_list), self.merge_config.n_process)
|
||||
|
||||
if use_gpu:
|
||||
start_idx = positions[rank]
|
||||
end_idx = positions[rank + 1] if rank + 1 < len(positions) else len(key_list)
|
||||
local_keys = key_list[start_idx:end_idx]
|
||||
# use index_gpu to preserve index
|
||||
index_gpu = index
|
||||
for i in range(len(positions) - 1):
|
||||
shard_file = (
|
||||
f"{self.merge_config.merge_prefix}-{i+1:05d}-of-{self.merge_config.n_process:05d}.safetensors"
|
||||
)
|
||||
for k in key_list[positions[i] : positions[i + 1]]:
|
||||
index_gpu["weight_map"][k] = shard_file
|
||||
shard_file = f"{self.merge_config.merge_prefix}-{rank+1:05d}-of-{dist.get_world_size():05d}.safetensors"
|
||||
for k in local_keys:
|
||||
index["weight_map"][k] = shard_file
|
||||
|
||||
key_list = list(index_list[0]["weight_map"].keys())
|
||||
positions = divide_positions(len(key_list), num)
|
||||
if not self.is_cpu:
|
||||
i = dist.get_rank()
|
||||
local_keys = key_list[positions[i] : positions[i + 1]]
|
||||
shard_file = f"{self.merge_config.merge_prefix}-{i+1:05d}-of-{dist.get_world_size():05d}.safetensors"
|
||||
if self.merge_config.tensor_type == "np":
|
||||
ValueError(f"Tensor type '{self.merge_config.tensor_type}' should be 'pd' when using GPU.")
|
||||
self.shard_merge_np(local_keys, index_list, shard_file)
|
||||
else:
|
||||
self.shard_merge_pd(local_keys, index_list, shard_file)
|
||||
if dist.get_world_size() > 1:
|
||||
dist.barrier()
|
||||
|
||||
for i in range(len(positions) - 1):
|
||||
shard_file = f"{self.merge_config.merge_prefix}-{i+1:05d}-of-{dist.get_world_size():05d}.safetensors"
|
||||
for k in key_list[positions[i] : positions[i + 1]]:
|
||||
index["weight_map"][k] = shard_file
|
||||
else:
|
||||
threads = []
|
||||
for i in range(len(positions) - 1):
|
||||
@@ -299,17 +351,11 @@ class MergeModel:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
# Save safe index file
|
||||
if not use_gpu or (use_gpu and rank == 0):
|
||||
if paddle.distributed.get_rank() == 0:
|
||||
save_index_file = os.path.join(self.merge_config.output_path, self.safe_index_name())
|
||||
with open(save_index_file, "w", encoding="utf-8") as f:
|
||||
if not use_gpu:
|
||||
content = json.dumps(index, indent=2) + "\n"
|
||||
else:
|
||||
content = json.dumps(index_gpu, indent=2) + "\n"
|
||||
f.write(content)
|
||||
self.merge_config.save_pretrained(self.merge_config.output_path)
|
||||
f.write(json.dumps(index, indent=2) + "\n")
|
||||
|
||||
def shard_merge_np(
|
||||
self,
|
||||
@@ -320,14 +366,13 @@ class MergeModel:
|
||||
merge_state_dict = {}
|
||||
for k in key_list:
|
||||
tensor_list = []
|
||||
|
||||
for i, model_path in enumerate(self.merge_config.model_path_list):
|
||||
with fast_safe_open(os.path.join(model_path, index_list[i]["weight_map"][k]), framework="np") as w:
|
||||
tensor = w.get_tensor(k)
|
||||
dtype = tensor.dtype
|
||||
# dtype==bfloat16: numpy(uint16) -> paddle(bfloat16) -> paddle(float32) -> numpy(float32)
|
||||
if tensor.dtype == np.uint16:
|
||||
tensor = paddle.to_tensor(tensor, dtype="bfloat16").astype("float32").numpy()
|
||||
tensor = paddle.Tensor(tensor, zero_copy=True).astype("float32").numpy()
|
||||
tensor_list.append(tensor)
|
||||
if self.merge_config.base_model_path is not None:
|
||||
with fast_safe_open(
|
||||
@@ -336,25 +381,19 @@ class MergeModel:
|
||||
) as w:
|
||||
base_tensor = w.get_tensor(k)
|
||||
if base_tensor.dtype == np.uint16:
|
||||
base_tensor = paddle.to_tensor(base_tensor, dtype="bfloat16").astype("float32").numpy()
|
||||
base_tensor = paddle.Tensor(base_tensor, zero_copy=True).astype("float32").numpy()
|
||||
tensor_list = [tensor - base_tensor for tensor in tensor_list]
|
||||
merge_state_dict[k] = self.merge_method.merge(tensor_list)
|
||||
if self.merge_config.base_model_path is not None:
|
||||
merge_state_dict[k] += base_tensor
|
||||
# dtype==bfloat16: numpy(float32) -> paddle(float32) -> paddle(bfloat16) -> numpy(uint16)
|
||||
if dtype == np.uint16:
|
||||
merge_state_dict[k] = paddle.to_tensor(merge_state_dict[k], dtype="float32").astype("bfloat16").numpy()
|
||||
del tensor_list
|
||||
if self.merge_config.base_model_path is not None:
|
||||
|
||||
del base_tensor
|
||||
merge_state_dict[k] = paddle.Tensor(merge_state_dict[k], zero_copy=True).astype("bfloat16").numpy()
|
||||
save_file(
|
||||
merge_state_dict,
|
||||
os.path.join(self.merge_config.output_path, shard_file),
|
||||
metadata={"format": "np"},
|
||||
)
|
||||
del merge_state_dict
|
||||
gc.collect()
|
||||
|
||||
def shard_merge_pd(
|
||||
self,
|
||||
@@ -365,81 +404,85 @@ class MergeModel:
|
||||
merge_state_dict = {}
|
||||
for k in key_list:
|
||||
tensor_list = []
|
||||
is_bf16 = False
|
||||
for i, model_path in enumerate(self.merge_config.model_path_list):
|
||||
with fast_safe_open(os.path.join(model_path, index_list[i]["weight_map"][k]), framework="np") as w:
|
||||
tensor = w.get_tensor(k)
|
||||
tensor = paddle.to_tensor(tensor)
|
||||
if tensor.dtype == paddle.bfloat16 and self.merge_config.device == "cpu":
|
||||
is_bf16 = True
|
||||
tensor = tensor.astype("float32")
|
||||
tensor_list.append(tensor)
|
||||
tensor_list.append(w.get_tensor(k))
|
||||
if self.merge_config.base_model_path is not None:
|
||||
with fast_safe_open(
|
||||
os.path.join(self.merge_config.base_model_path, index_list[-1]["weight_map"][k]),
|
||||
framework="np",
|
||||
) as w:
|
||||
base_tensor = w.get_tensor(k)
|
||||
base_tensor = paddle.to_tensor(base_tensor)
|
||||
tensor_list = [tensor - base_tensor for tensor in tensor_list]
|
||||
|
||||
merge_tensor = self.merge_method.merge(tensor_list)
|
||||
|
||||
if self.merge_config.base_model_path is not None:
|
||||
merge_tensor += base_tensor
|
||||
if is_bf16:
|
||||
merge_tensor = merge_tensor.astype("bfloat16")
|
||||
merge_state_dict[k] = merge_tensor.numpy()
|
||||
|
||||
del tensor_list
|
||||
paddle.device.cuda.empty_cache()
|
||||
if self.merge_config.base_model_path is not None:
|
||||
del base_tensor
|
||||
paddle.device.cuda.empty_cache()
|
||||
|
||||
tensor_list.append(w.get_tensor(k))
|
||||
is_bf16 = str(tensor_list[0].dtype) == "uint16"
|
||||
tensor_mem = int(np.prod(tensor_list[0].shape) * self.numpy_dtype_map[str(tensor_list[0].dtype)]) / (
|
||||
1024**3
|
||||
)
|
||||
if tensor_mem > self.merge_config.max_tensor_mem:
|
||||
tensor_split_list = [
|
||||
np.array_split(tensor, self.merge_config.split_pieces, axis=0) for tensor in tensor_list
|
||||
]
|
||||
merge_split = []
|
||||
for sp in range(self.merge_config.split_pieces):
|
||||
tensor_list = [tensor_split[sp] for tensor_split in tensor_split_list]
|
||||
if is_bf16:
|
||||
tensor_list = [
|
||||
paddle.Tensor(tensor, zero_copy=True).astype("float32") for tensor in tensor_list
|
||||
]
|
||||
else:
|
||||
tensor_list = [paddle.Tensor(tensor, zero_copy=True) for tensor in tensor_list]
|
||||
if self.merge_config.base_model_path is not None:
|
||||
base_tensor = tensor_list.pop()
|
||||
tensor_list = [tensor - base_tensor for tensor in tensor_list]
|
||||
merge_tensor = self.merge_method.merge(tensor_list)
|
||||
if self.merge_config.base_model_path is not None:
|
||||
merge_tensor += base_tensor
|
||||
if is_bf16:
|
||||
merge_split.append(merge_tensor.astype("bfloat16").numpy())
|
||||
else:
|
||||
merge_split.append(merge_tensor.numpy())
|
||||
merge_state_dict[k] = np.concatenate(merge_split, axis=0)
|
||||
else:
|
||||
if is_bf16:
|
||||
tensor_list = [paddle.Tensor(tensor, zero_copy=True).astype("float32") for tensor in tensor_list]
|
||||
else:
|
||||
tensor_list = [paddle.Tensor(tensor, zero_copy=True) for tensor in tensor_list]
|
||||
if self.merge_config.base_model_path is not None:
|
||||
base_tensor = tensor_list.pop()
|
||||
tensor_list = [tensor - base_tensor for tensor in tensor_list]
|
||||
merge_tensor = self.merge_method.merge(tensor_list)
|
||||
if self.merge_config.base_model_path is not None:
|
||||
merge_tensor += base_tensor
|
||||
if is_bf16:
|
||||
merge_state_dict[k] = merge_tensor.astype("bfloat16").numpy()
|
||||
else:
|
||||
merge_state_dict[k] = merge_tensor.numpy()
|
||||
save_file(
|
||||
merge_state_dict,
|
||||
os.path.join(self.merge_config.output_path, shard_file),
|
||||
metadata={"format": "np"},
|
||||
)
|
||||
del merge_state_dict
|
||||
paddle.device.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
def check_model_path(self, model_path):
|
||||
def check_model_path(self, model_path, lora_merge=False):
|
||||
if os.path.exists(os.path.join(model_path, self.safe_index_name())):
|
||||
with open(os.path.join(model_path, self.safe_index_name()), "r", encoding="utf-8") as f:
|
||||
index = json.load(f)
|
||||
safe_file_list = list(set(index["weight_map"][k] for k in index["weight_map"]))
|
||||
for i in range(len(safe_file_list)):
|
||||
if os.path.exists(os.path.join(model_path, safe_file_list[i])):
|
||||
continue
|
||||
else:
|
||||
ValueError(f"Not found {os.path.join(model_path, safe_file_list[i])}.")
|
||||
file_type = ["safetensors", False]
|
||||
file_type = "safetensors"
|
||||
elif os.path.exists(os.path.join(model_path, self.safe_weight_name())):
|
||||
file_type = ["safetensors_without_index", False]
|
||||
file_type = "safetensors_without_index"
|
||||
elif os.path.exists(os.path.join(model_path, self.weight_name())):
|
||||
file_type = ["pdparams", False]
|
||||
|
||||
# lora
|
||||
elif os.path.exists(os.path.join(model_path, SAFE_PEFT_WEIGHTS_INDEX_NAME)):
|
||||
with open(os.path.join(model_path, SAFE_PEFT_WEIGHTS_INDEX_NAME), "r", encoding="utf-8") as f:
|
||||
index = json.load(f)
|
||||
safe_file_list = list(set(index["weight_map"][k] for k in index["weight_map"]))
|
||||
for i in range(len(safe_file_list)):
|
||||
if os.path.exists(os.path.join(model_path, safe_file_list[i])):
|
||||
continue
|
||||
else:
|
||||
ValueError(f"Not found {os.path.join(model_path, safe_file_list[i])}.")
|
||||
file_type = ["safetensors", True]
|
||||
elif os.path.exists(os.path.join(model_path, SAFE_PEFT_WEIGHTS_NAME)):
|
||||
file_type = ["safetensors_without_index", True]
|
||||
elif os.path.exists(os.path.join(model_path, LORA_WEIGHTS_NAME)):
|
||||
file_type = ["pdparams", True]
|
||||
file_type = "pdparams"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Please check path {model_path} is correct. Support safetensors and pdparams in complete parameter format (not TP or PP format) only."
|
||||
f"Please check path {model_path} is correct. Support safetensors and pdparams only in complete parameter format (not TP or PP format) only."
|
||||
)
|
||||
return file_type
|
||||
|
||||
def check_lora_model_path(self, model_path):
|
||||
if os.path.exists(os.path.join(model_path, SAFE_PEFT_WEIGHTS_INDEX_NAME)):
|
||||
file_type = "lora_safetensors"
|
||||
elif os.path.exists(os.path.join(model_path, LORA_WEIGHTS_NAME)):
|
||||
file_type = "lora_pdparams"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Please check lora path {model_path} is correct. Support safetensors and pdparams only in complete parameter format (not TP or PP format) only."
|
||||
)
|
||||
return file_type
|
||||
|
||||
@@ -460,3 +503,185 @@ class MergeModel:
|
||||
return SAFE_WEIGHTS_INDEX_NAME
|
||||
else:
|
||||
return SAFE_MASTER_WEIGHTS_INDEX_NAME
|
||||
|
||||
def merge_lora_model(self):
|
||||
# Check model file type
|
||||
file_type_list = []
|
||||
file_type_list.append(self.check_lora_model_path(self.merge_config.lora_model_path))
|
||||
file_type_list.append(self.check_model_path(self.merge_config.base_model_path))
|
||||
# Merge model (distinguish between safetensors and pdparams)
|
||||
if "safetensors" in file_type_list[-1]:
|
||||
self.merge_safetensor_lora_model(file_type_list)
|
||||
else:
|
||||
self.merge_pdparams_lora_model(file_type_list)
|
||||
|
||||
def shard_lora_merge(self, key_list, base_index, shard_file, lora_config, file_type_list):
|
||||
merge_state_dict = {}
|
||||
base_state_dict = self.get_model_state_dict(self.merge_config.base_model_path, file_type_list[1], key_list)
|
||||
lora_state_dict = self.get_model_state_dict(self.merge_config.lora_model_path, file_type_list[0], key_list)
|
||||
if not lora_config.rslora:
|
||||
scaling = lora_config.lora_alpha / lora_config.r
|
||||
else:
|
||||
scaling = lora_config.lora_alpha / math.sqrt(lora_config.r)
|
||||
|
||||
for k in key_list:
|
||||
if lora_state_dict is not None and k in lora_state_dict.keys():
|
||||
tensor = lora_state_dict[k]
|
||||
else:
|
||||
tensor = base_state_dict[k]
|
||||
if "weight" in k:
|
||||
lora_A_key, lora_B_key = k.replace("weight", "lora_A"), k.replace("weight", "lora_B")
|
||||
lora_A_tensor = None
|
||||
if lora_state_dict is not None and lora_A_key in lora_state_dict.keys():
|
||||
lora_A_tensor, lora_B_tensor = lora_state_dict[lora_A_key], lora_state_dict[lora_B_key]
|
||||
is_bf16 = tensor.dtype == np.uint16
|
||||
tensor = paddle.Tensor(tensor, zero_copy=True)
|
||||
lora_A_tensor = paddle.Tensor(lora_A_tensor, zero_copy=True)
|
||||
lora_B_tensor = paddle.Tensor(lora_B_tensor, zero_copy=True)
|
||||
if self.is_cpu and is_bf16:
|
||||
tensor = tensor.astype("float32")
|
||||
lora_A_tensor = lora_A_tensor.astype("float32")
|
||||
lora_B_tensor = lora_B_tensor.astype("float32")
|
||||
tensor += lora_A_tensor @ lora_B_tensor * scaling
|
||||
tensor = tensor.astype("bfloat16").numpy()
|
||||
else:
|
||||
tensor += lora_A_tensor @ lora_B_tensor * scaling
|
||||
tensor = tensor.numpy()
|
||||
merge_state_dict[k] = tensor
|
||||
save_file(
|
||||
merge_state_dict,
|
||||
os.path.join(self.merge_config.output_path, shard_file),
|
||||
metadata={"format": "np"},
|
||||
)
|
||||
|
||||
def merge_safetensor_lora_model(self, file_type_list):
|
||||
# Load index
|
||||
base_index = self.get_safetensor_index(self.merge_config.base_model_path, file_type_list[-1])
|
||||
lora_config = LoRAConfig.from_pretrained(self.merge_config.lora_model_path)
|
||||
|
||||
# Initialize new index
|
||||
index = {}
|
||||
index["metadata"] = base_index["metadata"]
|
||||
index["metadata"]["total_size"] = int(index["metadata"]["total_size"])
|
||||
index["weight_map"] = {}
|
||||
|
||||
# LoRA Merge
|
||||
key_list = list(base_index["weight_map"].keys())
|
||||
if not self.is_cpu:
|
||||
rank = dist.get_rank()
|
||||
divided_key_list = divide_lora_key_list(key_list, dist.get_world_size(), lora_config)
|
||||
local_keys = divided_key_list[rank]
|
||||
shard_file = f"{self.merge_config.merge_prefix}-{rank+1:05d}-of-{dist.get_world_size():05d}.safetensors"
|
||||
self.shard_lora_merge(local_keys, base_index, shard_file, lora_config, file_type_list)
|
||||
|
||||
for i in range(len(divided_key_list)):
|
||||
shard_file = f"{self.merge_config.merge_prefix}-{i+1:05d}-of-{dist.get_world_size():05d}.safetensors"
|
||||
for k in divided_key_list[i]:
|
||||
index["weight_map"][k] = shard_file
|
||||
else:
|
||||
divided_key_list = divide_lora_key_list(key_list, self.merge_config.n_process, lora_config)
|
||||
threads = []
|
||||
for i in range(len(divided_key_list)):
|
||||
shard_file = (
|
||||
f"{self.merge_config.merge_prefix}-{i+1:05d}-of-{self.merge_config.n_process:05d}.safetensors"
|
||||
)
|
||||
t = Process(
|
||||
target=self.shard_lora_merge,
|
||||
args=(
|
||||
divided_key_list[i], # key_list
|
||||
base_index, # base index
|
||||
shard_file, # shard_file name
|
||||
lora_config,
|
||||
file_type_list,
|
||||
),
|
||||
)
|
||||
threads.append(t)
|
||||
for k in divided_key_list[i]:
|
||||
index["weight_map"][k] = shard_file
|
||||
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
# Save safe index file
|
||||
if paddle.distributed.get_rank() == 0:
|
||||
save_index_file = os.path.join(self.merge_config.output_path, self.safe_index_name())
|
||||
with open(save_index_file, "w", encoding="utf-8") as f:
|
||||
f.write(json.dumps(index, indent=2) + "\n")
|
||||
self.merge_config.save_pretrained(self.merge_config.output_path)
|
||||
|
||||
def merge_pdparams_lora_model(self, file_type_list):
|
||||
# Load & check state dict
|
||||
lora_state_dict = self.get_model_state_dict(self.merge_config.lora_model_path, file_type_list[0])
|
||||
base_state_dict = self.get_model_state_dict(self.merge_config.base_model_path, file_type_list[1])
|
||||
for key in lora_state_dict.keys():
|
||||
if "lora_A" in key:
|
||||
if key.replace("lora_A", "lora_B") not in lora_state_dict.keys():
|
||||
raise ValueError(f"{key} is not paired with {key.replace('lora_A', 'lora_B')}")
|
||||
if key.replace("lora_A", "weight") not in base_state_dict.keys():
|
||||
raise ValueError(f'{key.replace("lora_A", "weight")} does not exist in base model.')
|
||||
|
||||
# Load lora config
|
||||
lora_config = LoRAConfig.from_pretrained(self.merge_config.lora_model_path)
|
||||
if not lora_config.rslora:
|
||||
scaling = lora_config.lora_alpha / lora_config.r
|
||||
else:
|
||||
scaling = lora_config.lora_alpha / math.sqrt(lora_config.r)
|
||||
|
||||
# Create index
|
||||
merge_state_dict = {}
|
||||
index = {"metadata": {"total_size": 0}, "weight_map": {}}
|
||||
key_list = list(base_state_dict.keys())
|
||||
positions = divide_positions(len(key_list), dist.get_world_size())
|
||||
for ii in range(len(positions) - 1):
|
||||
shard_file = f"{self.merge_config.merge_prefix}-{ii+1:05d}-of-{dist.get_world_size():05d}.safetensors"
|
||||
for key in key_list[positions[ii] : positions[ii + 1]]:
|
||||
index["weight_map"][key] = shard_file
|
||||
index["metadata"]["total_size"] += int(
|
||||
np.prod(base_state_dict[key].shape) * self.numpy_dtype_map[str(base_state_dict[key].dtype)]
|
||||
)
|
||||
|
||||
# Merge state dict
|
||||
rank = dist.get_rank()
|
||||
local_keys = key_list[positions[rank] : positions[rank + 1]]
|
||||
for k in local_keys:
|
||||
if k in lora_state_dict.keys():
|
||||
tensor = lora_state_dict[k]
|
||||
else:
|
||||
tensor = base_state_dict[k]
|
||||
if "weight" in k:
|
||||
lora_A_key, lora_B_key = k.replace("weight", "lora_A"), k.replace("weight", "lora_B")
|
||||
if lora_A_key in lora_state_dict.keys():
|
||||
lora_A_tensor = lora_state_dict[lora_A_key]
|
||||
lora_B_tensor = lora_state_dict[lora_B_key]
|
||||
is_bf16 = tensor.dtype == np.uint16
|
||||
tensor = paddle.Tensor(tensor, zero_copy=True)
|
||||
lora_A_tensor = paddle.Tensor(lora_A_tensor, zero_copy=True)
|
||||
lora_B_tensor = paddle.Tensor(lora_B_tensor, zero_copy=True)
|
||||
if self.is_cpu and is_bf16:
|
||||
tensor = tensor.astype("float32")
|
||||
lora_A_tensor = lora_A_tensor.astype("float32")
|
||||
lora_B_tensor = lora_B_tensor.astype("float32")
|
||||
tensor += lora_A_tensor @ lora_B_tensor * scaling
|
||||
tensor = tensor.astype("bfloat16")
|
||||
else:
|
||||
tensor += lora_A_tensor @ lora_B_tensor * scaling
|
||||
tensor = tensor.numpy()
|
||||
merge_state_dict[k] = tensor
|
||||
|
||||
# Save safetensor file
|
||||
save_file(
|
||||
merge_state_dict,
|
||||
os.path.join(
|
||||
self.merge_config.output_path,
|
||||
f"{self.merge_config.merge_prefix}-{rank+1:05d}-of-{dist.get_world_size():05d}.safetensors",
|
||||
),
|
||||
metadata={"format": "np"},
|
||||
)
|
||||
# Save index file & merge config file
|
||||
if paddle.distributed.get_rank() == 0:
|
||||
save_index_file = os.path.join(self.merge_config.output_path, self.safe_index_name())
|
||||
with open(save_index_file, "w", encoding="utf-8") as f:
|
||||
f.write(json.dumps(index, indent=2) + "\n")
|
||||
self.merge_config.save_pretrained(self.merge_config.output_path)
|
||||
|
||||
@@ -11,6 +11,7 @@
|
||||
# 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 re
|
||||
|
||||
|
||||
def divide_positions(m, n):
|
||||
@@ -29,3 +30,44 @@ def divide_positions(m, n):
|
||||
positions.append(positions[-1] + base_value)
|
||||
positions.append(m)
|
||||
return positions
|
||||
|
||||
|
||||
def divide_lora_key_list(key_list, n, lora_config):
|
||||
lora_key = []
|
||||
other_key = []
|
||||
for module_name in key_list:
|
||||
if (
|
||||
any(re.fullmatch(target_module, module_name) for target_module in lora_config.target_modules)
|
||||
and "weight" in module_name
|
||||
):
|
||||
lora_key.append(module_name)
|
||||
else:
|
||||
other_key.append(module_name)
|
||||
lora_positions = divide_positions(len(lora_key), n)
|
||||
other_positions = divide_positions(len(other_key), n)
|
||||
divided_key_list = []
|
||||
for i in range(len(lora_positions) - 1):
|
||||
divided_key = (
|
||||
lora_key[lora_positions[i] : lora_positions[i + 1]]
|
||||
+ other_key[other_positions[i] : other_positions[i + 1]]
|
||||
)
|
||||
divided_key_list.append(divided_key)
|
||||
return divided_key_list
|
||||
|
||||
|
||||
def divide_safetensor_key_list(weight_map, n):
|
||||
file_map = {}
|
||||
for key in weight_map:
|
||||
if weight_map[key] in file_map:
|
||||
file_map[weight_map[key]].append(key)
|
||||
else:
|
||||
file_map[weight_map[key]] = [key]
|
||||
file_list = list(file_map.keys())
|
||||
p = divide_positions(len(file_list), n)
|
||||
key_list = []
|
||||
positions = [0]
|
||||
for i in range(n):
|
||||
for file in file_list[p[i] : p[i + 1]]:
|
||||
key_list += file_map[file]
|
||||
positions.append(len(key_list))
|
||||
return key_list, positions
|
||||
|
||||
@@ -33,24 +33,20 @@ class SparsifyMethod:
|
||||
|
||||
def dare(self, tensor):
|
||||
if self.merge_config.tensor_type == "np":
|
||||
mask = np.random.binomial(1, self.merge_config.reserve_p, size=tensor.shape).astype(tensor.dtype)
|
||||
tensor *= mask
|
||||
tensor *= (np.random.rand(*tensor.shape) < self.merge_config.reserve_p).astype(tensor.dtype)
|
||||
if self.merge_config.rescale:
|
||||
tensor /= self.merge_config.reserve_p
|
||||
return tensor
|
||||
elif self.merge_config.tensor_type == "pd":
|
||||
mask = paddle.bernoulli(paddle.full(tensor.shape, self.merge_config.reserve_p, dtype=tensor.dtype))
|
||||
|
||||
tensor *= mask
|
||||
if self.merge_config.rescale:
|
||||
tensor /= self.merge_config.reserve_p
|
||||
mode = "upscale_in_train" if self.merge_config.rescale else "downscale_in_infer"
|
||||
tensor = paddle.nn.functional.dropout(tensor, p=1 - self.merge_config.reserve_p, mode=mode, training=True)
|
||||
return tensor
|
||||
else:
|
||||
raise ValueError(f"Unkonwn tensor type {self.merge_config.tensor_type}")
|
||||
|
||||
def magprune(self, tensor):
|
||||
if self.merge_config.tensor_type == "np":
|
||||
if np.all(tensor == 0):
|
||||
if not np.any(tensor != 0):
|
||||
return tensor
|
||||
drop_p = 1 - self.merge_config.reserve_p
|
||||
# 1: ranking(descending)
|
||||
@@ -72,7 +68,7 @@ class SparsifyMethod:
|
||||
tensor /= 1 - probs
|
||||
return tensor
|
||||
elif self.merge_config.tensor_type == "pd":
|
||||
if paddle.all(tensor == 0):
|
||||
if not paddle.any(tensor != 0):
|
||||
return tensor
|
||||
drop_p = 1 - self.merge_config.reserve_p
|
||||
abs_tensor = paddle.abs(tensor)
|
||||
@@ -84,8 +80,7 @@ class SparsifyMethod:
|
||||
probs = probs * self.merge_config.epsilon / tensor.numel()
|
||||
p_min = drop_p - self.merge_config.epsilon / 2
|
||||
probs += p_min
|
||||
|
||||
mask = paddle.cast(paddle.bernoulli(1 - probs), tensor.dtype)
|
||||
mask = paddle.bernoulli(1 - probs).astype(tensor.dtype)
|
||||
tensor *= mask
|
||||
if self.merge_config.rescale:
|
||||
tensor /= 1 - probs
|
||||
@@ -96,7 +91,6 @@ class SparsifyMethod:
|
||||
def trim(self, tensor):
|
||||
if self.merge_config.tensor_type == "np":
|
||||
shape = tensor.shape
|
||||
org_sum = np.sum(np.abs(tensor))
|
||||
tensor = tensor.flatten()
|
||||
abs_tensor = np.abs(tensor)
|
||||
threshold = np.quantile(abs_tensor, 1 - self.merge_config.reserve_p)
|
||||
@@ -111,14 +105,11 @@ class SparsifyMethod:
|
||||
tensor[abs_tensor < threshold] = 0
|
||||
return tensor.reshape(shape)
|
||||
elif self.merge_config.tensor_type == "pd":
|
||||
shape = tensor.shape
|
||||
org_sum = paddle.sum(paddle.abs(tensor))
|
||||
abs_tensor = paddle.abs(tensor)
|
||||
threshold = paddle.quantile(abs_tensor, 1 - self.merge_config.reserve_p)
|
||||
mask = paddle.cast(abs_tensor >= threshold, tensor.dtype)
|
||||
tensor = tensor * mask
|
||||
|
||||
tensor = paddle.where(abs_tensor < threshold, paddle.zeros_like(tensor), tensor)
|
||||
if self.merge_config.rescale:
|
||||
org_sum = paddle.sum(abs_tensor)
|
||||
new_sum = paddle.sum(paddle.abs(tensor))
|
||||
if org_sum >= 1e-8 and new_sum >= 1e-8:
|
||||
tensor *= org_sum / new_sum
|
||||
|
||||
@@ -20,12 +20,11 @@ from parameterized import parameterized
|
||||
|
||||
from paddlenlp.mergekit import MergeConfig, MergeModel
|
||||
from paddlenlp.transformers import AutoModel
|
||||
from tests.testing_utils import require_gpu
|
||||
|
||||
|
||||
class TestMergeModel(unittest.TestCase):
|
||||
@parameterized.expand([("slerp",), ("della",), ("dare_linear",), ("ties",)])
|
||||
def test_merge_model(self, merge_method):
|
||||
def test_merge_model_np(self, merge_method):
|
||||
with TemporaryDirectory() as tempdir:
|
||||
model = AutoModel.from_pretrained("__internal_testing__/tiny-random-bert", dtype="bfloat16")
|
||||
pd_path = os.path.join(tempdir, "pd_model")
|
||||
@@ -68,20 +67,8 @@ class TestMergeModel(unittest.TestCase):
|
||||
mergekit = MergeModel(merge_config)
|
||||
mergekit.merge_model()
|
||||
|
||||
# test safetensor only with pd
|
||||
merge_config = MergeConfig(
|
||||
merge_method=merge_method,
|
||||
model_path_list=[safe_path, safe_path],
|
||||
output_path=tempdir,
|
||||
n_process=2,
|
||||
tensor_type="pd",
|
||||
)
|
||||
mergekit = MergeModel(merge_config)
|
||||
mergekit.merge_model()
|
||||
|
||||
@parameterized.expand([("slerp",), ("della",), ("dare_linear",), ("ties",)])
|
||||
@require_gpu(2)
|
||||
def test_merge_model_gpu(self, merge_method):
|
||||
def test_merge_model_pd(self, merge_method):
|
||||
with TemporaryDirectory() as tempdir:
|
||||
model = AutoModel.from_pretrained("__internal_testing__/tiny-random-bert", dtype="bfloat16")
|
||||
pd_path = os.path.join(tempdir, "pd_model")
|
||||
@@ -89,13 +76,41 @@ class TestMergeModel(unittest.TestCase):
|
||||
safe_path = os.path.join(tempdir, "safe_model")
|
||||
model.save_pretrained(safe_path, safe_serialization="safetensors")
|
||||
|
||||
# test safetensor only with pd and gpu
|
||||
# test mix
|
||||
merge_config = MergeConfig(
|
||||
merge_method=merge_method, model_path_list=[safe_path, pd_path], output_path=tempdir, tensor_type="pd"
|
||||
)
|
||||
mergekit = MergeModel(merge_config)
|
||||
mergekit.merge_model()
|
||||
|
||||
# test mix with base model
|
||||
merge_config = MergeConfig(
|
||||
merge_method=merge_method,
|
||||
model_path_list=[safe_path, safe_path],
|
||||
model_path_list=[safe_path, pd_path],
|
||||
output_path=tempdir,
|
||||
device="gpu",
|
||||
base_model_path=safe_path,
|
||||
tensor_type="pd",
|
||||
)
|
||||
mergekit = MergeModel(merge_config)
|
||||
mergekit.merge_model()
|
||||
|
||||
# test safetensor only
|
||||
merge_config = MergeConfig(
|
||||
merge_method=merge_method,
|
||||
model_path_list=[safe_path, safe_path],
|
||||
output_path=tempdir,
|
||||
tensor_type="pd",
|
||||
)
|
||||
mergekit = MergeModel(merge_config)
|
||||
mergekit.merge_model()
|
||||
|
||||
# test safetensor only with base model
|
||||
merge_config = MergeConfig(
|
||||
merge_method=merge_method,
|
||||
model_path_list=[safe_path, safe_path],
|
||||
output_path=tempdir,
|
||||
tensor_type="pd",
|
||||
base_model_path=safe_path,
|
||||
)
|
||||
mergekit = MergeModel(merge_config)
|
||||
mergekit.merge_model()
|
||||
|
||||
@@ -511,7 +511,7 @@ def require_paddle_up_to_2_gpus(test_case):
|
||||
def require_gpu(min_gpus: int = 1):
|
||||
def actual_decorator(func):
|
||||
gpu_count = paddle.device.cuda.device_count()
|
||||
|
||||
print("gpu count: ", gpu_count)
|
||||
if gpu_count < min_gpus:
|
||||
return unittest.skip(f"test requires {min_gpus} GPUs")(func)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user