Skip to content

fix: abort process groups on rank-local training failure - #781

Draft
maocheng23 wants to merge 1 commit into
maocheng/colocate-1-capture-rowsfrom
maocheng/colocate-2-teardown-abort
Draft

fix: abort process groups on rank-local training failure#781
maocheng23 wants to merge 1 commit into
maocheng/colocate-1-capture-rowsfrom
maocheng/colocate-2-teardown-abort

Conversation

@maocheng23

Copy link
Copy Markdown
Collaborator

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_group in the CLI's finally hangs 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-collective ProcessGroup.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 without ProcessGroup.abort fall back to destroy.
  • The CLI _train marks the exceptional path with a failed flag and passes abort=failed to teardown.

Related Issues

Splits #766. Stack: #1 capture-rows ← #2 (this)#3 rank0-tracker ← #4 colocated-core ← #5 hybrid-shard.

Accuracy Test

  • New test_failure_teardown_aborts_each_distinct_process_group verifies abort is called once per distinct group and collective destroy is not used on the failure path.
  • Existing teardown test keeps covering the success path.

Checklist

  • Format your code according to the Code Formatting with Pre-Commit (black --check and isort --check-only).
  • Add unit tests.
  • Update documentation as needed (none needed).

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant