Batching: serving a GPT to many users at once

Intro

0:00 Hello and welcome back to episode fourteen, where we're going to serve our model to a crowd, such as the players of a game, through an API that writes the stories for requests arriving together, as a batch. Let's begin with a quick look at where we're going to be by the end of the episode.

0:17 By the end of today, with sixty-four players asking our model for a story at once, the server writes about thirty-nine hundred tokens a second, against about six hundred and ninety when it writes one request at a time. And the median player waits about one point seven seconds for a story, instead of about nine point two.

0:35 On the website of the series, you can find more information about the project. The code is open source on GitHub, tagged at each step of each episode, so that you can follow along, and the docs get a page on serving the model today, with all of today's measurements. And each episode has its own page here, with its chapters and its code.

A game is a crowd

0:55 Now, before we get to today's work, let's quickly go over what we did last time. We wrote flash attention, which computes attention a block at a time without keeping its table of scores, first in PyTorch on the Mac and then as a Triton kernel on a rented GPU, and we trained the model with it. So far, the model has been served by the Gradio app from episode five, in Docker on a server, and by the page that runs it in the browser, and both write a story for one person at a time. A game is different, as many of its players can ask for something at once, and that's who we're going to serve today.

1:29 The game comes from my other series, nanoMMO, where we build an MMORPG from scratch, both the client and the server. Its focus is the engineering around the game, such as Docker, benchmarking, monitoring with Prometheus and Grafana, MySQL and Redis. Its name takes after Andrej Karpathy's nanochat, and like mvGPT, it keeps its code minimal, while covering an MMO from one end to the other. Say its world has a storyteller, and players can walk up to it and ask for a story with a few words, which our model writes. Then the game's server sends a request for each player who asks, and many of them can ask at once.

2:07 The players talk to the game's server, as they do today, and the game's server sends ours an HTTP request for each story, in the shape of OpenAI's completions API, which most clients and servers for language models speak. Our server loads the fine-tuned model from the Hub, by its tag, and writes one story at a time at first, which we're going to measure against a crowd. Then it writes the stories for the requests that arrive together as a batch, and we measure it again. And last, it goes into a Docker image of its own, so that a game can run it next to its other services. All right, let's get started with the server and its completions API.

A completions API

