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

[0z5a]

Thanks for picking this up, it fills an important gap for the folding work. I went through the converter path for folded ZeRO-3: routed experts aren't TP-sliced under folding (expert TP = 1), fragments don't collide since each global rank has a unique (ep_rank, edp_rank), ZeRO padding is stripped via ds_zero_partition_groups, and TP-replicated params go through replicated_map.rebuild. A synthetic folded grid over 7 layouts with odd sizes (so EDP padding occurs) reassembles fp32 / exp_avg / exp_avg_sq bitwise for me. The engine change is a no-op for non-folded runs, and comparing against the source engine's live state in the GPU test is the right reference.

1. [blocking] Existing unit test now fails with KeyError: 0 — ds_to_universal.py:709, test_autoep_unit.py:941

_autoep_expert_param_names_by_rank now keys metadata_by_rank by (tp_rank, dp_rank), but test_universal_converter_detects_zero3_partitioned_autoep_model_state still indexes metadata_by_rank[0]. It fails on CPU (DS_ACCELERATOR=cpu), and it's the one failure in the cpu-torch-latest check on 9362efc. Could you update it to metadata_by_rank[(0, 0)] and note the key shape in the helper's docstring, since other tests import it?

2. [should-fix] Docs promise module-only loads that aren't tested — config-json.md:1147

The new text promises weights-only / module-only universal loads at a different autotp_size, but TestZero3FoldedUniversalCheckpoint only exercises the default load_checkpoint. Module-only goes through load_module_checkpoint_state_from_checkpoint_dir_stage3, which never sees a folded source here. Could you add one load_module_only=True target that compares fp32 against the recorded state, or narrow the docs? Also, autoep.rst around L597 still says only autoep_size / DP size can change via universal checkpointing; worth mentioning autotp_size for folded runs.

3. [should-fix] CI covers a narrow set of layout changes — test_autoep_autotp_sharded_attention.py:529-535

On the 4-GPU runners (8, 4, 2) skips, so no CI-run target grows the dense DP degree; only (4, 4, 2) changes it, down to 1. A 4-rank (4, 1, 4) target would move dense DP 2→4, EDP 2→1 and TP 2→1 in one case. Also, with ATTENTION_ONLY_TP_SPECS only q/k/v (+bias) and o_proj are TP-sliced; the default-spec layouts (MLP gate/up column, down row, row-parallel bias, vocab-parallel embedding/head) never reach this converter path. One use_default_specs target, or a note in the description that they're out of scope?

4. [question] Experts are published as TP-replicated in the AutoTP UC metadata — ds_to_universal.py:608-611

Under folding, collect_autotp_universal_checkpoint_info publishes ...experts.w[123] as TP-replicated with a replicated_map over the EP-local shape, but TP peers hold different experts (with TP=EP=2 the groups coincide). Filtering those patterns out of the strict check works for this converter, but the metadata itself stays wrong. Would it be cleaner to skip AutoEP expert params when collecting AutoTP UC info? Anything that later plans from AFFINE_MAP in UC info (e.g. the shard-to-shard planner / executor) would otherwise trust a wrong map.

5. [nit] k_proj.bias exclusion — test_autoep_autotp_sharded_attention.py:577

The key bias only has a zero true gradient on its non-rotary half. With partial_rotary_factor=0.5 the rotary half is rotated by position-dependent RoPE; in fp64 on this tiny GLM-4-MoE I get max |grad| ≈ 1e-4 on the rotary dims (same order as q_proj.bias) vs ≈ 4e-11 on the pass-through dims. Could you exclude only [..., int(head_dim * partial_rotary_factor):] of each head instead of dropping the whole tensor?

6. [nit] Fallback partition values — ds_to_universal.py:870-871

If ds_zero_partition_groups were missing, this falls back to the dense DP size / rank, which is wrong for expert sub-groups (partitioned over EDP). Every partition-native AutoEP ZeRO-3 checkpoint writes it, so raising here seems safer than guessing.