[jinyouzhi]

Fixes an intermittent ZeRO-3 hang when parameters live in more than one partition group (e.g. AutoEP). all_gather_coalesced now sorts its parameters by ds_id before splitting them by group, so every rank issues the per-group all-gathers in the same order.

Related: #8060 (added per-group gathering to all_gather_coalesced). Found while validating deepspeedai/DeepSpeedExamples#1020 (AutoEP HF → ZeRO training → rollout E2E).

Bug

• Symptom: AutoEP + ZeRO-3 training hangs from step 3 on XPU (XCCL). It is intermittent and depends on per-process object addresses. It did not reproduce on CUDA/NCCL (see Validation), but issuing collectives on different communicators in a rank-dependent order is not safe in general, and the per-group order is rank-dependent regardless of backend.
• Affected configuration: ZeRO-3 with parameters partitioned over two or more process groups. Example: AutoEP, where dense params are on the DP group and experts on the expert-data-parallel (EDP) group. If all params share one group, nothing changes.
• Reproduce: the new unit test (4 ranks), or run_autoep_rollout_e2e.sh from DeepSpeedExamples#1020 with ZERO_STAGES=3.

Root cause

1. The ZeRO-3 prefetcher hands params_to_prefetch (a set) to all_gather_coalesced. A set's iteration order follows object addresses, which differ per process.
2. all_gather_coalesced buckets params into params_by_group in that iteration order. It then issues one gather per group, so the first group seen is gathered first.
3. So rank A can issue DP→EDP while rank B issues EDP→DP. Each blocks in a collective the other has not reached yet, a circular wait.
4. The existing sorted(params, key=lambda p: p.ds_id) ran only after bucketing, inside each per-group call. It fixed the order of params within a group, not the order of the groups.

It starts at step 3 because prefetching begins only once the module trace is complete. While the E2E run was hung, py-spy showed rank 2 in _all_gather_sequential with a single expert param (EDP) and the other ranks in _all_gather_coalesced for the dense params (DP).

Fix

• Move the ds_id sort to the entry of all_gather_coalesced, before bucketing. ds_id comes from a counter (Init.param_id) assigned in construction order, so it is identical on every rank. Both the group order and the order within each group are now rank-independent.
• Remove the later per-group sort. Each group is now a slice of an already sorted list, so it was redundant.
• Complexity: low (+8/−8 in partition_parameters.py). No API, config or checkpoint change.

Compatibility

Axis	Effect
ZeRO stage	Only ZeRO-3 uses this function.
Single partition group (default)	Same gather order and values as before. The sort now also runs before the NVMe swap-in and the world-size-1 early return, and neither depends on order.
Multiple groups (AutoEP / MoE EDP)	Groups are now gathered in ds_id order on every rank (the fix).
ZeRO++ hpZ / qwZ	The coordinator already splits by secondary tensor and by quantization before calling, and each call is sorted as before.
Param offload (CPU/NVMe)	The swap-in receives the sorted list; order-insensitive.
Determinism	All-gather has no reduction, so values are unchanged; only the issue order of the per-group collectives is fixed.
Cost	One sort per call over the params of a single fetch (tens to hundreds). The recursive per-group calls re-sort already-sorted input, which takes linear time.

Test

TestMultiGroupGatherOrder in tests/unit/runtime/zero/test_zero_stale_handle_wait.py (world size 2):

• One param partitioned over the world group and one over a separate group of the same two ranks (dist.new_group([0, 1])). The two groups are distinct communicators, so the gather is split into two per-group collectives.
• Rank 1 passes the params in reverse order.
• Asserts the gathered values.
• Without the fix it hangs.

Validation

Hardware and software: 8x Intel Arc Pro B70; driver 1.15.39122; Ubuntu 24.04; Python 3.12; torch 2.14.0.dev20260804+xpu; XCCL backend.

Command	Without fix	With fix
DS_UNITTEST_TIMEOUT=120 pytest --forked tests/unit/runtime/zero/test_zero_stale_handle_wait.py -k TestMultiGroupGatherOrder (2 ranks)	hangs, 2/2 runs	pass
pytest --forked tests/unit/runtime/zero/test_zero_stale_handle_wait.py tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py	—	5 passed
AutoEP E2E, ZeRO-3, 8 ranks, EP4, Qwen3-30B-A3B truncated to 4 layers (DeepSpeedExamples#1020)	hangs at training step 3	all steps complete; HF / training / rollout weights bitwise equal

Same test on CUDA: 8x NVIDIA L20, torch 2.14.1+cu130, NCCL 2.30.7.

Command	Without fix	With fix
TestMultiGroupGatherOrder, 2 ranks	pass, 3/3 runs	pass
TestMultiGroupGatherOrder, 4-rank variant ({0,1}/{2,3} subgroups)	pass, 3/3 runs	—

NCCL does not hang here, likely because ProcessGroupNCCL runs each communicator on its own internal stream, so the two per-group gathers progress concurrently. On NCCL the fix only removes the rank-dependent issue order.

Not validated: ZeRO++ hpZ/qwZ and NVMe param offload were not run. They are covered only by the reasoning in the compatibility table.

CI coverage

• cpu-torch-multi-latest (nightly, LOCAL_SIZE=4) runs the new test on CPU over gloo. Locally (10 runs) (LOCAL_SIZE=4 DS_ACCELERATOR=cpu) it passes with and without the fix: gloo appears to progress each process group independently, so issuing the groups in opposite order does not block there. On CPU and on NCCL the test therefore only guards the pass path; the hang reproduces on XCCL (see Validation).
• Plain cpu-torch-latest advertises one device and skips it. With 2 ranks the test runs on any 2-GPU runner. Modal's selector only scopes tests/unit/v1. xpu-max1100 / hpu-gaudi2 are triggered by the path but run fixed lists that do not include this file.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-21T13:30:37.446069Z	41168b1	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[pengdurice]

@jinyouzhi , please fix the CI issue.