Demystifying the Communication Characteristics for Distributed Transformer Models

TL;DR

This work investigates the communication bottlenecks in distributed transformer training for large language models, tackling how data, pipeline, and tensor parallelism—and ZeRO variants—shape inter-process data movement. It blends empirical profiling on the Frontier supercomputer with analytical models to quantify collective volumes, message sizes, and frequencies across sequence lengths, providing system-agnostic and interconnect-specific insights. Key findings include formalized volume formulas such as Allreduce_volume = $2 m$ and Allgather_volume = $m$, the importance of small-message optimization, and the observation that ZeRO-3 can incur roughly 1.5× more communication than DP/ZeRO-1/2 due to extra Allgather. The results offer practical guidance for MPI/communication library design and HPC middleware to better support transformer-scale workloads on current and future exascale systems.

Abstract

Deep learning (DL) models based on the transformer architecture have revolutionized many DL applications such as large language models (LLMs), vision transformers, audio generation, and time series prediction. Much of this progress has been fueled by distributed training, yet distributed communication remains a substantial bottleneck to training progress. This paper examines the communication behavior of transformer models - that is, how different parallelism schemes used in multi-node/multi-GPU DL Training communicate data in the context of transformers. We use GPT-based language models as a case study of the transformer architecture due to their ubiquity. We validate the empirical results obtained from our communication logs using analytical models. At a high level, our analysis reveals a need to optimize small message point-to-point communication further, correlations between sequence length, per-GPU throughput, model size, and optimizations used, and where to potentially guide further optimizations in framework and HPC middleware design and optimization.

Figures (15)

• Figure 1: 13-billion parameter model breakdown of communication and computation using ZeRO-1 and 8 tensor-parallel stages (single iteration)
• Figure 2: 20-billion parameter model breakdown of communication and computation using ZeRO-1 and 8 tensor-parallel stages (single iteration)
• Figure 3: A non-exhaustive list of what must be considered when characterizing LLM performance, scalability, and communication behavior.
• Figure 4: An illustration of 3D parallelism with 2 Data-Parallel ranks, 2 Pipeline-Parallel stages and 2 Tensor-Parallel ranks. Each Pipeline-Parallel stage holds half of the total layers.
• Figure 5: An illustration of ZeRO-3 with 4 Data-Parallel ranks and $N$ layers. Between each layer, an Allgather is needed to collect the parameters from all the workers.