LoRA微调7B模型从数据到推理：损失曲线与效果对比

{
  "instruction": "查询2023年销售金额超过100万的客户名称",
  "output": "SELECT customer_name FROM sales WHERE year = 2023 AND amount > 1000000"
}

from datasets import load_dataset
from transformers import AutoTokenizer

dataset = load_dataset("json", data_files="sql_data.jsonl")
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True)

def format_chat(example):
    # ChatGLM3使用特殊token： 和 
    messages = [
        {"role": "user", "content": example["instruction"]},
        {"role": "assistant", "content": example["output"]}
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False)
    return {"text": text}

dataset = dataset.map(format_chat)
# 切分训练/验证集
dataset = dataset["train"].train_test_split(test_size=0.05, seed=42)

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer

# 4bit量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",       # 使用nf4量化，比fp4精度更稳定
    bnb_4bit_compute_dtype=torch.bfloat16,  # 计算时使用bf16加速
    bnb_4bit_use_double_quant=True    # 双量化进一步减少显存
)

# 加载基座模型（注意：ChatGLM3需要trust_remote_code）
model = AutoModelForCausalLM.from_pretrained(
    "THUDM/chatglm3-6b",
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

# 为kbit训练准备模型
model = prepare_model_for_kbit_training(model)

# LoRA配置：只调整注意力层的线性投影
lora_config = LoraConfig(
    r=8,                          # 秩，8在效果和参数间平衡
    lora_alpha=16,                # 缩放系数，通常设为r的2倍
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # ChatGLM3的注意力模块名
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, lora_config)

# 训练参数
training_args = TrainingArguments(
    output_dir="./lora_sql_output",
    per_device_train_batch_size=4,   # 3090上4是安全值，梯度累积后等效batch=16
    gradient_accumulation_steps=4,
    num_train_epochs=3,
    learning_rate=2e-4,              # LoRA通常用1e-4到5e-4
    fp16=True,                       # 混合精度加速
    logging_steps=10,
    save_strategy="epoch",
    evaluation_strategy="steps",
    eval_steps=50,
    save_total_limit=2,
    report_to="tensorboard",
    remove_unused_columns=False      # ChatGLM3需要保留原始字段
)

# SFTTrainer自动处理padding和attention_mask
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    tokenizer=tokenizer,
    max_seq_length=512,              # 我们的SQL最长不超过300token
    dataset_text_field="text",
)

trainer.train()

from peft import PeftModel

# 先加载基座模型（需要量化配置与训练时一致）
base_model = AutoModelForCausalLM.from_pretrained(
    "THUDM/chatglm3-6b",
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)
# 加载LoRA adapter
model = PeftModel.from_pretrained(base_model, "./lora_sql_output/checkpoint-5400")

# 推理示例
inputs = tokenizer.apply_chat_template([
    {"role": "user", "content": "查询2024年没有下订单的客户ID"}
], return_tensors="pt").to("cuda")
outputs = model.generate(inputs, max_new_tokens=128, temperature=0.1)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)
# 输出：SELECT customer_id FROM customers WHERE customer_id NOT IN (SELECT customer_id FROM orders WHERE order_date LIKE '2024%')

# 合并与保存LoRA权重（可选）
model = model.merge_and_unload()
model.save_pretrained("./merged_lora_model")