Skip to content

Add AutoEP HF checkpoint to rollout precision check - #1020

Draft
jinyouzhi wants to merge 3 commits into
deepspeedai:masterfrom
jinyouzhi:autoep-rollout-e2e
Draft

jinyouzhi wants to merge 3 commits into
deepspeedai:masterfrom
jinyouzhi:autoep-rollout-e2e

Conversation

@jinyouzhi

@jinyouzhi jinyouzhi commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

Adds training/deepspeed_finetune_demo/autoep_rollout_e2e.py and run_autoep_rollout_e2e.sh, which check precision along HF checkpoint -> AutoEP ZeRO training -> universal checkpoint or AffineIR transfer -> rollout.

  1. Load a truncated HF checkpoint (Qwen3-30B-A3B, 4 layers) into an AutoEP training engine (EP4). Every weight is compared bitwise with the safetensors tensors.
  2. Train 3 steps.
  3. Move the trained weights to NaN-initialized rollout engines at EP8:
    • 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.
  4. Generate greedily with 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)

Check ZeRO-2 ZeRO-3
HF weights vs training engine bitwise, 10.4e9 elements bitwise, 10.4e9 elements
UCP rollout weights vs trained bitwise, 8.0e9 elements not run (1)
Affine rollout weights vs trained bitwise, 8.0e9 elements bitwise, 8.0e9 elements
UCP vs Affine logits bitwise -
Rollout vs training engine logits max diff 0.68, 16/1240 argmax differ max diff 0.67, 32/1248 argmax differ

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.

Step UCP, ZeRO-2 Affine, ZeRO-2 Affine, ZeRO-3
save / plan 7.1 s 3.2 s 0.05 s
convert / execute 107.7 s 0.27 s 0.23 s
load / dense copy 25.1 s 0.01 s 0.18 s

(1) A universal checkpoint converted from ZeRO-3 keeps ZeRO-3 model file names, which the ZeRO-1 rollout engine cannot load.

@jinyouzhi
jinyouzhi requested a review from tjruwase as a code owner October 21, 2026 11:12
jinyouzhi and others added 2 commits October 21, 2026 11:12
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant