MIT 6.S191: Secrets of Massively Parallel Training

0:05 [music] >> Hi everyone. And welcome back. So it's great to it's my pleasure to introduce Matthias Lechner who is the co-founder and chief chief technology officer of Liquid AI. Matthias previous to co-founding Liquid AI, he was a research researcher here at MIT CSAIL where he co-invented Liquid Neural Networks and then he spun off and co-founded Liquid AI. I've had the privilege of working with Matthias for the past 9 years from when he was at MIT all the way to present day at Liquid where I also was a co-founder together with Matthias. So I'm really excited by today's lecture.

0:51 Matthias is a leader in the field and it's a it's a pleasure for me to get to work with him every day. So I think you'll all enjoy this lecture greatly and I'll pass it over to Matthias. Thank you for joining us. Thank you so much. Thank you so much Alexander. Um It's great to be back at MIT. And today I'm going to talk about massively parallel training. Specifically how we at Liquid AI scale training runs up to thousands of GPUs.

1:20 So quickly to talk about what we cover in this talk. A quick intro to Liquid AI. We are MIT spin-off as Alexander mentioned around 3 years old and we optimize on foundation for efficient foundation model. That's what we are building and we just recently released half a year ago this LFM2 series of on-device optimized um language models and yeah, continue on that. So in today's talk, first of all I want to convince you or I want to answer the question of why we scale machine learning models in the first place.

1:54 Then I also want to talk about what we scale in a way, the different methods and the different strategies that we apply to shard and and partition our workload across a cluster. And finally, also how we scale. Specifically about the frameworks and and implementation of the methods that we talk about in different PyTorch training frameworks. And finally, I'm also giving like a quick case study in a way like of what methods in a way are typically used in some example training runs that we did in in in our company.

2:28 So first of all, why do we train neural networks on a GPU in the first place? And to answer this question, I want to go back a bit in time to 20 2011 in a way when AlexNet came out. And AlexNet really was a image recognition uh neural network. It was fully was one of the first fully trained end-to-end, you know, on a GPU and it really sparked the deep learning revolution. And specifically, this model was trained on two Nvidia GTX 580 GPUs for around 2 weeks.

3:00 And if you recap what Alexander and Ava taught us before in the previous days in a way, training of a neural network is essentially just a bunch of matrix multiplications. There's a lot of other operations as well involved, but like the heavy portion is like matrix multiplications. So let's look at an example of like a input X consisting of 256 size in a way times a feature of 4096 like 4K multiplied by weight matrix of 4K times 4K.

3:29 And this operation in a way requires around eight billion floating point operations, so additions and multiplications. And if you go back like to 2011 in a way, what was available there? Like a CPU around that time a typical one had like around 0.1 teraflops like trillion floating point operations that it could execute per second. Compare that to a like a GPU like this gaming GPUs that I mentioned before, they could do 16x more in a way. And um as I mentioned this model was trained for 2 weeks on the GPU. So um on CPU this would take basically a year because like you have two uh two of them times 16x.

4:08 Um so yeah, this is all um the key takeaway of that message is like we are training neural networks on GPU because um uh they're very compute intense and the type of compute uh the type of computation in a way that is you needed for training uh very fits fits well with GPUs in a way with the type of operations that are going to be accelerated. So this is a long time ago anyway. What has changed since then in a way to give you a more modern example of a H100 a GPU that's still used today for inference and training a lot um that has seen right uh around 1,000x more uh than the the gaming GPU here. So it sees really like this deep learning revolution that really started like um there are a lot of uh hardware improvements uh that have been made on that level.

4:51 So now that we can we answer the question of why we train on a GPU, I also want to convince you uh why uh we want to train on larger data sets. And uh to answer this question let's look at a Llama 2 uh 7B which is an open source a large language model trained by Meta in uh 20 uh 23. And uh it's open source so they share the training curve in a way. So on the x-axis we see the number of training the data set size basically the number of training tokens that this model has seen in training. On the on the on the y-axis the training perplexity which is kind of how well the model fits to the training data. And like there's a clear trend that with more uh data the loss gets lower anyway.

5:30 And um yeah. So basically the takeaway here is like more training data means a better loss means a better fit to the data. Of obviously there's some assumptions that the training data is fresh and so on and high quality um uh but assuming plentiful of data is available uh more data means uh better model. And this is not just me, you know, showing you the plot and like eyeballing like this correlation away, but this has been systematically and scientifically analyzed in like a thing called scaling laws um that really like capture this kind of loss dependency on the training set size as a as a power law.

6:03 So now that we know that hey, we should train on more data to get a better model, I also want to convince you that we should train a larger model on this on the same data. So again, Llama 2 uh they didn't like Meta didn't just train like one model variant, but they trained like four different variations of the same architecture with uh different number of parameters in the model. The 7B, the 13B, the 30 uh uh 34B, and the 70 billion parameter model. And again, they shared the training uh curves of uh those training runs.

6:31 And one thing that we see is like if you take any point on the x-axis and look at the uh loss of all of the four models, we see a monotonic decrease of basically a larger model at the same data uh gives us a better fit to the data. So yeah, larger model on same data means a lower loss, means a more capable model. Again, same story as before. This is not just me uh looking at the plot and and making this up, but this has been systematically and scientifically analyzed in in papers.

6:57 So just to recap the first portion, um the scaling imperative in a way like uh means like we want to train larger models with more parameters, and we want to train them on more data on more tokens in a way in order to get um a better performance to kind of drive uh from the scaling laws uh uh a better model quality. To give you some example from industry like the more on famous example, GPT-3 um really pushed that hard like with a 175 billion parameter model language model trained for 300 billion tokens.

