File size: 7,743 Bytes
e7f92e1 |
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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 |
from __future__ import annotations
import math
import warnings
from typing import Any, Optional, Union, List
import torch
import torch.nn as nn
from peft.tuners.lora import LoraLayer
class MultiAdapterLinear(nn.Module, LoraLayer):
"""
Custom LoRA module supporting multiple adapters for a linear layer.
This module extends the standard LoRA implementation to support multiple task-specific
adapters that can be dynamically selected during the forward pass. The task_label
parameter passed to the forward function determines which LoRA adapter(s) to use:
- If task_label is a string, all examples in the batch use the same adapter
- If task_label is a list of strings, each example can use a different adapter
This enables efficient multi-task inference where all task-specific LoRA adapters
are loaded in memory simultaneously and dynamically selected per example, eliminating
the need to switch adapter states between tasks and allowing optimal throughput
for mixed-task batches.
Derived from peft.tuners.lora.Linear.
"""
def __init__(
self,
base_layer,
adapter_name: str,
task_names: List[str],
r: int = 0,
lora_alpha: int = 1,
lora_dropout: float = 0.0,
fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
is_target_conv_1d_layer: bool = False,
init_lora_weights: Union[bool, str] = True,
use_rslora: bool = False,
use_dora: bool = False,
lora_bias: bool = False,
**kwargs,
) -> None:
super().__init__()
LoraLayer.__init__(self, base_layer, **kwargs)
self.fan_in_fan_out = fan_in_fan_out
self.task_names = task_names
self._active_adapter = adapter_name
self.update_layer(
adapter_name,
r,
lora_alpha=lora_alpha,
lora_dropout=lora_dropout,
init_lora_weights=init_lora_weights,
use_rslora=use_rslora,
use_dora=use_dora,
lora_bias=lora_bias,
)
self.is_target_conv_1d_layer = is_target_conv_1d_layer
def forward(self, x: torch.Tensor, task_label: Union[str, List[str]], *args: Any, **kwargs: Any) -> torch.Tensor:
self._check_forward_args(x, *args, **kwargs)
if self.disable_adapters:
if self.merged:
self.unmerge()
result = self.base_layer(x, *args, **kwargs)
elif self.merged:
result = self.base_layer(x, *args, **kwargs)
else:
result = self.base_layer(x, *args, **kwargs)
torch_result_dtype = result.dtype
lora_A_keys = self.lora_A.keys()
for active_adapter in self.active_adapters:
if active_adapter not in lora_A_keys:
continue
if isinstance(task_label, str):
lora_A = self.lora_A[active_adapter][task_label]
lora_B = self.lora_B[active_adapter][task_label]
dropout = self.lora_dropout[active_adapter]
scaling = self.scaling[active_adapter]
x = self._cast_input_dtype(x, lora_A.weight.dtype)
result = result + lora_B(lora_A(dropout(x))) * scaling
else:
unique_tasks = list(set(task_label))
lora_output = torch.zeros_like(result)
for task in unique_tasks:
task_indices = [i for i, t in enumerate(task_label) if t == task]
task_x = x[task_indices]
lora_A = self.lora_A[active_adapter][task]
lora_B = self.lora_B[active_adapter][task]
dropout = self.lora_dropout[active_adapter]
scaling = self.scaling[active_adapter]
task_x = self._cast_input_dtype(task_x, lora_A.weight.dtype)
task_lora_value = lora_B(lora_A(dropout(task_x))) * scaling
for i, idx in enumerate(task_indices):
lora_output[idx] = task_lora_value[i]
result = result + lora_output
result = result.to(torch_result_dtype)
return result
def __repr__(self) -> str:
rep = super().__repr__()
return "lora." + rep
def update_layer(
self,
adapter_name,
r,
lora_alpha,
lora_dropout,
init_lora_weights,
use_rslora,
use_dora: bool = False,
lora_bias: bool = False,
):
# This code works for linear layers, override for other layer types
if r <= 0:
raise ValueError(f"`r` should be a positive integer value but the value passed is {r}")
self.r[adapter_name] = r
self.lora_alpha[adapter_name] = lora_alpha
if lora_dropout > 0.0:
lora_dropout_layer = nn.Dropout(p=lora_dropout)
else:
lora_dropout_layer = nn.Identity()
self.lora_dropout.update(nn.ModuleDict({adapter_name: lora_dropout_layer}))
# Actual trainable parameters
self.lora_A[adapter_name] = nn.ModuleDict({
task_name: nn.Linear(self.in_features, r, bias=False)
for task_name in self.task_names
})
self.lora_B[adapter_name] = nn.ModuleDict({
task_name: nn.Linear(r, self.out_features, bias=lora_bias)
for task_name in self.task_names
})
self.lora_bias[adapter_name] = lora_bias
if use_rslora:
self.scaling[adapter_name] = lora_alpha / math.sqrt(r)
else:
self.scaling[adapter_name] = lora_alpha / r
self.reset_lora_parameters(adapter_name, init_lora_weights)
self._move_adapter_to_device_of_base_layer(adapter_name)
self.use_dora[adapter_name] = False
self.set_adapter(self.active_adapters)
def reset_lora_parameters(self, adapter_name, init_lora_weights):
if init_lora_weights is False:
return
if init_lora_weights is True:
# initialize A the same way as the default for nn.Linear and B to zero
# https://github.com/microsoft/LoRA/blob/a0a92e0f26c067cf94747bdbf1ce73793fa44d19/loralib/layers.py#L124
for task_name in self.task_names:
nn.init.kaiming_uniform_(self.lora_A[adapter_name][task_name].weight, a=math.sqrt(5))
elif init_lora_weights.lower() == "gaussian":
for task_name in self.task_names:
nn.init.normal_(self.lora_A[adapter_name][task_name].weight, std=1 / self.r[adapter_name])
else:
raise ValueError(f"Unknown initialization {init_lora_weights=}")
for task_name in self.task_names:
nn.init.zeros_(self.lora_B[adapter_name][task_name].weight)
if self.lora_bias[adapter_name]:
for task_name in self.task_names:
nn.init.zeros_(self.lora_B[adapter_name][task_name].bias)
def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None:
"""
Merge the active adapter weights into the base weights
"""
raise NotImplementedError("Merge operation is not supported")
def unmerge(self) -> None:
"""
This method unmerges all merged adapter layers from the base weights.
"""
raise NotImplementedError("Unmerge operation is not supported")
|