300 lines
9.6 KiB
Python
300 lines
9.6 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
'''
|
||
角色对话LoRA微调训练脚本
|
||
基于test.jsonl数据微调Qwen 8B模型生成游戏NPC对话
|
||
'''
|
||
|
||
import json
|
||
import os
|
||
import torch
|
||
from peft import LoraConfig, PeftModel, TaskType, get_peft_model
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
from transformers import TrainingArguments, Trainer, DataCollatorForSeq2Seq
|
||
from datasets import Dataset
|
||
import platform
|
||
import swanlab
|
||
from swanlab.integration.transformers import SwanLabCallback
|
||
|
||
# Windows multiprocessing兼容性修复
|
||
if platform.system() == "Windows":
|
||
import multiprocessing
|
||
multiprocessing.set_start_method('spawn', force=True)
|
||
|
||
os.environ['VLLM_USE_MODELSCOPE'] = 'True'
|
||
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
|
||
os.environ["TORCH_USE_CUDA_DSA"] = "1"
|
||
|
||
|
||
def process_func(example, tokenizer):
|
||
"""数据预处理函数"""
|
||
MAX_LENGTH = 1024
|
||
|
||
# 构建对话模板 - 专门针对角色对话
|
||
system_prompt = f"你是一个游戏中的NPC角色。{example['character']}"
|
||
instruction = example['instruction']
|
||
user_input = example['input']
|
||
|
||
# 定义输入部分
|
||
instruction = tokenizer(
|
||
f"<s><|im_start|>system\n{system_prompt}<|im_end|>\n"
|
||
f"<|im_start|>user\n{instruction + user_input}<|im_end|>\n"
|
||
f"<|im_start|>assistant\n<think>\n\n</think>\n\n",
|
||
add_special_tokens=False
|
||
)
|
||
|
||
# 定义输出部分
|
||
response = tokenizer(f"{example['output']}", add_special_tokens=False)
|
||
|
||
# 合并输入输出
|
||
input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.pad_token_id]
|
||
attention_mask = instruction["attention_mask"] + response["attention_mask"] + [1]
|
||
|
||
# 标签：只对输出部分计算损失
|
||
labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.pad_token_id]
|
||
|
||
# 截断处理
|
||
if len(input_ids) > MAX_LENGTH:
|
||
input_ids = input_ids[:MAX_LENGTH]
|
||
attention_mask = attention_mask[:MAX_LENGTH]
|
||
labels = labels[:MAX_LENGTH]
|
||
|
||
return {
|
||
"input_ids": input_ids,
|
||
"attention_mask": attention_mask,
|
||
"labels": labels
|
||
}
|
||
|
||
def load_model_and_tokenizer(model_path):
|
||
"""加载模型和分词器"""
|
||
print(f"Loading model from: {model_path}")
|
||
|
||
# 加载分词器
|
||
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True)
|
||
if tokenizer.pad_token is None:
|
||
tokenizer.pad_token = tokenizer.eos_token
|
||
|
||
# 加载模型
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
model_path,
|
||
device_map="auto",
|
||
torch_dtype=torch.bfloat16,
|
||
trust_remote_code=True
|
||
)
|
||
|
||
return model, tokenizer
|
||
|
||
def create_lora_config():
|
||
"""创建LoRA配置"""
|
||
config = LoraConfig(
|
||
task_type=TaskType.CAUSAL_LM,
|
||
target_modules=["q_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj", "o_proj"],
|
||
inference_mode=False,
|
||
r=8, # 增加rank以提高表达能力
|
||
lora_alpha=16, # alpha = 2 * r
|
||
lora_dropout=0.1,
|
||
modules_to_save=["lm_head", "embed_tokens"]
|
||
)
|
||
return config
|
||
|
||
def prepare_dataset(data_path, tokenizer):
|
||
"""准备数据集"""
|
||
print(f"Loading dataset from: {data_path}")
|
||
|
||
# 加载JSON数据
|
||
with open(data_path, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
|
||
print(f"Total samples: {len(data)}")
|
||
|
||
# 转换为Dataset格式
|
||
dataset = Dataset.from_list(data)
|
||
|
||
# 应用预处理函数
|
||
def tokenize_function(examples):
|
||
return process_func(examples, tokenizer)
|
||
|
||
tokenized_dataset = dataset.map(
|
||
tokenize_function,
|
||
remove_columns=dataset.column_names,
|
||
batched=False
|
||
)
|
||
|
||
return tokenized_dataset
|
||
|
||
def train_lora_model(model_path, data_path, output_dir):
|
||
"""训练LoRA模型"""
|
||
|
||
# 1. 加载模型和分词器
|
||
model, tokenizer = load_model_and_tokenizer(model_path)
|
||
|
||
# 2. 创建LoRA配置
|
||
lora_config = create_lora_config()
|
||
|
||
# 3. 应用LoRA
|
||
model = get_peft_model(model, lora_config)
|
||
|
||
# 4. 启用梯度计算
|
||
for param in model.parameters():
|
||
if param.requires_grad:
|
||
param.requires_grad_(True)
|
||
|
||
model.config.use_cache = False # 关闭缓存以节省显存
|
||
|
||
# 5. 准备数据集
|
||
train_preparedataset = prepare_dataset(data_path, tokenizer)
|
||
|
||
# 6. 配置训练参数 - 针对3080显卡优化
|
||
training_args = TrainingArguments(
|
||
output_dir=output_dir,
|
||
per_device_train_batch_size=2, # 减小batch size
|
||
gradient_accumulation_steps=4, # 增加梯度累积
|
||
logging_steps=10,
|
||
num_train_epochs=3, # 增加训练轮数以充分学习角色特征
|
||
save_steps=50,
|
||
learning_rate=5e-5, # 稍微提高学习率
|
||
warmup_ratio=0.1,
|
||
max_grad_norm=1.0,
|
||
save_on_each_node=True,
|
||
gradient_checkpointing=True,
|
||
gradient_checkpointing_kwargs={"use_reentrant": True},
|
||
dataloader_pin_memory=False, # 减少内存使用
|
||
remove_unused_columns=False,
|
||
report_to="none",
|
||
#bf16=True,
|
||
#fp16=True, # 使用混合精度训练
|
||
save_total_limit=3, # 只保留最新的3个检查点
|
||
)
|
||
|
||
#添加swan监测
|
||
swanlab_callback = SwanLabCallback(
|
||
project = "QwenLora_Learn",
|
||
experiment_name="Qwen3-8B-LoRA-experiment"
|
||
)
|
||
swanlab.login(api_key="pAxFTROvv3aspmEijax46")
|
||
# 7. 创建训练器
|
||
trainer = Trainer(
|
||
model=model,
|
||
args=training_args,
|
||
train_dataset=train_preparedataset,
|
||
data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
|
||
callbacks=[swanlab_callback] # 传入之前的swanlab_callback
|
||
)
|
||
|
||
# 8. 开始训练
|
||
print("Starting training...")
|
||
trainer.train()
|
||
|
||
# 9. 保存最终模型
|
||
final_output_dir = os.path.join(output_dir, "final_model")
|
||
trainer.save_model(final_output_dir)
|
||
tokenizer.save_pretrained(final_output_dir)
|
||
|
||
print(f"Training completed! Model saved to: {final_output_dir}")
|
||
return final_output_dir
|
||
|
||
def test_trained_model(model_path, lora_path):
|
||
"""测试训练后的模型"""
|
||
print("Testing trained model...")
|
||
|
||
# 加载基础模型
|
||
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True)
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
model_path,
|
||
device_map="auto",
|
||
torch_dtype=torch.bfloat16,
|
||
trust_remote_code=True
|
||
)
|
||
|
||
# 加载LoRA权重
|
||
model = PeftModel.from_pretrained(model, lora_path)
|
||
|
||
# 测试对话
|
||
test_cases = [
|
||
{
|
||
"system": "你是克莱恩，一位神秘学专家和侦探。",
|
||
"user": "请告诉我一些关于神秘学的知识。"
|
||
},
|
||
{
|
||
"system": "你是阿兹克，经验丰富的神秘学导师。",
|
||
"user": "学生遇到了危险，你会给出什么建议？"
|
||
},
|
||
{
|
||
"system": "你是塔利姆，一个有礼貌的普通人，遇到了困难。",
|
||
"user": "你最近怎么样？"
|
||
}
|
||
]
|
||
|
||
for i, case in enumerate(test_cases):
|
||
messages = [
|
||
{"role": "system", "content": case["system"]},
|
||
{"role": "user", "content": case["user"]}
|
||
]
|
||
|
||
inputs = tokenizer.apply_chat_template(
|
||
messages,
|
||
add_generation_prompt=True,
|
||
tokenize=True,
|
||
return_tensors="pt",
|
||
return_dict=True,
|
||
enable_thinking=False
|
||
).to('cuda')
|
||
|
||
inputs = {k: v.to(model.device) for k, v in inputs.items()}
|
||
|
||
with torch.no_grad():
|
||
logits = model(**inputs).logits
|
||
probs = torch.softmax(logits, dim=-1)
|
||
# 检查非法值
|
||
if torch.isnan(probs).any():
|
||
print("概率张量包含NaN！")
|
||
if torch.isinf(probs).any():
|
||
print("概率张量包含Inf！")
|
||
if (probs < 0).any():
|
||
print("概率张量包含负数！")
|
||
outputs = model.generate(
|
||
**inputs,
|
||
max_new_tokens=200,
|
||
do_sample=True,
|
||
temperature=0.7,
|
||
top_p=0.8,
|
||
pad_token_id=tokenizer.eos_token_id
|
||
)
|
||
|
||
response = outputs[0][inputs['input_ids'].shape[1]:]
|
||
decoded_response = tokenizer.decode(response, skip_special_tokens=True)
|
||
|
||
print(f"\n--- 测试用例 {i+1} ---")
|
||
print(f"系统提示: {case['system']}")
|
||
print(f"用户输入: {case['user']}")
|
||
print(f"模型回复: {decoded_response}")
|
||
|
||
def main():
|
||
# 配置路径
|
||
model_path = '/mnt/e/AI/Project02/AITrain/Qwen/Qwen3-8B-AWQ' # 基础模型路径
|
||
data_path = './npc_dialogue_dataset.json' # 训练数据路径
|
||
output_dir = './output/NPC_Dialogue_LoRA' # 输出目录
|
||
|
||
# #####test
|
||
# final_model_path = os.path.join(output_dir, "final_model")
|
||
# test_trained_model(model_path, final_model_path)
|
||
# 确保数据文件存在
|
||
if not os.path.exists(data_path):
|
||
print(f"数据文件不存在: {data_path}")
|
||
print("请先运行 prepare_dialogue_data.py 生成训练数据")
|
||
return
|
||
|
||
try:
|
||
# 训练模型
|
||
final_model_path = train_lora_model(model_path, data_path, output_dir)
|
||
|
||
# 测试模型
|
||
test_trained_model(model_path, final_model_path)
|
||
|
||
except Exception as e:
|
||
print(f"训练过程中出现错误: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
if __name__ == '__main__':
|
||
main() |