File size: 4,130 Bytes
87b732d | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """Copyright (c) Microsoft Corporation. Licensed under the MIT license."""
import math
from typing import Literal
import torch
from torch import nn
__all__ = ["LoRA", "LoRARollout", "LoRAMode"]
LoRAMode = Literal["single", "from_second", "all"]
class LoRA(nn.Module):
"""LoRA adaptation for a linear layer."""
def __init__(
self,
in_features: int,
out_features: int,
r: int = 4,
alpha: int = 1,
dropout: float = 0.0,
) -> None:
"""Initialise.
Args:
in_features (int): Number of input features.
out_features (int): Number of output features.
r (int, optional): Rank. Defaults to `4`.
alpha (int, optional): Alpha. Defaults to `1`.
dropout (float, optional): Drop-out rate. Defaults to `0.0`.
"""
super().__init__()
assert r > 0, "The rank must be strictly positive."
self.lora_alpha = alpha
self.r = r
self.lora_dropout = nn.Dropout(dropout)
self.lora_A = nn.Parameter(torch.empty((r, in_features)))
self.lora_B = nn.Parameter(torch.empty((out_features, r)))
self.scaling = self.lora_alpha / self.r
self.init_weights()
def init_weights(self) -> None:
"""Initialise weights."""
# Initialise A the same way as the default for `nn.Linear` and set B to zero.
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
nn.init.zeros_(self.lora_B)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Compute the LoRA adaptation.
Args:
x (torch.Tensor): Input to the linear layer.
Returns:
torch.Tensor: Additive correction for the output of the linear layer.
"""
x = self.lora_dropout(x) @ self.lora_A.transpose(0, 1) @ self.lora_B.transpose(0, 1)
return x * self.scaling
class LoRARollout(nn.Module):
"""Per-roll-out-step LoRA finetuning."""
def __init__(
self,
in_features: int,
out_features: int,
r: int = 8,
alpha: int = 8,
dropout: float = 0.0,
max_steps: int = 40,
mode: LoRAMode = "single",
) -> None:
"""Initialise.
Args:
in_features (int): Number of input features.
out_features (int): Number of output features.
r (int, optional): Rank. Defaults to `4`.
alpha (int, optional): Alpha. Defaults to `1`.
dropout (float, optional): Drop-out rate. Defaults to `0.0`.
max_steps (int, optional): Maximum number of roll-out steps. Defaults to `40`.
mode (str, optional): Mode. `"single"` uses the same LoRA for all roll-out steps,
`"from_second"` uses the same LoRA from the second roll-out step on, and `"all"`
uses a different LoRA for every roll-out step. Defaults to `"single"`.
"""
super().__init__()
self.mode = mode
self.max_steps = max_steps
lora_layers = max_steps if mode == "all" else 1
self.loras = nn.ModuleList(
[
LoRA(in_features, out_features, r=r, alpha=alpha, dropout=dropout)
for _ in range(lora_layers)
]
)
def forward(self, x: torch.Tensor, step: int) -> torch.Tensor:
"""Compute the LoRA adaptation.
Args:
x (torch.Tensor): Input to the linear layer.
step (int): Roll-out step, starting at zero.
Returns:
torch.Tensor: Additive correction for the output of the linear layer.
"""
assert step >= 0, f"Step must be non-negative, found {step}."
if step >= self.max_steps:
return 0
if self.mode == "single":
return self.loras[0](x)
elif self.mode == "from_second":
if step == 0:
return 0
else:
return self.loras[0](x)
elif self.mode == "all":
return self.loras[step](x)
else:
raise ValueError(f"Invalid mode: {self.mode}")
|