diff --git a/skyrl/backends/skyrl_train/workers/megatron/adapter_store.py b/skyrl/backends/skyrl_train/workers/megatron/adapter_store.py index ecdf99ccc3..4a38a1496e 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/adapter_store.py +++ b/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()) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_adapter_store_swap.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_adapter_store_swap.py new file mode 100644 index 0000000000..776bfcf169 --- /dev/null +++ b/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)