[fix][megatron] Make multi-LoRA adapter swaps safe against colocation offload

skyrl/backends/skyrl_train/workers/megatron/adapter_store.py

@@ -37,10 +37,26 @@ def _iter_buffers(model_chunks) -> Iterable[Tuple[int, int, Any]]:
 
 
 def _new_pinned_like(t: torch.Tensor) -> torch.Tensor:
-    """Allocate a pinned-CPU tensor with the same shape/dtype as t."""
+    """Allocate a pinned-CPU tensor with the same shape/dtype as t.
+
+    Safe to call on an offloaded buffer: only shape/dtype metadata is read,
+    which survives ``storage().resize_(0)``.
+    """
     return torch.empty_like(t, device="cpu").pin_memory()
 
 
+def _is_resident(t: Optional[torch.Tensor]) -> bool:
+    """True when ``t`` still owns storage we can copy to/from.
+
+    Megatron's offload path frees GPU buffers with ``storage().resize_(0)``
+    (see ``_ParamAndGradBuffer.offload_to_cpu``) and leaves the tensor object
+    — full shape and all — pointing at a zero-sized storage. Copying from it
+    raises ``cudaErrorInvalidValue``, so every read/write of a DDP buffer has
+    to be gated on this.
+    """
+    return t is not None and t.untyped_storage().size() > 0
+
+
 def _expected_lora_param_check(model_chunks) -> None:
     """Sanity-check: every trainable param under DDP buffers is a LoRA adapter param.
 
@@ -149,6 +165,10 @@ def __init__(self) -> None:
         self._pristine: Optional[AdapterSlot] = None
         self._current_id: Optional[str] = None
         self._signature: Optional[LoraSignature] = None
+        # True while the DDP grad buffers are offloaded, i.e. each adapter's
+        # grads live only in its CPU slot. Set by park_grads, cleared by
+        # unpark_grads.
+        self._grads_parked: bool = False
 
     @property
     def current_id(self) -> Optional[str]:
@@ -223,12 +243,34 @@ def _allocate_empty_slot(self, model_chunks, optimizer) -> AdapterSlot:
             slot.cpu_param_group_state.append(group_state)
         return slot
 
+    @staticmethod
+    def _require_param_residency(buf, mc_idx: int, buf_idx: int) -> None:
+        """Fail loudly when a swap is attempted with model params offloaded.
+
+        Callers must backload the model before swapping (the dispatch does
+        this via ``_ensure_on_gpu(..., need_model=True)``). Grads are allowed
+        to be offloaded — see :meth:`park_grads` — but params are not, because
+        the CPU mirror Megatron keeps for them (``param_data_cpu``) is not
+        per-adapter and would hand the next backload the wrong tenant's
+        weights.
+        """
+        if not _is_resident(buf.param_data):
+            raise RuntimeError(
+                f"AdapterStore: DDP buffer {mc_idx}/{buf_idx} param_data is offloaded; "
+                f"backload the model before swapping adapters."
+            )
+
     @torch.no_grad()
     def _snapshot(self, slot: AdapterSlot, model_chunks, optimizer) -> None:
         """Copy live GPU state into `slot` (CPU)."""
         for mc_idx, buf_idx, buf in _iter_buffers(model_chunks):
+            self._require_param_residency(buf, mc_idx, buf_idx)
             slot.cpu_param_data[mc_idx][buf_idx].copy_(buf.param_data, non_blocking=True)
-            slot.cpu_grad_data[mc_idx][buf_idx].copy_(buf.grad_data, non_blocking=True)
+            # Skip grads while the grad buffers are offloaded: their GPU
+            # storage is freed, and the slot's copy — parked by park_grads()
+            # just before the offload — is already the authoritative one.
+            if _is_resident(buf.grad_data):
+                slot.cpu_grad_data[mc_idx][buf_idx].copy_(buf.grad_data, non_blocking=True)
         for opt_idx, _opt in enumerate(iter_opts(optimizer)):
             groups = getattr(_opt, "shard_fp32_from_float16_groups", None) or []
             for g, group in enumerate(groups):
@@ -254,8 +296,13 @@ def _snapshot(self, slot: AdapterSlot, model_chunks, optimizer) -> None:
     def _restore(self, slot: AdapterSlot, model_chunks, optimizer) -> None:
         """Copy `slot` (CPU) into live GPU state."""
         for mc_idx, buf_idx, buf in _iter_buffers(model_chunks):
+            self._require_param_residency(buf, mc_idx, buf_idx)
             buf.param_data.copy_(slot.cpu_param_data[mc_idx][buf_idx], non_blocking=True)
-            buf.grad_data.copy_(slot.cpu_grad_data[mc_idx][buf_idx], non_blocking=True)
+            # Offloaded grad buffers have no storage to restore into; the
+            # incoming adapter's grads stay parked in its slot and are
+            # re-materialised by unpark_grads() on the next backload.
+            if _is_resident(buf.grad_data):
+                buf.grad_data.copy_(slot.cpu_grad_data[mc_idx][buf_idx], non_blocking=True)
         for opt_idx, _opt in enumerate(iter_opts(optimizer)):
             groups = getattr(_opt, "shard_fp32_from_float16_groups", None) or []
             for g, group in enumerate(groups):
@@ -372,6 +419,66 @@ def _copy_slot(self, src: AdapterSlot, dst: AdapterSlot) -> None:
                     else:
                         dst_pg[k] = v
 
+    # ------------------------------------------------------------------
+    # Grad parking across CPU offload
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def park_grads(self, model_chunks) -> None:
+        """Save the live adapter's grads into its slot before a grad offload.
+
+        Megatron frees ``grad_data`` on offload and zero-fills it on reload,
+        so grads accumulated by a ``forward_backward`` that hasn't reached its
+        ``optim_step`` yet are destroyed outright. Under colocation another
+        tenant's request (a sample, a forward) offloads in exactly that gap,
+        so the grads have to be parked per-adapter instead.
+
+        Call immediately before the strategy's offload. No-op when the grad
+        buffers are already offloaded or no adapter is live.
+        """
+        if self._current_id is None or self._current_id not in self._slots:
+            return
+        slot = self._slots[self._current_id]
+        parked = False
+        for mc_idx, buf_idx, buf in _iter_buffers(model_chunks):
+            if not _is_resident(buf.grad_data):
+                continue
+            slot.cpu_grad_data[mc_idx][buf_idx].copy_(buf.grad_data, non_blocking=True)
+            parked = True
+        if parked:
+            torch.cuda.current_stream().synchronize()
+            self._grads_parked = True
+
+    @torch.no_grad()
+    def unpark_grads(self, model_chunks) -> None:
+        """Re-materialise the live adapter's grads after a grad backload.
+
+        Mirror of :meth:`park_grads`. Restores the adapter that is live *now*,
+        which need not be the one that was live at park time — a swap in the
+        offload window only moves CPU slots around, and this is where the
+        result lands back on the GPU.
+
+        Call immediately after the strategy's backload. No-op unless grads
+        were parked and the buffers are resident again.
+        """
+        if not self._grads_parked:
+            return
+        if self._current_id is None or self._current_id not in self._slots:
+            # Live adapter was deleted while offloaded: nothing to restore,
+            # and the reloaded buffers are already zeroed.
+            self._grads_parked = False
+            return
+        slot = self._slots[self._current_id]
+        restored = False
+        for mc_idx, buf_idx, buf in _iter_buffers(model_chunks):
+            if not _is_resident(buf.grad_data):
+                continue
+            buf.grad_data.copy_(slot.cpu_grad_data[mc_idx][buf_idx], non_blocking=True)
+            restored = True
+        if restored:
+            torch.cuda.current_stream().synchronize()
+            self._grads_parked = False
+
     @torch.no_grad()
     def delete(self, model_id: str) -> None:
         """Drop the slot for `model_id`.
