First Attentions Last: Better Exploiting First Attentions for Efficient Transformer Training

TL;DR

The paper targets the communication bottleneck of tensor-parallel training for large transformers by reconfiguring MHA-MLP connections. It introduces FAL, which reuses the first MHA output across blocks to feed the MLP, eliminating intra-block all-reduces and enabling concurrent MHA/MLP computation; FAL+ further augments this signal to improve model quality. Empirical results show substantial multi-GPU speedups (up to 44%) and single-GPU throughput gains (up to 1.18x), with perplexity improvements over GPT baselines; FAL+ provides additional perplexity reductions without extra training time and generalizes across variants and modalities. The approach offers a practical, architecture-level path to more scalable transformer training and inference, with broad applicability beyond language models.

Abstract

As training billion-scale transformers becomes increasingly common, employing multiple distributed GPUs along with parallel training methods has become a standard practice. However, existing transformer designs suffer from significant communication overhead, especially in Tensor Parallelism (TP), where each block's MHA-MLP connection requires an all-reduce communication. Through our investigation, we show that the MHA-MLP connections can be bypassed for efficiency, while the attention output of the first layer can serve as an alternative signal for the bypassed connection. Motivated by the observations, we propose FAL (First Attentions Last), an efficient transformer architecture that redirects the first MHA output to the MLP inputs of the following layers, eliminating the per-block MHA-MLP connections. This removes the all-reduce communication and enables parallel execution of MHA and MLP on a single GPU. We also introduce FAL+, which adds the normalized first attention output to the MHA outputs of the following layers to augment the MLP input for the model quality. Our evaluation shows that FAL reduces multi-GPU training time by up to 44%, improves single-GPU throughput by up to 1.18x, and achieves better perplexity compared to the baseline GPT. FAL+ achieves even lower perplexity without increasing the training time than the baseline. Codes are available at: https://github.com/CASL-KU/FAL

Figures (20)

• Figure 1: Transformer block designs highlighting our proposed modifications. (a) Pre-LN architecture with Layer Normalization (LN), Multi Head Attention (MHA) and Multi Layer Perceptron (MLP). (b) FAL blocks with reconfigured connections. (c) FAL+ blocks with augmented connections. (d) End-to-end training time and perplexity comparison of Pre-LN architecture, FAL, and FAL+.
• Figure 2: Forward pass during Transformer training with tensor parallelism across 2 GPUs. (a) Standard training involves Broadcast, All-Reduce, and Aggregate communication. (b) FAL training requires only Broadcast and Aggregate, reducing communication overhead.
• Figure 3: Motivation for reconfiguring MHA-MLP connections. (a) CKA similarity analysis across successive GPT-2 blocks. The x-axis indicates the block index, and the y-axis shows the similarity between consecutive MHA outputs, Residual+MHA outputs (i.e., MLP inputs), and MLP outputs. (b) Connection ablation results measured by perplexity. Original denotes the unaltered model, All MHA removes all MHA layer, and All Connect removes all direct MHA-MLP connections.
• Figure 4: The crucial role of the first attentions. (a) Normalized gradient magnitude of the MHA outputs across Transformer blocks in GPT-2 for different datasets. The x-axis represents the block index. (b) Layer ablation results measured by perplexity. The x-axis indicates the index of the transformer block from which the MHA is omitted.
• Figure 5: Parallel kernel launch of FAL.