7:26 Then the Llama 2 that I just showed you, this largest variant with 70 billion parameters trained for 2 trillion tokens, so like uh roughly 10x uh scaling data. And then finally, um we just released this week Llama 2.5 uh 1.2B, which was trained for uh 28 trillion tokens. So another like uh more than 10x scale up from Lama 2. Um so, yeah, this shows a clear trend in the industry like from um scaling the dataset size. But, interestingly, I just told you um that we should train larger models, but the examples that I shared you um actually the model becomes smaller.

7:58 So, this is a bit uh actually a counterintuitive. Like, why is this the case? I just told you the exact opposite. And the reason for that is uh that the scaling laws that I just showed you, they only concern the training. But, after we train a model, we actually want to run it. We want to inference. And those add up the cost actually like uh if you let's say chat GPT and so on. Like, there's millions of users. You train it once, and then you have like millions of users in a way.

8:19 So, actually the inference cost dominate the compute like the the lifetime. If you look at the lifetime of a model, the training portion is actually only a small portion. Like, you need to think also of inference. So, that's why like there's a big push into more efficient models. To give you an example, like the the GPT-3 uh model, like the weights I cannot even download on my laptop. Uh I will run off out of storage. Whereas our LFM 2.5B it runs on a phone um much faster uh obviously. And uh it actually runs on a Raspberry. So, like the Raspberry Pi. So, the efficiency is is is really important. But, the quality like uh of LFM 2 is like much better than like the GPT-3.

8:55 Okay. So, now we know we should train uh more data and bigger models in a way. And uh there's one way that we can very easily scale the dataset size um by leveraging and exploiting effect from the training algorithm from gradient descent. And um basically the idea of data parallelism is um that each sample in a in a gradient um back descent update stochastic gradient descent update is kind of independent from each other in a way. And only at the end we kind of compute when you commit the gradients, we do a mean over the entire over all of the gradients for each individual sample, right?

9:28 Um and we can use this fact kind of to uh scale uh uh can scale the datasets by by using more than one TPU you the following way. We just replicate the model with optimizer states and everything across all the devices. Then we just distribute kind of a batch across the all of the devices in a way. Each of them individually computes in a way the gradients. And then at the end we synchronize and and have like only one communication step at the end where we kind of exchange information. And the advantage is that yeah, the most heavy compute of the forward pass and the backward pass of the model all run independently and like there's very very little communication necessary.

10:04 So uh in concrete figure in a way like a let's assume we have a batch size for 256 samples. We just distribute them on four cheap use. Each of them gets 40 uh 64 uh samples in a way. Each of them computes locally in a way uh the the gradients and and then at the end we do just the all reduce mean in a way where we just have one communication step anyway. So most of the time the cheap use are kind of independently working and the only like one uh synchronization synchronization step is necessary.

10:31 And this scales linearly. We can just, you know, stack in and add more cheap use like infinitely in a way like this uh yeah, we can in increases basically increases the batch the effective batch size in a way. And this is also in a way the uh bottleneck of this method in a way because um at some point the batch size becomes so large that actually the generalization of the model suffers. So uh if you train like a model one with a small batch size and one with a gigantic one uh on the same data on the same test uh sorry, training data and then evaluate it on the test set, uh the model trained on a smaller batch size uh will outperform the one. It has to do with like the loss landscape of the optimization and so on. There's a lot of papers actually written about this topic in a way. But in a way that uh just in a way the take home uh message is just you know, hey, batch size scaling is the most naive way to scale the dataset, but there's a limitation. And we need to think about if you want to scale in a way that the throughput of the model of the training run without scaling the batch sizes becomes much more challenging.

11:25 So this is about dataset size. So there's this naive approach at least up to certain um threshold in a way um of of scaling dataset. But what about memory? So, let's look a bit into exactly like how much GPU memory do we need to train a model? And let's take an example of 1 billion parameter model that has like a that we want to train on a batch size of 16 uh with like uh uh 256 layers and 1K feature size or hidden hidden dimension.

11:50 So, first of all, the parameters itself, like we usually store them in 16-bit precision in a way. So, it adds up to 2 GB. Then, we also need to have a gradient buffer because we need to compute the gradient. So, there's also a gradient that we need to allocate on the GPU. Uh and finally, we have the optimizer states. So, the optimizer I think Alexander and Albert talked about a lot about them. Basically, like we want to have different learning rate like different parameters learn different fast in a way. So, we want to we want not just want to apply gradient descent, but actually precondition the weight updates. And we do this uh with with with Adam in our case for instance. And this also requires some parameters in a way that are exactly like uh uh a multiplier on the on the parameters.

12:28 In this case, a 32-bit we store them. It's just an example. You can actually also quantize them to some extent. And then, finally, activations. And the activations uh why we need them is uh it has to do with backprop because like in back propagation, if you remember, um we first do the forward propagation in a way. And for the backward propagation, when when we get the backward signal in a way, we need to multiply it by the activations in a way to kind of for the computation of the gradients.

12:51 So, actually, we need to we need to remember them. So, in total, like if you want to train a 1 billion parameter model in this simple um calculation here, it adds up to 28 GB in a way. So, the message here is again like um even though the model is small, we need a lot of GPU memory just to train those this model even. So, if you extrapolate this to 70 billion parameter model, this would uh fit uh that would require 2 TB. And like the largest GPU single like single device memory that we have at the moment like 200 300 400 GB at the moment.

