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

2024年6月3日，昆仑万维宣布开源 2 千亿稀疏大模型 Skywork-MoE ， 性能强劲， 同时推理成本更低。Skywork-MoE 基于之前昆仑万维开源的 Skywork-13B 模型中间 checkpoint 扩展而来，是首个完整将 MoE Upcycling 技术应用并落地的开源千亿 MoE大模型，也是首个支持用单台 4090 服务器推理的开源千亿 MoE大模型。

模型架构:

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

天工 3.0 还训练了 75B 和 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和 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， 是均匀切分成 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 是目前能在 8x4090 服务器上推理的最大的开源 MoE 模型。8x4090 服务器一共有 192GB 的 GPU 显存，在 FP8 量化下，使用我们首创的非均匀 Tensor Parallel 并行推理方式，Skywork-MoE 可以在合适的 batch size 内达到 2200 tokens/s 的吞吐。天工团队完整开源了相关的推理框架代码和安装环境。

结语

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