LoRA微调7B模型全记录：从显存优化到指令遵循能力跃升

{
  "instruction": "解释什么是机器学习",
  "input": "",
  "output": "机器学习是一种人工智能技术..."
}

from datasets import load_dataset
from transformers import AutoTokenizer

def preprocess_function(examples):
    """将alpaca格式转为模型输入"""
    MAX_LENGTH = 512  # 限制序列长度，防止OOM

    # Qwen的chat模板格式
    prompts = []
    for inst, inp in zip(examples["instruction"], examples["input"]):
        if inp:
            prompt = f"user\n{inst}\n{inp}\nassistant\n"
        else:
            prompt = f"user\n{inst}\nassistant\n"
        prompts.append(prompt)

    # 对输出添加结束标记
    outputs = [out + "" for out in examples["output"]]

    # 拼接并tokenize
    full_texts = [p + o for p, o in zip(prompts, outputs)]

    tokenized = tokenizer(
        full_texts,
        max_length=MAX_LENGTH,
        truncation=True,
        padding="max_length",
        return_tensors="pt"
    )

    # 将输出部分的label保留，输入部分设为-100（忽略损失）
    labels = tokenized["input_ids"].clone()
    for i, (prompt, full) in enumerate(zip(prompts, full_texts)):
        prompt_len = len(tokenizer(prompt)["input_ids"])
        labels[i, :prompt_len] = -100  # label为-100的位置不参与loss计算

    tokenized["labels"] = labels
    return tokenized

# 加载数据集
dataset = load_dataset("yahma/alpaca-cleaned", split="train")
tokenized_dataset = dataset.map(
    preprocess_function,
    batched=True,
    remove_columns=dataset.column_names
)

from peft import LoraConfig, get_peft_model, prepare_for_kbit_training
from transformers import BitsAndBytesConfig
import torch

# 4bit量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True  # 双量化进一步省显存
)

# LoRA配置
lora_config = LoraConfig(
    r=16,  # 秩，增大可提升容量但增加显存
    lora_alpha=32,  # 缩放因子，过大会导致不稳定
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # 注意力层的4个投影
    lora_dropout=0.05,  # 防止过拟合
    bias="none",  # 不训练bias
    task_type="CAUSAL_LM"
)

# 加载4bit模型
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2-7B-Instruct",
    quantization_config=bnb_config,
    device_map="auto"  # 自动分配层到不同设备
)

# 将LoRA注入模型
model = prepare_for_kbit_training(model)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# 输出: trainable params: 8,388,608 || all params: 7,742,005,248 || trainable%: 0.1084

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./qwen2-lora",
    per_device_train_batch_size=4,  # 根据显存调整
    gradient_accumulation_steps=8,   # 等效batch size = 4*8 = 32
    num_train_epochs=3,
    learning_rate=2e-4,  # LoRA常用2e-4 ~ 5e-4
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    logging_steps=20,
    save_steps=500,
    eval_strategy="steps",
    eval_steps=200,
    fp16=False,  # 使用bf16更稳定
    bf16=True,
    optim="paged_adamw_8bit",  # 使用8bit优化器省显存
    gradient_checkpointing=True,  # 用计算换显存
    max_grad_norm=0.3,
    report_to="tensorboard"  # 记录loss曲线
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset.select(range(40000)),  # 取4万条训练
    eval_dataset=tokenized_dataset.select(range(40000, 45000)),
)
trainer.train()