[Perf] Free up temporary memory allocated for LoRA tensors for in-memory weight sync

[SumanthRH]

What does this PR do?

Moves an in-memory LoRA adapter (trainer.policy.model.lora.sync_mode=memory) off the GPU once vLLM has copied it into its LoRA slots, so the built adapter and the received tensors stop occupying GPU memory through generation and training.

TLDR: after add_adapter activates a GPU-built in-memory adapter, the registered LoRAModel and SkyRL's staged tensors move to pinned host memory, which is where vLLM keeps the LoRAModel for an adapter loaded from a directory. The build itself stays on the GPU. No new config.

Why it happens

1. With sync_mode=memory, the receiver stages the adapter tensors on the GPU, and patch_lora_in_memory builds the LoRAModel there with from_lora_tensors(..., device=manager.device). Building on the GPU keeps deduplicated aliases sharing storage and keeps vLLM's MoE packing fast.
2. Inside add_adapter, vLLM packs every expert's factors into stacked tensors (PackedLoRALayerWeights.pack_moe, torch.stack). For MoE models this packed copy holds every local expert's factors, including one copy per expert of factors that were sent once and aliased. vLLM then copies this worker's slice into the LoRA slots it allocated at startup (activate_adapter -> set_lora).
3. vLLM keeps the registered LoRAModel as the source for re-activating the adapter after an eviction, and SkyRL keeps the staged tensors to rebuild it after an LRU eviction. Both stayed on the GPU. They are allocated outside vLLM's sleep-mode pool, so sleeping the engine does not free them, and they stay allocated through generation and through training in colocated setups. Generation only reads the slots.

The directory path (sync_mode=disk) does not have this cost: WorkerLoRAManager._load_adapter builds the LoRAModel on pinned CPU memory (from_local_checkpoint(..., device="cpu")) and only the slots are on the GPU.

Size of the packed copy per vLLM worker: per MoE layer and local expert, 3 * r * (hidden + moe_intermediate) elements in the LoRA dtype (gate and up: A [r, hidden], B [moe_intermediate, r]; down: A [r, moe_intermediate], B [hidden, r]). With vLLM tensor parallelism and no expert parallelism every expert is local and the LoRAModel is not TP-sliced. For Qwen3-30B-A3B (48 layers x 128 experts, hidden 2048, moe_intermediate 768) at rank 64 that is 6.2 GiB per GPU in bf16, plus the staged tensors. In a colocated run on 8 GPUs, this PR lowers the peak GPU memory of every phase by about 6.4 GiB (see Validation). The cost grows with the rank, the number of experts and the number of MoE layers.

More details on large-scale tests to follow

Solution

patch_lora_in_memory.py:

