Support configurable parameter, gradient, and optimizer sharding in experimental FSDP v2 by Autumn1998 · Pull Request #6137 · NVIDIA/Megatron-LM

Autumn1998 · 2026-07-30T02:43:54Z

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

• # What does this PR do?

Adds no_shard, optim, and optim_grads sharding-strategy support to experimental MFSDP v2, including microbatch gradient accumulation and deterministic post-optimizer weight synchronization.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a [feature request](https://github.com/NVIDIA/Megatron-LM/issues/new?
template=feature_request.md) 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.

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-07-30T02:43:57Z

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.

github-actions · 2026-07-30T02:44:02Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

wujingyue

Is this ready for review? It's shown as draft currently.

Autumn1998 · 2026-07-30T05:52:12Z

+                fsdp_parameter_groups[parameter_group] = None

        for parameter_group in fsdp_parameter_groups:
            parameter_group.sync_model_weight_from_main_weight()


In the optimizer post hook, we convert the main weights into model weights. In ZeRO-1, the optimizer-facing weights are stored as flat shards, while the model weights are replicated, so there is an all-gather here. We would like to move this part into the first microbatch so that we can overlap between different layers. @wujingyue Do we have a corresponding design change for this? 🤔

Good question. sync_main_weight_to_model_weight, invoked by the post-optimizer-step hook, does a cast followed by a redistribute today:

Megatron-LM/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py

Lines 209 to 211 in 77c8772

self.main_weight.cast(self.model_weight.dtype).redistribute(

self.model_weight.placements, out=self.model_weight

)

.

Do we prefer to

keep cast in post-optimizer-step but redistribute (usually an allgather) in the first microbatch, or

keep both in the first microbatch, or

keep both in post-optimizer-step as is today?

My gut feeling is (1), because:

compute, which is on the critical path, should be batched for shorter latency

communication, which should be hidden behind compute, should be pipelined and exposed as little as possible

As the last resort, we could also introduce a user knob.

I also think option (1) is the better choice. In addition, we should find a clean way to determine whether it is the first microbatch. Perhaps it would make sense to first open a PR specifically to identify the “first microbatch.” 🤔

identify the "first microbatch."

Remember #5652 tried to add is_first_microbatch but got pivoted to is_last_microbatch because FSDP2 doesn't annotate first microbatch?

However, given it's only about redistribution, maybe something like

Megatron-LM/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py

Lines 322 to 323 in ff6b92a

if self.main_grad.placements != self._accumulation_placements:

self.main_grad = self.main_grad.redistribute(self._accumulation_placements)

will work without the user having to annotate?

We actually don’t need the user to specify the first microbatch; we can just determine it ourselves within the param group.

Autumn1998 · 2026-07-30T05:59:33Z

+            if not is_last_microbatch:
+                self.main_grad.local_buffer.add_(partial_grad.local_buffer)
+                return
+            partial_grad.local_buffer.add_(self.main_grad.local_buffer)


For no_shard and optim, MFSDP v2 currently accumulates partial gradients into main_grad and performs the final RS on the last microbatch.
However, many kernels fuse gradient accumulation by writing directly into main_grad, bypassing FSDP’s explicit accumulation step. The current implementation is not compatible with this fused path.

Yes, I'm aware of this unimplemented optimization. I believe it can be added by allocating reduce_scatter input before calling the fused backward+copy_in kernel. In fact, for exactly this extension, I chose to launch backward and copy_in on the same stream: http://nv/mfsdp-schedule-design => alternatives considered => FSDP2.

Autumn1998 · 2026-07-30T06:04:31Z

+            optimizer_grad = partial_grad.redistribute(self.main_weight.placements)
+            if partial_grad.placements[finalize_axis].reduce_op == dist.ReduceOp.SUM:
+                optimizer_grad.local_buffer.div_(self.mesh.size(finalize_axis))
+            self._install_sharded_grads(optimizer_grad)


Here, we install optimizer_grad instead of main_grad because main_grad has a Partial placement, while the optimizer requires a Flat shard.
The accumulation-gradient layout differs from the optimizer-gradient layout, so we may need a separate accumulation_grad buffer.

#6041 should make this unnecessary.

let me refine this PR with #6041 🤔

#6041 has been merged FYI

This is still fairly subtle. Take ZeRO-1 as an example. During gradient accumulation, main_grad uses a Partial placement. On the last microbatch, it is finalized with a reduce-scatter and becomes the optimizer-facing gradient with a Flat placement. Therefore, the lifecycle contains two transitions: Partial -> Flat on the last microbatch, followed by Flat -> Partial on the first microbatch of the next accumulation cycle.

These transitions should be handled differently. Partial -> Flat requires a real reduce-scatter and allocates a new output DBuffer. In contrast, Flat -> Partial should not perform communication: the Flat gradient belongs to the previous optimizer step and can be discarded after the optimizer has consumed it, so we can directly allocate a fresh Partial accumulation buffer.

The storage dependencies must still be maintained correctly. We need a record stream:

Partial buffer A │ │ enqueue Reduce-Scatter (asynchronous) ▼ Flat buffer B self.main_grad = B │ └── A loses its last reference and is released by the allocator │ ▼ A's storage is reused │ ▼ Reduce-Scatter is still reading A → corrupted gradients

