Fix FSDP double-buffer overflow in lazy main_grad_getter during backward#146
Open
sudhu2k wants to merge 1 commit into
Open
Fix FSDP double-buffer overflow in lazy main_grad_getter during backward#146sudhu2k wants to merge 1 commit into
sudhu2k wants to merge 1 commit into
Conversation
- Expose the reduce-scatter pipeline on the param_and_grad_buffer to enable back-pressure management, preventing overflow in the double-buffer pool during gradient allocation. - Implement a mechanism in main_grad_getter to drain the oldest pending reduce-scatter before allocating new grad buckets, ensuring efficient memory usage. - Adjust the wait_for_previous_grad_reduce method to operate on the reduce-scatter stream, improving synchronization during backward passes.
|
Oh, I just noticed this PR while looking at a PR that fixes this double buffer leak while working on an issue. It was an insanely lousy bug that survived for a long time untested. Sorry for the inconvenience! 🙏🏻 |
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
2 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.
What does this PR do?
Fixes an assertion failure (
No buffer found for bucket_id) when training withfsdp_double_buffer=True. During backward, TE callsget_main_grad()lazily viamain_grad_getter, which could allocate a grad bucket before draining the fixed double-buffer pool (size 2). Once two FSDP units' buckets were live, a third allocation would assert inFixedPoolAllocator.allocate.This PR adds double-buffer back-pressure in the lazy allocation path: expose
grad_reduce_pipelineonParamAndGradBuffer, call_enforce_double_buffer_limitbeforefetch_bucket()inmain_grad_getter, and run the drain on the reduce-scatter stream so buffer frees order correctly against in-flight reduce-scatters.Ports the fix from NVIDIA/Megatron-LM#5222 (commit
55638bc4), adapted for this fork's FSDP layout.Changes
megatron_fsdp.py: Setparam_and_grad_buffer.grad_reduce_pipelinesomain_grad_gettercan reach the reduce-scatter pipeline.param_and_grad_buffer.py—main_grad_getter: Enforce the double-buffer limit before allocating; remove the prematurefetch_bucket()call.param_and_grad_buffer.py—_enforce_double_buffer_limit: Issuewait_for_previous_grad_reduceonrs_streaminstead of the caller's compute stream.Test plan
fsdp_double_buffer=True, TE fused wgrad — training completes past first backward step (previously failed atbucket_id: 31withusing_buffer: {33, 32}).fsdp_double_buffer=False.