Repository navigation
Cover symmetric memory in CUDA graph test - #5987
Merged
Merged
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. |
wujingyue
changed the title
wujingyue
force-pushed
the
cuda-graph
branch
4 times, most recently
from
August 10, 2026 22:38
2b60bdc to
0deb195
Compare
wujingyue
marked this pull request as ready for review
August 10, 2026 22:38
Parametrize test_captures_full_iteration over use_symm_mem so CUDA graph capture and replay are verified for both the default and NCCL symmetric-memory FSDP collective paths, which produce identical loss trajectories. With symmetric memory the FSDP path allocates its all-gather and reduce-scatter staging buffers from a symmetric-memory MemPool and calls symm_mem.rendezvous(); the capture-stream warmup installs and rendezvouses those buffers so the captured region reuses already-registered windows. Warm the data-parallel communicator with a barrier before capture, since NCCL symmetric-memory window registration fails when a rendezvous is the first collective on a communicator. Size the bf16 layers so NCCL selects its symmetric-memory kernels, and use a zero regression target so the loss decreases by a robust margin across replays. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
shjwudp
approved these changes
Aug 12, 2026
cuichenx
approved these changes
Aug 12, 2026
jaredcasper
approved these changes
Aug 12, 2026
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31679739540 |
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Aug 13, 2026
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31723217131 |
ksivaman
pushed a commit
to ksivaman/Megatron-LM
that referenced
this pull request
Sep 1, 2026
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This branch was previously deployed
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
Extend the MFSDP v2 CUDA-graph unit test to cover two cases:
Only the symmetric-memory parametrization is marked
flakyandflaky_in_dev; the default path continues to run normally.Why
The symmetric-memory case tracks #6154. Its dedicated
MemPoolis separate from the CUDA-graph private pool, so allocation addresses and lifetimes are not yet guaranteed across capture and replay. Keep that coverage present but excluded from the default and LTS test runs until the required PyTorch allocator and buffer-registration support is available.Testing
ruff check tests/unit_tests/distributed/mfsdp_v2/test_cuda_graph.pypytest --collect-only -m "not flaky and not flaky_in_dev" tests/unit_tests/distributed/mfsdp_v2/test_cuda_graph.py(collects the default case and deselects the symmetric-memory case)