[Paper Club] Upcycling Large Language Models into Mixture of Experts

Sure.

Okay. Pointer works too.

Okay, cool. Uh, today I'm, I'm going to present a Mixture of Expert. I'm Ethan from NVIDIA. I work on scaling, uh, LLMs, transformers. Um, I, I might sound a bit muffled 'cause I, I'm having a cold, so apologies for that.

Yeah, that's-

So to this topic, uh, I, I'm going to give a brief introduction on Mixture of Experts, or MoE, first. Then I'm gonna talk about Megatron-Core MoE, like how we, uh, salary these, uh, MoEs and, uh, train and do inference efficiently.

Uh, finally going to talk about the upcycling M- LLM into MoE, which is, uh, of a recent paper.

So the AI models are growing larger and larger. This is a rather old picture from twenty twenty-one. Uh, Switch Transformer is the first model that surpassed the one trillion model size. Before that it's, uh, it's only hundreds billion parameters.

I think the-- At that time, it's growing ten times each year. Looks like it's slowing down recently. Um, so the, the question is, like, uh, uh, we only have so much compute. Uh, how, how can we make the model better without increasing the compute?

Uh, from, from Noam, he said, like, uh, "My unsubstantiated theory is that parameters are good for knowledge and compute or the flops is good for intelligence," uh, whatever this term means. So MoE is a good way of growing the parameters or growing knowledge without increasing the compute.

So what is MoE? Um, here is this, uh, very simple diagram from Switch Transformer. Like, uh, for the traditional, uh, LLMs, uh, transformers, you have the self-attention, the residual norm, uh, layer norm. And then you have a FFN layer, which is just a two linear layer, and then the residual.

Um, MoEs transform the FFN layer into multiple copy of them. You see here there are like a four FFN layers, and then each token selectively activate a few experts, which is selected by router. Um, the router is simply a matrix multiplication, a learnable matrix to, to select one of the expert based on the input.

The model size increase enhancing its capability, while the compute roughly remain the same as the original model.

Um, if we look into the MoE layer, it actually consists of, uh, several steps. Uh, it's more complicated than the original FFN layer. You can think of the original FFN layer as the, the third step computation. There is a, uh, there is a expert layer, which was the original FFN layer.

It's applied to the input token. So the first step here is, uh, is routing. So given, given the input token, for example, here you have, like, uh, six tokens, "The quick brown fox jumped over."

It went over the router. Router is simply a matrix that is applied on these tokens. It generates the probabilities and, uh, given we, we will take the highest probability, uh, as a router selected. So here is a rou- the expert indices are the experts which are, have the highest probability on these tokens.

And then the second step is the permutation. So given the, given the tokens selected expert, you need to align these, uh, input features with the expert. Like, uh, all of the tokens that select ex- expert zero will go to, uh, uh, you need to arrange them into a single matrix that are, uh, only for expert z- zero, the expert one, expert two.

Depending on the capacity factor, some of the token might be dropped. And the third step is the same as the original FFN layer, where you do the computation. After the computation, there is, uh, there, there will be a unpremier step where you need to arrange these tokens back to the original shape.

Then the router probability is applied as a scaling factor on all of these, uh, uh, all of these, uh, expert features.

So scaling this, uh, scaling this MoE tr- during training is very challenging. Um, like, first, the, the models are massive scales. Usually, for example, Mixral-eighty by seventy, you, you increase the model parameters roughly by, like, six X or seven X.

This puts substantial pressure on the mo- memory usage. And the router dispatching also has overhead Like pre, in the previous slide, you will notice there's a permute on, on permute operator that essentially increase, uh, activation memory by two times because all of those need to be stored.

If you have like top K routing, uh, it would also increase the memory of the activation by K because, uh, uh, the, the hidden states need to go to each expert, and you need to duplicate it. Uh, and also reduce the gemm efficiency because you need to do a loop over all the experts to, to do the gemm, uh, separately.

There's also like a imbalance issue if all of the tokens go to one expert, other, other experts on other GPU would be idle.

So, um, now, uh, let me introduce the Megatron-Core MoE, uh, which is how we accelerate these MoE models given these challenges.

Um, so Megatron, Megatron-LM and Megatron-Core is an open source library on GitHub available. Um, we accelerate not only MoE and also all of the LLMs, including like GPT, BERT, T5. Uh, not sure if anyone is still using those now, but primarily GPT models.

