
This repo shows how you can utilize AWS-designed silicon to run inference on Phind's version of Codellama-34B! I ran this model on HumanEval locally and was getting 53.757082802180754 tokens per second running on an [inf2.48xlarge](https://instances.vantage.sh/aws/ec2/inf2.48xlarge).  

The example below shows a single sample.

#### 509 tokens in 9.498470783233643 seconds

Human: You are an intelligent programming assistant. How do I implement a linked list in Rust?

To implement a linked list in Rust, you need to define a structure (or struct) that will represent your linked list nodes. Each node should have some data and a link to the next node in the list. Here's a basic implementation:

```rust
// Define a Node of the LinkedList
pub struct Node<T> {
    data: T,
    next: Option<Box<Node<T>>>,
}

// Define the LinkedList
pub struct LinkedList<T> {
  head: Option<Box<Node<T>>>,
  size: usize,
}

// Implement LinkedList methods
impl<T> LinkedList<T> {
pub fn new() -> Self {
LinkedList { head: None, size: 0 }
}

pub fn push(&mut self, data: T) {
  let new_node = Box::new(Node { data, next: self.head.take() });
  self.head = Some(new_node);
  self.size += 1;
}

pub fn pop(&mut self) -> Option<T> {
  let node = self.head.take()?;
  self.head = node.next;
  self.size -= 1;
Some(node.data)
}

pub fn len(&self) -> usize {
  self.size
}

pub fn is_empty(&self) -> bool {
  self.size == 0
}
    
}
```

In the code above, we defined a `Node` struct and a `LinkedList` struct. The `Node` struct contains some generic data of type `T` and a pointer to the next `Node`. The `LinkedList` struct contains a pointer to the head of the `LinkedList` and a size variable to keep track of the current size of the list. \n\nThe methods of the `LinkedList` are implemented to perform basic linked list operations like `push`, `pop`, `len`, and `is_empty`. The `push` method adds a new node to the front of the list, `pop` removes the node from the front of the list, `len` returns the size of the list and `is_empty` checks if the list is empty.

# Steps to run this on Inferentia/Trainium:

1. Launch an inf2.48xlarge instance using Amazon EC2. Use the HuggingFace Neuron [DLAMI](https://aws.amazon.com/marketplace/pp/prodview-gr3e6yiscria2).

2. Use the commands below to install the following packages or create a bash script. You can run the following commands in your terminal.
```
sudo apt-get update -y \
 && sudo apt-get install -y --no-install-recommends \
    aws-neuronx-dkms=2.15.9.0 \
    aws-neuronx-collectives=2.19.7.0-530fb3064 \
    aws-neuronx-runtime-lib=2.19.5.0-97e2d271b \
    aws-neuronx-tools=2.16.1.0 


pip3 install --upgrade \
    neuronx-cc==2.12.54.0 \
    torch-neuronx==1.13.1.1.13.0 \
    transformers-neuronx==0.9.474 \
    --extra-index-url=https://pip.repos.neuron.amazonaws.com
```

3. Clone the repo.
```
git lfs clone https://huggingface.co/VishaalY/Phind-CodeLlama-34B-v2-neuron-seqlen-2048-bs-1
```

4. Run the following commands to ensure that the script can read the files.

5. Move all the .neff files into a folder named 'neuron_artifacts'

![image/png](https://cdn-uploads.huggingface.co/production/uploads/63e2b95d282ee5f9624d4e6c/vojqFszTW0RdtIEiA3j_f.png)

6. Move all the other files into a folder named 'phind-34b-split'
   
![image/png](https://cdn-uploads.huggingface.co/production/uploads/63e2b95d282ee5f9624d4e6c/ZTDI0lOXl3xzgTV942jFv.png)

8. Create a file named inference.py in the same dir as the other two folders.

9. Run inference.py
```
import torch
from transformers_neuronx.module import save_pretrained_split
from transformers import LlamaForCausalLM
from transformers_neuronx.config import NeuronConfig
from transformers_neuronx import constants

import time
from transformers import AutoTokenizer
from transformers_neuronx.llama.model import LlamaForSampling

import os

print("Construct a tokenizer and encode prompt text.")
tokenizer = AutoTokenizer.from_pretrained('./Phind-CodeLlama-34B-v2')
prompt = "How do I reverse a linked list?"
input_ids = tokenizer.encode(prompt, return_tensors="pt")


print("Load from Neuron Artifacts")
neuron_model = LlamaForSampling.from_pretrained('./phind-34b-split', batch_size=1, tp_degree=24, amp='f16')
neuron_model.load('neuron_artifacts') # Load the compiled Neuron artifacts
neuron_model.to_neuron() # will skip compile
# ---------------------------------------------------------------------------------------------------------------

print("Run Inference.")

with torch.inference_mode():
    start = time.time()
    generated_sequences = neuron_model.sample(input_ids, sequence_length=2048, top_k=50)
    elapsed = time.time() - start
# ---------------------------------------------------------------------------------------------------------------

generated_sequences = [tokenizer.decode(seq) for seq in generated_sequences]
print(f'generated sequences {generated_sequences} in {elapsed} seconds')
```

9. If interested in how the artifacts are generated, here is a script to compile a llama based model.

```
import torch
from transformers_neuronx.module import save_pretrained_split
from transformers import LlamaForCausalLM
from transformers_neuronx.config import NeuronConfig
from transformers_neuronx import constants

import time
from transformers import AutoTokenizer
from transformers_neuronx.llama.model import LlamaForSampling

import os

# ---------------------------------------------------------------------------------------------------------------

print("Loading Model from HuggingFace. Converting to single bin / safetensor file for neuronx.")
model = LlamaForCausalLM.from_pretrained('Phind/Phind-CodeLlama-34B-v2')
save_pretrained_split(model, './phind-34b-split')

# Compiler flag -O1 is a workaround for “Too many instructions after unroll” in SDK 2.14
os.environ['NEURON_CC_FLAGS'] = '-O1'

# ---------------------------------------------------------------------------------------------------------------

print("Loading Model from single bin / safetensor. Convert to .neff format.")
neuron_model = LlamaForSampling.from_pretrained('./phind-34b-split', batch_size=1, tp_degree=24, amp='f16')
neuron_model.to_neuron()

# # ---------------------------------------------------------------------------------------------------------------

print("construct a tokenizer and encode prompt text")
tokenizer = AutoTokenizer.from_pretrained('./Phind-CodeLlama-34B-v2')
prompt = "How do I reverse a linked list?"
input_ids = tokenizer.encode(prompt, return_tensors="pt")

# # ---------------------------------------------------------------------------------------------------------------

print("run inference with top-k sampling")
with torch.inference_mode():
    start = time.time()
    generated_sequences = neuron_model.sample(input_ids, sequence_length=2048, top_k=50)
    elapsed = time.time() - start

generated_sequences = [tokenizer.decode(seq) for seq in generated_sequences]
print(f'generated sequences {generated_sequences} in {elapsed} seconds')

# ---------------------------------------------------------------------------------------------------------------

print("Saving Neuron Artifacts as .neff files.")
neuron_model.save('./neuron_artifacts') # can be copied and used on a different neuron instance
del neuron_model
```

This should result in the phind-split-34b dir and the neuron_artifacts dir being populated with the same files in the repo above.
{"license": "llama2"}
transformers pytorch llama text-generation license:llama2 autotrain_compatible endpoints_compatible text-generation-inference region:us