在实际工作中，经常有人问，7B、14B 或 70B 的模型需要多大的显存才能推理？如果微调它们又需要多大的显存呢？为了回答这个问题，本文整理了一份关于训练或推理所需显存的计算方式。如果对具体细节不感兴趣，可以直接参考经验法则评估资源需求。
数据精度
在开始介绍之前，先明确一个重要的概念——数据精度。数据精度指的是信息表示的精细程度，在计算机中由数据类型和其位数决定。计算显存时，从'原子'层面来看，必须知道使用数据的精度，因为它决定了数据存储的方式以及占用多少 bit。
目前常见的精度类型如下：
- 32-bit (4 Bytes): FP32 / float32
- 16-bit (2 Bytes): FP16 / float16, bfloat16
- 8-bit (1 Byte): int8
- 4-bit (0.5 Bytes): int4
经验法则
推理：参数量 × 精度。例如，假设模型都是 16-bit 权重发布的，即一个参数消耗 2 Bytes 内存。若模型参数量为 70B，基于上述经验法则，推理最低显存需求为 70B × 2 Bytes = 140 GB。
训练：通常需要 4 - 6 倍的推理资源。
推理阶段显存分析
在模型推理阶段，需要的资源主要有三部分：模型权重、KV Cache 和激活（在推理过程中创建的张量）。
1. 模型权重
加载模型权重（即模型大小）占用的资源主要依赖于模型的参数量和精度。其中，参数量基本不变，精度可以通过模型量化技术进行优化。尽管量化可能会轻微影响模型性能，但相比于选择更高精度的小模型来说，量化技术更受青睐。
公式：
模型大小 = 参数量 × 精度字节数
2. KV Cache
在 Transformer 的解码阶段，每次推理生成一个 token，依赖于之前的 token 结果。如果每次都对所有 token 重新计算一次，代价非常大。为了避免重复计算，通过 KV Cache 技术将其缓存到 GPU 内存中。
公式：
KV Cache = 2 × Batch Size × Sequence Length × Number of Layers × Hidden Size × Precision
*注意：第一个因子 2 解释了 K 和 V 矩阵。通常，Hidden Size 和 Number of Layers 的值可以在模型相关的配置文件中找到。
3. 激活内存
在模型的前向传播过程中，必须存储中间激活值。这些激活值代表了神经网络中每层的数据在向前传播时的输出。它们通常保持为 FP32 格式，以避免数值爆炸并确保收敛。
公式：
Activation Memory = Batch Size × Sequence Length × Hidden Size × (34 + (5 × Sequence Length × Number of attention heads) / Hidden Size)
训练阶段显存分析
训练阶段所需的资源，除了上述介绍的模型权重、KV Cache 和激活内存之外，还需要存储优化器和梯度状态，因此训练比推理需要更多的资源。
1. 优化器内存
优化器需要资源来存储参数和辅助变量。这些变量包括诸如 Adam 或 SGD 等优化算法使用的动量和方差等参数。这取决于优化状态的数量及其精度。例如，AdamW 优化器是最流行的微调 LLM 的方法，它为模型的每个参数创建并存储 2 个新参数。如果我们有一个 70B 的模型，优化器将创建 140B 的新参数！假设优化器的参数为 float32，即每个参数占用 4 字节的内存。优化器至少需要 140B × 4 Bytes = 560 GB 的资源。
不同优化器的状态数量如下：
- AdamW: 2 states (8 Bytes per parameter)
