This is especially important if you are changing any paramters that require a recompile because Optimum-neuron will let you take advantage of the compilation cache.

Neuronx model for Mistral

This repository contains AWS Inferentia2 and neuronx compatible checkpoints for mistralai/Mistral-7B-Instruct-v0.1.

However, this file includes an example of how to compile various versions of Mistral. Support isn’t available yet (as of 1/3/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 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 mistralai/Mistral-7B-Instruct-v0.1, you can run it directly using the steps below. Otherwise, jump to the Compilation of other Mistral 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 inference isn't supported yet.

From python:

# using python instead of git clone because I know this supports lfs on the DLAMI image
from huggingface_hub import Repository
repo = Repository(local_dir="Mistral-neuron", clone_from="aws-neuron/Mistral-neuron")

This should put a local copy in Mistral-neuron. This process should take a 5-10 minutes. If it completes in a few seconds the first time you run it, you are 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.module import save_pretrained_split
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("Mistral-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("Mistral-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.  Note that this points to the original model
tokenizer = AutoTokenizer.from_pretrained('mistralai/Mistral-7B-Instruct-v0.1')
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: (most of the time with amp=‘bf16’, the answer is ketchup. However, if I compiled with amp=f32, the answer was soy sauce. This was for a sample size of one, so let me know what you see —@jburtoft)

2024-Jan-03 15:59:21.0510 1486:2057 [0] nccl_net_ofi_init:1415 CCOM WARN NET/OFI aws-ofi-nccl initialization failed
2024-Jan-03 15:59:21.0510 1486:2057 [0] init.cc:138 CCOM WARN OFI plugin initNet() failed is EFA enabled?
['<s> [INST] What is your favourite condiment? [/INST] My favorite condiment is probably ketchup. It adds a perfect balance of sweet, tangy, and slightly spicy flavor to dishes, and is versatile enough to go with a wide variety of foods.</s>']

Compilation of other Mistral versions

If you want to use a different version of Mistral 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

# Load and save the CPU model with bfloat16 casting.  This also gives us a local copy
# change the Hugging Face model name (mistralai/Mistral-7B-Instruct-v0.1) 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('mistralai/Mistral-7B-Instruct-v0.1')
save_pretrained_split(model_cpu, 'mistralai/Mistral-7B-Instruct-v0.1-split')

# 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('mistralai/Mistral-7B-Instruct-v0.1-split', 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("mistralai/Mistral-7B-Instruct-v0.1-split")

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 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
)
("Mistral-neuron", batch_size=1, tp_degree=2, n_positions=256, amp='bf16', neuron_config=neuron_config)