[BUG] Fail to Resume From Checkpoint with Different GPU Number(Huggingface Trainer + Deepspeed)

Describe the bug
Hello, I encountered a problem when trying to resume from a previous checkpoint when I used Transformers Trainer + ZeRO 3 strategy to train a TinyLlama.
My previous run is conducted on 2 nodes * 2 A100 40GB on each node. The structure of the previous checkpoint is showed below.

Now I want to resume from this checkpoint with 1 node with 4 * A100 40GB, and the error below occurred.
I guess it may related to the different checkpoint format(e.g. the RNG states and model/optim states). Is there any method to consolidate the checkpoint format?
Any help will be really appreciated! For us students using HPC clusters, sometimes it's hard for us to always get the same GPU numbers or specifications. Thanks in advance!

    return inner_training_loop(
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/transformers/trainer.py", line 1708, in _inner_training_loop
    deepspeed_load_checkpoint(self.model_wrapped, resume_from_checkpoint)
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/transformers/integrations/deepspeed.py", line 402, in deepspeed_load_checkpoint
    deepspeed_load_checkpoint(self.model_wrapped, resume_from_checkpoint)
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/transformers/integrations/deepspeed.py", line 402, in deepspeed_load_checkpoint
    load_path, _ = deepspeed_engine.load_checkpoint(
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2752, in load_checkpoint
    load_path, _ = deepspeed_engine.load_checkpoint(
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2752, in load_checkpoint
        load_path, client_states = self._load_checkpoint(load_dir,load_path, client_states = self._load_checkpoint(load_dir,

  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2804, in _load_checkpoint
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2804, in _load_checkpoint
    sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine)
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py", line 43, in get_sd_loader
    sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine)
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py", line 43, in get_sd_loader
    return MegatronSDLoader(ckpt_list, version, checkpoint_engine)
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py", line 193, in __init__
    super().__init__(ckpt_list, version, checkpoint_engine)
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py", line 55, in __init__
    return MegatronSDLoader(ckpt_list, version, checkpoint_engine)
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py", line 193, in __init__
    self.check_ckpt_list()
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py", line 168, in check_ckpt_list
    super().__init__(ckpt_list, version, checkpoint_engine)
  File "/home/users/nus/miniconda3/envs/gpu_dev/lib/python3.10/site-packages/deepspeed/runtime/state_dict_factory.py", line 55, in __init__
    assert len(self.ckpt_list) > 0
AssertionError

To Reproduce
Steps to reproduce the behavior:

1. Go to '...'
2. Click on '....'
3. Scroll down to '....'
4. See error

Expected behavior
A clear and concise description of what you expected to happen.

ds_report output
Please run ds_report to give us details about your setup.

Screenshots

System info (please complete the following information):

• OS: [e.g. Ubuntu 18.04]
• GPU count and types [e.g. two machines with x8 A100s each]
• Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB]
• Python version
• Any other relevant info about your setup

Launcher context
Are you launching your experiment with the deepspeed launcher, MPI, or something else?

Docker context
Are you using a specific docker image that you can share?

Additional context
Add any other context about the problem here.