技術指南

ZeRO 和分片優化器

ZeRO(零冗餘優化器)透過跨 GPU 分片優化器狀態、梯度和權重,消除了資料並行性中浪費的記憶體重複。

概述

ZeRO(零冗餘優化器)透過跨 GPU 分片優化器狀態、梯度和權重,消除了資料並行性中浪費的記憶體重複。它可以讓您透過簡單的資料並行性來訓練巨大的模型,但只需要一小部分 GPU 記憶體。

ZeRO 和分片優化器是一個技術構建塊,會大規模影響模型品質、基礎設施成本、延遲和可靠性。

深入探討

在普通資料並行中,每個 GPU 都會儲存優化器狀態、梯度和參數的冗餘完整副本,這是極大的浪費,特別是對於 Adam 來說,其中優化器狀態可能是模型本身大小的幾倍。 DeepSpeed 中的 Microsoft 引入的 ZeRO 透過在 GPU 上對這些張量進行分區來消除這種冗餘,這樣每個設備只擁有一個切片。 ZeRO 分為三個漸進階段:第 1 階段將優化器狀態分片,第 2 階段新增梯度分片,第 3 階段將參數本身分片。根據需要,GPU 透過通訊、計算收集丟失的切片,然後釋放它們。結果是每個 GPU 的記憶體顯著降低,支援數十億到兆參數的訓練,同時保持資料並行的簡單程式設計模型。

技術洞察

ZeRO 用额外的通信来节省内存。在第 3 階段,在某個層的前向傳遞之前,全收集將該層的完整參數收集到每個 GPU 上;之後,非擁有的切片將被丟棄以回收記憶體。梯度是減少分散的,因此每個 GPU 僅保留與其擁有的參數相符的梯度切片。 PyTorch 的 FSDP(完全分片資料並行)原生實現了相同的想法,將模組包裝為動態分片和重新分片。

掌握 ZeRO 和分片优化器

ZeRO(零冗餘優化器)透過跨 GPU 分片優化器狀態、梯度和權重,消除了資料並行性中浪費的記憶體重複。它可以讓您透過簡單的資料並行性來訓練巨大的模型,但只需要一小部分 GPU 記憶體。 ZeRO 和分片優化器是一個技術構建塊,會大規模影響模型品質、基礎設施成本、延遲和可靠性。為了建立深入的理解,請將 ZeRO 和分片優化器視為一種操作模型,而不是單一功能:定義期望的結果,澄清假設,並將系統可以可靠地完成的任務與仍需要專家判斷的任務分開。

在實踐中,強大的團隊使用 ZeRO 和分片優化器根據可靠性和成本優化架構、資料和基礎架構選擇。他們記錄明確的成功標準,根據實際數據和工作流程進行測試,並根據觀察到的失敗模式而不是一次性基準測試勝利進行迭代。這就是理論理解轉變為跨產品、政策和營運的持久能力的地方。

多年來,架構決策決定著效能和營運成本。同時,優化一個基準測試可以隱藏更廣泛的系統弱點。最具彈性的方法是將實驗速度與治理規則結合:運行試點、捕獲證據、發布決策日誌,並隨著模型行為、使用者期望和監管要求的發展不斷更新保障措施。

戰略影響

多年來,架構決策決定著效能和營運成本。

多年來,架構決策決定著效能和營運成本。在高品質部署中,這會轉化為可衡量的操作規則、所有權邊界和定期審查儀式,以便團隊可以增強信心,而不是擴大模糊性。

技術教育幫助團隊選擇正確的堆疊,而不僅僅是最新的堆疊。

技術教育幫助團隊選擇正確的堆疊,而不僅僅是最新的堆疊。在高品質部署中,這會轉化為可衡量的操作規則、所有權邊界和定期審查儀式,以便團隊可以增強信心,而不是擴大模糊性。

更好的工程選擇可以減少生產中的可靠性事故。

更好的工程選擇可以減少生產中的可靠性事故。在高品質部署中,這會轉化為可衡量的操作規則、所有權邊界和定期審查儀式,以便團隊可以增強信心,而不是擴大模糊性。

ZeRO 和分片优化器的未来

