2024年6月3日，昆侖萬維宣布開源2000億稀疏大模型Skywork-MoE，性能強勁，同時推理成本更低。Skywork-MoE基于之前昆侖萬維開源的Skywork-13B模型中間checkpoint擴展而來，是首個完整將MoE Upcycling技術(shù)應(yīng)用并落地的開源千億MoE大模型，也是首個支持用單臺4090服務(wù)器推理的開源千億MoE大模型。

開源地址：

Skywork-MoE的模型權(quán)重、技術(shù)報告完全開源，免費商用，無需申請：

•模型權(quán)重下載：

https://huggingface.co/Skywork/Skywork-MoE-base

https://huggingface.co/Skywork/Skywork-MoE-Base-FP8

•模型開源倉庫：https://github.com/SkyworkAI/Skywork-MoE

•模型技術(shù)報告：https://github.com/SkyworkAI/Skywork-MoE/blob/main/skywork-moe-tech-report.pdf

•模型推理代碼：（支持8x4090服務(wù)器上8 bit量化加載推理）https://github.com/SkyworkAI/vllm

模型架構(gòu)：

本次開源的Skywork-MoE模型隸屬于天工3.0的研發(fā)模型系列，是其中的中檔大小模型（Skywork-MoE-Medium），模型的總參數(shù)量為146B，激活參數(shù)量22B，共有16個Expert，每個Expert大小為13B，每次激活其中的2個Expert。

天工3.0還訓(xùn)練了75B（Skywork-MoE-Small）和400B（Skywork-MoE-Large）兩檔MoE模型，并不在此次開源之列。

模型能力：

基于目前各大主流模型評測榜單評測了Skywork-MoE，在相同的激活參數(shù)量20B（推理計算量）下，Skywork-MoE能力在行業(yè)前列，接近70B的Dense模型。使得模型的推理成本有近3倍的下降。同時Skywork-MoE的總參數(shù)大小比DeepSeekV2的總參數(shù)大小要小1/3，用更小的參數(shù)規(guī)模做到了相近的能力。

技術(shù)創(chuàng)新：

為了解決MoE模型訓(xùn)練困難，泛化性能差的問題，相較于Mixtral-MoE，Skywork-MoE設(shè)計了兩種訓(xùn)練優(yōu)化算法：

1.Gating Logits歸一化操作

昆侖萬維在Gating Layer的token分發(fā)邏輯處新增了一個normalization操作，使得Gating Layer的參數(shù)學(xué)習(xí)更加趨向于被選中的top-2 experts，增加MoE模型對于top-2的置信度：

2.自適應(yīng)的Aux Loss

有別于傳統(tǒng)的固定系數(shù)（固定超參）的aux loss，在MoE訓(xùn)練的不同階段讓模型自適應(yīng)地選擇合適的aux loss超參系數(shù)，從而讓Drop Token Rate保持在合適的區(qū)間內(nèi)，既能做到expert分發(fā)的平衡，又能讓expert學(xué)習(xí)具備差異化，從而提升模型整體的性能和泛化水平。在MoE訓(xùn)練的前期，由于參數(shù)學(xué)習(xí)不到位，導(dǎo)致Drop Token Rate太高（token分布差異太大），此時需要較大的aux loss幫助token load balance；在MoE訓(xùn)練的后期，昆侖萬維希望Expert之間仍保證一定的區(qū)分度，避免Gating傾向為隨機分發(fā)Token，因此需要較低的aux loss降低糾偏。

訓(xùn)練Infra

如何對MoE模型高效地進(jìn)行大規(guī)模分布式訓(xùn)練是一個有難度的挑戰(zhàn)，目前社區(qū)還沒有一個最佳實踐。Skywork-MoE提出了兩個重要的并行優(yōu)化設(shè)計，從而在千卡集群上實現(xiàn)了MFU 38%的訓(xùn)練吞吐，其中MFU以22B的激活參數(shù)計算理論計算量。

1.Expert Data Parallel

區(qū)別于Megatron-LM社區(qū)已有的EP（Expert Parallel）和ETP（Expert Tensor Parallel）設(shè)計，昆侖萬維提出了一種稱之為Expert Data Parallel的并行設(shè)計方案，這種并行方案可以在Expert數(shù)量較小時仍能高效地切分模型，對Expert引入的all2all通信也可以最大程度地優(yōu)化和掩蓋。相較于EP對GPU數(shù)量的限制和ETP在千卡集群上的低效，EDP可以較好地解決大規(guī)模分布式訓(xùn)練MoE的并行痛點，同時EDP的設(shè)計簡單、魯棒、易擴展，可以較快地實現(xiàn)和驗證。

一個最簡單的EDP的例子，兩卡情況下TP=2，EP=2，其中Attention部分采用Tensor Parallel，Expert部分采用Expert Parallel

2.非均勻切分流水并行

由于first stage的Embedding計算和last stage的Loss計算，以及Pipeline Buffer的存在，流水并行下均勻切分Layer時的各stage計算負(fù)載和顯存負(fù)載均有較明顯的不均衡情況。昆侖萬維提出了非均勻的流水并行切分和重計算Layer分配方式，使得總體的計算/顯存負(fù)載更均衡，約有10%的端到端訓(xùn)練吞吐提升。

比較均勻切分和非均勻切分下的流水并行氣泡：對于一個24層Layer的LLM，(a)是均勻切分成4個stage，每個stage的layer數(shù)量是：[6,6,6,6].(b)是經(jīng)過優(yōu)化后的非均勻切分方式，切成5個stage，每個stage的layer數(shù)量是：[5,5,5,5,4]，在中間流水打滿的階段，非均勻切分的氣泡更低。

MoE Know-how

此外，Skywork-MoE還通過一系列基于Scaling Laws的實驗，探究哪些約束會影響Upcycling和From Scratch訓(xùn)練MoE模型的好壞。

一個可以遵循的經(jīng)驗規(guī)則是：如果訓(xùn)練MoE模型的FLOPs是訓(xùn)練Dense模型的2倍以上，那么選擇from Scratch訓(xùn)練MoE會更好，否則的話，選擇Upcycling訓(xùn)練MoE可以明顯減少訓(xùn)練成本。

4090推理

Skywork-MoE是目前能在8x4090服務(wù)器上推理的最大的開源MoE模型。8x4090服務(wù)器一共有192GB的GPU顯存，在FP8量化下（weight占用146GB），使用昆侖萬維首創(chuàng)的非均勻Tensor Parallel并行推理方式，Skywork-MoE可以在合適的batch size內(nèi)達(dá)到2200 tokens/s的吞吐。天工團隊完整開源了相關(guān)的推理框架代碼和安裝環(huán)境，詳情參見：https://github.com/SkyworkAI/Skywork-MoE

結(jié)語

希望本次開源的Skywork-MoE模型、技術(shù)報告和相關(guān)的實驗結(jié)果可以給開源社區(qū)貢獻(xiàn)更多的MoE訓(xùn)練經(jīng)驗和Know-how，包括模型結(jié)構(gòu)、超參選擇、訓(xùn)練技巧、訓(xùn)練推理加速等各方面，探索用更低的訓(xùn)練推理成本訓(xùn)更大更強的模型，在通往AGI的道路上貢獻(xiàn)一點力量。