Issue with Falcon/Falcoder while trying to use it on AWS EC2 Inferentia 2.8xlarge Instance

We are facing issues while using this model on the aforementioned machine. We were able to run the same experiment on G5 instance successfully but we are observing that the same code is not working on Inf2 machine instance. We are aware that it has Accelerator instead of NVIDIA GPU. Hence we tried the neuron-core's capability and added required helper code for using the capability of neuron-cores of the instance by using the torch-neuronx library. The code changes and respective error screenshots are provided below for your reference:

Code without any torch-neuronx usage - Generation code snippet:

generation_output = model.generate( input_ids = input_ids, attention_mask = attention_mask, generation_config = generation_config, return_dict_in_generate = True, output_scores = False, max_new_tokens = max_new_tokens, early_stopping = True ) #print("generation_output") #print(generation_output) s = generation_output.sequences[0] output = tokenizer.decode(s)

Code using torch-neuronx - helper function code snippet:

def generate_sample_inputs(tokenizer, sequence_length): dummy_input = "dummy" embeddings = tokenizer(dummy_input, max_length=sequence_length, padding="max_length",return_tensors="pt") return tuple(embeddings.values())

def compile_model_inf2(model, tokenizer, sequence_length, num_neuron_cores):

use only one neuron core os.environ["NEURON_RT_NUM_CORES"] = str(num_neuron_cores) import torch_neuronx payload = generate_sample_inputs(tokenizer, sequence_length) return torch_neuronx.trace(model, payload)

model = compile_model_inf2(model, tokenizer, sequence_length=512, num_neuron_cores=1)

Can this github issue address our specific problems mentioned above? https://github.com/oobabooga/text-generation-webui/issues/2260

My queries are basically:

1.Is it even feasible to do inference on this machine or should we go for G4dn.8xlarge as we are facing so many issues in Inf2? 2. Can we try Llama 2 on Inferentia 2 8xlarge machine or this is not supported? If not, which machine instance we should try considering cost-effectiveness?