[megatron] Skip frozen-master backload for adapter-only LoRA weight sync - #2157
Draft
avigyabb wants to merge 5 commits into
Draft
[megatron] Skip frozen-master backload for adapter-only LoRA weight sync#2157avigyabb wants to merge 5 commits into
avigyabb wants to merge 5 commits into
Conversation
…ers resident On TB-scale colocated models (Kimi K2.7: ~1.19TB BF16 masters) the CPU offload held ~1.34TB/node of pinned/anonymous RAM next to the vLLM engines' own ~1.3TB level-1 sleep backup — two full CPU copies of the model per 2.9TB node, leaving <100GB of margin. Every bring-up was a coin flip that ended in NUMA-local OOM kills (CONSTRAINT_MEMORY_POLICY, one socket full with TBs free on the other) or, with swap present, compress/disk thrash until Ray's keepalive declared the workers dead. Two changes to the offload layer: - Frozen (requires_grad=False, non-adapter) params now offload to file-backed mmap tensors, written once per rank on first offload (SKYRL_FROZEN_OFFLOAD_DIR, default /data/skyrl/frozen-offload, "0" disables and restores pinned RAM). The masters are immutable for the whole run, so their pages become clean, evictable page cache instead of unreclaimable RAM; re-offloads are free and backloads read 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). Keeping them GPU-resident lets the adapter-only weight sync export without ever backloading the masters. Measured on ghost002+ghost013 (2x8xB300): peak-phase MemAvailable went from 80-380GB (repeated OOM kills across 8 consecutive bring-up attempts) to ~1.5TB, with bring-up reliably passing. Co-authored-by: Cursor <cursoragent@cursor.com>
Replace the per-chunk _chunk_has_lora_adapters name heuristic (a third string-matching variant of LoRA detection) with the authoritative is_lora flag the strategy already receives from the worker. Also avoids walking named_parameters() twice per sleep/wake cycle on non-LoRA runs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
Under colocate_all, _prepare_for_weight_sync unconditionally backloaded the full policy model to GPU before every sync. For adapter-only LoRA sync (megatron, merge_lora=False) the broadcast exports straight from the fused adapter buffers, which stay GPU-resident while the model is offloaded — so the sync round-tripped the TB-scale frozen base weights disk -> GPU -> discarded for nothing, at least once per run (the first save_weights_for_sampler after build) and on any fully-async sync that overlaps an offloaded phase. Gate need_model on the sync path: adapter-only LoRA skips the model backload while keeping the offload-others side effect and the optimizer offload. merge_lora=True, full-parameter runs, and FSDP keep the backload (their sync extracts full weights on GPU). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Under
colocate_all,_prepare_for_weight_syncunconditionally backloads the full policy model to GPU before every weight sync. For adapter-only LoRA sync (megatron,merge_lora=False), the broadcast exports straight from the fused adapter buffers — which stay GPU-resident while the model is offloaded (#2062) — so the sync round-trips the TB-scale frozen base weights disk → GPU → discarded for nothing.This provably happens at least once per colocated run: models are marked offloaded right after build, so the first
save_weights_for_samplerpays the full backload just to broadcast a rank-32 adapter. It can recur in fully-async flows whenever a sync overlaps a phase where the model is offloaded. (Steady-state synchronous syncs run right after training with the model already on GPU, so they were unaffected.)How
_prepare_for_weight_syncnow gatesneed_modelon a new_weight_sync_needs_model_on_gpu()check:merge_lora=False→need_model=False._ensure_on_gpustill runs, so the offload-others side effect and optimizer offload are preserved — only the model backload is skipped.merge_lora=True, full-parameter runs, FSDP →need_model=Trueunchanged (their sync extracts full weights on GPU; FSDP has no resident-adapter guarantee, so it stays conservative)._finish_weight_syncis intentionally unchanged: offloading an already-offloaded model is a cheap worker-side no-op (empty-storage params are skipped) and keeps_gpu_stateconsistent.Notes
save_weights_for_sampler's non-colocated branch.Tests
Added
TestPrepareForWeightSync(CPU, no Ray needed) covering: adapter-only LoRA skips the backload;merge_lora=True, FFT, and FSDP+LoRA keep it; non-colocated is untouched. All 12 tests in the file pass.🤖 Generated with Claude Code