散文網(wǎng) » 科技 »學(xué)習(xí) » Transformer 估算 101

Transformer 估算 101

2023-05-19 15:45 作者:HuggingFace 0人讀過 | 我要投稿

本文主要介紹用于估算 transformer 類模型計算量需求和內(nèi)存需求的相關(guān)數(shù)學(xué)方法。

引言

其實，很多有關(guān) transformer 語言模型的一些基本且重要的信息都可以用很簡單的方法估算出來。不幸的是，這些公式在 NLP 社區(qū)中鮮為人知。本文的目的是總結(jié)這些公式，闡明它們是如何推導(dǎo)出來的及其作用。

注意:?本文主要關(guān)注訓(xùn)練成本，該成本主要由 GPU 的 VRAM 主導(dǎo)。如果你想知道有關(guān)推理成本 (通常由推理延遲主導(dǎo)) 的信息，可以讀讀 Kipply 寫的?這篇精彩博文。

算力要求

下式可用于估算訓(xùn)練一個 transformer 模型所需的算力成本:

$C \approx τ T = 6 P D$

這里:

$C$ ?是訓(xùn)練 transformer 模型所需的計算量，單位為總浮點運算數(shù) (FLOP)
$C = C^{前向} + C^{后向}$
$C^{前向} \approx 2 P D$
$C^{后向} \approx 4 P D$
$τ$ ?是訓(xùn)練集群的實際總吞吐量:? $τ = GPU? 數(shù) \times 每 ?GPU? 的實際每秒浮點運算數(shù) ?(實際 ?FLOPs)?$ ，單位為 FLOPs
$T$ ?是訓(xùn)練模型所花費的時間，以秒為單位
$P$ ?是 transformer 模型的參數(shù)量
$D$ ?是數(shù)據(jù)集大小，表示為數(shù)據(jù)集的總詞元數(shù)

該式由?OpenAI 的縮放定律論文?和?DeepMind 的縮放定律論文?提出并經(jīng)其實驗驗證。想要獲取更多信息，可參閱這兩篇論文。

下面，我們討論一下? $C$ ?的單位。? $C$ ?是總計算量的度量，我們可以用許多單位來度量它，例如:

FLOP - 秒，單位為? $每秒浮點運算數(shù) \times 秒$
GPU - 時，單位為? $GPU? 數(shù) \times 小時$
縮放定律論文傾向于以?PetaFLOP - 天?為單位，即單位為? $1 0^{15} \times 24 \times 3600$

這里需要注意? $實際 ?FLOPs$ ?的概念。雖然 GPU 的規(guī)格書上通常只宣傳其理論 FLOPs，但在實踐中我們從未達(dá)到過這些理論值 (尤其在分布式訓(xùn)練時！)。我們將在計算成本這一小節(jié)列出分布式訓(xùn)練中常見的? $實際 ?FLOPs$ ?值。

請注意，上面的算力成本公式來自于?這篇關(guān)于 LLM 訓(xùn)練成本的精彩博文。

參數(shù)量與數(shù)據(jù)集的權(quán)衡

嚴(yán)格來講，你可以隨心所欲地使用任意數(shù)量的詞元來訓(xùn)練 transformer 模型，但由于參與訓(xùn)練的詞元數(shù)會極大地影響計算成本和最終模型的性能，因此需要小心權(quán)衡。

我們從最關(guān)鍵的部分開始談起: “計算最優(yōu)” 語言模型?！癈hinchilla 縮放定律”，得名于提出 “計算最優(yōu)” 語言模型論文中所訓(xùn)練的模型名，指出計算最優(yōu)語言模型的?參數(shù)量?和?數(shù)據(jù)集大小?的近似關(guān)系滿足:? $D = 20 P$ 。該關(guān)系成立基于一個前提條件: 使用 1,000 個 GPU 1 小時和使用 1 個 GPU 1,000 小時成本相同。如果你的情況滿足該條件，你應(yīng)該使用上述公式去訓(xùn)練一個性能最優(yōu)且?GPU - 時?成本最小的模型。

但?我們不建議在少于 200B 詞元的數(shù)據(jù)集上訓(xùn)練 LLM。?雖然對于許多模型尺寸來說這是 “Chinchilla 最優(yōu)” 的，但生成的模型通常比較差。對于幾乎所有應(yīng)用而言，我們建議確定你可接受的推理成本，并訓(xùn)練一個滿足該推理成本要求的最大模型。

