Increase Megatron-FSDP overlap test dim to 8192 for reliable overlap - #5770
Merged
Conversation
test_overlaps_all_gather_and_compute asserts that at least num_children - 1 forward all-gathers overlap GEMM compute. At dim = 4096 the GEMM kernels are short enough that kernel launch and dispatch latency dominate, so the independent all-gather and GEMM kernels do not reliably co-reside on the GPU. On a single DGX H100 node (8xH100) with default NCCL settings the test fails nearly every run, because any one of the eight ranks missing an overlap fails the whole test. Raising the GEMM size to dim = 8192 makes each kernel long enough to amortize launch latency, so the forward all-gathers reliably overlap the preceding child's GEMM. Verified on 8xH100 under default NCCL: 25/25 runs pass, versus 0/12 at dim = 4096. The flakiness does not surface in CI because tests/unit_tests/run_ci_test.sh exports NCCL_MAX_NCHANNELS=1 to reduce NCCL memory. That serializes the all-gather onto a single channel, producing a longer, lighter-weight all-gather kernel that happens to overlap the small dim = 4096 GEMM. Isolated on 8xH100: NCCL_MAX_NCHANNELS=1 alone passes 10/10, while default NCCL fails 0/12. The test was therefore implicitly depending on a memory-reduction environment variable unrelated to the feature under test. Increasing dim removes that hidden dependency so the test passes under both default and CI NCCL configurations. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
|
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. |
Autumn1998
approved these changes
Jul 13, 2026
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29263458844 |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29287044324 |
github-merge-queue
Bot
removed this pull request from the merge queue due to no response for status checks
Jul 14, 2026
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29310102013 |
chochowski
pushed a commit
to chochowski/Megatron-LM
that referenced
this pull request
Jul 20, 2026
…VIDIA#5770) Signed-off-by: Jingyue Wu <jingyuew@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: mchochowski <mchochowski@nvidia.com>
terminator123
pushed a commit
to 021ai/Megatron-LM
that referenced
this pull request
Aug 3, 2026
…VIDIA#5770) Signed-off-by: Jingyue Wu <jingyuew@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
svcnvidia-nemo-ci
pushed a commit
to dimapihtar/Megatron-LM
that referenced
this pull request
Aug 4, 2026
…VIDIA#5770) Signed-off-by: Jingyue Wu <jingyuew@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.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
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.
Summary
test_overlaps_all_gather_and_computeasserts that at leastnum_children - 1forward all-gathers overlap GEMM compute. At
dim = 4096the GEMM kernels areshort enough that kernel launch/dispatch latency dominates, so the independent
all-gather and GEMM kernels do not reliably co-reside on the GPU. On a single
DGX H100 node (8×H100) with default NCCL settings the test fails nearly
every run — any one of the eight ranks missing an overlap fails the whole test.
Raising the GEMM size to
dim = 8192makes each kernel long enough to amortizelaunch latency so the forward all-gathers reliably overlap the preceding
child's GEMM.
Validation (8×H100, default NCCL)
Why the flakiness doesn't show up in CI
The unit-test harness
tests/unit_tests/run_ci_test.shexportsNCCL_MAX_NCHANNELS=1(to reduce NCCL memory). That serializes the all-gatheronto a single channel, producing a longer, lighter-weight all-gather kernel
that happens to overlap even the small
dim = 4096GEMM. Isolated on 8×H100:NCCL_MAX_NCHANNELS=1onlySo the test was implicitly relying on a memory-reduction env var unrelated to
the feature under test. Bumping
dimremoves that hidden dependency so itpasses under both default and CI NCCL configs.
🤖 Generated with Claude Code