第 7 章 — 优化器与数值稳定

开篇段落

如果说模型架构是训练的“骨架”，数据是“血液”，那么优化器与数值稳定性策略就是整个训练过程的“心脏与循环系统”。一个微小的超参失当，或是一次未被察觉的数值溢出，都可能导致数千 GPU 小时的努力付诸东流。本章，我们将深入这间至关重要的“引擎室”，对 LLM 训练中的核心组件进行精细解剖。我们将从默认王者 AdamW 出发，不仅探讨其超参的“what”，更深究其“why”，特别是 β₂ 和解耦权重衰减的现代选择。随后，我们将审视 Lion、Adafactor 等备选方案的独特设计与适用界，并详解 DeepSpeed Paged AdamW 如何打破显存壁垒。更重要的是，我们将系统性地构建一套数值稳定性的“纵深防御体系”，涵盖梯度累积、bf16 混合精度的奥秘、损失缩放的动态机制以及梯度裁剪的必要性。学完本章，你将不仅掌握一套“能用”的配置，更能理解其背后的原理，从而在面对训练不稳时，具备庖丁解牛般的诊断与调试能力。

文字论述

1. AdamW: 精调默认的王者

AdamW (Adam with Decoupled Weight Decay) 凭借其鲁棒性和高效性，已成为大规模语言模型训练的黄金标准。然而，“默认”不代表“无需理解”。恰恰相反，对其关键超参的精细调整是通往 SOTA 性能的第一步。

1.1 解耦权重衰减 (Decoupled Weight Decay) 的威力

要理解 AdamW，首先要明白它与 Adam + L2 正则化的区别。

  • 传统 L2 正则化: 在损失函数中加入 (λ/2) * ||w||² 项。这导致梯 g_t 变为 g_t' = g_t + λ * w_t。这个“正则化梯度”随后被送入 Adam 的动量和二阶矩估计中，影响了自适应学习率的计算。对于梯度较大（通常是信息量大）的权重，λ * w_t 的影响被 sqrt(v_t) 项削弱；对于梯度较小的权重，其影响反而被放大。这并非我们想要的正则化效果。

  • 解耦权重衰减: AdamW 将权重衰减直接从梯度更新中分离出来。在计算完基于梯度的更新步长 Δw_t 后，直接从权重中减去一个与其自身大小成正比的量。

    • 更新公式（概念上）：w_{t+1} = w_t - η * Δw_t - η * λ * w_t
    • 这意味着权重衰减的效果与学习率 η 和衰减系数 λ 直接相关，且独立于梯度的历史信息。它更像是一种朝向原点的“匀速”收缩，是更纯粹、更可控的正则化。

经验法则: 对于 LLM 预训练，权重衰减系数 wd (即 λ) 的一个非常稳健的起始点是 0.1。

1.2 键超参的现代理解

  • β₁ (一阶矩衰减率): 0.9 几乎是所有任务的通用标准，它控制了动量的“记忆长度”。修改它的风险高，收益不确定，因此我们通常固定此值。

  • β₂ (二阶矩衰减率): 这是 LLM 训练中最关键的超参之一。

    • 传统值 0.999: 意味着优化器对过去 1000 步的梯度平方信息有很长的“记忆”。这在相对平滑的损失曲面上表现良好。但在 LLM 训练这种高度非凸、梯度充满噪声的环境中，过长的记忆就像一艘笨重的油轮，转向缓慢，无法快速适应新的梯度信息。
    • 现代值 0.95: LLaMA 等工作的实践表明，更短的记忆（约 20 步）能让优化器像一艘灵活的快艇。当遇到一小段损失平稳期（plateau）或梯度分布突然变化时，v_t 能更快地“忘记”旧信息，采纳新信息，从而调整每个参数的学习率，帮助模型更快地逃离困境。对于从零训练，β₂=0.95 是强烈推荐的默认值。

  • ε (Epsilon): 分母稳定项。

    • 其值应与数值精度相匹配。在 fp32 下，1e-8 是标准值。但在 bf16/fp16 混合精度下，由于数值表示范围变窄，一个稍大的 ε 如 1e-6 或 1e-5 能提供额外的保护，防止 sqrt(v_t) 接近零时导致分母过小，更新步长爆炸。

