Spaces:
Running on Zero
Running on Zero
| """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 | |
| def in_features(self) -> int: | |
| return self.base.in_features | |
| 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 | |