大模型训练与推理优化1. GPU 时间计算基础1.1 GPU 计算中的基本概念hhidden size隐藏维度LTransformer Block 层数V词表大小参数量、FLOPs、显存和训练时间之间存在直接联系1.2 Transformer 模型参数计算1.2.1 Self-Attention 参数Q、K、V 投影矩阵Wq / Wk / Wv输出投影矩阵Wo总参数量4h² 4h当 hidden sizeh固定时多头数量通常不会改变总参数量1.2.2 MLP 参数经典结构h → 4h → h第一层参数4h² 4h第二层参数4h² h总参数量8h² 5h经典结构下MLP 参数量约为 Attention 的 2 倍1.2.3 LayerNorm 与单个 Transformer Block一个 LayerNorm 包含γ和β参数量为2h一个 Block 通常包含两个 LayerNorm共4hSelf-Attention4h² 4hMLP8h² 5hLayerNorm4h因此单个 Transformer Block12h² 13h当h较大时可以近似为12h²1.2.4 整个 Transformer 模型参数L个 Transformer BlockL × (12h² 13h)Embedding MatrixV × h因此总参数量可以近似表示为N ≈ Vh L(12h² 13h)进一步简化N ≈ Vh 12Lh²核心关系可以记成Attention ≈ 4h²MLP ≈ 8h²一个 Block ≈ 12h²因此 Transformer 参数量主要随着 hidden sizeh的平方增长。1.2.5 用12Lh²快速估算 LLaMA 参数量上一节得到一个 Transformer Block ≈ 12h²如果模型共有L个 Transformer Block那么主体参数量可以近似为N ≈ 12Lh²其中hhidden sizeLTransformer 层数NTransformer 主体参数量以 LLaMA 7B 为例h 4096L 32代入12 × 32 × 4096² ≈ 6.44B而 LLaMA 7B 的实际参数规模约为6.7B说明这个公式可以比较方便地估算模型大概属于多少 B。不同规模可以快速估算为模型hidden sizeh层数L12Lh²估算实际规模LLaMA 7B4096326.44B约 6.7BLLaMA 13B51204012.58B约 13.0BLLaMA 33B66566031.90B约 32.5BLLaMA 65B81928064.42B约 65.2B可以看到估算值与实际参数规模已经非常接近。1.2.6 为什么估算值和实际参数量不完全一致12Lh²主要估算 Transformer Block 的主体参数并不是完整模型的精确参数公式。实际模型还包括Token EmbeddingLM Head最后的 RMSNormSwiGLU 的具体中间维度Attention 的具体实现例如 LLaMA 7BEmbedding ≈ V × h当V ≈ 32000h 4096则Embedding ≈ 32000 × 4096 ≈ 131M因此12Lh²更适合用来快速判断模型规模而不是精确计算模型的每一个参数。1.2.7 参数规模的增长规律由N ≈ 12Lh²可以直接看出L × 2→ 参数量大约× 2h × 2→ 参数量大约× 4也就是说层数L对参数量是线性影响而 hidden sizeh对参数量是平方影响。核心关系可以记成一个 Block ≈ 12h²L 个 Block ≈ 12Lh²代入 h 和 L → 快速判断模型大概是几 B12Lh²是 Transformer 主体参数的快速估算公式不是 LLaMA 的精确参数公式。1.3 参数量与显存占用模型参数本身需要占用 GPU 显存FP32、FP16、BF16 每个参数占用空间不同训练阶段还需要保存 Gradient 和 Optimizer State参数量并不等于实际训练显存占用1.4 参数量与计算量1.4.1 FLOP 与 FLOPsFLOP一次浮点运算一次乘法约为1 FLOP一次加法约为1 FLOPFLOPs通常用来表示完成某个计算需要多少浮点运算GPU FLOPSGPU 每秒能够完成多少浮点运算因此训练时间可以先粗略理解为训练时间 ≈ 总 FLOPs / GPU 实际计算速度1.4.2 向量点积的计算量假设两个长度为m的向量x (x1, x2, ..., xm)y (y1, y2, ..., ym)它们做点积时需要m次乘法m - 1次加法因此FLOPs m (m - 1) 2m - 1当m很大时可以近似为向量点积 ≈ 2m FLOPs1.4.3 矩阵乘法的计算量假设Am × nBn × r那么C A × B结果矩阵C的大小为m × r其中每一个元素都是一次长度为n的向量点积一个元素 ≈ 2n FLOPs而C一共有m × r个元素因此矩阵乘法 FLOPs ≈ (m × r) × 2n整理得到FLOPs ≈ 2mnr这是后面计算 Transformer FLOPs 最重要的基础公式。1.4.4 参数量和 FLOPs 的区别以 Transformer 中的 Q 投影为例Q XWq假设Xs × hWqh × h其中sSequence Length / Token 数量hhidden sizeWq本身包含h × h h²个参数因此参数量 h²但是实际计算时(s × h) × (h × h)根据矩阵乘法公式FLOPs ≈ 2 × s × h × h所以Q 投影 FLOPs ≈ 2sh²可以看到参数量h²计算量2sh²多出来的s是因为同一组权重参数需要被s个 Token 重复使用。1.4.5 核心关系这一部分可以记住两条公式长度为 m 的向量点积 ≈ 2m FLOPs(m × n) × (n × r) 矩阵乘法 ≈ 2mnr FLOPs后面计算 Transformer 时Q、K、V 投影Attention输出投影WoMLP本质上都可以继续使用这个矩阵乘法公式。因此整体逻辑就变成参数量 → 矩阵乘法 → FLOPs → GPU FLOPS → 训练时间1.4.6 Self-Attention 的输入与 Q、K、V 计算量设输入X ∈ R^(b × s × h)其中bbatch sizessequence lengthhhidden sizeSelf-Attention 首先计算Q XWqK XWkV XWv其中Wq、Wk、Wv ∈ R^(h × h)以 Q 为例[b,s,h] × [h,h]根据矩阵乘法FLOPs ≈ 2mnr得到Q ≈ 2bsh² FLOPs因此 Q、K、V 三个投影一共QKV ≈ 6bsh² FLOPs1.4.7 Attention Score 与 V 的计算量Q、K 的形状都是[b,s,h]计算QK^T相当于[s,h] × [h,s]因此QK^T ≈ 2bs²h FLOPs这里出现s²是因为每一个 Token 都需要与其他 Token 计算相关性。所以标准 Self-Attention 对序列长度具有O(s²)的计算复杂度。经过 Softmax 后Attention 权重矩阵形状为A[b,s,s]再计算A × V即[s,s] × [s,h]因此Attention × V ≈ 2bs²h FLOPs严格来说Scaled Dot-Product Attention 的缩放因子是1 / √d_head不是1 / √h。1.4.8 输出投影与 Self-Attention 总 FLOPsAttention 的输出仍然为[b,s,h]最后还需要经过输出投影Wo ∈ R^(h × h)因此Wo ≈ 2bsh² FLOPs把 Self-Attention 的主要计算加起来Q、K、V6bsh²QK^T2bs²hAttention × V2bs²hWo2bsh²因此Self-Attention FLOPs ≈ 8bsh² 4bs²h通常这里暂时忽略 Softmax、缩放、Bias 和残差连接等较小计算。1.4.9 Self-Attention FLOPs 的核心规律Self-Attention 的计算量可以拆成两部分8bsh²来自Q、K、V、Wo也就是线性投影计算。而4bs²h来自QK^T Attention × V也就是 Token 之间的 Attention 计算。因此可以记成Self-Attention ≈ 8bsh² 4bs²h其中h²主要来自线性层s²主要来自 Token 与 Token 之间的 Attention所以h越大 → 线性层越贵s越长 → Attention 的平方项越贵这也为后面理解 Flash Attention、KV Cache 和长上下文优化做准备。1.4.10 MLP 的 FLOPs 计算经典 Transformer 的 MLP 一般采用h → 4h → h输入X[b,s,h]第一层[b,s,h] × [h,4h]根据矩阵乘法FLOPs ≈ 2mnr得到第一层 MLP ≈ 8bsh²第二层[b,s,4h] × [4h,h]同样得到第二层 MLP ≈ 8bsh²因此整个 MLPMLP FLOPs ≈ 16bsh²这里通常忽略 GELU、Bias、LayerNorm 和残差相加等较小计算。1.4.11 单个 Transformer Block 的 FLOPs上一节得到Self-Attention ≈ 8bsh² 4bs²h而MLP ≈ 16bsh²所以一个 Transformer BlockBlock FLOPs≈ (8bsh² 4bs²h) 16bsh²整理得到Block FLOPs ≈ 24bsh² 4bs²h其中24bsh²Q、K、V、Wo 和 MLP 等线性层计算4bs²hQK^T 和 Attention × V 的 Token-Token 计算因此一个 Block ≈ 24bsh² 4bs²h1.4.12 L 个 Transformer Block 的计算量如果模型共有L个 Transformer BlockL 个 Block FLOPs≈ L(24bsh² 4bs²h)这里bbatch sizessequence lengthhhidden sizeLTransformer Block 数量1.4.13 输出到词表的计算量最后一个 Transformer Block 输出[b,s,h]模型还需要通过 LM Head 将 hidden state 映射到整个词表[b,s,h] × [h,V]其中V vocabulary size因此LM Head FLOPs ≈ 2bshV所以整个模型一次前向传播的主要计算量可以写成Forward FLOPs ≈ L(24bsh² 4bs²h) 2bshV1.4.14 参数量与前向 FLOPs 的关系前面计算参数量时得到一个 Transformer Block ≈ 12h²而一个 Block 中主要线性层的前向计算量约为24bsh²可以写成24bsh² 2bs × 12h²也就是说前向 FLOPs ≈ 2 × Token 数量 × 参数量这里一个 batch 的 Token 数量为b × s因此可以粗略理解为每个参数在处理每个 Token 时通常会参与一次乘加运算而一次乘加约对应 2 FLOPs。这也是后面推导Training FLOPs ≈ 6ND的重要基础。1.4.15 这一部分的核心关系从前面的计算可以串成矩阵乘法 ≈ 2mnr↓Self-Attention ≈ 8bsh² 4bs²h↓MLP ≈ 16bsh²↓一个 Block ≈ 24bsh² 4bs²h↓L 个 Block ≈ L(24bsh² 4bs²h)↓整个模型前向 ≈ L(24bsh² 4bs²h) 2bshV下一步就可以继续从“前向传播 FLOPs”推到“训练 FLOPs”也就是经典的Training FLOPs ≈ 6ND1.4.16 从参数量推到前向 FLOPs前面已经得到一个 Transformer Block 参数量 ≈ 12h²如果模型共有L个 BlockN ≈ 12Lh²其中N模型参数量LTransformer Block 数量hhidden size对于一个 batchToken 数量 b × s前面又得到一个 Block 的主要线性计算约为24bsh²因此L个 BlockForward FLOPs ≈ 24Lbsh²注意24Lbsh² 2 × (12Lh²) × bs也就是Forward FLOPs ≈ 2 × 参数量 × Token 数量如果整个训练过程一共处理D个 TokenForward FLOPs ≈ 2ND其中N模型参数量D训练 Token 总数1.4.17 为什么前向传播是2ND可以把一个参数参与计算理解成一次乘加x × w → 1 次乘法结果累加 → 1 次加法因此1 次乘法 1 次加法 ≈ 2 FLOPs所以一个参数 × 一个 Token ≈ 2 FLOPs进一步得到N 个参数 × D 个 Token ≈ 2ND FLOPs因此Forward FLOPs ≈ 2ND1.4.18 为什么训练变成6ND训练不仅包含 Forward还需要 Backward。对于一个线性层Y XW前向传播主要计算Y XW可以看成1份矩阵乘法计算。反向传播主要还需要计算dX dY · W^T以及dW X^T · dY因此可以粗略理解成Forward ≈ 1 份Backward ≈ 2 份所以完整训练Forward Backward ≈ 3 × Forward而Forward ≈ 2ND因此Training FLOPs ≈ 3 × 2ND最终得到Training FLOPs ≈ 6ND1.4.196ND的核心直觉对于一个参数 × 一个 Token可以近似理解为Forward ≈ 2 FLOPsBackward输入梯度 ≈ 2 FLOPsBackward参数梯度 ≈ 2 FLOPs所以完整训练 ≈ 2 2 2 6 FLOPs最终可以记成推理 / Forward FLOPs ≈ 2ND训练 FLOPs ≈ 6ND这里的6ND是大模型训练计算量的快速估算公式并不是严格恒等式。它主要统计参数矩阵乘法带来的主要计算而像QK^TAttention × VSoftmax激活函数LayerNorm等额外计算没有被完整包含。因此它更适合用于快速估算大模型预训练需要多少计算量而不是精确计算每一次运算。1.5 从训练 FLOPs 到训练时间1.5.1 训练时间的基本公式前面已经得到训练总计算量Training FLOPs ≈ 6ND其中N模型参数量D训练 Token 总数课程中有时也会使用Training FLOPs ≈ 8ND作为更保守的经验估算把部分额外计算也考虑进去。训练时间本质上就是训练时间 总计算量 / 实际计算速度因此可以写成Time ≈ C / (G × P × U)其中C训练总 FLOPsGGPU 数量P单张 GPU 峰值 FLOP/sUGPU 有效利用率如果按照课程中的8NDTime ≈ 8ND / (G × P × U)1.5.2 GPU FLOPS 与实际计算速度GPU 标称的 TFLOPS 表示理论峰值计算速度。例如312 TFLOPS表示312 × 10^12 FLOP/s但真实训练时通常无法持续达到峰值因为还会受到GPU 间通信数据加载显存访问梯度同步非矩阵计算等因素影响。因此实际有效算力可以近似写成Effective FLOPS G × P × U这里的U更接近训练中的MFUModel FLOPs Utilization而不是简单的 GPU 使用率。1.5.3 一个 7B 模型的训练时间估算假设模型参数N 7B 7 × 10^9训练数据D 2T 2 × 10^12 TokenGPU 数量G 1000单卡峰值算力P 312 TFLOP/s有效利用率U 45%按照课程中的C ≈ 8ND得到C ≈ 8 × 7×10^9 × 2×10^12≈ 1.12 × 10^23 FLOPs实际总计算速度G × P × U 1000 × 312×10^12 × 0.45≈ 1.404 × 10^17 FLOP/s因此Time ≈ 1.12×10^23 / 1.404×10^17≈ 7.98 × 10^5 s换算后≈ 9.2 天所以在这些理想化假设下7B 2T Tokens 1000 张 GPU训练时间大约是9 天左右1.5.4 为什么增加 GPU 不一定线性加速理论上GPU × 2 → 计算速度 × 2 → 时间 ÷ 2但实际训练还存在All-Reduce 通信GPU 同步等待Pipeline Bubble网络带宽限制数据加载与 Checkpoint因此实际加速比 GPU 数量增长比例这也是后面需要学习 Distributed Training 的原因。1.5.5 常见单位换算1B 10^91T Token 10^12 Token1 TFLOP/s 10^12 FLOP/s因此做训练时间估算时一定先统一单位。1.5.6 从参数量到训练时间的完整链路整套逻辑可以记成参数量 N↓训练 Token 数 D↓训练计算量 ≈ 6ND或课程中使用 8ND↓实际算力 GPU 数 × 单卡峰值 FLOPS × 有效利用率↓训练时间 总 FLOPs / 实际 FLOPS也就是N、D 决定有多少活G、P、U 决定每秒能干多少活最后Time ≈ C / (G × P × U)1.5.7 训练时间实战175B 300B Tokens假设模型参数量N 175B 175 × 10^9训练 TokenD 300B 300 × 10^9GPU 数量G 1024单张 A100 峰值算力P 312 TFLOP/s有效利用率U 0.45课程继续采用Training FLOPs ≈ 8ND因此总训练计算量C ≈ 8 × 175×10^9 × 300×10^9得到C ≈ 4.2 × 10^23 FLOPs1.5.8 计算 GPU 集群的实际算力单张 A100 峰值算力312 TFLOP/s考虑 45% 的有效利用率312 × 0.45 ≈ 140.4 TFLOP/s1024 张 GPU1024 × 312×10^12 × 0.45得到≈ 1.44 × 10^17 FLOP/s因此整个 GPU 集群每秒实际能够完成约1.44 × 10^17次浮点运算。1.5.9 计算最终训练时间根据训练时间 总计算量 / 实际总算力代入Time ≈ 4.2×10^23 / 1.44×10^17得到≈ 2.92 × 10^6 秒换算成天2.92 × 10^6 / 86400 ≈ 33.8 天因此训练时间 ≈ 34 天也就是说在这些假设条件下175B 参数 300B Tokens 1024 张 A100大约需要34 天1.5.10 GPU 数量与利用率对训练时间的影响训练时间公式Time ≈ C / (G × P × U)因此GPU 数量G增加 → 训练时间下降单卡算力P增加 → 训练时间下降有效利用率U提高 → 训练时间下降例如利用率从45% → 90%在其他条件不变时理论训练时间大约可以减半34 天 → 17 天但现实中 GPU 数量增加后会带来更多通信梯度同步参数同步数据传输因此 GPU 翻倍并不一定意味着训练时间严格减半。1.5.116ND和8ND的区别前面推导得到的经典估算是Training FLOPs ≈ 6ND而课程这里使用Training FLOPs ≈ 8ND可以这样理解6ND主要估算参数矩阵乘法带来的训练计算量。而8ND属于课程中采用的更保守估算把部分额外计算开销也考虑进去。因此训练时间估算时必须先确认使用的是6ND还是8ND因为两者会直接影响最终结果。这一节最核心可以记成模型参数量 × Token 数 → 总工作量GPU 数 × 单卡算力 × 利用率 → 实际计算速度最终训练时间 总工作量 / 实际计算速度1.5.12 什么是 GPU HoursGPU Hours 表示GPU Hours GPU 数量 × 实际运行小时数例如1 张 GPU × 10 小时 10 GPU Hours10 张 GPU × 1 小时 10 GPU Hours因此 GPU Hours 更像是在描述总 GPU 资源消耗而不是人真正等待了多久。反过来在理想情况下实际训练时间 ≈ GPU Hours / GPU 数量1.5.13 GPU Hours 与实际训练时间假设某次训练总共消耗3,100,000 GPU Hours如果只使用1 张 GPU那么训练时间 3,100,000 小时换算成年3,100,000 / (24 × 365) ≈ 354 年也就是大约355 年如果使用更多 GPU100 张 GPU→ 理想约3.54 年1000 张 GPU→ 理想约129 天10000 张 GPU→ 理想约13 天所以GPU 数量 ↑ → Wall-clock Time ↓这也是为什么大模型训练必须依赖分布式训练。1.5.14 GPU 数量增加为什么不能完全线性加速理想情况下GPU × 2 → 训练时间 ÷ 2但现实中 GPU 越多还会增加GPU 间通信Gradient Synchronization参数同步网络传输Pipeline Bubble负载不均衡因此实际情况通常是实际训练时间 GPU Hours / GPU 数量也就是说GPU 数量增加以后并不能始终保持完美的线性加速。1.5.15 FLOPs、GPU FLOPS、GPU Hours 与训练时间这几个概念可以这样区分FLOPs表示整个训练任务总共需要做多少计算例如4.2 × 10^23 FLOPsGPU FLOPS表示GPU 每秒能够完成多少计算例如312 TFLOP/sGPU Hours表示总共消耗了多少 GPU × 小时例如3,100,000 GPU HoursWall-clock Training Time表示实际从训练开始到结束人真正等待了多久理想情况下Wall-clock Time ≈ GPU Hours / GPU 数量但现实中还需要考虑分布式效率、通信和利用率。1.5.16 大模型训练计算成本的完整链路现在可以把前面的内容完整串起来模型参数量 N训练 Token 数 D↓估算训练总 FLOPs↓结合 GPU 峰值算力与利用率↓估算 GPU Hours↓根据 GPU 数量↓估算 Wall-clock Training Time可以简单记成FLOPs 总共有多少活GPU FLOPS 每秒能干多少活GPU Hours 一共用了多少 GPU 资源Wall-clock Time 实际等了多久因此大模型训练真正解决的问题就是把巨大的总计算量通过大量 GPU 和分布式训练尽可能高效地并行完成。
大模型训练与推理优化（一）
大模型训练与推理优化1. GPU 时间计算基础1.1 GPU 计算中的基本概念hhidden size隐藏维度LTransformer Block 层数V词表大小参数量、FLOPs、显存和训练时间之间存在直接联系1.2 Transformer 模型参数计算1.2.1 Self-Attention 参数Q、K、V 投影矩阵Wq / Wk / Wv输出投影矩阵Wo总参数量4h² 4h当 hidden sizeh固定时多头数量通常不会改变总参数量1.2.2 MLP 参数经典结构h → 4h → h第一层参数4h² 4h第二层参数4h² h总参数量8h² 5h经典结构下MLP 参数量约为 Attention 的 2 倍1.2.3 LayerNorm 与单个 Transformer Block一个 LayerNorm 包含γ和β参数量为2h一个 Block 通常包含两个 LayerNorm共4hSelf-Attention4h² 4hMLP8h² 5hLayerNorm4h因此单个 Transformer Block12h² 13h当h较大时可以近似为12h²1.2.4 整个 Transformer 模型参数L个 Transformer BlockL × (12h² 13h)Embedding MatrixV × h因此总参数量可以近似表示为N ≈ Vh L(12h² 13h)进一步简化N ≈ Vh 12Lh²核心关系可以记成Attention ≈ 4h²MLP ≈ 8h²一个 Block ≈ 12h²因此 Transformer 参数量主要随着 hidden sizeh的平方增长。1.2.5 用12Lh²快速估算 LLaMA 参数量上一节得到一个 Transformer Block ≈ 12h²如果模型共有L个 Transformer Block那么主体参数量可以近似为N ≈ 12Lh²其中hhidden sizeLTransformer 层数NTransformer 主体参数量以 LLaMA 7B 为例h 4096L 32代入12 × 32 × 4096² ≈ 6.44B而 LLaMA 7B 的实际参数规模约为6.7B说明这个公式可以比较方便地估算模型大概属于多少 B。不同规模可以快速估算为模型hidden sizeh层数L12Lh²估算实际规模LLaMA 7B4096326.44B约 6.7BLLaMA 13B51204012.58B约 13.0BLLaMA 33B66566031.90B约 32.5BLLaMA 65B81928064.42B约 65.2B可以看到估算值与实际参数规模已经非常接近。1.2.6 为什么估算值和实际参数量不完全一致12Lh²主要估算 Transformer Block 的主体参数并不是完整模型的精确参数公式。实际模型还包括Token EmbeddingLM Head最后的 RMSNormSwiGLU 的具体中间维度Attention 的具体实现例如 LLaMA 7BEmbedding ≈ V × h当V ≈ 32000h 4096则Embedding ≈ 32000 × 4096 ≈ 131M因此12Lh²更适合用来快速判断模型规模而不是精确计算模型的每一个参数。1.2.7 参数规模的增长规律由N ≈ 12Lh²可以直接看出L × 2→ 参数量大约× 2h × 2→ 参数量大约× 4也就是说层数L对参数量是线性影响而 hidden sizeh对参数量是平方影响。核心关系可以记成一个 Block ≈ 12h²L 个 Block ≈ 12Lh²代入 h 和 L → 快速判断模型大概是几 B12Lh²是 Transformer 主体参数的快速估算公式不是 LLaMA 的精确参数公式。1.3 参数量与显存占用模型参数本身需要占用 GPU 显存FP32、FP16、BF16 每个参数占用空间不同训练阶段还需要保存 Gradient 和 Optimizer State参数量并不等于实际训练显存占用1.4 参数量与计算量1.4.1 FLOP 与 FLOPsFLOP一次浮点运算一次乘法约为1 FLOP一次加法约为1 FLOPFLOPs通常用来表示完成某个计算需要多少浮点运算GPU FLOPSGPU 每秒能够完成多少浮点运算因此训练时间可以先粗略理解为训练时间 ≈ 总 FLOPs / GPU 实际计算速度1.4.2 向量点积的计算量假设两个长度为m的向量x (x1, x2, ..., xm)y (y1, y2, ..., ym)它们做点积时需要m次乘法m - 1次加法因此FLOPs m (m - 1) 2m - 1当m很大时可以近似为向量点积 ≈ 2m FLOPs1.4.3 矩阵乘法的计算量假设Am × nBn × r那么C A × B结果矩阵C的大小为m × r其中每一个元素都是一次长度为n的向量点积一个元素 ≈ 2n FLOPs而C一共有m × r个元素因此矩阵乘法 FLOPs ≈ (m × r) × 2n整理得到FLOPs ≈ 2mnr这是后面计算 Transformer FLOPs 最重要的基础公式。1.4.4 参数量和 FLOPs 的区别以 Transformer 中的 Q 投影为例Q XWq假设Xs × hWqh × h其中sSequence Length / Token 数量hhidden sizeWq本身包含h × h h²个参数因此参数量 h²但是实际计算时(s × h) × (h × h)根据矩阵乘法公式FLOPs ≈ 2 × s × h × h所以Q 投影 FLOPs ≈ 2sh²可以看到参数量h²计算量2sh²多出来的s是因为同一组权重参数需要被s个 Token 重复使用。1.4.5 核心关系这一部分可以记住两条公式长度为 m 的向量点积 ≈ 2m FLOPs(m × n) × (n × r) 矩阵乘法 ≈ 2mnr FLOPs后面计算 Transformer 时Q、K、V 投影Attention输出投影WoMLP本质上都可以继续使用这个矩阵乘法公式。因此整体逻辑就变成参数量 → 矩阵乘法 → FLOPs → GPU FLOPS → 训练时间1.4.6 Self-Attention 的输入与 Q、K、V 计算量设输入X ∈ R^(b × s × h)其中bbatch sizessequence lengthhhidden sizeSelf-Attention 首先计算Q XWqK XWkV XWv其中Wq、Wk、Wv ∈ R^(h × h)以 Q 为例[b,s,h] × [h,h]根据矩阵乘法FLOPs ≈ 2mnr得到Q ≈ 2bsh² FLOPs因此 Q、K、V 三个投影一共QKV ≈ 6bsh² FLOPs1.4.7 Attention Score 与 V 的计算量Q、K 的形状都是[b,s,h]计算QK^T相当于[s,h] × [h,s]因此QK^T ≈ 2bs²h FLOPs这里出现s²是因为每一个 Token 都需要与其他 Token 计算相关性。所以标准 Self-Attention 对序列长度具有O(s²)的计算复杂度。经过 Softmax 后Attention 权重矩阵形状为A[b,s,s]再计算A × V即[s,s] × [s,h]因此Attention × V ≈ 2bs²h FLOPs严格来说Scaled Dot-Product Attention 的缩放因子是1 / √d_head不是1 / √h。1.4.8 输出投影与 Self-Attention 总 FLOPsAttention 的输出仍然为[b,s,h]最后还需要经过输出投影Wo ∈ R^(h × h)因此Wo ≈ 2bsh² FLOPs把 Self-Attention 的主要计算加起来Q、K、V6bsh²QK^T2bs²hAttention × V2bs²hWo2bsh²因此Self-Attention FLOPs ≈ 8bsh² 4bs²h通常这里暂时忽略 Softmax、缩放、Bias 和残差连接等较小计算。1.4.9 Self-Attention FLOPs 的核心规律Self-Attention 的计算量可以拆成两部分8bsh²来自Q、K、V、Wo也就是线性投影计算。而4bs²h来自QK^T Attention × V也就是 Token 之间的 Attention 计算。因此可以记成Self-Attention ≈ 8bsh² 4bs²h其中h²主要来自线性层s²主要来自 Token 与 Token 之间的 Attention所以h越大 → 线性层越贵s越长 → Attention 的平方项越贵这也为后面理解 Flash Attention、KV Cache 和长上下文优化做准备。1.4.10 MLP 的 FLOPs 计算经典 Transformer 的 MLP 一般采用h → 4h → h输入X[b,s,h]第一层[b,s,h] × [h,4h]根据矩阵乘法FLOPs ≈ 2mnr得到第一层 MLP ≈ 8bsh²第二层[b,s,4h] × [4h,h]同样得到第二层 MLP ≈ 8bsh²因此整个 MLPMLP FLOPs ≈ 16bsh²这里通常忽略 GELU、Bias、LayerNorm 和残差相加等较小计算。1.4.11 单个 Transformer Block 的 FLOPs上一节得到Self-Attention ≈ 8bsh² 4bs²h而MLP ≈ 16bsh²所以一个 Transformer BlockBlock FLOPs≈ (8bsh² 4bs²h) 16bsh²整理得到Block FLOPs ≈ 24bsh² 4bs²h其中24bsh²Q、K、V、Wo 和 MLP 等线性层计算4bs²hQK^T 和 Attention × V 的 Token-Token 计算因此一个 Block ≈ 24bsh² 4bs²h1.4.12 L 个 Transformer Block 的计算量如果模型共有L个 Transformer BlockL 个 Block FLOPs≈ L(24bsh² 4bs²h)这里bbatch sizessequence lengthhhidden sizeLTransformer Block 数量1.4.13 输出到词表的计算量最后一个 Transformer Block 输出[b,s,h]模型还需要通过 LM Head 将 hidden state 映射到整个词表[b,s,h] × [h,V]其中V vocabulary size因此LM Head FLOPs ≈ 2bshV所以整个模型一次前向传播的主要计算量可以写成Forward FLOPs ≈ L(24bsh² 4bs²h) 2bshV1.4.14 参数量与前向 FLOPs 的关系前面计算参数量时得到一个 Transformer Block ≈ 12h²而一个 Block 中主要线性层的前向计算量约为24bsh²可以写成24bsh² 2bs × 12h²也就是说前向 FLOPs ≈ 2 × Token 数量 × 参数量这里一个 batch 的 Token 数量为b × s因此可以粗略理解为每个参数在处理每个 Token 时通常会参与一次乘加运算而一次乘加约对应 2 FLOPs。这也是后面推导Training FLOPs ≈ 6ND的重要基础。1.4.15 这一部分的核心关系从前面的计算可以串成矩阵乘法 ≈ 2mnr↓Self-Attention ≈ 8bsh² 4bs²h↓MLP ≈ 16bsh²↓一个 Block ≈ 24bsh² 4bs²h↓L 个 Block ≈ L(24bsh² 4bs²h)↓整个模型前向 ≈ L(24bsh² 4bs²h) 2bshV下一步就可以继续从“前向传播 FLOPs”推到“训练 FLOPs”也就是经典的Training FLOPs ≈ 6ND1.4.16 从参数量推到前向 FLOPs前面已经得到一个 Transformer Block 参数量 ≈ 12h²如果模型共有L个 BlockN ≈ 12Lh²其中N模型参数量LTransformer Block 数量hhidden size对于一个 batchToken 数量 b × s前面又得到一个 Block 的主要线性计算约为24bsh²因此L个 BlockForward FLOPs ≈ 24Lbsh²注意24Lbsh² 2 × (12Lh²) × bs也就是Forward FLOPs ≈ 2 × 参数量 × Token 数量如果整个训练过程一共处理D个 TokenForward FLOPs ≈ 2ND其中N模型参数量D训练 Token 总数1.4.17 为什么前向传播是2ND可以把一个参数参与计算理解成一次乘加x × w → 1 次乘法结果累加 → 1 次加法因此1 次乘法 1 次加法 ≈ 2 FLOPs所以一个参数 × 一个 Token ≈ 2 FLOPs进一步得到N 个参数 × D 个 Token ≈ 2ND FLOPs因此Forward FLOPs ≈ 2ND1.4.18 为什么训练变成6ND训练不仅包含 Forward还需要 Backward。对于一个线性层Y XW前向传播主要计算Y XW可以看成1份矩阵乘法计算。反向传播主要还需要计算dX dY · W^T以及dW X^T · dY因此可以粗略理解成Forward ≈ 1 份Backward ≈ 2 份所以完整训练Forward Backward ≈ 3 × Forward而Forward ≈ 2ND因此Training FLOPs ≈ 3 × 2ND最终得到Training FLOPs ≈ 6ND1.4.196ND的核心直觉对于一个参数 × 一个 Token可以近似理解为Forward ≈ 2 FLOPsBackward输入梯度 ≈ 2 FLOPsBackward参数梯度 ≈ 2 FLOPs所以完整训练 ≈ 2 2 2 6 FLOPs最终可以记成推理 / Forward FLOPs ≈ 2ND训练 FLOPs ≈ 6ND这里的6ND是大模型训练计算量的快速估算公式并不是严格恒等式。它主要统计参数矩阵乘法带来的主要计算而像QK^TAttention × VSoftmax激活函数LayerNorm等额外计算没有被完整包含。因此它更适合用于快速估算大模型预训练需要多少计算量而不是精确计算每一次运算。1.5 从训练 FLOPs 到训练时间1.5.1 训练时间的基本公式前面已经得到训练总计算量Training FLOPs ≈ 6ND其中N模型参数量D训练 Token 总数课程中有时也会使用Training FLOPs ≈ 8ND作为更保守的经验估算把部分额外计算也考虑进去。训练时间本质上就是训练时间 总计算量 / 实际计算速度因此可以写成Time ≈ C / (G × P × U)其中C训练总 FLOPsGGPU 数量P单张 GPU 峰值 FLOP/sUGPU 有效利用率如果按照课程中的8NDTime ≈ 8ND / (G × P × U)1.5.2 GPU FLOPS 与实际计算速度GPU 标称的 TFLOPS 表示理论峰值计算速度。例如312 TFLOPS表示312 × 10^12 FLOP/s但真实训练时通常无法持续达到峰值因为还会受到GPU 间通信数据加载显存访问梯度同步非矩阵计算等因素影响。因此实际有效算力可以近似写成Effective FLOPS G × P × U这里的U更接近训练中的MFUModel FLOPs Utilization而不是简单的 GPU 使用率。1.5.3 一个 7B 模型的训练时间估算假设模型参数N 7B 7 × 10^9训练数据D 2T 2 × 10^12 TokenGPU 数量G 1000单卡峰值算力P 312 TFLOP/s有效利用率U 45%按照课程中的C ≈ 8ND得到C ≈ 8 × 7×10^9 × 2×10^12≈ 1.12 × 10^23 FLOPs实际总计算速度G × P × U 1000 × 312×10^12 × 0.45≈ 1.404 × 10^17 FLOP/s因此Time ≈ 1.12×10^23 / 1.404×10^17≈ 7.98 × 10^5 s换算后≈ 9.2 天所以在这些理想化假设下7B 2T Tokens 1000 张 GPU训练时间大约是9 天左右1.5.4 为什么增加 GPU 不一定线性加速理论上GPU × 2 → 计算速度 × 2 → 时间 ÷ 2但实际训练还存在All-Reduce 通信GPU 同步等待Pipeline Bubble网络带宽限制数据加载与 Checkpoint因此实际加速比 GPU 数量增长比例这也是后面需要学习 Distributed Training 的原因。1.5.5 常见单位换算1B 10^91T Token 10^12 Token1 TFLOP/s 10^12 FLOP/s因此做训练时间估算时一定先统一单位。1.5.6 从参数量到训练时间的完整链路整套逻辑可以记成参数量 N↓训练 Token 数 D↓训练计算量 ≈ 6ND或课程中使用 8ND↓实际算力 GPU 数 × 单卡峰值 FLOPS × 有效利用率↓训练时间 总 FLOPs / 实际 FLOPS也就是N、D 决定有多少活G、P、U 决定每秒能干多少活最后Time ≈ C / (G × P × U)1.5.7 训练时间实战175B 300B Tokens假设模型参数量N 175B 175 × 10^9训练 TokenD 300B 300 × 10^9GPU 数量G 1024单张 A100 峰值算力P 312 TFLOP/s有效利用率U 0.45课程继续采用Training FLOPs ≈ 8ND因此总训练计算量C ≈ 8 × 175×10^9 × 300×10^9得到C ≈ 4.2 × 10^23 FLOPs1.5.8 计算 GPU 集群的实际算力单张 A100 峰值算力312 TFLOP/s考虑 45% 的有效利用率312 × 0.45 ≈ 140.4 TFLOP/s1024 张 GPU1024 × 312×10^12 × 0.45得到≈ 1.44 × 10^17 FLOP/s因此整个 GPU 集群每秒实际能够完成约1.44 × 10^17次浮点运算。1.5.9 计算最终训练时间根据训练时间 总计算量 / 实际总算力代入Time ≈ 4.2×10^23 / 1.44×10^17得到≈ 2.92 × 10^6 秒换算成天2.92 × 10^6 / 86400 ≈ 33.8 天因此训练时间 ≈ 34 天也就是说在这些假设条件下175B 参数 300B Tokens 1024 张 A100大约需要34 天1.5.10 GPU 数量与利用率对训练时间的影响训练时间公式Time ≈ C / (G × P × U)因此GPU 数量G增加 → 训练时间下降单卡算力P增加 → 训练时间下降有效利用率U提高 → 训练时间下降例如利用率从45% → 90%在其他条件不变时理论训练时间大约可以减半34 天 → 17 天但现实中 GPU 数量增加后会带来更多通信梯度同步参数同步数据传输因此 GPU 翻倍并不一定意味着训练时间严格减半。1.5.116ND和8ND的区别前面推导得到的经典估算是Training FLOPs ≈ 6ND而课程这里使用Training FLOPs ≈ 8ND可以这样理解6ND主要估算参数矩阵乘法带来的训练计算量。而8ND属于课程中采用的更保守估算把部分额外计算开销也考虑进去。因此训练时间估算时必须先确认使用的是6ND还是8ND因为两者会直接影响最终结果。这一节最核心可以记成模型参数量 × Token 数 → 总工作量GPU 数 × 单卡算力 × 利用率 → 实际计算速度最终训练时间 总工作量 / 实际计算速度1.5.12 什么是 GPU HoursGPU Hours 表示GPU Hours GPU 数量 × 实际运行小时数例如1 张 GPU × 10 小时 10 GPU Hours10 张 GPU × 1 小时 10 GPU Hours因此 GPU Hours 更像是在描述总 GPU 资源消耗而不是人真正等待了多久。反过来在理想情况下实际训练时间 ≈ GPU Hours / GPU 数量1.5.13 GPU Hours 与实际训练时间假设某次训练总共消耗3,100,000 GPU Hours如果只使用1 张 GPU那么训练时间 3,100,000 小时换算成年3,100,000 / (24 × 365) ≈ 354 年也就是大约355 年如果使用更多 GPU100 张 GPU→ 理想约3.54 年1000 张 GPU→ 理想约129 天10000 张 GPU→ 理想约13 天所以GPU 数量 ↑ → Wall-clock Time ↓这也是为什么大模型训练必须依赖分布式训练。1.5.14 GPU 数量增加为什么不能完全线性加速理想情况下GPU × 2 → 训练时间 ÷ 2但现实中 GPU 越多还会增加GPU 间通信Gradient Synchronization参数同步网络传输Pipeline Bubble负载不均衡因此实际情况通常是实际训练时间 GPU Hours / GPU 数量也就是说GPU 数量增加以后并不能始终保持完美的线性加速。1.5.15 FLOPs、GPU FLOPS、GPU Hours 与训练时间这几个概念可以这样区分FLOPs表示整个训练任务总共需要做多少计算例如4.2 × 10^23 FLOPsGPU FLOPS表示GPU 每秒能够完成多少计算例如312 TFLOP/sGPU Hours表示总共消耗了多少 GPU × 小时例如3,100,000 GPU HoursWall-clock Training Time表示实际从训练开始到结束人真正等待了多久理想情况下Wall-clock Time ≈ GPU Hours / GPU 数量但现实中还需要考虑分布式效率、通信和利用率。1.5.16 大模型训练计算成本的完整链路现在可以把前面的内容完整串起来模型参数量 N训练 Token 数 D↓估算训练总 FLOPs↓结合 GPU 峰值算力与利用率↓估算 GPU Hours↓根据 GPU 数量↓估算 Wall-clock Training Time可以简单记成FLOPs 总共有多少活GPU FLOPS 每秒能干多少活GPU Hours 一共用了多少 GPU 资源Wall-clock Time 实际等了多久因此大模型训练真正解决的问题就是把巨大的总计算量通过大量 GPU 和分布式训练尽可能高效地并行完成。
