Fix FSDP2 reduce-scatter memory & performance regression on release/2.12#3501
Open
anatoliylitv wants to merge 3 commits into
Open
Conversation
Fixes large amount of stream appearing in the FSDP trace when using CUDA graphs with FSDP2. Notes: 1. This fix will only work on CUDA >= 13.2. 2. This patch relieves the issue but doesn't fully resolve it. There's a lot of edge cases that don't fully work currently (tp, ep, micro-batching) but it does reduce the number of streams. In a follow-up, I can try to test and clean-up remaining cases and land fixes incrementally to avoid things from breaking. 3. Not sure what's a good way to test this. Will add a test if the changes look good. 4. It makes the trace much more readable but still has some quirks like the copies, all-gather, and reduce-scatter spread out across different streams even tho the capture has the same stream. Haven't thought about how to resolve that. For a 16-layer linears model with simple FSDP: Before: <img width="1550" height="872" alt="image" src="https://github.com/user-attachments/assets/4dfd0a0d-9b84-466c-9d7a-f9ea06ae7512" /> After: <img width="1328" height="250" alt="image" src="https://github.com/user-attachments/assets/01b67f3f-197f-40e1-8147-f52a0a718881" /> Related issue pytorch#155679 - torch.compile creating a lot of streams Pull Request resolved: pytorch#183983 Approved by: https://github.com/ngimel, https://github.com/weifengpy Co-authored-by: Cursor <[email protected]>
…catter blocking backward compute (pytorch#186000) FSDP keeps **one** reduce-scatter input buffer in flight: the compute stream must wait on the previous reduce-scatter before the next copy-in (`chunk_cat`) can reuse that buffer (`FSDP::post_backward_rs_wait`). When the reduce-scatter is exposed, that recycle wait stalls backward compute every layer — e.g. 37.6 ms/step of compute-stream idle, 100% gated by a reduce-scatter finishing, on an omnifm_v5 trace. There's no compute→reduce-scatter data dependency (backward never reads the reduced gradient), so the coupling is pure buffer-reuse bookkeeping. This exposes the in-flight buffer count — hardwired to 1 — as a per-module tunable: ```python FSDPModule.set_reduce_scatter_max_input_buffers(max_input_buffers: int, *, recurse: bool = True) ``` Raising it lets the next copy-in write a **fresh** buffer instead of waiting, removing the stall. `1` (default) keeps today's behavior; `2` clears the stall once it's `≥` the reduce-scatter pipeline depth while bounding peak memory; larger caps overlap deeper at higher memory. **Design:** generalize the existing recycle to keep at most `max_input_buffers` buffers — reclaim the oldest (`current_stream.wait_event(oldest_rs)`, then drop the keepalive ref; no `record_stream`) before this layer appends one. `1` drains to zero and reuses the freed buffer → **byte-identical to upstream**, and the reclaim is a no-op once `≥` the pipeline depth. The copy-in stays on the compute stream (`foreach_reduce` unchanged) and grads free promptly, so only the retained input buffers cost memory. The cap is a positive int — no unbounded mode (a large cap already means "retain all"). Co-authored-by: Lei Tian <[email protected]> Pull Request resolved: pytorch#186000 Approved by: https://github.com/anshul-si Co-authored-by: Lei Tian <[email protected]>
…ytorch#186335) By default FSDP2 runs all-gather and reduce-scatter on separate CUDA streams but through the same process group -- one NCCL communicator, which processes one collective at a time and so serializes them on the wire. This adds an opt-in FSDPModule API to give reduce-scatter its own communicator: FSDPModule.set_separate_reduce_scatter_group(enable=True, *, recurse=True) When enabled, FSDP creates a dedicated process group over the shard ranks (dist.new_group), one per distinct set of shard ranks (typically one communicator), so reduce-scatter and all-gather can progress concurrently when the network can sustain it; enable=False resets to the shared group. The default is unchanged -- reduce-scatter shares the shard process group -- so this is purely opt-in and creates no extra communicators unless requested. This redesigns the approach explored in (closed) PR pytorch#177015, which made the separate communicator the unconditional default (and created one even for post-forward meshes); here it is an opt-in toggle. Test Plan: ``` python test/distributed/_composable/fsdp/test_fully_shard_overlap.py \ TestFullyShardOverlap.test_set_separate_reduce_scatter_group \ TestFullyShardOverlap.test_fully_shard_backward_comm_overlap ``` Both pass on 4xH100: - test_set_separate_reduce_scatter_group: default shares the shard PG; enabling creates one dedicated PG shared across same-rank-set meshes; disabling resets to the shared PG. - test_fully_shard_backward_comm_overlap: real backward AG/RS overlap (large matmuls + collectives) is no slower than a serialized single-communicator reference. Authored with Claude. Co-authored-by: Lei Tian <[email protected]> Pull Request resolved: pytorch#186335 Approved by: https://github.com/anshul-si ghstack dependencies: pytorch#186000 Co-authored-by: Lei Tian <[email protected]>
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
3 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.
Fix FSDP2 reduce-scatter memory & performance regression on release/2.12
https://amd-hub.atlassian.net/browse/ROCM-27061
[Regression] PyTorch 2.12 causes ~12% TFLOP/s drop for Llama-3.1-405B BF16 (torchtitan, 4-node) on MI325X vs PyTorch 2.11 — same ROCm 7.14.0 build
Problem
Training with FSDP2 on per-param-mesh / MoE-style models (e.g. DeepSeek V3) regressed
after commit
49a3f54d2d27("[FSDP2] overlap reduce-scatter with compute for per-param-meshpytorch#177319") was integrated into
release/2.12:buffer-reuse race corrupting in-flight reduce-scatter inputs
Root cause
#177319changedFSDPCommContext.reduce_scatter_state: ReduceScatterState | Noneto
reduce_scatter_states: list[ReduceScatterState]. For modules with more than oneparam group (per-param-mesh models with separate meshes for dense vs. expert params),
only the param group whose
_param_group_index == _num_param_groups - 1drains theshared list. This relies on an explicitly documented but unenforced assumption about
autograd hook firing order. When the assumption is violated, reduce-scatter input buffers
accumulate across param groups without being recycled, increasing peak memory and
potentially causing buffer-reuse races before in-flight collectives complete.
This feature had already been reverted once in upstream
pytorch/pytorchfor consistentlyfailing tests before being re-landed, confirming its fragility.
Fix
Cherry-picked three upstream
pytorch/pytorchfixes that were already present inROCm/pytorch'srelease/2.13anddevelopbranches but missing fromrelease/2.12:dbf467438698d97f0cf156set_reduce_scatter_max_input_buffersb5ff92918f4set_separate_reduce_scatter_groupThe key fix is pytorch#186000: it replaces the unbounded
list.clear()with a boundedpop(0)recycle loop gated onreduce_scatter_max_input_buffers(default1).At the default, this is byte-identical to the pre-regression behavior — at most one
reduce-scatter input buffer in flight at any time — eliminating both the memory
accumulation and the race condition. The opt-in APIs
(
set_reduce_scatter_max_input_buffers,set_separate_reduce_scatter_group) areavailable for deliberately enabling overlap once validated on target hardware.
Files changed
torch/distributed/fsdp/_fully_shard/_fsdp_param_group.pytorch/distributed/fsdp/_fully_shard/_fsdp_state.pytorch/distributed/fsdp/_fully_shard/_fsdp_collectives.pytorch/distributed/fsdp/_fully_shard/_fully_shard.pytest/distributed/_composable/fsdp/test_fully_shard_overlap.pytest/distributed/_composable/fsdp/test_fully_shard_comm.pytest/distributed/_composable/fsdp/test_fully_shard_training.pytest/distributed/_composable/fsdp/test_fully_shard_init.pyTest Plan
Run full build and tested model with regression.
Test Result
Run successfully. No performance drop and memory retry.
Submission Checklist