13:21 So, yeah, there's no way we can fit this. So, we need to have better strategies to kind of reduce the memory uh requirements of training. So, one method that you that is used heavily in industry is called activation checkpointing. And it's basically not a free lunch, but like a trade-off compute to memory in a way. And the idea is like if you recall without gradient checkpointing sorry without activation checkpointing for each layer that we compute 1 2 3 4 5 6 and so on, we need to store the activations because for the backward pass we need them.

13:55 The idea for for activation checkpointing is let's only store let's say that every fourth layer the activation and then so it means like for layer one and layer five will will keep the activations in memory, but for the other in between we just discard them. During the forward pass and then the gradient the backward pass in a way we just recompute like that portion in a way that's currently needed for that portion. So, let's say in a backward in a forward pass we just as I mentioned discard them and then in the backward path if if we arrive for instance layer number five we just recompute the activations from two three and four and do that and and finish in a way the the the back propagation run.

14:32 So, in this scenario here we see like four x reduction memory in practice it's it's it's it's a lot so that's why it's used heavily. Extra compute it sounds extremely expensive, but if you think about it the first forward pass we have to do anyway. The backward pass we have to do as well and backward pass also we have to do back propagation to the input and to the gradient itself. And then we just have one additional forward pass in a way.

14:53 So, it's like it sounds like it's twice the effort like the forward pass is only like a small portion of it. So, it doesn't add it doesn't blow up the compute. And this is also one of the reasons why it's used in practice a lot. Another reason it's used in practice a lot is because typical GPUs like the compute tends to be cheaper than memory in a way. So, so it's it's used as I mentioned before heavily in industry.

15:17 Another approach to reduce memory is offloading in a way offloading to CPU. And this idea is like general to computer science like in in in our laptops and so on and everything like we have a cache where we offload them to the main memory and this offloads to the disk in a way when when it becomes um um uh the idea is basically like we that we push infrequently used data to slower but larger memory keys. Like it's a standard concept in in computer science.

15:41 Um and uh we could apply this as well to training of neural networks. Specifically, the optimizer states are only used at the very end in a way to precondition the weight updates. We also have inactive layers like when we compute a layer 1 to 10, we don't need layer 100 and so on in memory maybe. And activations as I just mentioned as well are very heavily. And maybe this amortizes over the entire training run or one one great one update step in a way because um the data is not needed all the time.

16:07 However, like there's one big issue with that is the bandwidth. There's like a huge gap in a way between um the data that the GPU has access to versus the one that the host like the CPU and so on has access to. It needs to as well go through this PCI bus and then also through the another bus to get to the host memory and so on. So that's why it's it's it's quite um slow from the perspective of the GPU.

16:28 And this is why it's it is used but it's used more like for niche applications. So for instance like fine-tuning of a large language model. If you want to train if you want to fine-tune let's say a a 70 billion parameter model on one GPU, there's no other strategy that you can use. You just have to offload parts to the to the host memory. Another example is like context length. There's like a thing called a DeepSpeed Ulysses that basically say for very long context length uh you you also offload like part of the sequence in a way to the CPU memory and so on. So in in general like it's scenarios where you need to train only a little bit. So that's why like offloading is used maybe more for smaller training runs but if you really want to utilize a cluster, there are there are better strategies.

17:05 So we saw that hey, this offloading idea looks nice in a way but it's not quite there in a way. Um so let's look a bit more closely in a how actually our GPU cluster that we that we get um that is getting built and so on, like this gigawatt clusters that's all over the place at the moment. Um how they look like uh inside anyway. So, typically like we have uh node that has like a CPU that's connected to several GPUs. In practice, for for Nvidia it's like uh eight uh GPUs on on a single host anyway.

17:33 Each of the GPUs has its own cache and so on, its own HBM, like high bandwidth memory, as well uh and uh the CPU kind of the controls it that has its own memory pool and so on. And then what happens is like we connect the GPUs to each other with like a very high uh in- uh interconnect with like typically this is called NVLink in Nvidia domain. It has different name for names for different providers anyway.

17:54 And this is called scale up. So, basically we have like one node anyway where we don't just have one GPU in there and so on, but we really stack up up to like the maximum GPUs that we can fit on a single node anyway. This in in in typically it means eight. And then also we uh yeah, they are connected to this CPU with like a kind of slow bus anyway. And then we just also do a scale out anyway where we just add more machines uh of the same basically of the replica of the machine anyway to the network. And then we connect them together with like a network networking interface. In the Nvidia case uh InfiniBand is used here. And and this InfiniBand is you can think of very similar to kind of a normal networking that we use every day anyway to connect our our machines. The only difference is that it needs to support like this thing called uh remote direct memory access.

18:38 Because I mentioned before, like if you move everything that typical network stack anyway moves uh data around on the CPU and so on that controls it. But this uh remote uh direct memory access anyway is um if you want to two GPUs on two different nodes want to communicate and exchange data, you don't want to go through the host, you know, the CPU on both machines, but you just directly, you know, pipe it from the um uh HBM memory into kind of the the networking interface.

19:03 So, equipped with this knowledge of like how does actually a cluster look like that we have this um eight GPUs, each of them, you know, is is uh very fast in a way. The interconnect on the node is much faster than the interconnect on the scale out fabric and so on. Um so the next option that we have in a way to train faster and also larger models in a way is a sharding in a way. And again, this is a concept very generic to computer science. Basically, like the idea is to partition data across multiple locations in a way.

19:31 It's used heavily for instance databases for meta and so on. They are like also sharded across the globe and so on cuz you need the data not everywhere in one location. And specifically for neural network training, there are a lot of options that we have for sharding. Like we could like basically shard the model weights, we could shard the optimizer states, we could shard the activations. There's also data and gradients and other stuff that we could shard.

