Fix AutoEP ZeRO-1/2 universal conversion (#8198)
Fixes #8147
## Problem
AutoEP checkpoints saved with ZeRO stages 1 and 2 store expert optimizer
state in ZeRO-sharded fragments, but Universal Checkpoint conversion
currently merges those expert fragments through the generic path before
AutoEP consolidation.
With source EP size greater than one, conversion can combine different
EP ranks and fail while reshaping the result. The fallback path can also
produce expert tensors without repartition metadata and use
mixed-precision model weights where FP32 optimizer masters are required.
## Approach
- Identify fused AutoEP expert parameters from checkpoint metadata
before the generic merge, then reconstruct their FP32 masters and Adam
states from the ZeRO fragments in EP-rank order.
- Use the saved topology setting to handle both expert-before-data and
data-before-expert rank layouts.
- Save the existing expert metadata with every expert state so Universal
Checkpoint load can repartition to a different EP size.
- Restrict pipeline layer discovery to complete pipeline checkpoint
filenames so per-expert files are not misclassified.
---------
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>