File size: 2,733 Bytes
9ede8c0 | 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 | """
Generic LoRA (Low-Rank Adaptation, Hu et al. 2021) wrapper -- module-agnostic,
works on any nn.Module built from nn.Linear layers. Used here to fine-tune
the value head under self-play with a capacity-constrained update instead of
a full-parameter one, which empirically fine-tunes more reliably on a small,
self-generated data distribution without regressing.
"""
import torch
from torch import nn
class LoRALinear(nn.Module):
"""Wraps an existing nn.Linear, freezing its original weight/bias and
adding a trainable low-rank delta:
output = frozen_linear(x) + scaling * (x @ A^T) @ B^T
`B` is initialized to ZERO, so the wrapped layer's output is byte-
identical to the original before any training happens."""
def __init__(self, linear, rank=4, alpha=1.0):
super().__init__()
assert isinstance(linear, nn.Linear), f"LoRALinear only wraps nn.Linear, got {type(linear)}"
self.linear = linear
for p in self.linear.parameters():
p.requires_grad = False
self.rank = rank
self.scaling = alpha / rank
device = linear.weight.device
dtype = linear.weight.dtype
self.lora_A = nn.Parameter(torch.randn(rank, linear.in_features, device=device, dtype=dtype) * 0.01)
self.lora_B = nn.Parameter(torch.zeros(linear.out_features, rank, device=device, dtype=dtype))
def forward(self, x):
base = self.linear(x)
delta = (x @ self.lora_A.t()) @ self.lora_B.t()
return base + self.scaling * delta
def lora_parameters(self):
return [self.lora_A, self.lora_B]
def merge_into_base(self):
"""Fold the current LoRA delta into the frozen base weight and
reset B to zero -- used to "commit" a trained adapter back into a
plain nn.Linear-equivalent state so the saved checkpoint is an
ordinary state_dict, loadable with zero LoRA-awareness downstream."""
with torch.no_grad():
delta_w = self.scaling * (self.lora_B @ self.lora_A)
self.linear.weight.add_(delta_w)
self.lora_B.zero_()
def apply_lora(module, rank=4, alpha=1.0):
"""Recursively replace every nn.Linear submodule of `module` with a
LoRALinear wrapper. Returns the new LoRA parameters for the caller's
optimizer. Modifies `module` in place."""
lora_params = []
for name, child in list(module.named_children()):
if isinstance(child, nn.Linear):
wrapped = LoRALinear(child, rank=rank, alpha=alpha)
setattr(module, name, wrapped)
lora_params.extend(wrapped.lora_parameters())
else:
lora_params.extend(apply_lora(child, rank=rank, alpha=alpha))
return lora_params
|