模型训练 | 大模型高效微调 DoRA: LoRA升级版

import torchimport torch.nn as nnimport torch.nn.functional as F
class LoRALayer(nn.Module):def __init__(self, in_dim, out_dim, rank, alpha):super().__init__()std_dev = 1 / torch.sqrt(torch.tensor(rank).float())# 初始化参数 A 和 B，其中 A 是输入特征和排名之间的参数，B 是排名和输出特征之间的参数self.A = nn.Parameter(torch.randn(in_dim, rank) * std_dev)# 参数 Aself.B = nn.Parameter(torch.zeros(rank, out_dim))# 参数 Bself.alpha = alpha# LoRA 中的 alpha 参数
def forward(self, x):# LoRA 的前向传播过程x = self.alpha * torch.matmul(torch.matmul(x, self.A), self.B)# 使用参数 A 和 B 进行计算return x
class LinearWithLoRAMerged(nn.Module):def __init__(self, linear, rank, alpha):super().__init__()self.linear = linear# 原始的线性层self.lora = LoRALayer(linear.in_features, linear.out_features, rank, alpha)# LoRA 层
def forward(self, x):lora = torch.matmul(self.lora.A, self.lora.B)# 组合 LoRA 中的参数 A 和 Bcombined_weight = self.linear.weight + self.lora.alpha * lora.T# 将 LoRA 和原始权重结合起来return F.linear(x, combined_weight, self.linear.bias)# 使用结合后的权重进行线性变换

import torchimport torch.nn as nnimport torch.nn.functional as F
torch.manual_seed(0)

class DoRALayer(nn.Module):def __init__(self, d_in, d_out, rank=4, weight=None, bias=None):super().__init__()
# 初始化权重和偏置if weight is not None:self.weight = nn.Parameter(weight, requires_grad=False)else:self.weight = nn.Parameter(torch.Tensor(d_out, d_in), requires_grad=False)
if bias is not None:self.bias = nn.Parameter(bias, requires_grad=False)else:self.bias = nn.Parameter(torch.Tensor(d_out), requires_grad=False)
# 计算输出维度上的权重矩阵的列的幅度self.m = nn.Parameter(self.weight.norm(p=2, dim=0, keepdim=True))# 初始化 LoRA 的参数 A 和 Bstd_dev = 1 / torch.sqrt(torch.tensor(rank).float())self.lora_A = nn.Parameter(torch.randn(d_out, rank) * std_dev)self.lora_B = nn.Parameter(torch.zeros(rank, d_in))
def forward(self, x):# 计算 LoRA 矩阵lora = torch.matmul(self.lora_A, self.lora_B)# 微调权重adapted = self.weight + lora# 对调整后的权重进行归一化column_norm = adapted.norm(p=2, dim=0, keepdim=True)norm_adapted = adapted / column_norm# 计算最终权重calc_weights = self.m * norm_adapted# 使用最终权重进行线性变换return F.linear(x, calc_weights, self.bias)