Deploying Llama 2 on AWS Inferentia2 with optimum-neuron

Hugging Face has integrated optimum-neuron with the AWS Neuron SDK, enabling the deployment of Llama 2 models on AWS Inferentia2 accelerators. This integration allows users to leverage specialized hardware to improve text generation performance in terms of encoding time, latency, and throughput.

Deployment Workflow and Setup

Deploying Llama 2 on AWS Inferentia2 requires specific environment setups and a model compilation step to convert the model into a serialized format compatible with Neuron devices.

Environment Configuration

Users can set up their Inferentia2 instances using three primary methods:

  • Hugging Face Neuron Deep Learning AMI (DLAMI): The recommended approach, providing pre-packaged libraries including Optimum Neuron, Neuron Drivers, Transformers, Datasets, and Accelerate.
  • Hugging Face Neuron SDK DLC: Used for deployments on Amazon SageMaker.
  • Manual Installation: Following the optimum-neuron installation instructions on a fresh instance.

Model Export and Compilation

Because Neuron devices require static shapes, models must be compiled before execution. Using the NeuronModelForCausalLM API, users specify the following parameters during export:

  • Compiler Arguments: Defines the number of cores (each neuron device has two cores) and the precision (e.g., float16).
  • Input Shapes: Sets static dimensions for batch_size and sequence_length. The sequence_length is critical as it constrains the input context, the KV cache, and the maximum output length.

Once compiled, models can be saved locally or pushed to the Hugging Face Hub for reuse.

Text Generation Capabilities

optimum-neuron supports standard text generation using the transformers library or simplified optimum-neuron pipelines.

Generation Strategies

Supported generation strategies include:

  • Greedy search
  • Multinomial sampling with top-k and top-p (including temperature)
  • Most logits pre-processing filters, such as repetition penalty

Implementation Options

For streamlined deployment, the pipeline API from optimum-neuron allows users to load a pre-compiled model from the Hub and generate text with a single function call.

Performance Benchmarks

Benchmarks were conducted using Llama 2 7B and 13B models with various configurations on inf2.xlarge (for budget models) and inf2.48xlarge (for latency and throughput optimized models). All models used a maximum sequence length of 2048.

Encoding Time

Encoding time—the duration to process input tokens and generate the first output token—is a key metric for perceived user latency. For 256 input tokens (typical Q&A), encoding times ranged from 0.3s (Llama2 7B-B) to 0.9s (Llama2 7B-T). For 768 input tokens (typical RAG), times ranged from 0.5s (Llama2 7B-B) to 5.2s (Llama2 13B-T).

End-to-End Latency

End-to-end latency measures the total time to reach a sequence length of 1024 tokens. On high-end inf2.48xlarge instances, Llama2 7B-L achieved a latency of 6.2s for 768 new tokens, while Llama2 13B-L took 10.2s. The "budget" model on inf2.xlarge exhibited significantly higher latency, reaching 47.3s for 768 new tokens.

Throughput

Throughput is calculated as total tokens per second (end-to-end latency divided by batch_size * sequence_length).

  • High-Performance: The Llama2 7B-T configuration achieved up to 750 tokens/second for 256 new tokens.
  • Budget: The Llama2 7B-B model achieved between 22 and 32 tokens/second, which Hugging Face notes is sufficient for streaming use-cases given average human reading speeds.

Technical Limitations and Future Work

While performance is strong, Hugging Face identifies two primary areas for improvement:

  1. Throughput Scaling: Currently, throughput is primarily increased by increasing batch size, which is limited by device memory. Pipelining is being integrated as an alternative.
  2. Context Length: The requirement for static sequence lengths limits the ability to handle very long contexts. The team is exploring attention sinks as a potential solution.

Sources