第2章：优化器 - AIInfraGuide

📖 本章概述

优化器是训练的”舵手”——模型算出梯度之后，怎么更新参数全靠它。本章不讲优化器的数学推导细节，而是聚焦于 AI Infra 工程师最需要关心的两个问题：优化器内部到底存了什么（直接决定显存占用）、分布式场景下优化器有什么特殊要求（大 Batch 收敛问题）。这两个问题是理解后续 ZeRO 显存优化和混合精度训练的直接前提。

📑 章节结构

0. 什么是优化器

神经网络训练的本质是最小化损失函数——模型对训练数据的预测结果与真实标签之间的差距。反向传播（backward）算出的梯度告诉我们损失函数在当前参数位置的”下坡方向”，但往哪走只是问题的一半，怎么走、走多远才是决定训练成败的关键。优化器就是回答后半个问题的组件。

优化器在训练循环中的位置

for batch in dataloader:
    loss = model(batch)          # 前向：算预测 → 算损失
    loss.backward()              # 反向：算梯度
    optimizer.step()             # 优化器：用梯度 + 内部状态更新参数
    optimizer.zero_grad()        # 清零梯度，准备下一轮

optimizer.step() 这一行看似简单，背后做了三件事：（1）读取每个参数的梯度；（2）结合优化器内部状态（动量、二阶矩等）计算更新量；（3）更新参数。优化器的内部状态在 step() 之间持久保存，这就是为什么它会持续占用显存。

1. 优化器演进：从 SGD 到 AdamW

SGD：最朴素的优化器，参数更新规则 $\theta_{t+1} = \theta_t - \eta \cdot g_t$ ，无额外状态，每个参数只需存一份梯度。好比沿着当前最陡的方向迈一步，简单直接但容易在山谷里来回震荡
Momentum SGD：引入动量 $m_t = \beta m_{t-1} + g_t$ ，每个参数多存一个动量缓冲区（一阶动量）。好比给下山的球加了惯性，不容易被小坑绊住，但需要额外存储动量状态
Adam：同时维护一阶动量（梯度的指数移动平均 $m_t$ ）和二阶动量（梯度平方的指数移动平均 $v_t$ ），相当于同时记住”往哪走”和”路有多颠”。自适应学习率让不同参数有不同的更新步长
AdamW：修正了 Adam 中权重衰减（Weight Decay）的实现方式——Adam 将 Weight Decay 混入梯度中，与自适应学习率耦合导致正则化效果不对；AdamW 将其解耦为独立的参数缩放步骤。这是目前大模型训练的事实标准

2. 优化器状态的显存开销（本章重点）

理解优化器的显存开销是理解 ZeRO 的关键。以下分析以混合精度训练为前提（前向/反向用 FP16/BF16，优化器内部用 FP32 保持精度）：

各优化器的每参数显存开销

优化器	额外状态	每参数显存	7B 模型总开销
SGD	无	4B（FP32 参数副本）	~28 GB
Momentum SGD	一阶动量	8B（FP32 参数副本 + FP32 动量）	~56 GB
Adam / AdamW	一阶动量 + 二阶动量	12B（FP32 参数副本 + FP32 $m$ + FP32 $v$ ）	~84 GB

显存占比分析

以 AdamW + BF16 混合精度训练 7B 模型为例：

BF16 参数： $2 \times 7 \times 10^9$ = 14 GB
BF16 梯度： $2 \times 7 \times 10^9$ = 14 GB
优化器状态： $12 \times 7 \times 10^9$ = 84 GB（FP32 参数副本 28GB + 一阶动量 28GB + 二阶动量 28GB）
总计：~112 GB，其中优化器状态占 75%