1.3 Fused AdamW: 榨干硬件性能

在 H100 规模的训练中，每一步的耗时都至关重要。优化器步骤（optimizer.step）可能成为一个不可忽视的瓶颈，因为它涉及大量对参数和优化器状态的读写操作。

  • 瓶颈分析: 这个过程是内存带宽密集型 (Memory-Bandwidth Bound) 而非计算密集型。标准的 PyTorch 实现会为每个操作（加法、乘法、开方等）启动一个独立的 CUDA kernel。频繁的 kernel 启动开销和全局内存的反复读写会严重拖慢执行速度。
  • 融合 (Fusing): Fused AdamW 将整个更新逻辑（包括梯度缩放、动量计算、二阶矩更新、权重衰减和最终的参数更新）合并成一个或少数几个高性能的 CUDA kernel。这使得所有计算尽可能在 GPU 的寄存器和 L1/L2 缓存中完成，一次性从全局内存（HBM）读取数据，计算完毕后再写回，极大减少了内存I/O，从而将优化器步骤提速数倍。DeepSpeed 和 Apex 都提供了此类实现，在我们的技术栈中应默认启用。

2. 备选方案：特定场景下的利器

  • Lion (EvoLved Sign Momentum):

    • 核心机制: update = sign(β₁*m_{t-1} + (1-β₁)g_t)。它完全抛弃了二阶矩 v_t，更新方向仅由动量的符号决定，更新步长则由全局学习率 η 控制。
    • 优劣分析:
      • 优点: 内存占用减半（只需存 m_t），在某些基准测试中展现出比 AdamW 更强的性能。其恒定的更新幅度可能有助于跨越一些尖锐的局部最小值。
      • 缺点: 超参其敏感。学习率 η 通常需要设为 AdamW 的 1/5 到 1/3，而 wd 也需要相应放大 5-10 倍。β₁ 和 β₂ 也需要重新调整。它是一种高风险高回报的选择，不适合作为初次尝试的基线。

  • Adafactor:

    • 核心机制: 为了节省内存，它不存储完整的二阶矩 v_t，而是将其近似为一个低秩分解矩阵。它也不存储一阶动量，除非 β₁ 不为零。
    • 历史地位与现状: 在 ZeRO 技术出现之前，Adafactor 是训练超大模型（如 T5）的唯一选择。但在 DeepSpeed ZeRO 普及的今天，其近似带来的潜在性能损失，使其吸引力大大下降。在我们的 H100 集群上，有 Paged AdamW 和 ZeRO-3，几乎没有理由再选择 Adafactor。

  • DeepSpeed Paged AdamW: 内存的无限游戏

    • 工作原理: 它将 CPU 的系统内存（DRAM）视为 GPU 高速缓存（HBM）的扩展。绝大多数优化器状态（FP32 权重副本、动量、方差）都存放在 CPU 内存中。当需要为某一层参数更新时，DeepSpeed 会异步地、提前地将这些状态从 CPU “分页(page in)”到 GPU 的一块预留缓存区中。计算完成后，更新后的状态再被“分页(page out)”回 CPU 内存。
    • ASCII 流程图:
CPU DRAM (Hundreds of GBs)             GPU HBM (80GB)
+----------------------------+         +-----------------+
| Optimizer States (Layer N) | ------> | Paging Buffer   | --(compute)--> Updated States
| Optimizer States (Layer N+1)|         +-----------------+       |
| ...                        | <------ (write back) <-------------+
+----------------------------+

*   **关键**: **异步预取 (Asynchronous Prefetching)** 是隐藏延迟的关键。在 GPU 计算第 N 层的反向传播时，DeepSpeed 的后台线程已经在通过 PCIe 总线拉取第 N+1 层的优化器状态了。
*   **适用性**: 对于 13B 及以上规模的模，即便有 80GB HBM，优化器状态依然会占据巨大空间。Paged AdamW 配合 ZeRO-3 是解锁更大模型、更大批次训练能力的**核心技术**。

