Lancet：通过全图计算通信重叠加速 Mixture-of-Expert 训练
分布式、并行与集群计算
2024-05-01 v1 机器学习
摘要
Mixture-of-Expert (MoE) 技术在扩大深度神经网络模型参数规模方面发挥着关键作用。然而，它面临训练过程中 all-to-all 通信延迟显著的挑战。现有方法试图通过重叠 all-to-all 与专家计算来缓解此问题，但常未能实现充分的重叠，从而限制了性能提升的潜力。在本研究中，我们从更广泛的训练图层面视角挑战，考虑在更宏观的层面上实现计算通信重叠。在前向传播期间，通过谨慎的划分与流水线，使非 MoE 计算得以与 all-to-all 重叠。在反向传播期间，通过调度梯度权重计算，实现与 all-to-all 的重叠。我们在 Lancet 系统中实现了这些技术，该系统采用基于编译器的优化自动提升 MoE 模型训练。我们的广泛评估表明，Lancet 将非重叠通信所需时间显著减少，最高可降低 77%。此外，它相较于现有最先进解决方案实现了最高达 1.3 倍的端到端加速。
引用
@article{arxiv.2404.19429,
title = {Lancet: Accelerating Mixture-of-Experts Training via Whole Graph Computation-Communication Overlapping},
author = {Chenyu Jiang and Ye Tian and Zhen Jia and Shuai Zheng and Chuan Wu and Yida Wang},
journal= {arXiv preprint arXiv:2404.19429},
year = {2024}
}
备注
11 pages, 16 figures. Published in MLSys'24