Accelerate BERT inference with Hugging Face Transformers and AWS Inferentia

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.

Related reading

• huggingface.co · Jul 8, 2021

  Deploy Hugging Face models easily with Amazon SageMaker

  Hugging Face and Amazon introduced an inference solution integrating Hugging Face Transformers with Amazon SageMaker. The SageMaker Hugging Face Inference Toolkit and specialized Deep Learning Containers enable deployment of trained models or publicly available Hub models to managed production endpoints. Users configure endpoints with minimal code using the SageMaker Python SDK, standard pipelines, or custom inference scripts that override default behaviors. The architecture supports models stored in Amazon S3 as well as direct references via HF_MODEL_ID and HF_TASK environment variables. Deployments benefit from native AWS infrastructure features, including built-in monitoring, Identity and Access Management permission controls, and Virtual Private Cloud connectivity.

  Philipp Schmid

• huggingface.co · Jan 11, 2022

  Deploy GPT-J 6B for inference using Hugging Face Transformers and Amazon SageMaker

  Deploying EleutherAI's 6 billion parameter GPT-J model for production inference presents latency hurdles due to large memory footprints and lengthy startup times. Loading the model via standard methods takes up to several minutes, conflicting with strict real-time response limits such as Amazon SageMaker's 60-second threshold. To overcome this limitation, the model is serialized using PyTorch's native save mechanisms, packaged into a compressed archive with supporting assets, and stored on Amazon S3. This alternative loading workflow reduces GPT-J load times down to 7.7 seconds. An Amazon SageMaker real-time endpoint is then deployed on an NVIDIA T4 GPU instance using the Hugging Face Inference Toolkit.