研究摘要：
专家混合模型（Mixture-of-Experts, MoE）成为扩展预训练模型规模的最热门技术之一。通过动态激活专家，实现条件计算，大幅增加神经网络的参数数量，这对于吸收深度学习领域中广泛的知识至关重要。然而，即使在现有的系统和算法优化下，通信效率低下和内存消耗问题仍然是需要解决的重大挑战。本文介绍了一种名为 MPipeMoE 的研究工作，通过自适应且内存高效的流水线并行技术，加速了 MoE 的训练。并通过拓展自适应流水线并行方法，实现内存复用策略，通过消除内存冗余来减少内存需求，并开发了自适应选择组件，在运行时根据硬件能力和模型特性动态选择最优策略。在常见的 MoE 模型训练中，与现有最先进的方法相比，MPipeMoE 的训练速度最高提升了 2.8 倍，同时将内存占用减少了 47%，展现了在训练大规模模型中的卓越性能。
研究背景：
在深度学习领域，特别是自然语言处理 (NLP)，通过扩展模型规模来提升模型性能已成为趋势。近年来，Mixture-of-Experts (MoE) 架构成为了一种流行的技术，如图1所示，MoE可以通过稀疏激活来显著提高参数数量而不成比例地增加计算成本。然而，MoE 模型的训练面临着两个主要挑战：高昂的通信成本和显著的内存消耗，特别是在模型规模达到数百亿甚至数万亿参数时，这些问题尤为突出。这篇研究工作以此为出发点，探索如何以高效的方式解决通信和内存瓶颈，进一步推动 MoE 的大规模应用。
图 1 MoE 模型架构示意图
研究现状
目前的 MoE 系统（如 Switch Transformer 和 BASE Layer）在优化通信和计算效率方面已经取得了一些进展。例如，FasterMoE 通过引入专家影子机制和流水线并行训练减少了通信开销。然而，这些方法的粒度通常是固定的，缺乏动态适应能力，难以满足训练过程中的变化需求。此外，内存管理问题长期以来被忽视，特别是激活和临时缓冲区的内存消耗仍然是优化训练超大规模模型的主要障碍。
研究方法
为了应对现有方法在通信和内存管理方面的局限，论文提出了一种名为 MPipeMoE 的创新方案，以自适应流水线并行和高效内存重用策略为核心，解决了稀疏专家模型（MoE）训练中的关键挑战。MPipeMoE 的自适应流水线并行策略通过将输入批次细分为多个微批次，同时优化计算与通信的重叠，最大程度地提升了 GPU 和网络带宽的利用率。与传统固定粒度的流水线方案不同，该方法引入了在线算法，可以根据硬件性能和任务需求动态调整流水线的粒度，从而实现更灵活、更高效的并行配置。
在内存管理方面，论文对 MoE 模型训练中的内存消耗进行了深入分析，识别出模型状态、激活和临时缓冲区是主要的内存开销来源。针对这一问题，MPipeMoE 提出了内存共享机制，通过在不同阶段复用内存缓冲区，大幅降低激活和临时缓冲区的内存需求，最多可减少到原来的 1/n。此外，为进一步优化内存使用，该方法结合了重计算技术和 CPU 数据卸载策略，既能有效节省 GPU 内存，又在一定程度上降低了计算和数据传输的开销。
为了增强系统的适应性，MPipeMoE 还构建了一个运行时性能模型，根据硬件配置和具体任务的需求动态选择最佳的内存管理策略。该模型在运行时对不同的内存分配和计算方案进行性能预测，从而实现资源利用的优化。这种全方位的设计不仅提高了 MoE 模型训练的效率，还在保持模型性能的前提下显著降低了内存消耗和计算开销，为大规模稀疏专家模型的训练提供了一种切实可行且高效的解决方案。
图 2 流水并行示意图
研究成果
实验表明，MPipeMoE 在多种大规模稀疏模型上均展现出了显著的优势。如图3所示，在 8 台 NVIDIA DGX A100 服务器（共 64 块 A100 GPU）的测试环境中，MPipeMoE 在训练 MoE-GPT3 和 MoE-BERT 时，相较于现有的框架（如 FastMoE 和 FasterMoE），实现了最高 3.7 倍的训练速度提升，同时将内存占用降低了 47%。自适应流水线并行的引入使得 MPipeMoE 在各种批量大小下都表现稳定，无论是小批量高通信开销还是大批量高内存占用的场景，都能显著提升训练效率。
图 3 模型对比实验。左轴表示相对显存占用量，右轴表示加速比。
相关研究成果获得了高性能计算领域的广泛认可，并发表在以下顶级会议和期刊中：
1. 论文《MPipeMoE: Memory Efficient MoE for Pre-trained Models with Adaptive Pipeline Parallelism》 发表在高性能计算顶会 IPDPS 2023（39th IEEE International Parallel &
Distributed Processing Symposium）（CCF-B）。
2. 论文《MPMoE: Memory Efficient MoE for Pre-trained Models with Adaptive Pipeline Parallelism》 发表在 CCF-A 类顶级期刊 IEEE TPDS（Transactions on Parallel and Distributed Systems）上。