Loading a HF Model in Multiple GPUs and Run Inferences Using deepspeed.init_inference() on those GPUs

Hi,
I loaded this model from HF using the below code on 2xA100s:

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf", use_fast=True)
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", torch_dtype=torch.float16, device_map='balanced')

Then I loaded it using deepspeed.init_inference() as:

ds_model = deepspeed.init_inference(
    model=model,
    mp_size=2,
    dtype=torch.float16,
    replace_method="auto",
    replace_with_kernel_inject=True)

and then for inferences I used:

import time
start = time.time()
prompt = 'What is the Capital of France? '
inputs = tokenizer.encode("<human>: prompt \n<bot>:", return_tensors='pt').to(model.device)
outputs = ds_model.generate(inputs, max_new_tokens=500)
output_str = tokenizer.decode(outputs[0])
print(output_str)
end = time.time()
print('Inference Time is:', end - start)

Unfortunately, I got this error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/__init__.py", line 342, in init_inference
    engine = InferenceEngine(model, config=ds_inference_config)
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/inference/engine.py", line 116, in __init__
    self._create_model_parallel_group(config)
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/inference/engine.py", line 222, in _create_model_parallel_group
    self.mp_group = dist.new_group(ranks)
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/comm/comm.py", line 185, in new_group
    return cdb.new_group(ranks)
  File "/usr/local/lib/python3.10/dist-packages/deepspeed/comm/torch.py", line 325, in new_group
    return torch.distributed.new_group(ranks)
  File "/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py", line 3313, in new_group
    raise RuntimeError(
RuntimeError: the new group's world size should be less or equal to the world size set by init_process_group

I tried the mp_size as 1 too, but in vain.

Any guidance/help would be highly appreciated, thanks in anticipation!

Start at the deepspeed.init_inference() call and follow the traceback through deepspeed/inference/engine.py::_create_model_parallel_group and torch.distributed.new_group. Reproduce the setup with the shown Hugging Face model and two GPUs, then determine how the process-group configuration relates to mp_size. Done means the supported launch/configuration and any required changes are documented or covered by a regression test.