7B模型LoRA/QLoRA微调实录：显存从24G降到6G的工程化实践

# 核心依赖版本
torch==2.1.0+cu118
transformers==4.36.2
peft==0.7.1
bitsandbytes==0.41.3
datasets==2.16.1
accelerate==0.26.1

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# 关键参数说明
lora_config = LoraConfig(
    r=16,                 # 低秩矩阵维度，我最终选16（8欠拟合，32过拟合且显存+1.2G）
    lora_alpha=32,        # 缩放因子，一般设为r/2或2*r，这里取2*r
    target_modules=["query_key_value"],  # ChatGLM2的注意力层是合并的QKV
    lora_dropout=0.05,    # 防止过拟合，0.1也行
    bias="none",          # 不训练bias
    task_type="CAUSAL_LM",
)

from transformers import BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,                # 4-bit量化
    bnb_4bit_quant_type="nf4",        # NormalFloat，比fp4精度更高
    bnb_4bit_use_double_quant=True,   # 双量化：量化常数的二次量化
    bnb_4bit_compute_dtype=torch.bfloat16,  # 计算时反量化为bf16
)

def build_prompt(sample):
    # 业务工单样本：sample["content"]是工单描述，sample["category"]是意图类别
    prompt = f"""对以下工单内容进行意图分类，类别：报销咨询、账号问题、故障报修、投诉建议、其他。

工单内容：{sample['content']}

分类结果：{sample['category']}"""
    return prompt

# 使用tokenizer处理，注意padding和truncation
def tokenize_function(examples):
    texts = [build_prompt({"content": c, "category": l}) for c, l in zip(examples["content"], examples["label"])]
    tokenized = tokenizer(
        texts,
        padding="max_length",
        truncation=True,
        max_length=512,
        return_tensors="pt"
    )
    # 关键：把label设置成input_ids的拷贝，且将prompt部分的loss屏蔽
    labels = tokenized["input_ids"].clone()
    # 这里简单起见，不屏蔽prompt部分（实际业务中建议屏蔽，否则模型会学习生成prompt）
    tokenized["labels"] = labels
    return tokenized

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./chatglm2-lora-16",
    num_train_epochs=3,
    per_device_train_batch_size=8,      # 24G显存下，QLoRA+gradient_ckpt可以到16
    gradient_accumulation_steps=2,      # 实际batch_size = 8*2 = 16
    gradient_checkpointing=True,        # 必须开启，省约40%显存
    optim="paged_adamw_8bit",           # 8-bit优化器，省显存且稳定
    logging_steps=50,
    save_steps=500,
    learning_rate=2e-4,                 # LoRA一般用1e-4~5e-4，比全参大10倍
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    bf16=True,                          # 如果显卡不支持bf16，用fp16
    max_grad_norm=1.0,
    report_to="tensorboard",
)

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import prepare_model_for_kbit_training
from datasets import load_dataset

# 加载4-bit模型
model = AutoModelForCausalLM.from_pretrained(
    "THUDM/chatglm2-6b",
    quantization_config=bnb_config,
    device_map="auto",          # 自动分配层到GPU
    trust_remote_code=True,     # ChatGLM2需要
)

# 冻结+准备kbit训练
model = prepare_model_for_kbit_training(model)
model.gradient_checkpointing_enable()

# 注入LoRA
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # 输出：trainable params: 8,388,608 || all params: 6,245,275,648 || trainable%: 0.1343

# 数据集加载与训练
dataset = load_dataset("json", data_files="train.jsonl")
train_dataset = dataset["train"].map(tokenize_function, batched=True)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)
trainer.train()

model = prepare_model_for_kbit_training(model)
model.gradient_checkpointing_enable()  # 必须在prepare之后

原模型（未微调）：
工单内容：我的报销单被驳回两次了，系统提示发票模糊，但我重新传了三次还是一样。
分类结果：其他（错误）

LoRA r=16微调后：
工单内容：我的报销单被驳回两次了，系统提示发票模糊，但我重新传了三次还是一样。
分类结果：报销咨询（正确）