Add experimental Megatron-FSDP fully_shard implementation by wujingyue · Pull Request #4976 · NVIDIA/Megatron-LM

wujingyue · 2026-05-26T06:38:31Z

Stacked on top of #4835.

What does this PR do ?

Adds an experimental per-module Megatron-FSDP fully_shard path that uses DBuffer primitives to shard parameters, materialize full weights for compute, and reduce gradients back into sharded optimizer state.

Meta-parameter materialization is intentionally split out to the follow-up draft PR at #5369.

Issue tracking

Linked issue: N/A

Summary

add the experimental fully_shard(...) entry point plus FsdpModule, ParameterGroup, and Placements runtime state
group parameters by dtype and requires_grad, manage sharded main/model weight buffers, and sync optimizer-updated weights into compute weights
allocate model/gradient buffers from DBuffer.device instead of reaching through local_buffer.device
install forward/backward hooks for unshard, reshard, and parameter-completion-based gradient reduction, including accumulation across backward calls
extend distributed coverage for DBuffer layout/redistribution and the experimental FSDP path, including nested ownership, frozen parameters, optimizer-step visibility, CPU-initialized parameter sharding, autograd storage reuse, loss parity, and peak-memory reduction

Testing

BASE_REF=pull-request/4835 CHECK_ONLY=true bash tools/autoformat.sh
git diff --check origin/pull-request/4835..HEAD

copy-pr-bot · 2026-05-26T06:38:35Z

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.

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Summarize today's FSDP work: - Add DBuffer storage release/reallocate support and an in-place fully_allgather_into path for materializing replicated buffers. - Simplify ParameterGroup and FsdpModule around sharded DTensor parameters, reused unsharded Parameters, meta materialization, and default-stream unshard/reshard/reduce behavior. - Remove unused optimizer/offload/state/helper surface area from the minimal path and keep version-counter preservation scoped to unsharded model-weight materialization. - Expand DBuffer and experimental FSDP tests for layouts, storage lifecycle, DTensor contracts, meta reset, nested ownership, train-step parity, and peak-memory reduction. Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Add out= support to DBuffer redistribution and primitive communication ops, keeping axis inference in redistribute only. Use preallocated model and gradient buffers in the minimal FSDP path where possible, including direct first-gradient reduce-scatter into main_grad. Update DBuffer and experimental FSDP tests for AVG reductions, explicit primitive axes, storage reuse, and gradient accumulation behavior. Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Use ordered parameter tuples for FSDP parameter swapping and keep the sharded data path aligned with main_weight storage. Set grad_dtype for FSDP-managed parameters so BF16 main gradients can be reduced without pre-reduce casts, and update tests to verify sharded parameter data and grad backing buffers. Clean up hook naming, local gradient accumulation handling, and memory/test assertions for the minimal experimental path. Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Add DBuffer casting for dtype conversion before model-weight sync, refresh model weights from main weights before unshard, and cover the next-forward optimizer update path with FP32 main weights and default BF16 main grads on SGD's non-foreach path. Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

wujingyue · 2026-06-17T00:40:08Z

