如何計算transformer模型的參數量

1. 前言

最近，OpenAI推出的ChatGPT展現出了卓越的性能，引發了大規模語言模型(Large Language Model,LLM)的研究熱潮。大規模語言模型的“大”體現在兩個方面：模型參數規模大，訓練數據規模大。以GPT3為例，GPT3的參數量為1750億，訓練數據量達到了570GB。進而，訓練大規模語言模型面臨兩個主要挑戰：顯存效率和計算效率。

現在業界的大語言模型都是基于transformer模型的，模型結構主要有兩大類：encoder-decoder（代表模型是T5）和decoder-only，具體的，decoder-only結構又可以分為Causal LM（代表模型是GPT系列）和PrefixLM（代表模型是GLM）。歸因于GPT系列取得的巨大成功，大多數的主流大語言模型都采用Causal LM結構。因此，針對decoder-only框架，為了更好地理解訓練訓練大語言模型的顯存效率和計算效率，本文分析采用decoder-only框架transformer模型的模型參數量、計算量、中間激活值、KV cache。

為了方便分析，先定義好一些數學符號。記transformer模型的層數為?，隱藏層維度為?，注意力頭數為?。詞表大小為?，訓練數據的批次大小為?，序列長度為?。

2. 模型參數量

transformer模型由個相同的層組成，每個層分為兩部分：self-attention塊和MLP塊。

self-attention塊的模型參數有?的權重矩陣和偏置，輸出權重矩陣??和偏置，4個權重矩陣的形狀為?，4個偏置的形狀為?。self- attention塊的參數量為?。

MLP塊由2個線性層組成，一般地，第一個線性層是先將維度從?映射到，第二個線性層再將維度從映射到。第一個線性層的權重矩陣?的形狀為?，偏置的形狀為?。第二個線性層權重矩陣?的形狀為?，偏置形狀為?。MLP塊的參數量為?。

self-attention塊和MLP塊各有一個layer normalization，包含了2個可訓練模型參數：縮放參數?和平移參數?，形狀都是?。2個layernormalization的參數量為??。

總的，每個transformer層的參數量為?。

除此之外，詞嵌入矩陣的參數量也較多，詞向量維度通常等于隱藏層維度，詞嵌入矩陣的參數量為?。最后的輸出層的權重矩陣通常與詞嵌入矩陣是參數共享的。

關于位置編碼，如果采用可訓練式的位置編碼，會有一些可訓練模型參數，數量比較少。如果采用相對位置編碼，例如RoPE和ALiBi，則不包含可訓練的模型參數。我們忽略這部分參數。

綜上，層transformer模型的可訓練模型參數量為。當隱藏維度??較大時，可以忽略一次項，?模型參數量近似為?。

接下來，我們估計不同版本LLaMA模型的參數量。

2.1 訓練過程中的顯存占用分析

在訓練神經網絡的過程中，占用顯存的大頭主要分為四部分：模型參數、前向計算過程中產生的中間激活、后向傳遞計算得到的梯度、優化器狀態。這里著重分析參數、梯度和優化器狀態的顯存占用，中間激活的顯存占用后面會詳細介紹。訓練大模型時通常會采用AdamW優化器，并用混合精度訓練來加速訓練，基于這個前提分析顯存占用。

在一次訓練迭代中，每個可訓練模型參數都會對應1個梯度，并對應2個優化器狀態（Adam優化器梯度的一階動量和二階動量）。設模型參數量為?，那么梯度的元素數量為?，AdamW優化器的元素數量為。float16數據類型的元素占2個bytes，float32數據類型的元素占4個bytes。在混合精度訓練中，會使用float16的模型參數進行前向傳遞和后向傳遞，計算得到float16的梯度；在優化器更新模型參數時，會使用float32的優化器狀態、float32的梯度、float32的模型參數來更新模型參數。因此，對于每個可訓練模型參數，占用了。使用AdamW優化器和混合精度訓練來訓練參數量為?的大模型，?模型參數、梯度和優化器狀態占用的顯存大小為?。

2.2 推理過程中的顯存占用分析