And in the, inside the transformer layers, the attention is, is accelerated with all kinds of parallelism, including pipeline parallel, tensor parallel, uh, also the parallelism and the, the MoEs are also accelerated. This is what we are primarily talking about today, Megatron-Core MoE.

And, uh, on top, you, you have two, uh, customizable training loop. Megatron-LM provides simple barebone training loop. Uh, you can, you can ha-- easily hack, and Nemo pro-provide a high-level interface where you can, you can just provide pass onic configuration to train these models.

In the Megatron-Core MoE, we provide different approaches to accelerate these models. For the router, there, there OxLos and Synkorn. Basically, OxLos is a token, token tries MoE, and Synkorn, it's, uh, without token dropping, it can usually use in, uh, expert tries.

And the tokens dispatcher, they're in permute, unpermute for efficient memory saving. And the expert, for the expert, we have grouped MLP to accelerate these.

So first, it's, uh, expert model parallel. This is, uh, this is available in Megatron-Core MoE now. You can... Uh, usually, you would put all of the experts on one single GPUs, and then you do a far loop over all of the expert to compute the result.

But instead, we can put one expert on each GPU. Uh, this will release a lot of memory and also accelerate the training and the inference.

So token dropping, uh, the default we use is dropless, meaning, uh, all of the tokens can, can go to one expert, and no tokens are dropped. But we also support, uh, drop, token dropping with padding, uh, meaning given a set capacity factor, for example, four here, uh, each expert can at max accept four tokens.

Tokens beyond that are going to be dropped. So accuracy-wise and, uh, efficiency-wise, there are a lot of discussion around here. Like, uh, uh, a lot of, a lot of the pre-training experiment shows that token dropping is very efficient, and it doesn't impact performance.

But in some of the, like, the downstream fine-tuning, people realize dropless is better. Like, uh, maybe it's because of the domain shifts, the tokens are no longer balanced. So we provide both of the options.

So recently, there are a lot of new MoEs that have increasing number of experts. This will co- this will cause a very huge overhead. For example, the DeepSeek V2 MoE, it has one sixty experts, and eight of, eight of them are active.

If you think about, uh, the memory overhead, let's say you, uh, first you would have the tokens, uh, and this token need to go to different expert. And at this step, you'd have a scatter and a copy operator.

Depending on the number for top K, since the top K is eight here, each of the hidden state would be copy eight times, which is a very huge, uh, overhead. And after, after the expert operation is done, uh, there's another eighth copy of it here.

So we, we have a fused permutation operation available in Megatron or MoE, where all of these operations are fused. Because, uh, the copy operation here, it's, uh, it has zero compute, but it could cause duplicate in memory. If you have these fuse operations, it, uh, you can easily recompute these, uh, features during backward pass, uh, while saving a lot of memory.

Uh, let's also look at the implementation of Mixture eight by seven on Hugging Face Transformer. You also notice the, in the expert operation there, you would iterate over all of the experts and compute each of the gem operation one by one.

We found that, we found that this is very inefficient. Instead, uh, we provide a interface to CatlistGroupGem, where the, the CatlistGroupGem group all of the, like, looping over, uh, experts and calculate gem into a single operation. Uh, given any number of the ex- any number of experts and any number of the tokens, you can efficiently compute the output in one operation, and it's, it's very efficient.

Okay. Uh, so that's, that's pretty much all of the optimizations for MoE in Megatron-Core MoE. If, if any one of you have questions, we-- I can first answer those question and then go to MoE apps secondly.

Um, just curious, like is, are these available to, for everyone to use in the way that flash attention is? Or are these specific to the Megatron implementation?

So,

um, you can, you can import them as a standalone module and apply them to any, any of your module.

Core import library, right? Yeah.

Yeah. Yeah, you can import Megatron-Core as a library and ju-just use it in your network. But I, I think there are some caveats. So, for example, if you want to use expert parallelism, you will need to also use Megatron's, Megatron's, uh, uh, parallelism strategy to initialize the strategy first.

But if you're using, for example, the GroupGem, I, I think you can just get away with, uh, any kind of network. You can combine it with Hugging Face Transformer without any problem. This, this is just a, it, it's just a PyTorch layer with a fused operator.

