@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 300 lines 11.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Define LoRA functionality mixin for model runners."""
4
5 from collections.abc import Callable
6 from contextlib import contextmanager
7
8 import numpy as np
9 import torch
10 import torch.nn as nn
11
12 from vllm.config import VllmConfig
13 from vllm.config.lora import LoRAConfig
14 from vllm.logger import init_logger
15 from vllm.lora.layers import LoRAMapping, LoRAMappingType
16 from vllm.lora.request import LoRARequest
17 from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
18 from vllm.model_executor.models import supports_lora
19 from vllm.v1.worker.gpu_input_batch import InputBatch
20
21 logger = init_logger(__name__)
22
23
24 # Defined as a mixin for GPUModelRunner
25 class LoRAModelRunnerMixin:
26 lora_config: LoRAConfig | None
27 get_model: Callable[[], nn.Module]
28
29 def reset_lora_state(self) -> None:
30 """Invalidate LoRA state after base weights are replaced."""
31 if not self.lora_config:
32 return
33
34 from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA
35
36 self.lora_manager.remove_all_adapters()
37 for module in self.get_model().modules():
38 if isinstance(module, LogitsProcessorWithLoRA):
39 module.reset_sharded_to_full_mapping()
40
41 def load_lora_model(
42 self,
43 model: nn.Module,
44 vllm_config: VllmConfig,
45 device: torch.device,
46 ) -> nn.Module:
47 if not supports_lora(model):
48 raise ValueError(f"{model.__class__.__name__} does not support LoRA yet.")
49
50 # Add LoRA Manager to the Model Runner
51 self.lora_manager = LRUCacheWorkerLoRAManager(
52 vllm_config,
53 device,
54 model.embedding_modules,
55 )
56 return self.lora_manager.create_lora_manager(model, vllm_config)
57
58 def _set_active_loras(
59 self,
60 prompt_lora_mapping: tuple[int, ...],
61 token_lora_mapping: tuple[int, ...],
62 lora_requests: set[LoRARequest],
63 mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,
64 ) -> None:
65 self._ensure_lora_enabled()
66
67 # Set is_prefill to True, so we always use the SGMV kernels on
68 # non-cuda platforms.
69 # On cuda platforms we use the same kernels for prefill and
70 # decode and this flag is generally ignored.
71 lora_mapping = LoRAMapping(
72 token_lora_mapping,
73 prompt_lora_mapping,
74 is_prefill=True,
75 type=mapping_type,
76 )
77 self.lora_manager.set_active_adapters(lora_requests, lora_mapping)
78
79 def _ensure_lora_enabled(self) -> None:
80 if not hasattr(self, "lora_manager"):
81 raise RuntimeError("LoRA is not enabled. Use --enable-lora to enable LoRA.")
82
83 def set_active_loras(
84 self,
85 input_batch: InputBatch,
86 num_scheduled_tokens: np.ndarray,
87 num_sampled_tokens: np.ndarray | None = None,
88 mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,
89 ) -> None:
90 if num_sampled_tokens is None:
91 num_sampled_tokens = np.ones_like(num_scheduled_tokens, dtype=np.int32)
92
93 prompt_lora_mapping: tuple[int, ...] # of size np.sum(num_sampled_tokens)
94 token_lora_mapping: tuple[int, ...] # of size np.sum(num_scheduled_tokens)
95 lora_requests: set[LoRARequest]
96 prompt_lora_mapping, token_lora_mapping, lora_requests = (
97 input_batch.make_lora_inputs(num_scheduled_tokens, num_sampled_tokens)
98 )
99 return self._set_active_loras(
100 prompt_lora_mapping, token_lora_mapping, lora_requests, mapping_type
101 )
102
103 @contextmanager
104 def maybe_setup_dummy_loras(
105 self, lora_config: LoRAConfig | None, remove_lora: bool = True
106 ):
107 if lora_config is None:
108 yield
109 else:
110 # __enter__ code
111 assert self.lora_manager is not None, "LoRA is not enabled"
112
113 num_loras = lora_config.max_loras
114 lora_warmup_rank: int = (
115 lora_config.max_lora_rank if lora_config.max_lora_rank < 8 else 8
116 )
117 lora_warmup_rank = self.lora_manager.get_dummy_lora_warmup_rank(
118 lora_warmup_rank
119 )
120 # Make dummy lora requests
121 lora_requests: set[LoRARequest] = {
122 LoRARequest(
123 lora_name=f"warmup_{lora_id}",
124 lora_int_id=lora_id,
125 lora_path="/not/a/real/path",
126 )
127 for lora_id in range(1, num_loras + 1)
128 }
129
130 with self.lora_manager.dummy_lora_cache():
131 # Add the dummy LoRAs here so _set_active_loras doesn't try to
132 # load from disk.
133 for lr in lora_requests:
134 self.lora_manager.add_dummy_lora(lr, rank=lora_warmup_rank)
135
136 yield
137
138 # __exit__ code
139 if remove_lora:
140 self.lora_manager.remove_all_adapters()
141
142 @contextmanager
143 def maybe_select_dummy_loras(
144 self,
145 lora_config: LoRAConfig | None,
146 num_scheduled_tokens: np.ndarray,
147 mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,
148 num_sampled_tokens: np.ndarray | None = None,
149 num_active_loras: int = 0,
150 ):
151 """Context manager to select dummy LoRAs for capture/warmup.
152
153 Args:
154 lora_config: LoRA configuration, or None if LoRA is disabled.
155 num_scheduled_tokens: Array of scheduled token counts per request.
156 mapping_type: Which LoRA mapping to build (language or encoder).
157 num_sampled_tokens: Array of sampled token counts per request.
158 num_active_loras: Number of distinct active LoRAs to use.
159 - 0: No LoRA active (set up zero mappings).
160 - >0: Use exactly this many distinct LoRAs.
161
162 """
163 if num_sampled_tokens is None:
164 num_sampled_tokens = np.ones_like(num_scheduled_tokens, dtype=np.int32)
165
166 # Skip LoRA setup entirely only if no LoRA config
167 if lora_config is None:
168 yield
169 else:
170 # __enter__ code
171 assert self.lora_manager is not None, "LoRA is not enabled"
172
173 num_reqs = len(num_scheduled_tokens)
174 max_loras = lora_config.max_loras
175
176 # Determine how many distinct LoRAs to use and whether to include
177 # no-LoRA tokens (-1 entries).
178 # When num_active_loras > max_loras (e.g., max_loras + 1), we need
179 # to include -1 entries to simulate batches with both LoRA and
180 # no-LoRA tokens. This ensures prepare_tensors computes the correct
181 # num_active_loras that matches the cudagraph capture key.
182 if num_active_loras == 0:
183 # No LoRA active - use 0 mappings like the original code
184 effective_num_loras = 0
185 include_no_lora = False
186 elif num_active_loras > max_loras:
187 # num_active_loras > max_loras means we want max_loras adapters
188 # PLUS no-LoRA tokens (-1). This is the max_loras + 1 case.
189 effective_num_loras = max_loras
190 include_no_lora = True
191 else:
192 # Specific number of active LoRAs requested
193 effective_num_loras = min(num_active_loras, max_loras)
194 include_no_lora = False
195
196 # Make prompt lora mapping
197 # Assign LoRA IDs cyclically to simulate a worst-case scenario.
198 # LoRA IDs are 1-indexed (1 to max_loras) as required by LoRARequest.
199 # convert_mapping() will convert these to 0-indexed slot indices.
200 if effective_num_loras > 0:
201 if include_no_lora:
202 # Include -1 (no-LoRA) entries by cycling through
203 # -1, 1, 2, ..., effective_num_loras
204 # This ensures prepare_tensors sees both LoRA and no-LoRA
205 # tokens, computing num_active_loras = effective_num_loras+1
206 cycle_values = np.array(
207 list(range(1, effective_num_loras + 1)),
208 dtype=np.int32,
209 )
210 prompt_lora_mapping = cycle_values[
211 np.arange(num_reqs, dtype=np.int32) % len(cycle_values)
212 ]
213 else:
214 # Use 1 to effective_num_loras (1-indexed lora IDs)
215 prompt_lora_mapping = (
216 np.arange(num_reqs, dtype=np.int32) % effective_num_loras
217 ) + 1
218 else:
219 # No LoRA active - use 0 for all tokens (original behavior)
220 prompt_lora_mapping = np.zeros(num_reqs, dtype=np.int32)
221
222 # Make sample lora mapping
223 sample_lora_mapping = np.repeat(prompt_lora_mapping, num_sampled_tokens)
224
225 # Make token lora mapping
226 token_lora_mapping = np.repeat(prompt_lora_mapping, num_scheduled_tokens)
227
228 # Make dummy lora requests (only for the active LoRAs)
229 lora_requests: set[LoRARequest] = {
230 LoRARequest(
231 lora_name=f"warmup_{lora_id}",
232 lora_int_id=lora_id,
233 lora_path="/not/a/real/path",
234 )
235 for lora_id in range(1, effective_num_loras + 1)
236 }
237
238 self._set_active_loras(
239 tuple(sample_lora_mapping),
240 tuple(token_lora_mapping),
241 lora_requests,
242 mapping_type,
243 )
244
245 yield
246
247 @contextmanager
248 def maybe_dummy_run_with_lora(
249 self,
250 lora_config: LoRAConfig | None,
251 num_scheduled_tokens: np.ndarray,
252 num_sampled_tokens: np.ndarray,
253 remove_lora: bool = True,
254 num_active_loras: int = 0,
255 mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,
256 ):
257 """Context manager for dummy runs with LoRA.
258
259 Args:
260 lora_config: LoRA configuration.
261 num_scheduled_tokens: Array of scheduled token counts per request.
262 num_sampled_tokens: Array of sampled token counts per request.
263 remove_lora: Whether to remove LoRAs after the context exits.
264 num_active_loras: Number of distinct active LoRAs to use.
265 LoRA is activated when num_active_loras > 0.
266 mapping_type: Which LoRA mapping to build (language or encoder).
267
268 """
269 with (
270 self.maybe_setup_dummy_loras(lora_config, remove_lora),
271 self.maybe_select_dummy_loras(
272 lora_config,
273 num_scheduled_tokens,
274 mapping_type,
275 num_sampled_tokens,
276 num_active_loras,
277 ),
278 ):
279 yield
280
281 def maybe_remove_all_loras(self, lora_config: LoRAConfig | None):
282 if lora_config is None:
283 return
284 self.lora_manager.remove_all_adapters()
285
286 def add_lora(self, lora_request: LoRARequest) -> bool:
287 self._ensure_lora_enabled()
288 return self.lora_manager.add_adapter(lora_request)
289
290 def remove_lora(self, lora_id: int) -> bool:
291 self._ensure_lora_enabled()
292 return self.lora_manager.remove_adapter(lora_id)
293
294 def pin_lora(self, lora_id: int) -> bool:
295 self._ensure_lora_enabled()
296 return self.lora_manager.pin_adapter(lora_id)
297
298 def list_loras(self) -> set[int]:
299 self._ensure_lora_enabled()
300 return self.lora_manager.list_adapters()