在神經網絡的推理階段，沒有優化器狀態和梯度，也不需要保存中間激活。少了梯度、優化器狀態、中間激活，模型推理階段占用的顯存要遠小于訓練階段。模型推理階段，占用顯存的大頭主要是模型參數，如果使用float16來進行推理，推理階段模型參數占用的顯存大概是?。如果使用KVcache來加速推理過程，?KV cache也需要占用顯存，KVcache占用的顯存下文會詳細介紹。此外，輸入數據也需要放到GPU上，還有一些中間結果（推理過程中的中間結果用完會盡快釋放掉），不過這部分占用的顯存是很小的，可以忽略。

3. 計算量FLOPs估計

FLOPs，floating point operations，表示浮點數運算次數，衡量了計算量的大小。

如何計算矩陣乘法的FLOPs呢？

對于?，計算??需要進行??次乘法運算和??次加法運算，共計??次浮點數運算，需要??的FLOPs。對于??，計算??需要的浮點數運算次數為??。

在一次訓練迭代中，假設輸入數據的形狀為?。我們?先分析self-attention塊的計算，計算公式如下：

1. 計算?：矩陣乘法的輸入和輸出形狀為。計算量為。

2.?矩陣乘法的輸入和輸出形狀為

。計算量為??。

3. 計算在?上的加權??，矩陣乘法的輸入和輸出形狀為。計算量為??。

4. attention后的線性映射，矩陣乘法的輸入和輸出形狀為。計算量為??。

接下來分析MLP塊的計算，計算公式如下：

1. 第一個線性層，矩陣乘法的輸入和輸出形狀為。計算量為??。

2. 第二個線性層，矩陣乘法的輸入和輸出形狀為。計算量為??。

將上述計算量相加，得到每個transformer層的計算量大約為?。

此外，另一個計算量的大頭是logits的計算，將隱藏向量映射為詞表大小。矩陣乘法的輸入和輸出形狀為，計算量為??。

因此，對于一個?層的transformer模型，輸入數據形狀為?的情況下，一次訓練迭代的計算量為。

3.1 計算量與參數量的關聯

當隱藏維度?比較大，且遠大于序列長度?時，我們可以忽略一次項，計算量可以近似為?。前面提到當模型參數量為?，輸入的tokens數為?，存在等式。我們可以近似認為：?在一次前向傳遞中，對于每個token，每個模型參數，需要進行2次浮點數運算，即一次乘法法運算和一次加法運算。

一次訓練迭代包含了前向傳遞和后向傳遞，后向傳遞的計算量是前向傳遞的2倍。因此，前向傳遞 + 后向傳遞的系數。一次訓練迭代中，對于每個token，每個模型參數，需要進行?次浮點數運算。

接下來，我們可以估計訓練GPT3-175B所需要的計算量。對于GPT3，每個token，每個參數進行了6次浮點數運算，再乘以參數量和總tokens數就得到了總的計算量。GPT3的模型參數量為?，訓練數據量為??tokens。

3.2 訓練時間估計

模型參數量和訓練總tokens數決定了訓練transformer模型需要的計算量。給定硬件GPU類型的情況下，可以估計所需要的訓練時間。給定計算量，訓練時間（也就是GPU算完這么多flops的計算時間）不僅跟GPU類型有關，還與GPU利用率有關。計算端到端訓練的GPU利用率時，不僅要考慮前向傳遞和后向傳遞的計算時間，還要**考慮CPU加載數據、優化器更新、多卡通信和記錄日志的時間。一般來講，GPU利用率一般在之間。

上文講到一次前向傳遞中，對于每個token，每個模型參數，進行2次浮點數計算。使用激活重計算技術來減少中間激活顯存（下文會詳細介紹）需要進行一次額外的前向傳遞，因此前向傳遞+ 后向傳遞 + 激活重計算的系數=1+2+1=4。使用激活重計算的一次訓練迭代中，對于每個token，每個模型參數，需要進行?次浮點數運算。在給定訓練tokens數、硬件環境配置的情況下，訓練transformer模型的計算時間為：

以GPT3-175B為例，在1024張40GB顯存的A100上，在300Btokens的數據上訓練175B參數量的GPT3。40GB顯存A100的峰值性能為312TFLOPS，設GPU利用率為0.45，則所需要的訓練時間為34天，這與[7]中的訓練時間是對得上的。

以LLaMA-65B為例，在2048張80GB顯存的A100上，在1.4TBtokens的數據上訓練了65B參數量的模型。80GB顯存A100的峰值性能為624TFLOPS，設GPU利用率為0.3，則所需要的訓練時間為21天，這與[4]中的實際訓練時間是對得上的。

