Parallelism 1

Framing: from one GPU to the datacenter

00:04 All right. So today's going to be the second of the basic systems lectures. And now, we're going to move on to multi-machine optimization. And so the focus today is going to be all about parallelism across machines. And so the goal today is going to move from optimizing a single GPU's throughput to being able to understand the complexities and the details that are required to train really large models. And when models get large, they no longer fit on a single GPU. So you've got to split up your models across different machines, but also, you've got to be able to leverage all of the different servers that you have in order to train these models quickly. So we've got both compute and memory concerns that we're going to have to deal with. And communication across different machines, it's going to be quite heterogeneous. We have different kinds of communication across GPUs at different levels of hierarchy. And so this is going to lead to different parallelization paradigms. People use many different parallelization strategies all together at once. And we're going to talk through each one of the very popular ones. And then we'll talk about how you combine them together in order to efficiently train a very large model. And then I'm going to end the lecture with looking at some examples of how people are actually using these parallelization strategies to run their large scale distributed training runs. And so that's going to roughly map to the different parts of this lecture. We're just going to talk about the basics of networking first, and then we're going to talk about how do each of these networking hardware concepts map to different parallelization strategies

01:46 and then finally, some case studies to close off with to show you how it all comes together. So I told you about GPU scaling last week, and it's quite impressive seeing this super exponential curve of flops per GPU going way, way up. But if we want to rapidly scale out both our compute and memory, a single GPU isn't enough. We're going to have to wait for another couple years for this curve to continue going upwards and upwards and upwards. So if we want to train a really powerful language model here and now today, well, we have to rely on multi-machine parallelism. So if we look at the world's fastest supercomputers, that's what's being shown on the right here, the fastest supercomputers have exaflops and exaflops of compute. Those are the green lines that you see over there. That's what you're really going to have to rely on if you're going to try to train the biggest and baddest language models today. And so that's the compute side of why you want to think about multi-machine parallelism. But we've also got a memory angle for thinking about the same thing. So these two are really the core resources and the core concerns that you're going to have to think about. So in terms of memory, many of the models are getting quite big. And of course, memory on GPUs is also growing but not quite as quickly. And a single GPU is not going to be able to fit these models. Maybe eventually in the distant future, we won't have to worry about a lot of these. But we've got billions and billions of parameters, they're not going to fit very nicely into a single GPU. So we have to be very respectful of the memory constraints

The hardware hierarchy

03:22 that we have. So those are the realities that we have to deal with. And what are the tools that we have to have to be able to handle these? Well, GPUs, I'm sure you've noticed in the class cluster, don't come in singletons. A single machine will have multiple GPUs within the same physical rack. And so here's an example. I took this, I think, from the GPT-NeoX paper. But this is an old example, but the same lesson applies to the H100 machines that you have in class. So here, there's eight different GPUs. They're connected to the various CPUs through fast interconnects. Within each GPUs, you see this NVSwitch thing at the bottom. This is very, very fast connections across these eight GPUs. But if these eight GPUs want to talk to GPUs on a different machine, they're going to have to go through a networking switch. And you see this purple line that says, HDR InfiniBand. That's a much slower connection compared to the NVLink connection. You can see the difference in the throughput that's about eight times slower per lane. And so this hardware hierarchy that we have is going to have big implications for how we're going to end up parallelizing our models in practice. And so you can keep this mental model with you as I talk through these things. We have very, very fast connections within a single machine. And then when we go across machines, it's going to get slower. And then depending on the kind of hardware we're using, there might even be another level of slowness once we go beyond, let's say, 256 GPUs networked together. Many of you may already know this having taken systems or networking classes,

Collectives, and the identity that pays for everything

05:03 but here's a very, very brief refresher on collective communication operations. And the reason why I'm going to bring this up is there's one particular important identity or equivalence that you will need to know to really understand some of the finer points of the performance characteristics of the parallelization algorithms. So I'll talk through these, and then I'll talk through one important performance implication. So the first one, which all of you probably have heard of, is all reduce. So you have four machines, four ranks in this case, each one having its own piece of data. And what you'd like to do is perform some of reduction operation. Let's say, I want to sum all of these inputs, and then I want the output to be copied over to every single machine. And this is going to have roughly the cost of two times the total number of things that you're all reducing. You have a broadcast operation. And here, I'm taking a single input from rank two, and I'd like to copy it out to all of the remaining ranks. And this is going to have roughly on the order of one times the total number of outputs in terms of the communication costs. And then we've got reduction where we got different inputs, and that's going to be summed up and then sent only to one machine. And then the two that are quite important, even though these may not be quite as common, is going to be the all gathering scatter. So all gather is an operation where I'm taking a single subcomponent of, let's say, my parameters from rank zero and I'm copying it over to all the ranks. Same thing with rank 1, 2, 3. So each of these are handling different parts, let's say, the parameters and they're copied over to the rest of the machines. So that's copying what I have to everyone else and then reduce scatter, which is I'm taking each of the rows,

06:47 let's say, I'm summing them up, and then I'm sending the result only to rank zero. So this is a partial version of an all reduce. And hopefully, this diagram makes it clear how reduce scatter works. And so all gather and reduce scatter are quite important because in some sense, they are the primitive by which many of the parallelization algorithms are going to be built. And so this is an important equivalence or an identity. I will refer to it one or two times as key points in this lecture. If you want to do an all reduce, let's say, I've got different GPUs, A, B, C, D, and each of the GPUs are handling a different data point. And so I've got different gradients for each of these data points, and I'm going to need to sum those gradients, and then I need to pass all of those gradients back to the GPUs. This is a classic data parallel operation that I might need to do across my four GPUs. So that would be an all reduce. One important thing, though, is this could be replaced with two operations, a reduce scatter and all gather, where a reduce-scatter is going to sum each of the rows and then leave the result of the rows in, let's say, GPU 0, 1, 2, 3 respectively. And then I'm going to do an all gather to copy those back out to the remaining GPUs. So each GPU now is getting a full sum of a part of the parameters, and then it's going to copy it back to the remaining workers. And in the bandwidth limited regime, this is basically the best that you can do. All reduce, the best that you can do is roughly matching the bandwidth that you can get out of a reduce scatter and all gather. And you can convince yourself of this by writing out how many communication operations happen in both all

GPU all-to-all vs TPU toroidal mesh

08:28 reduce and the right hand side. The final thing that I want to briefly touch on before I move on to talking about the parallelization algorithms, and this is the one place I'll talk about GPU versus TPU. Most of the discussion today can actually abstract out the underlying hardware, but there is actually one important thing that I'll mention up front, so that I can refer to it later as I talk through this. How do we network together different machines or different accelerators in GPUs? Well, as I showed you in the GPT-NeoX slide here, how in the GPU world this generally works is you've got nodes, single machines that contain, let's say, 8 GPUs, and then you've got these switches that connect fairly quickly to each other. And these machines are connected all to all up to about 256 GPUs. So that's an important threshold up until which you have very fast, arbitrary communication between machines. And then above that, you're actually going to need much more slow communication, these leaf switches and spine switches, once you go beyond roughly a single rack's worth of GPU. On the other hand, if you look at TPU design from Google, they actually take a very different approach to networking their machines. You've got a single TPU chip, and they all talk to their neighbors very, very quickly. And so this is a very easily expandable what they call toroidal mesh. But you can only talk to your neighbors. And the reason why I'm talking about this right after the all reduce slide is if you think about doing these kinds of collective communications, like allreduce or reduce scatter, you can implement them just as efficiently on a toroidal mesh

