Skip to content

Add affine AutoEP checkpoint placement - #8544

Merged
delock merged 8 commits into
deepspeedai:masterfrom
jinyouzhi:autoep-affine-placement
Sep 24, 2026
Merged

Add affine AutoEP checkpoint placement#8544
delock merged 8 commits into
deepspeedai:masterfrom
jinyouzhi:autoep-affine-placement

Conversation

@jinyouzhi

@jinyouzhi jinyouzhi commented Sep 16, 2026

Copy link
Copy Markdown
Contributor

Follow up #8385

This pull request introduces support for flexible expert placement in DeepSpeed's AutoEP (Automatic Expert Placement) system, enabling non-uniform, non-contiguous, and replicated expert layouts. The main changes add a versioned expert placement descriptor, validation logic, and integration into checkpoint consolidation and metadata validation. This lays the groundwork for more advanced expert scheduling and model parallelism strategies.

The most important changes are:

AutoEP Expert Placement Descriptor and Affine Map Lowering

  • Added a new module autoep_affine.py that defines the expert placement descriptor, validation, legacy uniform descriptor synthesis, and lowering to affine maps for sharded tensor reconstruction. This enables flexible, versioned expert placement beyond the legacy uniform contiguous layout.

Integration into Checkpoint Consolidation and Metadata

  • Updated autoep_universal.py to:
    • Accept and validate the new expert placement descriptor in layer metadata, relaxing the requirement that num_local_experts * ep_size == num_experts when a placement is provided.
    • Use the placement descriptor and affine map for reconstructing full expert tensors during checkpoint consolidation, supporting arbitrary expert layouts. [1] [2]
    • Validate placement consistency during expert file consolidation.
    • Import and use the new placement logic. [1] [2]

Metadata Validation Enhancements

  • Updated autoep_zero3_metadata.py to:
    • Import and use placement validation and legacy descriptor synthesis.
    • Track and validate placements for each layer entry, and check runtime layer consistency. [1] [2] [3]
    • Handle the presence or absence of the placement descriptor during partitioned metadata validation.

Documentation Updates

  • Expanded the affine IR specification (affine_ir_spec.md) to document the new AutoEP placement descriptor, its semantics, and its integration into the IR and runtime, clarifying the distinction between placement provenance and scheduling.

Bugfixes and Robustness

  • Fixed a potential bug in affine.py by skipping empty piece lists during tensor rebuilding, preventing errors when a rank has no assigned pieces.

Related: #8252, #8230.

@jinyouzhi
Add affine AutoEP checkpoint placement
e88ffcf 
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
@delock
delock self-requested a review September 17, 2026 06:40
@jinyouzhi
jinyouzhi marked this pull request as ready for review September 18, 2026 16:06
@delock

delock commented Sep 21, 2026

Copy link
Copy Markdown
Collaborator

Hi @jinyouzhi thank you for your PR. This PR overall looks good to me, I have two questions:

  1. did you tested that the new checkpoint format can be correctly converted to UC during training?
  2. I see that UC info in checkpoint contains dict rather than affine map. Lower to affine map happens during convertion to UC checkpoint. For AutoTP, affine map is written in UC info in checkpoint. Is this difference a design choice?

@0z5a 0z5a mentioned this pull request Sep 22, 2026
Draft
@jinyouzhi

jinyouzhi commented Sep 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Thank you for comments.
For 1, I am design e2e test to validate loss curve consistency ( 50 steps -> save/load -> 50 steps vs 100 steps)
For 2, aggree with your observation, AutoEP should align with AutoTP by persisting a versioned per-parameter affine map at save time. I will retain the placement descriptor for provenance and legacy fallback.
@delock

jinyouzhi and others added 2 commits September 22, 2026 07:25
@jinyouzhi
Merge remote-tracking branch 'upstream/master' into pr-8544-add-affin…
d471be8 
…e-autoep-checkpoint-placement

Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