@@ -402,6 +509,11 @@ def swap_to(self, model_id: str, model_chunks, optimizer) -> None:
         agree on the live adapter before the next collective. TP/PP/EP groups
         do not need barriers because the swap is identical-shape on all
         ranks within those groups (LoRA signature is fixed).
+
+        Model params must be GPU-resident; the DDP grad buffers and the
+        DistributedOptimizer state need not be (colocation offloads them
+        between requests). Offloaded grads are left parked in their slots —
+        see :meth:`park_grads`.
         """
         if model_id not in self._slots:
             raise KeyError(f"AdapterStore: unknown adapter '{model_id}'")

skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py

@@ -1552,6 +1552,32 @@ def swap_to_adapter(self, model_id: str) -> None:
             return  # FFT path: no-op
         self.adapter_store.swap_to(model_id, self.actor_module, self.optimizer)
 
+    def offload_to_cpu(self, offload_optimizer: bool = True, offload_model: bool = True):
+        """Offload worker state, parking the live adapter's grads first.
+
+        The optimizer half of the offload frees the DDP grad buffers outright
+        (Megatron's reload zero-fills them), which would silently drop grads a
+        tenant accumulated in a ``forward_backward`` whose ``optim_step``
+        hasn't arrived yet — under colocation another tenant's sample lands in
+        exactly that gap. Parking moves them into the adapter's CPU slot,
+        where they survive both the offload and any swap that happens while
+        offloaded.
+        """
+        if offload_optimizer and self.adapter_store is not None and self.actor_module is not None:
+            self.adapter_store.park_grads(self.actor_module)
+        super().offload_to_cpu(offload_optimizer=offload_optimizer, offload_model=offload_model)
+
+    def backload_to_gpu(self, backload_optimizer: bool = True, backload_model: bool = True):
+        """Backload worker state, re-materialising the live adapter's grads.
+
+        Counterpart to :meth:`offload_to_cpu`. Restores whichever adapter is
+        live at this point, which may differ from the one parked if a swap
+        happened while offloaded.
+        """
+        super().backload_to_gpu(backload_optimizer=backload_optimizer, backload_model=backload_model)
+        if backload_optimizer and self.adapter_store is not None and self.actor_module is not None:
+            self.adapter_store.unpark_grads(self.actor_module)
+
     def adapter_store_state(self) -> dict:
         """Diagnostic: return current_id + registered model_ids. Cheap; useful
         for tests."""

skyrl/backends/skyrl_train/workers/worker_dispatch.py

@@ -86,11 +86,17 @@ def ensure_active_adapter(self, role: str, model_id: Optional[str]) -> None:
         No-op when ``model_id is None`` (single-tenant / FFT path) or when
         the workers don't have an AdapterStore (non-LoRA strategies).
 
-        Must be called *after* ``_ensure_on_gpu(role, ...)`` so the model
-        and optimizer storages are live before we tensor.copy_() into them.
+        The swap copies the DDP param buffers, so the model must be resident
+        before it runs. Callers generally arrive here right after their own
+        ``_ensure_on_gpu``; we repeat it (a no-op against dispatch-local
+        state) so paths that only need the optimizer — ``set_lr``, say —
+        can't hand the AdapterStore freed param storage. The optimizer and
+        grad buffers are allowed to stay offloaded: the store copies the
+        optimizer's CPU tensors in place and leaves parked grads alone.
         """
         if model_id is None or role not in self._actor_groups:
             return
+        self._ensure_on_gpu(role, need_optimizer=False, need_model=True)
         ray.get(self._actor_groups[role].async_run_ray_method("pass_through", "swap_to_adapter", model_id))
 
     def register_adapter(self, role: str, model_id: str) -> None:
@@ -417,7 +423,14 @@ def optim_step(self, model: str, model_id: Optional[str] = None) -> Optional[flo
         """Run optimizer step. For single-tenant training, the model should already be on GPU from forward_backward.
 
         For multi-tenant LoRA training, ``model_id`` is used to ensure the correct adapter is used.
+
+        The residency check is not redundant with ``forward_backward``'s under
+        multi-tenancy: another tenant's request (e.g. a sample, which offloads
+        the optimizer) can land between this model's forward_backward and its
+        optim_step, so the state this step writes to may have been offloaded
+        in the gap.
         """
+        self._ensure_on_gpu(model, need_optimizer=True, need_model=True)
         self.ensure_active_adapter(model, model_id)
         refs = self._actor_groups[model].async_run_ray_method("pass_through", "optim_step")
         grad_norms = ray.get(refs)