【0410 LLM新知】Mixture of Expert，2024年LLM的新標竿
前一篇有點卡文卡太久XD 先來插播另外一個最近很重要的進展，Mixture of Expert model。MoE最近有多熱，我快速講幾件事各位就能知道了。
- Elon Musk開源的Grok-1是一個314B的MoE model。
- Mistral AI在4/10偷偷release 8x22b的MoE model。
- JetMoE展現只要不到10萬美金就可以訓練出超過LLaMa2–7b的模型。
- Jamba mixture了transformer跟mamba block，在效果上堪比Mixtral 8x7b與LLaMa-2 70B，在效率上遠勝兩者。
- DS-MoE提出了新的訓練MoE model的方法，得到比Dense model更好的scaling law，並且inference也更快。
- Qwen1.5 MoE展現出MoE可以減少75% training cost達到一樣的效果，並且inference效果也更快。
大量的MoE發展讓我最近也積極探索MoE新的可能性，以及思考在業界我們應該怎麼應用MoE模型。
本文我就會淺談一下MoE模型，專注在介紹新的進展，希望幫助大家快速建立2024年對於MoE模型的認識。
本文會分3個段落，分別是
1. 回顧2023年我是怎麼理解MoE模型的
2. MoE模型最近的進展，為甚麼我們應該更多關注MoE模型
3. 其他我們應該稍微關注的MoE研究進展1. 回顧2023年我是怎麼理解MoE模型的
Mixture of Expert model（MoE model）第一次被大眾重視恐怕是去年GPT-4的parameter leak，謠傳GPT-4是8x220B的Mixture of Expert Model。大家開始發現MoE這條技術路線。而後Mistral AI release Mixtral 8x7b達到gpt3.5-turbo的能力更是加強大家對MoE的信心。
但在2023，筆者其實對MoE是相對沒有信心的，基於自己淺淺對MoE發展歷史的認識，我在MoE早期跟同事分享了兩個MoE可能最重要的思考問題。
第一個問題：Mixture of Expert model可以跟多大Dense model comparable？
一般MoE model會呈現以下形式，NxMb模型。
- N代表你總共有幾個Experts。
- M代表你每個Expert多大（會activate多少parameters）。
像是Mixtral 8x7b代表的就是，總共有8個experts，每個expert要activate 7b parameters，而因為Mixtral 8x7b每次會activate 2個experts，所以每次forward都要activate 12b parameters。
為甚麼是12b而不是14b，是因為即便是MoE，中間也有共用的parameters，因此不會是直接2*7。
因此我們可以得知不考慮模型細節的話，最少就有2個參數決定著MoE模型的大小、能力，就是Expert count（N）及Base model size（Ｍ）。
接著就要提到MoE這領域最重要的論文之一，由DeepMind與Google Research提出的 <Unified Scaling Laws for Routed Language Models> [1]
這篇論文提到不論各種routing機制，隨著Expert Count（Ｎ）增加，模型效果可以越來越好（如下圖），並且如果我們專注看Expert Count = 8的情形，可以看到8x15M的model大約等於1x25M的model，8x25M的model大約等於1x55M的model以此類推。
基於此我們可以得到第一個最重要的結論：8xMb的MoE模型可以comparable with "2*M"b的Dense模型。
因此Mixtral 8x7b的能力應該會跟Mistral 14B接近（如果真的有Mistral 14B model）。看到這裡有些人可能馬上就會想到問題了，Mixtral 8x7b的總參數量大約46.7B，但只能跟一個14B的模型能力相近，這樣不是虧大了嗎？
這就帶來MoE第二個重要結論：空間上虧了，但時間上賺了。
一樣以Mixtral 8x7b為例，雖然總模型參數量46.7B，但每次只要activate 12B的parameters，所以總計算量比14B模型更少，因此可以更快速的得到output。
因此MoE在2023年，在我心中其實就是一種典型"用空間換時間"的技巧。
第二個問題：我們要如何有效得到MoE model
目前已知訓練MoE模型有兩種主要的方法，
- 從0開始訓練一個MoE model，早期的switch transformer [2]以及最近的OpenMoE [3] 都是採用從0 pretrain一個MoE model。這種方法就如大家訓練一個Dense model一樣，耗時耗力，台灣大多公司沒有能力做，有能力的公司也很難做好。
- 基於一個訓練好的Dense model來"MoE化"，這種技巧我們叫做Sparse Upcycling [4]。
雖然想法很美好，但是Sparse Upcycling實際上不比從開始訓練容易，參考其論文在ViT上的實驗，我們可以看到當訓練量不多的時候，Sparse Upcycling根本效果出不來。而想要真的有效果，動輒就是幾百天的TPU days。
因此我們得出第三個結論：MoE model得來不易XD
2023年，我主要就基於這三個結論
- 8xMb的MoE模型可以comparable with 2*Mb的Dense模型。
- 空間上虧了，但時間上賺了。MoE是空間換時間的做法
- MoE model得來不易XD
快速做出判斷，MoE短期內都不會是我專注的目標，因為在台灣業界大部分時候部屬model都是memory bound而不是computation bound，同時效果也沒有非常出奇，再加上訓練的成本也沒有下降。
2. MoE模型最近的進展，為甚麼我們應該更多關注MoE模型
前一個段落講了2023年我主要的2個思考與3個結論，而這個段落我就想強調一件事：2024年這3個結論都或多或少被打破了。
參考論文 <Dense Training, Sparse Inference: Rethinking Training of Mixture-of-Experts Language Models> [5]，作者提出了一種新的MoE訓練方法，後續統稱DS-MoE，在結果上DS-MoE 6B可以跟Dense 6B得到接近的效果，總參數量相近，並且僅僅只需activate不到2B的parameters。而SMoE-5B（舊的MoE訓練方法，activate 1.2B）則會遠遠若一節，跟Dense 3B相近而已（如下圖）
DS-MoE藉由一個很簡單的方法達成這件事，在訓練時同步訓練所有Expert。（如下圖，以往的Sparse Training在訓練時執會update top-k expert的weight，大幅浪費訓練效率以及參數效率）
再加上Mutual information loss來控制expert間的load balance。（如下圖）
而因為是MoE model，所以速度上的優勢依舊非常明顯，Throughtput接近Dense的2倍。
方法上筆者並不想深究，但是可要先知道3個新的結論
- DS-MoE可以在一樣的空間下得到跟Dense model comparable的效果
- MoE不用再用"空間"來換時間，空間依舊，時間加快。
- MoE訓練效率可以大幅增加
而在上週JetMoE [6]的出現，更同步證明了3。JetMoE只用不到0.1M就得到了跟LLaMa 2 7B一樣的結果，而其之所以有效，正是因為MoE架構，他的dataset都使用opensource dataset，因此data scaling law上並沒有優勢。
JetMoE在文章中寫他們只使用96 H100 * 2 weeks，與LLaMa2 7b使用了總共184,320 A100 GPU hours相比，考量到H100跟A100的效果差，JetMoE大概只使用30%的訓練資源就達到了跟LLaMa2 7b一樣的效果，這應該是目前看到最有效的訓練優化方法。
類似的結果在Qwen 1.5 MoE [7]也有看到，Qwen 1.5 MoE說他們只用25%的訓練資源訓練MoE model就達到了跟他們Dense model一樣的結果。
同時我們回頭看效率（第二個結論）的部分，除了原本MoE對比Dense model的提升以外，Jamba [8] 更進一步藉由Mixture transformer block及Mamba block，進而大幅提升了他的throughput與Context length。（具體Mamba為甚麼能提升效率我之後也會寫文介紹）
Jamba架構圖（中間結合Transformer跟Mamba的架構）
Throughput比較（可以看到batchsize變大後優勢非常明顯，特別適合業界場景）
Context Length比較（Jamba可以輕鬆容納100k context with 1 A100 GPU）
從Jamba的結果我們可以發現，Mixture of transformer and Mamba甚至超過了單純的MoE model，效率上大幅提升，效果上也維持或是更好（可以參考Jamba論文）。
因此2023我不關注MoE的3個主要理由都或多或少被打破，且被多方驗證，2024我們應該要記得的新結論如下。
- MoE可以在一樣的空間下得到跟Dense model comparable的效果（DS-MoE）
- MoE不用再用”空間”來換時間，空間依舊，時間加快，且結合Mamba後效率還可以再飆升。（DS-MoE, Jamba）
- MoE訓練效率可以大幅增加（DS-MoE, Jet-MoE, Qwen 1.5 MoE）
基於這三個結論，MoE很有可能會慢慢取代Dense model的發展，變成主要開源閉源模型的技術標準。因此做為相關領域從業者的我們也需要花更多時間關注MoE的發展。
3. 其他我們應該稍微關注的MoE研究進展
第一、MoE的inference efficiency技巧，雖然MoE model activate較少experts，可以加速inference，但是目前主流還是把整個模型load到GPU memory來進行運算，其實空間上還是有優化的空間。這邊目前有一些研究在探討加速的可能性，不論是加速還是減少Memory使用，都能帶給serving MoE model大量的好處。
而這方面已經有論文開始討論如何加速MoE模型的運算，藉由更高效、聰明的方法來load expert into memory。[9]
第二、MoE model訓練出來的能力偏重跟Dense model會不會不同？雖然我們都說某個大小的MoE model可以跟Dense model comparable，但講的是overall performance comparable，會不會MoE model跟Dense model的能力側重不同。
一樣參考JetMoE [6]，可以看到JetMoE-8b的math reasoning跟coding能力是遠大於LLaMa2-7b的，是否MoE model可以具備更強的reasoning能力？
第三、MoE model跟VLM是否存在更有效的結合方法？目前已有的MoE VLM model就是MoE LLaVa。
雖然MoE LLaVa效果很好，但它還是較為Naive的結合了MoE跟VLM（在LLM的部分做MoE，前段Vision encoder依舊是走LLaVa那套）
而我們已知VLM最少就需要兩種截然不同的能力：1. Vision perception，2. Language reasoning。這種能力的分野是遠大於LLM中不同領域的能力。因此是否可以藉由MoE來分散不同能力所訓練到的expert，並且強化各個能力，這件事我沒有把握，但是一個可行的方向。
同時對VLM如果想深入了解的話，推薦另一篇文章，可以快速了解VLM。
結論
這篇我快速回顧了2023年研究界怎麼認識MoE以及2024的大進展，並且強調MoE在2024年的3個新的結論。
- DS-MoE可以在一樣的空間下得到跟Dense model comparable的效果（DS-MoE）
- MoE不用再用”空間”來換時間，空間依舊，時間加快，且結合Mamba後效率還可以再飆升。（DS-MoE, Jamba）
- MoE訓練效率可以大幅增加（DS-MoE, Jet-MoE, Qwen 1.5 MoE）
筆者也認為基於這三個結論，2024年MoE會慢慢取代Dense model的發展，並在2025年變成LLM的主流，因此做為相關從業者、研究者都應該花更多精力關注MoE相關系列發展。
Reference
- Clark, Aidan, et al. “Unified scaling laws for routed language models.” International conference on machine learning. PMLR, 2022.
- Fedus, William, Barret Zoph, and Noam Shazeer. “Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.” Journal of Machine Learning Research 23.120 (2022): 1–39.
- Xue, Fuzhao, et al. “Openmoe: An early effort on open mixture-of-experts language models.” arXiv preprint arXiv:2402.01739 (2024).
- Komatsuzaki, Aran, et al. “Sparse upcycling: Training mixture-of-experts from dense checkpoints.” arXiv preprint arXiv:2212.05055 (2022).
- Dense Training, Sparse Inference: Rethinking Training of Mixture-of-Experts Language Models https://arxiv.org/abs/2404.05567
- JetMoE https://research.myshell.ai/jetmoe
- Qwen 1.5 MoE https://qwenlm.github.io/blog/qwen-moe/
- Jamba https://www.ai21.com/blog/announcing-jamba
- Eliseev, Artyom, and Denis Mazur. “Fast inference of mixture-of-experts language models with offloading.” arXiv preprint arXiv:2312.17238 (2023).
- Lin, Bin, et al. “Moe-llava: Mixture of experts for large vision-language models.” arXiv preprint arXiv:2401.15947 (2024).
