通过编译器优化,让多核AI芯片高效运行大模型。
ELK: Exploring the Efficiency of Inter-core Connected AI Chips with Deep Learning Compiler Techniques
- 设计编译框架Elk,统筹计算、通信与存储的性能权衡。
- 在真实芯片模拟中实现平均94%的理想性能,提升大模型运行效率。
- 适合芯片架构设计与深度学习编译器研究者参考。
为满足深度学习模型日益增长的需求,AI芯片同时采用片外内存(如HBM)和高带宽低延迟的核间互连,实现核间直接数据交换。然而,由于计算(单核执行)、通信(核间数据交换)和输入输出(片外数据访问)之间存在根本性矛盾,难以有效探索这类核间互联AI(ICCA)芯片的性能效率。本文提出Elk,一种深度学习编译器框架,通过协同权衡三者性能因素,最大化ICCA芯片效率。Elk将性能因素建模为可配置参数,构建编译器内全局权衡空间,并引入新型归纳操作调度策略与成本感知片上内存分配算法,生成全局最优执行计划,充分重叠片外数据加载与片上执行。为验证其效果,我们基于真实ICCA芯片IPU-POD4构建全功能模拟器,并开发了用于敏感性分析的芯片仿真器。实验表明,Elk在平均情况下达到ICCA芯片理想屋顶线性能的94%,证明其在支持大模型方面的优势。此外,还展示了Elk在新ICCA芯片架构设计空间探索中的能力。
原文摘要 · Abstract (English)
To meet the increasing demand of deep learning (DL) models, AI chips are employing both off-chip memory (e.g., HBM) and high-bandwidth low-latency interconnect for direct inter-core data exchange. However, it is not easy to explore the efficiency of these inter-core connected AI (ICCA) chips, due to a fundamental tussle among compute (per-core execution), communication (inter-core data exchange), and I/O (off-chip data access). In this paper, we develop Elk, a DL compiler framework to maximize the efficiency of ICCA chips by jointly trading off all the three performance factors discussed above. Elk structures these performance factors into configurable parameters and forms a global trade-off space in the DL compiler. To systematically explore this space and maximize overall efficiency, Elk employs a new inductive operator scheduling policy and a cost-aware on-chip memory allocation algorithm. It generates globally optimized execution plans that best overlap off-chip data loading and on-chip execution. To examine the efficiency of Elk, we build a full-fledged emulator based on a real ICCA chip IPU-POD4, and an ICCA chip simulator for sensitivity analysis with different interconnect network topologies. Elk achieves 94% of the ideal roofline performance of ICCA chips on average, showing the benefits of supporting large DL models on ICCA chips. We also show Elk's capability of enabling architecture design space exploration for new ICCA chip development.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。