19:53 And um the key trade-off here is like in a way that we reduce in a way the memory that each device in a way must hold every time in a way all the time. But also it adds a communication overhead. Like there's you know, for the different strategies of the different things that we want to shard. Um actually communicating with the rest of the network might look very different. The simplest case is optimizer sharding because as I mentioned before, the optimizer is only used at the end kind of the precondition the weight update.

20:24 And so it's actually used not too often. All of the methods that I'm I'm going to speak about from now on and so on, they assume data parallel that we talked before in a way is it the default setting. So like data data parallel in distributed training, data parallel is kind of the the standard setting. Not just single GPU training, but like data parallel kind of is the default, which is nice. So in the default here without sharding data parallel, we have replicas on each of the two devices here, GPU 1 and GPU 0. Each of them needs to have like 12 GB in a way in this example scenario. And with optimizer sharding, we just split in a way the the the optimizer states in a way into two partitions. Um one on GPU 1, one on GPU 0. And reduces in a way the memory from 12 to 8 GB.

21:07 So this doesn't come for free, but it adds communication overhead and specifically the update process on the of the weights is slightly changed. So, first of all, the forward and backward path is completely independent anyway. So, there's no communication needed anyway. Um then at the end anyway, we do a reduced scatter where basically like each a GPU they exchange the information on the gradients. So, they we send the gradients for the for the specific shards to the to the shard anyway to the to the device where the shard is.

21:32 Then [snorts] we compute the the weight update in a way locally. Uh each shard on their portion only. And then at the end we kind of all gather kind of we exchange again the the the weight updates or the or the updated weights um across the entire um sharding group anyway. Again here like a lot of memory savings anyway because optimizer states uh typically like two variables per parameter like in the case of Adam for instance.

21:59 And yeah. Next up uh next strategy is pipeline parallelism. So, this reduces both memory, but also is used to improve increase the throughput um of our training run. And the idea is basically if you remember like a deep neural network, the the the name deep comes from that we have many layers. And the idea is, "Okay, let's just split the model, you know, and put like let's say layer one a layer zero to 31 on GPU one a GPU zero, sorry. The next couple of layers on GPU one and so on. Next layers on the next GPU and so on. So, we kind of split in depth anyway the model and shard it there."

22:32 There's one major issue with this approach. It's called pipeline bubbling pipeline bubble basically because if you think about training, if you remember again, we have to do the forward pass and the backward pass. And when you do the forward pass and compute like layer one zero one and so on, all of them sit on GPU zero. So, GPU zero is like running with 100% but the others are idling. So, they're waiting basically for for any workload anyway.

22:55 So, nothing comes. So, they're basically there's a lot of waste. The backward path, the same story but the reverse order anyway. A huge waste, basically. The idea to solve this problem is micro batching in a way where you split the entire batch into smaller units in a way. So, the the batch size and then you basically only GPU zero, you compute the first micro batch. Then you forward pass it to GPU one who does it then the the next layers and so on. At the same time, GPU zero, you know, it does already the second micro batch and so on. And the same in the way for the backward pass So, it's a a kind of a simple idea. And pipeline parallelism is used heavily in industry and specifically like for instance two examples of the in industry is not just use the micro batching, but they have more sophisticated methods.

23:39 So, for instance, Llama 3, they do actually some bidirectional stuff where kind of they overlap a kind of forward and backward pass more uh uh um interleave it much more, anyway. For instance, like chief the last chief you win a pipeline does prioritize for instance more backward pass early on. You know, in order to kind of say, "Okay, the the the previous chief use can already like back prop that portion." Like there's another algorithm from Deep Seek which was used in this Deep Seek V3, anyway. It's called 12 pipe that actually also further separates like more granular. Like instead of just micro batches forward backward pass, it also like separates backward to input with from backward to weight in a way into two separate um uh smaller units kind of that can be scheduled.

24:27 So, this is just, you know, splitting the network by depth in a way. Um but we can also split layer or within a layer, anyway. And this is called a tensor parallelism. So, remember like we do a lot of matrix multiplication, X * W = Y. And the idea with tensor parallelism is just that we split uh W, the weight matrix in the half. We can do this uh by column in a way with column parallelism in a way.

24:51 Um and each GPU has like um only a certain portion of the partition. Uh and we can do this also in the other way. Like we can split it by row anyway. And sounds the same at the beginning. Like when I first saw this in a way like I looks the same. Like what's the difference? The difference is like nuanced. It's specifically column parallelism in a way that X the input needs to be replicated on both devices.

25:14 But the output will be partitioned on both on basically separate devices. And for the row parallelism, the inputs can be partitioned on two different devices. And the outputs will also be partitioned on two devices. So we can actually leverage this fact for a thing let's call it a transformer MLP and so on. So you don't have just one matrix but two in a way you would stack two layers after another in a way. You add a non-linearity in the middle.

25:35 Otherwise it would be point like pointless basically. Um with the with a relu or in this case gelu activation function. In this case in a way you basically start with a replicated input. You would do the first uh column parallelism. It's uh partitioned the intermediate states. And then the second um uh you apply because the second layers in a row parallel. So basically the inputs can be partitioned. And then at the end uh you produce uh partitioned outputs again.

26:00 And then you do one all reduce in a way to kind of aggregate them again and make them replicated on the devices in a way the the the output in a way. So the idea is in with like stacking one column and one row after another, we can get away with sharding two layers with only one communication step at the end to synchronize again and and and exchange with the output information. So that's why it's used in practice typically like in that format in a way.

