fix(megatron): gate device_id process-group init (AIMA-227) - #939
Draft
gphuang wants to merge 1 commit into
Draft
Conversation
The device_id injection forces eager RCCL communicator creation for the world PG and every Megatron sub-group. On llama3-70B BF16 FSDP2 (1 node x 8 MI355X, TP/PP/EP=1, ROCm 26.5) that costs +16.0 GiB/GPU peak VRAM and +9.9 s of startup for no throughput change (-0.07 %, within noise over two interleaved repeats per arm). The memory is invisible to the torch caching allocator -- max reserved is byte-identical across arms -- so it is communicator-side. The deadlock the patch guards against did not reproduce on that recipe with the patch disabled, on ROCm 26.5 or on stock 26.3 (whose bundled Primus predates the patch). It is kept rather than removed because multi-node and TP/PP/EP > 1 were not tested, and gated behind enable_init_process_group_device_id (default false) so no recipe pays 16 GiB unknowingly. The existing FSDP2 and ODC guards are preserved. Also corrects the docstring's ~768 MiB estimate, which understated the measured cost by ~21x. Adds a test pinning the full condition truth table plus the shipped default, so the default cannot silently flip back.
Contributor
Author
Note: original repro question folded into PR summaryThe original repro question has been folded into the PR summary; the current request is tracked in the historical Flux-schnell question. |
Contributor
Author
Note: Llama result folded into PR summaryThe detailed Llama patch-OFF result has been folded into the PR summary; it does not cover the historical Flux-schnell workload tracked in the current question. |
Contributor
Author
Question: historical Flux-schnell repro@amital-amd — could you share the exact image/commit, config, node count, parallelism, and hang point from the original Flux-schnell Megatron/TE deadlock? We need one v26.5 patch ON/OFF liveness check before merging. |
gphuang
changed the title
|
@gphuang I did not have those details. Could not find them from logs, probably deleted them during cleanup of dir to save space. |
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
2 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
Tracking: AIMA-227. Introduced through PR #807 and merged to
mainthrough PR #856.Makes the
device_idinjection intotorch.distributed.init_process_groupopt-in throughenable_init_process_group_device_id(defaultfalse).Passing
device_idforces eager RCCL communicator creation and adds about 16 GiB/GPU outside the torch caching allocator. This exhausts the remaining VRAM on Llama-70B FP8 and causes the AIMA-227 throughput regression.An exact same-HEAD FP8 check recovered the 16 GiB, matched the old memory baseline, and completed 10/10 iterations with the patch disabled. Llama also completed without the patch on stock 26.3 and 26.5.
The patch is gated instead of deleted because the historical failure involved Flux-schnell on Megatron with TransformerEngine. If that workload still deadlocks on v26.5 with the patch disabled, the affected Flux recipes should opt in explicitly.
Test plan