Universal Checkpointing: read AutoEP ZeRO-3 optimizer files one at a time

[pengdurice]

Universal Checkpointing: read AutoEP ZeRO-3 optimizer files one at a time

Problem

ds_to_universal converts a ZeRO-3 AutoEP checkpoint in two parts. The dense parameters go through the usual
extract-and-merge path, which reads one optimizer file per worker and writes fragments to tmp/. The routed experts
go through _consolidate_zero3_autoep_expert_states, which first loads every rank's optimizer file into one dict
(optim_states_by_rank) and then keeps a copy of every expert fragment in memory until all of them are merged.

Together those files hold every optimizer state of the model: fp32 (or bf16) master weights plus both Adam moments.
For a large MoE model that is more than one host has:

Model	Parameters	Optimizer state, fp32 master + fp32 Adam (12 B/param)	bf16 master + bf16 Adam (6 B/param)
MiniMax-M3	426.2B	4.65 TiB	2.33 TiB
GLM-5.2	743.4B	8.11 TiB	4.06 TiB

The H200 nodes this was found on have 2 TiB of host memory, so the conversion cannot finish for either model.

Change

_consolidate_zero3_autoep_expert_states reads one rank's optimizer file at a time. Each expert fragment it finds is
written to <output>/tmp/autoep_expert_fragments/ (one file per parameter, optimizer state, EP rank and partition
rank) and the rank's file is released before the next one is read. The merge then loads, for one parameter and one
state at a time, only that parameter's fragments. Peak host memory becomes one optimizer file plus one full expert
tensor, instead of all optimizer files.

The checks are unchanged: missing ranks, duplicate or missing expert-data-parallel partitions, shape mismatches and
placement disagreements raise the same errors. The temporary fragments live under <output>/tmp, which the converter
already uses for the dense fragments and removes at the end.

Tests

Test	Result
tests/unit/checkpoint/test_autoep_affine_universal.py (calls the function directly; CPU)	20 passed (master 1c862f7 + this change); master 5f266e2 + this change, both files together: 63 passed
tests/unit/v1/moe/test_autoep_checkpoint.py (AutoEP ZeRO-3 save, convert, load at other layouts; GPU)	43 passed

Hardware for the GPU tests: 1 node of 8 H200 (driver 580.126.09), torch 2.8.0+cu126, pytest 9.1.1.
pre-commit run --files passes on the changed file.

Not measured: the peak host memory of a full-size conversion before and after. The bound follows from the code:
before, every rank's optimizer file is loaded at once (for MiniMax-M3 with fp32 master weights and fp32 Adam states,
4.65 TiB in total, more than the 2 TiB of the H200 hosts this was found on); after, one rank's file plus one
parameter's fragments.

Duplicate check

Searched open and recent pull requests on 2026-10-15 ("universal checkpoint", "ds_to_universal", "AutoEP
checkpoint"): #8800 plans live expert transfers between placements and #8637 links ZeRO-1/2 fragments; neither
changes how ds_to_universal reads the ZeRO-3 AutoEP optimizer files.

Compatibility

Cell	Validated
ZeRO-3 + AutoEP checkpoint conversion, save / convert / load at other layouts	yes (tests)
Dense (non-AutoEP) ZeRO-3 conversion	path unchanged
ZeRO-1 / ZeRO-2 conversion	path unchanged

AI assistance

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