| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Sequence as GenericSequence |
| 5 | from dataclasses import dataclass |
| 6 | |
| 7 | import torch |
| 8 | import torch.types |
| 9 | |
| 10 | from vllm.lora.peft_helper import PEFTHelper |
| 11 | from vllm.utils.torch_utils import PIN_MEMORY |
| 12 | |
| 13 | |
| 14 | class LoRALayerWeights: |
| 15 | """LoRA weights for a layer composed of two low rank matrixes.""" |
| 16 | |
| 17 | def __init__( |
| 18 | self, |
| 19 | module_name: str, |
| 20 | rank: int, |
| 21 | lora_alpha: int, |
| 22 | lora_a: torch.Tensor, |
| 23 | lora_b: torch.Tensor, |
| 24 | scaling: float | None = None, |
| 25 | ) -> None: |
| 26 | self.module_name = module_name |
| 27 | self.rank = rank |
| 28 | self.lora_alpha = lora_alpha |
| 29 | self.lora_a = lora_a |
| 30 | self.lora_b = lora_b |
| 31 | |
| 32 | if scaling is None: |
| 33 | self.scaling = self.lora_alpha / self.rank |
| 34 | else: |
| 35 | self.scaling = scaling |
| 36 | |
| 37 | def optimize(self) -> "LoRALayerWeights": |
| 38 | """Optimize the LoRA by merging the scaling into lora_b.""" |
| 39 | if self.scaling == 1: |
| 40 | return self |
| 41 | self.lora_b *= self.scaling |
| 42 | self.scaling = 1 |
| 43 | return self |
| 44 | |
| 45 | @property |
| 46 | def input_dim(self) -> int: |
| 47 | return self.lora_a.shape[1] |
| 48 | |
| 49 | @property |
| 50 | def output_dim(self) -> int: |
| 51 | return self.lora_b.shape[0] |
| 52 | |
| 53 | @property |
| 54 | def is_packed(self) -> bool: |
| 55 | return False |
| 56 | |
| 57 | @classmethod |
| 58 | def from_config( |
| 59 | cls, |
| 60 | module_name: str, |
| 61 | peft_helper: PEFTHelper, |
| 62 | ) -> "LoRALayerWeights": |
| 63 | # lora_a and lora_b are set to None for config-based construction |
| 64 | return cls( |
| 65 | module_name, |
| 66 | peft_helper.r, |
| 67 | peft_helper.lora_alpha, |
| 68 | None, |
| 69 | None, |
| 70 | peft_helper.vllm_lora_scaling_factor, |
| 71 | ) |
| 72 | |
| 73 | @classmethod |
| 74 | def create_dummy_lora_weights( |
| 75 | cls, |
| 76 | module_name: str, |
| 77 | input_dim: int, |
| 78 | output_dim: int, |
| 79 | rank: int, |
| 80 | dtype: torch.dtype, |
| 81 | device: torch.types.Device, |
| 82 | ) -> "LoRALayerWeights": |
| 83 | pin_memory = str(device) == "cpu" and PIN_MEMORY |
| 84 | lora_a = torch.zeros( |
| 85 | [rank, input_dim], dtype=dtype, device=device, pin_memory=pin_memory |
| 86 | ) |
| 87 | lora_b = torch.zeros( |
| 88 | [output_dim, rank], dtype=dtype, device=device, pin_memory=pin_memory |
| 89 | ) |
| 90 | |
| 91 | return cls( |
| 92 | module_name, |
| 93 | rank=rank, |
| 94 | lora_alpha=1, |
| 95 | lora_a=lora_a, |
| 96 | lora_b=lora_b, |
| 97 | ) |
| 98 | |
| 99 | |
| 100 | class PackedLoRALayerWeights(LoRALayerWeights): |
| 101 | """LoRA used for packed layers (eg. qkv_proj).""" |
| 102 | |
| 103 | def __init__( |
| 104 | self, |
| 105 | module_name: str, |
| 106 | rank: int, |
| 107 | lora_alphas: list[int | None], |
| 108 | lora_a: list[torch.Tensor | None], |
| 109 | lora_b: list[torch.Tensor | None], |
| 110 | scaling: list[float] | None = None, |
| 111 | ) -> None: |
| 112 | super().__init__( |
| 113 | module_name=module_name, |
| 114 | rank=rank, |
| 115 | lora_alpha=0, |
| 116 | lora_a=lora_a, |
| 117 | lora_b=lora_b, |
| 118 | scaling=scaling, # type: ignore |
| 119 | ) |
| 120 | self.lora_alphas = lora_alphas |
| 121 | if scaling is None: |
| 122 | self.scaling = [ # type: ignore |
| 123 | lora_alpha / self.rank # type: ignore # noqa |
| 124 | for lora_alpha in self.lora_alphas |
| 125 | ] |
| 126 | |
| 127 | @classmethod |
| 128 | def pack( |
| 129 | cls, loras: GenericSequence["LoRALayerWeights | None"] |
| 130 | ) -> "PackedLoRALayerWeights": |
| 131 | """Pack a list of LoRAs into a single LoRA. |
| 132 | |
| 133 | If LoRA is None, it signifies that the submodule does not have a LoRA. |
| 134 | """ |
| 135 | first_lora = next(lora for lora in loras if lora is not None) |
| 136 | for lora in loras: |
| 137 | if lora is None: |
| 138 | continue |
| 139 | lora.optimize() |
| 140 | rank = first_lora.rank |
| 141 | module_name = first_lora.module_name |
| 142 | obj = cls( |
| 143 | module_name, |
| 144 | rank, |
| 145 | [lora.lora_alpha if lora is not None else None for lora in loras], |
| 146 | [lora.lora_a if lora is not None else None for lora in loras], |
| 147 | [lora.lora_b if lora is not None else None for lora in loras], |
| 148 | scaling=[ |
| 149 | 1 if lora is not None else None # type: ignore |
| 150 | for lora in loras |
| 151 | ], |
| 152 | ) |
| 153 | return obj |
| 154 | |
| 155 | @classmethod |
| 156 | def pack_moe( |
| 157 | cls, |
| 158 | loras: GenericSequence["LoRALayerWeights | None"], |
| 159 | module_name: str, |
| 160 | is_non_gated_moe: bool = False, |
| 161 | ) -> "PackedLoRALayerWeights": |
| 162 | """Pack a list of LoRAs into a single LoRA. |
| 163 | |
| 164 | If LoRA is None, it signifies that the submodule does not have a LoRA. |
| 165 | """ |
| 166 | first_lora = next(lora for lora in loras if lora is not None) |
| 167 | assert first_lora is not None |
| 168 | rank = first_lora.rank |
| 169 | lora_alpha = first_lora.lora_alpha |
| 170 | assert len(loras) % 3 == 0 |
| 171 | w1_lora_a_lst = [] |
| 172 | w2_lora_a_lst = [] |
| 173 | w3_lora_a_lst = [] |
| 174 | w1_lora_b_lst = [] |
| 175 | w2_lora_b_lst = [] |
| 176 | w3_lora_b_lst = [] |
| 177 | # TODO: Consider the case where some experts don't have LoRA added. |
| 178 | for eid in range(len(loras) // 3): |
| 179 | w1_lora = loras[eid * 3] |
| 180 | w2_lora = loras[eid * 3 + 1] |
| 181 | w3_lora = loras[eid * 3 + 2] |
| 182 | # For non-gated MoE, w3 is not used, so we use w1's LoRA weights |
| 183 | # This is determined by checking the expert mapping (get_expert_mapping) |
| 184 | # which indicates when ckpt_up_proj_name is empty. |
| 185 | if w3_lora is None and is_non_gated_moe: |
| 186 | w3_lora = w1_lora |
| 187 | assert w1_lora is not None |
| 188 | assert w2_lora is not None |
| 189 | assert w3_lora is not None |
| 190 | |
| 191 | w1_lora_a_lst.append(w1_lora.lora_a) |
| 192 | w2_lora_a_lst.append(w2_lora.lora_a) |
| 193 | w3_lora_a_lst.append(w3_lora.lora_a) |
| 194 | |
| 195 | w1_lora_b_lst.append(w1_lora.lora_b) |
| 196 | w2_lora_b_lst.append(w2_lora.lora_b) |
| 197 | w3_lora_b_lst.append(w3_lora.lora_b) |
| 198 | |
| 199 | w1_lora_a = torch.stack(w1_lora_a_lst, dim=0) # (num_experts,rank,input_size) |
| 200 | w2_lora_a = torch.stack(w2_lora_a_lst, dim=0) |
| 201 | w1_lora_b = torch.stack(w1_lora_b_lst, dim=0) # (num_experts,output_size,rank) |
| 202 | w2_lora_b = torch.stack(w2_lora_b_lst, dim=0) |
| 203 | |
| 204 | # All w1, w2, w3 have the same scaling factor. Use the per-adapter |
| 205 | # scaling (e.g. alpha/sqrt(r) for rsLoRA) instead of alpha/rank. |
| 206 | scaling = first_lora.scaling |
| 207 | last_scaling = scaling |
| 208 | |
| 209 | if is_non_gated_moe: |
| 210 | # For non-gated MoE, reuse w1 tensors for w3 to avoid memory waste |
| 211 | # w3_lora_a_lst and w3_lora_b_lst are not relevant in this case |
| 212 | w3_lora_a = w1_lora_a |
| 213 | w3_lora_b = w1_lora_b |
| 214 | |
| 215 | # For non-gated MoE, avoid double-scaling by setting w3's scaling to 1. |
| 216 | last_scaling = 1.0 |
| 217 | else: |
| 218 | w3_lora_a = torch.stack(w3_lora_a_lst, dim=0) |
| 219 | w3_lora_b = torch.stack(w3_lora_b_lst, dim=0) |
| 220 | |
| 221 | obj = cls( |
| 222 | module_name, |
| 223 | rank, |
| 224 | [lora_alpha, lora_alpha, lora_alpha], |
| 225 | [w1_lora_a, w2_lora_a, w3_lora_a], |
| 226 | [w1_lora_b, w2_lora_b, w3_lora_b], |
| 227 | scaling=[scaling, scaling, last_scaling], |
| 228 | ) |
| 229 | return obj |
| 230 | |
| 231 | @classmethod |
| 232 | def pack_moe_stacked( |
| 233 | cls, |
| 234 | loras: GenericSequence["LoRALayerWeights | None"], |
| 235 | module_name: str, |
| 236 | ) -> "PackedLoRALayerWeights": |
| 237 | """Pack pre-stacked (3D) w1/w2/w3 expert LoRAs into a single LoRA. |
| 238 | |
| 239 | Unlike :meth:`pack_moe`, which stacks one 2D tensor per expert, this |
| 240 | expects each of the three input LoRAs to already carry the expert |
| 241 | dimension (``experts.w{1,2,3}``). It is used for "shared-outer" |
| 242 | adapters where some factors are shared across experts (expert-dim 1) |
| 243 | and cannot be reconstructed by stacking per-expert tensors. The |
| 244 | produced ``lora_a``/``lora_b`` are the ``[w1, w2, w3]`` lists that |
| 245 | ``FusedMoEWithLoRA.set_lora`` consumes directly. |
| 246 | """ |
| 247 | assert len(loras) == 3, ( |
| 248 | "shared-outer expert LoRA expects exactly w1/w2/w3 stacked tensors" |
| 249 | ) |
| 250 | w1_lora, w2_lora, w3_lora = loras |
| 251 | assert w1_lora is not None and w2_lora is not None and w3_lora is not None |
| 252 | rank = w1_lora.rank |
| 253 | lora_alpha = w1_lora.lora_alpha |
| 254 | scaling = w1_lora.scaling |
| 255 | return cls( |
| 256 | module_name, |
| 257 | rank, |
| 258 | [lora_alpha, lora_alpha, lora_alpha], |
| 259 | [w1_lora.lora_a, w2_lora.lora_a, w3_lora.lora_a], |
| 260 | [w1_lora.lora_b, w2_lora.lora_b, w3_lora.lora_b], |
| 261 | scaling=[scaling, scaling, scaling], |
| 262 | ) |
| 263 | |
| 264 | def optimize(self) -> "PackedLoRALayerWeights": |
| 265 | """Optimize the LoRA by merging the scaling into lora_b.""" |
| 266 | for i in range(len(self.lora_b)): |
| 267 | if self.scaling[i] == 1 or self.lora_b[i] is None: # type: ignore |
| 268 | continue |
| 269 | self.lora_b[i] *= self.scaling[i] # type: ignore |
| 270 | self.scaling[i] = 1 # type: ignore |
| 271 | return self |
| 272 | |
| 273 | @property |
| 274 | def input_dim(self) -> int: |
| 275 | raise NotImplementedError() |
| 276 | |
| 277 | @property |
| 278 | def output_dim(self) -> int: |
| 279 | raise NotImplementedError() |
| 280 | |
| 281 | @property |
| 282 | def is_packed(self) -> bool: |
| 283 | return True |
| 284 | |
| 285 | |
| 286 | @dataclass |
| 287 | class LoRAFullModuleWeights: |
| 288 | """LoRA weights for classification layers.""" |
| 289 | |
| 290 | module_name: str |
| 291 | weight: torch.Tensor |
| 292 | bias: torch.Tensor | None = None |
| 293 | |
| 294 | def pin_memory(self) -> None: |
| 295 | self.weight = self.weight.pin_memory() |
| 296 | if self.bias is not None: |
| 297 | self.bias = self.bias.pin_memory() |