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.

Takeaways

1. 1

   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.

2. 2

   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.

3. 3

   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.