 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 slash GPTpt2. Now when we talk about reproducing gpt2 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 miniseries when these releases are made. So there are the gpt2 miniseries 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 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 four models in the GPT-2 miniseries, 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, etc. 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 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 gonna 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-124M model. Now, previously when they were working on this, this is already five 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 bucks if you want to do this on the cloud 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-124M. 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-124M 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 to the code base 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, let's just say not used as much anymore. We 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. 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, GPT-2, Modeling, GPT-2.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. 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 pre-trained 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 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 1,024, we have up to 1,024 positions that each token can be attending to in the past. And every one of those positions in GPD2 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. 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 And it has structure because 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 three 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 one from zero 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 zero. 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 under-trained 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 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 gonna 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 can not 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 five 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 dot dot dot so you can read through these if you like but basically these are five different completions of the same prefix from this gpt2124m 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 the modeling GPT-2.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 on 24M 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 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 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 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 the encoder. GPT-2 is a decoder-only transformer, as we call it. So this entire encoder here is missing. 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 GPT NN 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 nnModuleDict. 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 nnEmbedding. And the weights of the position embeddings which is also just an NN embedding and if you remember an 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 the 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 dot h and then there's a this is indexed using numbers instead of indexed using strings so there's a dot h dot 0 1 2 etc all the way up till dot h dot 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 dot zero dot one two etc and the module list has uh 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 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 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 dot h is the all the blocks in gray and the lnf is this new layer that gets added here by the g-2 model and lmhead 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 feed forward 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 micro grad 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 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 Rx first goes through the layer normalization and then the attention and then goes back out to go to the layer normalization number two and the multi-layer perceptron, sometimes also referred to as a feedforward 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 there's 1,024 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 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 ends up just being a repeated application of map reduce, 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 non-linearity. So nn.gelu approximate is 10h. Now when we swing over to the PyTorch documentation, this is nn.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 10H. So as you can see, just as a preview here, GELU is basically like a ReLU, except there's no 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, at least to the specific formulation. It has to do with stochastic regularizers and the expectation of a modification to it after dropout, so you can read through all of that if you'd like here. And there's a little bit of the history as to why there isn't an approximate version of GELU and that comes from this issue here as far as I can tell. And in this issue Daniel Hendrix mentions that at the time when they developed this non-linearity 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, so they ended up basically developing this approximation. And this approximation 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 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 RELU 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 LAMA3 and so on this non-linearity also further changes to SWGLU and other variants like that, but for GPT-2 they use this approximate GELV. 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 the previous video, we had this multi-headed attention module. And this implementation made it obvious that these heads are not actually that complicated. There's basically in parallel inside every attention block, there's multiple heads and they're all functioning in parallel and 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 impact torch. But fundamentally and algorithmically, nothing is different from the implementation we saw before in this GIF repository. So to remind you very briefly and I don't want to go in this into this in too many 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 three 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 while 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 one 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 one always. And then recall from the previous video that doing the attention matrix multiplied 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 just reassembling all of that again. And this actually performs a 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, cAten is the same as c-atten. 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 2000 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 hybrid parameters, 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 5250257, 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 Phase 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 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 four models in this mini series of GPT-2. This is some of the Jupyter 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. 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 the 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 auto-aggressive 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 hard-coded 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 GPT-2.py. And fingers crossed. 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 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 as 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 zero 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 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 lmhead. 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 plus 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 five times up to 30 tokens with the prefix of hello on the language model, and these are the completions that we achieved. So we're going to try to replicate that on the left here. So number of turn 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 tech token library from OpenAI So here's the code that I've written. We're going to import the Tiktokin 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 T tokenizer 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 eight tokens. And then we're going to replicate these eight tokens for five times to get five rows of eight 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 pin to forward to get our logits so that we know what comes as the sixth token uh sorry as the ninth token in every one of these five rows. 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 this x, which is of size b by t, 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. 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 case sampling of 50, and I'm doing that because this is the HuggingFace default. So just looking at the HuggingFace docs here of a pipeline, there's a bunch of quarks that go into HuggingFace. And I mean, it's kind of a lot, honestly, but I guess the important one that I noticed is that they're using topk 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 zero 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 it doesn't get lost, and it 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 y loop gets tripped up and then finally we have an entire x of size um 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 tiktokennizer to get back the string, which we can print. And so terminal, new terminal, and let me python train gpt2. Python train GPT-2. Okay, so these are the generations that we're getting. Hello, I'm a language model, not a program. 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 HuggingFace model so this is the HuggingFace 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 Phase 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 initialize the GPT-2, this is the exact opening of 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 the 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 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 Javier 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 Now the results here of course are total garbage carbel 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 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 a CUDA. And then if you don't have a CUDA, do you at least have NPS? NPS is the backend for Apple Silicon. So if you have a MacBook that is fairly new, you probably have Apple Silicon on the inside. And then that has a GPU that is actually fairly capable, depending on which MacBook you have. And so you can use NPS, which will be potentially faster than CPU. And so we can print the device here. Now, once we have 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 create a pose, we are 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 you can't combine those. But here we are carefully initializing on the correct device as indicated by the input to this model. 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 gonna be too much slower. So if I override device here, oops. 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 one two three four five six okay about six 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 orice or where you take a TorchTensor and do a.toDevice 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 double you get it or you can just search tiny Shakespeare dataset and so I have in my file system is just lsinput.txt so I already downloaded it and here I'm reading the data set, 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 to 1, so 1000 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 data set 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 one byte, and so this is the same number, roughly a 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 tiktokin encoding for GPT-2, encode the data, the first 1000 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 slash in character. So that is a new line. And then here, for example, we have two two new lines 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 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 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 dot view of, for example, 4 by 6, which multiply to 24. And so it's just a two-dimensional rearrangement of these tokens. And you'll notice that when you view this one-dimensional sequence as two-dimensional 4 by 6 here, the first six tokens up to here end up being the first row, the next six tokens here end up being the second row, and so on. And so basically it's just going to stack up every six 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 three tokens and is going to try to predict that 198 comes next. So in this way we are able to create this two-dimensional batch 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 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 plus one 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 4x6. everything up to the last token, not including, and view it as 4x6. 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 target 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 plus one 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 one, 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 data loader 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 to token. And all of this should look familiar. We're loading a thousand characters. I'm setting BT to just be four 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 create the model and get the logits. And so here, as you see, I already ran this. It 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 zero, 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 logit 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 that cross entropy does not like multidimensional inputs. It can't take a B by T by vocab size. So what's happening here is that we are flattening out this three-dimensional tensor into just two 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 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 three dimensional tensor of logits to just be two-dimensional b times t all individual examples and vocab size on in terms of the length of each row and then it's also flattening out the targets which are also two-dimensional at this stage but we're gonna 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. Here we see that we printed 11 roughly. Notice that this is the tensor of a single element, which is this number 11. Now, we also want to be able to calculate a reasonable starting of starting point for a randomly initialized 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 we're hoping is that the probability of any arbitrary token is roughly 1 over 50,000 to 157 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 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. and do an optimization so let's get to that okay so let's do the optimization now so here we have the loss is 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 atom optimizer which is an alternative to the Stochastic Gradient Descent Optimizer SGT that we're using so SGT is a lot simpler Atom is a bit more involved and I actually specifically like the Atom W variation because in my opinion it kind of just like fixes a bug so Atom W is a bug fix of Atom is what i would say when we go to the documentation for atom w oh my gosh we see um 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, something that looks a bit like RMS prop, 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 this 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. and then go through the optimization. The first thing to always make sure, the copilot did not forget to zero the gradients. So always remember that you have to start with a zero gradient. Then when you get your loss and you do a dot backward, dot backward adds to gradients. So it deposits gradients. It always does a plus equals on whatever the gradients are which is why you must set them to zero 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 dot item is used here because loss is a tensor with a single element dot item will actually convert that to a single float and this float will live not will live on the cpu so this gets to some of the internals again of the devices but loss is a is a tensor with a single element and it lives on gpu for me because i'm using gpus when you call that item pytorch behind the scenes will take that one-dimensional tensor ship it back to the cpu uh 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 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. Expected all tensors to be on the same device but found at least two devices, CUDA 0 and CPU. So CUDA0 is the zeroth GPU because I actually have eight GPUs on this box. So the zeroth GPU on 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 of device it's not stateful it doesn't convert it to be a device it instead returns pointer to a new memory which is on the device so you see how we can just do model.to a device that does not apply to tensors you have to do buff equals buff.toDevice 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 in these that is roughly what we're seeing here so 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 XY batches and create a little data loader 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 what I came up with, and I wrote a little data loader light. So what this data loader light. So what this data loader does is we're importing the token 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 the beginning of the document and start reading it again? So we start off at position zero, 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 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 one and we need that plus one because remember we need the target token for the last token in the current patch and so that way we can do the x y exactly as we did it before 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 zero. So this is one way to write a very, very simple data loader that simply just goes through the file in chunks and is good enough for us for current purposes. And we're gonna complexify it later. And now we'd like to come back around here and we'd like to actually use our data loader. So the import tick token 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 four. So batch size was four and time was 32. And then here we 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 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 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 that 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 data loader simple CPU 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 of these crazy unit codes 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 have 338,000 tokens, which makes sense with our 3 to 1 compression ratio because there are 1 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. 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 Phase 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. So this one here is the token embedding at the bottom of the transformer. 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,000 to 57 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,000 to 57 to get the logis 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 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 wait time scheme. So what's happening here is that this is a common wait time scheme that actually comes 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 30. 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 30 reference is this paper. So this came out in 2017. And you can read the full paper, but basically it argues for this weight-signing scheme. And I think intuitively the idea for why you might want to do this comes from this paragraph here. And basically you can observe that 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 the 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 should have similar embeddings or similar weights. And so this is what motivates their exploration here. And 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. 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 on meat 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 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 to 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 bottom of it, flowing again into the WTE tensor. So we are currently not sharing WTE in our code, but we want to do that. So 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 lmhead. So this basically copies the data pointer, right? It copies the reference. And now the wte.wait 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 sensor 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 of here it's 768 times 50,257 so this is 40 million parameters and this is a 124 million parameter model so 40 divide 124 so this is like 30 percent 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 wait time 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 wait time scheme and we've saved a ton of parameters and we expect our model to work slightly better because of the 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 so this is a normal distribution for the weights and the standard deviation is 0.02 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 an end module and that iterates all the sub modules 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 gonna make sure to initialize the weight using a normal with the standard deviation of 0.02 if there's a bias in this layer we will make sure to initialize that to zero. Note that zero 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 zero 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 one and the offset in the layer norm to be zero. 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 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 follow the 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 D model sizes inside these transformers for GPT-2 are roughly 768, 1,600, etc. So 1 over the square root of, for example, 768 gives us 0.03. If we plug in 1,600 1600 we get 0.02 if we plug in three 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 hard coding 0.02 here, but you'd like typically 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 one over square root of n, where n is the number of residual layers. So this is what 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 the residual stream grows. So here's a small example if we start at zero and then we for 100 times we have sort of this residual stream of 768 zeros and then 100 times we add random which is a normal distribution zero mean one 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 we're always adding 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 layer 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. We'll do in the initialization, see that special nano GPT scale init is one. 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 this std should be 0.02 by default but then if hasAtter module of this thing then std times equals Copilot is not guessing correctly so we want one 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 negative 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 two times number of layers comes from the fact that the two times number of layers comes from the fact that every single one of our layers in the transformer actually has two blocks that add to the residual pathway, right? We have the attention and then the MLP. So that's where the two 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 of our old 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 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 gpt2 initialization in the way we've implemented it right now so this looks reasonable to me okay so at this point we have the gpt2 model we have some confidence that it's correctly implemented we've initialized it properly and we have a data loader 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 eight GPUs, and each one of those GPUs is an A100 SXM 80 GB. 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 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 VSCO to it. And that's how I develop. Now, when we look at the A100s that are available here, A100 80 gigabyte 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 laws. And the interesting thing I'd like you to note is when I do logits.dtype, this prints a 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 precisions. Not all computational workloads can tolerate small precision. So for example, if we go back to the data sheet, 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. 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 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 float 16 or be float 16, 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 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. Typical utilizations of your hardware, if you're getting 60 percent 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 tensor float 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 any time we have any operations that require matrix multiplication, they get broken up 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 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 galu non-linearities there's some layer norms sandwiched in between. So there's some additions in residuals, there's some Gellu 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 lot faster which are hidden inside our linear layers and they're accelerated through TensorCourse. Now the best reference I would say for TensorCourse is basically just go to the 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 tensor float 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 precisions 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 A-tax improvement as I mentioned to the hops that we get so tf32 specifically we're looking at this row here. And the way this works is normally fb32 has 32 bits. tf32 is the exact same bits. We have one sign bit. We have eight exponent bits, except the mantissa bits get cropped in the float and so basically we end up with just 19 bits instead of 32 bits because the last 13 bits get truncated they get dropped and all 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 in 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 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 go 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 codes 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. request and then it continues running. And so 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 to time but actually the GPU is still running because it takes a 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 And so actually if you need to, you want to wait for Chatku to synchronize, 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 zero we're not using the gpu and then by default pytorch will use gpu zero so we see that it gets filled up and we're using 35 gigabytes out of 80 gigabytes available and then 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 a thousand milliseconds 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. And so this will run for roughly 50 seconds as we have it right now. So that's our baseline in float 32. 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 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 the GPU. And we're going to see that a bit later as well. So for now, let's just stick with 16, 1024. 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. OK, 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, matrix multiplications will now use tensorflow32 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? 8x. And it also came from looking at it here, 156 tflops instead of 19.5. Okay, so what actually happened? So we're seeing that our throughput roughly 3x'd not 8x'd. So we are going 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 a 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 not 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 to bfloat16. So we're only going to be maintaining 16 bits per float. And we're going to use the bfloat16, and I'll explain in a bit fp16 difference, and we're going to be in this row. So when we go back to the documentation here for the A100, 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 BF16 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 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. you have for your numbers. And so the range of numbers is identical, but we have fewer possibilities within that range because we are truncating the mantissa. So we have less precision in that range. 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 in 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 BF16 and the Ampere, and they made it much simpler because we're just truncating Matissa, we have the exact same range, and we do not need gradient scalers. So everything is much, much simpler. 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. This change is not just local to the operation itself. So let's see how that works. There's some documentation here. So I think this is probably the 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 five 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 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 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 float 16 because if we do that we need to start using gradient scalars as well so we are going to be using b float 16. 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 um let me first break in to here before we actually run this so right right after logins, I'd like to show you that different from the tf32 that we saw, this is actually going to impact our tensors. So this logins 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 pfloat16, but not everything has changed. So model.transformer.wte, this is the weight token embedding table, it has a dot weight inside it, and the dtype of this weight this parameter is still torch float 32. so our parameters seem to still be in float 32 but our activations the logits are now in b float 16. so clearly this is why we get the mixed precision some things pytorch is keeping in float 32 some things pytorch is converting to lower precision and what gets converted at what point is not super clear. I remember scrolling down... is it here? 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 b float 16 and and when so for example only these matrix multiply like operations get converted to b float 16 but a lot of operations remain in float 32 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 running 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. Major x 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 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. 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 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 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 BFloat 16 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 that precision decrease. So that's B416 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 cnc++ code 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 rights. So let me unpack that a little bit. Okay, here we are. Okay, so we went from 300 milliseconds. 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 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 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 GELU nonlinearity that we've been looking at. So here we use the nnGELU. Now this here is me basically just breaking up the nnGELU which you remember has this formula. So this here is the equivalent implementation to what's happening inside GELU. Algorithmically it's identical. Now by default if we just were using this instead of ending.gulu here, what would happen without TorchCompile? 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 course, 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 travel time that actually causes a lot of issues. So here, remember this memory bandwidth? We can communicate about two 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 an 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 ALUs and everything like that is all stored on the chip and 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-wide 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 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 onset 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 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. Because this is a complicated topic, but it's worth understanding on a high level what's happening here. And I could probably spend an entire video of like two 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 they're connected. But these are two separate chips basically. Now here, this is zoom-in of kind of this cartoon diagram of a GPU, and what 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 multi-processors. 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 four 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 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 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 one 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 of 10 megabytes of memory collectively throughout 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 during the calculations and saving to the memory, and we're doing those round trips many, many times. But if it's 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 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 and running it a lot faster. So flash attention will come here and we will take out these four lines and FlashAttention implements these four 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 four operations. They're including dropout but we are not using dropout here. So we just have these four 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 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 their load and store into HBM, and that's what they avoid. And so in particular, they do not ever materialize this end-by-end attention matrix, this ATT here. 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 has 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 normalizer calculation for softmax and remarkably it came out ofVIDIA and it came out of it like really early 2018. So this is four 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 four years later from Stanford. So I don't fully understand 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 flash attention 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 flashAttention. And the way to do that basically in PyTorch is we are going to comment out these four lines. And we're going to replace them with a single line. And here we are calling this compound operation in PyTorch called scaled.productAttention. in PyTorch called scaled.productAttention. And PyTorch will call FlashAttention when you use it in this way. I'm not actually 100% sure why Torch Compile doesn't realize that these four lines should just call FlashAttention in this exact way. We have to do it again for it, which in my opinion is a little bit odd, but here we are. So you have to use this compound up and 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 see 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 retention is just an algorithmic rewrite, and it's a faster kernel, but it doesn't actually change any of the computation, and we achieve 6.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 divide 130-ish so this is maybe 27-ish percent improvement so really interesting and that is flash retention 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 two inside them. You can divide by two 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 two and lots of kernels are written in terms of powers of two 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 three times is kind of ugly. I'm not 100% sure maybe this can be improved, but this is ugly and not ideal. That's nice. 1,024 is very nice. That's a power of 2. 12 is a little bit suspicious. Not too many powers of 2. 768 is great. 50,000 to 57 is a really, really ugly number. First of all, it's odd. And there's 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're trying to optimize some kernels to run this fast and required a bunch of special case handling. So basically, we have some ugly numbers, and some of them are easier to fix than others. And in particular, the vocab size being 50,000 to 57, 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 two 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, 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 50,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 book app size has powers of two 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 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 4% improvement here by doing more calculations. And the reason for this is we've made an ugly number into a nice number. 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 that's doing is we're predicting additional dimensions of 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 zero. 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, a thousand tokens out of 50,000 to 57 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. 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 two so calculations are done in like chunks of 64 or chunks of 32 and when your 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 truncate up your input and they will do the nice part first and then they have a whole second second phase where they come back to anything that like uh remains uh and then they process the remaining part and the kernels for that could be very inefficient and so you're basically um 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 four percent 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 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 from 50,000 to 57,000 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'd like to follow the hyperparameters that are mentioned in the GPT-2 or GPT-3 paper. Now sadly GPT-2 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 small details that go into the model training, as to all the 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. 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 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 Atom 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 you can see is the default is 1 negative 8 and this is also 1 negative 8 let's just put it in so that we're explicit 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 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. 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 length of it, 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 the upper bounded 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 a 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. warm-up and it's got a cosine decay to 10% over some horizon and so we're gonna 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 one. 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 warm-up and the way this looks is that the learning rate is basically starts right at around zero, 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 zero, 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 4 i in max steps. I just change it to step now so that we have the notion of a step is a single optimization step 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 look of a 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 five 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 three negative four 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. So for this one, basically 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-4 so we could actually go higher in fact we may want to try to follow that and just set the max lr here at six then the that's the maximum learning rate the min learning rate is uh 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 terribly inside floor interesting I'm just modulating based on the iteration number of 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 zero because that would be not useful to update with a learning rate of zero. That's why there's an it plus one, so that on the zeroth iteration, we are not using exactly zero. 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 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 this is okay for us and for our purposes right now. And we're just going to use this ourselves. I don't think it makes 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 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 in the early stages of the optimization, again, the model is in a very atypical setting. And mostly what you're learning is that you're mostly learning to ignore the tokens that don't come up in your training set very often. You're learning very simple biases 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 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 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 de-correlated per examples. And that's where the actual work starts, and that's where the gradients become more de-correlated 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 return the sequence to the pool. They are exhausting a pool. So when they draw a sequence, it's gone until the next epoch of training. So we're already doing that because our data loader 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. 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 addmw optimizer and the weight decay that's by default used in addmw here is 0.01 so it's it's 10 times lower than what's used in GPT-3 paper here so the weight decay basically ends up making its way into the addmw3 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. So the one-dimensional tensors are in the no decay params 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 the network to kind of like distribute the work across more channels because there's sort of like a pool 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 we're 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 Adam W 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 spec.signature which is basically checking if this fused quark 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 big time so by default they 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 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 AtomW update instead of iterating over all the tensors. So that's the ConfigureOptimizers 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 decay tensors is 50, and it's most of the parameters. And number of non-decay tensors is 98, and these are the biases and the layer norm parameters mostly. And there's only 100,000 of those. So most of it is decayed. And then we are using the fused implementation of AtomW, 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 Adam W optimizer. So in a single commit here, we are introducing fused Adam, 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 the linear. So that is this, and we can take this out, and 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 atom parameters, beta 1, beta 2, the epsilon, and so on, these are very complicated mathematical relationships in the optimization literature, These are very complicated mathematical relationships in the optimization literature. And for the most part, 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 a 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 data loader. 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've seen the pattern that the bigger networks are trained with slightly lower learning rates and we also see this batch size where in the small networks they use a smaller batch size and in the bigger networks they use a smaller batch size and in the bigger networks they use a bigger batch size. Now, the problem 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 is my B? B equals, where do I call the deal? 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 1,024 tokens. So 0.5 e6 1 million divide 1,024. 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 a 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 to the 19 is 524288. 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 over 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. Oops what am I doing? 2 to the 19 in total so the grad accum will be 32. 19 in total so the gradacum will be 32. So therefore gradacum 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 make every step roughly three seconds in just napkin math so that's graticum 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, gradicum steps, we do this. And remember that lost at backward always deposits gradients. So we're doing doing inside lost at backward there's always a plus equals on the gradients so in every single lost at backward gradients will add up on the gradient testers so we lost at 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 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 four individual examples. And we're just doing simple regression with the mean squared loss over those four examples. Now, when we calculate the loss and we loss that backward 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 four examples. So this is the exact loss objective and this is the average, the 1 over 4, because there are four independent examples here. And then we have the four examples and their mean squared error, the squared error, and then this makes it the mean squared error. So therefore, we calculate the squared error, and then we normalize it to make it the mean over the examples. And there's four examples here. So now when we come to the gradient accumulation version of it, this here is the gradient accumulation version of it where we have grad accrued steps of four and I reset the gradient with grad accrued steps of four 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 achieved from that. So basically now we forward our function, calculate the exact same loss, do a backward, and we do that four times. And when we look at the gradient, you'll notice that the gradients don't match. So here we did a single batch of four, and here we did four gradient accumulation steps of batch size one and the gradients are not the same and basically the reason that they're not the same is exactly because this mean squared error gets lost this one quarter in this loss gets lost because what happens here is the loss objective for every one of the loops is just a 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 zeroth iteration, same in the first, third, and so on. And then when you do the loss.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 quarter 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 four and what happens now is that we're introducing we're scaling our loss we're introducing a one-quarter in front of all of these places so all the individual losses are now scaled by one quarter and then when we backward all of these accumulate with the sum but now there's a one quarter inside every one of these accumulate with the sum but now there's a one quarter 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 MSC loss the loss that we're calculating here in the model is using a reduction of mean as well so where is the loss? f dot 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. 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 a copilot 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.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 Gratacom 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 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 of kum which I initialize at zero 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 them. So that's loss of kum and then we're printing that here instead of loss and then in addition to that I had to account for the the grad accum 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 grad accum steps to be 32. And we're getting about three 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 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 eight GPUs here, and so we should be putting all of them to work. And in particular, they are going to collaborate and 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 distributed data parallel from PyTorch. There's also a legacy data parallel, which I recommend you not use, and that's kind of like, you know, legacy. Distribute data parallel works in a very simple way. We have eight GPUs, so we're going to launch eight 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 eight 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 eight of them, we're not going to be launching our script anymore with just PyTorch traingpt2.py. We're going to be running it with a special command called torch run in PyTorch. We'll see that in a bit. And torch run, when it runs our Python script, will actually make sure to run eight 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, Torch Run 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 eight. That's the total number of processes running. There's a rank, which is each process will basically run the exact same code at the exact same time roughly, but the only difference between these processes is that they all have a different DDP rank so the GPU 0 will have DDP rank of 0 GPU 1 will have 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 a gpus and so local rank is the rank of the gpu on a single node so from zero to seven as an example but for us we're mostly going to be running on a single box so the things we care about our 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 colon 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 ddprank == 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 zero the world size is just one 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 Torch Run, which we'll see in a bit, there's going to be eight 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 gonna be eight processes running in parallel. So as you read the code now you have to imagine there's eight 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 micro batches and sequence lengths these are all just per GPU right so now there's gonna be num processes of them running in parallel so we have to adjust this right because the Graticum steps now is gonna be total batch size divided by V times T times DDP world size because each process will do B times T and there's this many of them and so in addition to that we want to make sure that this fits nicely into total batch size which for us it will because 16 times 1.4 times 8 GPUs is 131k. And so 5, 2, 4, 2, 8, 8. This means that our Gratacum will be 4 with the current settings. So there's going to be 16 times 124 process 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 eight 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 TDP, world size. Okay. So each GPU calculates this. Now this is where we started to get right into issues, right? So each process is going to come by a print, and they're all going to print. So we're going to have eight copies of these prints. So one way to deal with this is exactly this master process variable that we have. So 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 eight times before getting into the data loader and we're gonna have to refactor it obviously maybe at this point is we should do some prints and just take it out for a spin and exit at this point so import sys and sys.exit and print imgpu ddp rank and print bye gpt rank and print pi. 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.pi like this. Now we're going to run with torch run and this is what it looks like. So torch run standalone number of processes for example is eight for us because we have eight GPUs, and then train gpt to that pi. So this is what the command would look like. And Torch Run, again, we'll run eight of these. So let's just see what happens. So first, it 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 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 unprocess on 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 gpu0 ended then gpu3 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 GPUs 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 wanna call 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 that to be printing and next up let's erase this. Next up we want to make sure that when we create data loader light we need to now make it aware of this multi-process loader light we need to now make it aware of this a multi-process 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 data set 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 data loader. And then when we come up here we see that we now take rank and processes and we save them. Now the current position will not be zero because what we want is we want to stride out all the processes. So one way to do this is we basically take self.b times self.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 is 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 number of processes, and they all go off to a different rank, and the position has to advance by the entire chunk and then here at b times t times self.numProcesses plus one 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 data loader light and you can notice that if process rank is zero and then processes is one 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 let's do the data loader. Okay so next up once they've all initialized the data loader they come here and they all create a GPT model. So we create eight GPT models on eight 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. Because the models are identical, there are eight 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 within-step because all the changes we're making are within-step changes. Now, the important thing here is when we construct the model model we actually have a bit of work to do here. GetLogits is deprecated so createModel. 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. 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 allReduce. 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 involved in 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 the synchronization of them and the backward pass. And this is just more efficient 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 lost at 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 gratacrum steps minus one, only at that last step do we want to actually do the all reduce to average up the gradients. So to do that, we come here. And the official sanctioned way by the way is to do this no sync 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 no sync there will be no communication. So they are telling us to do with ddp no sync uh 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 backward grad sync and this is uh being toggled around and changed and this is the variable that basically uh 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 last 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 grad 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 lost up 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 lost up 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 loss of coom, well presumably we're only going to be printing on the master process rank zero 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 need 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 loss if we want loss so if tdp again then this is a pytTorch 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 Lossicum. And so this Lossicum tensor exists on all the ranks when we call all reduce 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 lossacum averaged up and so when we print here on the master process the lossacum 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 not 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 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 data set 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.configureoptimizers which now doesn't work because model now is a DDP model so instead this has to become rawModel.configureOptimizers where rawModel is something I create here so right after I wrap the model into DDP I have to create the rawModel which in the case of DDP is a model.module, is where it stores the raw NM 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 just running a single GPU, you'll notice that this is a single GPU run with 32 RADACUM. The numbers won't exactly match up. And that's kind of a boring reason for why that happens. The reason for that is that in the data loader, 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 is just make the total batch size much smaller and the B and a T and then um so I think I used uh four times 124 times eight so I used 32768 as a total batch size and then um so I made sure that the single GPU will do eight gradient accumulation steps and then at the multi-GPU and then you're reducing the boundary effects of the data loader 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 hybrid parameters and we have outgrown our tiny Shakespeare file and we want to upgrade it. So let's move 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 with at least three karma. And that was kind of like their starting point. And they collected all the webpages and all the text in them. And so this was least three 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 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 data set all by itself because it's 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 and 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 WebText2 from before, so 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, 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 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 FindWebEDU subset which is 1.3 trillion of educational and 5.4 trillion of high educational content. So basically they're trying to filter CommonCurl to very high quality educational subsets and this is the one that we will use. There's a long web page 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 into 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 FineWebEDU I think and it's basically educational content from the Internet they show that training on educational content 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 10 billion sample of the fineweb edu because empirically in my previous few experiments this actually suffices to really get close to gpt performance. And it's simple enough to work with. And so let's work with the sample 10 BT. So our goal will be to download it, process it, and make sure that our data loader can work with it. So let's get to that. Okay, so I introduced another file here that will basically download FindWebEDU from Hugging Face datasets. It will pre-process and pre-tokenize all of the data, and it will save data shards to a folder on a local disk. And so while this is running, 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 basically looks like it's working fairly well. Like it's talking about nuclear energy in France. It's talking about Mexican America, some Mac, PyJs, etc. So actually it seems like their filters are working pretty well. The filters here, by the way, were applied automatically using LLM370B, 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 data set which this is all Hugging Face code running this you're going to need to pip install datasets so it's downloading the data set then it is tokenizing all of the documents inside this data set now when we tokenize the documents you'll notice that 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 50,256 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 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 mp.un16. We're using un.16 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 multi-processing code and it's honestly not that exciting so I'm not gonna step through it but we're loading the data set we're tokenizing it and we're saving everything to shards and the shards are numpy files so just storing a numpy array and which is very very similar to torch tensors which is very, very similar to TorchTensors. And the first shard, 0, 0, 0, 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 a 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 data data set. We're doing lots of tokens per second. We have eight 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 find web, we see that there's now 100 charts in it. And that makes sense because each chart is 100 million tokens. So 100 charts of that is 10 billion tokens in total. Now swinging over to the main file, I made some adjustments to our data loader again. And that's because we're not running with Shakespeare anymore. We want to use the find web charts. And so you'll see some code here that additionally basically can load these charts. We load the UN1616 numpy file we convert it to a torch.long tensor 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 data loader light so we can load the split train but also the split val the zero split and then we can load the shards and 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 data loader 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 I set up some numbers so we are doing 2 to the 19 tokens 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 E6 tokens divide 2 to the 19 is 715 steps so that's why warm up steps is set to 715 so this will exactly match So that's why warmup steps is set to 7.15. So this will exactly match the warmup schedule that GPT-3 used. And I think 7.15, by the way, is very 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 fix that. And then that's pretty much it. 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, 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 back 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-trading 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 Ts and dotted our Is. But if we let this run for a while,'re gonna actually get a pretty good model and the model that might even be on par with or better than GPT-2 124M. Okay so it looks like everything is going great we're processing 1.5 million tokens per second everything here looks good. We're doing 330 milliseconds per iteration and we have to do a total of... where are we printing that? 1973. So 19073 times 0.33 is this many seconds, this many minutes, so this will run for 1.7 hours. So one and a half 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-4. 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 Ts and dot Ris. 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 equals val that will basically create a data loader just for the validation shard. The other thing I did is in the data loader, I introduced a new function reset, which is called an init, and it basically resets the data loader. And that is very useful because when we come to the main training loop now, so this is the code I've added, and basically every hundredth iteration, including the zeroth iteration, we put the model into evaluation mode we reset the val loader and then 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 that 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 hundredth iteration, including on the very first iteration. So that's nice. That will tell us a little bit about how much we're overfitting. That said, 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 can 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 an sorry to find web 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 validation, once in a while 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. So here are some samples on iteration 1000. Hello, I'm a language model and I'm not able to get more creative. I'm a language model and 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 1000 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 1100. Okay this is you know the model is still a young baby. Okay so basically all of this sampling code that I've put here everything should be familiar 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. 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 end 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 now and after this as well and so everything should be exactly documented in the git commit history and so I think that'll be nice. So hopefully by the time you go to github this is removed and it's working and I will have fixed debug. Okay so I have the optimization running here and it's stepping and we're on step 6000 or so so we're about 30% through training. Now while this is training I would like to introduce one evaluation that we're going to use to supplement the validation set and that is the HelloSwag eval. So HelloSwag comes from this paper back in 2019 so it's a five-year-old eval now. And the way Hellaswag works is there is 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-dried 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 choices 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 host 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 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 Wikikihouse. 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 completion. One more thing that's kind of interesting about HeLaSwag 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 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 learn more. The thing to point out though is that this is five years ago and since then what happened to HelloSwag is that it's been totally just solved. And so now the language models here are 96%. So basically 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 percent but this is how far things got but still the reason people like hella swag and it's not used by the way in gpt2 but in gpt3 there is hella swag eval and lots of people use hella swag and so for g3, 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, 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 four rows and t tokens whatever that t happens to be then the shared context that is basically the context for the four choices the tokens of that are shared across all of the rows and then we have the four options so we kind of like lay them out and then only one of the options is correct in this case label three option three and so this is the correct option and option one two and four 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 zero for these 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 gonna 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 gonna 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 hella swag may not do it this way. They may do it in a multiple choice format where you sort of give the context a single time and then the four completions. And so the model is able to see all the four 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 hellosw 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 four by T array of tokens, the mask, which tells us which parts are the options and everything else is zero, and the label that is the correct label. And so that allows us to then iterate the examples and everything else is zero 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 gpt2 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 valuing 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 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 HelloSwag. Now, if you go up here, I'm showing that for GPT-2124M, if you run this script, you're going to see that HelloSwag 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 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 were training a GPT-2124M from scratch in 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 HelloSwag and how it evolves over time and see when and if we cross this 29.55 sort of region. So let's now walk through some of the changes to train GPT to that pipe. The first thing I did here is 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 code base when I put it up on GitHub, we're gonna 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 LLSWG 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 Hella Swag, but only if we are not using compile because compile breaks it. So I'm going to come back to this code for evaluating Hella Swag 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 um these are if we're not after we validate sample and evaluate heliswag we actually do a training step here and 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 the way this works is I'm trying to get all the GPUs to collaborate on the Halaswag. And so we're iterating all the examples. And then each process only picks the examples that are 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 a 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 a 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.allReduceOn 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 HellasVag accuracy. So 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 language model uh let's see if i can find any fun one 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, vowel, and the heli-sway 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 HeLa, and the HeLa 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 FineWebEDU. On the left, we have the loss. And in blue, we have the training loss. In orange, we have the validation loss. And in red, as a horizontal line, we have the opening IGPT2-124M model checkpoint, when it's just evaluated on the validation set of this FineWebEDU. So you can see that we are surpassing. This orange is below the red. evaluated on the validation set of this FIneWebEDU. 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 look at. Now, we would ideally like something that is withheld and comparable and somewhat standard. And so for us, that is Helioswag. And so on here, we see the Helioswag progress we made from 25% all the way here. In red, we see the OpenAI GPT-2-124M model in red, so it achieves this helos flag 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. 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 templating training. So number one, it could be that OpenAI GPT-2 was trained on a much wider data distribution. So in particular, FWEDU 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 Helioswag eval is fairly old maybe five years or so it is possible that aspects of Helioswag in some way or even identically have made it into the training set of FindWeb. 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 we're not doing something completely wrong. 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 FindWebEDU, they use HeLaSwag as an eval. So I would hope that they make sure that they deduplicate and that there's no HeLaSwag 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 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 a 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-nano-gpt 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 videoata and I'm going to talk about some of the things that happened after the video and I expect that we will have fixed the 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 And possibly it could be also that the dataset may have improved. So the original GPT-2 dataset was webtext. 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 two hours, I just did it times four so that that would take eight hours while I was sleeping. And so we did four 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 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 FineWebEDU dataset and that is to be determined but otherwise we are seeing that the heli swag actually went up by a lot and we almost we almost made it to the gpt3 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 data loader, and this data loader 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 in the identical format and you're introducing 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 up the documents and shuffle them around because the order of the documents shouldn't matter and they shouldn't basically want to break up that dependence because it's a kind of a spurious correlation and so our data letter 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 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 changes 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 GPT-3 and in that case basically the models would be roughly identical 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 heliosmMagui cheat 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 they are a lot more coherent so and they're actually addressing the fact that it's a language model almost so hello I'm a 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 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 training process. But if you just want to save the model, this is how you would do it. And 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 LswaggyVal but you may want to use something nicer like for example the Luther Luther evaluation hardness evaluation 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 just predicts the next token you can't talk to it like you can talk to chat GPT 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 data set into a data set 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 would 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 gonna stop it 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 case because it is a direct optimized implementation. So trainGPT2.py in LLM.C is basically the nano GPT and when you scroll through this file you'll find a lot of things that very much look like things that we've built up in this lecture and then when you look at train-gpt2.cu this this is the C CUDA implementation. So there's a lot of MPI and NICL, GPU, CUDA, C, C++, and you have to be familiar with that. But 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, a nanoGPT looking thing. On the right I have the lmc call and here I'm gonna launch the two. Both of these are going to be running on a single GPU and here I'm putting the lm.c on GPU1 and this one will grab GPU0 by default and then we can see here that lm.c compiled and then allocate space and it's stepping so basically meanwhile pytorch is still compiling because torch compile is a bit slower here than the lm.c nvcc c 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. 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 align 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 CQt implementation and they're the same except this one runs faster. So that's kind of I wanted to show you also briefly Ln.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 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 two-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 patients 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 FineWebEDU data sampling. Why can't we turn on Torch Compile? It currently breaks Generation and HeliSwag. What's up with that? In the data loader, we should probably be permuting our data when we reach epoPUB boundaries. There's a few more issues like that, and I expect to be documenting some of those over time in the build.nano.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 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
