[albmalamd]

Linked issue or supporting maintainer

Supported by @jeffdaily

BC-breaking?

No

Summary (human written only)

This change fixes communicator teardown hangs and an abort-path resource leak. Since NCCL 2.31, graceful communicator teardown (ncclCommFinalize/ncclCommDestroy) waits in a barrier for all communicator ranks on the same node. Aborted communicators skip this barrier.

• ProcessGroupNCCL::shutdown() finalizes communicators one at a time. Ranks may process them in different orders and wait on different communicator barriers forever. This change finalizes them together in one NCCL group.
• _abort_process_group() did not finalize TorchComms communicators or release NVSHMEM team pools. This change releases these resources after aborting, avoiding collective communicator destruction. If a finalize or release fails, the error is logged and cleanup continues.

It also removes the ROCm 10.2 skips that pytorch#199316 added for five teardown tests, because pytorch#200371 and pytorch#200682 already fixed them.

Agent details (if applicable)

Authored with assistance from an AI coding assistant (Cursor).

Current scope after rebasing. NCCL 2.31 added a barrier among a
communicator's host-local ranks to graceful teardown
(bootstrapIntraNodeBarrier in commDestroySync). Only an aborted
communicator skips it. RCCL 2.32 contains the same behavior.

This PR contains three changes:

1. ProcessGroupNCCL::shutdown() finalizes a group's communicators one at a
   time, in unordered_map order. Ranks can hold different communicators
   (for example, lazily created P2P communicators or one communicator per
   device) and enter different teardown barriers. The fix finalizes all
   communicators inside one ncclGroupStart/End, as
   abortCommsFromMap() already does for abort.
2. _abort_process_group() did not release the TorchComms comms or NVSHMEM
   team pools released by destroy_process_group(). It now finalizes
   TorchComms and releases NVSHMEM teams after aborting, in both the
   single-group and WORLD paths. Finalizing an already-aborted TorchComm
   skips its collective communicator destroy and frees only local resources.
   Releasing NVSHMEM after the abort matters because the release may
   synchronize the device. The TorchComms finalize loop is shared with
   destroy_process_group() and still deduplicates comms by identity.
   Post-abort cleanup is best-effort: a TorchComms finalize failure (for
   example a timed-out work item) or an NVSHMEM release failure (for example
   a sticky CUDA error from the release's device synchronize) is logged per
   communicator or group, and the remaining cleanup still runs.
   destroy_process_group() stays fail-fast.
3. [ROCm][ciflow/rocm-preview] Update rocm wheel to 10.2.0a20261014 pytorch/pytorch#199316 skipped five teardown tests on ROCm >= 10.2 (ROCM-31513). The
   failures they cover were already fixed in the tests by [c10d][nccl2] Abort timed-out groups on every rank in timeout tests pytorch/pytorch#200371 (nccl2
   timeout tests) and [c10d][test] Abort nccl_pg in monitored_barrier allreduce-hang tests pytorch/pytorch#200682 (monitored_barrier allreduce-hang tests), so
   the skips are removed. With the skips removed, all five pass on 2 and
   8 MI300X GPUs, with or without the rest of this PR. Reverting the two
   test fixes brings the hang and the ncclRemoteError failure back.

Related upstream work by other contributors. While this PR was under
review, upstream/main gained the overlapping fixes:

• @fduwjj's #194775,
  landed as 32575604cf2, aborts the original group for SHRINK_ABORT and
  documents that excluded ranks must also abort or exit.
• @d4l3k's #200371,
  landed as e63e88f27ce, aborts timed-out nccl2 subgroups on every rank in
  the timeout tests.

The equivalent two commits were dropped from this PR during the rebase.
Consequently, this PR no longer introduces the SHRINK_ABORT behavior
change and is not BC-breaking.

Validation. Before the rebase, the same changes were source-built for
gfx942 and tested on 8x MI300X with RCCL 2.32.3.

• Without grouped finalize, the three test_send_recv_nccl variants timed
  out after 300 seconds and test_nccl_backend_4gpu_module timed out after
  300 seconds. With it, they passed in 75 seconds total and 42 seconds,
  respectively.
• The new
  test_abort_finalizes_comms_and_releases_nvshmem_teams verifies subgroup
  and WORLD cleanup order: abort, then TorchComms finalize, then NVSHMEM
  release. It fails if cleanup is reverted or moved before the abort.
• test_abort_cleanup_continues_when_finalize_and_release_raise verifies
  that a "Work timed out during finalize" exception and a
  "CUDA error: unspecified launch failure" from the NVSHMEM release are
  logged without preventing subgroup or WORLD state cleanup.
• After rebasing, all seven focused TorchComms destroy/abort unit tests pass.

cd test
TEMP_DIR=$T BACKEND=nccl WORLD_SIZE=8 python -m pytest distributed/test_distributed_spawn.py -v -s -k test_send_recv_nccl
python -m pytest distributed/test_c10d_nccl.py -v -s -k test_nccl_backend_4gpu_module
python -m pytest distributed/test_c10d_torchcomms.py -v -k TestC10dTorchCommsDestroyDedup

An RCCL-only reproducer confirms the grouped-finalize mechanism: two ranks
finalizing two communicators in opposite orders hang, while the same calls
inside ncclGroupStart/End complete.

Not covered: no CUDA run here. spin fixlint was unavailable in the
agent environment; IDE diagnostics and git diff --check report no errors.

Pull Request resolved: pytorch#200673
Approved by: https://github.com/jeffdaily

(cherry picked from commit adc4a41)