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

　　2024年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 模型隶属于天工 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% 左右的端到端训练吞吐提升。