Let's reproduce GPT-2 (124M)
We reproduce the GPT-2 (124M) from scratch. This video covers the whole process: First we build the GPT-2 network, then we optimize its training to be really fast, then we set up the training run following the GPT-2 and GPT-3 paper and their hyperparameters, then we hit run, and come back the next morning to see our results, and enjoy some amusing model generations. Keep in mind that in some places this video builds on the knowledge from earlier videos in the Zero to Hero Playlist (see my channel). You could also see this video as building my nanoGPT repo, which by the end is about 90% similar.

Links:
- build-nanogpt GitHub repo, with all the changes in this video as individual commits: https://github.com/karpathy/build-nanogpt
- nanoGPT repo: https://github.com/karpathy/nanoGPT
- llm.c repo: https://github.com/karpathy/llm.c
- my website: https://karpathy.ai
- my twitter: https://twitter.com/karpathy
- our Discord channel: https://discord.gg/3zy8kqD9Cp

Supplementary links:
- Attention is All You Need paper: https://arxiv.org/abs/1706.03762
- OpenAI GPT-3 paper: https://arxiv.org/abs/2005.14165 - OpenAI GPT-2 paper: https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf- The GPU I'm training the model on is from Lambda GPU Cloud, I think the best and easiest way to spin up an on-demand GPU instance in the cloud that you can ssh to: https://lambdalabs.com 

