7B模型LoRA微调实录：显存占用直降68%与推理效果平衡

transformers==4.36.2
peft==0.7.1
bitsandbytes==0.41.3
accelerate==0.26.1
torch==2.1.2+cu118
datasets==2.16.1

def format_instruction(sample):
    return f"""### 指令：
{sample['instruction']}

### 输入：
{sample['input']}

### 回答：
{sample['output']}"""

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,                       # 4bit量化
    bnb_4bit_quant_type="nf4",               # NF4浮点量化，比int4更稳
    bnb_4bit_use_double_quant=True,          # 双重量化，减少显存
    bnb_4bit_compute_dtype=torch.bfloat16    # 计算时反量化为bf16，避免精度塌方
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf",
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

# 冻结全部参数
for param in model.parameters():
    param.requires_grad = False
    if param.ndim == 1:
        param.data = param.data.to(torch.float32)  # 防止LayerNorm的fp32被量化

# LoRA配置
from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    r=16,                    # 秩：16在效果和参数量间最平衡
    lora_alpha=32,           # 缩放因子：alpha/r=2，经验上最优
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],  # 覆盖所有线性层
    lora_dropout=0.1,        # 防止过拟合
    bias="none",             # 不训练bias，节省显存
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # 输出：trainable params: 4,194,304 || all params: 6,742,732,800 || trainable%: 0.0622

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./lora_results",
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,  # 实际batch=32
    learning_rate=2e-4,
    warmup_steps=100,
    num_train_epochs=3,
    logging_steps=50,
    save_steps=500,
    fp16=True,
    gradient_checkpointing=True,    # 用计算换显存
    optim="paged_adamw_8bit",       # bitsandbytes的8bit优化器
    lr_scheduler_type="cosine",
    max_grad_norm=0.3,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=data_collator,
)
trainer.train()