10:11 than you can on all to all connection. And so if you're optimizing purely for collective communications, it makes sense to think about things like TPU networking rather than GPU networking. I'll talk a little bit about pros and cons of this later as I go through different parallelization operations. So just to put this together, now we're going to start talking about a new unit of compute. Instead of the GPU, the new unit is the data center. The whole data center is going to be the thing that we're going to be doing. And now, we're going to try to come up with algorithms and sharding strategies that get us two different things. The first one is linear memory scaling. So as I scale up the number of GPUs, the biggest model that I can train is going to scale linearly with that. So I can train bigger and bigger models if I really want to. I also want linear compute scaling. As I get more and more GPUs, the useful computation that I'm doing to train the model scales linearly. And then finally, a lot of this, these algorithms are going to be implemented by just calling these very simple collective communications primitives in various ways. And so when we think about the performance characteristics of these parallel algorithms, it suffices to reason about basically counting the collective communications primitives. So that's kind of an important way to think about these. We don't go all the way down to the low level implementation of these algorithms here. And questions on part one? Yes. Sorry, but from the previous slide, does it mean that it's better to do reduce scatter as all gather rather than the algorithms or algorithm? Right. So this slide, right? Yeah. So the conclusion of this slide is that they're equivalent. And I think if you think about something like doing gradient descent in parallel, all reduce

11:55 is a very natural operation to do because you'll scatter your-- sorry, you'll distribute your data to different machines. And then you'll have to all reduce your gradients together. But what I'm saying is this very natural thing to do of all reduce can actually be written as a sum of two different operations and they're equivalent. So there's no performance hit by going from this left representation to this right one, at least in bandwidth. And that's going to have important implications in maybe five slides. So you can wait a little bit to see why I mentioned this. OK.

Three axes: data, model, activation parallelism

12:29 Any other questions? Good. OK. So now, we're going to get started. In some sense, this is the exciting algorithmic meat of the lecture. And there are three kinds of parallelism strategies, parallelism things that we should really be thinking about. So the first one is data parallelism. So data parallelism at a high level is the idea of I'm going to roughly copy the parameters across my different GPUs. I'm not going to worry about splitting my parameters up, but I will take my batch, and I will split my batch up. And different GPUs or different machines will get different slices of my batch. So that's data parallelism. There's lots of subtleties in how we execute that. Model parallelism now is starting to say, OK, I don't want all my GPUs to have all the different parts of my model. As my models get bigger, that's going to be a very big problem. So I need to cut up my model in very clever ways, and I need my GPU to handle different parts of my model. So that's going to be model parallelism. And then the final piece is activation parallelism. We don't really think too much about activations in our day to day lives lies because the PyTorch handles it very transparently. But as the models get bigger and the sequence lengths get longer, the activation memory starts to be a really big problem. So if you want to train these really big models with big, big batch sizes, you have to somehow manage the memory footprint of your activations. And so we have to split those up too. So there's some ways to handle that. And when we put all these together, we will have all the tools we need in order to scale up both compute and memory gracefully as we have lots and lots of machines.

Naive data parallel and the memory bill

14:09 So these are the core conceptual objects. And now, we're going to talk about implementing each of these ideas efficiently. So the starting point of data parallelism is just SGD. If we're doing very naive batch stochastic gradient descent, the formula for doing this looks like this equation that I have right here on the slide right here. I'm taking a batch size capital B, and I'm going to sum up all those gradients, and I'm going to update my parameters. So naive data parallelism is just saying, all right, take your batch size B, split that up, and send that to different machines. Each machine will compute some part of the sum. And then I will exchange all of my gradients together to synchronize after each-- before each gradient step, I will synchronize my gradients, and then I will take a parameter update. So now, I've been talking to you about compute and memory scaling and all these things. So let's just talk through what it looks like for each of these. So for compute scaling, data parallelism is pretty great. Each machine, each GPU is going to get B over M examples. And if my batch size is big enough, each GPU is going to get a pretty decent batch size, micro batch size. And it's able to hopefully saturate its compute. So that's good. What's the communication overhead? Well, I'm going to have to transmit twice the number of my parameters, every batch. Remember, on all reduce is going to roughly be twice the amount of stuff that you're all reducing in terms of communication cost. And so this is OK if the batch size is big. If my batch sizes are really big, I can mask the communication overhead of having to synchronize my gradients every now and then.

15:53 Memory scaling, I'm not touching this at all right. Every GPU needs to replicate the number of parameters. It needs to replicate the optimizer state. It's pretty bad for memory scaling. So if we didn't have to worry about memory at all, this is an OK strategy. But I think in practice memory is a problem. I think every one of you sitting here has experienced trying to put a big model onto a GPU and PyTorch telling you, oh, you're out of memory. And this is really a problem with your training as well because if you can fit more and more batch sizes, that's going to make the data parallel more efficient. And so ideally, you'd like to save on memory. So let's take a closer look at the memory usage of naive data parallel. And the memory situation is actually worse than it looks. It's actually quite terrible, because you've done this in assignment one, but we can think about how many copies of our model we need to store. And it's very large. Depending on the precision by which we're doing some of our training, you're going to need to store something like 16 bytes of data per parameter. And in fact, you need to store something like five copies of your weights. And this is really quite bad. Because if you just want to think about your model parameters, technically you only need two bytes. So where did that factor of 8 come from? Well, at least, you need gradients. And if you're computing your gradients in BF16, that's another 2 bytes. But then your optimizer state shows up, and that's a really big problem, because you've got 4 bytes of master weights, the things that you're accumulating into SGD, like these intermediate sums that you're doing. You need 4 or 2 bytes for Adam's first moment estimates,

17:39 because remember Adam keeps track of historical gradients. And then Adam also needs second moment estimates, the variance of the gradients that you've gotten in the past. And that's going to need another 4 or 2 bytes. And so what originally looked fine is actually now looking quite grim. And so the 16X, if I just draw it as a picture, you realize that most of your memory usage, at least in terms of parameter memory, is really being dominated by the optimizer states of your Adam optimizer. So your memory consumed is going to be a function of how many bytes are being used for your optimizer state. And that's generally going to be even more than the core parameter and gradient memory usage. And so for a simple example of a 7.5B model distributed over 64 accelerators, you're using a ton of memory. And this memory scales linearly upwards, total memory at least scales linearly upwards with the number of GPUs. So that's no good at all. But if once we look at this picture,

ZeRO stage 1: sharding the optimizer state

18:47 we get some very simple ideas. You might wonder clearly or maybe not clearly, I need the parameters and gradients to be copied across devices. That seems necessary to do data parallel. But do I really need all the optimizer states to be on every single machine? And once you ask that question, you can maybe get to the second row here. And this is going to be called optimizer state sharding. And if we could do that, then at least in this case, we can go from 120 gigabytes of total memory usage down to 31.4. And then maybe we can start sharding the gradients, and then now, we can get to 16.6 gigabytes of memory usage. And then if we also shard the parameters, we can go all the way down to 1.9 gigabytes of memory usage. And that would be a pretty good place to be, because now, we've fully sharded out all the optimizer state and parameter and gradient memory that we need. Yes. So why would we shard the optimizer state if we're doing, I guess, the gradient computation on each of them and then [INAUDIBLE] reducing that? Then how can we have a [INAUDIBLE]? That is a very good question. And the question is, how can we shard the optimizer state? When we're doing data parallel, GPU 0 has to be responsible for data point 1. So clearly, it needs to know about all the parameters and update it. So how can I possibly shard the optimizer state? And in a way, I think, zero, which is what this is, this is ZeRO, the zero-overhead data parallel optimizer, this is a very, in some ways, clever idea. Because it shows you that even when you're doing data parallel, you don't actually need to copy everything onto every machine. You can be really clever about how you do communications

