ds_inference success but OOM when use tp_presharded_mode=True[BUG]

Describe the bug
A clear and concise description of what the bug is.
I use transformers-bloom-inference to inference bloom 176b with 16 * v100 (32G) GPUs, fp16. Following is the transformers-bloom-inference:
https://github.com/huggingface/transformers-bloom-inference/blob/main/bloom-inference-scripts/bloom-ds-inference.py#L181

The core code in my try is :

if os.path.isfile(os.path.join(repo_root, "ds_inference_config.json")):
    tp_presharded_mode = True
else:
    tp_presharded_mode = False
if tp_presharded_mode:
    # tp presharded repos come with their own checkpoints config file
    checkpoints_json = os.path.join(repo_root, "ds_inference_config.json")
else:
    # for normal bloom repo we need to write the checkpoints config file
    write_checkpoints_json()
    kwargs["save_mp_checkpoint_path"] = os.path.normpath(repo_root) + "_tp_presharded"
    dist.barrier()

model = deepspeed.init_inference(
    model,
    mp_size=world_size,
    base_dir=repo_root,
    dtype=getattr(torch, infer_dtype),
    checkpoint=checkpoints_json,
    **kwargs,
)

I run bloom-ds-inference.py to inference bloom 176b. Firstly, I run with normal bloom repo (tp_presharded_mode = False). It successed and saved a presharded checkpoint. When I run with the presharded checkpoint (tp_presharded_mode = True), the loading is faster but it OOM in deepspeed.init_inference.

How to reduce memory usage with fp16?

related issue:
https://github.com/huggingface/transformers-bloom-inference/issues/92

Start with bloom-ds-inference.py around the referenced line 181 and compare the normal-checkpoint and tp_presharded paths using ds_inference_config.json. Trace the call into deepspeed.init_inference and reproduce the fp16 run with 16 V100 GPUs. Done means the presharded path no longer runs out of memory and the inference initialization succeeds.