[BUG] Unable to run my model with Tensor Parallelism

Describe the bug
I'm trying to run a simple FFN with the deepspeed inference engine in a tensor parallel manner. But when I monitor the GPU usage, i see that my model only runs on 1 GPU.
How do i use the inference engine to run this in a tensor parallel manner?

To Reproduce
Steps to reproduce the behavior:

import deepspeed
import torch
import torch.nn as nn
import os
from transformers import pipeline
from transformers.models.t5.modeling_t5 import T5Block
from deepspeed.inference.config import DeepSpeedTPConfig
tp = {
    "enabled": True,
    "tensor_parallel_size": 4,
}
tp_config = DeepSpeedTPConfig(tp)

local_rank = int(os.getenv('LOCAL_RANK', '0'))
world_size = int(os.getenv('WORLD_SIZE', '4'))

# Init translator
# translator = pipeline("translation_en_to_fr", model="t5-base", tokenizer="t5-base", device=local_rank)
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(768, 12288)
        self.fc2 = nn.Linear(12288, 256)
        self.fc3 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x
model = MyModel()
input = torch.randn(1, 768).half().to('cuda')
# DeepSpeed init_inference API
model = deepspeed.init_inference(model,
    tp = tp_config,
    dtype=torch.half,
    #injection_policy={T5Block: ('SelfAttention.o', 'EncDecAttention.o', 'DenseReluDense.wo')}
)

output = model(input)
print(output.shape)

run with

deepspeed --num_gpus 4 ds.py
Expected behavior
A clear and concise description of what you expected to happen.

ds_report output

root@50e47578e691:/ramyapra# ds_report
[2023-09-29 08:57:50,537] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
mp_size default=1 extra={'deprecated': True, 'new_param': 'tensor_parallel.tp_size'}
--------------------------------------------------
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]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.0
 [WARNING]  using untested triton version (2.0.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 ............... ['/usr/local/lib/python3.8/dist-packages/torch']
torch version .................... 2.0.1+cu117
deepspeed install path ........... ['/usr/local/lib/python3.8/dist-packages/deepspeed']
deepspeed info ................... 0.10.2+d7e2ede, d7e2ede, rnp-moe
torch cuda version ............... 11.7
torch hip version ................ None
nvcc version ..................... 11.8
deepspeed wheel compiled w. ...... torch 2.0, cuda 11.7
shared memory (/dev/shm) size .... 220.45 GB

Screenshots
If applicable, add screenshots to help explain your problem.

Docker context
Containerised version of nvcr.io/nvidia/pytorch:22.11-py3

Start with the ds.py reproduction, especially DeepSpeedTPConfig, deepspeed.init_inference, and the deepspeed --num_gpus 4 launch command. Read the inference tensor-parallel entry points and related configuration handling, then verify the model uses all four GPUs and document the required behavior or configuration.