DeepSpeed

12ffb990 - [AutoEP] Ensure consistent order of partition groups in all-gather when ZeRO-3 (#8876)

[AutoEP] Ensure consistent order of partition groups in all-gather when ZeRO-3 (#8876) 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. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: pengdurice <pengduhit@gmail.com>