2. LoRA与全量微调对比

任务目标

在(124M参数)上分别执行全量微调(Full Fine-tuning)与LoRA微调,对比两者在显存占用、可训练参数量、训练速度及下游任务效果上的差异。本案例使用transformers、peft与datasets库,在E2E NLG数据集的一个小子集上完成训练与评估,给出可复现的对比结论。

环境依赖

  • Python ≥ 3.9
  • torch ≥ 2.0
  • transformers ≥ 4.30
  • peft ≥ 0.5.0
  • datasets ≥ 2.12
  • nvidia-smi(用于显存观测)

安装命令:pip install torch transformers peft datasets

核心实现

以下代码完成:加载与E2E数据子集,分别构建全量微调模型和LoRA模型(r=8),各训练200步,记录显存峰值与训练耗时,并在验证集上计算BLEU分数作为效果指标。

import os
import time
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
from peft import LoraConfig, get_peft_model, TaskType
from datasets import load_dataset
from evaluate import load as load_metric
import numpy as np

os.environ["TOKENIZERS_PARALLELISM"] = "false"
device = "cuda" if torch.cuda.is_available() else "cpu"
MODEL_NAME = ""
MAX_LEN = 128
TRAIN_STEPS = 200
BATCH_SIZE = 4

# ---------- 数据准备 ----------
tokenizer = GPT2Tokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_ = tokenizer.eos_

def tokenize_fn(examples):
    texts = [t + " => " + r for t, r in zip(examples["input_text"], examples["target_text"])]
    enc = tokenizer(texts, truncation=True, max_length=MAX_LEN, padding="max_length")
    enc["labels"] = enc["input_ids"].copy()
    return enc

dataset = load_dataset("e2e_nlg", split="train").select(range(500))
split_ds = dataset.train_test_split(test_size=0.1, seed=42)
train_ds = split_ds["train"].map(tokenize_fn, batched=True, remove_columns=dataset.column_names)
eval_ds = split_ds["test"].map(tokenize_fn, batched=True, remove_columns=dataset.column_names)

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

# ---------- 模型构建 ----------
def build_full_model():
    return GPT2LMHeadModel.from_pretrained(MODEL_NAME).to(device)

def build_lora_model():
    base = GPT2LMHeadModel.from_pretrained(MODEL_NAME)
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=8,
        lora_alpha=16,
        lora_dropout=0.1,
        target_modules=["c_attn", "c_proj"]
    )
    model = get_peft_model(base, lora_config)
    model.print_trainable_parameters()
    return model.to(device)

# ---------- 训练与评估 ----------
def train_and_evaluate(model, run_name):
    train_args = TrainingArguments(
        output_dir=f"./tmp_{run_name}",
        max_steps=TRAIN_STEPS,
        per_device_train_batch_size=BATCH_SIZE,
        per_device_eval_batch_size=BATCH_SIZE,
        logging_steps=50,
        evaluation_strategy="steps",
        eval_steps=100,
        save_strategy="no",
        report_to=[],
        remove_unused_columns=False,
    )
    trainer = Trainer(
        model=model,
        args=train_args,
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        data_collator=data_collator,
    )
    # 显存峰值监控
    torch.cuda.reset_peak_memory_stats()
    start = time.time()
    trainer.train()
    elapsed = time.time() - start
    peak_mem = torch.cuda.max_memory_allocated() / 1024**3 if torch.cuda.is_available() else 0.0

    # 验证集BLEU
    model.eval()
    preds, refs = [], []
    for batch in torch.utils.data.DataLoader(eval_ds, batch_size=4, collate_fn=data_collator):
        input_ids = batch["input_ids"].to(device)
        with torch.no_grad():
            out = model.generate(input_ids, max_new_tokens=40, pad__id=tokenizer.eos__id)
        for i, o in enumerate(out):
            pred_text = tokenizer.decode(o[input_ids.shape[1]:], skip_special_tokens=True)
            ref_text = tokenizer.decode(batch["labels"][i], skip_special_tokens=True)
            if pred_text.strip():
                preds.append(pred_text.strip())
                refs.append([ref_text.strip()])
    bleu = load_metric("bleu")
    score = bleu.compute(predictions=preds, references=refs)["bleu"] if preds else 0.0

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f"\n=== {run_name} ===")
    print(f"可训练参数量: {trainable:,} / 总参数量: {total:,} ({100*trainable/total:.2f}%)")
    print(f"训练耗时: {elapsed:.1f}s | 峰值显存: {peak_mem:.2f} GB | BLEU: {score:.4f}")
    return {"trainable": trainable, "total": total, "time": elapsed, "mem": peak_mem, "bleu": score}

if __name__ == "__main__":
    print(">>> 全量微调")
    full_model = build_full_model()
    full_result = train_and_evaluate(full_model, "full_ft")

    print("\n>>> LoRA微调")
    lora_model = build_lora_model()
    lora_result = train_and_evaluate(lora_model, "lora_ft")

    print("\n========== 对比汇总 ==========")
    print(f"{'指标':<12}{'全量微调':<18}{'LoRA':<18}")
    print(f"{'可训练参数':<10}{full_result['trainable']:<18,}{lora_result['trainable']:<18,}")
    print(f"{'参数量占比':<10}{100*full_result['trainable']/full_result['total']:<18.2f}{100*lora_result['trainable']/lora_result['total']:<18.2f}")
    print(f"{'训练耗时(s)':<10}{full_result['time']:<18.1f}{lora_result['time']:<18.1f}")
    print(f"{'峰值显存(GB)':<10}{full_result['mem']:<18.2f}{lora_result['mem']:<18.2f}")
    print(f"{'BLEU':<10}{full_result['bleu']:<18.4f}{lora_result['bleu']:<18.4f}")

运行方式

将上述代码保存为compare_ft.py,在具备NVIDIA GPU(建议≥8GB显存)的环境中执行:

python compare_ft.py

若GPU显存不足,可将BATCH_SIZE降为2或1,并将TRAIN_STEPS降为100。

验证方法与判定标准

程序运行结束后,观察终端输出的对比汇总表。满足以下条件即判定实验成功:

  • 参数量对比:LoRA的可训练参数量应显著小于全量微调(通常低于1%),而全量微调为100%。
  • 显存对比:LoRA的峰值显存占用应低于全量微调(在相同batch size下通常降低30%~50%)。
  • 速度对比:LoRA的训练耗时通常短于全量微调。
  • 效果对比:LoRA的BLEU分数应与全量微调接近(差距通常在0.05以内),证明低秩适配未显著牺牲性能。

注意:由于仅训练200步且数据量有限,BLEU值会偏低(约0.1~0.3),但两者之间的相对差距才是本案例的观察重点。

常见问题

  • 显存不足(OOM):降低BATCH_SIZE或MAX_LEN,或使用gradient_accumulation_steps保持等效batch size。
  • BLEU为0:检查生成文本是否为空——若pred_text.strip()为空则被过滤,可尝试增加max_new_tokens或使用do_sample=True。
  • 无GPU环境:本案例依赖CUDA进行显存对比,纯CPU环境无法观测显存指标,建议在GPU环境运行。