`ZeRO-3 + MP8` Universal Checkpoint

Is it possible to convert a model trained using ZeRO-3 and MP=8 to a universal checkpoint?

Tracing through the universal checkpointing conversion tool (ds_to_universal), the model states remained unmerged, with 8 model parallel shards per each data parallel rank. E.g., with world_size = 2048, there are 2048 model state files,zero_pp_rank_{0-255}_{0-7} before and after the conversion.

When converting a model with ZeRO <= 2, MP > 1, the model state files are merged into a single file through merge_tp_slices.

If this is not possible, how would one extract and merge only the Z3 / MP checkpointed model states (along both z3 and model parallel partitions) to a single file?

The zero_to_fp32 script does not work since it only handles ZeRO-{2,3} without model parallelism.

Start with the universal checkpointing conversion tool, ds_to_universal, and trace how merge_tp_slices handles ZeRO and model-parallel partitions. Compare this with zero_to_fp32, reproducing the ZeRO-3 and MP=8 case described in the issue. Done means the model states can be merged into a single universal checkpoint, or the supported extraction limitation and path are documented.