AutoEP + AutoTP folding: Universal Checkpointing under ZeRO-3

[pengdurice]

AutoEP + AutoTP folding: Universal Checkpointing under ZeRO-3

Part of: #8672 (Q4 2026 roadmap), "Mixture-of-Expert (MoE) Training Scaling and Kernel Development": "AutoEP Performance Optimization & Parallel Folding (AutoEP + AutoTP)" (PoC @tohtana @pengdurice). Follows #8670.

Problem

#8670 makes ds_to_universal refuse a folded ZeRO-3 checkpoint: "Universal checkpoint conversion of an AutoEP +
AutoTP folded ZeRO-3 checkpoint is not supported yet. Load it with the same topology." So a folded run can only resume
with the same TP size, EP size and number of GPUs.

Three things stood in the way: two in the conversion, one in the load.

1. File keys. The AutoEP expert step of the ZeRO-3 conversion keyed checkpoint files by the ZeRO data-parallel rank
   alone (_zero3_rank_from_file). ZeRO-3 files are named zero_pp_rank_{dp}_mp_rank_{tp}. Under folding the files of
   different TP ranks share a data-parallel rank, so they collided. That is why AutoEP + AutoTP folding under ZeRO-3 #8670 refused.
2. The TP merge's strict check. _merge_tp_slice_files asserts that every pattern in AutoTP's universal info was
   used. AutoTP lists the AutoEP expert weights (experts.w1/w2/w3) among tp_replicated_parameter_patterns. Those
   weights are merged by the AutoEP step, not by the TP merge, so their patterns are never used and the assertion
   fails.
3. The layout check on load. _load_checkpoint compares the folding layout recorded in the model-states file with
   the runtime's (_validate_autoep_folding_checkpoint_metadata). This rejects any other TP size, EP size or an
   unfolded runtime, although a universal load re-cuts every tensor for the runtime's layout.

Change

• deepspeed/checkpoint/ds_to_universal.py:
  • The AutoEP ZeRO-3 files are keyed by (TP rank, DP rank), with _zero3_tp_dp_ranks_from_file, which the dense
    merge already uses. _zero3_rank_from_file is no longer used and is removed, along with the folding refusal.
  • _merge_tp_slice_files no longer counts a pattern as unused when it only matches parameters excluded from this
    merge. The ZeRO-3 path now passes the AutoEP expert names as excluded; the ZeRO-1/2 path already did.
  • The expert reassembly itself is unchanged. It already uses each file's EP rank and ZeRO partition rank. Under
    folding each expert slice exists on exactly one rank, and the existing duplicate and completeness checks still
    apply.
• deepspeed/runtime/engine.py _load_checkpoint: skip the folding-layout check when checkpoint.load_universal is
  set. Rigid loads keep it.
• Loading tensors needs no change. For each parameter, load_hp_checkpoint_state in stage3.py cuts the EP slice,
  then the TP shard, then the ZeRO partition in the parameter's own group (dense data-parallel or expert
  data-parallel). That is what folding needs.
• Docs: docs/code-docs/source/autoep.rst and docs/_pages/config-json.md.
• Tests, in tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py:
  • TestZero3FoldedCheckpoint no longer expects the conversion to be refused.
  • New TestZero3FoldedUniversalCheckpoint. A distributed fixture trains a folded TP 2 x EP 2 run on 4 ranks for one
    AdamW step, saves, converts, and records every parameter and both Adam states at full size, then the losses and
    parameters of the next step. Each target layout loads the universal checkpoint and must restore every recorded
    tensor exactly. Targets with the source's 2 data-parallel groups read the same batches, so they must also repeat
    the next step. Targets: the same layout, EP 4, TP 4, 8 ranks with TP 4, and AutoEP without AutoTP; each in fp32
    and bf16.

ZeRO-1/2 conversion of folded checkpoints is not covered: it stores experts in per-expert files, one per TP rank,
and is still refused in autoep_universal.py (resolve_expert_ckpt_path).

Tests

