| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | from collections.abc import Callable |
| 6 | from typing import TypeVar |
| 7 | |
| 8 | import torch |
| 9 | from torch import nn |
| 10 | |
| 11 | from vllm.config import VllmConfig |
| 12 | from vllm.config.lora import LoRAConfig |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.lora.layers import ( |
| 15 | BaseLayerWithLoRA, |
| 16 | ClassificationHeadWithLoRA, |
| 17 | FusedMoE3DWithLoRA, |
| 18 | FusedMoEWithLoRA, |
| 19 | LoRAMapping, |
| 20 | LoRAMappingType, |
| 21 | ) |
| 22 | from vllm.lora.lora_model import LoRAModel, MoEEPLoadSpec |
| 23 | from vllm.lora.lora_weights import ( |
| 24 | LoRAFullModuleWeights, |
| 25 | LoRALayerWeights, |
| 26 | PackedLoRALayerWeights, |
| 27 | ) |
| 28 | from vllm.lora.punica_wrapper import PunicaWrapperBase, get_punica_wrapper |
| 29 | from vllm.lora.utils import ( |
| 30 | from_layer, |
| 31 | from_layer_classification, |
| 32 | from_layer_logits_processor, |
| 33 | get_supported_lora_modules, |
| 34 | is_in_target_modules, |
| 35 | is_moe_model, |
| 36 | is_supported_lora_module, |
| 37 | process_packed_modules_mapping, |
| 38 | replace_submodule, |
| 39 | ) |
| 40 | from vllm.model_executor.layers.fused_moe import MoERunner |
| 41 | from vllm.model_executor.models import ( |
| 42 | SupportsLoRA, |
| 43 | SupportsMultiModal, |
| 44 | is_pooling_model, |
| 45 | supports_multimodal, |
| 46 | ) |
| 47 | from vllm.model_executor.models.module_mapping import MultiModelKeys |
| 48 | from vllm.model_executor.models.utils import PPMissingLayer |
| 49 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 50 | from vllm.multimodal.encoder_budget import MultiModalBudget |
| 51 | from vllm.utils.cache import LRUCache |
| 52 | from vllm.utils.torch_utils import PIN_MEMORY |
| 53 | |
| 54 | logger = init_logger(__name__) |
| 55 | |
| 56 | T = TypeVar("T") |
| 57 | DEFAULT_LANGUAGE_WRAPPER_KEY = "language_model" |
| 58 | |
| 59 | |
| 60 | class SupportsLoRAModel(nn.Module, SupportsLoRA): ... |
| 61 | |
| 62 | |
| 63 | class SupportsLoRAMultiModalModel(SupportsLoRAModel, SupportsMultiModal): ... |
| 64 | |
| 65 | |
| 66 | class AdapterLRUCache(LRUCache[int, T]): |
| 67 | def __init__(self, capacity: int, deactivate_fn: Callable[[int], object]): |
| 68 | super().__init__(capacity) |
| 69 | self.deactivate_fn = deactivate_fn |
| 70 | |
| 71 | def _on_remove(self, key: int, value: T | None): |
| 72 | logger.debug("Removing adapter int id: %d", key) |
| 73 | self.deactivate_fn(key) |
| 74 | return super()._on_remove(key, value) |
| 75 | |
| 76 | |
| 77 | class LoRAModelManager: |
| 78 | """A manager that manages multiple LoRA-fine-tuned models.""" |
| 79 | |
| 80 | def __init__( |
| 81 | self, |
| 82 | model: SupportsLoRAModel, |
| 83 | max_num_seqs: int, |
| 84 | max_num_batched_tokens: int, |
| 85 | vocab_size: int, |
| 86 | lora_config: LoRAConfig, |
| 87 | device: torch.device, |
| 88 | vllm_config: VllmConfig, |
| 89 | ): |
| 90 | """Create a LoRAModelManager and adapter for a given model. |
| 91 | |
| 92 | Args: |
| 93 | model: the model to be adapted. |
| 94 | max_num_seqs: the maximum number of sequences model can run in a |
| 95 | single batch. |
| 96 | max_num_batched_tokens: the maximum number of tokens model can run |
| 97 | in a single batch. |
| 98 | vocab_size: the vocab size of the model. |
| 99 | lora_config: the LoRA configuration. |
| 100 | device: the device the LoRA tensors are placed on. |
| 101 | vllm_config: the vLLM config for this engine. |
| 102 | |
| 103 | """ |
| 104 | self.model: SupportsLoRAModel = model |
| 105 | self.supported_lora_modules = get_supported_lora_modules(self.model) |
| 106 | assert self.supported_lora_modules, ( |
| 107 | f"No supported LoRA modules found in {self.model.__class__.__name__}." |
| 108 | ) |
| 109 | |
| 110 | self.adapter_type = "LoRA" |
| 111 | self.lora_config = lora_config |
| 112 | self.device = device |
| 113 | self.max_num_seqs = max_num_seqs |
| 114 | assert self.capacity >= self.lora_slots |
| 115 | self._registered_adapters: AdapterLRUCache[LoRAModel] = AdapterLRUCache( |
| 116 | self.capacity, self.deactivate_adapter |
| 117 | ) |
| 118 | self._active_adapters: AdapterLRUCache[None] = AdapterLRUCache( |
| 119 | self.lora_slots, self._deactivate_adapter |
| 120 | ) |
| 121 | self.max_num_batched_tokens = math.ceil(max_num_batched_tokens / 8) * 8 |
| 122 | self.lora_index_to_id: list[int | None] = [None] * self.lora_slots |
| 123 | self.vocab_size = vocab_size |
| 124 | |
| 125 | self.is_pooling_model = is_pooling_model(self.model) |
| 126 | self._pooling_task = ( |
| 127 | vllm_config.model_config.get_pooling_task( |
| 128 | tuple(self.model.pooler.get_supported_tasks()) |
| 129 | ) |
| 130 | if self.is_pooling_model |
| 131 | else None |
| 132 | ) |
| 133 | self.packed_modules: dict[str, list[str]] = {} |
| 134 | self.modules: dict[str, BaseLayerWithLoRA] = {} |
| 135 | self._last_mapping: LoRAMapping | None = None |
| 136 | self._last_slot_layout: tuple[int | None, ...] | None = None |
| 137 | is_moe = is_moe_model(self.model) |
| 138 | self._is_moe = is_moe |
| 139 | |
| 140 | self.supported_modules_to_save = ( |
| 141 | {"score", "classifier"} |
| 142 | if self.is_pooling_model |
| 143 | and vllm_config.model_config.score_type == "cross-encoder" |
| 144 | else set() |
| 145 | ) |
| 146 | |
| 147 | # When the engine is started with enable_mixed_moe_lora_format=True |
| 148 | # we force the universal 2D wrapper (FusedMoEWithLoRA) regardless of |
| 149 | # the model's 3D flag, so 2D and 3D adapters can coexist. |
| 150 | self._enable_mixed_moe_lora_format = ( |
| 151 | is_moe and lora_config.enable_mixed_moe_lora_format |
| 152 | ) |
| 153 | self._is_3d_moe_model = ( |
| 154 | self._is_moe |
| 155 | and self.model.is_3d_moe_weight |
| 156 | and not self._enable_mixed_moe_lora_format |
| 157 | ) |
| 158 | # Shared MoE adapters: w13 lora_A / w2 lora_B shared across experts, |
| 159 | # stored as pre-stacked experts.w{1,2,3} tensors (startup opt-in). |
| 160 | self._enable_moe_shared_loras = is_moe and lora_config.enable_moe_shared_loras |
| 161 | self.packed_modules_mapping = process_packed_modules_mapping( |
| 162 | self.model, |
| 163 | force_2d_moe=self._enable_mixed_moe_lora_format, |
| 164 | enable_moe_shared_loras=self._enable_moe_shared_loras, |
| 165 | ) |
| 166 | self._is_non_gated_moe = is_moe and self.model.is_non_gated_moe |
| 167 | self._use_ep = bool( |
| 168 | vllm_config and vllm_config.parallel_config.enable_expert_parallel |
| 169 | ) |
| 170 | self._init_punica_wrapper(max_num_batched_tokens, vllm_config) |
| 171 | self._create_lora_modules() |
| 172 | |
| 173 | self._classification_head: tuple[str, ClassificationHeadWithLoRA] | None = None |
| 174 | if self.is_pooling_model: |
| 175 | classification_heads = [ |
| 176 | (module_name, module) |
| 177 | for module_name, module in self.modules.items() |
| 178 | if isinstance(module, ClassificationHeadWithLoRA) |
| 179 | ] |
| 180 | if classification_heads: |
| 181 | assert len(classification_heads) == 1, ( |
| 182 | "Expected 1 classification head, but found " |
| 183 | f"{len(classification_heads)}." |
| 184 | ) |
| 185 | self._classification_head = classification_heads[0] |
| 186 | |
| 187 | self.moe_ep_load_spec: MoEEPLoadSpec | None = self._build_moe_ep_load_spec() |
| 188 | |
| 189 | self.model.lora_manager = self |
| 190 | |
| 191 | def _init_punica_wrapper( |
| 192 | self, max_num_batched_tokens: int, vllm_config: VllmConfig |
| 193 | ) -> None: |
| 194 | # Used to indicate whether the model is a multimodal model |
| 195 | self.supports_mm: bool = ( |
| 196 | supports_multimodal(self.model) |
| 197 | # In case the model only supports LoRA for |
| 198 | # text modules (e.g. ChatGLM) |
| 199 | and hasattr(self.model, "get_mm_mapping") |
| 200 | ) |
| 201 | self.punica_wrapper_mapping: dict[str, PunicaWrapperBase] = {} |
| 202 | if self.supports_mm: |
| 203 | self._maybe_init_mm(vllm_config, max_num_batched_tokens) |
| 204 | else: |
| 205 | llm_punica_wrapper = get_punica_wrapper( |
| 206 | max_num_batched_tokens, |
| 207 | max_batches=self.max_num_seqs, |
| 208 | device=self.device, |
| 209 | lora_config=self.lora_config, |
| 210 | ) |
| 211 | |
| 212 | self.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY] = ( |
| 213 | llm_punica_wrapper |
| 214 | ) |
| 215 | |
| 216 | def _maybe_init_mm( |
| 217 | self, |
| 218 | vllm_config: VllmConfig, |
| 219 | max_num_batched_tokens: int, |
| 220 | ) -> None: |
| 221 | mm_registry = MULTIMODAL_REGISTRY |
| 222 | |
| 223 | self.supports_tower_connector_lora = False |
| 224 | self.mm_mapping: MultiModelKeys = self.model.get_mm_mapping() |
| 225 | |
| 226 | # Only one language model can be included in the model. |
| 227 | assert len(self.mm_mapping.language_model) == 1 |
| 228 | |
| 229 | # Language model punica wrapper |
| 230 | llm_punica_wrapper = get_punica_wrapper( |
| 231 | max_num_batched_tokens, |
| 232 | max_batches=self.max_num_seqs, |
| 233 | device=self.device, |
| 234 | lora_config=self.lora_config, |
| 235 | ) |
| 236 | |
| 237 | lm_prefix = self.mm_mapping.language_model[0] |
| 238 | self.punica_wrapper_mapping[lm_prefix] = llm_punica_wrapper |
| 239 | |
| 240 | # First, determine if the model supports tower connector LoRA. |
| 241 | self.supports_tower_connector_lora = ( |
| 242 | self.supports_mm and self.model.supports_tower_connector_lora |
| 243 | ) |
| 244 | |
| 245 | # Then, handle the case where the feature is disabled in the config. |
| 246 | if not self.lora_config.enable_tower_connector_lora: |
| 247 | if self.supports_tower_connector_lora: |
| 248 | logger.info( |
| 249 | "%s supports adding LoRA to the tower modules. If needed, " |
| 250 | "please set `enable_tower_connector_lora=True`.", |
| 251 | self.model.__class__.__name__, |
| 252 | ) |
| 253 | self.supports_tower_connector_lora = False |
| 254 | return |
| 255 | |
| 256 | # After this point, the feature is enabled in the config. |
| 257 | # Now check if it's supported by the model. |
| 258 | if not self.supports_tower_connector_lora: |
| 259 | # Enabled but not supported: log warning and return. |
| 260 | logger.warning( |
| 261 | "LoRA with tower connector is enabled, but the model %s " |
| 262 | "does not support it. This will be ignored.", |
| 263 | self.model.__class__.__name__, |
| 264 | ) |
| 265 | return |
| 266 | |
| 267 | # Check if initialize the language model only. |
| 268 | if ( |
| 269 | vllm_config.model_config.multimodal_config |
| 270 | and vllm_config.model_config.multimodal_config.language_model_only |
| 271 | ): |
| 272 | logger.warning( |
| 273 | "Disabling `enable_tower_connector_lora` because the multimodal " |
| 274 | "model is configured to initialize the language model only." |
| 275 | ) |
| 276 | self.supports_tower_connector_lora = False |
| 277 | return |
| 278 | |
| 279 | logger.warning( |
| 280 | "LoRA for the tower and connector of multimodal models is " |
| 281 | "experimental and may contain bugs. Please report any related issues on " |
| 282 | "GitHub if you encounter them." |
| 283 | ) |
| 284 | |
| 285 | mm_budget = MultiModalBudget(vllm_config, mm_registry) |
| 286 | limit_per_prompt = max(mm_budget.mm_max_items_per_prompt.values()) |
| 287 | max_lora_tokens = mm_budget.get_encoder_budget() |
| 288 | lora_token_counts_by_modality = [ |
| 289 | self.model.get_mm_lora_token_counts( |
| 290 | modality=modality, |
| 291 | mm_kwargs=None, |
| 292 | num_mm_embeds=max_lora_tokens, |
| 293 | ) |
| 294 | for modality in mm_budget.mm_max_toks_per_item |
| 295 | ] |
| 296 | num_encoder_tokens = max( |
| 297 | tower_tokens for tower_tokens, _ in lora_token_counts_by_modality |
| 298 | ) |
| 299 | |
| 300 | # Tower wrappers |
| 301 | tower_punica_wrapper = get_punica_wrapper( |
| 302 | num_encoder_tokens, |
| 303 | max_batches=self.max_num_seqs * limit_per_prompt, |
| 304 | device=self.device, |
| 305 | lora_config=self.lora_config, |
| 306 | ) |
| 307 | for prefix in self.mm_mapping.tower_model: |
| 308 | self.punica_wrapper_mapping[prefix] = tower_punica_wrapper |
| 309 | |
| 310 | # Use wrapper for connector if present. |
| 311 | if self.mm_mapping.connector: |
| 312 | connector_tokens = max( |
| 313 | ( |
| 314 | connector_tokens |
| 315 | for _, connector_tokens in lora_token_counts_by_modality |
| 316 | if connector_tokens is not None |
| 317 | ), |
| 318 | default=None, |
| 319 | ) |
| 320 | if connector_tokens is not None: |
| 321 | connector_punica_wrapper = get_punica_wrapper( |
| 322 | connector_tokens, |
| 323 | max_batches=self.max_num_seqs * limit_per_prompt, |
| 324 | device=self.device, |
| 325 | lora_config=self.lora_config, |
| 326 | ) |
| 327 | for prefix in self.mm_mapping.connector: |
| 328 | self.punica_wrapper_mapping[prefix] = connector_punica_wrapper |
| 329 | else: |
| 330 | logger.warning_once( |
| 331 | "Connector LoRA support disabled: model does not implement " |
| 332 | "get_num_mm_connector_tokens(). This method is required to " |
| 333 | "determine the connector's token budget for LoRA operations." |
| 334 | ) |
| 335 | |
| 336 | def __len__(self) -> int: |
| 337 | return len(self._registered_adapters) |
| 338 | |
| 339 | @property |
| 340 | def capacity(self) -> int: |
| 341 | assert self.lora_config.max_cpu_loras is not None |
| 342 | return self.lora_config.max_cpu_loras |
| 343 | |
| 344 | @property |
| 345 | def lora_slots(self) -> int: |
| 346 | return self.lora_config.max_loras |
| 347 | |
| 348 | @property |
| 349 | def adapter_slots(self) -> int: |
| 350 | return self.lora_slots |
| 351 | |
| 352 | def activate_adapter( |
| 353 | self, |
| 354 | lora_id: int, |
| 355 | ) -> bool: |
| 356 | """Move LoRA into a GPU buffer to be used in the forward pass.""" |
| 357 | if lora_id in self._active_adapters: |
| 358 | return False |
| 359 | first_free_slot = next( |
| 360 | ( |
| 361 | (i, lora_id) |
| 362 | for i, lora_id in enumerate(self.lora_index_to_id) |
| 363 | if lora_id is None |
| 364 | ), |
| 365 | None, |
| 366 | ) |
| 367 | if first_free_slot is None: |
| 368 | raise ValueError("No free lora slots") |
| 369 | index, _ = first_free_slot |
| 370 | self._active_adapters[lora_id] = None |
| 371 | lora_model = self._registered_adapters[lora_id] |
| 372 | logger.debug( |
| 373 | "Activating LoRA. int id: %d, slot index: %d", lora_model.id, index |
| 374 | ) |
| 375 | self.lora_index_to_id[index] = lora_model.id |
| 376 | num_lora_weights_applied = 0 |
| 377 | for module_name, module in self.modules.items(): |
| 378 | if isinstance(module, ClassificationHeadWithLoRA): |
| 379 | full_module = self._get_module_to_save_weights(lora_model, module_name) |
| 380 | if full_module is None: |
| 381 | module.reset_module_to_save(index) |
| 382 | else: |
| 383 | module.set_module_to_save( |
| 384 | index, full_module.weight, full_module.bias |
| 385 | ) |
| 386 | |
| 387 | module_lora = self._get_lora_layer_weights(lora_model, module_name) |
| 388 | if not module_lora: |
| 389 | module.reset_lora(index) |
| 390 | logger.debug( |
| 391 | "No LoRA weights found for module %s, skipping.", module_name |
| 392 | ) |
| 393 | continue |
| 394 | |
| 395 | module.set_lora( |
| 396 | index, |
| 397 | module_lora.lora_a, |
| 398 | module_lora.lora_b, |
| 399 | ) |
| 400 | num_lora_weights_applied += 1 |
| 401 | logger.debug("Successfully loaded LoRA weights for module %s.", module_name) |
| 402 | if num_lora_weights_applied == 0: |
| 403 | logger.debug_once( |
| 404 | "No LoRA weights were applied for adapter %s on this worker. " |
| 405 | "Requests may use the base model; check --lora-target-modules. " |
| 406 | "This may be expected with pipeline or expert parallelism.", |
| 407 | lora_model.id, |
| 408 | ) |
| 409 | return True |
| 410 | |
| 411 | def _deactivate_adapter(self, lora_id: int): |
| 412 | try: |
| 413 | index = self.lora_index_to_id.index(lora_id) |
| 414 | self.lora_index_to_id[index] = None |
| 415 | except ValueError: |
| 416 | pass |
| 417 | |
| 418 | def _add_adapter(self, lora: LoRAModel): |
| 419 | self._create_merged_loras_inplace(lora) |
| 420 | self._registered_adapters[lora.id] = lora |
| 421 | |
| 422 | def pin_adapter(self, lora_id: int) -> bool: |
| 423 | """Pin a LoRAModel in the manager cache.""" |
| 424 | raise NotImplementedError( |
| 425 | "Pinning is not supported in LoRAModelManager. " |
| 426 | "Use LRUCacheLoRAModelManager for pinning" |
| 427 | ) # type: ignore |
| 428 | |
| 429 | def _set_adapter_mapping(self, mapping: LoRAMapping) -> None: |
| 430 | # Default to the main language model wrapper |
| 431 | if not (self.supports_mm and self.supports_tower_connector_lora): |
| 432 | target_prefix = ( |
| 433 | self.mm_mapping.language_model[0] |
| 434 | if self.supports_mm |
| 435 | else DEFAULT_LANGUAGE_WRAPPER_KEY |
| 436 | ) |
| 437 | elif mapping.type == LoRAMappingType.TOWER and self.mm_mapping.tower_model: |
| 438 | target_prefix = self.mm_mapping.tower_model[0] |
| 439 | elif mapping.type == LoRAMappingType.CONNECTOR and self.mm_mapping.connector: |
| 440 | target_prefix = self.mm_mapping.connector[0] |
| 441 | else: |
| 442 | target_prefix = self.mm_mapping.language_model[0] |
| 443 | |
| 444 | punica_wrapper = self._get_punica_wrapper(target_prefix) |
| 445 | assert punica_wrapper is not None |
| 446 | |
| 447 | punica_wrapper.update_metadata( |
| 448 | mapping, |
| 449 | self.lora_index_to_id, |
| 450 | self.lora_slots + 1, |
| 451 | self.vocab_size, |
| 452 | ) |
| 453 | |
| 454 | if self._classification_head is not None: |
| 455 | _, classification_head = self._classification_head |
| 456 | if classification_head.punica_wrapper is punica_wrapper: |
| 457 | classification_head.set_output_mapping( |
| 458 | tuple( |
| 459 | self.lora_index_to_id.index(lora_id) if lora_id > 0 else -1 |
| 460 | for lora_id in mapping.prompt_mapping |
| 461 | ) |
| 462 | ) |
| 463 | |
| 464 | def remove_all_adapters(self): |
| 465 | """Remove all LoRAModels from the manager.""" |
| 466 | self._registered_adapters.clear() |
| 467 | self.lora_index_to_id = [None] * self.lora_slots |
| 468 | self._active_adapters.clear() |
| 469 | self._last_mapping = None |
| 470 | self._last_slot_layout = None |
| 471 | |
| 472 | def _create_lora_modules(self): |
| 473 | def _parent_module(module_name: str) -> str: |
| 474 | # module name is a dot separated name. |
| 475 | # for example: |
| 476 | # - given an input 'x.y.z' return 'x.y' |
| 477 | # - given an input 'x' return '' |
| 478 | return module_name.rpartition(".")[0] |
| 479 | |
| 480 | wrapped_by_id: dict[int, BaseLayerWithLoRA] = {} |
| 481 | |
| 482 | for module_name, module in self.model.named_modules(remove_duplicate=False): |
| 483 | if isinstance(module, PPMissingLayer): |
| 484 | continue |
| 485 | |
| 486 | is_classifier_head = module_name in self.supported_modules_to_save |
| 487 | |
| 488 | if self.lora_config.target_modules is None: |
| 489 | if not is_supported_lora_module( |
| 490 | module_name, |
| 491 | module, |
| 492 | self.supported_lora_modules, |
| 493 | ): |
| 494 | continue |
| 495 | elif not self._match_target_modules(module_name) and not is_classifier_head: |
| 496 | # Full classification heads bypass the A/B LoRA target filter. |
| 497 | continue |
| 498 | |
| 499 | punica_wrapper = self._get_punica_wrapper(module_name) |
| 500 | if punica_wrapper is None: |
| 501 | logger.warning( |
| 502 | "Regarding %s, no matching PunicaWrapper " |
| 503 | "is found; %s will be ignored.", |
| 504 | self.model.__class__.__name__, |
| 505 | module_name, |
| 506 | ) |
| 507 | continue |
| 508 | |
| 509 | # TODO: Remove this restriction |
| 510 | # peft error when generating LoRA adapter with "gate" module: |
| 511 | # "Target module NemotronHTopkRouter() is not supported." |
| 512 | # Working LoRA adapter was created using peft with: |
| 513 | # LoraConfig(target_modules="all-linear", ...) |
| 514 | if self._is_non_gated_moe and module_name.endswith("mixer.gate"): |
| 515 | logger.debug_once( |
| 516 | "LoRA is not supported for non-gated MoE gate module." |
| 517 | " %s will be ignored.", |
| 518 | module_name, |
| 519 | ) |
| 520 | continue |
| 521 | |
| 522 | existing_wrapper = wrapped_by_id.get(id(module)) |
| 523 | if existing_wrapper is not None and "lm_head" not in module_name: |
| 524 | # Same underlying module was already wrapped under another |
| 525 | # path (e.g. a MoE gate held both directly on the block and |
| 526 | # inside the MoE runner). The adapter targets the canonical |
| 527 | # path (`mlp.gate`); rewire the alias attribute |
| 528 | # (`runner.gate`) to the SAME wrapper so the forward path |
| 529 | # through the alias still applies LoRA, but do NOT add a |
| 530 | # second entry to self.modules — otherwise `activate_adapter` |
| 531 | # would call `reset_lora` on the alias and wipe the weights |
| 532 | # just set under the canonical name, because the alias can't |
| 533 | # load LoRA weights due to name mismatch. |
| 534 | parent = self.model.get_submodule(_parent_module(module_name)) |
| 535 | # reference |
| 536 | setattr(parent, module_name.rpartition(".")[-1], existing_wrapper) |
| 537 | continue |
| 538 | |
| 539 | parts = module_name.split(".")[-1] |
| 540 | packed_moduled_lst = self.packed_modules_mapping.get(parts, []) |
| 541 | if isinstance(module, MoERunner): |
| 542 | # packed_moduled_lst is used here to just determine whether to |
| 543 | # instantiate FusedMoE3DWithLoRA or FusedMoEWithLoRA, and the |
| 544 | # difference between these two LoRA layers is whether the |
| 545 | # LoRA weights of w1 and w3 have already been fused on disk. |
| 546 | |
| 547 | packed_moduled_lst = ["w13"] if self._is_3d_moe_model else ["w1", "w3"] |
| 548 | if is_classifier_head: |
| 549 | new_module = from_layer_classification( |
| 550 | module, |
| 551 | self.lora_slots, |
| 552 | self.lora_config, |
| 553 | self.model.config, |
| 554 | ) |
| 555 | else: |
| 556 | new_module = from_layer( |
| 557 | module, |
| 558 | self.lora_slots, |
| 559 | self.lora_config, |
| 560 | packed_moduled_lst, |
| 561 | self.model.config, |
| 562 | ) |
| 563 | new_module = replace_submodule(self.model, module_name, new_module) |
| 564 | |
| 565 | if is_classifier_head: |
| 566 | self.model.pooler.replace_classifier(new_module) |
| 567 | if isinstance(new_module, BaseLayerWithLoRA): |
| 568 | wrapped_by_id[id(module)] = new_module |
| 569 | wrapped_by_id[id(new_module)] = new_module |
| 570 | |
| 571 | # (yard1): TODO make this more robust |
| 572 | if "lm_head" in module_name: |
| 573 | logits_processor_module_name = "logits_processor" |
| 574 | parent_module = _parent_module(module_name) |
| 575 | if parent_module: |
| 576 | logits_processor_module_name = ( |
| 577 | f"{parent_module}.{logits_processor_module_name}" |
| 578 | ) |
| 579 | |
| 580 | logits_processor_module = self.model.get_submodule( |
| 581 | logits_processor_module_name |
| 582 | ) |
| 583 | |
| 584 | new_module = replace_submodule( |
| 585 | self.model, |
| 586 | logits_processor_module_name, |
| 587 | from_layer_logits_processor( |
| 588 | logits_processor_module, |
| 589 | module, |
| 590 | self.lora_slots, |
| 591 | self.lora_config, |
| 592 | self.model.config, |
| 593 | ), |
| 594 | ) |
| 595 | |
| 596 | # Some matched modules can be unsupported by LoRA wrappers |
| 597 | # (e.g. subclasses with specialized forward behavior). |
| 598 | if not isinstance(new_module, BaseLayerWithLoRA): |
| 599 | error_msg = ( |
| 600 | "LoRA target module " |
| 601 | f"{module_name} ({type(module).__name__}) matched the " |
| 602 | "deployment configuration but could not be wrapped by any " |
| 603 | "LoRA layer implementation." |
| 604 | ) |
| 605 | if self.lora_config.target_modules is not None: |
| 606 | raise ValueError( |
| 607 | f"{error_msg} target_modules=" |
| 608 | f"{sorted(self.lora_config.target_modules)}" |
| 609 | ) |
| 610 | logger.warning_once("%s It will be ignored.", error_msg) |
| 611 | continue |
| 612 | self.register_module(module_name, new_module) |
| 613 | |
| 614 | self._register_packed_modules(module_name) |
| 615 | # All lora layers share the same punica_wrapper based on reference. |
| 616 | new_module.set_mapping(punica_wrapper) |
| 617 | |
| 618 | def register_module(self, module_name: str, module: "BaseLayerWithLoRA"): |
| 619 | assert isinstance(module, BaseLayerWithLoRA), ( |
| 620 | f"Module {module_name} must be a BaseLayerWithLoRA instance, " |
| 621 | f"got {type(module)}" |
| 622 | ) |
| 623 | self.modules[module_name] = module |
| 624 | |
| 625 | @staticmethod |
| 626 | def _pad_lora_pairs_to_triplets( |
| 627 | loras: list[LoRALayerWeights | None], |
| 628 | ) -> list[LoRALayerWeights | None]: |
| 629 | """Pad LoRA weight pairs to triplets for non-gated MoE. |
| 630 | |
| 631 | For non-gated MoE, each expert has 2 entries (w1, w2) that need to be |
| 632 | padded to triplets (w1, w2, None) to match pack_moe expectations. |
| 633 | """ |
| 634 | assert len(loras) % 2 == 0, "Expected pairs of LoRA weights for non-gated MoE." |
| 635 | padded: list[LoRALayerWeights | None] = [] |
| 636 | for i in range(0, len(loras), 2): |
| 637 | padded.extend(loras[i : i + 2]) |
| 638 | padded.append(None) |
| 639 | return padded |
| 640 | |
| 641 | def create_dummy_lora( |
| 642 | self, |
| 643 | lora_id: int, |
| 644 | rank: int, |
| 645 | embedding_modules: dict[str, str] | None = None, |
| 646 | ) -> LoRAModel: |
| 647 | """Create zero-initialized LoRAModel for warmup.""" |
| 648 | model = LoRAModel(lora_id, rank, {}) |
| 649 | for module_name, module in self.model.named_modules(): |
| 650 | if ( |
| 651 | not is_supported_lora_module( |
| 652 | module_name, |
| 653 | module, |
| 654 | self.supported_lora_modules, |
| 655 | ) |
| 656 | or not self._match_target_modules(module_name) |
| 657 | or not isinstance(module, BaseLayerWithLoRA) |
| 658 | or self._get_punica_wrapper(module_name) is None |
| 659 | ): |
| 660 | continue |
| 661 | parts = module_name.split(".") |
| 662 | if module_name not in self.packed_modules: |
| 663 | assert embedding_modules is not None |
| 664 | if parts[-1] in embedding_modules: |
| 665 | # Special-case lm_head: wrapped by LogitsProcessorWithLoRA. |
| 666 | # LoRA input dim is hidden_size, output dim is vocab size. |
| 667 | # LogitsProcessorWithLoRA handles extra vocab size directly. |
| 668 | if parts[-1] == "lm_head": |
| 669 | input_dim = module.lora_a_stacked[0].shape[-1] |
| 670 | output_dim = module.lora_b_stacked[0].shape[-2] |
| 671 | else: |
| 672 | input_dim = ( |
| 673 | module.base_layer.org_vocab_size |
| 674 | if hasattr(module.base_layer, "org_vocab_size") |
| 675 | else module.base_layer.weight.shape[1] |
| 676 | ) |
| 677 | output_dim = ( |
| 678 | module.base_layer.embedding_dim |
| 679 | if hasattr(module.base_layer, "embedding_dim") |
| 680 | else module.base_layer.weight.shape[0] |
| 681 | ) |
| 682 | lora = LoRALayerWeights.create_dummy_lora_weights( |
| 683 | module_name, |
| 684 | input_dim, |
| 685 | output_dim, |
| 686 | rank, |
| 687 | module.lora_a_stacked[0].dtype, |
| 688 | "cpu", |
| 689 | ) |
| 690 | model.loras[module_name] = lora |
| 691 | elif module.__class__.__name__ == "FusedMoE3DWithLoRA": |
| 692 | # Case for 3D moe model |
| 693 | # w2 |
| 694 | lora = LoRALayerWeights.create_dummy_lora_weights( |
| 695 | module_name, |
| 696 | module.w2_input_size, |
| 697 | module.w2_output_size, |
| 698 | rank * module.w2_lora_a_stacked[0].shape[1], # rank*num_experts |
| 699 | module.w2_lora_a_stacked[0].dtype, |
| 700 | "cpu", |
| 701 | ) |
| 702 | model.loras[module_name] = lora |
| 703 | # w13 |
| 704 | lora = LoRALayerWeights.create_dummy_lora_weights( |
| 705 | module_name, |
| 706 | module.w13_input_size, |
| 707 | module.w13_output_size, |
| 708 | rank |
| 709 | * module.w13_lora_a_stacked[0].shape[1], # rank*num_experts |
| 710 | module.w13_lora_a_stacked[0].dtype, |
| 711 | "cpu", |
| 712 | ) |
| 713 | model.loras[module_name + ".base_layer"] = lora |
| 714 | else: |
| 715 | lora = LoRALayerWeights.create_dummy_lora_weights( |
| 716 | module_name, |
| 717 | module.lora_a_stacked[0].shape[-1], |
| 718 | module.lora_b_stacked[0].shape[-2], |
| 719 | rank, |
| 720 | module.lora_a_stacked[0].dtype, |
| 721 | "cpu", |
| 722 | ) |
| 723 | model.loras[module_name] = lora |
| 724 | else: |
| 725 | parts = module_name.split(".") |
| 726 | replacements = self.packed_modules_mapping[parts[-1]] |
| 727 | n_slices = getattr(module, "n_slices", len(replacements)) |
| 728 | if module.__class__.__name__ == "FusedMoEWithLoRA": |
| 729 | replacements = replacements[ |
| 730 | : len(module.lora_a_stacked) // self.lora_slots |
| 731 | ] |
| 732 | subloras: list[LoRALayerWeights | None] = [] |
| 733 | # HACK: overrides replacements for qkvz = qkv + z case. |
| 734 | # Any better methods to handle this case? |
| 735 | if n_slices != len(replacements): |
| 736 | replacements = [f"slice_{i}" for i in range(n_slices)] |
| 737 | for i, r in enumerate(replacements): |
| 738 | lora = LoRALayerWeights.create_dummy_lora_weights( |
| 739 | module_name + "." + r, |
| 740 | module.lora_a_stacked[i].shape[-1], |
| 741 | module.lora_b_stacked[i].shape[-2], |
| 742 | rank, |
| 743 | module.lora_a_stacked[i].dtype, |
| 744 | "cpu", |
| 745 | ) |
| 746 | subloras.append(lora) |
| 747 | if module.__class__.__name__ == "FusedMoEWithLoRA": |
| 748 | # For non-gated MoE, pad subloras to 3 elements per expert |
| 749 | # to match pack_moe expectations (w1, w2, None for w3) |
| 750 | if self._is_non_gated_moe and len(subloras) > 0: |
| 751 | subloras = self._pad_lora_pairs_to_triplets(subloras) |
| 752 | lora = PackedLoRALayerWeights.pack_moe( |
| 753 | subloras, module_name, is_non_gated_moe=self._is_non_gated_moe |
| 754 | ) |
| 755 | else: |
| 756 | lora = PackedLoRALayerWeights.pack(subloras) |
| 757 | model.loras[module_name] = lora |
| 758 | return model |
| 759 | |
| 760 | def get_dummy_lora_warmup_rank(self, default_rank: int) -> int: |
| 761 | """Return a dummy LoRA rank compatible with wrapped modules. |
| 762 | |
| 763 | Dummy LoRAs keep warmup memory low by using a small rank. Fully |
| 764 | sharded MoE wrappers additionally require the dummy rank to be divisible |
| 765 | by tensor parallel size because they shard W13 along the rank axis. |
| 766 | """ |
| 767 | if not self.lora_config.fully_sharded_loras: |
| 768 | return default_rank |
| 769 | |
| 770 | required_multiple = 1 |
| 771 | for module in self.modules.values(): |
| 772 | if not getattr(module, "fully_sharded", False): |
| 773 | continue |
| 774 | required_multiple = math.lcm(required_multiple, module.tp_size) |
| 775 | |
| 776 | if required_multiple == 1 or default_rank % required_multiple == 0: |
| 777 | return default_rank |
| 778 | |
| 779 | adjusted_rank = ( |
| 780 | (default_rank + required_multiple - 1) // required_multiple |
| 781 | ) * required_multiple |
| 782 | if adjusted_rank > self.lora_config.max_lora_rank: |
| 783 | raise ValueError( |
| 784 | "Unable to choose a dummy LoRA warmup rank compatible with " |
| 785 | "fully sharded MoE modules: " |
| 786 | f"default_rank={default_rank}, " |
| 787 | f"required_multiple={required_multiple}, " |
| 788 | f"max_lora_rank={self.lora_config.max_lora_rank}" |
| 789 | ) |
| 790 | return adjusted_rank |
| 791 | |
| 792 | def _match_target_modules(self, module_name: str) -> bool: |
| 793 | """Check if a module passes the deployment-time target filter. |
| 794 | |
| 795 | Args: |
| 796 | module_name: Full dot-separated module name (e.g., |
| 797 | "model.layers.0.self_attn.o_proj") |
| 798 | |
| 799 | Returns: |
| 800 | True if the module passes the filter, False otherwise. |
| 801 | |
| 802 | """ |
| 803 | return is_in_target_modules( |
| 804 | module_name, |
| 805 | self.lora_config.target_modules, |
| 806 | self.packed_modules_mapping, |
| 807 | ) |
| 808 | |
| 809 | def _get_punica_wrapper(self, module_name: str) -> PunicaWrapperBase | None: |
| 810 | """Determine whether this module supports LoRA and which wrapper to use.""" |
| 811 | # For language model (early return) |
| 812 | if not self.supports_mm: |
| 813 | return self.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY] |
| 814 | |
| 815 | if module_name in self.supported_modules_to_save: |
| 816 | lm_wrapper = self.punica_wrapper_mapping.get( |
| 817 | self.mm_mapping.language_model[0] |
| 818 | ) |
| 819 | return lm_wrapper |
| 820 | |
| 821 | # For multimodal model |
| 822 | # NOTE Sort by prefix length (descending) to match the longest prefix first |
| 823 | # e.g., 'visual.merger' should match 'visual.merger' instead of 'visual.' |
| 824 | for prefix in sorted(self.punica_wrapper_mapping.keys(), key=len, reverse=True): |
| 825 | if module_name.startswith(prefix): |
| 826 | return self.punica_wrapper_mapping[prefix] |
| 827 | |
| 828 | return None |
| 829 | |
| 830 | def _register_packed_modules(self, module_full_name: str) -> None: |
| 831 | parts = module_full_name.split(".") |
| 832 | module_name = parts[-1] |
| 833 | replacements = self.packed_modules_mapping.get(module_name, []) |
| 834 | # When replacements is less than or equal to 1, it indicates that this |
| 835 | # module is not a packed module. |
| 836 | if len(replacements) <= 1: |
| 837 | return |
| 838 | prefix = ".".join(parts[:-1]) |
| 839 | self.packed_modules[module_full_name] = [ |
| 840 | prefix + "." + r if prefix else r for r in replacements |
| 841 | ] |
| 842 | |
| 843 | def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None: |
| 844 | for module_name, new_module_names in self.packed_modules.items(): |
| 845 | # For 2D RoutedExperts modules with EP, narrow the per-expert |
| 846 | # sub-module list to this rank's owned experts so pack_moe |
| 847 | # produces a tensor sized to local_num_experts directly. |
| 848 | packed_module_names = new_module_names |
| 849 | if module_name.endswith(".experts"): |
| 850 | new_module_names = self._restrict_to_local_experts( |
| 851 | module_name, new_module_names |
| 852 | ) |
| 853 | replacement_loras: list[LoRALayerWeights | None] = [] |
| 854 | has_replacement = False |
| 855 | for r in new_module_names: |
| 856 | lora = self._get_lora_layer_weights(lora_model, r) |
| 857 | replacement_loras.append(lora) |
| 858 | if lora: |
| 859 | has_replacement = True |
| 860 | if not has_replacement: |
| 861 | continue |
| 862 | for i in range(len(replacement_loras)): |
| 863 | if replacement_loras[i]: |
| 864 | continue |
| 865 | replacement_loras[i] = None |
| 866 | # HACK Temporary solution for the pool model. |
| 867 | if self.is_pooling_model and not lora_model.check_lora_name(module_name): |
| 868 | replaced_module_name = module_name.removeprefix("model.") |
| 869 | if lora_model.check_lora_name(replaced_module_name): |
| 870 | module_name = replaced_module_name |
| 871 | if module_name.endswith(".experts"): |
| 872 | if self._enable_moe_shared_loras: |
| 873 | lora_model.loras[module_name] = ( |
| 874 | PackedLoRALayerWeights.pack_moe_stacked( |
| 875 | replacement_loras, |
| 876 | module_name, |
| 877 | ) |
| 878 | ) |
| 879 | for module in packed_module_names: |
| 880 | lora_model.loras.pop(module, None) |
| 881 | continue |
| 882 | if self._is_non_gated_moe and len(replacement_loras) > 0: |
| 883 | replacement_loras = self._pad_lora_pairs_to_triplets( |
| 884 | replacement_loras |
| 885 | ) |
| 886 | lora_model.loras[module_name] = PackedLoRALayerWeights.pack_moe( |
| 887 | replacement_loras, |
| 888 | module_name, |
| 889 | is_non_gated_moe=self._is_non_gated_moe, |
| 890 | ) |
| 891 | else: |
| 892 | lora_model.loras[module_name] = PackedLoRALayerWeights.pack( |
| 893 | replacement_loras |
| 894 | ) |
| 895 | # Drop every candidate sub-module, including non-local expert |
| 896 | # entries that were loaded but did not contribute to the |
| 897 | # packed result. Without this they would keep extra CPU |
| 898 | # memory alive after pack_moe. |
| 899 | for module in packed_module_names: |
| 900 | lora_model.loras.pop(module, None) |
| 901 | |
| 902 | for lora in lora_model.loras.values(): |
| 903 | lora.optimize() |
| 904 | |
| 905 | for module_name, module in self.modules.items(): |
| 906 | if isinstance(module, FusedMoE3DWithLoRA): |
| 907 | self._stack_moe_lora_weights(lora_model, module, module_name) |
| 908 | elif isinstance(module, FusedMoEWithLoRA): |
| 909 | # When mixed mode is enabled the universal 2D wrapper has to |
| 910 | # absorb both 2D and 3D-format adapters. 3D-format adapters |
| 911 | # need to be split into per-(w1, w2, w3) tensors before the |
| 912 | # 2D set_lora can copy them into the stacked buffers. |
| 913 | if self._enable_mixed_moe_lora_format and getattr( |
| 914 | lora_model, "is_3d_lora_weight", False |
| 915 | ): |
| 916 | self._convert_3d_to_2d_moe_lora(lora_model, module, module_name) |
| 917 | else: |
| 918 | self._slice_moe_lora_ep(lora_model, module, module_name) |
| 919 | |
| 920 | first_lora: LoRALayerWeights = next(iter(lora_model.loras.values())) |
| 921 | assert first_lora.lora_a is not None |
| 922 | if isinstance(first_lora.lora_a, list): |
| 923 | lora_device = next(iter(first_lora.lora_a)) |
| 924 | else: |
| 925 | lora_device = first_lora.lora_a.device |
| 926 | # Execute pin_memory after LoRA weight merging, mainly because: |
| 927 | # 1. Some MoE models have a large number of LoRA weights. If we |
| 928 | # perform # pin_memory immediately after loading weights, the |
| 929 | # overhead is significant. |
| 930 | # 2. The weight packing above (e.g., pack_moe) may invalidate the |
| 931 | # pin_memory allocation, so we execute it after packing. |
| 932 | |
| 933 | pin_memory = str(lora_device) == "cpu" and PIN_MEMORY |
| 934 | if pin_memory: |
| 935 | for lora in lora_model.loras.values(): |
| 936 | if isinstance(lora.lora_a, list): |
| 937 | for index in range(len(lora.lora_a)): |
| 938 | if lora.lora_a[index] is None: |
| 939 | continue |
| 940 | lora.lora_a[index] = lora.lora_a[index].pin_memory() |
| 941 | lora.lora_b[index] = lora.lora_b[index].pin_memory() |
| 942 | else: |
| 943 | lora.lora_a = lora.lora_a.pin_memory() |
| 944 | lora.lora_b = lora.lora_b.pin_memory() |
| 945 | |
| 946 | if lora_model.modules_to_save: |
| 947 | for full_module in lora_model.modules_to_save.values(): |
| 948 | full_module.pin_memory() |
| 949 | |
| 950 | def _stack_moe_lora_weights( |
| 951 | self, lora_model: LoRAModel, module: FusedMoE3DWithLoRA, module_name: str |
| 952 | ): |
| 953 | module_lora = self._get_lora_layer_weights(lora_model, module_name) |
| 954 | |
| 955 | # Note (gnovack) - If MOE lora weights are not split into |
| 956 | # num_experts chunks, we split them here |
| 957 | if module_lora and torch.is_tensor(module_lora.lora_a): |
| 958 | # Handle PEFT file format where experts.base_layer is the |
| 959 | # gate_up_proj and experts is the down_proj |
| 960 | gate_up_proj_lora = self._get_lora_layer_weights( |
| 961 | lora_model, module_name + ".base_layer" |
| 962 | ) |
| 963 | down_proj_lora = module_lora |
| 964 | # FIXME Edge case where LoRA is not added to gate_up_proj |
| 965 | # or down_proj |
| 966 | assert gate_up_proj_lora is not None |
| 967 | assert down_proj_lora is not None |
| 968 | if self._is_3d_moe_model: |
| 969 | local_num_experts = module.w13_lora_a_stacked[0].shape[1] |
| 970 | # The checkpoint holds weights for all global experts, but |
| 971 | # each EP rank owns only local_num_experts. Reshape against |
| 972 | # the adapter's actual expert count, then slice this rank's |
| 973 | # owned expert range before it gets copied into the local |
| 974 | # stacked buffer. For non-EP (local == global) this is a |
| 975 | # no-op slice. |
| 976 | global_num_experts = module.global_num_experts |
| 977 | ep_rank = module.ep_rank |
| 978 | expert_start = ep_rank * local_num_experts |
| 979 | expert_end = expert_start + local_num_experts |
| 980 | |
| 981 | # (num_experts,rank,input_size) |
| 982 | gate_up_proj_lora.lora_a = gate_up_proj_lora.lora_a.reshape( |
| 983 | global_num_experts, -1, gate_up_proj_lora.lora_a.shape[-1] |
| 984 | )[expert_start:expert_end].contiguous() |
| 985 | down_proj_lora.lora_a = down_proj_lora.lora_a.reshape( |
| 986 | global_num_experts, -1, down_proj_lora.lora_a.shape[-1] |
| 987 | )[expert_start:expert_end].contiguous() |
| 988 | |
| 989 | # (output_size,rank,num_experts) |
| 990 | gate_up_proj_lora.lora_b = gate_up_proj_lora.lora_b.reshape( |
| 991 | gate_up_proj_lora.lora_b.shape[0], -1, global_num_experts |
| 992 | )[..., expert_start:expert_end] |
| 993 | down_proj_lora.lora_b = down_proj_lora.lora_b.reshape( |
| 994 | down_proj_lora.lora_b.shape[0], -1, global_num_experts |
| 995 | )[..., expert_start:expert_end] |
| 996 | |
| 997 | # (num_experts,output_size,rank) |
| 998 | gate_up_proj_lora.lora_b = gate_up_proj_lora.lora_b.permute( |
| 999 | 2, 0, 1 |
| 1000 | ).contiguous() |
| 1001 | down_proj_lora.lora_b = down_proj_lora.lora_b.permute( |
| 1002 | 2, 0, 1 |
| 1003 | ).contiguous() |
| 1004 | |
| 1005 | module_lora.lora_a = [ |
| 1006 | gate_up_proj_lora.lora_a, |
| 1007 | down_proj_lora.lora_a, |
| 1008 | ] |
| 1009 | module_lora.lora_b = [ |
| 1010 | gate_up_proj_lora.lora_b, |
| 1011 | down_proj_lora.lora_b, |
| 1012 | ] |
| 1013 | else: |
| 1014 | # Some 3D MoE models haven't added the `is_3d_moe_weight` |
| 1015 | # attribute yet, so fallback here |
| 1016 | num_experts = module_lora.lora_a.shape[0] // module_lora.rank |
| 1017 | |
| 1018 | gate_proj_a = gate_up_proj_lora.lora_a.chunk(num_experts, dim=0) |
| 1019 | up_proj_a = gate_up_proj_lora.lora_a.chunk(num_experts, dim=0) |
| 1020 | |
| 1021 | gate_proj_b = gate_up_proj_lora.lora_b[::2, ...].chunk( |
| 1022 | num_experts, dim=-1 |
| 1023 | ) |
| 1024 | up_proj_b = gate_up_proj_lora.lora_b[1::2, ...].chunk( |
| 1025 | num_experts, dim=-1 |
| 1026 | ) |
| 1027 | |
| 1028 | down_proj_a = down_proj_lora.lora_a.chunk(num_experts, dim=0) |
| 1029 | down_proj_b = down_proj_lora.lora_b.chunk(num_experts, dim=-1) |
| 1030 | |
| 1031 | lora_a = [] |
| 1032 | lora_b = [] |
| 1033 | for i in range(num_experts): |
| 1034 | lora_a.append(gate_proj_a[i]) |
| 1035 | lora_a.append(down_proj_a[i]) |
| 1036 | lora_a.append(up_proj_a[i]) |
| 1037 | |
| 1038 | lora_b.append(gate_proj_b[i]) |
| 1039 | lora_b.append(down_proj_b[i]) |
| 1040 | lora_b.append(up_proj_b[i]) |
| 1041 | |
| 1042 | module_lora.lora_a = lora_a |
| 1043 | module_lora.lora_b = lora_b |
| 1044 | |
| 1045 | def _convert_3d_to_2d_moe_lora( |
| 1046 | self, |
| 1047 | lora_model: LoRAModel, |
| 1048 | module: FusedMoEWithLoRA, |
| 1049 | module_name: str, |
| 1050 | ) -> None: |
| 1051 | """Convert a 3D-format MoE LoRA checkpoint into the 2D pack layout |
| 1052 | that `FusedMoEWithLoRA.set_lora` expects. |
| 1053 | |
| 1054 | On disk the 3D PEFT layout stores two flat tensor pairs per layer: |
| 1055 | - `{module}.base_layer.lora_{A,B}`: gate_up_proj, with shapes |
| 1056 | `(rank * num_experts, hidden)` / `(intermediate * 2, |
| 1057 | rank * num_experts)` |
| 1058 | - `{module}.lora_{A,B}`: down_proj, with shapes |
| 1059 | `(rank * num_experts, intermediate)` / `(hidden, |
| 1060 | rank * num_experts)` |
| 1061 | The 2D wrapper expects three stacked per-expert tensors, |
| 1062 | `[w1, w2, w3]`, with `(num_experts, rank, in)` for lora_a and |
| 1063 | `(num_experts, out, rank)` for lora_b. In the 3D layout w1 |
| 1064 | (gate_proj) and w3 (up_proj) share the rank-r intermediate |
| 1065 | representation, so both halves use the same lora_a tensor. |
| 1066 | |
| 1067 | Only invoked when `enable_mixed_moe_lora_format=True` and the |
| 1068 | source LoRARequest declares `is_3d_lora_weight=True`. |
| 1069 | """ |
| 1070 | gate_up_proj_lora = self._get_lora_layer_weights( |
| 1071 | lora_model, module_name + ".base_layer" |
| 1072 | ) |
| 1073 | down_proj_lora = self._get_lora_layer_weights(lora_model, module_name) |
| 1074 | if gate_up_proj_lora is None or down_proj_lora is None: |
| 1075 | # Either the adapter omits the experts entirely or the file |
| 1076 | # layout differs from what this path supports; leave the entry |
| 1077 | # untouched so set_lora can raise a clear error if needed. |
| 1078 | return |
| 1079 | |
| 1080 | local_num_experts = module.local_num_experts |
| 1081 | global_num_experts = module.global_num_experts |
| 1082 | ep_rank = module.ep_rank |
| 1083 | expert_start = ep_rank * local_num_experts |
| 1084 | expert_end = expert_start + local_num_experts |
| 1085 | |
| 1086 | # Reshape and EP-slice into per-expert 3D tensors. This mirrors |
| 1087 | # `_stack_moe_lora_weights`; for non-EP runs the slice is a no-op. |
| 1088 | gate_up_a = gate_up_proj_lora.lora_a.reshape( |
| 1089 | global_num_experts, -1, gate_up_proj_lora.lora_a.shape[-1] |
| 1090 | )[expert_start:expert_end].contiguous() |
| 1091 | gate_up_b = ( |
| 1092 | gate_up_proj_lora.lora_b.reshape( |
| 1093 | gate_up_proj_lora.lora_b.shape[0], -1, global_num_experts |
| 1094 | )[..., expert_start:expert_end] |
| 1095 | .permute(2, 0, 1) |
| 1096 | .contiguous() |
| 1097 | ) |
| 1098 | down_a = down_proj_lora.lora_a.reshape( |
| 1099 | global_num_experts, -1, down_proj_lora.lora_a.shape[-1] |
| 1100 | )[expert_start:expert_end].contiguous() |
| 1101 | down_b = ( |
| 1102 | down_proj_lora.lora_b.reshape( |
| 1103 | down_proj_lora.lora_b.shape[0], -1, global_num_experts |
| 1104 | )[..., expert_start:expert_end] |
| 1105 | .permute(2, 0, 1) |
| 1106 | .contiguous() |
| 1107 | ) |
| 1108 | |
| 1109 | # Split the fused gate_up_proj output dim into separate w1 / w3 |
| 1110 | # halves. GPT-OSS interleaves them along the output dim, all other |
| 1111 | # 3D MoE checkpoints we know about concatenate them. |
| 1112 | intermediate_x2 = gate_up_b.shape[1] |
| 1113 | if intermediate_x2 % 2 != 0: |
| 1114 | raise ValueError( |
| 1115 | "Expected gate_up_proj LoRA-B output dim to be 2 * intermediate, " |
| 1116 | f"got {intermediate_x2}." |
| 1117 | ) |
| 1118 | intermediate = intermediate_x2 // 2 |
| 1119 | base_arch = self.model.config.architectures[0] |
| 1120 | if base_arch == "GptOssForCausalLM": |
| 1121 | w1_b = gate_up_b[:, ::2, :].contiguous() |
| 1122 | w3_b = gate_up_b[:, 1::2, :].contiguous() |
| 1123 | else: |
| 1124 | w1_b = gate_up_b[:, :intermediate, :].contiguous() |
| 1125 | w3_b = gate_up_b[:, intermediate:, :].contiguous() |
| 1126 | |
| 1127 | # In the 3D layout w1 and w3 share the same rank-r mid |
| 1128 | # representation, so they reuse the same lora_a tensor. The 2D |
| 1129 | # wrapper's set_lora copies whatever it gets here into independent |
| 1130 | # per-slice buffers, so the sharing is purely a CPU-side memory |
| 1131 | # optimization and does not affect numerics. |
| 1132 | down_proj_lora.lora_a = [gate_up_a, down_a, gate_up_a] |
| 1133 | down_proj_lora.lora_b = [w1_b, down_b, w3_b] |
| 1134 | # Drop the redundant base_layer entry to avoid double pin_memory |
| 1135 | # and to keep the activation path looking up only the wrapper key. |
| 1136 | lora_model.loras.pop(module_name + ".base_layer", None) |
| 1137 | |
| 1138 | def _slice_moe_lora_ep( |
| 1139 | self, |
| 1140 | lora_model: LoRAModel, |
| 1141 | module: FusedMoEWithLoRA, |
| 1142 | module_name: str, |
| 1143 | ) -> None: |
| 1144 | """Slice cached LoRA tensors down to this rank's local experts. |
| 1145 | |
| 1146 | Shared factors have expert dimension one and remain available on every |
| 1147 | rank. Per-expert factors are narrowed to the local expert block. |
| 1148 | """ |
| 1149 | if not module.use_ep: |
| 1150 | return |
| 1151 | module_lora = self._get_lora_layer_weights(lora_model, module_name) |
| 1152 | if module_lora is None or not isinstance(module_lora.lora_a, list): |
| 1153 | return |
| 1154 | |
| 1155 | local_num_experts = module.local_num_experts |
| 1156 | global_num_experts = module.global_num_experts |
| 1157 | ep_rank = module.ep_rank |
| 1158 | expert_start = ep_rank * local_num_experts |
| 1159 | expert_end = expert_start + local_num_experts |
| 1160 | |
| 1161 | def _slice_local(t: torch.Tensor | None) -> torch.Tensor | None: |
| 1162 | if t is not None and t.shape[0] == global_num_experts: |
| 1163 | return t[expert_start:expert_end].contiguous() |
| 1164 | return t |
| 1165 | |
| 1166 | module_lora.lora_a = [_slice_local(a) for a in module_lora.lora_a] |
| 1167 | module_lora.lora_b = [_slice_local(b) for b in module_lora.lora_b] |
| 1168 | |
| 1169 | def _restrict_to_local_experts( |
| 1170 | self, module_name: str, new_module_names: list[str] |
| 1171 | ) -> list[str]: |
| 1172 | """Narrow a flat expert-major sub-module list to this rank's experts. |
| 1173 | |
| 1174 | ``new_module_names`` is produced by |
| 1175 | ``fused_moe_make_expert_params_mapping`` and is ordered |
| 1176 | ``[e=0,w1, e=0,w2, e=0,w3, e=1,w1, ...]`` (non-gated MoE has 2 |
| 1177 | entries per expert instead of 3). When the module is a 2D |
| 1178 | ``FusedMoEWithLoRA`` with EP enabled, we slice the list to the |
| 1179 | contiguous block of experts owned by this rank so the downstream |
| 1180 | ``pack_moe`` call only consumes local weights and produces a |
| 1181 | tensor sized to ``local_num_experts`` directly. |
| 1182 | |
| 1183 | Returns the original list unchanged for non-MoE modules, the 3D |
| 1184 | MoE path (handled separately by ``_stack_moe_lora_weights``), |
| 1185 | modules without EP, or layouts we cannot cleanly partition. |
| 1186 | """ |
| 1187 | module = self.modules.get(module_name) |
| 1188 | if not isinstance(module, FusedMoEWithLoRA): |
| 1189 | return new_module_names |
| 1190 | if isinstance(module, FusedMoE3DWithLoRA): |
| 1191 | return new_module_names |
| 1192 | if not getattr(module, "use_ep", False): |
| 1193 | return new_module_names |
| 1194 | global_num_experts = module.global_num_experts |
| 1195 | local_num_experts = module.local_num_experts |
| 1196 | ep_rank = module.ep_rank |
| 1197 | if global_num_experts <= 0 or len(new_module_names) % global_num_experts != 0: |
| 1198 | return new_module_names |
| 1199 | per_expert = len(new_module_names) // global_num_experts |
| 1200 | start = ep_rank * local_num_experts * per_expert |
| 1201 | end = start + local_num_experts * per_expert |
| 1202 | return new_module_names[start:end] |
| 1203 | |
| 1204 | def _build_moe_ep_load_spec(self) -> MoEEPLoadSpec | None: |
| 1205 | """Per-rank slicing metadata for 2D RoutedEXperts LoRA modules.""" |
| 1206 | if not self._use_ep or not self._is_moe: |
| 1207 | return None |
| 1208 | module = next( |
| 1209 | ( |
| 1210 | m |
| 1211 | for m in self.modules.values() |
| 1212 | if isinstance(m, FusedMoEWithLoRA) |
| 1213 | and not isinstance(m, FusedMoE3DWithLoRA) |
| 1214 | ), |
| 1215 | None, |
| 1216 | ) |
| 1217 | if module is None: |
| 1218 | return None |
| 1219 | return MoEEPLoadSpec( |
| 1220 | ep_rank=module.ep_rank, |
| 1221 | local_num_experts=module.local_num_experts, |
| 1222 | global_num_experts=module.global_num_experts, |
| 1223 | ) |
| 1224 | |
| 1225 | def _get_lora_layer_weights( |
| 1226 | self, lora_model: LoRAModel, module_name: str |
| 1227 | ) -> LoRALayerWeights | None: |
| 1228 | org_module_name = module_name |
| 1229 | if self.is_pooling_model and not lora_model.check_lora_name(module_name): |
| 1230 | # If it's a pool model, and the layer name is not found, |
| 1231 | # remove the prefix 'model.' and search again. |
| 1232 | module_name = module_name.removeprefix("model.") |
| 1233 | if lora_model.check_lora_name(module_name): |
| 1234 | org_module_name = module_name |
| 1235 | logger.info_once( |
| 1236 | "For the pool model, successfully loaded the LoRA weights " |
| 1237 | "after removing the prefix 'model.'." |
| 1238 | ) |
| 1239 | return lora_model.get_lora(org_module_name) |
| 1240 | |
| 1241 | def _get_module_to_save_weights( |
| 1242 | self, lora_model: LoRAModel, module_name: str |
| 1243 | ) -> LoRAFullModuleWeights | None: |
| 1244 | weights = lora_model.get_module_to_save(module_name) |
| 1245 | if weights is not None: |
| 1246 | return weights |
| 1247 | |
| 1248 | unprefixed_module_name = module_name.removeprefix("model.") |
| 1249 | weights = lora_model.get_module_to_save(unprefixed_module_name) |
| 1250 | if weights is not None and unprefixed_module_name != module_name: |
| 1251 | logger.info_once( |
| 1252 | "For the pool model, successfully loaded the full module weights " |
| 1253 | "after removing the prefix 'model.'." |
| 1254 | ) |
| 1255 | return weights |
| 1256 | |
| 1257 | def _validate_modules_to_save(self, lora_model: LoRAModel) -> None: |
| 1258 | if not lora_model.modules_to_save: |
| 1259 | return |
| 1260 | saved_module_name = next(iter(lora_model.modules_to_save)) |
| 1261 | if self._classification_head is None: |
| 1262 | raise ValueError( |
| 1263 | f"Cannot load full module {saved_module_name!r}: the model does " |
| 1264 | "not expose a unique ClassificationHeadWithLoRA." |
| 1265 | ) |
| 1266 | |
| 1267 | module_name, wrapper = self._classification_head |
| 1268 | full_module = self._get_module_to_save_weights(lora_model, module_name) |
| 1269 | if full_module is None: |
| 1270 | raise ValueError( |
| 1271 | f"Full module {saved_module_name!r} does not match the model's " |
| 1272 | f"classification head {module_name!r}." |
| 1273 | ) |
| 1274 | |
| 1275 | received_weight_shape = tuple(full_module.weight.shape) |
| 1276 | if ( |
| 1277 | full_module.weight.ndim != 2 |
| 1278 | or full_module.weight.size(0) < 1 |
| 1279 | or full_module.weight.size(0) > wrapper.max_lora_cls_labels |
| 1280 | or full_module.weight.size(1) != wrapper.input_size |
| 1281 | ): |
| 1282 | raise ValueError( |
| 1283 | f"Full module {saved_module_name!r} for {module_name!r} has " |
| 1284 | "an incompatible weight shape: expected " |
| 1285 | f"(1..{wrapper.max_lora_cls_labels}, {wrapper.input_size}), " |
| 1286 | f"received {received_weight_shape}." |
| 1287 | ) |
| 1288 | |
| 1289 | if full_module.bias is None: |
| 1290 | return |
| 1291 | expected_bias_shape = (full_module.weight.size(0),) |
| 1292 | received_bias_shape = tuple(full_module.bias.shape) |
| 1293 | if received_bias_shape != expected_bias_shape: |
| 1294 | raise ValueError( |
| 1295 | f"Full module {saved_module_name!r} for {module_name!r} has " |
| 1296 | "an incompatible bias shape: expected " |
| 1297 | f"{expected_bias_shape}, received {received_bias_shape}." |
| 1298 | ) |
| 1299 | |
| 1300 | def _validate_token_classification_lora(self, lora_model: LoRAModel) -> None: |
| 1301 | if self._pooling_task != "token_classify": |
| 1302 | return |
| 1303 | if self._classification_head is None: |
| 1304 | return |
| 1305 | |
| 1306 | module_name, _ = self._classification_head |
| 1307 | has_full_module = ( |
| 1308 | self._get_module_to_save_weights(lora_model, module_name) is not None |
| 1309 | ) |
| 1310 | has_lora_weights = ( |
| 1311 | self._get_lora_layer_weights(lora_model, module_name) is not None |
| 1312 | ) |
| 1313 | if not has_full_module and not has_lora_weights: |
| 1314 | return |
| 1315 | |
| 1316 | raise ValueError( |
| 1317 | f"LoRA adapter {lora_model.id} contains weights for classification " |
| 1318 | f"head {module_name!r}, but token_classify only supports LoRA on " |
| 1319 | "the model backbone." |
| 1320 | ) |
| 1321 | |
| 1322 | def deactivate_adapter(self, adapter_id: int) -> bool: |
| 1323 | if adapter_id not in self._active_adapters: |
| 1324 | return False |
| 1325 | self._deactivate_adapter(adapter_id) |
| 1326 | self._active_adapters.pop(adapter_id, None) |
| 1327 | return True |
| 1328 | |
| 1329 | def add_adapter(self, adapter: LoRAModel) -> bool: |
| 1330 | logger.debug("Adding lora. Model id: %d, int id: %d", adapter.id, adapter.id) |
| 1331 | if adapter.id in self._registered_adapters: |
| 1332 | return False |
| 1333 | if len(self._registered_adapters) >= self.capacity: |
| 1334 | raise RuntimeError("No free adapter slots.") |
| 1335 | self._add_adapter(adapter) |
| 1336 | return True |
| 1337 | |
| 1338 | def set_adapter_mapping(self, mapping: LoRAMapping) -> None: |
| 1339 | # The punica metadata derives from the slot layout as well as the |
| 1340 | # mapping: an out-of-band add_lora() can LRU-evict and reassign slots |
| 1341 | # while the running batch, and thus the mapping, is unchanged. |
| 1342 | slot_layout = tuple(self.lora_index_to_id) |
| 1343 | if self._last_mapping != mapping or self._last_slot_layout != slot_layout: |
| 1344 | self._set_adapter_mapping(mapping) |
| 1345 | self._last_mapping = mapping |
| 1346 | self._last_slot_layout = slot_layout |
| 1347 | |
| 1348 | def remove_adapter(self, adapter_id: int) -> bool: |
| 1349 | self.deactivate_adapter(adapter_id) |
| 1350 | if adapter_id not in self._registered_adapters: |
| 1351 | return False |
| 1352 | self._registered_adapters.pop(adapter_id, None) |
| 1353 | return True |
| 1354 | |
| 1355 | def list_adapters(self) -> dict[int, LoRAModel]: |
| 1356 | return dict(self._registered_adapters.cache) |
| 1357 | |
| 1358 | def get_adapter(self, adapter_id: int) -> LoRAModel | None: |
| 1359 | return self._registered_adapters.get(adapter_id) |
| 1360 | |
| 1361 | |
| 1362 | class LRUCacheLoRAModelManager(LoRAModelManager): |
| 1363 | """A model manager that manages multiple LoRAs with LRU cache.""" |
| 1364 | |
| 1365 | def add_adapter(self, lora: LoRAModel) -> bool: |
| 1366 | """Add a LoRAModel to the manager.""" |
| 1367 | logger.debug("Adding lora. Model id: %d, int id: %d", lora.id, lora.id) |
| 1368 | if lora.id not in self._registered_adapters: |
| 1369 | self._add_adapter(lora) |
| 1370 | was_added = True |
| 1371 | else: |
| 1372 | # We always touch to update the LRU cache order |
| 1373 | self._registered_adapters.touch(lora.id) |
| 1374 | was_added = False |
| 1375 | return was_added |
| 1376 | |
| 1377 | def activate_adapter( |
| 1378 | self, |
| 1379 | lora_id: int, |
| 1380 | ) -> bool: |
| 1381 | if ( |
| 1382 | lora_id not in self._active_adapters |
| 1383 | and len(self._active_adapters) >= self.lora_slots |
| 1384 | ): |
| 1385 | self._active_adapters.remove_oldest() |
| 1386 | result = super().activate_adapter(lora_id) |
| 1387 | # We always touch to update the LRU cache order |
| 1388 | self._active_adapters.touch(lora_id) |
| 1389 | return result |
| 1390 | |
| 1391 | def remove_oldest_adapter(self) -> bool: |
| 1392 | if len(self._registered_adapters) > 0: |
| 1393 | self._registered_adapters.remove_oldest() |
| 1394 | return True |
| 1395 | return False |
| 1396 | |
| 1397 | def pin_adapter(self, lora_id: int) -> bool: |
| 1398 | """Pin a LoRAModel in the manager cache.""" |
| 1399 | self._pin_lora_in_cpu_cache(lora_id) |
| 1400 | self._pin_lora_in_gpu_cache(lora_id) |
| 1401 | return True |
| 1402 | |
| 1403 | def _pin_lora_in_cpu_cache(self, lora_id: int): |
| 1404 | try: |
| 1405 | self._registered_adapters.pin(lora_id) |
| 1406 | except ValueError as err: |
| 1407 | raise ValueError( |
| 1408 | f"Pinning failed. LoRA {lora_id} is not registered." |
| 1409 | ) from err |
| 1410 | |
| 1411 | def _pin_lora_in_gpu_cache(self, lora_id: int): |
| 1412 | if lora_id not in self._active_adapters: |
| 1413 | # move lora to gpu if not already active |
| 1414 | self.activate_adapter(lora_id) |
| 1415 | |
| 1416 | self._active_adapters.pin(lora_id) |
| 1417 | |
| 1418 | |
| 1419 | def create_lora_manager( |
| 1420 | model: SupportsLoRAModel, |
| 1421 | max_num_seqs: int, |
| 1422 | max_num_batched_tokens: int, |
| 1423 | vocab_size: int, |
| 1424 | lora_config: LoRAConfig, |
| 1425 | vllm_config: VllmConfig, |
| 1426 | device: torch.device, |
| 1427 | lora_manager_cls: type[LoRAModelManager] = LoRAModelManager, |
| 1428 | **kwargs, |
| 1429 | ) -> LoRAModelManager: |
| 1430 | """Create a LoRA adapter for a given model.""" |
| 1431 | if not isinstance(model, SupportsLoRA): |
| 1432 | raise ValueError(f"Model {type(model)} is not supported for LoRA.") |
| 1433 | lora_manager = lora_manager_cls( |
| 1434 | model=model, |
| 1435 | max_num_seqs=max_num_seqs, |
| 1436 | max_num_batched_tokens=max_num_batched_tokens, |
| 1437 | vocab_size=vocab_size, |
| 1438 | lora_config=lora_config, |
| 1439 | vllm_config=vllm_config, |
| 1440 | device=device, |
| 1441 | **kwargs, |
| 1442 | ) |
| 1443 | return lora_manager |