[BUG] Excessive CPU and GPU Memory Usage with Multi-GPU Inference Using DeepSpeed

I am experiencing excessive CPU and GPU memory usage when running multi-GPU inference with DeepSpeed. Specifically, the memory usage does not scale as expected when increasing the number of GPUs. Below is the code I am using for inference:

import os
import torch
import deepspeed
import time
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
from deepspeed.inference.config import DeepSpeedTPConfig
from deepspeed.runtime.utils import see_memory_usage

local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))

model_dir = "/mnt/sgnfsdata/tolo-03-97/pretrained_models/internlm2-chat-20b"
trust_remote_code = True
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=trust_remote_code)
config = AutoConfig.from_pretrained(model_dir, trust_remote_code=trust_remote_code)
model = AutoModelForCausalLM.from_pretrained(model_dir,
                                             torch_dtype=torch.bfloat16,
                                             trust_remote_code=trust_remote_code
                                             )

model = model.eval()
see_memory_usage("After load model", force=True)

tp_config = DeepSpeedTPConfig(tp_size=world_size)
zero_config = DeepSpeedZeroConfig(stage=3,
                                  model_persistence_threshold=0,
                                  max_live_parameters=0,
                                  mics_shard_size=world_size
                                  )
ds_engine = deepspeed.init_inference(model=model,
                                     tensor_parallel=tp_config,
                                     dtype=torch.bfloat16,
                                     zero=zero_config,
                                     max_out_tokens=1024,
                                     replace_method="auto",
                                     replace_with_kernel_inject=True)

see_memory_usage("After DS-inference init", force=True)

model = ds_engine.module
print("device: ", model.device)
prompt = "what is deepspeed?"
t0 = time.time()
response = model.chat(tokenizer=tokenizer,
                      query=prompt,
                      history=[],
                      max_new_tokens=1024,
                      do_sample=True,
                      temperature=0.8,
                      top_p=0.8
                      )
t1 = time.time()
print(response)
print('=' * 100)
print("inference time: ", t1 - t0)
print('=' * 100)

Steps to Reproduce:

1. Run the script with 2 GPUs:

deepspeed --num_gpus 2 main.py --ds_inference

2. Run the script with 4 GPUs:

deepspeed --num_gpus 4 main.py --ds_inference

Expected Behavior:
I expected that using 4 GPUs would reduce the memory usage per GPU, ideally halving the GPU memory usage compared to running with 2 GPUs.

Actual Behavior:

With 2 GPUs:
    CPU virtual memory: 92.87GB
    Each GPU memory: 37.74GB

With 4 GPUs:
    CPU virtual memory: 162.92GB (significantly higher than expected)
    Each GPU memory: 37.74GB (no reduction)

Questions:

Why does the CPU virtual memory usage increase significantly when using more GPUs?
How can I reduce the memory usage per GPU when scaling up the number of GPUs?

System Info:

DeepSpeed version: 0.14.4
PyTorch version: 2.3.1
Transformers version: 4.42.3
Python version: 3.10
OS: ubuntu 24.04

Additional Context:
Any insights or suggestions on how to optimize the memory usage for multi-GPU inference with DeepSpeed would be greatly appreciated. Thank you!

Start by reproducing main.py with the two- and four-GPU commands, recording the memory reports around model loading and deepspeed.init_inference. Read the DeepSpeedTPConfig and DeepSpeedZeroConfig settings used in the script and compare how tensor parallelism and ZeRO-3 handle model memory. Done means explaining the scaling behavior and identifying a tested way to reduce per-GPU memory usage.