Drain predecessor reduce-scatter at dispatch time - #4940
Merged
deepakn94 merged 5 commits intoMay 27, 2026
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
Author
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
Author
|
/ok to test a817431 |
Contributor
Author
|
/ok to test fdfea86 |
deepakn94
force-pushed
the
dnarayanan/drain_predecessor_grad_reduce
branch
from
May 24, 2026 14:52
fdfea86 to
1dc5cae
Compare
deepakn94
marked this pull request as ready for review
May 24, 2026 16:13
yashaswikarnati
approved these changes
May 26, 2026
yaoyu-33
approved these changes
May 26, 2026
kvareddy
approved these changes
May 27, 2026
svcnvidia-nemo-ci
added
Approved
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/26494868163 |
Victarry
pushed a commit
to yanring/Megatron-LM
that referenced
this pull request
May 27, 2026
* origin/main: (50 commits) Drain predecessor reduce-scatter at dispatch time (NVIDIA#4940) ci: Add allow_failure flag to gpt and moe recipes that are failing in nightlies (NVIDIA#4905) fix(tests): initialize num_microbatches calculator in vision cudagraph tests (NVIDIA#4986) test: re-enable test_pp2_create_cudagraphs_first_stage on TE 2.15+ (NVIDIA#4985) ci: Add support for MBridge job gating based on PR labels (NVIDIA#4926) test(ci): re-enable 8experts2parallel_multi_dist_optimizer_instances_1node (NVIDIA#4984) test: re-enable paged stashing MoE tests (NVIDIA#4978) Fix elastification unwrap_model import (NVIDIA#4972) Avoid offsetting functional test master port (NVIDIA#4973) test: enable NVTE_CUTEDSL_FUSED_GROUPED_MLP via pytest fixture (NVIDIA#4931) chore(beep boop 🤖): Bump (main) (2026-05-25) test(release): add release goldens for deepseekv3/nemotron3 and set tp2pp2 exit-interval (NVIDIA#4932) Fix `get_batch` return order to ignore BlendedDataset provenance fields (NVIDIA#4952) ci: restore perf test torchrun logs (NVIDIA#4951) Various training utils (NVIDIA#4872) ci: Update training script paths in BERT and T5 (NVIDIA#4939) [MXFP8/FP4-param-gather] Post processing after forced param AG in eval (NVIDIA#4562) Fix mxfp8 param gather numerical issue when DP overlap is off (NVIDIA#4800) Add TEFusedDenseMLP for Dense+Grouped GEMM fusion on SM100+ (NVIDIA#4318) (NVIDIA#4786) Fix paged stashing test submodules lookup (NVIDIA#4925) ... # Conflicts: # megatron/training/training.py
janEbert
pushed a commit
to janEbert/Megatron-LM
that referenced
this pull request
Jun 2, 2026
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
mathemakitten
pushed a commit
to mathemakitten/Megatron-LM
that referenced
this pull request
Jun 12, 2026
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.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
5 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.
Summary
When
reduce_scatter_with_fp32_accumulation=True, each in-flight reduce-scatter pins an intermediate all-to-all output tensor until.wait()runs. Withoverlap_grad_reduce=Trueand multiple bucket groups, those buffers stayed live until end-of-step, inflating peak memory by O(num_buckets * bucket_size).next_param_gather_bucket_groupchaining: link each bucket group to its immediate predecessor in dispatch order (forward bucket order during backward) viaprevious_grad_reduce_bucket_group.start_grad_sync, drain the predecessor so its all-to-all buffer is freed before the new one is allocated.finish_grad_syncidempotent (underoverlap_grad_reduce) via a per-iterationgrad_reduce_finishedflag so the end-of-step finalize loop is a no-op for buckets already drained by their successor (and the lazy first-batch path is safe under repeated invocation).Linkage is gated on
overlap_grad_reduce && reduce_scatter_with_fp32_accumulation && num_distributed_optimizer_instances == 1— the only config in which the extra buffer exists.Results
8B model, 8 nodes / 64 H100s (PP=1),
overlap_grad_reduce=True, 7 grad-reduce buckets. Memory measured on rank 0 at steady state.Compared to
bf16rs (one-shot, no drain), the drain reclaims 2.2 GB of reserved memory at TP=4 and 4.2 GB at TP=2 — the saving is larger at TP=2 because each transientall_to_all_output_tensoris correspondingly larger. Throughput also improves (+2.6% at TP=4, +9.3% at TP=2), since less allocator churn means less iteration-level overhead.Loss curves also match well:
Test plan
tests/unit_tests/distributed/test_reduce_scatter_with_fp32_accumulation.py,test_param_and_grad_buffer.py,test_distributed_data_parallel.py,test_grad_sync_with_expert_parallel.py.🤖 Generated with Claude Code