You, you can just import it as a library. It's very standalone.

Cool.

Um, do you have any intuition as far as, like, the knowledge contained in each expert? Um, I guess previous to seeing this paper, I had assumed that, like, maybe one expert was good at economics, another one was good at physics and so forth.

But this makes it seem like it's more token by token. Um, so do you have any-

Yeah.

-intuition around that?

Yeah. So we haven't studied in our research, but I saw a lot of research on interpretability of the MoE models. So, uh, unfortunately, people didn't find like a significant, uh, interpretability inside this expert. Say one expert focus on math, the other focus on literature.

Uh, I think the problem is that, uh, neural network hidden states are already very entangled. So one hidden states are in this kinda superposition where it can represent multiple different features. So it's very hard to tell which expert focus on which area.

But there are some, uh, some evidence of specializations on early layers of the expert. For example, fir-first and second layers. You'll find some expert focus on like multiple tokens, some expert focus on single tokens, um, things like that.

And there's also one pretty interesting research from Facebook where the, they do a, a, a specialized training of dense model first, then combine those, uh, dense experts specialize into diff-different domain into a MoE. In that case, uh, it still preserves some of the specializations.

Thank you.

Um, just a quick question. We talked about top K sampling to select the expert throughout. Um, is there any benefit to using maybe a top P sampling approach similar to how you would use a top P sampling for selecting an expert as compared to top K?

Uh, top P? Uh, what do you mean top P?

Um, considering like, uh, a list of experts until they exceed like a given probability cumulatively.

Uh-huh.

Does that make sense?

I s- I see. Yeah. So, uh, in the top, uh, it, it will be dynamic. Let's say sometimes it can, it selects some more, sometime it selects less.

I, I think that promotes a little more diversity, I, I've heard.

Yeah. Yeah. I, I think that makes sense. That will create some difficulty in, uh, optimization, so.

Okay.

But I, I think another thing pretty exciting is the expert choice.

Okay.

You see in the expert choice model, uh, the selection is reversed.

Okay.

So here we, we talk about all of them. Usually this is like, uh, uh, uh, this is token choice, which means the token selects K experts.

Yeah.

So each token always have, for example, two experts applied on it. But expert choice is another way where the experts select tokens. The... Each expert only selects the K tokens constantly. So this, uh, even though this is fixed, but, uh, from the token perspective, each token can have zero expert applied to it, or more than zero, or all of the expert applied to it.

Yeah. In that case, you're either overloading the expert or not considering all the tokens. It's like a trade-off.

Yeah. In fact, the expert choice applies pretty well to vision models because vision models do not have causal mask.

Okay. Okay. Thank you so much.

Okay. Uh, I, I guess I can go to the, the next section, upcycling MoEs.

So,

so if you're going to remember one thing, uh, remember this. So you can upcycle your dense models into Mixture of Experts. By training these upcycled models, you can achieve better accuracy than simply training the dense model further for the same number of flops.

The context is we have so many, like, big dense models. For example, Llama, uh, four or 5B, and from NVIDIA we have NeMo-Trons through 40B. These models are, are huge. It's very expensive to retrain MoE variant of it.

If you, we want to further improve these models, uh, you can upcycle these models into MoE to achieve better accuracy. On our scaling experiments, we tried on 15B models, uh, upcycling and applied on one trillion tokens and achieved roughly about five percent improvement in terms of the validation loss and four percent improvement on MMLU.

Uh, it's, it's exciting because the theoretical sparse upcy- cycling paper, uh, found it difficult to scale beyond one billion parameters. We found there are several key factors to, to go beyond one billion.

Uh, so this is re- this is a concept of, uh, MoE. Let's say you have, you have the MLP in the original plain old transformer. And, uh, here this is a mixture of two experts. One of the expert is activated so the flop is the same as the original model while the parameters is increased.

To upcycle such model, you, you do two things. First is, uh, you copy the MLP layer into, into, uh, number of expert copies. Here is just two copies. And then you randomly initialize the router weights. Then you just train this model.

Pretty straightforward, right? Um, there's one caveat here. Like, uh, this model needs to perform the same as the original model in the first forward pass. Otherwise, it will lead to catastrophe forgetting. Like if, uh...

