[jinyouzhi]

A universal checkpoint converted from a ZeRO-3 run cannot be loaded into ZeRO-1/2, and one converted from a ZeRO-1/2 run cannot be loaded into ZeRO-3. This PR makes both directions load, with the fp32 master weights, Adam states and step restored exactly.

Reproduction

Train under one stage, save_checkpoint, run ds_to_universal, then load_checkpoint(..., load_optimizer_states=True) with checkpoint.load_universal: true under the other stage. On master a75ab89 each gap below surfaces once the previous one is worked around:

#	Direction	Error	Root cause
1	both	assert len(self.ckpt_list) > 0	_get_all_ckpt_names globs only the current stage's model-states name: zero_pp_rank_0_mp_rank_*_model_states.pt for ZeRO-3, mp_rank_*_model_states.pt for ZeRO-1/2. A universal load reads only metadata from this file, so either name works.
2	3 → 1/2	AssertionError: Empty ds_version	_save_optimizer_state_stage3 writes the inner optimizer dict, but ZeRO-3 keeps ds_version beside it. DeepSpeedZeroOptimizer._load_global_state asserts on it.
3	3 → 1/2	IndexError: too many indices for tensor of dimension 1	The ZeRO-3 converter (merge_zero3_slices) writes each state as a bare flat tensor; load_hp_checkpoint_state indexes ckpt_dict[PARAM]. The stage-3 loader already accepts both forms.
4	3 → 1/2	step silently not restored	A ZeRO-3 source writes no per-parameter step.pt; its step is in optimizer_state_dict.state.
5	1/2 → 3	KeyError: 'optimizer_state_dict'	The stage-3 loader reads the step and param_groups from optimizer_state_dict, which a ZeRO-1/2 source does not have (it has top-level param_groups and per-parameter step.pt).

Changes

• engine._get_all_ckpt_names: on a universal load with no match, glob the other stage's name, built by _get_ckpt_name (new optional partition_weights argument).
• ds_to_universal._save_optimizer_state_stage3: copy ds_version into zero/optimizer_state.pt.
• universal_checkpoint.py:
  • load_hp_checkpoint_state views a bare flat tensor as the parameter's shape. A flat tensor carries no TP layout, so if it does not fill the whole parameter (ZeRO-3 source → ZeRO-1/2 with TP>1) it raises a ValueError saying so, instead of an IndexError.
  • New load_universal_optimizer_step reads the step from either layout. It is used by the ZeRO-1/2 path (ZeROOptimizer.load_hp_checkpoint_state_from_checkpoint_dir, when no step.pt exists) and by the ZeRO-3 path.
• stage3.load_hp_checkpoint_state_from_checkpoint_dir_stage3: take param_groups from the top level when present (the ZeRO-3 converter has always written it there too).

Same-stage loads take the same paths as before: every new branch is reached only when the old lookup finds nothing.

Sibling paths checked: BF16_Optimizer shares load_hp_checkpoint_state_from_checkpoint_dir and its _load_global_state is a no-op, so it gets fixes 3 and 4 and is unaffected by 2. stage_1_and_2 and stage3 are both patched.

Not a duplicate

No existing issue or PR covers cross-stage universal loading. #9069 also touches bare flat tensors but fixes the ZeRO-3 + AutoTP load path, not ZeRO-1/2 targets. #7546, #8848 and #8898 are different bugs.

Tests

New TestUniversalCheckpointAcrossZeROStages in tests/unit/checkpoint/test_universal_checkpoint.py (world size 2, fp32, torch.optim.Adam), parametrized over (1→3), (2→3), (3→1), (3→2). It trains 2 steps, saves, converts, loads into the other stage, then checks:

• fp32 master, exp_avg and exp_avg_sq of every parameter are torch.equal to the source (the conversion only moves fp32 data, so the comparison is exact);
• the step is 2;
• one more training step runs.

Hardware: Intel Arc Pro B70 (XPU), Linux 6.19.10 in-tree xe driver, compute runtime 26.27.39122.11, torch 2.14.0.dev20260804+xpu. Run from the repo root.

pytest -q tests/unit/checkpoint/test_universal_checkpoint.py::TestUniversalCheckpointAcrossZeROStages

Tree	B70	NVIDIA L20 (driver 580.173.02, torch 2.14.1+cu130)
master a75ab89	4 failed (gap 1: assert len(self.ckpt_list) > 0)	4 failed, same assert
this PR	4 passed	4 passed

Regression on B70, master and this PR in the same environment:

pytest -q tests/unit/checkpoint/test_universal_checkpoint.py tests/unit/checkpoint/test_autotp_uc_checkpoint.py tests/unit/checkpoint/test_autoep_affine_universal.py

Tree	Result (219 tests)
master a75ab89	30 failed, 189 passed
this PR	25 failed, 194 passed

• The 4 extra failures on master are the new TestUniversalCheckpointAcrossZeROStages cases.
• All other failures occur on both trees and are unrelated to this change:
  • TestZeROUniversalCheckpointDP with ZeRO-1 and load_optim=True: the bit-exact optimizer-state check fails intermittently on this XPU setup. Which cases fail changes from run to run; re-running 7 of them gave 5/7 failures with this PR and 4/7 and 5/7 on master.
  • TestRealCheckpointUniversalConversionTPxPP: No backend type associated with device type xpu.

End to end on B70 (2 ranks): Qwen3-0.6B, bf16, torch.optim.AdamW. Train 2 steps under one stage, save, ds_to_universal, load under the other stage, and compare safe_get_full_fp32_param / safe_get_full_optimizer_state before and after:

Direction	Parameters	Mismatched (fp32, exp_avg, exp_avg_sq; exact)	Step after one more step
ZeRO-3 → ZeRO-1	310	0	3
ZeRO-1 → ZeRO-3	310	0	3

Validated matrix cells

• Validated cross-stage: ZeRO 1/2 ↔ 3 with fp32 (unit test, XPU and CUDA) and ZeRO 1 ↔ 3 with bf16 (Qwen3-0.6B on XPU); DP=2; no offload; no TP/PP.
• Same-stage universal loads: covered by the regression suite above (DP resizing, AutoTP, AutoEP).
• Not validated cross-stage: fp16, CPU/NVMe offload, TP/PP, MoE. The change has no device-specific code.

Limitations

• A checkpoint converted from ZeRO-3 before this PR lacks ds_version, so loading it into ZeRO-1/2 still fails the version assert; re-run ds_to_universal on the original checkpoint.
• ZeRO-3 source → ZeRO-1/2 target with TP>1 is not supported (the ZeRO-3 converter records no TP layout); it now fails with a clear error.

AI assistance

This PR was prepared with AI assistance (GitHub Copilot CLI), covering the root-cause analysis, the code, the tests and this description. I reviewed the changes and ran the tests reported above on the hardware listed.