1.4 资源核算:参数量、显存与计算量#
训练或部署一个模型之前,需要先定量回答三个问题,模型有多少参数、占用多少显存、消耗多少算力(FLOPs)。这三本账决定了单卡能否装下、训练一遍需时多久、推理为何受限。本节是斯坦福 CS336 第 2 讲的核心主题[^cs336],也是 第 3 章 · 扩展律 与 第 9 章 · 训推系统 的共同基础。
核心问题
单层参数≈$12d^2$、总参数≈$12Ld^2$ 中的系数12如何由 $Q,K,V,O$ 与 FFN 的各权重矩阵累加得到;
训练算力 $C\approx 6ND$ 中的系数6如何拆为前向 2 与反向 4,推理为何每 token 仅需 $2N$;
训练显存16 字节/参数的经典分解(fp16 参数 + fp32 主权重 + Adam 的 $m,v$)从何而来,为何优化器状态是主要瓶颈;
激活显存随 batch 与序列长度如何增长,梯度检查点如何以重算换取显存;
推理显存 = 参数 + KV-Cache,KV-Cache 的大小公式与算例,以及长上下文下它为何可能超过权重本身。
参数量决定模型自身的规模;算力(FLOPs)度量完成一次前向或训练所需的浮点运算总量,直接对应训练时长与硬件成本;显存则约束了同一时刻能容纳的参数、梯度、优化器状态与激活的总量。三者相互牵制:参数越多,所需算力越大,占用的显存也越高。本节的目标是为这三者各建立一条估算公式,使读者拿到任一模型的结构配置即可推断其规模、训练时长与硬件需求。
需要先界定本节范围:这里只处理静态核算,给定模型结构与硬件,推导参数量、算力、显存三个数值。
至于给定算力预算,如何在模型规模与数据量之间分配属于 3.1 · 幂律 与 3.2 · Chinchilla 计算最优 的范畴;
单卡装不下时如何切分到多卡则由 9.6 · 训练期并行 讨论。本节负责把这三本账的底层公式推导清楚,为后续章节提供依据。
一、参数量:主要开销在何处#
问题与整体结构#
一个 Transformer 由 $L$ 个结构相同的层堆叠而成。要估算总参数量,只需算清单层参数再乘以层数,最后补上嵌入层与输出层。单层参数几乎全部集中在两处:注意力的四个投影矩阵与前馈网络(Feed-Forward Network, FFN)的两个矩阵;LayerNorm 的增益与各处偏置只占极小比例。
逐个矩阵拆解#
设隐藏维度为 $d$(部分文献记作 $E$ 或 $D$,含义相同)。逐块拆解单层如下。
注意力块。每个 token 是一个 $d$ 维向量,需分别投影为 Query、Key、Value:
$$W_Q,\;W_K,\;W_V \in \mathbb{R}^{d\times d}\quad\Rightarrow\quad 3d^2$$多头注意力只是沿维度将这 $d$ 维切分给各头分别计算,不改变矩阵总大小($h$ 个头,每头 $d/h$ 维,拼接后仍为 $d$)。注意力计算之后还有一个输出投影 $W_O\in\mathbb{R}^{d\times d}$,再加 $d^2$。四个矩阵合计:
$$\underbrace{d^2}_{Q}+\underbrace{d^2}_{K}+\underbrace{d^2}_{V}+\underbrace{d^2}_{O}=4d^2$$FFN 块。标准 Transformer 的 FFN 先升维再降维,中间宽度取 $4d$(这一 4 倍膨胀自 GPT/BERT 以来成为惯例):
$$W_1\in\mathbb{R}^{d\times 4d}\;(=4d^2),\qquad W_2\in\mathbb{R}^{4d\times d}\;(=4d^2)\quad\Rightarrow\quad 8d^2$$此处需澄清一个常见误解:多头注意力并不增加参数量。它并非为每个头各配一套 $d\times d$ 的矩阵,而是把同一套 $W_Q,W_K,W_V$ 的输出沿维度切成 $h$ 段,每段 $d/h$ 维交给一个头独立计算注意力,最后拼回 $d$ 维。以 $d=512$、$h=8$ 为例,每头维度 $d/h=64$,八个头的投影合计仍是同一个 $512\times 512$ 矩阵,总参数不变。
多头带来的是不同头关注不同子空间的表达力,而非参数量的增加,这一区分在阅读论文与核对参数时都很重要。
单层合计为 $4d^2+8d^2=12d^2$。乘以层数 $L$,即得到常见的估计式:
$$\boxed{\,N\approx 12\,L\,d^2\,}$$单层参数由六个权重矩阵累加为 $12d^2$:注意力块的 $W_Q,W_K,W_V,W_O$ 各为 $d\times d$、合 $4d^2$,FFN 的 $W_1,W_2$ 各 $4d^2$、合 $8d^2$,其中 FFN 占 $\tfrac{8}{12}=\tfrac{2}{3}$。
为什么 FFN 占主要部分#
注意力占单层 $4d^2$、FFN 占 $8d^2$,后者是前者的 2 倍,占注意力 + FFN之和的 $\tfrac{8}{12}=\tfrac{2}{3}$。也就是说,每层约三分之二的参数与乘加运算发生在 FFN 中。JAX 的《Scaling Book》进一步指出,只要序列长度 $T\lt 8d$,MLP(即 FFN)的参数与 FLOPs 就主导整个前向1。
这正是 1.6 · 混合专家 MoE 选择在 FFN 上做文章的根本动机:若要在不成比例增加计算的前提下堆叠参数,就应针对占算力比重最大的部分。
沿此逻辑可进一步推得 MoE 的核心算术。既然 FFN 占每层约三分之二的算力,而目标又是参数量大、但单 token 实际计算量不随之激增,则自然的做法是把 FFN 复制为多份专家,每个 token 由路由器选择其中一两份激活。
如此总参数可堆叠至十余倍,而单 token 的 FLOPs 仅按被激活的专家数增长,这即是 MoE稀疏激活的实质。反之,若 FFN 只占单层约一成算力,针对它的改造收益便有限,MoE 也不会成为现今的主流形态。把参数量的分布核算清楚,是理解此类结构创新的前提。
具体数值:GPT-2 Small 的精确参数#
近似式 $12Ld^2$ 会遗漏偏置与嵌入。若要精确核算 GPT-2 Small($d=768$、$L=12$、词表 $V=50257$、上下文 $P=1024$),完整公式为2:
$$N = \underbrace{d(V+P)}_{\text{token+位置嵌入}} + L\cdot(\underbrace{12d^2+13d}_{\text{每层}}) + \underbrace{2d}_{\text{末层 LN}}$$其中每层可拆得更细:注意力的 c_attn 为 $3d^2+3d$(含偏置)、c_proj 为 $d^2+d$,合 $4d^2+4d$;FFN 的两个矩阵合 $8d^2+5d$;相加恰为 $12d^2+13d$。代入数字:嵌入项为 $768\times(50257+1024)=39{,}383{,}808$;
主体项为 $12\times(12\times768^2+13\times768)=12\times7{,}087{,}872=85{,}054{,}464$;末层项为 $2\times768=1{,}536$。三者相加为 $124{,}439{,}808$,即 GPT-2 Small1.24 亿参数的来源。
可见 $12Ld^2\approx 85\text{M}$ 的近似仅对应主体,嵌入层在小模型中占据近三分之一;但随模型增大($d$ 上千、$L$ 上百),主体项以 $d^2$ 增长,而嵌入项 $2Vd$ 仅线性增长,主体迅速占据主导,近似随之愈发精确。
嵌入是否计入参数:输入嵌入表 $V\times d$ 与输出(unembedding)$d\times V$ 各占 $Vd$。许多模型(含 GPT-2)采用权重共享(tied embeddings),两者为同一张表,只计一次 $Vd$。
在 GPT-3 175B 这类大模型中嵌入占比不足 1%,因此扩展律讨论算力时常将 $N$ 近似为非嵌入参数,细节见 3.1 · 幂律。
用 12Ld² 反推常见模型#
将 $N\approx 12Ld^2$ 作为估算工具,可对任一开源模型快速估重。方法是查出其隐藏维 $d$ 与层数 $L$ 代入即得大致参数量(若计嵌入还需补 $2Vd$,但对中大模型影响很小)。下表并列几款代表性配置,读者可自行验证,并从中体会 $d^2$ 这一平方项的作用:
| 模型档位 | 隐藏维 $d$ | 层数 $L$ | $12Ld^2$ 估算 | 典型场景 |
|---|---|---|---|---|
| 小(GPT-2 Small 量级) | 768 | 12 | 约 0.085B | 单卡即可运行 |
| 中(GPT-2 XL 量级) | 1600 | 48 | 约 1.47B | 单张消费卡可推理 |
| 7B 档 | 4096 | 32 | 约 6.4B | 开源主力,单卡可 LoRA 微调 |
| 70B 档 | 8192 | 80 | 约 64B | 需多卡,量化后可单机推理 |
由此表可得三点观察。其一,宽度比深度的成本更高:$d$ 出现在平方项中而 $L$ 仅出现在一次项中,因此将 $d$ 翻倍带来约 4 倍参数,将 $L$ 翻倍仅带来约 2 倍参数。这解释了为何现代大模型多为宽而不过深,单纯堆叠层数性价比较低,且深网络还伴随梯度传播与训练稳定性的额外困难。
其二,估算值与官方标称通常相差百分之几到百分之十几,差异源于被省略的嵌入、偏置、LayerNorm,以及不同实现的 FFN 膨胀倍数并非严格 4 倍(LLaMA 系采用 SwiGLU,中间宽度约 $\tfrac{8}{3}d$ 且需乘 3 个矩阵,须另行核算)。
其三,该工具的价值在于数量级而非精确位数:它使读者仅凭一篇论文或一个 config.json 即可判断模型规模与单卡可容性,无需下载权重逐一清点。
SwiGLU 与现代 FFN:LLaMA 的核算差异#
上述 $8d^2$ 对应经典两矩阵 FFN($W_1$ 升维、$W_2$ 降维,中间为 GELU/ReLU)的核算。
但 2023 年之后的主流开源模型(LLaMA、Qwen、Mistral 等)多改用门控线性单元(Gated Linear Unit)家族的 SwiGLU,它包含三个权重矩阵,门控投影 $W_{\text{gate}}$、上投影 $W_{\text{up}}$、下投影 $W_{\text{down}}$。
为使总参数量与经典 FFN 大致持平,这些模型将中间宽度从 $4d$ 收窄至约 $\tfrac{2}{3}\times 4d=\tfrac{8}{3}d$。如此三矩阵各约 $\tfrac{8}{3}d^2$,合计仍约 $8d^2$,单层依旧约 $12d^2$,因此 $12Ld^2$ 对 SwiGLU 模型同样适用,只是三分之二参数在 FFN的结论应修正为门控 FFN 仍是单层参数量最大的部分。
因此,精确核对 LLaMA 参数量时不应套用 $4d^2+8d^2$,而须按其实际的三矩阵结构与实际中间宽度逐项计算。
二、算力:训练 $C\approx 6ND$、推理 $\approx 2N$/token#
每个参数都要参与一次运算#
模型的计算几乎全部为矩阵乘法。一个 token 从输入流至输出,会与每一个权重参数发生一次乘、再累加一次的运算。因此处理一个 token 所需的算力正比于参数量。训练比推理昂贵,是因为训练在前向之外还需反向传播以计算梯度。将这两点定量化,即得到本节最重要的两个数值。
系数2与系数6的由来#
一个 $(m\times k)$ 矩阵乘 $(k\times n)$ 矩阵,输出含 $m\times n$ 个元素,每个元素是长度 $k$ 的点积,$k$ 次乘加即 $k$ 次乘加共 $2k$ 次浮点运算。故总量为:
$$\text{FLOPs} = 2\,m\,k\,n$$CS336 的 lecture_02.py 将其写作 2 * B * D * K(一乘一加)3。
前向为 $2N$/token。一个 token 流过所有权重矩阵,将上式对所有层、所有矩阵求和,乘加总量约等于$2\times$ 参数量。每个参数参与一次乘与一次加,$N$ 个参数即 $2N$ 次运算。
反向为 $4N$/token。反向传播需对每个矩阵乘 $Y=WX$ 计算两组梯度:对权重 $W$ 的梯度,以及对输入 $X$ 的梯度(用于将误差继续前传)。这两组各是一次与前向同规模的矩阵乘,故反向约为前向的 2 倍,即 $4N$。CS336 明确记为 forward = 2 * tokens * params、backward = 4 * tokens * params3。
训练总量。前向 2 加反向 4 为 6,对全部 $D$ 个训练 token:
$$\boxed{\,C\approx 6\,N\,D\,}$$推理只含前向。生成时没有反向,每个 token 约 $2N$ FLOPs。
$6ND$ 中为何不含注意力自身的算力:注意力中的 $QK^\top$ 与 $\text{softmax}\cdot V$ 两步矩阵乘并不经过权重参数。其 FLOPs 正比于 $\text{seq}^2$,在短序列、大模型时相对于权重矩阵乘(正比于 $N$)可忽略,故 $6ND$ 将其省略。
但当序列长度增至与 $d$ 同量级甚至更长时,注意力的平方项便会反超、不可再忽略,这既是长上下文训练同时消耗显存与算力的原因,也是 FlashAttention 等工作专攻注意力两步的背景。日常估算中,$6ND$ 是权重相关算力,已够用;精确核算长上下文时须补回注意力的 $\text{seq}^2$ 项。
这两条账的用途:有了 $C\approx 6ND$,可将算力预算换算为可训练的模型规模与数据量,这正是 3.2 · Chinchilla约 20 token/参数结论的推导起点。
有了推理 $\approx 2N$/token,则可理解 9.1 · 推理为什么慢、贵 中所述推理访存受限而非算力受限的原因:每生成一个 token 的算力很小,瓶颈在于将数十 GB 权重反复搬运进出。
前向、反向与训练/推理总量的拆分关系见图 1。
具体数值:GPT-3 训练的算力#
GPT-3($N\approx 175\text{B}$)训练约 $D\approx 3\times 10^{11}$ token,代入 $6ND$:
$$C\approx 6\times 1.75\times 10^{11}\times 3\times 10^{11}\approx 3.14\times 10^{23}\ \text{FLOPs}$$CS336 给出的正是这一量级4。换算到硬件:A100 的 bf16 峰值算力约 312 TFLOP/s $=3.12\times 10^{14}$ FLOP/s。若能 100% 利用,单卡跑完需 $\tfrac{3.14\times 10^{23}}{3.12\times 10^{14}}\approx 10^9$ 秒,约 32 年,因此必须数千张卡并行。
将并行规模与利用率一并代入,即可估算实际训练时长。假设采用 1024 张 A100、MFU 取现实的 40%,则集群有效算力约为 $1024\times 3.12\times 10^{14}\times 0.4\approx 1.28\times 10^{17}$ FLOP/s。
总算力除以有效算力得 $\tfrac{3.14\times 10^{23}}{1.28\times 10^{17}}\approx 2.45\times 10^6$ 秒,约合 28 天。这即是大模型训练动辄一个月的算术依据,它是 $6ND$ 除以卡数 × 峰值 × MFU的直接结果。
反之,该式也是制定训练预算的有效工具:给定模型规模 $N$、数据量 $D$ 与卡数,即可估出墙钟时间与能耗量级,避免开训后才发现无法在预期内完成。
关键指标:MFU(模型 FLOPs 利用率)#
真实训练无法达到峰值算力。MFU 定义为有效 FLOPs÷峰值 FLOPs $\times$ 用时,度量硬件被利用的充分程度:
$$\text{MFU} = \frac{6\,N\,D}{\text{峰值 FLOP/s}\times \text{GPU 数}\times \text{墙钟时间}}$$此处有一个需区分的近似概念:MFU 分子采用模型本身必须完成的 $6ND$;若开启梯度检查点,硬件还会额外执行重算的那部分 FLOPs,将硬件实际执行的全部 FLOPs置于分子,得到的是另一指标 HFU(硬件 FLOPs 利用率),其数值高于 MFU。性能对比时须辨明报告的是何者:HFU 高仅表明硬件未闲置,MFU 高才表明有效工作占比高。
日常判断管线优劣一般以 MFU 为准。
其余算力被通信、访存与流水线气泡消耗。业界公开报告中,训练 MFU 大多落在 35%–55% 区间[^mfu]:早期 GPT-3 约 20% 出头,此后 PaLM、Llama-3、字节跳动 MegaScale 等通过更好的计算通信重叠与调度将其推高至 40%–55%。
通常将 35%–45% 视作合格、50% 以上视作优秀(这些数值随硬件代际与模型结构浮动,只需把握其量级)。因此,拿到一份训练报告可先据此估算 MFU,以判断其管线尚有多少优化空间,若算得仅个位数百分比,多半是通信未重叠好,或 batch 过小、算子过碎。
推理侧的算力:prefill 与 decode 的差异#
推理 $\approx 2N$/token中隐含一个重要的二分。一次带上下文的生成分为两个阶段,算力特征迥异:
- prefill 预填充:将用户输入的整段 prompt 一次性送入,为其中每个 token 计算并缓存 K/V。设 prompt 含 $P$ 个 token,这一步约做 $2NP$ FLOPs,且这 $P$ 个 token 可并行置入一个大矩阵乘,因此 prefill 是算力受限(compute-bound)的,能充分利用 GPU。
- decode 逐 token 生成:此后每输出一个新 token,只处理这一个 token,约 $2N$ FLOPs,却需将整套数十 GB 权重从显存搬入计算单元一遍。算力小而搬运多,因此 decode 是访存受限(memory-bound)的,GPU 算力大量闲置。
两阶段在并行度、FLOPs 规模与瓶颈类型上均存在显著差异。
prefill 计算受限、decode 访存受限的分裂正是推理系统各项优化的出发点:连续批处理(continuous batching)将多个请求的 decode 步合并,摊薄权重搬运成本;投机解码(speculative decoding)用小模型先行预测若干 token、由大模型一次性验证,将多个 decode 步折叠为一次前向。
其成因与带宽优化方法是 9.1 · 推理为什么慢、贵 的核心议题。本节只需记住:训练看 FLOPs,推理 decode 看显存带宽,两者的瓶颈并不在同一处。
三、训练显存:优化器状态是主要瓶颈#
参数只是训练显存的一部分#
一个常见的误解是参数占多大空间、显存就占多少。实际上,训练时每个参数还伴随其梯度、优化器中的两个动量,以及一份高精度主权重。在混合精度加 Adam 的配置下,这些附加量之和是参数本身的数倍。
16 字节/参数的经典分解#
现代训练采用混合精度:前向/反向以 fp16(或 bf16,2 字节)计算以提升速度,但为保数值稳定,优化器中维护一份 fp32(4 字节)的主权重。Adam/AdamW 还需为每个参数存储两个统计量:一阶动量 $m$ 与二阶动量 $v$。将每个参数的开销列全(这正是 ZeRO 论文的分解5):
| 项目 | 精度 | 字节/参数 |
|---|---|---|
| 参数副本(前反向用) | fp16 | 2 |
| 梯度 | fp16 | 2 |
| 主权重(优化器用) | fp32 | 4 |
| Adam 一阶动量 $m$ | fp32 | 4 |
| Adam 二阶动量 $v$ | fp32 | 4 |
| 合计 | 16 字节 / 参数 |
其中后三项(主权重 + $m$ + $v$ = 12 字节)合称优化器状态,占 $\tfrac{12}{16}=75\%$。
其中最易被忽略、也最反直觉的是那份 fp32 主权重。既然前向反向均用 fp16 已足够快且省,为何还需额外保留一份 4 字节的高精度副本?原因在于 fp16 的表示精度不足:当学习率乘梯度得到的更新量极小(大模型训练后期,单步更新常在 $10^{-5}$ 量级)时,将如此小的增量加到 fp16 权重上会因舍入而丢失,参数不再变化、训练停滞。
解决方法是在 fp32 中累加更新、保留完整精度的主权重,每步再将其降精度为 fp16 供计算使用。因此这份 fp32 主权重并非冗余,而是混合精度训练能够收敛的前提。同理,Adam 的两个动量也采用 fp32,因为它们需长期累积、对精度更敏感,若以 fp16 存储会使优化过程失真。
$$\text{训练模型状态} \approx 16N\ \text{字节}$$每参数 16 字节在五个项目间的分配比例见图 2,其中优化器状态(后三项)独占 $\tfrac{12}{16}=75\%$。
优化器状态为何是瓶颈:ZeRO 的切分#
优化器状态($12N$)比参数本身($2N$)大 6 倍,是训练显存的绝对主体。微软 ZeRO 的核心观察是:数据并行时每张卡各存一份完整优化器状态,构成巨大冗余。ZeRO 将模型状态按三个阶段逐级切分到各卡5:Stage 1 切分优化器状态(节省最多,故最先做);Stage 2 进一步切分梯度;Stage 3 连参数也切分(各卡仅在用到时临时聚合)。
切分带来的节省很直接:设有 $G$ 张卡做数据并行,Stage 1 将那 $12N$ 的优化器状态平摊到 $G$ 张卡,每卡只留 $12N/G$;Stage 2 连梯度也摊、每卡再省一块;Stage 3 将参数本身也摊,理论上每卡的模型状态趋于 $16N/G$。
代价是通信量增加,参数被切分后,前向反向需临时从其他卡聚合所需分片(all-gather)、计算后再释放,因此阶段越高省显存越多、但通信越重。工程上常见的取舍是:显存充裕则停在 Stage 1/2(通信较轻),显存紧张才启用 Stage 3。
这也解释了同为数据并行,ZeRO 何以能训练比普通 DDP 大得多的模型,它消除了每卡各存一份完整状态的冗余。
其与张量/流水线并行的配合详见 9.6 · 训练期并行;此处只需记住,正是优化器状态是瓶颈这一显存账,直接催生了 ZeRO 的分层切分策略。
激活显存:随 batch 与序列长度膨胀#
除模型状态($16N$,与 batch 无关)外,还有第二大来源:激活值,前向时每层的中间输出,反向计算梯度时需要,故须先行保存。Korthikanti 等给出单层激活显存的量化式6:
$$\text{单层激活} \approx s\,b\,h\left(34 + 5\,\frac{a\,s}{h}\right)\ \text{字节}$$其中 $s$ 为序列长、$b$ 为 batch、$h$ 为隐藏维、$a$ 为注意力头数。此式有两个要点:其一,激活随 batch $b$ 与序列长 $s$ 线性增长(第一项);其二,括号内的 $\tfrac{as}{h}$ 项对应注意力矩阵 $s\times s$,含对序列长度的平方项,这即长上下文训练激活显存急剧膨胀的原因。
激活之所以必须保存,是因为反向传播计算某层权重的梯度时需用到该层前向时的输入(链式法则中的对应项)。前向自输入算至输出时保留每层中间输出,反向才能逐层回用。它与模型状态 $16N$ 的本质区别在于:模型状态是每参数维度的固定开销,与单次喂入的数据量无关;激活是每 token维度的开销,batch 越大、序列越长,需保存的中间结果越多。
因此同一模型将 batch 从 1 增至 32 时,参数、梯度、优化器状态一字节不变,激活却增长 32 倍。理解这一固定 vs 可变的二分,是判断增大 batch 是否会 OOM的关键。
节省激活:梯度检查点(gradient checkpointing)#
激活并非必须全部保存。梯度检查点(Chen 等 2016,亚线性内存7)的思路是:前向时只保存少数几个 checkpoint 层的激活,反向需用到中间激活时,从最近的 checkpoint 临时重算一遍前向。
每隔约 $\sqrt{L}$ 层存一个,即可将激活显存从 $O(L)$ 降至 $O(\sqrt{L})$,代价是多一次前向(算力约增加 33%,因多出 $2N$/token 的重算)。这是以算力换显存的经典权衡。
更精细的选择性重计算加序列并行可将激活显存降约 5 倍,使 530B GPT-3 在 2240 张 A100 上达到 54.2% MFU(相较全量重计算的 42.1%)6。
具体算例:训练一个 7B 模型的显存#
以 7B($N=7\times 10^9$)为例,混合精度加 AdamW:
7B 训练显存账(不含激活)
- fp16 参数:$2N=14\text{GB}$
- fp16 梯度:$2N=14\text{GB}$
- fp32 主权重:$4N=28\text{GB}$
- Adam 动量 $m$ + 方差 $v$:$8N=56\text{GB}$
按完整 $16N$ 口径约 112GB;即便不重复存 fp16 梯度、按bf16 参数 14GB + 优化器状态 fp32 约 84GB计,也接近百 GB 量级。
这一数量级已超过单张 A100 的 80GB 显存,再叠加激活,单卡无法容纳,从而必然导向多卡切分:9.6 · 训练期并行 中的 ZeRO/张量并行/流水线并行。
若将激活一并计入,显存需求进一步上升。假设序列长 2048、单卡 batch 取 4,按激活式 $sbh(34+5as/h)$ 逐层累加(不开重计算),一个 7B 模型仅激活即可再占用数十 GB,与模型状态处于同一量级8。
这正是实际训练中决定单卡能装多大 batch的往往不是参数、而是激活的原因:参数状态是固定成本,激活则随 batch 与序列长线性乃至平方增长。工程上因此有一条排查顺序:出现 OOM 时先考虑缩小 batch 或开启梯度检查点(削减激活),再考虑并行切分(削减模型状态),最后才是更换更大显存的卡。
对照推理侧:同一个 7B 以 fp16 只需存参数 14GB($2N$),一张 24GB 消费级卡即可容纳,训练与推理的显存相差近一个数量级,根源即在这 $16N$ 的优化器状态。换言之,能推理远不等于能训练:本地可运行的 7B,若从头训练则需一整台八卡服务器。这也是开源社区中能使用模型者远多于能训练模型者的原因,推理的门槛是 $2N$,训练的门槛是十余个 $N$。
常见误区辨析:其一,勿遗漏激活,大 batch/长序列下激活可能超过 $16N$,OOM 多发生于激活,应优先开启梯度检查点或缩小 batch。其二,
8-bit Adam/Adafactor等优化器可将 $m,v$ 压缩至 1–2 字节,从而削减大部分优化器状态。其三,LoRA 等参数高效微调只为极小一部分参数建立优化器状态,这正是其省显存的根本原因(详见 第 5 章 · SFT)。其四,勿将模型状态与激活混为一谈随手相加,前者随 $N$ 固定、后者随 batch·seq 浮动,两者的削减方式(切分 vs 重计算)不同,OOM 时须先分辨是哪一部分溢出再对症处理。
四、推理显存:参数 + KV-Cache#
推理去除了训练开销,但引入了缓存#
推理没有反向、没有梯度、没有优化器,训练时那三部分($14N$ 的优化器相关开销)全部不存在。显存账简化为两块:模型参数与 KV-Cache。但 KV-Cache 是新的开销来源:自回归生成时,为避免重复计算历史 token 的 Key/Value,将其缓存下来,这块缓存会随生成长度与并发量持续增大,长上下文下甚至可能超过模型权重本身。
先说明 KV-Cache 缓存的内容与缓存的必要性。自回归生成逐个产生 token:生成第 $t$ 个 token 时,注意力需以它与前 $t-1$ 个 token 的 Key 做匹配、再以匹配权重加权它们的 Value。这些历史 token 的 K/V 在各自被生成时即已算出,若每步重算,成本将随位置线性上升、总成本随序列长成平方增长,不可接受。
因此将每个 token 的 K/V 计算一次即存入缓存,后续步骤直接查表,以空间换取时间。这一交易在算力上极为划算,代价是显存:缓存随生成进行只增不减,直至本轮对话结束才释放。
核心公式#
每生成一个新 token,注意力需访问全部历史 token 的 K 与 V。将其缓存,缓存大小为9:
$$\text{KV-Cache} = 2\times L\times n_{\text{kv}}\times d_{\text{head}}\times \text{seq}\times \text{batch}\times \text{bytes}$$逐项说明:2 为 K 与 V 两份;$L$ 层各存一份;$n_{\text{kv}}\times d_{\text{head}}$ 为每层每 token 的 K(或 V)维度;再乘序列长 $\text{seq}$、并发 $\text{batch}$、每元素字节数。
关键优化点在于:GQA(分组查询注意力,Grouped-Query Attention)令多个 Query 头共享同一组 K/V 头,使 $n_{\text{kv}}\ll n_{\text{head}}$,KV-Cache 按倍数缩小,这是现代大模型(Llama-2/3、Qwen 等)普遍采用 GQA 的核心动机之一。
具体算例#
一个 80 层、KV 头数 8、head_dim 128 的模型,fp16(2 字节),序列 4096、batch 89:
$$2\times 80\times 8\times 128\times 4096\times 8\times 2\ \text{字节}\approx 10.7\text{GB}$$10.7GB 的 KV-Cache 已与一个 7B 模型的权重(14GB)处于同一量级。若将 batch 或序列再翻数倍,KV-Cache 即可超过权重本身,成为推理显存的主体。
此处有一条对推理部署重要的判断:权重是一次性、与并发无关的固定成本,KV-Cache 是每条并发请求各占一份、且随对话变长而增长的可变成本。假设以一张 80GB 卡部署上述模型,权重(如量化至 8-bit 后约 20GB)先占一块,剩余约 60GB 可用于放置 KV-Cache。
以剩余空间除以单请求单 token 的 KV 大小,即得该卡在给定序列长下最多可同时服务的请求数,这正是推理服务并发上限的物理来源。序列越长,每路请求占用的 KV 越多,可容纳的并发越少;这也是长上下文服务尤其消耗显存、成本较高的原因。
据此可反向做容量规划:欲支持更长上下文,或增加显存,或用 GQA/MQA 压小 $n_{\text{kv}}$,或将 KV 量化至低比特。
补充说明 GQA 的两个极端配置。最省的做法是多查询注意力(Multi-Query Attention, MQA),令所有 Query 头共享唯一一组 K/V($n_{\text{kv}}=1$),KV-Cache 除以头数、缩至最小,但表达力略有损失。
折中的 GQA 则将 Query 头分为若干组、每组共享一组 K/V(如 32 个 Query 头配 8 组 K/V),在显存与质量间取平衡。上述算例中 KV 头数取 8、而 Query 头很可能为 32 或 64,正是典型 GQA 配置,若退回每个 Query 头独享 K/V的原始多头注意力,同一模型的 KV-Cache 将再增数倍。
与后续章节的衔接#
正因 KV-Cache 随并发与长度线性膨胀且难以预测,它成为推理系统优化的主战场:PagedAttention 以类操作系统分页的方式按块管理 KV 内存、消除碎片(见 9.2 · 优化方法 · PagedAttention);
KV 量化/驱逐将其压至低比特或丢弃不重要的历史(见 16.1 · 量化)。整个 第 10 章 · KV-Cache 优化 均围绕这块缓存展开,本节只建立其大小公式,后续章节详述优化方法。
三本账速查表
场景 显存构成 每参数 / 关键量 训练(混合精度 + Adam) 参数 + 梯度 + 优化器状态 + 激活 $16N$ + 激活(随 batch·seq) 推理 参数 + KV-Cache $2N$(fp16)+ KV-Cache(随 seq·batch)
五、把三本账连起来用:核算流程#
单看每本账都不复杂,实际运用的要点在于:拿到一个具体需求时,知道应查哪本账、按何顺序计算。下面将典型场景串成一条决策流程。
首先确定模型规模:查 config 中的 $d$、$L$,套用 $12Ld^2$(含嵌入再加 $2Vd$),SwiGLU 模型按三矩阵与实际中间宽度修正,得到 $N$,其后所有核算以此为起点。其次估算训练时长:定下数据量 $D$,算 $C=6ND$,再除以卡数 × 单卡峰值 × MFU(现实取 0.35–0.5)得墙钟时间;若时间不可接受,回头削减 $D$ 或增加卡数。
第三判断单卡训练可行性:先算模型状态 $16N$(字节),再估激活 $sbh(34+5as/h)$ 逐层累加,两者之和若超单卡显存,则须并行,先以 ZeRO 切分优化器状态(节省最多),再张量/流水线并行,或开梯度检查点削减激活。
第四估算推理所需卡数:权重按精度算(fp16 为 $2N$、8-bit 为 $N$),再按目标序列长与并发算 KV-Cache($2Ln_{\text{kv}}d_{\text{head}}\cdot\text{seq}\cdot\text{batch}\cdot\text{bytes}$),两者之和除以单卡显存即最少卡数;反之,以单卡剩余显存 ÷ 单请求 KV得并发上限。
最后定位性能瓶颈:训练慢看 MFU(低于 30% 多为通信未重叠、算子过碎);推理慢分 prefill(算力受限)与 decode(访存受限)分别看,长上下文的高成本基本来自 KV-Cache。
这套流程的价值在于把一个含糊的工程问题拆成几个可用基本算术回答的子问题,从而在动手之前即预判可行性。这正是资源核算的意义:在编写代码之前,先在纸上把可行性算清楚。
主流配置怎样落到三本账#
| 配置 | 参数账 | 算力或显存账 | 具体例子 |
|---|---|---|---|
| GPT-2 Small | 12 层、宽度 768,主干近似 $12Ld^2$ | 权重共享使输入与输出嵌入只计一次 | 代入主干式得约 85M,再加词表与归一化后接近公开的 124M |
| GPT-3 175B | $N=175\text{B}$ | 训练近似 $6ND$,其中 $D=300\text{B}$ | 代入得约 $3.15\times10^{23}$ FLOPs,量级与论文训练账一致 |
| Llama 类 7B | SwiGLU 使用三张 FFN 矩阵 | fp16 权重本体约需 14GB | AdamW 全状态按约 $16N$ 估算为 112GB,训练需切分状态 |
| GQA 推理 | 查询头多于 KV 头 | KV-Cache 按 KV 头数线性缩放 | 128 个 Q 头配 8 个 KV 头时,缓存头维度成本降至 MHA 的 $1/16$ |
以一个完整算例贯通全流程。假设有 8 张 80GB 的卡,欲从头训练一个 13B 模型、喂入 2 万亿 token。第一步套用 $12Ld^2$ 确认其确为 13B 量级;
第二步算算力 $C=6\times 1.3\times 10^{10}\times 2\times 10^{12}\approx 1.56\times 10^{23}$ FLOPs,除以8 卡 × 单卡峰值 × MFU 0.4得墙钟时间,若算得需近半年,则可判定 8 卡不足,须减少数据或增加卡数。
第三步算显存:模型状态约 $16\times 1.3\times 10^{10}\approx 208$GB,单卡 80GB 显然不足,须 ZeRO 切分乃至叠加张量并行,再计激活还须开梯度检查点。走完这三步,即可在编写第一行训练代码之前,判明该任务在既有硬件上是否可行、大致需时多久、应如何切分。
一个常被提出的收尾问题是:这些公式是否会随硬件更新而过时。
会过时的只是具体数字(某代卡的峰值 TFLOP/s、显存容量、某模型的 MFU),结构性的核算不会变,$12Ld^2$、$6ND$、$16N$ 与 KV-Cache 公式均由Transformer 有几个矩阵、每步做几次乘加、每个参数拖几份状态推得,只要架构仍是 Transformer,这些比例关系即成立。
因此本节要读者记住的是推导骨架而非某一具体 GB 数:数字随代际更新时查阅最新值,骨架则长期有效。
小结#
- 资源核算包含三本账
- 参数量 $N\approx 12Ld^2$,单层为注意力 $4d^2$ 加 FFN $8d^2$,FFN 占三分之二,MoE 因此针对 FFN
- 算力方面,训练 $C\approx 6ND$(前向 2 加反向 4)、推理约 $2N$/token(无反向、访存受限),并以 MFU(现实约 35–55%)判断管线优劣
- 训练显存每参数约 16 字节(fp16 参数 2 + fp16 梯度 2 + fp32 主权重 4 + Adam $m,v$ 各 4),其中 12 字节的优化器状态是瓶颈,催生 ZeRO 分层切分
- 激活随 batch·seq 膨胀,可用梯度检查点以算力换显存
- 推理显存 = 参数 + KV-Cache($2Ln_{\text{kv}}d_{\text{head}}\cdot\text{seq}\cdot\text{batch}\cdot\text{bytes}$,GQA 靠 $n_{\text{kv}}\ll n_{\text{head}}$ 大幅压缩)
- 这三本账把扩展律、并行系统与推理优化串成一条线
掌握本节后,即可拿一款模型的配置在脑中完成多大、多久、几张卡的核算。这套看似朴素的算术是后续章节的公共基础:扩展律讨论给定算力如何分配给参数与数据,训推系统讨论装不下如何切分、算得慢如何加速,推理优化讨论如何压缩 KV-Cache,均建立在本节这三条公式之上。将 $12Ld^2$、$6ND$、$16N$ 与 KV-Cache 公式熟记,有助于顺畅理解后续内容。
上一节:1.3 · 注意力与 Transformer | 下一节:1.5 · 从零训一个 Transformer | 本章总览:第 1 章 · 大模型基础
参考文献#
行业好为 35–45%、50%+ 为优秀), 2025. https://debjitpaul.github.io/blog/2025/compute/
按 optimizer state / gradient / parameter 三阶段切分模型状态). https://arxiv.org/abs/1910.02054
选择性重计算 + 序列并行降激活约 5×,530B GPT-3 在 2240 A100 上 MFU 54.2% vs 全量重计算 42.1%), 2022-05-10. https://arxiv.org/abs/2205.05198
算例 $2\cdot80\cdot8\cdot128\cdot4096\cdot8\cdot2\approx 10.7$GB. https://lyceum.technology/magazine/kv-cache-memory-calculation-llm/
“All the Transformer Math You Need to Know”(JAX Scaling Book):总参数式 $L(3DF+4DNH+2D)+2DV$,$F=4D$ 时 $\approx 28LD^2+2DV$;MLP 参数与 FLOPs 在 $T\lt 8D$ 时主导;单矩阵乘前向 $2NPM$、反向对 A/B 各 $2NPM$ 故训练 $6NPM$. https://jax-ml.github.io/scaling-book/transformers/ ↩︎
Michael Wornow, “Transformer Math (Part 1): Counting Model Parameters”(GPT-2 精确参数式 $C=E(V+P)+L(12E^2+13E)+2E$,$E=768$、$L=12$ 得 124,439,808 参数), 2024-01-18. https://michaelwornow.net/2024/01/18/counting-params-in-transformer ↩︎
stanford-cs336/spring2024-lectures,
lecture_02.py:矩阵乘 FLOPs =2·B·D·K(一乘一加),forward = 2·tokens·params、backward = 4·tokens·params、训练总计6·tokens·params. https://github.com/stanford-cs336/spring2024-lectures/blob/main/lecture_02.py ↩︎ ↩︎Stanford CS336 Spring 2025 Lecture 2 transcript(资源核算分显存 GB 与算力 FLOPs 两类;A100 峰值 312 TFLOP/s,GPT-3 训练约 $3.14\times 10^{23}$ FLOPs). https://www.textpurr.com/transcript/stanford-cs336-language-modeling-from-scratch-spring-2025-lec-2-pytorch-resource ↩︎
Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models”, arXiv:1910.02054(混合精度 + Adam 每参数 16 字节。 ↩︎ ↩︎
Korthikanti et al., “Reducing Activation Recomputation in Large Transformer Models”, arXiv:2205.05198(单层激活 $\approx sbh(34+5as/h)$ 字节。 ↩︎ ↩︎
Chen et al., “Training Deep Nets with Sublinear Memory Cost”, arXiv:1604.06174(以一次额外前向为代价,将 $n$ 层网络激活显存降到 $O(\sqrt{n})$), 2016. https://arxiv.org/abs/1604.06174 ↩︎
“Breaking the memory barrier: how ZeRO revolutionizes large model training”(7B 模型 fp32 参数 28GB、AdamW 状态 56GB,bf16 参数 14GB + 状态 56GB ≈ 70GB,未含激活). https://medium.com/the-algorithmic-minds/breaking-the-gmemory-barrier-how-zero-revolutionizes-large-language-model-training-8e00d2e2f30a ↩︎
“KV Cache Memory Calculation for LLMs”(Lyceum Technology):KV-Cache = $2\cdot L\cdot n_{\text{kv}}\cdot d_{\text{head}}\cdot \text{seq}\cdot \text{batch}\cdot\text{bytes}$,GQA 用较小 $n_{\text{kv}}$ 降内存。 ↩︎ ↩︎