+        else:
+            self.model_weight = DBuffer(
+                mesh=self.mesh,
+                placements=model_weight_placements,


From @Autumn1998: seems similar to how we create unsharded_model_weight. Consider dedup.

wujingyue · 2026-06-17T00:53:40Z

+
+        return grad_hook
+
+    def pre_forward(self) -> None:


From @Autumn1998 : pre_forward/forward/post_forward may be called multiple times when activation recomputation is on.

@Autumn1998 to provide the code pointer

The user should probably call module.forward instead of module(...) when the intention is to avoid triggering hooks.

wujingyue · 2026-06-17T04:29:06Z

Recovered as replacement draft PR #5387 after this PR was closed due to the deleted base ref. The branch has been rebased directly on main: #5387

wujingyue mentioned this pull request May 26, 2026

Add minimal DBuffer implementation #4835

Merged

wujingyue commented May 26, 2026

View reviewed changes

Comment thread tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py

wujingyue mentioned this pull request Jun 2, 2026

Add FSDP stream context delayed release #5124

Closed

wujingyue commented Jun 8, 2026

View reviewed changes

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

wujingyue commented Jun 8, 2026

View reviewed changes

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

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

wujingyue force-pushed the fsdp/minimal branch from 1335461 to 3438a25 Compare June 10, 2026 22:51

wujingyue changed the base branch from main to pull-request/4835 June 10, 2026 22:58

wujingyue changed the title ~~Add minimal experimental FSDP path~~ Add experimental Megatron-FSDP fully_shard path Jun 11, 2026

wujingyue changed the title ~~Add experimental Megatron-FSDP fully_shard path~~ Add experimental Megatron-FSDP fully_shard implementation Jun 11, 2026

wujingyue force-pushed the fsdp/minimal branch from e2ba61f to fc06298 Compare June 11, 2026 04:10

wujingyue marked this pull request as ready for review June 11, 2026 04:26

wujingyue requested review from a team as code owners June 11, 2026 04:26

wujingyue added the Run tests label Jun 11, 2026

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

wujingyue force-pushed the fsdp/minimal branch from fc06298 to 1de42fe Compare June 11, 2026 04:27

copy-pr-bot Bot temporarily deployed to public June 11, 2026 04:27 Inactive

wujingyue requested review from Autumn1998, cspades, shjwudp and skyw June 11, 2026 04:27

copy-pr-bot Bot temporarily deployed to test June 11, 2026 04:28 Inactive

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

copy-pr-bot Bot temporarily deployed to public June 11, 2026 04:30 Inactive

copy-pr-bot Bot temporarily deployed to public June 11, 2026 04:31 Inactive

copy-pr-bot Bot temporarily deployed to public June 16, 2026 21:36 Inactive

copy-pr-bot Bot force-pushed the pull-request/4835 branch from c01bf24 to bed9193 Compare June 16, 2026 22:14

wujingyue added 13 commits June 16, 2026 22:22

WIP: add experimental minimal FSDP path

8b43c26

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Require matching FSDP main grad dtype

948caff

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Reuse FSDP model weights for matching main weights

5217e08

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Preserve autograd grad dtype in minimal FSDP

6ad5fbc

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Compare minimal FSDP loss curve with baseline

de302ea

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Adapt minimal FSDP to split DBuffer API

878b9a3

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Split minimal FSDP runtime modules

345e96e

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Split minimal FSDP module mixin

2521db9

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Remove experimental FSDP meta parameter support

19a814e

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

wujingyue force-pushed the fsdp/minimal branch from 86f98dd to 19a814e Compare June 16, 2026 22:25

copy-pr-bot Bot temporarily deployed to public June 16, 2026 22:25 Inactive

wujingyue requested a review from Autumn1998 June 16, 2026 22:26

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

copy-pr-bot Bot temporarily deployed to public June 16, 2026 22:38 Inactive

wujingyue commented Jun 17, 2026

View reviewed changes

shjwudp approved these changes Jun 17, 2026

View reviewed changes

copy-pr-bot Bot deleted the branch NVIDIA:pull-request/4835 June 17, 2026 04:18

copy-pr-bot Bot closed this Jun 17, 2026

wujingyue mentioned this pull request Jun 17, 2026

Add experimental Megatron-FSDP fully_shard implementation #5387

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add experimental Megatron-FSDP fully_shard implementation#4976

Add experimental Megatron-FSDP fully_shard implementation#4976
wujingyue wants to merge 13 commits into
NVIDIA:pull-request/4835from
wujingyue:fsdp/minimal

wujingyue commented May 26, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented May 26, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wujingyue Jun 17, 2026

Uh oh!

Uh oh!

Uh oh!

wujingyue Jun 17, 2026

Uh oh!

wujingyue Jun 17, 2026

Uh oh!

wujingyue Jun 17, 2026

Uh oh!

wujingyue commented Jun 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

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

Uh oh!

What does this PR do ?

Issue tracking

Summary

Testing

Uh oh!

copy-pr-bot Bot commented May 26, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wujingyue Jun 17, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

wujingyue Jun 17, 2026

Choose a reason for hiding this comment

Uh oh!

wujingyue Jun 17, 2026

Choose a reason for hiding this comment

Uh oh!

wujingyue Jun 17, 2026

Choose a reason for hiding this comment

Uh oh!

wujingyue commented Jun 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

wujingyue commented May 26, 2026 •

edited

Loading