[merceod]

Why

cosmos3_nano_tp2_sp2 (TP2 x SP2 on 4 GPUs) never finishes starting. Two workers sit in ProcessGroupGloo construction (connectFullMesh, waiting on the store for a peer) and the other two sit in ncclCommSplit from eagerConnectSingleDevice. TP-only and SP-only configs start fine.

The cause is how torch names a members-only group. new_group(..., use_local_synchronization=True) names it sha1("<ranks>_<len(_world.pg_names)>") (_hash_ranks_to_str in torch/distributed/distributed_c10d.py, same in torch 2.9.1 and 2.12.1). len(_world.pg_names) counts the groups this process has registered so far, and a rank that isn't in a group registers nothing. So two members of the same gloo group can get there with different counts, pick different names and store prefixes, and wait for each other forever.

The 2x2 mesh has rank tuples (0, 1), (0, 2), (1, 3) and (2, 3). With the current per-tuple loop, rank 0 reaches the (0, 2) gloo group after making 4 groups and rank 2 after making 2, so they hash 0_2_4 and 0_2_2. Ranks 1 and 3 have moved on to the (1, 3) ncclCommSplit, which needs ranks 0 and 2 to join as non-members, and they never do. A one-axis mesh puts every rank in one group at the same point, which is why this never showed up before. Nothing here is ROCm-specific, so a 2-D mesh on CUDA should hang the same way (not run on NVIDIA).

Making every rank create the groups in one global order isn't enough, even though the torch docstring suggests it is. All device subgroups first and then all gloo groups still hangs (rank 0 hashes 0_2_4, rank 2 0_2_3).

What changed

• WorkerParallelGroups.init_dist (mstar/distributed/communication.py) makes the gloo twins first, in the sorted order every rank already uses. A rank outside a tuple makes a gloo group of just itself at that step, so every rank's count moves in step and the members of each twin compute the same name.
• The device subgroups come after, unchanged. They're named from torch's shared counter, which members-only groups don't touch, so their names and order are the same as before.
• On XPU the device group still doubles as the cpu group and no gloo groups are made.
• Cost: one throwaway one-rank gloo group per tuple a rank isn't in (2 per rank on a 2x2 mesh).

Verification

On 4 x MI325X with our ROCm branch, which carries these commits:

	without the fix	with the fix
cosmos3_nano_tp2_sp2 startup	not healthy after 421 s, stacks as above	healthy after 71 s
Cosmos3-Nano text-to-video, 5 prompts vs the diffusers reference	n/a	5/5 videos, per-video median frame PSNR 16.7 / 22.3 / 22.4 / 12.2 / 12.4 dB

For scale, the same 5 prompts on 1 x MI300X give 16.5 / 22.3 / 22.3 / 12.8 / 12.6 dB, and the diffusers reference against itself with math SDPA gives a median of 19.3 dB. This checkpoint's videos drift a lot under any attention change, so the TP x SP numbers matching the 1-GPU ones is the check here.

A CPU-only gloo script with 4 ranks reproduces the hang without any mstar code. It makes the four groups above in the same order on every rank with use_local_synchronization=True, and ranks 0 and 2 enter (0, 2) with len(_world.pg_names) 2 and 1.

Tests

• test/modular/test_subgroup_creation_order.py (new, added to test/cpu-core.txt). A fake dist that names groups the way torch does drives init_dist on every rank of 2x2, 2x1, 1x2, 4x1, 2x4 and 4x2 meshes. It checks that the members of each gloo twin agree on its name, that every rank makes the same all-rank calls in the same order, that each twin is made by its members in one global order, and that the groups land on the right comm groups. The name check fails on main for the 2x2, 2x4 and 4x2 meshes.
• test_init_dist_uses_mixed_groups_only_on_xpu now expects the new call order (gloo twins, then device groups) and still checks the XPU path.
• pytest test/modular/test_parallel_groups.py test/modular/test_subgroup_creation_order.py: 40 passed.
• python -m pytest @test/cpu-core.txt (ROCm torch 2.12.1 env, GPUs hidden): 2963 passed, 236 skipped, 4 failed. The same 4 fail on main in that env (3 Wan VAE decode tolerance checks and one Rust graph-runtime test that needs a rebuilt extension), so nothing new.
• ruff check .: clean.

Not in this PR

• An upstream torch issue for the naming. Counting earlier groups with the same rank set would keep names unique and agree across members.

[vasilevklart]

Checked on CPU gloo with 4 ranks (torch 2.12.1), the 2x2 mesh hangs on main and works with this fix, same with [(0,1),(0,1,2,3)] and [(0,1,2),(0,3),(1,2,3)]

[vasilevklart]

The field comment (L298) and the docstring (L350) still say one new_group per tuple with non-members discarding the handle, now it's two passes with a one-rank stand-in

[vasilevklart]

This is len(_world.pg_names), the groups currently registered, so it stays equal only inside this loop and only while nothing is destroyed; worth saying here + AGENTS.md invariant 11 still says use_local_synchronization=True is enough

[vasilevklart]

We could guard with a real 4-rank gloo run

patch

--- a/test/modular/test_subgroup_creation_order.py
+++ b/test/modular/test_subgroup_creation_order.py
@@ -11,10 +11,14 @@
 
 from __future__ import annotations
 
