DeepSpeed
eec237ee - [AutoTP] Fix ZeRO-3 checkpoint consolidation to gather across TP and DP (#8168)

Commit
2 days ago
[AutoTP] Fix ZeRO-3 checkpoint consolidation to gather across TP and DP (#8168) AutoTP + ZeRO-3 silently produced incomplete checkpoints: both export paths handled only the ZeRO data-parallel dimension and dropped the tensor-parallel shards. - ds_to_universal.py: stage3 conversion recovers the (tp,dp) grid from checkpoint file names, extracts shards under the real tp_index, and reuses the stage<=2 TP-aware merge when tp_degree>1 (DP-only path preserved for tp_degree==1 -> no regression for plain ZeRO-3). - engine.py: _zero3_consolidated_16bit_state_dict nests GatherReplacedLayerParams inside GatheredParameters so save_16bit_model gathers both DP and TP; remove the blanket autotp+zero3 training block now that checkpoint consolidation is implemented. - stage3.py: load_hp_checkpoint_state resolves the TP shard before the ZeRO-DP partition, so universal checkpoint restore round-trips. Add end-to-end universal conversion tests and update existing tests for the refactored merge_tp_slices / extract_zero_shards_stage3 signatures. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Author
Parents
Loading