Repository navigation
Support full-iteration CUDA graphs with MFSDP v2 - #7075
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
force-pushed
the
mfsdp-v2-cuda-graph-probe
branch
from
September 4, 2026 03:32
4429266 to
3684f27
Compare
wujingyue
changed the title
wujingyue
marked this pull request as ready for review
September 4, 2026 03:32
wujingyue
commented
Sep 4, 2026
wujingyue
force-pushed
the
mfsdp-v2-cuda-graph-probe
branch
2 times, most recently
from
September 4, 2026 04:14
f9b9814 to
59ee1fc
Compare
wujingyue
force-pushed
the
mfsdp-v2-cuda-graph-probe
branch
from
September 4, 2026 04:23
59ee1fc to
f43b0c0
Compare
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
wujingyue
changed the title
The test only exercised --cuda-graph-impl full_iteration, which captures the forward and backward but not the optimizer step. --optimizer-cuda-graph captures the step itself, so the two cover disjoint parts of the iteration and the optimizer half went untested at scale. Enabling it required dropping --log-num-zeros-in-grad. MegatronOptimizer.step calls count_zeros() when that is set, and count_zeros() ends in a .item() host sync that is illegal inside the captured step (cudaErrorStreamCaptureUnsupported). That is not specific to MFSDP v2: the shared count_zeros_fp32() syncs the same way for every optimizer, and nothing in validate_args rejects the combination. Golden values are regenerated because capturable=True routes FusedAdam through multi_tensor_adam_capturable instead of multi_tensor_adam, which shifts the loss slightly: lm loss at step 50 moves 7.00104 -> 7.00118, a max relative deviation of 4.25e-4 over 50 steps with no compounding. Memory metrics move by 2.5-3.5 KiB against a 4 GiB peak. A fresh run reproduces the new values under the exact-match check, so the captured optimizer step is still reproducible run-to-run. Also drops the num-zeros series from the golden file. With log_num_zeros_in_grad off the optimizer reports a literal 0, so the recorded series was 50 zeros that measured nothing. It was never in this test's METRICS. Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
wujingyue
commented
Sep 7, 2026
|
|
||
| assert fully_shard_context_calls == [True] | ||
|
|
||
| @pytest.mark.parametrize("optimizer_cuda_graph", [False, True], ids=["eager", "cuda_graph"]) |
Contributor
Author
There was a problem hiding this comment.
optimizer_cuda_graph is covered by the unit test added.
Phlip79
approved these changes
Sep 7, 2026
Phlip79
left a comment
Phlip79
left a comment
Member
There was a problem hiding this comment.
What about GB200 functional test?
wujingyue
changed the title
yaoyu-33
approved these changes
Sep 8, 2026
Contributor
Author
I don't think we need it before #5615 |
jaredcasper
approved these changes
Sep 9, 2026
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
wujingyue
force-pushed
the
mfsdp-v2-cuda-graph-probe
branch
from
September 9, 2026 19:36
0590b47 to
3183151
Compare
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
balasaajay
approved these changes
Sep 10, 2026
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/34446913438 |
This was referenced Sep 10, 2026
2 of 5 tasks
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
10 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.
Closes #6152.
Adds full-iteration CUDA-graph support for Megatron-FSDP v2.