# Accelerate BERT inference with Hugging Face Transformers and AWS Inferentia

huggingface.co · Philipp Schmid · Mar 16, 2022

**Type:** Tutorial

## Summary

Production deployments of BERT and Transformer architectures often face cost and latency challenges because these models are significantly larger and more computationally intensive than traditional algorithms. To optimize text classification workloads, developers can compile vanilla PyTorch models for AWS Inferentia using the AWS Neuron SDK and its tracing utilities. Because the Neuron SDK requires static tensor dimensions, the model is traced with fixed input lengths and packaged alongside a custom inference script configuring one Neuron Core per worker. Deploying the resulting artifacts to an Amazon SageMaker ml.inf1.xlarge endpoint yields an average latency of 5 to 6 milliseconds for a sequence length of 128 across 10,000 synchronous evaluation requests.

## Context

As transformer models transition from research to large-scale production workloads, BERT and similar architectures remain relatively slow, large, and complex compared to traditional machine learning algorithms. Accelerating inference while managing infrastructure costs presents a key operational challenge.

## Approach / What changed

The implementation converts a Hugging Face sequence classification model to AWS Neuron using torch.neuron.trace with static input shapes, writes a custom inference.py handler setting NEURON\_RT\_NUM\_CORES=1, packages the artifacts to Amazon S3, and deploys a real-time SageMaker endpoint on an ml.inf1.xlarge instance using HuggingFaceModel.

## Takeaways

- AWS Inferentia chips contain four Neuron Cores, enabling configurations that load one model per core for high throughput or span one model across all cores for lower latency.
- The AWS Neuron SDK requires static input shapes for model compilation and inference, meaning compiled models only accept inputs matching the traced batch size and sequence length.
- Compiling DistilBERT with AWS Neuron and deploying on a SageMaker ml.inf1.xlarge instance achieved an average model latency of 5 to 6 milliseconds for sequence lengths of 128.

**Tags:** [AWS](https://yomu.fyi/topic/aws), [Deployment](https://yomu.fyi/topic/deployment), [Machine Learning](https://yomu.fyi/topic/machine-learning), [Performance](https://yomu.fyi/topic/performance), [Python](https://yomu.fyi/topic/python)

[Read original post](https://huggingface.co/blog/bert-inferentia-sagemaker)