大语言模型大 Batch 训练临界 Batch Size 经验测量与优化策略

3. 实验发现与分析

发现一：CBS 随训练过程动态演变

• CBS 在初始化时接近 0。
• 在训练初期（前 50B tokens）迅速增长。
• 随后进入平台期（Plateau），稳定在约 4096 (documents) 左右。
• 结论: CBS 不是一个静态常数，而是一个动态变化的量。这意味着在训练初期应该用小 BS，后期可以用大 BS。

发现二：模型规模不影响 CBS

• 对比 1B 和 7B 的曲线：
  • 两条曲线的形态和数值惊人地相似。
  • 启示: 我们可以用小模型（1B）测出的 CBS 规律，去指导大模型（7B 甚至更大）的训练配置，节省昂贵的探索成本。

发现三：梯度噪声尺度不可靠

• 实测的梯度噪声尺度（红色虚线）与本文测量的真实 CBS（蓝色实线）完全对不上。
• 噪声尺度严重低估了真实的 CBS（差了几个数量级）。
• 趋势也不匹配（特别是在 7B 模型上）。
• 打击: 证明了 McCandlish (2018) 的理论在 LLM/Adam 场景下失效。

4. 应用与结果

4.1 提出的策略：Batch Size Warmup

• 策略逻辑: 既然 CBS 随时间增长，我们应该动态调整 BS。
• 具体算法:
  1. 从小 BS 开始。
  2. 定期检测 CBS 是否超过了当前 BS 的两倍。
  3. 如果是，则将 BS 翻倍，并根据 $2\sqrt{2}$ 规则调整学习率。
• 实施: 在 OLMo 1B 训练中，BS 经历了两次翻倍（1024 -> 2048 -> 4096）。

4.2 核心结果对比 (Table 1)

• 三组对比实验:
  1. Small-Batch Control (一直用小 BS，1024): 理论上 Loss 最好，但慢。
  2. Large-Batch Control (一直用大 BS，4096): 跑得快，但初期 BS > CBS，导致 Loss 差。
  3. BS Warmup (Ours): 动态调整。
• 结果炸裂:
  • Loss: Warmup 方法的最终 Loss 甚至略优于 Small-Batch (2.5433 vs 2.5486)。
  • 效率: 相比 Small-Batch，节省了 43% 的梯度步数。

对比 Large-Batch：Large-Batch 虽然也快，但最终 Loss 明显变差，且无法恢复。

4.3 下游任务性能 (Downstream Tasks)

• 不仅仅看 Training Loss，还看了下游任务（C4, The Pile, BPB 等）。
• BS Warmup 在各项指标上均与 Small-Batch 持平或略优。
• 结论: 这种加速方法是'免费的午餐'，没有副作用。

5. 总结与个人思考

5.1 论文总结

1. 方法论: 提出了基于分支训练的 CBS 测量法，不依赖强理论假设。
2. 科学发现: CBS 随训练过程增长并趋于平稳；CBS 与模型大小无关。
3. 工程价值: 证明了 Gradient Noise Scale 在 LLM 上的失效；提出了 BS Warmup，实现了 40%+ 的训练加速。

5.2 局限性分析

• 成本问题: 虽然比从头跑多次要省钱，但'分支训练'本身依然需要额外的计算资源（每次 Checkpoint 都要分叉跑 2B tokens）。
• 参数敏感性: $\Delta$ (窗口大小) 和 epsilon（容忍度）的选择比较经验主义。如果 $\Delta$ 选小了，可能误判 CBS。
• 通用性验证: 目前只验证了 OLMo 架构和 Dolma 数据集，是否适用于 MoE 或其他架构还需验证。

5.3 工程实践启示

• 如果我们要训练新模型：不要迷信 OpenAI 提到的 Gradient Noise Scale。
• 调参策略：不要从头到尾用固定的 Batch Size。可以尝试手动实现简单的 Warmup（比如在前 10% step 用小 BS，后面翻倍），这可能在资源受限的情况下提升效果。
• Proxy 的使用：可以用小模型（如 1B）先跑一遍确定 CBS 曲线，直接套用到大模型训练中。

6. Q&A

(1) 'Noise Scale 是错的'? 我们以前可能一直以为梯度的方差能告诉我们 Batch Size 设多少，但这篇论文说在用 Adam 的时候这完全是误导。

(2) 关于 Methodology：'这样测量是不是很贵？' 确实有开销，但相比于盲目用小 BS 浪费的时间，或者用大 BS 导致模型训练坏了重跑，这个开销是划算的。而且作者发现 1B 的规律可以用在 7B 上，这大大降低了测量成本。

(3) 关于公式：Paper 中提到了 k√k scaling rule for Adam。为什么不是线性缩放（Linear Scaling） 因为 Adam 的 update rule 里分母有二阶矩估计，导致梯度的 scaling 变成了平方根关系（引用 Malladi et al., 2022）。

(4) Appendix D，作者试图推导 T√T 的 scaling law，但最后认为经验测量比强行推导理论更靠谱。这显示了作者严谨的态度。