张量并行 vs 流水线并行：多GPU上的AI模型并行化
多GPU上的AI模型并行化：张量并行 vs 流水线并行
现代AI模型越来越大，常常超过单个GPU的内存（VRAM）。为了训练或运行这些模型，多GPU并行化至关重要。目前主要有两种策略：张量并行（TP）和流水线并行（PP）。两者都分配计算，但方式根本不同，对延迟、吞吐量和硬件要求有不同的影响。本文解释每种方法的工作原理、涉及的额外开销以及每种方法的最佳用例。
流水线并行（PP）：按层划分
在流水线并行中，模型被垂直分割：神经网络的层被分配到不同的GPU上。例如，GPU 0处理前10层，GPU 1处理接下来的10层，依此类推。数据像流水线一样流经GPU。GPU 0计算其层并将激活传递给GPU 1。这引入了一个固有问题：流水线气泡。由于每个GPU必须等待前一个GPU完成，因此会出现空闲时间。在简单配置中，GPU经常空闲等待数据，降低了整体效率。然而，当处理批次（同时处理多个请求）时，这个缺点得到了缓解。当许多请求按顺序到达时，流水线可以被填满：当GPU 1仍在处理请求A时，GPU 0已经在处理请求B。这掩盖了空闲时间，显著提高了吞吐量。因此，PP适用于高批次吞吐量场景，其中单个请求的延迟不那么关键。
张量并行（TP）：按矩阵划分
相比之下，张量并行通过将每个层内的单个矩阵运算分布到多个GPU上来水平分割模型。每个GPU只持有每层权重的一部分（分片）。要计算一层的最终结果，所有GPU必须协作：它们执行部分计算，然后通过称为all-reduce的同步操作组合它们。这需要GPU之间持续通信——几乎在每一层。由此产生的通信开销是TP性能的关键因素。这种通信的速度很大程度上取决于所使用的互连技术。NVLink是GPU之间的专用高速链路，可以在几微秒内传输数据，使开销可以忽略不计。然而，当使用较慢的PCIe总线时，数据传输成为瓶颈。PCIe的带宽远低于GPU内部内存的带宽。等待数据交换所花费的时间可能占用总计算时间的很大一部分——在极端情况下通常为25%到50%。这可能会抵消并行化的好处。通过PCIe的典型all-reduce可能需要几毫秒，而实际计算只需要几微秒。没有NVLink，TP通常效率低下。
比较与开销
| 属性 | 流水线并行（PP） | 张量并行（TP） |
|---|---|---|
| 划分方式 | 按层（垂直） | 按矩阵（水平） |
| 通信频率 | 仅在层块之间（低频） | 每一层（非常频繁） |
| 通信量 | 大的激活数据（每块一次） | 小到中等数据（但非常频繁） |
| 主要开销 | 流水线气泡（空闲时间） | all-reduce同步（延迟） |
| 延迟（单次请求） | 较高（由于顺序步骤） | 较低（并行计算） |
| 吞吐量（批次） | 高（气泡被掩盖） | 中等（通信开销随规模增长） |
| 硬件要求 | 在PCIe上表现良好 | 需要NVLink或类似的快速互连 |
| 理想场景 | 大批次、多服务器、PCIe环境 | 配备NVLink的单服务器、延迟关键型应用 |
结论
张量并行和流水线并行之间的选择很大程度上取决于可用的硬件和性能目标。如果GPU通过NVLink连接，并且需要最小化单次请求的延迟，那么TP是更好的选择。它实现了细粒度并行化和低延迟，但需要非常快速的互连。如果只有PCIe连接可用，或者模型分布在多个服务器上，那么PP是更稳健、更高效的选择。流水线气泡通过批处理得到缓解，通信开销仍然很低。在实践中，这两种技术经常结合使用（例如，在节点内使用NVLink的TP，在节点间通过网络使用PP），以利用两者的优势。理解这些差异对于在AI中高效利用多GPU系统至关重要。