4. QLoRA微调Mistral-7B用于指令跟随

任务目标

使用QLoRA技术对Mistral-7B模型进行参数高效微调,在Alpaca格式的指令数据上训练,使模型能够遵循简单指令(如生成回答、执行文本任务)。本案例使用HuggingFace transformers、peft、bitsandbytes 库,在单张24GB显存GPU上可完成训练。

环境依赖

pip install transformers==4.38.2 peft==0.10.0 bitsandbytes==0.43.0 accelerate==0.27.2 datasets==2.17.1

实现代码

import torch
from transformers import (
    AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig,
    TrainingArguments, Trainer, DataCollatorForLanguageModeling
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import Dataset

# 1. 构造最小Alpaca格式训练数据(3条示例,实际使用请替换为完整数据集)
train_data = [
    {"instruction": "将以下句子翻译成英文:", "input": "今天天气很好。", "output": "The weather is nice today."},
    {"instruction": "给出一句鼓励的话。", "input": "", "output": "相信自己,你一定能做到!"},
    {"instruction": "列出三种水果。", "input": "", "output": "苹果、香蕉、橙子。"},
]

# 2. 格式化函数:将Alpaca格式转为模型输入
def format_alpaca(example):
    if example["input"]:
        text = f"### Instruction:\n{example['instruction']}\n\n### Input:\n{example['input']}\n\n### Response:\n{example['output']}"
    else:
        text = f"### Instruction:\n{example['instruction']}\n\n### Response:\n{example['output']}"
    return {"text": text}

dataset = Dataset.from_list(train_data).map(format_alpaca)

# 3. 4-bit量化配置(QLoRA核心)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# 4. 加载模型与分词器
model_name = "mistralai/Mistral-7B-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_ = tokenizer.eos_

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)

# 5. 准备QLoRA训练
model = prepare_model_for_kbit_training(model)
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)

# 6. 数据预处理
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, max_length=512, padding="max_length")

tokenized_dataset = dataset.map(tokenize_function, remove_columns=["text", "instruction", "input", "output"])

# 7. 训练配置
training_args = TrainingArguments(
    output_dir="./mistral-qlora-checkpoints",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=8,
    num_train_epochs=3,
    learning_rate=2e-4,
    fp16=True,
    logging_steps=1,
    save_strategy="epoch",
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
)

# 8. 启动训练
trainer.train()

# 9. 保存LoRA适配器
model.save_pretrained("./mistral-qlora-adapter")
tokenizer.save_pretrained("./mistral-qlora-adapter")

运行方式

  1. 将上述代码保存为 train_qlora.py,在GPU环境执行:python train_qlora.py
  2. 训练完成后,LoRA适配器保存在 ./mistral-qlora-adapter 目录。

推理验证

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

base_model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("./mistral-qlora-adapter")
model = PeftModel.from_pretrained(base_model, "./mistral-qlora-adapter")
model.eval()

def generate_response(instruction, input_text=""):
    if input_text:
        prompt = f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n"
    else:
        prompt = f"### Instruction:\n{instruction}\n\n### Response:\n"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs, max_new_tokens=128, temperature=0.7)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("### Response:\n")[-1]

# 测试示例
print(generate_response("将以下句子翻译成英文:", "今天天气很好。"))

通过条件

  • 训练过程无报错,loss随步数下降(初始约1.5-2.0,3轮后降至0.5以下)。
  • 推理时,模型能对训练集中见过的指令生成与标注输出语义一致的回复。
  • 对未见过但同类型的简单指令(如翻译另一句中文),模型能生成合理回答,说明指令跟随能力已习得。

关键说明

  • 本案例使用3条示例数据仅用于验证训练流程可跑通;实际应用需替换为数百至数千条Alpaca格式数据。
  • QLoRA通过4-bit量化将Mistral-7B的显存占用降至约6-8GB,LoRA只训练约0.1%参数,单卡24GB可训练。
  • 若显存不足,可降低 max_length 至256或减少 gradient_accumulation_steps。