Loading DeepSpeed sharded weights

[prabhuteja12]

Hello!

I'm saving a model that is being trained with Deepspeed stage 3 using the ModelCheckpoint callback. It saves multiple .pt files in a folder. However, load those weights, I can't seem to understand how to load them in a sharded fashion ie my code currently calls convert_zero_checkpoint_to_fp32_state_dict using from pytorch_lightning.utilities.deepspeed import convert_zero_checkpoint_to_fp32_state_dict to convert the checkpoint to a single model and loads it using MyModule.load_from_checkpoint and then pass to Trainer to continue training. My understanding is that Deepspeed allows for loading those shards directly without going through this conversion framework. Can someone help me with how to do this?