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")
运行方式
- 将上述代码保存为
train_qlora.py,在GPU环境执行:python train_qlora.py - 训练完成后,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。