Profiling a GPT with PyTorch's profiler

Intro

0:00 Hello and welcome back to episode six, where we're going to look at where our model spends its time, with PyTorch's profiler and Google's Model Explorer, and measure what our decisions in its code cost. 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, we can look at our model as a graph of the operations it runs, each coloured by how long it takes on the CPU, and the slowest is the head, which scores all ten thousand tokens of the vocabulary.

0:30 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 profiling today. And each episode has a page here as well, with its chapters and its code.

Last time

0:50 Now, before we get to today's work, let's quickly go over what we did last time. Our model has GPT-2's architecture, so we saved it under GPT-2's names, and pushed it to the Hugging Face Hub, with a tag per episode. A Gradio app writes stories with it on a server, and a single page runs it in the browser. Both write a token at a time with a KV cache, like ours from episode four. There, we saw that a model this small writes one story faster on the CPU than on the GPU. Today, we're going to see why, operation by operation, with PyTorch's profiler.

Measure, then decide

1:24 First, we're going to look at the model as PyTorch runs it, a graph of operations, in Model Explorer, a tool from Google. Then PyTorch's profiler times those operations, on the CPU and on the GPU, while the model reads a full context window, and while it writes. And with that, we're going to measure a few decisions in our code, such as writing GELU out, against the alternatives, on each device. Guessing where the time goes is easy to get wrong, and here the answer depends on the device, so we measure before we decide. All right, let's get started with what PyTorch runs when it runs our model.

The model as a graph

2:00 We write explore dot py, to show the model in Model Explorer. It imports argparse, Model Explorer, which Google made to look at models as graphs, torch, and load checkpoint. Main makes an argument parser, with a line on what the script does, and takes a checkpoint, and a port for the web page, eighty-eighty by default. It loads the model from the checkpoint, on the CPU, and puts it in evaluation mode. X is a full context window of random token ids, two hundred and fifty-six of them. Torch dot export needs an example input, x, as it traces the model's code with tensors of that shape, and records each operation it calls, with its shapes, as a graph. The graph is the model as PyTorch runs it, with no Python left in between: each linear layer, each layer norm, and each of the operations in our GELU. We count the nodes of the graph that call a function, which are the operations, and print how many there are. Then we hand the program to Model Explorer, under the name mvGPT, and it serves the graph as a web page, on the port we gave it, for us to open in the browser. And Model Explorer joins the requirements, next to the others.

3:10 We run it on our fine-tuned model, and a pass over two hundred and fifty-six tokens is a hundred and sixty-three operations. Model Explorer converts the graph, and serves it on port eighty-eighty, with an address that names the model to open.

3:27 Model Explorer starts with the model folded into a single layer, between its inputs and its output. Expanded, it's a long column of operations, the embeddings at the top, then the six blocks, and the head at the bottom, too small to read from here. The search box finds operations by name, such as erf, the error function in our GELU, and around it, the operations GELU is written with: a division, the error function, an addition and two multiplications, between the MLP's two linear layers. So our GELU is five operations, where PyTorch's own GELU is one, which we're going to come back to.

PyTorch's profiler

4:04 Next, we write profiler dot py, which imports argparse, json, os, statistics and time, torch, and from PyTorch's profiler, its activities, profile, and record function, which labels a stretch of the run, along with sync from bench dot py. The prompt is once upon a time, and we write a hundred tokens. Profiled makes a profiler for a device, which watches the CPU, and on CUDA, the GPU's kernels too. Above it, we note that on a Mac's GPU, the profiler sees the CPU's side: handing each operation over to the GPU, and we're going to keep that in mind when we read its numbers. Summary finds the run we labelled, and as its comment says, it keeps the operations the run called, not the ones they call in turn, such as add m m inside linear, which would count the time twice. Each event records its parent on the CPU, so the calls are the events whose parent is our labelled run. For each operation's name, we add up how often it was called, and its time, and sort them, slowest first, as a share of the run. Summary returns the run's time in milliseconds, how many calls it made, and a row for each operation. Profile pass runs the model on a full window once, to warm up, as a first run sets things up. Then it runs it once more inside the profiler, labelled pass, with a sync at the end, so the profiler waits for the GPU's work, which runs in the background, after the CPU has handed it over. Profile writing does that for generate, writing a hundred tokens with the cache, labelled writing. Ms per token times writing without the profiler, as the profiler adds a little time to each operation it records, so its own times run a little slow. It writes a hundred tokens five times, each between two syncs, so the GPU has finished, with perf counter around it, and keeps the milliseconds per token. Then it takes the median of the five, which ignores a run that something else on the machine slowed down. Print ops prints a table of the operations: each one's name, how many times it was called, and its share of the run's time, for the slowest eight. Main takes the checkpoint, the device, which defaults to the one pick device chooses, and a folder for the results. We seed PyTorch, load the model on the device, and put it in evaluation mode. X is two hundred and fifty-six random token ids, a full context window, so that each operation has enough work for its time to be mostly arithmetic, and not the overhead of calling it. We profile the pass, summarise it, and print its time and how many operations it ran, then its operations, slowest first, with how often each was called and its share of the time. Writing is a token at a time with the cache, which means lots of small operations. We encode the prompt, profile writing a hundred tokens, and time it again without the profiler, which we call plain. Then we print the operations per token, which is the writing run's calls divided by a hundred, and the time per token with the profiler and without it, and the five slowest operations. With dash dash out, we make the folder, and write both summaries as JSON, with the time without the profiler. And the pass's trace goes in a file too, a timeline of each operation, which Perfetto, a viewer from Google, opens at u i dot perfetto dot dev.

