微调太贵了

你花了两周在 8 张 GPU 上完成了 SFT 微调。现在想针对医疗领域再微调一次。

但全量微调意味着 64M 个参数全部要重新训练。每个新领域都重来一遍？

全量 vs LoRA

# 64M 参数全部更新
optimizer = AdamW(model.parameters())
# 10 个领域 = 10 × 64M

# 仅 ~0.4M 参数可训练
optimizer = AdamW(lora_params)
# 10 个领域 = 10 × 0.4M

Part 1：核心思想

低秩分解——用两个小矩阵近似一个大矩阵的变化量。

低秩分解旁路

不修改原始权重，而是在旁边加一条"小路"

定义：LoRA 在原始权重矩阵 W 旁添加低秩分解矩阵 B·A，等效权重变为 W' = W + B·A。

说白了，原始权重像一条主路，LoRA 像一条绕行的辅路。两条路的输出加在一起就是最终结果。训练时只修辅路，主路不动。

这个比喻只是为了帮你建立直观感受。在计算机内部，LoRA 实际上是两个线性变换的串联——先把维度从 d 压缩到 r（A 矩阵），再从 r 还原到 d（B 矩阵），没有真正的"路"。

参数量对比

rank=16 时，单个被注入权重的 LoRA 仅为 4.2%

Part 2：65 行代码的秘密

整个 LoRA 实现只有 65 行代码——一个类加四个函数。我们逐一理解它们各自解决什么问题。

LoRA 类定义

两个无偏置的线性层串联，仅此而已

class LoRA(nn.Module):
    def __init__(self, in_features, out_features, rank):
        super().__init__()
        self.rank = rank  # 秩，控制低秩矩阵大小
        self.A = nn.Linear(in_features, rank, bias=False)  # (d, r) 降维
        self.B = nn.Linear(rank, out_features, bias=False)  # (r, d) 升维
        # A 高斯初始化
        self.A.weight.data.normal_(mean=0.0, std=0.02)
        # B 全零初始化
        self.B.weight.data.zero_()

    def forward(self, x):
        return self.B(self.A(x))  # x → A(x) → B(A(x))

初始化策略

B=0 保证了训练起点的一致性

这是精心设计的细节：让 LoRA 从"什么都不做"开始，逐渐学会"需要做什么"。

注入 LoRA

遍历所有方形 Linear 层，逐个注入旁路

def apply_lora(model, rank=16):
    for name, module in model.named_modules():
        # 只对方形线性层注入（Q/K/V/O 投影、FFN 层）
        if isinstance(module, nn.Linear) and \
           module.weight.shape[0] == module.weight.shape[1]:
            lora = LoRA(module.weight.shape[0],
                        module.weight.shape[1],
                        rank=rank).to(model.device)
            setattr(module, "lora", lora)
            original_forward = module.forward

            def forward_with_lora(x, layer1=original_forward, layer2=lora):
                return layer1(x) + layer2(x)

            module.forward = forward_with_lora  # monkey-patch

为什么只选方形层？因为 Q/K/V/O 投影和 FFN 层是 (d, d) 的方形矩阵。Embedding 和 lm_head 不是方形，不注入。

Monkey-Patch 机制

不修改原始类定义，运行时替换 forward

被注入的层一览

8 层 Transformer 中哪些层会获得 LoRA？

保存与加载

只保存 LoRA 参数，文件极小（~50KB）

# 保存——只提取 lora 参数
def save_lora(model, path):
    raw_model = getattr(model, '_orig_mod', model)
    state_dict = {}
    for name, module in raw_model.named_modules():
        if hasattr(module, 'lora'):
            lora_state = {
                f'{name}.lora.{k}': v.cpu().half()
                for k, v in module.lora.state_dict().items()
            }
            state_dict.update(lora_state)
    torch.save(state_dict, path)

# 加载——找到对应 LoRA 恢复权重
def load_lora(model, path):
    state_dict = torch.load(path, map_location=model.device)
    for name, module in model.named_modules():
        if hasattr(module, 'lora'):
            lora_state = {
                k.replace(f'{name}.lora.', ''): v
                for k, v in state_dict.items()
                if f'{name}.lora.' in k
            }
            module.lora.load_state_dict(lora_state)

合并权重

推理时把 LoRA 溶入原始权重，零额外开销

def merge_lora(model, lora_path, save_path):
    load_lora(model, lora_path)
    raw_model = getattr(model, '_orig_mod', model)
    state_dict = {}
    for name, module in raw_model.named_modules():
        if isinstance(module, nn.Linear) and '.lora.' not in name:
            state_dict[f'{name}.weight'] = module.weight.data.clone().cpu().half()
            if hasattr(module, 'lora'):
                # W' = W + B @ A
                state_dict[f'{name}.weight'] += \
                    (module.lora.B.weight.data @ module.lora.A.weight.data).cpu().half()
    torch.save(state_dict, save_path)

Part 3：训练流程

LoRA 训练与全量 SFT 共享数据格式和循环结构，但有三处关键差异。

与全量 SFT 的差异

冻结与收集参数

train_lora.py 第 128-145 行

# 加载 SFT 模型，注入 LoRA
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
apply_lora(model)

# 统计参数量
total_params = sum(p.numel() for p in model.parameters())
lora_count = sum(p.numel() for n, p in model.named_parameters() if 'lora' in n)
print(f"总参数: {total_params/1e6:.3f}M | LoRA: {lora_count/1e6:.3f}M | "
      f"占比: {lora_count/total_params*100:.2f}%")

# 冻结非 LoRA 参数
lora_params = []
for name, param in model.named_parameters():
    if 'lora' in name:
        param.requires_grad = True
        lora_params.append(param)
    else:
        param.requires_grad = False

optimizer = optim.AdamW(lora_params, lr=args.learning_rate)

训练循环精讲

与全量 SFT 几乎相同，只改了两处

def train_epoch(epoch, loader, iters, lora_params, start_step=0):
    for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
        # 学习率调度、前向、loss 计算与全量 SFT 一致
        with autocast_ctx:
            res = model(input_ids, labels=labels)
            loss = (res.loss + res.aux_loss) / args.accumulation_steps

        scaler.scale(loss).backward()

        if step % args.accumulation_steps == 0:
            scaler.unscale_(optimizer)
            # ⬇️ 区别 1：梯度裁剪只作用于 LoRA 参数
            torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip)
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad(set_to_none=True)

        if step % args.save_interval == 0 and is_main_process():
            # ⬇️ 区别 2：只保存 LoRA 权重
            save_lora(model, lora_save_path)

超参数对比

LoRA vs 全量 SFT 的默认配置

Part 4：使用场景

什么时候用 LoRA，什么时候全量微调？

LoRA 适用场景

LoRA 数据格式

与 SFT 相同的 JSONL 格式

{"conversations": [
    {"role": "user", "content": "头痛应该吃什么药？"},
    {"role": "assistant", "content": "头痛的原因很多，建议先就医确诊。如果是普通感冒引起的..."}
]}

完整训练命令

# LoRA 微调（医疗领域）
$ python trainer/train_lora.py \
    --from_weight full_sft \
    --data_path ../dataset/lora_medical.jsonl \
    --lora_name lora_medical \
    --epochs 10 --learning_rate 1e-4

# 合并 LoRA 到基础模型（推理用）
$ python -c "
from model.model_lora import merge_lora
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
config = MiniMindConfig()
model = MiniMindForCausalLM(config)
merge_lora(model, '../out/lora_medical_768.pth', '../out/merged_768.pth')
"

Part 5：排错与练习

LoRA 训练中你可能遇到的坑。

常见错误诊断

动手练习

知识全景图