本文主要介紹用於估算 transformer 類模型計算量需求和記憶體需求的相關數學方法。
引言
其實，很多有關 transformer 語言模型的一些基本且重要的資訊都可以用很簡單的方法估算出來。不幸的是，這些公式在 NLP 社群中鮮為人知。本文的目的是總結這些公式，闡明它們是如何推匯出來的及其作用。
注意: 本文主要關注訓練成本，該成本主要由 GPU 的 VRAM 主導。如果你想知道有關推理成本 (通常由推理延遲主導) 的資訊，可以讀讀 Kipply 寫的 這篇精彩博文。
算力要求
下式可用於估算訓練一個 transformer 模型所需的算力成本:
這裡:
- \(C\) 是訓練 transformer 模型所需的計算量，單位為總浮點運算數 (FLOP)
- \(C=C_{\text {前向}}+C_{\text {後向}}\)
- \(C_{\text {前向}} \approx 2PD\)
- \(C_{\text {後向}} \approx 4PD\)
- \(\tau\) 是訓練叢集的實際總吞吐量: \(\tau=\text {GPU 數} \times \text {每 GPU 的實際每秒浮點運算數 (實際 FLOPs) }\)，單位為 FLOPs
- \(T\) 是訓練模型所花費的時間，以秒為單位
- \(P\) 是 transformer 模型的引數量
- \(D\) 是資料集大小，表示為資料集的總詞元數
該式由 OpenAI 的縮放定律論文 和 DeepMind 的縮放定律論文 提出並經其實驗驗證。想要獲取更多資訊，可參閱這兩篇論文。
下面，我們討論一下 \(C\) 的單位。 \(C\) 是總計算量的度量，我們可以用許多單位來度量它，例如:
FLOP - 秒，單位為 \({\text {每秒浮點運算數}} \times \text {秒}\)GPU - 時，單位為 \(\text {GPU 數}\times\text {小時}\)- 縮放定律論文傾向於以
PetaFLOP - 天為單位，即單位為 \(10^{15} \times 24 \times 3600\)
這裡需要注意 \(\text {實際 FLOPs}\) 的概念。雖然 GPU 的規格書上通常只宣傳其理論 FLOPs，但在實踐中我們從未達到過這些理論值 (尤其在分散式訓練時！)。我們將在計算成本這一小節列出分散式訓練中常見的 \(\text {實際 FLOPs}\) 值。
請注意，上面的算力成本公式來自於 這篇關於 LLM 訓練成本的精彩博文。
引數量與資料集的權衡
嚴格來講，你可以隨心所欲地使用任意數量的詞元來訓練 transformer 模型，但由於參與訓練的詞元數會極大地影響計算成本和最終模型的效能，因此需要小心權衡。
我們從最關鍵的部分開始談起: “計算最優” 語言模型。“Chinchilla 縮放定律”，得名於提出 “計算最優” 語言模型論文中所訓練的模型名，指出計算最優語言模型的 引數量 和 資料集大小 的近似關係滿足: \(D=20P\)。該關係成立基於一個前提條件: 使用 1,000 個 GPU 1 小時和使用 1 個 GPU 1,000 小時成本相同。如果你的情況滿足該條件，你應該使用上述公式去訓練一個效能最優且 GPU - 時 成本最小的模型。
但 我們不建議在少於 200B 詞元的資料集上訓練 LLM。 雖然對於許多模型尺寸來說這是 “Chinchilla 最優” 的，但生成的模型通常比較差。對於幾乎所有應用而言，我們建議確定你可接受的推理成本，並訓練一個滿足該推理成本要求的最大模型。
計算成本的經驗值
Transformer 模型的計算成本通常以 GPU - 時 或 FLOP - 秒 為單位。
- GPT-NeoX 的
實際 TFLOP/s在正常注意力機制下達到 150 TFLOP/s/A100，在 Flash 注意力機制下達到 180 FLOP/s/A100。這與其他高度最佳化的大規模計算庫一致，例如 Megatron-DS 的值是在 137 和 163 TFLOP/s/A100 之間。 - 一個通用的經驗法則是
實際 TFLOP/s可至 120 TFLOP/s/A100 左右。如果你得到低於 115 TFLOP/s/A100 的值，可能是你的模型或硬體配置有問題。 - 藉助 InfiniBand 等高速互連裝置，你可以在資料並行維度上實現線性或亞線性擴充套件 (即增加資料並行度應該以近乎線性的方式增加整體吞吐量)。下圖顯示了在橡樹嶺國家實驗室 (Oak Ridge National Lab) 的 Summit 超級計算機上測試出的 GPT-NeoX 庫的擴充套件性。請注意，這張圖用的是 V100，而本文中的其他大多數例子都是基於 A100 的。
記憶體需求
Transformer 模型通常由其 引數尺寸 來描述。但是，根據給定一組計算資源確定要訓練哪些模型時，你需要知道 該模型將佔用多少空間 (以位元組為單位)。這不僅需要考慮你的本地 GPU 可以推理多大的模型，還需要考慮給定訓練叢集中的總可用 GPU 記憶體可供訓練多大的模型。
推理
模型權重
大多數 transformer 模型都使用 混合精度進行訓練，可以是 fp16 + fp32 或是 bf16 + fp32。混合精度降低了模型訓練和推理所需的記憶體量。推理時，我們還可以將語言模型從 fp32 轉換為 fp16 甚至 int8，而沒有實質性的精度損失。下面我們看看在不同的資料型別下，模型所需記憶體有什麼不同 (以位元組為單位):
- 對 int8 而言，\(\text {模型記憶體}=1 \text { 位元組} /\text {引數}\cdot \text {引數量}\)
- 對 fp16 和 bf16 而言，\(\text {模型記憶體}=2 \text { 位元組} /\text {引數} \cdot \text {引數量}\)
- 對 fp32 而言，\(\text {模型記憶體}=4 \text { 位元組} /\text {引數}\cdot \text {引數量}\)
推理總記憶體
除了儲存模型權重所需的記憶體外，實際中前向傳播過程中還會有少量額外開銷。根據我們的經驗，此開銷在 20% 以內，該比例通常與模型無關。
總的來說，回答 “這個模型是否適合推理” 這一問題，可以用下述公式來獲得不錯的估計:
\(\text {推理總記憶體} \approx 1.2 \times \text {模型記憶體}\)
本文不會深究該開銷的來源，留待後面的文章來闡述。在本文的後續部分，我們將主要關注模型訓練的記憶體。如果你有興趣瞭解更多有關推理計算需求的資訊，請檢視 這篇深入介紹推理的精彩博文。現在，我們要開始訓練了！
訓練
除了模型權重之外，訓練還需要在裝置記憶體中儲存最佳化器狀態和梯度。這就是為什麼當你問 “我需要多少記憶體來執行模型？”，別人會立即回答 “這取決於是訓練還是推理”。訓練總是比推理需要更多的記憶體，通常多得多！
模型權重
首先，可以使用純 fp32 或純 fp16 訓練模型:
- 純 fp32，\(\text {模型記憶體}=4 \text { 位元組} /\text {引數} \cdot \text {引數量}\)
- 純 fp16，\(\text {模型記憶體}=2 \text { 位元組} /\text {引數} \cdot \text {引數量}\)
除了推理中討論的常見模型權重資料型別外，訓練還引入了 混合精度 訓練，例如 AMP。該技術尋求在保持收斂性的同時最大化 GPU 張量核的吞吐量。現代 DL 訓練領域經常使用混合精度訓練，因為: 1) fp32 訓練穩定，但記憶體開銷高且不能利用到 NVIDIA GPU 張量核、2) fp16 訓練穩定但難以收斂。更多混合精度訓練相關的內容，我們建議閱讀 tunib-ai 的 notebook。請注意，混合精度要求模型的 fp16/bf16 和 fp32 版本都儲存在記憶體中，而模型需要的記憶體如下:
- 混合精度 (fp16/bf16 + fp32), \(\text {模型記憶體}=2 \text { 位元組} /\text {引數} \cdot \text {引數量}\)
正如上面所講，這個僅僅是模型記憶體，還需要額外加上 \(4 \text { 位元組 / 引數} \cdot \text {引數量}\) 的 用於最佳化器狀態計算 的模型副本，我們會在下面的 最佳化器狀態 一節中算進去。
最佳化器狀態
Adam 有奇效，但記憶體效率非常低。除了要求你有模型權重和梯度外，你還需要額外保留三個梯度引數。因此，
-
對於純 AdamW，\(\text {最佳化器記憶體}=12 \text { 位元組}/\text {引數}\cdot \text {引數量}\)
- fp32 主權重: 4 位元組 / 引數
- 動量 (momentum): 4 位元組 / 引數
- 方差 (variance): 4 位元組 / 引數
-
對於像 bitsandbytes 這樣的 8 位最佳化器，\(\text {最佳化器記憶體} =6 \text { 位元組} /\text {引數} \cdot \text {引數量}\)
- fp32 主權重: 4 位元組 / 引數
- 動量: 1 位元組 / 引數
- 方差: 1 位元組 / 引數
-
對於含動量的類 SGD 最佳化器，\(\text {最佳化器記憶體} =8 \text { 位元組} /\text {引數} \cdot \text {引數量}\)
- fp32 主權重: 4 位元組 / 引數
- 動量: 4 位元組 / 引數
梯度
梯度可以儲存為 fp32 或 fp16 (請注意，梯度資料型別通常與模型資料型別匹配。因此，我們可以看到在 fp16 混合精度訓練中，梯度資料型別為 fp16)，因此它們對記憶體開銷的貢獻為:
- 對於 fp32，\(\text {梯度記憶體} = 4 \text { 位元組} /\text {引數} \cdot \text {引數量}\)
- 對於 fp16，\(\text {梯度記憶體} = 2 \text { 位元組} /\text {引數} \cdot \text {引數量}\)
啟用和 Batch Size
對於 LLM 訓練而言，現代 GPU 通常受限於記憶體瓶頸，而不是算力。因此，啟用重計算 (activation recomputation，或稱為啟用檢查點 (activation checkpointing) ) 就成為一種非常流行的以計算換記憶體的方法。啟用重計算 / 檢查點主要的做法是重新計算某些層的啟用而不把它們存在 GPU 記憶體中，從而減少記憶體的使用量。記憶體的減少量取決於我們選擇清除哪些層的啟用。舉個例子，Megatron 的選擇性重計算方案的效果如下圖所示:
圖中，紅色虛線表示 A100-80GB GPU 的記憶體容量，“present work” 表示使用選擇性啟用重計算後的記憶體需求。請參閱 Reducing Activation Recomputation in Large Transformer Models 一文，瞭解更多詳細資訊以及下述公式的推導過程。
下面給出儲存 transformer 模型啟用所需記憶體的基本公式:
其中:
- \(s\) 是序列長度，即序列中詞元的個數
- \(b\) 是每個 GPU 的 batch size
- \(h\) 是每個 transformer 層的隱含維度
- \(L\) 是 transformer 模型的層數
- \(a\) 是 transformer 模型中注意力頭 (attention heads) 的個數
- \(t\) 是張量並行度 (如果無張量並行，則為 1)
- 我們假設沒有使用序列並行
- 我們假設啟用資料型別為 fp16
由於重計算的引入也會引起計算成本的增加，具體增加多少取決於選擇了多少層進行重計算，但其上界為所有層都額外多了一次前向傳播，因此，更新後的前向傳播計算成本如下:
訓練總記憶體
至此，我們得到了回答 “這個模型是否適合訓練” 這一問題的一個很好的估算公式:
分散式訓練
分片最佳化器 (sharded optimizer)
巨大的最佳化器記憶體開銷促使大家設計和實現了分片最佳化器，目前常用的分片最佳化器實現有 ZeRO 和 FSDP。該分片策略可以使單 GPU 的最佳化器記憶體開銷隨 \(\text {GPU 個數}\) 線性下降，這就是為什麼你會發現某個模型及其訓練配置可能在大規模叢集上能跑，但到小規模叢集時就 OOM (Out Of Memory，記憶體耗盡) 了。下圖來自於 ZeRO 論文，它形象地說明了不同的 ZeRO 階段及它們之間的區別 (注意 \(P_{os}\)、\(P_{os+g }\) 和 \(P_{os+g+p}\) 通常分別表示為 ZeRO-1、ZeRO-2、ZeRO-3。ZeRO-0 通常表示 “禁用 ZeRO”):
下面，我們總結一下 ZeRO 各階段的記憶體開銷公式 (假定我們使用混合精度及 Adam 最佳化器):
- 對於 ZeRO-1，
- 對於 ZeRO-2，
- 對於 ZeRO-3，
其中，在訓練過程沒有使用流水線並行或張量並行的條件下，\(\text {GPU 數}\) 即為 \(\text {DP 並行度}\)。更多詳細資訊，請參閱 Sharded Optimizers + 3D Parallelism 一文。
請注意，ZeRO-3 引入了一組實時引數。這是因為 ZeRO-3 引入了一組配置項 ( stage3_max_live_parameters, stage3_max_reuse_distance, stage3_prefetch_bucket_size, stage3_param_persistence_threshold) 來控制同一時刻 GPU 記憶體中可以放多少引數 (較大的值佔記憶體更多但需要的通訊更少)。這些引數會對總的 GPU 記憶體使用量產生重大影響。
請注意，ZeRO 還可以透過 ZeRO-R 在資料並行 rank 間劃分啟用，這樣 \(\text {啟用記憶體}\) 還可以再除以張量並行度 \(t\)。更詳細的資訊，請參閱相關的 ZeRO 論文 及其 配置選項 (注意，在 GPT-NeoX 中，相應的配置標誌為 partition_activations)。如果你正在訓練一個大模型，啟用放不下記憶體而成為一個瓶頸，你可以使用這個方法用通訊換記憶體。把 ZeRO-R 與 ZeRO-1 結合使用時，記憶體消耗如下:
3D 並行
LLM 主要有 3 種並行方式:
資料並行: 在多個模型副本間拆分資料
流水線或張量 / 模型並行: 在各 GPU 之間拆分模型引數，因此需要大量的通訊開銷。它們的記憶體開銷大約是:
請注意，這是個近似公式，因為 (1) 流水線並行對降低啟用的記憶體需求沒有幫助、(2) 流水線並行要求所有 GPU 儲存所有正在進行的 micro batch 的啟用，這對大模型很重要、(3) GPU 需要臨時儲存並行方案所需的額外通訊緩衝區。
分片最佳化器 + 3D 並行
當 ZeRO 與張量並行、流水線並行結合時，由此產生的 GPU 間的並行策略如下:
值得一提的是，資料並行度對於計算訓練的全域性 batch size 至關重要。資料並行度取決於你想在訓練叢集中保持幾份完整模型副本:
雖然流水線並行和張量並行與所有 ZeRO 階段都相容 (例如，張量並行疊加上 ZeRO-3 後，我們會首先對張量進行切片，然後在每個張量並行單元中應用 ZeRO-3)，但只有 ZeRO-1 與張量和 / 或流水線並行相結合時會效果才會好。這是由於梯度劃分在不同並行策略間會有衝突 (如流水線並行和 ZeRO-2 都會對梯度進行劃分)，這會顯著增加通訊開銷。
把所有東西打包到一起，我們可以得到一個典型的 3D 並行 + ZeRO-1 + 啟用分割槽 方案:
總結
EleutherAI 的工程師經常使用上述估算方法來高效規劃、除錯分散式模型訓練。我們希望澄清這些經常被忽視的但很有用的實現細節，如果你想與我們討論或認為我們錯過了一些好的方法，歡迎你透過 contact@eleuther.ai 聯絡我們！
請使用如下格式引用本文:
@misc {transformer-math-eleutherai,
title = {Transformer Math 101},
author = {Anthony, Quentin and Biderman, Stella and Schoelkopf, Hailey},
howpublished = \url {blog.eleuther.ai/},
year = {2023}
}
英文原文: https://blog.eleuther.ai/transformer-math/
原文作者: Quentin Anthony，Stella Biderman，Hailey Schoelkopf，作者均來自 EleutherAI
譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。
審校/排版: zhongdongy (阿東)