Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models

By Shibo Wang, Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake A. Hechtman, Dehao Chen, K. Murthy, Marcello Maggioni, Qiao Zhang and 4 more. 2022

117 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Hardware and Architecture

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Hardware and Architecture

Citations in the corpus

Scaling Vision Transformers to 22 Billion Parameters (2023, methodology): “The general case of this technique is presented in Wang et al. (2022a), who also introduce the XLA operations we leverage here.”
Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve (2024, methodology): “We filter out-liers of these datasets by removing requests with total length more than 8192 and 16384 tokens, respectively.”
Ring Attention with Blockwise Transformers for Near-Infinite Context (2023, background): “…of near-infinite context introduces a vast array of exciting opportunities, such as large video-audio-language models, learning from extended feedback and trial-and-errors, understanding and generating codebase, and adapting AI models to understand scientific data such as gene sequences.”