`cp_spmd_types` hangs on ROCm CI: CP redistribution collectives race FSDP on `default_pg`

[rishisinhanj]

Spmd_types issues its redistribution collectives as raw blocking calls, with no async_op and no stream context. With async_op=False PyTorch runs the NCCL kernel on the caller's current stream without fencing it, so when the CP axis resolves to default_pg, those collectives race FSDP's on the same communicator and the group deadlocks. Passing async_op=True fixes it: 28/28 runs pass with identical losses, but we were not sure if it is a valid fix. This seems to be a genuine SPMD bug rather than something flaky on the ROCm end.

This is currently blocking re-enabling the ROCm leg of torchtitan CI (#4039). The failing test is cp_spmd_types in the "8 GPU Feature Tests" workflow (tests/integration_tests/features.py), and it is the only failure in the entire suite — the other 32 of 33 spmd tests pass. Failing job: https://github.com/pytorch/torchtitan/actions/runs/30675245769/job/91301019973

While this issue persists, we will make a PR to skip that test one test

Issue

The test hangs in a collective and gets killed by the NCCL watchdog:

WorkNCCL(SeqNum=211, OpType=_ALLGATHER_BASE, Timeout(ms)=100000) ran for 100015 ms
[PG ID 0 PG GUID 0(default_pg) Rank 0] failure detected by watchdog at work sequence id: 211
PG status: last enqueued work: 242, last completed work: 210
Fatal Python error: Aborted

The group is default_pg. In another CI run on the same test, work 318 completed while 317 did not, on that same PG, so device execution order is diverging from issue order.

cp_spmd_types has failed on ROCm since it was added in #3763 (June 25).
It went unnoticed because the ROCm legs were failing earlier in the job for unrelated reasons.

Repro

NGPU=4 ./run_train.sh --parallelism.spmd_backend spmd_types \
  --parallelism.context_parallel_degree=4 activation-checkpoint:none

llama3_debugmodel, 4 GPUs. Step 1 completes, then a later step wedges.

Cause

spmd_types issues its redistribution collectives as raw blocking calls:

pg = axis
_dist.dist.all_gather_into_tensor(out, x, group=pg)   # _collectives.py:299

No async_op, no stream context. With async_op=False, ProcessGroupNCCL launches the NCCL kernel on
the caller's current stream and does not fence it (ProcessGroupNCCL.cpp:3786-3789):

auto ncclStream = asyncOp ? ncclStreams_.at(key)
                          : at::cuda::getCurrentCUDAStream(device.index());
if (asyncOp) { syncStream(device, ncclEvents_[key], ncclStream); }

When cp_degree == world_size, DeviceMesh returns the default group itself rather than a new one
(device_mesh.py:559-576), so spmd and FSDP end up on the same communicator. Logging
(group, caller, async_op, stream) per collective, identical on all 4 ranks:

PG desc=default_pg ranks=[0,1,2,3]
  FSDP  all_gather_single      async_op=False  stream=0x1b697910
  FSDP  reduce_scatter_single  async_op=False  stream=0x401e6610
  SPMD  all_gather_into_tensor async_op=False  stream=0x0        <- compute stream, unfenced

Three streams driving one communicator with nothing ordering them. Host issue order is the same on
every rank, but device order isn't, and the collectives deadlock.

Why it's spmd_types and not the mesh

An earlier CI job ran both backends at cp=4/world=4, same GPUs, 7 minutes apart
(run 28282189799 / job 83799919203):

test	backend	result
cp	full_dtensor	pass
cp_spmd_types	spmd_types	fail

Same mesh construction, so both have the same default_pg aliasing, and only one of them hangs. DTensor goes through torch.ops._c10d_functional.* where asyncOp=true is the default, so its kernels run on ncclStreams_ behind a fence. We reproduced this locally too: full_dtensor passes 10/10.

Fix

Pass async_op=True and wait immediately, so the kernel lands on the per-communicator stream behind the fence. All 8 call sites in _collectives.py (lines 59, 64, 240, 299, 352, 597, 639, 677) need it; patching some but not others mismatches across ranks.

-        _dist.dist.all_gather_into_tensor(out, x, group=pg)
+        _w = _dist.dist.all_gather_into_tensor(out, x, group=pg, async_op=True)
+        if _w is not None:
+            _w.wait()

arm, cp=4/world=4	pass/total
unpatched	5/28
async_op=True	28/28
dedicated CP communicator (alternative)	10/10
create the group but never use it (control)	2/8
force FSDP onto one stream (diagnostic)	8/8

The create-but-don't-use arm rules out the patch just perturbing timing. Forcing a single stream also fixes it, which points at the same mechanism. This is in the spmd package and not torchtitan.

Losses are unchanged. With --debug.seed 42 --debug.deterministic --debug.deterministic_warn_only, patched and unpatched give the same step 1-10 traces
(8.12977 7.82932 7.07333 6.24817 5.27155 4.75987 4.41952 4.21944 4.64082 4.04201).

It's a race, so the hit rate moves around a lot between machines and runs, and
--debug.deterministic hides it completely. The arms above were interleaved within each batch, so drift hits control and patched runs equally.

Environment

We emulated the CI environment to reproduce this bug locally.

spmd_types 0.2.1, torch 2.14.0.dev20260731+rocm7.2, RCCL 2.27.7, 8x MI350X (gfx950), torchtitan
fc1d4600. The pattern is unchanged on main and v0.2.2, so bumping the pin won't help.

Side note

ProcessGroupNCCL.hpp:284-289 says "each NCCL call is scheduled on a separate CUDA stream that is different from the current CUDA stream". That isn't true for the async_op=False default and contradicts the code above.

[vidushi8]

Evaluation shows this SPMD fix has no perf impact and add more robustness.

Baseline (no fix): Pass rate 8/10 trial runs - Average step-10 TPS: **48,362.38** 
With SPMD async_op=True fix: Pass rate 10/10 trial runs - Average step-10 TPS: **48481.20**