Author: kaiyuan 編輯: 丁師兄大模型 Link: https://zhuanlan.zhihu.com/p/687226668

AI 算法在服務器中運行時，一個常見問題“單張 GPU 能承載多少模型參數？”，該問題跟模型結構、引擎框架、驅動版本、GPU 硬件相關。

本文圍繞大模型的訓練/推理場景，介紹 Transformer 類模型的顯存計算公式，幫助讀者能更好的了解全局顯存的組成以及如何優化顯存。

文中涉及的主要問題：

• 如何有效估算一個模型加載后的顯存值？

• 計算值與實際 GPU 中的最大值的差距可以有多大？

• 大模型切分策略是如何降低顯存的？計算公式怎么構建？

• 優化顯存的方法和常見的優化思路？

01

模型顯存內容分析

在模型訓練/推理時，顯存（顯卡的全局內存）分配一部分是給 AI 框架，另一部分給了系統（底層驅動）。

總的顯存消耗量可以通過 API 查詢，比如在 NVIDIA-GPU 上通過 nvidia-smi 指令能夠打印出各個進程的顯存消耗量。

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    1   N/A  N/A     67321      C         .../anaconda3/envs/py/bin/python    23646MiB |
|    1   N/A  N/A     71612      C         .../anaconda3/envs/py/bin/python      848MiB |
|    2   N/A  N/A     67321      C         .../anaconda3/envs/py/bin/python    25776MiB |
+---------------------------------------------------------------------------------------+

其中系統層的顯存消耗一般由驅動控制，用戶不可控；框架側的顯存消耗用戶可控，也是本文分析的重點。以 PyTorch 框架為例通過顯存可視化工具，看一下訓練過程中顯存的消耗。

如下圖是一個模型訓練過程中已用顯存的數值隨時間的變化：

注意：數據是具體的消耗值，不等于 cudaMalloc 創建的顯存值。

顯存消耗的內容包括：

• 模型參數（parameter）

• 優化器狀態值（optimizer_state）

• 激活值（activation）

• 梯度值（gradient）

• 輸出數據（input）

• 臨時變量（temporary）

• 自動梯度（autograd_detail）

• 未知數量（unknown）

從用戶側可以將這些數據進行一個分類：

• 可估算值：模型參數（parameter）、優化器狀態值（optimizer_state）、激活值（activation）、梯度值（gradient）、輸出數據（input）

• 未命名數據：臨時變量（temporary）、未知數據（unknown)

• 其他（框架）：自動梯度（autograd_detail）

其中“未命名數據”來源可能是用戶創建的一些臨時變量，這些變量未參與圖的計算過程，所以未被統計；或者是一些未被框架跟蹤（tracing）到的數據。“自動梯度數據"是在反向傳播求解梯度時產生的一些變量；

我們在顯存計算時會發現“為什么有時顯存估算值和實際測量值相差較大？”

其中一個可能的原因是：未知的數據太大。即顯存中可估算值占比相對較小，其它不可估算值的數據占比較大，導致計算值和實際值差距較大（誤差可超過 30%），比如估算得到的顯存消耗為 50GB，而實際測試達到了 75GB。

如下圖是運行一個 LLM 模型采集的一些過程數據，可以看到 unknown 占比有時能達到 30%。

不同時刻顯存的占比變化

02

計算公式

2.1 訓練場景

訓練顯存消耗（可估算部分）主要包括：模型參數（Model）+ 優化器狀態（Optimizer status）+梯度值（Gradient）+激活值（Activation）。

根據數值的變化，可將顯存消耗分為靜態/動態值。訓練過程中，模型參數、優化器狀態一般不會變化，這兩部分歸屬于靜態值；激活值、梯度值會隨著計算過程發生變化，將它們歸類到動態值。

下面主要來看一下這四種類型值的估算方法：

2.1.1 模型顯存（Model Memory)

模型自身所占用的顯存大小與參數量、參數類型相關。常見類型 fp32、fp16/bf16、還有 int8、fp8 等。