3. 数值稳定性的纵深防御体系

3.1 梯度累积 (Gradient Accumulation, A)

梯度累积是在不牺牲内存的前提下，达成大批量训练目标的基础。 Global batch size (tokens) = micro_batch_size * num_gpus (D) * gradient_accumulation_steps (A) 从优化器的视角看，它在每 A 个 micro-step 之后，看到的是一个由 D * A 个 micro-batch 累加而成的梯度，这等效于用一个巨大的 batch size 进行了一次更新。这对于稳定训练、利用 large-batch scaling law 至关重要。

3.2 混合精度训练：bf16 的统治

选择正确的数值格式是速度与稳定性的核心权衡。

| 格式 | 符号位 | 指数位 | 尾数位 | 动态范围 (近似) | 精度 | H100 Tensor Core 支持 | 备注 |

格式 符号位 指数位 尾数位 动态范围 (近似) 精度 H100 Tensor Core 支持 备注
FP32 1 8 23 10⁻³⁸ to 10³⁸ 高 (基线) (通过 TF32 模拟) 稳定但慢，内存占用大
TF32 1 8 10 10⁻³⁸ to 10³⁸ 中等 是 Hopper/Ampere 默认 Matmul 格式，对用户透明
BF16 1 8 7 10⁻³⁸ to 10³⁸ 低 是 LLM 训练首选，动态范围同 FP32
FP16 1 5 10 6x10⁻⁸ to 65504 中等 是 动态范围窄，极易溢出，需强依赖损失缩放

结论: bf16 的 8 位指数位使其拥有与 fp32 相同的动态范围，这意味着它能表示极大和极小的数值，从根本上解决了 fp16 的上溢（overflow）和下溢（underflow）问题。虽然其 7 位尾数牺牲了精度，但大量研究表明，对于神经网络随机梯度下降的本质来说，这种精度损失几乎不影响最终收敛效果。在 H100 平台上，bf16 是毫无疑问的最佳选择。

3.3 动态损失缩放 (Dynamic Loss Scaling)

即便在使用 bf16 时，损失缩放依然是一道重要的、几乎零成本的保险。

  • 机制详解:
    1. 初始化: 设置一个初始缩放因子 S_init (e.g., 2^16) 和一个检查周期 N_growth (e.g., 1000 steps)。
    2. 缩放: loss_scaled = loss * S。
    3. 反向传播: loss_scaled.backward()。梯度也被放大了 S 倍。
    4. 梯度反缩放与检查: 在 optimizer.step() 之前，框架会遍历所有梯度，检查是否存在 Inf 或 NaN。
      • 若存在溢出: 跳过此次权重更新。将 S 除以 2（S /= 2），重置检查周期计数器。
      • 若无溢出: 将所有梯度除以 S 恢复原值，然后执行 optimizer.step()。检查周期计数器加一。如果计数器达到 N_growth，则将 S 乘以 2（S *= 2），以尝试保留更小的梯度细节。
  • 作用: 这个动态调节机制确保了 S 始终处在一个“既能防止下溢，又不会导致上溢”的甜点区。在 PyTorch Lightning/DeepSpeed 中，只需开启相应配置，整个过程完全自动化。

3.4 梯度裁剪 (Gradient Clipping)

在训练初期，或者当遇到不稳定的数据批次时，梯度可能会发生爆炸，导致更新步长过大，瞬间摧毁模型的已有学习成果。梯度裁剪是限制这一破坏的“安全阀”。

  • 机制: 在优化器更新前，计算所有模型参数梯度的 L2 范数（grad_norm）。
    • grad_norm = sqrt(sum(g_i² for g in all_gradients))
  • 裁剪: 如果 grad_norm 超过一个预设的阈值 max_norm，则对所有梯度进行缩放：
    • g_clipped = g * (max_norm / grad_norm)
  • 这保证了梯度的“方向”不变，但其“度”被限制在 max_norm 以内。
  • 经验法则: 对于 LLM 训练，一个常见的 max_norm 值是 1.0。在训练日志中监控 grad_norm 的实际值非常重要。如果它频繁触及裁剪阈值，可能意味着学习率过高或存在其他不稳定性因素。