Where the time goes

7:16 On the CPU, a pass over two hundred and fifty-six tokens, profiled, is one hundred and sixty-three operations. Linear is about fifty-nine percent of the time, as the linear layers are where the arithmetic is, and attention about nine. Our GELU's error functions, multiplications and divisions add up to about eleven percent, plus a share of the additions, which is a lot for a function that does no matrix multiplication. When it writes, the model runs about one hundred and eighty-nine operations per token. It takes about one point three one milliseconds per token without the profiler, and about one point four nine with it, which is the profiler's own cost.

7:53 As a chart, the linear layers take most of a pass, and that's the arithmetic we need, as they multiply the channels by the model's weights. Our GELU comes next, a little ahead of attention, and well ahead of the layer norms. Its arithmetic works on each number on its own, which is cheap per number, but its five operations read and write all the MLP's channels five times over.

8:14 On the GPU, the profiled pass takes about one point six milliseconds, which is the time to hand its operations over, as the GPU runs them after that, in the background. When the model writes, the numbers look odd: about forty percent of the time goes to an operation called to, which copies a tensor to a device, and it's called once per token. The copy is a single number, and it's the first moment in each step where the CPU needs the GPU to be done, so the CPU waits there for the GPU's work, and the profiler counts that wait as the copy's time. I checked, and if we remove that copy, the wait moves to the next operation that needs the GPU, in top k's masking. So when we profile a GPU from the CPU's side, the time shows up where the CPU waits, and not where the GPU does the work.

9:02 To see the time on the graph, explore dot py gets dash dash times. It imports Model Explorer's node data builder, to colour the graph's nodes by a number, then record function, and profiled from profiler dot py. Labelled is a torch f x interpreter, which steps through a graph and runs each node, and we wrap each node's run in a record function with its name. An interpreter runs a graph slowly, with a Python call per node, which is fine here, as we want the time of each node, and not the fastest pass. Node times runs the exported graph that way once, to warm up, and once in the profiler, then keeps the time of each event whose name is a node's, in milliseconds. The profiler's averages add up its events by name, so each node's label collects the time of that node alone. In main, with dash dash times, we time the nodes on the CPU, print the five slowest with their milliseconds, so the terminal shows them too, and give the times to Model Explorer as node data, with a gradient from the slate for the quickest operations to coral for the slowest.

10:10 The slowest node is linear thirty-six, at about eight-tenths of a millisecond, which is the head, and the next four are linear layers in the MLPs, at about a fifth of a millisecond each.

10:22 In Model Explorer, the panel on the right now has our node data, with its sum and its largest value. Expanded, the column of operations has colour where the time goes: the linear layers glow, while the views, transposes and additions stay dark. And at the bottom, the head is the brightest. It multiplies each of the two hundred and fifty-six positions' three hundred and eighty-four channels by a matrix with ten thousand columns, one per token of the vocabulary, which is more arithmetic than any other layer.

Decisions, measured