20:29 to avoid all of this. So I will talk through exactly this. This is a great question. So what we're going to do is we're going to split up the optimizer states, as I said. So the first and second moments are now split up across all the GPUs. But everyone has the parameters and the gradients. So why is this important? If I have the parameters and gradients, let's say, I'm GPU zero, I have the parameters and gradients for everything. That's enough information for me to compute the full gradient. The full gradient update for this example can be computed. The only thing I can't do is I can't take that gradient and take an Adam step. I can't update my parameters unless I see all of the optimizer states. So that's the key idea. And so now, what's going to happen is GPU 0 is going to compute the gradients for everything, but GPU 0 is now only responsible for updating the parameters for the shard that they own. And that's the key idea. We're going to distribute the work of updating the parameters. And then we're going to synchronize the parameters back. So let me show you in much more gory detail how this works and the reason why it's called zero-overhead. So step one, every GPU gets a different data point let's say. I'm just going to simplify all this batch computation. I have GPU 0 through, let's say, 4. And every GPU gets a single example, and they compute a full gradient on the example that they own. Now, what I'm going to do next is I'm going to reduce scatter the gradients. So I'm going to send the gradients that-- I'm going to collect in some sense the gradients that each GPU owns. So GPU 0, let's say, is responsible for this first quarter of the parameters.

22:11 So the parameters are the Y-axis here and the X-axis here is GPUs. And so what we're going to do is we're going to reduce scatter to make sure that GPU 0 has all of the gradient information from all of the other GPUs for the subset of parameters that it is responsible for. So now, it gets this gradient information from GPU 1 and GPU 2 and GPU 3, and that's all reduced into GPU 0. Hopefully, that's clear. Now, GPU 0 has all the information it needs to update its own parameters, because it has the optimizer state corresponding to this first part. It has a full summed gradient for this first part. And now, so it's going to take a gradient update on their part of the parameters using gradient and state. And so that now I have the full updated parameters for this subset in my GPU 0, and all I need to do is all gather all of the updated parameters back in to all the ranks. So there's many questions here. I'll start here. Yes. Would you say [INAUDIBLE] the number of parameters that's permission, right? Or is that [INAUDIBLE]. Sorry. Say that again. and speak up. [INAUDIBLE] the number of parameters, that's per machine, right? Or is that total? So the question was whether the number of params communication cost was per machine or its total. Here, it's going to be total, because-- so this is going to be 1/4 of the parameter is going to be sent 3 times to this machine. And then you repeat that 4 times. In slide one, you said it was two times the number of parameters. Was that per machine or-- That was also total. Oh. OK. Yeah. 2 times number of parameters is total, because each block is going to have to be sent

23:50 to every other kind of machine. OK. Yes. So this question is not unique to what you're showing here, but you made me think of it. So the Adam optimizer that you showed seems to assume, largely assume independence of parameters. We've drawn all these diagrams that show the opposite. We have connected nodes and all that. And it seems, especially pressing when we're trying to split these and update them separately, does that create any issues? OK. So the question was AdamW seems to assume parameters operate independently. I'm assuming because you're saying, we track gradient sums and then we diagonally update the parameters. But we know that that's not fully diagonal, and so is there a problem. There have been better attempts at improving AdamW to not just be diagonal. There's things like K-FAC and all these other second order style optimizers that people have come up with. They haven't dethroned Adam, even though they do have their advantages. And there's some really interesting things that you can do with these kinds of improved second order preconditioning methods. Yes. What was the rows that we're reducing over? What is the rows that we're reducing over? So you're asking, what is the rows of this picture? Yeah. So imagine this is parameters here in the rows. So GPU 0 is responsible for some number of parameters. So this is a block of parameters up top. And so when we do reduce scatter, we're saying, take the gradients for example 0 for this block of parameters.

25:27 Take the gradients for example 1 for this same block of parameters. And then sum them all and put them in rank 0. That's what we're saying here. Cool. OK. And the key thing here is we're doing a reduce scatter and an all gather. And if you remember what I was saying before, well, a reduce scatter and an all gather has the same cost as an all reduce. And so there is a little bit of a surprising magic thing that happened here, which is that, well, we were doing an all reduce before on all the gradients to make sure everyone's gradients were synchronized. And that cost us 2 times the number of parameters. But if we're clever about how we're doing the updates, well, we can do a reduce scatter and an all gather, and in between the two steps, we can do some computation. And that gives us the same amount of compute communication cost but now, at least for the optimizer state, we fully sharded the optimizer state across the model. So ZeRO stage 1 is in some sense free in the bandwidth limited regime and gives you memory wins. Yes. [INAUDIBLE] can suppress the memory contribution of the higher notes, do people modify Adam to include higher moments? What do you mean by you can suppress the higher order contributions? Right, so for first and second moments, the amount of memory per GPU is divided by the total [INAUDIBLE]. Yes. That's right. So it seems like you might as well shard [INAUDIBLE]. I see.

27:04 So you're roughly saying, you could track way more optimizer state-- to rephrase what you're saying, you could have even more complicated optimizer state because you can divide that by the number of GPUs. While this is true, what we're going to do next is we're actually going to make the other components scale with GPUs. So that's going to make things in some sense not free anymore. Optimizer state will continue to be the bottleneck if we can divide everything by the number of GPUs. So hopefully that's a reasonable convincing answer. So we're going to build up stage by stage to ZeRO stage 3, which is more complicated.

ZeRO stages 2 and 3 (FSDP)

27:37 ZeRO stage 2 is still relatively simple. So now hopefully that optimizer state sharding trick made sense. I think that's very cool. So now we want to shard even more stuff. So I want to shard the gradients across the machines. So roughly, we can do the same kinds of trick as stage 1, but there is one additional complexity. And so what's the additional complexity? Well, we can never instantiate a full gradient vector. If I ever do the full backwards pass, and I try to compute a full gradient vector, I might go out of memory. So I want my maximum memory usage to basically be bounded by this, which is like full parameters, sharded gradient, sharded optimizer state. And so what we're going to have to do is when we do the backwards pass, as we're computing the gradient vector, we can't instantiate the full gradient first and then do communication. What we have to do is as we compute the gradients backwards, as soon as we compute a layer's worth of gradient, we're going to have to send that over to the corresponding GPU that it belongs to. So this is how it works. It's roughly the same idea. So now everyone has their own batch component. Everyone incrementally goes backwards on the computation graph. And let's say we're going to operate layer by layer. So layers are sharded atomically to different GPUs. So what we're going to do then is as we go backwards on the computation graph, after we compute a layer's gradients, immediately call a reduction operation to send this to the right worker. So a layer belongs to some worker. Maybe it's like GPU number 2 in this case. So we're just going to immediately reduce that, send that to the worker at that point.

29:16 And gradients are now no longer needed. I don't need to store the gradients on rank 0, 1, and 3. So I can immediately free that. And then now we continue this process. And so all the machines have their fully updated gradients. And now they have a full gradient for their share of the parameters, they have a full optimizer state for their share of the parameters, each machine can update their parameters, and it all gather the parameters back together. This looks like it's maybe more communication because you're doing this kind of reduction operation every layer. But this is only for a small amount of parameters. It's sharded. And so the full communication remains the same. So ZeRO stage 2 has some more overhead because we have to synchronize layer by layer and make sure that the gradients are properly sent to the right workers. But the overhead is pretty minimal. It's still very simple, fairly straightforward. Now the last one of these ZeRO stage 3 is more complicated for sure, but it allows you the greatest win of all, which is now essentially everything is divided by the number of GPUs that you have. And you can get the maximum savings possible. And if you've heard of FSDP, you've probably used that in some aspect of your life in the past. FSDP is exactly ZeRO stage 3. So now you'll kind of hopefully today know how FSDP works. So the same idea applies. We're going to shard everything, including the parameters. We're going to do the same thing as ZeRO stage 2, which is we're going to incrementally communicate and compute things so that we don't keep these big vectors of gradients lying around. And we're going to send and request parameters on demand while we're stepping through the compute graph, both