26:26 There's other dimensions that we can split and partition our data sequence and context parallelism. So um typically we train on uh like context lengths let's say in this case 8K. And we can say okay the first kind of tokens the 0 to 2000 are on GPU zero. The next 2000 tokens are on GPU uh one and so on. And I got a lot of layers in a neural network like MLPs layer norms and so on that they're kind of they don't mix time. So they're time independent in a way. So there's no communication actually needed between the GPUs in this case.

26:54 Um there are also like layers that are a bit more tricky to um paralyze in in in in length in in sequence, sorry. Specifically, attention layer that's used in transformer um model um requires a bit more care. Like there's a thing called context parallelism in a way that there's essentially like partitions the KV cache. Um and there are a lot of like algorithms on how to use efficiently ring attention and so on. So, it's a bit more basically in in complex version of of sequence parallelism.

27:21 So, [clears throat] yeah, again, like sequence parallelism for like the independent layers, context parallelism for the for the transformer attention layers in a way. And then this is typically used in practice kind of for long context training. So, when you train a a large neural network model, you train them typically like on short context like 8K and so on. And then after the main training phase, you extend the context length. So, you train just a couple of like trillion like a few billion hundred of billion to be trillion of tokens with like a longer context length.

27:45 So, it doesn't need to be trained from scratch really with with a 1 million context. Um another or the last um parallelism strategy in a way is about mixture of experts. So, the idea with mixture of experts is like a new type of neural network architecture on new type of layer in a way that tries to scale model capacity. So, try to really use the scaling laws by scaling model capacity without adding extra compute. Because in a traditional sense in a way if you scale in a way um the model in size, you know, our weight matrix becomes bigger and so on and that means the inference, the training, everything becomes more expensive.

28:18 But the idea with mixture of experts like let's basically slice up our weight matrix into like rows. Um they're called experts. And there's a router component. And the router component in a way uh tells in a way which one we should activate it. And the rest in a way is assumed to be zero in a way. So, it's kind of a way of sparsity is a way of structured sparsity in a network. And this used heavily because it's yeah, as I mentioned, scales the model capacity without adding additional compute.

28:45 And this leads to a very natural way to par to shard and partition in a way uh the mixture of expert layer by just putting different experts on different devices. So, let's say in this example here, we have expert zero and expert one on device zero, expert two and three on device one and so on. So, it's kind of a very natural way to kind of parallelize. Um there's also again like communication necessary in a way to kind of make sure the router the tokens are distributed and also the the results are aggregated.

29:11 An example of mixture of expert architecture that we trained, you know, um is like LFM-2 8B uh one uh A1B, which is like a 8.3 billion parameter model that has only 1.5 uh active. So, it's like it runs with the inference speed of a 1.5 billion model, but has the kind of the capacity of a 8 billion model. Um yeah, as I mentioned like it basically scales the params, um but the compute to much less extent.

29:37 And like there's a major challenge always the ME architecture is about load balancing and load balancing in two aspects. The first aspect is for the training stability. So, if you want to if you want to train ME layer actually, you need to make sure that the router in a way doesn't overfit in a way and just routes same like all tokens to one expert and then you actually train a small network because the rest is never activated. So, there's like uh making sure that the training is stable.

30:01 There's a lot of different um methods that's still being developed and so on on how to how to kind of balance in a way the load of the network during training properly. The second problem with load balancing is that exactly with the has to do with the expert parallelism is because like imagine we have the expert parallelism and the load balancing doesn't work and all the tokens get routed to uh expert uh zero and expert one, it means the GPU uh like one has all the load and the the rest of the GPUs are idling kind of. So, you have like a lot of like um it slows down the training a lot. So, you need to be really careful with load balancing. As I mentioned, there's still still actually actually free search going on into how to do this properly.

30:41 So, these are all different methods and strategies kind of to separate and partition the data and the compute and scale the training throughput. But, now I want to also quickly talk about frameworks and and and tools in a way that implement some of these strategies that we talked about. The first example is DeepSpeed Zero. The Zero stands for Zero Redundancy Optimizer. And it's basically about sharding the memory. So, again that the baseline, as I mentioned before, baseline in distributed training is always data parallel. So, data parallel training.

31:06 Um there's a stage one of Zero um where basically it starts optimizer sharding exactly as we learned before. Um there's also stage two, which shards also the gradients. Has a additional effort, but more on that later. And then also stage three, where you partition everything that the optimizers states, the gradients, and also the parameters. The communication pattern for Zero 1, which is the optimizer sharding that we showed before, it's basically this reduce scatter to kind of send the gradients out to the specific shard and then they're all gathered to kind of get the updated uh weights again. Updated parameters uh are synchronized in a way.

31:41 Uh for Zero 2, we do additional reduce scatter by sending in a way that the gradients anyway to the corresponding shard that has the gradients. So, that's one more step level of communication. And then Zero 3 anyway has actually much more than that because like you for the forward and backward path, the parameters the parameters have the problem that we use them much more frequently than the optimizer states. So, there's actually much more um um communication necessary anyway if you want to shard also the parameters.

32:06 In practice, it's still used um because you you block them together anyway. You don't do every individual weight, basically. And yeah, so basically Zero 3 allows you really to scale like to multiple GPUs uh where the model is so big that it doesn't fit on a single device uh on the single like memory. So, this is DeepSpeed uh DeepSpeed. Not to be confused with DeepSeek, which also provides a nice architectures and so on. Um uh there's another method uh like another way to implement sharding is Fully Sharded Data Parallel. It's a basically PyTorch native implementation of a zero-three style sharding that we learned before anyway.

