fix(mfsdp): support fsdp v2 zero-1 and zero-2 by Autumn1998 · Pull Request #10 · shjwudp/Megatron-LM

Autumn1998 · 2026-05-25T10:39:28Z

What does this PR do ?

⚠️ 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 the @mcore-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
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

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

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

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

… develop/mfsdp-refactor-main-stage2 # Conflicts: # megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py

shjwudp

This PR fixes several existing issues.

I'd prefer a more complete solution that addresses both optim and optim_grads in a single PR. We should ensure that both shard strategies are fully validated with the expected behaviors:

optim: performs one all-gather (AG) and one reduce-scatter (RS) per iteration
optim_grads: performs one AG and num_GA reduce-scatter (RS) operations

Please also include unit tests. These strategies are relatively complex, so a corresponding design doc would be highly valuable—particularly to clarify how the optim and optim_grads sharding strategies are implemented within the current code framework and pitfalls.

shjwudp · 2026-05-26T01:51:22Z

+                        + main_shard_meta.size
+                    ]
+                )
+                sync_replicated_model_buffer(model_weight_buffer)


This unshard can be executed during the forward pass, which is more favorable for overlap. The communication workflow from ZeRO-3 can be reused (only first micro-batch), and for optim and optim_grads sharding strategy, it is not necessary to register post-forward or post-backward hooks.

Additionally, how are optim_grads and optim gradients handle distinguished?

shjwudp · 2026-05-26T01:51:44Z

-        TracePoolAllocator() if enable_trace_pool else StorageFreeingBucketAllocator()
+    use_trace_pool = enable_trace_pool and sharding_strategy in (
+        "optim_grads",
+        "optim_grads_params",


Curious, why "optim" can not use trace pool allocator?

optim does not use temporary unshard/reduce-scatter buckets in the current flow: weights and grads used by compute are replicated, grad reduction isin-place all-reduce, and replicated weight refresh writes directly into the persistent model weight buffer. Therefore TracePoolAllocator has no allocation trace to plan and can hit the empty-trace assertion

I will fix this case

I see, make sense!

shjwudp · 2026-05-26T01:52:33Z

            start_offsets.append(start_offset)
            model_param_shards.append((model_shard, transpose_shard))

        quantize_main_weights_to_fp8(


Is fp8 param need special handle for optim and optim_grads sharding strategy?

shjwudp

Great to see ZeRO-1/2 introduced here (along with critically important unit tests) — this will be a key step toward unlocking Hybrid FSDP. That said, the implementation details will need several rounds of discussion to finalize. Please refer to my review comments. Thank you!

shjwudp · 2026-05-27T10:07:14Z

-    # --- Reduce-grad overlap tracking ---
-    reduce_grad_buckets: Dict[int, List[Tuple[torch.cuda.Event, ParameterGroup]]]
+    # --- Reduce-scatter grad overlap tracking ---
+    reduce_scatter_grad_buckets: Dict[int, List[Tuple[torch.cuda.Event, ParameterGroup]]]


We need potentially support no_shard, so reduce_grad make more sense here.

shjwudp · 2026-05-27T13:09:07Z

-        module.reduce_grad(async_op=ctx.enable_async_reduce_grad)
+        module.reduce_scatter_grad(
+            async_op=ctx.enable_async_reduce_grad,
+            allowed_sharding_strategies=("optim_grads", "optim_grads_params"),


Can we externalize allowed_sharding_strategies to reduce the complexity of the reduce_scatter_grad function?

shjwudp · 2026-05-27T13:11:41Z


+    @torch.no_grad()
+    def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
+        """Finish optimizer-facing gradient synchronization for this iteration."""


Moving the ZeRO-1's gradient handling to the end looks good for now, but we may need to revisit this in the future if ZeRO-1 overlap-grad-reduce support is required.

shjwudp · 2026-05-27T13:12:44Z

        """
        work = None
+        if not bwd_pass and self._needs_replicated_weight_buffer_refresh:
+            sync_replicated_buffer_from_shard(


Would be better merge sync_replicated_buffer_from_shard into L220 weight_buffer.unshard.

shjwudp

LGTM

Autumn1998 added 2 commits May 25, 2026 03:35

fix(mfsdp): support fsdp v2 zero-1 and zero-2

6b23d34

Merge remote-tracking branch 'origin/mfsdp_refactor_main_stage2' into…

434fa2e

… develop/mfsdp-refactor-main-stage2 # Conflicts: # megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py

shjwudp reviewed May 26, 2026

View reviewed changes

Autumn1998 added 7 commits May 26, 2026 01:34

fix(mfsdp): sync replicated weights from shards

6a33a56

fix(mfsdp): allow trace pool for optim strategy

ee6a812

fix(mfsdp): refresh replicated weights during unshard

494922a

fix(mfsdp): delay optim grad reduce-scatter

e62913d

Fix FSDP v2 ZeRO-1/2 grad sync

3664634

refactor(mfsdp): merge FSDP v2 grad sync traversal

d072fbf

docs(mfsdp): align FSDP v2 grad sync docs

54ef471

shjwudp reviewed May 27, 2026

View reviewed changes

Autumn1998 force-pushed the tongliu/fsdp-v2-mixed-precision branch from bf91245 to acd7241 Compare May 28, 2026 08:32

Refactor FSDP v2 unshard and reduce_grad

00507fd

Autumn1998 force-pushed the tongliu/fsdp-v2-mixed-precision branch from acd7241 to 00507fd Compare May 28, 2026 08:35

Refine FSDP v2 replicated buffer refresh

123d7f9

shjwudp approved these changes May 28, 2026

View reviewed changes

Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md

Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py Outdated

Autumn1998 force-pushed the tongliu/fsdp-v2-mixed-precision branch from c70e5b0 to 123d7f9 Compare May 28, 2026 11:12

Autumn1998 added 3 commits May 28, 2026 22:31

Fix FSDP v2 async grad install

090e26e

Merge mfsdp_refactor_main_stage2 into PR branch

152802f

Allow FSDP v2 ZeRO-1 and ZeRO-2 strategies

1670574

shjwudp merged commit 0ca4a59 into shjwudp:mfsdp_refactor_main_stage2 May 29, 2026

Conversation

Autumn1998 commented May 25, 2026

What does this PR do ?

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

shjwudp left a comment

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shjwudp left a comment

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!

shjwudp left a comment • 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!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

(Step 1): Add PR label `Expert Review`

shjwudp left a comment •

edited

Loading