+import multiprocessing as mp
+import time
 from types import SimpleNamespace
 
 import pytest
 import torch
+import torch.distributed as dist
+from test_rendezvous_store import _allow_cpu_device_id, _free_port
 
 from mstar.distributed import communication
 from mstar.distributed.communication import GlobalParallelConfig
@@ -146,3 +150,34 @@
         assert sp.device_group.ranks == tuple(sp.group_members) == sp.cpu_group.ranks
         assert tp.cpu_group.backend == "gloo" and tp.device_group.backend == "nccl"
         assert rank in tp.group_members and rank in sp.group_members
+
+
+def _real_rank(rank: int, port: int, out) -> None:
+    _allow_cpu_device_id()
+    view = _mesh(2, 2).per_worker_config[f"worker-{rank}"]
+    view.init_dist(init_method=f"tcp://localhost:{port}", device=torch.device("cpu"))
+    sums = []
+    for group in (view.node_to_tp_group["dit"], view.node_to_sp_group["dit"]):
+        value = torch.ones(1)
+        dist.all_reduce(value, group=group.cpu_group)
+        sums.append(value.item())
+    out.put((rank, sums))
+    dist.destroy_process_group()
+
+
+def test_a_2x2_mesh_starts_on_real_gloo():
+    """The fake above follows torch's naming; this runs torch's own."""
+    ctx = mp.get_context("spawn")
+    out = ctx.Queue()
+    port = _free_port()
+    procs = [ctx.Process(target=_real_rank, args=(rank, port, out)) for rank in range(4)]
+    for p in procs:
+        p.start()
+    deadline = time.monotonic() + 120
+    for p in procs:
+        p.join(timeout=max(0.0, deadline - time.monotonic()))
+    hung = [rank for rank, p in enumerate(procs) if p.is_alive()]
+    for p in procs:
+        p.kill()
+    assert not hung, f"ranks {hung} still in init_dist"
+    assert sorted(out.get(timeout=5) for _ in procs) == [(rank, [2.0, 2.0]) for rank in range(4)]

[merceod]

/review

[mstar-review]

Automated review

🟡 No findings on 3eca928, light review

Description

• The Tests section says the new module was "added to test/cpu-core.txt" and that the suite was run as python -m pytest @test/cpu-core.txt, but that file doesn't exist in this tree and the diff doesn't touch it; CPU Core now selects every module under test/modular/ via CPU_TEST_PATHS, so the description should give the current command and its results on this head.
• The startup and PSNR results were measured on the author's ROCm branch "which carries these commits", not on this head, and nothing ran on CUDA/NCCL even though this path is the default for every TP/SP deployment; say whether a TP-only or SP-only CUDA config was started on this head to show the reordering (gloo twins before device subgroups) doesn't regress the 1-D meshes that work today.

What the review checked

• torch names a members-only group from len(_world.pg_names), and a non-member of an NCCL split returns NON_GROUP_MEMBER without registering. The new gloo pass registers exactly one group per tuple on every rank, so each member hashes the same <ranks>_<1+i> for tuple i.
• A one-rank stand-in name r_c has two _-separated fields, and a real twin's name always has at least three, so a stand-in can't collide with a twin's store prefix.
• Device subgroups still use use_local_synchronization=False, so they're named from _world.group_count, which hashed-name groups don't advance. Their names and the all-rank split order match the old code.
• The gloo pass waits only on members, tuple by tuple in the shared sorted order, so it can't form a wait-for cycle. The NCCL splits then run on all ranks in the same order (invariant 11).
• On XPU, the device group is still reused as the cpu group and no gloo groups are created, as before.
• init_dist is the only new_group call site under mstar/, so the reordering affects no other code.
• The new test lives under test/modular/, so CPU Core picks it up automatically. It runs on CPU only: device=cpu skips set_device_index, resolve_dist_timeout(None) returns {}, and AllReduce.new(None) builds no collectives.
• The existing inline threads already cover the stale field comment and docstring, the pg_names caveat in the comment, invariant 11 wording in AGENTS.md, and the real 4-rank gloo test, so this review doesn't repeat them.

_{Advisory, not a merge gate. Rules: AGENTS.md. Reply to discuss.}

[Gaurav-Shah05]

The fix is right. In my runs on CPU gloo with 8 ranks, main hangs on both a 4x2 and a 2x4 mesh, and this head starts both with the right all-reduce sums.

A shape that might be simpler: on CUDA, initialize the default group as cpu:gloo,cuda:nccl, as init_dist already does on XPU, and have every rank make one new_group per tuple. torch names those from group_count, which every rank advances, and non-members join the members' ncclCommSplit. So the members-only gloo loop and the one-rank groups go away, every group is named from the counter torch shares across ranks, and cpu_group is the device group on CUDA as it is on XPU. I haven't run it on NCCL. What do you think?

[merceod]

Thanks for running both meshes. I like the mixed cpu:gloo,cuda:nccl default group, it is what XPU already does and it would drop the members-only loop and the one-rank stand-ins. It also changes how every group gets created at startup, including the NCCL ones, and I have not run that on NCCL yet, so I would rather keep this PR as the small fix for the deadlock and try your shape as a follow-up with NCCL numbers. Fine either way if you prefer it here.