Add experimental Megatron-FSDP fully_shard implementation - #5387
Merged
Conversation
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>
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>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
|
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. |
shjwudp
approved these changes
Jun 17, 2026
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Phlip79
approved these changes
Jun 22, 2026
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
deepakn94
approved these changes
Jun 23, 2026
Contributor
Author
|
/ok to test |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28079891589 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Recovered replacement for #4976, which GitHub closed after its base ref
pull-request/4835was deleted. This branch is now rebased directly onmain.What does this PR do ?
Adds an experimental per-module Megatron-FSDP
fully_shardpath 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
fully_shard(...)entry point plusFsdpModule,FsdpParameterGroup, andPlacementsruntime staterequires_grad, manage sharded main/model weight buffers, and sync optimizer-updated weights into compute weightsDBuffer.deviceinstead of reaching throughlocal_buffer.deviceTesting
git diff --check origin/main..HEADgit diff --checkfor the latest comment-only updateuv run --no-sync python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/megatron_fsdp/test_dbuffer.py tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py(27 passed, 6 skipped)