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