单个4090可推理，2000亿稀疏大模型「天工MoE」开源

发布时间:2024-06-05 09:02:54 · 责任编辑:字母汇 · 浏览量:65 次

在大模型浪潮中，训练和部署最先进的密集 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 倍的下降。

 

值得注意的是， Skywork-MoE 的总参数大小比 DeepSeekV2 的总参数大小要小 1/3，用更小的参数规模做到了相近的能力。

技术创新

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

Gating Logits 归一化操作

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

 

自适应的 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 降低纠偏。

 

训练 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 的设计简单、鲁棒、易扩展，可以较快的实现和验证。

 

一个最简单的 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% 左右的端到端训练吞吐提升。

 

比较均匀切分和非均匀切分下的流水并行气泡：对于一个 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 模型的好坏。

 

 

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

原文始发于微信公众号（机器之心）：单个4090可推理，2000亿稀疏大模型「天工MoE」开源