Episode 10 of a series that builds mvGPT, a minimum viable GPT, from scratch and takes it through the training pipeline of a language model. Our model learned a table of positions, a vector for each of its 256 places, which stops at the context window and keeps its KV cache from sliding past it. We replace it with rotary position embeddings, from the RoFormer paper: each query and key is turned by an angle that grows with its position, so attention depends on how far apart two tokens are. We add them to the model as a config option, with a cache that slides, pretrain for ten minutes next to episode two's model, and measure both past their context window, on the CPU and the GPU of a laptop.
Topics: Architecture
The code
Each step on screen is a tag in the code's repository. The code at the end of the episode is ep010/end.
Transcript
Intro
0:00 Hello and welcome back to episode ten, where we're going to give our model rotary position embeddings, which put the positions of the tokens into attention, so that its KV cache can slide along a story past the context window. Let's begin with a quick look at where we're going to be by the end of the episode.
0:19 By the end of today, past its context window, our model predicts each token in about two point two milliseconds on my laptop's CPU, with a KV cache that slides along the story, where the model with learned positions takes about eight point six, running the whole window again for each token.
0:34 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 rotary position embeddings today. And each episode has a page here as well, with its chapters and its code.
0:54 Now, before we get to today's work, let's quickly go over what we did last time. We fine-tuned our model on instructions with LoRA, where the pretrained weights stay frozen, and each linear layer gets a small update of low rank that we train, about two percent of the weights, saved as a file of its own and merged back in. Today, we go back to how the model writes. In episode four, our KV cache kept the keys and values of the tokens so far, so that each new token runs through the model alone. That holds up to the context window, two hundred and fifty-six tokens. Past it, generate runs the whole window again for each new token, as each token's position is mixed into its keys and values, and a token that moves back in the window gets new ones.
1:37 Today, we're going to replace the table of positions with rotary position embeddings, which turn each query and each key by an angle that depends on its position. With those, a cached key stays valid as the window moves on, so the cache can slide, dropping its oldest token. We pretrain a model with them for ten minutes, from the seed and the data of the model from episode two, and compare the two on held-out stories. And we measure both past the context window, for how well they predict and how long each token takes. All right, let's get started with how our model knows where each token is.
Learned positions
2:13 Right now, the model knows where each token is from a table of position embeddings, a row of three hundred and eighty-four numbers for each of the two hundred and fifty-six positions, learned in training like the token embeddings. The two are added before the first block, so from there on, the position is mixed into each token's channels, and into the keys and values each block computes from them. That costs us in two ways, once the stories get long. The table has no row past position two hundred and fifty-five, so the model can't place a token further than that, and once a story outgrows the window, each token moves back a position at each step, which changes its channels in all six blocks, along with its cached keys and values.
2:54 A table like ours is how GPT-2 does it, and how Andrej Karpathy does it in his video, Let's build GPT, from scratch, in code, spelled out, where he writes a small GPT in a notebook, one piece at a time, with a table of position embeddings next to the table of tokens. I'd recommend it alongside our first episode, as he explains each piece as he types it, and trains it on Shakespeare.
Rotary position embeddings
3:17 Rotary position embeddings come from RoFormer, Enhanced Transformer with Rotary Position Embedding, by Jianlin Su and others at Zhuiyi Technology, in Shenzhen, from 2021. Its abstract says that RoPE encodes each token's absolute position with a rotation, and that attention then depends on the relative position of two tokens, how far apart they are. Figure one shows how: the query and the key are split into pairs of channels, and each pair is turned by an angle that grows with the position, faster for some pairs and slower for others. And equation sixteen shows why it works: the query rotated for position m, times the key rotated for position n, comes out as the original query and key with a single rotation between them, by n minus m. Llama uses it, and so do most of the open models since.
4:05 Take two of a query's channels as the two coordinates of a point on a plane, an arrow from the origin. Rotary embeddings turn that arrow by an angle, m times theta, where m is the token's position. So the token at position zero isn't turned, the next one is turned by theta, the one after by two theta, and so on. The arrow keeps its length, and its direction now carries the position. A head's sixty-four channels make thirty-two pairs, and each pair gets a theta of its own, ten thousand to the minus two i over sixty-four. The first pair turns by one radian per token, which takes it once around in about six tokens, and the last pair takes about forty-seven thousand tokens to go once around. Between them, the pairs tell near and far positions apart. Keys are turned like the queries, by their own positions. The score of a query with a key is their dot product, which for two arrows is their lengths times the cosine of the angle between them. Turn the query by m theta and the key by n theta, and the angle between them changes by n minus m, times theta, wherever the two tokens are. So a score depends on what the two tokens are and how far apart they are, and not on where they sit in the window. And a cache can slide because of it, as a key rotated for its position gives the correct score against each later query, however far the window has moved.
Rotary embeddings in our model
5:24 In model dot py, the config gets a new field, rope, false by default: rotary position embeddings, instead of a learned embedding per position. As an option, it keeps our older checkpoints working, as their configs have no rope in them, so they load with their table of positions, as before. In the KV cache, a new field, position, starts at zero. Up to now, the next token's position was how many tokens the cache held, and once the cache drops tokens, the two come apart. So position holds the next token's place in the story, which keeps counting up once the cache starts dropping tokens. Then a new method, slide, takes how many tokens to keep, and drops the older ones in each block: a slice from the end of the tokens' dimension, first of the keys, and then of the values. Next, a function of its own, rotary, which takes the first position, how many tokens, the head size and the device. Speeds is the theta of each pair, ten thousand to the minus two i over sixty-four, where a range from zero to sixty-four in steps of two gives the thirty-two values of two i. Angles multiplies each position by each speed, which gives a table of T rows by thirty-two angles, and we return their cosines and sines, so each block can reuse them. Above it, we write what it computes: the angles for positions start to start plus T minus one, one per pair of channels, as cosines and sines, turning faster for the first pairs and slower for the last ones. Rotate then turns a query or a key by those angles. It splits the channels into two halves, x one and x two, so each pair of channels, i and i plus thirty-two, is turned by its position's angle. It returns the rotation from the formula, x one cos minus x two sin, and x one sin plus x two cos, for all the pairs at once. The paper pairs neighbouring channels, while Llama's code pairs the two halves, which is that rotation with the channels in another order, and as the projections are learned, either order works as well. In attention, forward takes the angles, and after the projections, when there are angles, the queries and the keys are rotated. Rotated by their positions, a query and a key score by how far apart they are, not where. That happens before the cache's update, so the cache keeps rotated keys, and the values aren't rotated at all, as they carry what a token says, not where it is. The block passes the angles on to attention, alongside the cache. In the GPT, with rotary embeddings, attention handles positions, and there's no table of them, which saves two hundred and fifty-six rows of three hundred and eighty-four, about a hundred thousand parameters. In forward, the first position comes from the cache's position, and x starts as the token embeddings. With rope, the angles are computed once, for all the blocks, and without it, we add the positions' embeddings, as before. After the blocks, the cache's position moves on by T. And generate now slides the cache along the story. The window is the block size, two hundred and fifty-six tokens. The whole window still runs again when we're not using the cache, or when the story is past the window with learned positions, and its comment gets a sentence more: with rotary embeddings, the cache can slide instead, as a key's rotation stays valid as the window moves. Without the cache, the model reads the last window of the story, as before. On the first step, the prompt goes in cut to the window, so a prompt longer than the window works too, and from then on, once the cache holds a whole window, slide drops its oldest token before the newest goes in. Positions keep counting past two hundred and fifty-five, which is fine, as the model compares distances, and those stay inside the window.
Ten minutes of pretraining
8:55 To train a model with rotary embeddings, pretrain dot py gets a flag. The run folder isn't a constant any more, as main takes rope, and the folder depends on it: each kind of position embedding gets a run folder of its own. So with rope, the checkpoints go to pretrain dash rope, and the model with learned positions, in checkpoints slash pretrain, stays as episode two left it, for us to compare against. The config gets rope, the TensorBoard run is named after the folder, and the checkpoints are saved into it. At the bottom, dash dash rope on the command line turns it on. The seed, the data, the batches and the learning rate are the ones from episode two, so the position embeddings are the one difference between the two runs.
9:40 So, let's pretrain for ten minutes with rotary embeddings. The model has fourteen million five hundred thousand parameters, about ninety-eight thousand fewer than with learned positions, as there's no table of positions any more. It trains at about twenty-six thousand tokens a second, against about twenty-seven thousand in episode two, as rotating the queries and keys adds a few small operations to each block. So in ten minutes, it gets through one thousand one hundred and seventy-eight steps, where episode two's run got through twelve hundred and fifty, and it ends with a held-out loss of about two point one eight.
10:16 Step for step, the model with rotary embeddings learns faster. At step a thousand, its held-out loss is about two point one eight, against about two point three one with learned positions, and it ends lower too, despite its fewer steps. I'd put that down to the table of positions, whose two hundred and fifty-six rows start out random and have to be learned from the data. Rotary embeddings build the distances into attention from the first step.
10:41 Evaluate dot py measures the two on the held-out stories, like in episode two. The model with learned positions has a perplexity of about ten, and with rotary embeddings, about nine, and its likeliest token is right about fifty-one percent of the time, against about fifty. Both read the stories within their window, so this is how well each learned, and not yet what happens past it.
Past the context window
11:06 To measure past the window, we write context dot py. It imports argparse, os and time, NumPy, torch and functional, sync from our bench, to wait for the GPU, our checkpoint loader, the data folder, the KV cache and pick device. Length is seven hundred and sixty-eight tokens, three context windows of held-out text, and we take eight streams of them. Window losses is the way a model with learned positions has to write past its window. For each token past the first window, the model runs the window that ends just before it, and we keep the loss of predicting it, from the logits at the last position. So each token past the first window is predicted from the two hundred and fifty-six tokens before it, run in full each time, one window after another, as generate does when it writes past the window with learned positions. Cache losses goes a token at a time through the cache. For each token but the last, if the cache holds a whole window, slide drops the oldest one, and the token goes through the model alone, with the cache, where attention covers that token and the two hundred and fifty-five before it. So the next token is predicted from two hundred and fifty-six tokens, like with the window. From the end of the first window on, we keep each loss. Full losses runs all the tokens at once, each attending to all the ones before it, up to seven hundred and sixty-seven tokens back, which is further than the model ever saw in training. Reduction none keeps a loss per token, past the first window. Timed runs one of those, between two syncs, so the GPU has finished, and returns the losses and the milliseconds per token. In main, the arguments are the checkpoints and the device, which defaults to the GPU when there is one. We memory-map the validation tokens, and cut eight streams of seven hundred and sixty-eight from them, one after the other, each as a batch of one on the device. Then we print a header, with the tokens we measure, and the columns of a table: the model, the method, the loss and the milliseconds. For each checkpoint, we load the model, put it in evaluation mode, and name it after its folder. The sliding window applies to any model, while the sliding cache and all of the context are for rotary embeddings alone, and the comment over them says why: learned positions stop at two hundred and fifty-six, so the cache can't slide, and the model can't see further. For each method, we average the loss over the eight streams, along with the time per token, and print a row, so each method of each model gets a line of its own. And as in our other scripts, main runs when we call context dot py from the command line.
13:37 First on the CPU, which we ask for with dash dash device. With learned positions, the sliding window predicts the tokens past the first window with a loss of about two point three two, at about eight point six milliseconds a token. With rotary embeddings, the sliding window gives about two point one eight, and the sliding cache gives about two point one eight too, as it sees the tokens the window sees, at about two point two milliseconds a token, against about twelve for its sliding window, about five point three times faster. And all of the context at once does worse, with a loss of about two point three four. Then on the GPU of my laptop, where a window run in full takes about three point five milliseconds a token with learned positions, and the sliding cache about three point zero. With a single story, a step on this GPU is mostly the time it takes PyTorch to send it each operation, as we saw in episode six, so the cache saves a lot less here.
14:31 Side by side, the sliding cache matches the sliding window's loss, and the model with rotary embeddings predicts better than the one with learned positions, as it did on the held-out stories. All of the context at once is worse than both, even though it sees more of the story. Rotary embeddings make attention depend on distances, but this model saw distances up to two hundred and fifty-five in training, and it has no training to go on past that. So rotary embeddings alone don't stretch the context a model was trained on. Models that read longer contexts train on longer sequences, or scale their rotations to fit. On the CPU, the cache does a token's work once, where the window runs two hundred and fifty-six tokens through all six blocks for each new one, which is where the difference in time comes from.
15:17 Andrej Karpathy's nanochat, his small ChatGPT, uses rotary embeddings too, and the docstring of its model lists them, with no positional embeddings. Its apply rotary emb splits the channels into two halves, like ours, and returns the two halves turned by cosines and sines. It turns them the other way, by minus theta, which, as its comment notes, works as well, as the rotation between a query and a key is what counts. And its rotations use a base of a hundred thousand, where ours use ten thousand, so its slowest pairs turn more slowly still, which recent models do for their longer contexts.
15:55 Last, the docs get a page on rotary position embeddings, with the rotation and why attention then depends on distances, the code and the sliding cache, the two models side by side, and the measurements past the window, on the CPU and the GPU. It goes in the list on the docs' home page, after LoRA, and at the end of the navigation. And the README has the commands to pretrain with rotary embeddings and to measure past the window.
Wrap-up
16:20 So, today we replaced our model's table of positions with rotary position embeddings, where each query and key is turned by an angle that grows with its position, so a score depends on how far apart two tokens are. In ten minutes of pretraining, from the seed and the data of episode two, the model with them reached a held-out loss of about two point one eight, against about two point two five with learned positions, in fewer steps. Past the context window, its cache slides, and on my laptop's CPU, it predicts each token in about two point two milliseconds, where running the window again takes about eight point six. And all of a longer context at once did worse, as rotary embeddings let a model handle the distances it was trained on, and no further on their own.
17:03 Let me know in the comments how far past its window you'd like to see a model like ours read, and which parts of attention 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.