Hugging Face Text Generation Inference now supports AWS Inferentia2

TL;DR

Hugging Face has made Text Generation Inference (TGI) generally available on AWS Inferentia2 through Amazon SageMaker, offering a GPU‑free, production‑grade path for serving large language models (LLMs) with low latency and high concurrency.

What is Text Generation Inference (TGI)?

TGI is a purpose‑built inference server for LLMs that provides:

  • Tensor parallelism and continuous batching for high throughput.
  • Optimized support for popular open‑source models such as Llama, Mistral, and others.
  • Production‑ready deployment used by companies like Grammarly, Uber, and Deutsche Telekom.

Why Inferentia2 matters

AWS Inferentia2 chips deliver high‑performance inference at a lower cost than GPUs. By integrating TGI with Inferentia2 and SageMaker, Hugging Face offers:

  • A seamless, managed deployment experience.
  • Compatibility with the same backend technologies that power HuggingChat, OpenAssistant, and Hugging Face’s serverless endpoints.
  • An alternative compute stack for customers who prefer AWS‑specific hardware.

Deploying Zephyr 7B on Inferentia2 – Step‑by‑step guide

The blog post walks through deploying the Zephyr 7B model (a DPO‑fine‑tuned version of Mistral‑7B‑v0.1) on an ml.inf2.8xlarge instance. The steps are self‑contained and can be reproduced by any SageMaker user.

1. Set up the development environment

pip install transformers "sagemaker>=2.206.0" --upgrade --quiet

The script obtains an IAM execution role, creates a SageMaker session, and prints the role ARN and region.

2. Retrieve the TGI NeuronX container image

from sagemaker.huggingface import get_huggingface_llm_image_uri
llm_image = get_huggingface_llm_image_uri(
    "huggingface-neuronx",
    version="0.0.20"
)
print(f"llm image uri: {llm_image}")

At the time of writing the latest DLC version was not yet exposed via the helper, so the raw ECR URI is used.

3. Compile or fetch a cached model for Inferentia2

Inferentia2 does not support dynamic shapes, so sequence length and batch size must be fixed at compile time. Hugging Face provides a neuron model cache with pre‑compiled configurations (e.g., Mistral‑7B, Zephyr‑7B). If a required configuration is missing, users can compile it with the Optimum CLI:

optimum-cli export neuron -m HuggingFaceH4/zephyr-7b-beta \
    --batch_size 4 --sequence_length 2048 \
    --num_cores 2 --auto_cast_type bf16 ./zephyr-7b-beta-neuron

The compiled artifact is then pushed to the Hub under aws-neuron/zephyr-7b-seqlen-2048-bs-4-cores-2.

4. Define endpoint configuration and deploy

Key environment variables for the TGI NeuronX container include:

  • HF_MODEL_ID, HF_NUM_CORES, HF_BATCH_SIZE, HF_SEQUENCE_LENGTH, HF_AUTO_CAST_TYPE
  • MAX_BATCH_SIZE, MAX_INPUT_LENGTH, MAX_TOTAL_TOKENS
from sagemaker.huggingface import HuggingFaceModel
config = {
    "HF_MODEL_ID": "HuggingFaceH4/zephyr-7b-beta",
    "HF_NUM_CORES": "2",
    "HF_BATCH_SIZE": "4",
    "HF_SEQUENCE_LENGTH": "2048",
    "HF_AUTO_CAST_TYPE": "bf16",
    "MAX_BATCH_SIZE": "4",
    "MAX_INPUT_LENGTH": "1512",
    "MAX_TOTAL_TOKENS": "2048",
}
llm_model = HuggingFaceModel(role=role, image_uri=llm_image, env=config)
llm = llm_model.deploy(initial_instance_count=1, instance_type="ml.inf2.8xlarge", container_startup_health_check_timeout=1800)

Deployment typically takes 10–15 minutes.

5. Run inference and chat with Zephyr 7B

The model uses a chat template. The tokenizer’s apply_chat_template method converts OpenAI‑style message dictionaries into the required prompt format.

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("aws-neuron/zephyr-7b-seqlen-2048-bs-4-cores-2")
messages = [
    {"role": "system", "content": "You are the AWS expert"},
    {"role": "user", "content": "Can you tell me an interesting fact about AWS?"},
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
payload = {
    "do_sample": True,
    "top_p": 0.6,
    "temperature": 0.9,
    "top_k": 50,
    "max_new_tokens": 256,
    "repetition_penalty": 1.03,
    "return_full_text": False,
    "stop": ["</s>"]
}
chat = llm.predict({"inputs": prompt, "parameters": payload})
print(chat[0]["generated_text"][len(prompt):])

The response confirms the model is operational on Inferentia2.

6. Clean up resources

llm.delete_model()
llm.delete_endpoint()

Implications and future work

  • Cost‑effective scaling: Inferentia2 offers a cheaper compute option for high‑throughput LLM serving compared with GPU clusters.
  • Model coverage: The current cache includes Llama, Mistral, and Zephyr; Hugging Face plans to expand supported architectures.
  • Compilation workflow: Users can rely on pre‑cached builds or compile custom configurations (up to ~45 minutes) using Optimum.
  • Roadmap: Ongoing efforts aim to support more models, improve the caching system, and streamline the compilation pipeline.

For questions, reach out to Philipp Schmid on Twitter or LinkedIn.

Sources