• apply_lora_in_memory_patch also wraps add_adapter on WorkerLoRAManager and on LRUCacheWorkerLoRAManager, which overrides it without calling super().
• _load_in_memory_adapter records the lora_int_id of every adapter it builds on the GPU. After the original add_adapter returns (the adapter is packed, registered and activated), the wrapper calls offload_in_memory_adapter for those ids only. vLLM calls add_adapter on every scheduler step, so for an adapter that is already active the wrapper does one set lookup. The offload runs under vLLM's gpu_sync_allowed(), as the rest of add_adapter's device work does.
• offload_in_memory_adapter copies each lora_a / lora_b of the registered LoRAModel (single tensors and per-slice lists) and the staged tensors to pinned host memory with non_blocking=True. Copies are keyed by (data_ptr, dtype, shape, stride), so tensors that shared GPU storage (a dense module and its staged tensor, staged aliases, pack_moe's reuse of w1 for w3 in non-gated MoE) share one host copy, and two different views of the same storage get separate copies. It then synchronizes the stream and calls torch.cuda.empty_cache() so a colocated trainer can use the freed memory.
• _load_in_memory_adapter builds on the host when the staged tensors are on the host. That is the case for a rebuild after an LRU eviction, which then works like a load from a directory: from_lora_tensors(device="cpu") returns the already pinned tensors as they are, so aliases stay shared, and no GPU memory is allocated for the rebuild while generation is running.

After an eviction from the slots, vLLM re-activates the adapter with set_lora from the host tensors, the same as for an adapter read from a directory.

.agents/docs/weight_sync.md and the SkyRLLoraConfig.sync_mode docstring describe the new placement and the size formula.

Alternatives considered

• Stage and build on the host. Copy each received chunk to pinned host memory and build the LoRAModel there, as the directory path does. Lowest GPU memory during the sync, but pack_moe (torch.stack of every expert's factors) then runs on the host. For MoE adapters that is much slower than building on the GPU and moving the result.
• Keep the copy on the GPU, but TP-sliced. For MoE, the largest factors (w13 lora_a, w2 lora_b) are replicated across TP ranks in vLLM's slot layout, so slicing would save little. It would also need a path through every LoRA layer's set_lora that accepts pre-sliced tensors, since vLLM slices unconditionally.
• fully_sharded_loras. Already supported. It shrinks the LoRA slots, not the LoRAModel copy, so it is independent of this change and can be combined with it.
• A config switch to keep the copy on the GPU. The offload costs one device-to-host copy of the packed adapter per sync (about 1 s for Qwen3-30B-A3B at rank 64) and in return makes in-memory sync use the same GPU memory as disk sync after the sync. A switch can be added if a workload needs the extra seconds back.

Constraints

• Host memory: each vLLM worker holds the packed LoRAModel in pinned host memory per registered adapter, the same amount the directory path holds, plus the deduplicated staged tensors. Freed pinned blocks return to PyTorch's pinned-host cache, not to the OS.
• The GPU peak during the sync is not reduced by the offload itself: the adapter is still built and packed on the GPU.
• Rebuild after a CPU-cache eviction (multi-tenant, more adapters than max_cpu_loras): the rebuild now runs on the host, including pack_moe, inside a scheduler step. That is slower than the previous GPU rebuild for large MoE adapters, and matches what the directory path does in the same case. Single-tenant runs with max_cpu_loras >= 1 never rebuild.

Test plan

CPU, tests/backends/skyrl_train/weight_sync/test_lora_in_memory_load.py (vllm marker, vLLM's real LoRAModel). A meta-device stage stands in for GPU tensors:

• test_device_built_adapter_is_offloaded_once: a device-built adapter is offloaded once by the patched add_adapter, and not again on a repeat add_adapter call for the active adapter.
• test_disk_adapter_is_not_offloaded: a directory-backed request goes to vLLM's original loader and is never offloaded.
• test_patch_wraps_add_adapter_of_both_worker_managers: both managers' add_adapter are wrapped (LRUCacheWorkerLoRAManager overrides it without calling super()).

GPU, tests/backends/skyrl_train/gpu/gpu_ci/patches/vllm/test_lora_in_memory_offload.py (new; runs in the FSDP GPU CI job). It drives vLLM's LoRAModel.from_lora_tensors and PackedLoRALayerWeights.pack_moe on CUDA, with a stand-in for the adapter manager:

• test_offload_moves_built_adapter_and_stage_to_pinned_host: every lora_a / lora_b and staged tensor ends up pinned on the host, torch.cuda.memory_allocated() returns to its value before the sync, packed values are unchanged, and aliased staged tensors share one host copy.
• test_rebuild_from_offloaded_stage_stays_on_host: a rebuild from the offloaded stage (vLLM's LRU path) builds on the host and keeps aliases shared (no re-pinning copy).
• test_views_with_different_strides_get_their_own_host_copy: a tensor and its transpose get separate, correct host copies.

The existing GPU e2e rows no_colocate_nccl_megatron_adapter_memory and colocate_nccl_megatron_adapter_memory in tests/backends/skyrl_train/gpu/gpu_ci/test_lora.py cover sync and generation with the real vLLM managers after the offload.

Validation

CPU, on a CPU-only run (CUDA_VISIBLE_DEVICES=, RAY_ADDRESS=local):

uv run --isolated --extra fsdp --extra dev pytest tests/backends/skyrl_train/weight_sync/test_lora_in_memory_load.py \
  tests/backends/skyrl_train/weight_sync/test_lora_receive.py tests/backends/skyrl_train/weight_sync/test_lora_target.py \
  tests/backends/skyrl_train/gpu/gpu_ci/patches/vllm/test_lora_in_memory_offload.py
# 33 passed, 3 skipped (the GPU tests skip without CUDA)

uv run --isolated --extra skyrl-train --extra dev pytest tests/train/ tests/backends/skyrl_train/ --ignore=tests/backends/skyrl_train/gpu -m "not vllm"
# 2 failed, 2241 passed, 42 skipped, 10 deselected

uv run --isolated --extra fsdp --extra dev pytest tests/train/ tests/backends/skyrl_train/ --ignore=tests/backends/skyrl_train/gpu -m "vllm"
# 7 failed, 378 passed, 11 skipped, 2259 deselected

uv run --isolated --extra dev pre-commit run --files <changed files>
# ruff, black, detect-secrets: Passed

All 9 failures are ModuleNotFoundError: No module named 'megatron' (test_moe_release_dispatcher_probs.py, test_sharded_rdt_ownership.py, test_sharded_rdt_source.py) and fail the same way on main without this change.

GPU (B200, rebased on main with vLLM 0.31.0; the memory run below was on vLLM 0.30):

uv run --isolated --extra dev --extra fsdp pytest -s tests/backends/skyrl_train/gpu/gpu_ci/patches/vllm/test_lora_in_memory_offload.py
# 3 passed
uv run --isolated --extra dev --extra megatron pytest -s tests/backends/skyrl_train/gpu/gpu_ci/test_lora.py -k "megatron_adapter_memory or frees_packed_ipc_buffer"
# 3 passed

End to end on Qwen3-30B-A3B: colocated GRPO on GSM8K, 1 node with 8 B200s, Megatron TP4 EP8, vLLM TP8, LoRA rank 64 with merge_lora=false and sync_mode=memory, 4 steps (examples/train/megatron/run_megatron_lora_qwen3-30b-a3b.sh with a smaller batch). Peak GPU memory per phase from nvidia-smi (max over GPUs, steps 2-4), main vs this PR:

Phase	main	This PR	Change
Weight sync	137.2 GiB	130.8 GiB	-6.4 GiB
Generation	138.0 GiB	131.1 GiB	-6.9 GiB
Log-prob forward	43.7 GiB	37.3 GiB	-6.4 GiB
Training	45.3 GiB	38.9 GiB	-6.4 GiB

The weight sync takes about 1 s longer (8.2-8.7 s vs 9.3-9.6 s, the device-to-host copy); step time is unchanged (63-66 s vs 64-65 s). The rollout-vs-trainer log-prob difference is unchanged (mean absolute difference 0.019-0.020 in both runs).

🤖 Generated with Claude Code

---

Note

Medium Risk
Changes vLLM worker LoRA lifecycle and GPU memory during sync/generation; LRU rebuilds now run pack_moe on the host, but behavior is covered by new GPU/CPU tests and existing e2e LoRA memory rows.

Overview
In-memory LoRA weight sync (sync_mode=memory) no longer keeps the built LoRAModel and staged tensors on the GPU after activation.

After vLLM's patched add_adapter copies the adapter into LoRA slots, offload_in_memory_adapter moves the registered model weights and SkyRL's staged tensors to pinned host memory (deduplicating shared GPU storage), synchronizes copies, and calls torch.cuda.empty_cache() so colocated trainers can reclaim GPU memory. Offload runs once per adapter via _AWAITING_OFFLOAD, including on LRUCacheWorkerLoRAManager (which overrides add_adapter without super()). Rebuilds after LRU eviction detect CPU-staged tensors and build on the host, matching the disk path.

Docs (weight_sync.md, SkyRLLoraConfig.sync_mode) describe the new lifecycle. Tests cover one-time offload, disk adapters unchanged, both manager classes patched, and GPU pinned-host behavior.

^{Reviewed by Cursor Bugbot for commit ebed80f. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request implements memory offloading for in-memory LoRA adapters, moving activated adapters and staged tensors from GPU to pinned host memory once they are loaded into vLLM's LoRA slots to reduce GPU memory footprint. The review feedback suggests using torch.empty_strided instead of torch.empty to preserve the stride and layout of non-contiguous tensors, safely retrieving add_adapter using .get() to avoid potential KeyErrors during upstream refactoring, and adding test assertions to verify that offloaded tensor strides match their original GPU counterparts.

[gemini-code-assist]

Using torch.empty to allocate the host tensor creates a contiguous tensor, which discards the original stride/layout of non-contiguous tensors (e.g., transposed or sliced views). To preserve the exact stride and layout of the original GPU tensor, use torch.empty_strided instead.

Suggested change

	host = torch.empty(tensor.shape, dtype=tensor.dtype, device="cpu", pin_memory=True)
	host = torch.empty_strided(tensor.shape, tensor.stride(), dtype=tensor.dtype, device="cpu", pin_memory=True)

[gemini-code-assist]

Accessing cls.__dict__["add_adapter"] directly will raise a KeyError if a future vLLM version refactors either WorkerLoRAManager or LRUCacheWorkerLoRAManager to inherit add_adapter rather than defining it directly. Using .get("add_adapter") and checking for None makes the patching logic more robust against upstream changes.

    for cls in (WorkerLoRAManager, LRUCacheWorkerLoRAManager):
        original = cls.__dict__.get("add_adapter")
        if original is not None:
            _ORIGINAL_ADD_ADAPTER[cls] = original
            cls.add_adapter = _make_patched_add_adapter(original)

[gemini-code-assist]

With the change to preserve strides using torch.empty_strided, we should assert that the strides of the offloaded host tensors match the original GPU tensors exactly.

Suggested change

	assert torch.equal(host, square.cpu())
	assert torch.equal(host_t, square.t().cpu())
	assert torch.equal(host, square.cpu())
	assert host.stride() == square.stride()
	assert torch.equal(host_t, square.t().cpu())
	assert host_t.stride() == square.t().stride()

[greptile-apps]

Confidence Score: 4/5

[High impact] Changes how in-memory LoRA adapters manage GPU and host memory.

The PR appears safe to merge, with a non-blocking cleanup issue that delays returning some GPU memory.

Findings

1.  Some GPU memory stays reserved ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[GPU staged tensors] --> B[Build and pack adapter]
  B --> C[Activate GPU slots]
  C --> D[Move model and stage to pinned host]
  D --> E[Rebuild or reactivate after eviction]
  E --> C

Loading

_{Reviews (1) · Last reviewed commit: "[Perf] Free up temporary memory allocate..." · Reviewed by Greptile}