DeepSpeed
da066407 - Fix AutoEP ZeRO-1/2 universal conversion (#8198)

Commit
11 days ago
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>
Author
Parents
Loading