技術指南

DeepSpeed 和 Megatron 訓練堆疊

DeepSpeed (Microsoft) 和 Megatron-LM (NVIDIA) 是軟體堆疊,使數千個 GPU 的數十億參數的訓練模型變得切實可行。

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

概述

沒有這些,現今的前沿模型根本無法在記憶體中容納,也無法在合理時間內完成訓練。

深入探討

在一個 GPU 上訓練大型模型是不可能的,因為權重、梯度和優化器狀態不適合。這些堆疊將工作分配給許多 GPU。 Megatron-LM 開創了張量並行性,在 GPU 上的每一層內切片單獨的矩陣乘法,以及管道並行性,將不同的層放在不同的 GPU 上。 DeepSpeed 的標誌性貢獻是 ZeRO(零冗餘優化器),它將優化器狀態、梯度和參數跨 GPU 分片而不是複製它們,從而大幅削減每個 GPU 的記憶體。兩者經常結合(Megatron-DeepSpeed)來訓練 BLOOM-176B 和 Megatron-Turing NLG 等模型。他們還添加了混合精度、激活檢查點和卸載到 CPU 或 NVMe,因此大型模型可以在有限的硬體上進行訓練。

技術洞察

ZeRO 有三個增加記憶體節省的階段:第 1 階段對優化器狀態進行分片,第 2 階段還對梯度進行分片,第 3 階段對參數本身進行分片,在前向和後向傳遞過程中按需收集它們。與張量並行性(層內)和管道並行性(層間)結合,形成「3D並行性」。關鍵的壓力是通訊開銷:每個分片分割都會增加 GPU 到 GPU 的流量,因此工程師調整分割以保持快速 NVLink 和 InfiniBand 連結飽和。

戰略影響

成本與預算

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

更明確的決策

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

品質管控

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

DeepSpeed 和 Megatron 訓練堆疊的未來

預計與 PyTorch 的原生 FSDP(完全分片資料並行)進行更緊密的集成,它吸收了許多 ZeRO 思想,模糊了研究堆疊和核心框架之間的界限。編譯器驅動的方法和自動並行規劃器旨在消除手動調整。隨著訓練集群成長到數十萬個加速器,容錯、彈性擴展和與運算的重疊通訊成為主要的工程前沿,同時支援 NVIDIA Blackwell 和客製化訓練晶片等新硬體。

現實世界的實施

使用跨數百個 GPU 的組合 Megatron-DeepSpeed 堆疊來訓練開放式多語言 BLOOM-176B 模型。

Microsoft 和 NVIDIA 使用 3D 平行訓練 5,300 億參數的 Megatron-Turing NLG 模型。

ZeRO-Offload 讓研究人員透過將優化器狀態溢位到 CPU RAM 來微調單一工作站 GPU 上的數十億參數模型。

在這些堆疊中使用激活檢查點,透過重新計算激活而不是儲存全部激勵來適應更長的上下文視窗。

風險與防護欄

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

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

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

實施路線圖

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 DeepSpeed and Megatron Training Stacks 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

下一步指南

GPTQ 和 AWQ 訓練後量化

常見問題

什麼是深速與密卡登訓練疊加?

DeepSpeed (Microsoft) 和 Megatron-LM (NVIDIA) 是軟體堆疊,使數千個 GPU 的數十億參數的訓練模型變得切實可行。如果沒有它們,今天的前沿模型根本無法適應記憶體或在合理的時間內完成訓練。

DeepSpeed 的 ZeRO 優化器的主要目的是什麼?

ZeRO(零冗餘優化器)透過跨 GPU 劃分優化器狀態、梯度和參數而不是在每個裝置上複製它們來消除記憶體冗餘。

Megatron-LM 中首創的張量並行如何分割模型?

張量並行性將單一層內的數學(例如大型矩陣乘法)跨多個 GPU 進行劃分,這就是層內分割。

哪個 ZeRO 階段透過對模型參數本身進行分片來提供最大的記憶體節省?

除了梯度和優化器狀態之外,ZeRO Stage 3 還對參數進行分片,按需收集它們,從而最大程度地減少記憶體。

在訓練過程中,「啟動檢查點」會犧牲什麼來節省記憶體?

激活檢查點儲存較少的中間激活,並在反向傳播期間重新計算它們,用額外的計算換取減少的記憶體。

為什麼這些技術的組合通常被稱為「3D 並行性」?

3D並行性堆疊了三種正交策略:資料並行性、張量(層內)並行性和管道(層間)並行性。