4. 中間激活值分析

除了模型參數、梯度、優化器狀態外，占用顯存的大頭就是前向傳遞過程中計算得到的中間激活值了，需要保存中間激活以便在后向傳遞計算梯度時使用。這里的激活（activations）指的是：前向傳遞過程中計算得到的，并在后向傳遞過程中需要用到的所有張量。這里的激活不包含模型參數和優化器狀態，但包含了dropout操作需要用到的mask矩陣。

在分析中間激活的顯存占用時，只考慮激活占用顯存的大頭，忽略掉一些小的buffers。比如，對于layernormalization，計算梯度時需要用到層的輸入、輸入的均值?和方差?。輸入包含了?個元素，而輸入的均值和方差分別包含了?個元素。由于?通常是比較大的（千數量級），有??。因此，對于layernormalization，中間激活近似估計為??，而不是?。

大模型在訓練過程中通常采用混合精度訓練，中間激活值一般是float16或者bfloat16數據類型的。在分析中間激活的顯存占用時，假設中間激活值是以float16或bfloat16數據格式來保存的，每個元素占了2個bytes。唯一例外的是，dropout操作的mask矩陣，每個元素只占1個bytes。在下面的分析中，單位是bytes，而不是元素個數。

每個transformer層包含了一個self-attention塊和MLP塊，并分別對應了一個layer normalization連接。

先分析self-attention塊的中間激活。self-attention塊的計算公式如下：

1. 對于?，需要保存它們共同的輸入?，這就是中間激活。輸入?的形狀為?，元素個數為?，占用顯存大小為?。

2. 對于?矩陣乘法，需要保存中間激活?，兩個張量的形狀都是?，占用顯存大小合計為?。

3. 對于函數，需要保存函數的輸入??，占用顯存大小為?，這里的?表示注意力頭數。

?的形狀為：?

?的形狀為：

?的形狀為：，元素個數為??，占用顯存大小為?。

4. 計算完函數后，會進行dropout操作。需要保存一個mask矩陣，mask矩陣的形狀與?相同，占用顯存大小為?。

5. 計算在?上的attention，即??，需要保存?，大小為?；以及?，大小為?。二者占用顯存大小合計為?。

6. 計算輸出映射以及一個dropout操作。輸入映射需要保存其輸入，大小為?；dropout需要保存mask矩陣，大小為?。二者占用顯存大小合計為?。

因此，將上述中間激活相加得到，self-attention塊的中間激活占用顯存大小為?。

接下來看MLP塊的中間激活。MLP塊的計算公式如下：

1. 第一個線性層需要保存其輸入，占用顯存大小為?。

2. 激活函數需要保存其輸入，占用顯存大小為?。

3. 第二個線性層需要保存其輸入，占用顯存大小為?。

4. 最后有一個dropout操作，需要保存mask矩陣，占用顯存大小為?。

對于MLP塊，需要保存的中間激活值為?。

另外，self-attention塊和MLP塊分別對應了一個layer normalization。每個layer norm需要保存其輸入，大小為?。2個layer norm需要保存的中間激活為?。

綜上，每個transformer層需要保存的中間激活占用顯存大小為?。對于層transformer模型，還有embedding層、最后的輸出層。embedding層不需要中間激活?？偟亩?，當隱藏維度?比較大，層數?較深時，這部分的中間激活是很少的，可以忽略。因此，對于?層transformer模型，中間激活占用的顯存大小可以近似為。

4.1 對比中間激活與模型參數的顯存大小

在一次訓練迭代中，模型參數（或梯度）占用的顯存大小只與模型參數量和參數數據類型有關，與輸入數據的大小是沒有關系的。優化器狀態占用的顯存大小也是一樣，與優化器類型有關，與模型參數量有關，但與輸入數據的大小無關。而中間激活值與輸入數據的大?。ㄅ未笮?/strong>?和序列長度?）是成正相關的，隨著批次大小?和序列長度的增大，中間激活占用的顯存會同步增大。當我們訓練神經網絡遇到顯存不足OOM（Out OfMemory）問題時，通常會嘗試減小批次大小來避免顯存不足的問題，這種方式減少的其實是中間激活占用的顯存，而不是模型參數、梯度和優化器的顯存。