面向 AMD GPU 的 LLM 训练分层数据流水线调度

作者：Gengyu Rao、Ghadeer Ahmed H Alabandi、Su Ann Chong、Thomas Huber、Kailash Gogineni、Subhajit Dutta Chowdhury、Saleel Kudchadker、Suchita Pati、Zirui Zhang、Zhen Huang、Yanyuan Qin、Xiaoming Peng、Lihuan Zhang、Botao Hu、Wei Huang、Yang Bo、Wen Xie、Xiangyu Zhao、Ruibin Zhang、Xiaobo Chen、Vidushi Goyal、Anshu Raina、Claire Songhyun Lee、Jiakun Yan、Yao Fu、Nusrat Islam、Ravi Dwivedula、Jie Zhang、Maria Garzaran、Zhenyu Gu、Emad Barsoum

查看 PDF HTML（实验性）

摘要：大语言模型（LLM）一直是近期机器学习突破的核心焦点。随着模型规模呈指数级增长，训练所需的计算量已达到天文数字。一块 GPU 能完成海量算术运算。然而，其计算单元所需的数据并不总是在就绪时可用。这些数据可能仍在慢速内存中，或位于另一块 GPU 上。当芯片等待时，昂贵的计算资源处于空闲状态，而在大规模场景下，即便利用率小幅下降也会转化为可观的成本。在本文中，我们将这种等待时间视为主要瓶颈。我们聚焦于为显式数据移动而设计的硬件特性。这些特性将数据从片外内存加载到芯片上，在加载过程中将其重排为计算所需的布局，并在 GPU 之间传输数据。通过编排这些特性，我们保持数据稳定地流向计算单元，使其持续忙碌。
我们提出了利用三条关键数据移动指令和硬件特性来解决该瓶颈的内核与框架实现。张量数据搬运器（TDM）将张量数据从高带宽内存（HBM）移入片上内存。转置加载指令在张量数据加载到片上时，将其转置为 LLM 计算所需的布局。最后，系统 DMA（SDMA）和远程 DMA（RDMA）分别通过纵向扩展网络和横向扩展网络在 GPU 之间移动数据。我们的结果表明，TDM 和转置加载在内核层面分别平均带来 1.88 倍和 1.19 倍的加速。SDMA 和 RDMA 在计算与通信重叠时分别平均带来 1.16 倍和 1.33 倍的加速。这是首个利用全部 GPU 数据流水线特性的 AMD LLM 训练系统。我们的实现已集成到面向 AMD GPU 的开源框架和库中。此 https URL 此 https URL

提交历史

来自：Gengyu Rao [查看邮箱]
[v1] 2026年11月3日星期二 23:44:04 UTC（90 KB）