feat: dist_muon offloading in megatron - #2739
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
--stream-optimizer-state-to-disk binds to DistributedOptimizer, which Muon never uses: set_default_megatron_args clears use_distributed_optimizer for every non-Adam optimizer and Muon is built through LayerWiseDistributedOptimizer. The two asserts guarding the flag were unreachable for Muon by construction. Rather than port the DistOpt-bound store to a different optimizer -- reimplementing its bucket planner and main->param write-back, and re-deriving the master-weight residency protocol -- reuse the offloader Megatron already runs for Muon. Float16OptimizerWithFloat16Params dispatches its step through an _optimizer_state_offloader slot, and ChunkedOptimizerStateOffloader fills that slot today under --chunked-optimizer-state-offload: chunked restore, assert_master_weights_resident, prefetch hooks and checkpoint sync already work with Muon. Its only tie to host memory is one allocator, _new_cpu_buffer. So subclass it and override just that, returning tensors backed by files on node-local NVMe. copy_ moves data to and from an mmap tensor exactly as for a pinned one, and is_pinned() returns False so the inherited code takes its synchronous path. File-backed pages land in the cgroup's reclaimable `file` rather than in `anon`, so the kernel can evict them under pressure instead of OOM-killing the run. Measured on Qwen3.5-35B-A3B, 4x GB300, TP2/EP4, at --optimizer-state-offload-fraction 1.0 -- the setting at which the pinned-CPU path dies at step 6 on a host peak of 908GB against a 919GB limit -- the disk backend ran to completion with anon flat at 86GB and grad_norm 0.1166, inside the 0.088-0.170 band the same model spans without offload. Gates are split by optimizer family rather than removed: Adam keeps requiring the distributed optimizer, Muon requires the dist_ prefix plus the chunk knobs the offloader is driven by. The flag's help described only the Adam path and is updated to cover both. --optimizer-cpu-offload stays Adam-only and is not a gap for Muon: HybridDeviceOptimizer substitutes the tensors its inner optimizer steps on, which collides with LayerWise resharding group["params"] and then wrapping each child to create fp32 masters.
1396e57 to
1a3715a
Compare
Zhichenzzz
changed the title
Mirrors test_qwen3_4B_offload_disk_stream.py with --optimizer dist_muon. dist_muon is the only route to the layer-wise distributed optimizer that the chunked offloader hangs off, so plain --optimizer muon would not exercise this path. Finishing the run proves little on its own: if the rebind misses the consuming module or the plugin import is swallowed, training falls back to pinned host memory and still converges. The offloader now logs one line per step, letting the test assert all four ranks really stepped against file-backed state rather than merely that install() ran. At this rollout size every sample truncates, so the registered metric gates sit at zero and cannot catch a regression by themselves; the disk-backed-step assertion is what guards the feature. Verified on 4x GB300: 4/4 ranks armed, 4/4 file-backed, rc=0.
The Muon path was a standalone module beside nvme_stream.py, duplicating its directory scheme and reading nothing from it. It now lives in that file: setup_muon_state_on_disk sits next to setup_optimizer_state_streaming, takes args the same way, and shares _state_dir_root and _purge_rank_dir, so both optimizers land under the same rank-scoped layout and both clear stale state the same way. Two entry points remain because the timing differs -- Muon has to be installed before the optimizer is built, since that is when the offloader is constructed, while the Adam store needs the built optimizer. Passing args instead of a prebuilt path lets model.py drop its path join and the os import that existed only for it. The unit test moves to test_nvme_stream.py, mirroring the module it covers. Also trims comments that restated the code or repeated the PR description. What stays is the two notes that stop a plausible edit from breaking things silently: why both modules are rebound, and why the argument gate keys off the dist_ prefix rather than use_layer_wise_distributed_optimizer.
mkstemp gets uniqueness from the stdlib, so the module-level counter and the pid-in-the-filename scheme both go away along with the itertools import. The counter existed only because two threads sharing a pid could otherwise collide on a name. set_(storage, 0, shape) does in one call what set_(storage) plus a view did in two, and CPU is already the default device. Folding the zero-length guard into the nbytes line keeps it next to the arithmetic it protects. Checked against the previous implementation on torch 2.13 across shapes and dtypes, including a zero-element tensor: same shape, dtype, device, pinned state, bit-exact round trip, no file left behind.
yueming-yuan
left a comment
yueming-yuan
left a comment
There was a problem hiding this comment.
overall LGTM, just add small comments
There was a problem hiding this comment.
should we default --optimizer-state-offload-chunk-size-mb to be a non-0 value, or add a warning? The current default is 0, and in megatron it seems 0 -> non-streaming, so the peak memory cannot be saved. The user might now be aware of this
There was a problem hiding this comment.
Megatron has this warning already
| state_dir = rank_dir | ||
| _disk_bytes = 0 | ||
|
|
||
| def _new_cpu_buffer(self, tensor: torch.Tensor) -> torch.Tensor: # type: ignore[override] |
There was a problem hiding this comment.
codex comment, looks reasonable:
Preserve existing mmap state during checkpoint adoption — miles_plugins/optimizers/nvme_stream.py:503
Megatron’s synchronize_for_checkpoint() calls adopt_cpu_optimizer_state(), which treats every non-pinned CPU tensor as foreign and reallocates it through _new_cpu_buffer. Because these mmap tensors intentionally report is_pinned() == False, every checkpoint copies the entire offloaded optimizer state into new mappings and permanently inflates _disk_bytes. For the hundreds-of-GB workloads targeted here, this adds a full-state copy to every checkpoint. Track already-managed mappings or override the adoption behavior.
| nbytes = max(tensor.numel() * tensor.element_size(), 1) | ||
| fd, path = tempfile.mkstemp(dir=directory, suffix=".bin") | ||
| os.close(fd) | ||
| storage = torch.UntypedStorage.from_file(path, shared=True, nbytes=nbytes) |
There was a problem hiding this comment.
codex + claude comment, looks reasonable
UntypedStorage.from_file(..., shared=True) expands the file with ftruncate but does not reserve physical blocks. If optimizer state exceeds available NVMe capacity, mapping succeeds and later writes can terminate the process with SIGBUS instead of producing an actionable ENOSPC error. Preallocate with posix_fallocate, as the existing _allocate_file path already does, before mapping and unlinking the file.
|
@claude review |
There was a problem hiding this comment.
Beyond the inline findings, I also looked at the temp-file handling in _disk_backed_like (nvme_stream.py:102-103), which a prior reviewer flagged on this same line: os.unlink(path) only runs after torch.UntypedStorage.from_file succeeds, so a failure there (e.g. ENOSPC or a permission error) would leak the mkstemp'd file — but this is a narrow error-path resource leak, not a correctness bug in the normal run path.
Extended reasoning...
This run confirmed one new finding (the dist_muon branch logging Adam-only chunk-size/moment-dtype details that it never reads) which is being posted as an inline comment. It also re-examined the mkstemp/os.unlink ordering in _disk_backed_like — the same line a human reviewer (yueming-yuan) left an unresolved inline comment on in the latest review round, with no subsequent commit visibly addressing it. Investigation confirms the leak is real but only triggers on the from_file error path (disk full, permission failure, etc.), not during normal operation, so it doesn't rise to a blocking finding on its own. Given the outstanding third-party review thread on that exact line has not been addressed by a code change since it was posted, and a new confirmed finding exists, approval is not appropriate here; a brief informational note is warranted since the ruled-out item directly overlaps with the open human review thread.
Additional findings (outside the current diff — GitHub can't attach inline comments there):
-
🟡
miles/utils/arguments.py— miles_validate_args logs Adam-specific stream details (chunk size, moment dtype) unconditionally, but the new dist_muon branch never uses either field.Extended reasoning...
A user running --optimizer dist_muon --stream-optimizer-state-to-disk --stream-optimizer-state-moment-dtype bf16 sees 'Streaming optimizer state to disk, ... moments=bf16' at startup. In reality setup_muon_state_on_disk/_disk_backed_like (miles_plugins/optimizers/nvme_stream.py) never reads args.stream_optimizer_state_moment_dtype or args.offload_train_disk_chunk_mb - the disk-backed buffer always matches the tensor's native dtype. The operator believes moments were downcast and plans disk capacity accordingly, but the muon path silently uses full-precision buffers, risking under-provisioned disk / larger-than-expected NVMe traffic. Fix: gate this log's content on the _muon_disk_state branch so it reports the fields that actually apply (chunked-optimizer-state-offload / optimizer-state-offload-fraction) for Muon.
Verification: Severity: nit. arguments.py:3422-3424 logs, inside the shared
if args.stream_optimizer_state_to_disk:block,Streaming optimizer state to disk, dir=..., chunk={args.offload_train_disk_chunk_mb}MB, moments={args.stream_optimizer_state_moment_dtype}with no branch on optimizer family. The new dist_muon path is gated in the same block (arguments.py:3379-3391), so this line fires for it. The…
Three things Yueming's review surfaced. adopt_cpu_optimizer_state reallocates every non-pinned CPU tensor it finds in optimizer.state, and our mappings never report pinned, so each checkpoint was copying the whole offloaded state into fresh mappings. The stock pinned path never hits this. Buffers now carry a marker and the allocator hands them back unchanged, which the offloader's own copy-onto-itself then treats as a no-op. Checked against the real base class: identity preserved, byte count flat, and a genuinely foreign CPU tensor is still relocated into a managed buffer. from_file sizes its file with ftruncate, which leaves it sparse: a full filesystem surfaces as SIGBUS at first touch rather than ENOSPC at setup. The fallocate-with-fallback that _allocate_file already had is now a shared _reserve that both paths use. Costs nothing measurable -- 1GB reserves in 1ms. Checkpointing fsyncs its own files while the writeback queue carries our dirty pages, so synchronize_for_checkpoint now msyncs our mappings first. This makes that cost ours and attributable; msync over a clean mapping is free, so the repeat is cheap. I could not measure how much it helps, and do not claim it fixes anything on its own. The chunk-size knob needs no code: Megatron already warns at 0, under a condition strictly broader than what this flag requires. The help text now points at it rather than duplicating the warning. Verified with checkpoints enabled on 4x H200, three saves, rc=0: the offloaded total goes 3.38 -> 6.77 GB at the first save, where the fp32 masters are offloaded for the first time, and then holds at 6.77 across the second. CI does not cover that path -- three saves of this model take 19 minutes, past the 1800s per-file limit -- so it stays a manual check.
2 participants
--stream-optimizer-state-to-diskcurrently only works with Adam. This makes it work with Muon too.Why it did not work
The flag binds to
DistributedOptimizer, which Muon never uses:set_default_megatron_argsclearsuse_distributed_optimizerfor every non-Adam optimizer, and Muon is built throughLayerWiseDistributedOptimizer. The two asserts guarding the flag were unreachable for Muon by construction — one required the distributed optimizer that had just been cleared, the other requiredoptimizer == "adam".Approach
Rather than port the DistOpt-bound store to a different optimizer — which would mean reimplementing its bucket planner and its main→param write-back, and re-deriving the master-weight residency protocol — this reuses the offloader Megatron already runs for Muon.
Float16OptimizerWithFloat16Paramsdispatches its step through an_optimizer_state_offloaderslot, andChunkedOptimizerStateOffloaderfills that slot today under--chunked-optimizer-state-offload: chunked restore, theassert_master_weights_residentprotocol and the prefetch hooks already work with Muon. Checkpointing needed two fixes, below. Its only tie to host memory is one allocator,_new_cpu_buffer, returning pinned CPU tensors.So
setup_muon_state_on_disksubclasses that offloader and overrides just the allocator, returning tensors backed by files on node-local NVMe. It lives innvme_stream.pybeside the Adam store and shares its rank-scoped directory layout and stale-state purge; the two entry points stay separate only because Muon must be installed before the optimizer is built, which is when the offloader is constructed.copy_moves data to and from an mmap tensor exactly as for a pinned one, andis_pinned()returns False so the inherited code takes its synchronous-copy path. The other two CPU allocations in that file were checked and left alone: one targets the parameter's own device, the other is a bounded per-chunk staging buffer.The flag's gates are split by optimizer family rather than removed — Adam keeps requiring the distributed optimizer, Muon requires the
dist_prefix plus the chunk knobs the offloader is driven by.Why it matters
File-backed pages land in the cgroup's reclaimable
filerather than inanon, so the kernel can evict them under pressure instead of OOM-killing the run.Measured on Qwen3.5-35B-A3B, 4x GB300, TP2/EP4, 8k responses, at
--optimizer-state-offload-fraction 1.0— the setting at which the pinned-CPU path dies:--chunked-optimizer-state-offloadalone)rc=0With the disk backend the non-reclaimable
anonstayed flat at 86GB across 4 optimizer steps whilefilegrew to 619GB, andgrad_normwas 0.1166 — inside the 0.088–0.170 band the same model spans without offload.Testing
tests/fast/optimizers/test_nvme_stream.py, 8 cases, no GPU needed. They cover the silent-failure path: if the allocator override stops returning file-backed storage, the offloader keeps working against pinned host memory while the log line still claims the disk backend is installed. Each assertion was checked against a deliberately broken implementation to confirm it fails.tests/e2e/megatron/test_qwen3_4B_muon_offload_disk.pymirrors the Adam testtest_qwen3_4B_offload_disk_stream.pywith--optimizer dist_muon, which is the only route to the layer-wise distributed optimizer the offloader hangs off. Completing the run proves little on its own — a rebind that misses the consuming module falls back to pinned host memory and still converges — so the offloader now logs one line per step and the test asserts all four ranks stepped against file-backed state rather than merely thatinstall()ran. Verified on 4x GB300: 4/4 ranks armed, 4/4 file-backed,rc=0in 292s.At that rollout size every sample truncates, so the registered metric gates sit at zero and cannot catch a regression by themselves — the same holds for the Adam test this mirrors. The disk-backed-step assertion is what guards the feature.
Checkpointing
Review found two ways this went wrong against
adopt_cpu_optimizer_state, whichreallocates every non-pinned CPU tensor in
optimizer.state. Ours never report pinned, soevery checkpoint was copying the whole offloaded state into fresh mappings; the stock
pinned path never reaches that branch. Buffers now carry a marker and the allocator returns
them unchanged. Separately,
synchronize_for_checkpointmsyncs our mappings before thewriter fsyncs its own files, so that cost is ours and attributable rather than left in the
kernel's writeback queue.
Verified on 4x H200 with three saves,
rc=0. The offloaded total goes 3.38 -> 6.77 GB atthe first save, where the fp32 masters are offloaded for the first time, then holds at 6.77
across the second -- reallocation would have doubled it again. All four ranks agree.
Two things worth knowing. Saving is measurably slower with the disk backend: 381s per save
against 229s for the same run on pinned host memory, an A/B on one box, so roughly +66%
from reading the state back through the mappings. And CI does not cover this: three saves
of this model take 19 minutes, past
run_suite.py's 1800s per-file limit, so the e2e teststays checkpoint-free and the above is a manual check.
Notes
--optimizer-cpu-offloadis a separate flag and remains Adam-only. It is not a gap for Muon:HybridDeviceOptimizerwraps an optimizer and substitutes the tensors it steps on, which collides withLayerWiseDistributedOptimizerreshardinggroup["params"]and then wrapping each child inFloat16OptimizerWithFloat16Paramsto create fp32 masters — two layers substituting the same tensors. Megatron's own answer is the_optimizer_state_offloaderslot, which is what this change builds on.