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

[casper-hansen]

Part of the Kimi K2.6/K2.7 series. Stacked on #2026 (4/n), #2031 (9/n), #2062 (11/n), and #2063 (12/n) — their commits appear in this diff; review the last commit ("Adapter-only colocated weight sync for megatron LoRA") for this PR's own changes.

What

save_weights_for_sampler routes megatron + LoRA + merge_lora=false colocated runs through a new adapter-only path instead of the full-weight NCCL broadcast:

1. Collective adapter export + per-node write via a new worker method save_lora_adapters (the engine-load half is split out of the worker's fused save+load with send_load_request=False). The bridge export reads only the LoRA adapter tensors, which stay GPU-resident through offload ([megatron] 11/n towards Kimi K2.6: file-backed offload of frozen masters so Tinker API LoRA training fits host RAM #2062 exempts the LoRA DDP buffers), so the TB-scale frozen masters remain offloaded the whole time.
2. Offload masters/optimizer only if a preceding forward/optim phase left them resident, so the engines fit on the GPUs; a cold sync skips this entirely.
3. Wake the engines only if they were asleep (engines_asleep is threaded from the backend). The eager post-create sleep is also skipped on this path — it was a pure sleep→wake round-trip.
4. Reset the prefix cache and issue the adapter load from the engine process once the engines are awake.

Supporting changes: _prepare_for_weight_sync releases cached CUDA allocator blocks before engine wake (also benefits the classic path); persist=True sampler saves on this path tar the just-written PEFT adapter files instead of running a merged multi-TB save_hf_model export (the classic persist path keeps the HF export but sleeps the engines first); SKYRL_LORA_SLEEP_LEVEL lets TB-scale LoRA runs opt into level-2 (discard-and-reload) sleeps. The classic full-broadcast path is unchanged for FFT / merged-LoRA. Unit tests cover the cold/hot/wake orderings (test_worker_dispatch.py).

Why

With merge_lora=false, vLLM serves the released base checkpoint and hot-loads the adapter — yet the colocated sync still moved terabytes to deliver gigabytes: backload ~1.19TB of frozen masters to GPU, wake engine weights, broadcast, offload again, wake KV cache.

Measured on our Kimi K2.7-Code runs (2x8xB300, INT4 QAT + LoRA rank 32 through the Tinker API): a sync took 13–80 minutes when it survived at all. Run-start syncs re-faulted TBs of mmap'd frozen weights through a page cache the engine build had just evicted (554GB checkpoint read + ~830GB/node sleep backup), stalling save_weights_for_sampler 15–25 minutes — past the client's deadline, which abandoned and retried, rebuilding the engines and repeating the eviction in a loop that never converged. The eager post-create sleep alone cost 5–6 minutes of ~830GB/node CPU backup (sequential ~50s/worker) for engines that the first sync would immediately wake. And the trainer's freed-but-cached allocator blocks (tens of GB after a forward/optim step) pushed wake_up(tags=["weights"]) over GPU capacity.

With this path, a cold sampler sync exports + loads in seconds with the engines awake the whole time, and the first end-to-end Kimi K2.7 LoRA training run through the Tinker API completed with stable step times.

Made with Cursor

[gemini-code-assist]

Code Review

This pull request introduces significant optimizations and robustness fixes for colocated training and inference, particularly for Megatron LoRA configurations. Key changes include file-backed mmap offloading for frozen parameters to prevent host OOMs, an optimized adapter-only weight sync path that avoids backloading massive base weights, and critical fixes in AdapterStore to handle offloaded grad buffers and prevent state leakage from deleted adapters. The review feedback is highly constructive, pointing out a potential cache collision risk in the offload directory, a memory optimization for writing tensors to disk, a potential crash in non-distributed environments within _is_lora_sync_writer_rank, and duplicated path resolution logic in save_sampler_checkpoint.

[gemini-code-assist]

The current cache key generation in _frozen_offload_file only hashes the parameter name, shape, and dtype. If multiple training runs or different models (with identical layer names and shapes but different base weights) are executed on the same machine/cluster using the default /data/skyrl/frozen-offload directory, they will silently overwrite or reuse each other's cached files. This will lead to silent, catastrophic model corruption where incorrect weights are loaded.

To prevent this, consider incorporating a unique identifier (such as a hash of the model path, a run ID, or a quick checksum/sample of the actual tensor values) into the cache key or the directory path.

[gemini-code-assist]

Using data.contiguous().view(torch.uint8).flatten().cpu().numpy().tobytes() creates a temporary Python bytes object containing the entire tensor's data in memory. For large model parameters, this can cause a significant memory spike.

Using numpy.ndarray.tofile() writes the array's raw data directly to the file stream, avoiding the creation of a temporary bytes object in the Python heap.

Suggested change

	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)
	tmp = f"{path}.tmp{os.getpid()}"
	with open(tmp, "wb") as f:
	data.contiguous().view(torch.uint8).flatten().cpu().numpy().tofile(f)
	os.replace(tmp, path)

[gemini-code-assist]

_is_lora_sync_writer_rank calls torch.distributed.get_world_size() and torch.distributed.all_gather_object() unconditionally. If this code is executed in a non-distributed environment (e.g., during single-GPU testing or mock runs where torch.distributed is not initialized), it will crash.

Add a guard to check if distributed training is initialized, falling back to returning True (since there is only one rank).

        cached = getattr(self, "_lora_sync_writer_cache", None)
        if cached is None:
            if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
                cached = True
            else:
                import socket

                hostnames = [None] * torch.distributed.get_world_size()
                torch.distributed.all_gather_object(hostnames, socket.gethostname())
                cached = hostnames.index(hostnames[torch.distributed.get_rank()]) == torch.distributed.get_rank()
            self._lora_sync_writer_cache = cached
        return cached

[gemini-code-assist]

Currently, save_sampler_checkpoint reconstructs the adapter directory path locally using os.path.basename(model_id). This duplicates the path resolution logic from the worker's _resolve_lora_sync_target and introduces a risk of mismatch if the worker uses a different naming convention (e.g., if model_id contains slashes and is not basename-sanitized on the worker).

Since save_weights_for_sampler already triggers the sync which resolves the correct path, consider modifying save_weights_for_sampler to return the resolved lora_sync_path and use it directly here.

[cursor]

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

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

[cursor]

Persist tars adapters off the driver

Medium Severity

persist=True on the adapter-only path tars lora_sync_path from the backend process. Writers are the first rank on each GPU node, and that path is allowed to be node-local so engines can hot-load locally. The driver often cannot see those files, so sampler checkpoints fail or archive an empty/stale directory.

 

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