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