2525first, ``from_lora_tensors``'s per-key ``.to()`` returns the same tensor, so
2626aliases created by :func:`expand_lora_aliases` keep sharing storage. Loading on
2727CPU would pin a private copy per key and re-inflate a deduplicated adapter to
28- its public size.
28+ its public size, and vLLM's MoE packing (``PackedLoRALayerWeights.pack_moe``)
29+ is much slower on the host.
2930
3031``add_lora`` still flows API server -> engine core -> every worker, so the
3132adapter name in the OpenAI-compatible model registry stays correct and
3233``load_inplace=True`` replaces the previous generation exactly as the
33- directory path does. Staged tensors stay registered after the load: the
34- ``LoRAModel`` shares their storage (same device and dtype), so this costs
35- nothing, and it lets vLLM rebuild the adapter after an LRU eviction the way it
36- would re-read a directory. The next sync for the same name replaces them;
34+ directory path does. Staged tensors stay registered after the load, so vLLM
35+ can rebuild the adapter after an LRU eviction the way it would re-read a
36+ directory. The next sync for the same name replaces them;
3737``discard_in_memory_adapter`` frees them on unload.
3838
39+ ``WorkerLoRAManager.add_adapter`` (and the LRU variant's) is wrapped too. Once
40+ it has activated a GPU-built adapter, i.e. copied this worker's slice into
41+ vLLM's pre-allocated LoRA slots, :func:`offload_in_memory_adapter` moves the
42+ registered ``LoRAModel`` and the staged tensors to pinned host memory, where
43+ the directory path keeps the ``LoRAModel`` as well. Generation reads only the
44+ slots. The ``LoRAModel`` holds ``pack_moe``'s stacked copy of every local
45+ expert's factors, allocated outside vLLM's sleep-mode pool. vLLM re-activates an
46+ evicted adapter by copying from these host tensors (``set_lora``), and a
47+ rebuild from host-staged tensors builds (and runs ``pack_moe``) on the host,
48+ as the directory path does.
49+
3950Sharing storage has one consequence: vLLM applies ``lora_alpha / r`` by
4051multiplying ``lora_b`` *in place* on every load (``LoRALayerWeights.optimize``),
4152so a non-unit scale would compound on each LRU rebuild. Staged adapters must
5970
6071_PATCHED = False
6172_ORIGINAL_LOAD_ADAPTER = None
73+ _ORIGINAL_ADD_ADAPTER : Dict [type , Any ] = {}
74+
75+ # ``lora_int_id``s whose ``LoRAModel`` was built on the GPU and not yet moved
76+ # to the host. ``LRUCacheWorkerLoRAManager.add_adapter`` runs on every
77+ # scheduler step, so the wrapper only offloads ids listed here.
78+ _AWAITING_OFFLOAD : set [int ] = set ()
6279
6380
6481@dataclass
@@ -177,17 +194,23 @@ def _load_in_memory_adapter(manager, lora_request, staged: StagedLoRAAdapter):
177194 if moe_ep_spec is not None :
178195 tensors = {k : v for k , v in tensors .items () if not _is_remote_expert_key (k , moe_ep_spec )}
179196
197+ # Build where the staged tensors live: on the GPU right after a sync, on
198+ # the host once offload_in_memory_adapter has moved them (a rebuild after
199+ # an LRU eviction then works like a load from a directory).
200+ build_on_host = next (iter (staged .tensors .values ())).device .type == "cpu"
180201 lora = manager ._lora_model_cls .from_lora_tensors (
181202 lora_model_id = lora_request .lora_int_id ,
182203 tensors = tensors ,
183204 peft_helper = peft_helper ,
184- device = manager .device ,
205+ device = "cpu" if build_on_host else manager .device ,
185206 dtype = manager .lora_config .lora_dtype ,
186207 model_vocab_size = manager .vocab_size ,
187208 weights_mapper = weights_mapper ,
188209 skip_prefixes = skip_prefixes ,
189210 )
190211 lora .is_3d_lora_weight = lora_request .is_3d_lora_weight
212+ if not build_on_host :
213+ _AWAITING_OFFLOAD .add (lora_request .lora_int_id )
191214 return lora
192215
193216
@@ -205,8 +228,72 @@ def _patched_load_adapter(self, lora_request):
205228 return _load_in_memory_adapter (self , lora_request , staged )
206229
207230
231+ def _to_host (tensor , moved : Dict [tuple , torch .Tensor ]):
232+ """Returns a pinned host copy of a CUDA ``tensor``, else ``tensor`` unchanged.
233+
234+ ``moved`` maps each copied tensor's ``(data_ptr, dtype, shape, stride)`` to
235+ its host copy, so tensors that are the same view of one GPU storage share
236+ one host copy.
237+ """
238+ if not isinstance (tensor , torch .Tensor ) or not tensor .is_cuda :
239+ return tensor
240+ key = (tensor .data_ptr (), tensor .dtype , tuple (tensor .shape ), tensor .stride ())
241+ host = moved .get (key )
242+ if host is None :
243+ host = torch .empty (tensor .shape , dtype = tensor .dtype , device = "cpu" , pin_memory = True )
244+ host .copy_ (tensor , non_blocking = True )
245+ moved [key ] = host
246+ return host
247+
248+
249+ def offload_in_memory_adapter (manager , lora_request ) -> None :
250+ """Moves an activated in-memory adapter's tensors from the GPU to pinned host memory.
251+
252+ Covers the registered ``LoRAModel``'s ``lora_a`` / ``lora_b`` (single
253+ tensors or per-slice lists) and the staged tensors for
254+ ``lora_request.lora_name``. Tensors that shared GPU storage share the host
255+ copy. Releases the freed blocks with ``torch.cuda.empty_cache`` so that
256+ other processes on the GPU (a colocated trainer) can use them.
257+ """
258+ lora_model = manager ._adapter_manager .get_adapter (lora_request .lora_int_id )
259+ if lora_model is None :
260+ return
261+ moved : Dict [tuple , torch .Tensor ] = {}
262+ for weights in lora_model .loras .values ():
263+ for attr in ("lora_a" , "lora_b" ):
264+ value = getattr (weights , attr , None )
265+ if isinstance (value , list ):
266+ setattr (weights , attr , [_to_host (t , moved ) for t in value ])
267+ else :
268+ setattr (weights , attr , _to_host (value , moved ))
269+ staged = _STAGED .get (lora_request .lora_name )
270+ if staged is not None :
271+ staged .tensors = {k : _to_host (v , moved ) for k , v in staged .tensors .items ()}
272+ if moved :
273+ # Waits for the non_blocking copies before returning the GPU blocks to the driver.
274+ torch .cuda .current_stream ().synchronize ()
275+ moved .clear ()
276+ torch .cuda .empty_cache ()
277+
278+
279+ def _make_patched_add_adapter (original ):
280+ def _patched_add_adapter (self , lora_request ):
281+ loaded = original (self , lora_request )
282+ # An id stays listed if ``original`` raised; the next successful call offloads it.
283+ if lora_request .lora_int_id in _AWAITING_OFFLOAD :
284+ from vllm .utils .gpu_sync_debug import gpu_sync_allowed
285+
286+ _AWAITING_OFFLOAD .discard (lora_request .lora_int_id )
287+ # vLLM runs add_adapter's own device work under the same context.
288+ with gpu_sync_allowed ():
289+ offload_in_memory_adapter (self , lora_request )
290+ return loaded
291+
292+ return _patched_add_adapter
293+
294+
208295def apply_lora_in_memory_patch () -> None :
209- """Install the wrapper on ``WorkerLoRAManager`` (idempotent, worker-process local)."""
296+ """Install the wrappers on vLLM's worker LoRA managers (idempotent, worker-process local)."""
210297 global _PATCHED , _ORIGINAL_LOAD_ADAPTER
211298 if _PATCHED :
212299 return
@@ -215,10 +302,19 @@ def apply_lora_in_memory_patch() -> None:
215302 except ModuleNotFoundError :
216303 # Trainer processes import the worker-extension module too.
217304 return
305+ from vllm .lora .worker_manager import LRUCacheWorkerLoRAManager
306+
218307 _ORIGINAL_LOAD_ADAPTER = WorkerLoRAManager ._load_adapter
219308 WorkerLoRAManager ._load_adapter = _patched_load_adapter
309+ # LRUCacheWorkerLoRAManager overrides add_adapter without calling super().
310+ for cls in (WorkerLoRAManager , LRUCacheWorkerLoRAManager ):
311+ _ORIGINAL_ADD_ADAPTER [cls ] = cls .__dict__ ["add_adapter" ]
312+ cls .add_adapter = _make_patched_add_adapter (_ORIGINAL_ADD_ADAPTER [cls ])
220313 _PATCHED = True
221- logger .debug ("in-memory LoRA patch installed on vllm WorkerLoRAManager._load_adapter" )
314+ logger .debug (
315+ "in-memory LoRA patch installed on vllm WorkerLoRAManager._load_adapter and on "
316+ "add_adapter of WorkerLoRAManager and LRUCacheWorkerLoRAManager"
317+ )
222318
223319
224320def is_patched () -> bool :
0 commit comments