32:42 It's slightly different um on on a on a the the code required in a way to implement this. But basically you can think of because it's PyTorch native you define a nn.Module uh FSDP block in a way and the FSDP block in a way handles uh the wrapper in a way handles uh the sharding for you. So basically think of of this if you write your Python script in a way um in the default setting is that all the data is sharded fully. Like that's what fully sharded. And then when your Python program enters the layer in a way it does a all gather and kind of gets all the parameters in a way in the forward pass. So it like in your Python script in a way you have access to all the weights of the layer of the module.

33:17 And then when your Python execution exits the program uh exits the layer, sorry. Um you discard in a way the fully the the aggregated uh information again. Um and only have the partitioned uh um parameters. Backward pass similar story but here we need to maybe also communicate the the gradients and all those but you know what very similar. The the execution engine enters in a way the backward pass you kind of get the information that you need and so on.

33:42 And um yeah there are several uh settings of FSDP. You can actually operate it in loss of only gradient and on I'll say only optimizer sharding and so on. So you can basically implement also zero stage one and stage two with with FSDP. And yeah it's a used heavily also with activation checkpointing in a way to train on on on uh larger models. So I just introduced uh a quick about DeepSpeed and then talk about FSDP and at the end I said basically they're providing the same features in a way. So what's what's the difference? How do how do they differ? Um but before answering this question there's another thing is called Megatron-LM and it's Nvidia's uh framework for sharding and and paralyzing in a way large transformer train pre-training or training runs.

34:24 Um yeah distributed training framework. And the way or how these different frameworks differ is the um I think I think difference is like uh who made it. In a way, for DeepSpeed comes from Microsoft. Um FSDP comes from Meta, and uh Megatron comes from Nvidia anyway. So, it's like uh they all offer very similar feature sets, but uh uh yeah, each company wants to do their own stuff anyway. Uh there's one company missing here. It's Google.

34:50 It's because they are using JAX and not PyTorch. So, um yeah, they are very similar uh conceptually. Um They're slightly different aspect on like focus and also the design philosophy in these different mm um um packages anyway. Like uh for DeepSpeed zero and FSDP, it's mostly about memory, but the the Megatron LM also like the parallelization and the throughput aspect is much more uh focused area as well. So, for instance, like pipeline parallel and tensor parallel, for instance, they need to be added um externally to FSDP and DeepSpeed, but in Megatron, there's like layers already implementing those things.

35:27 Um yeah, as I said, the philosophy is slightly different. For for DeepSpeed, um the idea is you don't want to don't want to change the code. You just wrap it in like one you define your model, your optimizer, your data loader, and then you just have one line change where you wrap it into a DeepSpeed uh function anyway. There for FSDP, it's PyTorch native idea, so you define your own nn.Module basically. And for in Megatron, it's a bit more complex. You have this model parallel unit that keeps track of like the how what is exactly sharded where and so on, like the communication groups basically.

35:53 And like the ease of use basically like for DeepSpeed, it's just you know, have one line change in your code, but all of the the way to shard and how what sharding methods use is all defined in a YAML and a config file anyway. So, it separates anyway the code from the configuration. In PyTorch and FSDP, it's more like yeah, it the the the aspect is really like hey, this is this feels like PyTorch native anyway. And for Megatron, much more custom code anyway.

36:13 But at the at the at the at the advantage that it allows you more uh has a bit more features. So, putting it all together anyway, um when you want to have a let's say a cluster of 2,000 GPUs and want to maximize anyway um and want to train a big model. The first point is really all like on memory. Like we want to use heavily activation checkpointing because as I mentioned before, like memory is typically more precious than compute in a way. So we do heavily activation checkpointing.

36:40 Another point that I didn't mention before, but mixed precision. So we basically we don't need the same precision levels for all of different portions of network. So for instance, like the activations are much more sensitive, but the weights can be quantized and can be expressed kind of with very low precision level and things like that. I also mentioned like FSDP and zero shotting in a way that kind of shots parameters, gradients, and optimizers and only kind of exchange this when needed in a way.

37:07 There's also offloading that we discussed. And as I mentioned, it's more like about niche applications. There's also about the throughput. This parallelization strategies and they're often called 4D or more recent now 5D parallelism because actually you can use them at the same time in a way. So you have different axes. You can think of it like different dimensions in a way and you just scale the different methods on the different axes. So we have like data parallel where basically you have you have data parallel, you have the tensor parallel where we split the layer.

37:36 You can do pipeline parallelism. Sequence context parallel, expert parallelism, which is very similar to tensor parallelism, but for mixture of experts in a way. And there's other optimizations like fused kernels where you make sure that the the overhead between sending communication to between the CPU and the GPU doesn't add up. Also like stuff like flash attention is basically very cache aware kind of algorithms in a way that make the make use of the cache on the on the device very efficiently.

38:04 And when we when I mentioned before like this is 5D or these dimensions means basically like imagine that 2000 GPUs in a way. Like and then we have like these different dimensions you add them together. You don't run like one big data parallel tensor parallel that's it in a way, but you really add them together in a way. So for instance, like TP8, so tensor parallel 8, what does it mean? It means that we kind of split the entire cluster into eight partitions and each of them in a way has their own batch memory mini batch and they only communicate these partitions of the cluster and they only communicate to kind of exchange the gradient information at the end.

