一、背景
最近在看腾讯最新混元大模型的 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 天
四、参考链接
本文转载自，作者：