Need help with configuring data and tensor parallel inference

[sarathkondeti]

Hey guys, I'm trying to run vicuna 13b in 2d-parallel config.
I was able to do tensor parallelism using tp_size:2 on gpu(0,1).
How do I add data parallelism, so that both the gpu pairs (0,1 & 2,3) would have DP and each pair has TP.

Setup 4xV100(16gb)
Attaching my current script below.

Any help would be greatly appreciated.
Thanks.

import os
import deepspeed
from transformers import LlamaForCausalLM, LlamaTokenizer
import argparse

ds_config = {
  "dtype": "torch.float16",
  "replace_with_kernel_inject": False,
  "tensor_parallel": {
    "tp_size": 2
  }
}
model_name = "lmsys/vicuna-13b-v1.5"

deepspeed.init_distributed()

local_rank = int(os.environ.get('LOCAL_RANK', 0))
world_size = int(os.environ.get('WORLD_SIZE', 1))

tokenizer = LlamaTokenizer.from_pretrained(model_name)
model = LlamaForCausalLM.from_pretrained(model_name)
model = deepspeed.init_inference(
    model,
    config = ds_config,
)

batch = tokenizer(
    "The primary use of LLaMA is research on large language models, including",
    return_tensors="pt",
    add_special_tokens=False
)
# here
batch = {k: v.cuda(local_rank) for k, v in batch.items()}
generated = model.generate(batch["input_ids"], max_length=100)
print(tokenizer.decode(generated[0]))