Spaces:
Running on Zero
Running on Zero
File size: 1,984 Bytes
ccf7a99 | 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 | """LoRA adapter for a frozen nn.Linear."""
from __future__ import annotations
import math
import torch
import torch.nn as nn
class LoRALinear(nn.Module):
def __init__(
self,
base_linear: nn.Linear,
rank: int = 8,
alpha: float | None = None,
dropout: float = 0.0,
):
super().__init__()
if rank <= 0:
raise ValueError(f"LoRA rank must be positive; got {rank}")
if alpha is None:
alpha = float(rank)
self.rank = rank
self.alpha = float(alpha)
self.scale = self.alpha / rank
self.base = base_linear
in_f = base_linear.in_features
out_f = base_linear.out_features
# Match base's device/dtype so post-construction LoRA wrapping
# (e.g. apply_lora_and_freeze after the model is on GPU) doesn't leave
# the adapter params on CPU.
device = base_linear.weight.device
dtype = base_linear.weight.dtype
self.lora_A = nn.Parameter(torch.empty(rank, in_f, device=device, dtype=dtype))
self.lora_B = nn.Parameter(torch.zeros(out_f, rank, device=device, dtype=dtype))
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
self.lora_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
# When False the adapter is skipped and forward == base(x). Lets callers
# evaluate the pristine frozen base (e.g. a self-distillation teacher)
# without unwrapping. Defaults True — backward-compatible.
self.enabled = True
@property
def in_features(self) -> int:
return self.base.in_features
@property
def out_features(self) -> int:
return self.base.out_features
def forward(self, x: torch.Tensor) -> torch.Tensor:
out = self.base(x)
if not self.enabled:
return out
update = self.lora_dropout(x) @ self.lora_A.t() @ self.lora_B.t()
return out + update * self.scale
|