HunYuan MoE：聊一聊 LLM 參數(shù)量、計算量和 MFU 等
一、背景
最近在看騰訊最新混元大模型的 Paper 時（[2411.02265] Hunyuan-Large: An Open-Source MoE Model with 52 Billion Activated Parameters by Tencent [1]），看到了如下關(guān)于計算 Budget 的公式
由于我們的工作中也會經(jīng)常根據(jù)計算資源 Budget 評估 LLM 預(yù)訓(xùn)練時間，而該公式與我們平常的計算方式不太一致；此外，如下圖所示，我們也看到很多文章中將上述公式中的第二項理解為長序列情況下 Attention 的額外計算開銷，而將 6ND -> 9.6ND 看成 Router 引入的開銷，與我們的計算結(jié)論也非常不一致，因此本文中我們簡單回顧一下相關(guān)內(nèi)容。
本文中的很多內(nèi)容我們在之前的文章中已經(jīng)介紹過，具體可以參考：
二、LLM 參數(shù)量&計算量
2.1 Dense LLM 參數(shù)量
標(biāo)準(zhǔn) Transformer 模型不管是 Encoder Only 的 Bert 系列模型，還是 Decoder Only 的 GPT 系列模型，同配置下參數(shù)量和計算量都是類似的。其中的一個關(guān)鍵點是：標(biāo)準(zhǔn) Transformer Block 輸入、輸出以及中間 Attention Out 的 Hidden Dim 是保持不變的，也就是一直是 Token Embedding 的 Hidden Dim。所有的 Transformer Block 都非常規(guī)整。
如下圖 Figure 1 Encoder 所示（圖片來自 [2302.14017] Full Stack Optimization of Transformer Inference: a Survey [2]），主要參數(shù)都來自幾個矩陣乘的 Weight 矩陣，其中 d 表示 Token Embedding 的 Hidden Dim，l 表示 Token 數(shù)，h 表示 MHA 中的 Head 個數(shù)，dFFN 表示 FFN 層中間升維后的 Dim。具體來說：
- MHA紅框：Wq，Wk，Wv 的大小都是 d x d。當(dāng)然這里也可以從 h 個 Head 的角度去看，則每個 Head 的 Wq，Wk，Wv 為 d x d/h，實際上參數(shù)量和計算量都是等價的。
- MHA藍(lán)框：在 MHA 的最后還有一個矩陣乘操作，對應(yīng)的 Wout 維度依然為 d x d。
- FFN綠框：在標(biāo)準(zhǔn) Transformer 的 FFN 中有兩個 Linear 層，先升維再降維，對應(yīng)的 Weight 矩陣 W1 和 W2 的大小都是 dFFN x d，并且標(biāo)準(zhǔn)的 dFFN 為 4d，也就是說 FFN 處兩個權(quán)重矩陣的參數(shù)量為 8d x d。?
綜上，在標(biāo)準(zhǔn)的 Transformer 模型或者 LLaMA 系列（MHA）中，如果忽略詞表、Embedding、LayerNorm 等參數(shù)后，總參數(shù)量為（所有 Transformer Block）：
N = nlayer *(nmha + nffn) = nlayer *(3d*d + d*d + 8d*d) = 12*nlayer*d*d
2.2 LLaMA LLM 參數(shù)量
2.2.1 SwiGLU
在 LLaMA 等模型中在 FFN 中會使用 SwiGLU 激活，這也就導(dǎo)致其會額外多了一個權(quán)重矩陣，如下圖所示：
在 LLaMA-1 的 Paper（[2302.13971] LLaMA: Open and Efficient Foundation Language Models [3]）中作者提到，使用 SwiGLU 后將 dFFN 從 4d 降低到了 8d/3。這樣 3 個權(quán)重矩陣的參數(shù)量還是 8d，總的參數(shù)量依然可以使用 12*nlayer*d*d 預(yù)估。
2.2.2 GQA
之前公式對應(yīng)的是 MHA（Multi Head Attention），這也是 LLaMA-1 系列模型的標(biāo)準(zhǔn)實現(xiàn)。不過，LLaMA-2 的 30B 和 70B 模型以及 LLaMA-3 的全部模型都開始使用 GQA（Grouped Query Attention）。使用 GQA 時，多個 Attention Head 會共享一個 Key 和 Value，比如常見的配置是每 4 或 8 個 Head 共享相同的 Key 和 Value，此時 Wk，Wv 的大小會變?yōu)?d x d/g，其中 g 表示每 g 個 Head 共享相同的 Key 和 Value。
在 LLaMA 2 的 Paper（[2307.09288] Llama 2: Open Foundation and Fine-Tuned Chat Models [4]）中，作者提到，為了保持使用 GQA 和使用 MHA 的總參數(shù)量保持不變，對于 GQA 模型，會將 FFN Dim 維度乘以 1.3（PS：不確定 1.3 是怎么來的）：
比如 LLaMA 2 70B 的 d 為 8192，dFFN 為 28672，如下圖所示，大概計算方式為（512 對齊）：
int(int(8192*8/3*1.3+512-1)/512)*512=28672
比較奇怪的的是 CodeLLaMA 34B 中使用了 GQA，但并未乘以 1.3（雖然 LLaMA 2 34B 未開源，但 LLaMA 2 34B 與 CodeLLaMA 34B 一樣，也就是說 LLaMA 2 34B 中也未乘以 1.3）。如下圖所示，其中 d 為 8192，dFFN 為 22016，大概計算方式為（512 對齊）：
int(int(8192*8/3+512-1)/512)*512=22016
其實關(guān)于 dFFN 的數(shù)值在 LLaMA 3.1 早期的技術(shù)報告中也鬧過烏龍，如下圖 Table 3 所示，其直接將 dFFN 寫成為 3d/2：
如下圖 Table 3 所示，Meta 在 Arxiv 上的 Paper（[2407.21783] The Llama 3 Herd of Models [5]）中已經(jīng)修正了這個問題，比如對于 8B 模型，d 為 4096，dFFN 為 14336，大概計算方式為（512 對齊）：
int(int(4096*8/3*1.3+512-1)/512)*512=14336
2.2.3 LLaMA 3
經(jīng)過上述調(diào)整之后，LLaMA 3 不再是標(biāo)準(zhǔn)的 Transformer Block，此時使用 N=12*d*d 來預(yù)估參數(shù)量已經(jīng)不太準(zhǔn)確。但依舊可以將其按照（Wq，Wout）（Wk，Wv），WFFN 和 WEmb 4 個部分來統(tǒng)計。比如，對于 LLaMA 3 模型，我們可以按照下述方式估計其參數(shù)量：
N = nlayer *(2d*d + 2d*d*kv/h + 3d*dFFN)+2*Vocab*d
根據(jù)上述公式，可以推算出上述不同模型的參數(shù)量：
8B：32*(4096*4096*(2+2*8/32)+3*4086*14336)+2*128000*4096=8B
70B：80*(8192*8192*(2+2*8/64)+3*8192*28672)+2*128000*8192=70B
405B：126*(16384*16384*(2+2*8/128)+3*16384*53248)+2*128000*16384=405B
2.3 Dense LLM 計算量
其實整體的計算量（僅考慮 Forward 操作）也是集中在模型中的一些矩陣乘法，包括：
- MHA紅框：Wq，Wk，Wv 對應(yīng)的計算量都為 2 x (d x d x l)，其中 2 表示一個乘法和一個加法。
- MHA藍(lán)框：Wout 對應(yīng)的計算量為 2 x (d x d x l)。
- MHA Attention（綠色方塊）：2 x (l x d/h x l + l x d/h x l) x h = 4 x d x l x l。需要說明的是，如果是 Decoder（LLM），由于 Causal Mask 的存在，此處的計算量應(yīng)該減半，也就是 2 x d x l x l。
- FFN綠框：W1 和 W2 對應(yīng)的計算量為 2 x (dFFN x d x l) 和 2 x (d x dFFN x l)。LLaMA 的 SwiGLU 類似。?
也可參考 [2302.14017] Full Stack Optimization of Transformer Inference: a Survey 中的 Table 2，其中列出了相應(yīng)矩陣乘操作的 Dim，也就可以推出對應(yīng)的計算量：
綜上，在標(biāo)準(zhǔn)的 Transformer 模型或者 LLaMA 系列（MHA）中，如果忽略詞表、Embedding、LayerNorm 等模塊，總計算量為（所有 Transformer Block）：
C = nlayer *(cmha + cffn) = nlayer *2*(3d*d*l + d*d*l + d*l*l + 8d*d*l)
C = 24*nlayer*d*d*l + 2*nlayer*d*l*l = 2*N*l + 2*nlayer*d*l*l
由于大部分情況下預(yù)訓(xùn)練的序列長度 l 遠(yuǎn)小于 12d，比如：
- LLaMA-1 7B 的 d 為 4K，預(yù)訓(xùn)練的 l 為 2K。
- LLaMA-2 7B 的 d 為 4K，預(yù)訓(xùn)練的 l 為 4K。
- LLaMA-3 8B 的 d 為 4K，預(yù)訓(xùn)練的 l 為 8K。
同理，如果是 Encoder 模型，比如 Bert，總的計算量為：
C = nlayer *2*(3d*d*l + d*d*l + 2*d*l*l + 8d*d*l) = 2*N*l + 4*nlayer*d*l*l
PS：其實不管模型中是否使用 GQA，是否使用 SwiGLU 或者像 LLaMA 中那樣將 FFN Dim 再乘以 1.3，其計算量還都可以用 C=2*N*l + 2*nlayer*d*l*l 來近似。
在 Scaling Law 的原始 Paper（[2001.08361] Scaling Laws for Neural Language Models [6]）中也有類似的計算公式，如下圖 Table 1 所示，可以看出其推導(dǎo)結(jié)論與上述我們的計算方式也基本一致，只不過紅框中為每個 Token 的計算量：
假設(shè)數(shù)據(jù)集中總共包含 D 個 Token，對于序列不是特別長的場景，所有 Token Forward的計算量可以近似為 C=2*N*D。
2.4 LLM 訓(xùn)練計算量
模型的訓(xùn)練包含 Forward 和 Backward 過程，如下圖 Figure 1 所示，Backward 過程實際上包含兩部分，一部分是對輸入的梯度（鏈?zhǔn)椒▌t），一部分是對權(quán)重的梯度。其實這兩部分主要的計算量都是矩陣乘法，并且大小與 Forward中的大小一致，因此往往會直接近似 Backward 的計算量為 Forward 的 2 倍。
綜上，訓(xùn)練總的計算量（參數(shù)更新計算量比較小，這里先忽略）為：
C = Forward + Backward = 6*N*D。
需要說明的是，在 [2104.04473] Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM [7] 中 NVIDIA 使用了 Activation 重計算，所以需要多一次 Forward 操作，相應(yīng)總的計算量為：
C = 2*Forward + Backward = 8*N*D。
2.5 MoE LLM 計算量
MoE（Mixture of Experts）與標(biāo)準(zhǔn) Transformer 模型的不同在于：MoE 模型將 Transformer 模型的 FFN Block 替換為 MoE Block，其中 MoE Block 由多個 Expert 組成，每一個 Expert 都是標(biāo)準(zhǔn)的 FFN，實際計算時每個 Token 都會經(jīng) Router 選擇 1 個或多個 Expert，然后執(zhí)行類似 FFN 的計算。如下圖所示（圖片來自 A Visual Guide to Mixture of Experts (MoE) [8]）：
由于 Router 的計算量通常很小，常常可以忽略。如果不考慮 Router 引入的計算量，由于上述每個 Expert 實際上與非 MoE 模型的 FFN 一樣，因此也可以用同樣的方式推導(dǎo)出訓(xùn)練時每個 Token 的計算量依然為 C=6*N*D，只不過這里的 N 不再是整個模型的參數(shù)量，而是每個 Token 激活的參數(shù)量。
如下圖所示，論文 [2402.07871] Scaling Laws for Fine-Grained Mixture of Experts [9] 中，作者也推到過 MoE 中計算量的公式，為了普適性，作者考慮了細(xì)粒度專家的場景，具體來說：
- G（Granularity）：表示粒度，也就是將一個完整的 FFN 切分為多少個細(xì)粒度的專家，比如一個 FFN 可以切分為 8 個小專家，每個專家的參數(shù)量為原始 FFN 的 1/8。
- E（Expansion rate）：表示膨脹率，也就是所有專家的參數(shù)量擴(kuò)展為單個標(biāo)準(zhǔn) FFN 參數(shù)量的多少倍。比如，G=8，總共 64 個細(xì)粒度專家，則這里的 E=64/8=8，相當(dāng)于有 8 個同標(biāo)準(zhǔn)規(guī)模的 FFN。
- cf：表示除 Router 之外的計算量與參數(shù)量的比值，也就是我們上述得到的 6。
- cr：表示 Router 的計算量與參數(shù)量的比值，對于比較簡單的 Router，通常為 6-20。?
從上述介紹也可以看出，通常 E*G 遠(yuǎn)小于 dmodel，因此也就可以忽略 Router 處的計算量。
2.6 Hunyuan MoE 參數(shù)量
首先使用之前的公式來計算 Hunyuan MoE 的總參數(shù)量和激活參數(shù)量，還是按照（Wq，Wout）（Wk，Wv），WFFN 和 WEmb 4 個部分：
- Wq，Wout：依舊是 2 x d x d
- Wk，Wv：由于使用了 GQA + CLA，每層 80 個 Head，并且只有 8 個 Key/Value Head，而 CLA 每 2 層共享 KV，等效于每層的 WK 和 WV 為 d x d/ 20。
- WFFN：每層 17 個 Expert，但每個 Token 激活 2 個 Expert。相當(dāng)于 MoE Layer 的總參數(shù)量為 17 x 3 x d x dFFN，激活參數(shù)量為 2 x 3 x d x dFFN。
- WEmb：由于 tie_word_embeddings 為 True，因此共享 Embedding，參數(shù)量為 vocab_size x d。
根據(jù)上述 4 個部分，可以得出與論文一致的參數(shù)量（389B，52B）：
總參數(shù)量：
64*(2*6400*6400+2/20*6400*6400+3*17*6400*18304)+128512*6400=388.7B
激活參數(shù)量：
64*(2*6400*6400+2/20*6400*6400+3*2*6400*18304)+128512*6400=51.3B
2.7 Hunyuan MoE 計算量
我們重新回顧 Hunyuan MoE 中的計算量 C，如下圖 Table 1 所示為 Hunyuan MoE 的模型配置。如下圖所示，作者提到其序列長度可以達(dá)到 8K，32K，256K。
這里以最短的 8K 序列長度為例。由于我們之前計算過，每個 Token 在 Attention 處忽略的計算量為 2*nlayer*nctx*d，相應(yīng)可以得出 Attention 額外的計算量為：
2*64*8*1024*6400=6710886400=6.7*109 >> 2.3*108
可以看出，將 2.3x108D 理解為長序列時 Attention 的計算開銷是不合理的。
除此之外，我們之前還忽略了 Router 的計算量，那么如果將其看成 Router 的計算量呢？因為有 16 個 Specialized Expert，相當(dāng)于 EG=16。對于比較簡單的 Router，cr 為 6-20 之間，如下所示推導(dǎo)出 cr 為 35 也可以接受。
64*6400*16*cr=2.3*108 => cr =35
按照上述的推導(dǎo)，那么 9.59ND 更可能是除 Router 之外的計算量，為什么這里不再是 6ND 呢？作者特意提到了長序列場景，因此我們不再忽略額外的 Attention 計算，從頭推導(dǎo)一下對應(yīng)的 C：
- 上述激活參數(shù)對應(yīng)的計算量為：6 * 51.3B*D=6ND。
- 額外的 Attention 的計算量為：6*nctx*6400*nlayer*D4=3.59ND。
如果除了 Router 外總的計算量為 C=9.59ND，可以得出，相當(dāng)于序列長度為 73K：
nctx = 3.59N/6/6400/64 = 74938 = 73K
然而作者在論文中也提到，長文本訓(xùn)練并不需要特別多的 Token，在 32K 和 256K 階段也都只使用了 10B Token，遠(yuǎn)小于預(yù)訓(xùn)練階段的 7T Token。所以依然比較奇怪 9.59 是怎么來的。
2.8 Megatron-LM 計算量計算
如下圖所示為 Megatron-LM 中有關(guān)計算量的計算公式（megatron/training/training.py#L98-L151 [10]），需要說明的是，這里的 expansion_factor 為 12：
- 第一行的 3 表示 Forward + Backward。
- 第二行的 2 表示將 Attention Block 和 FFN Block 都分別看成兩組，比如 Attention 中 K 和 V 是一對，Q 和 Out 是一對，Attention 的 Q x K 和 Score x V 是一對。
- 第三行的 2 表示矩陣乘法中的乘和加，如果從單個 Token 單個矩陣乘的視角，計算量就是參數(shù)量的 2 倍。?
三、LLM 訓(xùn)練效率&訓(xùn)練預(yù)估
3.1 MFU & HFU
為了評估 LLM 訓(xùn)練時的效率，業(yè)界通常會使用 Model FLOPS Utilization（MFU） 和 Hardware FLOPS Utilization（HFU） 兩個關(guān)鍵指標(biāo)來評估模型的 Forward 和 Backward 傳播過程中（包括任何的網(wǎng)絡(luò)同步開銷和 DataLoader IO）硬件的利用率。
- MFU = 預(yù)估 FLOPS/硬件理論 FLOPS。其中，預(yù)估 FLOPS 就是模型訓(xùn)練時理論需要的計算量，并不包括各種優(yōu)化方案額外引入的計算量，比如 Gradient Checkpointing/Activation Recomputation 等引入的額外計算量。
- HFU = 實際 FLOPS/硬件理論 FLOPS。其中，實際 FLOPS 也就是理論上所有實際發(fā)生的計算量，包含 Gradient checkpointing/Activation Recompution 等引入的額外計算量，因此 HFU 應(yīng)該總是 >= MFU。
如下所示為 Maximizing training throughput using PyTorch FSDP [11] 中 Meta 在 LLM 訓(xùn)練時的 MFU 和 HFU。對于 LLM 訓(xùn)練任務(wù)而言，通常在 A100/A800 GPU 集群中，MFU 可以達(dá)到 50%+，甚至接近 60%；而在 H100/H800 GPU 集群中， MFU 往往不超過 50%。
3.2 175B 訓(xùn)練實際預(yù)估
有了以上的計算公式，根據(jù)每個 Token 計算量（6N）、總的 Token 數(shù)，以及計算資源就可以大概預(yù)估出訓(xùn)練的總時長。如下所示：
訓(xùn)練天數(shù) = Token 數(shù) * Ctoken / (GPU 數(shù) * GPU FLOPs * MFU * 3600 * 24)
根據(jù)以上公式可以預(yù)估使用 8192 H100-80G GPU（FLOPs 為 989T），10T Token 數(shù)據(jù)訓(xùn)練 175B 模型的天數(shù)為 30 天，其中假設(shè) MFU 為 50%：
10T*6*175B/(8192*989T*50%)/3600/24=30 天
當(dāng)然，上述是理想的估算，實際的訓(xùn)練中往往無法保證一直穩(wěn)定的訓(xùn)練，比如可能出現(xiàn)機(jī)器異常導(dǎo)致任務(wù)中斷、降速等，因此在上述理論預(yù)估的基礎(chǔ)上往往會留一定的 Buffer，比如 10% 的 Buffer 時長。
3.3 LLaMA 3 405B 訓(xùn)練實際預(yù)估
對于 LLaMA 3 405B 模型，使用 15T Token，16384 H100 GPU，MFU 為 41%，那么對應(yīng)的理想訓(xùn)練時間為：
15T*6*405B/(16384*989T*41%)/3600/24=42.3 天
然而，Meta 在 LLaMA 3 的技術(shù)報告（[2407.21783] The Llama 3 Herd of Models）中介紹 LLaMA 3 的預(yù)訓(xùn)練時間在 54 天左右，比 42.3 天多一些，可能得原因有幾個方面。
其中一個原因可能是其提到的多達(dá) 466 次的訓(xùn)練中斷，其中包括 47 次有計劃中斷，以及 419 次的預(yù)期外中斷。在非預(yù)期中斷中，78% 是硬件問題，例如 GPU 或物理機(jī)其他組件的異常，而 GPU 相關(guān)問題占到 58.7%。
另一個可能的原因是不同訓(xùn)練階段使用了不同規(guī)模的 GPU 和數(shù)據(jù)量。比如，如下圖 Table 4 所示涉及到 3 種不同的配置：
但上述配置并不能與作者介紹的 405B 模型預(yù)訓(xùn)練階段對應(yīng)。比如，作者提到，在訓(xùn)練的早期階段使用比較小的 Batch Size以提高訓(xùn)練穩(wěn)定性，然后會增加 Batch Size 以提高訓(xùn)練效率。具體來說：
- 初始時 4K 序列長度，一個 Batch 內(nèi) 4M Token，對應(yīng) Global Batch Size 為 1K。
- 訓(xùn)練 252M Token 后擴(kuò)展到 8K，共 8M Token，對應(yīng) Global Batch Size 為 1K。
- 訓(xùn)練到 2.87T Token 時進(jìn)一步 Double Batch Size，擴(kuò)展到 16M Token，對應(yīng) Global Batch Size 為 2K。
- 針對長序列訓(xùn)練：作者分為 6 個階段逐步增加上下文長度，將序列長度從最初的 8K擴(kuò)展 到最終的 128K，這個階段大約使用了 800B Token。
3.4 其他數(shù)據(jù)
其實很多 Paper 中也會簡單提到其訓(xùn)練模型規(guī)模、訓(xùn)練資源、時長、MFU 等指標(biāo)，但往往一筆帶過，如下所示為一些示例：
以上圖中 Nemotron-4 340B 的相應(yīng)配置可以預(yù)估出對應(yīng)的訓(xùn)練時長，可以看出與實際時長非常接近：
0.2T*6*340B/(1536*989T*42.4%)/3600/24=7.3 天
0.2T*6*340B/(3072*989T*42.3%)/3600/24=3.7 天
7.6T*6*340B/(6144*989T*41%)/3600/24=72 天
四、參考鏈接
- ??https://arxiv.org/abs/2411.02265??
- ??https://arxiv.org/abs/2302.14017??
- ??https://arxiv.org/abs/2302.13971??
- ??https://arxiv.org/abs/2307.09288??
- ??https://arxiv.org/abs/2407.21783??
- ??https://arxiv.org/abs/2001.08361??
- ??https://arxiv.org/pdf/2104.04473??
- ??https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mixture-of-experts??
- ??https://arxiv.org/abs/2402.07871??
- ??https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/training.py#L98-L151??
- ??https://pytorch.org/blog/maximizing-training/???
