| """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 |
| |
| 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 |
|
|
