Skip to content

Commit 0c30582

Browse files
SumanthRHclaude
andcommitted
[Perf] Free up temporary memory allocated for LoRA tensors
With `trainer.policy.model.lora.sync_mode=memory`, vLLM builds the received adapter on the GPU (`from_lora_tensors(device=GPU)`), packs MoE expert factors into stacked tensors (`PackedLoRALayerWeights.pack_moe`) and copies this worker's slice into its pre-allocated LoRA slots. The registered `LoRAModel` and SkyRL's staged tensors then stayed on the GPU, allocated outside vLLM's sleep-mode pool, through generation and training, although generation reads only the slots. vLLM's directory path keeps the same `LoRAModel` on pinned host memory. The in-memory LoRA patch now also wraps `add_adapter` of `WorkerLoRAManager` and `LRUCacheWorkerLoRAManager`. After a GPU-built adapter is activated, `offload_in_memory_adapter` moves the `LoRAModel`'s `lora_a` / `lora_b` and the staged tensors to pinned host memory (tensors that shared GPU storage share one host copy), synchronizes and calls `torch.cuda.empty_cache()`. Only ids built since the last activation are offloaded, so the per-step `add_adapter` call stays a set lookup. The offload runs under vLLM's `gpu_sync_allowed()`. vLLM re-activates an evicted adapter from the host tensors with `set_lora`, and a rebuild from the offloaded stage builds on the host, as the directory path does. Tests: CPU wiring tests in test_lora_in_memory_load.py, and GPU tests of the copy in gpu_ci/patches/vllm/test_lora_in_memory_offload.py. Signed-off-by: SumanthRH <sumanthrh99@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent a9fafae commit 0c30582

6 files changed

Lines changed: 405 additions & 25 deletions

File tree

‎.agents/docs/weight_sync.md‎

