[megatron] 12/n towards Kimi K2.6: harden AdapterStore swaps for Tinker API multi-LoRA session churn

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

@@ -41,6 +41,22 @@ def _new_pinned_like(t: torch.Tensor) -> torch.Tensor:
     return torch.empty_like(t, device="cpu").pin_memory()
 
 
+def _grad_data_live(buf) -> bool:
+    """False while the buffer's grad storage is offloaded (freed).
+
+    DDP.offload_grad_buffers() frees grad_data via storage().resize_(0) while
+    keeping the tensor view intact; touching the view in that state is
+    undefined behavior (cudaMemcpyAsync on the stale pointer fails with
+    `invalid argument`). restore_grad_buffers()/zero_grad_buffer() reallocate
+    it zero-filled, i.e. Megatron already treats offloaded grads as discarded
+    — so snapshot/restore skip the grad copy instead of crashing. Grads are
+    only offloaded post-optim-step (never mid-accumulation), so the pending
+    grads at that point have been consumed and are safe to drop.
+    """
+    gd = buf.grad_data
+    return not (gd.is_cuda and gd.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,11 @@ def __init__(self) -> None:
         self._pristine: Optional[AdapterSlot] = None
         self._current_id: Optional[str] = None
         self._signature: Optional[LoraSignature] = None
+        # True when the live GPU state no longer mirrors any registered slot
+        # (nor pristine): set when the current adapter is deleted, cleared by
+        # the next completed swap_to(). While dirty, create() must NOT adopt
+        # the live state as a new adapter — it belongs to a deleted tenant.
+        self._live_dirty: bool = False
 
     @property
     def current_id(self) -> Optional[str]:
@@ -227,7 +248,14 @@ 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):
             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)
+            if _grad_data_live(buf):
+                slot.cpu_grad_data[mc_idx][buf_idx].copy_(buf.grad_data, non_blocking=True)
+            else:
+                # Grad storage is offloaded/freed: the live grads were already
+                # consumed by the last optim step and Megatron zero-fills on
+                # reload, so record them as zero rather than reading a freed
+                # pointer.
+                slot.cpu_grad_data[mc_idx][buf_idx].zero_()
         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,7 +282,12 @@ 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):
             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)
+            if _grad_data_live(buf):
+                buf.grad_data.copy_(slot.cpu_grad_data[mc_idx][buf_idx], non_blocking=True)
+            # else: grad storage is offloaded/freed. zero_grad_buffer() /
+            # restore_grad_buffers() reallocate it zero-filled before the next
+            # forward_backward, and a slot swapped in while offloaded carries
+            # post-step (zero) grads anyway — skipping loses nothing.
         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):
@@ -304,17 +337,26 @@ def register_pristine(self, model_chunks, optimizer, signature: LoraSignature) -
         self._signature = signature
         self._pristine = self._allocate_empty_slot(model_chunks, optimizer)
         self._snapshot(self._pristine, model_chunks, optimizer)
+        # _snapshot issues non_blocking D2H copies into pinned memory; the
+        # pristine slot is read on the *CPU* by create()'s _copy_slot, so the
+        # DMA must have landed before this returns.
+        torch.cuda.current_stream().synchronize()
 
     @torch.no_grad()
     def create(self, model_id: str, model_chunks, optimizer, signature: LoraSignature) -> None:
         """Register a new adapter slot.
 
-        - First registration: this is also the live adapter; allocate a slot
-          but skip the pristine→slot copy because the live state already
-          equals pristine. `current_id` becomes `model_id`.
+        - First registration (live state still pristine): this is also the
+          live adapter; allocate a slot but skip the pristine→slot copy
+          because the live state already equals pristine. `current_id`
+          becomes `model_id`.
         - Subsequent registrations: allocate slot and copy pristine → slot.
           Live state is unchanged (no swap). The new adapter only becomes
           live when the next `swap_to(model_id)` is issued.
+        - After a delete of the current adapter (`current_id is None` but
+          live state is dirty), the new adapter is seeded from pristine like
+          any other registration: adopting the live state here would silently
+          inherit the deleted tenant's weights, fp32 masters and Adam state.
         """
         if self._signature is None:
             raise RuntimeError("AdapterStore.create called before register_pristine")
