一句话：提出将通信集合与依赖计算操作分解为细粒度操作并重叠执行的技术，隐藏数据传输延迟，提升大规模模型系统吞吐。
大规模深度学习模型在加速器上运行时，设备内存常不足以容纳整个模型。层内模型并行通过将单层或算子切分到多设备来缓解内存压力，但产生的数据通信会占据显著执行时间，损害计算效率。本文提出一种通过分解重叠通信与依赖计算的技术：将原始通信集合与依赖计算操作一起分解为一系列更细粒度操作，创造更多重叠机会，并让细粒度通信与计算并行执行，从而隐藏数据传输延迟、提升系统利用率。在 TPU v4 Pods 上，使用 100 亿至 1 万亿参数的不同类型大模型评估，该技术将系统吞吐率提升 1.14-1.38 倍；在 1024 个 TPU 芯片上运行 5000 亿参数大语言模型时，最高峰值 FLOPS 利用率达 72%。
Large deep learning models have shown great potential with state-of-the-art results in many tasks. However, running these large models is quite challenging on an accelerator (GPU or TPU) because the on-device memory is too limited for the size of these models. Intra-layer model parallelism is an approach to address the issues by partitioning individual layers or operators across multiple devices in a distributed accelerator cluster. But, the data communications generated by intra-layer model parallelism can contribute to a significant proportion of the overall execution time and severely hurt the computational efficiency. As intra-layer model parallelism is critical to enable large deep learning models, this paper proposes a novel technique to effectively reduce its data communication overheads by overlapping communication with computation. With the proposed technique, an identified original communication collective is decomposed along with the dependent computation operation into a sequence of finer-grained operations. By creating more overlapping opportunities and executing the newly created, finer-grained communication and computation operations in parallel, it effectively hides the data transfer latency and achieves a better system utilization. Evaluated on TPU v4 Pods using different types of large models that have 10 billion to 1 trillion parameters, the proposed technique improves system throughput by 1.14 - 1.38x. The achieved highest peak FLOPS utilization is 72% on 1024 TPU chips with a large language model that has 500 billion parameters.
@inproceedings{WangWSDIHCMMZKG23,
title = {Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models},
author = {Shibo Wang and Jinliang Wei and Amit Sabne and Andy Davis and Berkin Ilbeyi and Blake Hechtman and Dehao Chen and Karthik Srinivasa Murthy and Marcello Maggioni and Qiao Zhang and Sameer Kumar and Tongfei Guo and Yuanzhong Xu and Zongwei Zhou},
booktitle = {ASPLOS},
year = {2022},
doi = {10.1145/3567955.3567959}
}