關于模型保存的大小估算方法：存儲 checkpoint（ckpt）時僅考慮模型本身，只要將顯存上模型內容存儲到磁盤中。

舉例：以 1B（billion）模型為例，若采用 fp32 類型將其存儲在磁盤上，其大小為：

1B 模型需要 3.725GB 存儲空間，進一步近似認為 1B≈4GB，可方便作存儲的估算推導，如 LLama13b，大約需要 52GB 存儲空間。

注意：混合精度（Mixed-precision）最后存儲的類型也是 fp32，公式也適合混合精度。

2.1.2 優化器狀態（Optimizer status）

在 LLM 中常見的優化器是 Adam，優化器中每個參數需要一個 Momentum 和一個 Variance 狀態參數，在混合精度訓練中 Adam 還有一份模型參數副本。

Adam 參數器狀態值計算公式（單位 GB）：

其中（4+4+4）的內容：

• 模型副本 4 Bytes

• Momentum 參數 4 Bytes

• Variance 參數 4 Bytes

• 如果是 8 位優化器，則計算變為：

• 模型副本 4 Bytes

• Momentum 參數 1Byte

• Variance 參數 1Byte

2.1.3 梯度值（Gradient）

梯度值與模型數據類型保持一致，計算如下（單位 GB）：

2.1.4 激活值（Activation）

激活值的大小跟模型參數、重計算、并行策略等相關，這里我們參考 Megtron 論文里面給的計算公式，來求解激活值所占用的顯存大小。

2.2 訓練的并行計算公式

目前，單卡的物理顯存基本不能滿足大模型的訓練需求，一般會采用模型并行方式來降低單卡顯存消耗。

常見的幾種方法：TP/SP/PP/Zero/重計算，這些方法出現在 DeepSpeed、Megtron 等并行框架中，目標都是讓 GPU 能夠裝下更大的模型。

其中：

• TP（TensorParallel）：tensor 并行；

• SP（SequenceParallel）：序列并行；

• PP（PipelineParallel）：pipeline 并行；

• Zero：參數服務器，分為 Zero1/2/3，最早出現在 deepspeed 中

當沒有并行策略時，僅模型本身的顯存需求（單卡）計算如下：

經過并行策略的調整，顯存需求可變為（舉例，PP/TP/zero1）：

2.3.1 3D 并行

3D 并行主要是 TP(SP)/PP/DP，其中 DP 為數據并行主要用于提升 bs（batch size），DP 不降低單卡的顯存消耗，但 TP(SP)/PP/DP 存在一個耦合關系，DP 的設置一般滿足：

而 TP(SP)/PP 可降低模型、激活值、梯度的顯存占用大小。

3D 并行對顯存計算的影響計算：

注意：梯度顯存沒有除以 TP，主要是考慮到反向計算時需要 AllGather 出完整 gradient。

3D 對激活值顯存的消耗改變需要結合重計算公式進一步分析。另一個問題，當前比較流行的 MoE 方式也會改變模型的參數分布進而改變計算。

但認為 MoE 構造的是多個小模型，改變的是模型的結構，這里計算暫不展開。

考慮MoE時參數的變化

2.3.2 重計算（Recomputation）

一般而言，我們會把前向計算中的中間數據保存下來用于反向計算，從而避免反復計算。

而重計算是指為了降低顯存消耗先丟棄一些前向計算結果，在反向傳播時再重新計算得到。

結合論文[Reducing Activation Recomputation in Large Transformer Models]里面給的計算公式，激活值所占用的顯存的計算公式如下：

單位 GB，參數說明：

• s 序列長度（sequence length），tokens 的量

• b 微批量大小（microbatch size）

• h 隱藏層大小（hidden dimension size）

• a attention 的頭數 （number of attention heads）

• t tensor 并行數值（tensor parallel size）

• L transformer 模型的層數

• λ 比例系數，當為 fp16 時，值等于 1 /（1024 * 1024 * 1024）

假設我們選用 Tensor 和序列并行、不開重計算，則單卡的公式變為：

2.3.4 Zero 方法