30:58 for the forward and backward passes. As we go through, we're going to send things around on demand. And of course, the key is to do this with as low overhead as possible. I think the thing that's really surprising about FSDP is not that this is possible, but that this is possible with relatively low overhead. You'll see why it's low overhead in the next slide. I admit that this is maybe not the most friendly graphic to start with, but this is, I promise, the baby version of FSDP. The next slide is a little bit more involved. But conceptually, this actually explains everything. So what we're doing is we're going to have model weights, and we're going to be all gathering the model weights as we go. So for each layer, no single GPU is going to have all the parameters. So I can't do the normal thing of saying oh GPU 0, go ahead and run the forward pass. That's not possible. So GPU 0, let's say, only owns the bottom most layer. So it does that computation and then it stops and it says-- it requests the parameters from all the other workers. So it stops and it does all gather, which is right here. You see there's the all gather step. It gathers all the parameters. Now it has the parameters that it needs to do a forward. So it can step forward and compute the layer that it didn't have before. And then now it can free the weights. It doesn't need the weights anymore. Get rid of it. Now it can all gather the next layer, I can do another forward, free the weights and I can repeat this right. The activations have to be stored. So the activation memory here is growing. So that's going to be an eventual problem. But if we ignore activations for the moment,

32:34 this is great because I load a layer, I do a forward, I free it. The memory overhead is very low here. Once I get to the end, now I can do the same thing with the backward pass. I can call backwards, and every time I move backwards through the neural network, I all gather for the parameters that I need. I can do a reduce scatter to update after the gradients that have been computed, and now I can free the weights. So I can free both the gradients that I don't need and the parameters. And at the very end, I've got a fully updated model. And so we've got three different operations that we've got to worry about here. We've got to all gather, we've got another all gather, and then we've got another reduce scatter, basically to update the model after we take the gradient update step. So conceptually, this is just a single step beyond ZeRO stage 2, but you do kind see that there is more overhead. So the total communication cost is now higher. We were kind of-- before, we had two times the number of parameters. Everything was kind of free in some sense. Now it's not right. There's a total of three times number of parameter communication cost and there's going to be cost associated with waiting for these communication things to finish. But I think the really cool thing about FSDP is it's actually surprisingly low overhead. You might imagine that because we're doing this crazy thing of asking for and sending parameters back and forth all the time, that things will be really slow. We have to be communicating all the time. But you can do this core idea of overlapping communication and computation.

34:10 So you want your GPU to be working, while also the communication is happening in the background, almost like prefetching, so that by the time you need some piece of information, it's already loaded up, it's already been communicated to you, and you're good to go. And so I'll talk through this example at the bottom here. But this is kind of the key to making FSDP actually somewhat efficient. So let's imagine we have a computation graph that looks something like this-- W1 W0 plus W2 W0 times x. Some input let's say is y. So some very simple computation graph like this. And then you might run FSDP. And you will get actually computation and communication that looks like this block diagram at the very end here. So the CPU it's nice that we did the Nsight Systems example last week because hopefully this diagram will now be clear. The CPU is going to basically dispatch a bunch of commands, asking the communication part of the GPU to basically go and fetch some parameters. It's going to dispatch things to the GPU to say, OK, all right, do some matrix multiplies and it's going to run far ahead in some sense of the GPU. We've seen this when we were looking at the profiler last week. Now let's look at the sequence of both communication and computation that happens on device now. Remember that I need to gather things on demand. So at the very beginning, I have to make sure that everyone has the weights for layer 0, or W0 here. So I do all gather 0, and I'm going to wait for that to complete. And once that's completed, I can do a forward step on W0. I can compute x times W0, let's say.

35:51 At this point, all gather 1 starts at the same time that all gather 0 ends. So as I'm doing this matrix multiply, I'm basically already starting to load the next parameters that I need. Of course, my communication is slower. And so there is some gap, but I end much quicker than the initial load. So now forward 1 can happen. And in the background, once again, I've started to load parameter number 2. And this yellow slice here, I'm now freeing the parameters associated with forward 1. And then now the other thing here is I'm repeating computation W0 is used twice. And so I don't need to communicate this. Again this happens very quickly. And I can do this very quickly. I have forward 2 now already loaded before I needed it. And so there's no bubble here. And then I can free number 2. That's the entirety of the forward pass. And you see that the gaps are relatively small here. And we were able to do a lot of loads before the compute needed to happen. And so by doing this very clever thing of queuing the requests for waits before you actually need them, you can avoid a lot of the overhead associated with communication. And then now at this point of forward 2, I'm done with the forward pass. I can free weight number 2. And I start on the backward pass. And you see that all gather 2 for the backward pass is already done. And so I can start on backward 2, backward 0. Weight 0 is already stored. So that's done. And then the high overhead here happens in the backward pass because I need to do reduce scatters and then all gathers and so on and so forth. Hopefully, you see this picture and you say, wow, it's kind of surprising that even though we're doing this crazy sharding-- if you go back to this picture, we fully sharded the parameters gradients and optimizer states--

37:32 but the total bandwidth that we need is only 3 times rather than 2 times. So that doesn't seem too bad. And the actual bubbles that we see are not horrendous. The communication is almost being fully being utilized and the computation isn't stalling for very long. So we're actually making pretty efficient use of the resources that we do have, which is cool. OK, yes. Where do I get [INAUDIBLE]. It's like to my understanding, the GPU and memory is full, so where does the weights get prefetched? Yeah, so you need a buffer in which you can store these weights. And so this picture is not quite right. You will have some overhead that you need associated with reading these weights for the current layer. And also, the other big elephant in the room is I haven't talked at all about activation. That's going to be like a big chunk, because you've got a big set of activations for a full model that are of living here in some sense. Yeah, cool. So this is kind of distributed data parallel. ZeRO is, in some ways, the way that people do distributed data parallel efficiently. And so there's different stages. And stage 1 is it's basically free. It's doing the same communication pattern as naive data parallel, but you get to shard your optimizer state. That's great. You might as well always do it. ZeRO stage 2 is twice the number of parameters. So the total bandwidth consumption is the same, but there is additional overhead in having to do this incremental freeing of the gradients as you go backwards. ZeRO stage 3 is more involved. You do three times the number of parameter communication cost,

39:09 but it's not so bad. We did have some overhead in the diagram that we saw before. But if you really cleverly mask your communication patterns, it's actually pretty good. And so people use data parallel even for fairly slow links in your networking pattern. And this is also conceptually very simple. One of the advantages here is especially data parallel doesn't care too much about the architecture. I didn't talk at all about how we actually implement a transformer in any of this. It's all very abstracted. And so this is one of the reasons why, for example, FSDP is so popular, it's very easy to write a wrapper that parallelizes arbitrary neural networks without having deep knowledge or deep introspection of what the architecture is actually doing. And so here's some examples. I worked out some examples because I'm always running out of memory on my GPUs. And you can see what's the maximum size of the model that I can fit on an eight times A100 80 gig node. And so for baseline, you might end up with oh, I can fit barely 6 billion parameter model. Whereas I think if I use ZeRO stage 3, I'm able to fit something like a 50 billion parameter model. There's big savings in my ability to fit larger and larger models by doing things like FSDP to cleverly save on memory. Oh, sorry, there's a question. Yes. I guess I'm a little unclear as to where the difference then, once you shard the parameters, what's the difference between that [INAUDIBLE]? Yeah, so model parallelism is really fundamentally about making sure that the parameters just

40:48 like live in separate-- let me see if I can find-- [INAUDIBLE] Yeah, yeah, yeah. So in some ways, it's true that we have sharded the parameters. So you could call this a kind of parallelism. But the whole point of model parallelism is to make sure that the parameters just live entirely in one machine. We're not going to try to ship them across in various ways. Only the activations are going to get shipped across. And so you'll see very different discussions in the model parallelism section. The focus there will be on communicating activations rather than communicating parameters. And that will be a big difference. Yes. If the parameters [INAUDIBLE], Why are you performing at all [INAUDIBLE]? So you're asking about this step, why are we doing all gather to gather weights onto all the machines when they're only on one machine. Is that right? Yeah, so we need to basically put-- we need to take the weights that live on one machine and scatter-- or is it gather or scatter? Sorry, I want to make sure I get this right. The terminology is a little bit sketchy for me. So I want to make sure I get-- sorry. Yeah. So what we want to do is the same as this. So each machine is going to have some parameter that I want to gather across all of the machines in order to make sure that each layer is properly replicated across all the GPUs. Is that the right question that you're asking? Or are you saying is there a simpler primitive that we could have invoked?

