[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>