BigScience BLOOM | 3D Parallelism Explained | Large Language Models | ML Coding Series
❤️ Become The AI Epiphany Patreon ❤️
https://www.patreon.com/theaiepiphany

👨‍👩‍👧‍👦 Join our Discord community 👨‍👩‍👧‍👦
https://discord.gg/peBrCpheKE

In this 4th video of the Large Language Model series I walk you through the BigScience's BLOOM model codebase!

The main focus is on understanding the 3D parallelism:
* Pipeline parallelism
* Model parallelism
* Data parallelism

A set of beautiful engineering ideas that are behind all of the recent scaling efforts and ML success stories!

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
✅ BLOOM code: https://github.com/bigscience-workshop/Megatron-DeepSpeed
✅ Ultimate Guide to Scaling Video: https://www.youtube.com/watch?v=hc0u4avAkuM
▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 

⌚️ Timetable: 
00:00:00 Intro - focusing on the 3D parallelism!
00:02:00 Quick setup
00:05:00 Stepping through the eval script
00:11:13 3D paralellism - model construction
00:15:00 Sharding the embedding table (model parallelism)
00:20:09 Sharding the transformer layer
00:22:30 LayerNorm fused kernels
00:23:50 Sharding the attention layer
00:25:15 ColumnParallel and RowParallel sharding
00:31:30 Synchronizing input and output embedding tables
00:34:45 Building the dataset (data parallelism)
00:39:15 3D parallelism - forward pass
00:39:25 Pipeline parallelism communication
00:43:35 Pass through the sharded embedding table
00:52:15 Pass through the sharded transformer layer
01:01:36 Sharded logit and cross-entropy computation
01:05:30 Recap
01:11:15 Outro

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 
💰 BECOME A PATREON OF THE AI EPIPHANY ❤️ 

If these videos, GitHub projects, and blogs help you,
consider helping me out by supporting me on Patreon! 

The AI Epiphany - https://www.patreon.com/theaiepiphany
One-time donation - https://www.paypal.com/paypalme/theaiepiphany 

Huge thank you to these AI Epiphany patreons:
Eli Mahler
Petar Veličković

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 

💼 LinkedIn - https://www.linkedin.com/in/aleksagordic/ 
🐦 Twitter - https://twitter.com/gordic_aleksa 
👨‍👩‍👧‍👦 Discord - https://discord.gg/peBrCpheKE

📺 YouTube - https://www.youtube.com/c/TheAIEpiphany/
📚 Medium - https://gordicaleksa.medium.com/ 
💻 GitHub - https://github.com/gordicaleksa 
📢 AI Newsletter - https://aiepiphany.substack.com/

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬

