Phind-CodeLlama-34B-v2-neuron-seqlen-2048-bs-1

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.

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:

// 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.

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

1. Clone the repo.

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

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

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

1. Move all the other files into a folder named 'phind-34b-split'

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

2. 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')

1. 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.