分片正在成為大規模訓練的預設選擇,而不是一種奇特的選擇。期望與卸載進行更深入的整合(透過 ZeRO-Infinity 將切片推送到 CPU 或 NVMe),透過計算更好地重疊全收集和減少分散以隱藏其成本,以及與張量和管道並行性的組合。隨著模型不斷增長,記憶體高效的分片優化器對於將模型適應實際的硬體預算至關重要。

現實世界的實施

使用 DeepSpeed ZeRO Stage 2 微調數十億參數的語言模型,否則模型會溢出 GPU 記憶體。

使用 PyTorch FSDP 進行訓練,跨 GPU 分片參數、梯度和優化器狀態,並根據需要每層收集它們。

應用 ZeRO-Offload 將優化器狀態推送到 CPU 內存,讓單一 GPU 訓練比其 VRAM 大很多倍的模型。

當 GPU 和 CPU 記憶體耗盡時,透過從 NVMe 儲存串流傳輸參數分片,使用 ZeRO-Infinity 擴展萬億參數模型。

實施模式

ZeRO 和分片优化器的实践

使用 DeepSpeed ZeRO Stage 2 微調數十億參數的語言模型,否則模型會溢出 GPU 記憶體。

使用 DeepSpeed ZeRO Stage 2 微調數十億參數的語言模型,否則會溢位 GPU 記憶體。當團隊預先定義品質閾值、為邊緣情況保留人工升級路徑並隨著時間的推移追蹤生產力增益和錯誤成本時,通常會獲得更好的結果。

ZeRO 和分片优化器的实践

使用 PyTorch FSDP 進行訓練,跨 GPU 分片參數、梯度和優化器狀態,並根據需要每層收集它們。

使用 PyTorch FSDP 進行訓練,它跨 GPU 分片參數、梯度和優化器狀態,並按需按層收集它們。當團隊預先定義品質閾值、為邊緣情況保留人工升級路徑並隨著時間的推移追蹤生產力增益和錯誤成本時,通常會獲得更好的結果。

ZeRO 和分片优化器的实践

應用 ZeRO-Offload 將優化器狀態推送到 CPU 內存,讓單一 GPU 訓練比其 VRAM 大很多倍的模型。

應用 ZeRO-Offload 將優化器狀態推送到 CPU 內存,讓單一 GPU 訓練比其 VRAM 大很多倍的模型。當團隊預先定義品質閾值、為邊緣情況保留人工升級路徑並隨著時間的推移追蹤生產力增益和錯誤成本時,通常會獲得更好的結果。

ZeRO 和分片优化器的实践

當 GPU 和 CPU 記憶體耗盡時,透過從 NVMe 儲存串流傳輸參數分片,使用 ZeRO-Infinity 擴展萬億參數模型。

當 GPU 和 CPU 記憶體耗盡時,透過從 NVMe 儲存串流參數分片,使用 ZeRO-Infinity 擴展萬億參數模型 團隊在預先定義品質閾值、為邊緣情況保留人工升級路徑並隨著時間的推移追蹤生產力增益和錯誤成本時,通常會獲得更好的結果。

風險與防護欄

!

優化一項基準測試可以隱藏更廣泛的系統弱點。

!

基礎設施和維護成本常常被低估。

!

隨著系統變得更加複雜,安全性和可觀察性差距可能會擴大。

實施路線圖

1

在實施之前定義延遲、品質和成本目標。

在實施之前定義延遲、品質和成本目標。將每個步驟視為證據門:如果不符合標準,則暫停推出,縮小差距,然後再擴大使用。

2

在實際負載和資料條件下進行基準測試。

在實際負載和資料條件下進行基準測試。將每個步驟視為證據門:如果不符合標準,則暫停推出,縮小差距,然後再擴大使用。

3

儀器監控錯誤、漂移和使用者影響。

儀器監控錯誤、漂移和使用者影響。將每個步驟視為證據門:如果不符合標準,則暫停推出,縮小差距,然後再擴大使用。

4

在擴展之前準備回滾和事件回應路徑。

在擴展之前準備回滾和事件回應路徑。將每個步驟視為證據門:如果不符合標準,則暫停推出,縮小差距,然後再擴大使用。

不斷探索