优化器状态是深度学习训练中优化器为自适应更新而维护的辅助变量,是显存占用的核心来源之一。
核心定义
优化器状态:优化器迭代更新参数时保存的"历史梯度信息",与模型参数维度一致。
为什么需要优化器状态?
以 Adam 优化器为例,需要记录两个关键状态:
- 一阶矩 $m$:梯度的指数移动平均(梯度方向)
- 二阶矩 $v$:梯度平方的指数移动平均(梯度波动)
更新公式: $$\theta = \theta - \eta \cdot \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon}$$
显存开销
| 组件 | 相对参数量 |
|---|---|
| 模型参数 | 1份 |
| 梯度 | 1份 |
| Adam 状态($m+v$) | 2份 |
单卡仅"模型+梯度+优化器状态"就占 4 倍参数量,这也是 ZeRO 优先分片优化器状态的原因。
不同优化器对比
| 优化器 | 状态数量 | 显存占比 |
|---|---|---|
| SGD(无动量) | 0 | 0% |
| SGD+动量 | 1 | 100% |
| Adam/AdamW | 2 | 200% |
ZeRO 分片逻辑
- ZeRO-1:仅拆分优化器状态,10个GPU时单卡仅存 1/10 的 $m$ 和 $v$
- 核心行为:第 $i$ 个 GPU 仅维护自己参数分片对应的状态
张芷铭的个人博客
Comments