Skip to content

[megatron] 11/n towards Kimi K2.6: file-backed offload of frozen masters so Tinker API LoRA training fits host RAM #2062

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
casper-hansen wants to merge 4 commits into
base: main
Choose a base branch
from
Open

[megatron] 11/n towards Kimi K2.6: file-backed offload of frozen masters so Tinker API LoRA training fits host RAM #2062

Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c3a70d3
[megatron] Offload frozen masters to file-backed mmap; keep LoRA buff…
casper-hansen Aug 19, 2026
aa604a0
[megatron] Decide offload buffer residency from the is_lora flag
avigyabb Sep 4, 2026
c321e58
unlinking file so no other process can use other model weights
avigyabb Sep 4, 2026
41c8571
got rid of .tobytes() which created 2x copies of the weights on cpu
avigyabb Sep 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 91 additions & 8 deletions skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):

@avigyabb avigyabb Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If my understanding is correct, I believe there might be a bug in how we do file naming here. If we launch two jobs using models of the same architecture (lets say a fresh model and the SFT’d version of it), there is a chance that the second job may use the weights of the first job’s model since we skip writing to the file if it already exists.

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)
Comment thread
cursor[bot] marked this conversation as resolved.
Outdated
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

@gemini-code-assist gemini-code-assist Bot Aug 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Critical/High Issues Identified:

  1. Correctness (High Severity) - Silent Weight Corruption across Runs/Checkpoints:
    The current key generation in _frozen_offload_file only hashes the parameter name, shape, and dtype:
    hashlib.sha1(f"{name}|{tuple(tensor.shape)}|{tensor.dtype}".encode())
    If you run different training jobs (e.g., with different base model checkpoints of the same architecture, or different initializations) on the same node, they will map to the exact same file paths. Since the code checks os.path.exists(path), the second run will silently reuse the frozen weights of the first run, leading to silent correctness corruption.
    Solution: Hash a small prefix of the actual tensor values (e.g., the first 1024 elements). This is extremely fast (microseconds) and guarantees uniqueness across different checkpoints/weights.

  2. Efficiency (Medium Severity) - Double CPU Memory Allocation:
    Calling .tobytes() on the entire tensor in _offload_frozen_param_to_file creates a Python bytes object copy of the entire tensor. For multi-gigabyte parameters, this duplicates the CPU memory consumption (one for the numpy array, one for the bytes object), which can trigger OOM.
    Solution: Write the contiguous numpy array directly to the file, as Python's write() accepts any object supporting the buffer protocol.

  3. Usability (Medium Severity) - Log Flooding on Failure:
    If the offload directory is not writable or offloading fails, a warning is logged for every single parameter, completely flooding the console/logs.
    Solution: Use a global flag to log the warning only once.

_warned_offload_failed = False


def _frozen_offload_file(name: str, tensor) -> str:
    import hashlib

    rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
    # Hash a small prefix of the tensor to distinguish different checkpoints/weights
    # of the same shape and dtype, preventing silent corruption or collisions
    # across different runs or models.
    flat = tensor.view(-1)
    prefix_len = min(1024, flat.numel())
    prefix = flat[:prefix_len].cpu().numpy().tobytes() if prefix_len > 0 else b""

    h = hashlib.sha1(f"{name}|{tuple(tensor.shape)}|{tensor.dtype}".encode())
    h.update(prefix)
    key = h.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:
                # Write the contiguous numpy array directly to avoid creating a
                # temporary Python bytes object copy of the entire tensor.
                f.write(data.contiguous().view(torch.uint8).flatten().cpu().numpy())
            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:
        global _warned_offload_failed
        if not _warned_offload_failed:
            logging.getLogger(__name__).warning(
                "file-backed frozen offload failed for %s (%s); falling back to pinned RAM (future failures will be silent)", name, exc
            )
            _warned_offload_failed = True
        return False



@torch.no_grad()
def offload_megatron_model_to_cpu(models):
"""
Expand All @@ -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:
Expand All @@ -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()
Expand Down
Loading