DoRA the explora

gistfile1.txt

import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F

# This layer is dropped into your pre-trained PyTorch model where nn.Linear is used
class DoRALayer(nn.Module):
    def __init__(self, d_in, d_out, rank=4):
        super().__init__()

        # Match original layer weights
        self.weight = nn.Parameter(torch.Tensor(d_out, d_in), requires_grad=False)
        self.bias = nn.Parameter(torch.Tensor(d_out), requires_grad=False)

        # m = Magnitude column-wise across output dimension
        self.m = nn.Parameter(self.weight.norm(p=2, dim=0, keepdim=True))

        self.lora_A = nn.Parameter(torch.randn(d_out, rank))
        self.lora_B = nn.Parameter(torch.randn(rank, d_in))

    def forward(self, x):
        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)