| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from contextlib import contextmanager |
| 5 | from typing import Any, Literal |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.config.lora import LoRAConfig |
| 11 | from vllm.exceptions import LoRAAdapterNotFoundError |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.lora.lora_model import LoRAModel |
| 14 | from vllm.lora.model_manager import ( |
| 15 | LoRAModelManager, |
| 16 | LRUCacheLoRAModelManager, |
| 17 | create_lora_manager, |
| 18 | ) |
| 19 | from vllm.lora.peft_helper import PEFTHelper |
| 20 | from vllm.lora.request import LoRARequest |
| 21 | from vllm.lora.utils import get_adapter_absolute_path |
| 22 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 23 | |
| 24 | logger = init_logger(__name__) |
| 25 | |
| 26 | |
| 27 | class WorkerLoRAManager: |
| 28 | """WorkerLoRAManager that manages LoRA models on the worker side. |
| 29 | |
| 30 | Every request, the requested LoRAs will be loaded (unless they are already |
| 31 | loaded), and every other LoRA will be unloaded.""" |
| 32 | |
| 33 | _manager_cls: type[LoRAModelManager] = LoRAModelManager |
| 34 | |
| 35 | def __init__( |
| 36 | self, |
| 37 | vllm_config: VllmConfig, |
| 38 | device: torch.device, |
| 39 | embedding_modules: dict[str, str], |
| 40 | lora_model_cls: type[LoRAModel] = LoRAModel, |
| 41 | ): |
| 42 | self._lora_model_cls = lora_model_cls |
| 43 | self.embedding_modules = embedding_modules |
| 44 | self._cached_dummy_lora: None | Literal[False] | LoRAModel = False |
| 45 | self.max_num_seqs = vllm_config.scheduler_config.max_num_seqs |
| 46 | self.max_num_batched_tokens = ( |
| 47 | vllm_config.scheduler_config.max_num_batched_tokens |
| 48 | ) |
| 49 | self.vocab_size = vllm_config.model_config.get_vocab_size() |
| 50 | lora_config = vllm_config.lora_config |
| 51 | if lora_config is None: |
| 52 | raise ValueError("LoRA config must be set for WorkerLoRAManager.") |
| 53 | self.lora_config: LoRAConfig = lora_config |
| 54 | |
| 55 | # Use get_text_config() in case of multimodal models |
| 56 | text_config = vllm_config.model_config.hf_config.get_text_config() |
| 57 | |
| 58 | # For encoder-decoder models (e.g., Whisper), use max_target_positions |
| 59 | # instead of max_position_embeddings |
| 60 | # TODO: Generalize max_position_embeddings handling for |
| 61 | # out-of-tree (OOT) encoder-decoder models |
| 62 | if vllm_config.model_config.is_encoder_decoder: |
| 63 | self.max_position_embeddings = getattr( |
| 64 | text_config, "max_target_positions", None |
| 65 | ) |
| 66 | else: |
| 67 | self.max_position_embeddings = getattr( |
| 68 | text_config, "max_position_embeddings", None |
| 69 | ) |
| 70 | self.device = device |
| 71 | # Lazily initialized by create_lora_manager. |
| 72 | self._adapter_manager: LoRAModelManager |
| 73 | |
| 74 | @contextmanager |
| 75 | def dummy_lora_cache(self): |
| 76 | """Use this context manager to reuse the dummy lora model |
| 77 | to avoid creating it repeatedly.""" |
| 78 | self._cached_dummy_lora = None |
| 79 | yield |
| 80 | self._cached_dummy_lora = False |
| 81 | |
| 82 | @property |
| 83 | def is_enabled(self) -> bool: |
| 84 | return True |
| 85 | |
| 86 | def create_lora_manager( |
| 87 | self, |
| 88 | model: torch.nn.Module, |
| 89 | vllm_config: VllmConfig, |
| 90 | ) -> Any: |
| 91 | if vllm_config is None: |
| 92 | raise ValueError("vllm_config must be provided to create a LoRA manager.") |
| 93 | lora_manager = create_lora_manager( |
| 94 | model, |
| 95 | max_num_seqs=self.max_num_seqs, |
| 96 | max_num_batched_tokens=self.max_num_batched_tokens, |
| 97 | vocab_size=self.vocab_size, |
| 98 | lora_config=self.lora_config, |
| 99 | device=self.device, |
| 100 | lora_manager_cls=self._manager_cls, |
| 101 | vllm_config=vllm_config, |
| 102 | ) |
| 103 | self._adapter_manager = lora_manager |
| 104 | return lora_manager.model |
| 105 | |
| 106 | def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel: |
| 107 | try: |
| 108 | supported_lora_modules = self._adapter_manager.supported_lora_modules |
| 109 | packed_modules_mapping = self._adapter_manager.packed_modules_mapping |
| 110 | expected_lora_lst: list[str] = [] |
| 111 | for module in supported_lora_modules: |
| 112 | if module in packed_modules_mapping: |
| 113 | expected_lora_lst.extend(packed_modules_mapping[module]) |
| 114 | else: |
| 115 | expected_lora_lst.append(module) |
| 116 | if module == "experts": |
| 117 | expected_lora_lst.append(module) |
| 118 | expected_lora_modules = set(expected_lora_lst) |
| 119 | lora_path = get_adapter_absolute_path(lora_request.lora_path) |
| 120 | |
| 121 | peft_helper = PEFTHelper.from_local_dir( |
| 122 | lora_path, |
| 123 | self.max_position_embeddings, |
| 124 | lora_request.tensorizer_config_dict, |
| 125 | ) |
| 126 | |
| 127 | # Validates the LoRA configuration against requirements before |
| 128 | # loading weights, throwing an exception if validation fails. |
| 129 | peft_helper.validate_legal(self.lora_config) |
| 130 | |
| 131 | # For some models like Qwen2VL, we need to use hf_to_vllm_mapper to ensure |
| 132 | # correct loading of lora weights. We only need to know about renames for |
| 133 | # this, so we use get_rename_mapper() to ignore stacking and deletions. |
| 134 | model = self._adapter_manager.model |
| 135 | hf_to_vllm_mapper = getattr(model, "hf_to_vllm_mapper", None) |
| 136 | if hf_to_vllm_mapper is not None: |
| 137 | hf_to_vllm_mapper = hf_to_vllm_mapper.get_rename_mapper() |
| 138 | |
| 139 | # Get model-defined prefixes to skip during LoRA loading. |
| 140 | lora_skip_prefixes = getattr(model, "lora_skip_prefixes", None) |
| 141 | |
| 142 | lora = self._lora_model_cls.from_local_checkpoint( |
| 143 | lora_path, |
| 144 | expected_lora_modules, |
| 145 | peft_helper=peft_helper, |
| 146 | lora_model_id=lora_request.lora_int_id, |
| 147 | device="cpu", |
| 148 | dtype=self.lora_config.lora_dtype, |
| 149 | model_vocab_size=self.vocab_size, |
| 150 | tensorizer_config_dict=lora_request.tensorizer_config_dict, |
| 151 | weights_mapper=hf_to_vllm_mapper, |
| 152 | skip_prefixes=lora_skip_prefixes, |
| 153 | moe_ep_spec=self._adapter_manager.moe_ep_load_spec, |
| 154 | ) |
| 155 | # Stamp the on-disk MoE layout onto the loaded model so the |
| 156 | # adapter manager can route 3D-format checkpoints through the |
| 157 | # 3D->2D conversion when running under the universal 2D wrapper. |
| 158 | lora.is_3d_lora_weight = lora_request.is_3d_lora_weight |
| 159 | |
| 160 | # Validate classification-head weights. |
| 161 | self._adapter_manager._validate_modules_to_save(lora) |
| 162 | self._adapter_manager._validate_token_classification_lora(lora) |
| 163 | |
| 164 | except FileNotFoundError as e: |
| 165 | # FileNotFoundError should be raised if both |
| 166 | # - No adapter found to download from huggingface (or in |
| 167 | # offline mode) |
| 168 | # - No local adapter files found at `lora_request.lora_path` |
| 169 | # For NotFoundError |
| 170 | raise LoRAAdapterNotFoundError( |
| 171 | lora_request.lora_name, lora_request.lora_path |
| 172 | ) from e |
| 173 | except Exception as e: |
| 174 | raise e |
| 175 | |
| 176 | return lora |
| 177 | |
| 178 | def add_dummy_lora(self, lora_request: LoRARequest, rank: int) -> bool: |
| 179 | if lora_request.lora_int_id in self.list_adapters(): |
| 180 | return False |
| 181 | if isinstance(self._cached_dummy_lora, LoRAModel): |
| 182 | dummy_lora = self._cached_dummy_lora.clone(lora_request.lora_int_id) |
| 183 | else: |
| 184 | dummy_lora = self._adapter_manager.create_dummy_lora( |
| 185 | lora_request.lora_int_id, rank, self.embedding_modules |
| 186 | ) |
| 187 | if self._cached_dummy_lora is None: |
| 188 | self._cached_dummy_lora = dummy_lora |
| 189 | return self._adapter_manager.add_adapter(dummy_lora) |
| 190 | |
| 191 | def get_dummy_lora_warmup_rank(self, default_rank: int) -> int: |
| 192 | return self._adapter_manager.get_dummy_lora_warmup_rank(default_rank) |
| 193 | |
| 194 | def pin_adapter(self, adapter_id: int) -> bool: |
| 195 | return self._adapter_manager.pin_adapter(adapter_id) |
| 196 | |
| 197 | def set_active_adapters(self, requests: set[Any], mapping: Any | None) -> None: |
| 198 | self._apply_adapters(requests) |
| 199 | if mapping is not None: |
| 200 | self._adapter_manager.set_adapter_mapping(mapping) |
| 201 | |
| 202 | def supports_tower_connector_lora(self) -> bool: |
| 203 | return ( |
| 204 | self._adapter_manager.supports_mm |
| 205 | and self._adapter_manager.supports_tower_connector_lora |
| 206 | ) |
| 207 | |
| 208 | def _apply_adapters(self, adapter_requests: set[Any]) -> None: |
| 209 | existing_adapters = self.list_adapters() |
| 210 | models_map = { |
| 211 | adapter_request.adapter_id: adapter_request |
| 212 | for adapter_request in adapter_requests |
| 213 | if adapter_request |
| 214 | } |
| 215 | if len(models_map) > self._adapter_manager.adapter_slots: |
| 216 | raise RuntimeError( |
| 217 | f"Number of requested models ({len(models_map)}) is greater " |
| 218 | "than the number of GPU model slots " |
| 219 | f"({self._adapter_manager.adapter_slots})." |
| 220 | ) |
| 221 | requested_ids = set(models_map) |
| 222 | for adapter_id in existing_adapters - requested_ids: |
| 223 | self.remove_adapter(adapter_id) |
| 224 | for adapter_id in requested_ids - existing_adapters: |
| 225 | self.add_adapter(models_map[adapter_id]) |
| 226 | |
| 227 | def add_adapter(self, adapter_request: Any) -> bool: |
| 228 | if adapter_request.adapter_id in self.list_adapters(): |
| 229 | return False |
| 230 | # One-time per adapter |
| 231 | with gpu_sync_allowed(): |
| 232 | loaded_adapter = self._load_adapter(adapter_request) |
| 233 | loaded = self._adapter_manager.add_adapter(loaded_adapter) |
| 234 | self._adapter_manager.activate_adapter(loaded_adapter.id) |
| 235 | return loaded |
| 236 | |
| 237 | def remove_adapter(self, adapter_id: int) -> bool: |
| 238 | return self._adapter_manager.remove_adapter(adapter_id) |
| 239 | |
| 240 | def remove_all_adapters(self): |
| 241 | self._adapter_manager.remove_all_adapters() |
| 242 | |
| 243 | def list_adapters(self) -> set[int]: |
| 244 | return set(self._adapter_manager.list_adapters()) |
| 245 | |
| 246 | |
| 247 | class LRUCacheWorkerLoRAManager(WorkerLoRAManager): |
| 248 | """WorkerLoRAManager that manages LoRA models on the worker side. |
| 249 | |
| 250 | Uses an LRU Cache. Every request, the requested LoRAs will be loaded |
| 251 | (unless they are already loaded) and least recently used LoRAs will |
| 252 | be unloaded if the cache is above capacity.""" |
| 253 | |
| 254 | _manager_cls: type[LRUCacheLoRAModelManager] = LRUCacheLoRAModelManager |
| 255 | |
| 256 | def create_lora_manager( |
| 257 | self, |
| 258 | model: torch.nn.Module, |
| 259 | vllm_config: VllmConfig, |
| 260 | ) -> Any: |
| 261 | if vllm_config is None: |
| 262 | raise ValueError("vllm_config must be provided to create a LoRA manager.") |
| 263 | lora_manager = create_lora_manager( |
| 264 | model, |
| 265 | lora_manager_cls=self._manager_cls, |
| 266 | max_num_seqs=self.max_num_seqs, |
| 267 | vocab_size=self.vocab_size, |
| 268 | lora_config=self.lora_config, |
| 269 | device=self.device, |
| 270 | max_num_batched_tokens=self.max_num_batched_tokens, |
| 271 | vllm_config=vllm_config, |
| 272 | ) |
| 273 | self._adapter_manager = lora_manager |
| 274 | return lora_manager.model |
| 275 | |
| 276 | def _apply_adapters(self, lora_requests: set[LoRARequest]) -> None: |
| 277 | loras_map = { |
| 278 | lora_request.lora_int_id: lora_request |
| 279 | for lora_request in lora_requests |
| 280 | if lora_request |
| 281 | } |
| 282 | if len(loras_map) > self._adapter_manager.lora_slots: |
| 283 | raise RuntimeError( |
| 284 | f"Number of requested LoRAs ({len(loras_map)}) is greater " |
| 285 | "than the number of GPU LoRA slots " |
| 286 | f"({self._adapter_manager.lora_slots})." |
| 287 | ) |
| 288 | for lora in loras_map.values(): |
| 289 | self.add_adapter(lora) |
| 290 | |
| 291 | def add_adapter(self, lora_request: LoRARequest) -> bool: |
| 292 | # Note that this method is not thread-safe. It may be invoked multiple |
| 293 | # times for the same adapter when using multiple API servers. |
| 294 | # This is ok because it's currently only called from |
| 295 | # the single-threaded core engine loop. |
| 296 | |
| 297 | with gpu_sync_allowed(): |
| 298 | if ( |
| 299 | lora_request.lora_int_id not in self.list_adapters() |
| 300 | or lora_request.load_inplace |
| 301 | ): |
| 302 | # Load the new adapter first to ensure it is actually valid, before |
| 303 | # evicting any existing adapters. |
| 304 | # This may cause the # of loaded lora adapters to very temporarily |
| 305 | # exceed `--max-cpu-loras`. |
| 306 | lora = self._load_adapter(lora_request) |
| 307 | |
| 308 | # Remove the existing adapter if it exists |
| 309 | # Use case for LoRA inplace |
| 310 | self._adapter_manager.remove_adapter(lora.id) |
| 311 | |
| 312 | # Loading succeeded, now check if we will exceed cache capacity and |
| 313 | # evict if the oldest adapter if so |
| 314 | if len(self._adapter_manager) + 1 > self._adapter_manager.capacity: |
| 315 | assert isinstance(self._adapter_manager, LRUCacheLoRAModelManager) |
| 316 | self._adapter_manager.remove_oldest_adapter() |
| 317 | # Then add the new adapter to the cache |
| 318 | loaded = self._adapter_manager.add_adapter(lora) |
| 319 | else: |
| 320 | # If the lora is already loaded, just touch it to |
| 321 | # update its position in the caches |
| 322 | loaded = ( |
| 323 | self._adapter_manager.get_adapter(lora_request.lora_int_id) |
| 324 | is not None |
| 325 | ) |
| 326 | self._adapter_manager.activate_adapter(lora_request.lora_int_id) |
| 327 | return loaded |