[論文翻譯] 分佈式訓練 Parameter sharding 之 ZeRO
0x00 摘要
Parameter sharding 就是把模型參數等切分到各個GPU之上，以此達到使用較少GPU實現大規模模型訓練的目的。本系列會以 Google，微軟和Facebook的論文，博客以及代碼來對parameter sharding 進行分析，大約有 5～6篇文章。
本文以 ZeRO: Memory Optimizations Toward Training Trillion Parameter Models 和 DeepSpeed: Extreme-scale model training for everyone爲主來進行分析，這是微軟開發的一個可以高效利用顯存的優化器，其會將模型狀態量（優化器狀態，梯度和模型參數）分佈在多個並行 GPU 之上，目的是在不使用模型並行的情況下對讓數十億參數模型進行訓練。
ZeRO是ZeRO-DP和ZeRO-R兩種方法的組合。ZeRO-DP是一種增強數據並行機制，它使用動態通信策略來將優化器狀態、梯度和參數進行分區，以最小化通信量和避免模型狀態的冗餘。ZeRO-R則使用分區激活重計算、恆定大小緩衝區和動態內存碎片整理機制來優化剩餘狀態的內存消耗。
本文並不會逐字分析，而是選擇了一些重點，並且爭取加入筆者自己的理解。
本系列其他文章如下：
[源碼解析] PyTorch 分佈式之 ZeroRedundancyOptimizer
0x01 綜述
這部分主要翻譯自 DeepSpeed: Extreme-scale model training for everyone。
1.1 挑戰
首先，我們需要了解訓練巨大模型所帶來的的顯存和計算效率的挑戰。
1.1.1 顯存效率
訓練萬億參數模型所需的顯存遠遠超出了單張 GPU 的顯存大小。比如，在使用 Adam 優化器進行混合精度訓練時需要約 16TB 的顯存來存儲模型狀態量（參數、梯度和優化器狀態量）。僅僅爲了存儲模型狀態，就需要 400 張英偉達 A100 GPU（每張40 GB 的顯存）。
激活函數也需要佔據額外的顯存，其隨批量大小（batch size）而增加。batch size設置爲1的情況下，訓練萬億參數模型就會使用超過 1 TB 顯存來存儲激活。人們也嘗試用 checkpoint 來處理激活顯存，就是用計算來換顯存，這可以將該顯存減少到大約20 GB，但是對於訓練而言，這個顯存需求仍然過高。
所以，必須在多個 GPU 設備之間有效地劃分模型狀態和激活顯存，才能讓這種大模型在不耗盡顯存的情況下進行訓練。
1.1.2 計算效率
基於 OpenAI 的研究 law of scaling來估算，端到端訓練一個萬億參數的模型大約需要 5000 Zflops（即 5 後面帶有 24 個零）。訓練這樣一個模型需要 4000 張 A100 以 50％ 的計算效率運行大約 100 天。
儘管大型超級計算 GPU 集羣可以擁有超過 4000 個 GPU，但是由於 batch size 的限制，要在這種規模上實現高計算效率仍然很具有挑戰性。計算效率隨着計算時間對通信時間的比例的增加而增加。該比例與 batch size成正比。但是，模型可以訓練的 batch size有一個上限，如果超過這個上限，則收斂情況會迅速惡化。
1.2 權衡
我們接下來看看數據並行、模型並行和流水線並行之間的權衡。
1.2.1 數據並行
數據並行是深度學習中的十分常見的技術。在數據並行中，每批輸入的訓練數據都在數據並行的 worker 之間進行平分。反向傳播之後，我們需要進行通信來規約梯度，以保證優化器在各個 worker 上可以得到相同的更新。數據並行性具有幾個明顯的優勢，包括計算效率高和工作量小。但是，數據並行的 batch size 會隨 worker 數量提高，而我們難以在不影響收斂性的情況下無限增加 batch szie。
- 顯存效率：數據並行會在所有 worker 之間複製模型和優化器，因此顯存效率不高。
- 計算效率：隨着並行度的提高，每個 worker 執行的計算量是恆定的。數據並行可以在小規模上實現近乎線性擴展。但是，因爲在 worker 之間規約梯度的通信成本跟模型大小成正相關，所以當模型很大或通信帶寬很低時，計算效率會受到限制。梯度累積是一種常見的用來均攤通信成本的策略，它可以增加batch size，在本地使用 micro-batch 進行多次正向和反向傳播，在進行優化器更新之前再規約梯度，從而分攤通信成本。
1.2.2 模型並行
模型並行是另一大類技術。它可以在多個 worker 之間劃分模型的各個層。就其本質而言，模型並行的計算和通信因模型結構而異，因此需要很大的工作量來實現。DeepSpeed 利用了英偉達的 Megatron-LM 來構建基於 Transformer 的大規模模型並行語言模型。模型並行會根據 worker 數量成比例地減少顯存使用，這是這三種並行模式中顯存效率最高的。但是其代價是計算效率最低。
- 顯存效率：模型並行的顯存使用量可以根據 worker 數量成比例地減少。至關重要的是，這是減少單個網絡層的激活顯存的唯一方法。DeepSpeed 通過在模型並行 worker 之間劃分激活顯存來進一步提高顯存效率。
- 計算效率：因爲每次前向和反向傳播中都需要額外通信來傳遞激活，模型並行的計算效率很低。模型並行需要高通信帶寬，並且不能很好地擴展到通信帶寬受限的單個節點之外。此外，每個模型並行worker 都會減少每個通信階段之間執行的計算量，從而影響計算效率。模型並行性通常與數據並行性結合使用，以便在內存和計算效率之間進行權衡。
1.2.3 流水線並行
流水線並行將模型的各層劃分爲可以並行處理的階段。當一個階段完成一個 micro-batch 的正向傳播時，激活內存將被髮送給流水線的下一個階段。類似地，當下一階段完成反向傳播時，將通過流水線把梯度反向傳遞回來。爲了確保流水線的各個階段能並行計算，必須同時計算多個 micro-batch 。目前已經有了幾種用於權衡內存和計算效率以及收斂行爲的實現，例如PipeDream。DeepSpeed 通過梯度累積來實現並行，並在相同的總 batch size下，可以達到與傳統數據並行和模型並行訓練收斂情況相同。
- 顯存效率：流水線並行減少的顯存與流水線的階段數成正比，這允許使模型的大小可以隨 worker 的數量線性擴展。但是，流水線並行不會減少每一層的激活顯存佔用量。此外，每個 worker 必須存儲同時運行的各個 micro-batch 的激活值。這導致流水線第一階段的激活內存與單個 mirco batch 的總激活內存大致相同。
- 計算效率：因爲流水線的通信量只和階段邊界的各層的激活值大小成正比，所以流水線並行的通信量最低。但是，它不能無限擴展。像模型並行一樣，增加流水線大小會減少每個流水線階段的計算量，這會降低計算與通信的比率。爲了實現良好的計算效率，流水線並行還要求其每個階段都進行完美的計算負載均衡。
此外，流水線並行性會在每個 batch 的開始和結束時因爲需要重新填充或排空流水線而產生 bubble overhead。使用流水線階段數的 4 倍或 8 倍的梯度累積步驟（以及 batch 大小）進行訓練，相較於只有一個流水線階段分別達到了 81％ 和 90％ 的擴展性。
1.3 通過 3D 並行實現內存和計算效率
數據、模型和管道並行在提高內存和計算效率方面各自發揮特定作用，所以DeepSpeed 結合了這三項強大的技術，可以訓練數萬億規模的模型並擴展到數千個 GPU。這三者的共生並行同時解決了訓練萬億參數模型的兩個基本挑戰：顯存效率和計算效率，讓深度學習訓練的規模遠遠超出了單獨使用每種策略可以企及的高度。因此，DeepSpeed 可以在顯存中放下巨大的模型，而不會犧牲速度。
顯存效率：模型的各層首先被劃分到不同的流水線階段，其次，把每個階段的層通過模型並行進一步進行劃分。這種 2D 組合同時減少了模型、優化器和激活所消耗的內存。然而，我們無法在不引入通信開銷的情況下無限地劃分模型，而通信開銷勢必會限制計算效率。
計算效率：爲了在不犧牲計算效率的情況下將 worker 數量擴展至超出模型和流水線並行能支持的規模，我們使用了 ZeRO 支持的數據並行功能（ZeRO-DP）。ZeRO-DP 不僅可以通過劃分優化器狀態量進一步提高顯存利用效率，而且還可以通過利用基於通信拓撲的映射關係，以最小的通信開銷擴展到任意數量的 GPU。
圖 1：32 個worker的 3D 並行示例。神經網絡的層被分爲四個流水線階段。每個流水線階段內的層在四個模型並行worker之間進一步劃分。最後，每個流水線階段都跨兩個數據並行實例進行復制，ZeRO 在這兩個數據並行副本之間對優化器狀態進行分區。
下圖展示了通信拓撲感知的 3D 映射：通過利用兩個關鍵的架構屬性，我們將 3D 並行中的每個維度仔細地映射到 worker 之上，從而最大化計算效率。
- 優化節點內和節點間的通信帶寬：模型並行具有三種策略中最大的通信開銷，因此我們優先考慮將模型並行 worker 組放置在節點內以利用更大的節點內帶寬。這裏我們基於英偉達 Megatron-LM 進行了張量切分式的模型並行。當模型並行組無法佔滿節點內的所有 worker 時，我們選擇將數據並行組放置在同一個節點內。否則它們將跨節點進行數據並行。流水線並行具有最低的通信量，因此我們可以跨節點進行調度流水線的各個階段，而不受到通信帶寬的限制。
- 通過並行通信增大帶寬：每個數據並行組需要傳遞的梯度量隨着流水線和模型並行的規模而線性減小，因此3D總通信量少於單純使用數據並行的通信量。此外，每個數據並行組會在本地的一小部分 worker 內部獨立進行通信，組間通信可以相互並行。因此，通過減少通信量和增加局部性與並行性，我們可以有效的擴大數據並行通信有效帶寬。
圖 2：圖 1 中的 worker 到八個節點（每個節點有四個 GPU）的系統上的 GPU 的映射。同一顏色的 GPU 在同一節點上。
1.4 3D 並行如何利用每種並行性
一個萬億參數模型可以使用 8 路模型並行、64 路管道並行和 8 路數據並行在 4,096 個 NVIDIA A100 GPU 上進行擴展。
-
通過結合模型並行和流水線並行，3D 並行在多個節點之上實現了出色的內存效率和高效計算效率。模型並行性提高了節點內的激活和模型狀態的存儲效率，流水線並行（相較於僅使用模型並行）則可以在不犧牲計算效率的情況下，跨節點高效存儲模型狀態。
-
通過將模型並行性與流水線並行性相結合，即使在非常小的批量大小下，流水線並行性也能以最小的氣泡開銷實現高計算效率。在 8 路模型並行下，每個模型每使用 1 個微批次，將導致每個 GPU 的有效微批次大小爲1/8。因此，流水線並行可以使用 8 倍流水線並行度的梯度累積步驟來實現 90% 的計算效率，並且每個 GPU 的總累積batch size 僅爲 1。當與數據並行相結合時，這讓 4096 張 GPU 上的總有效 batch size 爲 4096，並依然可以達到 90％ 的流水線效率。
但是數據並行性會帶來怎樣的計算效率呢？數據並行性是否需要每個 GPU 擁有大批量才能保持高效？
模型並行可以將每張GPU上的有效 batch 大小減小到小於 1。這允許流水線並行即使在小 batch 下依然可以有效隱藏流水線 bubble 開銷。請注意，通過使用跨節點流水線並行性，我們就可以讓流水線每個階段的數據並行節點之間的通信獨立發生並與其他流水線階段並行。實際上，在高端 GPU 集羣中常見的全連接的網絡拓撲中，這對於數據並行訓練的可用有效通信帶寬具有重要意義。由於流水線階段中的每個節點都可以與其對應的數據並行節點並行通信，因此有效的通信帶寬與流水線階段數量成正比。設置 64 個流水線並行級之後，有效帶寬將是往返於單個節點的帶寬的 64 倍。憑藉如此大的流水線並行有效帶寬，即使在計算與通信比率非常低的小批量下，數據並行性也能有效擴展。
0x02 引論
我們接下來以 ZeRO: Memory Optimizations Toward Training Trillion Parameter Models 來進行學習。
2.1 原文摘要
大規模深度學習模型可以顯著提高accuracy，但訓練數十億到數萬億的參數是一項挑戰，因爲單個GPU無法容納如此大的存儲模型和狀態量。現有的解決方案比如跨GPU的數據和模型並行存在了很大侷限性：這些方案雖然獲得了計算、通信和開發效率，但都是在各種因素之間權衡，而且有一個最基本的問題：模型只能位於有限的設備內存中。
論文作者開發了一種新的解決方案，使用零冗餘優化器（Zero）來優化內存，這樣可以極大地提高了訓練速度，同時增加了可以有效訓練的模型大小。ZeRO在尋求數據並行和模型並行的一個適當中間點，其希望消除數據和模型並行訓練中的內存冗餘，同時保持了較低的通信量和較高的計算粒度，使我們能夠以持續的高效率按設備數量比例調整模型大小。因此，ZeRO 可以獲得數據並行性和模型並行性的優點。可以用同樣內存來來運行更大的模型，可以使用數據並行方式來訓練那些以前只能使用模型並行進行訓練的模型。
2.2 原文引論
常見的數據並行性（DP）並不會減少每個設備的內存，而其他現有解決方案，如管道並行（PP）、模型並行（MP）、CPU卸載（CPU-Offloading）等這些都是在功能性、可用性以及內存和計算/通信效率之間進行權衡。在訓練大型模型的各種現有解決方案中，MP可能是最有前途的，然而，MP無法擴展到更大的尺寸。MP 會把模型進行垂直分割，將每一層中的計算和參數劃分到多個設備上，這需要在每一層之間進行大量通信。因此，它們在GPU間通信帶寬較高的單個節點內工作良好，但在單個節點之外，效率會迅速下降，因此無法有效的擴展到單個節點之外。
那麼，我們如何克服現有解決方案的侷限性，更有效地訓練大型模型呢？爲了回答這個問題，我們首先分析了現有系統在模型訓練方面的全部內存消耗，並將其分爲兩部分：
- 對於大型模型來說，大部分內存被模型狀態佔用，其中包括優化器狀態（如Adam中的動量和方差）、梯度和參數。
- 剩餘的內存被：激活、臨時緩衝區和不可用的碎片佔據，我們統稱爲剩餘狀態。
因此，我們開發了Zero Redundancy Optimizer，在獲得高計算和通信效率的同時，可以優化這兩個方面的內存效率。
2.2.1 優化模型狀態
模型狀態通常在訓練過程中消耗最大的內存量，但是現有的方法，如DP和MP並不能提供令人滿意的解決方案。
DP具有良好的計算/通信效率，但內存效率較差，而MP的計算/通信效率較差。更具體地說，DP在所有數據並行進程中複製整個模型狀態，導致冗餘內存消耗；雖然MP對這些狀態進行分區以獲得較高的內存效率，但往往會導致過於細粒度的計算和昂貴的通信，從而降低了擴展效率。此外，這些方法靜態地維護整個訓練過程中所需的所有模型狀態，即使在訓練過程中並非始終需要所有模型狀態。
基於這些觀察結果，我們開發了ZeRO-DP，ZeRO-DP在保持DP的計算/通信效率基礎之上，同時實現了MP的內存效率。ZeRO-DP通過對模型狀態進行分區而不是複製來消除數據並行進程中的內存狀態冗餘，這樣每個GPU之上的內存消耗將會和數據並行度成反比，並通過在訓練期間使用動態通信調度來保留同DP基本一致的計算粒度和通信量，這樣可以保持計算/通信效率。
在模型訓練期間，大部分內存被以下三種情況之一消耗：
- 激活。
- OGP狀態，即由優化器狀態（O），參數梯度（G）和參數本身（P）組成的張量。
- 臨時緩衝區。
可能有人會問爲什麼不考慮輸入數據的內存消耗，其實，輸入數據所佔用的顯存其實並不大，這是因爲用戶基本使用迭代器讀取數據，這意味着數據並不是一次性全部讀入顯存，因此每次輸入所佔用的顯存與整個網絡參數相比其實是微不足道的。
ZeRO DP有三個主要的優化階段（如下圖1所示），它們對應於優化器狀態、梯度和參數的分區。當逐步啓用時：
1）優化器狀態分區（Pos）：內存減少4倍，通信量與DP相同，此階段也被稱爲 ZeRO-OS。
2）添加梯度分區（Pos+g）優化：內存減少8倍，通信量與DP相同；
3）添加參數分區（Pos+g+p）優化：內存減少與DP的並行度成線性關係。模型內存被平均分配到每個GPU之上，每個gpu上的內存消耗與數據並行度成反比，但是通信量只是適度增加。例如，跨64個GPU（Nd=64）拆分將產生64倍的內存縮減。通信量適度增加了50%。
內存消耗具體可以參見下圖：
圖1：ZeRO-DP優化的三個階段之中每個設備內存消耗比較。ψ表示模型大小（參數數量），K表示優化器狀態的內存乘數，Nd表示DP並行度，即Nd個GPU。在本例中，我們假設基於Adam優化器的混合精度訓練，模型大小爲ψ=7.5B，DP爲Nd=64，K=12。
2.2.2 優化殘餘狀態內存
在使用 ZeRO-DP 優化模型狀態對應的內存之後，殘餘內存（Residual State Memory）成爲次要內存瓶頸，剩餘內存包括：激活、臨時緩衝區和不可用內存片段。我們開發了ZeRO-R來分別優化這三個因素所消耗的剩餘內存。
-
對於激活（從前向傳播結果之中存儲，用來支持後向傳播），我們注意到優化檢查點會有幫助，但是對於大型模型不夠用。因此，ZeRO-R通過在現有MP方案中識別和刪除激活副本來優化激活內存。它還可以在適當的時候將激活卸載到CPU。
-
ZeRO-R爲臨時緩衝區定義了適當的大小，以實現內存和計算效率的平衡。
-
我們觀察到在訓練中，由於不同張量生命週期的變化而會導致一些內存碎片。由於這些碎片的存在，會導致即便即使有足夠的可用內存，也會因爲缺少連續內存而使得內存分配失敗。ZeRO-R根據張量的不同生命週期來主動管理內存，防止內存碎片。
ZeRO-DP和ZeRO-R結合在一起形成了一個強大的DL訓練內存優化系統，我們統稱爲ZeRO。
2.2.3 ZeRO和MP
因爲ZeRO消除了DP中的內存效率不足，所以很自然地會問：我們還需要MP嗎？什麼時候需要？ZeRO如何與MP一起工作？
使用ZeRO之後，MP對於大型模型就不太有吸引力了。ZeRO-DP在減少每設備內存佔用方面至少與MP一樣有效，或者在MP無法均勻劃分模型時更有效。它還具有相當或更好的縮放效率。此外，數據並行非常容易使用，因此它廣泛適用於不同的工作負載，而如今的MP方法通常需要模型開發人員的一些額外工作來修改其模型，比如現有的工作（如Megatron-LM）只支持有限的操作和模型集。
儘管如此，仍然有一些情況下我們希望利用MP:
i）當與ZeRO-R一起使用時，MP可以減少大型模型的激活內存佔用。
ii）對於激活內存不是問題的較小模型。當單獨使用DP時，可能會因爲聚合batch size太大而無法實現良好的收斂性，這時候MP也有好處。在這種情況下，可以將ZeRO與MP結合起來，使模型具有可接受的聚合batch size。
0x03 相關工作
3.1 數據，模型和流水線並行
並行化是大型模型訓練的關鍵策略。對於可以塞進設備內存的模型，數據並行（DP）用於將訓練擴展到多個設備。在DP中，模型參數複製到每個設備上。在每個步驟中，一個小批量被均勻地分發到所有數據並行進程中，這樣每個進程都會對不同的數據樣本子集執行正向和反向傳播，並使用進程間的平均梯度來局部更新模型。
當一個模型不適合設備內存時，模型並行性（MP）和流水線並行性（PP）分別以垂直和水平方式在進程之間分割模型。
PP在層之間水平拆分一個模型，在不同設備上運行不同的分區，並使用微批處理隱藏管道氣泡。由於水平拆分和micro-batching，所以某些功能（如tied-weight和batch-normalization）難以實現。
流行的PP實現（如G-pipe）同時對模型參數和總激活進行分區，但需要與管道分區數量成比例的batch size來隱藏管道氣泡。大batch size可能會影響收斂速度，PP同時也需要大量內存來存儲激活。
PipeDream是PP的另一種實現，其保留了過時參數的多個副本，以隱藏管道氣泡，而不會顯著增加batch size，從而可以降低內存效率。此外，該實現不等同於標準DL訓練，並且對訓練收斂有影響。
相比之下，ZeRO獲得了與PP相同或更好的內存效率，而不會有PP帶來的功能、性能和與收斂的限制。
3.2 非並行方面的工作
原小標題爲Non-parallelism based approach to reduce memory。
除了MP和PP之外，還有很多旨在減少DL訓練內存開銷的工作。
3.2.1 減少激活內存
目前，有很多工作集中在減少激活的內存佔用上，包括壓縮、激活檢查點或實時分析。這些努力是互補的，可以與ZeRO一起工作。事實上，ZeRO-R中的激活內存減少完全可以和激活檢查點並行工作。
3.2.2 CPU Offload
也有一些工作利用計算節點的異構性，分別通過算法設計或虛擬化內存將模型狀態轉移到CPU內存。但是這導致有50%的時間被浪費在GPU-CPU-GPU傳輸。ZeRO的不同之處在於，它顯著降低了內存消耗，而無需將模型狀態存儲到CPU內存中。在極少數情況下，ZeRO-R可能只針對非常大的模型才卸載激活檢查點，以提高性能。
3.2.3 內存高效(Efficient)優化器
另一些工作是通過獲取模型參數和梯度的粗粒度統計數據來減少自適應優化方法的內存消耗，這可能會對模型收斂保證產生影響。ZeRO與這些工作是正交的，它的優化不會改變模型優化方法或影響模型收斂，但會有效地減少每個設備的優化器狀態和梯度的內存佔用。
3.3 訓練優化器
對於大型模型，自適應優化（Adaptive）方法對於達到SOTA性能和精度至關重要。與SGD相比，它以顯著的內存佔用爲代價，維護每個模型參數和梯度的細粒度一階和二階統計信息。ZeRO可以將這些優化器的內存佔用減少幾個數量級，使這些複雜的優化方法對於在具有適度設備內存的硬件上訓練大型模型非常實用。它還讓人們可以開發和使用更復雜、內存消耗更大、收斂性更好的優化器。
0x04 模型內存都去哪裏了?
讓我們退一步來研究一下當前訓練系統的內存消耗。例如，一個1.5B參數的GPT-2模型需要3GB內存用於16位精度的權重（或參數），但是人們卻不能使用Tensorflow或PyTorch在一個32GB內存的GPU上進行訓練。人們可能想知道所有的內存都去了哪裏。在模型訓練期間，大部分內存被模型狀態消耗，即由optimizer狀態、梯度和參數組成的張量。除了這些模型狀態，其餘的內存被激活、臨時緩衝區和碎片化內存消耗，我們稱之爲剩餘狀態。我們將從這兩個方面詳細研究內存消耗。
4.1模型狀態：優化器狀態，梯度和參數
原小標題爲：Model States: Optimizer States, Gradients and Parameters
大多數設備內存在訓練期間由模型狀態消耗。例如，用Adam，DL訓練中最流行的優化器之一作爲例子。Adam需要存儲兩個優化器狀態，i）時間平均動量（time averaged momentum）和ii）梯度方差（variance of the gradients）來計算更新。因此，要使用ADAM訓練模型，必須有足夠的內存來保存梯度動量和方差的副本。此外，也需要有足夠的內存來存儲梯度和權重本身。在這三種類型的參數相關張量中，優化器狀態通常消耗最多的內存，特別是在應用混合精度訓練時。
4.1.1 混合精度訓練
在當前一代NVIDIA GPU上訓練大型模型的最先進方法是通過混合精度（fp16/32）訓練，在這個方法中，參數和激活存儲爲fp16，從而能夠在這些GPU上使用高吞吐的張量核心單元。在混合精度訓練期間，優化器使用fp16權重和激活執行正向和反向傳播。但是，爲了在反向傳播結束時有效地計算和應用權重更新，混合精度優化器必須保留參數的fp32副本以及所有其他優化器狀態的fp32副本。
讓我們以Adam優化器爲例。使用Adam對帶有ψ個參數的模型進行混合精度訓練需要足夠的內存來保存參數和梯度的fp16副本。其內存需求分別爲2ψ和2ψ字節。此外，它還需要保存優化器狀態，參數動量和方差的fp32副本，其內存需求分別爲4ψ，4ψ和4ψ字節。
讓我們使用K來表示優化器狀態的內存乘數（multiplier），即存儲它們所需的額外內存是Kψ字節。混合精度Adam的K=12。總的來說，這將產生2ψ+2ψ+Kψ=16ψ字節的內存需求。對於具有15億個參數的GPT-2這樣的模型，這至少需要24GB的內存，遠遠高於單獨保存fp16參數所需的3GB內存。
4.2 剩餘內存佔用
原標題爲 Residual Memory Consumption
4.2.1 激活
在訓練期間，激活會佔用大量的內存。作爲一個具體的例子，1.5B參數的GPT-2模型以1K的序列長度和32的batch size 進行訓練，需要大約60GB的內存。激活檢查點（或激活重新計算）是一種常用的方法，可將激活內存減少到總激活的平方根，但需花費33%的重新計算開銷。這將使此模型的激活內存消耗減少到約 8 GB。
儘管有顯著的減少，但對於更大的模型，即便使用激活檢查點，激活內存也會變得相當大。例如，一個具有1000億個參數的類GPT模型，對於32大小的batch size，則需要大約60 GB的內存，即使使用激活檢查點也是如此。
4.2.2 臨時緩衝區
對於大型模型，用於存儲中間結果的臨時緩衝區會消耗大量內存。有些操作，比如gradient all-reduce或者 gradient norm computation 會傾向於將所有梯度融合到單個平坦緩衝區中，以此來運行一個統一操作，這樣可以提高吞吐量。例如，所有設備的帶寬都會隨着消息的增大而降低。雖然梯度本身通常存儲爲fp16張量，但融合緩衝區可以是fp32張量（具體取決於操作類型）。當模型較大時，這些臨時緩衝區大小是非常重要的。例如，對於參數爲1.5B的模型，扁平fp32緩衝區需要6GB內存。
4.2.3 內存碎片
到目前爲止，我們已經討論了訓練期間的實際內存消耗。此外，即使有足夠的可用內存，也可能耗盡可用內存。內存碎片就可能導致這種情況。如果沒有足夠的連續內存來滿足對內存的請求，即使總可用內存大於請求的內存，對內存的請求也會失敗。在訓練非常大的模型時，我們觀察到明顯的內存碎片，這會導致內存不足問題，在某些極端情況下，即使超過30%的內存仍然可用，依然無法分配內存。
0x05 ZeRO: 感悟和概述
ZeRO有兩組優化：i）ZeRO DP旨在減少模型狀態的內存佔用，ii）ZeRO-R旨在減少剩餘內存消耗。
5.1 感悟和概述: ZeRO-DP
ZeRO powered DP 基於三個關鍵感悟：
-
DP比MP具有更好的擴展效率，因爲MP降低了計算粒度，同時也增加了通信開銷。超過某一點後，較低的計算粒度會降低每個GPU的效率，而增加的通信開銷會隱藏跨GPU的可伸縮性，特別是在跨越節點邊界時。相反，DP具有更高的計算粒度和更低的通信量，從而帶來更高的效率。
-
DP內存效率低下，因爲模型狀態在所有數據並行進程中冗餘存儲。相反，MP對模型狀態進行分區以獲得內存效率。
-
DP和MP都保留了整個訓練過程中所需的模型所有狀態，但並非所有情況都需要。例如，每個層的參數僅在層的正向傳播和反向傳播期間需要。
基於這些感悟，ZeRO DP 在保留DP的訓練效率的同時，也實現了MP的內存效率。ZeRO DP對模型狀態進行分區，而不是複製它們，並使用動態通信計劃，該計劃利用模型狀態固有的時間特性，同時最小化通信量。通過這樣做，ZeRO-DP隨着DP度的增加線性地減少了模型的每設備內存佔用，同時保持通信量接近默認DP，這樣就保持了效率。
5.2 感悟和概述: ZeRO-R
5.2.1 降低激活內存
兩個關鍵感悟是：
-
MP對模型狀態進行分區，但通常需要複製激活內存。例如，如果我們垂直分割一個線性層的參數並跨兩個GPU並行計算它們，那麼每個GPU都需要整個激活來計算其分區。
-
對於GPT-2或更大的模型，算術強度（每次迭代的計算量與每次迭代的激活檢查點量之比）非常大(≥ 10K），並隨着隱藏維度增加而線性增加，從而可以隱藏激活檢查點的數據移動成本，即使在帶寬較低的情況下也是如此。
ZeRO通過跨GPU劃分激活檢查點來消除MP中的內存冗餘，並使用allgather按需重建它們。激活內存的減少與MP程度成比例。對於非常大的模型，ZeRO甚至可以選擇將激活分區移動到CPU內存中，同時由於這些模型中的運算強度很大，因此仍然可以實現良好的效率。
5.2.2 管理臨時緩衝區
ZeRO-R使用恆定大小的緩衝區來避免臨時緩衝區隨着模型大小的增加而崩潰，同時使它們足夠大以保持效率。
5.2.3 管理內存碎片
內存碎片是短生命週期內存對象和長生命週期內存對象交錯分配的結果。在正向傳播期間，激活檢查點的壽命很長，但重新計算的激活壽命很短。同樣，在反向計算中，激活梯度的壽命很短，而參數梯度的壽命很長。基於這一認識，ZeRO通過將激活檢查點和梯度移動到預先分配的連續內存緩衝區來執行動態內存碎片整理。這不僅提高了內存可用性，還通過減少內存分配器查找可用連續內存所需的時間來提高效率。
0x06 深入瞭解 ZeRO-DP
雖然現有的DP方法在每個設備上覆制模型狀態並引入顯著的內存開銷，但ZeRO DP通過跨數據並行的進程對它們（優化器狀態、梯度和參數）進行分區來消除這種內存冗餘。圖1量化並可視化了有無ZeRO-DP的內存需求。該圖顯示了（1）優化器狀態（2）梯度和（3）參數累積冗餘 這三種參數在分區後的內存佔用。我們將其稱爲ZeRO DP的三個優化階段：Pos、Pg和Pp，我們將在下面詳細說明。這裏把圖一再次貼出來。
6.1 Pos : 優化器狀態分區
對於一個\(N_d\)並行度的DP來說，我們將優化器狀態分組到\(N_d\)個相等的分區中，這樣第i個數據並行進程只更新與第i個分區對應的優化器狀態。因此，每個數據並行過程只需要存儲和更新總優化器狀態 的$ \frac{1}{N_d}\(，然後只更新\) \frac{1}{N_d}$個參數。在每個訓練步驟結束時，我們會執行一個跨數據並行進程的all-gather操作，以獲得跨所有數據並行進程的完全更新的參數。
如圖1所示的具體示例，7.5 B參數模型，使用64路DP（\(N_d\)=64），其Pos需要31.4GB內存。而使用標準DP則需要120 GB內存。此外，當\(N_d\)較大時，模型狀態的內存需求從4ψ+12ψ=16ψ字節減少到4ψ+\(\frac{12ψ}{N_d}\)字節≈ 4ψ字節，導致4x倍數的減少。
6.2 Pg: 梯度分區
由於每個數據並行進程只負責更新其相應的參數分區，因此，每個節點僅僅對自己負責的那部分參數的梯度進行規約。在歸併之後，每個節點只需要自己參數分區對應的梯度，對於其他的梯度不再需要，所以它們的內存可以被釋放。這將梯度的內存佔用從2ψ字節縮減到 \(\frac{2ψ}{N_d}\)。
實際上，這是一種 Reduce-Scatter操作，不同參數的梯度被減少到不同的進程之中。爲了提高效率，我們使用了bucketization策略，其中我們將對應於特定分區的所有梯度bucketization，並立即對整個bucket執行規約。在我們的例子中，我們在分區邊界執行一個reduce而不是 all-reduce，以減少內存佔用，並重疊計算和通信。
內存節省：通過消除梯度和優化器狀態冗餘，我們將內存佔用進一步降低到2ψ+ \(\frac{14ψ}{N_d}\)≈ 2Ψ. 如圖1中的示例所示，7.5 B參數模型使用Pos+g和64路DP（Nd=64）時只需要16.6 GB內存，而使用標準DP時需要120 GB內存。當\(N_d\)較大時，模型狀態的內存需求從2ψ+14ψ=16ψ字節減少到2ψ+\(\frac{14ψ}{N_d}\)字節≈ 2ψ字節，減少8倍。
6.3 Pp: 參數分區
就像優化器狀態和梯度一樣，每個進程只存儲與其分區對應的參數。當正向和反向傳播需要其分區外的參數時，會通過broadcast操作從適當的數據並行進程接收這些參數。雖然乍一看，這可能會導致顯著的通信開銷，但我們發現，這種方法只會將基線DP系統的總通信量增加到1.5倍，同時實現與Nd成比例的內存減少。
內存節省：通過參數分區，我們將ψ個參數的內存佔用從16ψ降低到 \(\frac{16ψ}{N_d}\)。 如圖1中的示例所示，7.5 B參數模型使用\(P_{os+g+p}\)和64路DP（Nd=64）時只需要1.9 GB內存，而使用標準DP時需要120 GB內存。
這有着深刻的含義：只要有足夠數量的設備來共享模型狀態，ZeRO-DP就可以適合任意大小的模型。
6.4 對模型大小的影響
分區Pos、Pos+g和Pos+g+p的三個階段分別將模型狀態下每個數據並行進程的內存消耗減少了4倍、8倍和\(N_d\)倍。表1分析了幾個示例模型在不同DP程度下，ZeRO-DP 3個階段下的模型狀態內存消耗。
如果不使用ZeRO，則無論DP程度如何，內存消耗都等於表中的第一行。注意，當Nd=64時，ZeRO可以分別使用Pos、Pos+g和Pos+g+p來訓練參數高達7.5B、14B和128B的模型。當Nd=1024時，啓用所有優化的ZeRO（Pos+g+p）可以訓練具有1萬億個參數的模型！或者可能是任意大小的模型！如果沒有ZeRO，DP可以運行的最大模型的參數纔不到15億個。
0x07 深入 ZeRO-R
7.1 \(P_a\): 將 Activation Checkpointing 分區
正如前面所討論，MP 在設計上就要求複製激活，從而在模型並行GPU之間產生激活的冗餘副本。ZeRO通過對激活進行分區來消除這種冗餘，並且在激活用於計算之前，才只以一個激活層的副本形式將它們一次性具化。
更具體地說，一旦計算了模型中一個層的前向傳播，輸入激活將在所有模型並行過程中進行分區，直到在反向傳播中再次需要它。此時，ZeRO使用all gather操作重新具化激活的複製副本。我們將此優化稱爲Pa。它與激活檢查點一起工作，只存儲分區的激活檢查點，而不是複製副本。此外，在非常大的模型和非常有限的設備內存的情況下，這些分區的激活檢查點也可以卸載到CPU上，以額外的通信成本將激活內存開銷降低到幾乎爲零，我們稱之爲\(P_{a+cpu}\)。
通過分區激活檢查點，ZeRO將激活佔用空間減少了一個與MP程度成比例的因子。考慮訓練一個100B模型，其批大小爲32，序列長度爲1024，MP的度數爲16。如果我們爲每個轉換器層檢查一個激活，那麼僅存儲激活檢查點就需要每個GPU大約33 GB的內存。但如果Pa爲零，則每個GPU的容量可以減少到2GB左右。此外，這個2GB可以卸載到CPU上，從而將激活的內存佔用減少到幾乎爲零。
7.2 CB: 固定大小緩衝區
ZeRO仔細選擇臨時數據緩衝區的大小，以平衡內存和計算效率。在訓練期間，某些操作的計算效率可能高度依賴於輸入大小，輸入越大，效率越高。例如，一個大的all-reduce操作比一個小的操作獲得更高的帶寬。因此，爲了獲得更好的效率，高性能庫（如NVIDIA Apex或Megatron）在應用這些操作之前將所有參數融合到單個緩衝區中。然而，融合緩衝器的內存開銷與模型大小成正比。例如，對於3B參數模型，32位融合緩衝區將需要12GB內存。爲了解決這個問題，當模型變得太大時，我們只需使用一個高性能的固定大小融合緩衝區（Constant Size Buffers）。通過這樣做，緩衝區大小不依賴於模型大小，並且通過保持足夠大的緩衝區大小，我們仍然可以實現良好的效率。
7.3 MD: 內存碎片整理
模型訓練中的內存碎片是激活檢查點和梯度計算的結果。在帶有激活檢查點的前向傳播期間，只有選定的激活被存儲用於反向傳播，而大多數激活被丟棄，因爲它們可以在反向傳播期間重新計算。這將創建短期內存（丟棄的激活）和長期內存（檢查點的激活）的交錯，導致內存碎片。類似地，在反向傳播期間，參數梯度是長生命週期的，而激活梯度和計算參數梯度所需的任何其他緩衝區是短生命週期的。同樣，這種短期內存和長期內存的交錯也會導致記憶碎片。
當有足夠的內存可用時，有限內存碎片通常不是問題，但對於使用有限內存運行的大型模型訓練，內存碎片會導致兩個問題，i）由於缺乏連續內存，即使有足夠的可用內存也會導致OOM，ii）由於內存分配器花費大量時間搜索連續內存塊以滿足內存請求，導致效率低下。
ZeRO通過爲激活檢查點和漸變預先分配連續內存塊，並在生成時將它們複製到預先分配的內存中，動態地進行內存碎片整理。MD不僅使ZeRO能夠以更大的批量訓練更大的模型，而且還可以在內存有限的情況下提高訓練效率。
0x08 ZeRO-DP 通信量分析
由於ZeRO通過消除內存冗餘提高了可以訓練模型的大小，所以很自然地會有疑問，是否在用通信量換取內存效率的問題。換句話說，與基線DP方法相比，ZeRO-powered到DP方法的通信量是多少？答案分爲兩個部分：i）ZeRO-DP使用Pos和Pg的時候並不會產生額外的通信，而可實現8倍的內存縮減；ii）ZeRO-DP在使用Pos和Pg之外的Pp時，最多會產生1.5倍的通信，但同時進一步將內存佔用減少了Nd倍。
8.1 數據並行通信量
數據並行訓練期間，在反向傳播結束，而在計算下一步的更新之前，會對所有數據並行進程的梯度進行平均。平均操作是使用all-reduce來完成的。對於大型模型，all- reduce通信完全受通信帶寬的限制，因此，我們的分析僅限於每個數據並行進程之間發送和發送的總通信量。
all-reduce的最新實現一般採用兩步方法，第一步是reduce-scatter操作，它規約了不同進程上數據的不同部分。下一步是all gather操作，其中每個進程收集所有進程上規約的數據。這兩個步驟的結果就是一個all-reduce操作。“reduce-scatter”和“all-gather”都是使用流水線方法實現的，這會導致總共ψ個元素（假設數據包含ψ個元素）的數據移動。因此，標準DP在每個訓練步驟中產生2ψ個數據移動。
8.2 ZeRO-DP 通信量
8.2.1 使用 Pos+g 的通信量
使用梯度分區之後，每個進程只存儲梯度的一部分，這是更新其相應的參數分區所必需的。因此，與all-reduce不同，ZeRO只需要在梯度上進行scatter-reduce操作，從而產生ψ的通信量。在每個進程更新其負責的參數的分區後，將執行all-gather以從所有數據並行進程收集所有更新的參數。這也導致通信量爲ψ。因此，每個訓練步驟的總通信量爲ψ+ψ=2ψ，與基線DP完全相同。
8.2.2 使用 Pos+g+p 的通信量
在參數分區之後，每個數據並行進程只存儲它負責更新的參數。因此，在前向傳播期間，它需要接收所有其他分區的參數。但是，這可以通過流水線來避免內存開銷。在計算與特定分區對應的模型部分的前向傳播之前，負責該分區的數據並行進程可以將權重廣播給所有數據並行進程。一旦該分區的前向傳播完成，就可以丟棄這些參數。因此，總通信量爲\(\frac{ψ×N_d}{N_d}=ψ\)。換言之，我們依靠在整個正向傳播中傳播來重新安排參數的all-gather，並在使用參數後丟棄這些參數。但是請注意，對於反向傳播，需要再次進行此all-gather（但是以相反的順序）。
因此，總通信量是reduce-scatter和 all-gather所產生的通信量的總和，總體積爲3ψ，是基線的1.5倍。梯度和參數分區都利用了這樣一種洞察，即並非所有的梯度和參數狀態都是始終需要的，而是通過明智地傳遞狀態來優化內存。
0x09 ZeRO-R 通信分析
我們將ZeRO-R中的分區激活檢查點（Pa）的通信量與基線MP進行了比較，結果表明，Pa引起的通信量增加通常不到基線MP的十分之一。此外，我們分析了Pa的通信開銷與DP通信量的關係，以確定Pa通過允許更大的批量和減少DP通信來提高效率的情況。我們利用這種分析來決定是否以及何時應用Pa以及Pa+cpu。
分區激活檢查點的通信量權衡取決於模型大小、檢查點策略和MP策略。爲了分享具體的見解，我們在使用SOTA MP方法（Megatron-LM）實現的模型背景下進行分析。
在帶有激活檢查點的Megatron-LM中，每個transformer在正向傳播中執行兩個大小爲batch×seq_length×hidden _dim 的all-reduce操作，兩個all-reduce操作用於正向重新計算，另外兩個all-reduce操作用於反向傳播。每個塊的總通信量爲 12 × seq length × hidden dim，因爲all reduce的通信量爲2 × message size。
當ZeRO-R對激活檢查點進行分區時，需要在每個激活檢查點上向前重新計算反向傳播之前執行額外的all gather操作。通常，我們檢查每個transformer塊的輸入激活，每個轉換器塊需要一個all gather。因此，通信開銷Pa爲seq length ∗ hidden dim，因爲所有聚集的通信量爲message size。因此，Pa的總通信開銷小於模型並行原始通信量的10%。
當MP與DP結合使用時，Pa可用於將數據並行通信量減少一個數量級，而模型並行通信量增加10%，並在數據並行通信成爲性能瓶頸時顯著提高效率。請注意，Pa將激活內存消耗降低了MP並行度，從而允許按比例增加批處理大小。對於大型模型，MP可以大到16個（DGX-2節點上的#GPU），允許批量大小最多增加16倍。數據並行訓練的通信量與批量大小成反比。因此，由於Pa導致批量大小增加一個數量級可能導致數據並行通信量減少一個數量級。
最後，如果採用Pa+cpu，分區激活檢查點將卸載到cpu，激活內存需求將減少到幾乎爲零，與Pa相比，cpu內存之間增加了2倍的數據移動。在極端情況下，DP通信量是主要瓶頸，因爲即使使用Pa，批大小也很小，只要cpu數據傳輸開銷小於DP通信量開銷，Pa+cpu就可以通過增加批處理大小來提高效率，這通常適用於小批處理大小。
在給定模型和硬件特性的情況下，我們利用上述分析來決定是否以及何時應用Pa和Pa+cpu。
0xFF 參考
論文解讀系列第十三篇：ZeRO——面向萬億級參數的模型訓練方法