关于云计算:基于-MegatronCore-的稀疏大模型训练工具阿里云MoE-大模型最佳实践

概述

随着大模型技术的一直倒退，模型构造和参数量级疾速演变。大模型技术的利用层出不穷。大模型展示惊人成果，但训练和推理老本高，始终是微小挑战。模型稠密化能升高计算和存储耗费。近期以 Mixtral 为代表的 MoE（多专家混合）大模型证实了稠密 MoE 技术能大幅升高计算量、晋升推理速度，模型成果甚至超过同规模浓密模型。阿里云 PAI 和 NVIDIA 团队深刻单干，基于 Megatron-Core MoE 框架，解决了 MoE 大模型训练落地时会遇到的可拓展性、易用性、功能性以及收敛精度等外围问题，在上游工作上获得了很好的模型成果。PAI 团队将上述 MoE 训练框架和技术与阿里云 AI 平台产品深度整合，使云上大模型用户能不便地进行 MoE 大模型的训练和部署。

基于 Megatron-Core 的 MoE 训练工具链

这一章节中，咱们从 MoE 算法介绍，Megatron-Core MoE 训练框架、以及云平台工具三个方面理解撑持 MoE 落地背地的根底能力。

MoE 算法介绍

MoE 的全称是 Mixture of Experts，其中的 Expert 对应的是 Transfomrer 模型的 MLP 层，在训练的时候从多个 MLP 中选取一个 MLP 进行激活 [1]（如下图所示）。这意味着模型能够在不减少 FLOPs 的状况下，通过减少 MLP 模块的数量来减少模型参数量级，进而晋升模型在上游工作上的成果。采纳 MoE 后的稠密 Transformer 模型和等同品质（验证集 loss 以及 Zero-shot NLU 上游工作性能）的浓密模型相比有较大幅度的训练和推理吞吐性能晋升。MoE 次要由以下两个要害局部组成：

• 稠密 MoE 层: 这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层蕴含若干 experts 组成，每个 experts 自身是一个独立的神经网络。在理论利用中这些专家通常是前馈网络 (FFN)。
• Router 路由: 这个局部用于决定哪些 tokens 被发送到哪个专家。例如在下图中，“More”这个 tokens 可能被发送到第二个专家，而“Parameters”这个 tokens 被发送到第一个专家。

接着咱们探讨下 MoE 中的 token 负载平衡问题。如上图所示的 Top-1 Routing 算法，每个输出 hidden_states 被 router 调配了 4 个 experts 中的 1 个用于前向推理计算，而后其输入被加权求和汇总后送入 Transformer 的下一个处理单元。这种路由形式会呈现多数受欢迎的 experts 须要解决很多 token，而其余 experts 仅需解决极少数量的 token 的状况，这导致解决极少 token 的那些 experts 无奈取得足够多的信息，从而导致训练效率大幅升高。

为了缓解这种 token 到 expert 负载不平衡 (Load Imbalancing) 的问题，能够引入了一个辅助损失函数，旨在激励给予所有专家雷同的重要性。这个损失函数确保所有专家接管到大抵相等数量的训练样本，从而均衡了专家之间的抉择。另外也能够通过 drop tokens 的形式失去缓解。首先定义一个 expert capacity，即为一个 expert 被调配到的 token 的容量。如果调配到一个 expert 的 token 数超出了一个固定容量，则多余的 tokens 会被丢掉。这些被丢掉的 tokens 不参加和 experts 的矩阵乘运算，间接通过一个残差连贯进入到下一个处理单元。如果一个 expert 没有被调配到其容量下限内的足够多的 tokens，则须要采纳 padding 的形式填充到容量下限。

