本文试图以最清晰的方式手动推导 Transformers 每一步的参数量到显存、计算量问题。理解底层，才能更好的做训练和优化。本文内容包括最基本的模型训练和推理过程中的显存占用，以及对其中中间激活值的优化方案和显存占用。****

1 训练过程

训练中的显存占用分两块

1. 模型状态，参数、梯度和优化器状态
2. 剩余状态， 激活值、临时buffer

1-1 模型状态显存

写在前面的前置知识，

• 显存占用 = 参数数量 x 该参数精度占用的 bytes 数
• 换算关系：Int8 需1 bytes, fp16 / bf16 数需 2 bytes, fp32 需要 4 bytes
• 目前训练 LLM 主流配置： AdamW 优化器 + 混合精度训练 + activation checkpoint 策略

设模型参数量为 $\Phi$，模型参数（fp16）、模型梯度（fp16）和优化器状态（fp32），总参数量 =，总数为 $2 \Phi + 2 \Phi + (4+4+4) \Phi = 16 \Phi$。

1. 这部分比较固定，且在整个训练过程中都要存在显存中。
2. 一般只能通过并行切分（Tensor Parallelism/Pipeline Parallism）能减少。
3. 不同优化器的 K 值不同，算法的中间变量、框架的实现都有可能有一定区别。

1-2 中间激活值显存

激活（activations）指的是：前向传递过程中计算得到的，并在后向传递过程中需要用到的所有张量。

中间激活值占用显存分两个部分分析：Attention 和 MLP，Embedding 没有中间值

Attention 部分计算公式如下：