[REQUEST] Partitioning the model states and optimizer states separately when resuming from checkpoint

Is your feature request related to a problem? Please describe.
I'm currently working on fine-tuning an LLM using LoRA, but the deepspeed checkpoints save the partitioned model states and optimizer states, which are incredibly large (I would only like to store the adapter model + optimizer states). I want to be able to load a checkpoint using my frozen PLM + adapter weights, partition those weights, and then load my optimizer states and LR_scheduler states saved. Currently, if we try to load the checkpoint directory without *model_states.pt, we reach an assertion error since the model states aren't there: https://github.com/lm-sys/FastChat/pull/1280

Describe the solution you'd like
I would like to be able to load from checkpoint-x and global_stepx using only the zero_pp_rank_x_mp_rank_xx_optim_states.pt and providing a model path with the frozen PLM weights (not partitioned).

Start at the checkpoint-loading path that raises an assertion when *model_states.pt is absent, using the zero_pp_rank_x_mp_rank_xx_optim_states.pt files and global_stepx as inputs. Trace how the frozen PLM and adapter weights are partitioned, then verify that optimizer and LR_scheduler states can resume without partitioned model-state files.