Since main_grad is replaced by a new DBuffer every global batch—either returned by redistribute() or allocated directly—could its storage address become unstable? Would CUDA graph capture or fused gradient accumulation require this address to remain stable across iterations?

I think an alternative approach is to create an accumulate grad (which uses partial/replicate placement) and separate it from the main grad (which is the flat shard), but let them share the same storage — for example, the main grad would be a view of the accumulate grad. @wujingyue What do you think?

Reduce-Scatter is still reading A → corrupted gradients

I doubt this. All these operations in the above diagram are enqueued to the same, reduce-scatter stream. So it should be safe to free after last use. https://dev-discuss.pytorch.org/t/fsdp-cudacachingallocator-an-outsider-newb-perspective/1486#p-2441-v1-single-stream-no-nonsense-cudacachingallocator-3

We need a record stream

Try our best not to. The whole blogpost is pretty much to teach people not to record_stream for FSDP.

I doubt this.

I take this back. I think you might have found the root cause of #5956 (comment).

Megatron-LM/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py

Line 345 in 8339200

self.main_grad = self.main_grad.redistribute(self.main_weight.placements)

deallocates the old main_grad on a different stream than it was allocated, which invites race conditions. I'll figure out a way forward ASAP.

cc @Achyuthan-S as well for the HFSDP CI failure he's investigating

Autumn1998 · 2026-07-30T10:00:04Z

Is this ready for review? It's shown as draft currently.

It’s still a draft; I’ll include #6041 and revise it.

# Conflicts: # megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py

wujingyue · 2026-07-31T05:53:41Z

                parameter_group = get_containing_parameter_group(parameter)
                if parameter_group is None:
                    continue
-                fsdp_parameter_groups.add(parameter_group)


I'd do this instead:

if parameter_group in synced_parameter_groups: continue parameter_group.sync_model_weight_from_main_weight() synced_parameter_groups.add(parameter_group)

I still little confusing: how does this ensure that the order in which communication is issued on each rank is the same?

synced_parameter_groups is only used for deduplication and is never traversed. As long as the parameter order in optimizer.param_groups is uniform, we are good.

Autumn1998 · 2026-07-31T10:12:14Z

/ok to test 64a447a

Autumn1998 added 3 commits July 29, 2026 03:38

Add experimental FSDP sharding strategy support

974fc0e

Fix experimental FSDP parameter group ordering

38e25d3

Use sharding strategy names in experimental FSDP tests

25b2df8

Autumn1998 requested review from a team as code owners July 30, 2026 02:43

svcnvidia-nemo-ci marked this pull request as draft July 30, 2026 02:44

Autumn1998 commented Jul 30, 2026

View reviewed changes

Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py

wujingyue reviewed Jul 30, 2026

View reviewed changes

Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py

Fix experimental FSDP gradient placement semantics

d19fe48

Autumn1998 commented Jul 30, 2026

View reviewed changes

Autumn1998 added 3 commits July 30, 2026 19:40

Merge remote-tracking branch 'ADLR/main' into main-zero-1-2

f22758e

# Conflicts: # megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py

Refine experimental FSDP gradient accumulation

2d523ef

Preserve gradient buffer stream lifetime

bb9fa5d

wujingyue reviewed Jul 31, 2026

View reviewed changes

Autumn1998 added 4 commits July 31, 2026 02:05

Keep FSDP accumulation gradients persistent

93fcd77

Track first microbatch per FSDP parameter group

d4a8933

Use microbatch state for gradient accumulation

afa3a34

Update gradient accumulation microbatch test

64a447a

Autumn1998 self-assigned this Jul 31, 2026

Autumn1998 marked this pull request as ready for review July 31, 2026 10:11

copy-pr-bot Bot temporarily deployed to public July 31, 2026 10:12 Inactive

github-actions Bot added the complexity: low label Jul 31, 2026

copy-pr-bot Bot temporarily deployed to public July 31, 2026 10:16 Inactive

copy-pr-bot Bot temporarily deployed to public July 31, 2026 10:28 Inactive

wujingyue mentioned this pull request Aug 1, 2026

Allocate main_grad on the reduce-scatter stream #6187

Open

	self.main_weight.cast(self.model_weight.dtype).redistribute(
	self.model_weight.placements, out=self.model_weight
	)

	if self.main_grad.placements != self._accumulation_placements:
	self.main_grad = self.main_grad.redistribute(self._accumulation_placements)

Uh oh!

Conversation

Autumn1998 commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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 Jul 30, 2026

Uh oh!

github-actions Bot commented Jul 30, 2026

Uh oh!

Uh oh!

wujingyue left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wujingyue Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Autumn1998 Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Autumn1998 Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wujingyue Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Autumn1998 commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Autumn1998 Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Autumn1998 commented Jul 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

Autumn1998 commented Jul 30, 2026 •

edited

Loading

wujingyue Jul 30, 2026 •

edited

Loading

Autumn1998 Jul 30, 2026 •

edited

Loading

Autumn1998 Jul 30, 2026 •

edited

Loading

wujingyue Jul 31, 2026 •

edited

Loading

Autumn1998 commented Jul 30, 2026 •

edited

Loading

Autumn1998 Jul 31, 2026 •

edited

Loading