技術指南

優化器狀態卸載到 CPU 和 NVMe

節省記憶體的技巧,將繁重的訓練記錄(優化器狀態、梯度,有時是權重)存放在 CPU RAM 或 NVMe SSD 上,而不是稀缺的 GPU 記憶體中。

閱讀時間約2分鐘最後更新

概述

It lets people train far larger models than their GPU's memory would otherwise allow.

深入探討

當您使用 Adam 這樣的優化器訓練神經網路時,每個參數都會帶來額外的負擔:兩個運行統計資料(動量和方差),加上權重的全精度副本及其梯度。在混合精準度訓練中,每個參數總計約 16 個字節,使權重本身的 2 個位元組相形見絀。卸載將這些負擔從 GPU 上移除。 CPU 卸載透過 PCIe 匯流排將最佳化器狀態串流傳輸到普通系統 RAM,而 NVMe 卸載將它們一直推送到快速固態磁碟。這項技術因 DeepSpeed 的 ZeRO-Infinity 和 ZeRO-Offload 而流行,以原始速度換取容量,讓單一 GPU 或小型叢集能夠微調具有數十億個參數的模型。

技術洞察

關鍵是將資料移動與計算重疊。優化器狀態位於 CPU/NVMe 中;在向後傳遞期間,分區在需要之前透過 PCIe 預取,且優化器步驟本身通常在 CPU 上運行。 ZeRO-Offload 將 float32 主權重和 Adam 矩保留在 CPU 上,因此只有前向和後向數學保留在 GPU 上。 NVMe 新增了分層緩存,因此 TB 等級狀態會溢出到磁碟,而熱分割區則保留在 RAM 中。

戰略影響

成本與預算

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

更明確的決策

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

品質管控

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

優化器狀態卸載到 CPU 和 NVMe 的未來

隨著模型的 GPU 記憶體不斷增長,分層卸載正在成為標準而不是奇特的。預計與更快的互連(例如模糊 CPU-GPU 邊界的 NVLink-C2C 和 CXL 記憶體池)以及更聰明的調度程式(可以預測要預取的狀態)的更緊密整合。 Grace Hopper 等統一記憶體架構減少了 PCIe 損失,框架正在推動多層卸載幾乎透明,以便愛好者可以在適度的硬體上微調大型模型。

現實世界的實施

使用 DeepSpeed ZeRO-Offload 在單一 24 GB 消費級 GPU 上微調 130 億參數 LLM,將 Adam 狀態推送到 CPU RAM。

一個小型研究實驗室透過 ZeRO-Infinity 將優化器狀態溢位到 NVMe 驅動器,在幾個 GPU 上訓練數十億參數模型。

Hugging Face Accelerate 設定可啟用 CPU 卸載,以便使用者可以執行完整的微調作業,否則會引發記憶體不足錯誤。

注重成本的新創公司租用更便宜、記憶體較低的雲端 GPU 並卸載到附加的 NVMe,而不是購買頂級 80 GB 卡。

風險與防護欄

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

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

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

實施路線圖

1

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

2

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

3

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

4

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

不斷探索

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 Optimizer State Offloading to CPU and NVMe 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

下一步指南

Adam 和自適應優化器

常見問題

What is Optimizer State Offloading to CPU and NVMe?

節省記憶體的技巧,將繁重的訓練記錄(優化器狀態、梯度,有時是權重)存放在 CPU RAM 或 NVMe SSD 上,而不是稀缺的 GPU 記憶體中。它可以讓人們訓練比 GPU 記憶體允許的更大的模型。

將優化器狀態卸載到 CPU 或 NVMe 的主要目標是什麼?

卸載將繁重的優化器狀態從 GPU 移至 CPU RAM 或 NVMe,從而釋放 GPU 內存,以便可以在同一硬體上訓練更大的模型。

對於混合精度的 Adam 優化器,哪些資料通常每個參數消耗最多的記憶體?

Adam 儲存動量、變異數和全精確度主權重,每個參數總共約 16 個位元組,遠遠超過 2 位元組半精確度權重。

哪個框架功能普及了大規模優化器卸載?

DeepSpeed 的 ZeRO-Offload 和 ZeRO-Infinity 大規模引入並推廣了 CPU 和 NVMe 的卸載優化器狀態。

卸載到 CPU 或 NVMe 的主要效能成本是什麼?

將狀態移出 GPU 意味著透過 PCIe 或 NVMe 傳輸數據,這比 GPU 記憶體上的速度慢,因此除非與計算重疊,否則吞吐量會下降。

卸載系統如何隱藏大部分傳輸延遲?

當 GPU 仍在運算時,調度程式透過 PCIe 預取所需的分割區,將通訊與運算重疊以掩蓋延遲。