42:29 Are you saying broadcast is the right object rather than all gather? I think maybe it's written that way because of some exceptions about layers not living on individual GPUs, but I'm not 100% sure. I agree with you that broadcast should be able to do the same thing if the parameters live on only one machine. OK, cool. Alrighty. Let me make sure we're-- got it. So there is a key resource in data parallel. And this is actually an important idea

Batch size is a finite resource

43:04 that I want you to remember. With data parallel, batch size is actually a really critical resource in the sense that you can't parallelize greater than your number-- sorry-- than your batch size because you can have at most one example on each machine. You can't go to fractional examples per machine. And so this means that if there's limits to your batch size, you stop being able to use data parallel. And there's diminishing returns to batch sizes. So in your assignment one, you may have played with varying batch sizes. But you kind of know that as you crank up the batch size past a certain point, you start to see fairly rapid diminishing returns to your optimization rates. And there's lots of papers written on this. OpenAI has a really nice one on something called critical batch sizes, where they basically argue that past a certain point, you have very rapid diminishing returns in how much each example is contributing to your ability to optimize. Basically, the intuition is that below a certain point, you have a lot of gradient noise and reducing that is very valuable. But at a certain point, you're really fundamentally limited by the number of gradient steps you're taking rather than variance reduction. And so that basically means data parallel alone isn't going to get you to arbitrarily large parallelism. And the batch size thing is a really important resource. Essentially, you have a fixed maximum batch size, and you can spend it in different ways. And I'll talk about that later because other kinds of parallelism also benefit from having bigger batches. And so you use your batch size in certain parts. And issues are going to remain with data parallel.

44:43 ZeRO stages 1 and 2 don't let you scale memory. ZeRO stage 3 is nice in principle, but it can be slow. And maybe more importantly-- and this relates to the earlier question-- it does not reduce activation memory. Ideally, you want to cut up my model entirely and make them live totally separately, because then the activation memory would also be reduced. And so now I want better ways to split up the model so I can fit these really big models in these GPUs. And so that's going to bring us to model parallelism.

Pipeline parallel: cutting by depth

45:18 We want to scale up in memory without changing the batch size. And we want an alternative axis where we don't need to spend or basically have big batch sizes in order to parallelize. And so what we're going to do is it's going to split up the parameters across GPUs. And in some ways, that's like ZeRO stage 3. But we're not going to communicate parameters anymore. We're going to pass activations around. And that's going to be different. And sometimes activations are going to be much smaller than parameters. And that'll be very good for us. So we'll cover two different types of parallelism. I'm going to talk about pipeline parallel, which is conceptually simpler, but much more horrible implementation wise, and tensor parallel, which is conceptually maybe less obvious, but honestly much nicer to implement and more commonly used. And they're going to correspond to two different ways of cutting up the model. So I think pipeline parallel is maybe the most obvious way to cut up a neural network. You know that a deep neural network comes in layers. So if I have layers, a very natural place to cut a network is to cut it up at the layer boundaries. And so each GPU is going to handle some subset of the layers, and I'm going to pass activations around. In this case, each layer belongs to a GPU. And GPUs are going to pass activations from one to the other. And in the backwards case, it's going to pass the backwards gradients backwards from GPU 3 to 0. So that's cool. That's great. What's wrong with this picture? Well, I think you should see that most of your GPUs are idle most of the time. This is actually quite terrible utilization. And so if I do this naive kind of parallelism

46:58 that I described before-- so if I have, each layer having a forward and let's say I have a single example, that's going to result in a diagram that looks like this. So different rows in this picture are different layers and also different GPUs. And the x-axis here is time where I'm going from left to right. So what do you see? Well, I first compute my first layer at the very left here. And then the activations get past the second layer. GPU 2 wakes up and it's like, all right, it's my turn. It does its job, passes to GPU 3 and then GPU 4. And now the backwards passes can begin, and so on and so forth. And you see this gigantic What people call bubble. This is a big overhead where you're doing absolutely nothing. And you see that the GPUs are active 1 over n of the time. So in some sense, this is the worst possible parallelism of I've added 4 GPUs, but I get the throughput of a single GPU. And so one thing you can do is you can be a little bit more clever about what you do, and you say, all right, I'm going to have a pipeline. I'm not just going to cut things up in layers. I'm going to have a sequence of things that need to be processed by each GPU. So now let's say I have a microbatch. So each machine is going to handle four examples. And what I'm going to do is I can finish my first example, my first data point, and I can send off the activations for that to my second GPU as soon as I finish. And then I can then get started working on my second data point. And so now I've overlapped communication and computation. The second GPU can start working while the first GPU continues to work. And now the size of the bubble can potentially

48:38 be reduced by having bigger batch sizes. And you can hopefully see why I said before that batch sizes are a resource. If you have a finite batch size, and you have pipeline parallel, you can use that same batch size to make your pipeline bubble size smaller, for example, or you could use it to do data parallel. So there's many different ways that you can take your single batch size and then split it up in two different ways. So now your micro batch size can control the bubble time. And in fact, the amount of the ratio of your overhead to the useful compute that you have is the number of stages minus 1 over the number of micro batches. So if you have big big batch sizes, pipeline parallel could potentially be efficient. But as we said before, batch sizes are finite. We can't just crank that up to whatever value that we want. So in general, pipelines seem really horrible. Why do we do it? Why do we incur this cost of a bubble in order to parallelize? Well, there's a couple reasons. Pipelines help save memory compared to data parallel. I mean, ZeRO stage 3 will also shard the parameters, but this also shards the activations, which is nice. Pipelines can also have good communication properties. It only depends on activations. It's also point to point. So it's possible that depending on your topology and depending on what you have, pipelines might actually be very favorable for the slower parts of your network. And so pipeline parallel is often going to be used on your slower network links. So internode or even sometimes across different racks or across different data centers, you might do-- actually, not data

50:18 centers-- across different racks, you might do pipeline parallel. One of the examples of a thing that I was recently told by some Google folks is, they were saying, actually, one of the big advantages of TPUs is that we don't have to do pipeline parallel very much because all of our connections are much bigger. They have this big toroidal mesh. They don't have this limit at 256 GPUs, where they're suddenly going towards a slower network link, where you might want to switch to pipeline parallel. So that's a real world kind of example of when you would start to think about pipeline parallel. And so this is an example from an NVIDIA paper. I'll talk about this paper in much greater detail later. They've done some really nice work showing performance characteristics of different kinds of parallelism. But you see with batch size 8, as you increase the pipeline parallel size, the number of devices, your utilization per GPU starts to really drop off. Whereas if you have a big, big batch size of 128, you can get away with pretty good utilization for a reasonably sized pipeline parallel. So batch sizes are really key to hiding the size of the bubble. Otherwise, you have issues. Of course, you can do different kinds of pipeline strategies. So instead of having these the standard patterns for scheduling the bubble. You can cut things up into finer pieces where you're assigning different stages, assigning different sublayers to different devices, and you're doing different computations at different parts. You can then interleave the pipeline better. And an advanced version of this that I want to spend a moment talking about-- and this is very, very clever, is zero bubble pipelining-- or I think in DeepSeek lingo, I think they call it DualPipe.