38:35 TP8 Uh so tensor parallel means for instance like tensor parallel is usually very communication intense operations. So you want to basically put them on a single node where we have these eight GPUs. So let's say I want to run tensor parallel on on the on one node in a way. Um and then pipeline parallel some will say hey I want to use and split the layers the network in depth on four different nodes. So kind of you have like four nodes in a way that the depth of the network is sharded across each each layer kind of tensor parallel on the on the on the machine with eight GPUs and then you have like groups of eight for the tensor parallel so where you where you basically basically process different types of different different part different portions of the batch of the training batch in a way. Then again also expert parallel and so on and sequence parallelism is another dimension in a way that you can add there.

39:24 So in practice in a way you basically stack them together. So to summarize a bit the talk in a way first mentioned of why we scale in a way the scaling laws that bigger model means a better fit to the data more data means a better fit to the data. Also discuss like why we train on GPUs they are much the training you know much more suited for the training. Um and we want to train models that don't fit on a single GPU.

39:47 Um and also for training a model we need actually much more memory than just the model weights. Uh we also discussed memory solutions like activation checkpointing you know offloading the sharding that I mentioned also quickly touched on mixed precision. And also parallelization strategy that both benefit data like memory reducing memory but also increase throughput in a way. Data parallel is the naive one that has a limit on batch size tensor parallel which splits in a way the weight matrices pipeline parallel where you split in depth and then also expert parallel that is specifically kind of a tensor parallel for a movie in a way and also context and sequence parallels uh slightly. And the key key principle again like it's really the network topology of the cluster in a way like what device is connected to which other device with what kind of uh memory bandwidth and latency and so on in a way to to match in a way the the the the communication hierarchy. So like yeah, for kernel development, so if you develop a new if you use kernel and so on, you actually you don't even think about HBM memory, but you think about cache. You know, what's what you you you that's that's your your abstraction level that you work on with.

40:54 Then you have the HBM memory, then you have this NVLink like kind of the local group in a way um of GPUs that sit on the same machine and then you have the scale out network of InfiniBand and so on. And again like there's no sufficient uh there's no single strategy that you say like hey for this cluster I need X and that's it in a way, but the idea is really that in practice you combine multiple of them. You have this 4D 5D parallelism um uh and also what is interesting is like depending on model size you use slightly different way to shard, you know? Like let's say like our 8B model is trained very differently than let's say the 1B model is trained very differently than how we would train it a 1 trillion model and so on.

41:28 Uh really depends, you know, on the model size and on your cluster in a way. Um that uh concludes uh my talk, yeah. THANK YOU. >> [applause] >> AMAZING. OKAY, WE HAVE SOME TIME FOR QUESTIONS. I just have a tiny one. Uh back one slide. In your final example, you're like we'd do this here and this here and this here. >> [sighs and gasps] >> What would be that final sort of memory versus time trade-off for your last example there?

42:05 This example here or what? >> Yeah. Yeah, the 2048 GPU training run. What do you mean? >> Like what is your versus cost? Like Like, how much memory did we save versus how much extra compute did we have to do? Yeah, it's a good question. Um the this is just an illustrative example because I mentioned in the next slide that actually depends on the size of the model. So, this was actually chosen very simple to like I think a 70 billion parameter uh Lama style model in a way.

42:36 Um and yeah, the the the idea is really that you cannot fit this um on a on a uh uh typically like what what happens in practice if you let's say a train a 70 billion parameter model and so on. Like, you can use uh FSDP and zero to kind of train it on like uh First of all, it doesn't fit on a single device, but you need it maybe needs uh like four nodes, then you can actually fit it with FSDP or Steep Speed zero.

42:56 But, your throughput will be very slow. So, um you would need um like years to train it in a way, right? But, there's so much competition and so on. You need to train a model faster in a way. So, that's why you need to add um more GPUs in a way. Exactly. And that's why you need to uh think about like, "Hey, we need to add this pipeline parallel and so on." Ideally, we would just like to train with data parallel because it's the simplest one easily. But, like there's like uh there's the batch size becomes already like very large. So, like for for language models, typically like they use like uh things like 8 billion tokens and so on. And up there uh from that point on, people start to notice like degradation in performance.

43:39 Uh thank you for the um excellent lecture. Really really great chart. Um question I have is how does how does your architecture shift when you have to um train around um sort of like a like more like a decentralized training um where you you don't have a co-located servers, right? And you have a you know, unreliable workers and nodes that can drop and um and you know, which which which of these dots um two parts like which of these which of these boxes like are you thinking, "Hey, you pick out."

44:06 And then building on the work that's already been done, I'm thinking like a trainer at home deep mind you know what these people prime intellect. How what what do you think they're really missing that I think like what do you think is the bottleneck that really can unlock training large language models in any decentralized way that's really an area of research that should be considered. Thank you. [clears throat] Yeah, that's a that's good question as well. Um there's actually like some work already going on in distributed training and specifically the data parallel portion is actually very well suited for distributed training because like there you don't need it it's basically the parallelization strategy that requires the least communication. So there's only like very little happening. So for instance, I think Google for instance does actually data parallel across different data centers. Because the communication is so little in a way that you can actually move do data parallel across multiple data centers in a way that different locations. And I've seen also like distributed training runs in a way at least smaller scales that try to do data parallel.

45:04 It becomes there's also more research going on in asynchronous because like some of them maybe you you have distributed very fast and some of them are slow cheap use and so on. So there's there's a thing called asynchronous training. Um but it's I think it's it has been heavily used I think in the past but recently I've not seen so many bigger runs in a way but it's it's an active research area. You mentioned earlier that um like at a certain point if you increase the batch size you um you harm the generalization of the models. So I'm just curious like if you're trying to train a small size model uh and you want to improve the wall clock at a certain point you can't just improve the data parallelism. So what is the strategy? Like do you train a larger model, distill it or like what is the how do you improve the wall clock of like for example the 1 billion parameter model that you guys are training?

