[REQUEST] Does Universal Checkpoint supports for MoE Checkpoint？ about deepspeed HOT 3 OPEN

@tiggerwu Can you tell us which MoE implementation you used?

from deepspeed.

@tiggerwu Can you tell us which MoE implementation you used?

Can you also share the deepspeed version and the stacktrace of the error?

from deepspeed.

@tiggerwu Can you tell us which MoE implementation you used?

Can you also share the deepspeed version and the stacktrace of the error?

Hello @tohtana , many thanks for your response.

1. I have implemented a GPT-3 small 125M model with 2 experts using the script https://github.com/microsoft/Megatron-DeepSpeed/blob/main/examples_deepspeed/MoE/ds_pretrain_gpt_125M_MoE64.sh in Megatron-Deepspeed. I trained it in 2 GPUs with DP=2 and ZeRO1. When I removed the layer_x_expert_x_mp_rank_00_model_states.pt file from checkpoint folder (Because it doesn't match the file pattern defined in ds_to_universal.py and will cause error) and used ds_to_universal.py to convert the saved checkpoint to universal checkpoint, I found there is no param_shapes key in mp_rank_00_model_states.pt and the error KeyError: 'param_shapes' raised.
2. The deepspeed version is 0.14.3
3. The error stacktrace:

Traceback (most recent call last):
  File "/path/to/deepspeed/checkpoint/ds_to_universal.py", line 401, in <module>
    main(args)
  File "/path/to/deepspeed/checkpoint/ds_to_universal.py", line 368, in main
    slice_shapes += mp_sd[PARAM_SHAPES]
KeyError: 'param_shapes'

from deepspeed.