| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | from typing import TYPE_CHECKING |
| 6 | |
| 7 | from huggingface_hub.utils import HfHubHTTPError, HFValidationError |
| 8 | from torch import nn |
| 9 | from transformers import PreTrainedConfig |
| 10 | |
| 11 | from vllm import envs |
| 12 | from vllm.config.lora import LoRAConfig |
| 13 | from vllm.logger import init_logger |
| 14 | |
| 15 | # being imported for _all_lora_classes below |
| 16 | from vllm.lora.layers import ( |
| 17 | BaseLayerWithLoRA, |
| 18 | ClassificationHeadWithLoRA, |
| 19 | ColumnParallelLinearWithLoRA, |
| 20 | ColumnParallelLinearWithShardedLoRA, |
| 21 | FusedMoE3DWithLoRA, |
| 22 | FusedMoEWithLoRA, |
| 23 | LogitsProcessorWithLoRA, |
| 24 | MergedColumnParallelLinearVariableSliceWithLoRA, |
| 25 | MergedColumnParallelLinearWithLoRA, |
| 26 | MergedColumnParallelLinearWithShardedLoRA, |
| 27 | MergedQKVParallelLinearWithLoRA, |
| 28 | MergedQKVParallelLinearWithShardedLoRA, |
| 29 | QKVParallelLinearWithLoRA, |
| 30 | QKVParallelLinearWithShardedLoRA, |
| 31 | ReplicatedLinearWithLoRA, |
| 32 | RowParallelLinearWithLoRA, |
| 33 | RowParallelLinearWithShardedLoRA, |
| 34 | VocabParallelEmbeddingWithLoRA, |
| 35 | ) |
| 36 | from vllm.model_executor.custom_op import maybe_get_oot_by_class |
| 37 | from vllm.model_executor.layers.fused_moe import MoERunner |
| 38 | from vllm.model_executor.layers.linear import LinearBase |
| 39 | from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding |
| 40 | from vllm.model_executor.utils import get_moe_expert_mapping, get_packed_modules_mapping |
| 41 | from vllm.transformers_utils.repo_utils import hf_api |
| 42 | |
| 43 | if TYPE_CHECKING: |
| 44 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 45 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 46 | from vllm.model_executor.models.utils import WeightsMapper |
| 47 | |
| 48 | logger = init_logger(__name__) |
| 49 | |
| 50 | |
| 51 | def get_captured_lora_counts(max_loras: int, specialize: bool) -> list[int]: |
| 52 | """Returns num_active_loras values for cudagraph capture. |
| 53 | |
| 54 | When specialize=True: powers of 2 up to max_loras, plus max_loras + 1. |
| 55 | When specialize=False: just [max_loras + 1]. |
| 56 | |
| 57 | This is the single source of truth for LoRA capture cases, used by both |
| 58 | CudagraphDispatcher and PunicaWrapperGPU. |
| 59 | """ |
| 60 | if not specialize: |
| 61 | return [max_loras + 1] |
| 62 | |
| 63 | return [ |
| 64 | n for n in range(1, max_loras + 2) if (n & (n - 1)) == 0 or n == max_loras + 1 |
| 65 | ] |
| 66 | |
| 67 | |
| 68 | _GLOBAL_LORA_ID = 0 |
| 69 | |
| 70 | |
| 71 | def get_lora_id(): |
| 72 | global _GLOBAL_LORA_ID |
| 73 | _GLOBAL_LORA_ID += 1 |
| 74 | return _GLOBAL_LORA_ID |
| 75 | |
| 76 | |
| 77 | # Order matters here: more specific wrappers must be checked before generic |
| 78 | # merged/column-parallel wrappers in from_layer(). |
| 79 | _all_lora_classes: tuple[type[BaseLayerWithLoRA], ...] = ( |
| 80 | VocabParallelEmbeddingWithLoRA, |
| 81 | ColumnParallelLinearWithLoRA, |
| 82 | MergedColumnParallelLinearWithLoRA, |
| 83 | QKVParallelLinearWithLoRA, |
| 84 | MergedQKVParallelLinearWithLoRA, |
| 85 | RowParallelLinearWithLoRA, |
| 86 | ReplicatedLinearWithLoRA, |
| 87 | LogitsProcessorWithLoRA, |
| 88 | ColumnParallelLinearWithShardedLoRA, |
| 89 | QKVParallelLinearWithShardedLoRA, |
| 90 | MergedColumnParallelLinearWithShardedLoRA, |
| 91 | MergedColumnParallelLinearVariableSliceWithLoRA, |
| 92 | MergedQKVParallelLinearWithShardedLoRA, |
| 93 | RowParallelLinearWithShardedLoRA, |
| 94 | FusedMoEWithLoRA, |
| 95 | FusedMoE3DWithLoRA, |
| 96 | ) |
| 97 | |
| 98 | |
| 99 | def is_moe_model(model: nn.Module) -> bool: |
| 100 | """Checks if the model contains MoERunner layers and warns the user.""" |
| 101 | if any(isinstance(module, MoERunner) for module in model.modules()): |
| 102 | logger.info_once("MoE model detected. Using fused MoE LoRA implementation.") |
| 103 | return True |
| 104 | return False |
| 105 | |
| 106 | |
| 107 | def from_layer( |
| 108 | layer: nn.Module, |
| 109 | max_loras: int, |
| 110 | lora_config: LoRAConfig, |
| 111 | packed_modules_list: list, |
| 112 | model_config: PreTrainedConfig | None = None, |
| 113 | ) -> nn.Module: |
| 114 | for lora_cls in _all_lora_classes: |
| 115 | # specifying kwargs so they can be easily accessed in decorator |
| 116 | if lora_cls.can_replace_layer( |
| 117 | source_layer=layer, |
| 118 | lora_config=lora_config, |
| 119 | packed_modules_list=packed_modules_list, |
| 120 | model_config=model_config, |
| 121 | ): |
| 122 | instance_layer = lora_cls(layer) |
| 123 | instance_layer.create_lora_weights(max_loras, lora_config, model_config) |
| 124 | return instance_layer |
| 125 | return layer |
| 126 | |
| 127 | |
| 128 | def from_layer_logits_processor( |
| 129 | layer: "LogitsProcessor", |
| 130 | lm_head: "ParallelLMHead", |
| 131 | max_loras: int, |
| 132 | lora_config: LoRAConfig, |
| 133 | model_config: PreTrainedConfig | None = None, |
| 134 | ) -> LogitsProcessorWithLoRA: |
| 135 | ret = LogitsProcessorWithLoRA( |
| 136 | layer, |
| 137 | lm_head.embedding_dim, |
| 138 | lm_head.weight.dtype, |
| 139 | lm_head.weight.device, |
| 140 | lm_head.get_sharded_to_full_mapping(), |
| 141 | ) |
| 142 | ret.create_lora_weights(max_loras, lora_config, model_config) |
| 143 | return ret |
| 144 | |
| 145 | |
| 146 | def from_layer_classification( |
| 147 | layer: nn.Module, |
| 148 | max_loras: int, |
| 149 | lora_config: LoRAConfig, |
| 150 | model_config: PreTrainedConfig | None = None, |
| 151 | ) -> ClassificationHeadWithLoRA: |
| 152 | instance_layer = ClassificationHeadWithLoRA(layer) |
| 153 | instance_layer.create_lora_weights(max_loras, lora_config, model_config) |
| 154 | return instance_layer |
| 155 | |
| 156 | |
| 157 | def replace_submodule( |
| 158 | model: nn.Module, module_name: str, new_module: nn.Module |
| 159 | ) -> nn.Module: |
| 160 | """Replace a submodule in a model with a new module.""" |
| 161 | parent = model.get_submodule(".".join(module_name.split(".")[:-1])) |
| 162 | target_name = module_name.split(".")[-1] |
| 163 | setattr(parent, target_name, new_module) |
| 164 | return new_module |
| 165 | |
| 166 | |
| 167 | def parse_fine_tuned_lora_name( |
| 168 | name: str, weights_mapper: "WeightsMapper | None" = None |
| 169 | ) -> tuple[str, bool]: |
| 170 | """Parse the name of lora weights. |
| 171 | |
| 172 | Args: |
| 173 | name: the name of the fine-tuned LoRA, e.g. |
| 174 | base_model.model.dense1.weight |
| 175 | weights_mapper: maps the name of weight, e.g. |
| 176 | `model.` -> `language_model.model.`, |
| 177 | return: |
| 178 | tuple(module_name, is_lora_a): |
| 179 | module_name: the name of the module, e.g. model.dense1, |
| 180 | is_lora_a whether the tensor is lora_a or lora_b. |
| 181 | |
| 182 | """ |
| 183 | # LoRA weight qualified name usually starts with `base_model.model.`, |
| 184 | # so we remove the prefix `base_model.model.` to make the following |
| 185 | # mapping correctly. |
| 186 | if name.startswith("base_model.model."): |
| 187 | name = name.replace("base_model.model.", "") |
| 188 | if weights_mapper: |
| 189 | mapped_name = weights_mapper._map_name(name) |
| 190 | if mapped_name is None: |
| 191 | raise ValueError("Mapped LoRA weight name cannot be None.") |
| 192 | # recover the prefix `base_model.model.` |
| 193 | name = "base_model.model." + mapped_name |
| 194 | else: |
| 195 | if weights_mapper: |
| 196 | mapped_name = weights_mapper._map_name(name) |
| 197 | if mapped_name is None: |
| 198 | raise ValueError("Mapped LoRA weight name cannot be None.") |
| 199 | name = mapped_name |
| 200 | |
| 201 | # In some situations, we may not start with `base_model.model.`. |
| 202 | # If we don't (e.g., ibm-granite/granite-speech-3.3-8b), |
| 203 | # we should keep the prefix intact. |
| 204 | start_index = 2 if name.startswith("base_model.model.") else 0 |
| 205 | |
| 206 | parts = name.split(".") |
| 207 | if (parts[-1] == "weight" or parts[-1] == "bias") and len(parts) >= 2: |
| 208 | if parts[-2] in ["lora_A", "lora_B"]: |
| 209 | new_name = ".".join(parts[start_index:-2]) |
| 210 | return new_name, parts[-2] == "lora_A" |
| 211 | # For modules_to_save in classification. |
| 212 | elif parts[-2] in ["score", "classifier"]: |
| 213 | new_name = parts[-2] |
| 214 | return new_name, False |
| 215 | |
| 216 | if parts[-1] == "lora_embedding_A" or parts[-1] == "lora_embedding_B": |
| 217 | new_name = ".".join(parts[start_index:-1]) |
| 218 | return new_name, parts[-1] == "lora_embedding_A" |
| 219 | |
| 220 | raise ValueError(f"{name} is unsupported LoRA weight") |
| 221 | |
| 222 | |
| 223 | def is_base_embedding_weights(name: str) -> bool: |
| 224 | # hardcoded subfixes for input & output embedding weights |
| 225 | embedding_suffixes = ( |
| 226 | ".embed_tokens.base_layer.weight", |
| 227 | ".lm_head.base_layer.weight", |
| 228 | ) |
| 229 | return name.endswith(embedding_suffixes) |
| 230 | |
| 231 | |
| 232 | def get_supported_lora_modules(model: nn.Module) -> list[str]: |
| 233 | """In vLLM, all linear layers support LoRA.""" |
| 234 | supported_lora_modules: set[str] = set() |
| 235 | for name, module in model.named_modules(): |
| 236 | # get the embedding modules if the module's embedding_modules |
| 237 | # is not empty. |
| 238 | embedding_modules = getattr(module, "embedding_modules", None) |
| 239 | if embedding_modules is not None: |
| 240 | for name in embedding_modules: |
| 241 | supported_lora_modules.add(name) |
| 242 | |
| 243 | if isinstance(module, (LinearBase, MoERunner)): |
| 244 | supported_lora_modules.add(name.split(".")[-1]) |
| 245 | |
| 246 | return list(supported_lora_modules) |
| 247 | |
| 248 | |
| 249 | def is_supported_lora_module( |
| 250 | module_name: str, |
| 251 | module: nn.Module, |
| 252 | supported_lora_modules: list[str], |
| 253 | ) -> bool: |
| 254 | """Check if a module is in the model's supported LoRA modules. |
| 255 | |
| 256 | The module name must match a model-supported suffix, and the runtime |
| 257 | module must belong to a module family handled by LoRA. |
| 258 | |
| 259 | Args: |
| 260 | module_name: Full dot-separated module name. |
| 261 | module: Runtime module associated with ``module_name``. |
| 262 | supported_lora_modules: List of module suffixes supported by the |
| 263 | model. |
| 264 | |
| 265 | Returns: |
| 266 | True if the module is supported, False otherwise. |
| 267 | |
| 268 | """ |
| 269 | module_suffix = module_name.rsplit(".", 1)[-1] |
| 270 | if module_suffix not in supported_lora_modules: |
| 271 | return False |
| 272 | |
| 273 | return isinstance( |
| 274 | module, |
| 275 | ( |
| 276 | LinearBase, |
| 277 | MoERunner, |
| 278 | VocabParallelEmbedding, |
| 279 | maybe_get_oot_by_class(VocabParallelEmbedding), |
| 280 | BaseLayerWithLoRA, |
| 281 | ), |
| 282 | ) |
| 283 | |
| 284 | |
| 285 | def is_in_target_modules( |
| 286 | module_name: str, |
| 287 | target_modules: list[str] | None, |
| 288 | packed_modules_mapping: dict[str, list[str]] | None = None, |
| 289 | ) -> bool: |
| 290 | """Check if a module passes the deployment-time target_modules filter. |
| 291 | |
| 292 | When target_modules is None (no restriction), all modules pass. |
| 293 | Otherwise, the module's suffix must be in the target_modules list. |
| 294 | |
| 295 | Args: |
| 296 | module_name: Full dot-separated module name. |
| 297 | target_modules: Optional deployment-time restriction list from |
| 298 | LoRAConfig.target_modules. |
| 299 | packed_modules_mapping: Optional model-defined mapping from packed |
| 300 | runtime module names to their adapter-visible submodule names |
| 301 | (e.g. ``{"gate_up_proj": ["gate_proj", "up_proj"]}``). |
| 302 | |
| 303 | Returns: |
| 304 | True if the module passes the filter, False otherwise. |
| 305 | |
| 306 | """ |
| 307 | if target_modules is None: |
| 308 | return True |
| 309 | target_module_set = set(target_modules) |
| 310 | module_suffix = module_name.split(".")[-1] |
| 311 | if module_suffix in target_module_set or module_name in target_module_set: |
| 312 | return True |
| 313 | |
| 314 | if not packed_modules_mapping: |
| 315 | return False |
| 316 | |
| 317 | # Runtime packed parent matched by deployment-time child targets. |
| 318 | packed_children = packed_modules_mapping.get(module_suffix) |
| 319 | if packed_children and any(child in target_module_set for child in packed_children): |
| 320 | return True |
| 321 | |
| 322 | # Adapter-visible packed child matched by deployment-time parent target. |
| 323 | return any( |
| 324 | module_suffix in children and packed_parent in target_module_set |
| 325 | for packed_parent, children in packed_modules_mapping.items() |
| 326 | ) |
| 327 | |
| 328 | |
| 329 | def get_adapter_absolute_path(lora_path: str) -> str: |
| 330 | """Resolves the given lora_path to an absolute local path. |
| 331 | |
| 332 | If the lora_path is identified as a Hugging Face model identifier, |
| 333 | it will download the model and return the local snapshot path. |
| 334 | Otherwise, it treats the lora_path as a local file path and |
| 335 | converts it to an absolute path. |
| 336 | |
| 337 | Args: |
| 338 | lora_path (str): The path to the lora model, which can be an absolute |
| 339 | path, a relative path, or a Hugging Face model identifier. |
| 340 | |
| 341 | Returns: |
| 342 | str: The resolved absolute local path to the lora model. |
| 343 | |
| 344 | """ |
| 345 | # Check if the path is an absolute path. Return it no matter exists or not. |
| 346 | if os.path.isabs(lora_path): |
| 347 | return lora_path |
| 348 | |
| 349 | # If the path starts with ~, expand the user home directory. |
| 350 | if lora_path.startswith("~"): |
| 351 | return os.path.expanduser(lora_path) |
| 352 | |
| 353 | # Check if the expanded relative path exists locally. |
| 354 | if os.path.exists(lora_path): |
| 355 | return os.path.abspath(lora_path) |
| 356 | |
| 357 | # If the path does not exist locally. |
| 358 | if envs.VLLM_USE_MODELSCOPE: |
| 359 | # If using ModelScope, we assume the path is a ModelScope repo. |
| 360 | from modelscope.hub.snapshot_download import InvalidParameter, snapshot_download |
| 361 | from requests import HTTPError |
| 362 | |
| 363 | download_fn = lambda: snapshot_download(model_id=lora_path) |
| 364 | download_exceptions = (HTTPError, InvalidParameter) |
| 365 | error_log = "Error downloading the ModelScope model" |
| 366 | else: |
| 367 | # Otherwise, we assume the path is a Hugging Face Hub repo. |
| 368 | download_fn = lambda: hf_api().snapshot_download( |
| 369 | repo_id=lora_path, |
| 370 | ) |
| 371 | download_exceptions = (HfHubHTTPError, HFValidationError) |
| 372 | error_log = "Error downloading the HuggingFace model" |
| 373 | |
| 374 | try: |
| 375 | local_snapshot_path = download_fn() |
| 376 | except download_exceptions: |
| 377 | # Handle errors that may occur during the download. |
| 378 | # Return original path instead of throwing error here. |
| 379 | logger.exception(error_log) |
| 380 | return lora_path |
| 381 | |
| 382 | return local_snapshot_path |
| 383 | |
| 384 | |
| 385 | def process_packed_modules_mapping( |
| 386 | model: nn.Module, |
| 387 | force_2d_moe: bool = False, |
| 388 | enable_moe_shared_loras: bool = False, |
| 389 | ) -> dict[str, list[str]]: |
| 390 | if is_moe_model(model): |
| 391 | # This method generates and returns a dictionary mapping packed module |
| 392 | # names to lists of their corresponding submodule names. It includes |
| 393 | # both static mappings and dynamic mappings for expert layers, where |
| 394 | # the expert indices are expanded based on the configured number |
| 395 | # of routed experts. |
| 396 | packed_modules_mapping = get_packed_modules_mapping(model) |
| 397 | # The 2D mapping is needed when the model itself is 2D, or when |
| 398 | # the engine forces the universal 2D wrapper via |
| 399 | # enable_mixed_moe_lora_format (so 3D models can also load 2D |
| 400 | # adapters through FusedMoEWithLoRA). |
| 401 | if enable_moe_shared_loras: |
| 402 | # Shared MoE adapters store one pre-stacked tensor per |
| 403 | # expert-projection (experts.w1/w2/w3) rather than a per-expert |
| 404 | # tensor list, so the packed mapping references the three stack |
| 405 | # names directly (drives expected_lora_modules and the packing). |
| 406 | packed_modules_mapping["experts"] = [ |
| 407 | "experts.w1", |
| 408 | "experts.w2", |
| 409 | "experts.w3", |
| 410 | ] |
| 411 | elif (not model.is_3d_moe_weight) or force_2d_moe: |
| 412 | # Filter out malformed entries: non-gated MoE has empty |
| 413 | # ckpt_up_proj_name which results in weight_name containing ".." |
| 414 | # (e.g., "experts.0.." instead of "experts.0.layer_name.") |
| 415 | packed_modules_mapping["experts"] = [ |
| 416 | weight_name.rstrip(".") |
| 417 | for _, weight_name, _, _ in get_moe_expert_mapping(model) |
| 418 | if ".." not in weight_name |
| 419 | ] |
| 420 | |
| 421 | return packed_modules_mapping |
| 422 | else: |
| 423 | return get_packed_modules_mapping(model) |