Serving GPT-2 at scale
By: Caleb Kaiser
Tags: Artificial Intelligence, Data Science, Programming, Machine Learning, Deep Learning

Serving GPT-2 at scale

Lessons learned scaling large deep models in production

Over the last few years, the size of deep learning models has increased at an exponential pace (famously among language models):

And in fact, this chart is out of date. As of this month, OpenAI has announced GPT-3, which is a 175 billion parameter model—or roughly ten times the height of this chart.

As models grow larger, they introduce new infrastructure challenges. For my colleagues and I building Cortex (open source model serving infrastructure), these challenges are front and center, especially as the number of users deploying large models to production increases.

Below, I’ve outline some of the design decisions we made to support serving models like GPT-2 at scale.

1. GPUs for inference

Running large deep learning inference on CPUs can be prohibitively slow. GPUs, on the other hand, have significantly more bandwidth (useful for the large computations involved in inference). This makes a large difference in inference latency.

For example, below I’ve deployed GPT-2 as an API on CPUs. Using the Cortex CLI, I’m monitoring the performance of the API, which I’ve named generator :

$ cortex get generator status up-to-date requested last update avg inference

live 1 1 5m 925 ms endpoint: http://abc123.us-west-2.elb.amazonaws.com/text/generator

The average inference takes nearly a full second to serve, which has serious implications at scale. For applications that have thousands of concurrent users, it will take a large number of replicas to keep the request queue under control.

Deploying on a GPU, as shown below, improves latency quite a bit:

$ cortex get generator status up-to-date requested last update avg inference

live 1 1 3m 199 ms endpoint: http://abc123.us-west-2.elb.amazonaws.com/text/generator

That’s roughly a 80% decrease from CPUs. For any application running at production scale, this is a meaningful improvement—but even at small scale, it can be crucial for latency-sensitive applications.

For example, features like Gmail’s Smart Compose require latency as close to realtime as possible. If Smart Compose generates text predictions more slowly than an average person types, recommendations are useless.

Implementing GPU support required a little refactoring. Cortex uses Kubernetes to manage model serving containers, but Kubernetes does not have out of the box support for managing GPUs across nodes. In order to manage GPUs through Kubernetes, we integrated NVIDIA’s device plugin for Kubernetes. Despite being in beta, we’ve found it to be very stable.

2. Autoscaling based on concurrency

GPU inference solved most of our latency and performance issues, but it introduced new challenges, particularly around autoscaling.

The Kubernetes pod autoscaler works by measuring CPU, not GPU, utilization. Obviously, this is useless for GPU inference.

We brainstormed a number of other approaches to autoscaling GPU instances. First, we tried autoscaling based on target latency. In this approach, users set an ideal latency for their API, and Cortex spun up additional instances until that latency was achieved. There were two key issues here:

If a user set their target latency unrealistically low—as in their model couldn’t possibly hit it with the provisioned resources—the cluster would autoscale infinitely.

Latency is a poor signal for scaling down instances. The latency of a request only changes when a replica is oversubscribed. Whether a replica is barely being used, or is at 90% utilization, latency will be the same so long as there are no queued requests

Additionally, maintaining and debugging this signal would be difficult for users. Any update to the model that affected latency, or any changes to the prediction API that changed latency (calls to another API, file operations, etc.) would require them to tweak their target latency. This was suboptimal.

The second idea we had was to monitor GPU utilization, but currently, Kubernetes metrics server does not offer GPU monitoring. We considered building this ourselves, but ultimately passed for a third option.

Our third option, and the one which turned out to be the best, was to implement request-based autoscaling, in which the number of enqueued requests and the concurrency capacity of a replica are used to calculate the number of needed replicas.

This approach worked, but required some refactoring:

We needed to redesign Cortex to use FastAPI and Uvicorn, as they allow us to run an asynchronous event loop that tracks requests as they’re added to the queue.

We added a target_replica_concurrency field to Cortex’s configuration files, so that users can specify how many requests a replica should handle concurrently at a steady rate.

There’s a full write up of the project here, but the end result is that now, Cortex can quickly autoscale to the optimal number of replicas (i.e. the minimum required without affecting performance).

3. Managing costs with spot instances

One of our ongoing focuses is lowering the cost of inference at scale. AWS’s cheapest GPU instance is the g4dn.xlarge , which costs $0.526 per hour and has a single GPU. Running a 500 GPU cluster, which some users do, for a month costs about $190,000.

Our initial instinct here was to support spot instances for inference. If you’re unfamiliar, spot instances are unused instances that AWS sells at a steep—sometimes as much as 70%—discount.

The challenge in implementation, however, was around how spot instances are allocated within AWS autoscaling groups.

Our goal was to allow users to opt to receive spot instances by default if available, and otherwise receive on demand instances. But autoscaling groups don’t work that way. Instead, users specify what percentage of nodes in an autoscaling group they want to be spot, and EKS tries to accommodate. There is no way to configure the autoscaling group to fall back to on demand instances when spot instances are not available.

To work around this, Cortex actually defines two separate autoscaling groups: one for spot instances, and another which is 100% on demand (and is only used as a backup for the primary autoscaling group). If a user requests spot instances, Cortex will first request nodes from the spot instance group. If an instance cannot be allocated (the scale-up request times out after five minutes), Cortex will request an instance from the on-demand instance group, and will temporarily blacklist the spot instance group so that future requests avoid the timeout.

Adapting to changes in model development

As model sizes continue to increase exponentially, there are other related changes within machine learning that affect how we design infrastructure.

First, pretrained models are becoming more prevalent in production. There are many teams doing no model training at all, and simply deploying pretrained models. Having a serving platform that can stand on its own, as opposed to only integrating with their training pipeline, is important for them—as they don’t have a pipeline.

Secondly, ASICs are becoming more important, as they present a good option for handling the resource needs of large models. Google has rolled out cloud TPUs, which allow for faster training and inference, while Amazon recently released Inferentia, a chip built specifically for inference.

All of these changes are good things—state-of-the-art models are being released pretrained, more teams are putting models into production, and the hardware for running inference is improving—but in order to take advantage of them, our model serving infrastructure needs to be constantly improved and rearchitected.