2024年6月3日，昆侖萬維宣布開源 2 千億稀疏大模型 Skywork-MoE ， 性能強勁， 同時推理成本更低。Skywork-MoE 基于之前昆侖萬維開源的 Skywork-13B 模型中間 checkpoint 擴展而來，是首個完整將 MoE Upcycling 技術(shù)應(yīng)用并落地的開源千億 MoE大模型，也是首個支持用單臺 4090 服務(wù)器推理的開源千億 MoE大模型。
模型架構(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 模型高效的進行大規(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 計算負載和顯存負載均有較明顯的不均衡情況。我們提出了非均勻的流水并行切分和重計算 Layer 分配方式，使得總體的計算/顯存負載更均衡，約有 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 是目前能在 8×4090 服務(wù)器上推理的最大的開源 MoE 模型。8×4090 服務(wù)器一共有 192GB 的 GPU 顯存，在 FP8 量化下（weight 占用 146GB），使用我們首創(chuàng)的非均勻 Tensor Parallel 并行推理方式，Skywork-MoE 可以在合適的 batch size 內(nèi)達到 2200 tokens/s 的吞吐。
結(jié)語
我們希望本次開源的 Skywork-MoE 模型、技術(shù)報告和相關(guān)的實驗結(jié)果可以給開源社區(qū)貢獻更多的 MoE 訓(xùn)練經(jīng)驗和 Know-how，包括模型結(jié)構(gòu)、超參選擇、訓(xùn)練技巧、訓(xùn)練推理加速等各方面， 探索用更低的訓(xùn)練推理成本訓(xùn)更大更強的模型，在通往 AGI 的道路上貢獻一點力量。?
