T3: Transparent Tracking & Triggering for Fine-grained Overlap of Compute & Collectives

📄 arXiv: 2401.16677v1 📥 PDF

作者: Suchita Pati, Shaizeen Aga, Mahzabeen Islam, Nuwan Jayasena, Matthew D. Sinclair

分类: cs.AR, cs.DC, cs.LG

发布日期: 2024-01-30

备注: To appear at the International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS) 2024

💡 一句话要点

提出T3以解决大语言模型通信效率低下问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 分布式训练 通信效率 硬件软件协同 Transformer模型 资源竞争 性能提升

📋 核心要点

  1. 现有分布式训练技术在设备数量增加时，通信效率显著降低，影响模型性能。
  2. T3通过硬件与软件协同设计，透明地重叠序列化通信与计算，减少资源竞争。
  3. 在T-NLG等重要模型中，T3实现了30%的通信加速和22%的数据移动减少，效果显著。

📝 摘要（中文）

随着大语言模型训练和推理对分布式技术的依赖加深，设备间的通信成为了效率提升的瓶颈。现有的分布式技术如张量并行性（TP）在模型执行时会序列化通信，导致资源竞争和效率降低。为了解决这一问题，本文提出了T3，通过硬件与软件的协同设计，透明地重叠序列化通信与计算，减少资源竞争。T3通过简单配置生产者的输出地址空间，将生产者操作与后续通信融合，且在硬件层面增加轻量级的跟踪与触发机制。实验结果表明，T3在重要的Transformer模型上实现了30%的通信加速和22%的数据移动减少，且在大规模模型中效果依然显著。

🔬 方法详解

问题定义：本文旨在解决大语言模型训练中因设备间通信导致的效率低下问题。现有的分布式技术如张量并行性（TP）在执行模型时会序列化通信，造成资源竞争和性能下降。

核心思路：T3的核心思路是通过硬件与软件的协同设计，透明地将序列化的通信与计算重叠，从而减少资源竞争。具体而言，T3通过配置生产者的输出地址空间，将生产者操作与后续通信进行融合。

技术框架：T3的整体架构包括两个主要模块：软件层面通过简单的配置实现生产者操作与通信的融合，硬件层面则增加了轻量级的跟踪与触发机制，以协调计算与通信的执行。此外，T3还利用计算增强内存来支持通信相关的计算。

关键创新：T3的主要创新在于其透明的融合机制和轻量级的硬件支持，这与现有方法的显著区别在于，它能够在不显著增加资源开销的情况下，提升通信与计算的重叠效率。

关键设计：T3的设计中，生产者的输出地址空间配置是关键参数，此外，轻量级的跟踪与触发机制确保了计算与通信的高效协调。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，T3在重要的Transformer模型（如T-NLG）上实现了30%的通信加速（最大可达47%），同时减少了22%的数据移动（最大可达36%）。在约5000亿参数的模型中，T3的效果依然显著，平均提升达到29%。

🎯 应用场景

T3的研究成果在大语言模型的训练和推理中具有广泛的应用潜力，尤其是在需要高效通信的场景下，如自然语言处理、机器翻译和对话系统等。通过提高通信效率，T3能够显著提升模型的训练速度和推理性能，推动更大规模模型的实际应用。

📄 摘要（原文）

Large Language Models increasingly rely on distributed techniques for their training and inference. These techniques require communication across devices which can reduce scaling efficiency as the number of devices increases. While some distributed techniques can overlap, and thus, hide this communication with independent computations, techniques such as Tensor Parallelism (TP) inherently serialize communication with model execution. One approach to hide this serialized communication is to interleave it with the producer operation (of the communicated data) in a fine-grained manner. However, this fine-grained interleaving of communication and computation in software can be difficult. Furthermore, as with any concurrent execution, it requires compute and memory resources to be shared between computation and communication, causing resource contention that reduces overlapping efficacy. To overcome these challenges, we propose T3 which applies hardware-software co-design to transparently overlap serialized communication while minimizing resource contention with compute. T3 transparently fuses producer operations with the subsequent communication via a simple configuration of the producer's output address space and requires minor software changes. At the hardware level, T3 adds a lightweight track and trigger mechanism to orchestrate the producer's compute, and communication. It further uses compute-enhanced memories for communication's attendant compute. As a result, T3 reduces resource contention, and efficiently overlaps serialized communication with computation. For important Transformer models like T-NLG, T3 speeds up communication-heavy sublayers by 30% geomean (max 47%) and reduces data movement by 22% geomean (max 36%). Furthermore, T3's benefits persist as models scale: geomean 29% for sublayers in $\sim$500-billion parameter models, PALM and MT-NLG.