Transcription
Uh, today it's our pleasure to host Nuaman Tazzy from HuggingFace. Uh, he's the lead author of the Ultrascale Playbook and a core developer of Nanatron. Um, HuggingFace is an open-source distributed training library. Um, his work spans projects like StarCoder 2, SmallLM3, and Mixture of Experts scaling with several initiatives, and he is passionate about making large-scale training practical and accessible. So, without further ado, let me hand it off to him.
Thank you so much for the invitation. So, hello everyone. Uh, my name is Dante. Uh, I've been working with Hugging Face now for four years, and the talk for today is scaling training to thousands of GPUs. Uh, so today's talk is going to be mainly uh inspired by uh the book uh that Hugging Face made. It's called the Ultrascale Playbook. So make sure to check it out. It's available online, and there's a printed version if you need it. Uh, but I also added some elements since the book now dates from one year ago, uh, especially regarding scaling and uh some new things that we can uh take a look at.
Uh, so without further ado, let's start with the first question: Why does scaling matter? And if we take a look at the latest releases of the uh latest LLMs, uh, we can see that the trend is still to train larger and larger. And why is that? Uh, that's also because we can see that there's a correlation uh with intelligence. It seems that the more uh the big the bigger the LLMs are, the smarter they are. And there's just an example from this week, Kim K 2.6, which is also uh a one trillion uh parameter model.
Um, so to take just some numbers, uh, models of these days uh are one trillion parameters large or even larger. Uh, they get trained on 15 trillion uh training tokens, and for context lengths as far as 1 million.
Um, so how can we do that? Uh, and to take a quick uh overview about uh what does this entail from an uh infra uh point of view? So if you want to train this kind of uh model, first you need to load the data, which is like about 15 trillion, imagine that's stored uh on a storage somewhere, and you need to load that continuously in your training loop. And each training iteration, even if the model is as big as a one trillion parameter model, needs to take about 1 second. Uh, and you're limited by the uh memory within each GPU or TPU or whatever you're using. And you need to save checkpoints to that storage every now and then. So you're putting a lot of pressure uh into these uh uh GPUs.
So the talk of today is going to be mainly focused on uh training uh the training iteration or the the training. So let's start quickly uh on one GPU. So there is a model, we do forward, backward, we compute gradients, and then we do the optimizer step, and we have an updated model. But for our case, we usually have a global batch size of 1 million to 10 million, uh, even like 50 million tokens. So that, of course, doesn't fit uh on a single training iteration. So what we do usually is we can do gradient accumulation. So we're going to divide this batch size to smaller batch sizes, and we're going to do multiple forward and backwards, and we're going to accumulate gradients sequentially. Uh, the problem with this approach is we're going to assume that the model already fits in memory, which is not always the case. Uh, models now are so big, one trillion, and we said that we're still limited by the available VRAM in each accelerator you're going to use. So how can we solve this with more GPUs? The other thing is that if you train sequentially uh with this approach, you're going to need to wait for a long time. Uh, so if we have more GPUs, how can we speed this up?
And for this, there are multiple forms of parallelisms we can use, and each one of them is used in a specific use case. And hopefully, today you can get at least an intuition which parallelism you can use for which uh use case. And um, to help you, uh, so basically for today, I'm going to mainly focus on the first two parallelisms because they're going to explain just the intuition of how we can uh think about these parallelisms. And the final three, we're going uh just quickly skim over them because uh there's not enough time, and I think that the approach uh can be uh generalized.
So before starting, uh, quick few notes. Uh, so this approach can be applied to an accelerator, GPU, CPUs. So in the slides, I use GPUs, but the same can be applied for anything. And it's useful for pre-training, post-training, distillation, training, inference. So any type of uh workload uh can benefit from these parallelisms, and it's applicable. You don't need 1,000 GPUs uh to benefit from this. As long as you have two, at least, uh, you can benefit from these parallelisms.
So naively, what we want to do? So we said that the model doesn't fit. So ideally, we want to be able to shard the model, so either vertically or horizontally. And that means we're going to also shard the gradients and the optimizer, and shard the data because we don't want to wait too long for the to train sequentially on data.
So let's deep dive into scaling. And first of all, we're going to start with data parallelism. And I'm assuming this is the basic form that everyone uh is using. So I'm going to go quickly over it. So instead of training sequentially over multiple batches, I'm going to feed each batch to a GPU. So I'm going to have different batches per GPU. Thus, the name data parallelism. I'm going to shard the data. Uh, and of course, if I have different data per GPU, each GPU is going to compute different gradients. And to keep the same model duplicated on each GPU, I'm going to need to synchronize these gradients. And how am I going to do that? It's going to be the same as accumulating gradients, but it's going to be in a distributed fashion. Thus, we're going to need to use uh the all-reduce collective operation, uh, which is just a distributed uh sum over multiple uh GPUs.
So we end up with the same gradients, and thus we can do the same optimizer step over all GPUs. And if you want a quick look at the code, it's going to look like this. We can use the DDP class by PyTorch, and when you wrap your model with it, when you do the forward and backward, it automatically handles the all-reduce part uh in the backward.
So there is a small problem with this. Is uh when we do the all-reduce, do the GPUs stay idle in that time? And for that, there is a very nice tool uh called the profiler that helps you visualize uh what the timelines basically or the workflow for the CPU and GPU. And for example, if we apply it on a naive implementation of DDP, we can see that there's the CPU stream, there's forward and backward, and then there is the GPU streams, GPU computation, and communication.
Uh, so there's two things. So the the biggest thing uh that we want is we don't want GPUs to stay idle. So we definitely want the GPU computation stream to be always full, which is not the case uh in here, as we can see here. Um, the other thing is we need to make sure that the CPU is always ahead of the GPU so that we can always schedule kernels beforehand. And the third thing is, for example, in this case of DDP, we can see that the all-reduce is not overlapped with the computation. So ideally, we want to overlap this communication with computation.
So this is a simplified diagram of uh the previous. So how can we overlap this all-reduce? We can just start uh all reducing the gradient once a bucket uh of the gradients is computed. Basically, when you do the forward, you start doing the backward from the from the last layers. So you don't need to do the backward of the entire model to be able to all-reduce them. Right? So once, for example, you compute the gradients of the last layer, you can already start the all-reduce. And this is what we do in practice. So for example, if layer two is the last layer, once the the gradients of the last layer are computed, we can all-reduce it, and we can continue doing the the backward. And this bucket, we can control the size of this bucket in PyTorch via this uh argument, bucket_cap_mb. So you can control how big or how small the buckets should be so that you have a better overlap of communication and computation. And this is what uh goes through the mind of any engineer who tries to scale models. Basically, when you have a lot of GPUs, you definitely want to overlap efficiently uh your computation with communication so that you avoid any idle time here for the GPU computation and uh to benefit of the the tensor cores of the GPUs. So this is the end goal uh of this presentation.
So quickly, we can go over a recap of DP. So it's easy to implement. Uh, it's very efficiently overlapped. It's really good. Uh, and it's model-agnostic, right? So we can just wrap any model, and it's going to automatically uh reduce uh the gradients in the backward. But the problem is that the optimizer step is duplicated across GPUs, right? Because if you have the same model on each of your GPUs, all of them are going to do the same uh uh optimizer step, which is bad, right? It's duplicated work. The second thing is that uh if you scale more and more uh in the number of GPUs, you're also going to need more and more batches of data. And this is what we call global batch size. So your global batch size is also going to scale. But of course, you don't, you can't just scale indefinitely, right? We're limited by tens of millions of tokens. Uh, and so you can't just apply DP if you have, I don't know, 1 million GPUs, you can just use the DP over all of your GPUs. And the third thing, which is the biggest bottleneck, is that uh DP assumes that the training step fits in memory, and this is not always the case, unfortunately. And so we need other approaches to optimize uh memory.
And we can look at the first one, which is ZeRO-1. So from the name, ZeRO-1 uh tries to solve the the optimizer step. Basically, we said that in DP, we're going to duplicate the model, and all of them are going to do the same optimizer step. So how can we reduce the memory at least of the optimizer states so that they're not duplicated and each GPU can do uh a different part of the optimizer of the optimization? So the simplest approach is uh instead of uh having uh the optimizer states of all the models, I'm going to shard them across my GPUs. And when I do the backward step, I'm only going to reduce-scatter the gradients. Uh, so I'm not going to have all-reduce here. I'm only going to have reduce-scatter. And then I'm going to do the optimizer step for each shard. And this is what's really good about ZeRO-1. We're not going to have duplicated work. And uh, once I have my updated parameter shards, I need to all-gather them before the next step. So this is ZeRO-1 uh in brief. And there is a very important intuition uh that you can learn from this. So before, in vanilla DP, we had an all-reduce in the backward, right? In here, we had an all-reduce. And the the trick that we made for ZeRO-1 to work is that we replaced the all-reduce with a reduce-scatter and then all-gather. And this works great, and we're going to reuse this trick uh in the future. Why? Because all-reduce essentially is just a reduce-scatter plus an all-gather. Uh, so basically, we didn't add any communication. Like, it's the same communication, just that instead of doing the all-reduce, I did the reduce-scatter first so that to compute the the gradients first, and then I can do the optimizer step, and then I all-gather uh the the parameters.
The second thing is, uh, for the gradients, I didn't use all of them. Right? So when I did the reduce-scatter, I only uh used the the shards uh that I optimized later. So I can just throw the the parts that I don't need from the gradients. And this is uh ZeRO-2. So ZeRO-2 uh optimizes or saves memory even further by just discarding the gradients that I don't use. And this is the difference between ZeRO-1 and ZeRO-2. It's a small difference. It's just that I don't need to keep the gradients. So it's easier said than done, right? Because in the implementation, it's annoying to just discard uh things. So people usually just use ZeRO-1 and they keep the gradients. Um, and also, uh, an important note, since now there is uh, the Adam optimizer, optimizer requires the full tensors, grads. So previously, people would take their entire parameters and they would just shard them. They would flatten everything. So we don't keep the notion of a tensor, right? You just flatten all your parameters and then you would shard them along the number of GPUs you have. But the problem is for Adam, you can't just split your parameters, right? You want your tensors uh to be uh full so that you can do the Newton's uh optimization on them. So the thing is here, when you want to split your parameters or your gradients on your GPUs, you need to make sure that to keep your tensors full. So for example, I'm going to give the first, I don't know, five full tensors to the first GPU, and the five sec, the second five tensors to the second GPU. Yeah. So this is for Adam, but also for other optimizers.
And lastly, for ZeRO-3, I also want to save memory of the parameters. And this is the craziest part, because how can I save memory uh by by sharding the the model's memory when I need them in the forward and the backward? And there's a very nice trick that uh the ZeRO-3 engineers came up with. Is when I do the optimizer step, I don't need to all-gather right away, right? I can just keep the the shards of the model uh on each GPU, and whenever I need to do the forward of that layer, for example, for the first layer, I'm just going to uh ask for these other shards from the other GPUs. So for example, for the first layer, I'm going to ask for the other shards, I'm going to do the forward, and then I'm going to free the other parameters. And then for the other layer, I'm going to do the same thing. So this looks very slow, right? Because on each layer, you need to ask for all the, like, for the rest of your parameters from your other GPUs, do the forward, and then free them. But in practice, it works very well. And the reason is, by the, it's the same for the backward. So for the backward, we also just do the all-gather of the weights and then compute the backward and we free the weights. So the reason is, we prefetch uh the the next layer while computing the current layer. So this is the trick. Uh, so I had a previously an all-gather here right after the optimizer step, but now, uh, I just separated the all-gather on multiple layers. And for example, for the first layer, I'm going to all-gather it here, I'm going to do the compute, and while I'm computing the for the forward of the first layer, I'm already all-gathering the the the next layer, etc., and then I free it, etc. And so, concretely, I never have more than two, what we call FSDP modules or FSDP units, in my memory because I always free uh one when I all-gather uh the other one. And the same thing for the backward. Uh, and so I still do the distributed optimizer and I repeat. So this is basically the best form of DP that we can have. Why? Because I efficiently uh optimize my memory. I saved memory on optimizer states, on gradients, and on the model. And I also uh prefetch. So everything is very well overlapped. And besides that, similar to uh DP's bucket size, I can also control the the size of the buckets, which uh we call the size of the FSDP modules that we have. And so for the code, so there is two, by the way, ZeRO-3 is also called FSDP. FSDP is Fully Sharded Data Parallel. And there are two versions in PyTorch. So FSDP-1, we used to just have a class FSDP where we can wrap our model and we can define the define the wrap policy, which is similar to the bucket size. And for FSDP-2, there's now a function called `fully_shard`, and you just call it over your model, and it's going to automatically transform uh your modules to FSDP modules. The difference between the two versions: the first one, as I said, it flattens all of your parameters, and we already talked about the problems of that. For example, Adam and some optimizers, they need full tensors, and this is bad. So ideally, when you shard your parameters, you want to keep full tensors on each GPU. So FSDP-2 enables that. The second thing is FSDP-2 uses Densor, which is uh a new utility in uh in PyTorch, and it helps it uh being combined with other forms of parallelism, as we're going to see.
And so to recap, please compare ZeRO-2 with ZeRO-3. So in ZeRO-2, we had a big all-gather here to uh all-gather the parameters that we optimized. And in ZeRO-3, I just uh divided this big all-gather into smaller all-gathers of the different FSDP modules. It allows me to overlap this all-gather with the forwards, and I can control the size of the FSDP modules, and it's also model-agnostic because I don't care about the architecture of my model. Uh, it's just great. It works very good.
So the pros and cons of DP ZeRO-3. It trades memory for compute, right? So DP vanilla, I didn't uh optimize memory at all. I just assumed that the model fits in memory. But if I needed to save memory, I can just do ZeRO-1 or ZeRO-2 or ZeRO-3. And this is something to pay attention to. So a lot of people just throw FSDP-2 on all of their models, and this is not a good thing because if your model fits uh on GPUs with ZeRO-1, you don't need ZeRO-3, right? Because ZeRO-3 just, you're going to add more communications and you're going to make your training slower. So just use the ZeRO degree that you need to save memory because there is no other benefits. Like, if ZeRO-1 fits, it's going to be faster than ZeRO-3. So the second thing is the great overlap of compute with compute, as we've seen, everything is well overlapped. And the third thing is model-agnostic. So there is no headache of implementation. You can just, whatever the the PyTorch module is, you can just wrap it with FSDP. You don't care if it's a state model, you don't care if it's anything, it's going to work out of the box. Uh, the negatives though is that, um, at scale, since you need to all-gather and the parameters for each uh FSDP module, if you scale it too much, even if it's overlapped, you can't hide the communication anymore, right? Because it's not perfect. If you have very slow network, you cannot overlap that anymore. So there is a solution for that, which is called hybrid sharding or HSDP. So you would use FSDP on a slower, on a smaller circle, and on the larger circle, you can use vanilla DP. So this is something you can look up later. And, uh, the other thing is ZeRO-3 is, we're still in data parallel. So we're still scaling uh with batches. So every GPU has a different batch, and we still need, we still have this uh restriction. We can't scale indefinitely.
So what if, um, what if I reached uh my global batch size? So I have, I don't know, 10,000 GPUs, and at 1,000 GPUs, I already reached my global batch size. I need another form of parallelism. So let's take a look at TP, and the motivation behind tensor parallelism is that if I have uh a model, a transformer model, for example, I want to keep the same batch, so the same inputs, the same data, and shard the model so that I do the same computation as if I had one GPU. So basically, instead of doing the computation with one GPU, I would do it with two, and essentially, I'm going to uh shard the memory for my model, for my gradients, optimizer states, by the number of GPUs I have. So the question is, and this is a very, uh, fun question that people like to ask in interviews, etc. So do we keep, do we keep full or half activations in tensor parallel? And we're going to answer that after a bit.
So let's start with an easy example, which is uh matrix multiplication. So if if we take X multiplied by Y by W, here a simple approach, and we want to parallelize that on two GPUs, we can just split uh our matrix uh on columns, and then I can all-gather uh the outputs. Let's take the example of two matrix multiplications now. So I have Y1, which is X * W1, and Y2, which is Y2 multiplied by W2. So if I split my first weight by columns, and if I split the second one by rows, I can do the same computation, which is two multiplications, and then I do an all-reduce here, and I would get the same result and the same computation as if I did this with one GPU. So this is kind of the power of tensor parallelism, and this is the intuition behind it. Uh, if I had one GPU, I would do two matrix multiplications with the full matrices. So I managed to uh distribute this operation on two GPUs by splitting uh the weights. So each GPU only has half the weights, and we each GPU only does half the compute. But the caveat of this is that I need to do a a distributed communication on all-reduce to keep the correctness. How does this work for the backward though? It's the same thing. Uh, and again, I'm going to assume that I have the same upstream grad, and this is a big headache, we can talk about it in a bit. So if we assume that I have the same upstream grad, which is uh this dY2, and dY2, similarly, it's just going to be multiplied by the transpose of W2 and W1. This is of course to to compute uh dX, which is the the out, the the output gradient. Um, so again, it's, I'm going to do two multiplications, and again, it's just half the weights, and again, I just need to do an all-reduce. So to efficiently distribute two matrix multiplications on multiple GPUs, I'm going to need one all-reduce in the forward and one all-reduce in the backward, under the assumption that I have the same upstream gradients, which is the the outputs, and I had the same inputs.
So let's apply this on our MLP. So in MLP, we have, for example, Y activation of X * A, and then I multiply it by another matrix. So I always going to have two matrix multiplications. So I apply the same trick. The first one is going to be a column linear, the second one is a row linear, and I need an all-reduce to keep the math correct. And as we said, we assume, and this is a big assumption in our transformer, we assume that I'm going to always have the same activations, the same inputs, the same data, and I need to have the same output, and later the same upstream gradients. So all the operations after Z must be the same uh in the two GPUs. Um, and another thing to basically, what we said, we're going to uh divide the compute by half. So A1 and A2 are halves of the original matrix A, right? So the hidden dimension is sharded uh in here, and this is what we just said, we needed the same input and we needed the same upstream gradient to uh keep the math correct.
Let's take an example of attention. So for attention, similarly, we have two matrix multiplications, which is the QKV projections and the output projections. So the question would be, uh, should we shard along the hidden dimension or number of heads or both? And what does it uh mean? So the question you really want to ask, and this is a very generic question, if you have a model and you don't know how to implement tensor parallelism on it, or you want to check that the implementation is correct, you should really ask: would GPUs compute the same activations as without TP? Let's say I sharded along the head dimension, right? Then when I do the softmax, which is Q * K, it's going to have a reduced hidden dimension, and that would impact the softmax calculation. So it wouldn't give the same results as if I had TP=1. But if I shard along the number of heads, since the heads are independent, I would get uh the same activations as if I didn't use uh TP. So this is really the the whole thing. So we shard along the number of heads, since the uh attention heads, the attention heads are independent, to keep attention correct. And again, uh, we need the same inputs and we need the same uh upstream gradient.
So as a recap, the pros is that we shard the model and compute across GPUs, and we're memory and sample efficient. So at the same time, uh, I sharded memory, so model parameters, gradients, and uh, and optimizer states, and I didn't need to have uh multiple batches, right? So I used the same sample across my GPUs. The second thing, I didn't need uh an all-reduce for the gradients, because all GPUs, I'm assuming they're doing the same work, like I have one big GPU uh that I distributed uh on smaller GPUs. And the cons is that communication is heavy, because at each attention, at each MLP block, I need to do an all-reduce for the forward and for the backward. The second thing is, and this is something we can uh see later, the parameters that are outside the TP region, which is the the region where the hidden dimension is sharded, should stay synchronized. Otherwise, I won't have uh the same uh inputs and the same upstream gradients, right? Because this is a big assumption for the matrix multiplications to stay correct. So in practice, uh, the parameters that are outside the TP region, they need to stay the same. They need to stay duplicated. So in practice, we all-reduce uh their grads, and this is something that screws up training when you use TP, uh, and just due to numerical precision, because theoretically, you don't need to do this, but due to numerical uh imprecisions, they drift apart. So in practice, we all-reduce their grads. And the third thing is, of course, complex implementation, and it's not easy if you change the uh architecture. Let's say, uh, you want to use MLA or some other fancy attention, it's a headache uh to apply tensor parallelism on it.
Let's take a look at sequence parallelism. So what's the motivation behind sequence parallelism, which is uh an a flavor uh to TP? So in TP, we've seen that we have this TP region, which is the attention, and this TP region, which is the MLP, and we do the all-reduce here. So the second layers, which are the row layers, they have an all-reduce. Okay, cool. So in the TP region, the shape is uh sequence, batch, and we shard the hidden dimension, we said. And after the second matrix in each TP region, I'm going to get back the original shapes. So I'm going to get back the original hidden dimension. So the question is, these parts are duplicated across GPUs. Uh, and it's annoying, as we said, to keep them uh synchronized. As I we said, in practice, we need to uh already use this layer norm, and if I had some other operation, I also need to make sure it's uh replicated. So can we just distribute them as well? And this is what sequence parallelism tries to to solve. And we're going to use the trick that we talked about earlier. So we have an all-reduce, and we know that all-reduce is reduce-scatter and all-gather. So how can we apply it uh in this case? This is what sequence parallelism is. So this all-reduce, I'm going to, I'm going to replace it with a reduce-scatter, and I'm going to, this identity, can call this an identity. This identity, I'm going to transform it to all-gather. So in the TP region, uh, I'm going to shard my activations uh along the hidden dimension, and then I reduce-scatter along the sequence dimension. So in this uh sequence region, SP region, uh, I have my sequence which is sharded by TP, and then I all-gather it again to have uh H, which is uh sharded. So basically, I play between uh sharding along the hidden dimension and along the sequence dimension. So this is moving from TP to TP with SP. And what this helps me uh do is multiple things. So first of all, the activations that are stored. So before, I had to store very big activations here, which are SBH, but thanks to sequence parallelism, my uh size is always divided by TP. So this uh helps me avoid uh big activations. The second thing is, it's the same amount of communications, since all-reduce, as we've seen, it's just reduce-scatter and all-gather. So I didn't have to add any communication. Uh, and the last question is about layer norm. So since we uh did this operation, should we sync layer norm grads in this case? And the answer is no. And for this, let's take a look at the backward. Um, so similarly to what we did before with the matrix multiplications, you can uh verify later that the backward of a reduce-scatter is all-gather, and the backward of an all-gather is reduce-scatter. Basically, reduce-scatter and all-gather, they are inverse of each other. Um, so if this is the forward, the backward becomes like this. So the backward, I'm going to do all-gather here, and then reduce-scatter here, etc. So we notice here that before layer norms, we have reduce-scatter. And thanks to these reduce-scatters that I now added, they take care of synchronizing the gradients. So it's really beautiful that it works out uh like this. So thanks to the fact that uh now the GPUs they don't do the same work, they see different patches, and thanks to the fact that in the backward, uh, I have a reduce-scatter here now, instead of just assuming uh the perfect case that it's an identity and that the layer norm would always stay synchronized, uh, I just have this all-gather and the reduce-scatter, and they take care of uh synchronizing uh my gradients. So you can uh verify this, the inverse of uh reduce-scatter and all-gather later.
And the pros and cons. So let's take a recap. Uh, so the the pros is that now we're sharding the model and compute across GPUs, and we'll be in memory and sample efficient, especially sample efficient, which is what we didn't have in DP. And again, we don't need grad, grad all-reduce, because GPUs are essentially doing the same work. And thanks to SP, I don't need to all-reduce layer norms as well, because uh, they now have the reduce-scatter in the backward. And but we still have the same problem. We're communication-heavy, and the implementation is complex. And for this communication-heavy, this is why it's recommended, if you've taken a look at the different distributed training libraries, they recommend that you use TP within a single node. So if you have a node of eight GPUs, they say that uh, it's better to use TP lower than eight, so that the communications for TP stay within a node, so that when you do forwards and backwards, you don't get stuck waiting for uh the all-reduces or the reduce-scatter and all-gather uh in forwards and backwards.
Now we've seen two uh axes of parallelisms, and before generalizing to the others, let's see what it means to uh combine them. So we've seen the data parallel, which shards data. Uh, here we have a data parallel size three, and we've seen uh tensor parallelism, which shards the model uh in two uh in two dimensions. So what does it mean to combine these two? Well, the cool thing about the the forms of parallelisms that we've seen is that they are orthogonal to each other. So if batch is uh sharded along the data axis, it's replicated, so it's replicated here along the model axis, and vice versa. If the model is uh sharded along the TP axis, it's replicated along the data axis. So this helps us uh like make a mesh of all the parallelisms. Now, now we've only seen two, but later we can see five, or there's even six or seven, depends on uh how many parallelisms you you want. And in the code, if you look at TorchTitan or Megatron or Natron, you're going to see that at first, they initialize the process group, so they initialize all the world sizes. Let's say you have six GPUs here, so you're going to create your TP groups and your DP groups, and later you can just pick on which axis you want to do the communication operations. So for example, I want to all-reduce along the TP group. So I'm going to all-reduce like this, and all-reduce like this, uh, along all the other axes. So this is uh the beauty of uh these parallelisms. We make sure that they are orthogonal to each other.
And now we've uh gone through the first two uh big parallelisms. The the last three, I'm going to go fast over them. So for PP, it's much easier than TP. So we have our layers like this, and instead of sharding them uh horizontally, I'm going to shard them vertically. So I'm going to put some layers on the first GPU, some layers on the second GPU. But we can see here that the problem is, for GPU one, at the beginning, it doesn't have any activations. So it's just going to stay keep waiting for the first GPU. And for the first GPU, after uh it goes the, after the forward goes through the first four layers, it just waits uh for the other GPU. So how can we solve this? Um, and for pipeline parallelism, there is what we call PP schedules, pipeline parallelism schedulers. So the schedulers, they define how uh your GPUs are going to do, are going to load the data, and they're going to communicate activations and gradients with the other ranks. So let's start with a simple assumption. Each GPU has one layer of your model, and we're going to start with the simplest scheduler, the all-forward, all-backward scheduler. Um, so we're going to have the first batch. Is this one? So it starts at the first GPU, and then it sends the activations to the second GPU, and then third, and then fourth. And then I have another batch. So that's the GPUs don't wait. I'm going to schedule eight batches. So all of them, they're going to do the forwards, and then uh, they're going to do the backwards. So every time one GPU does the backward, it sends the gradients to the previous layer, etc. And I repeat this. So this is all-forward, all-backward. And I schedule multiple batches to uh minimize the idle time. So the idle time is this, and it's the biggest problem in pipeline parallelism. Um, so there is the best way to solve it uh is to have more complex uh schedulers. So this is the one-forward, one-backward scheduler. And the name comes from the fact that if you prioritize backwards over forwards, you're going to have at some point, like in the middle, you're going to have uh forward, backward. So they're interleaved, basically forward, backward, forward, backward. But we still didn't solve the pipeline bubble. This just helps with the memory a little bit. So the best way to solve it uh would be to start doing forwards from this side as well. And this is what DeepSeek has done in their DualPipe. So basically, if you have eight devices, uh, you're going to distribute the layers in a round-robin fashion. So the first layer, you're going to put it both on the first device and last device, so that both of them, they can start with data, etc. So there's a forward that's going like this, and there's a forward that's going like this, and then a backward like this, and then a backward like this, and you manage the the overlap in the health. So as we said, it's easier said than done, because the implementation is very complex, and you need to keep track of the batches and the forwards and backwards so that you send the correct activations and the uh and the correct gradients.
So in general, uh, the pros and cons. So the pros is that the communications are cheap. There's no all-reduce, nothing to to have a headache about. You just exchange the activations and the gradients, and uh, the sharding is very efficient because every GPU keeps track of just a shard of the model. It's model-agnostic. I don't care about uh the nature of the model, as long as I just have layers and I just send activations and gradients. Um, and the uh, the cons is that I need to save activations for multiple micro-batches, which is annoying. So for example, if in in this case, I had eight uh backwards, in here they schedule 20 micro-batches. So I need to save 20 times the activations for the different. So usually for PP and FSDP, since they rely on activations a lot, we use, we offload the activations. So we use activation checkpointing, or now there is even CPU offloading. If you have a good GPU, you can offload the activations to the CPU, and then you just load them back uh when you need them. The other annoying thing is if you want to hide the PP bubble, you need an advanced pipeline scheduler like DualPipe, and the implementation is complex, it's very hard, and it's tough to add uh to the library. And 13 is you need multiple micro-batches to hide the PP bubble. So in a sense, this is like the DP issue, is that you need to scale uh your global batch size with uh the PP size.
And the last two parallelisms, they are specific to two use cases. For example, for CP, the issue we want to solve is sequence parallelism. So once you scale your sequence length, the activations, they, activations they explode. So the question, even with TP, even with PP, the question is, how can we amortize the memory cost uh when doing long context training? So what we've seen previously in data parallelism, when we have a batch, we can shard it along the batch dimension. But the other, what we can do now is, let's shard it along the sequence uh dimension. So it's analogous to data parallelism. But if I do that, and if I give uh different sequences to the GPUs, what about attention, right? Because when I do the forward, each GPU has only a part of the sequence. So how can I compute attention over my entire uh sequence? This is where ring attention comes from. And the idea is similar to online softmax that is used in FlashAttention. So you don't need to compute softmax out of the box, you can uh compute it uh in an online fashion. So each GPU, they compute softmax locally, and then they exchange K and V with the other GPUs, then they update the softmax, etc. So this is the idea behind ring attention. So the problem with this, small problem, is that we need to communicate K and V inside the attention block. So inside each uh attention block, which is a similar problem to TP. So it's also communication-heavy. And for that, we can use uh send-receive operation uh for GPUs to to send the the other K's and V's. You can also use all-gather uh for that.
So the pros is, it's the only parallelism that efficiently, efficiently, efficiently partitions large sequences memory. Sorry. Uh, and because the other parallelisms, they don't really uh shard the sequence. So if you have a big sequence, there's a problem. The cons is that it's communication-heavy because in each attention block, you need to exchange K's and V's. And similar to DP, since each GPU sees different data, we need to all-reduce the gradients. And it doesn't help much for short sequences. So really, it's really used if and only if uh you're doing long context training.
And for the last one, EP, it's used for, we don't have time to go uh really through them, but how can we parallelize them? So when you have your router and you want to uh route your tokens to different experts, you can just shard your experts across GPUs, right? So the question would be, what communications uh would I add so that I can do the the routing? For routing, there is uh Switch Transformer and Top-K version. So let's talk about Top-K, which is the more generic one. For this, let's take a look at all-to-all operation. So all-to-all, it's the most generic uh way of exchanging messages. So basically, you have, for example, four processes, four GPUs. Each one of them, they have four different messages. So all-to-all, what it does is just uh each uh process, it's going to send uh their messages to the other four, and each process, it's going to get uh messages from the other four. So in a sense, all processes they communicate with each other. It's like an N-squared type of communication. It's the most complex one. And how does this apply uh in MoE? So let's say uh each GPU has different experts. So when you're doing the routing, each GPU is going to have some tokens that need to be routed to the first GPU. So this is why all-to-all helps. So the first GPU, the first GPU here, which has the first experts, needs to get the tokens associated to it from the other uh GPUs.
Um, so what's, how can we apply expert parallelism? First, we shard the experts across the GPUs, and only the experts. So expert parallelism is really only applied in MoE, it doesn't touch uh attention. And so to avoid duplicated work in attention, we need to shard data across GPUs, right? So EP is actually, it's not orthogonal to DP in this sense. So you also shard data across uh expert parallelism. And we need an all-to-all communication to route tokens to their respective experts, which is the operation that we call dispatch. And of course, we also need to combine them uh after after the experts to retrieve the original uh sequences.
So the problem, and the biggest problem, uh, that you're going to face when scaling. So how do we know what all-to-all dispatch to uh to use? So basically, in order to know how the tokens are routed, you need, for the router to do the uh to compute the scores, right? So the CPU needs to wait uh for the GPU to do the router calculation. What we call the dispatch pre-process. And the GPU needs to tell us what are the buffer sizes and what are the tokens per expert uh that we need. And to solve this, there is, for example, DeepEP by DeepSeek, and there is HybridEP by Nvidia. And the only way to solve this is to use recent hardware, so far, which has IBGDA and RDMA. So these allow us to, so they allocate specific SMs and specific uh memory uh to basically, when whenever they calculate or they know which tokens to route to each GPU, they have already allocated the the memory for that uh calculation. And this is a problem because, for example, if you have H100s or if you don't have uh InfiniBand network, you're stuck with this CPU-GPU sync. And this is why most labs who do trainings, they have very slow trainings just because of this problem of hardware. And even this like DeepEP that DeepSeek has open-sourced, you're going to find out that it only works on IBG GDA, which
is uh infin. So the other labs have a hard time uh catching up uh with this. So in summary, E expert parism is the only way to distribute efficiently and uh it's communication heavy. It needs an all to all uh every block and since it has different data it needs to all reduce gradients and there is a CPU GPU sync which slows down your training very much to launch this batch operation. So in summary it is only to be used for MOE trainings and or so the practical uh solution that people do to train with expert parism is to combine it with PP and how do they do that?
So we said that the dispatch and combine are slow and if you remember for our one forward one backward uh we can overlap the forward and the backward. So what we do here is so we have two batches which is the blue and the green here. So when we want to do the this batch which is for the blue uh for the blue batch in parallel we can do the backward for the other batch and this is nice right because so MLP here doesn't need uh this batch and we can do the dispatch here and once this is done we can do the dispatch uh for the other batch etc. And this is already implemented in some libraries. For example, in Megatron, uh you can use these two flags.
So in a nutshell, you can combine all five parisms and you're going to have different communications in different uh steps of your training. And we've seen that uh CP and EP. So EP only apply to it only touches the CP only touches self attention. self attention here and yeah in all paralisms are orthogonal so you can combine uh all five of them and uh yeah as a recap we've seen the five parallelisms and you can find a cheat sheet that we made in the ultras scale playbook on how to make decisions following the number of GPUs uh you have and finally if you're more interested In the infra side of things, uh you can take a look at the small training playbook in the infra section. We've made a lot of benchmarks on how infra can handle uh these types of workload. And you can find more references in the two books mentioned also the Jack scaling book for a CPU uh point of view and if you're interested in these uh kind of things contributions are very welcome. So, Nanotron, Titan, Megatron, and them and a lot of other libraries. Uh, so feel free to uh reach out if you want to contribute. And one last thing is as we scale to thousands of GPUs, let's be mindful of their energy impacts and use them uh responsibly. Thank you.
>> We'll also uh take any questions in person or over Zoom. Uh any in the room?
>> Yeah. Um thanks for the talk. Um it seems that for expert parallelism, you also have to have the tokens routed evenly across ranks because if all of your tokens end up routed to one GPU, then everything else is idle. Um what are ways to fix this in practice?
>> Yeah. Uh that's what the load balancing tries to fix, right? So we have so depending on which approach for example DC uh they initially introduce the load balancing loss. Oh sorry. Um so the load balancing adds basically in the loss uh a factor that penalizes uh an even distribution of tokens. uh and there's multiple ways to do it. There's the auxiliary loss free where you just add a bias term when you compute the router and it would automatically uh adjust uh the distribution so that it's even uh across GPUs also have some on slido.
Um one is do different parallelism strategies meaningfully affect scaling efficiency, model performance or convergence or do they just enable reaching larger scales?
>> The first one scaling convergence
>> uh efficiency
>> convergence efficiency
>> I think like scaling loss. Um okay so actually the difference parism do not change uh the forwards and backwards. So ideally uh if everything is implemented correctly it should have the same effect as if you had one GPU. So like the same scaling laws applies uh and like they you shouldn't find any difference in using PP over EP uh from like uh uh like computation uh point of view.
Another is how can we reduce GPU idle time during say RL training when CPU side checkpointing or environment data processing becomes sequential and a bottleneck.
>> Yeah. Uh so for that uh at least for text uh you can allocate a lot of workers uh before. So the PyTorch data loaders they allow you to specify the number of workers that you can use and they can pre-process the data before the training uh workloads. So that should at least uh help alleviate uh the pre-processing so that it's done before you even go through the training iterations.
>> Any more in the room? Also someone asked do we have a automatic way to decide the best parallelism?
>> Yeah. Uh there's been some papers for that. I think if you take a look at Jack's uh book uh they've played with that uh and like like they the tip you uh decide uh on what's the best parallelism but it really depends on your uh specific so how many data you have the global batch size and uh the network you have. So the biggest uh factor uh in what we've said when you scale of course is the network bottleneck. So if you have for example let's say envelink so in GPUs uh usually in link is either 48 or for the most recent one 32 or even 72. So it depends on on that uh you can even like use the more communicationheavy forms of parisms or the less uh communication heavy ones. And for example for TPUs since they're in pods of 32 usually people they allow themselves to use tensor parallelism uh over the 32 pots. So yeah it really depends on the hardware and the uh global web size.
>> All right. Thanks guys. Um thanks so much uh Numan for the amazing and insightful talk. So let's give it a hand again for our speaker.