fix: abort process groups on rank-local training failure - #781
Draft
maocheng23 wants to merge 1 commit into
Draft
Conversation
destroy_process_group is collective; when one rank dies inside the run while its peers are blocked in a CUDA/FSDP collective, teardown hangs in NCCL communicator destruction, the originating traceback never surfaces, and torchrun cannot reap the job. On the exceptional path the CLI now calls destroy_distributed(abort=True), which uses the non-collective ProcessGroup.abort() on every distinct cached group (including the default group) so elastic sees the real exception. The success path keeps the existing collective destroy. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
5 tasks
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
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.
Motivation
Stack 2/5 replacing draft #766. When one rank dies inside a run (a data or capture error) while its peers are blocked in a CUDA/FSDP collective, the collective
destroy_process_groupin the CLI'sfinallyhangs in NCCL communicator destruction: the originating traceback never surfaces and torchrun cannot reap the job. Colocated online training makes rank-local capture failures a realistic event, but this fix applies to every distributed topology.Modifications
destroy_distributed(abort=False): on the exceptional path, call the non-collectiveProcessGroup.abort()on every distinct cached group, including the default group, so elastic surfaces the real exception. The success path keeps the existing collective destroy; torch builds withoutProcessGroup.abortfall back to destroy._trainmarks the exceptional path with afailedflag and passesabort=failedto teardown.Related Issues
Splits #766. Stack: #1 capture-rows ← #2 (this) ← #3 rank0-tracker ← #4 colocated-core ← #5 hybrid-shard.
Accuracy Test
test_failure_teardown_aborts_each_distinct_process_groupverifies abort is called once per distinct group and collective destroy is not used on the failure path.Checklist
black --checkandisort --check-only).