arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2407.05784cs.AR
收藏

Hecaton:利用可扩展小芯片系统训练大语言模型

Hecaton: Training Large Language Models with Scalable Chiplet Systems

Zongle Huang, Shupei Fan, Chen Tang, Xinyuan Lin, Shuwen Deng, Yongpan Liu

AI总结:

针对LLM训练并行化带来的通信开销问题,本工作提出Hecaton小芯片系统,通过定制架构调度和分布式训练方法优化DRAM访问与NoP通信,实现弱扩展性,在Llama3.1-405B上性能提升5.29倍、能耗降低3.46倍。

摘要:

大语言模型(LLMs)在各个领域取得了显著成功,但其训练和微调需要大规模计算与内存资源,因此必须采用并行化,而这会引入巨大的通信开销。在封装技术进步的推动下,小芯片架构成为一种潜在解决方案,因为它可以整合计算能力,并利用具有更好信号完整性、更高带宽和更低能耗的封装内链路。然而,大多数现有小芯片相关工作聚焦于深度神经网络(DNN)推理。将它们直接移植到LLM训练会引入大量动态随机存取存储器(DRAM)访问和封装内网络(NoP)开销,导致最先进的小芯片设计失效,凸显了这一研究空白。本工作提出Hecaton,一种用于LLM训练的可扩展且具成本效益的小芯片系统。我们首先提供一种带有定制调度的小芯片架构,可大幅减少DRAM访问。我们进一步设计了一种高效的分布式训练方法,降低NoP通信复杂度,并缓解静态随机存取存储器(SRAM)容量和布局的约束。理论分析表明,整个系统实现了弱扩展性:当工作负载和硬件资源成比例增长时,计算与通信的比率几乎保持不变。针对各种工作负载和硬件配置的实验验证了这一特性,与Megatron中的张量并行相比,Hecaton在Llama3.1-405B上实现了5.29倍的性能提升和3.46倍的能耗降低。据我们所知,我们提出了首个专门用于LLM训练或微调的小芯片架构,无论问题规模如何,都能保证性能。

英文摘要:

Large Language Models (LLMs) have achieved remarkable success in various fields, but their training and finetuning require massive computation and memory, necessitating parallelism which introduces heavy communication overheads. Driven by advances in packaging, the chiplet architecture emerges as a potential solution, as it can integrate computing power, as well as utilize on-package links with better signal integrity, higher bandwidth, and lower energy consumption. However, most existing chiplet-related works focus on DNN inference. Directly porting them to LLM training introduces significantly large quantities of DRAM access and network-on-package (NoP) overheads which make state-of-the-art chiplet designs fail, highlighting a research gap. This work proposes Hecaton, a scalable and cost-effective chiplet system for LLM training. We first provide a chiplet architecture with tailored scheduling that can largely reduce DRAM accesses. We further design an efficient distributed training method that reduces NoP communication complexity and relieves constraints on SRAM capacity and layout. Theoretical analysis shows that the entire system achieves weak scaling: as the workload and hardware resources grow proportionally, the computation-to-communication ratio remains nearly constant. Experiments with various workloads and hardware configurations verify the property, and Hecaton achieves $5.29\times$ performance improvement and $3.46\times$ energy reduction on Llama3.1-405B, compared to the tensor parallelism in Megatron. To the best of our knowledge, we propose the first chiplet architecture specifically used for LLM training or finetuning, with guaranteed performance regardless of the problem scale.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