52:05 But the core single trick is the same. So here, if you think about it, let's say we're doing the backwards pass to compute gradients. You can split this up into two different components. The first part is about back propagating the activations. So this is as I go down the residual connections, I need to compute essentially the derivative with respect to the activations. And then as I get to a parameter, I also want to compute the gradient itself. How am I going to update the parameters, not just how do the activation change with respect to the previous layers. And so to give you a concrete example, let's look at this bottom left diagram over here. So in this diagram, you see the forward pass. This is a single MLP. So we've got multiply by a weight. I do a non-linearity. And then I'm just going to output the non-linearity. So this is a kind of a naive single part of the MLP. Now let's look at the backwards. I have the derivative with respect to the loss. It comes in. And then I can compute how that's going to change the x's, the inputs to my MLP. So this is, in some sense, the derivative with respect to the activations here. And then as I compute these, of course, I can use them to compute the gradients that I need to update my weights. But the important thing is this part, this part of computing the gradients for the weights. This can be done whenever. There's no dependence of this. And so I can rearrange the scheduling for this computation to any part of the computation graph. And so what you can do is you can do your standard pipeline parallel for the parts that are serially dependent. But any time you have to do these computations just for updating the parameters, you can reschedule them wherever.

53:49 And so the key idea is you start with a nice what's called 1F1B pipeline. This is a nice optimized reducing the bubble size schedule. And then you can take this. And what you can do is you can separate this B, which is this computation of the backwards part, and then W, which is the computation necessary to compute the gradients of the weights. And now I can do the computation of the weights, the W's, where I would have originally had a bubble. So the parts where I originally had these white idle utilization components, I can now fill them in with these W's. And so by thinking carefully about what the serial dependencies actually are, I can now have something really nice where I'm getting actually good utilization out of my GPUs. To be clear, this is horrendously complicated. If you actually want to implement pipeline parallel in this way, you're going to have to intervene in how your autodiff is actually calculating these things. You have to have a queue that can track where things go. I heard a funny anecdote in a conversation recently from someone in a frontier lab training LLMs, and they said, actually, there's two people in the group that understand how the pipeline parallel in our infra works. One person left. And so there's a single load bearing person in our training infra. There are stories like this. Pipeline parallel is infrastructurally very, very complicated. It looks simple here. If you're interested, I encourage you to try and implement it. It does get pretty hairy pretty fast. And I think that's a good note on which to switch to the other model parallelism

Tensor parallel: cutting by width

55:31 because this is much simpler. And this is often very cleanly utilized by a lot of frameworks and a lot of-- even people training really big models rely very, very heavily or primarily on this model parallelism. So what other way can we split up a model? So if we think about it, most of what we do is matrix multiplies. In a big model, most of the computation is matrix multiplies. Most of the parameters are matrix multiplies or matrices. And so what can we do? Well, if we can parallelize just the matmuls, that would be pretty good. And so tensor parallel is this idea that we can take a big matrix, multiply and split it up into a set of submatrices that can be multiplied. So if I have this matrix multiply at the top-- we have X and X times A equals Y. What I can do instead is I can cut up into half, and then I can also cut up X into half. And I can compute the submatrices, I can sum them up, and then I will get my answer at the end. So conceptually, pipeline parallel is cutting along the depth dimension like the layers. Tensor parallel, which is what this is, is cutting up along the width dimension of your matrix multiplies. And so we're going to decompose into submatrices and then do partial sums. So here's an example of what it might look like in an MLP. We have each GPU handling a different submatrix of let's say a big MLP matrix multiply. And then we're going to have collective communications to synchronize the activations as we need them. So what are we going to do? So this is an MLP.

57:10 And the top half and the bottom half-- there's two different paths. These are splitting up the matrices. So I want to do this operation Y equals GeLU X times A. I'm going to split up my matrix A into A1 and A2. And then on the right-hand side, I want to compute dropout YB right, and then I want to return the result as Z. So I'm going to also cut up B. So I've cut up both of my big parameter matrices into two parts, A and B. And in the forward pass, what I'm going to do is I'm going to take my inputs X, and I'm just going to copy them twice. So each GPU is going to get the same inputs. And they're going to operate on it with A1 and A2. They have the same kind of-- oh, sorry. They're the same row dimensions. So it's going to be fine operating on them. So XA1 and XA2 is going to give you some activations Y1 and Y2. Those are going to go into B1 and B2. And then I'm going to do an all reduce to sum them up. That's exactly the figure I showed you before. So you copy, and then you all reduce, and you get the answer Z. In the backwards pass, now it's actually the reverse as the gradients come backwards in the backward steps. This G is going to be the identity. So I'm going to copy the derivatives on both sides. And I'm going to do the backwards operation all the way through. And once I get to F, this is an all reduce, because I've got two derivatives coming in from both paths. And then I sum them back up. So this F and G are synchronization barriers. In the forward pass, I do a single allreduce. On the backwards pass, I do a single allreduce, just at two different places in the computation graph. So now you can hopefully see how this is a very nice way of wherever you have a matrix multiply,

58:49 you can just cut up the matrix multiply and parallelize them across different devices. And as you might imagine, this is actually somewhat expensive. We have a synchronization barrier that lives per layer. It needs to communicate an activation the residual activation worth of stuff twice in a forward backward pass. And so tensor parallel, this very simple idea, is going to require very high speed interconnects. And so there's a rule of thumb. It's a very simple rule of thumb to remember which is that the tensor parallel, is applied within a single node. So a single box of let's say NVIDIA GPUs is going to ship with eight different GPUs that live in that same box. And as I showed you at the beginning of lecture today, they're very, very high speed connected. Those eight GPUs can talk to each other very quickly. And so it makes sense to use something like tensor parallel that's very bandwidth hungry between those eight devices. So what you will typically see is that tensor parallel is applied up to eight GPUs where the eight GPUs live in the same machine, because that gives you the least drop in performance. And so this is an example from Hugging Face's parallelization tutorial showing you the throughput decreases of different levels of tensor parallelism. You see that there are hits, 10% and 12% hits to throughput as you do tensor parallelism. But up until 8, well, maybe this is manageable. This is kind of the price you pay for just being able to parallelize more nicely. But then you go to 16 devices, and you get this kind of astounding 42% drop in performance. You go to 32, and you see another 65% drop in throughput.

60:32 And so you see, hopefully visually here, that you really want to stop at 8 for tensor parallelism. That's really the sweet spot, because of the kinds of hardware interconnects you can get your hands on. So how do things now compare to pipelined parallel? Well, compared to pipeline parallel, we don't really have to deal with this bubble thing that we had before. We don't need to consume larger batch sizes in order to reduce the bubble, which is nice. And there's relatively-- I wouldn't say very-- there's relatively low complexity in applying tensor parallel. All you really need to know about are where are the big matrix multiplies, can I split them up and make them live on different devices. The forwards and backwards operations still remain the same. Compared to implementing something like zero overhead or dual pipeline parallel, you're going to be in much, much better shape doing this. So the con is that it's much larger communication overhead. You've got, in pipeline parallel, batch size times sequence length times residual dimension, point to point communications per micro-batch. In tensor parallel, you've got eight times that per layer, and you've got all reduced communication. It's potentially a very large amount of communication that needs to be done. So the rule of thumb, as I said before, is tensor parallel is used whenever you have low latency, high bandwidth interconnects. You're going to see 2 to 16, depending on what kinds of machines you have of tensor parallel out in the wild. And I'll show you examples as I talk through at the very end here of examples of tensor parallel. Any questions on pipeline or tensor parallel before we move on to the kind of third kind, sequence

62:13 parallel and activation sharding? Yes. [INAUDIBLE] Yeah, so the question was, can they be used simultaneously. The answer is that, yeah, you do use them both. So I think we'll get to examples later. But I think the typical thing that you see is for large scale runs, you very often see tensor parallel. Pipeline parallel is often used on top of that. I think the only example I know of that does pipeline, but not tensor parallel would be DeepSeek, V3, as far as I know. So within a single machine, I guess you have-- say if you have five different machines, maybe the first 20% of the parameters are across the reach of the first machine, which you use tensor parallel for. And then that pipeline parallels into the second machine where you have the next step. Yeah, so the question was there-- do you do tensor parallel within machine and pipeline parallel across machine for example. Yeah, so you would do something like tensor parallel within machine and a combination of data and pipeline parallel across machines for example. And I'll show you the rule of thumb later. But basically, you do pipeline parallel