Zero 方法對顯存的優化和原理參考其論文[https://arxiv.org/abs/1910.02054]，其中包含了三種策略，對顯存降低的效果不一樣。

zero策略下顯存消耗的計算變化

假設不考慮 3D 并行和重計算，開啟 Zero 的計算公式為：

其中 N 是 GPU 的數量；LiveParams 是 Zero3 引入的參數，這些參數用于控制模型中哪些參數需要加載在 GPU 中，本身的顯存占用不可忽視。

2.3.5 訓練的綜合計算列舉

當條件確定好后，我們可將上述的公式綜合起來求解總的顯存消耗。通過一個具體的示例來說明。

假設相關的運算條件：

• 采用混合精度訓練

• 開啟 TP/SP/PP

• 不開重計算

• 開啟 Zero2

混合精度的單層的數據配置一般如下圖所示，需要注意的是 master weights 只要算一次，要么在優化器中計算要么在模型中計算，這里默認在優化器中考慮。

混合精度數值類型

計算公式如下（單位 GB）：

其中：

相關參數說明：

• params：模型參數

• N：GPU 數量

• PP：pipeline 并行數值

• TP：Tensor + 序列并行數值

• s 序列長度（sequence length），tokens 的量

• b 微批量大小（microbatch size）

• h 隱藏層大小（hidden dimension size）

• a attention 的頭數（number of attention heads）

• t tensor 并行數值（tensor parallel size）

• L transformer 模型的層數

注意：公式計算得到是一個估算值，且只考慮了模型部分，實際運行中的總數還需要考慮框架、分布式通信庫、環境變量、算法產生副本數據。

2.3 推理場景

推理的顯存量組成成分比訓練簡單，有一個簡單的估算公式：

總顯存占用：

相關內容可參看這篇 blog：Transformer Inference Arithmetic | kipply's blog。

總之，通過綜合求解公式可以知道模型顯存消耗主要部分，能幫助我們確定顯存的優化的策略。

03

顯存優化

由于大模型的參數成倍數的增長，遠超出了單 GPU 物理顯存所能承載的范圍，大模型訓練必然需要進行顯存優化。

顯存優化要么是優化算法本身，降低模型算法的顯存消耗；要么是去擴大顯存，通過一些置換方式獲得“額外“空間，由于顯存物理大小一定，我們獲得額外空間的方式不外乎兩種：

• 時間換空間；如，重計算

• 空間轉移；如，多卡并行/offload

其中，時間換空間通常會消耗算力、帶寬；空間轉移主要是消耗 I/O 帶寬，有一定的時延，可能會降低吞吐。

顯存優化的過程一般是從模型算法本身到底層，可以參考的優化路徑：

多卡并行 -> 算子/數據類型 -> 消除框架副本 -> 顯存管理 -> 底層 API

1、多卡并行：該手段相對來說是使用頻率最高，且一般不會影響運算的精度，可以用 2 節中的計算公式為參考去設計新的 TP/PP/DP/Zero/重計算的相關參數來降低顯存消耗。缺點：這些方式可能會增加額外的帶寬消耗。

2、算子優化：選取精度相同但顯存消耗更低的算子/方案。缺點：一般情況下，算子優化的過程耗時較長。

3、數據類型修改：用低精度替換高精度數據。比如用 fp16 代替 fp32，或者用更低的 int8/int4。缺點：該方式可能影響訓練收斂性/推理性能。

4、消除框架副本：在 AI 框架（如 pytorch）中有些數據是一些由框架產生的中間副本，可以進行優化消除；缺點：游湖成本較大。

5、顯存管理：通過顯存管理的知識可知[PyTorch 顯存管理]，框架的顯存管理會產生顯存碎片，通過優化顯存管理來優化碎片；缺點：目前可用的手段較少。

6、底層 API：在 GPU 的驅動庫中/CUDA 算子庫中，不同 API 顯存消耗不一樣，我們可以用顯存消耗更小算子去替換大顯存消耗算子，比如FlashAttention；

有些默認的操作會產生額外系統顯存，也可以考慮替換更高版本優化后的 API。