在用户数量可能在数百到数十万之间波动，且输入序列长度随每个请求而变化的生产环境中，部署生成式 AI 工作负载会面临独特的挑战。要在这些环境中实现低延迟推理，无论 GPU 生成方式或显存容量如何，多 GPU 设置都是必需的。为了提高生产级设置中的推理性能，我们很高兴推出 TensorRT-LLM Multi-shot，这是一种新的多 GPU 通信协议，利用 NVIDIA NVLink Switch 可将通信速度大幅提升高达 3 倍。本博客概述了这一新功能，以及它如何帮助开发者和解决方案架构师克服传统多 GPU 通信方法的限制。
传统 AllReduce 算法面临的挑战
对于低延迟推理，无论单个 GPU 的显存容量如何，多 GPU 都至关重要。但是，在低并发情况下，GPU 花在交换数据上的时间可能超过花在计算上的时间。为了获得最佳性能， 高效的 AllReduce 操作 –结合每个参与其中的 GPU 的部分结果的集合操作–至关重要。
传统方法使用基于环的算法，其中部分值在环形的 GPU 之间传递。每个 GPU 都贡献其值并将结果传递给其邻居。该过程重复 2N-2 次，其中 N 是协同工作的 GPU 数量，在该过程结束时，每个 GPU 都具有相同的总和值。需要对环进行第二次传递，以将总和值从最后一个 GPU 传播到其他 GPU。
Ring 方法可在每个通信步骤中高效利用可用的 GPU 到 GPU 带宽，但随着 GPU 数量的增加，步骤数也会增加。这会增加延迟，因为所有 GPU 都需要在 Ring 的每个步骤中保持同步。这些同步延迟会显著增加延迟开销，并可能导致难以满足更严格的延迟目标。
Ring AllReduce 算法描述如下：
- 环形算法：GPU-1 → GPU-2 → … → GPU-N → GPU-1 → GPU-2 → … → GPU-(N-1)
- 2N-2 步长，每步具有完整的 Tensor send/recv
- 延迟：2N-2 通信步骤。(N：GPU 的数量)
- 流量：(4N-4)/N 张量的 send/recv 字节数
使用 TensorRT-LLM MultiShot 应对 AllReduce 通信挑战
TensorRT-LLM MultiShot 是一种新算法，可利用 NVSwitch 中的组播，将 Ring AllReduce 的 O(N) 延迟最多降低 3 倍。组播是 NVSwitch 中的硬件加速功能，允许一个 GPU 发送数据一次，并将该数据同时发送到所有其他 GPU，从而将通信步骤的数量减少到两个 GPU 间的同步，同时保持带宽效率。如果没有 NVSwitch，这将占用 N 倍的通信带宽。
TensorRT-LLM Multishot 将 AllReduce 分离为 ReduceScatter 操作，然后是 AllGather 操作（有关集合操作的更多详细说明，请参阅 此文档 ）。
每个 GPU 仅负责累积结果张量的一部分。
第一步（或“shot”）涉及每个 GPU 将张量的不同切片发送到负责累积该张量切片的相应 GPU。
在本地累加后，每个 GPU 现在都有正确的和累加器，用于其独特的输出切片。
在第二步 (或“shot”) 中，每个 GPU 使用 NVSwitch 组播功能将结果切片广播到所有其他 GPU。这可最大限度地减少 NVSwitch 本身执行数据放大所需的每个 GPU 带宽；每个 GPU 一步发送 1/N 数据并接收完整的结果张量。
无论执行张量并行推理的 GPU 数量如何，整个操作仅需两次通信步骤。
- TensorRT-LLM MultiShot 算法：GPU_N 发送切片、计算切片和、在单个组播运算中广播结果。
- 延迟：2 个通信步骤（与 GPU 数量无关）
- 流量：2 张量字节的 send/recv（与 GPU 数量无关）
为何如此重要
由于此算法只需要两个通信步骤，而不是 2N-2 (其中 N 表示 GPU 数量)，因此 MultiShot 的速度几乎是 Ring AllReduce 的 3 倍。这种算法的优势在消息大小较小且并行度高的情况下尤为明显，而这正是需要最低延迟以获得出色的用户体验的场景。
这可用于降低最小延迟，或在给定延迟下提高吞吐量。在具有更激进的延迟阈值的场景中，这可能会导致 GPU 数量的超线性扩展。
实现最佳推理性能需要仔细的工作负载分析和对性能瓶颈的深入了解。通过内部工程工作以及与外部开发者和研究人员的密切合作，我们可以快速、频繁地优化平台的许多方面，为用户提供出色的性能。
随着我们继续识别和实施新的性能优化（一些可能是广泛的，另一些可能范围较窄），我们将定期提供有关这些优化的更新，提供技术动机和量化效益。