Activation memory and sequence parallel

63:21 because your models won't fit. If you could fit your entire model, you just do data parallel plus tensor parallel or just maybe even data parallel. Great, OK. Excellent. So then we've been talking about memory. And memory is, in some sense, a very important part of parallelization because we're going to be training big models. And so when you look at your memory, you realize that, actually, activations are a really big part of your memory usage. So if you look at standard kind of forward backward pass, I think this was from one of the PyTorch tutorials, you see that memory usage is very dynamic. So I'll just talk through this, because I think it's an interesting plot in general. You always have your parameters as you're training because that's static. But in iteration zero, you don't still have optimizer state at all. So actually, you don't have that part of your memory use. But as you do your forward and backward, you see activation grows, grows, grows, grows as you accumulate all the activations. And as you start your backwards pass, your activation goes down because you're freeing it as you use up your activations, and then you're accumulating your gradients. So your gradient memory usage goes up. And the peak is actually somewhere partially through your backwards pass where you haven't freed all your activations yet, and you're still building up your gradients. So in iteration two, you see the same thing here. So the point of this diagram is to say, well, we've thought about all the other pieces. We thought about the parameters, we've thought about optimizer state, we've thought about the gradients. But we have not thought about, very deeply at least, the activations. And so let's do that. So the final complexity that I want to talk you through

65:01 is the activation memory. So tensor and pipeline parallel can linearly reduce basically most things. But it can't actually reduce all of the activation memory usage. And so this is an example from one of the NVIDIA papers that's talking about how do you reduce activation memory. And I think one thing that's really interesting to see is that you make your models bigger and bigger. So going from left to right. You see that parameter and optimizer state memory can remain the same if we parallelize aggressively. But activation memory just kind of continues to grow because some parts of it don't parallelize very cleanly. So no matter the number of devices you have, actually you can't really get rid of the growth of activation memory per device. And I'll show you why in a moment here. Whereas I think if you do some slightly more clever things like recomputation, you can keep the activation memory low. And that's really key to parallelizing some of the biggest models. So what's the activation memory per layer? You've kind of done some of this transformer math and calculus before. So hopefully, you're now familiar with all of this. But we can compute what's the amount of activation memory we need per layer. And there's a handy formula here. And this is the amount of memory you need. It's sbh times 34 plus 5 as over h. And some of these numbers are mystifying, but actually they're not so mystifying. You can very much see that there's a left term and then there's a right term. The left term comes from the MLP and other pointwise operations. That's where sbh times 34 comes from. These depend on the size of your residual stream, the h. On the right side, you have a term

66:41 that's actually, if you multiply this out, as squared b because the h is cancel. That's the memory that you need for the softmax term and other quadratic terms in your attention. Of course, if you use FlashAttention and use recomputation, we know that we can drastically reduce that second term. So then let's say we do tensor parallel. We do tensor parallel everywhere we can. So we do it in the MLPs, we do it in the KQ computations, in the attention computation. We will end up with something that looks like this. And this is looking pretty good, but not quite there. So activation memory per layer divided by t, which is the number of devices that were tensor paralleling over. So if we're dividing by 8, ideally, we would divide all the activation memory by 8. But you see there's this straggler term, sbh times 10, that has not been reduced down. And if you think about what these are, these are the non matmul components. So the LayerNorm, the Dropouts, the inputs to the attention, and the MLP-- all of these terms will unfortunately continue to grow with size and they will not be parallelized very nicely. And so the very last thing that we need to think about is to take those simple pointwise operations, which thus far, we have not parallelized, and we just need to split them up. And there's a very simple way to split them up, which is to say, well, if we're doing a layer norm, these layer norms across different positions in the sequence do not interact at all with each other. They just don't care about anything else.

68:23 And so what we're going to do is let's say we have a 1024-long sequence, we're going to cut that up, and then each device will handle a different part of that layer norm, or different part of that dropout. Those pointwise operations can now be completely split up across the sequence dimension. And because now we're cutting things up across the sequence dimension, we're going to have to do some synchronization to make sure the parallel computations that we did can get aggregated back again. And so in the forward pass these guys, they're going to be all gathers and g bars are going to be reduced scatters. And in the backwards pass, the two are reversed. In some sense, they're a duality here between the two. And what we're doing here is for the layer norm, we've kind of scattered things around. And so we're going to have to gather them back together so that we can do of our standard computation. And then now whenever we get to the dropout, we want to scatter them back out into the parallel component that we have. And in the backwards pass, we're kind of doing that in the reverse. So hopefully, that is clear. This is a very simple idea. We're just parallelizing the very last components that we failed to parallelize before. And so now, we can put all of these different pieces together and get to the end, which is we started up here, which is no parallelism at all. We did tensor parallel, which allows us to divide everything that's not a pointwise op by t. And then if we apply this sequence parallelism idea, we can divide this component by t once more. And then we can do things like activation recomputation which is the FlashAttention trick, to remove the second term. And the minimal memory that you can easily get away with is going to be this thing on the bottom,

70:02 which is sbh 34 over t. And this is often used if you're looking at different formulas for transformer arithmetic on how much activation memory do I use. You often see something like sbh 34 and then if you have t tensor parallel divided by t, because this is the easy minimum that you can get for that kind of a memory. Any questions on sequence parallel and activations? Yes. There's one more thing, as the transformers stack up each other, I suppose the computational graph will grow more and more like model [INAUDIBLE] computational graph as a [INAUDIBLE] ever become a bottleneck in the communication between the GPUs? You're saying if we have something that's a more complicated computation graph than a single linear chain, will that become a problem? It's a good question. I haven't thought about that. I would guess not. At least for tensor parallel, this operates purely layer wise. It doesn't really care about the dependencies. Maybe for pipeline parallel, there's opportunities for increased parallelization if there's more than one branch.

Ring attention, expert parallel, and the summary table

71:06 But I'm not too sure. [INAUDIBLE] Yes, right. OK, cool. All right. So there's a few other parallelism strategies that I'm not going to talk about just because in the interest of time and fatiguing you, because I think I've already dragged you through a whole bunch of low level details about how to do parallelization. So the first one I want to talk about is context parallel or ring attention. You may have heard the term ring attention before. This is a way of essentially splitting up both the computation and the activation cost of computing really large attention, where essentially, you're just going to pass keys and values around different machines. So each machine is responsible for a different query, and then keys and values are going to travel from machine to machine in a ring like fashion in order to compute your KQV inner products. And the cool thing here is you already know how to do this because you've done the tiling for FlashAttention. So you know that attention can be computed in this kind of online tile by tile way. And that's what's happening in ring attention. The other thing, which now that you know tensor parallel is pretty straightforward, is expert parallelism. Expert parallelism, you can think of as almost like tensor parallel in the sense that you're splitting up one big MLP into smaller expert MLPs, let's say, and then scattering them across different machines. The key difference with expert parallelism

72:45 is that the experts are sparsely activated. And so you have to think a little bit about routing. And the routing is not going to be as predictable, let's say, as the all to all communication that we had before in tensor parallel, because now maybe one expert is overloaded. Your networking is going to be a little bit more complicated. But otherwise, conceptually you're living in the same world as tensor parallel for expert parallelism. So just to recap all the things we talked about, I've made a little small table of the different kinds of strategies that we have. We have DDP and ZeRO stage 1. This is kind of the naive data parallelism thing that you do. Here, you have some overhead per batch, you have no memory scaling, reasonable bandwidth properties, but you consume batch size in order to be able to do this. You need big batch sizes to have big data parallelism. You have FSDP, which is kind of a nicer version of ZeRO stage 1 in the sense that you can get memory scaling, but you're going to pay overhead across different layers. And so now you've got higher communication costs, and you've got potentially synchronization barriers that lead to poor utilization. Pipeline parallel is nice in that we no longer have this dependence on this per batch aspects, and we can get linear memory scaling, but we have another issue, which is this also consumes batch size and it's horrendous to set up and use. And so a lot of people like to avoid pipeline parallelism if it's possible. And then finally, tensor parallelism is very high cost in terms of bandwidth and the amount of synchronization you need to do.

