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

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

開源地址：

Skywork-MoE 的模型權(quán)重、技術(shù)報(bào)告完全開源，免費(fèi)商用，無(wú)需申請(qǐng)：

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

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

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

? 模型開源倉(cāng)庫(kù)：https://github.com/SkyworkAI/Skywork-MoE

? 模型技術(shù)報(bào)告：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 個(gè) Expert，每個(gè) Expert 大小為 13B，每次激活其中的 2 個(gè) Expert。

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

模型能力：

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

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

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

1. Gating Logits 歸一化操作

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

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

有別于傳統(tǒng)的固定系數(shù)（固定超參）的 aux loss， 我們?cè)?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 分布差異太大），此時(shí)需要較大的 aux loss 幫助 token load balance；在 MoE 訓(xùn)練的后期，我們希望 Expert 之間仍保證一定的區(qū)分度，避免 Gating 傾向?yàn)殡S機(jī)分發(fā) Token，因此需要較低的 aux loss 降低糾偏。

訓(xùn)練 Infra

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

1. Expert Data Parallel

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

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

2. 非均勻切分流水并行

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

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

MoE Know-how

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

一個(gè)可以遵循的經(jīng)驗(yàn)規(guī)則是：如果訓(xùn)練 MoE 模型的 FLOPs 是訓(xùn)練 Dense 模型的 2 倍以上，那么選擇 from Scratch 訓(xùn)練 MoE 會(huì)更好，否則的話，選擇 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 的吞吐。天工團(tuán)隊(duì)完整開源了相關(guān)的推理框架代碼和安裝環(huán)境，詳情參見：https://github.com/SkyworkAI/Skywork-MoE

結(jié)語(yǔ)

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