大型模型的張量平行性
將單一神經網路層內的數學分割到多個 GPU 上的方法,這樣對於一個裝置來說太大的模型仍然可以運作。
概述
It matters because frontier models have hundreds of billions of parameters that no single GPU can hold or compute fast enough alone.
深入探討
張量並行性(也稱為層內模型並行性)在 GPU 上分割各個權重矩陣,而不是將整個層放在單獨的裝置上。在 Transformer 中,大矩陣乘法(注意力投影和前饋 MLP)被拆分:例如,MLP 的第一個權重矩陣按列分區,第二個權重矩陣按行分區,因此每個 GPU 計算一個切片,然後單個 all-reduce 組合結果。注意力分散在多個頭中,每個 GPU 處理一個子集。由於每個 GPU 同時執行每一層的一部分,因此張量並行性會減少每個 GPU 的記憶體並加快運算速度,但它需要每層 GPU 之間進行頻繁的高頻寬通訊。這就是為什麼它通常被限制在由 NVLink 連接的節點內,並與管道和資料並行性相結合,以實現非常大的訓練和服務作業。
技術洞察
由 Megatron-LM 推廣的技巧是選擇分區尺寸,從而最大限度地減少通訊。按列拆分第一個 MLP 矩陣可讓每個 GPU 在本地應用非線性,無需同步;逐行拆分第二行意味著輸出只需要一次全歸約即可對部分結果求和。因此,每一層大約會產生兩個全歸約(向前)和兩個(向後)。由於這些集合發生在每一層,因此延遲占主導地位,因此張量並行性存在於 NVLink 等快速節點內連結背後,而不是較慢的節點間網路背後。
戰略影響
成本與預算
多年來,架構決策決定著效能和營運成本。
更明確的決策
技術教育幫助團隊選擇正確的堆疊,而不僅僅是最新的堆疊。
品質管控
更好的工程選擇可以減少生產中的可靠性事故。
大型模型張量並行的未來
張量並行性仍然是基礎性的,但越來越多地融入「3D並行性」(張量+管道+數據),並與混合專家模型的專家並行性相結合。 Megatron-LM、DeepSpeed 和 vLLM 等框架可自動執行分片。隨著 GPU 互連(NVLink、NVSwitch)和光學結構變得更快,節點邊界限制放寬,允許更廣泛的張量並行組。期望更聰明的自動並行化,選擇分片尺寸和群組大小,以最大限度地減少給定集群拓撲的通訊。
現實世界的實施
使用 Megatron-LM 將每一層的權重矩陣分片到一個 NVLink 連接節點中的 8 個 GPU 上,從而訓練 175B 參數模型。
在 vLLM 中使用 tensor_parallel_size=4 提供 70B 參數的聊天模型,以便權重適合四個 GPU 並即時回應。
將 Transformer 注意力頭拆分到 GPU 上,以便每個裝置計算一個子集,然後連接下一層的輸出。
結合節點內的張量平行性和跨節點的管道並行性,在大型 GPU 叢集上訓練萬億參數模型。
風險與防護欄
優化一項基準測試可以隱藏更廣泛的系統弱點。
基礎設施和維護成本常常被低估。
隨著系統變得更加複雜,安全性和可觀察性差距可能會擴大。
實施路線圖
在實施之前定義延遲、品質和成本目標。
在實際負載和資料條件下進行基準測試。
儀器監控錯誤、漂移和使用者影響。
在擴展之前準備回滾和事件回應路徑。
不斷探索
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the Tensor Parallelism for Large Models quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
常見問題
What is Tensor Parallelism for Large Models?
將單一神經網路層內的數學分割到多個 GPU 上的方法,這樣對於一個裝置來說太大的模型仍然可以運作。這很重要,因為前沿模型擁有數千億個參數,沒有任何一個 GPU 能夠單獨保存或計算得足夠快。
GPU 之間的張量平行性有何不同?
張量(層內)並行性將層內的權重矩陣進行分片,因此每個 GPU 都計算同一層的一部分,這與管道並行性不同,管道並行性將整個層放置在不同的 GPU 上。
在威震天式 MLP 分割中,兩個權重矩陣如何劃分以最小化通訊?
按列分割第一個矩陣可以讓每個 GPU 在本地應用非線性;將第二個按行拆分意味著單一全歸約對部分輸出求和,從而最大限度地減少同步。
為什麼張量並行性通常保留在單一節點內?
每層都會觸發集體通訊(全歸約),因此對延遲敏感的流量需要快速的節點內連結(如 NVLink),而不是較慢的節點間網路。
在張量並行下,注意力機制通常是如何並行的?
注意力頭是獨立的,因此它們分佈在 GPU 上——每個設備計算一些頭,然後組合輸出。
在張量並行中,什麼集體操作通常會組合部分結果?
All-reduce 對每個 GPU 上計算的部分輸出求和,以便它們就該層的結果達成一致;在前向和後向傳遞中,每層都會發生多次這種情況。