@(Engineering Practice)

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

引言

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

注意： 本文主要關(guān)注訓練成本肛循，該成本主要由 GPU 的 VRAM 主導(dǎo)。如果你想知道有關(guān)推理成本（通常由推理延遲主導(dǎo)）的信息银择，可以讀讀 Kipply 寫的這篇精彩博文多糠。

算力要求

下式可用于估算訓練一個 transformer 模型所需的算力成本：

這里:

• 是訓練 transformer 模型所需的計算量，單位為總浮點運算數(shù)（FLOP）
• 
• 
• 
• 是訓練集群的實際總吞吐量：欢摄，單位為 FLOPs
• 是訓練模型所花費的時間熬丧，以秒為單位
• 是 transformer 模型的參數(shù)量
• 是數(shù)據(jù)集大小，表示為數(shù)據(jù)集的總詞元數(shù)

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

下面绿淋，我們討論一下 的單位闷畸。 是總計算量的度量，我們可以用許多單位來度量它吞滞，例如：

• FLOP - 秒佑菩，單位為
• GPU - 時盾沫，單位為
• 縮放定律論文傾向于以 PetaFLOP - 天 為單位，即單位為

這里需要注意 的概念殿漠。雖然 GPU 的規(guī)格書上通常只宣傳其理論 FLOPs赴精，但在實踐中我們從未達到過這些理論值（尤其在分布式訓練時！）绞幌。我們將在計算成本這一小節(jié)列出分布式訓練中常見的 值蕾哟。

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

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

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

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

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

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

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

• GPT-NeoX 的 實際 TFLOP/s 在正常注意力機制下達到 150 TFLOP/s/A100胆胰，在 Flash 注意力機制下達到 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)線性或亞線性擴展（即增加數(shù)據(jù)并行度應(yīng)該以近乎線性的方式增加整體吞吐量）沐兵。下圖顯示了在橡樹嶺國家實驗室（Oak Ridge National Lab）的 Summit 超級計算機上測試出的 GPT-NeoX 庫的擴展性别垮。請注意，這張圖用的是 V100扎谎，而本文中的其他大多數(shù)例子都是基于 A100 的碳想。

內(nèi)存需求

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

推理

模型權(quán)重

模型權(quán)重

大多數(shù) transformer 模型都使用混合精度進行訓練螟左，可以是 fp16 + fp32 或是 bf16 + fp32∶俟唬混合精度降低了模型訓練和推理所需的內(nèi)存量。推理時巷嚣，我們還可以將語言模型從 fp32 轉(zhuǎn)換為 fp16 甚至 int8喘先，而沒有實質(zhì)性的精度損失。下面我們看看在不同的數(shù)據(jù)類型下廷粒，模型所需內(nèi)存有什么不同（以字節(jié)為單位）：

• 對 int8 而言窘拯，
• 對 fp16 和 bf16 而言，
• 對 fp32 而言坝茎，

推理總內(nèi)存

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

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

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

訓練

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

模型權(quán)重

首先涡上，可以使用純 fp32 或純 fp16 訓練模型：

• 純 fp32，
• 純 fp16拒名，

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

• 混合精度 (fp16/bf16 + fp32),

正如上面所講禁偎，這個僅僅是模型內(nèi)存，還需要額外加上 的用于優(yōu)化器狀態(tài)計算 的模型副本阀坏，我們會在下面的 優(yōu)化器狀態(tài) 一節(jié)中算進去如暖。

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

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

• 對于純 AdamW妄迁，
  • fp32 主權(quán)重：4 字節(jié) / 參數(shù)
  • 動量（momentum）：4 字節(jié) / 參數(shù)
  • 方差（variance）：4 字節(jié) / 參數(shù)
• 對于像 bitsandbytes 這樣的 8 位優(yōu)化器，
  • fp32 主權(quán)重：4 字節(jié) / 參數(shù)
  • 動量：1 字節(jié) / 參數(shù)
  • 方差：1 字節(jié) / 參數(shù)
• 對于含動量的類 SGD 優(yōu)化器李命，
  • fp32 主權(quán)重：4 字節(jié) / 參數(shù)
  • 動量：4 字節(jié) / 參數(shù)

梯度

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

• 對于 fp32阔籽，
• 對于 fp16流妻，

激活和 Batch Size

對于 LLM 訓練而言，現(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 一文，了解更多詳細信息以及下述公式的推導(dǎo)過程先慷。

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

其中：

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

由于重計算的引入也會引起計算成本的增加论熙，具體增加多少取決于選擇了多少層進行重計算缎浇，但其上界為所有層都額外多了一次前向傳播，因此赴肚，更新后的前向傳播計算成本如下：

訓練總內(nèi)存

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

分布式訓練

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

巨大的優(yōu)化器內(nèi)存開銷促使大家設(shè)計和實現(xiàn)了分片優(yōu)化器二蓝，目前常用的分片優(yōu)化器實現(xiàn)有 ZeRO 和 FSDP誉券。該分片策略可以使單 GPU 的優(yōu)化器內(nèi)存開銷隨 線性下降，這就是為什么你會發(fā)現(xiàn)某個模型及其訓練配置可能在大規(guī)模集群上能跑刊愚，但到小規(guī)模集群時就 OOM（Out Of Memory踊跟，內(nèi)存耗盡）了。下圖來自于 ZeRO 論文鸥诽，它形象地說明了不同的 ZeRO 階段及它們之間的區(qū)別（注意 商玫、 和 通常分別表示為 ZeRO-1、ZeRO-2牡借、ZeRO-3拳昌。ZeRO-0 通常表示 “禁用 ZeRO”）：

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

• 對于 ZeRO-1钠龙，

• 對于 ZeRO-2炬藤，

• 對于 ZeRO-3，

其中碴里，在訓練過程沒有使用流水線并行或張量并行的條件下沈矿， 即為 。更多詳細信息咬腋，請參閱 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 間劃分激活蠢壹，這樣 還可以再除以張量并行度 嗓违。更詳細的信息，請參閱相關(guān)的 ZeRO 論文 及其 配置選項（注意图贸，在 GPT-NeoX 中蹂季，相應(yīng)的配置標志為 partition_activations）。如果你正在訓練一個大模型疏日，激活放不下內(nèi)存而成為一個瓶頸偿洁，你可以使用這個方法用通信換內(nèi)存。把 ZeRO-R 與 ZeRO-1 結(jié)合使用時沟优，內(nèi)存消耗如下：

3D 并行

LLM 主要有 3 種并行方式：

數(shù)據(jù)并行： 在多個模型副本間拆分數(shù)據(jù)

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

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

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

當 ZeRO 與張量并行隘谣、流水線并行結(jié)合時增拥，由此產(chǎn)生的 GPU 間的并行策略如下：

3D 并行

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

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

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

總結(jié)

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

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

@misc {transformer-math-eleutherai,
  title = {Transformer Math 101},
  author = {Anthony, Quentin and Biderman, Stella and Schoelkopf, Hailey},
  howpublished = \url {blog.eleuther.ai/},
  year = {2023}
}

英文原文: <url> https://blog.eleuther.ai/transformer-math/ </url>
原文作者：Quentin Anthony，Stella Biderman，Hailey Schoelkopf嫩舟，作者們均來自EleutherAI
譯者: Matrix Yao (姚偉峰)氢烘，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓練推理家厌。