概述

随着大模型技术的一直倒退，模型构造和参数量级疾速演变。大模型技术的利用层出不穷。大模型展示惊人成果，但训练和推理老本高，始终是微小挑战。模型稠密化能升高计算和存储耗费。近期以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:4train_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

原文链接

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