Repository navigation
Conversation
Load a truncated HF MoE checkpoint into an AutoEP training engine, train a few steps, and move the weights to rollout engines with a different EP size through a universal checkpoint and through AffineIR transfers. Check every step bitwise and report the time of each path. The README compares the weight flows. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Release ZeRO-3 prefetched parameters after each forward so the later gathers do not find them in flight, report the executor's rounds as the maximum over ranks, and run only the AffineIR path for ZeRO-3, whose universal checkpoint the ZeRO-1 rollout engine cannot load. Add the ZeRO-3 results to the README. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
jinyouzhi
marked this pull request as draft
October 22, 2026 08:40
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Oct 29, 2026
…en ZeRO-3 (deepspeedai#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: deepspeedai#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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
training/deepspeed_finetune_demo/autoep_rollout_e2e.pyandrun_autoep_rollout_e2e.sh, which check precision along HF checkpoint -> AutoEP ZeRO training -> universal checkpoint or AffineIR transfer -> rollout.ucp:save_checkpoint->ds_to_universal->load_checkpoint(load_universal);affine: AutoEP storage maps ->plan_transfer->execute_transfers(experts, device to device); dense weights copied on device.HybridEngineRollout; teacher-force the other engines on the same sequences.The README adds a section with the weight-flow diagrams for both paths and the in-place
HybridEngineRollout, a comparison table, and measured times.Depends on deepspeedai/DeepSpeed#8800 (AutoEP storage maps) and deepspeedai/DeepSpeed#8809 (AffineIR transfer executor).
Results (8x Intel Arc Pro B70, Qwen3-30B-A3B truncated to 4 layers, EP4 -> EP8)
The last row is reported, not required: EP4 and EP8 group tokens differently in the grouped GEMM, so the bf16 accumulation order differs even with identical weights.
(1) A universal checkpoint converted from ZeRO-3 keeps ZeRO-3 model file names, which the ZeRO-1 rollout engine cannot load.
all_gather_coalesced, which issues the dense and expert-group all-gathers in an order that can differ across ranks. The ZeRO-3 numbers above were taken with a local fix (sort byds_idbefore grouping); ([AutoEP] Ensure consistent order of partition groups in all-gather when ZeRO-3 DeepSpeed#8876).