計算成本的經(jīng)驗值

Transformer 模型的計算成本通常以?GPU - 時?或?FLOP - 秒?為單位。

GPT-NeoX 的?實際 TFLOP/s?在正常注意力機(jī)制下達(dá)到 150 TFLOP/s/A100，在 Flash 注意力機(jī)制下達(dá)到 180 FLOP/s/A100。這與其他高度優(yōu)化的大規(guī)模計算庫一致，例如 Megatron-DS 的值是在 137 和 163 TFLOP/s/A100 之間。
一個通用的經(jīng)驗法則是?實際 TFLOP/s?可至 120 TFLOP/s/A100 左右。如果你得到低于 115 TFLOP/s/A100 的值，可能是你的模型或硬件配置有問題。
借助 InfiniBand 等高速互連設(shè)備，你可以在數(shù)據(jù)并行維度上實現(xiàn)線性或亞線性擴(kuò)展 (即增加數(shù)據(jù)并行度應(yīng)該以近乎線性的方式增加整體吞吐量)。下圖顯示了在橡樹嶺國家實驗室 (Oak Ridge National Lab) 的 Summit 超級計算機(jī)上測試出的 GPT-NeoX 庫的擴(kuò)展性。請注意，這張圖用的是 V100，而本文中的其他大多數(shù)例子都是基于 A100 的。

內(nèi)存需求

Transformer 模型通常由其?參數(shù)尺寸?來描述。但是，根據(jù)給定一組計算資源確定要訓(xùn)練哪些模型時，你需要知道?該模型將占用多少空間 (以字節(jié)為單位)。這不僅需要考慮你的本地 GPU 可以推理多大的模型，還需要考慮給定訓(xùn)練集群中的總可用 GPU 內(nèi)存可供訓(xùn)練多大的模型。

推理

模型權(quán)重

大多數(shù) transformer 模型都使用?混合精度進(jìn)行訓(xùn)練，可以是 fp16 + fp32 或是 bf16 + fp32?；旌暇冉档土四Ｐ陀?xùn)練和推理所需的內(nèi)存量。推理時，我們還可以將語言模型從 fp32 轉(zhuǎn)換為 fp16 甚至 int8，而沒有實質(zhì)性的精度損失。下面我們看看在不同的數(shù)據(jù)類型下，模型所需內(nèi)存有什么不同 (以字節(jié)為單位):

對 int8 而言， $模型內(nèi)存 = 1 ? 字節(jié) / 參數(shù) ? 參數(shù)量$
對 fp16 和 bf16 而言， $模型內(nèi)存 = 2 ? 字節(jié) / 參數(shù) ? 參數(shù)量$
對 fp32 而言， $模型內(nèi)存 = 4 ? 字節(jié) / 參數(shù) ? 參數(shù)量$

推理總內(nèi)存

除了存儲模型權(quán)重所需的內(nèi)存外，實際中前向傳播過程中還會有少量額外開銷。根據(jù)我們的經(jīng)驗，此開銷在 20% 以內(nèi)，該比例通常與模型無關(guān)。

總的來說，回答 “這個模型是否適合推理” 這一問題，可以用下述公式來獲得不錯的估計:

$推理總內(nèi)存 \approx 1.2 \times 模型內(nèi)存$

本文不會深究該開銷的來源，留待后面的文章來闡述。在本文的后續(xù)部分，我們將主要關(guān)注模型訓(xùn)練的內(nèi)存。如果你有興趣了解更多有關(guān)推理計算需求的信息，請查看?這篇深入介紹推理的精彩博文?，F(xiàn)在，我們要開始訓(xùn)練了！

訓(xùn)練

除了模型權(quán)重之外，訓(xùn)練還需要在設(shè)備內(nèi)存中存儲優(yōu)化器狀態(tài)和梯度。這就是為什么當(dāng)你問 “我需要多少內(nèi)存來運行模型？”，別人會立即回答 “這取決于是訓(xùn)練還是推理”。訓(xùn)練總是比推理需要更多的內(nèi)存，通常多得多！

模型權(quán)重

首先，可以使用純 fp32 或純 fp16 訓(xùn)練模型:

純 fp32， $模型內(nèi)存 = 4 ? 字節(jié) / 參數(shù) ? 參數(shù)量$
純 fp16， $模型內(nèi)存 = 2 ? 字節(jié) / 參數(shù) ? 參數(shù)量$

