Make your llama generation time fly with AWS Inferentia2

Update (02/2024): Performance has improved even more! Check our updated benchmarks.

In a previous post on the Hugging Face blog, we introduced AWS Inferentia2, the second-generation AWS Inferentia accelerator, and explained how you could use optimum-neuron to quickly deploy Hugging Face models for standard text and vision tasks on AWS Inferencia 2 instances.

In a further step of integration with the AWS Neuron SDK, it is now possible to use 🤗 optimum-neuron to deploy LLM models for text generation on AWS Inferentia2.

And what better model could we choose for that demonstration than Llama 2, one of the most popular models on the Hugging Face hub.

Setup 🤗 optimum-neuron on your Inferentia2 instance

Our recommendation is to use the Hugging Face Neuron Deep Learning AMI (DLAMI). The DLAMI comes with all required libraries pre-packaged for you, including the Optimum Neuron, Neuron Drivers, Transformers, Datasets, and Accelerate.

Alternatively, you can use the Hugging Face Neuron SDK DLC to deploy on Amazon SageMaker.

Note: stay tuned for an upcoming post dedicated to SageMaker deployment.

Finally, these components can also be installed manually on a fresh Inferentia2 instance following the optimum-neuron installation instructions.

Export the Llama 2 model to Neuron

As explained in the optimum-neuron documentation, models need to be compiled and exported to a serialized format before running them on Neuron devices.

Fortunately, 🤗 optimum-neuron offers a very simple API to export standard 🤗 transformers models to the Neuron format.

>>> from optimum.neuron import NeuronModelForCausalLM

>>> compiler_args = {"num_cores": 24, "auto_cast_type": 'fp16'}
>>> input_shapes = {"batch_size": 1, "sequence_length": 2048}
>>> model = NeuronModelForCausalLM.from_pretrained(
        "meta-llama/Llama-2-7b-hf",
        export=True,
        **compiler_args,
        **input_shapes)

This deserves a little explanation:

>>> model.save_pretrained("a_local_path_for_compiled_neuron_model")