Mixtral-8x7b 模型提供了另外一种思路来缓解负载不平衡的问题，其采纳的 Megablocks [2]论文中提出的 dropless moe 算法。如下图所示，首先 tokens 通过 router 被调配到具体的 experts，token 到 expert 的分配关系被存储到 expert indices 中，调配置信度被存储到 Probabilities 中。其次依据事后设置好的 expert 的容量将 tokens 分组搁置，超出容量的 tokens 被丢掉。接着将分组中的 tokens 和 experts 进行批量矩阵乘运算。最初输入后果依照调配置信度进行加权求和（如下右边公式)，其调配置信度的计算采纳的是 normalized softmax（如下左边公式）。

Megablock 论文认为这个 drop token 计划会导致成果降落，于是在此基础上引入块稠密算子来实现 dropless 性能。如下图所示，dropable 的 moe 算法在 experts 之间并行批量矩阵乘（如下图 A 所示）能够用一种块对角矩阵乘来等价表白（如下图 B 所示）。Megablock 设计了一个新的算子容许变长的输出（如下图 C 所示），而后输入到一个稠密对角矩阵上。引入了这个变长输出就意味着不须要 drop 的形式就能解决超出容量的 tokens。相似的，Megatron-Core MoE 应用 GroupedGEMM 解决了多 expert 变长输出的问题。

Megatron-Core MoE 训练框架

Megatron-Core 是 NVIDIA 推出的一个成熟且轻量级的大规模 LLM 训练框架，它蕴含了训练大规模 LLM 模型所需的所有关键技术，例如各类模型并行的反对、算子优化、通信优化、显存优化以及 FP8 低精度训练等。Megatron-Core 不仅继承了前代 Megatron-LM 的优良个性，还在代码品质、稳定性、性能丰富性和测试覆盖率上进行了全面晋升。更重要的是，Megatron-Core 在设计上更加解耦和模块化，不仅进步了 LLM 训练的效率和稳定性，也为二次开发和摸索新的 LLM 架构提供了更大的灵活性。

基于 Megatron-Core，钻研人员能够在大规模集群上高效地训练各种大型 LLM 模型，其中也包含对 MoE 模型的摸索。MoE 的建模形式也带来了参数量的急剧回升、MoE 层简单的通信逻辑、Routing 等额定的开销，都给训练框架带来了不小的挑战。其中首当其冲的就是如何在多个 GPU 间高效的拓展 MoE 训练。

在最新公布的 Megatron-Core v0.5 中，Megatron-Core 正式推出了对大规模 MoE 模型训练的反对，也就是 Megatron-Core MoE，涵盖了并行性、路由和负载平衡、性能优化、Token 散发机制等多个 feature。以下是一些值得关注的个性：

并行性（Parallelism）

Megatron-Core MoE 反对专家并行（Expert Parallel），这是一种专门为 MoE 模型设计的并行办法。在这种并行化策略中，不同的 Rank 负责解决其中一个或多个专家的计算。

此外，Megatron-Core MoE 还反对 3D 并行（Data Parallel, Tensor Parallel, Pipeline Parallel, Sequence Parallel）。对于更大的 MoE 模型，Megatron-Core MoE 也反对将专家并行（EP）与其它的并行 DP/TP/PP/SP/Distributed Optimizer 联合应用。

将来，Megatron-Core MoE 还将反对 Context Parallel，以反对更长序列的训练。

Token 散发机制（Token Dispatch Mechanism）

Megatron-Core MoE 目前提供了对 Dropless MoE 的反对（不进行 Token 抛弃）。行将退出对 Token drop MoE 的反对。

路由和负载平衡（Router and Load Balancing）

Megatron-Core MoE 提供了多种路由类型，包含通用的 Top-K router 和行将推出的 Expert Choice router。

在负载平衡算法方面，反对 Sinkhorn（S-BASE）、Z-Loss，以及 Load balancing Loss。

Grouped GEMM

