2. QLoRA微调LLaMA-7B

任务目标

本案例使用 bitsandbytes 将 LLaMA-7B 模型量化为 4-bit NF4（NormalFloat4）格式，并基于 QLoRA 方法在对话数据集上完成微调，最终验证模型在对话生成任务上的效果。核心目标是跑通“4-bit 量化 → QLoRA 微调 → 对话生成测试”的完整闭环。

环境与依赖

需要以下环境（建议使用支持 CUDA 的 GPU，显存 ≥ 10GB）：

  • Python 3.10+
  • PyTorch 2.0+（CUDA 11.8 或更高）
  • transformers ≥ 4.35.0
  • peft ≥ 0.7.0
  • bitsandbytes ≥ 0.41.0
  • datasets ≥ 2.14.0
  • accelerate ≥ 0.25.0

安装命令：pip install torch transformers peft bitsandbytes datasets accelerate

实现代码

以下代码完整实现了 4-bit NF4 量化加载、QLoRA 配置与微调、以及微调后的对话生成测试。数据集使用 datasets 库内置的 HuggingFaceH4/ultrachat_200k 的子集（取前 500 条用于演示），训练 1 个 epoch。

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 load_dataset

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

# ========== 2. 加载模型与分词器 ==========
model_name = "huggyllama/llama-7b"  # 使用可公开访问的 LLaMA-7B 权重
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
)

# ========== 3. 准备 QLoRA 训练 ==========
model = prepare_model_for_kbit_training(model)

lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)

# ========== 4. 加载并预处理对话数据集 ==========
dataset = load_dataset("HuggingFaceH4/ultrachat_200k", split="train[:500]")

def format_chat(example):
    # 将对话格式化为单轮指令-回复文本
    messages = example["messages"]
    user_msgs = [m["content"] for m in messages if m["role"] == "user"]
    assistant_msgs = [m["content"] for m in messages if m["role"] == "assistant"]
    if not user_msgs or not assistant_msgs:
        return {"text": ""}
    # 取第一轮对话作为训练样本
    text = f"User: {user_msgs[0]}\nAssistant: {assistant_msgs[0]}{tokenizer.eos_}"
    return {"text": text}

dataset = dataset.map(format_chat, remove_columns=dataset.column_names)
dataset = dataset.filter(lambda x: len(x["text"]) > 0)

def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, max_length=512)

tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])

# ========== 5. 训练配置与执行 ==========
training_args = TrainingArguments(
    output_dir="./qlora-llama7b-checkpoints",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=8,
    num_train_epochs=1,
    learning_rate=2e-4,
    fp16=True,
    logging_steps=10,
    save_steps=100,
    report_to=[],
    remove_unused_columns=False
)

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

trainer.train()

# ========== 6. 保存微调后的模型 ==========
model.save_pretrained("./qlora-llama7b-lora")
tokenizer.save_pretrained("./qlora-llama7b-lora")

# ========== 7. 对话生成测试 ==========
def generate_response(prompt, max_new_tokens=100):
    model.eval()
    inputs = tokenizer(f"User: {prompt}\nAssistant:", return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad__id=tokenizer.eos__id
        )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # 提取 Assistant 部分
    if "Assistant:" in response:
        response = response.split("Assistant:", 1)[1].strip()
    return response

# 测试示例
test_prompts = [
    "What is the capital of France?",
    "Explain quantum computing in simple terms."
]
for prompt in test_prompts:
    print(f"Prompt: {prompt}")
    print(f"Response: {generate_response(prompt)}")
    print("-" * 50)

运行方式

  1. 将上述代码保存为 qlora_llama7b.py。
  2. 在 GPU 环境下执行：python qlora_llama7b.py。
  3. 首次运行会自动下载模型权重（约 13GB，需网络连接），训练过程约需 10-20 分钟（取决于 GPU 性能）。

验证方法与通过标准

  • 训练日志检查：训练过程中 loss 应逐步下降（例如从 2.0 左右降至 1.0 以下），且无 NaN 或报错。
  • 模型保存检查：训练完成后，./qlora-llama7b-lora 目录下应包含 adapter_config.json 和 adapter_model.bin 文件。
  • 生成质量测试：运行脚本末尾的测试提示词，观察输出。通过标准为：模型能生成与提示词相关的、语法通顺的英文回复（例如对 “What is the capital of France?” 应回答 “Paris” 或包含 Paris 的完整句子）。
  • 显存占用验证：训练时显存占用应低于 10GB（4-bit 量化后约 6-8GB），若显存溢出则说明量化未生效或 batch size 过大。

常见问题

  • 模型下载失败：若无法访问 huggyllama/llama-7b，可替换为其他 LLaMA-7B 镜像（如 decapoda-research/llama-7b-hf），但需确认该仓库允许访问。
  • bitsandbytes 不支持当前 GPU：需 CUDA 11.8+ 且 GPU 算力 ≥ 7.5（如 RTX 30/40 系列）。
  • 训练 loss 不下降：检查学习率是否过大（建议 1e-4 至 3e-4），或数据集预处理是否正确（确认 text 字段非空）。
← 上一章 📖 返回目录 下一章 →