昆侖萬(wàn)維開(kāi)源2千億稀疏大模型天工MoE，全球首創(chuàng )能用4090推理

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

　　開(kāi)源地址：

　　Skywork-MoE 的模型權重、技術(shù)報告完全開(kāi)源，免費商用，無(wú)需申請：

　　• 模型權重下載：

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

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

　　• 模型開(kāi)源倉庫：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

　　模型架構：

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

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

　　模型能力：

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

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

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

　　1. Gating Logits 歸一化操作

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

　　2. 自適應的 Aux Loss

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

　　訓練 Infra

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

　　1. Expert Data Parallel

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

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

　　2. 非均勻切分流水并行

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

　　比較均勻切分和非均勻切分下的流水并行氣泡：對于一個(gè) 24 層 Layer 的 LLM， （a） 是均勻切分成 4 個(gè) stage，每個(gè) stage 的 layer 數量是：［6， 6， 6， 6］。（b） 是經(jīng)過(guò)優(yōu)化后的非均勻切分方式，切成 5 個(gè) stage， 每個(gè) stage 的 layer 數量是：［5， 5， 5， 5， 4］ ， 在中間流水打滿(mǎn)的階段，非均勻切分的氣泡更低。

　　MoE Know-how

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

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

　　4090 推理

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

　　結語(yǔ)

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