Megatron-Core MoE 开发了 GroupedGEMM 来解决多 Experts 变长输出这一问题。当每个 Rank 有多个专家时，Megatron-Core MoE 利用自 CUTLASS 2.8 引入的 Grouped GEMM 个性，将多个部分（可能是较小的）GEMM 操作合并为单个 GroupedGEMM kernel，可能大幅度提高 SM 利用率和性能。

同时，Megatron-Core MoE 还将局部效率较低的操作替换为优化后的 CUDA Kernel，如 Sinkhorn、local token permutation/unpermutation 等。

行将推出的性能与优化

近期将公布的局部新个性包含：

• 上下文并行（Context Parallel）
• FP8 低精度训练
• FP8 Grouped GEMM
• Token-Drop MoE

通过这些个性，Megatron-Core MoE 为用户提供了一个弱小的 MoE 训练框架，以在大规模集群上高效地训练和摸索大型 MoE 模型。随着更多功能的推出，咱们期待 Megatron-Core MoE 将在将来的 LLM 钻研和利用中施展更大的作用。

MoE 平台工具

在阿里云 PAI 灵骏分布式集群上运行的基于 Megatron 的 MoE 训练工具由三局部组成（如下图所示）。从上到下顺次是 PAI 平台，PAI-Megatron-Patch（https://github.com/alibaba/Pai-Megatron-Patch）以及 NVIDIA Megatron-Core(https://github.com/NVIDIA/Megatron-LM)。PAI 平台的 DSW 产品是为 AI 开发者量身定制的云端机器学习交互式开发 IDE，随时随地开启 Notebook 疾速读取数据、开发算法、训练及部署模型。DLC 产品则为您提供灵便、稳固、易用和极致性能的多机多卡深度学习训练环境。Pai-Megatron-Patch 是各类开源大模型和 Megatron 训练减速引擎之间的“桥梁”，为用户提供用 Megatron 训练开源大模型的易用性以及 LLM 算法场景定制化的灵活性。NVIDIA 开发的 Megatron 以其良好的业界口碑博得了泛滥大模型用户的青眼，其算子拆分和流水并行等模型并行技术曾经成为行业标准。基于 DeepSpeed Zero- 1 改良的 Distributed Optimizer 在大规模分布式训练过程中的降显存方面施展巨大作用。集成最新的算子交融，Transformer Engine 以及 Flash Attention 等训练减速技术更是将大模型的吞吐性能 & 显卡利用率晋升到了极致[3].。

以代码生成这个 LLM 利用场景为例，研发团队首先在 DSW 中开发 & 调试用于该场景的定制化数据处理流程，应用定制化的分词器对原始代码数据进行 ID 化，同时调试模型库中对应的代码生成模型比方 deepseek。接着在 DSW 中跑通实用于多机多卡训练的 HuggingFace 到 Megatron 的权重转换流程并生成 Megatron 可加载的权重。而后在 DLC 中配置 Patch 提供的训练启动脚本，将训练所需的超参数比方学习率以及训练减速开关比方 Flash Attention 等通过脚本传入 Megatron 训练引擎。

MoE 训练成果展现

咱们以最近社区热门的 Mixtral-8x7B 稠密 MoE 大模型为例，从 HuggingFace 到 Megatron 的零样本 loss 对齐，训练 loss 收敛曲线，代码生成上游工作评测，训练吞吐性能这四个方面向用户展现训练工具的预期成果，以此来验证工具的可靠性和稳定性。为不便用户测试，咱们提供了曾经解决好的 idxmap 格局的 Wudao 数据集，转换前的 HuggingFace 模型以及转换后的 Megatron-Core 格局的模型供下载应用，具体文件大小和下载链接如下所示。