74:26 But this has this really nice property that has no impact on batch sizes. So it's kind one parallelism strategy you can use that has no cost in terms of your global batch size, which is nice. So we have to balance a number of limited resources. We have memory, which is one resource. We have bandwidth and compute, which is another resource. And then we have batch size, which is kind of an unconventional resource, but one that you should really think of as a limited thing that you can spend on different aspects of these to improve your efficiency. And there's a very nice TPU parallelism, or TPU book, let's call it, from Google that I referred to last week. But also actually, they have a really nice parallelism section. And they have this great figure that I wanted to show you before I moved on to some of the examples. So the key quantity, as I was saying before, is the batch size. And depending on the ratio of batch size to the number of GPUs you have, different kinds of parallelism become optimal. And so they use certain formula on how much communication and computation you end up doing for each of these models. So this is a simplified formula to generate this plot. And you can see if your batch size is too small, you have lots of GPUs and really tiny batch sizes, then there is no way for you to be efficient. You're always communication bound, which is this bottom half here. And in fact, you're spending most of your time on communication. As you get more and more batch size, eventually you can get to a point where if you mix both FSDP-- so ZeRO stage 3-- and MP, which in this case is tensor parallel, you can actually get basically to a place where your compute bound.

76:07 So now you're not wasting your flops waiting for communication. And then finally, if you get to a point where your batch sizes are big, then you can just get away with pure data parallel. Pure FSDP is going to get you into a regime where the time you spend doing computation is higher than the time you spend doing communication. So if your batch size is big enough, you can just get away with FSDP. So this is a cool illustration of this idea of why would you mix these, when would you mix these, why is batch size of resource.

3D parallelism, and what real models do

76:41 Hopefully, this kind of shows you in a very visual way what this is. And so when you put these all together, you end up with what people call 3D or 4D parallelism. I think I've heard the term 5D parallelism recently. I wasn't quite sure what the fifth dimension was yet. I'll have to read up on that. But now you can put it all together the different dimensions of parallelism. And this is a really simple rule of thumb. I originally looked it up and put this together last year, but turns out it's still the same this year. So you can follow this now. So the first thing you have to do is you have to fit your model and your activations in memory. If you don't do that, you just cannot train. So this is a requirement. So until your model fits in memory, we have to split up our model. So we're going to do tensor parallelism. And we know that up to the number of GPUs per machine that's very efficient, that's very fast. So we're going to do tensor parallel up to that point. Now after that, depending on things like your desire to deal with pipeline parallel and/or your bandwidth constraints, you're either going to use ZeRO stage 3 or pipeline parallel across the machines until you can fit your model in memory. Now, after that point, well, until you run out of GPUs, you can now run the whole thing, and your only goal is to increase the amount of total flops that you have on hand. So you're going to scale the rest of the way with data parallel, because data parallel is-- it works well on low bandwidth communication channels and it is very simple. And so that's going to give you a way of using all of your GPUs. Now if your batch size is really small, then there is a way of trading batch sizes for better communication efficiency.

78:22 If you haven't consumed all of your batch size as a resource, what you can do is you can use gradient accumulation on your devices. And that will let you basically have effectively larger batch sizes even if you're memory constrained. And that will let you trade your batch size for better communication efficiency since you're synchronizing less often across machines. Simple rule of thumb. This will let you train models with reasonable efficiency no matter what you're doing. And so to make this concrete, I'll talk through a few examples at the very end here. I'll flash through both this really lovely paper back in 2021 from Megatron-LM basically showing you exactly these things in pictures, and also a lot of ablations as well as some of the models from last year. So this is a big table of how they trained models going from 1.7 billion parameters to 1 trillion parameters. And they get great utilization on all of these. You see percentage of theoretical peak flops that they get and it ranges from 40 to 52%. It's pretty good. And so you can see tensor parallel starts at 1, and then they eventually go up to 8, and then it caps out at 8. And so they're using tensor parallelism first. And then pipeline parallel stays at 1. But once the models get big enough, they can't fit these big models. So pipeline parallel has to increase in order to compensate. And then the data parallel size basically starts out as big as possible and then slowly kind of goes down because as we increase the amount of pipeline parallel, this is now consuming, in some sense, the batch sizes. And so you can't have effectively

80:02 as big of a batch size if they're being used, in some sense, for pipeline parallel. So careful 3D parallelism is going to give you linear gains in aggregate flops. So you see, if you do careful 3D parallelism, you can see very flat overall achieved flops per GPU, which is giving you-- if you add more GPUs linear scaling in the total aggregate throughput, that's great. Tensor parallel 8 is often optimal. You see this is the pipeline parallel size and the tensor parallel size. You see going to 8, 8 with a batch size of 30-- or sorry, batch size of 128 is optimal. Even if you have a smaller batch size, tensor parallel size of 8 remains optimal. And activation recomputation enables larger batch sizes. And remember that larger batches can in turn help you mask overhead for pipeline parallel. So activation recomputation, even though it's more flops, can pay for itself. We've seen that story play out already in FlashAttention. All right. So the last part of this is recent language models, what do they do. So I've gone through a few papers to look at examples of what people's parallelization strategy is. OLMo and the Dolma paper, they do FSDP for 7 billion parameter model. DeepSeek, the first paper does ZeRO stage 1 with tensor sequence and pipeline parallel. This is the vanilla thing that I told you. V3 actually does something slightly different. They do 16 way pipeline parallel, 64 way expert parallel, which is kind of tensor parallel,

81:44 and then ZeRO stage 1 for their data parallelism strategy. Yi, which is another Chinese model, does, once again, ZeRO stage 1 tensor and pipeline parallel. And Yi-Lightning, because they're doing MoEs, replaces tensor parallelism with expert parallelism. The final thing, if you're interested in state of the art distributed training with lots of details, Llama 3's report is actually really interesting to read. They have a lot of detail about how they do their networking, what things happen. And you see, once again, the kinds of things I said before. You see a tensor parallel of 8, you see CP or this is context parallel. This is only relevant for long context training, which is this very last step. So you can ignore that. And you've got pipeline parallel and data parallel happening in these first two phases. You can also even ignore the first stage here because that's the small batch size training that they did in order to be stable. And if you look at of their rationale for how they do their parallelism strategy. You see exactly what I had said before of basically, all right, you want to do TP, CP, pipeline parallel, and DP in that order in terms of the amount of bandwidth that you need, where data parallel can tolerate these long network latencies, because you can do this asynchronous fetching of sharded model weights. And so they're using the strategy that I told you in order to train some of the biggest models. The funny side note about Llama 3-- and you may have heard this in, not rumors, but casual conversation with your friends, is there's lots of GPU failures when you train models at a huge scale. They had 148 interruptions from faulty GPUs, totaling about 30%

83:28 of the total interruptions that they had. They had things like unplanned maintenance of machines. And that was 32 different things, 32 instances of interruptions for their training. And so when you're training a model this big-- I've talked about the algorithms, but you also need kind of fault tolerant architectures to be able to deal with these kinds of things. And I've also heard various stories of people saying the even scarier thing is not actually explicit model failures, but actually data corruption. GPUs can silently fail on you and give you garbage data, completely ruining your run. And then the last one example is for Gemma 2. And I wanted to end on this because this is a TPU example. They do ZeRO stage 3, which is roughly FSDP. And then they do model parallelism and data parallelism. And so here, as I said before, the TPUs allows them to stretch model parallelism a little bit further. So putting it all together, scaling beyond a certain point is going to require multi-GPU, multi-node parallelism. There's no single solution. So you want to combine all three approaches to leverage strength. And then there's simple and interpretable rules of thumb for how you might execute this parallelism in practice. All right. Thank you.