Avoid unnecessary NCCL collective coalescing in distributed optimizer #1847
Conversation
This file contains 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
Signed-off-by: Tim Moon <tmoon@nvidia.com>
crcrpar
approved these changes
Sep 29, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
I've been experiencing some data corruption in distributed optimizer checkpoints because PyTorch is not properly synchronizing the NCCL stream with the main CUDA stream. All indications point to a bug in PyTorch's infrastructure for coalesced NCCL calls and I've isolated it down to cases where we enter PyTorch's
_coalescing_managerbut do not perform any NCCL collectives. The debugger suggests that_coalescing_managersets this flag when it enters the context and fails to unset it, resulting in weird behavior in later NCCL calls. I haven't fully bottomed out this bug, but this PR fixes the issue for me.