So the trick to maintain this, uh, this feature is through the swapping of the top K softmax operator. So this is introducing Mixtral 8x7B. Let's say you, you have the MLP copied and then, uh,

and then you, you do the top K first to select two experts. And then you do the softmax on top, on top of the, uh, logits from the top K router. In this way, because the softmax is applied to the top K output, it always sum up to one.

Here I give example zero point seven, zero point three. And then you add these two output together. Uh, in this way, because the MLP layer is, uh, exactly the same as the dense model, these two copies are just the same.

So the model output is the same as the original dense model. This is a very important feature in upcycling, uh, because, uh, the upcycled MoE model is actually, actually behave exactly the same as dense without any training. And this will help stabilize the model and avoid catastrophe forgetting.

But the problem here is, uh, we actually found the mixed approach didn't work as well as expected because the, the original model, the original switch transformer from Google uses a softmax and top-K for a reason. And, uh, because of upcycling, if you switch to top-K then softmax, it actually hurts some performance.

The difference, the difference is, uh, simply the swap of top-K and the softmax order. We've... I've already explained, uh, how mixture did the soft, uh, top-K than softmax. So in the original sphere transformer paper, you apply softmax first, so the, the probability on all of the experts would sum up to one.

So if you apply top-K, uh, the output no longer sum up to one. Here I... It's the example zero point four, zero point two. So this is smaller than the original output. If you just train this model naively on large scale, uh, the model would catastrophically forget.

Uh, but on smaller scale model, like one B or under, you can kind of get away with this problem 'cause small model adapt very fast. This is probably one of the reason the original sparse absent cycling paper didn't go beyond one billion parameters.

However, the, uh, so we found a very simple approach to solve this problem. Uh, because the output scale is smaller than the original model, we can simply scale up the MLP output, uh, by the number of expert divided by top-K.

For example, if, if it's mixed row eight by seven, we simply scale the MLP layer output by four X. This will solve the problem with the scale, and the, the model still behaves the same as the original model.

You can train, train the upcycle model, uh, normally. And then we found, uh, it, with this approach it consistently outperforms the, the mixed approach. Uh, so we can... Because the benefit of, uh, the original sphere transformer. A bit of intuition behind why, like, top, softmax and top-K is better because the, if you apply softmax to all of the experts, the, the probability to distribution is always measured on all of the experts.

However, in, in the swap the case, the probability to distribution is only on two expert, and it's dynamic. It's harder for the model to learn. Additionally, if you only use top one, uh, this method would not work because, uh, if the op, softmax on one expert is always one.

There won't, wouldn't be any gradient to learn.

So next let's go to, uh, fine-grain MoE. Uh, this is very popular in the most recent MoEs. For example, the, the Qwen, Qwen-V-two use sixty-four experts. And DeepSeek would use one twenty, one twenty-eight or one sixty something expert.

Uh, granularity use more expert, but smaller ones. So this gives the flexibility of, uh, more combinations of different experts, so more representation power.

For example, here, originally you have the big serve two expert and one is selected. Instead, you can, you can expand the number of expert into four. Now each expert is, uh, uh, is, uh, smaller than before, so the compute is, is the same, and the parameters is also the same.

Uh, here each, uh, sub-notation. Here E-two means the expansion factor is two. This is how many times the expert is copied and, uh, T-two is granularity, or how many times the expert is, uh, segmented. So you can think of this, like, the expert is, is first copied into two copies here, and then each copy is, uh, is segmented two times.

And T-two means, uh, how many experts we route to. Uh, here you route to two expert, so the, the flops is the same as the original one.

Uh, so you would soon notice a problem in if you upcycle such model. Let's say you segment, uh, these two segment are not the same anymore because, uh, you segment one Mo-- one expert into two copy, uh, two shard.

If for the top two case, if you selected one shard two times and the other shard zero time, the output is no longer the same as the original MoE. And we, uh, we mentioned that- It's very important for upcycling to maintain the same, uh, forward path as the original dense model.

The solution here is also rather straightforward. So instead of randomly initialize a router, we would initialize half of the router and then duplicate it. So this would ensure that the probability distribution would be the same in each, uh, virtual group or i-in each shard group.

Uh, because the weights are the same, the top case selection will be exactly the same. Well, you would, uh... The highest score would be the same for these two group. Let's say h-here the example is that I, I shard the MLP into two parts.