45:59 Yeah, that's um to scale in a way yeah, that's a big issue. Like once it's basically scaling is easy as long as you're in a data parallel stream at the batch size is fine in a way, but it becomes very challenging. And like yeah, there's like I think uh um pipeline parallelism um I think probably the first thing that people use in a way to scale further. Um the problem that small models have specifically um on that is that each kind of it doesn't have many layers for instance. So it has only like 32 layers or 16 layers and so on. So there's actually not much room to shard stuff around. And also the tensor parallel and so on because the matrices are actually quite small already. So kind of them splitting into multiple pieces in a way like it's very tricky to scale up. So there's a lot of algorithmically um advancements that we do like for instance um distillation you mentioned in a way where you basically you um want to accelerate the training in a way not by running faster more tokens in a way, but by adding more signal for more having a richer training signal.

46:56 Specifically distillation where you say hey, you don't just predict the next token, but you need to predict the next distribution of a teacher model anyway. So this is for instance how we scaled in a way LVM 2. So we used like a previous generation LVM 1 7B model um that gives us a distribution not just of the next one logic token, but like um 16 basically. Thanks. Thanks for the excellent talk. Uh I have a question on one of the early slides. I was very impressed by the 28 trillion token on a 1 billion model.

47:28 And my question is um how much quality check are you able to do on 28 trillion tokens? And um if quality really how much it matters. And what is the current status of being able to distill? Like I have a huge data set, can I just distill the data that really matters and train only on that? Evidently not that good. Yeah, like that's um Yeah, very challenging. I think um if you look at some of the open source data sets on Hugging Face for instance, like you can get um the first kind of trillion tokens and so on very high-quality data from archive scientific papers from Wikipedia and so on. But the further you scale actually on the on the token axis actually it's more and more difficult to find high-quality data in a way, right?

48:17 Because you you source okay, PubMed archive and so on is nice. We we can then then you do filtering like on Common Crawl. I think they only really what what what we're doing for instance like we we take Common Crawl and filter it. And then training like in a way smaller filter models in a way that just detect in a way which one could be useful for the for the for the for the training of the large models and so on. We do a lot of ablation runs on data and so on.

48:39 There's also a lot of work going on in in synthetic data. Um that but this is also careful because you you don't want to like model it collapses the training. Because if you train we also train some models in a way where we we're fully synthetic, but this in a way collapses a lot of the training dynamics. So it has people call it like the model has no soul. Because it's trained only on synthetic data and so on. So it's it's it's it's a really big issue in a way and like the I think the the point really that they the first first the first tokens like highest quality data, but like it is finding you know the the last like 30 trillion plus tokens of high-quality data is very challenging, yeah.

49:12 Amazing. Okay, maybe two more questions. Uh thank you. Uh it was really impressive. Uh my question would be for the Liquid AI. As uh these are the smaller models compared to the GPT so that these are like we can use it on our phone. But as we progress like Liquid AI progress on to the future uh to basically compete with the bigger models for their accuracy and the quality of the uh basically answers or anything.

49:42 Um uh don't you think that the parameters and the size of the model will increase and it would not be able to fit on the like phone phone's RAM or something? Uh is it like do you have any like solution for this problem or it will become like normal models for like GPTs and we are using on laptops and things? Yeah, great question. Um The answer is that we are still doing active research, you know, we are all scientists in a way. We are still doing active research because we believe we can still push much much more than what's possible. Like even LFM 2 what I showed before it's it's it's boring in a sense in a way. There's much more that we can do in a way on algorithmically architecturally the side in a way. There's challenges like with with scaling, you know, like the the parameters and so on because like as you said like the the the the chip shortage and this kind of memory shortage at the moment and so on.

50:32 But one other in a way thing that I'm that at least Liquid AI has like we have very strong hardware partners with like with AMD, with Qualcomm and so on where we have like very good ties to. And they're also improving the hardware. So what we what what how we think the future is going to happen is like we co-evolve together with the hardware providers because the hardware will also at more than say you actually need a lot of more flops with low precision. We can we can offer that on the on the device.

50:54 You need maybe different tiers of memory and so on for this mixture of experts architecture and so on so like I think on the phone today in a way there will be limits but like we believe that in the next 5 years and so on also the hardware will be much more equipped with better design for for training for for running generative AI. Awesome. Last question. Thank you. So I have a question, you know, like you told Megatron LM right from Nvidia one of the framework.

51:22 So what do you suggest like you know if someone want to do a domain specific not edge-based device but domain specific model training like which framework is good? Thank you. I would say like the ease of use in a way probably the easiest to use is the DeepSpeed because as I mentioned before, like you would just have one line change. And you you get like a lot of the memory benefits in a way from that. There's also like DeepSpeed has also like this offloading that I mentioned before already also implemented. So imagine you have only one GPU and therefore cannot use FSDP, you can actually use offloading with DeepSpeed. So like the it's kind of the entry point in a way probably for for for DeepSpeed. And then the more you go you customize more. So for instance, like we use some we use DeepSpeed for kind of the early experimentations on multi-modality specifically. But then if you that our bigger pre-training runs we go like with FSDP plus a fork of Megatron. We have actually a fork of Megatron um where we use some of the sharding sharding strategies from Megatron, but we also incorporate it into a FSDP for memory saving.

52:32 I think okay, one more bigger round of applause for Matthias. >> [applause]