2:44 We start a new file, server dot py, which imports argparse, for its flags, json, for the requests and the answers, and os, for the environment variables. Queue is Python's queue that threads can share safely, where requests are going to wait for the model. Threading gives us the model's thread, along with the events a request waits on, time is for each answer's timestamp, and uuid gives each answer an id. From dataclasses, dataclass and field, for a small class that holds a request. And from http dot server, the HTTP server in Python's standard library, the base class for handling requests, and the threading server, which gives each connection a thread of its own. Then torch, and from transformers, the auto classes that load the model and the tokenizer from the Hub, as the app did in episode five, and detokenize, from our own tokenizer dot py. The repo and the revision come from environment variables, like the app's, with the fine-tuned model from episode three as the default, which writes the story an instruction asks for. A request is a dataclass, like our model's config since episode one. By default, a dataclass compares two objects by their fields, so we turn that off, and the comment says why: two requests asking for one story are still two requests. A request holds the prompt's token ids, max tokens, the temperature and top k. Tokens is a list, and field with a default factory gives each request a new one, as dataclasses refuse a plain list as a default, which all the requests would share. From tokens on, the fields are filled in by the model's thread, which sets done when they're ready. The finish reason starts as length, and done is a threading event, which a thread can wait on until another thread sets it. Write takes the model, a request and the token id of the end of text, and no grad turns off the gradients, as we're not training. X is the prompt, as a batch of one, and past is the KV cache, empty to begin with. For up to max tokens, the model takes x and the cache, and with use cache, transformers' GPT-2 hands the cache back with the new keys and values in it, like ours has since episode four. We take the logits at the last position and divide them by the temperature. Then the logits below the k-th largest get minus infinity, which leaves the top k tokens to pick from. Multinomial picks the next token from the softmax's probabilities, and indexing with none makes it a batch of one again, for the next step. If it's the end of text, the finish reason is stop, and we return, and otherwise the token joins the story. Serve runs on the model's thread, in a loop. Get waits until there's a request in the queue, write writes it, and setting done wakes up the thread that's waiting for it. Over the loop, we write that this is the model's thread, one request at a time, in the order they came. Decode turns the tokens into text, and for our word-level tokenizer, we join the tokens with detokenize, as the app does, and the comment says why: it's saved as WordPiece, whose decode puts spaces around punctuation. The BPE model from episode twelve goes through the tokenizer's own decode. The server is a threading HTTP server, with a request queue size of a hundred and twenty-eight, the number of connections that can wait for the server to accept them. With the default of five, a crowd of players gets turned away. Handler is a function that makes the class answering the requests, so that the class's methods have the tokenizer, the context window and the queue in reach, without any globals. The server makes a handler object from it for each request. Do get handles the GET requests, and any path but slash health gets a four oh four, with the error in the shape OpenAI's API gives its errors. Slash health answers with the model's name and tag, so that whoever runs the server can check that it's up. Do post answers slash v one slash completions, the path of OpenAI's API. It reads as many bytes as the content length header says, parses the JSON, and the tokenizer turns the prompt into ids. The room is what's left of the context window after the prompt, as the prompt and what's written after it have to fit in the context window. So a prompt that leaves no room gets a four hundred, with a message saying why. The request takes max tokens, cut to the room left, then the temperature and top k, with the defaults we've used since episode two. It goes on the queue, and this thread waits on the request's event while the model's thread writes the story. The answer has an id, the kind of object, a text completion, and when it was created, as OpenAI's answers do. The choices hold the story, decoded, and why it stopped, and the usage counts the tokens, so that a client can tell what a request cost. Reply turns the body into JSON, as bytes, then sends the status, the headers for the content's type and length, and the bytes. And log message is left empty, as a line per request is too many under load. In main, argparse reads the server's flags, and for now there's a single one, the port, which defaults to eight thousand, as it does in vLLM, a server for language models that we're going to look at later. The tokenizer and the model load from the Hub, and eval, as before, puts the model in evaluation mode. Then the queue, where requests wait for the model, which runs on a thread of its own. The thread runs serve, and daemon means it doesn't keep the program running once the main thread ends. The host comes from an environment variable, like the app's, and as the comment says, in Docker the server has to listen on all interfaces for a published port to reach it. The server gets the handler, with the model's context window, two hundred and fifty-six tokens. Print says where it's serving, and flush writes the line straight away, as Python holds its output back when it isn't printing to a terminal, such as in Docker's logs. And serve forever answers requests until we stop it. So, the model runs on a single thread, which takes the requests one at a time, while each connection's thread waits on its own request's event. As such, the connections don't hold each other up, and the model doesn't run on two threads at once.

8:49 We start the server in the background, with an ampersand at the end, and it loads the model from the Hub and prints where it's listening. Then curl sends it a request, where dash H sets the content type and dash d is the JSON body, asking for sixty tokens with the words dragon, castle and brave. And jq prints the answer indented, with the story, why it stopped, and how many tokens it took.

A crowd of players

