Overlap experimental FSDP communication - #5416
Draft
wujingyue wants to merge 5 commits into
Draft
Conversation
wujingyue
force-pushed
the
single-comm-stream
branch
from
June 22, 2026 04:30
84cab1b to
b7328b5
Compare
wujingyue
changed the title
wujingyue
force-pushed
the
single-comm-stream
branch
from
June 22, 2026 04:58
b7328b5 to
c8706f9
Compare
wujingyue
force-pushed
the
single-comm-stream
branch
from
June 24, 2026 21:33
c8706f9 to
0d42783
Compare
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
wujingyue
force-pushed
the
single-comm-stream
branch
from
July 9, 2026 18:53
0d42783 to
09d82d8
Compare
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Mirror the delayed-release path for gradient reductions: replace PreparedReduction / launch_prepared_reductions with a DelayedReduction dataclass, a delayed_reductions queue, enqueue_reduction(), and drain_delayed_reductions(target_length) structured like drain_delayed_releases. enqueue_reduction records a consumer event on the compute stream where the gradients are packed, so each reduce-scatter waits only for its own packing instead of a blanket reduce_scatter_stream.wait_stream(current_stream), removing a synchronization bubble. Drain at target_length=0 in pre_backward (right after the all-gather unshard) and in the end-of-backward callback, so each reduction launches as soon as its module's backward finishes and overlaps the next module's compute. Drop the always-true consumer_event "is not None" guards (record_event never returns None) and tighten DelayedRelease/DelayedReduction consumer_event to torch.cuda.Event. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
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
1 participant
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.
Summary
FsdpContextshared by nested experimental FSDP units.FSDP:layers.0:forward_computeandFSDP:layers.0:backward_computefor nsys attribution.Why
This brings the doublebuffer communication-overlap changes onto the experimental FSDP path now merged in
main. The single communication stream intentionally serializes AG and RS relative to each other, but keeps temporary communication buffers owned by one PyTorch caching-allocator stream. The NVTX ranges make the resulting compute/communication overlap easier to inspect in nsys.Validation
BASE_REF=main CHECK_ONLY=true SKIP_DOCS=false bash tools/autoformat.shpython -m py_compile megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.pyCUDA_VISIBLE_DEVICES=0,1 MASTER_PORT=29569 uv run --no-sync python -m torch.distributed.run --nproc-per-node 2 -m pytest -q --experimental tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py::test_fsdp_nvtx_ranges_use_named_module_paths tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py::test_fully_sharded_root_with_child_units_overlaps_all_gather_and_computeCUDA_VISIBLE_DEVICES=0,1 MASTER_PORT=29570 uv run --no-sync python -m torch.distributed.run --nproc-per-node 2 -m pytest -q --experimental tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py