https://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/mistral-ckpts/Mixtral-8x7B-v0.1-to-mcore-tp4-ep4.tar.zsthttps://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/mistral-ckpts/Mixtral-8x7B-v0.1.tar.zsthttps://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/mistral-datasets/wudao_mistralbpe_content_document.binhttps://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/mistral-datasets/wudao_mistralbpe_content_document.idxhttps://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/mistral-ckpts/Mixtral-8x7B-v0.1-to-mcore-tp4-ep4.tar.zsthttps://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/mistral-ckpts/Mixtral-8x7B-v0.1.tar.zsthttps://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/mistral-datasets/wudao_mistralbpe_content_document.binhttps://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/mistral-datasets/wudao_mistralbpe_content_document.idx

测试镜像地址请采纳：http://dsw-registry.cn-wulanchabu.cr.aliyuncs.com/pai/pytorch…

HF 到 Megatron 模型权重转换

因为从头预训练大模型所耗费的算力老本较高，因而用户广泛采纳加载 HuggingFace 提供曾经训练好的模型权重持续预训练或者微调。这就须要先把 HuggingFace 的 Transformer 构造的权重转换到 Megatron 的 Transformer 可辨认并加载的格局。下图显示的两边的算子之间的对照关系。

这里须要留神的是当 TP>1 的时候，对权重矩阵的切分与合并操作应合乎 Megatron 的建模形式。以 mlp 层的 gate_proj 和 up_proj 为例，按 TP= 2 切分后两个算子的左半局部先组合到一起，而后右半局部再组合到一起，而后拼接造成 Megatron 可辨认的 dense_h_to_4h 操作，如下图所示。这种拆分 & 合并的程序一旦呈现谬误，会导致两边的前向推理无奈对齐。

另外波及到 experts 的调配的时候也须要分外小心。咱们比照 dense 的转换在 moe 模型的 convert 脚本中新增了两个参数（expert_model_parallel_size 和 world_size（总卡数））来对 experts 的散布进行精细化解决。比方 mixtral 的 experts 的总数是 8，那么在 16 张卡有三种转换状况, tp4ep4, tp8ep2, tp2ep8，依照每种状况来调配 moe 层的权重。以 tp4ep4 为例，权重转换后的文件夹命名如下表所示：

在每个文件里存储两个 experts 的 FFN 权重，每个 tp_rank 有 8 个 experts 来做路由，如下图所示：

最初咱们通过 zeroshot loss 来掂量转换后的模型是否正确，从下表中能够看到，转换前后的误差根本能够管制在十分小的范畴内。下表中箭头精度转换（fp16->bf16）的含意是 checkpoint 格局是 fp16, 运行的时候加载成 bf16。

基于 MoE 模型持续预训练 & 微调收敛

咱们摸索了三个应用场景的收敛状况别离是：从头预训练，加载 Mixtral-8x7B-v0.1 持续预训练以及加载 Mixtral-8x7B-Instruct-v0.1 微调。具体操作办法可参考咱们提供的 Mixtral-8x7B 稠密大模型最佳实际，外面具体介绍了如何启动训练的细节。在从头预训练的试验中，咱们采纳的训练配置是：A800 机型 80G 显存，两机 16 卡，global_size=256，lr=1e-4, seq_len=2048, TP=4。训练 24 个小时到 2.4k 个 step 后的 loss 能够到收敛到 1.9 的曲线如下所示：

* 仅做技术参考和探讨

在加载 Mixtral-8x7B-v0.1 持续预训练的试验中，咱们采纳的训练配置是：A800 机型 80G 显存，两机 16 卡，global_size=256，lr=5e-5, seq_len=2048, TP=4。训练 18 个小时到 2k 个 step 后的 loss 收敛曲线如下所示：

* 仅做技术参考和探讨

在加载 Mixtral-8x7B-Instruct-v0.1 微调的试验中，咱们在代码生成这个上游工作中设计了三种微调场景别离是：base 模型纯自回归训练（橙色曲线，计算全副 loss），instruct 模型纯自回归训练（蓝色曲线，计算全副 loss) 和 instruct 模型纯自回归训练（红色曲线，计算 answer loss）。从三种状况的收敛 loss 曲线比照能够看出根本都收敛到 0.5 高低。其中采纳 base 模型来微调的 loss 下载不如其余两个，这是合乎预期的。

