FASTMOE_A FAST MIXTURE-OF-EXPERT TRAININGSYSTEM

研究背景

1. 背景介绍:
   这篇文章的研究背景是混合专家（MoE）模型在扩大语言模型规模方面具有巨大潜力，能够将模型参数扩展到万亿级别。然而，训练万亿规模的MoE模型需要算法和系统的协同设计，以支持高性能的分布式训练系统。现有的唯一平台依赖于Google的硬件（TPU）和软件（Mesh TensorFlow）堆栈，并且不对外开放，特别是对GPU和PyTorch社区不可用。
2. 研究内容:
   该问题的研究内容包括开发一个基于PyTorch的分布式MoE训练系统FastMoE，提供灵活的模型设计和易于适应不同应用（如Transformer-XL和Megatron-LM）的接口。FastMoE通过高性能加速技术优化了训练速度，并支持在多个GPU和节点上扩展专家数量。
3. 文献综述:
   • 该问题的相关工作包括BERT、GPT-2/-3、XLNet、RoBERTa、T5、GShard和Switch Transformer等大规模语言模型的发展
   • MoE架构由Shazeer等人提出，能够在不显著增加计算量的情况下大幅扩展模型规模。
   • MOE的双刃剑
     1. 在不显着增加计算量（FLOPs）的情况下将模型大小扩大几个数量级(好处)
     2. 当扩展到数千名专家时，MoE的全对全通信模式的不平衡给算法和系统的协同设计带来了新的挑战(问题
   • 传统的深度学习库如PyTorch （Paszke et al., 2019）和TensorFlow （Abadi et al., 2016）不能直接支持MoE
   • 现有的MoE训练系统主要依赖于Google的TPU和Mesh TensorFlow，缺乏灵活性和开放性
   • GShard系统（Chen et al., 2020）实现了MoE模型的分布式版本。它在多达2048个TPU上训练语言模型，每个TPU上每层有一个专家。因此，MoE层包含的参数比非MoE层多2048倍。
   • Switch Transformer （Fedus et al., 2021）中，将模型进一步扩大到1.6万亿，显示出系统支持大规模训练模型的强大能力。不幸的是，这个系统还没有公开可用。它与TPU集群紧密绑定，这使得很难在商品设备上复制实验

研究方法

这篇论文提出了FastMoE系统。具体来说：

• 灵活的系统设计:
  • FastMoE支持使用任意网络作为专家，并提供了一个层次化的接口，允许用户自定义门网络和专家网络。FastMoE还包含一个高度优化的Transformer前馈神经网络（FFN）实现。
  • FastMoE支持将多个专家放在同一个worker上，使专家数量的配置空间更加灵活（即专家数量不必等于数据并行数），这与GShard的设计不同
• 模型并行方法:
  • FastMoE支持在多个GPU和节点上分布专家，通过模型并行方法扩展模型容量。输入数据交换的细节被隐藏在FMoE接口中，用户只需编写单个专家的代码
  • FastMoE中集成了一个插件式模块，可以用MoE网络快速替换原来Megatron-LM模型中的ffn
• 异构同步模块: FastMoE引入了数据并行通信组标签，以解决不同部分的网络在不同工作器之间复制时的梯度同步问题。

实验设计

• 单GPU训练速度: 在NVIDIA TESLA V100 PCIe GPU上测试FastMoE的性能，与纯PyTorch实现的基线进行比较。测试了不同的输入样本数、隐藏层大小和选择的专家数。
• 跨GPU和跨节点的可扩展性: 在一个由8个节点组成的集群上进行实验，每个节点配备1个NVIDIA Tesla V100 GPU。通过增加GPU数量来测试FastMoE的可扩展性。
• 端到端性能提升: 使用Megatron-LM在8个GPU上训练一个12层的GPT模型，比较基线模型和MoE模型的训练速度和损失收敛情况。

结果与分析

• 单GPU训练速度: FastMoE在单GPU上的表现优于基线实现，特别是在专家数量增加时，基线的训练时间显著增加，而FastMoE的延迟保持稳定。
• 跨GPU和跨节点的可扩展性: FastMoE在多个GPU和节点上表现出良好的可扩展性，整体吞吐量从2个GPU的10 TFLOPs增加到8个GPU的25 TFLOPs，显示出次线性扩展。
• 端到端性能提升: MoE模型在相同的训练时间内实现了更低的损失，表明FastMoE在训练大模型方面的有效性。

结论

这篇论文介绍了FastMoE，一个开源的MoE模型训练系统，基于流行的PyTorch框架，支持在GPU上高效训练。FastMoE在单GPU上的性能经过优化，并且在多个GPU和节点上具有良好的可扩展性，能够进一步扩展模型规模。实验结果表明，MoE模型在使用FastMoE进行训练时表现出更好的性能。未来的工作包括支持专家之间的负载均衡和优化系统功能，以提高用户体验。
这篇论文为大规模MoE模型的训练提供了有效的解决方案，并展示了其在实际应用中的潜力。

参考

1. https://laekov.com.cn/l/slides/fastermoe.pdf