# Conflicts:
#	deepspeed/module_inject/auto_ep_layer.py
@jinyouzhi
Align AutoEP checkpoint maps with AffineIR
2ef5093 
Persist versioned per-parameter affine maps in AutoEP checkpoint metadata and use them for conversion and restore, while retaining placement descriptors for provenance and legacy checkpoints.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
@jinyouzhi

jinyouzhi commented Sep 22, 2026

Copy link
Copy Markdown
Contributor Author

@delock

I completed an end-to-end loss-resume validation using the DeepSpeedExamples script (deepspeedai/DeepSpeedExamples#1014 ) training/deepspeed_finetune_demo/run_autoep_affine_ir_checkpoint_experiment.sh , with Moonlight-16B-A3B on 4 GPUs (AutoEP=4, ZeRO-3 with CPU parameter/optimizer offload).

The experiment trains an uninterrupted 100-step baseline, saves a native checkpoint at step 50, then resumes both directly from the native checkpoint and after  ds_to_universal.py  conversion. The converter completed the AutoEP ZeRO-3 expert-state consolidation successfully, and all runs had finite loss values. For every resumed step (51–100), the native-resume and Universal-checkpoint-resume losses are exactly identical. Their shared difference from the uninterrupted baseline is  max_abs_error=0.1374  and  mean_abs_error=0.01052 ; because the native and Universal paths match at every step, this difference is not introduced by the Universal/AffineIR conversion.

I attached the loss plot below.

loss_comparison

@jinyouzhi
Validate AutoEP affine placement metadata
4012bfc 
Use explicit placement rank IDs for local expert lookup and reject checkpoints whose persisted affine map conflicts with placement provenance.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
delock
delock reviewed Sep 23, 2026
View reviewed changes
Comment thread deepspeed/checkpoint/autoep_universal.py
@@ -288,19 +336,27 @@ def consolidate_autoep_zero12_expert_states(temp_dir, output_dir, expert_param_i

@delock delock Sep 23, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is consolidate_autoep_zero12_expert_states still needed?

@delock delock Sep 23, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

also if it is no longer called, does it mean this functionality is lost? Is it intentionally not called?

@jinyouzhi jinyouzhi Sep 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed. thanks

@jinyouzhi
Fix AutoEP ZeRO-1/2 universal optimizer states
a924776 
Consolidate ZeRO FP32 master expert weights and optimizer states after per-expert model files, and verify native and Universal loss-resume parity for ZeRO stages 1 and 2.

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

jinyouzhi commented Sep 24, 2026

Copy link
Copy Markdown
Contributor Author

fix ZeRO1/2 issues and also have a test with ZeRO-2 @delock

loss_zero2

@jinyouzhi
Avoid duplicate AutoEP FP32 checkpoint writes
b7df13e 
Remove the unreachable expp_rank optimizer consolidation helper; ZeRO-1/2 optimizer state is handled by the active shard merger. Keep expert-file FP32 fallbacks in the temporary directory and let the merger write each final tensor once.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
@delock
delock enabled auto-merge September 24, 2026 03:53
auto-merge was automatically disabled September 24, 2026 05:02

Head branch was pushed to by a user without write access

@jinyouzhi
fix null module state in universal conversion
87a00c9 
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
@jinyouzhi
jinyouzhi force-pushed the autoep-affine-placement branch from 1a6df6a to 87a00c9 Compare September 24, 2026 05:22
@delock
delock added this pull request to the merge queue Sep 24, 2026
Merged via the queue into deepspeedai:master with commit 3571027 Sep 24, 2026
13 checks passed
@Achyuthan-S Achyuthan-S mentioned this pull request Sep 24, 2026
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@delock delock delock approved these changes
@tjruwase tjruwase Awaiting requested review from tjruwase tjruwase is a code owner
@tohtana tohtana Awaiting requested review from tohtana tohtana is a code owner
@hwchen2017 hwchen2017 Awaiting requested review from hwchen2017 hwchen2017 is a code owner
@loadams loadams Awaiting requested review from loadams loadams is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@jinyouzhi @delock