[dev] Add experimental decoupled compact LayerWise DDP layout for Muon by Wohox · Pull Request #5388 · NVIDIA/Megatron-LM

Wohox · 2026-06-17T05:01:45Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do ?

dev counterpart of #5391 (which targets main). Same change, implemented on top of dev. See #5391 for the main-branch version.

Add an experimental compact, decoupled LayerWise DDP buffer layout for the Muon (layer-wise) distributed optimizer that removes the persistent dp_size * max(shard_load) padding from the long-lived param/grad buffers.

There is no new flag: this reuses the existing --no-use-layer-wise-param-layout. With a Muon layer-wise distributed optimizer under --use-distributed-optimizer, disabling the shard-aligned padded LayerWise layout (use_layer_wise_param_layout=False) selects the compact decoupled path. use_layer_wise_param_layout is mirrored onto both DistributedDataParallelConfig and OptimizerConfig (auto-populated by field name, default True = padded layout), so the configs read it directly — no derived switch.

In this mode use_distributed_optimizer becomes a per-buffer property:

LayerWise-managed (Muon 2D matrix) buffers use a compact no-padding DDP layout and locally disable DistributedOptimizer semantics: all-reduce gradients, legacy whole-param ping-pong ownership, and allgather_params param sync.
Sibling non-LayerWise buffers (embeddings, biases, layernorm) keep the standard byte-level DistributedOptimizer layout.

The effective flag is computed per _ParamAndGradBuffer / _ParamAndGradBucketGroup. partition_buckets splits a force-single bucket group (disable_bucketing / non-first VPP chunks) by the effective per-bucket use_distributed_optimizer, so Muon (all-reduce) and sibling (reduce-scatter) buckets never share a group; when all buckets agree this collapses to a single group, identical to the prior behavior.

get_model and wrap_model_chunks_with_ddp share a single if use_layer_wise_distributed_optimizer: branch — both the padded and compact cases force ddp_config.use_distributed_optimizer=True, tag params for buffer routing, and compute the LayerWise full_param_layout. The padded-vs-compact decision lives entirely in compute_full_param_layout / _ParamAndGradBuffer, which read ddp_config.use_layer_wise_param_layout.

Compatibility

The default padded LayerWise layout is unchanged (use_layer_wise_param_layout=True).
Blockwise/MXFP8 compute with fp8_param_gather=False (params persist in bf16) is supported.
FP8/FP4 parameter gather is rejected at arg-validation for any layer-wise distributed optimizer: it requires DistributedOptimizer param buffers that the layer-wise path does not provide.
The compact path requires num_distributed_optimizer_instances == 1 (the non-DistOpt Muon buffers only all-reduce within a single optimizer instance).

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

copy-pr-bot · 2026-06-17T05:01:48Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Wohox · 2026-06-17T09:52:26Z

/ok to test f615b2d

Wohox · 2026-06-23T01:57:12Z

/ok to test e965c83

Wohox · 2026-06-23T03:39:22Z

/ok to test 8c128c8

…uffer distopt Two changes from PR NVIDIA#5388 review (FDecaYed): 1. Make the compact decoupled LayerWise layout the DEFAULT. Replace the `--no-use-layer-wise-param-layout` (store_false) flag with `--use-layer-wise-param-layout` (store_true, default False), and flip the DistributedDataParallelConfig.use_layer_wise_param_layout default to False. The padded LayerWise layout is now opt-in (its dp_size*max(shard_load) padding is a large memory overhead and the reduce-scatter perf gain is not always significant). validate_args message updated to the new flag name. 2. Drop the separate `use_distributed_optimizer` argument threaded into _ParamAndGradBucketGroup. Instead, each _ParamAndGradBuffer bakes its per-buffer DistOpt decision into its own ddp_config ONCE via dataclasses.replace (a LayerWise/Muon buffer on the compact decoupled layout disables DistributedOptimizer; siblings keep the model-level setting). buffer.ddp_config is then the single source of truth and bucket groups simply inherit it in partition_buckets -- no per-site threading or per-call replace. Behavior-preserving. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: pingtianl <pingtianl@nvidia.com>

Wohox · 2026-06-29T03:29:03Z

/ok to test 505f3ed

Flip OptimizerConfig.use_layer_wise_param_layout from True to False so the compact decoupled LayerWise DDP layout is the default, and update the docstring so the (default) annotation tracks the False branch. Signed-off-by: Pingtian Li <pingtianl@nvidia.com>

FDecaYed

LGTM

Wohox · 2026-06-29T07:18:02Z

/ok to test 667002c

svcnvidia-nemo-ci · 2026-06-29T08:28:52Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28358870884

Wohox · 2026-06-30T12:26:22Z

/ok to test ccf1d1b

…layout) Regenerate golden_values_dev_dgx_gb200.json for gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon and its _1node variant. Removing --no-use-layer-wise-param-layout switches these tests from the compact decoupled LayerWise layout back to the default padded LayerWise layout, which changes the deterministic trajectory. Values regenerated in the GB200 CI container (mcore-pyt-dev-dgx_gb200, torch nv26.04) on GB200; an unchanged sibling test reproduced its committed golden bitwise, confirming CI fidelity. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pingtian Li <pingtianl@nvidia.com>

