However, this file includes an example of how to compile various versions of Zephyr. Support isn’t available yet (as of 1/9/2024) in the optimum neuron framework, so we use the base transformers library.

These instructions closely follow the Developer Guide. Look there for more detailed explanations, especially for the GQA settings.

This model has been compiled to run on an inf2.xlarge (the smallest Inferentia2 instance). You can run it on a bigger instance, but it will only use two cores no matter how many are available, unless you change the core number available in compilation. Remember that each Neuron processor has two cores.

Set up the environment

First, use the DLAMI image from Hugging Face. It has most of the utilities and drivers preinstalled. However, you will need to update transformers-neruonx from the source to get Mistral/Zephyr support.

python -m pip install git+https://github.com/aws-neuron/transformers-neuronx.git

Running inference from this repository

If you want to run a quick test or if the exact model you want to use is HuggingFaceH4/zephyr-7b-beta, you can run it directly using the steps below. Otherwise, jump to the Compilation of other Mistral/Zephyr versions section.

First, you will need a local copy of the library. This is because one of the nice things that the Hugging Face optimum library does is abstract local loads from repository loads. However, Mistral/Zephyr inference isn't supported yet.

# To speed up downloads we can use hf_transfer
pip install hf_transfer
HF_HUB_ENABLE_HF_TRANSFER=1

# use huggingface-cli to download model to local dir
huggingface-cli download aws-neuron/zephyr-7b-beta-neuron --local-dir zephyr-7b-beta-neuron

This should put a local copy in zephyr-7b-beta-neuron. This process should take a 5-10 minutes. If it completes in a few seconds the first time you run it, you are likely having problems with git-lfs. You can see this by using ls -al to check the size of the files downloaded. You will also notice it later when you get parsing errors.

Next, load the model and neff files from disk into the Neuron processors:

import torch
from transformers_neuronx import constants
from transformers_neuronx.mistral.model import MistralForSampling
from transformers_neuronx.config import NeuronConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

# Set sharding strategy for GQA to be shard over heads
neuron_config = NeuronConfig(
    grouped_query_attention=constants.GQA.SHARD_OVER_HEADS
)

# define the model.  These are the settings used in compilation.
# If you want to change these settings, skip to "Compilation of other Mistral versions"
model_neuron = MistralForSampling.from_pretrained("zephyr-7b-beta-neuron", batch_size=1, \
    tp_degree=2, n_positions=256, amp='bf16', neuron_config=neuron_config)

# load the neff files from the local directory instead of compiling
model_neuron.load("zephyr-7b-beta-neuron")

# load the neff files into the neuron processors.  
# you can see this process happening if you run neuron-top from the command line in another console.
# if you didn't do the previous load command, this will also compile the neff files
model_neuron.to_neuron()

Inference example

This points to the original model for the tokenizer because the tokenizer is the same.
If you are compiling your own and want to have a single reference for everything, you can copy the special_tokens_map.json and tokenizer* from the original model to your local copy.

# Get a tokenizer and example input. This points to original tokenizer. 
# tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-alpha")
# this refers to tokenizer from local copy
tokenizer = AutoTokenizer.from_pretrained("zephyr-7b-beta-neuron")
text = "[INST] What is your favourite condiment? [/INST]"
encoded_input = tokenizer(text, return_tensors='pt')

# Run inference
with torch.inference_mode():
    generated_sequence = model_neuron.sample(encoded_input.input_ids, sequence_length=256, start_ids=None)
    print([tokenizer.decode(tok) for tok in generated_sequence])

Example output:

["<s> [INST] What is your favourite condiment? [/INST]\nHere's a little script to test people's favorite condiment.\n\nYou can do this with paper cones and have people guess what's in it, but they need to write their guess on a piece of of paper and put it in a jar before they take a bite.\n\nIn this version, we have ketchup, mustard,mayonnaise,bbq sauce, and relish.\n\nThe script is straightforward, so as long as your bottle isn’t too tiny, you can add to the bottom of the script,or re-shape the form of the script a bit.\n\nIf you put their guesses in a jar before they take a bite,you can put all their guesses in the jar as soon as they're done,and show the container as they guess.\nAs for removing lines from the script,you'll probably be removing the ones from the bottom of the script,or adding lines to the top of of the script.\nIf for no matter reason your bottle is too tiny to set all the guesses in,you can write their guesses on cards or bits of paper,and set"]

Compilation of other Mistral versions

If you want to use a different version of Mistral or Zephyr from Hugging Face, use the slightly modified code below. It essentially removes the “load” command. When the “to_neuron()” command sees that the model object doesn’t include the neff files, it will kick off the recompile. You can save them at the end so you only have to do the compilation process once. After that, you can use the code above to load a model and the neff files from the local directory.

import torch
from transformers_neuronx import constants
from transformers_neuronx.mistral.model import MistralForSampling
from transformers_neuronx.module import save_pretrained_split
from transformers_neuronx.config import NeuronConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id="HuggingFaceH4/zephyr-7b-beta"

# Load and save the CPU model with bfloat16 casting.  This also gives us a local copy
# change the Hugging Face model name (HuggingFaceH4/zephyr-7b-beta) below to what you want
# You can update the other model names if you want, but they just reference a directory on the local disk.
model_cpu = AutoModelForCausalLM.from_pretrained(model_id)
save_pretrained_split(model_cpu, model_id)

# Set sharding strategy for GQA to be shard over heads
neuron_config = NeuronConfig(
    grouped_query_attention=constants.GQA.SHARD_OVER_HEADS
)

# Create and compile the Neuron model
model_neuron = MistralForSampling.from_pretrained(model_id, batch_size=1, \
    tp_degree=2, n_positions=256, amp='bf16', neuron_config=neuron_config)
model_neuron.to_neuron()

#save compiled neff files out to the same directory
model_neuron.save("HuggingFaceH4/zephyr-7b-beta")

Arguments passed during compilation

The settings use in compilation are the same as shown above in the code. If you want to change these, you will need to recompile. If you don’t want to pass them in each time, you could update the config.json file. This is another nice thing the Hugging Face optimum neuron framework does for us. You can see an example of the format by looking at one of the Llama model config.json files. For example.

neuron_config = NeuronConfig(
    grouped_query_attention=constants.GQA.SHARD_OVER_HEADS
)
("zephyr-7b-beta-neuron", batch_size=1, tp_degree=2, n_positions=256, amp='bf16', neuron_config=neuron_config)