[BUG] Distributed language model inference takes too much memory.

Describe the bug
A 20b model on huggingface, distributed 4-ways in fp16 precision, takes up to 24GB memory during inference.

To Reproduce
Steps to reproduce the behavior:

1. Simple inference script to reproduce
   `
   import torch
   import torch.distributed as dist
   import torch.nn.functional as F
   import numpy as np
   import deepspeed
   import os
   import time
   from transformers.deepspeed import HfDeepSpeedConfig
   from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM
   from datasets import load_dataset

if name == "main":
# print(torch.cuda.device_count())
local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
dist.init_process_group(backend='nccl', rank=local_rank, world_size=world_size)

# Load the TriviaQA dataset
dataset = load_dataset("trivia_qa", "rc.nocontext")
test_dataset = dataset["validation"]
train_dataset = dataset["train"]

train_pairs = []
count = 0
for example in train_dataset:
    train_pairs.append({'question': example['question'], 'answer': example['answer']['value']})
    count+=1
    if count >=5:
        break

# Extract and format question-answer pairs
qa_pairs = [{'question': example['question'], 'answer': example['answer']['value']} for example in test_dataset]

# Construct a prompt with 5 examples
prompt = ""
for i, pair in enumerate(train_pairs):
    prompt += f"Example {i+1}:\nQuestion: {pair['question']}\nAnswer: {pair['answer']}\n\n"

batch_size = 48

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-1.4b", padding_side="left", torch_dtype=torch.float16)
tokenizer.pad_token = tokenizer.eos_token

oracle_model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neox-20b", torch_dtype=torch.float16)


oracle_model = deepspeed.init_inference(
                oracle_model,
                replace_with_kernel_inject=False,
                tp={"tp_size": 4,},
                dtype=torch.float16,
                #checkpoint=checkpoint_dict,
               )

current_prompts = []
curr_count = 0

for pair in qa_pairs:
    # Constructing the prompt for each question
    current_prompts.append("Question: " + pair['question'])
    draft_input_ids = tokenizer.batch_encode_plus(current_prompts, padding='longest')
    curr_count += 1
    if curr_count == 48:
        curr_count = 0
        input_tensors = torch.tensor(draft_input_ids['input_ids']).cuda(local_rank)
        oracle_model.generate(input_tensors, max_new_tokens=10, pad_token_id=tokenizer.eos_token_id, use_cache=False)
        input_tensors.cpu()
        current_prompts = []
    else:
        continue

`
3. What packages are required and their versions
Python 3.10, Latest Transformers/deepspeed
5. How to run the script
deepspeed --num_gpus 4 test_load_deepspeed.py

Expected behavior
Peak GPU memory usage should be not much more than 10GB.

ds_report output
DeepSpeed C++/CUDA extension op report
NOTE: Ops not installed will be just-in-time (JIT) compiled at
runtime if needed. Op compatibility means that your system
meet the required dependencies to JIT install the op.
JIT compiled ops requires ninja
ninja .................. [OKAY]
op name ................ installed .. compatible
[WARNING] async_io requires the dev libaio .so object and headers but these were not found.
[WARNING] async_io: please install the libaio-dev package with apt
[WARNING] If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
[WARNING] Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
[WARNING] sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.1
[WARNING] using untested triton version (2.1.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
DeepSpeed general environment info:
torch install path ............... ['/users/myan/anaconda3/envs/deepspeed/lib/python3.10/site-packages/torch']
torch version .................... 2.1.0
deepspeed install path ........... ['/users/myan/anaconda3/envs/deepspeed/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.11.2, unknown, unknown
torch cuda version ............... 11.8
torch hip version ................ None
nvcc version ..................... 11.8
deepspeed wheel compiled w. ...... torch 2.1, cuda 11.8
shared memory (/dev/shm) size .... 251.78 GB

Screenshots
Peak memory usage monitored by Nvidia-smi is around 24GB.

System info (please complete the following information):

• OS: Ubuntu 20.04
• GPU count and types: One machine with four 40GB A100s.
• (if applicable) Hugging Face Transformers/Accelerate/etc. versions: 4.35.0 (Reproducible in older version such as 4.31 as well).
• Python version: 3.10

Start with the reproduced test_load_deepspeed.py script and the deepspeed.init_inference entry point using tensor parallelism. Run it with the reported Python, PyTorch, CUDA, Transformers, and DeepSpeed versions while monitoring peak GPU memory, then compare the result with the expected usage of not much more than 10GB per GPU.