9:12 To simulate the crowd, we write a new file, load dot py, which imports argparse, json, random, for the words each player asks for, and statistics, for the median. Threading runs the players at once, time is for the clock, and from urllib, the standard library's HTTP client, its errors and its requests, so the script runs on Python's standard library alone. The words are things from nanoMMO's world, such as sword, gate, skeleton and dragon, and as the comment says, they're what players in a game could ask a storyteller for, words from the game that the model knows. Merchant, healer and bard aren't among the model's ten thousand tokens, so they'd come in as unknown tokens, and the stories would come out without them. A client is one player, and its random generator gets a seed of its own, so each player asks for different words, and each run asks for those words again. Its comment says what a client does: a request, the wait for its story, then the next. For each of its requests, the prompt is an instruction with three of the words, and the body is the JSON for it, with max tokens, and perf counter starts the clock, before the request is built. urllib's request takes the URL, the body and the headers, here the content type. Then urlopen sends it, which is a POST, as it has a body, and json load reads the answer, where the usage says how many tokens the story took. A request that fails, with an error from urllib or a reset connection, goes in the errors and stays out of the times. We catch both, as when the server resets a connection while we wait for the answer, Python raises a plain connection error, without urllib's wrapping. And the time it took goes in the results, with its tokens. Run takes the URL, how many clients, and how many requests and tokens each of them asks for, and it starts with empty lists for the results and the errors. It makes a thread for each client, with its seed and both lists, which all the threads share. Appending to a list from several threads is safe in Python, as each append happens in one go. We start the clock, then each thread, and join waits for each of them to finish, so the clock stops once the last player has its last story. If no request got through, it stops with a message asking whether the server is running. The row has the clients, the requests that got through, the errors and the seconds the run took, and the tokens a second are all the tokens written, divided by those seconds. The median is the middle of the sorted times, and the ninety-fifth percentile is the wait that ninety-five percent of the requests beat, so one request in twenty waits longer than that. In main, argparse reads the flags, and the URL defaults to our server's completions, on this machine. Clients takes one or more numbers, with n args plus, and by default, it's one, four, sixteen and sixty-four, which run one after another. Each client sends four requests, one after another, of a hundred tokens each, so a crowd of sixty-four sends two hundred and fifty-six of them. One short request warms the server up first, as the first pass through the model takes longer than the rest. Then we print the table's header, with its columns padded to line up. For each number of clients, run simulates that crowd, and we keep the row it returns in a list and print it, with the clients, the requests and the errors, then the tokens a second, the median and the ninety-fifth percentile. And dash dash out, when it's given, writes the rows as JSON, for the charts.

12:39 Now we run the load script, with the server running in the background as before, and dash dash out to keep the measurements. With one player, a story of a hundred tokens takes about one hundred and fifty milliseconds, about six hundred and eighty tokens a second. With more players, the server writes about as many tokens a second, as it takes them one at a time, so each player waits for the stories of the players in front of them. With sixteen players, the median wait is about two point two seconds, and with sixty-four, it's about nine point two seconds, while the slowest five percent wait about nine point two seconds or more.

Batching

13:13 To write the stories for several requests together, their prompts go into one tensor, a row each, and as prompts can have different lengths, the shorter ones get padding tokens. We pad on the left, so that each prompt ends in the last column, and the token each row writes next goes in the next column, for all the rows at once. The attention mask has a one for each token and a zero for each pad, and attention leaves the pads out, so no token attends to them. And the positions count each row's own tokens, from the cumulative sum of the mask, minus one, so the first word of each prompt is at position zero, however much padding there is in front of it. Our model learned an embedding for each position, so without that, a padded prompt would read as if it started later, and the model would write something else.

