SGD、Momentum、Adam、AdamW 怎么更新？
各要多少额外显存？
先弄清楚：每个概念是什么
- SGD 与 learning rate
- 沿 gradient 的反方向迈一步，learning rate \(\eta\) 是步长。SGD 只用一个 mini-batch 估计 gradient，所以有噪声。
- Momentum
- 把过去的 gradient 累加成速度，像滚下山的球。方向一致的分量越滚越快，来回震荡的分量互相抵消。
- Adam
- momentum 加上每个参数自己的步长。m 记 gradient 的平均，v 记 g² 的平均，步子按 \(1/\sqrt{v}\) 缩放，gradient 一直很大的参数步子被压小。
- Weight decay vs L2
- L2 是在 loss 上加 \(\tfrac{\lambda}{2}\|w\|^2\)，gradient 多一项 λw。weight decay 是每步直接把 weight 缩一点。
- Optimizer state
- optimizer 为每个参数额外存的数，SGD 0 份、momentum 1 份、Adam 2 份。
30 秒回答
SGD steps against the minibatch gradient, w ← w − ηg. Momentum accumulates past gradients into a velocity, damping oscillation. Adam keeps moving averages of the gradient, m, and the squared gradient, v, bias-corrects both because they start at zero, and updates w by η·m/(√v + ε), so each parameter gets its own step size. AdamW decouples weight decay: instead of adding λw to the gradient, where Adam would rescale it, it shrinks the weights directly by ηλw. Memory-wise, SGD adds no state, momentum one copy of the parameters and Adam two; mixed-precision Adam totals 16 bytes per parameter, about 112 GB for a 7B model.
关键公式
SGD 与 Momentum
\[ b_t = \mu\, b_{t-1} + g_t,\qquad w_t = w_{t-1} - \eta\, b_t \]\(\mu = 0\) 就是 SGD；gradient 不变时 \(b \to g/(1-\mu)\)，\(\mu = 0.9\) 步长放大 10 倍。
Adam（默认 η = 1e-3、β₁ = 0.9、β₂ = 0.999、ε = 1e-8）
\[ m_t = \beta_1 m_{t-1} + (1-\beta_1)\, g_t,\qquad v_t = \beta_2 v_{t-1} + (1-\beta_2)\, g_t^2 \] \[ \hat m_t = \frac{m_t}{1-\beta_1^{t}},\qquad \hat v_t = \frac{v_t}{1-\beta_2^{t}} \] \[ w_t = w_{t-1} - \eta\,\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon} \]为什么要 bias correction
m、v 从 0 起步，前几步偏小。假设 gradient 分布不变，\(\mathbb E[v_t] = \mathbb E[g^2]\,(1-\beta_2^{t})\)，所以除以 \(1-\beta_2^{t}\)；\(t\) 大了修正自动消失。例：第 1 步 g = 2，\(m_1 = 0.2\)、\(v_1 = 0.004\)，修正后 \(\hat m_1 = 2\)、\(\hat v_1 = 4\)，正好走 \(\eta\)；不修正会走 \(0.2/\sqrt{0.004} \approx 3.16\,\eta\)。
AdamW vs Adam + L2
Adam + L2 把 λw 加进 g，它跟着被 \(\sqrt{\hat v}\) 除掉，所以 gradient 大的参数几乎不被 decay。AdamW 的 m、v 只用 loss 的 gradient，decay 单独做：
\[ w_t = w_{t-1} - \eta\lambda\, w_{t-1} - \eta\,\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon} \]所有参数按同一比例 ηλ 缩小。SGD 下两者完全一样，因为 λw 乘 η 正好是 ηλw。
显存：每种 optimizer 多存多少
每个参数存 weight、gradient、optimizer state，按 bytes 加起来再乘参数量（fp32 4 bytes，bf16 2 bytes；不含 activation）。7B 模型：
| optimizer | optimizer state | bytes / param | 7B 合计 |
|---|---|---|---|
| SGD | 无 | 4 + 4 = 8 | 56 GB |
| Momentum | b | 8 + 4 = 12 | 84 GB |
| Adam / AdamW | m、v | 8 + 8 = 16 | 112 GB |
| Adam，mixed precision | fp32 master、m、v | 2 + 2 + 12 = 16 | 112 GB |
只问 optimizer state 时，7B 的 Adam 多存 14B 个数，fp32 就是 56 GB。纯 fp32 的 Adam 是 4 倍参数量（w、g、m、v）。mixed precision 时 bf16 省下的被 fp32 master copy 吃回去，还是 16 bytes。低精度的 weight 直接更新会把小的更新量舍掉，所以要这份 master copy，见 mixed precision。一张卡放不下就用 ZeRO 切，见 DP、ZeRO、FSDP。
必答点
- SGD：\(w \leftarrow w - \eta g\)；momentum 累加历史 gradient，步长放大 \(1/(1-\mu)\)
- Adam = momentum（m）+ 按 \(\sqrt{v}\) 缩放，每个参数自己的步长
- bias correction：m、v 从 0 起步偏小，除以 \(1-\beta^t\)
- AdamW：decay 不进 gradient，直接 \(w \leftarrow w - \eta\lambda w\)
- 额外 state：SGD 0、momentum 1、Adam 2 份；mixed precision Adam 16 bytes/param，7B 约 112 GB
追问
Adam 和 SGD 怎么选？
LLM 基本都用 AdamW（LLaMA 取 β₂ = 0.95、weight decay 0.1）。图像分类上 SGD + momentum 过去常比 Adam 泛化好，AdamW 原文认为主因是 L2 在 Adam 里不起作用。
为什么要 warmup + cosine decay？
刚开始参数随机、gradient 不稳，直接用峰值 learning rate 容易发散，所以先线性升上去，再按 cosine 慢慢降。LLaMA 用 2000 步 warmup，最后降到峰值的 10%。
来源Kingma & Ba 2014（Adam）Alg. 1（p.2）、§2.1、§3（pp.2–3）· Loshchilov & Hutter 2017（AdamW）§2、Alg. 2（pp.2–3）· Goodfellow et al. §8.3.2（pp.292–296）· Murphy, PML1 §8.2.4.1（pp.287–288）、§8.4.6.3（p.300）· Rajbhandari et al. 2019（ZeRO）§3.1、§5（pp.7–11）· Micikevicius et al. 2017 §3.1 · LLaMA §2.3 · PyTorch AdamW