Dora原理代码讲解

ops/2024/10/15 22:12:09/

Dora原理及代码讲解

相关的完整训练应用及验证代码在：
https://github.com/mst272/LLM-Dojo 项目的dora部分：Dora部分

前验证知识：Lora

Lora估计大家已经很熟悉了，在这里就不详细介绍Lora的一些原理的，简而言之就是冻结已经训练好的模型权重，在其线性层或其它层中增加一些参数，只训练这些新增的参数层。
可以直接用这一个图来表示：
在这里插入图片描述

直接用代码来解释一下：

创建LoraLayer

python">import torch.nn as nn
import torch# 构建LoraLayer
class LoRALayer(nn.Module):def __init__(self, in_dim, out_dim, rank,  alpha):super().__init__()std_dev = 1 / torch.sqrt(torch.tensor(rank).float())self.A = nn.Parameter(torch.rand(in_dim, rank)*std_dev)self.B = nn.Parameter(torch.zeros(rank, out_dim))self.alpha = alphadef forward(self, x):x = self.alpha * (x @ self.A @ self.B)return x

将Lora合并到线性层

python">class LinearWithLoRA(nn.Module):def __init__(self, linear, rank, alpha):super().__init__()self.linear = linearself.lora = LoRALayer(linear.in_features,linear.out_features,rank,alpha)def forward(self,x):return self.linear(x) + self.lora(x)

Lora的变体—Dora

具体原理可以在代码中理解，实际运用中将LinearWithLoRA替换为LinearWithDoRA即可使用。

python">class LinearWithDoRA(nn.Module):def __init__(self, linear, rank, alpha):super().__init__()self.linear = linearself.lora = LoRALayer(linear.in_features, linear.out_features, rank, alpha)self.m = nn.Parameter(torch.ones(1, linear.out_features))def forward(self, x):linear_out = self.linear(x)lora_out = self.lora(x)lora_out_norm = lora_out / (lora_out.norm(p=2, dim=1, keepdim=True) + 1e-9)dora_modification = self.m * lora_out_normreturn linear_out + dora_modification

具体如何使用我们可以自己简单的建一个Model进行测试，例如：

python">class TestMLP(nn.Module):def __init__(self, num_features, num_hidden1, num_hidden2, num_class):super().__init__()self.layers = nn.Sequential(nn.Linear(num_features, num_hidden1),nn.ReLU(),nn.Linear(num_hidden1, num_hidden2),nn.ReLU(),nn.Linear(num_hidden2, num_class))def forward(self, x):x = self.layers(x)return x