import torch
import torch.nn as nn
import math
class LoRALinear(nn.Module):
def __init__(self, in_features: int, out_features: int, rank: int = 8, alpha: float = 16.0):
super().__init__()
# 1. Freeze base pretrained linear layer
self.linear = nn.Linear(in_features, out_features, bias=False)
self.linear.weight.requires_grad = False
self.rank = rank
self.scaling = alpha / rank
# 2. Low-rank matrices A and B
self.lora_A = nn.Parameter(torch.zeros(rank, in_features))
self.lora_B = nn.Parameter(torch.zeros(out_features, rank))
# 3. Initialization: A ~ Gaussian, B = 0
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
nn.init.zeros_(self.lora_B)
def forward(self, x: torch.Tensor) -> torch.Tensor:
base_out = self.linear(x)
lora_out = (x @ self.lora_A.T @ self.lora_B.T) * self.scaling
return base_out + lora_out
def merge_weights(self):
"""Merges B*A into base weight matrix W_0 for zero-latency inference."""
self.linear.weight.data += (self.lora_B @ self.lora_A) * self.scaling
print("LoRA weights merged successfully into base linear layer!")