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

模型架構：

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

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

模型能力：

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

技術創新：

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

1.?Gating Logits 歸一化操作

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

2.?自適應的 Aux Loss

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

訓練 Infra

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

1. Expert Data Parallel

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

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

2.?非均勻切分流水并行

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

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

MoE Know-how

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

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

4090 推理

Skywork-MoE 是目前能在 8×4090 服務器上推理的最大的開源 MoE 模型。8×4090 服務器一共有 192GB 的 GPU 顯存，在 FP8 量化下（weight 占用 146GB），使用我們首創的非均勻 Tensor Parallel 并行推理方式，Skywork-MoE 可以在合適的 batch size 內達到 2200 tokens/s 的吞吐。

結語

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