So the, the first, I need to select orange e-exactly once and the blue exactly once. By duplicating the router weights, this achieves the purpose.

And, uh, the formula for scaling the weights are also the same, except there's an additional granularity factor. They said, but it's, uh, complicated, so I, I don't... Uh, I'll skip this part. Uh, so for our experiment, we, we do our, uh, experiment on the eight trillion tokens Nemotron to train on.

The ablation study is on a smaller model, Nemotron-2B. And the bigger model, we do it on eight by fifteen B on one trillion tokens.

And we found that the learning rate is the most important hyperparameters, as always in machine learning. And this is true also for upcycling. Learning rate is the most important parameter. So in the original sparse upcycling paper, the learning rate is, is taken from the minimum, uh, the, the ending learning rate from pre-training, so sometimes it can be rather small.

So we found that if you, if you have a high learning rate, it could help the upcycling a lot. Here, the orange line is the lowest learning rate. Uh, basically, you continue fine-tune the model into MoE. Uh, this learning rate is typical, for example, for other tasks like alignment.

But for, uh, upcycling, the model need to adapt to a new local minima. Uh, we need larger learning rate for this. We found that the best is to use the original highest peak learning rate from pre-training, which works the best.

And if you dive into the, the base of this upcycled model, you'll find something really interesting.

So if you apply constant small learning rate, like in, in fine-tuning or alignment, the cosine similarity between the base model and the upcycled model would be almost one. This is true for most of the aligned model. For example, Llama chat versus Llama base.

If you use a high peak learning rate for upcycling, the, the cosine similarity would be much lower, around zero point seven. Also, we kinda... We also analyze the Mistral eight by seven base versus Mistral seven B. We found that the similarity is also around there, which means you need higher learning rate for upcycling.

An additional experiment on number of expert, we found sixty-four expert is kind of like the sweet spot. If you increase the number of expert beyond sixty-four, it provides dimish-diminishing return.

Finally, this is the large scale upcycling. So we upcycled the eight by fifteen B model on one trillion token.

So the, the validation loss, uh, so there are three models here. Let me explain the base model. Fifteen B is trained on eight trillion tokens. Uh, this is pre-training data. So validation loss one point six, and the MMLU fifty-nine.

So the continued training model, uh, it target more on the, the academic benchmarks and more ta- to obtain higher performance on the evaluations. So we, uh, this is roughly one trillion tokens. And the upcycling is performed on the same data for comparison.

Continued training is just as a dense model continued training. You will notice that actually the data plays the biggest factor. Uh, even the, the base model, uh, continued trained, it can have a huge boost on MMLU. Uh, and the upcycled model is another, like, a four to five percent improvement on top of that.

But continual training is, like, twenty, twenty percent improvement because of data. As this reminds us again, like, data is the most important in ML.

If, if you put the five percent improvement into, into the law, scaling law perspective, we can roughly gauge, like, how much the five percent improvement means. Uh, like the, the fine-grain MoE we upcycled, uh, this is, this has the same FLOP as the dense model.

And, uh, this has about four percent improvement in terms of the loss. Four percent improvement if you plug in the scaling law from OpenAI, it's roughly, uh, represent one point seven X bigger model. And, uh, the, the non-fine-grain model, uh, Val top two, this is the, the same config as, uh, eight by seven Mixtro.

It has increased FLOPs because of top two. Uh, that's roughly one point seven X more FLOPs, and, uh, it's roughly two times as powerful as the original model. Given that we only spent, like, one-eighth of the original pre-training compute, uh, this is a decent, uh, some saving compared to training these MoEs from scratch.

Okay. Uh, thank you for listening. I, I put the paper links, uh, Megatron-Core MoE GitHub there. And, uh, NeMo GitHub, where we provide high-level training interface. You can also, uh, follow me on LinkedIn and Twitter. I, I can take questions now.

Hey, nice, nice meeting, Ethan. Um, there are a whole bunch of questions. Uh, I think we'll stop the recording so we can, uh, op-open up for questions. But, uh, I think everyone is very-

Sure

... excited by the presentation.

Thank you.

There's a lot of questions right now, so people want your slides. Uh, so let people know-

Thank you

... where to get your slides.

Yeah, sure. I, I can share the slides