fix: avoid AccumulateGrad stream mismatch by using default stream for DDP init#21579
Open
s-zx wants to merge 1 commit intoLightning-AI:masterfrom
Open
fix: avoid AccumulateGrad stream mismatch by using default stream for DDP init#21579s-zx wants to merge 1 commit intoLightning-AI:masterfrom
s-zx wants to merge 1 commit intoLightning-AI:masterfrom
Conversation
… DDP init DDP was previously wrapped in torch.cuda.stream(torch.cuda.Stream()) which caused the AccumulateGrad node to be created on a non-default stream. This triggered PyTorch's stream mismatch warning when running backward, especially with Fabric and gradient accumulation (no_backward_sync). The fix removes the custom stream context so DDP initialization runs on the default stream, matching subsequent forwards/backwards as recommended by PyTorch's warning message. Fixes Lightning-AI#21567 Signed-off-by: s-zx <[email protected]>
s-zx
requested review from
ethanwharris,
justusschock,
lantiga and
tchaton
as code owners
github-actions
bot
added
fabric
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #21579 +/- ##
=======================================
- Coverage 87% 87% -0%
=======================================
Files 270 270
Lines 24078 24073 -5
=======================================
- Hits 20863 20855 -8
- Misses 3215 3218 +3 |
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
Fixes the AccumulateGrad stream mismatch warning when training with Fabric + DDP, especially with gradient accumulation (
no_backward_sync). The warning occurred because DDP was initialized insidetorch.cuda.stream(torch.cuda.Stream()), creating AccumulateGrad nodes on a non-default stream that did not match subsequent forwards/backwards.Root Cause
DDP setup used a custom CUDA stream context (from PR #17334) when wrapping the model. This caused the AccumulateGrad nodes to be created on a non-default stream. When backward ran on the default stream (or vice versa), PyTorch emitted:
Fix
Remove the custom stream context so DDP initialization runs on the default stream, matching subsequent forwards/backwards as recommended by PyTorch.
lightning_fabric/strategies/ddp.py: Removetorch.cuda.stream(torch.cuda.Stream())contextlightning/pytorch/strategies/ddp.py: Same changeFixes #21567
📚 Documentation preview 📚: https://pytorch-lightning--21579.org.readthedocs.build/en/21579/