LoRA微调7B模型：数据准备、训练配置与效果对比全记录

import json
from datasets import Dataset

def format_chatml(question, answer):
    # Qwen2.5-Instruct使用chatml格式
    conversation = [
        {"role": "system", "content": "你是一个专业的法律助手，回答要准确引用法条。"},
        {"role": "user", "content": question},
        {"role": "assistant", "content": answer}
    ]
    # 使用tokenizer的apply_chat_template，但这里先转为纯文本
    text = "system\n你是一个专业的法律助手...\n"
    text += f"user\n{question}\n"
    text += f"assistant\n{answer}"
    return text

# 加载数据
with open("legal_qa.jsonl", "r") as f:
    data = [json.loads(line) for line in f]

formatted = [format_chatml(d["question"], d["answer"]) for d in data]
dataset = Dataset.from_dict({"text": formatted})
dataset = dataset.train_test_split(test_size=0.1, seed=42)

from transformers import (
    AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
import bitsandbytes as bnb

# 1. 加载4bit量化模型
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct",
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True
)

# 2. 准备kbit训练（冻结原参数，激活gradient checkpointing）
model = prepare_model_for_kbit_training(model)

# 3. LoRA配置
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],  # 只微调这两个模块
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)

# 4. 训练参数
training_args = TrainingArguments(
    output_dir="./qwen-lora-legal",
    per_device_train_batch_size=4,      # A100 80G上可到8，但留余量
    gradient_accumulation_steps=4,       # 实际batch_size=4*4=16
    num_train_epochs=3,
    learning_rate=2e-4,                  # LoRA常用lr，比全参大10倍
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    logging_steps=10,
    save_steps=200,
    evaluation_strategy="steps",
    eval_steps=200,
    report_to="wandb",
    bf16=True,
    remove_unused_columns=False
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    tokenizer=tokenizer,
    data_collator=lambda data: tokenizer.pad(
        [{"input_ids": tokenizer.encode(d["text"], truncation=True, max_length=1024)} for d in data],
        return_tensors="pt",
        padding=True
    )
)

trainer.train()