Klotski通过智能调度专家计算与加载,大幅减少MoE模型推理中的等待空隙。
Klotski: Efficient Mixture-of-Expert Inference via Expert-Aware Multi-Batch Pipeline
- 设计专家感知的多批流水线,动态平衡计算与I/O时间。
- 在不同硬件和模型下实现最高85.12倍的吞吐量提升。
- 适合需要高效部署超大规模MoE模型的研究者与工程师。
Mixture of Experts(MoE)凭借其稀疏结构,使语言模型规模扩展至万亿参数而计算成本增长有限。然而,巨大的参数量给推理带来挑战,因GPU内存扩张速度跟不上参数增长。尽管离线技术利用CPU和磁盘内存,并行化I/O与计算以提高效率,但MoE中每个专家的计算时间常短于I/O时间,导致流水线中出现大量空闲间隙。为此,我们提出Klotski,一种高效的MoE推理引擎,通过新颖的专家感知多批流水线范式显著减少流水线空隙。该范式使用批量处理延长当前层计算时间,以重叠下一阶段加载时间。尽管此思想已在密集模型中有效应用,但更多批次可能激活更多专家,导致更长加载时间和更多空隙。因此,不同于传统方法,我们基于不同批大小下的异构计算与I/O需求及激活模式,优化推理顺序以平衡计算与I/O时间并最小化空隙。此外,为适应不同硬件环境与模型,我们设计了约束敏感的I/O-计算调度器与相关性感知的专家预取器,实现最小化流水线空隙的调度。实验表明,相较于现有最优技术,Klotski在吞吐率-延迟权衡上表现更优,吞吐量最高提升85.12倍。
原文摘要 · Abstract (English)
Mixture of Experts (MoE), with its distinctive sparse structure, enables the scaling of language models up to trillions of parameters without significantly increasing computational costs. However, the substantial parameter size presents a challenge for inference, as the expansion in GPU memory cannot keep pace with the growth in parameters. Although offloading techniques utilise memory from the CPU and disk and parallelise the I/O and computation for efficiency, the computation for each expert in MoE models is often less than the I/O, resulting in numerous bubbles in the pipeline. Therefore, we propose Klotski, an efficient MoE inference engine that significantly reduces pipeline bubbles through a novel expert-aware multi-batch pipeline paradigm. The proposed paradigm uses batch processing to extend the computation time of the current layer to overlap with the loading time of the next layer. Although this idea has been effectively applied to dense models, more batches may activate more experts in the MoE, leading to longer loading times and more bubbles. Thus, unlike traditional approaches, we balance computation and I/O time and minimise bubbles by orchestrating their inference orders based on their heterogeneous computation and I/O requirements and activation patterns under different batch numbers. Moreover, to adapt to different hardware environments and models, we design a constraint-sensitive I/O-compute planner and a correlation-aware expert prefetcher for a schedule that minimises pipeline bubbles. Experimental results demonstrate that Klotski achieves a superior throughput-latency trade-off compared to state-of-the-art techniques, with throughput improvements of up to 85.12x.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。