LLaMA-Factory 中微调算法的原理与显存优化思路

ds_config.json 里我通常这么写：

{
  "train_batch_size": 4,
  "gradient_accumulation_steps": 2,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 5e-5
    }
  },
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu"
    }
  }
}

ZeRO stage 2 把优化器状态和梯度分片，配合 CPU offload 能把显存压得更低，代价是通信多了点，单卡训练时几乎无感。混合精度 fp16 也能省下将近一半显存，前提是显卡支持。

一些特别的优化场景

处理特别长的文本时，梯度检查点（gradient checkpointing）是必备技能：

model.gradient_checkpointing_enable()

用计算换显存，前向过程中中间激活不保存，反向时重新计算。虽然训练慢了，但对于长序列是救命的。如果用了 Flash Attention，还能进一步降低注意力计算的显存：

from llama_factory import enable_flash_attention
enable_flash_attention(model)

多任务场景下，我习惯给每个任务挂一个独立的 LoRA 模块，而不是合在一起：

class MultiTaskLORA(nn.Module):
    def __init__(self, model, num_tasks):
        self.loras = nn.ModuleList([
            LoRAForModel(model, rank=8) for _ in range(num_tasks)
        ])
    def forward(self, x, task_id):
        return self.loras[task_id](x)

Adapter 也类似，可以这样添加：

model.add_adapter("task1", config=AdapterConfig())
model.add_adapter("task2", config=AdapterConfig())

这样切换任务时不用重新加载模型，直接换 adapter 权重就行。

折腾中遇到的坑

OOM 排查
最直接的方式是开一个终端跑 nvidia-smi -l 1 看着显存变化。一旦 OOM，先从 batch_size 开刀，减半再试；不行就降 max_length 和 lora_rank。梯度累积可以模拟大的 batch size 而不增加显存，所以 batch_size 设为 1 配合 4 步累积也是可行的。

训练不收敛
学习率是第一个要怀疑的。AdamW 的 lr 从 2e-5 到 5e-5 是比较安全的范围，配合 warmup 和线性衰减：

optimizer = AdamW(model.parameters(), lr=2e-5)
scheduler = get_linear_schedule_with_warmup(
    optimizer, num_warmup_steps=100, num_training_steps=1000
)

梯度裁剪也能防止 loss 飞掉：

torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

这些配置在 LLaMA-Factory 的 trainer 里基本都有现成参数，直接写到 train.json 就行。

说到底，理解算法的数学背景不是为了炫技，而是帮你在面对具体问题时能做出有理有据的调整。比如看到 lora_rank 就想到低秩分解，看到 gradient checkpointing 就想到用计算换空间。多实验、多观察 nvidia-smi，比背配置有用得多。