LoRA微调7B参数大模型全记录：从数据清洗到显存优化

import re
import json
from datasets import Dataset

def clean_legal_text(text):
    # 还原甲方乙方指代
    text = re.sub(r'甲方', '合同方A', text)
    text = re.sub(r'乙方', '合同方B', text)
    # 去除多余空白
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def format_llama3_prompt(instruction, input_text=None, output_text=None):
    # Llama-3的对话模板
    if input_text:
        user_content = f"指令：{instruction}\n上下文：{input_text}"
    else:
        user_content = f"指令：{instruction}"
    if output_text:
        return f"user\n{user_content}assistant\n{output_text}"
    return f"user\n{user_content}assistant\n"

# 读取原始数据
with open('raw_qa.json', 'r') as f:
    raw_data = json.load(f)

formatted_data = []
for item in raw_data:
    question = clean_legal_text(item['question'])
    answer = clean_legal_text(item['answer'])
    context = clean_legal_text(item.get('context', ''))
    prompt = format_llama3_prompt(question, context, answer)
    formatted_data.append({'text': prompt})

# 过滤过短或过长的样本
formatted_data = [d for d in formatted_data if 50 < len(d['text']) < 2048]
print(f"清洗后样本数: {len(formatted_data)}")

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# 量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True
)

# 加载基座模型
model_id = "meta-llama/Meta-Llama-3-8B"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

# 冻结基座参数，准备k-bit训练
model = prepare_model_for_kbit_training(model)

# LoRA配置
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.1,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# 输出：trainable params: 8,388,608 || all params: 8,127,834,112 || trainable%: 0.1032

from trl import SFTTrainer, SFTConfig

training_args = SFTConfig(
    output_dir="./lora-llama3-legal",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    learning_rate=2e-5,
    num_train_epochs=3,
    logging_steps=20,
    save_steps=500,
    eval_steps=500,
    evaluation_strategy="steps",
    warmup_ratio=0.05,
    lr_scheduler_type="cosine",
    bf16=True,
    max_seq_length=512,
    packing=False,  # 不打包，保持样本独立
    report_to="tensorboard",
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=Dataset.from_list(formatted_data[:16000]),
    eval_dataset=Dataset.from_list(formatted_data[16000:]),
    tokenizer=tokenizer,
)
trainer.train()