LoRA 低秩微调大模型

python import torch import torch.nn as nn import torch.nn.functional as F

### 定义一个 LoRALayer
我们对线性层进行包装：
```python
class LoRALinear(nn.Module):
    def __init__(self, in_features, out_features, r=4, lora_alpha=16, lora_dropout=0.0):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.r = r
        self.lora_alpha = lora_alpha
        self.scaling = lora_alpha / r

        # 原始权重层，冻结
        self.linear = nn.Linear(in_features, out_features, bias=False)
        self.linear.weight.requires_grad = False

        # LoRA 可训练参数
        self.lora_A = nn.Parameter(torch.zeros(r, in_features))
        self.lora_B = nn.Parameter(torch.zeros(out_features, r))
        self.lora_dropout = nn.Dropout(p=lora_dropout)

        # 初始化
        nn.init.kaiming_uniform_(self.lora_A, a=5 ** 0.5)
        nn.init.zeros_(self.lora_B)   # B 初始全零，让 Lora 从零开始

    def forward(self, x):
        result = self.linear(x)                      # 原始输出
        if self.r > 0:
            x = self.lora_dropout(x)
            lora_out = (x @ self.lora_A.T @ self.lora_B.T) * self.scaling
            result += lora_out
        return result

注意：初始化时 $A$ 使用随机分布，$B$ 初始化为零，确保训练开始时 LoRA 分支输出为零，不影响原始模型行为。

应用 LoRA 到线性层

在实际使用时，我们需要替换目标模型中的线性层。以 Hugging Face 的 transformers 库为例，可以定义一个替换函数：

def apply_lora_to_model(model, target_modules, r=8, lora_alpha=32):
    for name, module in model.named_children():
        if isinstance(module, nn.Linear) and name in target_modules:
            in_f, out_f = module.in_features, module.out_features
            lora_linear = LoRALinear(in_f, out_f, r=r, lora_alpha=lora_alpha)
            lora_linear.linear.weight.data = module.weight.data
            setattr(model, name, lora_linear)
        else:
            # 递归处理子模块
            apply_lora_to_model(module, target_modules, r, lora_alpha)

通常 target_modules 可设置为例如 ['q_proj', 'k_proj', 'v_proj', 'o_proj']（具体名称依模型而定）。

训练过程简介

训练循环与普通 PyTorch 训练完全一致：

optimizer = torch.optim.AdamW(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-3
)

# 训练循环
model.train()
for batch in dataloader:
    loss = model(batch)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()