#bloom #bigscience #3dparallelism #scaling
 What's cracking guys, Alex here. In this video, I'm doing the fourth video of the large language model series and I'll be walking you through the Bloom code base. So that's the 175 billion parameter model from the big science event. And I'm gonna focus on showing you how the 3D parallelism is implemented in code. So that means model parallelism, pipeline parallelism and data parallelism. So you're gonna see side by side comparisons of diagrams. I showed you in my previous video on the ultimate guide to scaling and you're gonna see the actual code. So as a quick recap, here's the playlist. As I said, the first video, the ultimate guide to scaling. In that one, I showed you a couple of techniques for scaling the models up to scales, up to those scales of hundreds of billions and even trillions of parameters. So techniques such as model parallelism, data parallelism, pipeline parallelism, mixed precision training and zero optimizer. In the second video, I showed you, I went through, I skimmed through the papers behind the most popular publicly shared large language models such as GPT-New, X1EB, Big Science Bloom and OpenPretrain Transformer from Meta. And then in the third video, I actually went through the code base behind the opt model and I focused on showing you the mixed precision training, the loss scaling, those details. And additionally, I focused on showing you how complex these code bases actually are if you want to get to these scales. So it's not just a matter of taking n, where n is the number of parameters and shifting it from five to 100 and expecting everything to work. Obviously, it's a bit more complex than that. Okay, in any case, as I said, I'm gonna show you the model parallelism from the ultimate guide to scaling video. I'm gonna show you the pipeline parallelism and I'm finally gonna show you the data parallelism approach where you replicate the weights across multiple devices, for example, GPUs and then you do the data replication magic. So you reduce the gradients using all reduce and then you broadcast the weights, et cetera, et cetera. Okay, enough rambling, let's get started. So first thing you have to do is obviously clone the repository. So that's the Megatron DeepSpeed one. I'm gonna obviously link the link down in the video description and just get started through this start fast with me. So just follow the instructions to set up the conda environment and it should be straightforward. Although I did experience some issues on Windows and so ultimately I had to shift to Windows subsystem for Linux and then everything worked out of the box. So it's kind of funny because DeepSpeed is actually from Microsoft folks and it's not working on Windows. I guess it's a bit ironic. In any case, go through this with me, download the data. We'll need these vocab and merges files for the tokenizer a bit later. Go through the training script potentially. I did go through it myself, but I decided to only show you the eval script because it contains all the necessary code for me to showcase the 3D parallelism approach. This one is complex because it additionally uses DeepSpeed which is basically built on top of Megatron and so there is a lot of a lot, the stack of function calls is too deep and it will be super confusing for me to show you that unless I wanna teach you DeepSpeed which is not the point of this video. Okay, long story short, basically go and take these parameters for the eval script. Let me just find where the main pie script, so this one. So, and get started from here. Okay guys, so now I'm just gonna open up my VS code. So you see here Megatron DeepSpeed, I have this Bloom 3 conda environment and I'm just gonna open up VS code. So that's my ID of choice. And let me first show you a couple of things I had to modify in the arguments here. So let me just see what's important. Again, you'll have to link this vocab file after you download it. Just create a random TXT file and you don't really have to download the Viki text dataset because this is going to work out of the box like this. And yeah, I mean, if your goal is to understand how the code works, if your goal is to actually do something with Viki text and obviously you wanna download the actual dataset. For me, I just literally like pasted some random text inside of a TXT file. Okay, the second thing, again, you'll have to download this file, the GPT-2 merges. I also kinda commented out this load function and then we'll just be loading random weights, which again, does not matter because I don't care about the actual result, I just care about these to step through the code and show you what's going on. Then I had this number of workers set to zero because it's easier again to step through the code. And finally, I commented out the checkpoints activations because this does not make any sense for the eval script. So I guess that's just a minor bug on their side. Having said all that, let's start digging into the code. So I'm gonna open up the main function here, the main script. I'm going to grab the eval Viki text. So that's the config for this file that I've just showed you. And let me run this thing. Okay, guys, let me quickly go through the initialize Megatron function. Let's see what's going on there, although it's not that important for what I wanna show you. Just make sure that we have CUDA, so that we have a GPU device. We set some global variables here. This is kinda interesting. So there is the, whoops, don't show this again, please. Okay, so they first parse the arguments and let me kinda dig into that function first. So just to show you how it looks like. So the idea here is to, they have these global variables. So they parse the arguments, they set the, they assign those arguments to this global variable, and then it's gonna be visible across the code base. So that's kinda how they went about designing this thing. I'm gonna skip through the, I'm gonna skip, yeah. I'm just gonna maybe quickly show you how complex this is. There is a lot of arguments that go into this, but I'm gonna skip all of those, and we're just gonna see how that looks like when we start stepping through the code. So as you can see, a lot of error checking, and I'm gonna hit F5 and hopefully get here, yeah. Okay, so let's continue here. So that's the parsing of arguments. Then I'm gonna skip this part. We don't care about that one. And then we build a tokenizer. Again, that's just going to be GPT tokenizer. You can see it here. So we're gonna kinda go through this GPT-2BPE tokenizer, and that's a common tokenizer used throughout many code bases. And just doing this sample coding series, I realized how pervasive this tokenizer is, and I first covered it in my OpenAI clip video, so do check that one out. I'm gonna link it somewhere above if you are curious to learn more about the tokenizers. But for now, I'm just gonna basically disable the breakpoints and skip over the actual tokenizer. And let's continue. So we just do some padding of the vocab such that it's basically suitable to model parallelism. Otherwise, you'd have problems. And that's it. Now we just set up some TensorBoard. This is kinda cool, the code carbon tracker. So they did have this functionality where they track how much are they impacting the environment by just measuring the amount of CO2 emission. But they did comment it out, so as you can see here, they don't actually use it, but it's a cool idea, I guess. Because they say it's very unstable, blah, blah, blah, whatever. Okay, let's step out of this, and let's continue. And that's it. Okay, let me toggle on the breakpoints again. I'm gonna skip most of these things. So get our, so now this is the cool thing. Wherever you are in the code base, you'll literally just grab that global variable and you're golden. You don't have to pass the arguments through the function argument list. Well, it's debatable whether that's the best thing to do, but yeah, I think it's nice. There is a couple of interesting things in this function, so I'm gonna quickly step through it. There is some initialization of the, so I'm using the NCCL, which is NVIDIA Collective Communication Library for all of those old reviews and other operations that you use to communicate and synchronize the GPUs. So it's gonna happen in this function. I'm not gonna step into it because it's fairly complex. Just library code. And then let me just see what else is interesting here, which set the random seed just for reproducibility purposes, obviously. And yeah, that's it. And then in the compile dependencies, there is a couple of interesting things going on. Again, we are going to, so first there's some, this part is not that interesting, but there's one part where they load certain CUDA kernels, and we're gonna later actually use them and leverage that C++ code for extra performance. So I'm gonna skip through most of these. We actually saw something similar, if not the same. Actually, it's the same code because Opt was also using Megatron, so my previous video was also using Megatron in the background. So we've seen the same code here. So what we do is we load these kernels. I'm gonna step into it quickly and just show you how it looks like, but it's nothing super crucial. I mean, not related to the 3D parallelism in that sense. So here you can see we're just loading various kernels for like scaled up or for softmax and then for some other softmax and layer norm, et cetera, et cetera. So I'm gonna put the breakpoint here, hit F5 and just exit here. So it was more of FYI than anything else. Okay, there is some barrier function call which basically synchronizes across. So all of the, it basically waits, all of the different devices, GPUs will wait at that line until the other GPUs hit that same line and then the barrier is gonna release the execution across all of the devices. That's what the barrier does. But I have a single GPU, so that's not interesting. And even though I have a single GPU, I'm still going to be able to show you the 3D parallelism as we'll soon see. Okay, let's just exit here and that's pretty much everything I wanted to show you. Now let's get into the main script. So because I'm using Viki Text 103, I'm gonna take this particular main function and then let's now enter here. Okay, so we fetched the arguments and let's see what's going on. So because we are using Viki Text, we're gonna set the eval metric to loss. That's not that important. And this is where stuff becomes interesting. We're now going to see how the parallelism is implemented in the code here. Okay, so let me re-enable the breakpoints and let's continue stepping through this. So here is the model provider. So we enter this get model function. Let's now see what's going on. Okay, so we hit this else branch. And now, because I have a single device and because of that, my pipeline is only of length one. I have a single stage. And that's why both the pre-process as well as the post-process here are gonna be set to true, as you can see here. But if this was on a, if you were running on a, let's say a two by two topology setup, so we have four GPUs in total and I'm gonna use that topology as a running example. Let me kind of open up the OneNote and show you what I mean. So let me go here. And this is going to be the running topology I'll be using. So basically, I'm gonna have something like this. So we have a GPU here. We have a GPU here. We have another GPU here. And we have yet another GPU here. So we're gonna have three types of parallelism, but I'm gonna show you only the model and the pipeline here on the image. So these ones here are going to be stage one. So stage or stage zero, and this is gonna be stage one, okay? So basically that means we'll have half of the layers of our transformer being here. Maybe if we have 24 layers, 12 will be here. So we'll have 12 layers here. And we'll have 12 layers here. Next up, I'm gonna do something like this. So we're gonna have model parallelism. So across this dimension, we have MP. And across this dimension here, we have pipeline parallelism PP, okay? So this one is gonna be zero and one. So this is gonna be the topology that I want you to have in your mind. And we'll see how it's gonna be useful in a couple of minutes. Okay, again, keep in mind that in the setup I just showed you, the devices from the stage zero of the pipeline parallelism will have, this will be true and this will be false. And then the second stage will have, this will be true and this will be false. And if you have even more stages, then some of these devices will have both of these pre and post process flags set to false because they are neither the first stage nor the last stage of the pipeline. Okay, hopefully that makes sense. Now let's enter this function here. So yeah, we just set this to true because we're dealing with loss metric. Let's continue. And here is the GPT model. So I'm gonna be using the GPT transformer as the underlying architecture. So let's enter here and let's see what's going on. So it's gonna inherit from Megatron module, which is a light wrapper around torch modules. I'm not gonna step through the constructor. Let's fetch the arguments here and let's just kind of skip across all of these details. This is the interesting part, the language model. So now we're gonna enter that function. And here, okay, again, we just fetch the arguments. We continue here. Now there is a transformer language model. So that's gonna be interesting one. So let me hit F5, let's get here. So it's inheriting again from the Megatron module. Okay, we can step across all of these. I'm gonna focus only on the parts that are interesting. So here's the embedding. This is the interesting part. So let me kind of put a break point here, hit F5. Let's see how this is going to work. Okay, so the embedding table is gonna have a couple of details. So here is the Voke parallel embedding. That's what they're using in the background. Parallel suggests that we are basically gonna do a model parallelism of this embedding table. And by the way, let me just go back a step up here. You can see that we only enter here if the pre-process is set to true. And that's only true for the first stage in the pipeline or zeroth index, as I kind of made the notation. But like the first devices in the pipeline will only have this embedding table, which makes sense, because we wanna split the model layers across different stages. Okay, so let's get back here. So here is the Voke parallel embedding. Let's enter here. Let's gonna hit F5. Whoops, let me hit F5 and enter there. Okay, so again, let's see what's gonna happen here. So this is the interesting part. This get tensor model parallel world size is going to in the topology that I showed you a couple of minutes ago, it will return two. So for my particular setup, it's obviously gonna return one because they have a single GPU, but in general, it's gonna return two. And now let's see what's the repercussion of that. So here, what happens is we fetch the rank. So the tensor model parallel rank, and that's going to be zero for the basically devices here. So it's gonna be zero for these devices. It's gonna be one for these devices here, okay? So let me go back to code. Now with that in mind and having this thing, so this is going to be two in the example I'm showing you. Let's now enter this function and let me show you what it does. So it takes the global vocab size, which is this big 50K slots in the embedding table. And we're going to divide it by two because our model parallelism world size is actually of size two in the running example I'm using. Okay, so it will be something like 25K here. Okay, and then what we do is we enter this function here and you can see here, so for rank zero, we would have this thing would be zero because we multiply it by zero 25K. And then this would be 25K plus, no, sorry, this would be zero plus 25K. So that would be 25K. So that means we would return zero and 25,000 here as a result. Obviously, in my case, I'm returning 50K, but bear with me. Okay, so that's the first thing we need to understand. Now those variables are gonna be used to split the embedding table across multiple devices. That's the whole point. So now you can see here, the number of embeddings would be 25,000 if we had the topology I'm using. We can skip this, we can skip this. Those are just some bits and bytes or something, basically where they use integer eight precision, et cetera, et cetera. We're gonna ignore that in this video. So here is the interesting part. So now we create the weight of the embedding table. So we can see here, we are passing this variable, which would be 25,000 again, and we're passing the embedding dimension, which is 1,024, blah, blah, blah. That's it. So the weight would be, as you can see, 25,000 times 1,000 on one device and on those devices that have a rank zero, and then it will be additionally 25,000 for those devices that have rank one, again, talking about model parallelism here. Okay, so we just now initialize it and that's it. So now you've seen for the first time how the model parallelism looks like. So let me just exit here. Let's now continue. The second thing we want to do here is create some positional embeddings. Now these, because they are much smaller than the vocab embedding table, they're just going to allocate the whole thing across all of the different devices. So we have like complete, we can see here 1,024, because that's gonna be our sequence and that's how many positionally encodings we need. Okay, let's continue. We just initialize those and let's see now what's happening here. We can skip all of those. We just add some dropout, blah, blah, blah. I keep entering this function, I'm not sure why, but we now have the embedding and let's continue. Now let me just update the diagram here. I think this is gonna be useful. So let me just take some different color. So we have the embedding table currently placed like this. So we have embedding table here and we have embedding table here. And the size here is 25K and the size here is 25K as well. Okay, let's not go back to decode. Now we have the parallel transformer. Let's see how the transformer is going to be split. Okay, so let's enter the parallel transformer and here it is. Okay, again, skipping these details until we get here. So here is an interesting assert. So the number of layers has to be divisible by the world size for the pipeline parallelism. In the running example I'm using, it's gonna be two. So that means you have to be like, the number of layers has to be divisible by two. In my actual setup, I have, this will be one, so yeah. Okay, now this is the interesting line. So we see here that we divide the number of layers by the number of stages. And again, so because we have 24, that means we would now have 12 layers being allocated on this particular device, okay? So again, keep in mind, this code, so the thing I'm just stepping through, this same thing would be happening in parallel on four devices. And each of those devices have different contexts because they lie in a different part of the topology and because of that they will be entering different types, different code, they will be going through the different code paths, that's the main idea. Hopefully you understand that. Let's continue here. Okay, so now we create the offset. So we grab the rank and we multiply it by the number of layers. So that means that the offset for these two devices here is gonna be zero, for these two is gonna be zero. And for the two ones here is gonna be 12, right? So let's go back here. So it's gonna be one, the rank will be one for those and then times the number of layers, 12. So offset is gonna be 12 for them. So let's see how offset plays a role here. So we now form 12 layers and we have the offset here. So you can see that we are passing that index information into this build layer function. So let's enter the build layer function. Here it is. And there we start building the actual transformer layers. We'll do this 12 times. So let me just enter this layer here. Let me hit F5 and we are here. So again, let me just focus on the important parts, ignoring many details obviously, otherwise it will be too bothersome. Okay, so first things first, we want to instantiate a layer norm. Okay, let's see how the layer norm looks like. The layer norm has some fused components, which is interesting. So that means we'll be using optimized CUDA kernels, fancy word for saying functions, implemented in CUDA library that are running super fast on GPUs obviously, on Nvidia GPUs. So here you can see we import this fused mixed precision layer norm CUDA. So let me show you where that was constructed. And I actually showed you that in the beginning of the video. So here, let me just move this thing here. So if I were to copy paste this here, so you can see here we are loading that precise variable is loaded using the CPP extension load helper. So we kind of passed, I kind of showed you this function before, and that's where it was initialized. Now let's get back to the actual code. Let me go to the top of the stack. We import that kernel and yeah, we're gonna then use it a bit later. So everything else is not that interesting. I'm not gonna focus on the layer norm because this one is not going to be split across devices. Okay, let me exit this. And let me just kind of move this a bit to the left. And we are out. Now let's see the parallel attention. This is gonna be the interesting part. So let's enter the parallel attention. I'm gonna hit F5. Let's go to the part that we care about. And that's the part where we split. So here we can see, we can step over all of these. The magic starts happening right here. So again, model parallel world size is gonna be equal to, so it's gonna be two. And now we divide the projection size by two, blah, blah, blah. So just the devising number of attention heads to two by two. And now we get to the interesting part. So this is the column parallel linear feed forward layer. That might be a mouthful, but let me remind you what that is. So this is what we are actually building right now. This is the self attention layer. And what we are building right now are these, as you can see, column parallel layers. So we take the weights and we split them in this column wise fashion. And then we're gonna use, so we're gonna use these here. The Q1, K1 and V1 are going to be used here, as you can see. And then the Q2, K2 and V2 are gonna be used on the different device. And that's what we are currently building. That's how this code, that's what this code is implementing. So let me enter the actual column parallel linear. And you can see again, the weight matrix is being split in this column parallelism fashion. So let me hit five, let me enter here. Let's see where that's gonna happen. So here is the world size. Again, this is gonna be equal to two. And then the output size is gonna be divided by two. Okay, so let's continue. And here is the actual, so we are now allocating the weights. This is going to, so this would, if we had the topology I'm using as a running example, this would be, as you can see here, divided by two. So that means we're literally splitting it column wise. That's the whole point, okay? And then once we have those weights, we just initialize the weights and there's some bias. Let's keep that part. It's not that interesting. So let's exit here, again, entering this thing. Okay, so that's the first part. Let's see what else. So now we have to create some softmax stuff that's going to be this part, the softmax. The only difference is this is gonna be optimized. So obviously there are no learnable weights in this one. So I'm gonna just skip it. We're not sharding it. So now we create the dropout, blah, blah, blah. And there is the row parallel linear. So this is the second component. So let me go back again to the OneNote. And here is the second part. So here you can see the other part of the self-attention layer. We have this linear layer and we split it row wise this time, B1 and B2. So you can see here. So let me go back to the code and let me enter the row parallel linear. So if I enter here, you can see again the diagram and let me put a break point here. Let me hit that five. So here we are. Again, world size, it's gonna be two. And let me continue. And again, this time you can see that the input size per partition is the one that's being divided. So here the input will be divided, not the output size. And because of that, we are doing row wise splitting, right? That's it. That's the difference. Again, I'm gonna skip everything else here. Let me kind of step out of this function, shift F11 does that on Windows keyboard. And here we are. So we have the row parallel linear instantiated. Let's continue. We don't have row three embeddings. So we can skip this part as well. And here we are. So that was the parallel attention. Okay. Again, let me just update the drawing. We have the parallel attention being sharded across two devices in this running example. So that means we'll have half of the weights here and half of the weights here. Obviously we'll need 12 of these. So I kind of didn't do it justice here, but yeah, you can imagine we'll have 12 of these transformer self-attention layers. So let's go back to the code. Let me exit here. So we have layer norm again. I won't bother. I'm just gonna exit this layer norm. Let's see what else we skip the decoder part. So now there is the parallel MLP. So again, how the transformer block looks like, we have the self-attention module and then we have the MLP part. So now here we'll see first the column-wise splitting and then the row-wise splitting. So this is the diagram we'll implement. Let's go back to the code. Let me show you this quickly. And then this is the last thing I'll show you. And then we are pretty much done with this model construction. So let me hit F5. Let me enter here. So again, column parallel. And I'm just gonna skip. I'm just going to disable the breakpoints because it's the same logic as before, right? We've seen this. This is the same logic as before. So I'm just gonna skip over all of that. Let's get here. Now some activation stuff. Here is the row parallel linear and that's it. And that's the one that has the 4H usually is the common convention. So the hidden layer will have 4X more weights than the input and output, dimension, sorry. Okay. Again, let's exit here. And that's our basically transformer layer. Again, you can always follow here where we are. We're currently implementing the transformer layer. So I'm gonna exit here and we do this 12 times. So I'm gonna, again, let me just make sure it's disabled. I'm gonna enable the breakpoint here, hit F5 and we have 12 layers. So again, let me update the drawing. Here is the drawing. We now have 12 layers here and here. So we have 12 layers here and here. And obviously, because this is running in parallel, these two devices would also have 12 layers here. So 12 layers and 12 layers. The difference is these two, because they are in the rank zero of the pipeline parallelism, they additionally have the embedding table sharded. Cool. Guys, let me know how you like the video so far. Leave any feedback if you have any. That would be super cool. So again, post-process. That's going to be active only on the last stage of the pipeline. So that means only those devices will have the layer norm, but I'm gonna skip the, I won't enter the layer norm because we already seen it. That's the optimized layer norm. And that's it. Again, that was the parallel transformer. Let's continue. We don't have the decoder part. We don't have the pooler, so we can exit here. We can exit here. And we have our language model being constructed. Again, I'm not sure why I'm entering this thing. So I just, once I enter it, I just do step out and that's how I exit in case you're wondering. Okay, guys, here we are, GPT model. We just constructed it. Now there is this initialized word embedding. So because I have, in my actual setup, I have a single GPU, so I'll just hit return. But in general, if we had the actual topology I'm using as the running example, we'd be hitting this code here. And let me show you why this is interesting. So they mentioned it here. So parameters are shared between the word embeddings layer. So that's the common thing you always do. So you have the embedding table being shared between, so mapping the tokens into the embedding vectors. And then at the output, you also use that embedding table to map the embedding vectors into the vocab. So basically distribution across the various words, subtokens from your, or tokens from your vocabulary. So those are tied. And you can see here, they have a nice trick that they do. So if this was running on the last stage of the pipeline, then they would create embedding table. Again, it's gonna be a parallel embedding table. So that means it's gonna be shorted across two devices of the last stage of the pipeline. And then this is the nice trick. So basically what they do is they do all reduce. So ensure that first and last stages have the same initial parameter values. So because this was not independently allocated and because they are lying on different devices, because of that, they have to do all reduce. So you can see here, if we are hitting either the first stage of the pipeline or the last stage, apply the all reduce method on this word embedding, or word embedding's weight. If you're not familiar with these collective operations, I did not explain them so far, just check out the Viki page. They have a nice visual explanation. It's fairly simple. Roughly what's gonna happen is assume that the first stage will literally receive the weights from the last stage, and then they're gonna divide it by two. And if the last stage does the same thing, you'll end up with the same result. And that means you'll have the same weights both in the first stage and in the last stage, and thus you've synchronized the weights. As simple as that. Okay, guys, hopefully that was interesting. And hopefully that makes sense with the drawing I'm using in OneNote. So let me exit here. So here is the, now some stuff that's not that vital. I'm gonna skip all of this, put the break point here, hit F5. Let's see what else. So we push all of these models onto the GPU. We use this float, we convert them to float 16 weights. So I'm going to just, basically let me enable the break points now. So I'm gonna hit this one a bit later, so we will see how that is going to be used. Okay, local DDP data distributed parallelism. We're gonna ignore that for the time being. And let's just return the model. Okay, so again, because I have load set to none, I will not be loading the checkpoints, which means I'll basically be using random weights. Okay, now we build the data set. Let me quickly walk you through this one. Although that's not that relevant, so tokenizer. So again, I just load this random dummy file I created. And you can see what it has. Literally I copy pasted. I actually took a single paragraph from WikiText data set and just copy pasted a couple of times just to have something there. So now we just do some processing, the tokenization, blah, blah, blah. Okay, let me just exit here. This is not interesting. So we then tokenize the text. Okay, I'm just hitting the breakpoints. So this is some processing that they do. As you can see, a lot of details that I'm trying to unsuccessfully ignore. So we now do the tokenization and you can see here, now we have the, let me show you how that's gonna look like. Tokenized data is just a list of tokens of our text. So the text was tokenized and we have 448 tokens, the smallest data set ever. Okay, now we use that to construct this LM data set. Again, I'm not going to dig into a lot of details because that's just the language modeling detail and we are trying to focus on the 3D parallelism in this video. Okay, let's exit here. We have the data set. Now we build the data loader. The reason I'm gonna show you this one is because it involves the data parallelism approach. So here it is. In PyTorch, it's fairly trivial. So here you can see get data parallel world size. So now additionally, I did not show you this so far, but imagine having this same group of devices, like four devices like this. So I'm just gonna change the color. So imagine having this same group duplicated maybe twice or even four times, okay? And we're doing the same thing in each of these groups until now. So now it's gonna be a bit different for each of these groups. Okay, so let's say this is group zero and this is group one. I'm gonna ignore, we're gonna have two groups, okay? So let's go back to the code. We would get world size would be two. And let me show you how this is implemented in the background. It literally calls this distributed get world size and then you specify the particular group. So here they specify the get data parallel group and that's why basically they get number two. Okay, so some constants there, I will not dig deeper than that. Let's continue, let's get the rank. So now where the rank is zero for the left group, so rank would be zero again for this group here, it will be one for these devices here. So that's how we'd have a differentiation. Let's continue. We formed the distributed sampler, okay? And that one is going to be shifting a part of our batch to one group and the second part to the other group. So that's gonna handle in the background by the PyTorch. Okay, so now the data loader, blah, blah, blah, pinned memory, again, a lot of details, I'll have to skip those. And now we enter the evaluate and print results. So here is the actual function where we'll be doing a feed forward pass through the model. So we just set it to eval mode. We entered the TorchNodeGrad context because we're doing inference, we're not trying to train anything here. And now we're going to fetch our first batch. So I'm going to ignore how that works. So I'm gonna disable the breakpoints, otherwise I would hit some breakpoints I left there. So let me hit F10 and we have a batch. So a batch is simply going to be, as you can see here, we have text and we have pet mask keys. So text is going to be, I guess, eight times 1,024 or something. Yeah, so eight because we have a batch of size eight and 1,025 because our sequence that we're gonna pass to the transformer is gonna be 1,024, but it's 1,025 because we'll be using one as the target. And so you're gonna see that in a second. So yeah, hopefully that makes sense. Like you literally take the, it's 1,025 because you take the first 1,024 as the input and then you're trying to predict shifted by one to the right, you try to predict that sequence. And that's why you have 1,025 here. Okay, let's continue. Let's enter the forward pass. Let's see what's going on here. So some batch processing, again, I'm gonna ignore all of that because that's just a language modeling standard stuff. This is the interesting part. So this is the receive forward. So what's going on here? So let's enter there. Let me now enable the breakpoints. So here what happens is the following. So if we are in the first stage, then the input tensor is none, okay? But if we were in some other stage, then we call this communicate method. And that means we're waiting for the output of the previous stage to pass that output as the input of this stage. And if that happens, we're going to override the batch that we just loaded from the data loader. So that might be a mouthful. It's gonna hopefully make sense in a couple of minutes. So here is the input tensor. Let's continue here. Okay, some unwrapping of the model, not that interesting. So here's the part where we set the input tensor. And let me kind of dig into this. So blah, blah, blah, we call the GPT model set input tensor function. And then, okay, let me exit this one. So then we enter the transformer language model set input tensor. Again, a lot of levels of indirection, and I keep hitting this random function. Okay, so here we are. So finally we get to parallel transformer, and we set the input tensor to none, because we are now in the first stage. I mean, at least in the actual setup I have. But let me show you where this one is used. So it's going to be used here a bit later in the forward pass. And we'll see that if preprocess is true, that means we're in the first stage. And in that case, we'll actually be taking the data from the batch. Otherwise, if we are not in the first stage, we'll be passing the actual input tensor that's passed from the previous stage. So again, we'll see that a bit later. So let me now exit all of this. Let me exit this, let me exit this. And let's continue here. So here is the actual forward pass. Again, just so that we are on the same page, let me show you what's going to happen here. So what's happening is we are having our distributed data loader fetching a batch. And I'm just going to simply draw it as a simple 2D square. And we're going to basically divide it into two pieces. This one is going to be passed to this group. And the first chunk is gonna be passed into the zeroth group here, okay? And now what happens is this is going to be fed here. And while that's happening, these two guys are blocked. So the thing is that communicate method I showed you is a blocking method. So that means until this has done its feed, so until this data has passed through the first stage, these guys are waiting. And then, only then do they start doing the actual computation. That's how this is gonna look like. And obviously the same happens for the rank one, for the second group, second data parallel group. So I will not be drawing it there, obviously. Okay, let's enter the feed forward pass. Okay, we have to enter this forward call. So distributed blah, blah, blah. There is a lot of levels of indirection again here. Let me just exit here. Oh my God, oh my God, let me enter here. Okay, so finally we are in the forward pass of the float 16 module. So let's see what's happening here. So as you can see here, if we're in the first stage of the pipeline, we convert our data to float 16, because we're gonna do mixed precision. So now inputs are gonna be converted to float 16. And now we call the actual transformer with the float 16 data. Okay, let's enter here. Let me hit F11. I think I can just hit F5 and I'll hit the, yeah. So here we are. We are in the GPT model forward pass, and let's start going through it. Okay, so again, I think I can hit F5 because they have a break point there. Yeah, so we are in the transformer language model. Okay, so as you can see here, only the first stages, because pre-process is true for first stages of the pipeline, are going to do the embedding part. Okay, so we do the embedding part. Let me kinda enter this. Whoops, I hate this. Let me enter it. So here we are in the embedding forward. Let me see whether it's useful to enter this one as well. So here we are in the vocab parallel embedding. Let's see how the forward would look like. So again, keep in mind that we have sharded that embedding table here and here. So we only have 25,000 of our embedding vectors here, and the other 25,000 is here. So they'll somehow have to communicate between each other, such that we can end up with the correct results. So let's see how that's gonna work. So let's enter here. So we would actually be entering this branch if I had the setup that I'm using as running examples, because the model parallel size would be two. So what would happen is the input mask, as you can see here, it will be true for all of those indices that are smaller or bigger than the index here, so then the boundary index here. And these would be like, these will be zero and 25,000 for this one, for this one here, and 25,000 to 50,000 for this one, okay? And that's why this is gonna mask those indices that do not belong to this device, okay? So there is this subtraction here that we have to do so that we have correct results, because imagine we are now on this device, because its start index is 25,000, we want to subtract 25,000 such that we end up with zero, because we are indexing in its table and the zero index is the first slot in that table on that device, if that makes sense. Sorry for being a bit confusing here. And you can see here, for all of the other indices, we just set them to zero. Maybe let me try and draw this quickly. So here is what's going on. Okay, we have basically input that contains indices from the whole table, so that means it might have numbers. Okay, let's quickly check what's the actual size. So let me use this input as the running example. So we have shape 8024, okay? So let's imagine we just have 1024, because we have batch size of one, it's gonna be easier. Okay, so here we are. So we have a vector that's like 1024. So this is gonna be 1024, okay? And it might contain in any of these slots, in any of the dimensions, let's say this is dimension I, it might contain numbers from zero to 50K. I'm just gonna run the numbers because it's easier, okay? So now let's assume we are on this device here. Okay, so that device only has the following indices. So it has, so we have two tables. This one contains from zero to 25K, and this one contains from 25K to 50K, 25K to 50K, okay? So what this device is going to do is it's going to mask all of those values here. So let me take a different color. So all of those values that have numbers that are zero to 25K, because it does not contain those numbers. So it's going to mask them out by putting them to zero. So let's see, this one is also put to zero, this one is also put to zero, blah, blah, blah, okay? And then it's going to take this number here. Let's say this one is 26K. It's going to subtract the 25K because that's the offset, and we're going to end up with 1K. And 1K, when you index with 1K into this table, you end up with some vector, embedding vector here. And that vector is, if you take a look at the whole table, is at the 26K, okay? So that's why we do all of this offsetting. Hopefully this now makes a bit more sense. Okay, let's continue. So now we do the embedding and we get the results. Okay, so let's check out the shapes. Shape is going to be 8024,024. This is a sequence length, this is the embedding dimension. And now what we would have to do is we would have to set those vectors that were masked out to zero, otherwise they would have incorrect results, okay? And finally, what we do is we call this reduced from tensor model parallel region. Let me quickly enter this one, and then I'm going to show you what is going on here. So here we are. Let me just go like this. So let me actually enter here. Let me enter here, reduce. And here is what's going to happen. In case that the model that we have multiple stages, sorry, this is model parallelism. In case that we have multiple shards across, if we have model parallelism of size two, as the running example is using, then we would not enter this line. We would do the all reduce across the group, across the model parallel group. So let me break that down, what we have done right there. So let's again, get back to the drawing and let's see what happened. Okay, so, okay, because these guys here were masked out, so that means they were, they maybe, this one was maybe 14K. And we've artificially set its value to zero. Because it was zero, that means it was indexing the zero vector from this table, which is incorrect because we want the 14,000, right? So this one here, again, let me just change the color and let me explain. So this one here was maybe 14K. So that's like the token ID, okay? So that one was indexing somewhere here, but because we've set it to zero, it was actually indexing this vector here. And that's why we have, oops, it's glitching. That's why we had to do the zeroing out because we don't want to have this value here. We don't want to have this embedding vector for this particular token. That's why we do the zeroing. That's why we do this part here. And then what the all reduce does is the following. Let me go back here. So this is going to happen. Let me see whether I can explain this nicely. Basically, this device here mapped this set of tokens into obviously various bunch of embedding vectors. So the dimensionality here is the same as the dimensionality up there. So it's 1024, but here we now have the hidden dimension. Okay, so this thing, this vector here is gonna be valid for this device. So that vector is gonna be valid. So let me change the color to green. So this one here, this token is gonna be valid on this particular device. Okay, so it's valid here. But as we just saw, it's going to be invalid for this one. Okay, so let me draw another table. So let me draw the results. So this is the output from this device. So this is the output from the second device. And it contains a corrupt data, actually all zeros, because we've zeroed it out. So it's gonna be all zeros. And so now what the all reduce does is basically, you sum these up. If you sum them up, you'll end up with a correct result both here and here, because if you just add zeros onto this device, you keep the same value. But this device will now have the correct value of the embedding vector. And that's the magic of the all reduce. So what's going to happen is these GPUs here are going to do that type of a communication and update the embedding vectors. So that's a beautiful thing, like if you ask me. So hopefully I managed to explain it nicely. And now finally, let's get out of here. Okay, again, just to update our mental model. So where we are at the moment is in the state where basically these two devices here are waiting for the information. And we have the embedding vectors on this device and on this device being the same because they're synchronized because of this all reduce we've just done. Okay, so we are now there. We have that and now let's get back to the code. Okay, so after this, we're gonna have some positional embeddings being added that's less interesting because it's not sharded. Let's continue we have some drop out, blah, blah, blah. And we end up with the output. Okay, so now let's see what's going on. So we now do the call to the encoder. And that's just going to be our transformer basically. So let's enter there. Let me hit F5. So we are now in the parallel transformer. Let's step over all of these. If self pre process, so again, if we're in the first stage we'll be doing this. So if we're in the first stage, we'll be grabbing the embedding vectors we just created. So let me show you the shape, 8024 to 1024. And we just grab those, we do some transpose operation and we do continuous and we continue. But if you recall this line here, this was from the, if we were, if this was not the first stage of the pipeline, but instead if this was maybe second or third or whatever stage, then instead we'd have hidden states being equal to input tensor, which came from the previous stage. Okay, so we would bypass this part here. Hopefully that makes sense. Let's not continue. We're not doing checkpointing. Okay, now we are just going to do a feed forward through all of the transformer layers. Okay, so let's do that. Let's grab the, so here is the layer. I can hit F5, I think. So we are now in the parallel transformer layer. And the first thing we do is we pass the data through the layer norm. Let me quickly show you how that's gonna look like because it's fused. Oh my God, this sucks. Okay, so let's just step over here. Okay, here we are. And here is the fused layer norm, F5M function we call this apply method. So let me enter here. So here you can see how it looks like. Oops. So what it does is ultimately it calls this CUDA kernel, which is a global variable, which we loaded in the beginning of this program. And then we just call the forward F5. So we cannot actually enter here because it's just a compiled C++ code. So I cannot enter with this current debugger. Maybe there is some, maybe it can be tweaked so we can enter and see the C++ code, but I'm not sure how that will work. Okay, so in any case, I just wanted to show you how the kernels look like. Now let's exit from the layer norm. And now let's hit the actual attention part. Okay, let me hit F5. Here is the attention part. So we are in the parallel attention part. Again, remember the weights are sharded such that each of these devices in the rank zero and rank one of the model parallelism only have a portion of the weights. Okay, so we end up with some... Okay, so we first hit the column parallel layer. Let me go back up again here up the stack. So we hit the query key value is actually the column parallel layer for a function. So that's again, let me show you here. So we are hitting the column parallel. This is the part I was showing you. So let me go back to the code now and let's see what's going on here. So we call this weird function copy to tensor. So what is it doing exactly? Let's maybe enter. So let's enter here. Let's enter here. And you can see here in the forward pass, it's just doing identity function. So why are we doing this? Well, that's because in the backward we'll be doing the reduce method. We'll be calling the all reduce. So let's see why that makes sense. Let me show you the diagram here. It's kind of clear once you see it. So what we're doing here is this. We are implementing F. F is, as you can see here, it's just the identity function in the forward pass because we just passed the data, the same data X to one device as well as to the other device. But once we go backward, we want to all reduce the information. We want to all reduce the gradients that came from this portion and this portion. We want to basically average them and then pass them to the previous layers. Hopefully that makes sense. That's why we have it implemented in such a way. Let's now exit this part. Now this is less interesting. We just apply the actual linear layer and then blah, blah, blah, bias. We can exit there. That's less interesting. Okay, so now we just have some, what's this, reshaping. We have some view, blah, blah, blah. We're splitting into three parts, the query key values. I'm not gonna focus on the actual transformer logic because I do assume you either know that or even if you don't know, just focus, let's focus on the 3D parallelism. So again, just some transformer stuff. We're changing the view so we can do multiplication between the keys and the queries and then basically use that to, okay. So here we just preallocate a tensor that's gonna contain the scores for the, that we'll then apply the softmax. We'll then apply a softmax. As you can see here, so we first do the multiplication of queries and keys here so we can skip all of that. And then we do, we now have attention scores. And in order to convert scores into actual probabilities, we have to apply the softmax. So again, I'm gonna exit this softmax layer. It's not that interesting. We've already seen the CUDA kernels. So I'm gonna just skip that. Okay, attention. We just do the dropout, blah, blah, blah. Let's continue. My computer is kind of struggling. You can see it's kind of lagging behind. So I cannot go faster than this. We just do some, again, view manipulation. And then we do the batch matrix multiply, which is going to basically sum up the value vectors. And we end up with the final result here. And then we just change the view. We do pre-mute, blah, blah, blah. Nothing new there, hopefully. And finally, we hit the dense part. And the dense part is, let me show you here. The dense part is this part. So that should be the row parallel linear layer, right? Because we're splitting the, we're now here. We're now here. Let me go back to the code. Let me enter here. Whoops. Let me enter here. Oh my God. So we are in the row parallel linear layer, as I said. So let's quickly see how it's going to work. Again, a couple of details there. So you can see we first applied the linear layer, and then we call this reduce from tensor model parallel region. So let's see how that thing looks like. That one in the forward pass is doing the all reduce. And in the backward, it's just doing the identity function. So it's just the, literally the polar opposite from what the column-wise was doing, right? So let's go back here. And why is that? Let me show you the diagram again. So here you can see in the forward, we want to all reduce these weights such that we end up with the correct result on both devices. But in the backward, we're just copy pasting the gradients. That's the reason why that was implemented in such a manner. Okay. So I'm going to skip all of this because it's not relevant. And we are done with the parallel tension. We are exiting the attention right now. And okay, now we have some residual stuff, blah, blah, blah, biases. We don't care about those. Some of those are JIT compiled so we wouldn't be able to step through it either way. Post attention layer norm. I think we're going to hit the break point here. So I'm going to remove it, exit this function. Let's exit the function here. And we're back. So we don't have the coder. Now we hit the MLP. So MLP is going to be this part. Again, we're now past the data that came out of this part. We're going to pass it here and we're going to repeat the same thing. So I'm going to skip it. So we're going to do the, as you can see here, column wise and then row wise, but it's the same logic. So we can safely skip this part. Okay. So let me do disable break points. We do the same thing and that's it. We are done here. Let me just exit this part here and let's return back the value. And that was the layer one. We obviously repeat that 12 times. So I'm going to put the break point here, hit F5. And now we're here. So let's see where we are now here with the diagram. So the data has successfully propagated all the way. Let me change. I'm not sure what color is the best one here. So maybe let's do it like this. So we are now here. We successfully propagated our activations all the way to here. And now we'll be passing them slowly to the second stage of the pipeline. So let's go back to the code. In case we are in the last stage of the pipeline, we'd be doing this additionally, passing it through the layer norms. And that's it. We are now here. We can step over all of this and we return back the results from our transformer language model. Guys, this is pretty much it. Like this is, we've seen the main concepts already. There is now a couple of additional details here. So like post language model processing was this one. Let me enable the break points. Let me hit F5. Here we are. So again, there's gonna be some synchronization between the devices in these ultimate parts. So yeah, assume now we're gonna project those weights into the vocab space. So let's see what's going on. So we have input parallel. That's gonna be, I guess, 8,024,024. Yep. And then once we do step over, we're gonna end up with the vocab. So eight, yeah, 1,024,50,304. So that's it. Again, there was some communication happening here between the devices such that we have consistent results on both devices. And that's it. Let's continue here. Whoops. We're now almost at the bottom of the stack. So we are now in this float 16 module. What we now do is because we're in the last stage, if we're in the last stage, we convert from float 16 to float 32, and then we return back the results. Okay. So let's exit here and we are back. So again, now we're gonna have this send forward function. And you can see here, if it's not the last stage, then communicate. So that means all of the stages will have to pass the data to some other device, unless you're at the last stage of the pipeline, which makes sense. Again, communicate is a blocking method. So that means if you're waiting for the data, you'll be blocked. That's why you have the concept of micro batches so that you can minimize the waiting, the idle time of the devices. Hopefully that makes sense. Okay, so in the last stage of the pipeline, you do this additional vocab parallel cross entropy. And I can quickly maybe enter this one, although I think it's gonna be fairly detailed. There's a lot of details here, but like the concepts are the same, like nothing new there. So I'm gonna slowly step out of all of this. So yeah, as you can see, they have devices you have to communicate such that we end up with the correct logits and then blah, blah, blah. You can go through this one at your own pace. I think it should be fairly easy to understand it. Okay, let's exit that part. Let me exit it. And that's it. We just do some summation. We return back the loss. We do some additional overviews here. And now we keep on doing this. Ultimately, you'll end up with the evaluation metrics. I'm gonna disable the breakpoints, hit that five here and exit here. And I'll let's say it's from here and that's it. Now, if we are on the last rank, we just, and by the way, here's how the last rank is calculated in this particular example. So here you can see that we just get the rank without any group, which means that now, however many devices we have like in total, we just find its absolute rank. So that means that for example, this one will be ranked zero, this one will be ranked one, two, three, then five, six, seven, eight. We will be returning those results. So basically we only on the last rank, so only on a single device in our whole cluster, we do this operation. So basically, yeah, just some blah, blah, blah perplexity. And then we print the results, but we don't care about the actual results. So yeah, so you can see here, average loss perplexity adjusted perplexity, blah, blah, blah. Nothing interesting there. And now finally we're done. Okay, guys, we've seen a lot so far. Let me do just a quick recap and then we can wrap this video up. So we have a couple of differences depending on whether we are dealing with the last stage of the pipeline, with the first stage of the pipeline, or the intermediate stages of the pipeline. So let me go and run the eval script again. So I'm gonna hit run here. And I have a breakpoint set inside of this forward step function. There's a couple of things I wanna explain here. So first things first is this receive forward function. Then I'm gonna explain the differences in the forward pass through the model and then the send forward part. So first things first, here, if we enter, you can see that the first stage in the pipeline will have the input tensor set to none, as you can see here, which means that they'll actually be fetching the data from the data loader. So they'll grab the chunk of data from the data loader and pass that as the input into the embedding table and then the transformer layers. So again, let me show you the diagram here. So that means those guys, the first stage, will grab the chunk of data from the batch and will pass that data through the, as you can see here, embedding table and only they have the embedding table, if you recall from the explanation back then. And then they'll pass the data through the transformer layers and ultimately they end up with this, with the output activations denoted in gray color here. Okay. And then all of the other stages will block on the communicate method. And that means they'll be waiting until the previous stage sends the activations to this stage. And then it will continue on the execution. So that means for these guys here, they'll be blocked until the gray activations are sent to the second stage here, and then they'll do the forward pass. Okay, so that's first thing. Then let's go back to the forward step. We have the forward pass through the model. So I'm gonna now just enable all of the breakpoints and let me hit that five just to enter into the actual forward pass. So I want to get, whoops, let me just do this. I'm gonna basically disable this and put a breakpoint here. So this is where the difference happens. So basically this is, as you can see here, GPT model forward, depending on whether we're dealing with the last stage or any other stage of the pipeline, we'll be entering this branch here. And as you can see here, if I enter this one, this thing does converts the activations into the logits. So that means we'll end up with the vocab space dimensionality. So let me digest what I just said there. So if you take a look at the LM output, all of the stages will be just returning the pure activations. So that means 8024,024. So those are the raw activations from the transformer block. But the last stage will be returning this processed information because it maps the tokens into the vocab space. So let me show you what I mean by that. So if we exit here, the outputs for the last stage will be, whoops, I think I have to do step over here and then let's do output shape. So you can see here 50K. It will be 25K in the case of having sharded, the model sharding. And yeah, that's it. So let me go outside of this thing. And now we're here. Okay, so now, again, remember output is one of two things. It's either the raw activations from the transformer layer or it's the tokens mapped into the vocab space for the last stage of the pipeline. And then the send forward, as you can see here, all of the stages will be communicating the activations to the next stage, but the last stage will just skip. As you can see here, if not the last stage, then communicate. If it is the last stage, we'll just ignore this. This is a no-op operation, no-op for the last stage. Okay, and finally, as you can see here, we have the last stage branch here, which means that the last stage will pass the logits into this function here and calculate the cross entropy, and then it will return back the loss. Whereas all of the other stages first communicated information to the next stage, and then they return none. And then if we go a level up from this function, so this is the forward step, you can see that the output, again, for the last stage will be all reduced and then just added to this total output, which will finally return and used to calculate the perplexity, et cetera, et cetera. So hopefully that makes sense. Again, let me go through the diagram. I think it's fairly straightforward. So again, the first stage here will be grabbing the data from the data loader, who will be passing that data through the embedding table and then through the transformer layers until we have the raw activations, and then they'll be passing those activations onto the next stage. Whereas the next stage will be waiting for the activations and then grabbing the activations from the previous stage, passing those activations through the 12 transformer layers, and finally calculating the logits, calculating the cross entropy loss, and then we'll be accumulating that information. If we were to have additionally intermediate stages in this diagram, obviously, then those would be receiving the activations from the previous stage and then just passing the activations to the next stage. So we kind of have three different branches in this pipeline logic. Cool, that was a long video. Hopefully you liked it. If you did, I would really appreciate any feedback. This will be the last video from the large language model series. I might create one more on the GPT-NewX20B if I have any interest. If you'd like me to go through the GPT-NewX20B code base, do comment down below and I'll be monitoring the comments. In any case, if you found this video useful, consider subscribing, sharing it out. That's the best way you can support the channel. And until next time, bye bye. Abū La mumṭaffāṭ 탓 rather le,