14:00 Write now takes a batch of requests, along with the ids of the padding token and of the end of text. The width is the longest prompt in the batch, and with left padding, each prompt ends in the last column, so each row's next token goes after it. So x has a row per request, with the padding it needs and then its prompt. The mask has a row per request too, a zero for each pad and a one for each token, and its comment says what it's for: attention leaves the padding out. The temperature becomes a column, a row per request, so it divides each row's logits by that request's temperature, and top k is a column too. Past is the cache, and writing is the list of requests whose stories aren't finished yet. The loop runs for the longest max tokens in the batch. The positions are the cumulative sum of the mask along each row, minus one, and clamp turns the pads' minus ones into zeros. On the first step that's the whole prompt, and after that the last column, the new token's, which is what the slice keeps. So a row's positions count its own tokens, not the padding in front of them. The model gets the attention mask and the position ids, with the cache as before. The logits are each row's last position, divided by its temperature. For top k, topk takes as many logits from each row as the largest top k in the batch, sorted from the largest down, and gather picks from each row the value at its own k, minus one, as the index starts at zero. And the comment reads, each row keeps its own top k, its k-th largest logit, among the largest k of the batch. Multinomial now picks a token for each row, as a column, and the mask grows by a column of ones, for the new tokens. Then each request gets its token, and a request that's finished skips it, as its row still runs with the others, and what it writes is dropped. The end of text sets the finish reason to stop and takes the request off the list, and otherwise the token joins its story, until it has max tokens. Once that list is empty, the batch is done, so a row that finished early still takes its share of each step until then. Serve now takes two more arguments, the most requests in a batch, and the wait, in seconds. The batch starts with the first request, where get waits for one, as before. The deadline is the wait from now, on the monotonic clock, which keeps going forward when someone sets the computer's clock. While the batch has room, get waits for another request, for whatever time is left until the deadline, and queue empty means the time is up. The comment over the loop says why we wait: without it, the first of a crowd would go alone, with the rest queued behind it. Then write writes the whole batch, and each request's event is set. In main, dash dash batch is the most requests written at once, as its help says, sixteen by default, which we're going to come back to once we've measured it. Dash dash wait takes a float, how long a batch waits to fill, in milliseconds, five by default, which is short next to the time a story takes to write. The thread gets both, and as the flag is in milliseconds, the wait is divided by a thousand, into the seconds the queue's timeout takes. And the line it prints says up to how many requests it writes at once. Then we write a test, in tests slash test server dot py. It imports transformers' config and model classes for GPT-2, to build a tiny one with random weights, so the test needs no download, and request and write from our server. The pad's id is zero, and the end of text's is one. Requests makes a request for each prompt, asking for twelve tokens at a temperature of one. As the comment says, a top k of one always picks the likeliest token, so what's written doesn't depend on chance. The test's name says what it checks, that a batch writes what each request writes alone. It seeds torch's generator, so that each run gets those random weights and prompts again, and builds a model with a vocabulary of fifty tokens, sixty-four positions, thirty-two channels, and two blocks of two heads. The prompts have three, nine and five random tokens, and the comment says why: they have different lengths, so the batch pads some of them. Each prompt is written alone first, as a batch of one, and then the three together, and the stories have to match. If the mask or the positions were wrong, the padded rows would write something else, so this test pins down both.

18:22 We run the test with pytest, and it passes, so a batch writes what each of its prompts writes alone. I also tried it without the mask, and without the positions, and it fails both times, so it catches either mistake.

18:35 We start the server again, and it now writes up to sixteen requests at once. With one player, a story takes about one hundred and fifty milliseconds, a little longer than before, partly as a lone request waits its five milliseconds before it's written. With sixteen players, the server writes about thirty-one hundred tokens a second instead of about six hundred and ninety, and the median wait goes from about two point two seconds to about five hundred and ten milliseconds. With sixty-four players, the batches fill up and the rest queue behind them, so the median wait is about one point nine seconds.

19:08 With dash dash batch sixty-four, a batch can take all of the crowd at once. Up to sixteen players, it's about as fast as batches of sixteen, as no more than sixteen requests are waiting at any time. With sixty-four players, it writes about thirty-nine hundred tokens a second, against about thirty-two hundred with batches of sixteen, and the median wait is about one point seven seconds. These numbers move by about a tenth from one run to the next, and with sixty-four players, batches of sixty-four write about twenty-one percent more than batches of sixteen, so the gain is larger than that variation.

19:42 On a chart, one request at a time stays flat along the bottom, while batches write more as more players ask, up to about thirty-nine hundred tokens a second. And one at a time, the median wait grows in a straight line with the crowd, to about nine point two seconds at sixty-four players, while with batches, it's about one point seven.

