[BUG] Universal checkpointing doesn't work when changing model parallel size (pp and dp change are ok)

Hello,

I am using Deepspeed 0.14.4, and converting checkpoints for a model. I can change the data and pipeline parallel size ok, and it will convert, load and resume training ok. However, when I try to change the model parallel size, it complains that the numel do match in my layernorms. In the documentation, it says I should be able to change the mp size. I go from mp=2, universal checkpoint convert, then use mp=1. I am using zero-1 optimizer.

I am using a forked version of GPT-neox also.

Here's the error I get. Does that seem right?

[rank2]: Traceback (most recent call last):                                                                                                            
[rank2]:   File "/home/etnguyen/savanna/train.py", line 18, in <module>                                                                                
[rank2]:     pretrain(global_config=global_config)                                                                                                     
[rank2]:   File "/home/etnguyen/savanna/savanna/training.py", line 172, in pretrain                                                                    
[rank2]:     model, optimizer, lr_scheduler = setup_model_and_optimizer(global_config=global_config, use_cache=False)                                  
[rank2]:                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                  
[rank2]:   File "/home/etnguyen/savanna/savanna/training.py", line 894, in setup_model_and_optimizer                                                   
[rank2]:     global_config.iteration = load_checkpoint(                                                                                                
[rank2]:                               ^^^^^^^^^^^^^^^^                                                                                                
[rank2]:   File "/home/etnguyen/savanna/savanna/checkpointing.py", line 224, in load_checkpoint                                                        
[rank2]:     checkpoint_name, state_dict = model.load_checkpoint(                                                                                      
[rank2]:                                   ^^^^^^^^^^^^^^^^^^^^^^                                                                                      
[rank2]:   File "/home/etnguyen/miniconda3/envs/bh_evo2/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 2770, in load_checkpoint       
[rank2]:     success = self._load_zero_checkpoint(load_dir, tag, load_optimizer_states=load_optimizer_states)                                          
[rank2]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                          
[rank2]:   File "/home/etnguyen/miniconda3/envs/bh_evo2/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 2959, in _load_zero_checkpoint 
[rank2]:     self.optimizer.load_state_dict(state_dict_list=zero_sd_list,                                                                              
[rank2]:   File "/home/etnguyen/miniconda3/envs/bh_evo2/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 2291, in load_state
_dict                                                                                                                                                  
[rank2]:     self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights)                                         
[rank2]:   File "/home/etnguyen/miniconda3/envs/bh_evo2/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 2296, in _load_univ
ersal_checkpoint                                                                                                                                       
[rank2]:     self.load_hp_checkpoint_state_from_checkpoint_dir("bit16_groups", checkpoint_folder)                                                      
[rank2]:   File "/home/etnguyen/miniconda3/envs/bh_evo2/lib/python3.11/site-packages/deepspeed/runtime/base_optimizer.py", line 47, in load_hp_checkpoi
nt_state_from_checkpoint_dir                                                                                                                           
[rank2]:     step = lp.load_hp_checkpoint_state(os.path.join(checkpoint_dir, self.param_names[lp]), tp_rank,                                           
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                           
[rank2]:   File "/home/etnguyen/miniconda3/envs/bh_evo2/lib/python3.11/site-packages/deepspeed/checkpoint/universal_checkpoint.py", line 82, in load_hp
_checkpoint_state                                                                                                                                      
[rank2]:     assert full_param_numel == tp_world_size * tp_slice_numel, \                                                                              
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                 
[rank2]: AssertionError: Loading /scratch/hielab/etnguyen/checkpoints/evo2/og2_v2/100m_6h_3m_3s_2a_10k_524k_mp_test3/global_step4_universal/zero/9.inpu
t_layernorm.scale/exp_avg_sq.pt full param numel 1536 != tensor slice numel 768 * tp_world_size 1   

Start with deepspeed/checkpoint/universal_checkpoint.py, especially load_hp_checkpoint_state, and trace its call from deepspeed/runtime/zero/stage_1_and_2.py. Reproduce the mp=2 to mp=1 conversion with the reported Zero-1 configuration and compare the checkpoint tensor sizes; done means the converted checkpoint loads and resumes training without the layernorm assertion.