本章小结

  • 优化器基线: AdamW 是最可靠的选择。使用 Fused 实现，并设置现代超参：β₁=0.9, β₂=0.95, ε=1e-6, wd=0.1。
  • 内存管理: 面对 13B+ 模型，DeepSpeed Paged AdamW 结合 ZeRO-3 是打破显存墙、实现大规模训练的关键。
  • 数值精度: bf16 是 H100 平台的标准配置。它提供了与 fp32 相同的动态范围，极大地提升了训练稳定性，同时带来显著的性能和内存优势。
  • 稳定“铁三角”:
    1. 梯度累积: 实现大批量训练的基石，用以满足 Scaling Law 对 GB_tok 的要求。
    2. 动态损失缩放: 混合精度训练的自动化“保丝”，防止梯度因数值范围问题而丢失或爆炸。
    3. 梯度裁剪: 限制梯度爆炸的“安全阀”，max_norm=1.0 是一个稳健的起点。

常见陷阱与错误 (Gotchas)

  1. β₂ 沿用旧值 0.999:

    • 症状: 模型训练初期 loss 下降极其缓慢，或者在训练中期轻易地陷入一个平稳期（plateau）无法自拔。
    • 诊断: 将 β₂ 想象成优化器的“惯性”。0.999 的惯性巨大，难以转向。0.95 则灵活得多，能快速响应梯度变化。
    • 补救: 如果训练已经开始，可以考虑在 loss 平稳期“热切换”β₂ 到 0.95，有时能激活训练。但最好的做法是从一开始就使用 0.95。

  2. Lion/其他新优化器与学习率不匹配:

    • 症状: 切换到 Lion 优化器后，训练在第一个 step 就 loss=NaN。
    • 诊断: Lion 的更新机制对学习率的尺度要求与 AdamW 完全不同。直接沿用 AdamW 的 LR schedule 是灾难性的。
    • 补救: 严格遵循 Lion 论文或官方实现的建议，将 AdamW 的峰值学习率除以 3-10 作为 Lion 的起始点，并对权重衰减进行相应放大。对任何新优化器，都要假定其超参空间不兼容，必须从小规模实验开始重新搜索。

  3. 梯度范数持续触顶或为 Inf/NaN:

    • 症状: 日志中 grad_norm 持续等于 max_norm (e.g., 1.0)，或者直接报告 grad_norm 是 Inf。损失缩放因子 S 不断下降。
    • 诊断 checklist:
      1. 学习率过高?: 这是最常见的原因。尝试降低峰值学习率 30%-50%。
      2. 数据问题?: 检查当前批次的数据，是否存在异常值或损坏的样本。
      3. 模型模块不稳定?: 某个自定义的算子或数值计算在 bf16 下不稳定？可以尝试用 PyTorch hooks 打印出每一层的梯度范数，定位到第一个出现 Inf 的层。
      4. 初始化问题?: 检查权重初始化否合理，不当的初始化可能导致早期激活值或梯度爆炸。
      5. RoPE Scaling 伪影?: 在使用 RoPE scaling 扩展上下文时，不当的 base 或 theta 值也可能引入数值问题。

  4. 忽视 Paged AdamW 的 I/O 瓶颈:

    • 症状: 开启 Paged AdamW 后，GPU 利用率显著下降，tokens/s 不升反降。
    • 诊断: 使用 nvidia-smi dmon 或 dcgm 监控 PCIe 带宽使用率。如果带宽持续被打满，说明 CPU-GPU 的数据拷贝成了瓶颈。
    • 补救:
      • 确认使用的是高速的 PCIe Gen4/Gen5 连接。
      • 在多 NUMA 节点的服务器上，确保训练进程和其使用的 GPU 绑定在同一个 NUMA 节点上，避免跨节点内存访问带来的高延迟。
      • 增加 DeepSpeed 配置中用于 pining 和 prefetching 的 buffer 大小，但这会消耗更多 CPU 内存。
      • 最终，如果硬件瓶颈无法解决，可能需要接受较低的吞吐，或者调整并行策以减少需要 offload 的数据量。