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
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).