[albmalamd]

Linked issue or supporting maintainer

Supported by @jeffdaily

BC-breaking?

No

Summary (human written only)

Follow-up to pytorch#198299, pytorch#198301 and pytorch#198302. Some FSDP/HSDP/TP distributed tests started failing with ENOSPC error again on 10 GiB /dev/shm after RCCL 2.32.3 doubled the per-communicator /dev/shm segment to 82.6 MB per rank. This change reuses or frees the groups they create (MoE expert group, TP and HSDP meshes, hybrid-shard auto groups) and runs each tp_ep configuration in its own processes. Peak /dev/shm drops from ~10.6 GB to at most 5.3 GB.

Agent details (if applicable)

I used Cursor (Claude) to root-cause the failures, write the change and run the validation below. I reviewed the code and the results.

Agent-generated details

Root cause. RCCL's proxy-ops pool is sized as ops[MAX_OPS_PER_PEER * NCCL_MAX_LOCAL_RANKS] and ncclShmOpen fallocates all of it. The NCCL v2.32.3 sync raised NCCL_MAX_LOCAL_RANKS from 72 to 144, so each communicator reserves 82,576,040 B per rank, up from ~41 MB. Only about 16 communicators per rank fit in 10 GiB on 8 GPUs.

Changes.

• MixtureOfExperts shares one size-1 expert group per parent group instead of creating one per model build.
• TestTPFSDPIntegration and TestHSDPWithDeviceMeshAndDTensor build their mesh once per class, as in [ROCm][test]Reuse device meshes across ComposabilityTest cases pytorch/pytorch#198299.
• test_fsdp_hybrid_shard_basic_setup destroys the groups FSDP auto-creates in its use_device_mesh=False subtests. Reusing them would skip the setup path under test.
• test_shard_placement_fn_tp_ep is parametrized over (tp_degree, dp_replicate), because it still needs ~27 communicators per rank with the [ROCm][test] Reuse device meshes across FSDP2 ND-mesh subtests pytorch/pytorch#198301 reuse.

Rejected alternative. Freeing tp_ep's groups when its configuration changes fails with Could not resolve the process group registered under the name 3. DeviceMesh.__eq__/__hash__ ignore process groups, so DTensor caches return an equal mesh whose groups were destroyed. A generic per-test group cleanup in MultiProcContinuousTest would hit the same hazard.

Validation. 8x MI300X, 10 GiB /dev/shm, torch 2.13.0+rocm10.2.0a20261008 (RCCL 2.32.3). test_fully_shard_training.py on main needs a newer torch than this wheel, so its identical change was validated on release/2.13.

python -m pytest distributed/fsdp/test_fsdp_core.py -k mixture_of_experts -v --timeout 900
python -m pytest distributed/fsdp/test_fsdp_hybrid_shard.py -v --timeout 900
python -m pytest distributed/fsdp/test_fsdp_tp_integration.py -v --timeout 900
python -m pytest distributed/fsdp/test_hsdp_dtensor_state_dict.py -v --timeout 900
python -m pytest "distributed/_composable/fsdp/test_fully_shard_training.py::TestFullyShardNDTraining" -v --timeout 900

Scope	Before	After	Peak /dev/shm after
test_fsdp_core.py -k mixture_of_experts	ENOSPC	12 passed	1.3 GB
test_fsdp_hybrid_shard.py	ENOSPC	7 passed	5.3 GB
test_fsdp_tp_integration.py	ENOSPC	3 passed	2.0 GB
test_hsdp_dtensor_state_dict.py	ENOSPC	8 passed	2.0 GB
TestFullyShardNDTraining (incl. 4x tp_ep)	ENOSPC	5 passed	5.3 GB

test_mixture_of_experts_offload_true_no_shard_cuda also has a pre-existing, unrelated flaky FSDP-vs-DDP parity mismatch on ROCm, seen with RCCL 2.30.7 too.

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

(cherry picked from commit 605ec3f)

Cherry-picked to release/2.12 branch via #3821

Cherry-picked to release/2.13 branch via #3822

Cherry-picked to release/2.15 branch via #3844

[albmalamd]

!cherry-pick --onto release/2.12

[albmalamd]

!cherry-pick --onto release/2.13

[rocm-repo-management-api-6]

Created branch autogenerated/release/2.12_cherry-pick_pr-3808 and #3821

Comment processed by Build

[rocm-repo-management-api-6]

Created branch autogenerated/release/2.13_cherry-pick_pr-3808 and #3822

Comment processed by Build

[albmalamd]

!cherry-pick --onto release/2.15

[rocm-repo-management-api-6]

Created branch autogenerated/release/2.15_cherry-pick_pr-3808 and #3844

Comment processed by Build