Unable to inference Falcon-40b with deepspeed.

Hi Team,

I am trying to inference falcon-40b on my g5-48xlarge (8gpu - 23 gb each) machine. But unable to proceed.
Tried to follow with Auto-TP test-falcon tutorial (provided by @RezaYazdaniAminabadi ), but ending up with below error on deepspeed.init_inference:

Error:
RuntimeError: the new group's world size should be less or equal to the world size set by init_process_group

Code:

from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
import transformers
import torch
import deepspeed
import time
from deepspeed.accelerator import get_accelerator

model = "tiiuae/falcon-40b"

tokenizer = AutoTokenizer.from_pretrained(model, use_fast=True)


model = AutoModelForCausalLM.from_pretrained(model, trust_remote_code=True).bfloat16()
model = deepspeed.init_inference(model, mp_size=4)

Environment:

Python - 3.10.11
Deepspeed - 0.10.2
Torch - 2.0.1+cu117
Transformers - 4.33.1

Can someone please advice ?

Begin with the linked Auto-TP test-falcon tutorial and the deepspeed.init_inference(model, mp_size=4) call; reproduce the failure with the listed Python, DeepSpeed, PyTorch, and Transformers versions. Done would require a documented reproducible path for Falcon-40b inference or a precise compatibility diagnosis, neither of which the issue currently specifies.