LoRA 适用结构与代码实现

LoRA 最初就是围绕线性层提出的。对于一个线性映射 $W_0 \in \mathbb{R}^{d_{out}\times d_{in}}$，LoRA 用两个低秩矩阵来表示更新：

这样可训练参数从 $d_{out}d_{in}$ 变成 $r(d_{out}+d_{in})$。这也是为什么 Transformer 中的大量投影层天然适合 LoRA：它们本来就是规模很大的线性变换 #Hu et al., 2021。

在实践中，并不是所有 Attention 子矩阵都同样重要。原始 LoRA 论文里常见的是对 attention 的投影矩阵做适配，而后续经验通常发现，v_proj 与 o_proj 是非常常见的起点组合；如果资源允许，再把 FFN 相关层一起纳入，效果往往更稳 #Hu et al., 2021。

Amazon Science 关于 target module 选择的研究进一步说明，只用部分输出投影也可能获得很高性价比，但精度与延迟之间存在明显 trade-off #Amazon-Science-LoRA-targets。

很多人第一次接触 LoRA 时，只会把它看成“attention 适配器”。但从模型功能上看，Attention 更像信息路由器，FFN/MLP 更像语义变换器；后者对知识重写和任务适配同样关键。Sebastian Raschka 也强调，如果只盯着 Key / Value 矩阵，常常会白白损失一部分性能 #Raschka-LoRA-DoRA。

在卷积层上做 LoRA，本质上仍然是在大权重张量上寻找低维更新子空间，只不过卷积核是四维张量而不是二维矩阵。工程上常见的做法，是把卷积更新拆成更小的低秩卷积/线性组合，从 kernel 维度或 channel 维度近似原更新 #Modi-2024-Conv-LoRA。

对扩散模型来说，这一点尤其重要，因为 U-Net 里不只有 attention，还有大量卷积与残差块。也正因此，Stable Diffusion 社区的 LoRA 训练与 NLP LoRA 在“目标模块选择”上会显著不同。

loralib 和 labml.ai 的实现都覆盖了 Embedding 版本，这对词表领域迁移、多语言词表偏移等场景是有意义的 #microsoft-LoRA #labml-LoRA。

但 LayerNorm / RMSNorm 往往参数量本来就极小，直接放开训练比给它套 LoRA 更直接。因此从参数效率角度看，它通常不是 LoRA 的优先目标层。

一个最小可运行的 LoRA 线性层其实只需要四个核心元素：

• 冻结的原始权重 weight
• 低秩矩阵 lora_A
• 低秩矩阵 lora_B
• 缩放系数 alpha / r

下面的实现保留了最关键的工程细节：A 随机初始化、B 零初始化、训练态按残差方式计算、推理态可选择合并权重。

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional


class LoRALinear(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        r: int = 8,
        alpha: Optional[int] = None,
        dropout: float = 0.0,
        bias: bool = True,
    ):
        super().__init__()
        self.r = r
        self.alpha = alpha if alpha is not None else r
        self.scaling = self.alpha / self.r

        self.weight = nn.Parameter(
            torch.empty(out_features, in_features),
            requires_grad=False,
        )
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features), requires_grad=False)
        else:
            self.register_parameter("bias", None)

        self.lora_A = nn.Parameter(torch.empty(r, in_features))
        self.lora_B = nn.Parameter(torch.empty(out_features, r))
        self.lora_dropout = nn.Dropout(dropout)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.lora_A, a=5 ** 0.5)
        nn.init.zeros_(self.lora_B)

    def forward(self, x: torch.Tensor, merge: bool = False):
        if merge:
            delta_w = self.lora_B @ self.lora_A * self.scaling
            return F.linear(x, self.weight + delta_w, self.bias)
        base_out = F.linear(x, self.weight, self.bias)
        lora_out = (x @ self.lora_A.T @ self.lora_B.T) * self.scaling
        return base_out + self.lora_dropout(lora_out)

从训练视角看，LoRA 是“主干路径 + 低秩残差路径”；但从部署视角看，只要把 $BA$ 合并回原始权重，就能得到与显式适配器等价的线性层，从而避免推理时多走一条分支。这也是 LoRA 常被说成“几乎没有额外推理延迟”的原因之一 #Hu-et-al.-2021。

在真实项目里，自己手写 LoRA 更多是为了理解原理；真正上训练任务时，最常见的选择是 HuggingFace PEFT。它把 target modules、rank、alpha、dropout、bias 策略、分层 rank pattern 等都封装到了 LoraConfig 里 #HF-PEFT-LoraConfig。

from peft import LoraConfig, get_peft_model, TaskType

config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM,
)

from transformers import AutoModelForCausalLM
from peft import LoraConfig, get_peft_model, TaskType

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    device_map="auto",
    torch_dtype="bfloat16",
)

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM,
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

实际训练时，如果你再叠加 4-bit 量化，就会自然进入 QLoRA 路线；这时 PEFT 往往与 bitsandbytes 协同使用 #HF-PEFT-LoraConfig。