昆仑万维开源 2 千亿稀疏大模型天工 MoE,能用 4090 推理

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

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

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 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 设计了两种训练优化算法:

在 Gating Layer 的 token 分发逻辑处新增了一个 normalization 操作,使得 Gating Layer 的参数学习更加趋向于被选中的 top-2 experts,增加 MoE 模型对于 top-2 的置信度:

有别于传统的固定系数(固定超参)的 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 降低纠偏。

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

区别于 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

由于 first stage 的 Embedding 计算和 last stage 的 Loss 计算,以及 Pipeline Buffer 的存在, 流水并行下均匀切分 Layer 时的各 stage 计算负载和显存负载均有较明显的不均衡情况。项目团队提出了非均匀的流水并行切分和重计算 Layer 分配方式,使得总体的计算/显存负载更均衡,约有 10% 左右的端到端训练吞吐提升。

比较均匀切分和非均匀切分下的流水并行气泡:对于一个 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 Laws 的实验,探究哪些约束会影响 Upcycling 和 From Scratch 训练 MoE 模型的好坏。

一个可以遵循的经验规则是:如果训练 MoE 模型的 FLOPs 是训练 Dense 模型的 2 倍以上,那么选择 from Scratch 训练 MoE 会更好,否则的话,选择 Upcycling 训练 MoE 可以明显减少训练成本。

Skywork-MoE 是目前能在 8×4090 服务器上推理的最大的开源 MoE 模型。8×4090 服务器一共有 192GB 的 GPU 显存,在 FP8 量化下(weight 占用 146GB),使用我们首创的非均匀 Tensor Parallel 并行推理方式,Skywork-MoE 可以在合适的 batch size 内达到 2200 tokens/s 的吞吐。

天工团队完整开源了相关的推理框架代码和安装环境,详情参见:https://github.com/SkyworkAI/Skywork-MoE

未经允许不得转载:岩猫星空网 » 昆仑万维开源 2 千亿稀疏大模型天工 MoE,能用 4090 推理