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

[casper-hansen]

Part of the Kimi K2.6/K2.7 series (previous: #2024, #2025, #2026, #2027, #2029, #2031, #2032). Standalone; no dependency on the other open PRs, but #2026 (merge_lora=false adapter sync) and the upcoming adapter-only weight sync build on it.

What

Two changes to the Megatron CPU-offload layer in megatron_utils.py, for LoRA training through the Tinker API on models whose frozen base does not fit host RAM twice:

• Frozen (requires_grad=False, non-adapter) params offload to file-backed mmap tensors instead of pinned/anonymous RAM. Written once per rank on first offload under SKYRL_FROZEN_OFFLOAD_DIR (default /data/skyrl/frozen-offload; set to 0 to disable and restore the previous pinned-RAM behavior). The masters are immutable for the whole run, so after the one-time write their pages are clean, evictable page cache instead of unreclaimable RAM; re-offloads are free and backloads read from hot cache.
• LoRA models skip offloading the fused DDP buffers entirely: Megatron's param/grad buffers only hold grad-requiring params, which for LoRA is the adapters (a few GB). They stay GPU-resident across offload/backload cycles, which is also what lets an adapter-only weight sync export the adapter without ever backloading the masters (follow-up PR).

Why

Measured on our 2x8xB300 (sm103) cluster training Kimi K2.7-Code (INT4 QAT + LoRA, colocated vLLM engines) through the Tinker API: the trainer's CPU offload held ~1.34TB/node of pinned/anonymous RAM (~1.19TB BF16 masters + optimizer/misc) next to the vLLM engines' own ~1.3TB/node level-1 sleep backup — two full CPU copies of the model on a 2.9TB node, leaving <100GB of margin. Bring-up was a coin flip that ended in NUMA-local OOM kills (CONSTRAINT_MEMORY_POLICY: one socket full while the other had TBs free), or with swap enabled, compress/disk thrash until Ray's keepalive declared the workers dead.

With this change, peak-phase MemAvailable on the same nodes went from 80–380GB (repeated OOM kills across 8 consecutive bring-up attempts) to ~1.5TB, and bring-up passes reliably.

Made with Cursor

[gemini-code-assist]

Code Review

This pull request introduces file-backed mmap CPU offloading for frozen weights in Megatron-based distributed training to reduce RAM usage and prevent OOM issues. It also skips offloading fused buffers when training LoRA adapters. The review feedback highlights three key areas for improvement: a high-severity correctness issue where the hash key for the offload file only uses the parameter name, shape, and dtype (which can lead to silent weight corruption across different runs or checkpoints on the same node); an efficiency issue where calling .tobytes() on the entire tensor duplicates CPU memory allocation; and a usability issue where failures during offloading flood the logs with warnings for every parameter.

[gemini-code-assist]

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

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit c3a70d3. Configure here.}

[avigyabb]

Hey Casper, thanks again for all your contributions! Left a comment regarding some of the naming logic. If you're strapped on time, just lmk and I can take this over the line if you'd like!

[avigyabb]

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.

[avigyabb]

Also, if my understanding is correct, when we first prepare_for_weight_sync we call _ensure_on_gpu which calls backload_to_gpu which loads the entire model from disk to GPU and then when we weight sync, we discard the GPU copy of the frozen weights. For the adapter only sync path, we never needed to bring these onto the GPU to begin with. Maybe fixing this logic can be a follow up PR from our end?

SkyRL/skyrl/backends/skyrl_train/workers/worker_dispatch.py

Line 169 in b0f74ac

self._actor_groups[model].backload_to_gpu(