在用戶數量可能在數百到數十萬之間波動，且輸入序列長度隨每個請求而變化的生產環境中，部署生成式 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 數量的超線性擴展。
實現最佳推理性能需要仔細的工作負載分析和對性能瓶頸的深入了解。通過內部工程工作以及與外部開發者和研究人員的密切合作，我們可以快速、頻繁地優化平臺的許多方面，為用戶提供出色的性能。
隨著我們繼續識別和實施新的性能優化（一些可能是廣泛的，另一些可能范圍較窄），我們將定期提供有關這些優化的更新，提供技術動機和量化效益。
?
