Support HSDP deferred DP-outer gradient reduction by Achyuthan-S · Pull Request #5743 · NVIDIA/Megatron-LM

Achyuthan-S · 2026-07-10T15:52:56Z

Summary

Adds HSDP DP-outer gradient synchronization to the experimental Megatron-FSDP path, building on the microbatch context from #5652. Partially addresses #5616.

Relax the main_grad/main_weight placement guard so main_grad may rest at a DP-outer-Partial accumulation state (Partial on an axis where main_weight is Replicate).
Gate the gradient reduction on is_last_microbatch: reduce-scatter DP-inner and accumulate into main_grad every backward; reduce DP-outer and install the sharded parameter gradients only on the last microbatch.
Compose the reduction one mesh axis at a time so a 2-D DP mesh can reduce-scatter DP-inner and then reduce DP-outer.
Thread is_last_microbatch from the FSDP context through post_backward.
Add a 2-D-mesh HSDP parity test driven via microbatch(...).

Scope and design notes (draft — feedback welcome)

HSDP only: parameters and optimizer stay DP-outer-replicated ([Replicate, Flat]); gradients rest [Partial, Flat]. HFSDP (optimizer sharded across DP-outer) is a follow-up.
The last-microbatch DP-outer reduction produces a per-step reduced-gradient buffer for the optimizer; a persistent/reused buffer could replace it.
The gradient placement's DP-outer Partial reduce op must match the collective op (AVG). Symmetric-memory reduction across more than one axis is not yet supported (raises NotImplementedError).
Single stream; overlapping DP-outer communication on a separate stream is orthogonal (see Overlap experimental FSDP communication #5416 / Add forward all-gather overlap #5513).

Testing

Run on a 2×2 mesh (4 ranks):

python -m torch.distributed.run --standalone --nproc-per-node 4 -m pytest -q
tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py

Covers HSDP parity against single-rank SGD for num_microbatches in {1, 3}; the existing 1-D FSDP path is unchanged.
cc @wujingyue @ahmadki

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

copy-pr-bot · 2026-07-10T15:52:59Z

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-10T15:53:05Z

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.

Copilot

Pull request overview

This PR extends the experimental Megatron-FSDP (“MFSDP v2”) path to support HSDP-style deferred DP-outer gradient synchronization, using the microbatch context to delay DP-outer reduction until the last microbatch.

Changes:

Allow main_grad to differ from main_weight placements in the HSDP-compatible case (grad Partial where weight Replicate) and introduce deferred DP-outer gradient reduction logic.
Thread is_last_microbatch through post_backward to gate the final DP-outer reduction and installation of sharded parameter grads.
Add a 2-D device-mesh parity test validating HSDP loss equivalence to single-rank SGD across multiple microbatch counts.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
`tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py`	Adds an HSDP parity test on a 2-D DP mesh and introduces HSDP placement definitions.
`megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py`	Implements deferred DP-outer gradient reduction/accumulation and relaxes gradient/weight placement constraints.
`megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py`	Passes `context.is_last_microbatch` into gradient reduction to enable last-microbatch gating.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Achyuthan-S · 2026-07-10T16:12:01Z

+    if len(grad_placements) != len(weight_placements):
+        return False
+    for grad_placement, weight_placement in zip(grad_placements, weight_placements):
+        if grad_placement == weight_placement:
+            continue
+        if isinstance(grad_placement, Partial) and isinstance(weight_placement, Replicate):
+            continue
+        return False
+    return True


Restricted _grad_placements_reduce_to_weight to at most one deferred (Partial→Replicate) axis, matching the single-axis redistribute. Multi-axis deferral is a follow-up.

Achyuthan-S · 2026-07-10T16:12:19Z

+        if self._defers_grad_reduction:
+            self._accumulate_deferred_grad(partial_grad, partial_op, is_last_microbatch)
+            for parameter in self.unsharded_parameters:
+                parameter.grad = None
+            return


_reduced_grad is now cleared at the first microbatch of each step, so the previous step's buffer is released after zero_grad().

…ad buffer Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

wujingyue

Thanks for the PR!

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

Achyuthan-S · 2026-07-11T06:35:36Z

Thanks for the PR!

Anytime @wujingyue , you can always tag me in the future for any other issues , I find them very interesting and love solving them .

Rest main_grad at main_weight placements so it backs .grad and reuses the zero_grad/set_to_none reset; reduce-scatter the sharded axes every backward and all-reduce the replicated DP-outer axis in place only on the last microbatch. Drop the per-group state and multi-axis reduction helper; cover both set_to_none modes in the HSDP parity test. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

main_grad rests at its DP-outer-Partial gradient placement between microbatches and is all-reduced to main_weight's placements on the last microbatch, per HSDP semantics. Extend DBuffer.get_dtensor to represent Partial so main_grad backs .grad throughout, reusing the existing zero_grad/set_to_none accumulation. Cover both set_to_none modes in the HSDP parity test. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

wujingyue

LGTM overall!

wujingyue · 2026-07-15T15:21:59Z

/ok to test

copy-pr-bot · 2026-07-15T15:22:02Z

/ok to test

@wujingyue, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/

wujingyue · 2026-07-15T15:22:39Z

/claude fix

Signed-off-by: svcnvidia-nemo-ci <svcnvidia-nemo-ci@nvidia.com>

svcnvidia-nemo-ci · 2026-07-15T15:24:59Z

🛠️ Claude fix commit 434a4d9d63b9 (attempt 1)

⚠️ This explanation is AI-generated and may be inaccurate; the exact commit is authoritative.

What changed
No edits were made. The pinned base merge already reports all conflicts resolved (no unmerged files; the only ’=======’ matches are Markdown underlines and reStructuredText table borders, not git conflict markers), and all 79 changed Python files compile cleanly.

Files changed by Claude

No additional file edits; the pinned base was merged.

Why
No repairable signal was present: the CI log directory ／home／runner／work／＿temp／claude-fix-ci／ is empty (no lint or unit failure logs), the steering file is empty, and git reported no remaining merge conflicts to resolve. With no observed conflict or terminal failure in scope, there was nothing to fix.

DCO
@Achyuthan-S, please fix any DCO failures on your commits before merge. DCO does not block this workflow.

View exact commit

Sanitized and posted by svcnvidia-nemo-ci.

svcnvidia-nemo-ci · 2026-07-15T15:25:03Z

/ok to test 434a4d9

ericharper · 2026-07-17T20:56:30Z

/ok to test c7c7183

@wujingyue

Filter on activity_type == 'kernel' so the profiler's gpu_user_annotation duplicates aren't double-counted, and format the boolean (per @wujingyue's suggestion). Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

Achyuthan-S · 2026-07-17T21:18:16Z

@wujingyue @ericharper Signed off and black-formatted the activity_type change in 96764cd — DCO is green now. The force-push cleared CI authorization and disabled auto-merge, so it needs one more /ok to test 96764cd and auto-merge re-enabled (or a manual merge once green). Thanks!

wujingyue · 2026-07-17T21:19:15Z

/claude fix

Signed-off-by: svcnvidia-nemo-ci <svcnvidia-nemo-ci@nvidia.com>

svcnvidia-nemo-ci · 2026-07-17T21:21:48Z

🛠️ Claude fix commit 89b5296dc5be (attempt 1)

⚠️ This explanation is AI-generated and may be inaccurate; the exact commit is authoritative.

What changed
No edits made. The pinned base merge in pr-head／ is already fully resolved ("All conflicts fixed", no unmerged files, no genuine conflict markers). All 27 changed Python files parse cleanly.

Files changed by Claude

No additional file edits; the pinned base was merged.

Why
No actionable conflict or failure was present: git reports all conflicts resolved with zero unmerged paths; the CI log directory ／home／runner／work／＿temp／claude-fix-ci／ is empty (no lint or unit-test failures to clear); and the steering file is empty. With no reported failure and no lint tooling available, and policy forbidding broadening the change, no local repair applies.

DCO
@Achyuthan-S, please fix any DCO failures on your commits before merge. DCO does not block this workflow.

View exact commit

Sanitized and posted by svcnvidia-nemo-ci.

svcnvidia-nemo-ci · 2026-07-17T21:21:52Z

/ok to test 89b5296

svcnvidia-nemo-ci · 2026-07-18T02:02:06Z

❌ Claude fix stopped because exact-SHA CI did not complete in time. View exact-SHA CI.

wujingyue · 2026-07-18T15:35:48Z

/claude fix

Signed-off-by: svcnvidia-nemo-ci <svcnvidia-nemo-ci@nvidia.com>

svcnvidia-nemo-ci · 2026-07-18T15:38:56Z

🛠️ Claude fix commit 66b592a2cb3c (attempt 1)

⚠️ This explanation is AI-generated and may be inaccurate; the exact commit is authoritative.

What changed
Resolved the one merge conflict in tests／unit＿tests／distributed／mfsdp＿v2／test＿fully＿shard.py by keeping the PR (HEAD) side, which filters NCCL events via the shared ＿nccl＿events() helper. The base side had inlined equivalent kernel-only filtering; that logic is already inside ＿nccl＿events(), and HEAD matches the sibling block at lines 324-326 using the same helper. cuda＿events stays defined for the downstream gemm＿events block. Verified no markers remain and ast.parse passes.

Files changed by Claude

tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py

Why
A both-modified text conflict: PR HEAD refactored CUDA NCCL event filtering into a ＿nccl＿events() helper, while the pinned base merge (b4ad280) independently inlined equivalent kernel-only filtering (activity＿type == "kernel") with comments. Both express the same intent, so I kept HEAD's helper-based form for consistency with the rest of the file.

DCO
@Achyuthan-S, please fix any DCO failures on your commits before merge. DCO does not block this workflow.

View exact commit

Sanitized and posted by svcnvidia-nemo-ci.

svcnvidia-nemo-ci · 2026-07-18T15:39:00Z

/ok to test 66b592a

svcnvidia-nemo-ci · 2026-07-18T20:19:13Z

❌ Claude fix stopped because exact-SHA CI did not complete in time. View exact-SHA CI.

wujingyue · 2026-07-18T23:30:55Z

/ok to test 66b592a

svcnvidia-nemo-ci · 2026-07-20T22:23:29Z

🔄 Merge queue validation started!

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

Support HSDP grad reduce in experimental mFSDP fully_shard

6095687

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

Achyuthan-S requested review from a team as code owners July 10, 2026 15:52

Copilot AI review requested due to automatic review settings July 10, 2026 15:52

Achyuthan-S requested a review from a team as a code owner July 10, 2026 15:52

svcnvidia-nemo-ci marked this pull request as draft July 10, 2026 15:53

github-actions Bot added the community-request label Jul 10, 2026

Copilot started reviewing on behalf of Achyuthan-S July 10, 2026 15:53 View session

Copilot AI reviewed Jul 10, 2026

View reviewed changes

Restrict deferred-grad guard to one axis and release stale reduced-gr…

7248685

…ad buffer Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

wujingyue requested changes Jul 10, 2026

View reviewed changes

Remove over-defensive main_grad/main_weight placement guard

3d271f3

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

Achyuthan-S added 2 commits July 13, 2026 16:00

wujingyue reviewed Jul 15, 2026

View reviewed changes

wujingyue marked this pull request as ready for review July 15, 2026 15:21

svcnvidia-nemo-ci requested a review from a team July 15, 2026 15:21

wujingyue added the Run tests label Jul 15, 2026

svcnvidia-nemo-ci added the complexity: low label Jul 15, 2026

Apply Claude fix for PR NVIDIA#5743

434a4d9

Signed-off-by: svcnvidia-nemo-ci <svcnvidia-nemo-ci@nvidia.com>

copy-pr-bot Bot temporarily deployed to public July 15, 2026 15:25 Inactive

copy-pr-bot Bot temporarily deployed to test July 15, 2026 15:26 Inactive

ericharper enabled auto-merge July 17, 2026 20:56

copy-pr-bot Bot temporarily deployed to public July 17, 2026 20:57 Inactive

copy-pr-bot Bot temporarily deployed to public July 17, 2026 21:00 Inactive

copy-pr-bot Bot temporarily deployed to public July 17, 2026 21:01 Inactive

yaoyu-33 approved these changes Jul 17, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to public July 17, 2026 21:09 Inactive

Count only NCCL kernel events in _nccl_events

96764cd

Filter on activity_type == 'kernel' so the profiler's gpu_user_annotation duplicates aren't double-counted, and format the boolean (per @wujingyue's suggestion). Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

auto-merge was automatically disabled July 17, 2026 21:15
Head branch was pushed to by a user without write access

Achyuthan-S force-pushed the fsdp/hsdp-grad-reduce branch from c7c7183 to 96764cd Compare July 17, 2026 21:15

Apply Claude fix for PR NVIDIA#5743

89b5296

Signed-off-by: svcnvidia-nemo-ci <svcnvidia-nemo-ci@nvidia.com>

copy-pr-bot Bot temporarily deployed to public July 17, 2026 21:22 Inactive

copy-pr-bot Bot had a problem deploying to test July 17, 2026 21:22 Error

Apply Claude fix for PR NVIDIA#5743

66b592a

Signed-off-by: svcnvidia-nemo-ci <svcnvidia-nemo-ci@nvidia.com>

jiemingz approved these changes Jul 20, 2026

View reviewed changes

wujingyue mentioned this pull request Jul 20, 2026

MFSDP v2 should support HSDP and HFSDP placements #5616

Open

Uh oh!

Conversation

Achyuthan-S commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Scope and design notes (draft — feedback welcome)

Testing

Uh oh!

copy-pr-bot Bot commented Jul 10, 2026

Uh oh!

github-actions Bot commented Jul 10, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Achyuthan-S Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Achyuthan-S Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

wujingyue left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Achyuthan-S commented Jul 11, 2026

Uh oh!

wujingyue left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wujingyue commented Jul 15, 2026

Uh oh!

copy-pr-bot Bot commented Jul 15, 2026

Uh oh!

wujingyue commented Jul 15, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 15, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 15, 2026

Uh oh!

ericharper commented Jul 17, 2026

Uh oh!

Achyuthan-S commented Jul 17, 2026

Uh oh!

wujingyue commented Jul 17, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 17, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 17, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 18, 2026

Uh oh!

wujingyue commented Jul 18, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 18, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 18, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 18, 2026

Uh oh!

wujingyue commented Jul 18, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Achyuthan-S commented Jul 10, 2026 •

edited

Loading