On master d371f33cc with this change, one node of 8 NVIDIA H200 GPUs, torch 2.8.0+cu126 (and master
without the change).

File	Result
tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py (includes the 10 new universal-checkpoint cases)	43 passed, 2 failed: TestTiedOutputHead::test_row_parallel_tied_head_is_refused (2 cases) fails the same on master d371f33cc without this change: the test expects AutoTP to refuse a row-parallel output head tied to the embedding, and master no longer does
tests/unit/v1/moe/test_autoep_checkpoint.py	43 passed
tests/unit/checkpoint/test_autotp_uc_checkpoint.py	38 passed, 1 failed, 6 errors: the same tests fail on master without this change (see below)
tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py	33 passed
tests/unit/checkpoint/test_autoep_affine_universal.py	20 passed
tests/unit/checkpoint/test_universal_checkpoint.py	144 errors: the same on master without this change (see below)

The failures in these two files do not come from this change. Master d371f33cc without this change gives the same
failing tests:

• 6 errors in test_autotp_uc_checkpoint.py and all 144 in test_universal_checkpoint.py are "fixture ... not found"
  for fixtures defined as DistributedFixture classes (affine_resume_checkpoint, baseline_ws2, baseline_ws4).
  This environment has pytest 9.1.1, which does not recognize the _pytestfixturefunction attribute that
  tests/unit/common.py sets to register those classes as fixtures.
• 1 failure in test_autotp_uc_checkpoint.py (TestRealCheckpointUniversalConversionTPxPP::test): gloo
  "Connection closed by peer" in the conversion's worker pool.

On master 5f266e2 (2026-10-15; 1 node of 8 H200, driver 580.126.09, torch 2.8.0+cu126), the failing
tests are the same with and without this change:

File	With this change	Master 5f266e2 without it
test_autoep_autotp_sharded_attention.py	43 passed, 2 failed (TestTiedOutputHead, 2 cases)	33 passed, 2 failed (the same 2)
test_autoep_checkpoint.py	43 passed	not rerun
test_autotp_uc_checkpoint.py	38 passed, 1 failed, 6 errors	38 passed, 1 failed, 6 errors (the same tests)
test_autotp_universal_checkpoint.py	33 passed	not rerun
test_autoep_affine_universal.py	20 passed	not rerun

The 6 errors are "fixture ... not found" (affine_resume_checkpoint 4, uneven_vocab_checkpoint 2) and the failure is
gloo "Connection closed by peer" in TestRealCheckpointUniversalConversionTPxPP::test, on both sides.

Further checks (not part of the PR's tests)

A wider version of the new test: 7 target layouts. Besides the 5 above: 8 ranks with TP 2 x EP 4, and AutoEP EP 4
without AutoTP. Run in fp32 and bf16, with the optimizer-offload PR for folding applied, and in 3 offload
combinations: saved and loaded without offload, saved and loaded with offload, saved without and loaded with offload.
All 42 loads restore every parameter and both Adam states exactly. Every target with the source's 2 data-parallel
groups repeats the next step: losses equal, and parameters within 1e-4 relative, leaving out the attention key bias.
In bf16 only the same layout repeats it.

The key bias is left out of the next-step comparison because its true gradient is zero. Its computed gradient is
rounding noise that another layout adds up in another order, and AdamW turns that noise into steps of about lr
(observed differences 6e-4 to 1.5e-3 relative).

Duplicate check

Searched open and recent pull requests on 2026-10-15 ("universal checkpoint", "ds_to_universal", "folding", "AutoEP"): none converts folded ZeRO-3 checkpoints. #8800 plans live expert transfers between placements; #8637 links ZeRO-1/2 fragments.

Proposal

New functionality: per CONTRIBUTING.md it needs an accepted proposal issue first. Draft:
descriptions/r3/peng-autoep-folding-ucp-v1-proposal-issue.md. TODO(user): open the issue, wait for acceptance, then link it here.

AI assistance

This change, its tests and this description were prepared with AI assistance (Claude).