[{"start": 0.0, "end": 2.52, "text": " What's cracking guys, Alex here."}, {"start": 2.52, "end": 4.34, "text": " In this video, I'm doing the fourth video"}, {"start": 4.34, "end": 6.16, "text": " of the large language model series"}, {"start": 6.16, "end": 9.32, "text": " and I'll be walking you through the Bloom code base."}, {"start": 9.32, "end": 12.32, "text": " So that's the 175 billion parameter model"}, {"start": 12.32, "end": 14.120000000000001, "text": " from the big science event."}, {"start": 14.120000000000001, "end": 15.780000000000001, "text": " And I'm gonna focus on showing you"}, {"start": 15.780000000000001, "end": 19.080000000000002, "text": " how the 3D parallelism is implemented in code."}, {"start": 19.080000000000002, "end": 20.72, "text": " So that means model parallelism,"}, {"start": 20.72, "end": 22.84, "text": " pipeline parallelism and data parallelism."}, {"start": 22.84, "end": 26.52, "text": " So you're gonna see side by side comparisons of diagrams."}, {"start": 26.52, "end": 28.32, "text": " I showed you in my previous video"}, {"start": 28.32, "end": 30.04, "text": " on the ultimate guide to scaling"}, {"start": 30.04, "end": 31.88, "text": " and you're gonna see the actual code."}, {"start": 31.88, "end": 34.8, "text": " So as a quick recap, here's the playlist."}, {"start": 34.8, "end": 37.72, "text": " As I said, the first video, the ultimate guide to scaling."}, {"start": 37.72, "end": 39.88, "text": " In that one, I showed you a couple of techniques"}, {"start": 39.88, "end": 41.9, "text": " for scaling the models up to scales,"}, {"start": 41.9, "end": 43.92, "text": " up to those scales of hundreds of billions"}, {"start": 43.92, "end": 45.5, "text": " and even trillions of parameters."}, {"start": 45.5, "end": 48.0, "text": " So techniques such as model parallelism,"}, {"start": 48.0, "end": 50.120000000000005, "text": " data parallelism, pipeline parallelism,"}, {"start": 50.120000000000005, "end": 53.24, "text": " mixed precision training and zero optimizer."}, {"start": 53.24, "end": 55.120000000000005, "text": " In the second video, I showed you,"}, {"start": 55.120000000000005, "end": 57.120000000000005, "text": " I went through, I skimmed through the papers"}, {"start": 57.12, "end": 61.64, "text": " behind the most popular publicly shared large language models"}, {"start": 61.64, "end": 64.56, "text": " such as GPT-New, X1EB, Big Science Bloom"}, {"start": 64.56, "end": 67.84, "text": " and OpenPretrain Transformer from Meta."}, {"start": 67.84, "end": 69.2, "text": " And then in the third video,"}, {"start": 69.2, "end": 74.0, "text": " I actually went through the code base behind the opt model"}, {"start": 74.0, "end": 77.02, "text": " and I focused on showing you the mixed precision training,"}, {"start": 77.02, "end": 78.84, "text": " the loss scaling, those details."}, {"start": 78.84, "end": 80.67999999999999, "text": " And additionally, I focused on showing you"}, {"start": 80.67999999999999, "end": 84.28, "text": " how complex these code bases actually are"}, {"start": 84.28, "end": 86.12, "text": " if you want to get to these scales."}, {"start": 86.12, "end": 87.92, "text": " So it's not just a matter of taking n,"}, {"start": 87.92, "end": 89.60000000000001, "text": " where n is the number of parameters"}, {"start": 89.60000000000001, "end": 92.48, "text": " and shifting it from five to 100"}, {"start": 92.48, "end": 93.94, "text": " and expecting everything to work."}, {"start": 93.94, "end": 96.24000000000001, "text": " Obviously, it's a bit more complex than that."}, {"start": 96.24000000000001, "end": 98.12, "text": " Okay, in any case, as I said,"}, {"start": 98.12, "end": 100.12, "text": " I'm gonna show you the model parallelism"}, {"start": 100.12, "end": 103.4, "text": " from the ultimate guide to scaling video."}, {"start": 103.4, "end": 106.96000000000001, "text": " I'm gonna show you the pipeline parallelism"}, {"start": 106.96000000000001, "end": 109.84, "text": " and I'm finally gonna show you the data parallelism approach"}, {"start": 109.84, "end": 113.44, "text": " where you replicate the weights across multiple devices,"}, {"start": 113.44, "end": 117.0, "text": " for example, GPUs and then you do the data replication magic."}, {"start": 117.0, "end": 120.12, "text": " So you reduce the gradients using all reduce"}, {"start": 120.12, "end": 122.8, "text": " and then you broadcast the weights, et cetera, et cetera."}, {"start": 122.8, "end": 127.8, "text": " Okay, enough rambling, let's get started."}, {"start": 127.88, "end": 129.6, "text": " So first thing you have to do"}, {"start": 129.6, "end": 131.88, "text": " is obviously clone the repository."}, {"start": 131.88, "end": 134.32, "text": " So that's the Megatron DeepSpeed one."}, {"start": 134.32, "end": 136.44, "text": " I'm gonna obviously link the link down"}, {"start": 136.44, "end": 138.8, "text": " in the video description"}, {"start": 138.8, "end": 142.72, "text": " and just get started through this start fast with me."}, {"start": 142.72, "end": 144.2, "text": " So just follow the instructions"}, {"start": 144.2, "end": 146.04, "text": " to set up the conda environment"}, {"start": 146.04, "end": 147.72, "text": " and it should be straightforward."}, {"start": 147.72, "end": 150.56, "text": " Although I did experience some issues on Windows"}, {"start": 150.56, "end": 152.44, "text": " and so ultimately I had to shift"}, {"start": 152.44, "end": 154.2, "text": " to Windows subsystem for Linux"}, {"start": 154.2, "end": 156.76, "text": " and then everything worked out of the box."}, {"start": 156.76, "end": 158.44, "text": " So it's kind of funny because DeepSpeed"}, {"start": 158.44, "end": 159.78, "text": " is actually from Microsoft folks"}, {"start": 159.78, "end": 161.46, "text": " and it's not working on Windows."}, {"start": 161.46, "end": 162.84, "text": " I guess it's a bit ironic."}, {"start": 162.84, "end": 166.64, "text": " In any case, go through this with me, download the data."}, {"start": 166.64, "end": 169.28, "text": " We'll need these vocab and merges files"}, {"start": 169.28, "end": 171.44, "text": " for the tokenizer a bit later."}, {"start": 171.44, "end": 173.32, "text": " Go through the training script potentially."}, {"start": 173.32, "end": 175.32, "text": " I did go through it myself,"}, {"start": 175.32, "end": 177.84, "text": " but I decided to only show you the eval script"}, {"start": 177.84, "end": 180.07999999999998, "text": " because it contains all the necessary code"}, {"start": 180.07999999999998, "end": 183.68, "text": " for me to showcase the 3D parallelism approach."}, {"start": 183.68, "end": 186.48, "text": " This one is complex because it additionally uses DeepSpeed"}, {"start": 186.48, "end": 188.8, "text": " which is basically built on top of Megatron"}, {"start": 188.8, "end": 190.6, "text": " and so there is a lot of a lot,"}, {"start": 190.6, "end": 193.56, "text": " the stack of function calls is too deep"}, {"start": 193.56, "end": 196.16, "text": " and it will be super confusing for me to show you that"}, {"start": 196.16, "end": 198.12, "text": " unless I wanna teach you DeepSpeed"}, {"start": 198.12, "end": 200.0, "text": " which is not the point of this video."}, {"start": 200.0, "end": 201.72, "text": " Okay, long story short,"}, {"start": 202.8, "end": 207.76, "text": " basically go and take these parameters for the eval script."}, {"start": 207.76, "end": 211.76, "text": " Let me just find where the main pie script, so this one."}, {"start": 211.76, "end": 214.72, "text": " So, and get started from here."}, {"start": 214.72, "end": 218.68, "text": " Okay guys, so now I'm just gonna open up my VS code."}, {"start": 218.68, "end": 220.72, "text": " So you see here Megatron DeepSpeed,"}, {"start": 220.72, "end": 223.52, "text": " I have this Bloom 3 conda environment"}, {"start": 223.52, "end": 225.28, "text": " and I'm just gonna open up VS code."}, {"start": 225.28, "end": 228.12, "text": " So that's my ID of choice."}, {"start": 228.12, "end": 230.32, "text": " And let me first show you a couple of things"}, {"start": 230.32, "end": 233.84, "text": " I had to modify in the arguments here."}, {"start": 233.84, "end": 236.8, "text": " So let me just see what's important."}, {"start": 236.8, "end": 238.92000000000002, "text": " Again, you'll have to link this vocab file"}, {"start": 238.92000000000002, "end": 240.16, "text": " after you download it."}, {"start": 241.12, "end": 243.92000000000002, "text": " Just create a random TXT file"}, {"start": 243.92000000000002, "end": 248.6, "text": " and you don't really have to download the Viki text dataset"}, {"start": 248.6, "end": 252.12, "text": " because this is going to work out of the box like this."}, {"start": 252.12, "end": 256.52, "text": " And yeah, I mean, if your goal is to understand"}, {"start": 256.52, "end": 257.36, "text": " how the code works,"}, {"start": 257.36, "end": 259.56, "text": " if your goal is to actually do something with Viki text"}, {"start": 259.56, "end": 261.24, "text": " and obviously you wanna download the actual dataset."}, {"start": 261.24, "end": 264.52000000000004, "text": " For me, I just literally like pasted some random text"}, {"start": 264.52000000000004, "end": 265.96000000000004, "text": " inside of a TXT file."}, {"start": 266.8, "end": 267.8, "text": " Okay, the second thing,"}, {"start": 267.8, "end": 271.76, "text": " again, you'll have to download this file, the GPT-2 merges."}, {"start": 271.76, "end": 274.88, "text": " I also kinda commented out this load function"}, {"start": 274.88, "end": 277.44, "text": " and then we'll just be loading random weights,"}, {"start": 277.44, "end": 278.88, "text": " which again, does not matter"}, {"start": 278.88, "end": 280.6, "text": " because I don't care about the actual result,"}, {"start": 280.6, "end": 283.04, "text": " I just care about these to step through the code"}, {"start": 283.04, "end": 285.04, "text": " and show you what's going on."}, {"start": 285.04, "end": 287.40000000000003, "text": " Then I had this number of workers set to zero"}, {"start": 287.40000000000003, "end": 290.40000000000003, "text": " because it's easier again to step through the code."}, {"start": 290.40000000000003, "end": 293.04, "text": " And finally, I commented out the checkpoints activations"}, {"start": 293.04, "end": 296.16, "text": " because this does not make any sense for the eval script."}, {"start": 296.16, "end": 299.20000000000005, "text": " So I guess that's just a minor bug on their side."}, {"start": 299.20000000000005, "end": 302.04, "text": " Having said all that, let's start digging into the code."}, {"start": 302.04, "end": 304.52000000000004, "text": " So I'm gonna open up the main function here,"}, {"start": 304.52000000000004, "end": 306.08000000000004, "text": " the main script."}, {"start": 306.08000000000004, "end": 309.76, "text": " I'm going to grab the eval Viki text."}, {"start": 309.76, "end": 312.96000000000004, "text": " So that's the config for this file"}, {"start": 312.96000000000004, "end": 314.44, "text": " that I've just showed you."}, {"start": 314.44, "end": 317.16, "text": " And let me run this thing."}, {"start": 317.16, "end": 319.36, "text": " Okay, guys, let me quickly go through"}, {"start": 319.36, "end": 321.52, "text": " the initialize Megatron function."}, {"start": 321.52, "end": 322.56, "text": " Let's see what's going on there,"}, {"start": 322.56, "end": 326.12, "text": " although it's not that important for what I wanna show you."}, {"start": 327.04, "end": 328.6, "text": " Just make sure that we have CUDA,"}, {"start": 328.6, "end": 330.84, "text": " so that we have a GPU device."}, {"start": 330.84, "end": 332.64, "text": " We set some global variables here."}, {"start": 332.64, "end": 333.46, "text": " This is kinda interesting."}, {"start": 333.46, "end": 338.04, "text": " So there is the, whoops, don't show this again, please."}, {"start": 338.04, "end": 340.92, "text": " Okay, so they first parse the arguments"}, {"start": 340.92, "end": 344.04, "text": " and let me kinda dig into that function first."}, {"start": 344.04, "end": 346.32, "text": " So just to show you how it looks like."}, {"start": 346.32, "end": 349.5, "text": " So the idea here is to, they have these global variables."}, {"start": 349.5, "end": 352.08000000000004, "text": " So they parse the arguments, they set the,"}, {"start": 352.08000000000004, "end": 355.72, "text": " they assign those arguments to this global variable,"}, {"start": 355.72, "end": 357.76000000000005, "text": " and then it's gonna be visible across the code base."}, {"start": 357.76000000000005, "end": 361.44, "text": " So that's kinda how they went about designing this thing."}, {"start": 361.44, "end": 365.26, "text": " I'm gonna skip through the, I'm gonna skip, yeah."}, {"start": 365.26, "end": 368.38, "text": " I'm just gonna maybe quickly show you how complex this is."}, {"start": 368.38, "end": 370.8, "text": " There is a lot of arguments that go into this,"}, {"start": 370.8, "end": 372.08000000000004, "text": " but I'm gonna skip all of those,"}, {"start": 372.08, "end": 374.15999999999997, "text": " and we're just gonna see how that looks like"}, {"start": 375.24, "end": 377.21999999999997, "text": " when we start stepping through the code."}, {"start": 377.21999999999997, "end": 379.97999999999996, "text": " So as you can see, a lot of error checking,"}, {"start": 379.97999999999996, "end": 382.71999999999997, "text": " and I'm gonna hit F5 and hopefully get here, yeah."}, {"start": 382.71999999999997, "end": 384.91999999999996, "text": " Okay, so let's continue here."}, {"start": 384.91999999999996, "end": 386.76, "text": " So that's the parsing of arguments."}, {"start": 386.76, "end": 389.91999999999996, "text": " Then I'm gonna skip this part."}, {"start": 389.91999999999996, "end": 391.64, "text": " We don't care about that one."}, {"start": 391.64, "end": 393.64, "text": " And then we build a tokenizer."}, {"start": 393.64, "end": 398.64, "text": " Again, that's just going to be GPT tokenizer."}, {"start": 399.08, "end": 400.76, "text": " You can see it here."}, {"start": 400.76, "end": 405.2, "text": " So we're gonna kinda go through this GPT-2BPE tokenizer,"}, {"start": 405.2, "end": 407.76, "text": " and that's a common tokenizer used"}, {"start": 407.76, "end": 409.52, "text": " throughout many code bases."}, {"start": 409.52, "end": 412.36, "text": " And just doing this sample coding series,"}, {"start": 412.36, "end": 415.8, "text": " I realized how pervasive this tokenizer is,"}, {"start": 415.8, "end": 419.56, "text": " and I first covered it in my OpenAI clip video,"}, {"start": 419.56, "end": 420.52, "text": " so do check that one out."}, {"start": 420.52, "end": 422.15999999999997, "text": " I'm gonna link it somewhere above"}, {"start": 422.15999999999997, "end": 425.15999999999997, "text": " if you are curious to learn more about the tokenizers."}, {"start": 425.15999999999997, "end": 428.8, "text": " But for now, I'm just gonna basically disable"}, {"start": 428.8, "end": 433.04, "text": " the breakpoints and skip over the actual tokenizer."}, {"start": 433.04, "end": 434.64, "text": " And let's continue."}, {"start": 434.64, "end": 437.8, "text": " So we just do some padding of the vocab"}, {"start": 437.8, "end": 442.28000000000003, "text": " such that it's basically suitable to model parallelism."}, {"start": 442.28000000000003, "end": 445.12, "text": " Otherwise, you'd have problems."}, {"start": 445.12, "end": 446.36, "text": " And that's it."}, {"start": 446.36, "end": 448.72, "text": " Now we just set up some TensorBoard."}, {"start": 448.72, "end": 450.84000000000003, "text": " This is kinda cool, the code carbon tracker."}, {"start": 450.84000000000003, "end": 453.28000000000003, "text": " So they did have this functionality"}, {"start": 453.28000000000003, "end": 456.24, "text": " where they track how much are they impacting"}, {"start": 456.24, "end": 460.32, "text": " the environment by just measuring the amount of CO2 emission."}, {"start": 460.32, "end": 462.84000000000003, "text": " But they did comment it out, so as you can see here,"}, {"start": 462.84000000000003, "end": 467.0, "text": " they don't actually use it, but it's a cool idea, I guess."}, {"start": 467.0, "end": 468.76, "text": " Because they say it's very unstable,"}, {"start": 468.76, "end": 469.96000000000004, "text": " blah, blah, blah, whatever."}, {"start": 469.96000000000004, "end": 473.16, "text": " Okay, let's step out of this, and let's continue."}, {"start": 473.16, "end": 474.24, "text": " And that's it."}, {"start": 475.2, "end": 478.2, "text": " Okay, let me toggle on the breakpoints again."}, {"start": 478.2, "end": 479.92, "text": " I'm gonna skip most of these things."}, {"start": 479.92, "end": 481.86, "text": " So get our, so now this is the cool thing."}, {"start": 481.86, "end": 484.0, "text": " Wherever you are in the code base,"}, {"start": 484.0, "end": 486.96, "text": " you'll literally just grab that global variable"}, {"start": 486.96, "end": 488.84, "text": " and you're golden."}, {"start": 488.84, "end": 491.08, "text": " You don't have to pass the arguments"}, {"start": 491.08, "end": 493.68, "text": " through the function argument list."}, {"start": 494.72, "end": 497.08, "text": " Well, it's debatable whether that's the best thing to do,"}, {"start": 497.08, "end": 499.68, "text": " but yeah, I think it's nice."}, {"start": 499.68, "end": 501.88, "text": " There is a couple of interesting things in this function,"}, {"start": 501.88, "end": 504.4, "text": " so I'm gonna quickly step through it."}, {"start": 504.4, "end": 507.08, "text": " There is some initialization of the,"}, {"start": 507.08, "end": 508.44, "text": " so I'm using the NCCL,"}, {"start": 508.44, "end": 511.36, "text": " which is NVIDIA Collective Communication Library"}, {"start": 511.36, "end": 515.72, "text": " for all of those old reviews and other operations"}, {"start": 515.72, "end": 519.52, "text": " that you use to communicate and synchronize the GPUs."}, {"start": 519.52, "end": 520.72, "text": " So it's gonna happen in this function."}, {"start": 520.72, "end": 523.88, "text": " I'm not gonna step into it because it's fairly complex."}, {"start": 523.88, "end": 525.84, "text": " Just library code."}, {"start": 525.84, "end": 529.1, "text": " And then let me just see what else is interesting here,"}, {"start": 530.04, "end": 531.36, "text": " which set the random seed"}, {"start": 531.36, "end": 534.44, "text": " just for reproducibility purposes, obviously."}, {"start": 534.44, "end": 537.16, "text": " And yeah, that's it."}, {"start": 537.16, "end": 539.78, "text": " And then in the compile dependencies,"}, {"start": 539.78, "end": 541.72, "text": " there is a couple of interesting things going on."}, {"start": 541.72, "end": 544.04, "text": " Again, we are going to,"}, {"start": 544.04, "end": 545.72, "text": " so first there's some,"}, {"start": 545.72, "end": 547.6, "text": " this part is not that interesting,"}, {"start": 547.6, "end": 552.52, "text": " but there's one part where they load certain CUDA kernels,"}, {"start": 552.52, "end": 554.3, "text": " and we're gonna later actually use them"}, {"start": 554.3, "end": 558.92, "text": " and leverage that C++ code for extra performance."}, {"start": 558.92, "end": 561.16, "text": " So I'm gonna skip through most of these."}, {"start": 561.16, "end": 565.48, "text": " We actually saw something similar, if not the same."}, {"start": 565.48, "end": 566.76, "text": " Actually, it's the same code"}, {"start": 566.76, "end": 568.9399999999999, "text": " because Opt was also using Megatron,"}, {"start": 568.94, "end": 570.8800000000001, "text": " so my previous video was also using Megatron"}, {"start": 570.8800000000001, "end": 572.44, "text": " in the background."}, {"start": 572.44, "end": 574.6800000000001, "text": " So we've seen the same code here."}, {"start": 574.6800000000001, "end": 577.6800000000001, "text": " So what we do is we load these kernels."}, {"start": 577.6800000000001, "end": 579.36, "text": " I'm gonna step into it quickly"}, {"start": 579.36, "end": 581.1800000000001, "text": " and just show you how it looks like,"}, {"start": 581.1800000000001, "end": 584.2600000000001, "text": " but it's nothing super crucial."}, {"start": 585.7600000000001, "end": 589.0, "text": " I mean, not related to the 3D parallelism in that sense."}, {"start": 589.0, "end": 592.6, "text": " So here you can see we're just loading various kernels"}, {"start": 594.58, "end": 597.5400000000001, "text": " for like scaled up or for softmax"}, {"start": 597.54, "end": 600.76, "text": " and then for some other softmax and layer norm,"}, {"start": 600.76, "end": 601.5999999999999, "text": " et cetera, et cetera."}, {"start": 601.5999999999999, "end": 603.8, "text": " So I'm gonna put the breakpoint here,"}, {"start": 603.8, "end": 606.0799999999999, "text": " hit F5 and just exit here."}, {"start": 606.0799999999999, "end": 609.28, "text": " So it was more of FYI than anything else."}, {"start": 609.28, "end": 612.7199999999999, "text": " Okay, there is some barrier function call"}, {"start": 612.7199999999999, "end": 614.88, "text": " which basically synchronizes across."}, {"start": 614.88, "end": 617.68, "text": " So all of the, it basically waits,"}, {"start": 617.68, "end": 622.5999999999999, "text": " all of the different devices, GPUs will wait at that line"}, {"start": 622.5999999999999, "end": 625.0, "text": " until the other GPUs hit that same line"}, {"start": 625.0, "end": 629.0, "text": " and then the barrier is gonna release the execution"}, {"start": 629.0, "end": 630.28, "text": " across all of the devices."}, {"start": 630.28, "end": 631.68, "text": " That's what the barrier does."}, {"start": 631.68, "end": 635.08, "text": " But I have a single GPU, so that's not interesting."}, {"start": 635.08, "end": 636.98, "text": " And even though I have a single GPU,"}, {"start": 636.98, "end": 640.22, "text": " I'm still going to be able to show you the 3D parallelism"}, {"start": 640.22, "end": 641.46, "text": " as we'll soon see."}, {"start": 641.46, "end": 643.76, "text": " Okay, let's just exit here"}, {"start": 643.76, "end": 646.36, "text": " and that's pretty much everything I wanted to show you."}, {"start": 646.36, "end": 648.6, "text": " Now let's get into the main script."}, {"start": 648.6, "end": 652.4, "text": " So because I'm using Viki Text 103,"}, {"start": 652.4, "end": 656.04, "text": " I'm gonna take this particular main function"}, {"start": 656.04, "end": 658.06, "text": " and then let's now enter here."}, {"start": 658.06, "end": 661.1, "text": " Okay, so we fetched the arguments"}, {"start": 661.1, "end": 662.8, "text": " and let's see what's going on."}, {"start": 662.8, "end": 666.0799999999999, "text": " So because we are using Viki Text,"}, {"start": 666.0799999999999, "end": 668.88, "text": " we're gonna set the eval metric to loss."}, {"start": 668.88, "end": 670.64, "text": " That's not that important."}, {"start": 670.64, "end": 674.02, "text": " And this is where stuff becomes interesting."}, {"start": 674.02, "end": 676.5799999999999, "text": " We're now going to see how the parallelism"}, {"start": 676.5799999999999, "end": 679.0, "text": " is implemented in the code here."}, {"start": 679.0, "end": 682.26, "text": " Okay, so let me re-enable the breakpoints"}, {"start": 682.26, "end": 685.26, "text": " and let's continue stepping through this."}, {"start": 685.26, "end": 687.4, "text": " So here is the model provider."}, {"start": 687.4, "end": 689.3199999999999, "text": " So we enter this get model function."}, {"start": 689.3199999999999, "end": 691.28, "text": " Let's now see what's going on."}, {"start": 691.28, "end": 694.16, "text": " Okay, so we hit this else branch."}, {"start": 694.16, "end": 697.24, "text": " And now, because I have a single device"}, {"start": 697.24, "end": 701.48, "text": " and because of that, my pipeline is only of length one."}, {"start": 701.48, "end": 703.2, "text": " I have a single stage."}, {"start": 703.2, "end": 705.48, "text": " And that's why both the pre-process"}, {"start": 705.48, "end": 708.38, "text": " as well as the post-process here"}, {"start": 708.38, "end": 711.2, "text": " are gonna be set to true, as you can see here."}, {"start": 711.2, "end": 713.44, "text": " But if this was on a,"}, {"start": 713.44, "end": 715.0400000000001, "text": " if you were running on a,"}, {"start": 715.0400000000001, "end": 718.5600000000001, "text": " let's say a two by two topology setup,"}, {"start": 718.5600000000001, "end": 720.4000000000001, "text": " so we have four GPUs in total"}, {"start": 720.4000000000001, "end": 722.98, "text": " and I'm gonna use that topology as a running example."}, {"start": 722.98, "end": 725.6, "text": " Let me kind of open up the OneNote and show you what I mean."}, {"start": 725.6, "end": 727.4000000000001, "text": " So let me go here."}, {"start": 727.4000000000001, "end": 731.3000000000001, "text": " And this is going to be the running topology I'll be using."}, {"start": 731.3000000000001, "end": 734.32, "text": " So basically, I'm gonna have something like this."}, {"start": 734.32, "end": 737.3000000000001, "text": " So we have a GPU here."}, {"start": 738.34, "end": 740.84, "text": " We have a GPU here."}, {"start": 740.84, "end": 743.12, "text": " We have another GPU here."}, {"start": 743.12, "end": 745.94, "text": " And we have yet another GPU here."}, {"start": 745.94, "end": 748.6, "text": " So we're gonna have three types of parallelism,"}, {"start": 748.6, "end": 750.12, "text": " but I'm gonna show you only the model"}, {"start": 750.12, "end": 751.88, "text": " and the pipeline here on the image."}, {"start": 751.88, "end": 755.76, "text": " So these ones here are going to be stage one."}, {"start": 755.76, "end": 757.8000000000001, "text": " So stage or stage zero,"}, {"start": 757.8000000000001, "end": 760.2, "text": " and this is gonna be stage one, okay?"}, {"start": 760.2, "end": 763.76, "text": " So basically that means we'll have half of the layers"}, {"start": 763.76, "end": 765.44, "text": " of our transformer being here."}, {"start": 765.44, "end": 768.76, "text": " Maybe if we have 24 layers, 12 will be here."}, {"start": 768.76, "end": 770.6, "text": " So we'll have 12 layers here."}, {"start": 770.6, "end": 772.64, "text": " And we'll have 12 layers here."}, {"start": 773.64, "end": 776.28, "text": " Next up, I'm gonna do something like this."}, {"start": 776.28, "end": 777.82, "text": " So we're gonna have model parallelism."}, {"start": 777.82, "end": 780.3000000000001, "text": " So across this dimension, we have MP."}, {"start": 781.22, "end": 783.16, "text": " And across this dimension here,"}, {"start": 783.16, "end": 786.52, "text": " we have pipeline parallelism PP, okay?"}, {"start": 786.52, "end": 789.32, "text": " So this one is gonna be zero and one."}, {"start": 789.32, "end": 791.64, "text": " So this is gonna be the topology"}, {"start": 791.64, "end": 794.0400000000001, "text": " that I want you to have in your mind."}, {"start": 794.0400000000001, "end": 796.1600000000001, "text": " And we'll see how it's gonna be useful"}, {"start": 796.1600000000001, "end": 797.62, "text": " in a couple of minutes."}, {"start": 797.62, "end": 802.62, "text": " Okay, again, keep in mind that in the setup I just showed you,"}, {"start": 803.94, "end": 808.58, "text": " the devices from the stage zero of the pipeline parallelism"}, {"start": 808.58, "end": 811.58, "text": " will have, this will be true and this will be false."}, {"start": 811.58, "end": 814.16, "text": " And then the second stage will have,"}, {"start": 814.16, "end": 816.34, "text": " this will be true and this will be false."}, {"start": 816.34, "end": 817.94, "text": " And if you have even more stages,"}, {"start": 817.94, "end": 821.98, "text": " then some of these devices will have both of these"}, {"start": 821.98, "end": 824.78, "text": " pre and post process flags set to false"}, {"start": 824.78, "end": 826.98, "text": " because they are neither the first stage"}, {"start": 826.98, "end": 829.14, "text": " nor the last stage of the pipeline."}, {"start": 829.14, "end": 830.44, "text": " Okay, hopefully that makes sense."}, {"start": 830.44, "end": 833.04, "text": " Now let's enter this function here."}, {"start": 834.1, "end": 836.66, "text": " So yeah, we just set this to true"}, {"start": 836.66, "end": 839.2, "text": " because we're dealing with loss metric."}, {"start": 839.2, "end": 840.04, "text": " Let's continue."}, {"start": 840.04, "end": 841.34, "text": " And here is the GPT model."}, {"start": 841.34, "end": 844.1, "text": " So I'm gonna be using the GPT transformer"}, {"start": 844.1, "end": 847.02, "text": " as the underlying architecture."}, {"start": 847.02, "end": 850.1, "text": " So let's enter here and let's see what's going on."}, {"start": 850.1, "end": 852.22, "text": " So it's gonna inherit from Megatron module,"}, {"start": 852.22, "end": 855.26, "text": " which is a light wrapper around torch modules."}, {"start": 855.26, "end": 858.14, "text": " I'm not gonna step through the constructor."}, {"start": 858.14, "end": 859.5, "text": " Let's fetch the arguments here"}, {"start": 859.5, "end": 862.9, "text": " and let's just kind of skip across all of these details."}, {"start": 862.9, "end": 864.58, "text": " This is the interesting part, the language model."}, {"start": 864.58, "end": 868.0, "text": " So now we're gonna enter that function."}, {"start": 869.54, "end": 873.38, "text": " And here, okay, again, we just fetch the arguments."}, {"start": 873.38, "end": 874.78, "text": " We continue here."}, {"start": 874.78, "end": 876.84, "text": " Now there is a transformer language model."}, {"start": 876.84, "end": 878.46, "text": " So that's gonna be interesting one."}, {"start": 878.46, "end": 881.86, "text": " So let me hit F5, let's get here."}, {"start": 881.86, "end": 884.3, "text": " So it's inheriting again from the Megatron module."}, {"start": 884.3, "end": 888.26, "text": " Okay, we can step across all of these."}, {"start": 888.26, "end": 890.8199999999999, "text": " I'm gonna focus only on the parts that are interesting."}, {"start": 890.8199999999999, "end": 891.66, "text": " So here's the embedding."}, {"start": 891.66, "end": 892.8199999999999, "text": " This is the interesting part."}, {"start": 892.8199999999999, "end": 897.14, "text": " So let me kind of put a break point here, hit F5."}, {"start": 897.14, "end": 900.06, "text": " Let's see how this is going to work."}, {"start": 900.06, "end": 903.3399999999999, "text": " Okay, so the embedding table"}, {"start": 903.3399999999999, "end": 905.3, "text": " is gonna have a couple of details."}, {"start": 905.3, "end": 907.3399999999999, "text": " So here is the Voke parallel embedding."}, {"start": 907.3399999999999, "end": 909.28, "text": " That's what they're using in the background."}, {"start": 909.28, "end": 914.28, "text": " Parallel suggests that we are basically gonna do"}, {"start": 915.5, "end": 917.86, "text": " a model parallelism of this embedding table."}, {"start": 917.86, "end": 920.78, "text": " And by the way, let me just go back a step up here."}, {"start": 920.78, "end": 923.42, "text": " You can see that we only enter here"}, {"start": 923.42, "end": 925.9399999999999, "text": " if the pre-process is set to true."}, {"start": 925.9399999999999, "end": 929.8399999999999, "text": " And that's only true for the first stage in the pipeline"}, {"start": 929.8399999999999, "end": 933.4399999999999, "text": " or zeroth index, as I kind of made the notation."}, {"start": 933.4399999999999, "end": 936.36, "text": " But like the first devices in the pipeline"}, {"start": 936.36, "end": 939.06, "text": " will only have this embedding table, which makes sense,"}, {"start": 939.06, "end": 941.4599999999999, "text": " because we wanna split the model layers"}, {"start": 941.4599999999999, "end": 944.14, "text": " across different stages."}, {"start": 944.14, "end": 945.9399999999999, "text": " Okay, so let's get back here."}, {"start": 945.9399999999999, "end": 948.3399999999999, "text": " So here is the Voke parallel embedding."}, {"start": 948.3399999999999, "end": 949.6199999999999, "text": " Let's enter here."}, {"start": 949.6199999999999, "end": 951.54, "text": " Let's gonna hit F5."}, {"start": 951.54, "end": 953.8399999999999, "text": " Whoops, let me hit F5 and enter there."}, {"start": 953.8399999999999, "end": 957.8199999999999, "text": " Okay, so again, let's see what's gonna happen here."}, {"start": 957.8199999999999, "end": 959.26, "text": " So this is the interesting part."}, {"start": 959.26, "end": 962.6999999999999, "text": " This get tensor model parallel world size"}, {"start": 962.6999999999999, "end": 964.7199999999999, "text": " is going to in the topology"}, {"start": 964.7199999999999, "end": 966.7399999999999, "text": " that I showed you a couple of minutes ago,"}, {"start": 966.7399999999999, "end": 968.16, "text": " it will return two."}, {"start": 968.16, "end": 969.54, "text": " So for my particular setup,"}, {"start": 969.54, "end": 970.74, "text": " it's obviously gonna return one"}, {"start": 970.74, "end": 972.42, "text": " because they have a single GPU,"}, {"start": 972.42, "end": 974.24, "text": " but in general, it's gonna return two."}, {"start": 974.24, "end": 977.74, "text": " And now let's see what's the repercussion of that."}, {"start": 977.74, "end": 982.74, "text": " So here, what happens is we fetch the rank."}, {"start": 982.86, "end": 985.66, "text": " So the tensor model parallel rank,"}, {"start": 985.66, "end": 990.66, "text": " and that's going to be zero for the basically devices here."}, {"start": 991.5, "end": 993.3399999999999, "text": " So it's gonna be zero for these devices."}, {"start": 993.3399999999999, "end": 996.4599999999999, "text": " It's gonna be one for these devices here, okay?"}, {"start": 996.46, "end": 998.82, "text": " So let me go back to code."}, {"start": 998.82, "end": 1003.22, "text": " Now with that in mind and having this thing,"}, {"start": 1003.22, "end": 1008.22, "text": " so this is going to be two in the example I'm showing you."}, {"start": 1009.62, "end": 1011.14, "text": " Let's now enter this function"}, {"start": 1011.14, "end": 1012.38, "text": " and let me show you what it does."}, {"start": 1012.38, "end": 1015.86, "text": " So it takes the global vocab size,"}, {"start": 1015.86, "end": 1020.86, "text": " which is this big 50K slots in the embedding table."}, {"start": 1021.1, "end": 1023.5400000000001, "text": " And we're going to divide it by two"}, {"start": 1023.54, "end": 1026.82, "text": " because our model parallelism world size"}, {"start": 1026.82, "end": 1030.02, "text": " is actually of size two in the running example I'm using."}, {"start": 1030.02, "end": 1033.82, "text": " Okay, so it will be something like 25K here."}, {"start": 1033.82, "end": 1036.62, "text": " Okay, and then what we do is we enter this function here"}, {"start": 1037.7, "end": 1040.6599999999999, "text": " and you can see here, so for rank zero,"}, {"start": 1040.6599999999999, "end": 1043.62, "text": " we would have this thing would be zero"}, {"start": 1043.62, "end": 1046.18, "text": " because we multiply it by zero 25K."}, {"start": 1046.18, "end": 1049.54, "text": " And then this would be 25K plus,"}, {"start": 1050.94, "end": 1053.42, "text": " no, sorry, this would be zero plus 25K."}, {"start": 1053.42, "end": 1054.8600000000001, "text": " So that would be 25K."}, {"start": 1054.8600000000001, "end": 1058.5800000000002, "text": " So that means we would return zero and 25,000 here"}, {"start": 1058.5800000000002, "end": 1059.8400000000001, "text": " as a result."}, {"start": 1061.5, "end": 1063.7, "text": " Obviously, in my case, I'm returning 50K,"}, {"start": 1063.7, "end": 1064.7, "text": " but bear with me."}, {"start": 1064.7, "end": 1069.0, "text": " Okay, so that's the first thing we need to understand."}, {"start": 1069.0, "end": 1070.98, "text": " Now those variables are gonna be used"}, {"start": 1070.98, "end": 1074.66, "text": " to split the embedding table across multiple devices."}, {"start": 1074.66, "end": 1076.38, "text": " That's the whole point."}, {"start": 1076.38, "end": 1077.66, "text": " So now you can see here,"}, {"start": 1077.66, "end": 1079.94, "text": " the number of embeddings would be 25,000"}, {"start": 1079.94, "end": 1082.98, "text": " if we had the topology I'm using."}, {"start": 1082.98, "end": 1085.1, "text": " We can skip this, we can skip this."}, {"start": 1085.1, "end": 1088.74, "text": " Those are just some bits and bytes or something,"}, {"start": 1088.74, "end": 1092.02, "text": " basically where they use integer eight precision,"}, {"start": 1092.02, "end": 1092.8600000000001, "text": " et cetera, et cetera."}, {"start": 1092.8600000000001, "end": 1095.3, "text": " We're gonna ignore that in this video."}, {"start": 1095.3, "end": 1096.6200000000001, "text": " So here is the interesting part."}, {"start": 1096.6200000000001, "end": 1099.3, "text": " So now we create the weight of the embedding table."}, {"start": 1099.3, "end": 1102.78, "text": " So we can see here, we are passing this variable,"}, {"start": 1102.78, "end": 1105.06, "text": " which would be 25,000 again,"}, {"start": 1105.06, "end": 1107.0, "text": " and we're passing the embedding dimension,"}, {"start": 1107.0, "end": 1109.34, "text": " which is 1,024, blah, blah, blah."}, {"start": 1109.34, "end": 1110.18, "text": " That's it."}, {"start": 1110.18, "end": 1113.0600000000002, "text": " So the weight would be, as you can see,"}, {"start": 1113.0600000000002, "end": 1116.94, "text": " 25,000 times 1,000 on one device"}, {"start": 1116.94, "end": 1121.0600000000002, "text": " and on those devices that have a rank zero,"}, {"start": 1121.0600000000002, "end": 1123.98, "text": " and then it will be additionally 25,000"}, {"start": 1123.98, "end": 1127.74, "text": " for those devices that have rank one,"}, {"start": 1127.74, "end": 1130.02, "text": " again, talking about model parallelism here."}, {"start": 1130.02, "end": 1132.8200000000002, "text": " Okay, so we just now initialize it and that's it."}, {"start": 1132.8200000000002, "end": 1135.5800000000002, "text": " So now you've seen for the first time"}, {"start": 1135.5800000000002, "end": 1138.1200000000001, "text": " how the model parallelism looks like."}, {"start": 1138.1200000000001, "end": 1140.0600000000002, "text": " So let me just exit here."}, {"start": 1140.06, "end": 1141.74, "text": " Let's now continue."}, {"start": 1141.74, "end": 1143.86, "text": " The second thing we want to do here"}, {"start": 1143.86, "end": 1146.98, "text": " is create some positional embeddings."}, {"start": 1146.98, "end": 1149.1399999999999, "text": " Now these, because they are much smaller"}, {"start": 1149.1399999999999, "end": 1151.8, "text": " than the vocab embedding table,"}, {"start": 1151.8, "end": 1154.86, "text": " they're just going to allocate the whole thing"}, {"start": 1154.86, "end": 1157.46, "text": " across all of the different devices."}, {"start": 1157.46, "end": 1161.8999999999999, "text": " So we have like complete, we can see here 1,024,"}, {"start": 1161.8999999999999, "end": 1163.46, "text": " because that's gonna be our sequence"}, {"start": 1163.46, "end": 1166.72, "text": " and that's how many positionally encodings we need."}, {"start": 1166.72, "end": 1169.3600000000001, "text": " Okay, let's continue."}, {"start": 1170.44, "end": 1172.4, "text": " We just initialize those"}, {"start": 1172.4, "end": 1175.4, "text": " and let's see now what's happening here."}, {"start": 1175.4, "end": 1177.0, "text": " We can skip all of those."}, {"start": 1177.0, "end": 1179.2, "text": " We just add some dropout, blah, blah, blah."}, {"start": 1179.2, "end": 1181.52, "text": " I keep entering this function, I'm not sure why,"}, {"start": 1181.52, "end": 1184.68, "text": " but we now have the embedding and let's continue."}, {"start": 1184.68, "end": 1186.72, "text": " Now let me just update the diagram here."}, {"start": 1186.72, "end": 1188.26, "text": " I think this is gonna be useful."}, {"start": 1188.26, "end": 1191.16, "text": " So let me just take some different color."}, {"start": 1191.16, "end": 1194.38, "text": " So we have the embedding table currently placed like this."}, {"start": 1194.38, "end": 1196.7600000000002, "text": " So we have embedding table here"}, {"start": 1196.7600000000002, "end": 1198.48, "text": " and we have embedding table here."}, {"start": 1199.72, "end": 1201.88, "text": " And the size here is 25K"}, {"start": 1202.92, "end": 1205.8000000000002, "text": " and the size here is 25K as well."}, {"start": 1205.8000000000002, "end": 1208.42, "text": " Okay, let's not go back to decode."}, {"start": 1210.48, "end": 1212.5200000000002, "text": " Now we have the parallel transformer."}, {"start": 1212.5200000000002, "end": 1215.64, "text": " Let's see how the transformer is going to be split."}, {"start": 1215.64, "end": 1219.3200000000002, "text": " Okay, so let's enter the parallel transformer"}, {"start": 1220.2, "end": 1221.4, "text": " and here it is."}, {"start": 1221.4, "end": 1225.96, "text": " Okay, again, skipping these details until we get here."}, {"start": 1225.96, "end": 1227.8200000000002, "text": " So here is an interesting assert."}, {"start": 1227.8200000000002, "end": 1231.2800000000002, "text": " So the number of layers has to be divisible"}, {"start": 1231.2800000000002, "end": 1235.44, "text": " by the world size for the pipeline parallelism."}, {"start": 1235.44, "end": 1238.38, "text": " In the running example I'm using, it's gonna be two."}, {"start": 1238.38, "end": 1240.76, "text": " So that means you have to be like,"}, {"start": 1240.76, "end": 1243.0800000000002, "text": " the number of layers has to be divisible by two."}, {"start": 1243.0800000000002, "end": 1246.16, "text": " In my actual setup, I have, this will be one, so yeah."}, {"start": 1246.16, "end": 1248.0400000000002, "text": " Okay, now this is the interesting line."}, {"start": 1248.0400000000002, "end": 1251.2, "text": " So we see here that we divide the number of layers"}, {"start": 1251.2, "end": 1253.42, "text": " by the number of stages."}, {"start": 1253.42, "end": 1255.68, "text": " And again, so because we have 24,"}, {"start": 1255.68, "end": 1260.4, "text": " that means we would now have 12 layers"}, {"start": 1260.4, "end": 1263.8400000000001, "text": " being allocated on this particular device, okay?"}, {"start": 1263.8400000000001, "end": 1266.0, "text": " So again, keep in mind, this code,"}, {"start": 1266.0, "end": 1268.02, "text": " so the thing I'm just stepping through,"}, {"start": 1268.02, "end": 1270.16, "text": " this same thing would be happening in parallel"}, {"start": 1270.16, "end": 1271.44, "text": " on four devices."}, {"start": 1271.44, "end": 1273.74, "text": " And each of those devices have different contexts"}, {"start": 1273.74, "end": 1276.48, "text": " because they lie in a different part of the topology"}, {"start": 1276.48, "end": 1279.1200000000001, "text": " and because of that they will be entering different types,"}, {"start": 1279.12, "end": 1281.1599999999999, "text": " different code, they will be going through the different"}, {"start": 1281.1599999999999, "end": 1283.0, "text": " code paths, that's the main idea."}, {"start": 1283.0, "end": 1284.9599999999998, "text": " Hopefully you understand that."}, {"start": 1284.9599999999998, "end": 1286.6799999999998, "text": " Let's continue here."}, {"start": 1286.6799999999998, "end": 1289.2399999999998, "text": " Okay, so now we create the offset."}, {"start": 1289.2399999999998, "end": 1293.08, "text": " So we grab the rank and we multiply it"}, {"start": 1293.08, "end": 1294.4399999999998, "text": " by the number of layers."}, {"start": 1294.4399999999998, "end": 1298.08, "text": " So that means that the offset for these two devices here"}, {"start": 1298.08, "end": 1301.6, "text": " is gonna be zero, for these two is gonna be zero."}, {"start": 1301.6, "end": 1309.6, "text": " And for the two ones here is gonna be 12, right?"}, {"start": 1309.6799999999998, "end": 1310.6799999999998, "text": " So let's go back here."}, {"start": 1310.6799999999998, "end": 1313.5, "text": " So it's gonna be one, the rank will be one for those"}, {"start": 1313.5, "end": 1315.6, "text": " and then times the number of layers, 12."}, {"start": 1315.6, "end": 1317.5, "text": " So offset is gonna be 12 for them."}, {"start": 1317.5, "end": 1319.8799999999999, "text": " So let's see how offset plays a role here."}, {"start": 1319.8799999999999, "end": 1324.5, "text": " So we now form 12 layers and we have the offset here."}, {"start": 1324.5, "end": 1327.6599999999999, "text": " So you can see that we are passing that index information"}, {"start": 1327.6599999999999, "end": 1328.84, "text": " into this build layer function."}, {"start": 1328.84, "end": 1331.0, "text": " So let's enter the build layer function."}, {"start": 1331.0, "end": 1332.6, "text": " Here it is."}, {"start": 1332.6, "end": 1335.52, "text": " And there we start building the actual transformer layers."}, {"start": 1335.52, "end": 1337.64, "text": " We'll do this 12 times."}, {"start": 1337.64, "end": 1341.04, "text": " So let me just enter this layer here."}, {"start": 1341.04, "end": 1344.8, "text": " Let me hit F5 and we are here."}, {"start": 1344.8, "end": 1348.34, "text": " So again, let me just focus on the important parts,"}, {"start": 1348.34, "end": 1350.32, "text": " ignoring many details obviously,"}, {"start": 1350.32, "end": 1352.76, "text": " otherwise it will be too bothersome."}, {"start": 1352.76, "end": 1354.1, "text": " Okay, so first things first,"}, {"start": 1354.1, "end": 1356.8, "text": " we want to instantiate a layer norm."}, {"start": 1356.8, "end": 1358.48, "text": " Okay, let's see how the layer norm looks like."}, {"start": 1358.48, "end": 1361.76, "text": " The layer norm has some fused components,"}, {"start": 1361.76, "end": 1362.8600000000001, "text": " which is interesting."}, {"start": 1362.8600000000001, "end": 1366.3, "text": " So that means we'll be using optimized CUDA kernels,"}, {"start": 1366.3, "end": 1368.84, "text": " fancy word for saying functions,"}, {"start": 1368.84, "end": 1371.28, "text": " implemented in CUDA library"}, {"start": 1371.28, "end": 1373.64, "text": " that are running super fast on GPUs obviously,"}, {"start": 1373.64, "end": 1375.56, "text": " on Nvidia GPUs."}, {"start": 1375.56, "end": 1377.92, "text": " So here you can see we import"}, {"start": 1377.92, "end": 1381.24, "text": " this fused mixed precision layer norm CUDA."}, {"start": 1381.24, "end": 1383.9, "text": " So let me show you where that was constructed."}, {"start": 1383.9, "end": 1385.1, "text": " And I actually showed you that"}, {"start": 1385.1, "end": 1386.0, "text": " in the beginning of the video."}, {"start": 1386.0, "end": 1388.72, "text": " So here, let me just move this thing here."}, {"start": 1388.72, "end": 1391.12, "text": " So if I were to copy paste this here,"}, {"start": 1391.12, "end": 1393.8, "text": " so you can see here we are loading"}, {"start": 1393.8, "end": 1396.4, "text": " that precise variable is loaded"}, {"start": 1396.4, "end": 1398.84, "text": " using the CPP extension load helper."}, {"start": 1398.84, "end": 1399.76, "text": " So we kind of passed,"}, {"start": 1399.76, "end": 1402.16, "text": " I kind of showed you this function before,"}, {"start": 1402.16, "end": 1404.82, "text": " and that's where it was initialized."}, {"start": 1404.82, "end": 1407.32, "text": " Now let's get back to the actual code."}, {"start": 1407.32, "end": 1409.68, "text": " Let me go to the top of the stack."}, {"start": 1409.68, "end": 1412.08, "text": " We import that kernel"}, {"start": 1412.08, "end": 1415.36, "text": " and yeah, we're gonna then use it a bit later."}, {"start": 1415.36, "end": 1417.28, "text": " So everything else is not that interesting."}, {"start": 1417.28, "end": 1419.6799999999998, "text": " I'm not gonna focus on the layer norm"}, {"start": 1419.6799999999998, "end": 1423.78, "text": " because this one is not going to be split across devices."}, {"start": 1425.6399999999999, "end": 1427.8, "text": " Okay, let me exit this."}, {"start": 1427.8, "end": 1431.26, "text": " And let me just kind of move this a bit to the left."}, {"start": 1431.26, "end": 1432.82, "text": " And we are out."}, {"start": 1432.82, "end": 1434.1599999999999, "text": " Now let's see the parallel attention."}, {"start": 1434.1599999999999, "end": 1435.76, "text": " This is gonna be the interesting part."}, {"start": 1435.76, "end": 1437.9599999999998, "text": " So let's enter the parallel attention."}, {"start": 1437.9599999999998, "end": 1439.28, "text": " I'm gonna hit F5."}, {"start": 1440.1799999999998, "end": 1442.04, "text": " Let's go to the part that we care about."}, {"start": 1442.04, "end": 1444.28, "text": " And that's the part where we split."}, {"start": 1444.28, "end": 1446.46, "text": " So here we can see,"}, {"start": 1448.3999999999999, "end": 1451.26, "text": " we can step over all of these."}, {"start": 1451.26, "end": 1453.12, "text": " The magic starts happening right here."}, {"start": 1453.12, "end": 1456.7, "text": " So again, model parallel world size is gonna be equal to,"}, {"start": 1456.7, "end": 1457.8799999999999, "text": " so it's gonna be two."}, {"start": 1457.8799999999999, "end": 1460.36, "text": " And now we divide the projection size by two,"}, {"start": 1461.44, "end": 1463.32, "text": " blah, blah, blah."}, {"start": 1463.32, "end": 1468.32, "text": " So just the devising number of attention heads to two by two."}, {"start": 1468.3999999999999, "end": 1470.74, "text": " And now we get to the interesting part."}, {"start": 1470.74, "end": 1475.74, "text": " So this is the column parallel linear feed forward layer."}, {"start": 1475.84, "end": 1477.04, "text": " That might be a mouthful,"}, {"start": 1477.04, "end": 1479.6, "text": " but let me remind you what that is."}, {"start": 1479.6, "end": 1482.6, "text": " So this is what we are actually building right now."}, {"start": 1482.6, "end": 1485.1200000000001, "text": " This is the self attention layer."}, {"start": 1485.1200000000001, "end": 1488.48, "text": " And what we are building right now are these,"}, {"start": 1488.48, "end": 1492.76, "text": " as you can see, column parallel layers."}, {"start": 1492.76, "end": 1495.1, "text": " So we take the weights and we split them"}, {"start": 1495.1, "end": 1497.74, "text": " in this column wise fashion."}, {"start": 1497.74, "end": 1502.74, "text": " And then we're gonna use, so we're gonna use these here."}, {"start": 1502.98, "end": 1507.7, "text": " The Q1, K1 and V1 are going to be used here,"}, {"start": 1507.7, "end": 1509.36, "text": " as you can see."}, {"start": 1509.36, "end": 1513.72, "text": " And then the Q2, K2 and V2 are gonna be used"}, {"start": 1513.72, "end": 1515.48, "text": " on the different device."}, {"start": 1515.48, "end": 1516.92, "text": " And that's what we are currently building."}, {"start": 1516.92, "end": 1520.1200000000001, "text": " That's how this code, that's what this code is implementing."}, {"start": 1520.1200000000001, "end": 1524.04, "text": " So let me enter the actual column parallel linear."}, {"start": 1524.04, "end": 1527.16, "text": " And you can see again, the weight matrix is being split"}, {"start": 1527.16, "end": 1532.16, "text": " in this column parallelism fashion."}, {"start": 1532.3200000000002, "end": 1535.5600000000002, "text": " So let me hit five, let me enter here."}, {"start": 1535.5600000000002, "end": 1536.78, "text": " Let's see where that's gonna happen."}, {"start": 1536.78, "end": 1537.76, "text": " So here is the world size."}, {"start": 1537.76, "end": 1539.88, "text": " Again, this is gonna be equal to two."}, {"start": 1539.88, "end": 1543.1200000000001, "text": " And then the output size is gonna be divided by two."}, {"start": 1544.48, "end": 1546.72, "text": " Okay, so let's continue."}, {"start": 1546.72, "end": 1550.5600000000002, "text": " And here is the actual, so we are now allocating the weights."}, {"start": 1552.0800000000002, "end": 1554.88, "text": " This is going to, so this would,"}, {"start": 1554.88, "end": 1557.8000000000002, "text": " if we had the topology I'm using as a running example,"}, {"start": 1557.8000000000002, "end": 1560.96, "text": " this would be, as you can see here, divided by two."}, {"start": 1560.96, "end": 1563.3200000000002, "text": " So that means we're literally splitting it column wise."}, {"start": 1563.3200000000002, "end": 1565.24, "text": " That's the whole point, okay?"}, {"start": 1565.24, "end": 1567.16, "text": " And then once we have those weights,"}, {"start": 1567.16, "end": 1570.4, "text": " we just initialize the weights and there's some bias."}, {"start": 1570.4, "end": 1571.5400000000002, "text": " Let's keep that part."}, {"start": 1571.5400000000002, "end": 1573.98, "text": " It's not that interesting."}, {"start": 1573.98, "end": 1577.0, "text": " So let's exit here, again, entering this thing."}, {"start": 1577.0, "end": 1580.0600000000002, "text": " Okay, so that's the first part."}, {"start": 1580.0600000000002, "end": 1581.4, "text": " Let's see what else."}, {"start": 1581.4, "end": 1586.4, "text": " So now we have to create some softmax stuff"}, {"start": 1587.2800000000002, "end": 1590.8000000000002, "text": " that's going to be this part, the softmax."}, {"start": 1590.8000000000002, "end": 1593.22, "text": " The only difference is this is gonna be optimized."}, {"start": 1593.22, "end": 1597.44, "text": " So obviously there are no learnable weights in this one."}, {"start": 1597.44, "end": 1599.3200000000002, "text": " So I'm gonna just skip it."}, {"start": 1599.3200000000002, "end": 1601.92, "text": " We're not sharding it."}, {"start": 1601.92, "end": 1604.96, "text": " So now we create the dropout, blah, blah, blah."}, {"start": 1604.96, "end": 1606.68, "text": " And there is the row parallel linear."}, {"start": 1606.68, "end": 1608.2, "text": " So this is the second component."}, {"start": 1608.2, "end": 1610.92, "text": " So let me go back again to the OneNote."}, {"start": 1610.92, "end": 1612.52, "text": " And here is the second part."}, {"start": 1612.52, "end": 1614.2, "text": " So here you can see the other part"}, {"start": 1614.2, "end": 1615.8000000000002, "text": " of the self-attention layer."}, {"start": 1615.8000000000002, "end": 1618.3200000000002, "text": " We have this linear layer and we split it"}, {"start": 1618.3200000000002, "end": 1620.6000000000001, "text": " row wise this time, B1 and B2."}, {"start": 1620.6000000000001, "end": 1622.3200000000002, "text": " So you can see here."}, {"start": 1622.3200000000002, "end": 1624.16, "text": " So let me go back to the code"}, {"start": 1624.16, "end": 1626.24, "text": " and let me enter the row parallel linear."}, {"start": 1626.24, "end": 1629.1200000000001, "text": " So if I enter here, you can see again the diagram"}, {"start": 1629.1200000000001, "end": 1631.16, "text": " and let me put a break point here."}, {"start": 1631.16, "end": 1632.68, "text": " Let me hit that five."}, {"start": 1632.68, "end": 1634.1200000000001, "text": " So here we are."}, {"start": 1634.1200000000001, "end": 1636.22, "text": " Again, world size, it's gonna be two."}, {"start": 1637.3600000000001, "end": 1639.44, "text": " And let me continue."}, {"start": 1639.44, "end": 1641.96, "text": " And again, this time you can see"}, {"start": 1641.96, "end": 1644.8, "text": " that the input size per partition"}, {"start": 1644.8, "end": 1647.04, "text": " is the one that's being divided."}, {"start": 1647.04, "end": 1651.68, "text": " So here the input will be divided, not the output size."}, {"start": 1651.68, "end": 1655.1200000000001, "text": " And because of that, we are doing row wise splitting, right?"}, {"start": 1655.1200000000001, "end": 1655.96, "text": " That's it."}, {"start": 1655.96, "end": 1657.3200000000002, "text": " That's the difference."}, {"start": 1657.3200000000002, "end": 1659.76, "text": " Again, I'm gonna skip everything else here."}, {"start": 1659.76, "end": 1661.96, "text": " Let me kind of step out of this function,"}, {"start": 1661.96, "end": 1665.3600000000001, "text": " shift F11 does that on Windows keyboard."}, {"start": 1667.16, "end": 1668.0, "text": " And here we are."}, {"start": 1668.0, "end": 1671.68, "text": " So we have the row parallel linear instantiated."}, {"start": 1671.68, "end": 1672.64, "text": " Let's continue."}, {"start": 1672.64, "end": 1674.08, "text": " We don't have row three embeddings."}, {"start": 1674.08, "end": 1676.4, "text": " So we can skip this part as well."}, {"start": 1676.4, "end": 1677.4, "text": " And here we are."}, {"start": 1677.4, "end": 1679.2, "text": " So that was the parallel attention."}, {"start": 1679.2, "end": 1680.32, "text": " Okay."}, {"start": 1680.32, "end": 1682.44, "text": " Again, let me just update the drawing."}, {"start": 1682.44, "end": 1686.52, "text": " We have the parallel attention being sharded"}, {"start": 1686.52, "end": 1689.68, "text": " across two devices in this running example."}, {"start": 1689.68, "end": 1693.36, "text": " So that means we'll have half of the weights here"}, {"start": 1694.64, "end": 1696.9, "text": " and half of the weights here."}, {"start": 1696.9, "end": 1699.88, "text": " Obviously we'll need 12 of these."}, {"start": 1699.88, "end": 1702.92, "text": " So I kind of didn't do it justice here,"}, {"start": 1702.92, "end": 1705.52, "text": " but yeah, you can imagine we'll have 12 of these"}, {"start": 1705.52, "end": 1708.0, "text": " transformer self-attention layers."}, {"start": 1708.88, "end": 1711.4, "text": " So let's go back to the code."}, {"start": 1711.4, "end": 1713.02, "text": " Let me exit here."}, {"start": 1713.02, "end": 1714.96, "text": " So we have layer norm again."}, {"start": 1714.96, "end": 1716.0400000000002, "text": " I won't bother."}, {"start": 1716.0400000000002, "end": 1718.2, "text": " I'm just gonna exit this layer norm."}, {"start": 1719.8400000000001, "end": 1722.0400000000002, "text": " Let's see what else we skip the decoder part."}, {"start": 1722.0400000000002, "end": 1723.8400000000001, "text": " So now there is the parallel MLP."}, {"start": 1723.8400000000001, "end": 1726.24, "text": " So again, how the transformer block looks like,"}, {"start": 1726.24, "end": 1730.76, "text": " we have the self-attention module"}, {"start": 1730.76, "end": 1732.2, "text": " and then we have the MLP part."}, {"start": 1732.2, "end": 1735.54, "text": " So now here we'll see first the column-wise splitting"}, {"start": 1735.54, "end": 1737.2, "text": " and then the row-wise splitting."}, {"start": 1737.2, "end": 1739.28, "text": " So this is the diagram we'll implement."}, {"start": 1739.28, "end": 1741.8, "text": " Let's go back to the code."}, {"start": 1741.8, "end": 1743.04, "text": " Let me show you this quickly."}, {"start": 1743.04, "end": 1744.6, "text": " And then this is the last thing I'll show you."}, {"start": 1744.6, "end": 1746.08, "text": " And then we are pretty much done"}, {"start": 1746.08, "end": 1747.84, "text": " with this model construction."}, {"start": 1747.84, "end": 1749.52, "text": " So let me hit F5."}, {"start": 1749.52, "end": 1750.44, "text": " Let me enter here."}, {"start": 1750.44, "end": 1752.36, "text": " So again, column parallel."}, {"start": 1753.36, "end": 1754.6, "text": " And I'm just gonna skip."}, {"start": 1754.6, "end": 1758.1599999999999, "text": " I'm just going to disable the breakpoints"}, {"start": 1758.1599999999999, "end": 1760.36, "text": " because it's the same logic as before, right?"}, {"start": 1760.36, "end": 1761.48, "text": " We've seen this."}, {"start": 1761.48, "end": 1764.06, "text": " This is the same logic as before."}, {"start": 1764.06, "end": 1766.4599999999998, "text": " So I'm just gonna skip over all of that."}, {"start": 1766.4599999999998, "end": 1767.9599999999998, "text": " Let's get here."}, {"start": 1769.6799999999998, "end": 1771.48, "text": " Now some activation stuff."}, {"start": 1771.48, "end": 1774.6999999999998, "text": " Here is the row parallel linear and that's it."}, {"start": 1775.6, "end": 1777.56, "text": " And that's the one that has the 4H"}, {"start": 1777.56, "end": 1780.6399999999999, "text": " usually is the common convention."}, {"start": 1780.64, "end": 1785.64, "text": " So the hidden layer will have 4X more weights"}, {"start": 1785.8400000000001, "end": 1789.74, "text": " than the input and output, dimension, sorry."}, {"start": 1789.74, "end": 1790.5800000000002, "text": " Okay."}, {"start": 1791.64, "end": 1793.4, "text": " Again, let's exit here."}, {"start": 1793.4, "end": 1796.5200000000002, "text": " And that's our basically transformer layer."}, {"start": 1796.5200000000002, "end": 1798.68, "text": " Again, you can always follow here where we are."}, {"start": 1798.68, "end": 1801.0400000000002, "text": " We're currently implementing the transformer layer."}, {"start": 1801.0400000000002, "end": 1804.72, "text": " So I'm gonna exit here and we do this 12 times."}, {"start": 1804.72, "end": 1808.5800000000002, "text": " So I'm gonna, again, let me just make sure it's disabled."}, {"start": 1808.58, "end": 1812.1999999999998, "text": " I'm gonna enable the breakpoint here, hit F5"}, {"start": 1812.1999999999998, "end": 1813.76, "text": " and we have 12 layers."}, {"start": 1813.76, "end": 1815.6599999999999, "text": " So again, let me update the drawing."}, {"start": 1817.1999999999998, "end": 1818.3999999999999, "text": " Here is the drawing."}, {"start": 1818.3999999999999, "end": 1820.96, "text": " We now have 12 layers here and here."}, {"start": 1822.96, "end": 1825.62, "text": " So we have 12 layers here and here."}, {"start": 1825.62, "end": 1829.48, "text": " And obviously, because this is running in parallel,"}, {"start": 1829.48, "end": 1834.48, "text": " these two devices would also have 12 layers here."}, {"start": 1834.84, "end": 1838.08, "text": " So 12 layers and 12 layers."}, {"start": 1838.08, "end": 1839.6399999999999, "text": " The difference is these two,"}, {"start": 1839.6399999999999, "end": 1842.98, "text": " because they are in the rank zero of the pipeline parallelism,"}, {"start": 1842.98, "end": 1845.8, "text": " they additionally have the embedding table sharded."}, {"start": 1845.8, "end": 1846.6399999999999, "text": " Cool."}, {"start": 1846.6399999999999, "end": 1848.52, "text": " Guys, let me know how you like the video so far."}, {"start": 1848.52, "end": 1850.24, "text": " Leave any feedback if you have any."}, {"start": 1850.24, "end": 1851.1599999999999, "text": " That would be super cool."}, {"start": 1851.1599999999999, "end": 1852.6399999999999, "text": " So again, post-process."}, {"start": 1852.6399999999999, "end": 1855.0, "text": " That's going to be active only"}, {"start": 1855.0, "end": 1857.0, "text": " on the last stage of the pipeline."}, {"start": 1858.36, "end": 1861.28, "text": " So that means only those devices will have the layer norm,"}, {"start": 1861.28, "end": 1862.4399999999998, "text": " but I'm gonna skip the,"}, {"start": 1862.4399999999998, "end": 1865.1599999999999, "text": " I won't enter the layer norm because we already seen it."}, {"start": 1865.1599999999999, "end": 1867.32, "text": " That's the optimized layer norm."}, {"start": 1867.32, "end": 1868.84, "text": " And that's it."}, {"start": 1868.84, "end": 1871.6799999999998, "text": " Again, that was the parallel transformer."}, {"start": 1873.04, "end": 1873.8799999999999, "text": " Let's continue."}, {"start": 1873.8799999999999, "end": 1875.56, "text": " We don't have the decoder part."}, {"start": 1875.56, "end": 1878.1799999999998, "text": " We don't have the pooler, so we can exit here."}, {"start": 1878.1799999999998, "end": 1879.24, "text": " We can exit here."}, {"start": 1879.24, "end": 1881.6, "text": " And we have our language model being constructed."}, {"start": 1881.6, "end": 1884.4399999999998, "text": " Again, I'm not sure why I'm entering this thing."}, {"start": 1884.4399999999998, "end": 1887.6799999999998, "text": " So I just, once I enter it, I just do step out"}, {"start": 1887.6799999999998, "end": 1889.96, "text": " and that's how I exit in case you're wondering."}, {"start": 1889.96, "end": 1891.8799999999999, "text": " Okay, guys, here we are, GPT model."}, {"start": 1891.8799999999999, "end": 1894.04, "text": " We just constructed it."}, {"start": 1894.04, "end": 1896.9199999999998, "text": " Now there is this initialized word embedding."}, {"start": 1896.92, "end": 1901.28, "text": " So because I have, in my actual setup,"}, {"start": 1901.28, "end": 1904.24, "text": " I have a single GPU, so I'll just hit return."}, {"start": 1904.24, "end": 1907.04, "text": " But in general, if we had the actual topology I'm using"}, {"start": 1907.04, "end": 1910.64, "text": " as the running example, we'd be hitting this code here."}, {"start": 1910.64, "end": 1912.72, "text": " And let me show you why this is interesting."}, {"start": 1912.72, "end": 1913.88, "text": " So they mentioned it here."}, {"start": 1913.88, "end": 1918.88, "text": " So parameters are shared between the word embeddings layer."}, {"start": 1919.3200000000002, "end": 1921.0800000000002, "text": " So that's the common thing you always do."}, {"start": 1921.0800000000002, "end": 1925.04, "text": " So you have the embedding table being shared between,"}, {"start": 1925.04, "end": 1927.72, "text": " so mapping the tokens into the embedding vectors."}, {"start": 1927.72, "end": 1931.1599999999999, "text": " And then at the output, you also use that embedding table"}, {"start": 1931.1599999999999, "end": 1934.6, "text": " to map the embedding vectors into the vocab."}, {"start": 1934.6, "end": 1938.1599999999999, "text": " So basically distribution across the various words,"}, {"start": 1938.1599999999999, "end": 1941.54, "text": " subtokens from your, or tokens from your vocabulary."}, {"start": 1942.62, "end": 1943.72, "text": " So those are tied."}, {"start": 1943.72, "end": 1947.12, "text": " And you can see here, they have a nice trick that they do."}, {"start": 1947.12, "end": 1952.12, "text": " So if this was running on the last stage of the pipeline,"}, {"start": 1952.12, "end": 1954.9599999999998, "text": " then they would create embedding table."}, {"start": 1954.9599999999998, "end": 1957.1999999999998, "text": " Again, it's gonna be a parallel embedding table."}, {"start": 1957.1999999999998, "end": 1959.32, "text": " So that means it's gonna be shorted across two devices"}, {"start": 1959.32, "end": 1961.6399999999999, "text": " of the last stage of the pipeline."}, {"start": 1961.6399999999999, "end": 1963.7199999999998, "text": " And then this is the nice trick."}, {"start": 1963.7199999999998, "end": 1967.2399999999998, "text": " So basically what they do is they do all reduce."}, {"start": 1967.2399999999998, "end": 1968.84, "text": " So ensure that first and last stages"}, {"start": 1968.84, "end": 1970.6799999999998, "text": " have the same initial parameter values."}, {"start": 1970.6799999999998, "end": 1973.52, "text": " So because this was not independently allocated"}, {"start": 1973.52, "end": 1977.32, "text": " and because they are lying on different devices,"}, {"start": 1977.32, "end": 1980.52, "text": " because of that, they have to do all reduce."}, {"start": 1980.52, "end": 1982.56, "text": " So you can see here, if we are hitting"}, {"start": 1982.56, "end": 1986.16, "text": " either the first stage of the pipeline or the last stage,"}, {"start": 1986.16, "end": 1990.16, "text": " apply the all reduce method on this word embedding,"}, {"start": 1990.16, "end": 1992.32, "text": " or word embedding's weight."}, {"start": 1992.32, "end": 1994.92, "text": " If you're not familiar with these collective operations,"}, {"start": 1994.92, "end": 1998.56, "text": " I did not explain them so far, just check out the Viki page."}, {"start": 1998.56, "end": 2000.12, "text": " They have a nice visual explanation."}, {"start": 2000.12, "end": 2000.96, "text": " It's fairly simple."}, {"start": 2000.96, "end": 2003.8799999999999, "text": " Roughly what's gonna happen is assume that the first stage"}, {"start": 2003.8799999999999, "end": 2007.2, "text": " will literally receive the weights from the last stage,"}, {"start": 2007.2, "end": 2009.4, "text": " and then they're gonna divide it by two."}, {"start": 2009.4, "end": 2012.18, "text": " And if the last stage does the same thing,"}, {"start": 2012.18, "end": 2013.72, "text": " you'll end up with the same result."}, {"start": 2013.72, "end": 2015.5600000000002, "text": " And that means you'll have the same weights"}, {"start": 2015.5600000000002, "end": 2018.3600000000001, "text": " both in the first stage and in the last stage,"}, {"start": 2018.3600000000001, "end": 2020.6000000000001, "text": " and thus you've synchronized the weights."}, {"start": 2020.6000000000001, "end": 2021.64, "text": " As simple as that."}, {"start": 2021.64, "end": 2025.24, "text": " Okay, guys, hopefully that was interesting."}, {"start": 2025.24, "end": 2026.5600000000002, "text": " And hopefully that makes sense"}, {"start": 2026.5600000000002, "end": 2028.8000000000002, "text": " with the drawing I'm using in OneNote."}, {"start": 2028.8000000000002, "end": 2029.96, "text": " So let me exit here."}, {"start": 2029.96, "end": 2034.96, "text": " So here is the, now some stuff that's not that vital."}, {"start": 2035.88, "end": 2039.24, "text": " I'm gonna skip all of this, put the break point here,"}, {"start": 2039.24, "end": 2040.42, "text": " hit F5."}, {"start": 2041.32, "end": 2042.52, "text": " Let's see what else."}, {"start": 2042.52, "end": 2047.52, "text": " So we push all of these models onto the GPU."}, {"start": 2049.92, "end": 2054.24, "text": " We use this float, we convert them to float 16 weights."}, {"start": 2054.24, "end": 2056.92, "text": " So I'm going to just, basically let me enable"}, {"start": 2056.92, "end": 2059.24, "text": " the break points now."}, {"start": 2059.24, "end": 2060.84, "text": " So I'm gonna hit this one a bit later,"}, {"start": 2060.84, "end": 2063.62, "text": " so we will see how that is going to be used."}, {"start": 2063.62, "end": 2068.62, "text": " Okay, local DDP data distributed parallelism."}, {"start": 2069.94, "end": 2072.5, "text": " We're gonna ignore that for the time being."}, {"start": 2072.5, "end": 2074.2999999999997, "text": " And let's just return the model."}, {"start": 2075.22, "end": 2078.7799999999997, "text": " Okay, so again, because I have load set to none,"}, {"start": 2078.7799999999997, "end": 2080.2999999999997, "text": " I will not be loading the checkpoints,"}, {"start": 2080.2999999999997, "end": 2083.46, "text": " which means I'll basically be using random weights."}, {"start": 2083.46, "end": 2085.38, "text": " Okay, now we build the data set."}, {"start": 2085.38, "end": 2087.74, "text": " Let me quickly walk you through this one."}, {"start": 2087.74, "end": 2091.02, "text": " Although that's not that relevant, so tokenizer."}, {"start": 2091.02, "end": 2096.02, "text": " So again, I just load this random dummy file I created."}, {"start": 2097.54, "end": 2099.66, "text": " And you can see what it has."}, {"start": 2099.66, "end": 2100.82, "text": " Literally I copy pasted."}, {"start": 2100.82, "end": 2105.82, "text": " I actually took a single paragraph from WikiText data set"}, {"start": 2106.22, "end": 2107.86, "text": " and just copy pasted a couple of times"}, {"start": 2107.86, "end": 2109.32, "text": " just to have something there."}, {"start": 2110.34, "end": 2112.3, "text": " So now we just do some processing,"}, {"start": 2112.3, "end": 2114.3, "text": " the tokenization, blah, blah, blah."}, {"start": 2114.3, "end": 2115.56, "text": " Okay, let me just exit here."}, {"start": 2115.56, "end": 2117.54, "text": " This is not interesting."}, {"start": 2117.54, "end": 2119.86, "text": " So we then tokenize the text."}, {"start": 2119.86, "end": 2121.9, "text": " Okay, I'm just hitting the breakpoints."}, {"start": 2121.9, "end": 2123.42, "text": " So this is some processing that they do."}, {"start": 2123.42, "end": 2125.26, "text": " As you can see, a lot of details"}, {"start": 2125.26, "end": 2128.78, "text": " that I'm trying to unsuccessfully ignore."}, {"start": 2128.78, "end": 2132.1800000000003, "text": " So we now do the tokenization and you can see here,"}, {"start": 2132.1800000000003, "end": 2133.94, "text": " now we have the, let me show you"}, {"start": 2133.94, "end": 2136.02, "text": " how that's gonna look like."}, {"start": 2136.02, "end": 2140.2400000000002, "text": " Tokenized data is just a list of tokens of our text."}, {"start": 2140.2400000000002, "end": 2144.1, "text": " So the text was tokenized and we have 448 tokens,"}, {"start": 2144.1, "end": 2145.6800000000003, "text": " the smallest data set ever."}, {"start": 2145.68, "end": 2150.0, "text": " Okay, now we use that to construct this LM data set."}, {"start": 2151.1, "end": 2154.7, "text": " Again, I'm not going to dig into a lot of details"}, {"start": 2154.7, "end": 2157.7999999999997, "text": " because that's just the language modeling detail"}, {"start": 2157.7999999999997, "end": 2160.3999999999996, "text": " and we are trying to focus on the 3D parallelism"}, {"start": 2160.3999999999996, "end": 2161.3999999999996, "text": " in this video."}, {"start": 2161.3999999999996, "end": 2163.72, "text": " Okay, let's exit here."}, {"start": 2163.72, "end": 2164.8399999999997, "text": " We have the data set."}, {"start": 2164.8399999999997, "end": 2166.7999999999997, "text": " Now we build the data loader."}, {"start": 2166.7999999999997, "end": 2168.04, "text": " The reason I'm gonna show you this one"}, {"start": 2168.04, "end": 2171.14, "text": " is because it involves the data parallelism approach."}, {"start": 2171.14, "end": 2172.24, "text": " So here it is."}, {"start": 2172.24, "end": 2174.2, "text": " In PyTorch, it's fairly trivial."}, {"start": 2174.2, "end": 2177.7599999999998, "text": " So here you can see get data parallel world size."}, {"start": 2177.7599999999998, "end": 2180.8999999999996, "text": " So now additionally, I did not show you this so far,"}, {"start": 2180.8999999999996, "end": 2185.8999999999996, "text": " but imagine having this same group of devices,"}, {"start": 2186.4399999999996, "end": 2188.08, "text": " like four devices like this."}, {"start": 2188.08, "end": 2189.8999999999996, "text": " So I'm just gonna change the color."}, {"start": 2189.8999999999996, "end": 2194.12, "text": " So imagine having this same group duplicated maybe twice"}, {"start": 2194.12, "end": 2197.12, "text": " or even four times, okay?"}, {"start": 2197.12, "end": 2199.2, "text": " And we're doing the same thing"}, {"start": 2199.2, "end": 2201.56, "text": " in each of these groups until now."}, {"start": 2201.56, "end": 2203.7599999999998, "text": " So now it's gonna be a bit different"}, {"start": 2203.76, "end": 2205.1800000000003, "text": " for each of these groups."}, {"start": 2205.1800000000003, "end": 2208.6400000000003, "text": " Okay, so let's say this is group zero and this is group one."}, {"start": 2208.6400000000003, "end": 2211.7200000000003, "text": " I'm gonna ignore, we're gonna have two groups, okay?"}, {"start": 2212.6000000000004, "end": 2214.1600000000003, "text": " So let's go back to the code."}, {"start": 2215.44, "end": 2217.88, "text": " We would get world size would be two."}, {"start": 2217.88, "end": 2219.6400000000003, "text": " And let me show you how this is implemented"}, {"start": 2219.6400000000003, "end": 2220.5200000000004, "text": " in the background."}, {"start": 2220.5200000000004, "end": 2223.84, "text": " It literally calls this distributed get world size"}, {"start": 2223.84, "end": 2226.38, "text": " and then you specify the particular group."}, {"start": 2226.38, "end": 2229.1600000000003, "text": " So here they specify the get data parallel group"}, {"start": 2229.16, "end": 2234.16, "text": " and that's why basically they get number two."}, {"start": 2234.16, "end": 2235.52, "text": " Okay, so some constants there,"}, {"start": 2235.52, "end": 2237.8799999999997, "text": " I will not dig deeper than that."}, {"start": 2237.8799999999997, "end": 2239.8799999999997, "text": " Let's continue, let's get the rank."}, {"start": 2239.8799999999997, "end": 2242.52, "text": " So now where the rank is zero for the left group,"}, {"start": 2242.52, "end": 2245.7599999999998, "text": " so rank would be zero again for this group here,"}, {"start": 2245.7599999999998, "end": 2248.24, "text": " it will be one for these devices here."}, {"start": 2248.24, "end": 2250.3599999999997, "text": " So that's how we'd have a differentiation."}, {"start": 2251.66, "end": 2252.6, "text": " Let's continue."}, {"start": 2252.6, "end": 2255.58, "text": " We formed the distributed sampler, okay?"}, {"start": 2255.58, "end": 2258.18, "text": " And that one is going to be shifting"}, {"start": 2258.18, "end": 2261.3599999999997, "text": " a part of our batch to one group"}, {"start": 2261.3599999999997, "end": 2263.44, "text": " and the second part to the other group."}, {"start": 2263.44, "end": 2267.72, "text": " So that's gonna handle in the background by the PyTorch."}, {"start": 2267.72, "end": 2270.52, "text": " Okay, so now the data loader, blah, blah, blah,"}, {"start": 2270.52, "end": 2272.44, "text": " pinned memory, again, a lot of details,"}, {"start": 2272.44, "end": 2274.2799999999997, "text": " I'll have to skip those."}, {"start": 2274.2799999999997, "end": 2277.52, "text": " And now we enter the evaluate and print results."}, {"start": 2277.52, "end": 2280.08, "text": " So here is the actual function"}, {"start": 2280.08, "end": 2283.3399999999997, "text": " where we'll be doing a feed forward pass through the model."}, {"start": 2283.3399999999997, "end": 2286.68, "text": " So we just set it to eval mode."}, {"start": 2286.68, "end": 2288.96, "text": " We entered the TorchNodeGrad context"}, {"start": 2288.96, "end": 2290.12, "text": " because we're doing inference,"}, {"start": 2290.12, "end": 2292.3799999999997, "text": " we're not trying to train anything here."}, {"start": 2292.3799999999997, "end": 2296.9199999999996, "text": " And now we're going to fetch our first batch."}, {"start": 2296.9199999999996, "end": 2300.24, "text": " So I'm going to ignore how that works."}, {"start": 2300.24, "end": 2302.0, "text": " So I'm gonna disable the breakpoints,"}, {"start": 2302.0, "end": 2304.72, "text": " otherwise I would hit some breakpoints I left there."}, {"start": 2304.72, "end": 2306.56, "text": " So let me hit F10 and we have a batch."}, {"start": 2306.56, "end": 2310.48, "text": " So a batch is simply going to be,"}, {"start": 2310.48, "end": 2311.9199999999996, "text": " as you can see here, we have text"}, {"start": 2311.9199999999996, "end": 2314.6, "text": " and we have pet mask keys."}, {"start": 2314.6, "end": 2318.2, "text": " So text is going to be, I guess, eight times 1,024"}, {"start": 2318.2, "end": 2319.04, "text": " or something."}, {"start": 2319.7999999999997, "end": 2323.2, "text": " Yeah, so eight because we have a batch of size eight"}, {"start": 2323.2, "end": 2328.2, "text": " and 1,025 because our sequence that we're gonna pass"}, {"start": 2328.92, "end": 2331.02, "text": " to the transformer is gonna be 1,024,"}, {"start": 2331.02, "end": 2335.0, "text": " but it's 1,025 because we'll be using one as the target."}, {"start": 2335.0, "end": 2337.0, "text": " And so you're gonna see that in a second."}, {"start": 2338.0, "end": 2339.4, "text": " So yeah, hopefully that makes sense."}, {"start": 2339.4, "end": 2341.7599999999998, "text": " Like you literally take the, it's 1,025"}, {"start": 2341.76, "end": 2345.0, "text": " because you take the first 1,024 as the input"}, {"start": 2345.0, "end": 2347.6400000000003, "text": " and then you're trying to predict shifted by one"}, {"start": 2347.6400000000003, "end": 2349.96, "text": " to the right, you try to predict that sequence."}, {"start": 2349.96, "end": 2353.48, "text": " And that's why you have 1,025 here."}, {"start": 2353.48, "end": 2354.32, "text": " Okay, let's continue."}, {"start": 2354.32, "end": 2356.1200000000003, "text": " Let's enter the forward pass."}, {"start": 2356.1200000000003, "end": 2357.6000000000004, "text": " Let's see what's going on here."}, {"start": 2357.6000000000004, "end": 2360.5600000000004, "text": " So some batch processing, again,"}, {"start": 2360.5600000000004, "end": 2362.1600000000003, "text": " I'm gonna ignore all of that"}, {"start": 2362.1600000000003, "end": 2366.4, "text": " because that's just a language modeling standard stuff."}, {"start": 2366.4, "end": 2367.6000000000004, "text": " This is the interesting part."}, {"start": 2367.6000000000004, "end": 2369.6400000000003, "text": " So this is the receive forward."}, {"start": 2369.6400000000003, "end": 2370.92, "text": " So what's going on here?"}, {"start": 2370.92, "end": 2372.4, "text": " So let's enter there."}, {"start": 2373.7200000000003, "end": 2376.76, "text": " Let me now enable the breakpoints."}, {"start": 2376.76, "end": 2379.4, "text": " So here what happens is the following."}, {"start": 2379.4, "end": 2382.62, "text": " So if we are in the first stage,"}, {"start": 2382.62, "end": 2385.84, "text": " then the input tensor is none, okay?"}, {"start": 2385.84, "end": 2388.38, "text": " But if we were in some other stage,"}, {"start": 2388.38, "end": 2390.64, "text": " then we call this communicate method."}, {"start": 2390.64, "end": 2393.46, "text": " And that means we're waiting for the output"}, {"start": 2393.46, "end": 2396.76, "text": " of the previous stage to pass that output"}, {"start": 2396.76, "end": 2399.08, "text": " as the input of this stage."}, {"start": 2399.08, "end": 2401.7999999999997, "text": " And if that happens, we're going to override"}, {"start": 2401.7999999999997, "end": 2404.12, "text": " the batch that we just loaded from the data loader."}, {"start": 2404.12, "end": 2405.92, "text": " So that might be a mouthful."}, {"start": 2405.92, "end": 2408.2799999999997, "text": " It's gonna hopefully make sense in a couple of minutes."}, {"start": 2408.2799999999997, "end": 2409.96, "text": " So here is the input tensor."}, {"start": 2409.96, "end": 2411.04, "text": " Let's continue here."}, {"start": 2411.04, "end": 2414.24, "text": " Okay, some unwrapping of the model, not that interesting."}, {"start": 2414.24, "end": 2416.88, "text": " So here's the part where we set the input tensor."}, {"start": 2416.88, "end": 2419.18, "text": " And let me kind of dig into this."}, {"start": 2419.18, "end": 2421.68, "text": " So blah, blah, blah, we call the GPT model"}, {"start": 2421.68, "end": 2423.08, "text": " set input tensor function."}, {"start": 2423.08, "end": 2426.64, "text": " And then, okay, let me exit this one."}, {"start": 2426.64, "end": 2430.3599999999997, "text": " So then we enter the transformer language model"}, {"start": 2430.3599999999997, "end": 2431.6, "text": " set input tensor."}, {"start": 2431.6, "end": 2433.62, "text": " Again, a lot of levels of indirection,"}, {"start": 2433.62, "end": 2436.2799999999997, "text": " and I keep hitting this random function."}, {"start": 2436.2799999999997, "end": 2437.52, "text": " Okay, so here we are."}, {"start": 2437.52, "end": 2439.4, "text": " So finally we get to parallel transformer,"}, {"start": 2439.4, "end": 2442.7799999999997, "text": " and we set the input tensor to none,"}, {"start": 2442.7799999999997, "end": 2444.64, "text": " because we are now in the first stage."}, {"start": 2444.64, "end": 2447.2, "text": " I mean, at least in the actual setup I have."}, {"start": 2447.2, "end": 2449.56, "text": " But let me show you where this one is used."}, {"start": 2449.56, "end": 2453.42, "text": " So it's going to be used here a bit later in the forward pass."}, {"start": 2453.42, "end": 2456.24, "text": " And we'll see that if preprocess is true,"}, {"start": 2456.24, "end": 2458.0, "text": " that means we're in the first stage."}, {"start": 2458.0, "end": 2459.3599999999997, "text": " And in that case,"}, {"start": 2459.3599999999997, "end": 2461.8799999999997, "text": " we'll actually be taking the data from the batch."}, {"start": 2461.8799999999997, "end": 2464.0, "text": " Otherwise, if we are not in the first stage,"}, {"start": 2464.0, "end": 2466.7999999999997, "text": " we'll be passing the actual input tensor"}, {"start": 2467.72, "end": 2469.3599999999997, "text": " that's passed from the previous stage."}, {"start": 2469.3599999999997, "end": 2472.1, "text": " So again, we'll see that a bit later."}, {"start": 2472.1, "end": 2474.4799999999996, "text": " So let me now exit all of this."}, {"start": 2474.4799999999996, "end": 2477.12, "text": " Let me exit this, let me exit this."}, {"start": 2477.12, "end": 2478.8799999999997, "text": " And let's continue here."}, {"start": 2478.8799999999997, "end": 2480.72, "text": " So here is the actual forward pass."}, {"start": 2480.72, "end": 2482.7999999999997, "text": " Again, just so that we are on the same page,"}, {"start": 2482.7999999999997, "end": 2485.6, "text": " let me show you what's going to happen here."}, {"start": 2485.6, "end": 2489.36, "text": " So what's happening is we are having"}, {"start": 2489.36, "end": 2493.92, "text": " our distributed data loader fetching a batch."}, {"start": 2493.92, "end": 2498.92, "text": " And I'm just going to simply draw it as a simple 2D square."}, {"start": 2499.7999999999997, "end": 2504.44, "text": " And we're going to basically divide it into two pieces."}, {"start": 2504.44, "end": 2507.72, "text": " This one is going to be passed to this group."}, {"start": 2507.72, "end": 2511.2, "text": " And the first chunk is gonna be passed"}, {"start": 2511.2, "end": 2513.52, "text": " into the zeroth group here, okay?"}, {"start": 2513.52, "end": 2516.4, "text": " And now what happens is this is going to be fed here."}, {"start": 2517.84, "end": 2521.14, "text": " And while that's happening, these two guys are blocked."}, {"start": 2521.14, "end": 2523.88, "text": " So the thing is that communicate method I showed you"}, {"start": 2523.88, "end": 2525.18, "text": " is a blocking method."}, {"start": 2525.18, "end": 2529.04, "text": " So that means until this has done its feed,"}, {"start": 2529.04, "end": 2532.44, "text": " so until this data has passed through the first stage,"}, {"start": 2532.44, "end": 2533.6, "text": " these guys are waiting."}, {"start": 2533.6, "end": 2535.7599999999998, "text": " And then, only then do they start"}, {"start": 2535.7599999999998, "end": 2537.28, "text": " doing the actual computation."}, {"start": 2537.28, "end": 2538.6, "text": " That's how this is gonna look like."}, {"start": 2538.6, "end": 2542.44, "text": " And obviously the same happens for the rank one,"}, {"start": 2542.44, "end": 2545.96, "text": " for the second group, second data parallel group."}, {"start": 2545.96, "end": 2548.44, "text": " So I will not be drawing it there, obviously."}, {"start": 2548.44, "end": 2551.68, "text": " Okay, let's enter the feed forward pass."}, {"start": 2553.0, "end": 2558.0, "text": " Okay, we have to enter this forward call."}, {"start": 2558.44, "end": 2560.2400000000002, "text": " So distributed blah, blah, blah."}, {"start": 2560.2400000000002, "end": 2563.52, "text": " There is a lot of levels of indirection again here."}, {"start": 2563.52, "end": 2565.44, "text": " Let me just exit here."}, {"start": 2568.4, "end": 2571.26, "text": " Oh my God, oh my God, let me enter here."}, {"start": 2571.26, "end": 2574.2000000000003, "text": " Okay, so finally we are in the forward pass"}, {"start": 2574.2000000000003, "end": 2576.2000000000003, "text": " of the float 16 module."}, {"start": 2576.2000000000003, "end": 2577.88, "text": " So let's see what's happening here."}, {"start": 2577.88, "end": 2579.2000000000003, "text": " So as you can see here,"}, {"start": 2579.2000000000003, "end": 2581.46, "text": " if we're in the first stage of the pipeline,"}, {"start": 2581.46, "end": 2584.26, "text": " we convert our data to float 16,"}, {"start": 2584.26, "end": 2586.5600000000004, "text": " because we're gonna do mixed precision."}, {"start": 2586.5600000000004, "end": 2590.4, "text": " So now inputs are gonna be converted to float 16."}, {"start": 2590.4, "end": 2593.42, "text": " And now we call the actual transformer"}, {"start": 2593.42, "end": 2595.32, "text": " with the float 16 data."}, {"start": 2595.32, "end": 2596.5600000000004, "text": " Okay, let's enter here."}, {"start": 2598.76, "end": 2600.5600000000004, "text": " Let me hit F11."}, {"start": 2600.56, "end": 2603.4, "text": " I think I can just hit F5 and I'll hit the, yeah."}, {"start": 2603.4, "end": 2604.24, "text": " So here we are."}, {"start": 2604.24, "end": 2607.84, "text": " We are in the GPT model forward pass,"}, {"start": 2607.84, "end": 2610.32, "text": " and let's start going through it."}, {"start": 2610.32, "end": 2612.12, "text": " Okay, so again, I think I can hit F5"}, {"start": 2612.12, "end": 2614.2799999999997, "text": " because they have a break point there."}, {"start": 2614.2799999999997, "end": 2616.42, "text": " Yeah, so we are in the transformer language model."}, {"start": 2616.42, "end": 2620.04, "text": " Okay, so as you can see here, only the first stages,"}, {"start": 2620.04, "end": 2623.04, "text": " because pre-process is true for first stages of the pipeline,"}, {"start": 2623.04, "end": 2625.04, "text": " are going to do the embedding part."}, {"start": 2625.04, "end": 2626.7999999999997, "text": " Okay, so we do the embedding part."}, {"start": 2626.7999999999997, "end": 2629.04, "text": " Let me kinda enter this."}, {"start": 2629.04, "end": 2630.32, "text": " Whoops, I hate this."}, {"start": 2630.32, "end": 2631.28, "text": " Let me enter it."}, {"start": 2631.28, "end": 2633.36, "text": " So here we are in the embedding forward."}, {"start": 2633.36, "end": 2636.1600000000003, "text": " Let me see whether it's useful to enter this one as well."}, {"start": 2638.84, "end": 2641.44, "text": " So here we are in the vocab parallel embedding."}, {"start": 2641.44, "end": 2643.6000000000004, "text": " Let's see how the forward would look like."}, {"start": 2643.6000000000004, "end": 2647.1200000000003, "text": " So again, keep in mind that we have sharded"}, {"start": 2647.1200000000003, "end": 2648.84, "text": " that embedding table here and here."}, {"start": 2648.84, "end": 2653.06, "text": " So we only have 25,000 of our embedding vectors here,"}, {"start": 2653.06, "end": 2655.2400000000002, "text": " and the other 25,000 is here."}, {"start": 2655.2400000000002, "end": 2658.6400000000003, "text": " So they'll somehow have to communicate between each other,"}, {"start": 2658.64, "end": 2661.58, "text": " such that we can end up with the correct results."}, {"start": 2661.58, "end": 2663.7599999999998, "text": " So let's see how that's gonna work."}, {"start": 2663.7599999999998, "end": 2665.7999999999997, "text": " So let's enter here."}, {"start": 2667.0, "end": 2669.24, "text": " So we would actually be entering this branch"}, {"start": 2669.24, "end": 2672.3199999999997, "text": " if I had the setup that I'm using as running examples,"}, {"start": 2672.3199999999997, "end": 2676.3599999999997, "text": " because the model parallel size would be two."}, {"start": 2676.3599999999997, "end": 2679.48, "text": " So what would happen is the input mask,"}, {"start": 2679.48, "end": 2681.64, "text": " as you can see here, it will be true"}, {"start": 2681.64, "end": 2686.3199999999997, "text": " for all of those indices that are smaller or bigger"}, {"start": 2686.32, "end": 2690.96, "text": " than the index here, so then the boundary index here."}, {"start": 2690.96, "end": 2694.2000000000003, "text": " And these would be like, these will be zero and 25,000"}, {"start": 2694.2000000000003, "end": 2697.6000000000004, "text": " for this one, for this one here,"}, {"start": 2697.6000000000004, "end": 2700.92, "text": " and 25,000 to 50,000 for this one, okay?"}, {"start": 2700.92, "end": 2705.4, "text": " And that's why this is gonna mask those indices"}, {"start": 2705.4, "end": 2707.92, "text": " that do not belong to this device, okay?"}, {"start": 2707.92, "end": 2711.6400000000003, "text": " So there is this subtraction here that we have to do"}, {"start": 2711.6400000000003, "end": 2713.32, "text": " so that we have correct results,"}, {"start": 2713.32, "end": 2716.44, "text": " because imagine we are now on this device,"}, {"start": 2716.44, "end": 2719.0, "text": " because its start index is 25,000,"}, {"start": 2719.0, "end": 2723.32, "text": " we want to subtract 25,000 such that we end up with zero,"}, {"start": 2723.32, "end": 2725.1600000000003, "text": " because we are indexing in its table"}, {"start": 2725.1600000000003, "end": 2728.6000000000004, "text": " and the zero index is the first slot in that table"}, {"start": 2728.6000000000004, "end": 2730.56, "text": " on that device, if that makes sense."}, {"start": 2730.56, "end": 2732.98, "text": " Sorry for being a bit confusing here."}, {"start": 2734.02, "end": 2736.7400000000002, "text": " And you can see here, for all of the other indices,"}, {"start": 2736.7400000000002, "end": 2738.56, "text": " we just set them to zero."}, {"start": 2738.56, "end": 2740.6400000000003, "text": " Maybe let me try and draw this quickly."}, {"start": 2740.64, "end": 2743.44, "text": " So here is what's going on."}, {"start": 2743.44, "end": 2748.18, "text": " Okay, we have basically input that contains indices"}, {"start": 2748.18, "end": 2751.8599999999997, "text": " from the whole table, so that means it might have numbers."}, {"start": 2753.2799999999997, "end": 2755.58, "text": " Okay, let's quickly check what's the actual size."}, {"start": 2755.58, "end": 2758.68, "text": " So let me use this input as the running example."}, {"start": 2758.68, "end": 2762.3599999999997, "text": " So we have shape 8024, okay?"}, {"start": 2762.3599999999997, "end": 2764.48, "text": " So let's imagine we just have 1024,"}, {"start": 2764.48, "end": 2767.8599999999997, "text": " because we have batch size of one, it's gonna be easier."}, {"start": 2767.8599999999997, "end": 2769.68, "text": " Okay, so here we are."}, {"start": 2769.68, "end": 2773.44, "text": " So we have a vector that's like 1024."}, {"start": 2774.44, "end": 2778.68, "text": " So this is gonna be 1024, okay?"}, {"start": 2778.68, "end": 2781.7999999999997, "text": " And it might contain in any of these slots,"}, {"start": 2781.7999999999997, "end": 2784.9199999999996, "text": " in any of the dimensions, let's say this is dimension I,"}, {"start": 2784.9199999999996, "end": 2788.04, "text": " it might contain numbers from zero to 50K."}, {"start": 2788.04, "end": 2791.08, "text": " I'm just gonna run the numbers because it's easier, okay?"}, {"start": 2791.08, "end": 2795.16, "text": " So now let's assume we are on this device here."}, {"start": 2795.16, "end": 2798.6, "text": " Okay, so that device only has the following indices."}, {"start": 2798.6, "end": 2801.24, "text": " So it has, so we have two tables."}, {"start": 2802.12, "end": 2804.92, "text": " This one contains from zero to 25K,"}, {"start": 2804.92, "end": 2809.92, "text": " and this one contains from 25K to 50K, 25K to 50K, okay?"}, {"start": 2812.44, "end": 2814.2, "text": " So what this device is going to do"}, {"start": 2814.2, "end": 2818.04, "text": " is it's going to mask all of those values here."}, {"start": 2818.04, "end": 2819.4, "text": " So let me take a different color."}, {"start": 2819.4, "end": 2822.16, "text": " So all of those values that have numbers"}, {"start": 2822.16, "end": 2824.4, "text": " that are zero to 25K,"}, {"start": 2824.4, "end": 2826.52, "text": " because it does not contain those numbers."}, {"start": 2826.52, "end": 2828.84, "text": " So it's going to mask them out by putting them to zero."}, {"start": 2828.84, "end": 2830.92, "text": " So let's see, this one is also put to zero,"}, {"start": 2830.92, "end": 2833.16, "text": " this one is also put to zero, blah, blah, blah, okay?"}, {"start": 2833.16, "end": 2835.88, "text": " And then it's going to take this number here."}, {"start": 2835.88, "end": 2839.36, "text": " Let's say this one is 26K."}, {"start": 2839.36, "end": 2842.56, "text": " It's going to subtract the 25K because that's the offset,"}, {"start": 2843.72, "end": 2845.8, "text": " and we're going to end up with 1K."}, {"start": 2845.8, "end": 2848.84, "text": " And 1K, when you index with 1K into this table,"}, {"start": 2849.8, "end": 2853.32, "text": " you end up with some vector, embedding vector here."}, {"start": 2853.32, "end": 2856.5, "text": " And that vector is, if you take a look at the whole table,"}, {"start": 2856.5, "end": 2859.08, "text": " is at the 26K, okay?"}, {"start": 2859.08, "end": 2860.64, "text": " So that's why we do all of this offsetting."}, {"start": 2860.64, "end": 2862.64, "text": " Hopefully this now makes a bit more sense."}, {"start": 2862.64, "end": 2864.84, "text": " Okay, let's continue."}, {"start": 2864.84, "end": 2869.84, "text": " So now we do the embedding and we get the results."}, {"start": 2869.84, "end": 2872.08, "text": " Okay, so let's check out the shapes."}, {"start": 2872.96, "end": 2875.88, "text": " Shape is going to be 8024,024."}, {"start": 2875.88, "end": 2878.76, "text": " This is a sequence length, this is the embedding dimension."}, {"start": 2878.76, "end": 2881.08, "text": " And now what we would have to do"}, {"start": 2881.08, "end": 2886.08, "text": " is we would have to set those vectors that were masked out"}, {"start": 2886.88, "end": 2890.88, "text": " to zero, otherwise they would have incorrect results, okay?"}, {"start": 2890.88, "end": 2892.6, "text": " And finally, what we do is we call this"}, {"start": 2892.6, "end": 2894.98, "text": " reduced from tensor model parallel region."}, {"start": 2894.98, "end": 2896.3199999999997, "text": " Let me quickly enter this one,"}, {"start": 2896.3199999999997, "end": 2898.7599999999998, "text": " and then I'm going to show you what is going on here."}, {"start": 2898.7599999999998, "end": 2901.16, "text": " So here we are."}, {"start": 2901.16, "end": 2902.3199999999997, "text": " Let me just go like this."}, {"start": 2902.3199999999997, "end": 2904.12, "text": " So let me actually enter here."}, {"start": 2904.12, "end": 2907.12, "text": " Let me enter here, reduce."}, {"start": 2907.12, "end": 2908.7599999999998, "text": " And here is what's going to happen."}, {"start": 2908.76, "end": 2912.5600000000004, "text": " In case that the model that we have multiple stages,"}, {"start": 2912.5600000000004, "end": 2913.6800000000003, "text": " sorry, this is model parallelism."}, {"start": 2913.6800000000003, "end": 2917.96, "text": " In case that we have multiple shards across,"}, {"start": 2917.96, "end": 2920.1800000000003, "text": " if we have model parallelism of size two,"}, {"start": 2920.1800000000003, "end": 2922.84, "text": " as the running example is using,"}, {"start": 2922.84, "end": 2925.1600000000003, "text": " then we would not enter this line."}, {"start": 2925.1600000000003, "end": 2928.96, "text": " We would do the all reduce across the group,"}, {"start": 2928.96, "end": 2930.6800000000003, "text": " across the model parallel group."}, {"start": 2930.6800000000003, "end": 2933.92, "text": " So let me break that down, what we have done right there."}, {"start": 2933.92, "end": 2936.32, "text": " So let's again, get back to the drawing"}, {"start": 2936.32, "end": 2937.32, "text": " and let's see what happened."}, {"start": 2937.32, "end": 2941.92, "text": " Okay, so, okay, because these guys here were masked out,"}, {"start": 2941.92, "end": 2943.6800000000003, "text": " so that means they were, they maybe,"}, {"start": 2943.6800000000003, "end": 2945.32, "text": " this one was maybe 14K."}, {"start": 2945.32, "end": 2948.92, "text": " And we've artificially set its value to zero."}, {"start": 2948.92, "end": 2951.56, "text": " Because it was zero, that means it was indexing"}, {"start": 2951.56, "end": 2954.28, "text": " the zero vector from this table, which is incorrect"}, {"start": 2954.28, "end": 2956.32, "text": " because we want the 14,000, right?"}, {"start": 2956.32, "end": 2958.84, "text": " So this one here, again, let me just change the color"}, {"start": 2958.84, "end": 2959.8, "text": " and let me explain."}, {"start": 2959.8, "end": 2962.0, "text": " So this one here was maybe 14K."}, {"start": 2963.0, "end": 2965.4, "text": " So that's like the token ID, okay?"}, {"start": 2965.4, "end": 2967.92, "text": " So that one was indexing somewhere here,"}, {"start": 2967.92, "end": 2969.96, "text": " but because we've set it to zero,"}, {"start": 2969.96, "end": 2972.6800000000003, "text": " it was actually indexing this vector here."}, {"start": 2972.6800000000003, "end": 2975.14, "text": " And that's why we have, oops, it's glitching."}, {"start": 2975.14, "end": 2977.36, "text": " That's why we had to do the zeroing out"}, {"start": 2977.36, "end": 2980.76, "text": " because we don't want to have this value here."}, {"start": 2980.76, "end": 2982.92, "text": " We don't want to have this embedding vector"}, {"start": 2982.92, "end": 2985.4, "text": " for this particular token."}, {"start": 2985.4, "end": 2986.96, "text": " That's why we do the zeroing."}, {"start": 2986.96, "end": 2988.7200000000003, "text": " That's why we do this part here."}, {"start": 2988.7200000000003, "end": 2992.34, "text": " And then what the all reduce does is the following."}, {"start": 2992.34, "end": 2993.46, "text": " Let me go back here."}, {"start": 2993.46, "end": 2994.4, "text": " So this is going to happen."}, {"start": 2994.4, "end": 2996.76, "text": " Let me see whether I can explain this nicely."}, {"start": 2996.76, "end": 3001.76, "text": " Basically, this device here mapped this set of tokens"}, {"start": 3003.02, "end": 3008.02, "text": " into obviously various bunch of embedding vectors."}, {"start": 3009.76, "end": 3013.04, "text": " So the dimensionality here is the same"}, {"start": 3013.04, "end": 3014.7200000000003, "text": " as the dimensionality up there."}, {"start": 3014.7200000000003, "end": 3018.88, "text": " So it's 1024, but here we now have the hidden dimension."}, {"start": 3018.88, "end": 3022.32, "text": " Okay, so this thing, this vector here"}, {"start": 3022.32, "end": 3024.6000000000004, "text": " is gonna be valid for this device."}, {"start": 3024.6000000000004, "end": 3026.2400000000002, "text": " So that vector is gonna be valid."}, {"start": 3026.2400000000002, "end": 3028.0, "text": " So let me change the color to green."}, {"start": 3028.84, "end": 3031.1200000000003, "text": " So this one here, this token is gonna be valid"}, {"start": 3032.48, "end": 3033.88, "text": " on this particular device."}, {"start": 3033.88, "end": 3035.42, "text": " Okay, so it's valid here."}, {"start": 3035.42, "end": 3039.2200000000003, "text": " But as we just saw, it's going to be invalid for this one."}, {"start": 3039.2200000000003, "end": 3042.0800000000004, "text": " Okay, so let me draw another table."}, {"start": 3042.0800000000004, "end": 3043.7200000000003, "text": " So let me draw the results."}, {"start": 3043.7200000000003, "end": 3045.2000000000003, "text": " So this is the output from this device."}, {"start": 3045.2000000000003, "end": 3047.56, "text": " So this is the output from the second device."}, {"start": 3047.56, "end": 3052.56, "text": " And it contains a corrupt data, actually all zeros,"}, {"start": 3056.08, "end": 3057.6, "text": " because we've zeroed it out."}, {"start": 3057.6, "end": 3059.12, "text": " So it's gonna be all zeros."}, {"start": 3059.12, "end": 3063.12, "text": " And so now what the all reduce does is basically,"}, {"start": 3063.12, "end": 3064.7599999999998, "text": " you sum these up."}, {"start": 3064.7599999999998, "end": 3067.88, "text": " If you sum them up, you'll end up with a correct result"}, {"start": 3067.88, "end": 3071.08, "text": " both here and here, because if you just add zeros"}, {"start": 3071.08, "end": 3074.32, "text": " onto this device, you keep the same value."}, {"start": 3074.32, "end": 3077.04, "text": " But this device will now have the correct value"}, {"start": 3077.04, "end": 3078.2, "text": " of the embedding vector."}, {"start": 3078.2, "end": 3080.4, "text": " And that's the magic of the all reduce."}, {"start": 3080.4, "end": 3084.34, "text": " So what's going to happen is these GPUs here"}, {"start": 3084.34, "end": 3086.68, "text": " are going to do that type of a communication"}, {"start": 3086.68, "end": 3089.2799999999997, "text": " and update the embedding vectors."}, {"start": 3089.2799999999997, "end": 3091.64, "text": " So that's a beautiful thing, like if you ask me."}, {"start": 3091.64, "end": 3096.64, "text": " So hopefully I managed to explain it nicely."}, {"start": 3096.8, "end": 3098.7599999999998, "text": " And now finally, let's get out of here."}, {"start": 3098.7599999999998, "end": 3102.88, "text": " Okay, again, just to update our mental model."}, {"start": 3102.88, "end": 3106.48, "text": " So where we are at the moment is in the state"}, {"start": 3106.48, "end": 3111.48, "text": " where basically these two devices here are waiting"}, {"start": 3111.72, "end": 3113.2, "text": " for the information."}, {"start": 3113.2, "end": 3116.7400000000002, "text": " And we have the embedding vectors on this device"}, {"start": 3116.7400000000002, "end": 3119.04, "text": " and on this device being the same"}, {"start": 3119.04, "end": 3121.4, "text": " because they're synchronized because of this all reduce"}, {"start": 3121.4, "end": 3122.22, "text": " we've just done."}, {"start": 3122.22, "end": 3124.08, "text": " Okay, so we are now there."}, {"start": 3124.08, "end": 3126.48, "text": " We have that and now let's get back to the code."}, {"start": 3127.4, "end": 3130.32, "text": " Okay, so after this, we're gonna have some"}, {"start": 3130.32, "end": 3134.56, "text": " positional embeddings being added that's less interesting"}, {"start": 3134.56, "end": 3136.92, "text": " because it's not sharded."}, {"start": 3136.92, "end": 3139.7999999999997, "text": " Let's continue we have some drop out, blah, blah, blah."}, {"start": 3139.7999999999997, "end": 3142.4, "text": " And we end up with the output."}, {"start": 3142.4, "end": 3145.36, "text": " Okay, so now let's see what's going on."}, {"start": 3145.36, "end": 3148.4, "text": " So we now do the call to the encoder."}, {"start": 3148.4, "end": 3151.36, "text": " And that's just going to be our transformer basically."}, {"start": 3151.36, "end": 3153.36, "text": " So let's enter there."}, {"start": 3154.44, "end": 3155.7, "text": " Let me hit F5."}, {"start": 3155.7, "end": 3158.16, "text": " So we are now in the parallel transformer."}, {"start": 3158.16, "end": 3160.72, "text": " Let's step over all of these."}, {"start": 3160.72, "end": 3164.2, "text": " If self pre process, so again, if we're in the first stage"}, {"start": 3164.2, "end": 3165.7599999999998, "text": " we'll be doing this."}, {"start": 3165.7599999999998, "end": 3167.2799999999997, "text": " So if we're in the first stage,"}, {"start": 3167.2799999999997, "end": 3171.9199999999996, "text": " we'll be grabbing the embedding vectors we just created."}, {"start": 3171.9199999999996, "end": 3176.12, "text": " So let me show you the shape, 8024 to 1024."}, {"start": 3176.12, "end": 3179.08, "text": " And we just grab those, we do some transpose operation"}, {"start": 3179.08, "end": 3181.8199999999997, "text": " and we do continuous and we continue."}, {"start": 3181.8199999999997, "end": 3184.8799999999997, "text": " But if you recall this line here,"}, {"start": 3184.8799999999997, "end": 3187.3999999999996, "text": " this was from the, if we were,"}, {"start": 3187.3999999999996, "end": 3190.12, "text": " if this was not the first stage of the pipeline,"}, {"start": 3190.12, "end": 3193.0, "text": " but instead if this was maybe second or third"}, {"start": 3193.0, "end": 3194.96, "text": " or whatever stage,"}, {"start": 3194.96, "end": 3197.24, "text": " then instead we'd have hidden states"}, {"start": 3197.24, "end": 3199.38, "text": " being equal to input tensor,"}, {"start": 3199.38, "end": 3201.32, "text": " which came from the previous stage."}, {"start": 3201.32, "end": 3204.8, "text": " Okay, so we would bypass this part here."}, {"start": 3204.8, "end": 3206.16, "text": " Hopefully that makes sense."}, {"start": 3207.12, "end": 3208.6, "text": " Let's not continue."}, {"start": 3208.6, "end": 3211.2, "text": " We're not doing checkpointing."}, {"start": 3211.2, "end": 3214.52, "text": " Okay, now we are just going to do a feed forward"}, {"start": 3214.52, "end": 3216.96, "text": " through all of the transformer layers."}, {"start": 3216.96, "end": 3219.6, "text": " Okay, so let's do that."}, {"start": 3219.6, "end": 3221.58, "text": " Let's grab the, so here is the layer."}, {"start": 3221.58, "end": 3223.56, "text": " I can hit F5, I think."}, {"start": 3223.56, "end": 3226.08, "text": " So we are now in the parallel transformer layer."}, {"start": 3226.08, "end": 3229.24, "text": " And the first thing we do is we pass the data"}, {"start": 3229.24, "end": 3230.58, "text": " through the layer norm."}, {"start": 3230.58, "end": 3233.52, "text": " Let me quickly show you how that's gonna look like"}, {"start": 3233.52, "end": 3234.88, "text": " because it's fused."}, {"start": 3234.88, "end": 3236.72, "text": " Oh my God, this sucks."}, {"start": 3236.72, "end": 3239.44, "text": " Okay, so let's just step over here."}, {"start": 3239.44, "end": 3240.7999999999997, "text": " Okay, here we are."}, {"start": 3240.7999999999997, "end": 3244.3199999999997, "text": " And here is the fused layer norm,"}, {"start": 3244.3199999999997, "end": 3246.88, "text": " F5M function we call this apply method."}, {"start": 3246.88, "end": 3248.56, "text": " So let me enter here."}, {"start": 3248.56, "end": 3250.84, "text": " So here you can see how it looks like."}, {"start": 3250.84, "end": 3252.0, "text": " Oops."}, {"start": 3252.0, "end": 3256.6000000000004, "text": " So what it does is ultimately it calls this CUDA kernel,"}, {"start": 3256.6000000000004, "end": 3257.6800000000003, "text": " which is a global variable,"}, {"start": 3257.6800000000003, "end": 3260.2000000000003, "text": " which we loaded in the beginning of this program."}, {"start": 3261.08, "end": 3263.08, "text": " And then we just call the forward F5."}, {"start": 3263.08, "end": 3264.6800000000003, "text": " So we cannot actually enter here"}, {"start": 3264.6800000000003, "end": 3266.7200000000003, "text": " because it's just a compiled C++ code."}, {"start": 3266.7200000000003, "end": 3269.3, "text": " So I cannot enter with this current debugger."}, {"start": 3269.3, "end": 3271.92, "text": " Maybe there is some, maybe it can be tweaked"}, {"start": 3271.92, "end": 3273.84, "text": " so we can enter and see the C++ code,"}, {"start": 3273.84, "end": 3276.0, "text": " but I'm not sure how that will work."}, {"start": 3276.0, "end": 3277.32, "text": " Okay, so in any case,"}, {"start": 3277.32, "end": 3279.58, "text": " I just wanted to show you how the kernels look like."}, {"start": 3279.58, "end": 3282.66, "text": " Now let's exit from the layer norm."}, {"start": 3282.66, "end": 3286.16, "text": " And now let's hit the actual attention part."}, {"start": 3286.16, "end": 3288.08, "text": " Okay, let me hit F5."}, {"start": 3288.08, "end": 3289.2, "text": " Here is the attention part."}, {"start": 3289.2, "end": 3291.52, "text": " So we are in the parallel attention part."}, {"start": 3291.52, "end": 3293.88, "text": " Again, remember the weights are sharded"}, {"start": 3293.88, "end": 3296.54, "text": " such that each of these devices"}, {"start": 3296.54, "end": 3299.44, "text": " in the rank zero and rank one of the model parallelism"}, {"start": 3299.44, "end": 3301.52, "text": " only have a portion of the weights."}, {"start": 3301.52, "end": 3304.54, "text": " Okay, so we end up with some..."}, {"start": 3304.54, "end": 3307.48, "text": " Okay, so we first hit the column parallel layer."}, {"start": 3307.48, "end": 3309.84, "text": " Let me go back up again here up the stack."}, {"start": 3309.84, "end": 3311.88, "text": " So we hit the query key value"}, {"start": 3311.88, "end": 3316.2400000000002, "text": " is actually the column parallel layer for a function."}, {"start": 3316.2400000000002, "end": 3318.12, "text": " So that's again, let me show you here."}, {"start": 3318.12, "end": 3320.48, "text": " So we are hitting the column parallel."}, {"start": 3320.48, "end": 3322.32, "text": " This is the part I was showing you."}, {"start": 3322.32, "end": 3323.76, "text": " So let me go back to the code now"}, {"start": 3323.76, "end": 3324.96, "text": " and let's see what's going on here."}, {"start": 3324.96, "end": 3327.84, "text": " So we call this weird function copy to tensor."}, {"start": 3327.84, "end": 3330.32, "text": " So what is it doing exactly?"}, {"start": 3330.32, "end": 3331.26, "text": " Let's maybe enter."}, {"start": 3331.26, "end": 3332.2400000000002, "text": " So let's enter here."}, {"start": 3332.2400000000002, "end": 3333.32, "text": " Let's enter here."}, {"start": 3333.32, "end": 3335.84, "text": " And you can see here in the forward pass,"}, {"start": 3335.84, "end": 3337.56, "text": " it's just doing identity function."}, {"start": 3337.56, "end": 3339.0, "text": " So why are we doing this?"}, {"start": 3339.0, "end": 3341.2000000000003, "text": " Well, that's because in the backward"}, {"start": 3341.2000000000003, "end": 3343.1600000000003, "text": " we'll be doing the reduce method."}, {"start": 3343.1600000000003, "end": 3344.84, "text": " We'll be calling the all reduce."}, {"start": 3344.84, "end": 3346.28, "text": " So let's see why that makes sense."}, {"start": 3346.28, "end": 3347.52, "text": " Let me show you the diagram here."}, {"start": 3347.52, "end": 3349.78, "text": " It's kind of clear once you see it."}, {"start": 3349.78, "end": 3352.28, "text": " So what we're doing here is this."}, {"start": 3352.28, "end": 3354.08, "text": " We are implementing F."}, {"start": 3354.08, "end": 3356.1200000000003, "text": " F is, as you can see here,"}, {"start": 3356.1200000000003, "end": 3359.6400000000003, "text": " it's just the identity function in the forward pass"}, {"start": 3359.6400000000003, "end": 3361.48, "text": " because we just passed the data,"}, {"start": 3361.48, "end": 3364.32, "text": " the same data X to one device"}, {"start": 3364.32, "end": 3366.02, "text": " as well as to the other device."}, {"start": 3366.02, "end": 3368.04, "text": " But once we go backward,"}, {"start": 3368.04, "end": 3370.06, "text": " we want to all reduce the information."}, {"start": 3370.06, "end": 3371.78, "text": " We want to all reduce the gradients"}, {"start": 3371.78, "end": 3374.96, "text": " that came from this portion and this portion."}, {"start": 3374.96, "end": 3377.0800000000004, "text": " We want to basically average them"}, {"start": 3377.0800000000004, "end": 3380.2400000000002, "text": " and then pass them to the previous layers."}, {"start": 3380.2400000000002, "end": 3381.1800000000003, "text": " Hopefully that makes sense."}, {"start": 3381.1800000000003, "end": 3383.96, "text": " That's why we have it implemented in such a way."}, {"start": 3383.96, "end": 3385.84, "text": " Let's now exit this part."}, {"start": 3386.92, "end": 3387.84, "text": " Now this is less interesting."}, {"start": 3387.84, "end": 3390.96, "text": " We just apply the actual linear layer"}, {"start": 3390.96, "end": 3394.44, "text": " and then blah, blah, blah, bias."}, {"start": 3394.44, "end": 3395.28, "text": " We can exit there."}, {"start": 3395.28, "end": 3396.12, "text": " That's less interesting."}, {"start": 3396.12, "end": 3398.7200000000003, "text": " Okay, so now we just have some,"}, {"start": 3398.7200000000003, "end": 3400.88, "text": " what's this, reshaping."}, {"start": 3400.88, "end": 3402.9, "text": " We have some view, blah, blah, blah."}, {"start": 3402.9, "end": 3405.32, "text": " We're splitting into three parts, the query key values."}, {"start": 3405.32, "end": 3407.4, "text": " I'm not gonna focus on the actual transformer logic"}, {"start": 3407.4, "end": 3410.12, "text": " because I do assume you either know that"}, {"start": 3410.12, "end": 3411.28, "text": " or even if you don't know,"}, {"start": 3411.28, "end": 3414.46, "text": " just focus, let's focus on the 3D parallelism."}, {"start": 3414.46, "end": 3418.48, "text": " So again, just some transformer stuff."}, {"start": 3418.48, "end": 3420.06, "text": " We're changing the view"}, {"start": 3420.06, "end": 3422.56, "text": " so we can do multiplication between the keys"}, {"start": 3422.56, "end": 3427.56, "text": " and the queries and then basically use that to, okay."}, {"start": 3428.32, "end": 3431.92, "text": " So here we just preallocate a tensor"}, {"start": 3431.92, "end": 3434.32, "text": " that's gonna contain the scores"}, {"start": 3435.2, "end": 3439.2999999999997, "text": " for the, that we'll then apply the softmax."}, {"start": 3440.96, "end": 3442.84, "text": " We'll then apply a softmax."}, {"start": 3442.84, "end": 3443.68, "text": " As you can see here,"}, {"start": 3443.68, "end": 3447.6, "text": " so we first do the multiplication of queries and keys here"}, {"start": 3447.6, "end": 3450.68, "text": " so we can skip all of that."}, {"start": 3450.68, "end": 3453.68, "text": " And then we do, we now have attention scores."}, {"start": 3453.68, "end": 3457.46, "text": " And in order to convert scores into actual probabilities,"}, {"start": 3457.46, "end": 3459.18, "text": " we have to apply the softmax."}, {"start": 3460.68, "end": 3462.8399999999997, "text": " So again, I'm gonna exit this softmax layer."}, {"start": 3462.8399999999997, "end": 3464.16, "text": " It's not that interesting."}, {"start": 3465.12, "end": 3466.7599999999998, "text": " We've already seen the CUDA kernels."}, {"start": 3466.7599999999998, "end": 3469.0, "text": " So I'm gonna just skip that."}, {"start": 3469.0, "end": 3470.08, "text": " Okay, attention."}, {"start": 3470.08, "end": 3472.68, "text": " We just do the dropout, blah, blah, blah."}, {"start": 3472.68, "end": 3473.88, "text": " Let's continue."}, {"start": 3473.88, "end": 3477.7200000000003, "text": " My computer is kind of struggling."}, {"start": 3477.7200000000003, "end": 3480.12, "text": " You can see it's kind of lagging behind."}, {"start": 3480.12, "end": 3483.04, "text": " So I cannot go faster than this."}, {"start": 3483.04, "end": 3486.2400000000002, "text": " We just do some, again, view manipulation."}, {"start": 3486.2400000000002, "end": 3488.92, "text": " And then we do the batch matrix multiply,"}, {"start": 3488.92, "end": 3491.84, "text": " which is going to basically sum up the value vectors."}, {"start": 3491.84, "end": 3494.2400000000002, "text": " And we end up with the final result here."}, {"start": 3494.2400000000002, "end": 3495.88, "text": " And then we just change the view."}, {"start": 3495.88, "end": 3498.2000000000003, "text": " We do pre-mute, blah, blah, blah."}, {"start": 3498.2000000000003, "end": 3501.76, "text": " Nothing new there, hopefully."}, {"start": 3501.76, "end": 3504.46, "text": " And finally, we hit the dense part."}, {"start": 3504.46, "end": 3507.6000000000004, "text": " And the dense part is, let me show you here."}, {"start": 3507.6000000000004, "end": 3508.76, "text": " The dense part is this part."}, {"start": 3508.76, "end": 3512.48, "text": " So that should be the row parallel linear layer, right?"}, {"start": 3512.48, "end": 3514.44, "text": " Because we're splitting the, we're now here."}, {"start": 3514.44, "end": 3515.6000000000004, "text": " We're now here."}, {"start": 3515.6000000000004, "end": 3517.1200000000003, "text": " Let me go back to the code."}, {"start": 3518.5600000000004, "end": 3519.4, "text": " Let me enter here."}, {"start": 3519.4, "end": 3520.2200000000003, "text": " Whoops."}, {"start": 3520.2200000000003, "end": 3521.36, "text": " Let me enter here."}, {"start": 3522.6400000000003, "end": 3523.48, "text": " Oh my God."}, {"start": 3525.0200000000004, "end": 3528.1400000000003, "text": " So we are in the row parallel linear layer, as I said."}, {"start": 3528.1400000000003, "end": 3530.6400000000003, "text": " So let's quickly see how it's going to work."}, {"start": 3530.64, "end": 3533.7599999999998, "text": " Again, a couple of details there."}, {"start": 3533.7599999999998, "end": 3538.56, "text": " So you can see we first applied the linear layer,"}, {"start": 3538.56, "end": 3539.8799999999997, "text": " and then we call this reduce"}, {"start": 3539.8799999999997, "end": 3542.2799999999997, "text": " from tensor model parallel region."}, {"start": 3542.2799999999997, "end": 3544.94, "text": " So let's see how that thing looks like."}, {"start": 3544.94, "end": 3547.56, "text": " That one in the forward pass is doing the all reduce."}, {"start": 3547.56, "end": 3550.7999999999997, "text": " And in the backward, it's just doing the identity function."}, {"start": 3550.7999999999997, "end": 3554.64, "text": " So it's just the, literally the polar opposite"}, {"start": 3554.64, "end": 3558.0, "text": " from what the column-wise was doing, right?"}, {"start": 3558.0, "end": 3559.52, "text": " So let's go back here."}, {"start": 3559.52, "end": 3560.58, "text": " And why is that?"}, {"start": 3560.58, "end": 3562.2799999999997, "text": " Let me show you the diagram again."}, {"start": 3562.2799999999997, "end": 3563.92, "text": " So here you can see in the forward,"}, {"start": 3563.92, "end": 3566.14, "text": " we want to all reduce these weights"}, {"start": 3566.14, "end": 3569.92, "text": " such that we end up with the correct result on both devices."}, {"start": 3569.92, "end": 3571.12, "text": " But in the backward,"}, {"start": 3571.12, "end": 3573.7599999999998, "text": " we're just copy pasting the gradients."}, {"start": 3573.7599999999998, "end": 3576.72, "text": " That's the reason why that was implemented"}, {"start": 3576.72, "end": 3577.9, "text": " in such a manner."}, {"start": 3577.9, "end": 3578.74, "text": " Okay."}, {"start": 3578.74, "end": 3582.0, "text": " So I'm going to skip all of this"}, {"start": 3582.0, "end": 3584.12, "text": " because it's not relevant."}, {"start": 3584.12, "end": 3587.0, "text": " And we are done with the parallel tension."}, {"start": 3587.0, "end": 3589.4, "text": " We are exiting the attention right now."}, {"start": 3589.4, "end": 3592.2400000000002, "text": " And okay, now we have some residual stuff,"}, {"start": 3592.2400000000002, "end": 3594.6, "text": " blah, blah, blah, biases."}, {"start": 3594.6, "end": 3595.92, "text": " We don't care about those."}, {"start": 3595.92, "end": 3597.4, "text": " Some of those are JIT compiled"}, {"start": 3597.4, "end": 3600.0, "text": " so we wouldn't be able to step through it either way."}, {"start": 3600.96, "end": 3602.2200000000003, "text": " Post attention layer norm."}, {"start": 3602.2200000000003, "end": 3604.2400000000002, "text": " I think we're going to hit the break point here."}, {"start": 3604.2400000000002, "end": 3607.48, "text": " So I'm going to remove it, exit this function."}, {"start": 3607.48, "end": 3609.34, "text": " Let's exit the function here."}, {"start": 3611.12, "end": 3612.04, "text": " And we're back."}, {"start": 3612.04, "end": 3614.44, "text": " So we don't have the coder."}, {"start": 3614.44, "end": 3616.2000000000003, "text": " Now we hit the MLP."}, {"start": 3616.2000000000003, "end": 3618.2200000000003, "text": " So MLP is going to be this part."}, {"start": 3618.22, "end": 3620.6, "text": " Again, we're now past the data"}, {"start": 3620.6, "end": 3622.0, "text": " that came out of this part."}, {"start": 3622.0, "end": 3623.3999999999996, "text": " We're going to pass it here"}, {"start": 3623.3999999999996, "end": 3624.7599999999998, "text": " and we're going to repeat the same thing."}, {"start": 3624.7599999999998, "end": 3625.64, "text": " So I'm going to skip it."}, {"start": 3625.64, "end": 3627.9399999999996, "text": " So we're going to do the, as you can see here,"}, {"start": 3627.9399999999996, "end": 3629.68, "text": " column wise and then row wise,"}, {"start": 3629.68, "end": 3631.2, "text": " but it's the same logic."}, {"start": 3631.2, "end": 3633.64, "text": " So we can safely skip this part."}, {"start": 3633.64, "end": 3634.48, "text": " Okay."}, {"start": 3634.48, "end": 3636.6, "text": " So let me do disable break points."}, {"start": 3636.6, "end": 3640.12, "text": " We do the same thing and that's it."}, {"start": 3640.12, "end": 3641.16, "text": " We are done here."}, {"start": 3642.48, "end": 3644.3999999999996, "text": " Let me just exit this part here"}, {"start": 3644.3999999999996, "end": 3646.3999999999996, "text": " and let's return back the value."}, {"start": 3646.4, "end": 3649.44, "text": " And that was the layer one."}, {"start": 3649.44, "end": 3651.96, "text": " We obviously repeat that 12 times."}, {"start": 3651.96, "end": 3655.1600000000003, "text": " So I'm going to put the break point here, hit F5."}, {"start": 3655.1600000000003, "end": 3656.8, "text": " And now we're here."}, {"start": 3656.8, "end": 3660.14, "text": " So let's see where we are now here with the diagram."}, {"start": 3660.14, "end": 3663.44, "text": " So the data has successfully propagated all the way."}, {"start": 3663.44, "end": 3664.28, "text": " Let me change."}, {"start": 3664.28, "end": 3666.1600000000003, "text": " I'm not sure what color is the best one here."}, {"start": 3666.1600000000003, "end": 3668.52, "text": " So maybe let's do it like this."}, {"start": 3668.52, "end": 3669.58, "text": " So we are now here."}, {"start": 3671.46, "end": 3674.96, "text": " We successfully propagated our activations"}, {"start": 3674.96, "end": 3676.36, "text": " all the way to here."}, {"start": 3676.36, "end": 3678.36, "text": " And now we'll be passing them slowly"}, {"start": 3678.36, "end": 3680.6400000000003, "text": " to the second stage of the pipeline."}, {"start": 3680.6400000000003, "end": 3683.0, "text": " So let's go back to the code."}, {"start": 3683.0, "end": 3685.8, "text": " In case we are in the last stage of the pipeline,"}, {"start": 3685.8, "end": 3688.08, "text": " we'd be doing this additionally,"}, {"start": 3688.08, "end": 3690.6800000000003, "text": " passing it through the layer norms."}, {"start": 3690.6800000000003, "end": 3691.92, "text": " And that's it."}, {"start": 3691.92, "end": 3693.7200000000003, "text": " We are now here."}, {"start": 3693.7200000000003, "end": 3696.6, "text": " We can step over all of this"}, {"start": 3696.6, "end": 3699.2400000000002, "text": " and we return back the results"}, {"start": 3699.2400000000002, "end": 3701.2400000000002, "text": " from our transformer language model."}, {"start": 3701.2400000000002, "end": 3703.52, "text": " Guys, this is pretty much it."}, {"start": 3703.52, "end": 3704.58, "text": " Like this is,"}, {"start": 3704.58, "end": 3708.2999999999997, "text": " we've seen the main concepts already."}, {"start": 3708.2999999999997, "end": 3710.7799999999997, "text": " There is now a couple of additional details here."}, {"start": 3710.7799999999997, "end": 3714.2999999999997, "text": " So like post language model processing was this one."}, {"start": 3714.2999999999997, "end": 3716.18, "text": " Let me enable the break points."}, {"start": 3716.18, "end": 3717.58, "text": " Let me hit F5."}, {"start": 3717.58, "end": 3718.86, "text": " Here we are."}, {"start": 3718.86, "end": 3720.74, "text": " So again, there's gonna be some synchronization"}, {"start": 3720.74, "end": 3723.38, "text": " between the devices in these ultimate parts."}, {"start": 3723.38, "end": 3727.1, "text": " So yeah, assume now we're gonna project those weights"}, {"start": 3727.1, "end": 3729.68, "text": " into the vocab space."}, {"start": 3729.68, "end": 3730.74, "text": " So let's see what's going on."}, {"start": 3730.74, "end": 3731.9, "text": " So we have input parallel."}, {"start": 3731.9, "end": 3736.2200000000003, "text": " That's gonna be, I guess, 8,024,024."}, {"start": 3736.2200000000003, "end": 3737.06, "text": " Yep."}, {"start": 3737.06, "end": 3739.3, "text": " And then once we do step over,"}, {"start": 3739.3, "end": 3741.06, "text": " we're gonna end up with the vocab."}, {"start": 3742.34, "end": 3746.38, "text": " So eight, yeah, 1,024,50,304."}, {"start": 3746.38, "end": 3747.7000000000003, "text": " So that's it."}, {"start": 3747.7000000000003, "end": 3749.78, "text": " Again, there was some communication happening here"}, {"start": 3749.78, "end": 3752.86, "text": " between the devices such that we have consistent results"}, {"start": 3752.86, "end": 3754.26, "text": " on both devices."}, {"start": 3754.26, "end": 3755.56, "text": " And that's it."}, {"start": 3755.56, "end": 3756.88, "text": " Let's continue here."}, {"start": 3758.1800000000003, "end": 3759.02, "text": " Whoops."}, {"start": 3759.02, "end": 3761.9, "text": " We're now almost at the bottom of the stack."}, {"start": 3761.9, "end": 3764.2599999999998, "text": " So we are now in this float 16 module."}, {"start": 3764.2599999999998, "end": 3767.22, "text": " What we now do is because we're in the last stage,"}, {"start": 3767.22, "end": 3768.54, "text": " if we're in the last stage,"}, {"start": 3768.54, "end": 3771.14, "text": " we convert from float 16 to float 32,"}, {"start": 3771.14, "end": 3772.86, "text": " and then we return back the results."}, {"start": 3772.86, "end": 3773.78, "text": " Okay."}, {"start": 3773.78, "end": 3778.78, "text": " So let's exit here and we are back."}, {"start": 3778.9, "end": 3782.42, "text": " So again, now we're gonna have this send forward function."}, {"start": 3782.42, "end": 3786.62, "text": " And you can see here, if it's not the last stage,"}, {"start": 3786.62, "end": 3788.02, "text": " then communicate."}, {"start": 3788.02, "end": 3790.74, "text": " So that means all of the stages will have to pass the data"}, {"start": 3790.74, "end": 3792.9, "text": " to some other device,"}, {"start": 3792.9, "end": 3795.38, "text": " unless you're at the last stage of the pipeline,"}, {"start": 3795.38, "end": 3796.3, "text": " which makes sense."}, {"start": 3796.3, "end": 3798.64, "text": " Again, communicate is a blocking method."}, {"start": 3798.64, "end": 3800.54, "text": " So that means if you're waiting for the data,"}, {"start": 3800.54, "end": 3801.62, "text": " you'll be blocked."}, {"start": 3801.62, "end": 3803.82, "text": " That's why you have the concept of micro batches"}, {"start": 3803.82, "end": 3805.54, "text": " so that you can minimize the waiting,"}, {"start": 3805.54, "end": 3807.3, "text": " the idle time of the devices."}, {"start": 3807.3, "end": 3809.22, "text": " Hopefully that makes sense."}, {"start": 3809.22, "end": 3813.06, "text": " Okay, so in the last stage of the pipeline,"}, {"start": 3813.06, "end": 3816.94, "text": " you do this additional vocab parallel cross entropy."}, {"start": 3816.94, "end": 3821.5, "text": " And I can quickly maybe enter this one,"}, {"start": 3821.5, "end": 3824.34, "text": " although I think it's gonna be fairly detailed."}, {"start": 3824.34, "end": 3826.14, "text": " There's a lot of details here,"}, {"start": 3826.14, "end": 3828.1, "text": " but like the concepts are the same,"}, {"start": 3828.1, "end": 3829.58, "text": " like nothing new there."}, {"start": 3829.58, "end": 3832.26, "text": " So I'm gonna slowly step out of all of this."}, {"start": 3832.26, "end": 3835.1, "text": " So yeah, as you can see,"}, {"start": 3835.1, "end": 3836.7400000000002, "text": " they have devices you have to communicate"}, {"start": 3836.7400000000002, "end": 3838.78, "text": " such that we end up with the correct logits"}, {"start": 3838.78, "end": 3841.62, "text": " and then blah, blah, blah."}, {"start": 3841.62, "end": 3843.82, "text": " You can go through this one at your own pace."}, {"start": 3843.82, "end": 3848.1800000000003, "text": " I think it should be fairly easy to understand it."}, {"start": 3848.1800000000003, "end": 3851.2200000000003, "text": " Okay, let's exit that part."}, {"start": 3851.2200000000003, "end": 3852.4, "text": " Let me exit it."}, {"start": 3853.44, "end": 3854.28, "text": " And that's it."}, {"start": 3854.28, "end": 3855.82, "text": " We just do some summation."}, {"start": 3855.82, "end": 3858.1800000000003, "text": " We return back the loss."}, {"start": 3858.1800000000003, "end": 3860.6600000000003, "text": " We do some additional overviews here."}, {"start": 3860.6600000000003, "end": 3863.54, "text": " And now we keep on doing this."}, {"start": 3863.54, "end": 3866.6600000000003, "text": " Ultimately, you'll end up with the evaluation metrics."}, {"start": 3866.6600000000003, "end": 3868.6600000000003, "text": " I'm gonna disable the breakpoints,"}, {"start": 3868.6600000000003, "end": 3872.6400000000003, "text": " hit that five here and exit here."}, {"start": 3872.64, "end": 3876.2599999999998, "text": " And I'll let's say it's from here and that's it."}, {"start": 3876.2599999999998, "end": 3878.62, "text": " Now, if we are on the last rank,"}, {"start": 3878.62, "end": 3880.22, "text": " we just, and by the way,"}, {"start": 3880.22, "end": 3881.98, "text": " here's how the last rank is calculated"}, {"start": 3881.98, "end": 3883.3399999999997, "text": " in this particular example."}, {"start": 3883.3399999999997, "end": 3887.44, "text": " So here you can see that we just get the rank"}, {"start": 3887.44, "end": 3889.94, "text": " without any group, which means that now,"}, {"start": 3889.94, "end": 3893.56, "text": " however many devices we have like in total,"}, {"start": 3893.56, "end": 3895.98, "text": " we just find its absolute rank."}, {"start": 3895.98, "end": 3898.64, "text": " So that means that for example,"}, {"start": 3898.64, "end": 3900.2, "text": " this one will be ranked zero,"}, {"start": 3900.2, "end": 3902.9399999999996, "text": " this one will be ranked one, two, three,"}, {"start": 3902.9399999999996, "end": 3904.8399999999997, "text": " then five, six, seven, eight."}, {"start": 3904.8399999999997, "end": 3907.48, "text": " We will be returning those results."}, {"start": 3907.48, "end": 3910.8799999999997, "text": " So basically we only on the last rank,"}, {"start": 3910.8799999999997, "end": 3913.52, "text": " so only on a single device in our whole cluster,"}, {"start": 3913.52, "end": 3914.8799999999997, "text": " we do this operation."}, {"start": 3914.8799999999997, "end": 3919.8799999999997, "text": " So basically, yeah, just some blah, blah, blah perplexity."}, {"start": 3920.02, "end": 3922.06, "text": " And then we print the results,"}, {"start": 3922.06, "end": 3923.7799999999997, "text": " but we don't care about the actual results."}, {"start": 3923.7799999999997, "end": 3925.06, "text": " So yeah, so you can see here,"}, {"start": 3925.06, "end": 3928.5, "text": " average loss perplexity adjusted perplexity, blah, blah, blah."}, {"start": 3928.5, "end": 3929.64, "text": " Nothing interesting there."}, {"start": 3929.64, "end": 3931.52, "text": " And now finally we're done."}, {"start": 3931.52, "end": 3933.8599999999997, "text": " Okay, guys, we've seen a lot so far."}, {"start": 3933.8599999999997, "end": 3935.3599999999997, "text": " Let me do just a quick recap"}, {"start": 3935.3599999999997, "end": 3937.42, "text": " and then we can wrap this video up."}, {"start": 3937.42, "end": 3939.8399999999997, "text": " So we have a couple of differences depending on"}, {"start": 3939.8399999999997, "end": 3942.12, "text": " whether we are dealing with the last stage of the pipeline,"}, {"start": 3942.12, "end": 3943.62, "text": " with the first stage of the pipeline,"}, {"start": 3943.62, "end": 3946.2999999999997, "text": " or the intermediate stages of the pipeline."}, {"start": 3946.2999999999997, "end": 3948.96, "text": " So let me go and run the eval script again."}, {"start": 3948.96, "end": 3951.14, "text": " So I'm gonna hit run here."}, {"start": 3951.14, "end": 3952.7999999999997, "text": " And I have a breakpoint set"}, {"start": 3952.7999999999997, "end": 3955.56, "text": " inside of this forward step function."}, {"start": 3955.56, "end": 3957.2799999999997, "text": " There's a couple of things I wanna explain here."}, {"start": 3957.28, "end": 3960.7200000000003, "text": " So first things first is this receive forward function."}, {"start": 3960.7200000000003, "end": 3963.84, "text": " Then I'm gonna explain the differences in the forward pass"}, {"start": 3963.84, "end": 3966.7200000000003, "text": " through the model and then the send forward part."}, {"start": 3966.7200000000003, "end": 3969.44, "text": " So first things first, here, if we enter,"}, {"start": 3969.44, "end": 3972.6600000000003, "text": " you can see that the first stage in the pipeline"}, {"start": 3972.6600000000003, "end": 3976.3, "text": " will have the input tensor set to none, as you can see here,"}, {"start": 3976.3, "end": 3978.6400000000003, "text": " which means that they'll actually be fetching the data"}, {"start": 3978.6400000000003, "end": 3979.8, "text": " from the data loader."}, {"start": 3979.8, "end": 3982.6400000000003, "text": " So they'll grab the chunk of data from the data loader"}, {"start": 3982.6400000000003, "end": 3985.48, "text": " and pass that as the input into the embedding table"}, {"start": 3985.48, "end": 3986.96, "text": " and then the transformer layers."}, {"start": 3986.96, "end": 3989.2400000000002, "text": " So again, let me show you the diagram here."}, {"start": 3989.2400000000002, "end": 3991.52, "text": " So that means those guys, the first stage,"}, {"start": 3991.52, "end": 3994.4, "text": " will grab the chunk of data from the batch"}, {"start": 3994.4, "end": 3998.28, "text": " and will pass that data through the, as you can see here,"}, {"start": 3998.28, "end": 4000.88, "text": " embedding table and only they have the embedding table,"}, {"start": 4000.88, "end": 4004.7400000000002, "text": " if you recall from the explanation back then."}, {"start": 4004.7400000000002, "end": 4007.7200000000003, "text": " And then they'll pass the data through the transformer layers"}, {"start": 4007.7200000000003, "end": 4009.8, "text": " and ultimately they end up with this,"}, {"start": 4009.8, "end": 4013.12, "text": " with the output activations denoted in gray color here."}, {"start": 4013.12, "end": 4014.68, "text": " Okay."}, {"start": 4014.68, "end": 4017.72, "text": " And then all of the other stages will block"}, {"start": 4017.72, "end": 4019.0, "text": " on the communicate method."}, {"start": 4019.0, "end": 4022.96, "text": " And that means they'll be waiting until the previous stage"}, {"start": 4022.96, "end": 4026.1, "text": " sends the activations to this stage."}, {"start": 4026.1, "end": 4028.52, "text": " And then it will continue on the execution."}, {"start": 4028.52, "end": 4031.0, "text": " So that means for these guys here,"}, {"start": 4031.0, "end": 4035.44, "text": " they'll be blocked until the gray activations are sent"}, {"start": 4035.44, "end": 4039.96, "text": " to the second stage here, and then they'll do the forward pass."}, {"start": 4039.96, "end": 4041.72, "text": " Okay, so that's first thing."}, {"start": 4041.72, "end": 4046.24, "text": " Then let's go back to the forward step."}, {"start": 4046.24, "end": 4049.3199999999997, "text": " We have the forward pass through the model."}, {"start": 4049.3199999999997, "end": 4052.16, "text": " So I'm gonna now just enable all of the breakpoints"}, {"start": 4052.16, "end": 4053.72, "text": " and let me hit that five just to enter"}, {"start": 4053.72, "end": 4055.7799999999997, "text": " into the actual forward pass."}, {"start": 4055.7799999999997, "end": 4059.24, "text": " So I want to get, whoops, let me just do this."}, {"start": 4059.24, "end": 4062.64, "text": " I'm gonna basically disable this"}, {"start": 4062.64, "end": 4064.3799999999997, "text": " and put a breakpoint here."}, {"start": 4064.3799999999997, "end": 4066.24, "text": " So this is where the difference happens."}, {"start": 4066.24, "end": 4067.72, "text": " So basically this is, as you can see here,"}, {"start": 4067.72, "end": 4070.7999999999997, "text": " GPT model forward, depending on whether we're dealing"}, {"start": 4070.8, "end": 4074.1200000000003, "text": " with the last stage or any other stage of the pipeline,"}, {"start": 4074.1200000000003, "end": 4076.6000000000004, "text": " we'll be entering this branch here."}, {"start": 4076.6000000000004, "end": 4079.6600000000003, "text": " And as you can see here, if I enter this one,"}, {"start": 4079.6600000000003, "end": 4084.6600000000003, "text": " this thing does converts the activations into the logits."}, {"start": 4086.04, "end": 4088.0, "text": " So that means we'll end up"}, {"start": 4088.0, "end": 4090.76, "text": " with the vocab space dimensionality."}, {"start": 4090.76, "end": 4093.84, "text": " So let me digest what I just said there."}, {"start": 4093.84, "end": 4095.88, "text": " So if you take a look at the LM output,"}, {"start": 4095.88, "end": 4097.66, "text": " all of the stages will be just returning"}, {"start": 4097.66, "end": 4099.02, "text": " the pure activations."}, {"start": 4099.02, "end": 4102.64, "text": " So that means 8024,024."}, {"start": 4102.64, "end": 4106.160000000001, "text": " So those are the raw activations from the transformer block."}, {"start": 4106.160000000001, "end": 4108.4800000000005, "text": " But the last stage will be returning"}, {"start": 4109.320000000001, "end": 4112.56, "text": " this processed information because it maps"}, {"start": 4112.56, "end": 4117.56, "text": " the tokens into the vocab space."}, {"start": 4119.540000000001, "end": 4120.96, "text": " So let me show you what I mean by that."}, {"start": 4120.96, "end": 4125.52, "text": " So if we exit here, the outputs for the last stage will be,"}, {"start": 4125.52, "end": 4130.4400000000005, "text": " whoops, I think I have to do step over here"}, {"start": 4130.4400000000005, "end": 4132.02, "text": " and then let's do output shape."}, {"start": 4132.02, "end": 4133.400000000001, "text": " So you can see here 50K."}, {"start": 4134.52, "end": 4138.96, "text": " It will be 25K in the case of having sharded,"}, {"start": 4138.96, "end": 4140.5, "text": " the model sharding."}, {"start": 4141.4400000000005, "end": 4143.120000000001, "text": " And yeah, that's it."}, {"start": 4143.120000000001, "end": 4145.540000000001, "text": " So let me go outside of this thing."}, {"start": 4146.92, "end": 4148.0, "text": " And now we're here."}, {"start": 4148.0, "end": 4152.160000000001, "text": " Okay, so now, again, remember output is one of two things."}, {"start": 4152.160000000001, "end": 4154.96, "text": " It's either the raw activations from the transformer layer"}, {"start": 4154.96, "end": 4159.44, "text": " or it's the tokens mapped into the vocab space"}, {"start": 4159.44, "end": 4161.0, "text": " for the last stage of the pipeline."}, {"start": 4161.0, "end": 4164.92, "text": " And then the send forward, as you can see here,"}, {"start": 4164.92, "end": 4166.72, "text": " all of the stages will be communicating"}, {"start": 4166.72, "end": 4168.24, "text": " the activations to the next stage,"}, {"start": 4168.24, "end": 4169.76, "text": " but the last stage will just skip."}, {"start": 4169.76, "end": 4173.7, "text": " As you can see here, if not the last stage,"}, {"start": 4173.7, "end": 4174.88, "text": " then communicate."}, {"start": 4174.88, "end": 4176.84, "text": " If it is the last stage, we'll just ignore this."}, {"start": 4176.84, "end": 4180.96, "text": " This is a no-op operation, no-op for the last stage."}, {"start": 4180.96, "end": 4183.56, "text": " Okay, and finally, as you can see here,"}, {"start": 4183.56, "end": 4186.72, "text": " we have the last stage branch here,"}, {"start": 4186.72, "end": 4190.72, "text": " which means that the last stage will pass the logits"}, {"start": 4192.240000000001, "end": 4195.84, "text": " into this function here and calculate the cross entropy,"}, {"start": 4195.84, "end": 4198.240000000001, "text": " and then it will return back the loss."}, {"start": 4198.240000000001, "end": 4201.0, "text": " Whereas all of the other stages"}, {"start": 4201.0, "end": 4203.320000000001, "text": " first communicated information to the next stage,"}, {"start": 4203.320000000001, "end": 4205.52, "text": " and then they return none."}, {"start": 4205.52, "end": 4208.4400000000005, "text": " And then if we go a level up from this function,"}, {"start": 4208.4400000000005, "end": 4209.72, "text": " so this is the forward step,"}, {"start": 4209.72, "end": 4212.0, "text": " you can see that the output, again, for the last stage"}, {"start": 4212.0, "end": 4214.28, "text": " will be all reduced and then just added"}, {"start": 4214.28, "end": 4216.7, "text": " to this total output, which will finally return"}, {"start": 4216.7, "end": 4220.04, "text": " and used to calculate the perplexity, et cetera, et cetera."}, {"start": 4220.04, "end": 4221.56, "text": " So hopefully that makes sense."}, {"start": 4221.56, "end": 4223.28, "text": " Again, let me go through the diagram."}, {"start": 4223.28, "end": 4225.1, "text": " I think it's fairly straightforward."}, {"start": 4225.1, "end": 4228.46, "text": " So again, the first stage here will be grabbing the data"}, {"start": 4228.46, "end": 4231.6, "text": " from the data loader, who will be passing that data"}, {"start": 4231.6, "end": 4232.96, "text": " through the embedding table"}, {"start": 4232.96, "end": 4234.46, "text": " and then through the transformer layers"}, {"start": 4234.46, "end": 4236.12, "text": " until we have the raw activations,"}, {"start": 4236.12, "end": 4238.7, "text": " and then they'll be passing those activations"}, {"start": 4238.7, "end": 4239.96, "text": " onto the next stage."}, {"start": 4239.96, "end": 4243.58, "text": " Whereas the next stage will be waiting for the activations"}, {"start": 4243.58, "end": 4246.8, "text": " and then grabbing the activations from the previous stage,"}, {"start": 4246.8, "end": 4251.0, "text": " passing those activations through the 12 transformer layers,"}, {"start": 4251.9, "end": 4254.74, "text": " and finally calculating the logits,"}, {"start": 4254.74, "end": 4256.56, "text": " calculating the cross entropy loss,"}, {"start": 4256.56, "end": 4259.24, "text": " and then we'll be accumulating that information."}, {"start": 4259.24, "end": 4262.0, "text": " If we were to have additionally intermediate stages"}, {"start": 4262.0, "end": 4263.38, "text": " in this diagram, obviously,"}, {"start": 4263.38, "end": 4266.0, "text": " then those would be receiving the activations"}, {"start": 4266.0, "end": 4267.84, "text": " from the previous stage and then just passing"}, {"start": 4267.84, "end": 4269.26, "text": " the activations to the next stage."}, {"start": 4269.26, "end": 4272.12, "text": " So we kind of have three different branches"}, {"start": 4272.12, "end": 4274.400000000001, "text": " in this pipeline logic."}, {"start": 4274.400000000001, "end": 4277.06, "text": " Cool, that was a long video."}, {"start": 4277.06, "end": 4278.22, "text": " Hopefully you liked it."}, {"start": 4278.22, "end": 4281.2, "text": " If you did, I would really appreciate any feedback."}, {"start": 4281.2, "end": 4282.320000000001, "text": " This will be the last video"}, {"start": 4282.320000000001, "end": 4284.68, "text": " from the large language model series."}, {"start": 4284.68, "end": 4287.8, "text": " I might create one more on the GPT-NewX20B"}, {"start": 4287.8, "end": 4290.24, "text": " if I have any interest."}, {"start": 4290.24, "end": 4294.18, "text": " If you'd like me to go through the GPT-NewX20B code base,"}, {"start": 4294.18, "end": 4295.4400000000005, "text": " do comment down below"}, {"start": 4295.4400000000005, "end": 4298.42, "text": " and I'll be monitoring the comments."}, {"start": 4298.42, "end": 4300.36, "text": " In any case, if you found this video useful,"}, {"start": 4300.36, "end": 4302.6, "text": " consider subscribing, sharing it out."}, {"start": 4302.6, "end": 4305.0, "text": " That's the best way you can support the channel."}, {"start": 4305.0, "end": 4306.92, "text": " And until next time, bye bye."}, {"start": 4306.92, "end": 4331.62, "text": " Ab\u016b La mum\u1e6daff\u0101\u1e6d \ud0d3 rather le,"}]