GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (Paper Explained)
deep learning machine learning arxiv explained neural networks ai artificial intelligence paper nlp billion parameters float32 attention mechanism transformer scale gpt-3 google gshard xla sharding parallelism mixture of experts trillion tpus distributed m4 multilingual translation natural language processing
Google builds a 600 billion parameter transformer to do massively multilingual, massive machine translation. Interestingly, the larger model scale does not come from increasing depth of the transformer, but from increasing width in the feedforward layers, combined with a hard routing to parallelize computations on up to 2048 TPUs. A very detailed engineering paper!

OUTLINE:
0:00 - Intro & Overview
4:10 - Main Results
5:10 - Mixture-of-Experts
16:00 - Difference to Scaling Classic Transformers
18:50 - Backpropagation in Mixture-of-Experts
20:05 - MoE Routing Algorithm in GShard
38:20 - GShard Einsum Examples
47:40 - Massively Multilingual Translation
56:00 - Results
1:11:30 - Conclusion & Comments

ERRATA:
I said the computation of MoE scales linearly, but actually, it's sub(!)-linear.

Paper: https://arxiv.org/abs/2006.16668

Abstract:
Neural network scaling has been critical for improving the model quality in many real-world machine learning applications with vast amounts of training data and compute. Although this trend of scaling is affirmed to be a sure-fire approach for better model quality, there are challenges on the path such as the computation cost, ease of programming, and efficient implementation on parallel devices. GShard is a module composed of a set of lightweight annotation APIs and an extension to the XLA compiler. It provides an elegant way to express a wide range of parallel computation patterns with minimal changes to the existing model code. GShard enabled us to scale up multilingual neural machine translation Transformer model with Sparsely-Gated Mixture-of-Experts beyond 600 billion parameters using automatic sharding. We demonstrate that such a giant model can efficiently be trained on 2048 TPU v3 accelerators in 4 days to achieve far superior quality for translation from 100 languages to English compared to the prior art.

Authors:
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, Zhifeng Chen

