Let's reproduce GPT-2 (124M)

Hi everyone. So today we are going to be continuing our Zero to Hero series. And in particular, today we are going to reproduce the GPT-2 model, the 124 million version of it. So when OpenAI released GPT-2, this was 2019, and they released it with this blog post. On top of that, they released this paper, and on top of that, they released this code on GitHub. So openai/gpt-2. Now, when we talk about reproducing GPT-2, we have to be careful because in particular in this video we're going to be reproducing the 124 million parameter model. So the thing to realize is that there's always a mini-series when these are— releases are made. So there are the GPT-2 mini-series made up of models at different sizes, and usually the biggest model is called the GPT-2. But basically the reason we do that is because you can put the model sizes on on the x-axis of plots like this, and on the y-axis you put a lot of downstream metrics that you're interested in, like translation, summarization, question answering, and so on. And you can chart out these scaling laws. So basically, as the model size increases, you're getting better and better at downstream metrics. And so in particular for GPT-2, if we scroll down in the paper, there are 4 models in the GPT-2 mini-series starting at 124 million, all the way up to 1,558 million. Now, the reason my numbers, the way I say them, disagree with this table is that this table is wrong. If you actually go to the GPT-2 GitHub repo, they sort of say that there was an error in how they added up the parameters. But basically, this is the 124 million parameter model, et cetera. So the 124 million parameter had 12 layers in the transformer, and it had 768 channels in the transformer, 768 dimensions. And I'm going to be assuming some familiarity with what these terms mean because I covered all of this in my previous video, "Let's build GPT-2," uh, "Let's build GPT from scratch." So I covered that in the previous video in this playlist. Now if we do everything correctly and everything works out well, by the end of this video we're going to see something like this where we're looking at the validation loss, which basically measures how good we are at predicting the next token in a sequence on some validation data that the model has not seen during training. And we see that we go from doing that task not very well because we're initializing from scratch all the way to doing that task quite well by the end of the training. And hopefully we're going to beat the GPT-2 124M model. Now previously when they were working on this, this is already 5 years ago, so this was probably a fairly complicated optimization at the time and the GPUs and the compute was a lot smaller. Today you can reproduce this model in roughly an hour or probably less even, and it will cost you about $10 if you want to do this on the cloud, cloud compute, a sort of computer that you can all rent. And if you pay $10 for that computer, you wait about an hour or less, you can actually achieve a model that is as good as this model that OpenAI released. And one more thing to mention is unlike many other models, OpenAI did release the weights for GPT-2, so those weights are all available in this repository, but the GPT-2 paper is not always as good with all of the details of training. So in addition to the GPT-2 paper, we're going to be referencing the GPT-3 paper, which is a lot more concrete in a lot of the hyperparameters and optimization settings and so on, and it's not a huge departure in the architecture from the GPT-2 version of the model. So we're going to be referencing both GPT-2 and GPT-3 as we try to reproduce GPT-2 1.24M. So let's go. So the first thing I would like to do is actually start at the end or at the target. So in other words, let's load the GPT-2 1.24M model as it was released by OpenAI and maybe take it for a spin. Let's sample some tokens from it. Now the issue with that is when you go through the codebase of GPT-2 and you go into the source, and you click in on the model.py, you'll realize that actually this is using TensorFlow. So the original GPT-2 code here was written in TensorFlow, which is, you know, not— let's just say not used as much anymore. So we'd like to use PyTorch because it's a lot friendlier, easier, and I just personally like it a lot more. The problem with that is the initial code is in TensorFlow. We'd like to use PyTorch. So instead, to get the target, we're going to use the Hugging Face Transformers code, which I like a lot more. So when you go into the transformers/source/transformers/models/gpt2/modeling/gpt2.py, you will see that they have the GPT-2 implementation of that transformer here in this file, and it's like medium readable but not fully readable. But what it does is it did all the work of converting all those weights from TensorFlow to PyTorch friendly, and so it's much easier to load and work with. So in particular, we can look at the GPT-2 model here and we can load it using Hugging Face Transformers. So swinging over, this is what that looks like. From transformers import the GPT-2 LM head model and then from pretrained GPT-2. Now one awkward thing about this is that when you do GPT-2 as the model that we're loading, this actually is the 124 million parameter model. If you want the actual, the GPT-2, the 1.5 billion, then you actually want to do -xl. So this is the 124M, our target. Now what we're doing is when we actually get this, we're initializing the PyTorch nn.Module as defined here in this class. From it, I want to get just the state dict, which is just the raw tensors. So we just have the tensors of that file. And by the way, here this is a Jupyter Notebook, but this is Jupyter Notebook running inside VS Code, so I like to work with it all in a single sort of interface, so I like to use VS Code. So this is the Jupyter Notebook extension inside VS Code. So when we get the state_dict, this is just a dict, so we can print the key and the value, which is the tensor, and let's just look at the shapes. So these are sort of the different parameters inside the GPT-2 model and their shape. So the w weight for token embedding is of size 50,257 by 768. Where this is coming from is that we have 50,257 tokens in the GPT-2 vocabulary, and the tokens, by the way, these are exactly the tokens that we've spoken about in the previous video. On my tokenization series. So the previous videos just before this, I go into a ton of detail on tokenization. GPT-2 tokenizer happens to have this many tokens. For each token, we have a 768-dimensional embedding that is the distributed representation that stands in for that token. So each token is a little string piece, and then the 768 numbers are the vector that represents that token. And so this is just our lookup table for tokens. And then here we have the lookup table for the positions. So because GPT-2 has a maximum sequence length of 1024, we have up to 1024 positions that each token can be attending to in the past. And every one of those positions in GPT-2 has a fixed vector of 768 that is learned by optimization. And so this is the position embedding and the token embedding. And then everything here is just the other weights and biases and everything else of this transformer. So when you just take, for example, the positional embeddings and flatten it out and take just the 20 elements, you can see that these are just the parameters. These are weights, floats, just we can take and we can plot them. So these are the position embeddings. And we get something like this, and you can see that this has structure. And it has structure because what we, what we have here really is every row in this visualization is a different position, a fixed absolute position in the range from 0 to 1024. And each row here is the representation of that position. And so it has structure because these positional embeddings end up learning these sinusoids and cosines that sort of like represent each of these positions. And each row here stands in for that position and is processed by the transformer to recover all the relative positions and sort of realize which token is where and attend to them depending on their position, not just their content. So when we actually just look into an individual column inside these, and I just grabbed 3 random columns, you'll see that, for example, here we are focusing on every single channel, and we're looking at what that channel is doing as a function of position from 1— from 0 to 1023, really. And we can see that some of these channels basically, like, respond more or less to different parts of the position spectrum. So this green channel really likes to fire for everything after 200, up to 800, but not less, but a lot less, and has a sharp drop-off here near 0. So who knows what these embeddings are doing and why they are the way they are. You can tell, for example, that because they're a bit more jagged and they're kind of noisy, you can tell that this model was not fully trained. And the more trained this model was, the more you would expect to smooth this out. And so this is telling you that this is a little bit of an undertrained model. But in principle, actually, these curves don't even have to be smooth. This should just be totally random noise. And in fact, in the beginning of the optimization, it is complete random noise because this position embedding table is initialized completely at random. So in the beginning, you have jaggedness, and the fact that you end up with something smooth is already kind of impressive, that that just falls out of the optimization, because in principle you shouldn't even be able to get any single graph out of this that makes sense, but we actually get something that looks a little bit noisy, but for the most part looks sinusoidal-like. In the original transformer, in the original transformer paper, the Attention is All You Need paper, the positional embeddings are actually initialized and fixed, if I remember correctly, to sinusoids and cosines of different frequencies, and that's the positional encoding and it's fixed. But in GPT-2, these are just parameters and they're trained from scratch. Just like any other parameter, and that seems to work about as well. And so what they do is they kind of like recover these sinusoidal-like features during the optimization. We can also look at any of the other matrices here. So here I took the first layer of the transformer and looking at like one of its weights and just the first block of 300 by 300, and you see some structure, but like Again, like, who knows what any of this is? If you're into mechanistic interpretability, you might get a real kick out of trying to figure out, like, what is going on, what is this structure, and what does this all mean? But we're not going to be doing that in this video. But we definitely see that there's some interesting structure, and that's kind of cool. What we're mostly interested in is we've loaded the weights of this model that was released by OpenAI, and now using the Hugging Face Transformers, we cannot just get all the raw weights, but we can also get the, what they call pipeline, and sample from it. So this is the prefix, "Hello, I'm a language model," comma, and then we're sampling 30 tokens and we're getting 5 sequences. And I ran this and this is what it produced. "Hello, I'm a language model." But what I'm really doing is making a human-readable document. There are other languages, but those are So you can read through these if you like, but basically these are 5 different completions of the same prefix from this GPT-2 1.24M. Now if I go here, I took this example from here, and sadly even though we are fixing the seed, we are getting different generations from the snippet than what they got. So presumably the code changed, but What we see though at this stage that's important is that we are getting coherent text. So we've loaded the model successfully. We can look at all its parameters and the keys tell us where in the model these come from. And we want to actually write our own GPT-2 class so that we have full understanding of what's happening there. We don't want to be working with something like, uh, the modeling_gpt2.py because it's just too complicated. We want to write this from scratch ourselves. So we're going to be implementing the GPT model here in parallel. And as our first task, let's load the GPT-2 124M into the class that we're going to develop here from scratch. That's going to give us confidence that we can load the OpenAI model, and therefore there's a setting of weights that exactly is the 124 model. But then of course what we're going to do is we're going to initialize the model from scratch instead and try, try to train it ourselves. On a bunch of documents that we're going to get, and we're going to try to surpass that model. So we're going to get different weights and everything's going to look different, hopefully better even, but we're going to have a lot of confidence that because we can load the OpenAI model, we are in the same model family and model class, and we just have to rediscover a good setting of the weights, but from scratch. So let's now write the GPT-2 model and let's load the weights and make sure that we can also generate text that looks coherent. Okay, so let's now swing over to the Attention is All You Need paper that started everything, and let's scroll over to the model architecture, the original Transformer. Now, remember that GPT-2 is slightly modified from the original Transformer. In particular, we do not have, uh, the encoder. GPT-2 is a decoder-only Transformer, as we call it. So this entire encoder here is missing. And in addition to that, this cross attention here that was using that encoder is also missing. So we delete this entire part. Everything else stays almost the same, but there are some differences that we're going to sort of look at here. So there are two main differences. When we go to the GPT-2 paper under 2.3.model, we notice that first there's a reshuffling of the layer norms, so they change place. And second, an additional layer normalization was added here to the final self-attention block. So basically all the layer norms here, instead of being after the MLP or after the attention, they swing before it, and an additional layer norm gets added here right before the final classifier. So now let's implement some of the first sort of skeleton NN modules here in our GPTNN module. And in particular, we're going to try to match up this schema here that is used by Hugging Face Transformers, because that will make it much easier to load these weights from this state dict. So we want something that reflects this schema here. So here's what I came up with. Basically, we see that the main container here that has all the modules is called transformer. So I'm reflecting that with an nn.Module dict. And this is basically a module that allows you to index into the submodules using keys, just like a dictionary, strings. Within it, we have the weights of the token embeddings, wte, and that's an nn.Embedding, and the weights of the position embeddings, which is also just an nn.Embedding. And if you remember, nn.Embedding is really just a fancy little wrapper module around just a single, single array of numbers, a single block of numbers just like this. It's a single tensor, and an embedding is a glorified wrapper around a tensor that allows you to access its elements by indexing into the rows. Now in addition to that, we see here that we have a.h, and then there's a— this is indexed using numbers instead of indexed using strings. So there's a.h,.0, 1, 2, etc. All all the way up till.h.11. And that's because there are 12 layers here in this transformer. So to reflect that, I'm creating also an h, I think that probably stands for hidden. And instead of a module dict, this is a model list. So we can index it using integers exactly as we see here,.0,.1,.2, et cetera. And the module list has n layer blocks and the blocks are yet to be defined in a module in a bit. In addition to that, following the GPT-2 paper, we had— we need an additional final layer norm that we're going to put in there. And then we have the final classifier, the language model head, which projects from 768, the number of embedding dimensions in this GPT, all the way to the vocab size, which is 50,257. And GPT-2 uses no bias for this final layer norm. Sort of projection. So this is the skeleton, and you can see that it reflects this. So the WTE is the token embeddings. Here it's called output embedding, but it's really the token embeddings. The PE is the positional encodings. Those two pieces of information, as we saw previously, are going to add and then go into the transformer. The.h is the— all the blocks in gray, and the LNF is this new layer that gets added here. By the GPT-2 model, and lm_head is this linear part here. So that's the skeleton of the GPT-2. We now have to implement the block. Okay, so let's now recurse to the block itself. So we want to define the block. So I'll start putting them here. So the block, I like to write out like this. These are some of the initializations, and then this is the actual forward pass of what this block computes. And notice here that there's a change from the transformer again that is mentioned in the GPT-2 paper. So here the layer normalizations are after the application of attention or feedforward. In addition to that, note that the normalizations are inside the residual stream. You see how feedforward is applied and this arrow goes through and through the normalization. So that means that your residual pathway has normalizations inside them. And this is not very good or desirable. You actually prefer to have a single clean residual stream all the way from supervision, all the way down to the inputs, the tokens. And this is very desirable and nice because the gradients that flow from the top, if you remember from your micrograd, addition just distributes gradients during the backward stage to both of its branches equally. So addition is a branch in the gradients. And so that means that the gradients from the top flow straight to the inputs, the tokens, through the residual pathways unchanged. But then in addition to that, the gradient also flows through the blocks, and the blocks, you know, contribute their own contribution over time and kick in and change the optimization over time. But basically, clean residual pathway is desirable from an optimization perspective. And then this is the pre-normalization version. Where you see that our x first goes through the layer normalization and then the attention and then goes back out to go to the layer normalization number 2 and the multi-layer perceptron, sometimes also referred to as a feed-forward network or an FFN, and then that goes into the residual stream again. And the one more thing that is kind of interesting to note is that recall that attention is a communication operation. It is where all the tokens and and there's 1024 tokens lined up in a sequence. And this is where the tokens communicate. This is where they exchange information. So attention is an aggregation function. It's a pooling function. It's a weighted sum function. It is a reduce operation. Whereas MLP, this MLP here happens at every single token individually. There's no information being collected or exchanged between the tokens. So the attention is the reduce and the MLP is the map. And what you end up with is that the transformer just ends up just being a repeated application of MapReduce, if you want to think about it that way. So this is where they communicate and this is where they think individually about the information that they gathered. And every one of these blocks iteratively refines the representation inside the residual stream. So this is our block, slightly modified from this picture. Okay, so let's now move on to the MLP. So the MLP block I implemented as follows. It is relatively straightforward. We basically have two linear projections here that are sandwiched in between the GeLU nonlinearity. So nn.GeluApproximate is tanh. Now when we swing on— swing over to the PyTorch documentation, This is n.gelu, and it has this format, and it has two versions: the original version of Gelu, which we'll step into in a bit, and the approximate version of Gelu, which we can request using tenh. So as you can see, just as a preview here, Gelu is basically like a ReLU, except there's no flat, exactly flat tail here at exactly zero. But otherwise, it looks very much like a slightly smoother ReLU. It comes from this paper here, Gaussian Error Linear Units, and you can step through this paper and there's some mathematical kind of like reasoning that leads to an interpretation that leads to this specific formulation. It has to do with stochastic regularizers and the expectation of a modification to adaptive dropout. So you can read through all of that if you'd like here. And there's a little bit of a history as to why there is an approximate version of GELU. And that comes from this issue here, as far as I can tell. And in this issue, Daniel Hendricks mentions that at the time when they developed this nonlinearity, the erf function, which you need to evaluate the exact GELU, was very slow in TensorFlow. So they ended up basically developing this approximation, and this approximation that then ended up being picked up by BERT and by GPT-2, etc. But today there's no real good reason to use the approximate version. You'd prefer to just use the exact version because I— like, my expectation is that there's no big difference anymore, and this is kind of like a historical kind of quirk. But we are trying to reproduce GPT-2 exactly, and GPT-2 used the 10H approximate version, so we prefer to stick with that. Now, one other reason to actually just intuitively use Gelu instead of Verlu is previously in the In videos in the past, we've spoken about the dead ReLU neuron problem, where in this tail of a ReLU, if it's exactly flat at zero, any activations that fall there will get exactly zero gradient. There's no change, there's no adaptation, there's no development of the network if any of these activations end in this flat region. But the GeLU always contributes a local gradient, and so there's always going to be a change, always going to be an adaptation and sort of smoothing it out ends up empirically working better in practice, as demonstrated in this paper, and also as demonstrated by it being picked up by the BERT paper, GPT-2 paper, and so on. So for that reason, we adopt this non-linearity here in the GPT-2 reproduction. Now, in more modern networks, also like Llama 3 and so on, this non-linearity also further changes to Swigloo and other variants like that. But for GPT-2, they used this approximate galup. Okay, and finally we have the attention operation. So let me paste in my attention. So I know this is a lot, so I'm going to go through this a bit quickly, a bit slowly, but not too slowly because we have covered this in the previous video and I would just point you there. So this is the attention operation. Now in the previous video, you will remember this is not just attention. This is, multi-headed attention, right? And so in a previous video, we had this multi-headed attention module, and this implementation made it obvious that these heads are not actually that complicated. Uh, there's basically in parallel inside every attention block, there's multiple heads and they're all functioning in parallel and, uh, their outputs are just being concatenated and that becomes the output of the multi-headed attention. So the heads are just kind of like parallel streams and their outputs get concatenated. And so it was very simple and made the head be kind of like fairly straightforward in terms of its implementation. What happens here is that instead of having two separate modules and indeed many more modules that get concatenated, all of that is just put into a single self-attention module. And instead I'm being very careful and doing a bunch of transpose, split tensor gymnastics to make this very efficient in PyTorch. But fundamentally and algorithmically, nothing is different from the implementation we saw before in this Git repository. So to remind you very briefly, and I don't want to go into this in too much time, but we have these tokens lined up in a sequence and there's 1,020 of them. And then each token at this stage of the attention emits 3 vectors: the query, key, and the value. And first what happens here is that the queries and the keys have to multiply each other to get sort of the attention amount, like how interesting they find each other. So they have to interact multiplicatively. So what we're doing here is we're calculating the QKV, we're splitting it, and then there's a bunch of gymnastics as I mentioned here. And the way this works is that we're basically making the, number of heads, nh, into a batch dimension. And so it's a batch dimension just like B, so that in these operations that follow, PyTorch treats B and nh as batches, and it applies all the operations on all of them in parallel, in both the batch and the heads. And the operations that get applied are: number 1, the queries and the keys interact to give us our attention. This is the autoregressive mask, that makes sure that the tokens only attend to tokens before them and never to tokens in the future. The softmax here normalizes the attention, so it sums to 1 always. And then recall from the previous video that doing the attention matrix multiply with the values is basically a way to do a weighted sum of the values of the tokens that we found interesting at every single token. And then the final transpose_contiguous and view is is just reassembling all of that again, and this actually performs the concatenation operation. So you can step through this slowly if you'd like, but it is equivalent mathematically to our previous implementation. It's just more efficient in PyTorch, so that's why I chose this implementation instead. Now, in addition to that, I'm being careful with how I name my variables. So for example, c_attn is the same as c_attn, and so actually our keys should basically exactly follow the schema of the Hugging Face Transformers code, and that will make it very easy for us to now port over all the weights from exactly this sort of naming conventions because all of our variables are named the same thing. But at this point, we have finished the GPT-2 implementation, and what that allows us to do is we don't have to basically use this file from Hugging Face, which is fairly long. This is 2,000 lines of code. Instead, we just have less than 100 lines of code, and this is the complete GPT-2 implementation. So at this stage, we should just be able to take over all the weights, set them, and then do generation. So let's see what that looks like. Okay, so here I've also changed the gpt_config so that the numbers here, the hyperparameters, agree with the GPT-2 124M model. So The maximum sequence length, which I call block size here, is 124. The number of tokens is 50,257, which if you watch my tokenizer video know that this is 50,000 merges, BPE merges, 256-byte tokens, the leaves of the BPE tree, and one special end-of-text token that delimits different documents and can start generation as well. And there are 12 layers, there are 12 heads in the attention, and the dimension of the Transformers was 768. So here's how we can now load the parameters from Hugging Face to our code here and initialize the GPT class with those parameters. So let me just copy paste a bunch of code here. And I'm not going to go through this code too slow, too quickly, too slowly, because honestly, it's not that interesting. It's not that exciting. We're just loading the weights. So it's kind of dry. But as I mentioned, there are 4 models in this miniseries of GPT-2. This is some of the Jupyter code, code that we had here on the right. I'm just porting it over. These are the hyperparameters of the GPT-2 models. We're creating the config object and creating our own model. And then what's happening here is we're creating the state dict both for our model and for the Hugging Face model. And then what we're doing here is we're going over to hugging_face_model_keys and we're copying over those tensors. And in the process, we are kind of ignoring a few of the buffers. They're not parameters, they're buffers. So for example, attention.bias, that's just used for the autoregressive mask. And so we are ignoring some of those masks, and that's it. And then one additional kind of annoyance is that this comes from the TensorFlow repo, and I'm not sure how— this is a little bit annoying, but some of the weights are transposed. From what PyTorch would want. And so manually I hardcoded the weights that should be transposed, and then we transpose them if that is so. And then we return this model. So the from_pretrained is a constructor or a class method in Python that returns the GPT object if we just give it the model type, which in our case is GPT-2, the smallest model that we're interested in. So this is the code and this is how you would use it. And we can pop open the terminal here in VS Code and we can python train_gpt2.py and fingers crossed. Okay, so we didn't crash. And so we can load the weights and the biases and everything else into our NN module. But now let's also get additional confidence that this is working and let's try to actually generate from this model. Okay, now before we can actually generate from this model, we have to be able to forward it. We didn't actually write that code yet, so here's the forward function. So the input to the forward is going to be our indices, our tokens, uh, token indices, and they are always of shape B by T. And so we have batch dimension of B, and then we have the time dimension of up to T, and the T can't be more than the block size. The block size is the maximum sequence length. So B by T indices arranged in sort of like a two-dimensional layout. And remember that basically every single row of this is of size up to block size, and this is T tokens that are in a sequence. And then we have B independent sequences stacked up in a batch so that this is efficient. Now here we are forwarding the position embeddings and the token embeddings, and this code should be very recognizable from the previous lecture. So we basically use a range, which is kind of like a version of range but for PyTorch, and we're iterating from 0 to t and creating this positions sort of indices. And then we are making sure that they're on the same device as idx because we're not going to be training on only CPU. That's going to be too inefficient. We want to be training on GPU, and that's going to come in, in a bit. Then we have the position embeddings and the token embeddings. And the addition operation of those two. Now notice that the position embeddings are going to be identical for every single row of, of input, and so there's broadcasting hidden inside this plus where we have to create an additional dimension here, and then these two add up because the same position embeddings apply at every single row of our examples stacked up in a batch. Then we forward the transformer blocks, and finally the last layer norm and the LM head. So what comes out after forward is the logits. And if the input was B by T indices, then at every single B by T, we will calculate the logits for what token comes next in the sequence. So what is the token B(t+1)? The one on the right of this token. And vocab size here is the number of possible tokens. And so therefore this is the tensor that we're going to obtain. And these logits are just a softmax away from becoming probabilities. So this is the forward pass of the network, and now we can get logits. And so we're going to be able to generate from the model imminently. Okay, so now we're going to try to set up the identical thing on the left here that matches Hugging Face on the right. So here we've sampled from the pipeline and we sampled 5 times up to 30 tokens with a prefix of "Hello, I'm a language model," and these are the completions that we achieved. So we're going to try to replicate that on the left here. So num_return_sequences is 5, max_length is 30. So the first thing we do, of course, is we initialize our model. Then we put it into evaluation mode. Now, this is a good practice to put the model into eval when you're not going to be training it. You're just going to be using it. And I don't actually know if this is doing anything right now for the following reason. Our model up above here contains no modules or layers that actually have a different behavior at training or evaluation time. So for example, dropout, batch norm, and a bunch of other layers have this kind of behavior, but all of these layers that we've used here should be identical in both training and evaluation time. So potentially model.eval does nothing, but then I'm not actually sure if this is the case, and maybe PyTorch internals do some clever things depending on the evaluation mode inside here. The next thing we're doing here is we are moving the entire model to CUDA. So we're moving this, all of the tensors, to GPU. So I'm SSH'd here to a cloud box, and I have a bunch of GPUs on this box. And here I'm moving the entire model and all of its members and all of its tensors and everything like that. Everything gets shipped off to basically a whole separate computer that is sitting on the GPU. And the GPU is connected to the CPU and they can communicate, but it's basically a whole separate computer with its own computer architecture. And it's really well catered to parallel processing tasks like those of running neural networks. So I'm doing this so that the model lives on the GPU, a whole separate computer, and it's just going to make our code a lot more efficient because all of this stuff runs a lot more efficiently on the GPUs. So that's the model itself. Now the next thing we want to do is we want to start with this as the prefix when we do the generation. So let's actually create those prefix tokens. So here's the code that I've written. We're going to import the TikTok library from OpenAI and we're going to get the GPT-2 encoding. So that's the tokenizer for GPT-2. And then we're going to encode this string and get a list of integers, which are the tokens. Now these integers here should actually be fairly straightforward because we can just copy paste this string and we can sort of inspect what it is in tiktokenizer. So just pasting that in, these are the tokens that are going to come out. So this list of integers is what we expect tokens to become. And as you recall, if you saw my video, of course, all the tokens, they're just little string chunks, right? So these are— this is the chunkation of this string into GPT-2 tokens. So once we have those tokens, it's a list of integers, we can create a torch tensor out of it. In this case, it's 8 tokens. And then we're going to replicate these 8 tokens for 5 times to get 5 rows of 8 tokens. And that is our initial input x, as I call it here, and it lives on the GPU as well. So x now is this idx that we can put into forward to get our logits so that we know what comes as the 6th token, sorry, as the 9th token in every one of these 5 rows. Okay, and we are now ready to generate. So let me paste in one more code block here. So what's happening here in this code block is we have these x which is of size b by t, right? So batch by time. And we're going to be in every iteration of this loop, we're going to be adding a column of new indices into each one of these rows, right? And so these are the new indices and we're appending them to the sequence as we're sampling. So with each loop iteration, we get one more column into x. And all of the operations happen in the context manager of torch.nograd. This is just telling PyTorch that we're not going to be calling.backward on any of this. So it doesn't have to cache all the intermediate tensors. It's not going to have to prepare in any way for a potential backward later, and this saves a lot of space and also possibly some time. So we get our logits, we get the logits at only the last location, we throw away all the other logits, we don't need them, we only care about the last column's logits. So this is being wasteful, but this is just kind of like an inefficient implementation of sampling. So it's correct but inefficient. So we get the last column of logits, pass it through softmax to get our probabilities. Then here I'm doing top-k sampling of 50, and I'm doing that because this is the Hugging Face default. So just looking at the Hugging Face docs here of a pipeline, there's a bunch of quarks that go into Hugging Face, and I mean, it's kind of a lot honestly, but I guess the important one that I noticed is that they're using top-k by default. Which is 50. And what that does is that— so that's being used here as well. And what that does is basically we want to take our probabilities and we only want to keep the top 50 probabilities and anything that is lower than the 50th probability, we just clamp to 0 and renormalize. And so that way we are never sampling very rare tokens. The tokens we're going to be sampling are always in the top 50 of most likely tokens. And this helps keep the model kind of on track and it doesn't blabber on and doesn't get lost and doesn't go off the rails as easily. And it kind of like sticks in the vicinity of likely tokens a lot better. So this is the way to do it in PyTorch and you can step through it if you like. I don't think it's super insightful, so I'll speed through it. But roughly speaking, we get this new column of tokens. We append them on x and basically the columns of x grow. Until this while loop gets tripped up. And then finally, we have an entire X of size 5 by 30 in this case, in this example, and we can just basically print all those individual rows. So I'm getting all the rows, I'm getting all the tokens that were sampled, and I'm using the decode function from tiktokenizer to get back the string, which we can print. And so terminal, new terminal, And let me python train-gpt2. Okay, so these are the generations that we're getting. Hello, I'm a language model, not a program. Um, new line, new line, etc. Hello, I'm a language model, and one of the main things that bothers me when they create languages is how easy it becomes to create something that I mean, so this will just like blabber on, right? In all these cases. Now, one thing you will notice is that these generations are not the generations of Hugging Face here. And I can't find the discrepancy, to be honest. And I didn't fully go through all these options, but probably there's something else hiding in addition to the top-p. So I'm not able to match it up. But just for correctness, down here below in the Jupyter Notebook and using the Hugging Face model, So this is the Hugging Face model here. I was— I replicated the code, and if I do this and I run that, then I am getting the same results. So basically the model internals are not wrong, it's just I'm not 100% sure what the pipeline does in Hugging Face, and that's why we're not able to match them up. But otherwise the code is correct and we've loaded all the tensors correctly. So we're initializing the model correctly and everything here works. So long story short, we've ported all the weights, we initialized the GPT-2. This is the exact OpenAI GPT-2 and it can generate sequences and they look sensible. And now here, of course, we're initializing with GPT-2 model weights, but now we want to initialize from scratch from random numbers and we want to actually train a model that will give us sequences as good as or better than these ones in quality. And so that's what we turn to next. So it turns out that using the random model is actually fairly straightforward because PyTorch already initializes our model randomly and by default. So when we create the GPT model in the constructor, this is all— all of these layers and modules have random initializers that are there by default. So when these linear layers get created and so on, there's default constructors, for example, using the Xavier initialization that we saw in the past to construct the weights of these layers. And so creating a random model instead of a GPT-2 model is actually fairly straightforward. And we would just come here and instead we would create model equals GPT. And then we want to use the default config, gpt_config, and the default config uses the 124M parameters. So This is the random model initialization and we can run it and we should be able to get results. Now the results here of course are total garbage garble and that's because this is a random model and so we're just getting all these random token string pieces chunked up totally at random. So that's what we have right now. Now one more thing I wanted to point out by the way is In case you do not have CUDA available because you don't have a GPU, you can still follow along with, with what we're doing here to some extent, and probably not to the very end because by the end we're going to be using multiple GPUs and actually doing a serious training run. But for now you can actually follow along decently okay. So one thing that I like to do in PyTorch is I like to auto-detect the device that is available to you. So in particular you could do that like this. So here we are trying to detect the device to run on that has the highest compute capability. You can think about it that way. So by default, we start with CPU, which of course is available everywhere because every single computer will have a CPU. But then we can try to detect, do you have a GPU? If so, use CUDA. And then if you don't have CUDA, do you at least have MPS? MPS is the backend for Apple Silicon. So if you have a MacBook that is fairly new, you probably have Apple Silicon. Silicon on the inside, and then that has a GPU that is actually fairly capable, uh, depending on which MacBook you have. And so you can use MPS, which will be potentially faster than CPU. And so we can print the device here. Now once we have the device, we can actually use it in place of CUDA. So we just swap it in, and notice that here when we call model on x, if this x here is on CPU instead of GPU, then it will work fine because here in the forward, which is where PyTorch will come, when we created pos, we were careful to use the device of idx to create this tensor as well. And so there won't be any mismatch where one tensor is on CPU, one is on GPU, and that you can't combine those. But here we are carefully initializing on the correct device as indicated by the input to this model. So this will auto-detect device. For me, this will be, of course, GPU, so using device CUDA. But you can also run with, as I mentioned, another device, and it's not going to be too much slower. So if I override device here— oops— if I override device equals CPU, then we'll still print CUDA, of course, but now we're actually using CPU. 1, 2, 3, 4, 5, 6. Okay, about 6 seconds. And actually, we're not using torch.compile and stuff like that, which will speed up everything a lot faster as well. But you can follow along even on the CPU, I think, to a decent extent. So that's a note on that. Okay, so I do want to loop around eventually into what it means to have different devices in PyTorch and what it is exactly that PyTorch does in the background for you when you do something like module.to_device or where you take a torch.Tensor and do a.to_device and what exactly happens and how that works. But for now, I'd like to get to training and I'd like to start training the model. And for now, let's just say the device makes code go fast and let's go into how we can actually train the model. So to train the model, we're going to need some dataset. And for me, the best debugging, simplest dataset that I like to use is the Tiny Shakespeare dataset. And it's available at this URL, so you can wget it, or you can just search Tiny Shakespeare dataset. And so I have in my file system as just ls input.txt, so I already downloaded it. And here I'm reading the dataset, getting the first 1000 characters, and printing the first 100. Now remember that GPT-2 has roughly a compression ratio— the tokenizer has a compression ratio of roughly 3:1, So 1,000 characters is roughly 300 tokens here that will come out of this in the slice that we're currently getting. So this is the first few characters. And if you want to get a few more statistics on this, we can do word_count on input.txt. So we can see that this is 40,000 lines, about 200,000 words in this dataset, and about 1 million bytes in this file. And knowing that this file is only ASCII characters, there's no crazy Unicode here as far as I know, and so every ASCII character is encoded with 1 byte, and so this is the same number, roughly 1 million characters inside this dataset. So that's the dataset size by default, very small and minimal dataset for debugging to get us off the ground. In order to tokenize this dataset, we're going to get tiktoken encoding for GPT-2. Encode the data, the first 1,000 characters, and then I'm only going to print the first 24 tokens. So these are the tokens as a list of integers, and if you can read GPT-2 tokens, you will see that 198 here, you'll recognize that as the slashing character, so that is a newline. And then here, for example, we have two newlines, so that's 198 twice here. So this is just a tokenization of the first 24 tokens. So what we want to do now is we want to actually process these token sequences and feed them into a transformer. And in particular, we want them— we want to rearrange these tokens into this idx variable that we're going to be feeding into the transformer. So we don't want a single very long one-dimensional sequence. We want an entire batch where each sequence is up to— is basically t tokens, and t cannot be larger than the maximum sequence length. And then we have these t, t long sequences of tokens, and we have b independent examples of sequences. So how can we create a b by t tensor that we can feed into the forward out of these one-dimensional sequences? So here's my favorite way to, to achieve this. So if we take torch and then we create a tensor object out of this list of integers and just the first 24 tokens, My favorite way to do this is basically you do a.view of, for example, 4 by 6, which multiply to 24. And so it's just a 2-dimensional rearrangement of these tokens. And you'll notice that when you view this 1-dimensional sequence as 2-dimensional, 4 by 6 here, the first 6 tokens up to here end up being the first row. The next 6 tokens here end up being the second row. And so on. And so basically it's just going to stack up the, the every 6 tokens in this case as independent rows, and it creates a batch of tokens in this case. And so for example, if we are token 25 in the transformer, when we feed this in and this becomes the idx, this token is going to see these 3 tokens and is going to try to predict that 198 comes next. So in this way, we are able to create this 2-dimensional batch. That's, that's quite nice. Now, in terms of the label that we're going to need for the target to calculate the loss function, how do we get that? Well, we could write some code inside the forward pass because we know that the next token in a sequence, which is the label, is just to the right of us. But you'll notice that actually we, for this token at the very end, 13, we don't actually have the next correct token because we didn't load it. So, uh, we actually didn't get enough information here. So I'll show you my favorite way of basically getting these batches. And I like to personally have not just the input to the transformer, which I like to call x, but I also like to create the labels tensor, which is of the exact same size as x but contains the targets at every single position. And so here's the way that I like to do that. I like to make sure that I fetch +1 token because we need the ground truth for the very last token, for 13. And then when we're creating the input, we take everything up to the last token, not including, and view it as for Weissex. And when we're creating targets, we do the buffer but starting at index 1, not index 0. So we're skipping the first element and we view it in the exact same size. And then when I print this, here's what happens, where we see that basically, as an example, for this token 25, its target was 198, and that's now just stored at the exact same position in the targets tensor, which is 198. And also this last token 13 now has its label, which is 198, and that's just because we loaded this +1 here. So basically this is the way I like to do it. You take long sequences, you view them in two-dimensional terms so that you get batches of time. And then we make sure to load one additional token. So we basically load a buffer of tokens of B times T plus 1, and then we sort of offset things and view them. And then we have two tensors. One of them is the input to the transformer, and the other exactly is the labels. And so let's now reorganize this code. And create a very simple DataLoader object that tries to basically load these tokens and feed them to the transformer and calculate the loss. Okay, so I've reshuffled the code here accordingly. So as you can see here, I'm temporarily overriding to_run_on_cpu and importing tok_token, and all of this should look familiar. We're loading 1,000 characters. I'm setting bt to just be 4 and 32 right now, just because we're debugging, we just want to have a single batch that's very small. And all of this should now look familiar and follows what we did on the right. And then here we get the— we create the model and get the logits. And so here, as you see, I already ran this, only runs in a few seconds, but because we have a batch of 4 by 32, our logits are now of size 4 by 32 by 50,257. So those are the logits for what comes next at every position. And now we have the labels which are stored in y. So now is the time to calculate the loss and then do the backward pass and then the optimization. So let's first calculate the loss. Okay, so to calculate the loss, we're going to adjust the forward function of this nn.Module in the model. And in particular, we're not just going to be returning logits, but also we're going to return the loss. And we're going to not just pass in the input indices, but also the targets. In y. And now we will print not logits.shape anymore, we're actually going to print the loss function, and then sys.exit of 0 so that we skip some of the sampling logic. So now let's swing up to the forward function, which gets called there, because now we also have these optional targets. And when we get the targets, we can also calculate the loss. And remember that we want to basically return logits_loss, and loss by default is None. But let's put this here. If targets is not None, then we want to calculate the loss. And Copilot is already getting excited here and calculating the— what looks to be correct loss. It is using the cross-entropy loss, as is documented here. So this is a function in PyTorch under the functional. Now, what is actually happening here? Because it looks a little bit scary. Basically, the f.cross_entropy does not like multi-dimensional inputs. It can't take a B by T by vocab size. So what's happening here is that we are flattening out this 3-dimensional tensor into just 2 dimensions. The first dimension is going to be calculated automatically, and it's going to be B times T. And then the last dimension is vocab size. So basically this is flattening out this 3-dimensional tensor of logits to just be 2-dimensional, B times T, all individual examples and vocab size on, uh, in terms of the length of each row. And then it's also flattening out the targets, which are also 2-dimensional at this stage, but we're going to just flatten them out so they're just a single tensor of B times T. And this can then pass into cross_entropy to calculate a loss, which we return. So this should basically at this point run because it's not too complicated. So let's run it and let's see if we should be printing the loss. And here we see that we printed 11, roughly. And so, um, And notice that this is tensor of a single element, which is this number 11. Now we also want to be able to calculate a reasonable, uh, kind of starting point for a random initialization network. So we covered this in previous videos, but our vocabulary size is 50,257. At initialization of the network, you would hope that every vocab element is getting roughly a uniform probability. So that we're not favoring at initialization any token way too much. We're not confidently wrong at initialization. So what we're hoping is that the probability of any arbitrary token is roughly 1 over 50,257. And now we can sanity check the loss because remember that the cross-entropy loss is just basically the negative log likelihood. So if we now take this probability and we take it through the natural logarithm, and then we do the negative, that is the loss we expect at initialization. And we covered this in previous videos. So I would expect something around 10.82, and we're seeing something around 11. So it's not way off. This is roughly the probability I expect at initialization. So that tells me that at initialization, our probability distribution is roughly diffuse. It's a good starting point, and we can now perform the optimization and tell the network which elements you know, should follow correctly in what order. So at this point we can do a loss.backward, calculate the gradients, and do an optimization. So let's get to that. Okay, so let's do the optimization now. So here we have the loss— this is how we get the loss. But now basically we want a little for loop here. So for i in range, let's do 50 steps or something like that. Let's create an optimizer object. In PyTorch. And so here we are using the Adam optimizer, which is an alternative to the stochastic gradient descent optimizer SGD that we were using. So SGD is a lot simpler, Adam is a bit more involved. And I actually specifically like the AdamW variation because in my opinion, it kind of just like fixes a bug. So AdamW is a bug fix of Adam is what I would say. When we go to the documentation for AdamW, Oh my gosh. We see that it takes a bunch of hyperparameters and it's a little bit more complicated than the SGD we were looking at before because in addition to basically updating the parameters with the gradient scaled by the learning rate, it keeps these buffers around and it keeps two buffers, the m and the v, which it calls the first and the second moment. So something that looks a bit like momentum and something that looks a bit like RMSProp. If you're familiar with it, but you don't have to be. It's just kind of like a normalization that happens on each gradient element individually and speeds up the optimization, especially for language models. But I'm not going to go into the detail right here. We're going to treat it as a bit of a black box, and it just optimizes the objective faster than SGD, which is what we've seen in the previous lectures. So let's use it as a black box in our case. Create the optimizer object. And then go through the optimization. The first thing to always make sure the Copilot did not forget to zero the gradients. So, um, always remember that you have to start with a zero gradient. Then when you get your loss and you do a.backward,.backward adds to gradients, so it deposits gradients. It, it always does a on on whatever the gradients are, which is why you must set them to 0. So this accumulates the gradient from this loss. And then we call the step function on the optimizer to update the parameters and to decrease the loss. And then we print the step and the loss..item is used here because loss is a tensor with a single element..item will actually convert that to a single float. And this float will live on the CPU. So this gets to some of the internals again of the devices, but loss is a tensor with a single element and it lives on GPU for me because I'm using GPUs. When you call.item, PyTorch behind the scenes will take that one-dimensional tensor, ship it back to the CPU memory, and convert it into a float that we can just print. So this is the optimization and this should probably just work. Let's see what happens. Actually, sorry, let me— instead of using CPU override, let me delete that so this is a bit faster for me and it runs on CUDA. Oh, expected all tensors to be on the same device but found at least two devices, CUDA 0 and CPU. Cuda0 is the 0th GPU because I actually have 8 GPUs on this box. So the 0th GPU in my box and CPU. And model we have moved to device, but when I was writing this code I actually introduced a bug because buff we never moved to device. And you have to be careful because you can't just do buff.to_device. It's not stateful, it doesn't convert it to be a device. It didn't instead returns pointer to a new memory which is on the device. So you see how we can just do model.to_device, but it does not apply to tensors. You have to do buff equals buff.to_device, and then this should work. Okay, so what do we expect to see? We expect to see a reasonable loss in the beginning, and then we continue to optimize just a single batch. And so we want to see that we can overfit this single batch. We can we can crush this little batch and we can perfectly predict the indices on just this little batch. And indeed, that is roughly what we're seeing here. So, um, we started off at roughly 10.82, 11 in this case, and then as we continue optimizing on this single batch without loading new examples, we are making sure that we can overfit a single batch and we are getting to very, very low loss. So the transformer is memorizing this single individual batch. And one more thing I didn't mention is the learning rate here is 3e-4, which is a pretty good default for most optimizations that you want to run at a very early debugging stage. So this is our simple inner loop and we are overfitting a single batch and this looks good. So now what comes next is we don't just want to overfit a single batch. We actually want to do an optimization. So we actually need to iterate these epochs. Xy batches and create a little DataLoader that makes sure that we're always getting a fresh batch and that we're actually optimizing a reasonable objective. So let's do that next. Okay, so this is where I came up with and I wrote a little DataLoader lite. So what this DataLoader does is we're importing toktoken up here, reading the entire text file from this single input.txt, tokenizing it, and then we're just printing the number of tokens in total. And the number of batches in a single epoch of iterating over this dataset. So how many unique batches do we output before we loop back around to the beginning of the document and start reading it again? So we start off at position 0 and then we simply walk the document in batches of B times T. So we take chunks of B times T and then always advance by B times T. And, um, it's important to note that we're always advancing our position by exactly B times T. But when we're fetching the tokens, we're actually fetching from current position to B times T plus 1. And we need that plus 1 because remember, we need the target token for the last token in the current batch. And so that way we can do the XY exactly as we did it before. And if we are to run out of data, we'll just loop back around to 0. So this is one way to write a very, very simple DataLoader that simply just goes through the file in chunks and is good enough for us for current purposes. And we're going to complexify it later. And now we'd like to come back around here and we'd like to actually use our DataLoader. So the import tiktoken has moved up and actually all of this is now useless. So instead we just want a train_loader for the training data. And we want to use the same hyperparameters for 4. So batch size was 4 and time was 32. And then here we need to get the XY for the current batch. So let's see if Copilot gets it because this is simple enough. So we call the next batch and then we make sure that we have to move our tensors from CPU to the device. So here, when I converted the tokens, notice that I didn't actually move these tokens to the GPU. I left them on the CPU, which is the default. And that's just because I'm trying not to waste too much memory on the GPU. In this case, this is a tiny dataset and it would fit, but it's fine to just ship it to GPU right now for our purposes right now. So we get the next batch. We keep the DataLoaderSimpleCPU class. And then here we actually ship it to the GPU and do all the computation. And let's see if this runs. So python train_gpt2.py. And what do we expect to see before this actually happens? What we expect to see is now we're actually getting the next batch. So we expect to not overfit a single batch. And so I expect our loss to come down, but not too much. And that's because I still expect it to come down because In the 50,257 tokens, many of those tokens never occur in our dataset. So there are some very easy gains to be made here in the optimization by, for example, taking the biases of all the logits that never occur and driving them to negative infinity. And that would basically just— it's just that all these crazy unicodes or different languages, those tokens never occur, so their probability should be very low. And so the gains that we should be seeing are along the lines of basically deleting the usage of tokens that never occur. That's probably most of the loss gain that we're going to see at this scale right now, but we shouldn't come to zero because we are only doing 50 iterations, and I don't think that's enough to do an epoch right now. So let's see what we got. We, we have 338,000 tokens, which makes sense with our 3:1 compression ratio because there are one million characters. So one epoch with the current setting of B and T will take 2,600 batches, and we're only doing 50 batches of optimization in here. So we start off in a familiar territory as expected, and then we seem to come down to about 6.6. So basically things seem to be working okay right now with respect to our expectations. So that's good. Okay, next I want to actually fix a bug that we have in our code. It's not a major bug, but it is a bug with respect to how GPT-2 training should happen. So the bug is the following: we were not being careful enough when we were loading the weights from Hugging Face, and we actually missed a little detail. So if we come here, notice that the shape of these two tensors is the same. So this one here is the token embedding at the bottom of the transformer, right? So, and this one here is the language modeling head at the top of the transformer. And both of these are basically two-dimensional tensors and their shape is identical. So here the first one is the output embedding, the token embedding, and the second one is this linear layer at the very top, the classifier layer. Both of them are of shape 50,257 by 768. This one here is giving us our token embeddings at the bottom, and this one here is taking the 768 channels of the transformer and trying to upscale that to 50,257 to get the logits for the next token. So they're both the same shape, but more than that, actually, if you look at comparing their elements, In PyTorch, this is an element-wise equality, so then we use.all, and we see that every single element is identical. And more than that, we see that if we actually look at the data pointer, this is what— this is a way in PyTorch to get the actual pointer to the data and the storage. We see that actually the pointer is identical. So not only are these two separate tensors that happen to have the same shape and elements, they're actually pointing to the identical tensor. So what's happening here is that this is a common weight tying scheme that actually comes from the original, from the original Attention is All You Need paper and actually even the reference before it. So if we come here, Embeddings and Softmax in the Attention is All You Need paper, they mentioned that in our model, we shared the same weight matrix between the two embedding layers and the pre-softmax linear transformation similar to Thirty. Um, so this is an awkward way to phrase that these two are shared and they're tied and they're the same matrix. And the Thirty reference is this paper. Um, so this came out in 2017 and you can read the full paper, but basically it argues for this weight tying scheme. And I think Intuitively, the idea for why you might want to do this comes from this paragraph here. And basically, you, you can observe that, um, you actually want these two matrices to behave similar in the following sense. If two tokens are very similar semantically, like maybe one of them is all lowercase and the other one is all uppercase, or it's the same token in a different language or something like that, If you have similarity between two tokens, presumably you would expect that they are nearby in the token embedding space. But in the exact same way, you'd expect that if you have two tokens that are similar semantically, you'd expect them to get the same probabilities at the output of a transformer because they are semantically similar. And so both positions in the transformer at the very bottom and at the top have this property that similar tokens similar tokens should have similar embeddings or similar weights. And so this is what motivates their exploration here. And they, they kind of, you know, I don't want to go through the entire paper, and you can go through it, but this is what they observe. They also observe that if you look at the output embeddings, they also behave like word embeddings, um, if you, um, if you just kind of try to use those weights as word embeddings. So they kind of observe this similarity, they try to tie them, and they observe that they can get much better performance in that way. And so this was adopted in the Attention is All You Need paper, and then it was used again in GPT-2 as well. So I couldn't find it in the Transformers implementation. I'm not sure where they tie those embeddings, but I can find it in the original GPT-2 code introduced by OpenAI. So So this is OpenAI GPT-2 source model. And here where they are forwarding this model, and this is in TensorFlow, but that's okay. We see that they get the WTE token embeddings. And then here is the encoder of the token embeddings and the position. And then here at the bottom, they use the WTE again to do the logits. So when they get the logits, it's a matmul of this output from the transformer and the WTE tensor is reused. And so the WTE tensor basically is used twice on the bottom of the transformer and on the top of the transformer. And in the backward pass, we'll get gradients contributions from both branches, right? And these gradients will add up on the WTE tensor. So we'll get a contribution from the classifier layer, and then at the very end of the transformer, we'll get a contribution at the, at the bottom of it flowing again into the WTE tensor. So we want to— we are currently not sharing WTE in our code, but we want to do that. So weight sharing scheme, and one way to do this, let's see if Copilot gets it. Oh, it does. Okay, so this is one way to do it. Basically relatively straightforward. What we're doing here is we're taking the wte.weight and we're simply redirecting it to point to the elem_head. So this basically copies the data pointer, right? It copies the reference, and now the wte.weight becomes orphaned, the old value of it, and PyTorch will clean it up. Python will clean it up. And so we are only left with a single tensor and it's going to be used twice in the forward pass. And this is to my knowledge all that's required. So we should be able to use this and this should probably train. We're just going to basically be using this exact same tensor twice. And we weren't being careful with tracking the likelihoods. But according to the paper and according to the results, you'd actually expect slightly better results doing this. And in addition to that, one other reason that this is very, very nice for us is that this is a ton of parameters, right? What is the size here? It's 768 times 50,257. So this is 40 million parameters and this is a 124 million parameter model. So 40 divided by 124. So this is like 30% of the parameters are being saved using this weight time scheme. And so this might be one of the reasons that this is working slightly better. If you're not training the model long enough because of the weight tying, you don't have to train as many parameters. And so you become more efficient in terms of the training process because you have fewer parameters and you're putting in this inductive bias that these two embeddings should share similarities between tokens. So this is the weight tying scheme, and we've saved a ton of parameters, and we expect our model to work slightly better because of this scheme. Okay, next I would like us to be a bit more careful with the initialization and to try to follow the way GPT-2 initialized their model. Now unfortunately, the GPT-2 paper and the GPT-3 paper are not very explicit about initialization, so we kind of have to read between the lines. And instead of going to the paper, which is quite vague, there's a bit of information in the code that OpenAI released. So when we go to the model.py, we see that when they initialize their weights, they are using the standard deviation of 0.02. And that's how they— they— so this is a normal distribution for the weights, and the standard deviation is 0.02. For the bias, they initialize that with 0. And then when we scroll down here, Why is this not scrolling? The token embeddings are initialized at 0.02 and position embeddings at 0.01 for some reason. So those are the initializations and we'd like to mirror that in GPT-2 in our module here. So here's a snippet of code that I sort of came up with very quickly. So what's happening here is at the end of our initializer for the GPT module, we're calling the apply function of nn.Module, and that iterates all the submodules of this module and applies init_weights function on them. And so what's happening here is that we're in— we're iterating all the modules here. And if they are an nn.Linear module, then we're going to make sure to initialize the weight using a normal with a standard deviation of 0.02. If there's a bias in this layer, we will make sure to initialize that to 0. Note that the 0 initialization for the bias is not actually the PyTorch default. By default, the bias here is initialized with a uniform, so that's interesting. So we make sure to use 0, and for the embedding, we're just going to use 0.02 and keep it the same. So we're not going to change it to 0.01 for positional because it's about the same. And then if you look through our model, the only other layer that requires initialization and that has parameters is the layer norm. And the PyTorch default initialization sets the scale in the layer norm to be 1 and the offset in the layer norm to be 0. So that's exactly what we want. And so we're just going to keep it that way. And so this is the default initialization if we are following the— where is it— the GPT-2 source code that they released. I would like to point out, by the way, that Typically, the standard deviation here on this initialization, if you followed Javier initialization, would be 1 over the square root of the number of features that are incoming into this layer. But if you'll notice, actually, 0.02 is basically consistent with that because the, the model sizes inside these transformers for GPT-2 are roughly 768, 1600, etc. So 1 over the square root of, for example, 768 gives us 0.03. If we plug in 1600, we get 0.02. If we plug in 3 times that, 0.014, etc. So basically 0.02 is roughly in the vicinity of reasonable values for these initializations anyway. So it's not completely crazy to be hardcoding 0.02 here, but you'd like typically some something that grows with the model size instead, but we will keep this because that is the GPT-2 initialization per their source code. But we are not fully done yet on initialization because there's one more caveat here. So here, a modified initialization which accounts for the accumulation on the residual path with model depth is used. We scale the weight of residual layers initialization by a factor of 1 over square root of n, where n is the number of residual layers. So this is what the GPT-2 paper says. So we have not implemented that yet, and we can do so now. Now I'd like to actually kind of like motivate a little bit what they mean here, I think. So here's roughly what they mean. If you start out with zeros in your residual stream, remember that each residual stream is of this form where we continue adding to it. x is x plus something. Some kind of contribution. So every single block of the residual network contributes some amount and it gets added. And so what ends up happening is that the variance of the activations in residual stream grows. So here's a small example. If we start at 0 and then we, for 100 times, we have sort of this residual stream of 768 zeros and then 100 times we add, um, random, which is a normal distribution, 0 mean, 1 standard deviation. If we add to it, then by the end the residual stream has grown to have standard deviation of 10, and that's just because, um, we're always adding, um, these numbers. And so this scaling factor that they use here exactly compensates for that growth. So if we take n and we basically scale down every one of these contributions into the residual stream by 1 over the square root of n, so 1 over the square root of n is n to the negative 0.5, right, because n to the 0.5 is the square root and then 1 over the square root is n negative 0.5. If we scale it in this way, then we see that we actually get 1. So this is a way to control the growth of activations inside the residual stream in the forward pass. And so we'd like to initialize in the same way where these weights that are at the end of each block, so this cproj, uh, layer, uh, the GPT paper proposes to scale down those weights by 1 over the square root of the number of residual layers. So one crude way to implement this is the following. I don't know if this is PyTorch sanctioned, but it works for me. Is we all do in the initialization, see that self.special NanogPT scale init is 1. So we're setting kind of like a flag for this module. There must be a better way in PyTorch, right? But I don't know. Okay, so we're basically attaching this flag and trying to make sure that it doesn't conflict with anything previously. And then when we come down here, this std should be 0.02 by default, but then if has_attr module of this thing, then std times equals What? Copilot is not guessing correctly. So we want 1 over the square root of the number of layers. So the number of residual layers here is twice times self.config.layers, and then this times -0.5. So we want to scale down that standard deviation, and this should be correct and implement that. I should clarify, by the way, that the 2 times number of layers comes from the fact that every single one of our layers in the transformer actually has 2 blocks that add to the residual pathway, right? We have the attention and then the MLP. So that's where the 2 times comes from. And the other thing to mention is that what's slightly awkward, but we're not going to fix it, is that because we are weight sharing the WTE and the LM head, In this iteration over all submodules, we're going to actually come around to that tensor twice. So we're going to first initialize it as an embedding with 0.02, and then we're going to come back around it again in a linear and initialize it again using 0.02. And it's going to be 0.02 because the LM head is of course not, not scaled. So it's not going to come here. It's just, it's going to be basically initialized twice using the identical same initialization, but that's okay. And then scrolling over here, I added some code here so that we have reproducibility to set the seeds. And now we should be able to python train_gpt2.py and let this running. And as far as I know, this is the GPT-2 initialization in the way we've implemented it right now. So this looks reasonable to me. Okay, so at this point we have the GPT-2 model, we have some confidence that it's correctly implemented, we've initialized it properly, and we have a DataLoader that's iterating through data batches, and we can train. So now comes the fun part. I'd like us to speed up the training by a lot so we're getting our money's worth with respect to the hardware that we are using here, and we're going to speed up the training by quite a bit. Now you always want to start with what hardware do you have, what does it offer, and are you fully utilizing it. So in my case, if we go to NVIDIA SMI, we can see that I have 8 GPUs and each one of those GPUs is an A100 SXM 80GB. So this is the GPU that I have available to me in this box. Now, when I look— when I use to spin up these kinds of boxes, by the way, my favorite place to go to is Lambda Labs. They do sponsor my development and that of my projects, but I— this is my favorite place to go, and this is where you can spin up one of these machines and you pay per hour, and it's very, very simple. So I like to spin them up and then connect VS Code to it, and that's how I develop. Now when we look at the A100s that are available here, A100 80GB SXM is the GPU that I have here. And we have a bunch of numbers here for how many calculations you can expect out of this GPU. So when I come over here and I break in right after here, so python train_gpt, so I'm breaking in right after we calculate the logits and the loss. And the interesting thing I'd like you to note is when I do logits.dtype, this prints torch.float32. So by default in PyTorch, when you create tensors, and this is the case for all the activations and for the parameters of the network and so on, by default everything is in float32. That means that every single number, activation or weight and so on, is using a float representation that has 32 bits. And that's actually quite a bit of memory. And it turns out empirically that for deep learning as a computational workload, this is way too much. And deep learning and the training of these networks can tolerate significantly lower precision. Not all computational workloads can tolerate small precision. So for example, if we go back to the datasheet, you'll see that actually these GPUs support up to FP64, and this is quite useful, I understand, for a lot of scientific computing applications, and there they really need this, but we don't need that much precision for deep learning training. Training. So currently we are here, FP32, and with this code as it is right now, we expect to get at most 19.5 teraflops of performance. That means we're doing 19.5 trillion operations, floating point operations. So this is floating point multiply, add, most, most likely. And so these are the floating point operations. Now notice that if we are willing to go down in precision, so TF32 is a lower precision format we're going to see in a second, you can actually get an 8x improvement here, and if you're willing to go down to float16 or bfloat16, you can actually get times 16x performance, all the way to 312 teraflops. You see here that NVIDIA likes to cite numbers that have an asterisk here. This asterisk says with sparsity, but we are not going to be using sparsity in our code, and I don't know that this is very widely used in the industry right now. So most people look at this number here without sparsity, and you'll notice that we could have got even more here, but this is int8, and int8 is used for inference, not for training, because int8 has a— it basically has uniform spacing, and we actually require a float so that we get a better match to the normal distributions that occur during training of neural networks, where both activations and weights are distributed as a normal distribution. And so floating points are really important to, to match that representation. So we're not typically using int8 for training, but we are using it for inference. And if we bring down the precision, we can get a lot more teraflops out of the tensor cores available in the GPUs. We'll talk about that in a second. But in addition to that, if all of these numbers have fewer bits of representation, it's going to be much easier to move them around. And that's where we start to get into the memory bandwidth and the memory of the model. So not only do we have a finite capacity of the number of bits that our GPU can store, But in addition to that, there's a speed with which you can access this memory, and you have a certain memory bandwidth. It's a very precious resource, and in fact, many of the deep learning workloads for training are memory-bound. And what that means is actually that the tensor cores that do all these extremely fast multiplications, most of the time they're waiting around, they're idle, because we can't feed them with data fast enough. We can't load the data fast enough from memory. So typical utilization of your hardware, if you're getting 60% utilization, you're actually doing extremely well. So half of the time in a well-tuned application, your tensor cores are not doing multiplies because the data is not available. So the memory bandwidth here is extremely important as well. And if we come down in the precision for all the floats, all the numbers, weights, and activations, suddenly require less memory, so we can store more and we can access it faster. So everything speeds up and it's amazing. And now let's reap the benefits of it. And let's first look at the TensorFlow 32 format. Okay, so first of all, what are tensor cores? Well, tensor cores— tensor core is just an instruction in the A100 architecture, right? So what it does is it does basically a little 4x4 matrix multiply. So this is just matrix multiplication here of 4x4 matrices. And there are multiple configurations as to what precision any of these matrices are, in what precision the internal accumulate happens, and then what is the output precision, input precision, etc. So there's a few switches, but it's basically a 4x4 multiply. And then anytime we have any operations that require matrix multiplication, they get broken up into these— into this instruction of little 4x4 multiply. And so everything gets broken up into this instruction because it's the fastest way to multiply matrices. And it turns out that most of the computational work that we're doing up above, all of it really is matrix multiplication. Most of the work computationally happens in the linear layers. Linear, linear, etc. There's a few things sandwiched in between. So there's some additions in residuals, there's some GeLU nonlinearities, there's some layer norms, etc. But if you just time them, you'll see that these are nothing. Like, basically the entire transformer is just a bunch of matrix multiplications, really. And especially at this small scale, 124 million parameter model, actually the biggest matrix multiplication by far is the classifier layer at the top. That is a massive matrix multiply of going from 768 to 50,257. And that matrix multiply dominates anything else that happens in that network, roughly speaking. So it's matrix multiplies that become a lot faster, which are hidden inside our linear layers, and they're accelerated through Tensor Cores. Now, the best reference I would say for Tensor Cores is basically just go to the, um, A100 architecture white paper, and then it's pretty detailed, and but I think people— it's like relatively readable mostly if you half understand what's happening. So Figure 9, TensorFlow 32. So this is the explanation basically for TF32 and what happens here, and you see that there's many configuration options here available. So the input operands and what precision are they in, the accumulator and what basically the internal representation within the instruction when you do the accumulate of this matrix multiplication. So the intermediate plus equals of the intermediate little vector multiplies here, that all happens in FP32. And then this is an ATX improvement, as I mentioned, to the ops that we get. So TF32 specifically, we're looking at this row here. And the way this works is, um, normally FP32 has 32 bits. TF32 is the exact same bits. We have 1 sign bit, we have 8 exponent bits, except the mantissa bits get cropped in the float. And so basically, um, we end up with just 19 bits instead of 32 bits because the last 13 bits get truncated, they get dropped. And all of this is internal to the instruction, so none of it is visible to anything in our PyTorch. None of our PyTorch code will change, all of the numbers will look identical. It's just that when you call the Tensor Core instruction internally in the hardware, it will crop out these 13 bits and that allows it to calculate this little matrix multiply significantly faster, 8x faster. Now of course this speedup comes at a cost, and the cost is that we are reducing the precision. Our accumulate is still an FP32, our output is FP32, our inputs are FP32, but internally things get truncated in the operands to perform the operation faster, and so our results are starting to be a bit more approximate. Approximate, but empirically when you actually train with this, you basically can't tell the difference. So the reason I like TF32 is because if you can tolerate a little bit of a precision fudge, then this is free. Like, none of your code sees this. It's fully internal to the operation, and the operation to you just goes 8x faster, and it's a bit more approximate. And so it's a pretty sweet spot, I would say, in optimization. And let's see what that looks like first. So I've set up our code to just time the iterations. So import time. I changed the hyperparameters so that we have something a bit more that reflects kind of workload that we want to run because we want to do a fairly large run at the end of this. So let's use batch size 16 and let's now use the actual GPT-2 maximum sequence length of 1024 tokens. So this is the configuration. And then for 50 iterations. I'm just doing something very lazy here. I'm doing time.time to get the current time, and then this is the optimization loop. And now I want to time how long this takes. Now one issue with working with GPUs is that as your CPU— when your CPU runs, it's just scheduling work on GPU. It's ordering some work, right? And so it sends a request and then it continues running. And so we can actually— it can happen sometimes that we sort of speed through this and we queue up a lot of kernels to run on the GPU. And then the CPU sort of like gets here and takes time.time. But actually the GPU is still running because it takes it time to actually work through the work that was scheduled to run. And so you're just building up a queue for the GPU. And so actually, if you need to, you want to wait torch.cuda.syncronize(), and this will wait for the GPU to finish all the work that was scheduled to run up above here, and then we can actually take the time. So basically, we're waiting for the GPU to stop this iteration, take the time, and then we're going to just print it. So here I'm going to run the training loop, and here on the right, I'm watching NVIDIA SMI. So we start off at 0, we're not using the GPU, and then by default PyTorch will use GPU 0, so we see that it gets filled up and we're using 35 gigabytes out of 80 gigabytes available. And then here on the left we see that because we've cranked up the batch size, now it's only 20 batches to do a single epoch on our Tiny Shakespeare, and we see that we're seeing roughly 1,000 milliseconds per per iteration here, right? So the first iteration sometimes is slower, and that's because PyTorch might be doing a lot of initializations here on the very first iteration. And so it's probably initializing all these tensors and buffers to hold all the gradients. And I'm not 100% sure all the work that happens here, but this could be a slower iteration. When you're timing your logic, you always want to be careful with that. But basically we're seeing 1,000 milliseconds per iteration, um, and so this will run for roughly 50 seconds as we have it right now. So that's our baseline in float32. One more thing I wanted to mention is that if this doesn't fit into your GPU and you're getting out of memory errors, then start decreasing your batch size until things fit. So instead of 16, try 8 or 4 or whatever you need to fit, um, the batch into your GPU. And if you have a bigger GPU, you can actually potentially get away with 32 and so on. By default, you want to basically max out the batch size that fits on your GPU, and you want to keep it nice numbers. So use numbers that have lots of powers of 2 in them. So 16 is a good number, 8, 24, 32, 48. These are nice numbers, but don't use something like 17 because that will run very inefficiently on a GPU, and we're going to see that a bit later as well. So for now, let's just stick with 16, 1024. Before. And the one thing that I added also here, and I ran it again, is I'm calculating tokens per second throughput during training because we might end up changing the batch size around over time. But tokens per second is the objective measure that we actually really care about. How many tokens of data are we training on? And what is the throughput of tokens that we're getting in our optimization? So right now we're processing and training on 163,000 tokens per second roughly, and that's a bit more objective metric. Okay, so let's now enable TF32. Now luckily PyTorch makes this fairly easy for us, and to enable TF32 you just need to do a single line, and it's this. And when we go to the PyTorch documentation here for this function, basically this tells PyTorch what kind of kernels to run, and by default I believe it is highest, highest precision for matmul, and that means that everything happens in float32, just like it did before. But if we set it to high, as we do right now, matrix multiplications will now use tf32 when it's available. My GPU is the A100, so it's an Ampere series, and therefore tf32 is available. If you have an older GPU, this might not be available for you. But for my GPU it's available. And so what I expect PyTorch to do is that every single place where we see an nn.Linear, inside there, there's a matrix multiplication, and I expect that matrix multiplication now to be running on tensor cores utilizing the tf32 precision. So this is the single line of change that is, I believe, necessary. And let's rerun this. Now we saw that In terms of the throughput that is promised to us, we're supposed to be getting 8x roughly. So let's see what happens. And that 8x came from here, right? Um, 8x. And it also came from looking at it, um, here, 156 TFLOPS instead of 19.5. Okay, so what actually happened? So we're seeing that our throughput roughly 3x, not 8x. So we are going— we're from 1000 milliseconds, we're going down to 300 milliseconds, and our throughput is now about 50,000 tokens per second. So we have a roughly 3x instead of 8x. So what happened? And basically what's happening here is, again, a lot of these workloads are memory bound. And so even though the TF32 offers in principle a lot faster throughput, all of these numbers everywhere are still float32s, and it's float32 numbers that are being shipped all over the place through the memory system, and it's just costing us way too much time to shuttle around all this data. And so even though we've made the multiply itself much faster, we are memory bound, and we're not actually seeing the full benefit that would come from this napkin math here. That said, we are getting one— 3x faster throughput, and this is free. Single line of code in PyTorch. All your variables are still float32 everywhere. It just runs faster, and it's slightly more approximate, but we're not gonna notice it basically. So that's tf32. Okay, so let's now continue. So we've exercised this row, and we saw that we can crop out some of the precision inside the operation itself. But we saw that we're still memory bound, we're still moving around all these floats, right? Otherwise, and we're paying that cost because of this. So let's now decrease the amount of stuff that we're going to be moving around, and we're going to do that by dropping down to bfloat16. So we're only going to be maintaining 16 bits per float, and we're going to use the bfloat16— I'll explain in a bit, uh, fp16 difference— and, uh, we're going to be in this row. So when we go back to the documentation here for the A100, um, we see here the precisions that are available, and this is the original FP32. The TF32 crops out the precision, and then here in BF16 you see that it is very similar to TF32, but it's even more aggressive in cropping off of the precision, the mantissa of this float. So the important thing with bfloat16 is that the exponent bits and the sign bit, of course, remain unchanged. So if you're familiar with your float numbers, and I think this should, should probably be an entire video by itself, the exponent sets the range that you can represent of your numbers, and the precision is how much precision you have for your numbers. And so the range of numbers is identical, but we can— we have fewer possibilities within that range because we are truncating the mantissa, so we have less precision in that range. But what that means is that things are actually fairly nice because we have the original range of numbers that are representable in float, but we just have less precision for it. And the difference with FP16 is that they actually touch and change the range. So FP16 cannot represent the full range of FP32. It has a reduced range, and that's where you start to actually run into issues because now you need these gradient scalars and things like that. And I'm not going to go into the detail of that in this video because that's a whole video by itself. But FP16 actually historically came first. That was available in the Volta series before Ampere. And so FP16 came first and everyone started to train FP16. But everyone had to use all these gradient scaling operations, which are kind of annoying, and it's an additional source of state and complexity. And the reason for that was because the exponent range was reduced in FP16. So that's the IEEE FP16 spec. And then they came out with BFP16 and the Ampere, and they made it much simpler because we're just truncating mantissa, we have the exact same range, and we do not need gradient scalars. So everything is much, much simpler. Now, when we do use BF16, though, we are impacting the numbers that we might be seeing in our PyTorch code. These— this change is not just local to the operation itself. So let's see how that works. There's some documentation here that— so I think this is probably the best, best page to explain how to use mixed precision in PyTorch, because there are many other tutorials and so on. Even within PyTorch documentation, that are a lot more confusing. And so I recommend specifically this one because there's 5 other copies that I would not recommend. And then when we come here, ignore everything about everything, ignore everything about gradient scalars, and only look at torch.autocast. And basically also this comes to a single line of code at the end. So this is the context manager that we want. And we want to use that in our network. When you click into the torch.autocast, autocasting, it has a few more, a bit more guideline for you. So it's telling you do not call bfloat16 on any of your tensors, just use autocast and only surround the forward pass of the model and the loss calculation. And that's the only two things that you should be surrounding. Leave the backward and the optimizer step alone. So that's the guidance that comes from the PyTorch team. So we're going to follow that guidance. And for us, because the loss calculation is inside of the model forward pass for us, we are going to be doing this. And then we don't want to be using torch.float16 because if we do that, we need to start using gradient scalars as well. So we are going to be using bfloat16. This is only possible to do in Ampere. But this means that the changes are extremely minimal, like basically just this one line of code. Let me first break in to here before we actually run this. So right after logits, I'd like to show you that different from the tf32 that we saw, this is actually going to impact our tensors. So this logits tensor If we now look at this and we look at the dtype, we suddenly see that this is now bfloat16. It's not float32 anymore. So our activations have been changed. The activations tensor is now bfloat16, but not everything has changed. So model.transformer.wte— this is the weight token embedding table. It has a.weight inside it. And the dtype of this weight, this parameter, is still torch.float32. So our parameters seem to still be in float32, but our activations, the logits, are now in bfloat16. So clearly this is why we get the mixed precision. Some things PyTorch is keeping in float32, some things PyTorch is converting to lower precision. Um, and what gets converted at what point is not super super clear. I remember scrolling down. Is it here? Blah. Okay, I can't find it. I thought it was here. Okay, there we go. So there are a few docs on when you're using this autocast, what gets converted to bfloat16 and when. So for example, only these matrix multiply-like operations get converted to bfloat16, but a lot of operations remain in float32. So in particular, a lot of normalizations like layer norms and things like that, not all of those layers might be converted. So only some layers selectively would be run in bfloat16, but things like softmax, layer norms, log softmax, so loss function calculations, a lot of those things might remain in float32 because they are more susceptible to precision changes. Matrix multiplies are fairly robust to precision changes. So some parts of the network are impacted more or less by the precision change. So basically only some parts of the, of the model are running in reduced precision. Let's take it for a spin and let's actually see what kind of improvement we achieve here. Okay, so we used to be 333 milliseconds, we're now 300, and we used to be somewhere around 50,000 tokens per second, we're now 55. So we're definitely running faster, but maybe not a lot faster, and that's because there are still many, many bottlenecks in our GPT-2. We're just getting started, but we have dropped down the precision as far as we can with my current GPU, which is a100. We're using PyTorch Autocast. Unfortunately, I don't actually exactly know what PyTorch Autocast do— does. I don't actually know exactly what's in bfloat16, what's in float32. We could go in and we could start to scrutinize it, but these are the kinds of rules that PyTorch has internally, and unfortunately they don't document it very well. So we're not going to go into that into the— in too much detail. But for now, we are training in bfloat16. We do not need a gradient scaler, and the reason things are running faster is because we are able to run tensor cores in bfloat16 now. That means we are in this row, but we are also paying in precision for this, so we expect slightly less accurate results with respect to the original FP32, but empirically in many cases this is a worth it kind of trade-off because it allows you to run faster and you could, for example, train longer and make up for the, for that precision decrease. So that's the Bfloat16 for now. Okay, so as we can see, we are currently at about 300 milliseconds per iteration and we're now going to reach for some really heavy weapons in the PyTorch arsenal. And in particular, we're going to introduce torch.compile. So torch.compile is really quite incredible infrastructure from the PyTorch team. And it's basically a compiler for neural networks. Like it's almost like GCC for C and C++ code. And this is just this GCC of neural nets. So came out a while ago and extremely simple to use. The way to use torch.compile is to do this. It's a single line of code to compile your model and return it. Now this line of code will cost you compilation time. But as you might guess, it's going to make the code a lot faster. So let's actually run that because this will take some time to run. But currently, remember, we're at 300 milliseconds and we'll see what happens. Now, while this is running, I'd like to explain a little bit of what torch.compile does under the hood. So feel free to read this page of PyTorch. But basically, there's no real good reason for you to not use torch.compile in your PyTorch. I kind of feel like You should be using it almost by default if you're not, unless you're debugging and you want your code to run really fast. And there's one line here in torch.compile that I found that actually kind of like gets to why this is faster. Speedup mainly comes from reducing Python overhead and GPU read writes. So let me unpack that a little bit. Okay, here we are. Okay, so we went from 300 milliseconds. We're now running at 129 milliseconds. So this is a 300 divided by 129, about 2.3x improvement from a single line of code in PyTorch. So quite incredible. So what is happening? What's happening under the hood? Well, when you pass the model to torch.compile, what we have here in this nn.Module, this is really just the algorithmic description of what we'd like to happen in our network. And torch.compile will analyze the entire thing and it will look at what operations you'd like to use. And with the benefit of knowing exactly what's going to happen, it doesn't have to run in what's called the eager mode. It doesn't have to just kind of like go layer by layer. Like the Python interpreter normally would start at the forward, and the Python interpreter will go, okay, let's do this operation, and then let's do that operation. And it kind of materializes all the operations as it goes through. So these calculations are dispatched and run in this order, and the Python interpreter and this code doesn't know what kind of operations are going to happen later. But torch.compile sees your entire code at the same time, and it's able to know what operations you intend to run, and it will kind of optimize that process. The first thing it will do is it will, it will take out the Python interpreter from the forward pass entirely, and it will kind of compile this entire neural net as a single object with no Python interpreter involved. So it knows exactly what's going to run, it will just run that, and it's all going to be running in efficient code. The second thing that happens is this read-write that they mentioned very briefly. So a good example of that, I think, is the Gellu nonlinearity that we've been looking at. So here we use the NN_Gellu. Now this here is me basically just breaking up the NN_Gellu which you remember has this formula. So this here is the equivalent implementation to what's happening inside GILU. Algorithmically, it's identical. Now, by default, if we just were using this instead of nn.gilu here, what would happen without torch.compile? Well, the Python interpreter would make its way here, and then it would be, okay, well, there's an input. Well, let me first— let me raise this input to the third power. And it's going to dispatch a kernel that takes your input and raises to the third power, and that kernel will run. And when this kernel runs, what ends up happening is this input is stored in the memory of the GPU. So here's a helpful example of the layout of what's happening, right? You have your CPU— this is in every single computer— there's a few cores in there, and you have your RAM, your memory. And the CPU can talk to the memory, and this is all well known. But now we've added the GPU, and the GPU is a slightly different architecture, of course. They can communicate, and it's different in that it's got a lot more cores than a CPU. All of those cores are individually a lot simpler too, but it also has memory, right? This high bandwidth memory. I'm sorry if I'm botching it. HBM. I don't even know what that stands for. I'm just realizing now. But this is the memory, and it's very equivalent to RAM, basically, in the computer. And what's happening is that input is living in the memory, and when you do input cubed, this has to travel to the GPU, to the cores, and to all the caches and registers on the actual chip of this GPU. And it has to calculate all the elements to the third, and then it saves the result back to the memory. And it's this, uh, travel time that actually causes a lot of issues. So here, remember this memory bandwidth? We can communicate about 2 terabytes per second, which is a lot, but also we have to traverse this link and it's very slow. So here on the GPU, we're on-chip and everything is super fast within the chip, but going to the memory is extremely expensive, takes extremely long amount of time. And so we load the input, do the calculations, and load back the output. And this round trip takes a lot of time. And now right after we do that, we multiply by this constant. So what happens then is we dispatch another kernel, and then the result travels back, all the elements get multiplied by a constant, and then the results travel back to the memory. And then we take the result and we add back input. And so this entire thing again travels to the GPU, adds the inputs, and gets written back. So we're making all these round trips from the memory to actually where the computation happens, because all the tensor cores and the ALUs and everything like that is all stored on the chip in the GPU. So we're doing a ton of round trips, and PyTorch, without using torch.compile, doesn't know to optimize this because it doesn't know what kind of operations you're running later. You're just telling it, raise the power to the third, then do this, then do that, and it will just do that in that sequence. But torch.compile sees your entire code. It will come here and it will realize, wait, all of these are element-wise operations. And actually what I'm going to do is I'm going to do a single trip of input to the GPU. Then for every single element, I'm going to do all of these operations while that memory is on the GPU or chunks of it. Rather, and then I'm going to write back a single time. So we're not going to have these round trips. And that's one example of what's called kernel fusion, and is a major way in which everything is sped up. So basically, if you have your benefit of hindsight, and you know exactly what you're going to compute, you can optimize your round trips to the memory. And you're not going to pay the memory bandwidth cost. And that's fundamentally what makes some of these operations a lot faster. And what they mean by read writes here. So let me erase this because we are not using it. And yeah, we should be using torch.compile, and our code is now significantly faster, and we're doing about 125,000 tokens per second. But we still have a long way to go. Before we move on, I wanted to supplement the discussion a little bit with a few more figures, uh, because this is a complicated topic, but it's worth understanding on a high level, uh, what's happening here. And I could probably spend an entire video of like 2 hours on this, but just a preview of that basically. So this chip here, that is the GPU. This chip is where all the calculations happen mostly, but this chip also does have some memory in it. But most of the memory by far is here in the high bandwidth memory, HBM, and is connected— they're connected, but these are two separate chips basically. Now here, this is a zoom in of kind of this cartoon diagram of a GPU, and we're seeing here is number one, you see this HBM. I realize it's probably very small for you, but on the sides here it says HBM, and so that's the links to the HBM. Now the HBM is again off-chip. On the chip, there are a large number of these streaming multiprocessors, Every one of these is an SM. There's 120 of them in total. And this is where a lot of the calculations happen. And this is a zoom-in of a single individual SM. It has these 4 quadrants. And see, for example, tensor core, this is where a lot of the matrix multiply stuff happens. But there's all these other units to do all different kinds of calculations for FP64, FP32, and for integers and so on. Now, so we have all this logic here to do the calculations. But in addition to that, on the chip, there is memory sprinkled throughout the chip. So L2 cache is some amount of memory that lives on the chip. And then on the SMs themselves, there's L1 cache. I realize it's probably very small for you, but this blue bar is L1. And there's also registers. And so there is memory stored here. But the way this memory is stored is very different from the way memory is stored in HBM. This is a very different implementation using— just in terms of like what the silicon looks like, it's a very different implementation. So here you would be using transistors and capacitors, and here it's a very different implementation with SRAM and what that looks like. But long story short is there is memory inside the chip, but it's not a lot of memory. That's the critical point. So this is some cart— this is an example diagram of a slightly different GPU, just like here, where it shows that, for example, typical numbers for CPU DRAM memory, which is this thing here, you might have 1 terabyte of disk, right? But it would be extremely expensive to access, especially for a GPU. You have to go through the CPU here. Now next we have the HBM. So we have tens of gigabytes of HBM memory on a typical GPU. Here, but it's, as I mentioned, very expensive to access. And then on the chip itself, everything is extremely fast within the chip, but we only have a couple 10 megabytes of memory collectively throughout the chip. And so there's just not enough space because the memory is very expensive on the chip. And so there's not a lot of it, but it is lightning fast to access in relative terms. And so basically, whenever we have these kernels, The more accurate picture of what's happening here is that we take these inputs, which live by default on the global memory, and now we need to perform some calculation. So we start streaming the data from the global memory to the chip. We perform the calculations on the chip and then stream it back and store it back to the global memory, right? And so if we are— if we don't have torch.compile, We are streaming the data through the chip, doing the calculations, and saving to the memory, and we're doing those round trips many, many times. But if it torch compiled, then we start streaming the memory as before, but then while we're on the chip, we have a chunk of the data that we're trying to process. So that chunk now lives on the chip. While it's on the chip, it's extremely fast to operate on. So if we have kernel fusion, we can do all the operations right there in an element-wise fashion. and those are very cheap, and then we do a single round trip back to the global memory. So operator fusion basically allows you to keep your chunk of data on the chip and do lots of calculations on it before you write it back, and that gives huge savings, and that's why torch.compile ends up being a lot faster, or that's one of the major reasons. So again, just a very brief intro to the memory hierarchy and roughly what torch.compile does for you. Now torch.compile is amazing, but there are operations that torch.compile will not find, and an amazing example of that is flash attention, to which we turn next. So flash attention comes from this paper from Stanford in 2022, and it's this incredible algorithm for performing attention, so, and running it a lot faster. So flash attention will come here and we will take out these 4 lines and FlashAttention implements these 4 lines really, really quickly. And how does it do that? Well, FlashAttention is a kernel fusion operation. So you see here we have in this diagram, they're showing PyTorch and you have these 4 operations. They're including dropout, but we are not using dropout here. So we just have these 4 lines of code here. And instead of those, we are fusing them into a single fused kernel of flash attention. So it's a kernel fusion algorithm, but it's a kernel fusion that torch.compile cannot find. And the reason that it cannot find it is that it requires an algorithmic rewrite of how attention is actually implemented here in this case. And what's remarkable about it is that flash attention actually, if you just count the number of flops, Flash attention does more flops than this attention here, but flash attention is actually significantly faster. In fact, they cite 7.6 times faster potentially. And that's because it is very mindful of the memory hierarchy as I described it just now. And so it's very mindful about what's in high-bandwidth memory, what's in the shared memory, and it is very careful with how it orchestrates the computation. Such that we have fewer reads and writes to the high-bandwidth memory. And so even though we're doing more flops, the expensive part is the load and store into HBM, and that's what they avoid. And so in particular, they do not ever materialize this n×n attention matrix, this ATT here. A flash attention is designed such that this matrix never gets materialized at any point, and it never gets read or written to the HBM. And this is a very large matrix, right? So because this is where all the queries and keys interact, and we're sort of getting for each head, for each batch element, we're getting a T by T matrix of attention, which is a million numbers, even for a single head at a single batch index. So basically, this is a ton of memory, and this is never materialized. And the way that this is achieved is that basically the fundamental algorithmic rewrite here relies on this online softmax trick, which was proposed previously, and I'll show you the paper in a bit. And the online softmax trick coming from a previous paper shows how you can incrementally evaluate a softmax without having to sort of realize all of the inputs to the softmax to do the normalization. And you do that by having these intermediate variables m and l, and there's an update to them that allows you to evaluate the softmax in an online manner. Now, flash attention— actually, so recently Flash Attention 2 came out as well, so I have that paper up here as well that has additional gains to how it calculates flash attention. And the original paper that this is based on basically is this online normalized calculation for softmax. And remarkably, it came out of NVIDIA, and it came out of it like really early 2018. So this is 4 years before flash attention, and this paper says that we propose a way to compute the classical softmax with fewer memory accesses and hypothesize that this reduction in memory accesses should improve softmax performance on actual hardware. And so they are extremely correct in this hypothesis, but it's really fascinating to me that they're from NVIDIA and that they had this realization but they didn't actually take it to the actual flash attention. That had to come 4 years later from Stanford. So I don't fully understand the historical— how this happened historically, but they do basically propose this online update to the softmax right here, and this is fundamentally what they reuse here to calculate the softmax in a streaming manner. And then they realize that they can actually fuse all the other operations with the online softmax calculation into a single fused kernel, FlashAttention, and that's what we are about to use. So a great example, I think, of being aware of memory hierarchy, the fact that flops don't matter, the entire memory access pattern matters, and that torch.compile is amazing, but there are many optimizations that are still available to us that potentially torch.compile cannot find. Maybe one day it could, but right now it seems like a lot to ask. So here's what we're going to do. We're going to use use flash attention. And the way to do that basically in PyTorch is we are going to comment out these 4 lines and we're going to replace them with a single line. And here we are calling this compound operation in PyTorch called scaled.product attention. And PyTorch will call flash attention when you use it in this way. I'm not actually 100% sure why torch.compile doesn't realize that these 4 lines should just call flash attention in this exact way. We have to do it again for it, which in my opinion is a little bit odd, but, um, here we are. So you have to use this compound up, and, uh, let's wait for a few moments before torch.compile gets around to it. And then let's remember that we achieved 6.05661. I have it here. That's the loss we are expecting to and we took 130 milliseconds before this change. So we're expecting to see the exact same result by iteration 49, but we expect to see faster runtime because Flash Attention is just an algorithmic rewrite and it's a faster kernel, but it doesn't actually change any of the computation and we should have the exact same optimization. So okay, so we're a lot faster, we're at about 95 milliseconds and we achieved 6 0.058. Okay, so they're basically identical up to a floating point fudge factor. So it's the identical computation, but it's significantly faster going from 130 to roughly 96. And so this is 96 divided by 130-ish, so this is maybe 27-ish percent improvement. So really interesting. And that is flush attention. Okay, we are now getting to one of my favorite optimizations, and it is simultaneously the dumbest and the most brilliant optimization. And it's always a little bit surprising to me. Anyway, so basically I mentioned a few minutes ago that there are some numbers that are nice and some numbers that are ugly. So 64 is a beautiful, nice number. 128 is even nicer. 256 is beautiful. What makes these numbers beautiful is that there are many powers of 2 inside them. You can divide by 2 many times. And examples of ugly numbers are like 13 and 17 and something like that, prime numbers, numbers that are not even, and so on. And so pretty much you always want to use nice numbers in all of your code that deals with neural networks or CUDA, because everything in CUDA works in sort of like powers of 2, and lots of kernels are written in terms of powers of 2, and there are lots of blocks of sizes 16 and 64 and so on. So everything is written in those terms, and you always have special case handling for all kinds of logic that when your inputs are not made of nice numbers. So let's see what that looks like. Basically, scan your code and look for ugly numbers is roughly the heuristic. So 3 times is kind of ugly. I'm not 100% sure, maybe this can be improved, but this is ugly and not ideal. 4 times is nice. So that's nice. 1024 is very nice. That's a power of 2. 12 is a little bit suspicious. Not too many powers of 2. 768 is great. 50257 is a really, really ugly number. It's— first of all, it's odd, so, and there's no— not too many powers of 2 in there. So this is a very ugly number and it's highly suspicious. And then when we scroll down, all these numbers are nice. And then here we have mostly nice numbers except for 25. So in this configuration of GPT-2 XL, a number of heads is 25. That's a really ugly number. That's an odd number. And actually, this did cause a lot of headaches for us recently when we were trying to optimize some kernels. To run this fast and required a bunch of special case handling. So basically these numbers are— we have some ugly numbers and some of them are easier to fix than others. And in particular the vocab size being 50,257, that's a very ugly number, very suspicious, and we're going to fix it. Now when you fix these things, one of the easy ways to do that is you basically increase the number until it's the nearest power of 2 that that you like. So here's a much nicer number, it's 50,304. And why is that? Because 50,304 can be divided by 8 or by 16 or by 32, 64, it can even be divided by 128 I think. Yeah, so it's a very nice number. So what we're going to do here is— this is the gpt-config and you see that we initialize vocab_size to 50,257. Let's override just that element to be 3,304. Okay, so everything else stays the same. We're just increasing our vocabulary size. So we're adding— it's almost like we're adding fake tokens so that vocab size has powers of 2 inside it. Now actually, what I'm doing here, by the way, is I'm increasing the amount of computation that our network will be doing. If you just count the flops, on like, do the math of how many flops we're doing. We're going to be doing more flops. And we still have to think through whether this doesn't break anything. But if I just run this, let's see what we get. Currently, this ran in maybe 96.5 milliseconds per step. I'm just kind of like eyeballing it. And let's see what kind of a result we're going to get. While this is compiling, Let's think through whether our code actually works okay when we increase the vocab size like this. Let's look at where vocab size is actually used. So we swing up to the init and we see that it's used inside the embedding table, of course, so all the way at the bottom of the transformer, and it's used at the classifier layer all the way at the top of the transformer, so in two places. And let's take a look, and we're running at 93. So 93 milliseconds instead of 96.5. So we are seeing a roughly, yeah, 4% improvement here by doing more calculations. And the reason for this is we fixed— we've made an ugly number into a nice number. Let's— I'm going to come into the explanation for that a little bit again. But for now, let's just convince ourselves that we're not breaking anything when we do this. So first of all, we've made the WTE, the embedding table for the tokens, we've made it larger. It's almost like we introduced more tokens at the bottom, and these tokens are never used because the GPT tokenizer only has tokens up to 50,256. And so we'll never index into the rows that we've added. So we're wasting a little bit of space here by creating memory that's never going to be accessed, never going to be used, etc. Now that's not fully correct because this WTE weight ends up being shared and ends up being used in the classifier here at the end. So what is that doing to the classifier right here? Well, what, what that's doing is we're predicting additional dimensions at the classifier now, and we're predicting probabilities for tokens that will of course never be present in the training set. And so therefore the network has to learn that these probabilities have to be driven to 0. And so the logits that the network produces have to drive those dimensions of the output to negative infinity. But that's no different from all the other tokens that are already in our dataset, or rather that are not in our dataset. So Shakespeare only probably uses, let's say, 1,000 tokens out of 50,257 tokens. So most of the tokens are already being driven to zero probability by the optimization. We've just introduced a few more tokens now that in a similar manner will never be used and have to be driven to zero in probability. So functionally though, nothing breaks. We're using a bit more extra memory, but otherwise this is a harmless operation as far as I can tell. But, and we're adding calculation, but it's running faster. And it's running faster because, as I mentioned, in CUDA, so many kernels use block tiles, and these block tiles are usually nice numbers, so powers of 2. So calculations are done in like chunks of 64 or chunks of 32, And when your desired calculation doesn't neatly fit into those block tiles, there are all kinds of boundary kernels that can kick in to like do the last part. So basically in a lot of kernels, they will chunkate up your input and they will do the nice part first. And then they have a whole second phase where they come back to anything that like remains and then they, process the remaining part, and the kernels for that could be very inefficient. And so you're basically spinning up all this extra compute, and it's extremely inefficient. So you might as well pad your inputs and make it fit nicely, and usually that empirically ends up actually running faster. So this is another example of a 4% improvement that we've added, and this is something that also torch.compile did not find for us. You would hope that torch.compile at some point could figure an optimization like this out, but for now, this is it. And I also have to point out that we're using PyTorch nightly, so that's why we're only seeing 4%. If you're using PyTorch 2.3.1 or earlier, you would actually see something like 30% improvement just from this change, from changing it to— from 50,257 to 53,004. So again, one of my favorite examples also of having to understand the under the hood and how it all works, and to know what kinds of things to tinker with to push the performance of your code. Okay, so at this point we have improved the performance by about 11x, right? Because we started at about 1000 milliseconds per step and we're now down to like 93 milliseconds. So that's quite good. And we're doing a much better job of utilizing our GPU resources. So I'm going to now turn to more algorithmic changes. And improvements to the actual optimization itself. And what we would like to do is we would like to follow the hyperparameters that are mentioned in the GPT-2 or GPT-3 paper. Now sadly, GPT-2 is, uh, doesn't actually say too much. It's very nice of them that they released the model weights and the code, but the paper itself is extremely vague as to the optimization details. The code itself that they released as well, the code we've been looking at this is just the inference code. So there's no training code here and very few hyperparameters. So this doesn't also tell us too much. So for that, we have to turn to the GPT-3 paper. And in the appendix of the GPT-3 paper, they have a lot more hyperparameters here for us to use. And the GPT-3 paper in general is a lot more detailed as to all the, you know, small details that go into the model training. But GPT-3 models were never released. So GPT-2, we have the weights but no details, and GPT-3, we have lots of details but no weights. So, but roughly speaking, GPT-2 and GPT-3 architectures are very, very similar. And basically, there are very few changes. The context length was expanded from 1024 to 2048. And that's kind of like the major change. And some of the hyperparameters around the transformer have changed. But otherwise, they're pretty pretty much the same model. It's just that GPT-3 was trained for a lot longer on a bigger dataset and has a lot more thorough evaluations. And the GPT-3 model is 175 billion instead of 1.6 billion in the GPT-2. So long story short, we're going to go to GPT-3 paper to follow along some of the hyperparameters. So to train all the versions of GPT-3, we use Adam with beta 1, beta 2 of 0.9 and 0.95. So let's swing over here and make sure that the betas parameter, which you can see here defaults to 0.9 and 0.999, is actually set to 0.9 and 0.95. And then the epsilon parameter, uh, you can see is the default is 1e-8, and this is also 1e-8. Let's just put it in so that we're explicit. Uh, now next up they say we clip the grad global norm of the gradient at 1.0. So what this is referring to is that once we calculate the gradients right after loss.backward, um, we basically have the gradients at all the parameter tensors. And what people like to do is basically clip them to have some kind of a maximum norm. So in PyTorch, this is fairly easy to do. Uh, it's one line of code here that we have to insert right after we calculate the gradients. And what this utility function is doing is it's calculating the global norm of the parameters. So every single gradient on all the parameters, you square it and you add it all up and you take a big square root of that. And that's the norm of the parameter vector, basically. It's the, it's the length of it, if you, if you'd like to look at it that way. And we are basically making sure that its length is no more than 1.0, and we're going to clip it. And the reason that people like to use this is that sometimes you can get unlucky during the optimization. Maybe it's a bad data batch or something like that. And if you get very unlucky in the batch, you might get really high loss, and really high loss could lead to a really high gradient. And this could basically shock your model and shock the optimization. So people like to use a gradient norm clipping to prevent the model from basically getting too big of shocks in terms of the gradient magnitude. And they upper bound it in this way. It's a bit of a hacky solution. It's about like a patch on top of like deeper issues. But people still do it fairly frequently. Now, the clip_grad_norm returns the norm of the gradient, which I like to always visualize. Because it is useful information. And sometimes you can look at the norm of the gradient. And if it's well behaved, things are good. If it's climbing, things are bad, and they're destabilizing during training. Sometimes you could get a spike in the norm. And that means there's some kind of an issue or instability. So the norm here will be norm. And let's do a 0.4f or something like that. And I believe this is just a float, and so we should be able to print that. So that's global gradient clipping. Now they go into the details of the learning rate scheduler. So they don't just use a fixed learning rate like we do here for 3e^-4, but there's actually basically a cosine decay learning rate schedule. It's got a warm-up and it's got a cosine decay to 10% over some horizon. And so we're going to implement this in a second. I just like to see the norm printed here. Okay, there we go. So what happened here is the norm is actually really high in the beginning, 30 or so. And you see that as we continue training, it kind of like stabilizes at values below 1. And this is not that crazy uncommon for the norm to be high in the very first few stages. Basically what's happening here is the model is completely random, and so there's a ton of learning happening very early in the network, but that learning is kind of like, you know, it's mostly learning the biases of the output tokens. And so it's a bit of an unstable time, but the network usually stabilizes in a very few iterations. So this looks relatively reasonable to me. Except usually I would expect— this looks a little bit funky that we go from 28 to 6 to 2 and then to 10. It's not completely insane, but it's just kind of a little bit funky. Okay, so let's now get to the learning rate scheduler. So the learning rate schedule that's used here in GPT-3 is what's called a cosine decay learning schedule with warmup. And the way this looks is that the learning rate is basically starts starts right at around 0, linearly ramps up over some amount of time, and then comes down with this cosine sort of form and comes down to some kind of a minimum learning rate that's up to you. So here the minimum learning rate is 0, but here in the paper they said that they use cosine decay for learning rate down to 10% of its value over the first 260 billion tokens, and then training continues 10% after. And there's a linear warm-up over the first 375 million tokens. So that's about the learning rate. So let's now implement this. So I already implemented it here, and the way this works is— let me scroll down first here. I changed our training loop a little bit. So this was a for i in max_steps. I just changed it to step now so that we have the notion of a step is a single optimization step in the, in in the for loop. And then here I get the lr for this step of the optimization using a new function I call get_lr. And then in PyTorch, to set the learning rate, I think this is the way to set the learning rate. It's a little bit gnarly because you have to— basically there's a notion of different parameter groups that could exist in the optimizer, and so you actually have to iterate over them even though we currently have a single param group only. And you have to set the lr in this for loop kind of style. Is my impression right now. So we have this local variable lr, we set the learning rate, and then on the bottom, I'm also printing it. So that's all the changes I made to this loop. And then of course, the get_lr is my scheduler. Now it's worth pointing out that PyTorch actually has learning rate schedulers and you can use them. And I believe there's a cosine learning rate schedule in PyTorch. I just don't really love using that code because Honestly, it's like 5 lines of code. And I fully understand what's happening inside these lines. So I don't love to use abstractions where they're kind of inscrutable, and then I don't know what they're doing. So personal style. So the max learning rate here is, let's say, 3e^-4. But we're going to see that in GPT-3 here, they have a table of what the maximum learning rate is for every model size. Size. So for, for this one, basically 12, 12-layer 768 GPT-3. So the GPT-3 small is roughly like a GPT-2 124M. We see that here they use a learning rate of 6e����. So we could actually go higher. In fact, we may want to try to follow that and just set the max_lr here at 6. Then the— that's the maximum learning rate. The min learning rate is 10% of that per description in the paper, some number of steps that we're going to warm up over, and then the maximum steps of the optimization, which I now use also in the for loop down here. And then you can go over this code if you like. It's not, it's not terribly insightful or interesting. I'm just modulating based on the iteration number which learning rate there should be. So this is the warm-up region. This is the region after the optimization, and then this is the region sort of in between, and this is where I calculate the cosine learning rate schedule. And you can step through this in detail if you'd like, but this is basically implementing this curve. And I ran this already, and this is what that looks like. So when we now run, we start at some very low number, Now note that we don't start exactly at 0 because that would be not useful to update with a learning rate of 0. That's why there's an it plus 1 so that on the 0th iteration we are not using exactly 0, we're using something very, very low. Then we linearly warm up to maximum learning rate, which in this case was 3e-4 when I ran it, but now would be 6e-4. And then it starts to decay all the way down to, um, 3e-5. Which was at the time 10% of the original learning rate. Now, one thing we are not following exactly is that they mentioned that— let me see if I can find it again. We're not exactly following what they did because they mentioned that their training horizon is 300 billion tokens, and they come down to 10% of the initial learning rate at 260 billion, and then they train after 260 with 10%. So basically their decay time is less than the max_steps_time, whereas for us they're exactly equal. So it's not exactly faithful, but it's, um, it's an okay, um, it's— this is okay for us and for our purposes right now. And, um, we're just going to use this ourselves. I don't think it makes too, too big of a difference, honestly. I should point out that what learning rate schedule you use is totally up to you. There's many different types. Cosine learning rate has been popularized a lot by GPT-2 and GPT-3, but people have come up with all kinds of other learning rate schedules, and this is kind of like an active area of research as to which one is the most effective at training these networks. Okay, next up, the paper talks about the gradual batch size increase. So there's a ramp on the batch size that is linear, and and you start with very small batch size and you ramp up to a big batch size over time. We're going to actually skip this and we're not going to work with it. And the reason I don't love to use it is that it complicates a lot of the arithmetic because you are changing the number of tokens that you're processing at every single step of the optimization. And I like to keep that math very, very simple. Also, my understanding is that this is not like a major improvement. And also my understanding is that this is not like an algorithmic optimization improvement. It's more of a systems and speed improvement. And roughly speaking, this is because, uh, in the early stages of the optimization, uh, again, the model is in a very atypical setting, and mostly what you're learning is that, um, you're mostly learning to ignore the tokens, uh, that don't come up in your training set very often. You're learning very simple biases and, and that kind of a thing. And so Every single example that you put through your network is basically just telling you, use these tokens and don't use these tokens. And so the gradients from every single example are actually extremely highly correlated. They all look roughly the same in the, in the original parts of the optimization because they're all just telling you that these tokens don't appear and these tokens do appear. And so because the gradients are all very similar and they're highly correlated, then why are you doing batch sizes of like millions when if you do a batch size of 32K, you're basically getting the exact same gradient early on in the training. And then later in the optimization, once you've learned all the simple stuff, that's where the actual work starts. And that's where the gradients become more decorrelated per examples. And that's where they actually offer you sort of statistical power in some sense. So we're going to skip this just because it kind of complicates things. And we're going to go to data are sampled without replacement during training, so until an epoch boundary is reached. So without replacement means that they're not sampling from some fixed pool and then take a sequence, train on it, but then also like return the sequence to the pool. They are exhausting a pool, so when they draw a sequence, it's, it's gone until the next epoch of training. So we're already doing that because our DataLoader iterates over chunks of data, so there's no replacement. They don't become eligible to be drawn again until the next epoch. So we're basically already doing that. Um, all models use a weight decay of 0.1 to provide a small amount of regularization. So let's implement the weight decay. And you see here that I've already kind of made the changes. And in particular, instead of creating the optimizer right here, I'm creating a new configure_optimizers function inside the model, and I'm passing in some of the hyperparameters instead. So let's look at the configure_optimizers, which is supposed to return the optimizer object. Okay, so it looks complicated, but it's actually really simple, and it's just— we're just being very careful, and there's a few settings here to go through. The most important thing with respect to this line is that you see there's a weight decay parameter here, and I'm passing that into, well, I'm passing that into something called optim_groups that eventually ends up going into the AdamW optimizer. And the weight decay that's by default used in AdamW here is 0.01, so it's, it's 10 times lower than what's used in the GPT-3 paper here. So the weight decay basically ends up making its way into the AdamW3D optimizer groups. Now, what else is going on here in this function? So the two things that are happening here that are important is that I'm splitting up the parameters into those that should be weight decayed and those that should not be weight decayed. So in particular, it is common to not weight decay biases and any other sort of one-dimensional tensors. So the one-dimensional tensors are in the node DecayParams, and these are also things like layer norm, scales, and biases. It doesn't really make sense to weight decay those. You mostly want to weight decay the weights that participate in matrix multiplications, and you want to potentially weight decay the embeddings. And we've covered in previous video why it makes sense to decay the weights, because you can sort of think of it as a regularization. Because when you're pulling down all the weights, you're forcing the optimization to use more of the weights, and you're not allowing any one of the weights individually to be way too large. You're forcing, you're forcing the network to kind of like distribute the work across more channels because there's sort of like a pull of gravity on the weights themselves. So that's why we are separating it in those ways here. We're only decaying the embeddings and the matmul participating weights. We're printing the number of parameters that were decaying and not. Most of the parameters will be decayed. And then one more thing that we're doing here is I'm doing another optimization here. And previous AdamW did not have this option, but later parts of PyTorch introduced it. And that's why I'm guarding it with an inspect.signature, which is basically checking if this fused, um, kwarg is present inside atom-w. And then if it is present, I'm gonna end up using it and passing it in here because some earlier versions do not have fused equals. So here's atom-w fused equals. It did not used to exist and it was added later. And there's some docs here for what's happening. And basically they say that by default they do not use fused because it is relatively new and we want to give it sufficient bake time. So by default, I don't use fused, but fused is a lot faster when it is available and when you're running on CUDA. And what that does is instead of iterating in a for loop over all the parameter tensors and updating them, that would launch a lot of kernels, right? And so fused just means that it's a— all those kernels are fused into a single kernel. You get rid of a lot of overhead. And you a single time on all the parameters call a kernel that updates them. And so it's just basically a kernel fusion for the AdamW update instead of iterating over all the tensors. So that's the configure_optimizers function that I like to use. And we can rerun and we're not going to see any major differences from what we saw before, but we are going to see some prints coming from here. So let's just take a look at what they look like. So we see that number of decayed tensors is 50, and it's most of the parameters, and number of non-decayed tensors is 98, and these are the biases and the layer norm parameters mostly, and that's— there's only 100,000 of those. So most of it is decayed, and then we are using the fused implementation of AdamW, which will be a lot faster. So if you have it available, I would advise you to use it. I'm not actually 100% sure why they don't default to it. It seems fairly benign and harmless. And also because we are using the fused implementation, I think this is why we have dropped— notice that the running time used to be 93 milliseconds per step, and we're now down to 90 milliseconds per step because of using the fused Atom-W optimizer. So in a single commit here, we are introducing fused Atom, getting improvements on the time, and we're adding or changing the weight decay, but we're only weight decaying the two-dimensional parameters, the embeddings and the matrices that participate in linear. So that is this, and we can take this out. And, uh, yeah, that is it for this line. One more quick note before we continue here. I just want to point out that the relationship between weight decay, learning rate, batch size, The Adam parameters, beta 1, beta 2, the epsilon, and so on. These are very complicated mathematical relationships in the optimization literature. And for the most part, I'm in this video, I'm just trying to copy paste the settings that OpenAI used. But this is a complicated topic, quite deep. And yeah, in this video, I just want to copy the parameters because it's a whole different video to really talk about that in detail and give it proper justice. Instead of just high-level intuitions. Now the next thing that I want to move on to is that this paragraph here, by the way, we're going to turn back around to when we improve our DataLoader. For now, I want to swing back around to this table where you will notice that for different models, we of course have different hyperparameters for the transformer that dictate the size of the transformer network. We also have a different learning rate. So we're seeing the pattern that the bigger networks are trained with slightly lower learning rates. And we also see this batch size where in the, in the small networks they use a smaller batch size and in the bigger networks they use a bigger batch size. Now the problem with for us is we can't just use 0.5 million batch size because if I just try to come in here and I try to set this B, where's my B? B equals— where do I call the— okay, b equals 16. If I try to set— well, we have to be careful. It's not 0.5 million because this is the batch size in the number of tokens. Every single one of our rows is 1024 tokens. So 0.5e6, 1 million divided by 1024. This would need about a 488 batch size. So the problem is I can't come in here and set this to 488 because my GPU would explode. This would not fit for sure. And so, but we still want to use this batch size because again, as I mentioned, the batch size is correlated with all the other optimization hyperparameters and the learning rates and so on. So we want to have a faithful representation of all the hyperparameters and therefore we need to use a batch size of 0.5 million roughly. But the question is, how do we use 0.5 million if we only have a small GPU? Well, for that we need to use what's called gradient accumulation. So we're going to turn to that next, and it allows us to simulate in a serial way any arbitrary batch size that we set. And so we can do a batch size of 0.5 million, we just have to run longer and we have to process multiple sequences and basically add up all the gradients from them to simulate a batch size of 0.5 million. So let's turn to that next. Okay, so I started the implementation right here just by adding these lines of code. And basically what I did is first I set the total batch size that we desire. So this is exactly 0.5 million, and I used a nice number, a power of 2, because 2^19 is 524,288. So it's roughly 0.5 million. It's a nice number. Now our micro batch size, as we call it now, is 16. So this is going to be, we still have B by T indices that go into the transformer and do forward backward, but we're not going to do an update, right? We're going to do many forward backwards. We're going to, and those gradients are all going to plus equals on the parameter gradients. They're all going to add up. So we're going to do forward backward grad accum steps number of times. And then we're going to do a single update once all that is accumulated. So in particular, our micro batch size is just now controlling how many tokens, how many rows we're processing in a single go of a forward backward. So here we are doing 16 times 124. We're doing 16,384 tokens per forward backward, and we are supposed to be doing 2 to the 19— whoops, what am I doing— 2 to the 19 in total, so the grada cum will be 32. So therefore grada cum here will work out to 32, and we have to do 32 forward-backward and then a single update. Now we see that we have about 100 milliseconds for a single forward-backward, so doing 32 of them will be will make every step roughly 3 seconds, just napkin math. So that's Gradacom steps, but now we actually have to implement that. So we're going to swing over to our training loop because now this part here and this part here, the forward and the backward, we have to now repeat this 32 times before we do everything else that follows. So let's see how we can implement that. So let's come over here. And actually, we do have to load a new batch every single time. So let me move that over here. And now this is where we have the inner loop. So for micro_step in range, uh, gradaccum_steps, we do this. And remember that loss.backward always deposits gradients. So we're doing inside loss.backward, there's always a on the gradients. So in every single loss.backward, gradients will add up on the gradient tensors. So we loss.backward, and then we get all the gradients over there. And then we normalize and everything else should just follow. So we're very close. But actually, there's like subtle and deep issue here. And this is actually incorrect. So invite— I invite you to think about why this is not yet sufficient. And let me fix it then. Okay, so I brought back the Jupyter notebook so we can think about this carefully in a simple toy setting and see what's happening. So let's create a very simple neural net that takes a 16— vector of 16 numbers and returns a single number. And then here I'm creating some random examples x and some targets y, and then we are using the mean squared loss here to calculate the loss. So basically what this is, is 4 individual examples, and we're just doing simple regression with the mean squared loss over those 4 examples. Now when we calculate the loss and we loss.backpropagate and look at the gradient, this is the gradient that we achieve. Now the loss objective here, notice that in MSE loss, the default for the loss function is reduction is mean. So we're calculating the average mean loss, the mean loss here over the 4 examples. So this is the exact loss objective and this is the average, the 1 over 4, because there are 4 independent examples here. And then we have the 4 examples and their mean squared error, the squared error, and then this makes it the mean squared error. So therefore we are, we calculate the average calculate the squared error, and then we normalize it to make it the mean over the examples. And there's 4 examples here. So now when we come to the gradient accumulation version of it, this, this here is the gradient accumulation version of it where we have grad_accum_steps of 4, and I reset the gradient. We've grad_accum_steps of 4, and now I'm evaluating all the examples individually instead and calling loss.backward on them many times. And then we're looking at the gradient that we achieve from that. So basically now we forward our function, calculate the exact same loss, do it backward, and we do that 4 times. And when we look at the gradient, uh, you'll notice that the gradients don't match. So here we, uh, did a single batch of 4, and here we did, uh, 4 gradient accumulation steps of batch size 1, and the gradients are not the same. And basically the, the reason that they're not the same is exactly because this mean squared error gets lost. This 1/4 in this loss gets lost because what happens here is the loss objective for every one of the loops is just the mean squared error, which in this case, because there's only a single example, is just this term here. So that was the loss in the 0th iteration, the same in the 1st, 3rd, and so on. And then when you do the loss step backward, we're accumulating gradients. And what happens is that accumulation in the gradient is basically equivalent to doing a sum in the loss. So our loss actually here is this without the factor of 1/4 outside of it. So we're missing the normalizer, and therefore our gradients are off. And so the way to fix this, or one of them, is basically we can actually come here and we can say loss equals loss divide 4. And what happens now is that we're introducing— we're scaling our loss. We're introducing a 1/4 in front of all of these places. So all the individual losses are now scaled by 1/4. And then when we backward, all of these accumulate with a sum. But now there's a 1/4 inside every one of these components. And now our losses will be equivalent. So when I run this, you see that the gradients are now identical. So long story short, with this simple example, when you step through it, you can see that basically the reason that this is not correct is because in the same way as here in the MSE loss, the loss that we're calculating here in the model is using a reduction of mean as well. So where's the loss? f.cross_entropy. And by default, the reduction here in cross entropy is also— I don't know why they don't show it, but it's the mean, the mean loss at all the B by T elements, right? So there's a reduction by mean in there. And if we're just doing this gradient accumulation here, we're missing that. And so the way to fix this is to simply compensate for the number of gradient accumulation steps. And we can in the same way divide this loss. So in particular here, the number of steps that we're doing is loss equals loss divided gradient accumulation steps. So even Copilot, sorry, gets the modification. But in the same way exactly, we are scaling down the loss so that when we do loss.backward, which basically corresponds to a sum in the objective, we are summing up the already normalized loss. And therefore, when we sum up the losses divided by Gradacum steps, we are recovering the additional normalizer. And so now these two will be— now this will be equivalent to the original sort of optimization because the gradient will come out the same. Okay, so I had to do a few more touch-ups and I launched the optimization here. So in particular, one thing we want to do because we want to print things nicely is, well, first of all, we need to create like an accumulator over the loss. We can't just print the loss because we'd be printing only the final loss at the final micro step. So instead we have loss_accum, which I initialize at 0, and then I accumulate the loss into it. And I'm using detach so that I'm detaching the tensor from the graph, and I'm just trying to keep track of the values. So I'm making these leaf nodes when I add add them. So that's lossaccum, and then we're printing that here instead of loss. And then in addition to that, I had to account for the gradaccum steps inside the tokens processed, because now the tokens processed per step is B times T times gradient accumulation. So long story short, here we have the optimization. It looks reasonable, right? We're starting at a good spot. We calculated the gradaccum steps to be 32. And we're getting about 3 seconds here, right? And so this looks pretty good. Now if you'd like to verify that your optimization and the implementation here is correct and you're working on the side, well now because we have the total batch size and the gradient accumulation steps, our setting of B is purely a performance optimization kind of setting. So if you have a big GPU, you can actually increase this to 32. And you'll probably go a bit faster. If you have a very small GPU, you can try 8 or 4, but in any case, you should be getting the exact same optimization and the same answers up to like a floating point error, because the gradient accumulation kicks in and, and can handle everything serially as necessary. So that's it for gradient accumulation, I think. Okay, so now is the time to bring out the heavy weapons. You've noticed that so far we've only been using a single GPU for training, but actually I am paying for 8 GPUs here, and so we should be putting all of them to work. And in particular, they are going to collaborate and, you know, optimize over tokens at the same time and communicate so that they're all kind of collaborating on the optimization. For this, we are going to be using the DistributedDataParallel from PyTorch. There's also a legacy DataParallel, which I recommend you not use, and that's kind of like, you know, legacy. Distributed Data Parallel works in a very simple way. We have 8 GPUs, so we're going to launch 8 processes, and each process is going to be assigned a GPU. And for each process, the training loop and everything we've worked on so far is going to look pretty much the same. Each GPU, as far as it's concerned, is just working on exactly what we've built so far, but now secretly there's 8 of them. And they're all going to be processing slightly different parts of the data. And we're going to add one more part where once they all calculate their gradients, there's one more part where we do an average of those gradients. And so that's how they're going to be collaborating on the computational workload here. So to use all 8 of them, we're not going to be launching our script anymore with just PyTorch train GPT-2 .py, we're going to be running it with a special command called torchrun in PyTorch. We'll see that in a bit. And torchrun, when it runs our Python script, will actually make sure to run 8, 8 of them in parallel. And it creates these environmental variables where each of these processes can look up which, basically, which one of the processes it is. So for example, torchrun will set rank, local_rank, and world_size environmental variables. And so this is a bad way to detect whether DDP is running. So if we're using torch.run, if DDP is running, then we have to make sure that CUDA is available because I don't know that you can run this on CPU anymore or that that makes sense to do. This is some setup code here. The important part is that there's a world size, which for us will be 8. That's the total number of processes running. There's a rank, which is, um, each process will basically run the exact same code at the exact same time, roughly. But all the process— the only difference between these processes is that they all have a different DTP rank. So the GPU 0 will have DTP rank of 0, GPU 1 will have a rank of 1, etc. So otherwise they're all running the exact same script. It's just that ddp_rank will be a slightly different integer. And that is the way for us to coordinate that they don't, for example, run on the same data. We want to— we want them to run on different parts of the data and so on. Now local_rank is something that is only used in a multi-node setting. We only have a single node with 8 GPUs. And so local_rank is the rank of the GPU on a single node. So from 0 to 7 as an example. But for us, we're mostly going to be running on a single box, so the things we care about are rank and world size. This is 8, and this will be whatever it is depending on the GPU that, that this particular instantiation of the script runs on. Now here we make sure that according to the local rank, we are setting the device to be cuda: and colon indicates which GPU to use if there are more than one GPUs. So depending on the local rank of this process, it's going to use just the appropriate GPU, so there's no collisions on which GPU is being used by which process. And finally, there's a boolean variable that I like to create, which is the ddp_rank 0. So the master process is arbitrarily process number 0, and it does a lot of the printing, logging, checkpointing, etc. And the other processes are thought of mostly as compute processes that are assisting. And so master process 0 will have some additional work to do. All the other processes will, will mostly just be doing forward/backwards. And if we're not using DDP and none of these variables are set, we revert back to single GPU training. So that means that we only have rank 0, the world size is just 1, and we are the master process. And we try to auto-detect the device, and this is world as normal. So, so far all we've done is we've initialized DDP, and in the case where we're running with torchrun, which we'll see in a bit, there's going to be 8 copies running in parallel. Each one of them will have a different rank, and now we have to make sure that everything happens correctly afterwards. So the tricky thing with running multiple processes is you always have to imagine that there's going to be 8 processes running in parallel. So as you read the code now, you have to imagine there's 8, you know, 8 Python interpreters running down these lines of code. And the only difference between them is that they have a different DDP rank. So they all come here, they all pick the exact same seed, they all make all of these calculations completely unaware of the other copies running, roughly speaking, right? So they all make the exact same calculations. And now we have to adjust these calculations to take into account that there's actually like a certain world size and certain ranks. So in particular, these microbatches and sequence lengths, these are all just per GPU, right? So now there's going to be num processes of them running in parallel. So we have to adjust this, right? Because the Gradacum steps now, is going to be total batch size divided by B times T times DDP role size, because each process will do B times T and there's this many of them. And so in addition to that, we, we want to make sure that this fits nicely into total batch size, which for us it will, because 16 times 1.24 times 8, 8 GPUs is 131. K, and so 524288, this means that our gradaccum will be 4 with the current settings, right? So there's going to be 16 times 124 processed in each GPU, and then there's 8 GPUs, so we're going to be doing 131,000 tokens in a single forward-backward on the 8 GPUs. So we want to make sure that this fits nicely so that we can derive a nice gradient accumulation steps. And, uh, yeah, let's just adjust the comments here. Times DDP world size. Okay, so each GPU calculates this. Now this is where we start to get run into issues, right? So we are— each process is going to come by a print and they're all going to print. So we're going to have 8 copies of these prints. So one way to deal with this is exactly this master process variable that we have. So in if master_process, then guard this. And that's just so that we just print this a single time, because otherwise all the processes would have computed the exact same variables, and there's no need to print this 8 times. Before getting into the DataLoader, we're going to have to refactor it, obviously. Maybe at this point is, uh, we should do some prints and just take it out for a spin and exit at this point. So import sys, sys.exit() and print "I am GPU DDP rank" and print "bye." So, so now let's try to run this and just see how this works. So let's take it for a spin just so we see what it looks like. So normally we used to launch python train_gpt2.py like this. Now we're going to run with torchrun and this is what it looks like. So torchrun standalone, number of processes, for example, is 8 for us because we have 8 GPUs, and then train_gpt2.py. So this is what the command would look like. And torchrun again will run 8 of these. So let's just see what happens. So first, gets a little busy. So there's a lot going on here. So first of all, there's some warnings from distributed, and I don't actually know that these mean anything. I think this is just like the code is setting up and the processes are coming online, and we're seeing some preliminary failure to collect while the processes come up. I'm not 100% sure about that, but we start to then get into actual prints. So all all the processes went down, and then the first print actually comes from process 5, just by chance. And then it printed— so process 5 basically got here first. It said, "I'm process— I'm GPU 5, bye." And then these prints come from the master process. So process 5 just finished first for whatever reason. It just depends on how the operating system scheduled the processes to run. Then GPU 0 ended, then GPU 3 and 2, and then probably process 5 or something like that has exited. And DDP really doesn't like that because we didn't properly dispose of the multi-GPU setting. And so process group has not been destroyed before we destruct. So it really doesn't like that. And in an actual application, we would want to called destroy_process_group so that we clean up DDP properly. And so it doesn't like that too much. And then the last of the GPUs finish and that's it. So basically we can't guarantee when these processes are running, it's totally arbitrary, but they are running in parallel. We don't want them to be printing. And next up, let's erase this. Next up, we want to make sure that when we create DataLoaderLite, we need to now make it aware of this multiprocessing setting because we don't want all the processes to be loading the exact same data. We want every process to get its own chunk of data so that they're all working on different parts of the dataset, of course. So let's adjust that. So one particularly simple and a naive way to do this is we have to make sure that we pass in the rank and the size to the DataLoader. And then we come up here, we see that we now take rank and processes and we save them. Now the current position will not be 0, because what we want is we want to stride out all the processes. So one way to do this is we basically take salt.b times salt.t and then multiply it by the process rank. So process rank 0 will start at 0, but process rank 1 now starts at b times t. Process rank 2 starts at 2 times b times t, etc. So that is the initialization. Now we still— they still do this identically, but now when we advance, we don't advance by b times t, we advance by b times t times number of processes, right? So basically the total number of tokens that we're consuming is b times t times num processes. And they all go off to a different rank, and the position has to advance by the entire chunk. And then here, b times t times self.num processes plus 1 would be to exceed number of tokens, and then we're going to loop. And when we loop, we want to of course loop in the exact same way, so we sort of like reset back. So this is the simplest change that I can find for kind of a very simple distributed DataLoader lite. And you can notice that if process_rank is 0 and num processes is 1, then the whole thing will be identical to what we had before. But now we can have actually multiple processes running and this should work fine. So that's the DataLoader. Okay, so next up, once they've all initialized the DataLoader, they come here and they all create a GPT model. So we create 8 GPT models on 8 processes, but because the seeds are fixed here, they all create the same identical model. They all move it to the device of their rank, and they all compile the model. And because the models are identical, there are 8 identical compilations happening in parallel, but that's okay. Now, none of this changes because that is on a per-step basis, and we're currently working kind of within step because we need to just all the changes we're making are kind of like within-step changes. Now the important thing here is when we construct the model, we actually have a bit of work to do here. get_logits is deprecated, so create_model. We need to actually wrap the model into the distributed data parallel container. So this is how we wrap the model into the DDP container. And these are the docs for DDP, and they're quite extensive, and there's a lot of caveats and a lot of things to be careful with because everything complexifies times 10 when multiple processes are involved. But roughly speaking, this device IDs, I believe, has to be passed in. Now unfortunately, the docs for what device IDs is, is extremely unclear. So when you actually like come here, this comment for what device IDs is, is roughly nonsensical, but I'm pretty sure it's supposed to be the DDP local rank. So not the DDP rank, the local rank. So this is what you pass in here. This wraps the model. And in particular, what DDP does for you is in a forward pass, it actually behaves identically. So my understanding of it is nothing should be changed in the forward pass. But in the backward pass, as you are doing the backward pass, In the simplest setting, once the backward pass is over on each independent GPU, each independent GPU has the gradient for all the parameters. And what DDP does for you is once the backward pass is over, it will call what's called all-reduce, and it basically does an average across all the ranks of their gradients, and then it will deposit that average on every single rank. So every single rank will end up with the average on it. And so basically that's the communication. It just synchronizes and averages the gradients. And that's what DDP offers you. Now DDP actually is a little bit more, is a little bit more involved than that because as you are doing the backward pass through the layers of the transformer, it actually can dispatch communications for the gradient while the backward pass is still happening. So there's overlap of the communication of the gradients and and the synchronization of them and the backward pass. And this is just more efficient and to do it that way. So that's what DDP does for you. Forward is unchanged and backward is mostly unchanged. And we're tacking on this average as we'll see in a bit. Okay, so now let's go to the optimization. Nothing here changes. Let's go to the optimization here, the inner loop and think through the synchronization of these gradients in the DEP. So basically by default what happens, as I mentioned, is when you do loss.backward here, it will do the backward pass and then it will synchronize the gradients. The problem here is because of the gradient accumulation steps loop here, we don't actually want to do the synchronization after every single loss.backward because we are just depositing gradients and we're doing that serially and we just want them adding up and we don't want to synchronize every single time. That would be extremely wasteful. So basically we want to add them up. And then on the very last, it's only on the very last step when microstep becomes gradaccum_steps 1, only at that last step do we want to actually do the allreduce to average up the gradients. So to do that, we come here and the official sanctioned way, by the way, is to do, this nosync context manager. So PyTorch says this is a context manager to disable gradient synchronization across DDP processes. So within this context, gradients will be accumulated, and basically when you do nosync, there will be no communication. So they are telling us to do with DDP nosync, do the gradient accumulation, accumulate grads, and then they are asking us to do DDP again with another input and dot backward. And I just really don't love this. I just really don't like it. The fact that you have to copy paste your code here and use a context manager, and this is just super ugly. So when I went to the source code here, you can see that when you enter, you simply toggle this variable, this require_backwards_gradsync, and this is being toggled around and changed. And this is the variable that basically if you step through it, is being toggled to determine if the gradient is going to be synchronized. So I actually just kind of like to use that directly. So instead, what I like to do is the following. Right here before the loss.backward, if we are using DDP, then, then basically we only want to synchronize— we only want this variable to be true when it is the final iteration. In all the other iterations inside the microsteps, we want it to be false. So I just toggle it like this. So require_backward_graph_sync should only turn on when the microstep is the last step. And so I'm toggling this variable directly, and I hope that that impacts last_step_backward. And this is a naughty thing to do because, you know, they could probably change the DDP and this variable will go away. But for now, I believe this, this works, and it allows me to avoid the use of context managers and code duplication. I'm just toggling the variable, and then loss.backward will not synchronize most of the steps, and it will synchronize the very last step. And so once this is over and we come out, every single rank will suddenly magically have the average of all the gradients that were stored on all the ranks. So now we have to think through whether that is what we want, and also if this suffices, and whether— how it works with the loss, and what is loss_accum. So let's think through that now. And the problem I'm getting at is that we've averaged the gradients, which is great, but the loss_accum has not been impacted yet. And this is outside of the DDP container, so that is not being averaged. And so here when we are printing lossaccum, well, presumably we're only going to be printing on the master process, rank 0, and it's just going to be printing the losses that it saw on its process. But instead we want it to print the loss over all the processes and the average of that loss because we did average of gradients. So we want the average of loss as well. So simply here after this, this is the code that I've used in the past. And instead of lossf we want lossaccum. So if TDP again, then dist is a PyTorch distributed. I import it. Where do I import it? Oh gosh. So this file is starting to get out of control, huh? So if, so import torch.distributed as dist. So dist.allreduce, and we're doing the average on lossaccum. And so this lossaccum tensor exists on all the ranks. When we call allreduce of average, it creates the average of those numbers, and it deposits that average on all the ranks. So all the ranks after this call will now contain lossaccum averaged up. And so when we print here on the master process, the lossaccum is identical in all the other ranks as well. So here, if master_process— oops— we want to print like this. Okay, and finally we have to be careful because we're now processing even more tokens, so times ddp_world_size. That's the number of tokens that we've processed up above, and everything else should be fine. The only other thing to be careful with is, as I mentioned, you want to destroy the process group so that we are nice to Nickel and it's not going to— to DDP and it's not going to complain to us when we exit here. So that should be it. Let's try to take it for a spin. Okay, so I launched the script and it should be printing here imminently. We're now training with 8 GPUs at the same time, so the gradient accumulation steps is not 32, it is now divide 8 and it's just 4. So Otherwise, this is what the optimization now looks like. And wow, we're going really fast. So we're processing 1.5 million tokens per second now. So these are some serious numbers. And the Tiny Shakespeare dataset is so tiny that we're just doing like so many epochs over it, most likely. But this is roughly what it looks like. One thing that I had to fix, by the way, is that this was model.configure_optimizers. Which now doesn't work because model now is a DDP model. So instead this has to become raw_model.configure_optimizers where raw_model is something I create here. So right after I wrap the model into DDP, I have to create the raw model, which in the case of DDP is a model.module is where it stores the raw NN module of GPT-2 as we have it, which contains the configure_optimizers function that we want to call. So that's one thing that I have to fix. Otherwise, this seems to run. Now one thing you'll notice is that when you actually compare this run and the numbers in it to the just running a single GPU, you'll notice that this is a single GPU run with 32 Gradacum. The numbers won't exactly match up, and it's kind of a boring reason for why that happens. The reason for that is that in the DataLoader we're basically just iterating through batches in a slightly different way because now we're looking for an entire page of data. And if that page for all the GPUs, if that chunk exceeds the number of tokens, we just loop. And so actually the single GPU and the GPU process will end up resetting in a slightly different manner. And so our batches are slightly different, and so we get slightly different numbers. But one way to convince yourself that this is Okay, it just makes the total batch size much smaller and the B and the T. And then, so I think I used 4 times 124 times 8. So I used 32,768 as a total batch size. And then, so I made sure that the single GPU will do 8 gradient accumulation steps and then at the multi-GPU, and then you're reducing the boundary effects of the DataLoader and you'll see that the numbers match up. So long story short, We're now going really, really fast. The optimization is mostly consistent with GPT-2 and 3 hyperparameters, and we have outgrown our tiny Shakespeare file and we want to upgrade it. So let's move to next— to that next. So let's now take a look at what datasets were used by GPT-2 and GPT-3. So GPT-2 used this WebText dataset that was never released. There's an attempt at reproducing it called OpenWebText. So basically, roughly speaking, what they say here in the paper is that they scraped all outbound links from Reddit and then with at least 3 karma, and that was kind of like their starting point, and they collected all the web pages and all the text in them. And so this was 45 million links and this ended up being 40 gigabytes of text. So, so that's roughly what GPT-2 says about its dataset. So it's basically outbound links from Reddit. Now when we go over to GPT-3, there's a training dataset section and that's where they start to talk about Common Crawl, which is a lot more used. Actually, I think even GPT-2 talked about Common Crawl, but basically it's not a very high-quality dataset all by itself because it is extremely noisy. This is a completely random subset of the internet and it's much worse than you think. So people go into great lengths to filter Common Crawl. Because there's good stuff in it, but most of it is just like ad spam, random tables and numbers and stock tickers, and it's just a total mess. So that's why people like to train on these data mixtures that they curate and are careful with. So a large chunk of these data mixtures typically will be Common Crawl. Like for example, 50% of the tokens will be Common Crawl. But then here in GPT-3, they're also using WebText-2 from before. That's Reddit outbound, but they're also adding, for example, books, and they're adding Wikipedia. There's many other things you can decide to add. Now, this dataset for GPT-3 was also never released. So today, some of the datasets that I'm familiar with that are quite good and would be representative of something along these lines are number one, the Red Pajama dataset, or more specifically, for example, the Slim Pajama subset of the Red Pajama dataset, which is a cleaned and deduped deduplicated version of it. And just to give you a sense, again, it's a bunch of Common Crawl, C4, which is also, as far as I know, more Common Crawl but processed differently. And then we have GitHub, Books Archive, Wikipedia, Stack Exchange. These are the kinds of datasets that would go into these data mixtures. Now, specifically, the one that I like that came out recently is called FindWeb Dataset. So this is an attempt to basically collect really high-quality Common Crawl data and filter it, in this case, to 15 trillion tokens. And then in addition to that, more recently, Hugging Face released this Fineweb Edu subset, which is 1.3 trillion of educational and 5.4 trillion of high-educational content. So basically, they're trying to filter Common Crawl to very high-quality educational subsets, and this is the one that we will use. There's a long webpage here on Fineweb, and they go into a ton of detail about how they process the data, which is really fascinating reading, by the way. And I would definitely recommend, if you're interested in data mixtures and so on and how data gets processed at these scales, look at this page. And more specifically, we'll be working with the Fineweb Edu, I think, and it's basically educational content from the internet. They show that training on educational content in, in in their metrics works really, really well. And we're going to use this sample 10 billion tokens subsample of it because we're not going to be training on trillions of tokens. We're just going to train on a 10 billion sample of the Fineweb Edu because empirically in my previous few experiments, this actually suffices to really get close to GPT-2 performance. And it's, simple enough to work with. And so let's work with the sample10bt. So our goal will be to download it, process it, and make sure that our DataLoader can work with it. So let's get to that. Okay, so I introduced another, um, file here that will basically download FineWebEDU from Hugging Face datasets. It will preprocess and pre-tokenize all of the data, and it will save data shards to a folder on local disk. And so while this is running, I just wanted to briefly mention that you can kind of look through the dataset viewer here just to get a sense of what's in here. And it's kind of interesting. I mean, it's a— it basically looks like it's working fairly well. Like it's talking about nuclear energy in France. It's talking about Mexican America. Some MacPyJs, etc. So actually, it seems like their filters are working pretty well. The filters here, by the way, were applied automatically using Llama 3.70B, I believe. And so basically, LLMs are judging which content is educational, and that ends up making it through the filter. So that's pretty cool. Now, in terms of the script itself, I'm not going to go through the full script because it's not as interesting and not as LLM-centric. But when you run this, basically, number one, we're going to load the dataset, which this is all Hugging Face code running this. You're going to need to pip install datasets. So it's downloading the dataset, then it is tokenizing all of the documents inside this dataset. Now when we tokenize the documents, you'll notice that to tokenize a single document, we first start the tokens with the end-of-text token, and this is a special token in the GPT-2 tokenizer, as you know. So 50256 is the ID of the end-of-text, and this is what begins a document, even though it's called end-of-text. But this is the first token that begins a document. Then we extend with all of the tokens of that document. Then we create a NumPy array out of that. We make sure that all the tokens are between Oh, okay, let me debug this. Okay, so apologies for that, it just had to do with me using a float division in Python. It must be integer division so that this is an int and everything is nice. Okay, but basically the tokenization here is relatively straightforward, returns tokens in np.un16. We're using un16 to save a little bit of space because 2 to the 16 minus 1 is 65,000, so the GPT-2 max token ID is well below that. And then here there's a bunch of multiprocessing code, and it's honestly not that exciting, so I'm not going to step through it. But we're loading the dataset, we're tokenizing it, and we're saving everything to shards. And the shards are NumPy files, uh, so just storing a NumPy array, and, uh, which is very, very similar to torch tensors. And the first shard is 000 is a validation shard, and all the other shards are training shards. And as I mentioned, they all have 100 million tokens in them exactly. And that just makes it easier to work with, to shard the files, because if we just have a single massive file, sometimes they can be hard to work with on the disk. And so sharding it is just kind of nicer from that perspective. And yeah, so we'll just let this run. This will be probably 30-ish minutes or so, and then we're going to come back to actually train on this data, and we're going to be actually doing some legit pre-training in this case. This is a good dataset. We're doing lots of tokens per second. We have 8 GPUs. The code is ready, and so we're actually going to be doing a serious training run. So let's get back in a bit. Okay, so we're back. So if we ls edu fine-web, we see that there's now 100 shards. In it, and that makes sense because each shard is 100 million tokens, so 100 shards of that is 10 billion tokens in total. Now swinging over to the main file, I made some adjustments to our DataLoader again, and that's because we're not running with Shakespeare anymore. We want to use the Fineweb shards, and so you'll see some code here that additionally basically can load these shards. We load the UN16-9 NumPy file, we convert it to a torch.LongTensor, which is what a lot of the layers up top expect by default. And then here we're just enumerating all the shards. I also added a split to DataLoaderLite so we can load the split train but also the split val, the zero split. And then we can load the shards. And then here we also have not just the current position now but also the current shard. So we have a position inside a shard. And then when we run out of tokens in a single shard, we first advance the shard and loop if we need to, and then we get the tokens and readjust the position. So this DataLoader will now iterate all the shards as well. So I changed that, and then the other thing that I did while the data was processing is our train loader now has split_train, of course, and down here I set up some, some I set up some numbers. So we are doing 2 to the 19 tokens per, per, per step. And we want to do roughly 10 billion tokens because that's how many unique tokens we have. So if we did 10 billion tokens, then divide that by 2 to the 19, we see that this is 19,073 steps. So that's where that's from. And then the GPT-3 paper says that they warm up the learning rate over 375 million tokens. So I came here and 375,000,000 e6 tokens divide, uh, 2 to the 19 is 715 steps. So that's why warmup_steps is set to 715. So this will exactly match, um, the warmup schedule that GPT-3 used. And I think 715, by the way, is very, uh, mild. And this could be made significantly more aggressive, probably even like 100 is good enough, but it's okay, let's leave it for now so that we have the exact hyperparameters of GPT-3. So I fixed that, and then that's pretty much it. We can, we can run. So we have our script here, and we can launch. And actually, sorry, let me do one more thing. Excuse me. For my GPU, I can actually fit more batch size, and I believe I can fit 64 on my GPU as a micro batch size. So let me try that. I could be misremembering, but that means 64 times 124 per GPU, and then we have 8 GPUs. So that means we would not even be doing gradient accumulation if this fits, because this just multiplies out to the full total batch size. So no gradient accumulation, and that would run pretty quickly if that fits. Let's go, let's go. I mean, if this works, then this is basically a serious pre-training training run. We're not logging, we're not evaluating the validation split, we're not running any evaluations yet. So it's not— we haven't crossed our T's and dotted our I's. But if we let this run for a while, we're going to actually get a pretty good model, and the model that might even be on par with or better than GPT-2 1.24M. Okay, so it looks like everything is going great. We're processing 1.5 million tokens per second. Uh, everything here looks good. We're doing 330 milliseconds per iteration, and we have to do a total of, uh, where are we printing that? 1973. So 1973 times 0.33 is this many seconds, this many minutes. So this will run for 1.7 hours. So 1.5-hour run like this, and we don't even have to use gradient accumulation, which is nice. And you might not have that luxury in your GPU. In that case, just start decreasing the batch size until things fit, but keep it to nice numbers. So that's pretty exciting. We're currently warming up the learning rate, so you see that it's still very low, 1e����. So this will ramp up over the next few steps all the way to 6e������. -4 here. Very cool. So now what I'd like to do is let's cross the T's and dot our I's. Let's evaluate on the validation split and let's try to figure out how we can run evals, how we can do logging, how we can visualize our losses and all the good stuff. So let's get to that before we actually do the run. Okay, so I've adjusted the code so that we're evaluating on the validation split. So creating the val_loader just by passing in split split equals val, that will basically create a DataLoader just for the validation shard. The other thing I did is in the DataLoader, I introduced a new function reset, which is called at init, and it basically resets the DataLoader. And that is very useful because when we come to the main training loop now, so this is the code I have added, and basically every 100th iteration, including the 0th iteration, we put the model into evaluation mode, we reset the val_loader, and then, um, no gradients involved, we're going to basically accumulate the gradients over, say, 20 steps and then average it all up and print out the validation loss. And so that basically is the exact same logic as the training loop, roughly, but there's no loss.backward. It's only inference. We're just measuring the loss, we're adding it up, Everything else otherwise applies and is exactly as we've seen it before. And so this will print the validation loss every 100th iteration, including on the very first iteration. So that's nice. That will tell us some amount, some, a little bit about how much we're overfitting. That said, like, we have roughly infinity data, so we're mostly expecting our train and val loss to be about the same. But the other reason I'm kind of interested in this is because we're take the GPT-2 124M as OpenAI released it. We can initialize from it and we can basically see what kind of loss it achieves on the validation loss as well. And that gives us kind of an indication as to how much that model would generalize to 124M, but it's not— sorry, to Fineweb Edu validation split. That said, it's not a super fair comparison to GPT-2 because it was trained on a very different data distribution, but it's still kind of like an interesting data point. And in any case, you would always want to have a validation split in a training run like this so that you can make sure that you are not overfitting. And this is especially a concern if we were to make more epochs in our training data. So for example, right now we're just doing a single epoch, but if we get to a point where we want to train on 10 epochs or something like that, we would be really careful with— maybe we are memorizing that data too much. If we have a big enough model, and our validation split would be one way to tell whether that is happening. Okay, and in addition to that, if you remember at the bottom of our script, we had all of this orphaned code for sampling from way back when. So I deleted that code and I moved it up to here. So once in a while we sample a valid— validation. Once in a while we sample— we generate samples, and then we do that only every 100 steps, and we train on every single step. So that's how I have a structure right now, and I've been running this for 1,000 iterations. So here are some samples on iteration 1,000. "Hello, I'm a language model and I'm not able to get more creative." "I'm a language model and the languages file you're learning about here is— or is the beginning of a computer." Okay, so this is all like pretty— there's still a garble, but we're only at iteration 1,000. And we've only just barely reached the maximum learning rate. So this is still learning. We're about to get some more samples coming up in 1,100. Okay. Okay. This is, you know, the model is still, still a young baby. Okay. So basically all of this sampling code that I've put here, everything should be familiar with to you and came from before. The only thing that I did is I created a generator object in PyTorch so that I have a direct control over the sampling of the random numbers. We don't— because I don't want to impact the RNG state of the random number generator that is the global one used for training. I want this to be completely outside of the training loop. And so I'm using a special sampling RNG, and then I make sure to seed it, that every single rank has a different seed, and then I pass in here where we sort of consume random numbers in multinomial where the sampling happens, I make sure to pass in the generator object there. Otherwise, this is identical. Now, the other thing is you'll notice that we're running a bit slower. That's because I actually had to disable torch.compile to get this to sample. And so we're running a bit slower. So for some reason, it works with no torch.compile, but when I torch.compile my model, I get a really scary error from PyTorch and I have no idea how to resolve it right now. So probably by the time you see this code released or something like that, maybe it's fixed. But for now, I'm just going to do `and False`, and I'm going to bring back `torch.compile`, and you're not going to get samples. And I think I'll fix this later. By the way, I will be releasing all this code, and actually I've been very careful about making git commits every time we add something. And so I'm going to release the entire repo that starts completely from scratch all the way to, uh, now and after this as well. And so everything should be exactly documented in the Git commit history, um, and so I think that will be nice. So hopefully by the time you go to GitHub, uh, this is removed and it's working and I will have fixed the bug. Okay, so I have the optimization running here and it's stepping and we're on step 6,000 or so, so we're about 30% through training. Now while this is training, I would like to introduce, uh, one evaluation that we're going to use to to supplement the validation set, and that is the HELLOswag eval. So HELLOswag comes from this paper back in 2019, so it's a 5-year-old eval now. And the way HELLOswag works is that it's basically a sentence completion dataset. So it's a multiple choice. For every one of these questions, we have basically a shared context, like a woman is outside with a bucket and a dog. The dog is running around trying to avoid bath. She: A) rinses the bucket off with soap and blow-dries the dog's head, B) uses a hose to keep it from getting soapy, C) gets the dog wet and it runs away again, or D) gets into a bathtub with the dog. And so basically the idea is that these multiple choice are constructed so that one of them is a natural continuation of the sentence and the others are not. And the others might not make sense, like "uses the hose to keep it from getting soapy." That makes no sense. And so what happens is that models that are not trained very well are not able to tell these apart, but models that have a lot of world knowledge and can tell which— and can tell a lot about the world will be able to create these completions. And these sentences are sourced from ActivityNet and from wikiHow. And at the bottom of the paper, there's kind of like a cool chart of the kinds of domains in wikiHow. So there's a lot of sentences from computers and electronics and homes and garden, and it has kind of a broad coverage of the kinds of things you need to know about the world in order to find the most likely completion and the identity of that of that completion. One more thing that's kind of interesting about TELA-SWAG is the way it was constructed is that the incorrect options are deliberately adversarially sourced. So they're not just random sentences, they're actually sentences generated by language models, and they're generated in such a way that language models basically find them difficult but humans find them easy. And so they mentioned that humans have a 95% accuracy on this set, but at the time the state-of-the-art language models had only 48%, and so at the time this was a good benchmark. Now you can read the details of this paper to, to learn more. The thing to point out though is that this is 5 years ago, and since then what happened to Hello Swag is that it's been totally just solved, and so now the language models here are 96%, so basically the 4%, the last 4% is probably errors in the dataset or the questions are really, really hard. And so basically this dataset is kind of crushed with respect to language models. But back then the best language model was only at about 50%. But this is how far things got. But still the reason people like HellaSWAG and it's not used by the way in GPT-2, but in GPT-3 there is HellaSWAG eval and lots of people use HellaSWAG. And so for GPT-3, we have results here that are cited, so we know what percent accuracies GPT-3 attains at all these different model checkpoints for HellaSWAG eval. And the reason people like it is because HellaSWAG is a smooth eval, and it is an eval that offers quote-unquote early signal. So early signal means that even small language models are going to start at the random chance of 25%, but they're going to slowly improve and you're going to see 25, 26, 27, etc. And you can see slow improvement even when the models are very small and it's very early. So it's smooth, it has early signal, and it's been around for a long time. So that's why people kind of like this eval. Now the way that we're going to evaluate this is as follows. As I mentioned, we have a shared context, and this is kind of like a multiple choice task. But instead of giving the model a multiple choice question and asking it for A, B, C, or D, uh, we can't do that because these models, when they are so small as we are seeing here, the models can't actually do multiple choice. They don't understand the concept of associating a label to one of the options of multiple choice. They don't understand that. So we have to give it to them in a native form, and the native form is a token completion. So here's what we do. We construct a batch of 4 rows and, uh, T tokens, whatever that T happens to be. Then the shared context, that is basically the context for the 4 choices, the tokens of that are shared across all of the rows. And then we have the 4 options. So we kind of like lay them out and and then only one of the options is correct. In this case, label 3, option 3. And so this is the correct option and option 1, 2, and 4 are incorrect. Now these options might be of different lengths. So what we do is we sort of like take the longest length and that's the size of the batch, B by T. And then some of these here are going to be padded dimensions. So they're going to be unused. And so we need the tokens we need the correct label, and we need a mask that tells us which tokens are active, and the mask is then 0 for these, uh, padded areas. So that's how we construct these batches. And then in order to get the language model to predict A, B, C, or D, the way this works is basically we're just going to look at the tokens, their probabilities, and we're going to pick the option that gets the lowest or the highest average probability for the token, so for the tokens, because that is the most likely completion according to the language model. So we're just going to look at the probabilities here and average them up across the options and pick the one with the highest probability, roughly speaking. So this is how we're going to do hella swag. And this is, I believe, also how GPT-3 did it. This is how GPT-3 did it, as far as I know. But you should note that some of the other evals where you might see Ella Swag may not do it this way. They may do it in a multiple choice format where you sort of give the, the context a single time and then the 4 completions. And so the model is able to see all the 4 options before it picks the best possible option. And that's actually an easier task for a model because you get to see the other options when you're picking your choice. But unfortunately models at our size can't do that. Only models at a bigger size are able to do that. And so our models are actually slightly handicapped in this way that they are not going to see the other options. They're only going to see one option at a time and they just have to assign probabilities and the correct option has to win out in this metric. All right, so let's now implement this very briefly and incorporate it into our script. Okay, so what I've done here is I've introduced a new file called hello_swag.py that you can take a look into. And I'm not going to step through all of it because this is not exactly like deep code, deep code. It's kind of like a little bit tedious, honestly, because what's happening is I'm downloading Hello Swag from GitHub and I'm rendering all of its examples. And there are a total of 10,000 examples. I am rendering them into this format. And so here at the end of this render_example function, you can see that I'm returning the tokens, the tokens of this 4 by T array of tokens, the mask which tells us which parts are the options and everything else is 0, and the label that is the correct label. And so that allows us to then iterate the examples and render them. And I have an evaluate function here which can load a GPT-2 from Hugging Face, and it runs the eval here. And basically just calculates, just as I described, it predicts the option that has the lowest or the highest probability. And the way to do that actually is we can basically evaluate the cross-entropy loss So we're basically evaluating the loss of predicting the next token in the sequence, and then we're looking at the row that has the lowest average loss, and that's the option that we pick as the prediction. And then we do some stats and prints and stuff like that. So that is a way to evaluate TALASWAG. Now, if you go up here, I'm showing that for GPT-2 124M, if you run this script, you're going to see that the Hellaswag gets 29.55%. So that's the performance we get here. Now remember that random chance is 25%, so we haven't gone too far. And GPT-2 XL, which is the biggest, the GPT-2, gets all the way up to 49% roughly. So these are pretty low values considering that today's state of the art is more like 95%. So these are definitely older models by now. And then there's one more thing called Eleuther Harness. Which is a very common piece of infrastructure for running evals for language models, and they get slightly different numbers. And I'm not 100% sure what the discrepancy is for these. It could be that they actually do the multiple choice instead of just the completions, and that, that could be the discrepancy, but I'm not 100% sure about that. I'd have to take a look. But for now, our script reports 29.55, and so that is the number that we'd like to beat if we're training a GPT-2 1/24M from scratch and ourselves. So now I'm going to go into actually incorporating this eval into our main training script, and basically because we want to evaluate it in a periodic manner so that we can track how a swag and how it evolves over time and see when, when and if we cross this 29.55 sort of region. So let's now walk through some of the changes to train_gpt2.py. The first thing I did here is I actually made use_compile optional, kind of, and I disabled it by default. And the problem with that is— the problem with compile is that unfortunately it does make our code faster, but it actually breaks the evaluation code and the sampling code. It gives me a very gnarly message and I don't know why. So hopefully by the time you get to the codebase when I put it up on GitHub, we're going to fix that by then. But for now, I'm running without torch.compile, which is why you see this be a bit slower. So we're running without torch.compile. I also created a log directory log where we can place our log.txt, which will record the train loss, validation loss, and the Hello Swag accuracies. So a very simple text file, and we're going to open for writing. So that it sort of starts empty, and then we're going to append to it. I created a simple variable that helps tell us when we have a last step. And then basically periodically inside this loop, every 250th iteration or at the last step, we're going to evaluate the validation loss. And then every 250th iteration, we are going to evaluate how Hellaswag, but only if we are not using compile because compile breaks it. So I'm going to come back to this code for evaluating Hellaswag in a second. And then every 250th iteration as well, we're also going to sample from the model. And so you should recognize this as our ancient code from way back when we started the video, and we're just sampling from the model. And then finally here, these are if we're not After we validate, sample, and evaluate HellaSWAG, we actually do a training step here. And so this is one step of training, and you should be pretty familiar with all of what this does. And at the end here, once we get our training loss, we write it to the file. So the only thing that changed that I really added is this entire section for HellaSWAG eval. And the way this works is I'm trying to get all the GPUs to collaborate on the HellaSWAG. And so we're iterating all the examples, and then each process only picks the examples that assigned to it. So we sort of take i and mod it by the world size, and we have to make it equal to rank, otherwise we continue. And then we render an example, put it on the GPU, we get the logits. Then I create a helper function that helps us basically predict the option with the lowest loss. So this comes here, the prediction. And then if it's correct, we sort of keep count. And then if multiple processes were collaborating on all this, then we need to synchronize their stats. And so the way— one way to do that is to package up our statistics here into tensors, which we can then call this.already_sum and sum. And then here we sort of unwrap them from tensors so that we just have ints. And then here the master process will print and log the HelloSwag accuracy. So that's kind of the— that's kind of it. And that's what I'm running right here. So you see this optimization here and we just had a generation and this is step 10,000 out of about 20,000, right? So we are halfway done and these are the kinds of samples that we are getting at this stage. So let's take a look. Hello, I'm a language model, so I'd like to use it to generate some kinds of output. 'Hello, I'm a language model and I'm a developer for a lot of companies.' 'Hello, I'm a language model,' uh, let's see if I can find any fun one. Um, I don't know, you can go through this yourself, but certainly the predictions are getting less and less random. It seems like the model is a little bit more self-aware and using language that is a bit more specific to it being a language model. Hello, I'm a language model and like how the language is used to communicate. I'm a language model and are going to be speaking English and German. Okay, I don't know. So let's just wait until this optimization finishes and we'll see what kind of samples we get. And we're also going to look at the train, val, and the HELLA score accuracy and see how we're doing with respect to GPT-2. Okay, good morning. So focusing for a moment on the Jupyter notebook here on the right, I created a new cell that basically allows us to visualize the train, val, and HELLA, and the HELLA score. And you can step through this. It basically like parses the log file that we are writing and A lot of this is just like boring matplotlib code, but basically this is what our optimization looks like. So we ran for 19,073 steps, which is roughly 10 billion tokens, which is— whoops, oh my gosh— which is one epoch of the sample 10B of Fineweb B2. On the left we have the loss, and in the blue we have the training loss, in orange we have the validation loss, In red as a horizontal line, we have the OpenAI GPT-2 124M model checkpoint when it's just evaluated on the validation set of, of this Fineweb EDU. So you can see that we are surpassing— this orange is below the red— so we're surpassing the validation set of this dataset. And like I mentioned, the dataset distribution is very different from what GPT-2 trained on, so this is not an exactly fair comparison. But it's a good cross-check to, to look at. Now, we would ideally like something that is withheld and comparable and somewhat standard. And so for us, that is HellaSWAG. And so on here, we see the HellaSWAG progress we made from 25% all the way here. In red, we see the OpenAI GPT-2 1.24M model in red. So it achieves this HellaSWAG here. And the GPT-3 model 124M, which was trained on 300 billion tokens, achieves green. So that's over here. So you see that we basically surpassed the GPT-2 124M model right here, which is really nice. Now, interestingly, we were able to do so with only training on 10 billion tokens, while GPT-2 was trained on 100 billion tokens. Tokens. So for some reason we were able to get away with significantly fewer tokens for training. There are many possibilities as to why we could match or surpass this accuracy with only Temblion training. So number one, it could be that OpenAI GPT-2 was trained on a much wider data distribution. So in particular, Fineweb Edu is all English. It's not multilingual. And there's not that much math and code. And so math and code and multilingual could have been stealing capacity from the original GPT-2 model, and basically that could be partially the reason why this is not working out. There's many other reasons. So for example, the hellaSWAG eval is fairly old, maybe 5 years or so. It is possible that aspects of hellaSWAG in some way, or even identically, have made it into the training set of FinWeb. We don't know for sure, but if that was the case, then we are basically looking at the training curve instead of the validation curve. So long story short, this is not a perfect eval, and there's some caveats here, but at least we have some confidence that, that we're not doing something completely wrong. And, and it's probably the case that when people try to create these datasets, they try to make sure that test sets that are very common are not part of the training set. For example, when Hugging Face created the FinewebEDU, they use HellaSwag as an eval. So I would hope that they make sure that they deduplicate and that there's no HellaSwag in the training set, but we can't be sure. The other thing I wanted to address briefly is look at this loss curve. This looks really weird. This looks really wrong here. I don't actually know 100% what this is, and I suspect it's because the 10 billion sample of FinewebEDU was not properly shuffled. And there's some issue here with the data that I don't fully understand yet. And there's some weird periodicity to it. And because we are in a very lazy way sort of serializing all the tokens and just iterating on them from scratch without doing any permutations or any random sampling ourselves, I think we're inheriting some of the ordering that they have in the dataset. So this is not ideal. But hopefully by the time you get to this repo, some of these things, by the way, will hopefully be fixed. And I will release this build_nanoGPT repo. And right now it looks a little ugly and preliminary. So hopefully by the time you get here, it's nicer. But down here, I'm going to show errata. And I'm going to talk about some of the things that happened after the video. And I expect that we will have fixed this small issue. But for now, basically, this shows that our training is not completely wrong, and it shows that we're able to surpass the accuracy with only 10x the token budget. And possibly it could be also that the dataset may have improved. So the original GPT-2 dataset was web text. It's possible that not a lot of care and attention went into the dataset. This was very early in LLMs. Whereas now there's a lot more scrutiny on good practices around deduplication, filtering, quality filtering, and so on. And it's possible that the dataset we're training on is just of higher quality per token, and that could be giving us a boost as well. So a number of caveats to think about, but for now we're pretty happy with this. And yeah, now the next thing I was interested in is, as you see, it's a morning now, so there was an overnight, and I wanted to basically see how far I could push the result So to do an overnight run, I basically did instead of one epoch, which took roughly 2 hours, I just did it times 4 so that that would take 8 hours while I was sleeping. And so we did 4 epochs or roughly 40 billion tokens of training, and I was trying to see how far we could get. And so this was the only change, and I reran the script. And when I point and read the log file at the 40B, this is what the curves look like. Okay, so to narrate this, number one, we are seeing this issue here with the periodicity through the different epochs and something really weird with the FinWebEDU dataset, and that is to be determined. But otherwise, we are seeing that the HELLA-SWAG actually went up by a lot, and we almost, we almost made it to the GPT-3 124M accuracy up here. But not quite. So it's too bad that I didn't sleep slightly longer. And I think if this was a 5-epoch run, we may have gotten here. Now one thing to point out is that if you're doing multi-epoch runs, we're not actually being very careful in our DataLoader. And we're not— this DataLoader goes through the data in exactly the same format and exactly the same order. And this is kind of suboptimal. And you would want to look into extensions where you actually permute the data randomly, you permute the documents around in every single shard on every single new epoch, and potentially even permute the shards. And that would go a long way into decreasing the periodicity. And it's also better for the optimization, so that you're not seeing things ident— in the identical format, and you're introducing some of the, some randomness. In how the documents follow each other. Because you have to remember that in every single row, these documents follow each other, and then there's the end of text token, and then the next document. So the documents are currently glued together in the exact same identical manner. But we actually want to break, break up the documents and shuffle them around, because the order of the documents shouldn't matter. And they shouldn't— basically, we want to break up that dependence because it's kind of a spurious correlation. And so our data loader is not currently doing that, and that's one improvement you could think of making. The other thing to point out is we're almost matching GPT-3 accuracy with only 40 billion tokens. GPT-3 trained on 300 billion tokens. So again, we're seeing about a 10x improvement here with respect to learning efficiency. The other thing I wanted to— and I don't actually know exactly what to attribute this to— other than some of the things that I already mentioned previously for the previous one. The other thing I wanted to briefly mention is the max_lr here. I saw some people already play with this a little bit in a previous related repository, and it turns out that you can actually almost like 3x this. So it's possible that the maximum learning rate can be a lot higher, and for some reason, the GPT-3 hyperparameters that we are inheriting are actually extremely conservative, and you can actually get away with a high higher learning rate and it would train faster. So a lot of these hyperparameters are quite tunable and feel free to play with them and they're probably not set precisely correctly and it's possible that you can get away with doing this basically. And if you wanted to exactly be faithful to GPT-3 you would also want to make the following difference. You'd want to come here And the sequence length of GPT-3 is 2x, it's 2048 instead of 1024. So you would come here, change this to 2048 for T, and then if you want the exact same number of tokens, half a million per iteration or per step, you want to then decrease this to 32. So they still multiply to half a mil. So that would give your model sequence length equal to that of GPT-3, and in that case basically the the models would be roughly identical as far as I'm, as far as I'm aware, because again, GPT-2 and GPT-3 are very, very similar models. Now we can also look at some of the samples here from the model that was trained overnight. So this is the optimization, and you see that here we stepped all the way to 76,290 or so, and these are the— the highest FLOPS we achieved was 33.24 and these are some of the samples from the model. And you can see that if you read through this and pause the video briefly, you can see that there are a lot more coherent, so, and they're actually addressing the fact that it's a language model almost. So, hello, I'm language model and I try to be as accurate as possible. I'm a language model, not a programming language. I know how to communicate. I use Python. I don't know if you pause this and look at it and then compare it to the one, to the model that was only trained for 10 billion. You will see that these are a lot more coherent and you can play with this yourself. One more thing I added to the code, by the way, is this chunk of code here. So basically right after we evaluate the validation loss, if we are the master process, in addition to logging the validation loss every 5,000 steps, we're also going to save the checkpoint, which is really just the state dictionary of the model. And so checkpointing is nice just because you can save the model and later you can use it in some way. If you wanted to resume the optimization, then in addition to saving the model, we have to also save the optimizer state dict, because remember that the optimizer has a few additional buffers because of Adam. So it's got the m and v, and you need to also resume the optimizer properly. You have to be careful with your RNG seeds, random number generators, and so on. So if you wanted to exactly be able to resume optimization, you have to think through the state of the, of the training process. But if you just want to save the model, this is how you would do it. And one, one nice reason why you might want to do this is because you may want to evaluate the model a lot more carefully. So here we are only kind of like winging the HellaSwag eval, but you may want to use something nicer, like for example the Luther evaluation hardness— evaluation harness— hardness. So this is a way to also evaluate language models, and so it's possible that you may want to use basically different infrastructure to more thoroughly evaluate the models on different evaluations and compare it to the OpenAI GPT-2 model on many other tasks, like for example, that involve math, code, or different languages and so on. So this is a nice functionality to have as well. And then the other thing I wanted to mention is that everything we've built here, this is only the pre-training step. So the GPT here is a, it dreams documents. It just predicts the next token. You can't talk to it like you can talk to ChatGPT. If you wanted to talk to the model, we have to fine-tune it into the chat format. And it's not actually like that complicated. If you're looking at supervised fine-tuning or SFT, really what that means is we're just swapping out a dataset into a dataset that is a lot more conversational, and there's a user-assistant, user-assistant kind of structure, and we just fine-tune on it, and then we We basically fill in the user tokens and we sample the assistant tokens. It's not a lot more deeper than that, but basically we swap out the dataset and continue training. But for now, we're going to stop at pre-training. One more thing that I wanted to briefly show you is that, of course, what we've built up today was building towards NanoGPT, which is this repository from earlier. But also, there's actually another NanoGPT implementation, and it's hiding in a more recent project that I've been working on called llm.c. And llm.c is a pure C CUDA implementation of GPT-2 or GPT-3 training, and it just directly uses CUDA and is written as C CUDA. Now the NanogPT here acts as reference code in PyTorch to the C implementation, so we're trying to exactly match up the two, but we're hoping that the C CUDA is faster. And of course, currently that seems to be the the case because it is a direct optimized implementation. So train_gpt2.py in llm.c is basically the NanoGPT, and when you scroll through this file, you'll find a lot of things that very much look like, um, things that we've built up in this lecture. And then when you look at train_gpt2.cu, this is the CUDA implementation. So there's a lot of MPI, Nickel, GPU, CUDA, C, C++, and you have to be familiar with that. But, uh, um, when this is built up, we can actually run the two side by side, and they're going to produce the exact same results, but lm.c actually runs faster. So let's see that. So on the left, I have PyTorch, uh, NanoGPT-looking thing. On the right, I have the llmc call, and here I'm gonna launch to 2. Both of these are going to be running on a single GPU, and here I'm putting the lm.c on GPU 1, and this one will grab GPU 0 by default. And then we can see here that lm.c compiled and then allocates space, and it's stepping. So basically, uh, meanwhile PyTorch is still compiling because torch.compile is a bit slower here than the lm.c NVCC CUDA compile. And so this program has already started running and we're still waiting here for torch.compile. Now, of course, this is a very specific implementation to GPT-2 and 3. PyTorch is a very general neural network framework, so they're not exactly comparable. But if you're only interested in training GPT-2 and 3, lm.c is very fast. It takes less space. It's faster to start and it's faster per step. And so PyTorch started stepping here, and as you can see, we're running at about 223,000 tokens per second here and about 185,000 tokens per second here. Um, so quite a bit slower, but I don't have full confidence that I exactly squeezed out all the juice from the PyTorch implementation. But the important thing here is notice that if I line up the steps, you will see that the losses and the norms that are printed between these two are identical. So on the left we have the PyTorch and on the right this C code implementation, and they're the same except this one runs faster. So that's kind of— I wanted to show you also briefly lm.c, and this is a parallel implementation, and it's also something that you may want to play with or look at, and it's kind of interesting. Okay, so at this point I should probably start wrapping up the video because I think it's getting way longer than I anticipated, but we did cover a lot of ground and we built everything from scratch. So as a brief summary, we were looking at the GPT-2 and GPT-3 papers. We were looking at how you set up these training runs and all the considerations involved. We wrote everything from scratch and then we saw that over the duration of either a 2-hour training run or an overnight run, we can actually match the 124 million parameter checkpoints of GPT-2 and GPT-3 to a very large extent. In principle, the code that we wrote would be able to train even bigger models if you have the patience or the computing resources, and so you could potentially think about training some of the bigger checkpoints as well. There are a few remaining issues to address. What's happening with the loss here, which I suspect has to do with the fine-web-edu data sampling. Why can't we turn on torch.compile? It currently breaks generation and HellaSwag. What's up with that? In the DataLoader, we should probably be permuting our data when we reach epoch boundaries. So there's a few more issues like that, and I expect to be documenting some of those over time in the build-managed-gpt repository here, which I'm going to be releasing with this video. If you have any questions or would like to talk about anything that we covered, please go to discussions tab so we can talk here, or please go to issues or pull requests, pull requests, depending on what you'd like to contribute. Or also have a look at the Zero to Hero Discord, and I'm going to be hanging out here on NanoGPT. Otherwise, for now, I'm pretty happy about where we got, and I hope you enjoyed the video, and I will see you later.