HeterMoE：异构GPU上Mixture-of-Experts模型的高效训练
分布式、并行与集群计算
2025-04-08 v1 机器学习
摘要
Mixture-of-Experts(MoE)架构正变得越来越流行，用作放大大语言模型(LLM)的方法。为节省成本，已提出针对由新旧代GPU混合组成的GPU集群进行异构感知训练的解决方案。然而，现有解决方案对不同MoE模型组件(即注意力和专家)的性能特征不敏感，未充分利用每个GPU的计算能力。在本论文中，我们引入HeterMoE，一种在异构GPU上高效训练MoE模型的系统。我们的关键洞察是，新型GPU在注意力方面显著超越旧一代，而旧一代GPU在专家计算上仍相对高效。HeterMoE将注意力和专家计算分离，其中旧一代GPU仅被分配专家模块。通过提出的斑马并行，HeterMoE重叠不同GPU上的计算，同时采用非对称专家分配策略进行细粒度负载平衡，以最小化GPU空闲时间。我们的评估显示，HeterMoE相对于现有MoE训练系统实现了最高2.3倍加速，相对于最优异构感知解决方案实现1.4倍加速。即使将A40集群中一半GPU替换为V100，HeterMoE也能保持95%的训练吞吐量。
引用
@article{arxiv.2504.03871,
title = {HeterMoE: Efficient Training of Mixture-of-Experts Models on Heterogeneous GPUs},
author = {Yongji Wu and Xueshen Liu and Shuowei Jin and Ceyu Xu and Feng Qian and Z. Morley Mao and Matthew Lentz and Danyang Zhuo and Ion Stoica},
journal= {arXiv preprint arXiv:2504.03871},
year = {2025}
}