代码生成工作评测

代码生成是指通过自然语言形容生成相应的可执行代码，从而晋升开发人员的工作效率。随着 LLM 模型的一直倒退，在代码生成工作上也获得了突破性的停顿，如 github copilot 可能晋升 55% 的开发效率。咱们针对代码生成这个上游畛域工作，基于 Evol-Instruct-Code-80K 数据对 Mixtral-8x7B-Instruct-v0.1 模型进行微调，并测试了在微调前后模型 Human-eval 上的通过率，和 LLaMA2 等开源通用模型比照能够看到微调后的模型在代码生成的能力上有显著的劣势，如下表所示。

咱们在 A800 机型 80G 显存，两机 16 卡下采纳的微调参数如下，具体操作办法可参考咱们提供的 Mixtral-8x7B 稠密大模型最佳实际。

micro batch size: 1 
global batch size: 128 
lr: 1e-5 
min_lr: 1e-6 
seqlen:2048 
padlen:2048 
tp:4 
pp:1 
ep:4
train_iter:2500

* 仅做技术参考和探讨

在评测时，咱们应用如下的 prompt 来生成代码：

f"[INST] Create a Python script for this problem:{question} [/INST]"

训练吞吐性能

笔者注：Megatron-Core MoE 在短期内还会退出对于通信、计算以及 Token dispatch 的多项性能优化，以后的测试数据不代表最终的性能数据。

在吞吐速度评测环节，阿里 PAI 团队调研了 Megablocks(https://github.com/stanford-futuredata/Megablocks)中的 dmoe 实现，一方面是因为 Mixtral-8x7b 论文中说采纳的是 Megablocks 的框架训练的 MoE 模型，另一方面咱们也想摸索下在雷同 Megatron 平台底座上，哪个 MoE 实现形式对训练吞吐减速成果更好。Megablocks 论文中提供了同 Dense 模型的收敛速度的比拟，如下图所示的具备 2.4x 的收敛减速能力，咱们冀望基于 Megatron-Core 的实现计划也具备同样甚至更好的减速成果。

吞吐性能比照试验采纳的是阿里云 PAI 灵骏平台上 A800 机型 80G 显存，两机 16 卡运行环境。咱们发现采纳 Mixtral-8x7B 模型配置时，Megablocks 的 seqlen 如果等于 1024/2048 会呈现 OOM。因而咱们将 Megablocks 和 Megatron-Core 的 seqlen 都设置为 512。同时咱们关上 –moe-grouped-gemm 开启 GroupedGEMM 晋升多 Experts 时的 GPU 利用率，从下表的吞吐数据能够看出，以后 Megatron-Core 的吞吐速度比 Megablock 快 10% 左右。此外，因为咱们得悉 Megatron-Core 以后正在推动 MoE 的性能优化，因而以后的数据仅供参考。

总结

在基于 Megatron-Core 的稠密大模型训练工具：阿里云 MoE 大模型最佳实际开发过程中，咱们围绕稠密大模型训练测试了以下核心技术的性能：

• MoE 根底技术平台：基于 Megatron-Core MoE 的多重训练减速技术的可靠性。
• MoE 落地 Pipeline：HF 到 Megatron 的模型权重转换在持续预训练 & 微调以及代码生成上游工作中的成果。

后续在 Pai-Megatron-Patch 中还会陆续放出更多高质量的大模型最佳实际，敬请期待。

参考文献

[1]. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

[2]. Megablocks: Efficient Sparse Training with Mixture-of-Experts

[3]. Reducing Activation Recomputation in Large Transformer Models

作者：李鹏，颜子杰，王明，颜海强，刘振寰，黄俊

单位：阿里云人工智能平台 PAI，NVIDIA DevTech Team

原文链接

本文为阿里云原创内容，未经容许不得转载。