【建议收藏】一文讲透大模型核心资源消耗：参数量、计算量、中间激活、KV cache

最新推荐文章于 2025-11-24 19:11:26 发布

原创最新推荐文章于 2025-11-24 19:11:26 发布 · 934 阅读

28 ·

CC 4.0 BY-SA版权

文章标签：

#transformer #ai #深度学习 #人工智能 #程序员 #转行 #大模型

本文重点参考了文章分析 transformer 模型的参数量、计算量、中间激活、KV cache，同时加入了自己一些理解。

01 前提假设

GPT 系列模型及大多数 LLM 使用的基座模型都是 Transformer Decoder，这里以 GPT1 模型为代表来分析大模型。

假设字典大小为 V，查询向量的维度为 d，隐藏层向量大小为 d，Transformer层的层数为 N，FFN 层隐藏层向量大小为 4d，模型输入为 batch size（B）、sequence length（S）。

在这里插入图片描述

GPT1 模型

02 模型参数量

（1）参数量是模型本身可以经过反向传播更新的模型参数数据量大小，模型参数量大小和模型输入无关。

（2）Transformer 从前往后依次需要经历 embedding、positional embedding、masked multi-head attention、residual and layer normlization、FFN、linear 层。

计算分别如下：

embedding 层：Vd
positional embedding：无可训练的参数
masked multi-head attention：4dd+d(bias)。看代码 WQ、WK 和 WV 计算时无 bias，WO 计算时包含 bias；隐藏层维度和查询向量维度相同
两个 residual and layer normlization：4d
FFN：8dd+5d(bias)
末尾 linear 层：使用 embedding 层的转置，无新增参数

考虑层数 N，那么模型参数量为 Vd+N*(12dd+10d)。当 d 数值较大且远大于 V 时，此时参数量近似为 12Ndd。

03 前向传递计算量

（1）最开始理解计算量=输入*模型参数，后来了解到这一点并不完全对，如 self-attention 计算过程中没有模型参数，但确实产生了计算量。

（2）补充 3 点矩阵计算背景知识：

可以进行计算前提条件：第一个矩阵列数等于第二个矩阵行数；
计算过程：第一个矩阵行数依次分别乘第二个矩阵列数；
计算量：对于 1*N 和 N*1 的矩阵计算，计算量为 N 次加法和 N 次乘法，共 2N。

（3）这里从输入依次从前向后总结计算量：这里忽略了计算量小头-包括 layer normalization 等。

输入计算得到 QKV：6BSdd
QKT 计算：[B，num_heads，S，d_head] * [B，num_heads，d_head，S] = 2BSSd
SV 计算：[B，num_heads，S，S] * [B，num_heads，S，d_head] =2BSSd
PO 计算：[B，S，d] * [d，d] = 2BSdd
FFN 第一层：[B，S，d] * [d，4d] = 8BSdd
FFN 第二层：[B，S，4d] * [4d，d] = 8BSdd
隐藏层映射到词表：[B，S，d] * [d，V] = 2BSdV

考虑层数 N，那么计算量为 2BSdV+N*(24BSdd+4BSSd)。当 d 数值较大且远大于 V、S 时，此时参数两近似为 24NBSdd。

04 反向传播计算量

（1）反向传播的核心是利用梯度下降法求得梯度，然后优化器根据梯度更新模型参数。

（2）大模型的核心是矩阵计算，矩阵计算的反向传播需要对两个输入 X 和 W 分别求导，要经过两次矩阵计算。

以下述 self-attention 计算为例，对于 Q=X WQ 计算，前向传播只需计算一次，反向传播则需要对 X 和 WQ 分别求导，需要计算 2 次。因此，反向传播的计算量近似是前向传递计算量的 2 倍。

在这里插入图片描述

前向传递与对应反向传播

05模型参数量和计算量的关系

（1）训练神经网络的一次迭代分为三步：前向传递计算损失函数；反向传播计算梯度；优化器根据梯度更新模型参数。在实验中发现优化器更新的耗时几乎可以忽略。

（2）在一次迭代中，对于每个 token，参数量和计算量关系分别如下。

在前向传递中，上面提到参数量近似为 12Ndd，计算量近似为 24NBSdd，所以计算量/参数量=2BS。

所以，在一次前向传递中，对于每个 token，每个模型参数需要进行 2 次浮点数运算，即一次乘法法运算和一次加法运算。

在反向传播中，根据上述计算量分析，每个参数需要进行 4 次浮点数计算。

所以，如果在不考虑激活重计算情况下，对于每一个 token，需要进行 6 次浮点数计算。

如果考虑 flash attention 提出的激活重计算，需要进行一次额外的前向传递，那么在一次训练中，对于每一个 token，需要进行 8 次浮点数计算。

06 混合精度训练

（1）为什么在训练时使用混合精度训练的策略？

如果使用 float32 进行全参数训练，对硬件显存和存储要求极高；

如果使用 float16 进行全参数训练，则会出现浮点数精度限制导致的溢出风险，造成精度损失。

因此，大家一般会使用混合精度训练，即使用 float16 的模型参数进行前向传递和基于反向传播计算得到 float16 的梯度；

在优化器更新模型参数阶段，会使用 float32 的优化器状态、float32 的梯度、float32 的模型参数来更新模型参数。

（2）在混合精度训练中，具体的计算流程如下：

第一，float16 模型参数前向传递；
第二，反向传播得到 float16 的梯度；
第三，float16 的梯度转换成 float32 的梯度，与 float32 的优化器状态，更新 float32 模型参数；
第四，float32 的模型参数转换为 float16 的模型参数。

混合精度训练流程

07 模型训练显存

（1）背景知识：float16 数据类型的元素占 2 个 bytes，float32 数据类型的元素占 4 个 bytes。大模型在训练时通常使用混合精度训练策略，通常使用 AdamW 优化器。

（2）计算公式：模型训练显存 = 模型参数显存 + 梯度显存 + 优化器显存 + 中间激活显存

假如模型参数为 1，那么分别计算如下：

模型参数显存：float16 + float32 = 6
梯度显存：float32 = 4（假设float16梯度在经过优化器转换为float32精度时，会自动释放掉显存）
优化器显存：float32 + float32 = 8（对应优化器的一阶动量和二阶动量）
中间激活显存：下一节会详细介绍。flash attention、梯度检查点、混合精度训练等技术都可以降低中间激活显存占用，一般情况下中间激活显存占全部显存的 10%-30%。

（3）如果模型参数量为 N（float16），那么训练阶段的显存约为 18N bytes（不考虑中间激活）。