File size: 2,607 Bytes
fd5a7ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """LoRA, written out: a frozen linear layer plus a trained low-rank correction."""
import math
import torch
import torch.nn as nn
TARGETS = ("q_proj", "k_proj", "v_proj", "o_proj")
class LoRALinear(nn.Module):
"""y = base(x) + (alpha / r) * B A x
base is the pretrained layer, frozen. A is (r, in), B is (out, r), and only they train.
B starts at zero, so at step 0 the correction is zero and the model is the pretrained one.
"""
def __init__(self, base: nn.Linear, r: int, alpha: float):
super().__init__()
self.base = base
for p in self.base.parameters():
p.requires_grad = False
w = base.weight
# A and B stay float32 whatever the base is stored in: AdamW's small updates vanish in bf16
self.A = nn.Parameter(torch.empty(r, base.in_features, dtype=torch.float32, device=w.device))
self.B = nn.Parameter(torch.zeros(base.out_features, r, dtype=torch.float32, device=w.device))
nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
self.scale = alpha / r
def forward(self, x):
delta = (x.to(self.A.dtype) @ self.A.T @ self.B.T) * self.scale
return self.base(x) + delta.to(x.dtype)
def inject(model, r=16, alpha=32, targets=TARGETS):
"""Freeze every weight, then wrap the attention projections of every layer.
Returns {"layers.3.q_proj": LoRALinear, ...}."""
for p in model.parameters():
p.requires_grad = False
wrapped = {}
for i, layer in enumerate(model.model.layers):
for name in targets:
lora = LoRALinear(getattr(layer.self_attn, name), r, alpha)
setattr(layer.self_attn, name, lora)
wrapped[f"layers.{i}.{name}"] = lora
return wrapped
def lora_state(model):
"""Only the trained tensors: this is the whole adapter."""
return {k: v.detach().cpu() for k, v in model.state_dict().items() if k.endswith((".A", ".B"))}
def load_adapter(model, path, r=16, alpha=32):
"""Wrap the model and load trained A and B into it."""
inject(model, r, alpha)
if str(path).endswith(".safetensors"):
from safetensors.torch import load_file
state = load_file(path)
else:
state = torch.load(path, map_location="cpu")
missing, unexpected = model.load_state_dict(state, strict=False)
assert not unexpected and not [k for k in missing if k.endswith((".A", ".B"))]
return model
def count(model):
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
return trainable, total
|