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.randn(in_dim, rank) * std_dev)
self.B = nn.Parameter(torch.zeros(rank, out_dim))
self.alpha = alpha
def forward(self, x):
x = self.alpha * (x @ self.A @ self.B)
return x
class LinearWithLoRA(nn.Module):
def __init__(self, linear, rank, alpha):
super().__init__()
self.linear = linear
self.lora = LoRALayer(
linear.in_features, linear.out_features, rank, alpha
)
def forward(self, x):
return self.linear(x) + self.lora(x)
Here, x represents the pre-trained model layer. Now, when comparing it to the lit-gpt LoRA implementation
the model weights are being passed to a linear layer and then added with LoRA weights.
I'm confused about where this extra linear layer comes from in lit-gpt code?, because in your implementation, rather than passing the weights to any layer, we are adding LoRA directly to the model's linear layer. Could you please clarify the difference between these two approaches and explain the purpose of the additional linear layer in the lit-gpt implementation?