Links:
YouTube: https://www.youtube.com/c/yannickilcher
Twitter: https://twitter.com/ykilcher
Discord: https://discord.gg/4H8xxDF
BitChute: https://www.bitchute.com/channel/yannic-kilcher
Minds: https://www.minds.com/ykilcher
 OpenAI has a 175 billion parameter model. You thought that was large? That's cute. Check out Google's 600 billion parameter model, 600 billion floating point numbers doing things at the same time. This has absolutely become a body part measuring competitions between companies. Google be like, oh, GPT-3. I spit on you. I spit on you and your little tiny 175 billion. OK. Let's stop kidding. This is a giant model that Google has trained right here. The paper we're going to look at today is called G-shard Scaling Giant Models with Conditional Computation and Automatic Sharding by Dmitri Lepikhin et al. of Google. And this paper basically tells the story of how they built this 600 billion parameter model, how actually they attempted to build a model that had a trillion parameters but just didn't manage to quite train it. And this is all using this system called G-shard. So I haven't actually seen the code out for G-shard yet, but I'm going to maybe assume that this is something that they're going to release at some point. Who knows? Or maybe I just haven't seen it yet. So this is basically describing a system on how to train these giant models. So if you have watched my video on GPT-3, which, of course, was this 175 billion parameter model of OpenAI, which already was record breaking, the paper was very much like, oh, we built a model and look at what things it can do. So that was the OpenAI paper. This paper here is like the complete opposite. It basically says, oh, yeah, we do language model. But here is how we built the model, which is equally cool. So OpenAI basically just made everything bigger. And here they say to make everything even bigger, you need some tricks in how to build models. And then they've basically developed this entire framework to build these giant models. And this paper mainly describes that framework. And the actual task here, which is machine translation, is almost sort of a side thing in the paper. It's just a task to showcase what this system can do. So this is very much an engineering paper rather than that much than a machine learning paper. And that's how you have to look at it right here. That being said, the machine learning results are of course, quite impressive. If you look at this graph here, you have a quality gain. It's a difference in blur score. And this is a quality score for machine translation over the previous state of the art. So over their baseline, which, as you can see here, you have 37 billion weights, 150 billion weights, and 600 billion weights, which they only train. They train for, you know, 2000 and on 2048 TPUs for just four days. They stress this is very efficient because they just have to train it for four days on 2000 TPUs. Absolutely crazy. So let's have a look at what this paper does. If you enjoy this, if you enjoyed this at the end, consider, you know, sharing the video out if you like it and tell me what you think about this stuff in the comments. Alright, so we'll go through the abstract and then we'll go through highlighted sections of the paper because the paper is 23 pages long. So I won't be able to cover everything, just kind of give you the high level ideas and highlight a few things. Actually let's not go into the abstract. Let's go into these results first. So as you can see, they managed to continue the trend. The trend in NLP has always been, at least since, you know, transformers were invented, the bigger the better, like larger model, larger data, more compute means better performance. And this is sort of unbroken here, as you can see, if you increase the number of parameters in these models, you do get a very, very big gain in these blur score, though it sort of seems to be kind of a logarithmic scaling, like you have to keep doubling and doubling and doubling the number of weights, sort of like Moore's law in computation. You can see that at the same time, the training wall time is going down and the computational cost, the computational cost of these models, it doesn't scale quadratically like you would expect, it scales linearly. And that's the big difference here in how these authors scale their model, rather than how the open AI authors scale their model. So in a traditional transformer, it looks like this. So it has these blocks of attention. If you don't know what this is, I have a video called Attention is All You Need. I explain how the attention blocks in transformers work. So this is nothing different. These are just standard transformers. There is an encoder and a decoder. Everything works as you know. So you have these blocks, you have n blocks. These are the number of layers that you have. And in these blocks, you always have an attention layer, and then a feed forward layer that acts on the tokens. So without repeating too much what an attention mechanism does, basically, you have input tokens. So this is a sequence, it's technically a set processing unit, but we use it for sequences of text. So here you have six tokens, a sentence of maybe six words. And then you transform it with the attention layer by having this attention mechanism that routes information from tokens to from positions to other positions, maybe like this route is here, route this here. And then you have a feed forward network that is applied on a per token basis. So each of these tokens now goes through this feed forward network and is kind of transformed. So the embedding of that token is transformed by that feed forward network. Now every token does this. And it's always the same feed forward network. So this network here is the same as this network. Now usually, when we talk about scaling transformers, we talk about this part right here, we talk about the attention mechanism. And also we talk about this part, the number of layers. So you know, we talk about scaling the number of transformer layers, more layers, more layers, more layers. And if we want to scale the attention mechanism, what that basically means is we have we increase the context size of the text we can input. So transformers are very limited by the size of this context right here that they can take. Like the original transformer started with something like 512 tokens that they were able to take because this attention mechanism has quadratic complexity. This went up and the open AI GPT-3, I believe, had a context size of 2048 tokens, which if it scales quadratically, that's quite an achievement. And also it stacked the layers very, very deep. Now in this paper, they scale the transformers differently. They basically leave the context size and I believe that their context size is 1024. So significantly smaller than the open AI context size. And they don't scale the layers. So their largest transformer is 36 layers, whereas I believe GPT-3 was maybe correct me but I think it was like 90 or 100 layers or something like this, at least significantly larger than this. Instead, what they scale is this part right here, the feet forward layers. Now that might seem counterintuitive. But they basically, they basically say what if we didn't only have one feed forward network right here, but we had many, right? We don't always have the same. We have many, many feed forward networks, different ones that can do different things. So that's what they call experts. Each one of these feed forward layers is an expert. And then you have yet another routing mechanism, kind of like in attention, you have a routing mechanism that decides which tokens go where. Okay, so this token here, this token here, this token here, and the sort of the implication being that different tokens, different parts of the input you want to transform require a different kind of transformations here. And these different experts can sort of specialize in how they transform the input. Now their task here is going to be machine translation as a multitask setup. So what you'll have is you'll have all kinds of languages like French and German, what's the E, and maybe a lot of languages. I don't know any other languages. And you want to translate all of them to English and you want to do it using the same model. So these experts here, they might specialize in the individual languages. Like maybe you will have to handle a pronoun differently if it comes from German than if it comes from French. You want to do it with the same model at the same time. That means you maybe want to have the one expert specialize in German pronouns and one expert specialize in French pronouns. Also you can think of the experts as maybe one specializes in question words, it doesn't matter which language they're from. And the other one specializes in some sort of other kind of linguistic feature. In any case, this number of experts here is if you want to scale that up, then that becomes the bottleneck of the transformer. They go up to 2000, 2048 experts in parallel. So that doesn't fit into a single accelerator anymore. And that's why the entire system has to be sharded. And that's what they call G shard. So G shard, the main application here is going to be how can we build this giant model on many, many distributed computers where the attention mechanism isn't the problem. The attention mechanism we just distribute like we do data parallelism. The attention, it lives on all of the accelerators, it synchronizes and so on. But the experts here, there's only so this expert lives on machine A, this expert lives on machine B, this expert lives on machine C. And then we do a hard routing. So we don't do a soft routing like an attention, we do a hard routing where one token goes to one or at maximum two experts. So this is sent to these machines. And then after the machines, you kind of gather all the results back right here. So G shard is the system that enables this sharding of these experts. And the everything in between everything that is necessary, but it can also be applied to shard any computation. And that's why it's so cool. So here you see what what they do, they always they take these transformers. And they always consider a block of two transformer layers. So this is a block of two transformer layers, you can see there is twice the attention, and there's twice this feed forward. So in one point, this feed forward is just a regular everything, all the tokens go through the same network. So that's like a classic transformer. But here, you have a lot of these different experts and the tokens are routed to these experts. It's important that the tokens are hard routed, right? If the tokens were soft routed, you don't you don't gain anything, because every token has to go through every expert. But here, the tokens are hard routed to the expert, which means that you can if you if I have an input size of 1024 tokens, maybe only 10 go to this one, and maybe only 10 of that those go to this one. Now you also have a batch size, of course, I haven't actually looked at what the batch size here is, but you usually have quite a large batch size in these things like maybe a batch size of 1000 as well. So ultimately, what you'll end up is, you know, 1000 times 10 tokens going to the first expert and so on. But still, you can significantly parallelize this computation. Okay. So this this, if you use G chart, this is going to result in the following in the thing on the right, where you have two machines, this is machine one, and this is machine two, you can see that the machines will what happened here, or someone made the PowerPoint mistake. So you can see that the the attention, everything is shared between the machines. So this here and this here, these are synchronized, the weights are synchronized, right, you simply do a data sharing. But here, you can see that you have model parallelism, model parallel mixture of experts, where on the first machine, you have the first expert, and then you have e devices. And on the last one, you have the last expert. And then it's all routed out and routed in again. And then you can continue your transformer. And this is layer after layer. So what's the problem here, the problem is that an operation like this is going to come to incur significant sort of overhead in terms of communication, and so on, if you were to do it naively, and it's going to be a real pain to program this. And that's why G chart is made to do all of this automatically. And you don't, you don't incur much of a cost, because you distribute. So what's the difference to the old scaling? Why don't they just make transformers larger in number of layers? And that's because this this is, I guess, what opening into as well, if you make transformers simply larger in number of layers, sorry, if you make it transformers larger in the attention mechanism, it just won't fit into memory at some point. And you'll you'll have to share that somehow. And you can do this with G shard. If you scale it in number of layers, that incurs significant cost where you have to wait, because you have to forward propagate, and then you have to backward propagate in your training sequence. And if you have just too many layers, then a lot of the a lot of the frameworks get at their limit, where at some point they say, well, I still have to wait for the signal to come back in order to continue. And they explore this in this benchmark right here. You can see they say the largest model, the 600 billion parameter model that achieved the best translation quality was trained with 2000 TPU v3 cores for three days, a total cost of 22 TPU core years. In contrast, training all 100 bilingual baseline models would have required 29 core years. So the model here is faster than if you train them individually. But if you want to train a single transformer that is just very deep, and achieves reasonable performance, you have to invest a lot more. Our best quality dense single transformer model 2.3 billion parameters. So it's also significantly smaller. Achieving this was trained with G pipe, which is a previous framework. So G pipe is kind of a task runner that also distributes computation was trained with G pipe on 2048 TPU cores for six weeks or a total of 235 TPU core years. By the way, for if you if you have $1 per TPU hour, that'll only cause that'll only, I guess set you back about 2 million or so. It's easy peasy. Or even 200,000 just, you know, a tiny, tiny bit of of money. But you can see that this transformer model that is dense, which means that is a classic transformer where you stack the transformer layers, you stack them, you stack them, you stack them. It, in fact, it has 96 layers, their baseline 96 layer transformer model, that's sort of what opening I did, they just kept stacking the transformer layers. You get a model that has less parameters and trains for much longer. And its performance is only about this good. Whereas here, if you scale not into depth, but into width of these experts, and it's not dense, but it's shorted, which means it calculates this in a in a kind of sparsified way because it has this hard routing, you can scale up to a lot more parameters. So 600 billion parameters, over 200 times more parameters than the deep model, and you can get a much better performance. Okay, so this is what is different here, it scales into these experts rather than scaling into depth or, or size of the attention mechanism itself. All right, the question, I guess that you come up with if you're a machine learner is how do you back propagate if you route if here you route to these different experts, and you do a hard routing like here, how do you back propagate the signal because it seems like you need a soft routing. But this has been handled, in fact, these mixture of experts has been introduced previously, in a paper I think called outrageously large language models or something like this. And so they've introduced that, you know, it, it still works. So backprop still works through so basically you have a backprop path through here. And because you put a little bit of noise in this routing, every path gets explored a few times, and therefore you have enough backprop signal to make it work. It can it could technically fail, but they do observe generally that it does work if you do this kind of hard routing with a bit of noise. All right, so where do we go from here, as I said, this is an engineering paper, and it's a long engineering paper. So they, they set up their, they set up a lot of a lot of the details of engineering directly in the paper, which we're not used to in the machine learning world. They really detail how they shard things and so on, which is pretty cool. But I invite you to look at the paper yourself. If you really want to know what's going on right here. Suffice to say, they, as you can see right here, what they do is, this is the input right here, and then they have this weight matrix, which is a this routing, this is learned routing weights. Okay. So you have trainable weights that decide how to route the input, and that's dependent on the input. So you have a bunch of inputs that comes from the lower layer. And this matrix right here determines where to route them. Probably says, okay, the input is a vector like this. I know that must probably go to the expert number three. Okay. And you have a softmax across that. So it's a really, it's an assignment to, it's a soft assignment to the experts. So once you've done the soft assignment to the expert, you do a hard assignment by collecting the top two. For each token, you say you collect the top two experts, and you only send it to the top two experts and you ignore all else, which is not a lot right there. At times there are 2000 experts in the system. And yeah, you distribute and you have some noise. So with a random probability, you actually don't even send it to the second expert. You just leave it at the first one. And with some noise, you send it also to the second one. And I think that that noise is part of what if what makes the system work a bit. And then you also have this auxiliary loss right here that you add on top, which just makes sure that you distribute evenly. So this encourages the system to distribute the tokens evenly, because sorry, what it penalizes is a this here is the mean assignment to each expert. So it penalizes whenever the mean assignment is out of out of line, basically, so a distribution assignment to the expert or one expert gets a lot of tokens, because I don't know, it tends to be really good at something. So all the tokens are routed to it. And the other expert don't get a lot that's penalized. So you encourage the system to distribute tokens evenly between those experts. And then there are also like upper limits where you drop tokens and so on. They really build a system that is out for performance rather than machine learning correctness. So they demonstrate how to do this in in sort of code with their system. And the cool thing about their system is that you don't have to do much. What you'll have to do is just specify which tensors are sharded at along which dimensions and the system does the rest. So this is pretty cool. So this here is this mixture of experts, mixture of experts as you would write it in code. And they make use a lot of this Einstein, this Einstein some notation. If you don't know what the Einstein some notation is, it's a general notation to describe matrix or tensor multiplications. So a for example, if you were to multiply two matrices, you could have a string there, you describe it as a string and it comes from how Einstein wrote up the kind of tensor contractions in his work. So if you want to multiply two matrices, you can you could put the string a b b c goes to a c. So this and then you put two matrices right here. This will tell it, okay, I have a one matrix. I'm going to call the axis a and b, I have another matrix or tensor where I'm going to call the axis b and c. Now I have the resulting tensor, and I want the first axis to be a and the a is this one. And I want the last axis to be c and the c is this one. And b is nowhere b is not in the output, which means it should contract over b. So it should sum along b, sorry, it should multiply along b and then add such contract over b. So this here describes a regular matrix matrix multiplication. Now if I could do something else, I could do something like a just a element wise product, an element wise product would be something like this a b comma a b goes to a b, which means here, I have a in the first input, and here I have a again. And I'm so you already see that you can even though these are different tensors, you can call the axis the same, which means that they're going to somehow be multiplied together. Now if you leave it away here, it means that it's going to be contracted and therefore the axis no longer exists. But here we don't leave it away, which simply means that these axes are going to be multiplied together. And the same for b right here. So this describes an element wise. This describes an element wise product, you can go really funky with this. So this, this here would be a row wise dot product, where a is more it for all the a is it's element wise, but then over b, it's contracted. So you know, you can go, you can go wild with the Einstein some notation, you can describe a lot of things with it. So here is this algorithm to distribute the computation among these different experts. So you have the inputs and the weight matrix for the, they call this the gates function. That's the routing function to these experts. So what do we do? We first of all, we have these tensors, these, they have these grouping, these grouping dimension right here. So they come along to along groups, which in our case, we could maybe say these are batches or the batch dimension. So they come across groups, and there is the sequence length and there is this M right here. That's going to be the feature dimension, the M. And you can see the M is contracted. So the M is no longer here. So the gating function is going to route each input token right here to one of the experts for each thing in the group. So you can see, you can express this with an Einstein some notation. Then you have a top two gating, which selects the top two from each of the last, from each of the entries. And that gives you this dispatch mask and the sorry, and the weights that you have to use at the end to combine. You can use the dispatch mask in order to distribute the inputs. So you have reshaped inputs, and so on. So I'm not going to go through all of this right here, but you can express all of this in terms of the Einstein some notation. And you can express pretty much any sort of computation that is along the line. You can express the attention mechanism and so on. You can express the feed forward layers in terms of these Einstein some notations and the underlying the underlined dimensions here are the dimensions where we want to shard the computation. So here, because we have this G underlined, that means that we are interested in sharding the computation along this axis. So this, I said, this is the batch dimension. This is your classic data parallelism, which means that the first machine gets the first couple of data points, the second machine gets the second couple of data points, and so on. And you can see in the weight matrix, there is no sharding, which means that the weight matrix lives on every machine as a copy of one another. This is different from from here, where you can see that what we're now going to do is here it's still sharded according to the batch, but we now are going to shard this according to the different experts. So we're going to route whatever the inputs are in to these experts. And then we're going to execute the computations on the experts. So this is now sharded according to the experts. And at the end, right here, you can see this is still sharded according to the experts. We're going to put it back together. And now it's sharded according to the groups again. That's what we said, we have the input right here, the inputs, and the inputs are maybe distributed according to the according to machines, right, we have these go through the first machine, these the second, these the third, and so on. This is your classic data parallelism. But then we have all of these experts. And now all of a sudden, we're going to route these things to the individual experts. And we're going to execute the computation in parallel on the experts. And then after that, we're going to put back together from wherever we got them now have to. So this goes here again. And so this is just the reverse of what we did before. So right, like that. So you get all of the outputs again. I hope you kind of can imagine how this happens. So the first difference is, is that's sharded according to a different dimension. And the second difference is, is that when we shard in data parallelism, we execute the same computation on all the machines, which means that we have the same weight matrix. If we do x times w in a feet forward layer, and we shard this thing here in data parallelism, what we do is we send the x to different machines, we split the x, we send it to different machines, this is x1, next to x3, x4. But we always multiply it with the same weight matrix that weight matrix lives on all of the machines and is regularly synchronized, it's kept synchronous in some way. Whereas if we shard x to the experts, then the experts have individual functions. So the expert one is different from the expert two is different from the expert three, and so on. Which means that before it wasn't important where x was routed, because we would execute the same computation. So we can just, you know, sharded according to you know, the first 10 go there, the next 10 go there. But here, it's not crucially important where they are routed to to which expert. And that's why we learn the function that is going to route them. So this is learned, this is these first line here, these are the weights that we learn to route, then we route right here. And we calculate your your, we calculate the feet forward layers on the expert, you see that this wi and wo, they are the weight matrices of the feet forward layer, the feet forward layers are, you have your input, you multiply it by wi, you have a ReLU, ReLU, and then you multiply it by wo. So it's kind of a two layer feet forward network. So this two layer feet forward network, as you can see, this is sharded according to the experts. And then, and the important part is, of course, that here, the weight is also sharded according to the experts. And that's what makes each expert different. And then it's combined again down here. So I hope you kind of get the idea of what this algorithm does. But the fact that we shard according to these experts is in fact different than your regular sharding where you shard the data like the batch, the batches, but keep the model in parallel, keep the model synchronized. With their system right now, this is how easy this is. So before we simply stated our algorithm in Einstein's sumnotations, there is no way to underline code and that magically happened something that was simply for us to visualize. Now we want to apply their system in order to make this actually sharded. And with the Gshard system, and as I said, I don't know if the code is out or it will be out, but with the Gshard system, this is basically all that you have to do. So you have these functions, they're called split and replicate. What replicate does is it takes that weight tensor and it replicates it on all the machines and that keeps it synchronized. This is a computation where we simply want to shard out the different to the different machines but keep it synchronized. And you can see if you do this, this is the operation, then the system knows, ah, this here is replicated across the machines. So that means I'm going to distribute the data points according to this G dimension, according to the batch dimension and multiply it with this matrix according to this Einstein sum notation string on all of the machines. And I'm going to keep this tensor in sync. Okay, so the system knows as opposed to that you have you have the split tensor right here. So the split, what it does is it splits a computation here the dispatch expert inputs, it splits it according to a axis index onto D different machines or into D different parts. So you see here you calculate the how you should do the routing and the resulting tensors first dimension is this E dimension. And then you say that should be split, you know, according to this first dimension onto D different places and these D different places are now separate. They don't have the they don't have to be kept in sync. Everyone has their own weights. And now when you do this, you know, according to this dimension, you can see because we know Einstein sum notation now, you can see this E appears here, here and here. So this operation is going to be applied element wise, that means independent of each other in the direction of this dimension, the system understands that since this tensor is sharded according to that dimension, I have to execute this on each of these entries in separate with on each expert having their own weight matrix right here. I hope this is a bit clear that their system makes it super easy. You can basically do two things. You can say this thing here is my classic parallelism where I want to keep it in sync. And this thing here is where I want to split up and do different computation on the different parts. And then they have also a general function that is more powerful. Yeah, they and they you can auto partition and whatnot. So they have a a a they have this we implemented the partitioner in the XLA compiler, which means that anything that can translate to XLA is a target for the system. And that's, you know, TensorFlow and pytorch can do this. So technically, this can come to any of those systems. But of course, who has their 2000 TPUs lying around to make use of this? But no, I'm kidding. I mean, this, I they here use it for transformers. And I am very excited to to see what people can come up with for the system, I believe a system like this where it's super easy to to shard. And they have some, you know, they talk about, okay, we do the single machine compiler. So the compiler is also fast and so on. I don't even want to go into this. But this is very well engineered, it seems. And they, they, they basically implement this for all of the operators. So I'm very excited to see what people can come up with outside of the traditional applications. I think there can be new types of models developed simply because we have a system like this that makes it easier. So yeah, I'm excited. So here, they show a bit how this works on the example of this Einstein, some notation. So here, we want to do this thing here, which if you remember, this is the operation where we want to route the input to these experts. So we want to start with something that is sharded according to the batch dimension. That means that we, you know, we have different different parts of the batch on different machines. And we want to route this and finally end up with something that is sharded on the different experts. So this is what the system does is first you have these here are the different shards, right? You want to multiply this, as you can see, this and this right here means that these this routing table is also sharded according to the same machines. So you have the zero is all on the same machine, the one is all on the same machine, and so on. So what you want to do is you want to contract is there you want to contract according to this s dimension, right, which we have we have omitted right here. And if you multiply that, sorry, okay, we omit the s so this is not much of a this is not much of a graphic right here. But then they have this reshard operation where they do and you don't have to worry about this. So from here to here, there is this reshard operation that just shards it according to the according to E. Yep. I find this to be a bit more a bit more insightful. So if you have something like this, this which is a regular matrix multiplication, right? And you want to contract along B, this is exactly the example we had before. So here is a situation where our tensor is sharded according to the B dimension and this tensor is also sharded according to the B dimension. You want to do a matrix multiplication of the whole tensor. So what can you do, you're supposed to multiply these two matrices, but they are sharded on different machines. If you consider what you actually have to do is you have to multiply each row here with each column here. And that in an element wise fashion. So that distributes according to you have to multiply this by this plus this by this plus this by this plus the red by the red. So you can simply multiply the zero tensors together, the one tensors together, the two tensors together and the three tensors together. Each one will give you a full matrix and then you can simply add all of them in order to get your full results. This is illustrated down here. So what machine one does, it simply multiplies its shard by its own shard of the second matrix, which will give it this thing here. And by the nature of how matrix multiplication is constructed, you can simply do an all reduce, which means you reduce you sum across all of the machines, and that will give you the full result. So this is a this is a an example of how this works. This is, you know, pretty simple. And I believe you may have seen something like this already when you were looking at just parallelizing matrix multiplication, and so on. So this system handles this transparently, right? If you're shorted like this, this is what the system will do. However, if you are shorted differently, the system will act differently. So here is a system you want to do the same matrix multiplication, but the first tensor happens to be shorted according to the A dimension, the second tensor happens to be shorted according to the C dimension. And you want to end up with something that's shorted to the C dimension. Now we have an additional constraint here that here you can see, we kind of assume that this full thing here fits into memory, mainly because we want to obtain the full result you see here, a and c should not be shorted. So we assume that we can keep that in memory. But here we want the final result to be shorted according to C, which imposes the additional constraint that it might be that the full matrix never fits into memory. So how are we going to calculate all of that? We can't do the same trick anymore. Now this G short system apparently realizes itself when something is out of memory, and it can do a smart move around being out of memory using a loop, which basically means that it will compute entry by entry or block by block. So these are the matrices we have to multiply. And you can see that if I want to do multiply this by this, that's fine, I can do this on one machine. And that will give me the block up here. But if I want the block up here, I have to multiply this by this, which is across two different machines. So what this system does is it's going into a while loop because it realizes there's not enough memory. And it kind of sends around these different slices to the different parts, each time computing a little piece. So here, first, we do this by this, this is fine. But then we grab ourselves from the we we grab ourselves this one here, calculate the next little piece up here, and then we grab ourselves the number two, calculate the piece here. And then so this is from zero, this is from two, the one we already had, and then we grab ourselves piece three, and multiply that until here, until we have this final slice that we want. Okay, so this goes in a while loop in multiple rounds, the system gets knows itself when it has to do this, and when it can calculate the full thing at once because it fits into in memory. It's even smarter than that, and that it can do these halo exchanges. So if you have to do something like this, a convolution, now in a convolution, what you'll do if you think of a think of an image, and you want to do a convolution on it, but the image happens to be sharded. Let's say the image is so large, it's sharded across nine different machines like this. Now, if you want to do a convolution, that's pretty cool, you know, here, here, here, but here, all of a sudden, your convolution is across two different machines. So this system, G shard will adapt automatically, and do these halo exchanges where it kind of sends around from this machine, it'll send something to this machine such that it can do the convolution in that step, and vice versa. And then this can be padded accordingly, as you can see. This is I think this is this is this was like super ugly to implement. If you just imagine that for each of these operations, you have to think about, okay, how can you express this with these MPI primitives like dynamic slice and collective permute, and so on. It's just an absolute nightmare. And I'm very happy that other people have done this, and I will probably just get to use it. So there is a lot more to this system than I've just explained, I just try to give you a flavor of what building a system like this means and how easy it is to use it like this. In order to implement all of this mixture of experts things, you simply go from this, which is one single machine implementation, how you would write it to this, which is now the same, it's almost the same code. But this now you can run on however many machines, and if you compile it with the system, it will do what you expect it to do in this shorted way. Completely crazy. Okay, so they apply this to massively multilingual massive machine translation. So two things, it's massively multilingual, and it's massive machine, which means, I guess a lot of machines. And the reason here is twofold. So what they say is, we have massively multilingual translation. Why don't they just look at single machine translation? And it has a very specific reason. Namely, if you have massively multilingual translation, which means that you have a lot of different languages, and you all have to translate them, ideally, to all the other languages or, you know, every language pair. But in this case, they only look at all the languages to English. I don't exactly know why, but I guess there must be some kind of reason. If you do this, then you can make use of a thing where there are languages that you just don't have much data on. Like I don't know, Basque or something like this. There's not that many people speaking Basque or Swiss German. There's not even a written form, a standard written form of Swiss German. So you just don't have as many resources. And for other languages, you have giant amounts of resources. And what you can make use of is this phenomenon called positive language transfer, where it happens that, for example, Swiss German is very close to German. Now, they can't understand us, which is a giant advantage for us, but still it shares a lot of similarities with German. So if you learn a lot about German, you can sort of transfer learn to Swiss German pretty easily. So if you have a system that does German and Swiss German at the same time, you can perform better on both languages because the Swiss German part of your model, the part of your model that does Swiss German, profits from the German inputs as well. Now don't understand me wrong. There is not an individual part of your model that for each language, it's all done at the same time. But still you can imagine that, you know, some of these things will specialize in some of the languages. But the hope is that if you have German and Swiss German in the same training set, that if the model realizes what a question construct is in German, it will be able to apply that also to Swiss German with some minor modification. So there is a benefit of having these many languages, especially for the low resource languages. Okay. So as the number of languages, sorry, as the number of language pairs to be modeled within a single translation model increases, positive language transfer starts to deliver large gains for low resource languages. Given the number of languages considered, which I believe is a hundred here, M4 has a clear advantage on improving the low resource task. On the contrary, for high resource languages, the increased number of tasks limit per task capacity within the model, resulting in lower translation quality compared to a models, to a models trained on a single language pair. This capacity bottleneck for high resource languages can be relaxed by increasing the model size to massive scale in order to satisfy the need for additional capacity. So basically they're saying, if we train all of these languages together, that will help a lot for these low resource languages, but it might hurt the high resource languages because now we would have enough data technically to train a French to English model on this giant model. We could train that. And now that we have all these other languages in there, it just hurts us because we don't have enough parameters. And we can solve this, of course, by simply adding more parameters. So that's the solution. Add more parameters and you increase the capacity of the model and you still get the benefits of the positive language transfer. So their investigations is going to be into how much can we scale this? And is there like a sweet spot where because if you, if you increase the parameters too much, you counteract this positive language transfer again. So since, you know, since Swiss German and German can sort of benefit from each other. However, if we have too many parameters, so, and then we end up having all of these experts right here and the tokens are always routed to these experts and it always happens that all the Swiss German tokens are always routed to this expert and all the German tokens are always routed to that expert. There will be no sharing of weights. There will be this positive language transfer will not happen because we have too much capacity. So the goal is to find a sweet spot between positive language transfer and this capacity bottleneck. They do use an in-house data set, which we don't have access to, but they say the training corpus mined from the web contains parallel documents for a hundred languages to and from English adding up to a total of 25 billion training examples. However, they only use from 100 languages to English. This result in approximately 13 billion training examples to be used for model training. So that's a lot. It's a lot of data, especially for translation. It's kind of a noisy translation because it's mined from the web, but still it's a lot of data. They have baselines. So the baselines are first of all, in order to form our baselines, we trained separate bilingual neural machine translation models for each language pair. So that means a single model for each language to English, depending on the available training data per language. And then they also have a baseline where they try open AI style to build as deep as single transformer as possible. And by that, they mean we also include a variant of a dense 96 layer transformer encoder decoder network trained with G pipe pipeline parallelism on the same data set as another baseline. So the difference again here is that this 96 layer is a dense transformer, which means that all of the tokens go through the same computation and we don't shard the computation out to these experts, right? We do shard according to the batch, but all of them go through the same parameters. And that means we can we can only scale up the number of layers and that severely limits the that severely limits the computational efficiency. Even if we have, you know, your pipeline parallelism and so on that hurts. They say training to convergence took over six weeks on 2000 TPU course. That's crazy. But I guess, yeah, you know, I was saying earlier that that I always thought we were happy. I always thought we were happy in machine learning because kind of the hip science fields being biology, like genetics and machine learning. I was thought like, oh, but these biology people, they always need like million dollar grants from government to run their experiments. And we can just sit down with a laptop. This time is over. If you start a PhD now start applying for money to get TPUs. Yeah. Okay. In any case, here you can see what this does. So they compare a bunch of models right here. So this T, this is this big dense transformer that's going to be one of our baselines and the other baseline here is going to be the zero axis. The zero axis means this is the single model for that language pair. So only so for each language, they trained one model only on data from that language. And that's going to be the worst thing here because this multi language translation in one model will generally help you if you have enough parameters. You can see all the models here have enough parameters such that the difference here, this is difference in blue is positive including this baseline model right here. So the baseline model, as you can see, has 2.3 billion parameters, even though it takes that much longer to train and that's, as we said, a function of the fact that it's dense and deep, so that hurts in training efficiency. And then you have these mixture of expert models. They always consider two things. They consider different numbers of experts. You can see it goes from 128 to 2048 experts. And they consider a number, different number of layers from 12 layers to 36 layers, 36 layers still being way smaller than the 96 layer transformer here. And that's the reason why it trains faster. So it doesn't train faster. So the reason it trains faster is because it has less layers. And then the reason it has more parameters is because it has a lot of these experts. And the art here is to constrain how much these more experts hurt you. So you know, you could run into the same problem where if you scale up the experts, in fact, you do, it doesn't fit into memory anymore. And it's going to hurt you a lot in training efficiency, kind of like if you increase the number of layers. But the G shard system prevents that it lets you up the number of experts without incurring the cost. That being said, it does not let you up the number of layers, you're going to incur the same cost if you up the number of layers as you have with the dense transformers. So does this help? It helps a lot. As you can see right here, there's a general trend upwards. And what's the x axis, the x axis is low resource languages. So you can see that as we as we go to lower and lower resource languages, this multi task training, this multilingual translation improves significantly over the baseline where we only train a system for that language specifically. And these 10k examples, it's it's it's quite a bit, but it's not that much, especially since it's noisy data. So this is specifically good for low resource languages. But you can see also the high resource languages here benefit from the multilingual translation. And that's a function of the fact that we have, you know, large enough models. In fact, you can see the larger the models, the more the difference in blue is, and there's not really an end in sight. And they also see it say that they haven't seen convergence in training. So you can technically train this forever. Yeah, you can also see that the the lowest mixture of experts right here is almost on par with their big dense transformer that took so much longer to train. Right. So this lowest model right here, I believe it took I don't want to go back, but it took it took hours or so or few hours to train, whereas this 96 layer dense transformer took these six weeks to train. So has to be said, the number of TPUs is not to be neglected, but if you're Google, you know, you just have them laying around. What's also interesting here, and you can start seeing this two things. First of all, you can see that the difference between here in between the dense transformer and this baseline model is very low for high resource languages, but gets larger for low resource languages. This is an indication that the dense transformer, it does more to share parameters between the languages because it shares parameters between all the things because all the tokens go through the same computation. So it is going to be a bit better in low resource languages, but still the general trend upwards holds even for the mixture of experts. The second thing is that you see there is a crossover here in these in these big in these biggest models. And what are the big models? One, the blue one is the one with 2048 experts and the green one is the one with 500 experts. They're both as deep models. But all of a sudden, over here for the high resource languages, it's still true that if you up the number of parameters, you get a benefit. So up the number of experts as well, you get a benefit. But over here for the low resource languages, it's it you see, it actually hurts you to up the number of experts. And that's the phenomenon exactly we talked about before. If you have too many of these experts, and you do a hard routing, that means all the tokens go a different way. And that means you don't get any sharing benefit from the multilingual translation. And they investigate a lot. And they basically claim that their sweet spot of expert in their particular task appears to be somewhere in between these 2000 and this 500 expert number, where you can see it doesn't always help you to scale up the model. So I have to say maybe the transformers, maybe they need a ResNet moment. So I believe in computer vision, it was sort of the same problem that we try to build deeper models and why like, okay, this, this is more width. But yeah, I think there might be some breakthrough on the horizon where someone just figures out how to train these giant models, even more giant transformer models with deeper layers. And then there's a new era of transformers. However, this is not that effect. I'm sorry, I said this at the wrong place. This is not that effect. This is to show that in this case, we do benefit for the high resource languages because we increase capacity. But for the low resource languages, we suffer if we up the number of experts too much, because they don't share any parameters anymore between the languages or between the different parts. Like it's not a necessity that the different languages are going to be routed to different experts. But it's probably going to happen, right? There's no hard coded thing that says if it's this language, it needs to go there. It just probably is going to happen this way because the different languages are going to be needed to be treated differently and therefore the system learns to route first and foremost, those two different experts. Here you can see the model sizes, including this 60 layer models model with 2000 experts that they didn't manage to train. They said they had numerical instability, but that had one trillion parameters. And I'm pretty sure they're, they're cool. They must be quite mad about this, right? Like you have the trillion parameters, even though it's not that much bigger than the 600 billion, that the trillion, it would be cool to write a paper like a trillion parameter model. But for now they are at the 600 billion mark and they simply want to tell you that they have actually compiled a model that's that big, just didn't manage to train it. And yeah, that's here. Here is where I wanted to say that maybe we're waiting for the ResNet moment where all of a sudden someone figures something out that makes the training of basically infinitely deep transformers possible. Like we made the training for almost infinitely deep CNNs possible with ResNet. Okay, so they conclude this and so they, that's the investigation of what the number of experts and so on gives you. And here is a bit of a different investigation where they more care about training efficiency. So they ask themselves, how many billion tokens of input do we need to reach a given cross entropy? So here, the more tokens you need, the lower your efficiency is, right? You can see that the general trend is the following. If you up the number of layers, you get more efficient, you can see and just look at this column for now, this point seven column, you can see it already pretty clearly. So here you go from 12 layers to 36, you gain efficiency, here you gain here you gain pretty predictable. If you up the number of layers, you need to see fewer tokens to get to the same cross entropy. And in fact, you can get to a lower cross entropy altogether at the end. We've known this for language models already. The other effect is of course, what happens if we go not deeper, but wider, if we increase these number of experts, if we increase this sparse computation. So here you can see, let's just look at the 12 layers for now. Let's look at all the rows where there's 12 layers. So here you get a significant advantage by upping the number of experts from 100 to 500. But then you hurt upping the number of experts to 2000, right? So that's that's sort of your you're hurting efficiency by upping the number of experts too much. And the same if we look at the 36 layer, so you gain massive efficiency by upping the number of experts, but you lose that a fish part of that efficiency again, by increasing it even more. Now we saw that the this model is still the best model, but it's not as efficient as that model. And that gives you another indication that there is sort of a sweet spot between these two things between the positive transfer and the bottleneck capacity that appears to be somewhere in between right here. So that's pretty interesting. Because we know about depth that you can basically up and up and up and get more efficient, but with not that much. Yeah, the largest model can be trained in under four days to achieving the best quality. Yes, yes, yes, but this is just a yeah. So here, oh, you can see the batch size in in tokens is quite, quite a bit. So yeah, if you have a 1000, if you have a context window of 1000, that means the batch size here was about 4000. So as as expected. Yeah, this is just easy peasy 22 TPU core years. I've seen someone on Twitter saying this, this is the new measure for computer. It's no longer like flops. It's TPU core years. Just mad, mad. And yeah. So 42 days to train this thing right here. Crazy, crazy, crazy. All right. They also have a number of investigations in other parts of efficiency, like per device memory consumption. You can see here that as you up the as you up the number of experts, you can see here, here, here, your weights don't go up because as you up the number of experts, you can just up the number of machines and the per machine weight usage will be the same, right? Because the experts are independent of each other, each one has their own weight matrix. So you can just add machines and you keep your weight requirements the same. However, if you go deeper, then your weights increase because you're now deeper, you have more layers. You have your so also your transformer weights will be higher and so on. So you go deeper right here. You see 3660 layers, your memory consumption increases for the weight. And also, this is the other big part in transformers, right? The activations that you have to save, because as we said, if you have a transformer and I have layer, layer, layer, layer, I basically have to keep around each of these signals in order to do back propagation. And that's why also the activation here increases as I go deeper. Now you can see percentually, it decreases again here. So what's happening? Technically, you don't have to keep these things around. You can also once the signal comes back, you can recompute them from the beginning or from an intermediate point. Now this increases computation, but saves the need to store the activations. And apparently G shard, yet another thing it does is it will recompute as necessary the activations if it realizes that you don't have enough memory to store them. So all of this is pretty crazy, honestly. And they look at where the different computations go. And I don't want to go into this. And they have these micro benchmarks where they really show that the increase in complexity is really according to square root of n, because that's how long it takes to distribute along these actors, sorry, along these experts. There's a lot to this paper. And there's no time to go through all of it. I think this video is already way too long. I hope I have given you an impression of what's possible with this system. And as I said, I'm excited what people can come up with. Just to say that in the appendix here, they detail that they have done this for all the operations in XLA. So for example, convolution, this is so ugly, how you have to implement the convolution because you have to padding must be correct across these expert across the the sharded machine. So there are no experts anymore. This is just G shard, the padding has to be correct. The strides have to be correct. Data needs to be exchanged according to the machines, the window size needs to be correct, blah, blah, blah. So just thank you for doing this and not having to do it myself. Yeah, I'm excited as soon as as the codes out, if I get a hold of it, I'll you know, link it or you'll find it once it's out. If it's already out, I'm just too dumb to see it. I enjoyed reading this. It's different than a machine learning paper. It kind of shows you what goes into engineering a system like this, and how easy it can be if it's engineered well to then apply it. I think this is going to be extremely helpful to the community. And with that said, 23 pages later, see you next time. Bye bye.
 OpenAI has a 175 billion parameter model.
 You thought that was large?
 That's cute.
 Check out Google's 600 billion parameter model, 600 billion floating point numbers doing things
 at the same time.
 This has absolutely become a body part measuring competitions between companies.
 Google be like, oh, GPT-3.
 I spit on you.
 I spit on you and your little tiny 175 billion.
 OK.
 Let's stop kidding.
 This is a giant model that Google has trained right here.
 The paper we're going to look at today is called G-shard Scaling Giant Models with Conditional
 Computation and Automatic Sharding by Dmitri Lepikhin et al. of Google.
 And this paper basically tells the story of how they built this 600 billion parameter
 model, how actually they attempted to build a model that had a trillion parameters but
 just didn't manage to quite train it.
 And this is all using this system called G-shard.
 So I haven't actually seen the code out for G-shard yet, but I'm going to maybe assume
 that this is something that they're going to release at some point.
 Who knows?
 Or maybe I just haven't seen it yet.
 So this is basically describing a system on how to train these giant models.
 So if you have watched my video on GPT-3, which, of course, was this 175 billion parameter
 model of OpenAI, which already was record breaking, the paper was very much like, oh,
 we built a model and look at what things it can do.
 So that was the OpenAI paper.
 This paper here is like the complete opposite.
 It basically says, oh, yeah, we do language model.
 But here is how we built the model, which is equally cool.
 So OpenAI basically just made everything bigger.
 And here they say to make everything even bigger, you need some tricks in how to build
 models.
 And then they've basically developed this entire framework to build these giant models.
 And this paper mainly describes that framework.
 And the actual task here, which is machine translation, is almost sort of a side thing
 in the paper.
 It's just a task to showcase what this system can do.
 So this is very much an engineering paper rather than that much than a machine learning
 paper.
 And that's how you have to look at it right here.
 That being said, the machine learning results are of course, quite impressive.
 If you look at this graph here, you have a quality gain.
 It's a difference in blur score.
 And this is a quality score for machine translation over the previous state of the art.
 So over their baseline, which, as you can see here, you have 37 billion weights, 150
 billion weights, and 600 billion weights, which they only train.
 They train for, you know, 2000 and on 2048 TPUs for just four days.
 They stress this is very efficient because they just have to train it for four days on
 2000 TPUs.
 Absolutely crazy.
 So let's have a look at what this paper does.
 If you enjoy this, if you enjoyed this at the end, consider, you know, sharing the video
 out if you like it and tell me what you think about this stuff in the comments.
 Alright, so we'll go through the abstract and then we'll go through highlighted sections
 of the paper because the paper is 23 pages long.
 So I won't be able to cover everything, just kind of give you the high level ideas and
 highlight a few things.
 Actually let's not go into the abstract.
 Let's go into these results first.
 So as you can see, they managed to continue the trend.
 The trend in NLP has always been, at least since, you know, transformers were invented,
 the bigger the better, like larger model, larger data, more compute means better performance.
 And this is sort of unbroken here, as you can see, if you increase the number of parameters
 in these models, you do get a very, very big gain in these blur score, though it sort of
 seems to be kind of a logarithmic scaling, like you have to keep doubling and doubling
 and doubling the number of weights, sort of like Moore's law in computation.
 You can see that at the same time, the training wall time is going down and the computational
 cost, the computational cost of these models, it doesn't scale quadratically like you would
 expect, it scales linearly.
 And that's the big difference here in how these authors scale their model, rather than
 how the open AI authors scale their model.
 So in a traditional transformer, it looks like this.
 So it has these blocks of attention.
 If you don't know what this is, I have a video called Attention is All You Need.
 I explain how the attention blocks in transformers work.
 So this is nothing different.
 These are just standard transformers.
 There is an encoder and a decoder.
 Everything works as you know.
 So you have these blocks, you have n blocks.
 These are the number of layers that you have.
 And in these blocks, you always have an attention layer, and then a feed forward layer that
 acts on the tokens.
 So without repeating too much what an attention mechanism does, basically, you have input
 tokens.
 So this is a sequence, it's technically a set processing unit, but we use it for sequences
 of text.
 So here you have six tokens, a sentence of maybe six words.
 And then you transform it with the attention layer by having this attention mechanism that
 routes information from tokens to from positions to other positions, maybe like this route
 is here, route this here.
 And then you have a feed forward network that is applied on a per token basis.
 So each of these tokens now goes through this feed forward network and is kind of transformed.
 So the embedding of that token is transformed by that feed forward network.
 Now every token does this.
 And it's always the same feed forward network.
 So this network here is the same as this network.
 Now usually, when we talk about scaling transformers, we talk about this part right here, we talk
 about the attention mechanism.
 And also we talk about this part, the number of layers.
 So you know, we talk about scaling the number of transformer layers, more layers, more layers,
 more layers.
 And if we want to scale the attention mechanism, what that basically means is we have we increase
 the context size of the text we can input.
 So transformers are very limited by the size of this context right here that they can take.
 Like the original transformer started with something like 512 tokens that they were able
 to take because this attention mechanism has quadratic complexity.
 This went up and the open AI GPT-3, I believe, had a context size of 2048 tokens, which if
 it scales quadratically, that's quite an achievement.
 And also it stacked the layers very, very deep.
 Now in this paper, they scale the transformers differently.
 They basically leave the context size and I believe that their context size is 1024.
 So significantly smaller than the open AI context size.
 And they don't scale the layers.
 So their largest transformer is 36 layers, whereas I believe GPT-3 was maybe correct
 me but I think it was like 90 or 100 layers or something like this, at least significantly
 larger than this.
 Instead, what they scale is this part right here, the feet forward layers.
 Now that might seem counterintuitive.
 But they basically, they basically say what if we didn't only have one feed forward network
 right here, but we had many, right?
 We don't always have the same.
 We have many, many feed forward networks, different ones that can do different things.
 So that's what they call experts.
 Each one of these feed forward layers is an expert.
 And then you have yet another routing mechanism, kind of like in attention, you have a routing
 mechanism that decides which tokens go where.
 Okay, so this token here, this token here, this token here, and the sort of the implication
 being that different tokens, different parts of the input you want to transform require
 a different kind of transformations here.
 And these different experts can sort of specialize in how they transform the input.
 Now their task here is going to be machine translation as a multitask setup.
 So what you'll have is you'll have all kinds of languages like French and German, what's
 the E, and maybe a lot of languages.
 I don't know any other languages.
 And you want to translate all of them to English and you want to do it using the same model.
 So these experts here, they might specialize in the individual languages.
 Like maybe you will have to handle a pronoun differently if it comes from German than if
 it comes from French.
 You want to do it with the same model at the same time.
 That means you maybe want to have the one expert specialize in German pronouns and one
 expert specialize in French pronouns.
 Also you can think of the experts as maybe one specializes in question words, it doesn't
 matter which language they're from.
 And the other one specializes in some sort of other kind of linguistic feature.
 In any case, this number of experts here is if you want to scale that up, then that becomes
 the bottleneck of the transformer.
 They go up to 2000, 2048 experts in parallel.
 So that doesn't fit into a single accelerator anymore.
 And that's why the entire system has to be sharded.
 And that's what they call G shard.
 So G shard, the main application here is going to be how can we build this giant model on
 many, many distributed computers where the attention mechanism isn't the problem.
 The attention mechanism we just distribute like we do data parallelism.
 The attention, it lives on all of the accelerators, it synchronizes and so on.
 But the experts here, there's only so this expert lives on machine A, this expert lives
 on machine B, this expert lives on machine C. And then we do a hard routing.
 So we don't do a soft routing like an attention, we do a hard routing where one token goes
 to one or at maximum two experts.
 So this is sent to these machines.
 And then after the machines, you kind of gather all the results back right here.
 So G shard is the system that enables this sharding of these experts.
 And the everything in between everything that is necessary, but it can also be applied to
 shard any computation.
 And that's why it's so cool.
 So here you see what what they do, they always they take these transformers.
 And they always consider a block of two transformer layers.
 So this is a block of two transformer layers, you can see there is twice the attention,
 and there's twice this feed forward.
 So in one point, this feed forward is just a regular everything, all the tokens go through
 the same network.
 So that's like a classic transformer.
 But here, you have a lot of these different experts and the tokens are routed to these
 experts.
 It's important that the tokens are hard routed, right?
 If the tokens were soft routed, you don't you don't gain anything, because every token
 has to go through every expert.
 But here, the tokens are hard routed to the expert, which means that you can if you if
 I have an input size of 1024 tokens, maybe only 10 go to this one, and maybe only 10
 of that those go to this one.
 Now you also have a batch size, of course, I haven't actually looked at what the batch
 size here is, but you usually have quite a large batch size in these things like maybe
 a batch size of 1000 as well.
 So ultimately, what you'll end up is, you know, 1000 times 10 tokens going to the first
 expert and so on.
 But still, you can significantly parallelize this computation.
 Okay.
 So this this, if you use G chart, this is going to result in the following in the thing
 on the right, where you have two machines, this is machine one, and this is machine two,
 you can see that the machines will what happened here, or someone made the PowerPoint mistake.
 So you can see that the the attention, everything is shared between the machines.
 So this here and this here, these are synchronized, the weights are synchronized, right, you simply
 do a data sharing.
 But here, you can see that you have model parallelism, model parallel mixture of experts,
 where on the first machine, you have the first expert, and then you have e devices.
 And on the last one, you have the last expert.
 And then it's all routed out and routed in again.
 And then you can continue your transformer.
 And this is layer after layer.
 So what's the problem here, the problem is that an operation like this is going to come
 to incur significant sort of overhead in terms of communication, and so on, if you were to
 do it naively, and it's going to be a real pain to program this.
 And that's why G chart is made to do all of this automatically.
 And you don't, you don't incur much of a cost, because you distribute.
 So what's the difference to the old scaling?
 Why don't they just make transformers larger in number of layers?
 And that's because this this is, I guess, what opening into as well, if you make transformers
 simply larger in number of layers, sorry, if you make it transformers larger in the
 attention mechanism, it just won't fit into memory at some point.
 And you'll you'll have to share that somehow.
 And you can do this with G shard.
 If you scale it in number of layers, that incurs significant cost where you have to
 wait, because you have to forward propagate, and then you have to backward propagate in
 your training sequence.
 And if you have just too many layers, then a lot of the a lot of the frameworks get at
 their limit, where at some point they say, well, I still have to wait for the signal
 to come back in order to continue.
 And they explore this in this benchmark right here.
 You can see they say the largest model, the 600 billion parameter model that achieved
 the best translation quality was trained with 2000 TPU v3 cores for three days, a total
 cost of 22 TPU core years.
 In contrast, training all 100 bilingual baseline models would have required 29 core years.
 So the model here is faster than if you train them individually.
 But if you want to train a single transformer that is just very deep, and achieves reasonable
 performance, you have to invest a lot more.
 Our best quality dense single transformer model 2.3 billion parameters.
 So it's also significantly smaller.
 Achieving this was trained with G pipe, which is a previous framework.
 So G pipe is kind of a task runner that also distributes computation was trained with G
 pipe on 2048 TPU cores for six weeks or a total of 235 TPU core years.
 By the way, for if you if you have $1 per TPU hour, that'll only cause that'll only,
 I guess set you back about 2 million or so.
 It's easy peasy.
 Or even 200,000 just, you know, a tiny, tiny bit of of money.
 But you can see that this transformer model that is dense, which means that is a classic
 transformer where you stack the transformer layers, you stack them, you stack them, you
 stack them.
 It, in fact, it has 96 layers, their baseline 96 layer transformer model, that's sort of
 what opening I did, they just kept stacking the transformer layers.
 You get a model that has less parameters and trains for much longer.
 And its performance is only about this good.
 Whereas here, if you scale not into depth, but into width of these experts, and it's
 not dense, but it's shorted, which means it calculates this in a in a kind of sparsified
 way because it has this hard routing, you can scale up to a lot more parameters.
 So 600 billion parameters, over 200 times more parameters than the deep model, and you
 can get a much better performance.
 Okay, so this is what is different here, it scales into these experts rather than scaling
 into depth or, or size of the attention mechanism itself.
 All right, the question, I guess that you come up with if you're a machine learner is
 how do you back propagate if you route if here you route to these different experts,
 and you do a hard routing like here, how do you back propagate the signal because it seems
 like you need a soft routing.
 But this has been handled, in fact, these mixture of experts has been introduced previously,
 in a paper I think called outrageously large language models or something like this.
 And so they've introduced that, you know, it, it still works.
 So backprop still works through so basically you have a backprop path through here.
 And because you put a little bit of noise in this routing, every path gets explored
 a few times, and therefore you have enough backprop signal to make it work.
 It can it could technically fail, but they do observe generally that it does work if
 you do this kind of hard routing with a bit of noise.
 All right, so where do we go from here, as I said, this is an engineering paper, and
 it's a long engineering paper.
 So they, they set up their, they set up a lot of a lot of the details of engineering
 directly in the paper, which we're not used to in the machine learning world.
 They really detail how they shard things and so on, which is pretty cool.
 But I invite you to look at the paper yourself.
 If you really want to know what's going on right here.
 Suffice to say, they, as you can see right here, what they do is, this is the input right
 here, and then they have this weight matrix, which is a this routing, this is learned routing
 weights.
 Okay.
 So you have trainable weights that decide how to route the input, and that's dependent
 on the input.
 So you have a bunch of inputs that comes from the lower layer.
 And this matrix right here determines where to route them.
 Probably says, okay, the input is a vector like this.
 I know that must probably go to the expert number three.
 Okay.
 And you have a softmax across that.
 So it's a really, it's an assignment to, it's a soft assignment to the experts.
 So once you've done the soft assignment to the expert, you do a hard assignment by collecting
 the top two.
 For each token, you say you collect the top two experts, and you only send it to the top
 two experts and you ignore all else, which is not a lot right there.
 At times there are 2000 experts in the system.
 And yeah, you distribute and you have some noise.
 So with a random probability, you actually don't even send it to the second expert.
 You just leave it at the first one.
 And with some noise, you send it also to the second one.
 And I think that that noise is part of what if what makes the system work a bit.
 And then you also have this auxiliary loss right here that you add on top, which just
 makes sure that you distribute evenly.
 So this encourages the system to distribute the tokens evenly, because sorry, what it
 penalizes is a this here is the mean assignment to each expert.
 So it penalizes whenever the mean assignment is out of out of line, basically, so a distribution
 assignment to the expert or one expert gets a lot of tokens, because I don't know, it
 tends to be really good at something.
 So all the tokens are routed to it.
 And the other expert don't get a lot that's penalized.
 So you encourage the system to distribute tokens evenly between those experts.
 And then there are also like upper limits where you drop tokens and so on.
 They really build a system that is out for performance rather than machine learning correctness.
 So they demonstrate how to do this in in sort of code with their system.
 And the cool thing about their system is that you don't have to do much.
 What you'll have to do is just specify which tensors are sharded at along which dimensions
 and the system does the rest.
 So this is pretty cool.
 So this here is this mixture of experts, mixture of experts as you would write it in code.
 And they make use a lot of this Einstein, this Einstein some notation.
 If you don't know what the Einstein some notation is, it's a general notation to describe matrix
 or tensor multiplications.
 So a for example, if you were to multiply two matrices, you could have a string there,
 you describe it as a string and it comes from how Einstein wrote up the kind of tensor contractions
 in his work.
 So if you want to multiply two matrices, you can you could put the string a b b c goes
 to a c.
 So this and then you put two matrices right here.
 This will tell it, okay, I have a one matrix.
 I'm going to call the axis a and b, I have another matrix or tensor where I'm going to
 call the axis b and c.
 Now I have the resulting tensor, and I want the first axis to be a and the a is this one.
 And I want the last axis to be c and the c is this one.
 And b is nowhere b is not in the output, which means it should contract over b.
 So it should sum along b, sorry, it should multiply along b and then add such contract
 over b.
 So this here describes a regular matrix matrix multiplication.
 Now if I could do something else, I could do something like a just a element wise product,
 an element wise product would be something like this a b comma a b goes to a b, which
 means here, I have a in the first input, and here I have a again.
 And I'm so you already see that you can even though these are different tensors, you can
 call the axis the same, which means that they're going to somehow be multiplied together.
 Now if you leave it away here, it means that it's going to be contracted and therefore
 the axis no longer exists.
 But here we don't leave it away, which simply means that these axes are going to be multiplied
 together.
 And the same for b right here.
 So this describes an element wise.
 This describes an element wise product, you can go really funky with this.
 So this, this here would be a row wise dot product, where a is more it for all the a
 is it's element wise, but then over b, it's contracted.
 So you know, you can go, you can go wild with the Einstein some notation, you can describe
 a lot of things with it.
 So here is this algorithm to distribute the computation among these different experts.
 So you have the inputs and the weight matrix for the, they call this the gates function.
 That's the routing function to these experts.
 So what do we do?
 We first of all, we have these tensors, these, they have these grouping, these grouping dimension
 right here.
 So they come along to along groups, which in our case, we could maybe say these are
 batches or the batch dimension.
 So they come across groups, and there is the sequence length and there is this M right
 here.
 That's going to be the feature dimension, the M. And you can see the M is contracted.
 So the M is no longer here.
 So the gating function is going to route each input token right here to one of the experts
 for each thing in the group.
 So you can see, you can express this with an Einstein some notation.
 Then you have a top two gating, which selects the top two from each of the last, from each
 of the entries.
 And that gives you this dispatch mask and the sorry, and the weights that you have to
 use at the end to combine.
 You can use the dispatch mask in order to distribute the inputs.
 So you have reshaped inputs, and so on.
 So I'm not going to go through all of this right here, but you can express all of this
 in terms of the Einstein some notation.
 And you can express pretty much any sort of computation that is along the line.
 You can express the attention mechanism and so on.
 You can express the feed forward layers in terms of these Einstein some notations and
 the underlying the underlined dimensions here are the dimensions where we want to shard
 the computation.
 So here, because we have this G underlined, that means that we are interested in sharding
 the computation along this axis.
 So this, I said, this is the batch dimension.
 This is your classic data parallelism, which means that the first machine gets the first
 couple of data points, the second machine gets the second couple of data points, and
 so on.
 And you can see in the weight matrix, there is no sharding, which means that the weight
 matrix lives on every machine as a copy of one another.
 This is different from from here, where you can see that what we're now going to do is
 here it's still sharded according to the batch, but we now are going to shard this according
 to the different experts.
 So we're going to route whatever the inputs are in to these experts.
 And then we're going to execute the computations on the experts.
 So this is now sharded according to the experts.
 And at the end, right here, you can see this is still sharded according to the experts.
 We're going to put it back together.
 And now it's sharded according to the groups again.
 That's what we said, we have the input right here, the inputs, and the inputs are maybe
 distributed according to the according to machines, right, we have these go through
 the first machine, these the second, these the third, and so on.
 This is your classic data parallelism.
 But then we have all of these experts.
 And now all of a sudden, we're going to route these things to the individual experts.
 And we're going to execute the computation in parallel on the experts.
 And then after that, we're going to put back together from wherever we got them now have
 to.
 So this goes here again.
 And so this is just the reverse of what we did before.
 So right, like that.
 So you get all of the outputs again.
 I hope you kind of can imagine how this happens.
 So the first difference is, is that's sharded according to a different dimension.
 And the second difference is, is that when we shard in data parallelism, we execute the
 same computation on all the machines, which means that we have the same weight matrix.
 If we do x times w in a feet forward layer, and we shard this thing here in data parallelism,
 what we do is we send the x to different machines, we split the x, we send it to different machines,
 this is x1, next to x3, x4.
 But we always multiply it with the same weight matrix that weight matrix lives on all of
 the machines and is regularly synchronized, it's kept synchronous in some way.
 Whereas if we shard x to the experts, then the experts have individual functions.
 So the expert one is different from the expert two is different from the expert three, and
 so on.
 Which means that before it wasn't important where x was routed, because we would execute
 the same computation.
 So we can just, you know, sharded according to you know, the first 10 go there, the next
 10 go there.
 But here, it's not crucially important where they are routed to to which expert.
 And that's why we learn the function that is going to route them.
 So this is learned, this is these first line here, these are the weights that we learn
 to route, then we route right here.
 And we calculate your your, we calculate the feet forward layers on the expert, you see
 that this wi and wo, they are the weight matrices of the feet forward layer, the feet forward
 layers are, you have your input, you multiply it by wi, you have a ReLU, ReLU, and then
 you multiply it by wo.
 So it's kind of a two layer feet forward network.
 So this two layer feet forward network, as you can see, this is sharded according to
 the experts.
 And then, and the important part is, of course, that here, the weight is also sharded according
 to the experts.
 And that's what makes each expert different.
 And then it's combined again down here.
 So I hope you kind of get the idea of what this algorithm does.
 But the fact that we shard according to these experts is in fact different than your regular
 sharding where you shard the data like the batch, the batches, but keep the model in
 parallel, keep the model synchronized.
 With their system right now, this is how easy this is.
 So before we simply stated our algorithm in Einstein's sumnotations, there is no way
 to underline code and that magically happened something that was simply for us to visualize.
 Now we want to apply their system in order to make this actually sharded.
 And with the Gshard system, and as I said, I don't know if the code is out or it will
 be out, but with the Gshard system, this is basically all that you have to do.
 So you have these functions, they're called split and replicate.
 What replicate does is it takes that weight tensor and it replicates it on all the machines
 and that keeps it synchronized.
 This is a computation where we simply want to shard out the different to the different
 machines but keep it synchronized.
 And you can see if you do this, this is the operation, then the system knows, ah, this
 here is replicated across the machines.
 So that means I'm going to distribute the data points according to this G dimension,
 according to the batch dimension and multiply it with this matrix according to this Einstein
 sum notation string on all of the machines.
 And I'm going to keep this tensor in sync.
 Okay, so the system knows as opposed to that you have you have the split tensor right here.
 So the split, what it does is it splits a computation here the dispatch expert inputs,
 it splits it according to a axis index onto D different machines or into D different parts.
 So you see here you calculate the how you should do the routing and the resulting tensors
 first dimension is this E dimension.
 And then you say that should be split, you know, according to this first dimension onto
 D different places and these D different places are now separate.
 They don't have the they don't have to be kept in sync.
 Everyone has their own weights.
 And now when you do this, you know, according to this dimension, you can see because we
 know Einstein sum notation now, you can see this E appears here, here and here.
 So this operation is going to be applied element wise, that means independent of each other
 in the direction of this dimension, the system understands that since this tensor is sharded
 according to that dimension, I have to execute this on each of these entries in separate
 with on each expert having their own weight matrix right here.
 I hope this is a bit clear that their system makes it super easy.
 You can basically do two things.
 You can say this thing here is my classic parallelism where I want to keep it in sync.
 And this thing here is where I want to split up and do different computation on the different
 parts.
 And then they have also a general function that is more powerful.
 Yeah, they and they you can auto partition and whatnot.
 So they have a a a they have this we implemented the partitioner in the XLA compiler, which
 means that anything that can translate to XLA is a target for the system.
 And that's, you know, TensorFlow and pytorch can do this.
 So technically, this can come to any of those systems.
 But of course, who has their 2000 TPUs lying around to make use of this?
 But no, I'm kidding.
 I mean, this, I they here use it for transformers.
 And I am very excited to to see what people can come up with for the system, I believe
 a system like this where it's super easy to to shard.
 And they have some, you know, they talk about, okay, we do the single machine compiler.
 So the compiler is also fast and so on.
 I don't even want to go into this.
 But this is very well engineered, it seems.
 And they, they, they basically implement this for all of the operators.
 So I'm very excited to see what people can come up with outside of the traditional applications.
 I think there can be new types of models developed simply because we have a system like this
 that makes it easier.
 So yeah, I'm excited.
 So here, they show a bit how this works on the example of this Einstein, some notation.
 So here, we want to do this thing here, which if you remember, this is the operation where
 we want to route the input to these experts.
 So we want to start with something that is sharded according to the batch dimension.
 That means that we, you know, we have different different parts of the batch on different
 machines.
 And we want to route this and finally end up with something that is sharded on the different
 experts.
 So this is what the system does is first you have these here are the different shards,
 right?
 You want to multiply this, as you can see, this and this right here means that these
 this routing table is also sharded according to the same machines.
 So you have the zero is all on the same machine, the one is all on the same machine, and so
 on.
 So what you want to do is you want to contract is there you want to contract according to
 this s dimension, right, which we have we have omitted right here.
 And if you multiply that, sorry, okay, we omit the s so this is not much of a this is
 not much of a graphic right here.
 But then they have this reshard operation where they do and you don't have to worry
 about this.
 So from here to here, there is this reshard operation that just shards it according to
 the according to E.
 Yep.
 I find this to be a bit more a bit more insightful.
 So if you have something like this, this which is a regular matrix multiplication, right?
 And you want to contract along B, this is exactly the example we had before.
 So here is a situation where our tensor is sharded according to the B dimension and this
 tensor is also sharded according to the B dimension.
 You want to do a matrix multiplication of the whole tensor.
 So what can you do, you're supposed to multiply these two matrices, but they are sharded on
 different machines.
 If you consider what you actually have to do is you have to multiply each row here with
 each column here.
 And that in an element wise fashion.
 So that distributes according to you have to multiply this by this plus this by this
 plus this by this plus the red by the red.
 So you can simply multiply the zero tensors together, the one tensors together, the two
 tensors together and the three tensors together.
 Each one will give you a full matrix and then you can simply add all of them in order to
 get your full results.
 This is illustrated down here.
 So what machine one does, it simply multiplies its shard by its own shard of the second matrix,
 which will give it this thing here.
 And by the nature of how matrix multiplication is constructed, you can simply do an all reduce,
 which means you reduce you sum across all of the machines, and that will give you the
 full result.
 So this is a this is a an example of how this works.
 This is, you know, pretty simple. And I believe you may have seen something like this already
 when you were looking at just parallelizing matrix multiplication, and so on.
 So this system handles this transparently, right?
 If you're shorted like this, this is what the system will do.
 However, if you are shorted differently, the system will act differently.
 So here is a system you want to do the same matrix multiplication, but the first tensor
 happens to be shorted according to the A dimension, the second tensor happens to be shorted according
 to the C dimension.
 And you want to end up with something that's shorted to the C dimension.
 Now we have an additional constraint here that here you can see, we kind of assume that
 this full thing here fits into memory, mainly because we want to obtain the full result
 you see here, a and c should not be shorted.
 So we assume that we can keep that in memory.
 But here we want the final result to be shorted according to C, which imposes the additional
 constraint that it might be that the full matrix never fits into memory.
 So how are we going to calculate all of that?
 We can't do the same trick anymore.
 Now this G short system apparently realizes itself when something is out of memory, and
 it can do a smart move around being out of memory using a loop, which basically means
 that it will compute entry by entry or block by block.
 So these are the matrices we have to multiply.
 And you can see that if I want to do multiply this by this, that's fine, I can do this on
 one machine.
 And that will give me the block up here.
 But if I want the block up here, I have to multiply this by this, which is across two
 different machines.
 So what this system does is it's going into a while loop because it realizes there's not
 enough memory.
 And it kind of sends around these different slices to the different parts, each time computing
 a little piece.
 So here, first, we do this by this, this is fine.
 But then we grab ourselves from the we we grab ourselves this one here, calculate the
 next little piece up here, and then we grab ourselves the number two, calculate the piece
 here.
 And then so this is from zero, this is from two, the one we already had, and then we grab
 ourselves piece three, and multiply that until here, until we have this final slice that
 we want.
 Okay, so this goes in a while loop in multiple rounds, the system gets knows itself when
 it has to do this, and when it can calculate the full thing at once because it fits into
 in memory.
 It's even smarter than that, and that it can do these halo exchanges.
 So if you have to do something like this, a convolution, now in a convolution, what you'll
 do if you think of a think of an image, and you want to do a convolution on it, but the
 image happens to be sharded.
 Let's say the image is so large, it's sharded across nine different machines like this.
 Now, if you want to do a convolution, that's pretty cool, you know, here, here, here, but
 here, all of a sudden, your convolution is across two different machines.
 So this system, G shard will adapt automatically, and do these halo exchanges where it kind
 of sends around from this machine, it'll send something to this machine such that it can
 do the convolution in that step, and vice versa.
 And then this can be padded accordingly, as you can see.
 This is I think this is this is this was like super ugly to implement.
 If you just imagine that for each of these operations, you have to think about, okay,
 how can you express this with these MPI primitives like dynamic slice and collective permute,
 and so on.
 It's just an absolute nightmare.
 And I'm very happy that other people have done this, and I will probably just get to
 use it.
 So there is a lot more to this system than I've just explained, I just try to give you
 a flavor of what building a system like this means and how easy it is to use it like this.
 In order to implement all of this mixture of experts things, you simply go from this,
 which is one single machine implementation, how you would write it to this, which is now
 the same, it's almost the same code.
 But this now you can run on however many machines, and if you compile it with the system, it
 will do what you expect it to do in this shorted way.
 Completely crazy.
 Okay, so they apply this to massively multilingual massive machine translation.
 So two things, it's massively multilingual, and it's massive machine, which means, I guess
 a lot of machines.
 And the reason here is twofold.
 So what they say is, we have massively multilingual translation.
 Why don't they just look at single machine translation?
 And it has a very specific reason.
 Namely, if you have massively multilingual translation, which means that you have a lot
 of different languages, and you all have to translate them, ideally, to all the other
 languages or, you know, every language pair.
 But in this case, they only look at all the languages to English.
 I don't exactly know why, but I guess there must be some kind of reason.
 If you do this, then you can make use of a thing where there are languages that you just
 don't have much data on.
 Like I don't know, Basque or something like this.
 There's not that many people speaking Basque or Swiss German.
 There's not even a written form, a standard written form of Swiss German.
 So you just don't have as many resources.
 And for other languages, you have giant amounts of resources.
 And what you can make use of is this phenomenon called positive language transfer, where it
 happens that, for example, Swiss German is very close to German.
 Now, they can't understand us, which is a giant advantage for us, but still it shares
 a lot of similarities with German.
 So if you learn a lot about German, you can sort of transfer learn to Swiss German pretty
 easily.
 So if you have a system that does German and Swiss German at the same time, you can perform
 better on both languages because the Swiss German part of your model, the part of your
 model that does Swiss German, profits from the German inputs as well.
 Now don't understand me wrong.
 There is not an individual part of your model that for each language, it's all done at the
 same time.
 But still you can imagine that, you know, some of these things will specialize in some
 of the languages.
 But the hope is that if you have German and Swiss German in the same training set, that
 if the model realizes what a question construct is in German, it will be able to apply that
 also to Swiss German with some minor modification.
 So there is a benefit of having these many languages, especially for the low resource
 languages.
 Okay.
 So as the number of languages, sorry, as the number of language pairs to be modeled within
 a single translation model increases, positive language transfer starts to deliver large
 gains for low resource languages.
 Given the number of languages considered, which I believe is a hundred here, M4 has
 a clear advantage on improving the low resource task.
 On the contrary, for high resource languages, the increased number of tasks limit per task
 capacity within the model, resulting in lower translation quality compared to a models,
 to a models trained on a single language pair.
 This capacity bottleneck for high resource languages can be relaxed by increasing the
 model size to massive scale in order to satisfy the need for additional capacity.
 So basically they're saying, if we train all of these languages together, that will help
 a lot for these low resource languages, but it might hurt the high resource languages
 because now we would have enough data technically to train a French to English model on this
 giant model.
 We could train that.
 And now that we have all these other languages in there, it just hurts us because we don't
 have enough parameters.
 And we can solve this, of course, by simply adding more parameters.
 So that's the solution.
 Add more parameters and you increase the capacity of the model and you still get the benefits
 of the positive language transfer.
 So their investigations is going to be into how much can we scale this?
 And is there like a sweet spot where because if you, if you increase the parameters too
 much, you counteract this positive language transfer again.
 So since, you know, since Swiss German and German can sort of benefit from each other.
 However, if we have too many parameters, so, and then we end up having all of these experts
 right here and the tokens are always routed to these experts and it always happens that
 all the Swiss German tokens are always routed to this expert and all the German tokens are
 always routed to that expert.
 There will be no sharing of weights.
 There will be this positive language transfer will not happen because we have too much capacity.
 So the goal is to find a sweet spot between positive language transfer and this capacity
 bottleneck.
 They do use an in-house data set, which we don't have access to, but they say the training
 corpus mined from the web contains parallel documents for a hundred languages to and from
 English adding up to a total of 25 billion training examples.
 However, they only use from 100 languages to English.
 This result in approximately 13 billion training examples to be used for model training.
 So that's a lot.
 It's a lot of data, especially for translation.
 It's kind of a noisy translation because it's mined from the web, but still it's a lot of
 data.
 They have baselines.
 So the baselines are first of all, in order to form our baselines, we trained separate
 bilingual neural machine translation models for each language pair.
 So that means a single model for each language to English, depending on the available training
 data per language.
 And then they also have a baseline where they try open AI style to build as deep as single
 transformer as possible.
 And by that, they mean we also include a variant of a dense 96 layer transformer encoder decoder
 network trained with G pipe pipeline parallelism on the same data set as another baseline.
 So the difference again here is that this 96 layer is a dense transformer, which means
 that all of the tokens go through the same computation and we don't shard the computation
 out to these experts, right?
 We do shard according to the batch, but all of them go through the same parameters.
 And that means we can we can only scale up the number of layers and that severely limits
 the that severely limits the computational efficiency.
 Even if we have, you know, your pipeline parallelism and so on that hurts.
 They say training to convergence took over six weeks on 2000 TPU course.
 That's crazy.
 But I guess, yeah, you know, I was saying earlier that that I always thought we were happy.
 I always thought we were happy in machine learning because kind of the hip science fields
 being biology, like genetics and machine learning.
 I was thought like, oh, but these biology people, they always need like million dollar
 grants from government to run their experiments.
 And we can just sit down with a laptop.
 This time is over.
 If you start a PhD now start applying for money to get TPUs.
 Yeah.
 Okay.
 In any case, here you can see what this does.
 So they compare a bunch of models right here.
 So this T, this is this big dense transformer that's going to be one of our baselines and
 the other baseline here is going to be the zero axis.
 The zero axis means this is the single model for that language pair.
 So only so for each language, they trained one model only on data from that language.
 And that's going to be the worst thing here because this multi language translation in
 one model will generally help you if you have enough parameters.
 You can see all the models here have enough parameters such that the difference here,
 this is difference in blue is positive including this baseline model right here.
 So the baseline model, as you can see, has 2.3 billion parameters, even though it takes
 that much longer to train and that's, as we said, a function of the fact that it's dense
 and deep, so that hurts in training efficiency.
 And then you have these mixture of expert models.
 They always consider two things.
 They consider different numbers of experts.
 You can see it goes from 128 to 2048 experts.
 And they consider a number, different number of layers from 12 layers to 36 layers, 36
 layers still being way smaller than the 96 layer transformer here.
 And that's the reason why it trains faster.
 So it doesn't train faster.
 So the reason it trains faster is because it has less layers.
 And then the reason it has more parameters is because it has a lot of these experts.
 And the art here is to constrain how much these more experts hurt you.
 So you know, you could run into the same problem where if you scale up the experts, in fact,
 you do, it doesn't fit into memory anymore.
 And it's going to hurt you a lot in training efficiency, kind of like if you increase the
 number of layers.
 But the G shard system prevents that it lets you up the number of experts without incurring
 the cost.
 That being said, it does not let you up the number of layers, you're going to incur the
 same cost if you up the number of layers as you have with the dense transformers.
 So does this help?
 It helps a lot.
 As you can see right here, there's a general trend upwards.
 And what's the x axis, the x axis is low resource languages.
 So you can see that as we as we go to lower and lower resource languages, this multi task
 training, this multilingual translation improves significantly over the baseline where we only
 train a system for that language specifically.
 And these 10k examples, it's it's it's quite a bit, but it's not that much, especially
 since it's noisy data.
 So this is specifically good for low resource languages.
 But you can see also the high resource languages here benefit from the multilingual translation.
 And that's a function of the fact that we have, you know, large enough models.
 In fact, you can see the larger the models, the more the difference in blue is, and there's
 not really an end in sight.
 And they also see it say that they haven't seen convergence in training.
 So you can technically train this forever.
 Yeah, you can also see that the the lowest mixture of experts right here is almost on
 par with their big dense transformer that took so much longer to train.
 Right.
 So this lowest model right here, I believe it took I don't want to go back, but it took
 it took hours or so or few hours to train, whereas this 96 layer dense transformer took
 these six weeks to train.
 So has to be said, the number of TPUs is not to be neglected, but if you're Google, you
 know, you just have them laying around.
 What's also interesting here, and you can start seeing this two things.
 First of all, you can see that the difference between here in between the dense transformer
 and this baseline model is very low for high resource languages, but gets larger for low
 resource languages.
 This is an indication that the dense transformer, it does more to share parameters between the
 languages because it shares parameters between all the things because all the tokens go through
 the same computation.
 So it is going to be a bit better in low resource languages, but still the general trend upwards
 holds even for the mixture of experts.
 The second thing is that you see there is a crossover here in these in these big in
 these biggest models.
 And what are the big models?
 One, the blue one is the one with 2048 experts and the green one is the one with 500 experts.
 They're both as deep models.
 But all of a sudden, over here for the high resource languages, it's still true that if
 you up the number of parameters, you get a benefit.
 So up the number of experts as well, you get a benefit.
 But over here for the low resource languages, it's it you see, it actually hurts you to
 up the number of experts.
 And that's the phenomenon exactly we talked about before.
 If you have too many of these experts, and you do a hard routing, that means all the
 tokens go a different way.
 And that means you don't get any sharing benefit from the multilingual translation.
 And they investigate a lot.
 And they basically claim that their sweet spot of expert in their particular task appears
 to be somewhere in between these 2000 and this 500 expert number, where you can see
 it doesn't always help you to scale up the model.
 So I have to say maybe the transformers, maybe they need a ResNet moment.
 So I believe in computer vision, it was sort of the same problem that we try to build deeper
 models and why like, okay, this, this is more width.
 But yeah, I think there might be some breakthrough on the horizon where someone just figures
 out how to train these giant models, even more giant transformer models with deeper
 layers.
 And then there's a new era of transformers.
 However, this is not that effect.
 I'm sorry, I said this at the wrong place.
 This is not that effect.
 This is to show that in this case, we do benefit for the high resource languages because we
 increase capacity.
 But for the low resource languages, we suffer if we up the number of experts too much, because
 they don't share any parameters anymore between the languages or between the different parts.
 Like it's not a necessity that the different languages are going to be routed to different
 experts.
 But it's probably going to happen, right?
 There's no hard coded thing that says if it's this language, it needs to go there.
 It just probably is going to happen this way because the different languages are going
 to be needed to be treated differently and therefore the system learns to route first
 and foremost, those two different experts.
 Here you can see the model sizes, including this 60 layer models model with 2000 experts
 that they didn't manage to train.
 They said they had numerical instability, but that had one trillion parameters.
 And I'm pretty sure they're, they're cool.
 They must be quite mad about this, right?
 Like you have the trillion parameters, even though it's not that much bigger than the
 600 billion, that the trillion, it would be cool to write a paper like a trillion parameter
 model.
 But for now they are at the 600 billion mark and they simply want to tell you that they
 have actually compiled a model that's that big, just didn't manage to train it.
 And yeah, that's here.
 Here is where I wanted to say that maybe we're waiting for the ResNet moment where all of
 a sudden someone figures something out that makes the training of basically infinitely
 deep transformers possible.
 Like we made the training for almost infinitely deep CNNs possible with ResNet.
 Okay, so they conclude this and so they, that's the investigation of what the number of experts
 and so on gives you.
 And here is a bit of a different investigation where they more care about training efficiency.
 So they ask themselves, how many billion tokens of input do we need to reach a given cross
 entropy?
 So here, the more tokens you need, the lower your efficiency is, right?
 You can see that the general trend is the following.
 If you up the number of layers, you get more efficient, you can see and just look at this
 column for now, this point seven column, you can see it already pretty clearly.
 So here you go from 12 layers to 36, you gain efficiency, here you gain here you gain pretty
 predictable.
 If you up the number of layers, you need to see fewer tokens to get to the same cross
 entropy. And in fact, you can get to a lower cross entropy altogether at the end.
 We've known this for language models already.
 The other effect is of course, what happens if we go not deeper, but wider, if we increase
 these number of experts, if we increase this sparse computation.
 So here you can see, let's just look at the 12 layers for now.
 Let's look at all the rows where there's 12 layers.
 So here you get a significant advantage by upping the number of experts from 100 to 500.
 But then you hurt upping the number of experts to 2000, right?
 So that's that's sort of your you're hurting efficiency by upping the number of experts
 too much.
 And the same if we look at the 36 layer, so you gain massive efficiency by upping the
 number of experts, but you lose that a fish part of that efficiency again, by increasing
 it even more.
 Now we saw that the this model is still the best model, but it's not as efficient as that
 model.
 And that gives you another indication that there is sort of a sweet spot between these
 two things between the positive transfer and the bottleneck capacity that appears to be
 somewhere in between right here.
 So that's pretty interesting.
 Because we know about depth that you can basically up and up and up and get more efficient, but
 with not that much.
 Yeah, the largest model can be trained in under four days to achieving the best quality.
 Yes, yes, yes, but this is just a yeah.
 So here, oh, you can see the batch size in in tokens is quite, quite a bit.
 So yeah, if you have a 1000, if you have a context window of 1000, that means the batch
 size here was about 4000.
 So as as expected.
 Yeah, this is just easy peasy 22 TPU core years.
 I've seen someone on Twitter saying this, this is the new measure for computer.
 It's no longer like flops.
 It's TPU core years.
 Just mad, mad.
 And yeah.
 So 42 days to train this thing right here.
 Crazy, crazy, crazy.
 All right.
 They also have a number of investigations in other parts of efficiency, like per device
 memory consumption.
 You can see here that as you up the as you up the number of experts, you can see here,
 here, here, your weights don't go up because as you up the number of experts, you can just
 up the number of machines and the per machine weight usage will be the same, right?
 Because the experts are independent of each other, each one has their own weight matrix.
 So you can just add machines and you keep your weight requirements the same.
 However, if you go deeper, then your weights increase because you're now deeper, you have
 more layers.
 You have your so also your transformer weights will be higher and so on.
 So you go deeper right here.
 You see 3660 layers, your memory consumption increases for the weight.
 And also, this is the other big part in transformers, right?
 The activations that you have to save, because as we said, if you have a transformer and
 I have layer, layer, layer, layer, I basically have to keep around each of these signals
 in order to do back propagation.
 And that's why also the activation here increases as I go deeper.
 Now you can see percentually, it decreases again here.
 So what's happening?
 Technically, you don't have to keep these things around.
 You can also once the signal comes back, you can recompute them from the beginning or from
 an intermediate point.
 Now this increases computation, but saves the need to store the activations.
 And apparently G shard, yet another thing it does is it will recompute as necessary
 the activations if it realizes that you don't have enough memory to store them.
 So all of this is pretty crazy, honestly.
 And they look at where the different computations go.
 And I don't want to go into this.
 And they have these micro benchmarks where they really show that the increase in complexity
 is really according to square root of n, because that's how long it takes to distribute along
 these actors, sorry, along these experts.
 There's a lot to this paper.
 And there's no time to go through all of it.
 I think this video is already way too long.
 I hope I have given you an impression of what's possible with this system.
 And as I said, I'm excited what people can come up with.
 Just to say that in the appendix here, they detail that they have done this for all the
 operations in XLA.
 So for example, convolution, this is so ugly, how you have to implement the convolution
 because you have to padding must be correct across these expert across the the sharded
 machine.
 So there are no experts anymore.
 This is just G shard, the padding has to be correct.
 The strides have to be correct.
 Data needs to be exchanged according to the machines, the window size needs to be correct,
 blah, blah, blah.
 So just thank you for doing this and not having to do it myself.
 Yeah, I'm excited as soon as as the codes out, if I get a hold of it, I'll you know,
 link it or you'll find it once it's out.
 If it's already out, I'm just too dumb to see it.
 I enjoyed reading this.
 It's different than a machine learning paper.
 It kind of shows you what goes into engineering a system like this, and how easy it can be
 if it's engineered well to then apply it.
 I think this is going to be extremely helpful to the community.
 And with that said, 23 pages later, see you next time.
 Bye bye.