@@ -329,12 +371,14 @@ def create(self, model_id: str, model_chunks, optimizer, signature: LoraSignatur
             raise ValueError(f"AdapterStore: adapter '{model_id}' already registered")
 
         slot = self._allocate_empty_slot(model_chunks, optimizer)
-        if self._current_id is None:
+        if self._current_id is None and not self._live_dirty:
             # First adapter: live state IS pristine; slot will be filled on
             # the next snapshot (i.e. swap-away). Treat live as authoritative.
             self._current_id = model_id
         else:
-            # Seed the new slot from pristine.
+            # Seed the new slot from pristine. When current_id is None but
+            # live is dirty (current adapter was deleted), the next
+            # swap_to(model_id) restores this pristine copy into live.
             self._copy_slot(self._pristine, slot)
         self._slots[model_id] = slot
 
@@ -384,6 +428,9 @@ def delete(self, model_id: str) -> None:
         del self._slots[model_id]
         if self._current_id == model_id:
             self._current_id = None
+            # Live GPU state still mirrors the deleted adapter. Mark it dirty
+            # so create() won't adopt it; the next swap_to() overwrites it.
+            self._live_dirty = True
 
     @torch.no_grad()
     def swap_to(self, model_id: str, model_chunks, optimizer) -> None:
@@ -407,9 +454,8 @@ def swap_to(self, model_id: str, model_chunks, optimizer) -> None:
         if self._current_id == model_id:
             return  # no-op fast path
 
-        dp_group = mpu.get_data_parallel_group()
         if dist.is_available() and dist.is_initialized():
-            dist.barrier(group=dp_group)
+            dist.barrier(group=mpu.get_data_parallel_group())
 
         if self._current_id is not None:
             current_slot = self._slots[self._current_id]
@@ -421,6 +467,7 @@ def swap_to(self, model_id: str, model_chunks, optimizer) -> None:
         torch.cuda.current_stream().synchronize()
 
         self._current_id = model_id
+        self._live_dirty = False  # live now mirrors target_slot
 
         if dist.is_available() and dist.is_initialized():
-            dist.barrier(group=dp_group)
+            dist.barrier(group=mpu.get_data_parallel_group())

tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_adapter_store_swap.py

@@ -0,0 +1,138 @@
+"""Unit tests for AdapterStore swap semantics around offload and delete.
+
+Regressions covered (both hit in production on 2026-08-18):
+
+1. swap_to() while the DDP grad buffers are offloaded — Megatron frees
+   grad_data via storage().resize_(0) after every post-step offload, and the
+   H2D copy_ into the freed view failed with CUDA `invalid argument`.
+2. create() after delete-of-current — the store adopted the live GPU state
+   (the *deleted* tenant's weights, fp32 masters and Adam state) as the new
+   adapter instead of seeding from pristine.
+
+The tests fake the Megatron DDP buffers by monkeypatching _iter_buffers, so
+they run on a single GPU with KB-scale allocations and no distributed init.
+"""
+
+import pytest
+import torch
+from types import SimpleNamespace
+
+from skyrl.backends.skyrl_train.workers.megatron import adapter_store as astore
+from skyrl.backends.skyrl_train.workers.megatron.adapter_store import (
+    AdapterStore,
+    LoraSignature,
+)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+
+_SIG = LoraSignature(
+    rank=8,
+    alpha=16,
+    target_modules=("linear_proj",),
+    lora_type="lora",
+    tp_size=1,
+    pp_size=1,
+    ep_size=1,
+)
+
+
+class _FakeBuf:
+    """Stands in for a megatron _ParamAndGradBuffer (param_data + grad_data)."""
+
+    def __init__(self, n: int = 64):
+        self.param_data = torch.randn(n, device="cuda", dtype=torch.bfloat16)
+        self.grad_data = torch.zeros(n, device="cuda", dtype=torch.bfloat16)
+        self.params = []  # _expected_lora_param_check iterates this
+
+    def free_grad_storage(self):
+        """Mimic DDP.offload_grad_buffers(): free storage, keep the view."""
+        torch.cuda.synchronize()
+        self.grad_data.untyped_storage().resize_(0)
+
+
+def _fake_opt():
+    return SimpleNamespace(optimizer=SimpleNamespace(state={}, param_groups=[]))
+
+
+@pytest.fixture
+def store_env(monkeypatch):
+    """AdapterStore wired to one fake buffer; model_chunks is [buf] itself."""
+    buf = _FakeBuf()
+    monkeypatch.setattr(astore, "_iter_buffers", lambda chunks: [(0, 0, chunks[0])])
+    monkeypatch.setattr(astore, "_expected_lora_param_check", lambda chunks: None)
+    store = AdapterStore()
+    opt = _fake_opt()
+    store.register_pristine([buf], opt, _SIG)
+    return store, buf, opt
+
+
+def test_swap_with_offloaded_grads_does_not_crash(store_env):
+    """Prod crash: snapshot/restore must skip grad copies on freed storage."""
+    store, buf, opt = store_env
+    store.create("A", [buf], opt, _SIG)  # first adapter: adopts live
+    store.create("B", [buf], opt, _SIG)  # seeded from pristine
+    pristine_params = buf.param_data.clone()
+
+    # "Train" A, then offload grads exactly like offload_after_step does.
+    buf.param_data.add_(1.0)
+    buf.free_grad_storage()
+
+    # Previously: torch.AcceleratorError CUDA invalid argument in _restore.
+    store.swap_to("B", [buf], opt)
+
+    assert store.current_id == "B"
+    torch.cuda.synchronize()
+    # B was seeded from pristine, so live params must be back to pristine.
+    assert torch.equal(buf.param_data, pristine_params)
+    # A's slot still captured the trained params (param storage was live) ...
+    a_params = store._slots["A"].cpu_param_data[0][0]
+    assert torch.equal(a_params, (pristine_params + 1.0).cpu())
+    # ... while its grads were recorded as zero instead of reading freed memory.
+    a_grads = store._slots["A"].cpu_grad_data[0][0]
+    assert torch.count_nonzero(a_grads) == 0
+
+
+def test_create_after_delete_does_not_adopt_live_state(store_env):
+    """Prod bug: a new adapter created after delete-of-current inherited the
+    deleted tenant's live state instead of pristine."""
+    store, buf, opt = store_env
+    pristine_params = buf.param_data.clone()
+
+    store.create("A", [buf], opt, _SIG)  # adopts live (true first create)
+    buf.param_data.add_(2.0)  # "train" A
+    store.delete("A")  # session expiry: current cleared, live now stale
+
+    store.create("B", [buf], opt, _SIG)
+    # B must NOT adopt the deleted tenant's live state.
+    assert store.current_id is None
+
+    store.swap_to("B", [buf], opt)
+    torch.cuda.synchronize()
+    assert store.current_id == "B"
+    assert torch.equal(buf.param_data, pristine_params)
+
+
+def test_full_prod_sequence_create_before_expiry_delete(store_env):
+    """The exact 06:57 ordering: create B while A is current, then A expires
+    (delete -> current=None), then swap_to(B) with grads offloaded."""
+    store, buf, opt = store_env
+    pristine_params = buf.param_data.clone()
+
+    store.create("A", [buf], opt, _SIG)
+    buf.param_data.add_(3.0)  # overnight training
+    store.create("B", [buf], opt, _SIG)  # new session registers first...
+    store.delete("A")  # ...then the stale session expires
+    buf.free_grad_storage()  # trainer idle: grads offloaded
+
+    store.swap_to("B", [buf], opt)
+    torch.cuda.synchronize()
+    assert store.current_id == "B"
+    assert torch.equal(buf.param_data, pristine_params)
+
+    # A subsequent adapter created while B is live seeds from pristine and
+    # round-trips through a snapshot of B without touching freed grads.
+    store.create("C", [buf], opt, _SIG)
+    store.swap_to("C", [buf], opt)
+    torch.cuda.synchronize()
+    assert store.current_id == "C"
+    assert torch.equal(buf.param_data, pristine_params)