240 lines
9.1 KiB
Python
240 lines
9.1 KiB
Python
import os
|
||
import sys
|
||
__package__ = "scripts"
|
||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||
import torch
|
||
import warnings
|
||
from transformers import AutoTokenizer
|
||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||
import argparse
|
||
warnings.filterwarnings('ignore', category=UserWarning)
|
||
|
||
# MiniMind 配置（需与训练时一致）
|
||
lm_config = MiniMindConfig(
|
||
hidden_size=512,
|
||
num_hidden_layers=8,
|
||
max_position_embeddings=32768,
|
||
vocab_size=6400, # 确保与 tokenizer 一致
|
||
use_moe=False,
|
||
num_attention_heads=8,
|
||
num_key_value_heads=2,
|
||
rms_norm_eps=1e-5,
|
||
rope_theta=10000.0
|
||
)
|
||
|
||
|
||
def merge_lora_weights(model, lora_state_dict):
|
||
"""
|
||
将 LoRA 权重合并到基础模型中
|
||
LoRA 权重格式: {module_name}.lora.A.weight 和 {module_name}.lora.B.weight
|
||
合并公式: W_new = W_old + B @ A
|
||
|
||
注意: 根据 model_lora.py，LoRA 只应用于方阵线性层 (shape[0] == shape[1])
|
||
"""
|
||
# 清理 state_dict 中的 'module.' 前缀（如果有）
|
||
cleaned_lora_dict = {k.replace('module.', ''): v for k, v in lora_state_dict.items()}
|
||
|
||
# 收集所有 LoRA 权重
|
||
lora_weights = {}
|
||
for key in cleaned_lora_dict.keys():
|
||
if '.lora.' in key:
|
||
# 提取模块名称和权重类型
|
||
parts = key.split('.lora.')
|
||
if len(parts) == 2:
|
||
module_name = parts[0]
|
||
weight_type = parts[1] # 'A.weight' 或 'B.weight'
|
||
|
||
if module_name not in lora_weights:
|
||
lora_weights[module_name] = {}
|
||
lora_weights[module_name][weight_type] = cleaned_lora_dict[key]
|
||
|
||
print(f"找到 {len(lora_weights)} 个应用了 LoRA 的模块")
|
||
|
||
# 合并 LoRA 权重到对应的线性层
|
||
merged_count = 0
|
||
skipped_count = 0
|
||
|
||
for name, module in model.named_modules():
|
||
# 只处理线性层
|
||
if isinstance(module, torch.nn.Linear):
|
||
# 检查是否有对应的 LoRA 权重
|
||
if name in lora_weights:
|
||
lora_data = lora_weights[name]
|
||
if 'A.weight' in lora_data and 'B.weight' in lora_data:
|
||
A = lora_data['A.weight'] # shape: [rank, in_features]
|
||
B = lora_data['B.weight'] # shape: [out_features, rank]
|
||
|
||
# 验证形状匹配
|
||
original_weight = module.weight.data # shape: [out_features, in_features]
|
||
expected_out, expected_in = original_weight.shape
|
||
|
||
# A: [rank, in_features], B: [out_features, rank]
|
||
if A.shape[1] != expected_in or B.shape[0] != expected_out:
|
||
print(f" ⚠ 跳过 {name}: 形状不匹配 (A: {A.shape}, B: {B.shape}, 期望: [{expected_out}, {expected_in}])")
|
||
skipped_count += 1
|
||
continue
|
||
|
||
# 合并: W_new = W_old + B @ A
|
||
# B @ A: [out_features, rank] @ [rank, in_features] = [out_features, in_features]
|
||
lora_delta = torch.matmul(B, A) # shape: [out_features, in_features]
|
||
module.weight.data = original_weight + lora_delta.to(original_weight.device)
|
||
merged_count += 1
|
||
print(f" ✓ 合并 LoRA 权重到: {name} (形状: {original_weight.shape}, rank: {A.shape[0]})")
|
||
|
||
print(f"成功合并 {merged_count} 个 LoRA 权重")
|
||
if skipped_count > 0:
|
||
print(f"跳过 {skipped_count} 个不匹配的权重")
|
||
|
||
return model
|
||
|
||
|
||
def merge_lora_to_minimind(lora_model_path, base_model_path, output_path, config=None, dtype=torch.float16):
|
||
"""
|
||
将 MiniMind + LoRA 权重合并并保存为 safetensors 格式
|
||
|
||
Args:
|
||
lora_model_path: LoRA 权重文件路径 (.pth)
|
||
base_model_path: 基础模型文件路径 (.pth)
|
||
output_path: 输出目录路径
|
||
config: MiniMindConfig 配置对象
|
||
dtype: 输出模型的数据类型
|
||
"""
|
||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
|
||
# 使用提供的配置或默认配置
|
||
if config is None:
|
||
config = lm_config
|
||
|
||
print("=" * 60)
|
||
print("开始合并 LoRA 权重到 MiniMind 模型")
|
||
print("=" * 60)
|
||
|
||
# 1. 加载基础模型
|
||
print(f"\n1. 加载基础模型: {base_model_path}")
|
||
model = MiniMindForCausalLM(config)
|
||
|
||
if os.path.isfile(base_model_path):
|
||
state_dict = torch.load(base_model_path, map_location='cpu')
|
||
|
||
# 如果 checkpoint 包含 'model' 键，则提取模型权重
|
||
if isinstance(state_dict, dict) and 'model' in state_dict:
|
||
state_dict = state_dict['model']
|
||
print(" ℹ 检测到完整 checkpoint，提取模型权重")
|
||
|
||
# 清理 state_dict 中的 'module.' 前缀（如果有，用于 DDP 模型）
|
||
cleaned_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
|
||
|
||
# 过滤出模型相关的键（排除 optimizer、scheduler 等）
|
||
model_keys = set(model.state_dict().keys())
|
||
filtered_state_dict = {k: v for k, v in cleaned_state_dict.items() if k in model_keys}
|
||
|
||
missing_keys, unexpected_keys = model.load_state_dict(filtered_state_dict, strict=False)
|
||
print(" ✓ 基础模型加载完成")
|
||
if missing_keys:
|
||
print(f" ⚠ 警告: {len(missing_keys)} 个键未找到（已忽略）")
|
||
if unexpected_keys:
|
||
print(f" ℹ 信息: {len(unexpected_keys)} 个额外键（已忽略）")
|
||
else:
|
||
raise FileNotFoundError(f"基础模型文件不存在: {base_model_path}")
|
||
|
||
# 2. 加载 LoRA 权重
|
||
print(f"\n2. 加载 LoRA 权重: {lora_model_path}")
|
||
if not os.path.isfile(lora_model_path):
|
||
raise FileNotFoundError(f"LoRA 权重文件不存在: {lora_model_path}")
|
||
|
||
lora_state_dict = torch.load(lora_model_path, map_location='cpu')
|
||
print(f" ✓ LoRA 权重加载完成 (包含 {len(lora_state_dict)} 个键)")
|
||
|
||
# 3. 合并 LoRA 权重
|
||
print(f"\n3. 合并 LoRA 权重到基础模型...")
|
||
model = merge_lora_weights(model, lora_state_dict)
|
||
|
||
# 4. 转换精度
|
||
print(f"\n4. 转换模型精度为 {dtype}...")
|
||
model = model.to(dtype)
|
||
model.eval()
|
||
|
||
# 5. 保存模型为 safetensors 格式
|
||
print(f"\n5. 保存模型到: {output_path}")
|
||
os.makedirs(output_path, exist_ok=True)
|
||
|
||
# 注册模型类以便 transformers 能够识别
|
||
MiniMindConfig.register_for_auto_class()
|
||
MiniMindForCausalLM.register_for_auto_class("AutoModelForCausalLM")
|
||
|
||
# 保存模型（使用 safe_serialization=True 保存为 safetensors）
|
||
model.save_pretrained(output_path, safe_serialization=False)
|
||
print(" ✓ 模型已保存为 safetensors 格式")
|
||
|
||
# 6. 保存 tokenizer
|
||
print(f"\n6. 保存 tokenizer...")
|
||
tokenizer_path = os.path.join(os.path.dirname(__file__), '..', 'model')
|
||
if os.path.exists(tokenizer_path):
|
||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
||
tokenizer.save_pretrained(output_path)
|
||
print(" ✓ Tokenizer 保存完成")
|
||
else:
|
||
print(f" ⚠ 警告: Tokenizer 路径不存在: {tokenizer_path}")
|
||
|
||
# 7. 打印参数量
|
||
print(f"\n7. 模型信息:")
|
||
model_params = sum(p.numel() for p in model.parameters())
|
||
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
print(f" 总参数量: {model_params / 1e6:.2f} M = {model_params / 1e9:.2f} B")
|
||
print(f" 可训练参数: {trainable_params / 1e6:.2f} M")
|
||
|
||
print("\n" + "=" * 60)
|
||
print(f"✅ LoRA 已合并并保存到: {output_path}")
|
||
print("=" * 60)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# 配置路径
|
||
parser = argparse.ArgumentParser(description="Merge LoRA weights into base MiniMind model")
|
||
parser.add_argument(
|
||
"--lora_model_path",
|
||
type=str,
|
||
required=True,
|
||
# default="/home/dieu/minimind/out/lora/lora_identity_512.pth",
|
||
help="Path to the LoRA weights (.pth file)"
|
||
)
|
||
parser.add_argument(
|
||
"--base_model_path",
|
||
type=str,
|
||
required=True,
|
||
# default="/home/dieu/minimind/out/full_sft_512.pth",
|
||
help="Path to the base model (.pth file)"
|
||
)
|
||
parser.add_argument(
|
||
"--output_path",
|
||
type=str,
|
||
required=True,
|
||
# default="/home/dieu/minimind/out/merged_model_test",
|
||
help="Directory to save the merged model"
|
||
)
|
||
parser.add_argument(
|
||
"--dtype",
|
||
type=str,
|
||
default="float16",
|
||
choices=["float16", "float32", "bfloat16"],
|
||
help="Data type for the merged model (default: float16)"
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 映射 dtype 字符串到 torch.dtype
|
||
dtype_map = {
|
||
"float16": torch.float16,
|
||
"float32": torch.float32,
|
||
"bfloat16": torch.bfloat16,
|
||
}
|
||
dtype = dtype_map[args.dtype]
|
||
# 调用合并函数
|
||
merge_lora_to_minimind(
|
||
lora_model_path=args.lora_model_path,
|
||
base_model_path=args.base_model_path,
|
||
output_path=args.output_path,
|
||
config=lm_config,
|
||
dtype=dtype
|
||
)
|