fix(cuda-graphs): MB-928 align DDP initialization with capture stream - #6021
Conversation
Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
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. |
|
/ok to test 9a05740 |
81659ff to
6ec7ab0
Compare
6ec7ab0 to
78a1b64
Compare
78a1b64 to
f311a97
Compare
Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
Head branch was pushed to by a user without write access
|
/ok to test af504fe |
|
/ok to test 16e9c71 |
Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
16e9c71 to
2232cc0
Compare
|
/ok to test 2232cc0 |
| ddp_stream = get_shared_capture_stream() | ||
| else: | ||
| # Preserve a dedicated initialization stream for all other implementations. | ||
| ddp_stream = torch.cuda.Stream() |
There was a problem hiding this comment.
Would it make sense to use:
full_iteration: the shared capture streamlocal/transformer_engine: a dedicated side streamnone: the default stream
Without CUDA graphs, forward/backward normally run on the default stream. Constructing DDP on a separate stream can associate its retained AccumulateGrad nodes with that stream, while incoming gradients are produced on the default stream, resulting in Pytorch warnings like:
The AccumulateGrad node's stream does not match the stream of the node that produced the incoming gradient.
The FP8 failures seen at 5caed0b do not demonstrate that none requires a side stream. The initial TE CUDA-graph test failed during capture and leaked TE’s temporary class-level monkey-patch wrapper; the subsequent main_grad failures appear to be cascading failures from that leaked state.
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30348900592 |
7 participants
Background
Full-iteration CUDA graph capture can fail when DDP-retained AccumulateGrad nodes reference a different initialization stream.
What changed
full_iterationCUDA graphs.none,local, andtransformer_engine.Details
Both DDP construction paths select the shared stream only for full-iteration capture. Other implementations retain the side-stream behavior required by Transformer Engine FP8 gradient setup.
Tested
tools/autoformat.shtest_dist_utils.py— 68 passedtest_fp8_param.py— 9 passed, 10 skippedtest_full_cuda_graph.py— 2 passed per rankThe separate GB200 failure was a launcher API read timeout before tests started.