universal resume checkpoint from deepspeed

[pengzhangzhi]

Description & Motivation

One pain in training with deepspeed is that when resume from a checkpoint you have to use the same amount of gpus as the num of gpus that the checkpoint was trained on. otherwise, you will see the following error:

deepspeed.runtime.zero.utils.ZeRORuntimeException: The checkpoint being loaded used a DP world size of 32 but the current world size is 128. Automatic adjustment of ZeRO's optimizer state partitioning with a new world size is not currently supported.

see this issue. deepspeedai/DeepSpeed#3810

Also when the model is well trained, and we want to paly with the inference, it would be a problem to load a deepspeed ckpt for inference as it requires the same num of GPUs as the training?

But currently, Deepspeed proposes a universal checkpointing to convert the deepspeed ckpt to universal ckpt, which can be loaded in whatever many of gpus.
Please refer to the link
https://github.com/microsoft/Megatron-DeepSpeed/tree/main/examples_deepspeed/universal_checkpointing#zero-stage-2-training

Would lightning integrate this feature?

Also a weird usage, is that any way to load only the checkpoint of the model while ignoring other state ckpt like optimizer?

Pitch

No response

Alternatives

No response

Additional context

No response

cc @Borda @awaelchli

[awaelchli]

@pengzhangzhi Do you mean this?

https://lightning.ai/docs/pytorch/stable/advanced/model_parallel/deepspeed.html#collating-single-file-checkpoint-for-deepspeed-zero-stage-3

After converting a checkpoint like shown there, you should be able to load it on a different number of GPUs.