MoE HunYuan 计算量和 LLM 参数量 MFU 等 聊一聊
一、背景
最近在看腾讯最新混元大模型的 Paper 时([2411.02265] Hunyuan-Large: An Open-Source MoE Model with 52 Billion Activated Parameters by Tencent [1]),看到了如下关于计算 Budget 的公式
由于咱们的上班中也会经常依据计算资源 Budget 评价 LLM 预训练期间,而该公式与咱们平时的计算模式不太分歧;此外,如下图所示,咱们也看到很多文章中将上述公式中的第二项了解为长序列状况下 Attention 的额外计算开支,而将 6ND -> 9.6ND 看成 Router 引入的开支,与咱们的计算论断也十分不分歧,因此本文中咱们便捷回忆一下相关内容。
本文中的很多内容咱们在之前的文章中曾经引见过,详细可以参考:
二、LLM 参数量&计算量
2.1 Dense LLM 参数量
规范 Transformer 模型不论是 Encoder Only 的 Bert 系列模型,还是 Decoder Only 的 GPT 系列模型,同性能下参数量和计算量都是相似的。其中的一个关键点是:规范 Transformer Block 输入、输入以及两边 Attention Out 的 Hidden Dim 是坚持不变的,也就是不时是 Token Embedding 的 Hidden Dim。一切的 Transformer Block 都十分规整。
如下图 Figure 1 Encoder 所示(图片来自 [2302.14017] Full Stack Optimization of Transformer Inference: a Survey [2]),重要参数都来自几个矩阵乘的 Weight 矩阵,其中 d 示意 Token Embedding 的 Hidden Dim,l 示意 Token 数,h 示意 MHA 中的 Head 个数,dFFN 示意 FFN 层两边升维后的 Dim。详细来说:
综上,在规范的 Transformer 模型或许 LLaMA 系列(MHA)中,假设疏忽词表、Embedding、LayerNorm 等参数后,总参数量为(一切 Transformer Block):
N = nlayer *(nmha + nffn) = nlayer *(3d*d + d*d + 8d*d) = 12*nlayer*d*d
2.2 LLaMA LLM 参数量
2.2.1 SwiGLU
在 LLaMA 等模型中在 FFN 中会经常使用 SwiGLU 激活,这也就造成其会额外多了一个权重矩阵,如下图所示:
在 LLaMA-1 的 Paper([2302.13971] LLaMA: Open and Efficient Foundation Language Models [3])中作者提到,经常使用 SwiGLU 后将 dFFN 从 4d 降落到了 8d/3。这样 3 个权重矩阵的参数量还是 8d,总的参数量依然可以经常使用 12*nlayer*d*d 预估。
之前公式对应的是 MHA(Multi Head Attention),这也是 LLaMA-1 系列模型的规范成功。不过,LLaMA-2 的 30B 和 70B 模型以及 LLaMA-3 的所有模型都开局经常使用 GQA(Grouped Query Attention)。经常使用 GQA 时,多个 Attention Head 会共享一个 Key 和 Value,比如经常出现的性能是每 4 或 8 个 Head 共享相反的 Key 和 Value,此时 Wk,Wv 的大小会变为 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 的总参数量坚持不变,关于 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
其实关于 dFFN 的数值在 LLaMA 3.1 早期的技术报告中也闹过乌龙,如下图 Table 3 所示,其间接将 dFFN 写成为 3d/2:
如下图 Table 3 所示,Meta 在 Arxiv 上的 Paper([2407.21783] The Llama 3 Herd of Models [5])中曾经批改了这个疑问,比如关于 8B 模型,d 为 4096,dFFN 为 14336,大略计算模式为(512 对齐):
int(int(4096*8/3*1.3+512-1)/512)*512=14336
2.2.3 LLaMA 3
通过上述调整之后,LLaMA 3 不再是规范的 Transformer Block,此时经常使用 N=12*d*d 来预估参数量曾经不太准确。但照旧可以将其依照(Wq,Wout)(Wk,Wv),WFFN 和 WEmb 4 个局部来统计。比如,关于 LLaMA 3 模型,咱们可以依照下述模式预计其参数量:
N = nlayer *(2d*d + 2d*d*kv/h + 3d*dFFN)+2*Vocab*d
依据上述公式,可以推算出上述不同模型的参数量:
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 操作)也是集中在模型中的一些矩阵乘法,包括:
也可参考 [2302.14017] Full Stack Optimization of Transformer Inference: a Survey 中的 Table 2,其中列出了相应矩阵乘操作的 Dim,也就可以推出对应的计算量:
综上,在规范的 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
由于大局部状况下预训练的序列长度 l 远小于 12d,比如:
同理,假设是 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 所示,可以看出其推导论断与上述咱们的计算模式也基本分歧,只不过红框中为每个 Token 的计算量:
假定数据集中总共蕴含 D 个 Token,关于序列不是特地长的场景,一切 Token Forward的计算量可以近似为 C=2*N*D。
2.4 LLM 训练计算量
模型的训练蕴含 Forward 和 Backward 环节,如下图 Figure 1 所示,Backward 环节实践上蕴含两局部,一局部是对输入的梯度(链式规律),一局部是对权重的梯度。其实这两局部重要的计算量都是矩阵乘法,并且大小与 Forward中的大小分歧,因此往往会直凑近似 Backward 的计算量为 Forward 的 2 倍。
综上,训练总的计算量(参数降级计算量比拟小,这里先疏忽)为:
C = Forward + Backward = 6*N*D。
须要说明的是,在 [2104.04473] Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM [7] 中 NVIDIA 经常使用了 Activation 重计算,所以须要多一次性 Forward 操作,相应总的计算量为:
C = 2*Forward + Backward = 8*N*D。
2.5 MoE LLM 计算量
MoE(Mixture of Experts)与规范 Transformer 模型的不同在于:MoE 模型将 Transformer 模型的 FFN Block 交流为 MoE Block,其中 MoE Block 由多个 Expert 组成,每一个 Expert 都是规范的 FFN,实践计算时每个 Token 都会经 Router 选用 1 个或多个 Expert,而后口头相似 FFN 的计算。如下图所示(图片来自 A Visual Guide to Mixture of Experts (MoE) [8]):
由于 Router 的计算量通常很小,经常可以疏忽。假设不思考 Router 引入的计算量,由于上述每个 Expert 实践上与非 MoE 模型的 FFN 一样,因此也可以用雷同的模式推导出训练时每个 Token 的计算量依然为 C=6*N*D,只不过这里的 N 不再是整个模型的参数量,而是每个 Token 激活的参数量。
如下图所示,论文 [2402.07871] Scaling Laws for Fine-Grained Mixture of Experts [9] 中,作者也推到过 MoE 中计算量的公式,为了普适性,作者思考了细粒度专家的场景,详细来说:
从上述引见也可以看出,通常 E*G 远小于 dmodel,因此也就可以疏忽 Router 处的计算量。
2.6 Hunyuan MoE 参数量
首先经常使用之前的公式来计算 Hunyuan MoE 的总参数量和激活参数量,还是依照(Wq,Wout)(Wk,Wv),WFFN 和 WEmb 4 个局部:
依据上述 4 个局部,可以得出与论文分歧的参数量(389B,52B):
总参数量:
64*(2*6*6+2/20*6*6+3*17*6*18304)+128512*6=388.7B
激活参数量:
64*(2*6*6+2/20*6*6+3*2*6*18304)+128512*6=51.3B
2.7 Hunyuan MoE 计算量
咱们从新回忆 Hunyuan MoE 中的计算量 C,如下图 Table 1 所示为 Hunyuan MoE 的模型性能。如下图所示,作者提到其序列长度可以到达 8K,32K,256K。
这里以最短的 8K 序列长度为例。由于咱们之前计算过,每个 Token 在 Attention 处疏忽的计算量为 2*nlayer*nctx*d,相应可以得出 Attention 额外的计算量为:
2*64*8*1024*6=6710886=6.7*109 >> 2.3*108
可以看出,将 2.3x108D 了解为长序列时 Attention 的计算开支是不正当的。
除此之外,咱们之前还疏忽了 Router 的计算量,那么假设将其看成 Router 的计算量呢?由于有 16 个 Specialized Expert,相当于 EG=16。关于比拟便捷的 Router,cr 为 6-20 之间,如下所示推导出 cr 为 35 也可以接受。
64*6*16*cr=2.3*108 => cr =35
依照上述的推导,那么 9.59ND 更或许是除 Router 之外的计算量,为什么这里不再是 6ND 呢?作者特地提到了长序列场景,因此咱们不再疏忽额外的 Attention 计算,从头推导一下对应的 C:
假设除了 Router 外总的计算量为 C=9.59ND,可以得出,相当于序列长度为 73K:
nctx = 3.59N/6/6/64 = 74938 = 73K
但是作者在论文中也提到,长文本训练并不须要特地多的 Token,在 32K 和 256K 阶段也都只经常使用了 10B Token,远小于预训练阶段的 7T Token。所以依然比拟奇异 9.59 是怎样来的。
2.8 Megatron-LM 计算量计算
如下图所示为 Megatron-LM 中无关计算量的计算公式(megatron/training/training.py#L98-L151 [10]),须要说明的是,这里的 expansion_factor 为 12:
三、LLM 训练效率&训练预估
3.1 MFU & HFU
为了评价 LLM 训练时的效率,业界通常会经常使用 Model FLOPS Utilization(MFU) 和 Hardware FLOPS Utilization(HFU) 两个关键目的来评价模型的 Forward 和 Backward 流传环节中(包括任何的网络同步开支和>如下所示为 Maximizing training throughput using PyTorch FSDP [11] 中 Meta 在 LLM 训练时的 MFU 和 HFU。关于 LLM 训练义务而言,通常在 A100/A800 GPU 集群中,MFU 可以到达 50%+,甚至凑近 60%;而在 H100/H800 GPU 集群中, MFU 往往不超越 50%。
3.2 175B 训练实践预估
有了以上的计算公式,依据每个 Token 计算量(6N)、总的 Token 数,以及计算资源就可以大略预估出训练的总时长。如下所示:
训练天数 = Token 数 * Ctoken / (GPU 数 * GPU FLOPs * MFU * 3600 * 24)
依据以上公式可以预估经常使用 8192 H100-80G GPU(FLOPs 为 989T),10T Token 数据训练 175B 模型的天数为 30 天,其中假定 MFU 为 50%:
10T*6*175B/(8192*989T*50%)/3600/24=30 天
当然,上述是现实的预算,实践的训练中往往不可保障不时稳固的训练,比如或许出现机器意外造成义务终止、降速等,因此在上述实践预估的基础上往往会留必定的 Buffer,比如 10% 的 Buffer 时长。
3.3 LLaMA 3 405B 训练实践预估
关于 LLaMA 3 405B 模型,经常使用 15T Token,16384 H100 GPU,MFU 为 41%,那么对应的现实训练期间为:
15T*6*405B/(16384*989T*41%)/3600/24=42.3 天
但是,Meta 在 LLaMA 3 的技术报告([2407.21783] The Llama 3 Herd of Models)中引见 LLaMA 3 的预训练期间在 54 天左右,比 42.3 天多一些,或许得要素有几个方面。
其中一个要素或许是其提到的多达 466 次的训练终止,其中包括 47 次有方案终止,以及 419 次的预期外终止。在非预期终止中,78% 是配件疑问,例如 GPU 或物理机其余组件的意外,而 GPU 相关疑问占到 58.7%。
另一个或许的要素是不同训练阶段经常使用了不同规模的 GPU 和数据量。比如,如下图 Table 4 所示触及到 3 种不同的性能:
但上述性能并不能与作者引见的 405B 模型预训练阶段对应。比如,作者提到,在训练的早期阶段经常使用比拟小的 Batch Size以提高训练稳固性,而后会参与 Batch Size 以提高训练效率。详细来说:
3.4 其余数据
其实很多 Paper 中也会便捷提到其训练模型规模、训练资源、时长、MFU 等目的,但往往一笔带过,如下所示为一些示例:
以上图中 Nemotron-4 340B 的相应性能可以预估出对应的训练时长,可以看出与实践时长十分凑近:
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 天
四、参考链接
本文转载自,作者: