FlashMoE用单个GPU内核加速MoE模型,显著提升训练效率。
FlashMoE: Fast Distributed MoE in a Single Kernel
- 将专家计算与跨GPU通信融合到一个持久化内核中,实现细粒度流水线
- 在8张H100上实现9倍更高显卡利用率、6倍更低延迟、5.7倍更高吞吐
- 适合大规模分布式训练场景,尤其对资源受限的超大模型训练有帮助
Mixture-of-Experts(MoE)模型的计算稀疏性使得模型规模增大时计算成本呈亚线性增长,为训练超大规模神经网络提供了可扩展路径。然而,现有实现存在显卡利用率低、延迟高、无法利用任务局部性等问题,主要由于依赖CPU调度、主机发起通信和频繁启动内核。为此,我们开发了FlashMoE,一个完全驻留于GPU的MoE算子,将专家计算与跨GPU通信融合进单一持久化GPU内核。FlashMoE支持调度、计算、合并阶段的细粒度流水线,消除启动开销并减少空闲间隔。不同于以往工作,FlashMoE采用单边、设备发起的跨GPU(R)DMA传输,避免了批量同步集体操作,从而通过消除稀疏激活层中的冗余网络负载提升了数据传输效率。在配备8张H100 GPU的节点上,对包含最多128个专家和16K令牌序列的MoE模型进行评估,FlashMoE相比最先进基线实现高达9倍的显卡利用率、6倍的降低延迟、5.7倍的提升吞吐量以及4倍更好的重叠效率,尽管其使用FP32,而基线采用FP16。结果表明,精心设计的GPU内核与硬件协同是突破大规模分布式机器学习性能瓶颈的关键。代码已开源:https://github.com/osayamenja/FlashMoE。
原文摘要 · Abstract (English)
The computational sparsity of Mixture-of-Experts (MoE) models enables sub-linear growth in compute cost as model size increases, thus offering a scalable path to training massive neural networks. However, existing implementations suffer from low GPU utilization, significant latency overhead, and a fundamental inability to leverage task locality, primarily due to CPU-managed scheduling, host-initiated communication, and frequent kernel launches. To overcome these limitations, we develop FlashMoE, a fully GPU-resident MoE operator that fuses expert computation and inter-GPU communication into a single persistent GPU kernel. FlashMoE enables fine-grained pipelining of dispatch, compute, and combine phases, eliminating launch overheads and reducing idle gaps. Unlike existing work, FlashMoE eliminates bulk-synchronous collectives for one-sided, device-initiated, inter-GPU (R)DMA transfers, thereby unlocking payload efficiency by eliminating bloated or redundant network payloads in sparsely activated layers. When evaluated on an 8-H100 GPU node with MoE models comprising up to 128 experts and 16K token sequences, FlashMoE achieves up to 9x higher GPU utilization, 6x lower latency, 5.7x higher throughput, and 4x better overlap efficiency compared to state-of-the-art baselines, despite using FP32, whereas the baselines use FP16. FlashMoE shows that principled GPU kernel-hardware co-design is key to unlocking the performance ceiling of large-scale distributed ML. We provide code at https://github.com/osayamenja/FlashMoE.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。