Chapters:
00:00:00 intro: Let’s reproduce GPT-2 (124M)
00:03:39 exploring the GPT-2 (124M) OpenAI checkpoint
00:13:47 SECTION 1: implementing the GPT-2 nn.Module
00:28:08 loading the huggingface/GPT-2 parameters
00:31:00 implementing the forward pass to get logits
00:33:31 sampling init, prefix tokens, tokenization
00:37:02 sampling loop
00:41:47 sample, auto-detect the device
00:45:50 let’s train: data batches (B,T) → logits (B,T,C)
00:52:53 cross entropy loss
00:56:42 optimization loop: overfit a single batch
01:02:00 data loader lite
01:06:14 parameter sharing wte and lm_head
01:13:47 model initialization: std 0.02, residual init
01:22:18 SECTION 2: Let’s make it fast. GPUs, mixed precision, 1000ms
01:28:14 Tensor Cores, timing the code, TF32 precision, 333ms
01:39:38 float16, gradient scalers, bfloat16, 300ms
01:48:15 torch.compile, Python overhead, kernel fusion, 130ms
02:00:18 flash attention, 96ms
02:06:54 nice/ugly numbers. vocab size 50257 → 50304, 93ms
02:14:55 SECTION 3: hyperpamaters, AdamW, gradient clipping
02:21:06 learning rate scheduler: warmup + cosine decay
02:26:21 batch size schedule, weight decay, FusedAdamW, 90ms
02:34:09 gradient accumulation
02:46:52 distributed data parallel (DDP)
03:10:21 datasets used in GPT-2, GPT-3, FineWeb (EDU)
03:23:10 validation data split, validation loss, sampling revive
03:28:23 evaluation: HellaSwag, starting the run
03:43:05 SECTION 4: results in the morning! GPT-2, GPT-3 repro
03:56:21 shoutout to llm.c, equivalent but faster code in raw C/CUDA
03:59:39 summary, phew, build-nanogpt github repo

Corrections:
I will post all errata and followups to the build-nanogpt GitHub repo (link above)

SuperThanks:
I experimentally enabled them on my channel yesterday. Totally optional and only use if rich. All revenue goes to to supporting my work in AI + Education.