[论文解读] TACCL: Guiding Collective Algorithm Synthesis using Communication Sketches

Aashaka Shah, Vijay Chidambaram+7 more|arXiv (Cornell University)|Nov 8, 2021

Ferroelectric and Negative Capacitance Devices被引用 6

一句话总结

TACCL 是一种新颖的工具，通过利用人工提供的通信草图来引导自动合成高效 GPU 集体通信算法，显著缩小了搜索空间，并实现了在多节点 GPU 拓扑结构上的可扩展、高性能算法生成。其性能相比 NCCL 最高提升 6.7 倍，并将 BERT 和 Transformer-XL 训练速度提升 11%–2.3 倍。

ABSTRACT

Machine learning models are increasingly being trained across multiple GPUs and servers. In this setting, data is transferred between GPUs using communication collectives such as AlltoAll and AllReduce, which can become a significant bottleneck in training large models. Thus, it is important to use efficient algorithms for collective communication. We develop TACCL, a tool that enables algorithm designers to guide a synthesizer into automatically generating algorithms for a given hardware configuration and communication collective. TACCL uses a novel communication sketch abstraction to get crucial information from the designer to significantly reduce the search space and guide the synthesizer towards better algorithms. TACCL also uses a novel encoding of the problem that allows it to scale beyond single-node topologies. We use TACCL to synthesize algorithms for three collectives and two hardware topologies: DGX-2 and NDv2. We demonstrate that the algorithms synthesized by TACCL outperform the Nvidia Collective Communication Library (NCCL) by up to 6.7x. We also show that TACCL can speed up end-to-end training of Transformer-XL and BERT models by 11%--2.3x for different batch sizes.

研究动机与目标

• 解决在具有混合带宽和延迟特性的大规模异构 GPU 拓扑结构中设计高效集体通信算法的挑战。
• 克服在多节点系统中探索集体算法路由与调度选项完整搜索空间的计算不可行性。
• 通过一种新型抽象——通信草图，融入算法设计者的高层次直觉，从而减少合成时间并提升可扩展性。
• 实现针对特定硬件配置（如 DGX-2 和 NDv2）量身定制的 Allreduce、Allgather 和 Alltoall 算法的自动生成。
• 通过结合草图引导搜索与松弛的混合整数线性规划（MILP）公式化方法及启发式优化技术，实现优于现有库（如 NCCL）的性能表现。

提出的方法

• 引入通信草图——一种受程序草图启发的高层抽象，使算法设计者能够表达对集体通信中数据路由与调度结构的直觉。
• 将集体通信合成问题建模为一种松弛路由与连续时间调度的混合整数线性规划（MILP），以提升可扩展性。
• 应用三阶段合成流水线：松弛路由以探索可行路径，启发式排序以优先选择高带宽链路，以及带连续性约束的精确调度以最小化空闲时间。
• 利用对称性与拓扑感知编码减少状态空间，从而实现在复杂拓扑（如 2D-环形拓扑和分层集群，例如 DGX-2 和 Azure NDv2）上的合成。
• 使用经过测量的网络拓扑和输入大小信息作为输入，引导合成器生成面向硬件、兼顾延迟与带宽优化的算法。
• 通过后端将合成的算法映射到底层 GPU 通信原原子，确保与现有分布式深度学习框架的兼容性。

实验结果

研究问题

• RQ1 通信草图能否在大规模 GPU 拓扑结构中有效引导集体通信算法的自动合成，同时显著缩小搜索空间？
• RQ2 如何将高效集体算法的合成扩展至单节点系统之外，并适用于具有不同链路速度的异构网络拓扑？
• RQ3 采用直观草图的人机协同方法，在生成高性能通信算法方面，与完全自动化或基于启发式的方案相比，其性能优势有多大？
• RQ4 所提出的带松弛和启发式优化的 MILP 公式化方法，是否能在保持大规模 GPU 集群可处理性的同时，实现接近最优的性能？
• RQ5 在不同工作负载和硬件配置下，所合成的算法与最先进的库（如 NCCL）相比，在实际应用中表现如何？

主要发现

• TACCL 在不到 5 分钟内为 8 个 Azure NDv2 节点（64 个 GPU）合成出 Allgather 算法，在各种数据大小下实现的算法带宽最高达 NCCL 的 1.7 倍。
• 对于 8 个 DGX-2 节点（128 个 GPU），TACCL 在约 11 小时内生成了有效的 Allgather 算法，证明了其在大规模多节点系统中的可扩展性。
• 所合成的算法在测试的集体通信原语和拓扑结构上，端到端通信性能相比 NCCL 最高提升 6.7 倍。
• 由于通信瓶颈减少，BERT 和 Transformer-XL 模型的端到端训练速度提升了 11% 至 2.3%，具体取决于小批量大小。
• TACCL 通过利用通信草图中的对称性，成功为非分层拓扑（如 2D $6\times8$ 环形拓扑）生成了高效算法。
• 使用通信草图使得能够探索多种针对不同输入大小范围优化的算法变体，充分体现了开发者的设 计直觉。