單個4090可推理，2000億稀疏大模型「天工MoE」開源

發布于 2024-6-4 13:04

瀏覽

0收藏

在大模型浪潮中，訓練和部署最先進的密集 LLM 在計算需求和相關成本上帶來了巨大挑戰，尤其是在數百億或數千億參數的規模上。為了應對這些挑戰，稀疏模型，如專家混合模型（MoE），已經變得越來越重要。這些模型通過將計算分配給各種專門的子模型或「專家」，提供了一種經濟上更可行的替代方案，有可能以極低的資源需求達到甚至超過密集型模型的性能。

6 月 3 日，開源大模型領域又傳來重要消息：昆侖萬維宣布開源 2 千億稀疏大模型 Skywork-MoE，在保持性能強勁的同時，大幅降低了推理成本。

Skywork-MoE 基于此前昆侖萬維開源的 Skywork-13B 模型中間 checkpoint 擴展而來，是首個完整將 MoE Upcycling 技術應用并落地的開源千億 MoE 大模型，也是首個支持用單臺 4090 服務器推理的開源千億 MoE 大模型。

讓大模型社區更為關注的是，Skywork-MoE 的模型權重、技術報告完全開源，免費商用，無需申請。

模型權重下載地址：

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

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

模型開源倉庫：https://github.com/SkyworkAI/Skywork-MoE
模型技術報告：https://github.com/SkyworkAI/Skywork-MoE/blob/main/skywork-moe-tech-report.pdf
模型推理代碼：（支持 8x4090 服務器上 8 bit 量化加載推理） https://github.com/SkyworkAI/vllm

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

完整相關的推理框架代碼和安裝環境見：https://github.com/SkyworkAI/Skywork-MoE

Skywork-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 倍的下降。

單個4090可推理，2000億稀疏大模型「天工MoE」開源-AI.x社區

值得注意的是， Skywork-MoE 的總參數大小比 DeepSeekV2 的總參數大小要小 1/3，用更小的參數規模做到了相近的能力。

技術創新

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

Gating Logits 歸一化操作

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

單個4090可推理，2000億稀疏大模型「天工MoE」開源-AI.x社區

自適應的 Aux Loss

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

單個4090可推理，2000億稀疏大模型「天工MoE」開源-AI.x社區

訓練 Infra

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

Expert Data Parallel

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

單個4090可推理，2000億稀疏大模型「天工MoE」開源-AI.x社區

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

非均勻切分流水并行

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

單個4090可推理，2000億稀疏大模型「天工MoE」開源-AI.x社區

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

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

單個4090可推理，2000億稀疏大模型「天工MoE」開源-AI.x社區