10:52 Now we measure the decisions, in decisions dot py. It imports what it needs to time things, and mock, from Python's unittest, which swaps a function for another while a block runs. It imports our model module as m v g p t, along with the story from bench dot py, and from profiler dot py, the prompt, ms per token and the profiling. Attention by hand is scaled dot product attention, written out for the arguments model dot py passes it. It's the formula from episode one: q, k transposed, over the square root of d k, the causal mask where it applies, minus infinity where a token mustn't look, the softmax, and times v. Each decision is a set of patches: model dot py runs as it is, with a function swapped for another. As written swaps no function, while GELU as one operation swaps our gelu for PyTorch's, which computes GELU's formula in a single operation. And attention by hand swaps PyTorch's attention for ours. Ms per pass times a pass over the full window. It runs one pass first, to warm up, then ten more, each between two syncs, so the GPU has finished, with perf counter around it, and it keeps the milliseconds of each. It returns the median of the ten, like ms per token in profiler dot py, so a pass slowed by something else on the machine doesn't count. Measure is what we do for each decision, and it returns a dictionary with a pass's time, a token's time, the operations per token, from profiling writing, and the logits the model gives the story, which we're going to compare between decisions. Main makes an argument parser, with a line on what the script measures, and takes the checkpoint, the device and a folder for the results, like profiler dot py. It loads the model and its tokenizer on that device, puts it in evaluation mode, so dropout is off, and encodes the prompt and the story as tensors of token ids. X is two hundred and fifty-six random token ids, a full context window for timing a pass, while the short prompt is for writing a token at a time with the cache. We print a header, and go through the decisions. An exit stack is a context manager that holds others: we enter each of the decision's patches in it, and when the with block ends, it undoes them all, so the next decision starts from the code as written. Inside the with block, we measure the decision, with its patches in place. The first decision is the code as written, and its logits are the reference, so each other decision's largest difference from them should be rounding, as each of them computes our model's formula with different operations. We keep the row, and print the pass, the token, the operations and the difference. Then torch dot compile, which turns the pass into fused kernels the first time it runs, for one shape of input. It reads the graph, like torch dot export, and generates code that does several operations in one go. The first call compiles, which takes a few seconds, so we make it before timing. Then we time the compiled pass, check its logits against the reference, and print its row, with no token time, as writing with a growing cache changes the input's shape at each step, and the compiled code would be made again and again. Last comes the vocabulary, as the head scores each token of the vocabulary, so its size follows the vocabulary's. An untrained model is as fast as a trained one, so there's no need to train one per vocabulary. We print a header, and for vocabularies of five, ten and twenty thousand tokens, we seed PyTorch, and build an untrained model with that vocabulary, in evaluation mode. X modulo the vocabulary's size keeps the token ids inside it. We time a pass and a token, count the head's weights, which grow with the vocabulary, and print them. With dash dash out, everything goes in a JSON file, named after the device.

14:38 On the CPU, a pass takes about eight point eight milliseconds as written. With GELU as one operation, about seven point eight, with twenty-four fewer operations per token, and the logits differ by rounding. Attention by hand makes a pass slower, about ten point six, and torch dot compile changes little here, at about eight point nine. Writing a token, the three uncompiled versions take between about one point three and one point four milliseconds. A head four times as large, for twenty thousand tokens instead of five thousand, makes a pass about nine percent slower.

15:14 On the GPU, the picture changes, and a pass takes about four point six milliseconds as written, and about three point six with GELU as one operation, which saves a fifth of the pass, against a tenth on the CPU. Torch dot compile brings it to about three point eight, as fusing operations means fewer launches and fewer trips to memory, which cost the GPU more than the arithmetic does. And attention by hand, writing a token, takes about five point one milliseconds, against about two point one as written, more than twice as long.

15:44 Side by side, GELU as one operation is the clearest win on both devices, while torch dot compile helps on the GPU and not on the CPU. On the GPU, each operation has a cost of its own, of launching it and of reading and writing its tensors in memory, so fewer, fused operations pay off more there. Attention by hand is slower on both, as PyTorch's attention computes the scores in fused kernels, while ours writes the whole grid of scores to memory, and reads it back for each of the operations that follow.

What we keep

16:15 So, with the numbers in front of us, let's decide what we keep. Our GELU written out costs about a tenth of a pass on the CPU and a fifth on the GPU, and we keep it, as this series is about reading the code. On the GPU, torch dot compile wins most of that time back, though on the CPU it doesn't. We keep PyTorch's attention, as ours by hand is slower on both devices. And twice the vocabulary, twenty thousand tokens instead of ten, makes a pass about five percent slower on the CPU, and about twenty-one percent on the GPU. So the vocabulary is mostly a decision about coverage, as in episode two, with a cost in speed that's larger on the GPU. On a bigger model, on a big GPU, these numbers would come out differently again, so we measure each decision on the machine it runs on.

17:02 Last, the docs get a page on profiling, with Model Explorer, the profiler on the CPU and the GPU, the graph coloured by time, and the decisions with today's timings, and the docs' home page links to it. And the README has the commands to show the model as a graph, to profile it, and to time the decisions.

Wrap-up

17:22 So, today we looked at our model as a graph of a hundred and sixty-three operations, in Model Explorer, and timed them with PyTorch's profiler. On the CPU, the time goes where the arithmetic is, in the linear layers, with the head the slowest node. On the GPU, the number of operations matters, and the profiler times the CPU's side, so the time shows up where the CPU waits for the GPU. Of the decisions we measured, GELU as one operation saves about a tenth of a pass on the CPU and a fifth on the GPU, and attention by hand is slower on both, while torch dot compile helps on the GPU and not on the CPU.

18:00 Let me know in the comments which other decisions in our code you'd like us to measure, and on which machines. And if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.