Lines changed: 23 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,8 @@ worker start: no layerwise reload (the base model is untouched)
221221
rank 0 POST /skyrl/v1/load_lora_adapter {lora_name, in_memory: true}
222222
API server add_lora(LoRARequest(lora_path="skyrl-memory://<name>", load_inplace=True))
223223
worker (patched) WorkerLoRAManager._load_adapter -> from_lora_tensors(device=GPU) from the stage
224+
add_adapter: pack + activate into the LoRA slots, then
225+
offload_in_memory_adapter LoRAModel + stage -> pinned host
224226
```
225227

226228
Arming lasts exactly one round — `finish_weight_update` disarms — so a missing or failed
@@ -250,12 +252,13 @@ Memory and lifecycle, per inference GPU:
250252
tensor, not the public size. Duplicate groups are per `(module, EP-rank expert group,
251253
lora_A|lora_B)`, verified with `torch.equal`, so a wrong group size costs bandwidth only.
252254
- **Staged tensors** (`patch_lora_in_memory._STAGED`, keyed by adapter name) are the unique
253-
tensors in the inference dtype. They are kept after the load because the `LoRAModel`
254-
built from them shares their storage (same device, same dtype), and because vLLM rebuilds
255+
tensors in the inference dtype. They are kept after the load because vLLM rebuilds
255256
an LRU-evicted adapter by calling `_load_adapter` again: with a directory it re-reads the
256-
files, here it re-reads the stage. A resync for the same name replaces the stage;
257-
`RemoteInferenceClient.unload_lora_adapter` also issues `skyrl_discard_in_memory_lora` so
258-
an unloaded tenant frees its tensors.
257+
files, here it re-reads the stage. After activation they live on pinned host memory (see
258+
below), so a rebuild builds on the host like the directory path, including `pack_moe`,
259+
which is several times slower there than on the GPU. A resync for the same
260+
name replaces the stage; `RemoteInferenceClient.unload_lora_adapter` also issues
261+
`skyrl_discard_in_memory_lora` so an unloaded tenant frees its tensors.
259262
- **The stage must never be written by vLLM.** vLLM applies `lora_alpha / r` by multiplying
260263
`lora_b` *in place* on every load (`LoRALayerWeights.optimize`), and since the `LoRAModel`
261264
shares the staged storage, a non-unit scale would compound on every LRU rebuild (dense and
@@ -265,11 +268,21 @@ Memory and lifecycle, per inference GPU:
265268
`lora_alpha == r`, so vLLM's scale is exactly 1 and `optimize` is a no-op; the in-memory
266269
loader rejects a config where they differ. The published adapter is vLLM-shaped, not the
267270
trainer's raw tensors: a consumer reading `lora_B` directly sees it pre-scaled.
268-
- **The registered `LoRAModel` lives on the GPU**, unlike the directory path where it lives
269-
on pinned CPU memory. vLLM packs per-expert LoRA into stacked tensors sized to the local
270-
experts, so a MoE adapter costs roughly its *local un-deduplicated* size per registered
271-
adapter (about 1/EP of the public size). With multi-tenant `max_loras` / `max_cpu_loras`
272-
above 1, budget that per resident adapter.
271+
- **The registered `LoRAModel` lives on pinned host memory after activation**, as on the
272+
directory path. It is built on the GPU (fast `pack_moe`, aliases kept shared), and the
273+
patched `add_adapter` calls `offload_in_memory_adapter` once vLLM has copied it into its
274+
pre-allocated LoRA slots: `lora_a` / `lora_b` of every module and the staged tensors move
275+
to pinned host memory (tensors that shared GPU storage share one host copy), then
276+
`torch.cuda.empty_cache()`. Generation reads only the slots; an evicted adapter is
277+
re-activated by vLLM's `set_lora` from the host tensors. The GPU therefore holds the
278+
built adapter only during the sync. vLLM packs per-expert LoRA into stacked tensors sized
279+
to the local experts, so a MoE `LoRAModel` costs roughly its *local un-deduplicated* size
280+
-- per MoE layer and local expert, `3 * r * (hidden + moe_intermediate)` elements in the
281+
LoRA dtype. That copy plus the staged tensors is the transient GPU peak of each sync,
282+
and the pinned host memory each worker holds per registered adapter (freed host blocks
283+
return to PyTorch's pinned-host cache, not to the OS). The LoRA slots themselves are allocated at startup and
284+
scale with `max_loras`; with multi-tenant `max_loras` above 1, budget the slots per
285+
resident adapter.
273286
Measured on GLM-5.3 (`glm_moe_dsa`, 256 routed experts, Megatron EP8 -> vLLM TP8, B300
274287
268 GiB): **~22.4 GiB per resident adapter per inference GPU** — vLLM's weight-load
275288
memory fell from 220.94 to 198.56 GiB going `max_loras` 2 -> 1. On a 1.4 TB base model

‎skyrl/backends/skyrl_train/patches/vllm/patch_lora_in_memory.py‎

Lines changed: 104 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,17 +25,28 @@
2525
first, ``from_lora_tensors``'s per-key ``.to()`` returns the same tensor, so
2626
aliases created by :func:`expand_lora_aliases` keep sharing storage. Loading on
2727
CPU 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
3132
adapter 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+
3950
Sharing storage has one consequence: vLLM applies ``lora_alpha / r`` by
4051
multiplying ``lora_b`` *in place* on every load (``LoRALayerWeights.optimize``),
4152
so a non-unit scale would compound on each LRU rebuild. Staged adapters must
@@ -59,6 +70,12 @@
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+
208295
def 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

224320
def is_patched() -> bool:

‎skyrl/train/config/config.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -113,8 +113,9 @@ class SkyRLLoraConfig(BaseConfig):
113113
``"memory"`` ships the adapter tensors through the configured
114114
``generator.inference_engine.weight_sync_backend`` transport (``nccl``: NCCL broadcast when
115115
non-colocated, CUDA IPC when colocated) and vLLM builds the adapter from the received GPU
116-
tensors; nothing is written. Shared expert adapters are sent once and aliased on the receiver.
117-
Megatron only."""
116+
tensors; nothing is written. Once the adapter is copied into vLLM's LoRA slots, the built
117+
adapter and the received tensors move to pinned host memory, where ``"disk"`` keeps the built
118+
adapter. Shared expert adapters are sent once and aliased on the receiver. Megatron only."""
118119
target_modules: str = "all-linear"
119120
"""Modules to apply LoRA to.
120121
``"all-linear"`` targets every linear layer for FSDP/PEFT, and is remapped to a fixed module

‎tests/backends/skyrl_train/gpu/gpu_ci/patches/vllm/__init__.py‎

Whitespace-only changes.

0 commit comments

Comments
 (0)