Route non-Muon params through DistributedOptimizer - #4771
Merged
deepakn94 merged 10 commits intoMay 21, 2026
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
Author
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
deepakn94
force-pushed
the
dnarayanan/distopt_for_non_muon_params
branch
from
May 13, 2026 18:00
694aa11 to
34cd925
Compare
Contributor
Author
|
/claude review |
Contributor
Author
|
/ok to test 34cd925 |
deepakn94
force-pushed
the
dnarayanan/distopt_for_non_muon_params
branch
from
May 13, 2026 20:36
34cd925 to
d52da35
Compare
Contributor
Author
|
/ok to test d52da35 |
deepakn94
commented
May 15, 2026
Contributor
Author
|
/claude review |
FDecaYed
approved these changes
May 19, 2026
FDecaYed
left a comment
FDecaYed
left a comment
Contributor
There was a problem hiding this comment.
Other part looks ok but I have some concerns on the tagging layerwise managed param and DDP wrapping part. I think we'll need another PR to refactor it. current approach of hard coding muon split seems not correct in longer term
My understanding is optimizer setup needs to happen after DDP wrapping due to it changing param/grad buffer. But on the other hand, we want a "optimizer managed layout" that can make comm optimal. If done naturally, this means creating optimizer first.
So overall I think we should a single source of code dealing with adam/muon/soap/etc split, as well as producing layout. all run in the beginning and same information is feed into both DDP and optimizers init.
kvareddy
approved these changes
May 19, 2026
jaredcasper
approved these changes
May 19, 2026
lmcafee-nvidia
approved these changes
May 19, 2026
yaoyu-33
approved these changes
May 19, 2026
yashaswikarnati
approved these changes
May 19, 2026
Splits the DDP buffer for layer-wise optimizers so Muon-managed matrices stay in a shard-aligned LayerWiseDistributedOptimizer buffer, while everything else (embeddings, biases, LayerNorm weights) lands in a separate byte-level DistributedOptimizer buffer. The top-level optimizer becomes ChainedOptimizer([LayerWise(Muon), DistOpt(Adam)]). This drops the massive (dp_size - 1) * embedding_numel padding overhead that the all-LayerWise layout pays when there are unique-large params (typical for embeddings). What changed: - BufferKey gains an is_layer_wise_distributed_optimizer field; group_params_for_buffers reads it from a param attribute set by tag_params_for_buffer_routing before DDP wrap. - LayerWiseDistributedOptimizer.compute_full_param_layout dispatches per-buffer: LayerWise's shard-aligned layout for its buffers, DistributedOptimizer's byte-level layout for the rest. - LayerWiseDistributedOptimizer._shard_params_from_layout and set_bucket_layerwise_params_list filter to LayerWise-managed buckets only via the _bucket_is_layer_wise_managed helper. - _get_megatron_emerging_optimizer constructs a real DistributedOptimizer for non-emerging-optimizer (Adam) groups when the layer-wise + layout path is active, chained at the top level alongside LayerWise. - partition_buckets uniqueness check relaxed to only enforce on uint8 buffers so two bf16 buffers can coexist in one DDP wrapper. - setup_process_groups_for_optimizer takes use_gloo_process_groups threaded through. Guard assertions for not-yet-supported configurations: the split path fails fast on num_distributed_optimizer_instances > 1 (instance_id is hardcoded to 0) and on expert-parallel non-Muon param groups (would need a second DistOpt with the expert-DP group). Both fall back cleanly when use_layer_wise_param_layout=False (the default in production). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
When LayerWiseDistributedOptimizer and DistributedOptimizer are chained (layer-wise + DistOpt split path), each was calling ``model_chunk.start_param_sync()`` which iterates through *all* DDP bucket groups, so the same buckets ended up being all-gathered twice per step. Adds ``start_param_sync_for_bucket_group_subset`` as a method on both LayerWise and DistOpt — each implementation iterates only its own bucket groups: - LayerWise walks the DDP bucket groups and filters by ``_bucket_is_layer_wise_managed`` (default_for_untagged=True, since legacy LayerWise owns every bucket). - DistOpt filters with ``default_for_untagged=False`` so untagged buckets (pure non-LayerWise paths) still count as DistOpt-managed. To preserve DDP's FP8 / MXFP8 post-all-gather work, the per-bucket-group sync goes through a new ``DistributedDataParallel._start_bucket_group_param_sync`` helper, which both optimizers call. ``DDP.start_param_sync`` is refactored to delegate to the same helper so the FP8 logic lives in one place. The LayerWise param sync now lives on ``step_with_ready_grads`` (not ``step``) so it also runs when LayerWise is a child of an outer ``ChainedOptimizer`` (which calls ``step_with_ready_grads`` directly and bypasses ``step``). Also default ``use_layer_wise_param_layout=True`` at the production ``get_model`` call site so live layer-wise training runs the shard-aligned + DistOpt split path. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The previous size-matching algorithm grouped equal-sized parameters from arbitrary backprop positions into the same bucket, producing buckets whose backprop range was non-contiguous. _ParamAndGradBuffer.__init__'s backprop-order iteration then saw non-monotonic bucket ids and tripped its bucket_id == cur + 1 invariant. Reusing the buffer's monotonic emit order also stalled overlap_grad_reduce because a bucket only dispatches reduce-scatter once all its params have grads, and a non-contiguous bucket only completes well past the end of its backward segment. The new layout walks params in backprop order, accumulates them into a contiguous chunk until the chunk reaches bucket_size, then bin-packs the chunk into dp_size shards via greedy LPT (largest first, assign to least-loaded shard). Each bucket therefore spans a single contiguous backprop range -- restoring the assertion and the overlap_grad_reduce dispatch timing -- while LPT keeps shard sizes close to balanced and zero-padded when params_per_layer * num_layers divides dp_size. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
In the LayerWise + DistributedOptimizer split path, non-Muon params are managed by a real DistributedOptimizer instance even though the arg parser flips ``args.use_distributed_optimizer`` off (line 1462). With the previous condition, ``_build_sharded_state_dict_metadata`` skipped populating ``distrib_optim_sharding_type``, so when the DistOpt sub-optimizer's ``sharded_state_dict`` ran it fell through to the deprecated ``fully_sharded_model_space`` default, which sets ``flattened_range`` on the produced ShardedTensors. After commit 5ab481c (Dec 2025) removed flattened_range support from ``ShardedTensor.validate_metadata_integrity``, that path raises ``CheckpointingException("ShardedTensor.flattened_range is not supported.")``. Treat ``use_layer_wise_distributed_optimizer`` as equivalent to ``use_distributed_optimizer`` for metadata purposes, so the DistOpt sub-optimizer gets ``dp_reshardable`` (or ``fully_reshardable`` / ``fsdp_dtensor`` as configured) instead of falling through to the broken default. Tests: - CPU-only ``TestBuildShardedStateDictMetadata`` pins the truth table of ``_build_sharded_state_dict_metadata``; the new ``use_layer_wise_distributed_optimizer``-only rows are the regression guard. - End-to-end ``test_optimizer_common_state_dict_hybrid`` runs a full ``save_checkpoint``/``load_checkpoint`` roundtrip through the hybrid ChainedOptimizer (LayerWise wrapping Muon + a sibling DistOpt). Test util adjustments to support this: ``setup_model_and_optimizer`` now threads ``use_gloo_process_groups`` through to ``get_megatron_optimizer`` (the hybrid DistOpt setup uses an explicit ``pg_collection`` which is incompatible with the Gloo-on default), and the optimizer state-seeding loop recurses into nested ChainedOptimizers and skips optimizers without ``init_state_fn``. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- Drop ``use_gloo_process_groups`` from ``_get_megatron_emerging_optimizer``. The hybrid LayerWise + DistOpt path always supplies an explicit ``pg_collection`` to ``setup_process_groups_for_optimizer``, which unconditionally rejects Gloo groups in that mode, so the parameter only ever held one legal value. Hardcode ``use_gloo_process_groups=False`` at the call site and drop the matching plumbing from the ``setup_model_and_optimizer`` test util. - Collapse ``dtype_to_buffer_map`` in ``partition_buckets`` to a single ``fp8_buffer`` lookup. The map was only consulted for the unique fp8 (uint8) buffer; storing every non-uint8 buffer under its dtype was unused and obscured the uniqueness invariant. Direct scan + assert keeps the invariant and drops the dict. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The LPT-in-contiguous-backprop layout finalised each bucket as soon as the chunk's raw numel reached ``bucket_size``. At high dp this left each bucket with too few params to fill all ``dp_size`` shards, so the LPT pack ended with most shards empty and the bucket's largest param forcing ``dp_size * max_param`` of padding. On the 8B Llama-style config at dp=32 with ``--ddp-num-buckets 6``, every Muon-side bucket held ~21 params for 32 shards and burned ~414M elements of padding each — total 1.98B of padding for 1.74B real params (113% overhead), which manifested as a ~12 GB jump in rank-0 ``allocated`` versus the legacy LayerWise path. Extend the finalisation threshold to ``max(bucket_size, dp_size * chunk_max_param * PADDING_FLOOR)`` so a chunk grows until its raw numel can absorb most of the on-buffer cost imposed by its largest param. With ``PADDING_FLOOR = 0.9`` each bucket's overhead stays at or below ~11%. On the same 8B dp=32 config this collapses 6 buckets into 3 well-balanced ones: overhead drops 113% -> 21%, padding drops 1.98B -> 369M, and rank-0 ``allocated`` drops 24.4 GB -> 14.8 GB. 4n/8n behaviour is unchanged because the floor only fires when ``dp_size * max_param > bucket_size``. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…mode The model-chunk split into ``[[chunks[0]], chunks[1:]]`` and the matching per-chunk ``overlap_param_gather_with_optimizer_step`` attribute are only applied by ``get_megatron_optimizer`` outside the emerging-optimizer branch. ``_get_megatron_emerging_optimizer`` returns before that split, so the first-chunk post-optimizer-step param-gather dispatch (``ChainedOptimizer.step_with_ready_grads`` line 1289) never fires and the flag is silently a no-op in Muon mode. Fail fast at optimizer construction so the misconfiguration is obvious instead of an unexplained throughput loss. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The old ``is_layer_wise_distributed_optimizer`` field/attribute read as if the boolean answered ``is this a distributed optimizer``, when it actually answers ``is this buffer managed by the LayerWise optimizer``. Rename consistently across BufferKey, the ``torch.nn.Parameter`` attribute set by ``tag_params_for_buffer_routing``, and the ``is_layer_wise_managed_param`` / ``_bucket_is_layer_wise_managed`` helpers. Pure rename, no behaviour change. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/26130597558 |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/26138083006 |
The precomputed LayerWise param layout (and the DistOpt-for-non-Muon routing it enables) changes bf16 reduction ordering versus the legacy LayerWise ping-pong path, so loss curves differ bit-for-bit between the two. Add a ``--no-use-layer-wise-param-layout`` opt-out so existing runs that pinned golden values against the legacy path can keep their numerics; the precomputed layout stays the default for new runs. The two failing functional tests (``gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon`` and its ``_1node`` sibling) gain the opt-out flag so their existing golden values continue to compare bit-for-bit. New ``..._param_layout`` and ``..._param_layout_1node`` variants exercise the new default with golden values captured from the run that generated the divergence. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Member
|
/ok to test b6309df |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/26207656966 |
This was referenced May 28, 2026
This was referenced Jun 5, 2026
Closed
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
9 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.
Summary
Splits the DDP buffer for layer-wise optimizers so Muon-managed matrices stay in a shard-aligned
LayerWiseDistributedOptimizerbuffer, while everything else (embeddings, biases, LayerNorm weights) lands in a separate byte-levelDistributedOptimizerbuffer. The top-level optimizer becomesChainedOptimizer([LayerWise(Muon), DistOpt(Adam)]). This drops the massive(dp_size - 1) × embedding_numelpadding overhead that the all-LayerWise layout pays when there are unique-large params (typical for embeddings).What changed
BufferKeygains anis_layer_wise_distributed_optimizerfield;group_params_for_buffersreads it from a param attribute set bytag_params_for_buffer_routingbefore DDP wrap.LayerWiseDistributedOptimizer.compute_full_param_layoutdispatches per-buffer: LayerWise's shard-aligned layout for its buffers,DistributedOptimizer's byte-level layout for the rest.LayerWiseDistributedOptimizer._shard_params_from_layoutandset_bucket_layerwise_params_listfilter to LayerWise-managed buckets only (_bucket_is_layer_wise_managedhelper)._get_megatron_emerging_optimizerconstructs a realDistributedOptimizerfor non-emerging-optimizer (Adam) groups when the layer-wise + layout path is active, chained at the top level alongside LayerWise.partition_bucketsuniqueness check relaxed to only enforce on uint8 buffers (so two bf16 buffers can coexist in one DDP wrapper).start_param_sync_for_bucket_group_subsethelper: when LayerWise and DistOpt are chained, each callsstart_param_syncon only its own bucket groups (filtered via a predicate) so the same buckets aren't all-gathered twice per step.bucket_id == cur + 1invariant in_ParamAndGradBuffer.__init__and preservingoverlap_grad_reducedispatch timing.threshold = max(bucket_size, dp_size * chunk_max_param * 0.9)so high-dp configs don't emit narrow buckets where the chunk's largest param dominates shard size — keeps per-bucket padding overhead under ~11%._build_sharded_state_dict_metadatatreatsuse_layer_wise_distributed_optimizeras equivalent touse_distributed_optimizerfor sharding-type purposes, fixing the deprecatedfully_sharded_model_spacefallback that crashed checkpoint save post-Integrate LayerWiseDistributedOptimizer with DDP buffer infrastructure #4509.Convergence
Performance
8B Llama-style (TP=4 PP=1, mb=1, seq=8192),
--ddp-num-buckets 6, sweep over DP ∈ {8, 16, 32} × GA ∈ {1, 2, 4, 8}. Per-cell numbers are medians across the 4 GA points at that DP, hybrid minus legacy:Take-away: hybrid is throughput-neutral to slightly-better than legacy at every DP.
Guard assertions for not-yet-supported configurations
The split path fails fast (clear error message) on:
num_distributed_optimizer_instances > 1—distributed_optimizer_instance_idis hardcoded to 0 in this path; needs proper computation to lift.Both fall back cleanly when
use_layer_wise_param_layout=False(the legacy LayerWise ping-pong path).Test plan
test_optimizer_common_state_dict_hybridintests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py).🤖 Generated with Claude Code