主要观点总结
这篇文章介绍了PyTorch中新的异步张量并行（Async-TP）技术，该技术显著提升了大规模语言模型训练的性能。文章详细讨论了实现过程中的技术挑战和解决方案，并提供了关于如何在TorchTitan和PyTorch中使用Async-TP的指南。
关键观点总结
关键观点1: 异步张量并行（Async-TP）的概念和优势
Async-TP通过将通信和计算操作分解并重叠执行，提高了大规模语言模型训练的效率。它在Llama3系列模型上的测试表明，可以显著提升前向传播和端到端训练速度。
关键观点2: 技术挑战与解决方案
文章讨论了实现Async-TP过程中的技术挑战，包括通信开销和计算效率问题，并提出了相应的解决方案。例如，通过利用CUDA的P2P机制和SymmetricMemory抽象来优化通信，通过交替流方法解决放大的量化效率问题。
关键观点3: TorchTitan和PyTorch中的Async-TP使用指南
文章提供了关于如何在TorchTitan和PyTorch中使用Async-TP的详细指南，包括使用torch.compile自动检测和应用Async-TP，以及在eager模式下直接使用Async-TP算子的方法。同时，还讨论了局限性和未来工作计划。
正文
上半部分 (Original传统方式): 两个分区(Partition 0和1)按顺序执行, 先进行AllGather操作收集数据, 然后执行Einsum矩阵计算，这种方式计算和通信是串行的,存在等待时间。下半部分 (Overlapped异步方式): 通信被分解为异步Send和Recv操作，计算也被分解为更小的Einsum操作。两个分区可以同时进行: Partition 0发送A0到Partition 1,同时计算Einsum(A0, B0)，Partition 1发送A1到Partition 0,同时计算Einsum(A1, B1)，通过Dynamic Update更新计算结果。这种方法通信和计算可以重叠执行，减少了整体等待时间，如图右侧箭头所示,相比传统方式节省了时间。
Wang等人的图示展示了如何将此技术应用于all-gather后跟一个matmul。all-gather被分解为send和recv操作,而matmul被分为子matmuls。通过这种分解,可以同时计算一个子matmul,同时传输下一个子matmul所需的数据,从而有效地隐藏通信延迟。
性能挑战
尽管异步张量并行的概念在理论上很简单直观,但要在CUDA中实现高性能的实现却面临着一些挑战。在本节中,我们将探讨这些挑战,并讨论我们采用的解决方案。
致谢
: 这些挑战中的许多最初是由Luca Wehrstedt(https://discuss.pytorch.org/u/lcw/summary)探索的。PyTorch中的异步TP实现从他在xformers中的异步TP工作中获得了重要启发。
通信开销
在分解通信时,使用NCCL的send/recv操作可能很有诱惑力,因为它们易于使用。然而,NCCL的send/recv操作具有一些特性,使其不太适合异步张量并行:
-
重叠计算和通信之间的竞争
- 虽然人们普遍认为计算和通信是可以独立使用的两种资源,但实际上它们的独立性是有细微差别的,确实会发生竞争。在节点内设置(TP最常见的情况)中,NCCL的send/recv kernel 会利用SM通过NVLink传输数据,这减少了可用于重叠矩阵乘法 kernel 的SM数量,从而降低了速度。有趣的是,观察到的速度下降可能超过通信 kernel 消耗的资源百分比。由于cuBLAS试图选择以完整waves执行的 kernel ,通信 kernel 占用的SM可能会打破平衡,导致矩阵乘法 kernel 需要执行额外的wave。
-
双向同步
- NCCL的send/recv kernel 执行双向同步,这意味着发送方和接收方都会被阻塞直到操作完成。这种方法对于算子内并行中的数据传输并不总是最优的。根据具体情况,可能更适合对多个数据传输执行单次同步,或者在向远程节点推送数据和从远程节点拉取数据之间进行选择。
幸运的是,我们可以通过利用CUDA的P2P机制来避免前面提到的缺点。该机制允许设备通过将其映射到虚拟内存地址空间来访问对等设备上分配的内存。这种机制使得内存操作(加载/存储/原子等)可以通过NVLink执行(目前,PyTorch中的async-TP实现需要所有设备对之间都有NVLink连接(例如通过NVSwitch)才能实现加速。这是我们计划在未来解决的限制)。此外,当通过cudaMemcpyAsync在对等设备之间传输连续数据时,该操作由拷贝引擎(拷贝引擎是GPU上的专用硬件单元,用于管理不同内存位置之间的数据移动,独立于GPU的计算核心(SM)运行)处理,不需要任何SM,从而避免了前面讨论的竞争问题(通过拷贝引擎的数据传输仍然共享相同的内存带宽。然而,这不太可能造成显著的竞争,因为(1)传输速率受限于NVLink带宽,低到足以避免内存带宽竞争,(2)重叠的矩阵乘法是计算密集型的)。
为了在未来利用这种机制实现async-TP和类似用例,我们开发了一个名为SymmetricMemory的实验性抽象(https://github.com/pytorch/pytorch/blob/main/torch/csrc/distributed/c10d/SymmetricMemory.hpp)。从概念上讲,它表示一个在设备组之间对称分配的缓冲区,通过虚拟内存/多播地址为每个GPU提供对其对等设备上所有相应缓冲区的访问。使用async-TP不需要直接与SymmetricMemory交互,但用户可以利用它来创建类似于async-TP的自定义细粒度、节点内/算子内优化。
由于工作分解导致的Wave量化效率放大问题
