1import torch
2import torch.nn as nn
3
4
5class LoRALinear(nn.Module):
6 """A frozen linear layer with a trainable low-rank bypass."""
7
8 def __init__(self, base: nn.Linear, r: int = 8, alpha: int = 16):
9 super().__init__()
10 self.base = base
11 self.base.weight.requires_grad_(False) # W0 is FROZEN
12
13 d, k = base.weight.shape
14 self.A = nn.Parameter(torch.randn(r, k) * 0.01) # Gaussian init
15 self.B = nn.Parameter(torch.zeros(d, r)) # ZERO init:
16 self.scale = alpha / r # so BA = 0 at step one --
17 # training starts at W0 exactly.
18
19 def forward(self, x: torch.Tensor) -> torch.Tensor:
20 # frozen path + low-rank correction
21 return self.base(x) + self.scale * (x @ self.A.T @ self.B.T)
22
23 @torch.no_grad()
24 def merge(self):
25 """Fold the adapter in: zero-latency deployment."""
26 self.base.weight += self.scale * (self.B @ self.A)
27
28
29# --- The arithmetic that changed the economics ---
30d = k = 12_288 # one GPT-3 attention projection
31full = d * k # 150,994,944 trainable params
32lora = 2 * d * 8 # 196,608 at rank 8
33print(f"trainable: {full:,} -> {lora:,} ({full // lora}x fewer)")
34
35# Applied to W_q and W_v across all 96 layers, versus full fine-tuning
36# with Adam states for 175B params: ~10,000x fewer trainable parameters,
37# ~3x less GPU memory -- and the merged model runs at base-model speed.