面向 LLM 训练的层次化数据流水线：AMD GPU 上的数据调度研究
Scheduling Data to Compute: Hierarchical Data Pipelining for LLM Training on AMD GPUs
研究团队提出在 AMD GPU 上通过层次化数据流水线解决 LLM 训练中计算单元等待数据的瓶颈，利用 Tensor Data Mover（TDM）、Transpose-load、SDMA 和 RDMA 三类数据搬运特性保持数据持续供给计算单元。
Authors: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
Abstract:Large language models (LLM) have been the key focus of recent machine learning breakthroughs. As the model size grows exponentially, the computation needed for training is astronomical. A GPU can do a huge amount of arithmetic. However, the data its compute units need is not always available when they are ready. This data may still be in slow memory or on another GPU. While the chip waits, expensive compute sits idle, and at scale even a small drop in utilization translates into substantial cost. In this paper, we treat this waiting time as the primary bottleneck. We focus on hardware features designed for explicit data movement. These features load data from off-chip memory onto the chip, rearrange it into the layout the computation requires as it is loaded, and transfer it between GPUs. By orchestrating them, we keep a steady pipeline of data flowing to the compute units so that they remain busy.
We present kernel and framework implementations that address this bottleneck using three key data-movement instructions and hardware features. The Tensor Data Mover (TDM) moves tensor data from high-bandwidth memory (HBM) into on-chip memory. Transpose-load instructions transpose tensor data into the layout that LLM computations require as it is loaded on-chip. Finally, System DMA (SDMA) and Remote DMA (RDMA) move data between GPUs over the scale-up and scale-out networks, respectively. Our results show that TDM and transpose-load provide on average 1.88x and 1.19x speedup at the kernel level, respectively. The SDMA and RDMA lead to on average 1.16x and 1.33x speedup for computation overlapped with communication. This is the first AMD LLM training system to utilize all the GPU data pipelining features. Our implementation has been integrated into open source framework and library for AMD GPUs. this https URL this https URL
| Subjects: | Hardware Architecture (cs.AR); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2611.05751 [cs.AR] |
| (or arXiv:2611.05751v1 [cs.AR] for this version) | |
| https://doi.org/10.48550/arXiv.2611.05751 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Gengyu Rao [view email]
[v1]
Tue, 3 Nov 2026 23:44:04 UTC (90 KB)
来源：arXiv cs.AI · arxiv.org