除了推理中討論的常見模型權(quán)重數(shù)據(jù)類型外，訓(xùn)練還引入了?混合精度?訓(xùn)練，例如?AMP。該技術(shù)尋求在保持收斂性的同時最大化 GPU 張量核的吞吐量?，F(xiàn)代 DL 訓(xùn)練領(lǐng)域經(jīng)常使用混合精度訓(xùn)練，因為: 1) fp32 訓(xùn)練穩(wěn)定，但內(nèi)存開銷高且不能利用到 NVIDIA GPU 張量核、2) fp16 訓(xùn)練穩(wěn)定但難以收斂。更多混合精度訓(xùn)練相關(guān)的內(nèi)容，我們建議閱讀?tunib-ai 的 notebook。請注意，混合精度要求模型的 fp16/bf16 和 fp32 版本都存儲在內(nèi)存中，而模型需要的內(nèi)存如下:

混合精度 (fp16/bf16 + fp32),? $模型內(nèi)存 = 2 ? 字節(jié) / 參數(shù) ? 參數(shù)量$

正如上面所講，這個僅僅是模型內(nèi)存，還需要額外加上? $4 ? 字節(jié) ?/? 參數(shù) ? 參數(shù)量$ ?的?用于優(yōu)化器狀態(tài)計算?的模型副本，我們會在下面的?優(yōu)化器狀態(tài)?一節(jié)中算進(jìn)去。

優(yōu)化器狀態(tài)

Adam 有奇效，但內(nèi)存效率非常低。除了要求你有模型權(quán)重和梯度外，你還需要額外保留三個梯度參數(shù)。因此，

對于純 AdamW， $優(yōu)化器內(nèi)存 = 12 ? 字節(jié) / 參數(shù) ? 參數(shù)量$

fp32 主權(quán)重: 4 字節(jié) / 參數(shù)
動量 (momentum): 4 字節(jié) / 參數(shù)
方差 (variance): 4 字節(jié) / 參數(shù)

對于像?bitsandbytes?這樣的 8 位優(yōu)化器， $優(yōu)化器內(nèi)存 = 6 ? 字節(jié) / 參數(shù) ? 參數(shù)量$

fp32 主權(quán)重: 4 字節(jié) / 參數(shù)
動量: 1 字節(jié) / 參數(shù)
方差: 1 字節(jié) / 參數(shù)

對于含動量的類 SGD 優(yōu)化器， $優(yōu)化器內(nèi)存 = 8 ? 字節(jié) / 參數(shù) ? 參數(shù)量$

fp32 主權(quán)重: 4 字節(jié) / 參數(shù)
動量: 4 字節(jié) / 參數(shù)

梯度

梯度可以存儲為 fp32 或 fp16 (請注意，梯度數(shù)據(jù)類型通常與模型數(shù)據(jù)類型匹配。因此，我們可以看到在 fp16 混合精度訓(xùn)練中，梯度數(shù)據(jù)類型為 fp16)，因此它們對內(nèi)存開銷的貢獻(xiàn)為:

對于 fp32， $梯度內(nèi)存 = 4 ? 字節(jié) / 參數(shù) ? 參數(shù)量$
對于 fp16， $梯度內(nèi)存 = 2 ? 字節(jié) / 參數(shù) ? 參數(shù)量$

激活和 Batch Size

對于 LLM 訓(xùn)練而言，現(xiàn)代 GPU 通常受限于內(nèi)存瓶頸，而不是算力。因此，激活重計算 (activation recomputation，或稱為激活檢查點 (activation checkpointing) ) 就成為一種非常流行的以計算換內(nèi)存的方法。激活重計算 / 檢查點主要的做法是重新計算某些層的激活而不把它們存在 GPU 內(nèi)存中，從而減少內(nèi)存的使用量。內(nèi)存的減少量取決于我們選擇清除哪些層的激活。舉個例子，Megatron 的選擇性重計算方案的效果如下圖所示:

圖中，紅色虛線表示 A100-80GB GPU 的內(nèi)存容量，“present work” 表示使用選擇性激活重計算后的內(nèi)存需求。請參閱?Reducing Activation Recomputation in Large Transformer Models?一文，了解更多詳細(xì)信息以及下述公式的推導(dǎo)過程。