20:01 Batches of sixteen are the default because the KV cache grows with the batch. A token's keys and values take two times six blocks, times three hundred and eighty-four channels, of four bytes each, which is about eighteen kilobytes. So sixteen full rows, of two hundred and fifty-six tokens, hold about seventy-five megabytes, which a server has to spare. The model from episode twelve has twelve blocks of seven hundred and sixty-eight channels and a window of five hundred and twelve tokens, so a token takes about seventy-two kilobytes. And sixty-four full rows would hold about two point four gigabytes, which is more than a small server may have free.

In Docker

20:41 To run the server in Docker, the Dockerfile from episode five gets stages, each starting with a FROM line and a name, so that the app and the server share their first steps. The base stage starts from Python's slim image, installs PyTorch's build for the CPU as before, and sets the host to all interfaces, for both of them. The server's stage starts from base, and installs the server's requirements, from requirements dash server dot txt, which we write in a moment. Then it copies server dot py and tokenizer dot py, the two files the server needs, exposes port eight thousand, and runs the server. The comment above it says how to build it, with docker build dash dash target server. The app's stage is what the Dockerfile was in episode five, and as its comment says, it's the last stage, so it's what docker build makes without a target, which keeps deploy dot sh building the app. The docker ignore file leaves everything out, with a star, and its comment says why: the images need the app, the server and what they import, and nothing else from the repo. Then it lets those files in, with an exclamation mark in front of each, so the checkpoints and the data stay out of what Docker sends to the build. And the server's requirements are transformers, pinned, without Gradio, which the server doesn't use.

21:56 Docker build with dash dash target server builds the server's stage, and dash t names the image. Docker run dash d starts it in the background, publishing port eight thousand, and the volume keeps the model it downloads from the Hub, like the app's. Then we run the load script again, with one, sixteen and sixty-four players. In Docker on the Mac, a lone player waits about five hundred and fifty milliseconds, and the server writes about one hundred and eighty tokens a second, about three point eight times slower than outside Docker. Docker on the Mac runs Linux in a virtual machine, with PyTorch's Linux build for ARM, and our model runs slower on that build than on PyTorch for macOS. Batching still writes about five point zero times as much, at about nine hundred tokens a second.

vLLM

22:43 Our batches start together and end together, so a row that finishes early keeps running, with its tokens dropped, until the last one is done. And a request that arrives just after a batch starts waits for all of it. Servers built for language models schedule each step instead, which is called continuous batching: a finished row leaves straight away, and a waiting request takes its place. vLLM, which started at UC Berkeley, does that, and it keeps the KV cache in blocks of a fixed size, which it calls PagedAttention, so that sequences of different lengths share the GPU's memory without padding. It runs on the CPU too, on x86, with AVX-512 or with fewer features on AVX2, and on ARM, and it has Docker images for the CPU. It loads both kinds of model we have on the Hub, the fine-tuned GPT-2 from episode three and the GPT-NeoX from episode twelve, and it serves OpenAI's completions API on port eight thousand, as ours does, so a client written for one works with the other. Our server is about a hundred and fifty lines of Python that we can read line by line, and it joins our word-level tokens back together, where vLLM's decode would put spaces around the punctuation.

23:54 Last, the docs get a page on serving, with the API, how the server is built, the batches and the test, all of today's measurements, and vLLM as the alternative. The docs' index and their navigation link to the new page. And the README has the commands to run the server and to put it under load.

Wrap-up

24:12 So, today we served our model as an API in the shape of OpenAI's completions, for a crowd such as the players of nanoMMO. One request at a time, it wrote about six hundred and ninety tokens a second whatever the crowd, so with sixty-four players, the median wait was about nine point two seconds. Writing the stories for the waiting requests together, with left padding, an attention mask and each row's own positions, it wrote about thirty-nine hundred tokens a second, and the wait went down to about one point seven seconds. The server has a Docker image of its own, next to the app's. And vLLM is the alternative, with continuous batching and OpenAI's API, on a GPU or on the CPU.

24:55 Let me know in the comments what you'd like a model like ours to write in a game, and which parts of serving you'd like us to go deeper on. And if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.