[albmalamd]

Go through the following checklist:

• Passes lint (lintrunner -a)
• Added/updated tests

Summary (human written only)

Sort split ranks did not account for preserve_rank_order correctly and sorted subgroup ranks by default. This was happening in eager path

Agent details (AI-generated)

Root cause

The new_group DeviceMesh path sorts subgroup ranks by default and preserves mesh order when preserve_rank_order=True. The eager path instead passed mesh order directly to split_group, which intentionally preserves the supplied order. Default subgroup order therefore depended on the initialization path for every backend that supports splitting.

Fix

Sort ranks before every split_group call when preserve_rank_order=False. Keep mesh order when preserve_rank_order=True, matching the documented DeviceMesh contract and the existing new_group path.

Behavior change

Eager legacy NCCL, Gloo, XCCL, NCCL-lazy, TorchComms, and custom split backends now sort permuted subgroup ranks by default. This aligns them with non-eager initialization. Callers that need mesh order should pass preserve_rank_order=True.

Validation

The rank-order tests verify:

• Default subgroup ranks are sorted.
• preserve_rank_order=True follows mesh order.
• The local rank matches the corresponding mesh coordinate.

It covers the configured eager backend, explicit eager legacy NCCL, bound eager Gloo, and the existing eight-rank non-eager case. All four tests pass with the fix. Removing the sorting line makes both the legacy NCCL and Gloo regressions fail with reversed local ranks.

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 agent_space/device-mesh-nightly/bin/python test/distributed/test_device_mesh.py DeviceMeshTest.test_device_mesh_preserve_rank_order_eager DeviceMeshTest.test_device_mesh_preserve_rank_order_eager_nccl_legacy DeviceMeshTestGlooBackend.test_device_mesh_preserve_rank_order_eager DeviceMeshTestNDim.test_device_mesh_preserve_rank_order -v

.venv/bin/spin quicklint

Pull Request resolved: pytorch#201949
Approved by: https://github.com/d4l3k

(cherry picked from commit 020ab1c)