下面給出存儲 transformer 模型激活所需內(nèi)存的基本公式:

$無重計算的激活內(nèi)存 = s bh L (10 + \frac{24}{t} + 5 \frac{a ? s}{h ? t}) ? 字節(jié)$

$選擇性重計算的激活內(nèi)存 = s bh L (10 + \frac{24}{t}) ? 字節(jié)$

$全部重計算的激活內(nèi)存 = 2 ? s bh L ? 字節(jié)$

其中:

$s$ ?是序列長度，即序列中詞元的個數(shù)
$b$ ?是每個 GPU 的 batch size
$h$ ?是每個 transformer 層的隱含維度
$L$ ?是 transformer 模型的層數(shù)
$a$ ?是 transformer 模型中注意力頭 (attention heads) 的個數(shù)
$t$ ?是張量并行度 (如果無張量并行，則為 1)
我們假設(shè)沒有使用序列并行
我們假設(shè)激活數(shù)據(jù)類型為 fp16

由于重計算的引入也會引起計算成本的增加，具體增加多少取決于選擇了多少層進(jìn)行重計算，但其上界為所有層都額外多了一次前向傳播，因此，更新后的前向傳播計算成本如下:

$2 P D \leq C_forward \leq 4 P D$

訓(xùn)練總內(nèi)存

至此，我們得到了回答 “這個模型是否適合訓(xùn)練” 這一問題的一個很好的估算公式:

$訓(xùn)練總內(nèi)存 = 模型內(nèi)存 + 優(yōu)化器內(nèi)存 + 激活內(nèi)存 + 梯度內(nèi)存$

分布式訓(xùn)練

分片優(yōu)化器 (sharded optimizer)

巨大的優(yōu)化器內(nèi)存開銷促使大家設(shè)計和實現(xiàn)了分片優(yōu)化器，目前常用的分片優(yōu)化器實現(xiàn)有?ZeRO?和?FSDP。該分片策略可以使單 GPU 的優(yōu)化器內(nèi)存開銷隨? $GPU? 個數(shù)$ ?線性下降，這就是為什么你會發(fā)現(xiàn)某個模型及其訓(xùn)練配置可能在大規(guī)模集群上能跑，但到小規(guī)模集群時就 OOM (Out Of Memory，內(nèi)存耗盡) 了。下圖來自于?ZeRO?論文，它形象地說明了不同的 ZeRO 階段及它們之間的區(qū)別 (注意? $P^{os}$ 、 $P^{os + g}$ ?和? $P^{os + g + p}$ ?通常分別表示為 ZeRO-1、ZeRO-2、ZeRO-3。ZeRO-0 通常表示 “禁用 ZeRO”):

下面，我們總結(jié)一下 ZeRO 各階段的內(nèi)存開銷公式 (假定我們使用混合精度及 Adam 優(yōu)化器):

對于 ZeRO-1，

$訓(xùn)練總內(nèi)存 \approx 模型內(nèi)存 + \frac{優(yōu)化器內(nèi)存}{GPU? 數(shù)} + 激活內(nèi)存 + 梯度內(nèi)存$

對于 ZeRO-2，

$訓(xùn)練總內(nèi)存 \approx 模型內(nèi)存 + 激活內(nèi)存 + \frac{優(yōu)化器內(nèi)存 + 梯度內(nèi)存}{GPU? 數(shù)}$

對于 ZeRO-3，

$訓(xùn)練總內(nèi)存 \approx 激活內(nèi)存 + \frac{模型內(nèi)存 + 優(yōu)化器內(nèi)存 + 梯度內(nèi)存}{GPU? 數(shù)} + (ZeRO-3? 實時參數(shù)量)$

其中，在訓(xùn)練過程沒有使用流水線并行或張量并行的條件下， $GPU? 數(shù)$ ?即為? $DP? 并行度$ 。更多詳細(xì)信息，請參閱?Sharded Optimizers + 3D Parallelism?一文。

請注意，ZeRO-3 引入了一組實時參數(shù)。這是因為 ZeRO-3 引入了一組配置項 (?stage3_max_live_parameters, stage3_max_reuse_distance, stage3_prefetch_bucket_size, stage3_param_persistence_threshold) 來控制同一時刻 GPU 內(nèi)存中可以放多少參數(shù) (較大的值占內(nèi)存更多但需要的通信更少)。這些參數(shù)會對總的 GPU 內(nèi)存使用量產(chǎn)生重大影響。

