昆仑万维开源2千亿稀疏大模型天工MoE，全球首创能用4090推理

2024年6月3日，昆仑万维宣布开源2千亿MoE大模型，性能强劲，同时推断成本。Skywork-MoE基于昆仑万维开源的Skywork-13B模型中间检查点扩展到来， 是首个完整的MoE升级技术应用并落地的开源千亿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

• 模型推理代码：（支持8×4090服务器上8位计算加载推理） https://github.com/SkyworkAI/vllm

模型架构：

本次开源的Skywork-MoE模型属于天工3.0的研发模型系列，是其中的中档大小模型（Skywork-MoE-Medium），模型的总参数量为146B，激活参数量22B，共有16个专家，每个专家大小为13B，每次激活其中2个专家。

天工 3.0 还训练了 75B （Skywork-MoE-Small）和 400B （Skywork-MoE-Large）两档 MoE 模型，并不在涉外开源之列。

模型能力：

我们基于目前最主流模型得分得分了Skywork-MoE，在相同的激活参数量20B（推理计算量）下，Skywork-MoE的能力在行业前列，接近70B的密集模型。使得模型的推理成本有近3倍的下降。同时Skywork-MoE的总参数大小比DeepSeekV2的总参数大小要小1/3，用更小的参数规模做到了相近的能力。

技术创新：

为了解决 MoE 模型训练困难，性能化差的问题，泛借鉴于 Mixtral-MoE，Skywork-MoE 设计了两种训练优化算法：

1. Gating Logits 归一化操作

我们在门控层的令牌分发逻辑处新增了一个归一化操作，使得门控层的参数学习更加趋向于被选中的top-2专家，增加MoE对于模型top-2的置信度：

2. 自适应的 Aux 损耗

除了传统的固定系数（固定超参）的辅助损失，我们在教育部让训练的不同阶段让模型自适应的选择合适的辅助损失超参系数，从而导致代币率下降保持在合适的区间内，既能做到专家分散的平衡，又能让专家学习具备偏差化，从而提升模型整体的性能和泛化水平。在教育部让训练的前期，由于参数学习不够位，导致代币率下降太高（代币）分配差异较大），此时需要加大的aux损失有助于令牌负载平衡；在MoE训练的高级中，我们希望Expert之间仍保证一定的区分度避免，门控总线为随机分配Token，因此需要较低的aux损失降低纠偏。

训练基础设施

如何对 MoE 模型高效进行大规模训练是一个有难度的挑战，目前社区还没有一个最佳实践。Skywork-MoE 提出了两个重要的计算优化设计，从而在千卡集群上实现了 MFU 38% 的训练吞吐，其中 MFU 以 22B 的激活参数计算理论量。

1. 专家数据并行

区别于Megatron-LM社区已有的EP（Expert Parallel）和ETP（Expert Tensor Parallel）设计，我们提出了一种称为专家数据并行的玩具设计方案，这种玩具方案可以在专家数量较小时仍能的切分模型，对专家引入的all2all通信也可以最大程度的优化和屏蔽。初步于对EP GPU数量的限制和ETP在千卡集群上的低效率，EDP可以更好的解决大规模训练 MoE 的硬件痛点，同时 EDP 的设计简单、鲁棒、易扩展，可以较快的实现和验证。

一个最简单的 EDP 的例子，两种卡情况下 TP = 2， EP = 2，其中 Attention 部分采用 Tensor Parallel ， Expert 部分采用 Expert Parallel

2.非均匀切分流水工件

由于第一阶段的嵌入计算和最后阶段的损耗计算，以及管道缓冲区的均匀存在，流水负荷下切分层时的各阶段计算负载和显存负载负载较明显的不均衡情况。我们提出了非均匀的流水负载切分和重计算方式分配，使得总体的计算/显存负载更加均衡，继续10%左右的后续训练吞吐提升。

比较均匀切分和非均匀切分下的流水玩具：对于一个24层的LLM，(a)是均匀切分4个阶段，每个阶段的层数为：[6,6,6,6]。 (b) 是经过优化后的非均匀切分方式，雕塑5个阶段，每个阶段的层数量为：[5, 5, 5, 5, 4] ，在中间流水打满的阶段，非均匀切分的语音语音。

教育部技术诀窍

此外，Skywork-MoE 还通过一系列基于 Scaling Laws 的实验，探究哪些约束会影响 Upcycling 和 From Scratch 训练 MoE 模型的好坏。

一个可以遵循的经验规则是：如果训练MoE模型的FLOPs是训练密集模型的2倍以上，那么选择从头开始训练MoE会更好，否则的话，选择Upcycling训练MoE可以减少明显的训练成本。

第4090章 推理

Skywork-MoE 是目前能在 8×4090 服务器上推理的最大的开源 MoE 模型。8×4090 服务器一共有 192GB 的 GPU 显存，在 FP8 自定义下（权重占用 146GB），使用我们首创的非 Tensor Parallel 家具推理方式，Skywork-MoE 可以在合适的批量大小内达到 2200 tokens/s的吞吐。天工团队完整开源了相关的推理框架代码和安装环境，详情参见：https://github.com/SkyworkAI/Skywork-MoE

结语

我们希望本次开源的Skywork-MoE训练模型、技术报告和相关的实验结果可以给开源社区贡献更多的MoE训练经验和Know-how，包括模型结构、超参选择、技巧、训练推理加速等各方面，探索用前期的训练成本训练更大的模型，在走向AGI的道路上贡献一点力量。