Transcription
okay, so last lecture I gave an overview of language models and what it means to build them from scratch and why we want to do that. Also talked about tokenization, which is going to be the first half of the first assignment.
Um, today's lecture will be going through actually building a model. Um, we'll discuss the primitives in PyTorch that are needed. Um, we're going to start with tensors, build models, optimizers, and training loop. And we're going to place close attention to efficiency, in particular how we're using resources, both um memory and compute.
Okay. So to motivate things a bit, here's some here's some questions. Okay. Um, these questions are going to be answerable by napkin math. So get your napkins out. Um, so how long would it take to train a 70 billion parameter dense transformer model on 15 trillion tokens on 1,24 H100s?
Okay, so you know, I'm just going to sketch out the sort of the give you a flavor of the type of things that we want to do. Okay, so here's how you go about reasoning it. Um, you count the no total number of flops um needed to um you know, train. So that's six times the number of parameters times the number of tokens. Okay. And where does that come from? That will be uh what we'll talk about in this lecture. Um, you can look at the promised number of flops per second that H100 gives you. Um, the MFU, which is something we'll see later, let's just set it to 0.5. Um, and you can look at the number of flops uh per day that your hardware is going to give you at this particular MFU. So 1,024 of them um for um you know, one day, and then you just divide the total number of flops you need to train the model by the number of flops that you're supposed to get. Okay. And that gives you about 144.
Okay. So this is very simple uh calculations at the end of the day. Um, we're going to go through a bit more where uh these numbers uh come from and in particular where the six times number of um parameters times number of tokens comes from. Um, okay, so here's a question: What is the largest model you can train on eight H100s using atom W if you're not being uh too clever?
Okay, so um H100 has 80 gigabytes of HBM uh memory. Um, the number of bytes per parameter that you need for the parameters, the gradients, optimizer state is uh 16, and we'll talk more about where that comes from. Um, and the number of parameters uh is basically total amount of memory divide by number of bytes you need per parameter, and that gives you uh about 40 uh billion parameters.
Okay. And this is very rough because it it doesn't take you into activations which depends on batch size and sequence length, which I'm not really going to talk about but will be important for assignment one. Okay. So this is a rough back envelope calculation. And this is something that you're probably not used to doing. You just implement a model, you train it, and what happens happens. But remember that efficiency is the name of the game. And to be efficient you have to know exactly how many flops you're actually expending because when these numbers get large these directly translate into dollars and uh you want that to be as small as possible.
Okay, so we'll talk more about the the details of how these numbers uh arise. Um, you know, uh we will not actually go over the transformer. So Tatu is going to talk at a um over the conceptual overview of that next time. Um, and there's uh many ways you can learn about transformer if you haven't already uh looked at it. There's assignment one. If you do assignment one you'll definitely know what a transformer is. Um, and the handout actually does a pretty good job of walking through all the different pieces. There's a mathematical description. If you like pictures, there's pictures. There's a lot of stuff you can uh look on online. So but instead I'm going to work with simpler models and really talk about the primitives and the resource accounting piece.
Okay. So remember last time I said what kinds of knowledge can you learn? So mechanics um in this lecture it's going to be just PyTorch and understanding how PyTorch works at um a fairly primitive uh level. So that's will be pretty straightforward. Mindset is about resource accounting and it's not hard. It's just you just have to do it. And intuitions um unfortunately this is just going to be broad strokes uh you know, for now. Actually there's not really much intuition that I'm going to talk about in terms of how um anything we're doing translates to good models. Um, this is more about the mechanics and mindset.
Okay. So let's start with memory accounting. Um, and then I'll talk about compute accounting. Um, and then we'll build up bottom up. Okay. So the best place to start is a tensor. So tensors are the building block of for storing everything in deep learning: Parameters, gradients, optimizer state, uh data, activations. There's sort of these atoms. Um, you can read lots of documentation about them. Um, you're probably very familiar with how to create tensors. There's creating tensors different ways. Um, you can also create a tensor and not initialize it and use some special uh initialization um for the parameters um to um if you want.
Um, okay, so those are um you know, tensors. Um, so let's talk about memory and how much uh memory tensors take up. So every tensor that we'll probably be interested in is stored as a floating-point number. And so there's many ways to represent floating point. So the most default uh way um is float 32. And float 32 has 32 uh bits. They're allocated one for sign, eight for exponent, and 23 for the you know the fraction. So exponent gives you dynamic range and fraction gives you you know uh different basically specifies different values. Um, float 32 is also known as FP32 or single uh you know, precision um is sort of the gold uh standard um in you know, computing. Some people also refer to float 32 as full precision. That's a little bit confusing because full is really depending on who you're talking to. Um, if you're talking to a scientific computing person they'll kind of laugh at you when you say float 32 is uh is really full because they'll use float 64 or even more. Um, but if you're talking to a machine learning person float 32 is the max you'll ever probably need to go because deep learning is is kind of sloppy like that.
Okay. So let's look at the memory. So the memory is it's very simple. It's determined by the number of values you have in your tensor and the data type of each value. Okay, so if you create a torch tensor of 4x8 matrix um the default it will give you a type of uh float 32 um the size is 4 by 8 and the number of elements is 32. Each element size is four bytes. Um, 32 bits is four bytes. Um, and the memory usage is simply the number of elements times the number of um uh size of each element and that'll give you 128 bytes. Okay, so this is should be pretty um easy. And just to give some intuition, if you get the one matrix in the F4 layer of uh GPT3 is this number by this number and that gives you 2.3 gigabytes. Okay, so that's one matrix is you know, can these matrices can be pretty big. Okay, so uh float 32 is the default. Um, but of course these matrices get big. So um naturally you want to make them smaller so you use less memory. And also it turns out if you make them smaller you also uh make it go faster too.
Okay. Um, so another type of uh representation is called float 16 and as the name suggests it's uh 16 bits um where both the exponent and the fraction are um shrunk down from um 8 to 5 and 23 to 10. Okay, so this is known as half precision and it cuts down half the memory. Um, and that's all great except for the dynamic range uh for these float 16 isn't great. So for example, if you try to make a number um like uh 10e1e-8 in float 16 um it uh basically rounds down to zero and you get underflow.
Okay. So uh the float 16 is not great for representing um very small numbers or or very big numbers as a matter of fact. Um, so if you use float 16 u for training for small models it's probably going to be okay, but for large models when you're um having lots of matrices and you can get uh instability or underflow or overflow and bad things happen.
Um, okay, so um one um thing that has happened which is nice is there's been another representation of BF16 which stands for brain float. um, this was developed in 2018 to address the issue that for deep learning um we actually care about dynamic range more than uh we care about this fraction. So basically BF16 allocates more to the exponent and less to the fraction.
Okay. So it uses the same memory as um floating point 16 but it has a dynamic range of float 32. Okay. So that sounds really good. Um, and it actually uh the catch is that this resolution which is determined by the fraction is worse. Uh, but this doesn't matter as much uh for deep learning. So now if you try to create a a tensor with one E minus 8 in BF16 then you get something that's not zero.
Okay. So um you can dive into the details, I'm not going to go into this but you can stare at the actual full specs of all the different floating point you know operations. Okay. Okay, so BF16 is basically what uh you will typically use to do computations because um it's it's sort of you know good enough um for fear computations. It turns out that um for storing optimizer states and parameters you still need float 32 um for otherwise your training will go um haywire.
Um, so if you're bold so now we have um something called FP8 or uh 8-bit and as the name suggests um this was developed in 2022 um um by you know Nvidia so now they have um essentially if you look at FP and BF16 it's like this and FP wow you really don't have that many bits to store stuff right so it's very crude. Um, there's there's two sort of variants depending on if you want to have more resolution or more um dynamic range. Um, and um I'm not going to say too much about this but FP8 is you know supported by H100. It's not really available on previous generation. Um, but but at a high level you know training with float 32 which is I think you know is is what you would do if you're not trying to optimize you know too much and it's it's sort of safe. It requires more memory. Um, you can go down to FP8 or BF16. Um, and but you can get some instability. Um, basically I don't think you would probably want to use uh a float 16 at this point for for deep learning. And um you can become more sophisticated by looking at particular places in your pipeline either forward pass or backward pass or optimizers or gradient accumulation and really figure out what the minimum precision you need at this particular places. And that's called gets into kind of mixed precision training. So for example, some people like to use float 32 for the you know the attention to make sure that doesn't kind of uh you know get messed up of what four simple F4 passes with Matt Mall's BF-16 is is is fine.
Okay, pause a bit for for questions. So we talked about tensors and we looked at uh depending on how what representation how much storage they take. Yeah. Can you just clarify about the mix position like when you would use 32 in the be?
Yeah. So the question is when would you use um uh float 32 or BF16. Um, I don't have time to get into the you know the exact uh details and it sort of varies depending on the model size and everything but um generally for the the parameters and optimizer states you use float 32. You can think about BF16 as something that's more transitory like you you basically take your parameters you cast it to BF16 and you kind of run ahead with that model but then the thing that you're going to accumulate over time you want to have higher precision.
Yeah. Okay. So now let's talk about uh compute. So that was memory. Um, so uh compute obviously depends on what the hardware you know is. By default tensors are stored in CPU. So for example, if you just in PyTorch um say X equals torch.0 is 3232 then it'll put it on your CPU, it'll be in um the CPU memory. Um, now of course that's no good because if you're not using your GPU then you're going to be orders of magnitude too slow. So you need to explicitly say in PyTorch that you need to move it to the GPU, and this is uh you it's actually just to make it very clear in pictures there's a CPU it has RAM and that um has to be moved over um to the the GPU, there's a data transfer um which is is c um which takes some work um takes some time. okay, so whenever you have a tensor in PyTorch you should always keep in your mind where is this residing because just looking at the variable or just looking at the code you can't always tell and if you want to be careful about computation and data movement you have to really know where it is. um, you can probably do things like assert um where it is in various places of code just to um document or be sure.
Okay. So um so let's look at uh what hardware we have. So we have um in this case we have one GPU. Um, this was run on the H100 clusters that you guys have access to. And um this um GPU is a you know H100, 80 GB of high bandwidth memory um and there's a you know it gives you the cache size and so on. Okay. So um so if you have remember the x is on uh CPU um you can move it uh just by specifying you know two which is a kind of a general PyTorch function. Um, you can also create a tensor directly on a GPU so you don't have to move it at all. Um, and if everything goes well I'm looking at the memory allocated before and after. Uh, the difference should be exactly uh two 30 by 32x32 um matrices of four byte floats. Okay. So um it's a192. Okay. Okay, so this is just a sanity check that um the code is is uh is doing what is advertised. Okay, so now you have your answers on the GPU. Um, what do you uh do? So there's many operations um that you'll be needing for assignment one and in general to do any deep learning application. Um, and most tensors you just create by performing operations on other tensors and each operations has some memory and compute uh footprint. So let's make sure we understand uh that.
Um, so first of all um what is actually a tensor in PyTorch right? tensors are a mathematical object in PyTorch, they're actually pointers into some allocated memory. Okay. So if you have let's say um a 4x4 matrix um what it actually looks like is a long array and what the tensor has is metadata that specifies how to get to address into that array. And the metadata is going to be um you know two numbers, a stride for each or actually number of one number per dimension of the of the tensor in this case because there's um you know two dimensions it's uh you know stride zero and stride one. um, stride zero uh specifies if you were in dimension zero to get to the next row to increment that index how many do you have to skip and so going down the rows you skip four so stride uh zero is four and um to go to the next column um you skip one. So stride one is one. Okay. So with that to find an element let's say one two one comma two um it's simply just um multiply the indexes by the stride and you get to your index which is uh six here. So that would be here or or here. Okay. So that's basically what's going underneath the hood um for for for tensors.
Okay, so this is um relevant because um you can have multiple tensors that use the same storage and this is useful because you don't want to copy the tensor all over the all over the place. So imagine you have a 2x3 matrix here. Um, many operations um don't actually create a new tensor. They just create a different view and doesn't make a copy. So you have to keep make sure that um um your mutations uh if you start mutating one tensor it's going to cause the other one to mutate. Okay. So um for example if you just get um row zero okay, so remember y is uh this tensor and sorry x is 1 2 3 4 5 6 and y is x0 which is just the first row. Okay. And um you can sort of double check. There's this you know function I wrote that says if you look at the underlying storage whether these two tensors have the same storage or not. Okay. So this um definitely doesn't copy the tensor. It just creates a view. Um, you can get column one. Um, this also doesn't uh copy the tensor. Um, oops, don't need to do that. Um, you can call a view function which can take any tensor and uh ch look at it in terms of the different dimensions. Um, 2x3 uh um actually this should be maybe the other way around Uh as a 3x2 tensor um so that's also doesn't you know change uh do any copying. Um, you can transpose um that also doesn't copy. Um, and then you know like I said if you start mutating X um then Y actually gets mutated as well because X and Y are just pointers into the same underlying storage.
Okay. So things are um one thing that you have to be careful of is that some views are contiguous which means that if you run through the tensor it's like just slide going through the the uh array in your storage um but some are not. So in particular if you transpose it now you're you know what does it mean when you're you're transposing it you're you're sort of going down now so you're kind of if you imagine going through the tensor you're kind of skipping around and if you have a non-contiguous tensor um then if you try to further view it in a different way then this is not going to work. Okay. So in some cases if you have a non-contiguous tensor you can make it contiguous first and then uh you can apply whatever oper viewing operation you want to it and then in this case x and y um do not have the same storage because contiguous in this case uh makes a copy.
Okay. So um this is just ways of slicing and dicing a tensor. Um, views are free so feel free to use them define different variables um to make your it sort of easier to read your code um because they're not allocating any memory. Um, but you know remember that contiguous or reshape which is basically contiguous view um uh can uh create a copy and so just be careful what you're doing. Okay, questions before moving on?
All right. So um hopefully some a lot of this will be review for those of you have uh kind of done a lot of PyTorch before but it's uh helpful to just do it systematically make sure we're on the same page. So um here are some operations that uh do create new tensors and in particular elementwise operations all create new tensors obviously because you need a somewhere else to store the new value. Um, there's a you know triangular U is also element operation that uh comes in handy when you want to create a causal attention mask um which you'll need for your um your assignment. Um, but nothing is interesting that interesting uh here.
Um, okay, so let's talk about map malls. So the bread and butter of deep learning is matrix multiplications. And I'm sure all of you have done a matrix multiplication but just in case this is what it looks like. Um, you take a 16x32 times a 32x2 matrix you get a 16x2 matrix. Um, and but in general when we do you know our machine learning application um all operations are you want to do in a batch and in the case of language models this usually means um for every example in a batch and for every sequence in a batch you want to do something. Okay, so generally what you're going to have in instead of just a matrix is you're going to have a tensor where the dimensions are typically batch, sequence, and then whatever thing you're you're trying to do. In this case it's a matrix for every um token in your in your data set. And so you know PyTorch is nice enough to make this uh work well for you. So when you take this for um the dimension tensor and this matrix um what actually ends up happening is that for every batch every example in every token you're multiplying these two um matrices. Okay. And then the result is that you get your your resulting matrix for each of the first two elements. So this is just like a there's nothing uh fancy going on but this is just a pattern that um I think uh you is helpful to you know think about.
Okay, so I'm going to take a little bit of a a digression and talk about um inops. Um, and so the motivation fors is the following: so normally you in PyTorch you define some um you know tensors. And then you see stuff like this where you take x and multiply by y transpose minus 2 minus 1 and you're you kind of look at this and you say okay what is minus 2 um well I think that's uh this sequence and then minus 1 is this hidden because you're indexing backwards. Um, and it's really easy to mess this up because if you look at your code and you see minus one minus two um you're kind of if you're you're good you write a bunch of comments but then the comments are can get out of date with the code and then um you you have a bad time debugging. So the the solution is to use inops here. So this is inspired by Einstein's summation notation. Um, and the idea is that we're just going to name all the dimensions um instead of you know relying on indices essentially. Okay. So there's a library called Jack typing which um is is is helpful for um as a as a way to specify the dimensions in the types. So normally in PyTorch you would just define write your code and then you would comment oh here's what the dimensions would be. So if you use jacks typing then um you have this notation where as a string you just write down what the dimensions are. So this is a slightly kind of more natural way of documenting.
Now notice that there's no enforcement here right because PyTorch types are sort of a a little bit of a lie in in in PyTorch. Uh, so it can be enforced. You you can use a checker right? Yeah, you can write a check but um not by default. Yeah.
Um, okay. So let's look at uh the you know Ein sum. So Ein sum is basically matrix multiplication on steroids with good bookkeeping. Um, so here's our example here. We have X which is let's just think about this as uh you have a batch dimension, you have a sequence dimension, and you have uh four um hiddens, and Y is the same size. um, you originally had to do this thing. Um, and now what you do instead is you basically write down the the dimensions uh names of the dimensions of the two uh tensors. So batch sequence one hidden batch sequence two hidden and you just write what you dimension should appear in the output. Okay. So I write batch here because I just want to uh basically uh you know carry that over and then I write seek one and seek two and notice that I don't write hidden and any dimension that is not named in output is just summed over and any dimension that is um named is sort of just iterated over. Okay. So um once you get used to this this is actually very very uh you know helpful. It may maybe looks if you're seeing this for the first time it might seem a bit you know strange and long but u trust me once you get used to it'll be better than um doing minus 2 minus 1.
Um, if you're a little bit uh you know slicker you can um use dot dot dot to represent broadcasting over any number of dimensions. So in this case um uh instead of writing batch I can just write dot dot.
dot and this would handle the case where instead of maybe batch, I have batch one, batch two, or or some other uh arbitrary long sequence. Yeah, question: does for compile this like is it guaranteed to compile to like I guess so the question is: is it guaranteed to compile to something efficient? Um this uh I think the short answer is yes. Um I don't know if you have any, you know, nuances. We will figure out the best way to reduce the best order of uh dimensions to reduce and then use that. If you're using it within portra compile, only do that one time and then you know reuse the same implementation over and over again, better than anything designed by hand. Yeah. Okay.
So um so let's look at reduce. So reduce operates on one tensor and it basically aggregates some dimension or dimensions of the tensor. So you have this tensor uh before you would write mean to sum over the the the final dimension and now you basically say um actually okay so this replace this with sum. So reduce and um again you say hidden and hidden is disappear disappeared. So which means that you are um aggregating over that dimension. Okay. So you can check that this indeed kind of works over here. Okay. So so maybe one final example of uh this is um sometimes um in a you know tensor one dimension actually represents multiple dimensions and you want to unpack that and operate over one of them and pack it back. So um in this case uh let's say you have batch uh sequence and then this eight-dimensional vector is actually a flattened representation of number of heads times some you know hidden dimension. Okay, so and then you have a way a vector that is um needs to operate on that hidden dimension. So you can do this very elegantly uh using um inops uh by um re calling rearrange and this basically you can think about it, we saw view uh before, it's kind of like um kind of a you know fancier version which basically looks at the same data but uh um you know differently. Um so here uh it basically says this dimension is actually heads in hidden one, I'm going to explode that into two dimensions. Um and uh you have to specify the the number of heads here because there's multiple ways to split a number into you know two. Um let's see. This might be a little bit long. Um okay, maybe it's not worth looking at right now. Um and given that x you can uh perform your transformation using line sum. So this is something hidden one which corresponds to x and then hidden one hidden two which corresponds to w and that gives you um something hidden too. Okay. Um and then you can rearrange back. So this is the just the inverse of breaking up. So you have the your two dimensions and you group it into one. So that's just a flattening operation that's uh um you know with uh everything all the other dimensions kind of left alone. Okay. So there is a tutorial um for for this that I would recommend you go through and it gives you a bit more. So um on you don't have to use this um because you're building it from scratch. So you can kind of do anything you want but um in assignment one we do give you guidance and it's something probably to invest in. Okay.
Um so now let's talk about uh computation uh no cost of t tensor operations. So we introduce a bunch of operations um and you know how much do they cost? So a floating-point operation is uh any operation floating point like addition or multiplication. These are them and um these are kind of the you know main ones that are going to um I think matter in terms of flop count. Um one thing that is uh is sort of a pet peeve of mine is that when you say flops, it's actually unclear what you mean. So you could mean flops with a lowercase s which stands for number of floating operations. This is measures amount of computation that you've done or you could mean flops um also written with a uppercase s which means floating points per second which is used to measure the speed of hardware. So um we're not going to in this class use uppercase s because I find that very confusing and just write /s to denote that this is floating point per second. Okay. Okay. So just to give you uh some intuition about flops, Uh GBD3 took about uh 323 flops, GPD4 was 2 E25 flops um speculation and there was a US executive order that any foundation model with over 126 flops had to be reported to government um which now u has been revoked um but the EU has still they're going u still has a something that's hasn't the EU AI act which is 1e25 which hasn't been revoked um so uh you know some intuitions uh A100 has a peak performance of 312 teraflop per second. Um and um H100 has a peak performance of um 1979 teraflop per second with uh sparsity and approximately 50% without um and um if you look at uh you know the Nvidia has these specification sheets um so you can see that um the the flops actually depends on what you're trying to do. So if you're using BF or F FP32 it's actually really really bad. Like the if you run FP32 on H100 you're not getting um it's it's orders of magnitude worse than if you're doing um FP16 or um and if you're willing to go down to FP8 then it can be even faster. And you know for the for when I first read it I didn't realize but there's an asterisk here and this means with sparsity. So uh usually you're in a lot of the m matrices we have in this class are dense. So you don't actually get this. You get something like uh you know it's exactly half that number. Exactly half. Okay. Okay. So now you can do a bucket of envelope of calculations: eight H100s for two weeks is um just you know eight times the number of flops per second times the number of seconds in a in a in a week. Um actually this is this might be one week. Okay, so that's one week and that's uh 4.7 * e to the 21 um which is you know some number and you can kind of contextualize the the flop counts with other um model counts. Yeah. Sparity mean? So that means if so what does sparsity mean? That means if your matrices are uh sparse, it's a specific like structured sparity. It's like two out of four elements in each like group of four elements is zero. That's the only case if you get that that speed. No one uses it. Yeah, it's a marketing department uses it. Okay.
So um let's go through a simple example. So remember we're not going to touch the transformer but I think even a linear model gives us a lot of the building blocks and intuitions. So suppose we have n points, each point is d-dimensional and the linear model is just going to match map each dimensional vector to a k dimensional uh vector. Okay. So let's set some number of points um is B, dimension is D, K is the number of outputs. Um and let's create our data matrix X, um our weight matrix uh uh W and the linear model is just some map. So nothing you know too too interesting going on. And um you know the question is how many flops was that? And uh the way you would you look at this is you say well um when you do the matrix multiplication you have basically for every j k triple I have to multiply two numbers together and I also have to add that number to the the the total. Okay. So the answer is 2 times um the basically the product of all the dimensions involved. So the the left dimension, the middle dimension and the right dimension. Okay, so this is something that you should just kind of remember if you're doing a matrix multiplication. The number of flops is two times um the product of the three dimensions. Okay, so the flops of other operations are usually kind of linear in the size of um the the matrix or tensor. Um and in general no other operation you encounter in deep learning is expensive as matrix multiplication for large enough uh matrices. So this is why I think a lot of the napkin math is very simple because we're only looking at the matrix multiplications that are going are performed by the model. Um now of course there are regimes where if your matrices are um small enough then the cost of other things starts to dominate but generally that's not a good regime you want to be in because the hardware is designed for big much multiplication. So um sort of by it's a little bit circular but by kind of we end up in this regime where we only consider uh models where the mammals are the dominant um uh you know cost. Okay, any questions about this this number two times the product of the three dimensions? This is just a useful thing. Would the algorithm of matrix motivation always be the same because the chip might have optimized that they always the same? Yeah. So the question is like uh is the does this essentially does this depend on the the matrix multiplication you know algorithm? Um in general I guess we'll look at this the next week when we or the week after when we look at kernels. I mean actually there's a lot of optimization that goes underneath under the hood when it comes to matrix uh multiplications. Um and there's a lot of specialization depending on the shape. Um so this is I would say this is just a kind of a crude you know estimate um that is ba basically like the right order of magnitude. Okay. So um yeah uh additions and multiplications are considered equivalent. Yeah, additions and multiplications are considered are equivalent. So um one way I find helpful to interpret this so at the end of the day this is just a matrix multiplication but I'm going to try to give a little bit of meaning um to this um which is why I've set up this as kind of a little toy machine learning problem. So B is really stands for the number of data points and DK is the number of parameters. So for this particular model the number of flops that's required for forward pass is two times the number of tokens or number of data points times the number of parameters. Okay. So this turns out to actually generalize to transformers. There's an asterisk there because uh um there there's you know the sequence length and other stuff but this is roughly right for uh if your sequent length isn't isn't too large. Um so okay so now this is just a number of floating-point operations right um so how does this actually translate to a walk like time which is presumably the thing you actually care about how long do you have to wait for your run um so let's time this so I have this function that is just going to um um do it five times and I'm going to perform the matrix multiply operation. Um we'll talk a little bit later about this um um two weeks from now why the other code is here. Um but um for now we get an actual time. So that matrix uh took you know 0.16 seconds um and the actual flops per second which is how many flops did it do per second is uh 5.4 E13. Okay. So now you can compare this with you know the marketing materials and um for the A100 and H100 um and you know as we look at the spec sheet the the flops depends on the data type and we see that the the promise flops per second which um you know for H100 for I guess this is for float 32 is you know 67 you know teraflops as we looked and so that is the number of promise flops per second we we had um and now if you look at the uh there's a helpful notion called model flops utilization or MFU which is the actual number of flops divided by the promise flops okay so you take the actual number of flops remember which what you actually uh you know witnessed the number of floating-point operations that that are useful for your model um uh divided by the actual time it took divided by this promise flops per second which is you know from the glossy brochure um you can get a MFU of uh 0.8 eight. Okay. So usually you see people talking about their MFUs and um something greater than 0.5 is you usually considered to be good and if you're like you know 5% MFU that's uh considered to be really bad, usually can't get you know close to that close to you know uh you know 90 or 100 um because this is sort of ignoring all sort of communication and overhead. It's just like the literal uh computation of the flops. Okay. And usually MFU is much higher if the magic matrix multiplications dominate. Um Um okay. So that's MFU. Any any questions about uh this? Yeah. Uh you're using the promise flop per sec uh not considering the sparse. So this promised flop per sec is uh not considering this as much. Yeah. Um one a note is like uh this is actually a you know there's also something called hardware you know flops uization and the the inov the motivation here is that we are all we're trying to look at the it's it's called model because we're looking at the number of effective useful operations that the model is you know performing okay. And uh so it's a way of kind of standardizing. It's not the actual number of flops that are are done because you could have optimization in your code that cache a few things or redo um you know um recomputation of some things and in some sense you're still computing the same model. So what matters is that you're this is sort of trying to look at the model complexity and you shouldn't be penalized just because you were clever in your MFU if you were clever and you didn't actually do the flops um but you you said you did. Okay. So you can also do the same with BF16. Um Oops. And here we see that for uh BF the um the time is actually much uh better right so 03 um instead of .16 so the actual flops per second is higher um the even accounting for sparsity the the promise flops is uh still quite high so the MFU actually um lower for BF uh 16. Um this is you know maybe surprisingly low but um but sometimes the promise flops is a bit of a you know optimistic. So always benchmark your code uh and don't just uh kind of assume that you're going to get certain um levels of performance. Okay. So uh just to summarize um matrix multiplications u dominate the the compute and the general rule of the thumb is that it's two times the product of the dimensions flops. Um the flops per second floating points per second depends on you know the hardware and also the data type. So the fancier the hardware you have the higher it is, the the smaller the data type the usually the faster it is. Um and uh MFU is a useful notion to look at how well you've you're essentially squeezing you know your your hardware. Yeah. I I've heard that often uh in order to get like the maximum utilization you want to use these like tensor cores on the machine. And so like does PyTorch by default use these tensor cores and like are these uh accounting for that? Yeah. So the question is uh what about those tensor cores? So if you go to this um spec sheet um you'll see that um you know these are all on the tensor core. So the tensor core is basically you know you know specialized hardware to do um map models. Um so if you are um you know if you're so by default it should use it and if you especially if you're using PyTorch you know compile it will generate the the code that will um use the the hardware properly. Okay.
So um let's talk a little about you know gradients. So and the the reason is that we've only looked at matrix multiplication or in other words basically feed forward um forward passes and the number of flops. Um but there's also a computation that comes from computing gradients and we want to track down how much uh that is. Um okay so just to consider a simple example, simple linear model um where you take the uh the prediction of a linear model and you um you look at the MSE with respect to five though not a very interesting loss but I think it'll it's illustrative for uh looking at the the gradients. Okay. So remember in the forward pass you have your x you have your your w which um you want to compute the gradient with respect to you make a prediction um by taking a linear product and then you you have your loss. Okay. And in the backward pass uh you just call loss backwards. And um in this case the the gradient which is this variable attached to the tensor is um turns out to be what you what you want. Okay. So everyone has done uh you know gradients in PyTorch before. Um so let's look at how many flops uh are required for computing gradients. Um okay. So let's uh look at a slightly more complicated model. So now it's uh a two-layer linear model where you have X which is B by D times um W1 which is D by D. So that's the first layer and then you take the your hidden activations um H1 and you pass it through another linear layer W2 and to get a K dimensional vector and you do some compute some loss. Okay, so this is a two-layer linear network. And just as a kind of review if you look at the number of forward flops um what you had to do was um you have to multiply look at W1. You have to multiply X by W1 um and add it to your H1 and you have to take H1 and W2 and you have to add it to your your H2. Okay, so the no total number of flops again is 2 * the the product of all the dimensions in your map mall plus two times the product dimensions your m all for the second matrix. Okay, in other words two times the total number of parameters in this case. Okay, so so what about the backward pass? So this part will be um a little bit more involved. Um so we can recall the model uh x to h1 to h2 and the loss. So in the backward path you have to compute a bunch of gradients and the gradients that are relevant is you have to compute uh the gradient with respect to um h1 um you know h2 w1 and w2 of of the loss. So D loss D each of these variables. Okay. So how long does it take to compute that? Let's just look at W2 for now. Okay. So the things that touch W2 um you can compute by looking at the the chain rule. Um so W2 grad. So the gradient of d loss dw2 is um you you sum um h1 times uh the gradient of the loss with respect to h2. Okay. So that's just a chain rule for um for w2. And this is um so all the gradients are the you know the same size as the underlying you know of vectors. Um so this turns out to be um essentially looks like a matrix uh you know multiplication and so the same you know calculus holds which is that it's 2 times the number of uh the product of all the dimensions B * D * K okay but this is only the gradient with respect to uh W2. We also need to compute the gradient with respect to h1 because we have to keep on back propagating um to w1 and and so on. Okay. So that is uh going to be um you know the product of w2 um you know times uh you know h2 sorry I think this should be that grad of h2 h2.grad red. Um so that turns out to also be essentially you know um looks like the matrix multiplication and it's the same number of flops for for computing the gradient of uh each one. Okay. So when you add the two uh so that's just for w2. Um you do the same thing for W1 and that's uh which has D * D parameters and when you add it all uh up it's uh so for this um for W2 the amount of computation was 4 * B * D * K and for W1 it's also 4 * B * D * uh D because W1 is D by. Okay. So um I know there's a lot of symbols here I'm going to try also to uh give you a visual account for this. So this is from a a blog post that I think um may work uh better. We'll see. Um okay I have to wait for the animation to loop back. So so basically this is uh one layer of the neuron now. It has um you know the hiddens and then the weights to the next layer. And so I have to Okay the problem with this animation is I have to wait. Okay, ready set. Okay, so first I have to multiply W and A and I have to add it to this. That's a forward pass. And now I'm going to multiply this these two and then add it to to that. And I'm going to multiply and then add it to that. Okay. Any questions? Which a way to slow this down. Um but you know the details maybe I'll I'll let you kind of ruminate on but um the the high level is that there's two times the number of parameters for the forward pass and four times the number of parameters for the backward pass. And um we can just kind of work it out uh via the chain row here. Yeah. For the homeworks are we also using the you said some PyTorch implementation is allowed some isn't. Are we allowed to use docu or we are doing the like entirely by hand doing the gradient? Uh so the question is in the homework are you going to compute gradients by hand? And the answer is no. Uh you're going to just use pietorch gradient. This is uh just to break it down so we can do the the um counting flops. Okay. Any any questions about this before I move on? Okay. Just to summarize the forward pass is for this particular model is 2 times the number of data points times the number of parameters and backward is four times the number of data points times the number of parameters which means that total it's six times number of data points times parameters okay and that's explains why uh there was that six in the beginning when I asked the motivating uh question. So now this is for a simple you know linear model. Um it turns out that mo many models this is uh basically the bulk of a computation um when essentially every computation you do has uh you know touches essentially a new parameters roughly right and you know obviously this doesn't hold you can find models where this doesn't hold because you can have like one parameter through parameter sharing um and have you know a billion flops but that's generally what not what models look like. Okay, so let me um move on. Um so far I've I've basically finished talking about the the resource accounting. So we looked at tensors. We looked at some computation on tensors. We looked at how uh much tensors take to store and also how many flops takes tensors uh take when you do various operations on them. Um now let's start building up uh different you know models. Um I think this part isn't necessarily going to be um you know that uh you know conceptually you know interesting or challenging but it's more for maybe just you know completeness. Okay. So um so parameters in PyTorch are stored as these nn parameter um objects. Um let's talk a little bit about parameter initialization. Um so if you have let's say a uh um you know parameter that has uh okay so you you generate um a okay do your sorry your w parameter is an input dimension by hidden dimension matrix um you're still in the linear model case so let's just generate an input and let's feed it through the output okay so rand and unit no gausian is uh you know seems innocuous. Um what happens when you do this is that if you look at the output um you get some pretty large numbers right and this is because when you you know have the the number uh grows as essentially the square root of the the hidden dimension. Um and so when you have large models uh this is going to you know blow up and um and training can be very unstable. Um so so typically what you want to do is initialize in a way that's um you know invariant to hidden or at least you know when you you're guaranteed that it's not going to you know blow up. And one simple way to do this is just rescale
By the one over square root of number of, you know, of inputs. Um, so basically, let's redo this: w equals, you know, parameter where I simply divide by the square root of the input dimension. Um, and then now when you uh feed it through the output, now you get things that are stable around, you know, this will actually concentrate to um, you know, something like normal 0, 1. Okay, so this is basically, you know, this has been explored pretty extensively in deep learning literature; it's known uh up to a constant as Xavier initialization, and typically I guess it's uh fairly common if you want to be extra safe, uh you don't trust the normal because it doesn't have it has unbounded tails, and you just say I'm going to truncate to minus 3, 3, so I don't get any large values and I don't want any to to mess with that. Okay.
Um, Okay. So um, so let's build a, you know, just a simple model. Um, it's going to have um dimensions and two layers. There's this, you know, I just made up this name cruncher. It's a custom model which is a deep linear network which has um num layers layers, and each layer is a a linear um model which has essentially just a matrix multiplication. Okay.
Um, so um, this parameters of these this model is looks like um I have um um layers for the the first layer um which is a D by D matrix, uh the second layer which is also a D by D matrix, and then I have a um a head um or a final layer. Okay. So if I get the number of of parameters of this model, then it's going to be um d^2 + d^2 + d. Okay, so nothing too surprising there. Um, and I'm going to move it to the GPU because I run want this to run one fast. And um, I'm going to generate some random data and feed it through the data. And the forward pass is just going through um the layers um and then finally applying the head. Um, okay. So with that model, I let's try to I'm going to use this model and do some, you know, stuff with it. Um, but just one kind of general digression.
Um, randomness is something that uh is is sort of can be annoying in some cases if you're trying to reproduce a bug, for example. It shows up in many places: initialization, dropout, data ordering. Um, and just a best practice is I we recommend you always pass a fix a random seed. Um, so you can reproduce your your model, or at least as well as you can. Um, and in particular, having a difference random seed for every source of randomness is is you know nice because then you can, for example, fix initialization or fix the data ordering but vary other other things. Um, determinism is your your friend when you're debugging. Um, and you know in code unfortunately there's you know many places where you can use randomness and just be uh cognizant of you know which one you're using and just if you want to be safe just uh set the seed for all of them.
Um, data loading um I guess I'll go through this quickly. It's it's not um it'll be useful for your your assignment. Um, so in language modeling, data is typically just a sequence of integers because this is remember output by the tokenizer. Um, and you serialize them into um you can serialize them into NumPy arrays. Um, and one I guess thing that's maybe useful is that uh you don't want to load all your data into memory at at once because, for example, the Llama data is the 2.8 8 terabytes. Um, but you can sort of pretend to load it by using this uh handy function called memmap. So which gives you essentially a variable that is mapped to a file. So when you try to access um the data, it actually on the on demand loads the file, and then using that you can uh create a data loader that um you know is a you know samples data from your um batch. So I'm going to skip over that just in interest of time.
Um, let's talk a little bit about, you know, optimizer. So we've defined our model. Um, so there's many optimizers. Um, just kind of maybe going through the intuitions behind some of them. So of course there's stochastic gradient descent. You compute the gradient of your batch. You take a step in that direction. No questions asked. Um, there's an idea called momentum which dates back to classic optimization nest um where you uh have a running average of your um gradients and you update against the you know the running average instead of your instantaneous you know gradient. Um, and then you have um Adagrad which um you you know you scale the gradients by uh your um your the average over the the norms of your or I guess not the norms the the square of the gradients. You also have RMSprop which uh is an improved version of Adagrad which uses an exponential averaging rather than just like a flat average. And then finally Adam, which appeared in 2014, which is essentially combining RMSprop and momentum. So that's why you're maintaining both uh your running average of your gradients but also running average of your gradient squared. Okay. So since you're going to implement uh Adam um in homework one, I'm not going to do that; instead I'm going to implement um, you know, Adagrad. So so the way you implement an optimizer in in you know PyTorch is that you override the the optimizer class, and you have to um uh let's see maybe I'll and I'll get to the implementation once we step through it.
Um, so let's define some data. Um, compute the forward pass on the loss, and then you compute the gradients, and then you when you call optimizer.step, um this is where the optimizer uh actually is active. So what this looks like is um your parameters are grouped u by, for example, you have one for the layer zero, layer one, and then the final you know weights, um and you can access a state which is a a dictionary from parameters to um you know whatever you want to store as optimizer state, um the gradient of that parameter you assume is already uh calculated um by the the backward pass. Um, and now you can do things like um you know in in Adagrad you're storing the sum of the gradient uh squares. So you can get that G2 variable, um and you can update that based on the square of the gradient. So this is element-wise squaring of the um the gradient, and you put it back into the state. Okay. So then your obviously your optimizer is responsible for updating the parameters, and this is just the you know you update the learning rate times the gradient divided by this uh scaling. So now this state is kept over across multiple uh invocations of um you know the optimizer. Okay. So um, and then at the you know end of your optimizer step you can you know free up the memory uh just to um which is I think going to actually be more important when you look when we talk about model parallelism.
Okay, so let's talk about the memory requirements of uh the optimizer states and actually basically at this point everything. So um you need to uh the number of parameters in this model is d^2 times the number of layers plus d for the final uh head. Okay. Um, the number of activations. So this is something we didn't do before, but now for this simple model it's fairly easy to do. It's just um B times you know D times the the number of of layers you have for every layer for every data point for every dimension you have to hold the activations. Um, for the gradients um this is the same as the number of parameters, and the number of optimizer states in for Adagrad it's uh you remember we had to store um the the gradient squared. So that's another copy of the parameters. So putting all together we have um the total memory is assuming you know FP32 which means four bytes times the number of parameters, number of activations, number of gradients, and number of optimizer states. Okay. And that gives us uh you know some number which um is 496 here. Okay. So this is um a fairly simple calculation. In the assignment one, Um you're going to do this for the transformer which is a little bit more involved because you have to there's not just matrix multiplications but there's um many matrices, there's attention, and there's all these other things, but the general form of the calculation is the same. You have parameters, activations, gradients, and optimizer states. Okay. And the so and the FLOPS required again for this model is six times uh the number of tokens or number of data points times the number of parameters, and um you know that's basically concludes the resource uh accounting for this particular model. Um, and uh if for reference if you're curious about working this out for for transformers you can um consult um some of these articles.
Okay. So in the remaining time, I think um maybe I'll pause for questions. Um, and we talked about building up the tensors, and then we built a kind of a very small model, and you know we talked about optimization and how many how much memory and how much uh compute was required. Yeah. So the question is why do you need to store the activations? So naively you need to store the activations because when you're uh when you're doing the backward pass the gradients um of let's say the first layer depend on the activation. So the gradients of the I layer depends on the activation there. Um, now if you're smarter you don't have to store the activations or you don't have to store all of them. You can recompute them, and that's something a technique we'll uh called activation checkpointing which we're going to talk about later.
Okay, so let's just do this quick uh you know actually there's not much to say here, but you know here's your typical you know training loop um where you define the model, define the optimizer, and you get the data, um feed forward, backward, and take a step in a parameter space. Um, And um I guess it'll be more interesting I I guess next time I should show like actual one plug which I isn't available on this uh on this version, but um so one note about checkpointing. Um, so training language models takes a long time, and you certainly will crash at some point, so you don't want to lose all your progress. Um, so you want to periodically save your model to disk, and just to be very clear the thing you want to save is both the model and uh the optimizer and probably you know which iteration you're on. Um, I should add that, and then you can just load it up.
One maybe final note, and I'll uh end is um I alluded to kind of mixed precision your training. Um, you know, choice of the data type has the different trade-offs. If you have higher precision, it's uh more accurate and stable, um but it's more expensive, and lower precision vice versa. Um, and as we mentioned before, by default the recommendation is use float 32, um but try to use BF16 or even FP8 whenever possible. So you can use um lower precision for the forward pass um but float 32 for the rest. Um, and this is an idea that goes back to the you know 2017; there's exploring mixed precision training. Um, PyTorch has um some tools that automatically allow you to do you know mixed precision training, um because it can be sort of annoying to have to specify which parts of your model it needs to be you know what precision. Um, generally you define your model as you know sort of this clean modular thing, and and the specifying the precision is sort of like a you know something that needs to cut across that. Um, and one I guess maybe one kind of general comment is that people are pushing the envelope on what precision is needed; there's some you know papers that show you can actually use um FP8 you know all the way you know through. Um, there's um I guess one of the challenges is of course when you have lower precision it gets very numerically unstable, but then you can do various tricks to you know control the uh the numeric of your model during training so that you don't get into these you know bad regimes. So this is where I think um the systems and the model architecture design kind of are synergistic because you want to design models now that we have a lot of model design is just governed by hardware. So even the transformer as we mentioned last time is governed by the having GPUs, and now if we notice that you know Nvidia chips have the property that if um lower precision even like int4 for example is one thing. Now if you can make your model training actually work on int4 which is I think you know quite quite hard, then you can get massive you know speed ups and your model will be more you know efficient. Um, now there's another uh thing which we'll talk about later which is you know often you'll train your model using more sane you know floating point, but when it comes to inference you can go crazy and you take your trained model and then you can quantize it and get a lot of the gains from very very uh aggressive quantization. So somehow training is a lot um more difficult to do with low precision, but once you have a trained model it's much easier to um make it low precision. Okay.
So um I will wrap up there. Just to conclude, um we have talked about the different primitives to use to train a model, building up from tensors all the way to the training loop. We talked about uh memory accounting and FLOPS accounting um for these simple models. Um, hopefully once you go through assignment one, all these concepts will be really solid because you'll be um uh applying these ideas for the actual transformer. Okay, see you next time.