請注意，ZeRO 還可以通過?ZeRO-R?在數(shù)據(jù)并行 rank 間劃分激活，這樣? $激活內(nèi)存$ ?還可以再除以張量并行度? $t$ 。更詳細(xì)的信息，請參閱相關(guān)的?ZeRO 論文?及其?配置選項?(注意，在 GPT-NeoX 中，相應(yīng)的配置標(biāo)志為?partition_activations)。如果你正在訓(xùn)練一個大模型，激活放不下內(nèi)存而成為一個瓶頸，你可以使用這個方法用通信換內(nèi)存。把 ZeRO-R 與 ZeRO-1 結(jié)合使用時，內(nèi)存消耗如下:

$訓(xùn)練總內(nèi)存 \approx 模型內(nèi)存 + \frac{優(yōu)化器內(nèi)存}{GPU? 數(shù)} + \frac{激活內(nèi)存}{張量并行度} + 梯度內(nèi)存$

3D 并行

LLM 主要有 3 種并行方式:

數(shù)據(jù)并行:?在多個模型副本間拆分?jǐn)?shù)據(jù)

流水線或張量 / 模型并行:?在各 GPU 之間拆分模型參數(shù)，因此需要大量的通信開銷。它們的內(nèi)存開銷大約是:

$并行后模型內(nèi)存 \approx \frac{模型內(nèi)存}{流水線并行度 \times 張量并行度}$

$并行后梯度內(nèi)存 \approx \frac{梯度內(nèi)存}{流水線并行度}$

請注意，這是個近似公式，因為 (1) 流水線并行對降低激活的內(nèi)存需求沒有幫助、(2) 流水線并行要求所有 GPU 存儲所有正在進(jìn)行的 micro batch 的激活，這對大模型很重要、(3) GPU 需要臨時存儲并行方案所需的額外通信緩沖區(qū)。

分片優(yōu)化器 + 3D 并行

當(dāng) ZeRO 與張量并行、流水線并行結(jié)合時，由此產(chǎn)生的 GPU 間的并行策略如下:

值得一提的是，數(shù)據(jù)并行度對于計算訓(xùn)練的全局 batch size 至關(guān)重要。數(shù)據(jù)并行度取決于你想在訓(xùn)練集群中保持幾份完整模型副本:

$數(shù)據(jù)并行度 ?=? \frac{GPU? 數(shù)}{流水線并行度 \times 張量并行度}$

雖然流水線并行和張量并行與所有 ZeRO 階段都兼容 (例如，張量并行疊加上 ZeRO-3 后，我們會首先對張量進(jìn)行切片，然后在每個張量并行單元中應(yīng)用 ZeRO-3)，但只有 ZeRO-1 與張量和 / 或流水線并行相結(jié)合時會效果才會好。這是由于梯度劃分在不同并行策略間會有沖突 (如流水線并行和 ZeRO-2 都會對梯度進(jìn)行劃分)，這會顯著增加通信開銷。

把所有東西打包到一起，我們可以得到一個典型的?3D 并行 + ZeRO-1 + 激活分區(qū)?方案:

$訓(xùn)練總內(nèi)存 \approx \frac{模型內(nèi)存}{流水線并行度 \times 張量并行度} + \frac{優(yōu)化器內(nèi)存}{GPU? 數(shù)} + \frac{激活內(nèi)存}{張量并行度} + \frac{梯度內(nèi)存}{流水線并行度}$

總結(jié)

EleutherAI 的工程師經(jīng)常使用上述估算方法來高效規(guī)劃、調(diào)試分布式模型訓(xùn)練。我們希望澄清這些經(jīng)常被忽視的但很有用的實現(xiàn)細(xì)節(jié)，如果你想與我們討論或認(rèn)為我們錯過了一些好的方法，歡迎你通過?contact@eleuther.ai?聯(lián)系我們！

請使用如下格式引用本文:

英文原文:?https://blog.eleuther.ai/transformer-math/

原文作者: Quentin Anthony，Stella Biderman，Hailey Schoelkopf，作者均來自 EleutherAI

譯者: Matrix Yao (姚偉峰)，英特爾深度學(xué)習(xí)工程師，工作方向為 transformer-family 模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓(xùn)練推理。

審校/排版: zhongdongy (阿東)

標(biāo)簽：