[tinker] 13/n towards Kimi K2.6: adapter-only colocated weight sync for Tinker API LoRA (merge_lora=false)

skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py

@@ -20,6 +20,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+import os
 from typing import Any, Dict, List, Optional, Union
 
 import torch
@@ -222,6 +224,70 @@ def load_megatron_grads_to_gpu(models):
                     param.grad = param.grad.to(torch.cuda.current_device(), non_blocking=True)
 
 
+def _chunk_has_lora_adapters(model_chunk) -> bool:
+    """True when the chunk trains LoRA adapters (and only adapters).
+
+    Megatron's fused param/grad buffers only hold grad-requiring params, so
+    for a LoRA model they contain nothing but the adapters (a few GB). The
+    frozen base weights live outside the buffers and are offloaded
+    param-by-param below.
+    """
+    return any("adapter" in name for name, param in model_chunk.named_parameters() if param.requires_grad)
+
+
+# Frozen (requires_grad=False, non-adapter) weights are immutable for the
+# whole run, so their CPU offload copies can live in file-backed mmap storage
+# instead of RAM: the pages are then *clean page cache* the kernel can evict
+# and re-read freely, instead of ~1.3TB/node of anonymous/pinned memory that
+# competes with the vLLM engines for physical RAM (the source of repeated
+# NUMA OOM kills and compress-swap stalls on TB-scale colocated models).
+# Files are written once per rank on first offload and reused afterwards.
+# Set SKYRL_FROZEN_OFFLOAD_DIR=0 (or empty) to restore pinned-RAM offload.
+_FROZEN_OFFLOAD_DIR = os.environ.get("SKYRL_FROZEN_OFFLOAD_DIR", "/data/skyrl/frozen-offload")
+
+
+def _frozen_offload_enabled() -> bool:
+    return bool(_FROZEN_OFFLOAD_DIR) and _FROZEN_OFFLOAD_DIR != "0"
+
+
+def _frozen_offload_file(name: str, tensor) -> str:
+    import hashlib
+
+    rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+    key = hashlib.sha1(f"{name}|{tuple(tensor.shape)}|{tensor.dtype}".encode()).hexdigest()[:20]
+    rank_dir = os.path.join(_FROZEN_OFFLOAD_DIR, f"rank{rank}")
+    os.makedirs(rank_dir, exist_ok=True)
+    return os.path.join(rank_dir, f"{key}.bin")
+
+
+def _offload_frozen_param_to_file(name: str, param) -> bool:
+    """Move a frozen param's data to a file-backed mmap CPU tensor.
+
+    Returns True on success; False to let the caller fall back to pinned RAM.
+    """
+    try:
+        data = param.data.detach()
+        nbytes = data.numel() * data.element_size()
+        path = _frozen_offload_file(name, data)
+        if not (os.path.exists(path) and os.path.getsize(path) == nbytes):
+            tmp = f"{path}.tmp{os.getpid()}"
+            with open(tmp, "wb") as f:
+                f.write(data.contiguous().view(torch.uint8).flatten().cpu().numpy().tobytes())
+            os.replace(tmp, path)
+        mapped = (
+            torch.from_file(path, shared=False, size=nbytes, dtype=torch.uint8)
+            .view(data.dtype)
+            .view(data.shape)
+        )
+        param._offload_cpu_data = mapped
+        return True
+    except (OSError, RuntimeError) as exc:
+        logging.getLogger(__name__).warning(
+            "file-backed frozen offload failed for %s (%s); falling back to pinned RAM", name, exc
+        )
+        return False
+
+
 @torch.no_grad()
 def offload_megatron_model_to_cpu(models):
     """
@@ -233,22 +299,37 @@ def offload_megatron_model_to_cpu(models):
     """
     for model_chunk in models:
         if isinstance(model_chunk, DDP):
-            for buffer in model_chunk.buffers + model_chunk.expert_parallel_buffers:
-                # use megatron buffer built in function to offload to cpu
-                # https://github.com/NVIDIA/Megatron-LM/blob/core_v0.16.0/megatron/core/distributed/param_and_grad_buffer.py#L964
-                buffer.offload_to_cpu(move_params=True, move_grads=False)
+            # LoRA: keep the fused buffers (adapters only, a few GB) resident.
+            # The adapter-only weight sync exports straight from these GPU
+            # tensors, so the TB-scale frozen masters never need to round-trip
+            # through the GPU just to sync a rank-32 adapter.
+            if not _chunk_has_lora_adapters(model_chunk):
+                for buffer in model_chunk.buffers + model_chunk.expert_parallel_buffers:
+                    # use megatron buffer built in function to offload to cpu
+                    # https://github.com/NVIDIA/Megatron-LM/blob/core_v0.16.0/megatron/core/distributed/param_and_grad_buffer.py#L964
+                    buffer.offload_to_cpu(move_params=True, move_grads=False)
 
             # LoRA-aware offloading: offload non-lora base weights that live
             # outside the fused Megatron buffers (e.g. HF/bridge "to_wrap" weights).
+            # Frozen weights are immutable, so prefer file-backed mmap copies
+            # (clean, evictable page cache) over pinned RAM; see
+            # _offload_frozen_param_to_file.
+            use_file_offload = _frozen_offload_enabled()
             for name, param in model_chunk.named_parameters():
                 if (
                     param.is_cuda
                     and not param.requires_grad
                     and "adapter" not in name
                     and param.data.storage().size() > 0
                 ):
-                    cpu_tensor = param.data.detach().cpu().pin_memory()
-                    param._offload_cpu_data = cpu_tensor
+                    if hasattr(param, "_offload_cpu_data") and param._offload_cpu_data is not None:
+                        # Frozen data never changes: the existing CPU copy
+                        # (file-backed or pinned) is still valid; just free
+                        # the GPU side again.
+                        pass
+                    elif not (use_file_offload and _offload_frozen_param_to_file(name, param)):
+                        cpu_tensor = param.data.detach().cpu().pin_memory()
+                        param._offload_cpu_data = cpu_tensor
                     param._offload_cuda_numel = param.data.numel()
                     param.data = torch.empty(0, dtype=param.data.dtype, device=param.data.device)
         else:
@@ -260,8 +341,10 @@ def offload_megatron_model_to_cpu(models):
 def load_megatron_model_to_gpu(models):
     for model_chunk in models:
         if isinstance(model_chunk, DDP):
-            for buffer in model_chunk.buffers + model_chunk.expert_parallel_buffers:
-                buffer.reload_from_cpu(move_params=True, move_grads=False)
+            # LoRA buffers never offload (see offload_megatron_model_to_cpu).
+            if not _chunk_has_lora_adapters(model_chunk):
+                for buffer in model_chunk.buffers + model_chunk.expert_parallel_buffers:
+                    buffer.reload_from_cpu(move_params=True, move_grads=False)
 
             # Restore any LoRA-frozen base weights that were offloaded above.
             device_id = torch.cuda.current_device()

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())