从零开始编写 LoRA 代码：原理与实战指南

从零开始编写 LoRA 代码：原理与实战指南

LoRA（Low-Rank Adaptation）作为一种用于微调大语言模型（LLM）的流行技术，最初由微软研究团队在论文《LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS》中提出。不同于全量微调，LoRA 不调整神经网络的所有参数，而是专注于更新一小部分低秩矩阵，从而大幅减少了训练所需的计算量和显存。

由于 LoRA 的微调质量与全模型微调相当，它被视为微调神器。本文基于知名机器学习研究者 Sebastian Raschka 的实践教程，从头构建 LoRA 实现，帮助开发者深入理解其底层逻辑。

LoRA 层的设计原理

用代码表述一个 LoRA 层，核心在于将原始线性层的权重分解为两个低秩矩阵 A 和 B。设输入维度为 in_dim，输出维度为 out_dim，秩为 r。缩放因子 alpha 控制对模型行为的调整幅度。

初始化策略如下：

• 矩阵 A：使用随机分布中的较小值初始化。
• 矩阵 B：使用零初始化。

这使得初始状态下 LoRA 分支的输出为零，不影响预训练模型的原始行为。

PyTorch 实现示例

以下是结合原始线性层和 LoRA 层的 LinearWithLoRA 类实现：

import torch
import torch.nn as nn

class LinearWithLoRA(nn.Module):
    def __init__(self, in_features, out_features, r=8, alpha=1.0, dropout=0.0):
        super().__init__()
        self.r = r
        self.alpha = alpha
        self.scaling_factor = alpha / r
        
        # 冻结原始权重
        self.original_layer = nn.Linear(in_features, out_features)
        for param in self.original_layer.parameters():
            param.requires_grad = False
            
        # 初始化 LoRA 矩阵
        self.lora_A = nn.Parameter(torch.zeros(r, in_features))
        self.lora_B = nn.Parameter(torch.zeros(out_features, r))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        
        self.dropout = nn.Dropout(dropout)

     ():
        original_output = .original_layer(x)
        lora_output = (.lora_B @ .lora_A) @ x.T
         original_output + (lora_output * .scaling_factor).T