The compact decoupled LayerWise DDP layout (now the default) routes non-Muon (Adam) params through a real DistributedOptimizer (reduce-scatter) instead of the legacy in-LayerWise all-reduce path, shifting lm-loss/num-zeros within tolerance and lowering peak memory. Regenerated on dgx_h100 with the common_pile CI dataset; verified deterministic + resume-consistency pass. Signed-off-by: Pingtian Li <pingtianl@nvidia.com>

Wohox · 2026-06-30T13:25:58Z

/ok to test 99af691

svcnvidia-nemo-ci · 2026-06-30T23:50:25Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28483433802

Wohox force-pushed the pingtian/claude/muon-layerwise-compact-buffers branch from a2a4b68 to 6e726d3 Compare June 17, 2026 09:37

Wohox changed the title ~~Add experimental decoupled compact LayerWise DDP layout for Muon~~ [dev] Add experimental decoupled compact LayerWise DDP layout for Muon Jun 17, 2026

Wohox marked this pull request as ready for review June 17, 2026 09:50

Wohox requested review from a team as code owners June 17, 2026 09:50

svcnvidia-nemo-ci added the complexity: medium label Jun 17, 2026

copy-pr-bot Bot temporarily deployed to public June 17, 2026 09:53 Inactive

copy-pr-bot Bot temporarily deployed to public June 17, 2026 09:56 Inactive

copy-pr-bot Bot temporarily deployed to public June 17, 2026 09:57 Inactive

Wohox mentioned this pull request Jun 17, 2026

[main] Add experimental decoupled compact LayerWise DDP layout for Muon #5391

Closed

6 tasks

copy-pr-bot Bot temporarily deployed to public June 17, 2026 10:06 Inactive

Wohox force-pushed the pingtian/claude/muon-layerwise-compact-buffers branch from f615b2d to e965c83 Compare June 23, 2026 01:51

copy-pr-bot Bot temporarily deployed to public June 23, 2026 01:57 Inactive

copy-pr-bot Bot temporarily deployed to test June 23, 2026 01:58 Inactive

copy-pr-bot Bot temporarily deployed to public June 23, 2026 02:01 Inactive

copy-pr-bot Bot temporarily deployed to public June 23, 2026 02:10 Inactive

copy-pr-bot Bot temporarily deployed to public June 23, 2026 03:39 Inactive

copy-pr-bot Bot temporarily deployed to test June 23, 2026 03:40 Inactive

copy-pr-bot Bot temporarily deployed to public June 23, 2026 03:43 Inactive

copy-pr-bot Bot temporarily deployed to public June 26, 2026 08:17 Inactive

copy-pr-bot Bot temporarily deployed to public June 26, 2026 08:26 Inactive

Wohox force-pushed the pingtian/claude/muon-layerwise-compact-buffers branch from 0689947 to b21a44a Compare June 26, 2026 08:40

Merge branch 'dev' into pingtian/claude/muon-layerwise-compact-buffers

505f3ed

copy-pr-bot Bot temporarily deployed to public June 29, 2026 03:29 Inactive

copy-pr-bot Bot temporarily deployed to test June 29, 2026 03:29 Inactive

copy-pr-bot Bot temporarily deployed to public June 29, 2026 03:32 Inactive

copy-pr-bot Bot temporarily deployed to public June 29, 2026 03:33 Inactive

copy-pr-bot Bot temporarily deployed to public June 29, 2026 03:43 Inactive

FDecaYed approved these changes Jun 29, 2026

View reviewed changes

FDecaYed enabled auto-merge June 29, 2026 05:22

Merge branch 'dev' into pingtian/claude/muon-layerwise-compact-buffers

667002c

Wohox and others added 3 commits June 30, 2026 06:25

Merge branch 'dev' into pingtian/claude/muon-layerwise-compact-buffers

99af691

Wohox mentioned this pull request Jul 13, 2026

[Add] LayerWise reduce-scatter fallback for high-padding layouts. #5738

Open

FDecaYed mentioned this pull request Jul 29, 2026

chore: nightly sync main into dev (28_07_2026) #6112

Closed

lhb8125 mentioned this pull request Aug 18, 2026

Megatron-LM Monthly Activity Report — July 2026 #6607

Open

Navigation Menu

Uh oh!

[dev] Add experimental decoupled compact LayerWise DDP layout for Muon - #5388

[dev] Add experimental decoupled compact LayerWise DDP layout for Muon#5388
Wohox merged 8 commits into
NVIDIA:devfrom
Wohox:pingtian/claude/muon-layerwise-compact-buffers

Wohox commented Jun 17, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jun 17, 2026

Uh oh!

Wohox commented Jun 17, 2026

Uh oh!

Wohox commented Jun 23, 2026

Uh oh!

Wohox commented Jun 23, 2026

Uh oh!

Wohox commented Jun 29, 2026

Uh oh!

FDecaYed left a comment

Uh oh!

Wohox commented Jun 29, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 29, 2026

Uh oh!

Wohox commented Jun 30, 2026

Uh oh!

Wohox commented Jun 30, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

Wohox commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Jun 17, 2026

Uh oh!

Wohox commented Jun 17, 2026

Uh oh!

Wohox commented Jun 23, 2026

Uh oh!

Wohox commented Jun 23, 2026

Uh oh!

Wohox commented Jun 29, 2026

Uh oh!

FDecaYed left a comment

Choose a reason for hiding this comment

Uh oh!

Wohox commented Jun 29, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 29, 2026

Uh oh!

Wohox commented Jun 30, 2026

Uh oh!

Wohox commented Jun 30, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Wohox commented Jun 17, 2026 •

edited

Loading