Support DeepSpeed checkpoints with DeepSpeed Inference

As discussed it would be really cool if DeepSpeed trained models that have been saved via deepspeed_model.save_checkpoint(...) be plugged into DeepSpeed Inference! Currently it only supports Megatron LM models/models that are pre-loaded and then loaded into the engine (Like HF Transformer Models).

Ideally something like:

model = deepspeed.init_inference(
    model,
    mp_size=self.num_processes,
    dtype=torch.half
    replace_method='auto'
)

model.load_from_checkpoint('deepspeed.ckpt') # expose load_from_checkpoint function, the same as the DeepSpeedEngine!

cc @RezaYazdaniAminabadi

Let me know where I can help! I think this really boils down to providing split/merge functions under a new State Dict Factory: https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/state_dict_factory.py#L103-L137

Start with deepspeed/runtime/state_dict_factory.py at the referenced State Dict Factory section, then review the init_inference entry point and the existing DeepSpeedEngine checkpoint-loading behavior. The work is done when a model saved with deepspeed_model.save_checkpoint(...) can be loaded through the proposed inference-engine API, including the required split/merge behavior.