[albmalamd]

Linked issue or supporting maintainer

Supported by @pragupta

BC-breaking?

No

Summary (human written only)

This change fixes hangs of the tests test_wait_times_out_and_process_survives and test_timeout_without_process_abort in test_c10d_nccl2.py on release/2.14 with ROCm 10.2. The hang comes from RCCL, and 2.14 hits it because of nccl2's lazy initialization: new_group creates the subgroup as a standalone 8-rank communicator, and for such a communicator RCCL in ROCm 10.2 runs small all_reduces with its DDA kernel, which doesn't check abortFlag. As a result, after the deliberate timeout ncclCommAbort never returns, because RCCL's proxy thread waits in hipDeviceSynchronize for that kernel to finish. The change in _ProcessGroupNCCL2SubgroupTest._init_pg binds the default group to a device (init_process_group(device_id=...)), so the subgroup is created with ncclCommSplit, as on 2.15, and RCCL doesn't use DDA for split communicators, so the abort works. This works around the RCCL bug rather than fixing it; the bug remains for standalone 8-rank communicators until ROCm/rocm-systems#11245 ships.

Agent details (if applicable)

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

Failures. distributed/test_c10d_nccl2.py::ProcessGroupNCCL2BlockingWaitTest::test_wait_times_out_and_process_survives and ::ProcessGroupNCCL2WatchdogNoTearDownTest::test_timeout_without_process_abort hang, F F F at about 30 min per attempt, in ROCm/rockrel run 41069793579. test_c10d_nccl2 took 181 min, and the distributed shard 2/3 job was cancelled at its 6 h limit, so the rest of that shard never ran.

Root cause. release/2.14 initializes nccl2 lazily, so new_group creates the tests' subgroup as a standalone communicator instead of splitting the default one. For a standalone single-node 8-rank communicator, RCCL in ROCm 10.2 runs small all_reduces with its DDA-IPC kernel, which never polls abortFlag. After the deliberate timeout, ncclCommAbort never returns: commReclaim joins the RCCL proxy thread, which is blocked in hipDeviceSynchronize behind that kernel. Split communicators get no DDA resources, so their kernels exit on abort. release/2.15 passes because nccl2 binds the default group eagerly there (pytorch#192895). The RCCL fix is ROCm/rocm-systems#11245 (on develop, not in ROCm 10.2).

Fix. Test-only, +21 lines. _ProcessGroupNCCL2SubgroupTest._init_pg binds the default group to cuda:rank with init_process_group(device_id=...), the pattern ProcessGroupNCCL2EagerNewGroupTest already uses. The subgroup is then split, as on 2.15, and the tests still exercise the same nccl2 timeout, error and abort paths.

Validation (8x MI300X, ROCm 10.2.0rc1, the CI torch 2.14 wheel)

• The two tests: 4/4 runs hang before the change, 12/12 pass after (about 82 s per pair).
• Full file after the change: 42 passed, 3 skipped. ProcessGroupNCCLLegacyNonblockingTest is deselected because of a separate race ([c10d] Serialize ncclCommGetAsyncError between the watchdog and helper polls pytorch/pytorch#203369).
• Causality on the unchanged tests: both pass with RCCL_DDA_ENABLE=0 and with 4 ranks.
• An RCCL-only repro (no PyTorch) hangs ncclCommAbort on a standalone 8-rank communicator for more than 120 s; on a split communicator it returns in 0.64 s.

Not included. #3700 skips test_timeout_with_communicator_cleanup on ROCm for the same RCCL defect. With this change it passes (2/2), while the baseline hangs, so the skip can be dropped in a follow-up.