通过重用首层注意力输出，显著降低大模型训练通信开销。
First Attentions Last: Better Exploiting First Attentions for Efficient Transformer Training
- 将首层注意力输出直接用于后续层MLP输入，跳过通信密集的MHA-MLP连接
- 多卡训练时间减少44%，单卡吞吐提升1.18倍，且困惑度更低
- 适合大规模分布式训练场景，尤其对资源受限的高效训练有帮助
随着千亿参数级Transformer模型训练日益普遍，多GPU并行训练已成为标准做法。然而现有架构在张量并行（TP）中存在显著通信开销，每个模块的MHA-MLP连接需执行all-reduce操作。我们发现，可通过跳过这些连接实现效率提升，同时利用首层注意力输出作为替代信号。为此提出FAL架构：将首层MHA输出重定向至后续层MLP输入，消除每块的MHA-MLP连接，从而去除all-reduce通信，并支持单卡上MHA与MLP并行执行。进一步提出FAL+，在后续层MHA输出中加入归一化后的首层注意力输出，以增强输入信号。实验表明，FAL使多卡训练时间最多缩短44%，单卡吞吐最高提升1.18倍，且困惑度优于基线GPT；FAL+在不增加训练时间的前提下达到更低困惑度。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。