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

.agents/docs/weight_sync.md

@@ -221,6 +221,8 @@ worker               start:  no layerwise reload (the base model is untouched)
 rank 0               POST /skyrl/v1/load_lora_adapter {lora_name, in_memory: true}
 API server           add_lora(LoRARequest(lora_path="skyrl-memory://<name>", load_inplace=True))
 worker (patched)     WorkerLoRAManager._load_adapter -> from_lora_tensors(device=GPU) from the stage
+                     add_adapter: pack + activate into the LoRA slots, then
+                       offload_in_memory_adapter                 LoRAModel + stage -> pinned host
 ```
 
 Arming lasts exactly one round — `finish_weight_update` disarms — so a missing or failed
@@ -250,12 +252,13 @@ Memory and lifecycle, per inference GPU:
   tensor, not the public size. Duplicate groups are per `(module, EP-rank expert group,
   lora_A|lora_B)`, verified with `torch.equal`, so a wrong group size costs bandwidth only.
 - **Staged tensors** (`patch_lora_in_memory._STAGED`, keyed by adapter name) are the unique
-  tensors in the inference dtype. They are kept after the load because the `LoRAModel`
-  built from them shares their storage (same device, same dtype), and because vLLM rebuilds
+  tensors in the inference dtype. They are kept after the load because vLLM rebuilds
   an LRU-evicted adapter by calling `_load_adapter` again: with a directory it re-reads the
-  files, here it re-reads the stage. A resync for the same name replaces the stage;
-  `RemoteInferenceClient.unload_lora_adapter` also issues `skyrl_discard_in_memory_lora` so
-  an unloaded tenant frees its tensors.
+  files, here it re-reads the stage. After activation they live on pinned host memory (see
+  below), so a rebuild builds on the host like the directory path, including `pack_moe`,
+  which is several times slower there than on the GPU. A resync for the same
+  name replaces the stage; `RemoteInferenceClient.unload_lora_adapter` also issues
+  `skyrl_discard_in_memory_lora` so an unloaded tenant frees its tensors.
 - **The stage must never be written by vLLM.** vLLM applies `lora_alpha / r` by multiplying
   `lora_b` *in place* on every load (`LoRALayerWeights.optimize`), and since the `LoRAModel`
   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:
   `lora_alpha == r`, so vLLM's scale is exactly 1 and `optimize` is a no-op; the in-memory
   loader rejects a config where they differ. The published adapter is vLLM-shaped, not the
   trainer's raw tensors: a consumer reading `lora_B` directly sees it pre-scaled.
-- **The registered `LoRAModel` lives on the GPU**, unlike the directory path where it lives
-  on pinned CPU memory. vLLM packs per-expert LoRA into stacked tensors sized to the local
-  experts, so a MoE adapter costs roughly its *local un-deduplicated* size per registered
-  adapter (about 1/EP of the public size). With multi-tenant `max_loras` / `max_cpu_loras`
-  above 1, budget that per resident adapter.
+- **The registered `LoRAModel` lives on pinned host memory after activation**, as on the
+  directory path. It is built on the GPU (fast `pack_moe`, aliases kept shared), and the
+  patched `add_adapter` calls `offload_in_memory_adapter` once vLLM has copied it into its
+  pre-allocated LoRA slots: `lora_a` / `lora_b` of every module and the staged tensors move
+  to pinned host memory (tensors that shared GPU storage share one host copy), then
+  `torch.cuda.empty_cache()`. Generation reads only the slots; an evicted adapter is
+  re-activated by vLLM's `set_lora` from the host tensors. The GPU therefore holds the
+  built adapter only during the sync. vLLM packs per-expert LoRA into stacked tensors sized
+  to the local experts, so a MoE `LoRAModel` costs roughly its *local un-deduplicated* size
+  -- per MoE layer and local expert, `3 * r * (hidden + moe_intermediate)` elements in the
+  LoRA dtype. That copy plus the staged tensors is the transient GPU peak of each sync,
+  and the pinned host memory each worker holds per registered adapter (freed host blocks
+  return to PyTorch's pinned-host cache, not to the OS). The LoRA slots themselves are allocated at startup and
+  scale with `max_loras`; with multi-tenant `max_loras` above 1, budget the slots per
+  resident adapter.
   Measured on GLM-5.3 (`glm_moe_dsa`, 256 routed experts, Megatron EP8 -> vLLM TP8, B300
   268 GiB): **~22.4 GiB per resident adapter per inference GPU** — vLLM's weight-load
   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

@@ -25,17 +25,28 @@
 first, ``from_lora_tensors``'s per-key ``.to()`` returns the same tensor, so
 aliases created by :func:`expand_lora_aliases` keep sharing storage. Loading on
 CPU would pin a private copy per key and re-inflate a deduplicated adapter to
-its public size.
+its public size, and vLLM's MoE packing (``PackedLoRALayerWeights.pack_moe``)
+is much slower on the host.
 
 ``add_lora`` still flows API server -> engine core -> every worker, so the
 adapter name in the OpenAI-compatible model registry stays correct and
 ``load_inplace=True`` replaces the previous generation exactly as the
-directory path does. Staged tensors stay registered after the load: the
-``LoRAModel`` shares their storage (same device and dtype), so this costs
-nothing, and it lets vLLM rebuild the adapter after an LRU eviction the way it
-would re-read a directory. The next sync for the same name replaces them;
+directory path does. Staged tensors stay registered after the load, so vLLM
+can rebuild the adapter after an LRU eviction the way it would re-read a
+directory. The next sync for the same name replaces them;
 ``discard_in_memory_adapter`` frees them on unload.
 
+``WorkerLoRAManager.add_adapter`` (and the LRU variant's) is wrapped too. Once
+it has activated a GPU-built adapter, i.e. copied this worker's slice into
+vLLM's pre-allocated LoRA slots, :func:`offload_in_memory_adapter` moves the
+registered ``LoRAModel`` and the staged tensors to pinned host memory, where
+the directory path keeps the ``LoRAModel`` as well. Generation reads only the
+slots. The ``LoRAModel`` holds ``pack_moe``'s stacked copy of every local
+expert's factors, allocated outside vLLM's sleep-mode pool. vLLM re-activates an
+evicted adapter by copying from these host tensors (``set_lora``), and a
+rebuild from host-staged tensors builds (and runs ``pack_moe``) on the host,
+as the directory path does.
+
 Sharing storage has one consequence: vLLM applies ``lora_alpha / r`` by
 multiplying ``lora_b`` *in place* on every load (``LoRALayerWeights.optimize``),
 so a non-unit scale would compound on each LRU rebuild. Staged adapters must
@@ -59,6 +70,12 @@
 
 _PATCHED = False
 _ORIGINAL_LOAD_ADAPTER = None
+_ORIGINAL_ADD_ADAPTER: Dict[type, Any] = {}
+
+# ``lora_int_id``s whose ``LoRAModel`` was built on the GPU and not yet moved
+# to the host. ``LRUCacheWorkerLoRAManager.add_adapter`` runs on every
+# scheduler step, so the wrapper only offloads ids listed here.
+_AWAITING_OFFLOAD: set[int] = set()
 
 
 @dataclass
@@ -177,17 +194,23 @@ def _load_in_memory_adapter(manager, lora_request, staged: StagedLoRAAdapter):
     if moe_ep_spec is not None:
         tensors = {k: v for k, v in tensors.items() if not _is_remote_expert_key(k, moe_ep_spec)}
 
+    # Build where the staged tensors live: on the GPU right after a sync, on
+    # the host once offload_in_memory_adapter has moved them (a rebuild after
+    # an LRU eviction then works like a load from a directory).
+    build_on_host = next(iter(staged.tensors.values())).device.type == "cpu"
     lora = manager._lora_model_cls.from_lora_tensors(
         lora_model_id=lora_request.lora_int_id,
         tensors=tensors,
         peft_helper=peft_helper,
-        device=manager.device,
+        device="cpu" if build_on_host else manager.device,
         dtype=manager.lora_config.lora_dtype,
         model_vocab_size=manager.vocab_size,
         weights_mapper=weights_mapper,
         skip_prefixes=skip_prefixes,
     )
     lora.is_3d_lora_weight = lora_request.is_3d_lora_weight
+    if not build_on_host:
+        _AWAITING_OFFLOAD.add(lora_request.lora_int_id)
     return lora
 
 
@@ -205,8 +228,74 @@ def _patched_load_adapter(self, lora_request):
     return _load_in_memory_adapter(self, lora_request, staged)
 
 
+def _to_host(tensor, moved: Dict[tuple, torch.Tensor]):
+    """Returns a pinned host copy of a CUDA ``tensor``, else ``tensor`` unchanged.
+
+    ``moved`` maps each copied tensor's ``(data_ptr, dtype, shape, stride)`` to
+    its host copy, so tensors that are the same view of one GPU storage share
+    one host copy.
+    """
+    if not isinstance(tensor, torch.Tensor) or not tensor.is_cuda:
+        return tensor
+    key = (tensor.data_ptr(), tensor.dtype, tuple(tensor.shape), tensor.stride())
+    host = moved.get(key)
+    if host is None:
+        host = torch.empty(tensor.shape, dtype=tensor.dtype, device="cpu", pin_memory=True)
+        host.copy_(tensor, non_blocking=True)
+        moved[key] = host
+    return host
+
+
+def offload_in_memory_adapter(manager, lora_request) -> None:
+    """Moves an activated in-memory adapter's tensors from the GPU to pinned host memory.
+
+    Covers the registered ``LoRAModel``'s ``lora_a`` / ``lora_b`` (single
+    tensors or per-slice lists) and the staged tensors for
+    ``lora_request.lora_name``. Tensors that shared GPU storage share the host
+    copy. Releases the freed blocks with ``torch.cuda.empty_cache`` so that
+    other processes on the GPU (a colocated trainer) can use them.
+    """
+    lora_model = manager._adapter_manager.get_adapter(lora_request.lora_int_id)
+    if lora_model is None:
+        return
+    moved: Dict[tuple, torch.Tensor] = {}
+    for weights in lora_model.loras.values():
+        for attr in ("lora_a", "lora_b"):
+            value = getattr(weights, attr, None)
+            if isinstance(value, list):
+                setattr(weights, attr, [_to_host(t, moved) for t in value])
+            else:
+                setattr(weights, attr, _to_host(value, moved))
+            # Drops the reference to the GPU tensor so that empty_cache below can release it.
+            del value
+    staged = _STAGED.get(lora_request.lora_name)
+    if staged is not None:
+        staged.tensors = {k: _to_host(v, moved) for k, v in staged.tensors.items()}
+    if moved:
+        # Waits for the non_blocking copies before returning the GPU blocks to the driver.
+        torch.cuda.current_stream().synchronize()
+        moved.clear()
+        torch.cuda.empty_cache()
+
+
+def _make_patched_add_adapter(original):
+    def _patched_add_adapter(self, lora_request):
+        loaded = original(self, lora_request)
+        # An id stays listed if ``original`` raised; the next successful call offloads it.
+        if lora_request.lora_int_id in _AWAITING_OFFLOAD:
+            from vllm.utils.gpu_sync_debug import gpu_sync_allowed
+
+            _AWAITING_OFFLOAD.discard(lora_request.lora_int_id)
+            # vLLM runs add_adapter's own device work under the same context.
+            with gpu_sync_allowed():
+                offload_in_memory_adapter(self, lora_request)
+        return loaded
+
+    return _patched_add_adapter
+
+
 def apply_lora_in_memory_patch() -> None:
-    """Install the wrapper on ``WorkerLoRAManager`` (idempotent, worker-process local)."""
+    """Install the wrappers on vLLM's worker LoRA managers (idempotent, worker-process local)."""
     global _PATCHED, _ORIGINAL_LOAD_ADAPTER
     if _PATCHED:
         return
@@ -215,10 +304,19 @@ def apply_lora_in_memory_patch() -> None:
     except ModuleNotFoundError:
         # Trainer processes import the worker-extension module too.
         return
+    from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
+
     _ORIGINAL_LOAD_ADAPTER = WorkerLoRAManager._load_adapter
     WorkerLoRAManager._load_adapter = _patched_load_adapter
+    # LRUCacheWorkerLoRAManager overrides add_adapter without calling super().
+    for cls in (WorkerLoRAManager, LRUCacheWorkerLoRAManager):
+        _ORIGINAL_ADD_ADAPTER[cls] = cls.__dict__["add_adapter"]
+        cls.add_adapter = _make_patched_add_adapter(_ORIGINAL_ADD_ADAPTER[cls])
     _PATCHED = True
-    logger.debug("in-memory LoRA patch installed on vllm WorkerLoRAManager._load_adapter")
+    logger.debug(
+        "in-memory LoRA patch installed on vllm WorkerLoRAManager._load_adapter and on "
+        "add_adapter of WorkerLoRAManager and LRUCacheWorkerLoRAManager"
+    )
 
 
 def is_patched() -> bool:

skyrl/train/config/config.py

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