大规模预训练中批大小设置的理论与实践

1. 大规模预训练中批大小设置的核心挑战

在深度学习模型训练中，批大小(batch size)的选择一直是个微妙而关键的问题。传统认知中，增大批大小可以提升GPU利用率并减少训练时间，但会降低模型泛化能力。这种认知源于小规模实验的观察，但当模型规模扩展到千亿参数级别时，情况发生了根本性变化。

1.1 传统批大小理论的局限性

经典的E(S)公式描述了总数据消耗E与达到特定损失所需优化步骤数S之间的关系。这个公式基于恒定学习率的假设，形式如下：

E(S) = B × S

其中B是批大小。这个线性关系暗示着：只要总计算量(B×S)相同，使用不同批大小最终效果应该相近。然而在现代大规模预训练中，这个理论完全失效了——使用Warmup-Stable-Decay(WSD)学习率调度器时，实际训练动态呈现完全不同的模式。

关键发现：当使用WSD调度器时，E(S)关系呈现明显的三阶段特征，这与恒定学习率下的线性关系截然不同。

1.2 WSD调度器带来的新动态

论文通过理论推导和实验验证，建立了适用于WSD调度器的E(S)关系新公式。这个分段函数将训练过程划分为三个特征鲜明的阶段：

1. 初始阶段：E与S呈反比关系(E ∝ 1/S)
2. 过渡阶段：E是S的二次函数(E ∝ S²)
3. 渐近阶段：E与S恢复线性关系(E ∝ S)

这种非线性动态解释了为什么传统批大小理论在大规模预训练中失效——训练过程不再是简单的线性计算量累积，而是存在明显的阶段转变。

2. 批大小选择的关键理论突破

2.1 最小批大小阈值Bmin

论文首次定义了达到目标损失所需的最小批大小阈值Bmin。从几何角度看，它等于E(S)曲线渐近线的斜率。这个物理最小值意味着：

• 当B < Bmin时，无论如何增加训练步数S，模型都无法达到目标损失
• Bmin随着训练进行(目标损失降低)而单调递增

这个发现颠覆了"只要训练足够久，小批量也能达到好效果"的传统认知。在大模型训练中，某些阶段必须使用足够大的批大小才能突破训练瓶颈。

2.2 最优批大小Bopt

在考虑计算效率时，论文提出了最优批大小Bopt的概念，定义为使总数据消耗E最小的批大小。几何上，它是从原点到E(S)曲线最小值点连线的斜率。

Bopt的物理意义在于：

• 使用Bopt能在给定计算预算下最大化数据效率
• 与Bmin类似，Bopt也随训练进程单调递增
• 实际批大小应在Bmin和Bopt之间权衡选择

实操建议：在资源允许的情况下，批大小应尽可能接近Bopt以获得最佳训练效率。当显存受限时，也不应低于Bmin。

3. 动态批大小调度策略

3.1 固定批大小的局限性

基于Bmin和Bopt都随训练进程增加的特性，论文指出固定批大小策略存在根本缺陷：

1. 早期阶段：使用过大的批大小会浪费计算资源
2. 后期阶段：批大小不足会限制模型继续优化
3. 整体效率：无法适应不同训练阶段的需求变化

3.2 动态调度算法设计

论文提出的动态批大小调度器根据已消耗的数据总量分阶段调整批大小。以Qwen3实验为例，采用的策略是：

初始批大小 = 2M tokens 在每消耗125B tokens后： 批大小序列 = [2M, 4M, 5M, 6M]

这种设计的关键考量包括：

1. 阶段划分基于数据消耗量而非训练步数
2. 批大小增幅遵循Bopt的增长曲线
3. 阶段过渡平滑避免训练不稳定

3.3 实际效果验证

在Qwen3 Dense和MoE模型上的实验表明，动态调度策略相比固定批大小：

1. 训练效率提升：达到相同损失所需总计算量减少15-20%
2. 模型质量提高：MMLU准确率提升1.2%，CMMLU提升0.8%
3. 训练稳定性：损失曲线更平滑，梯度噪声控制更好

4. 缩放定律与批大小的关系

4.1 神经缩放定律概述

缩放定律描述了大模型性能随规模增长的统计规律，主要有三种形式：

1. 模型规模缩放：性能 ∝ N^α (N为参数量)
2. 数据量缩放：性能 ∝ D^β (D为训练数据量)
3. 计算量缩放：性能 ∝ C^γ (C为计算预算)

批大小选择与这三种缩放形式都密切相关，特别是计算量缩放。计算预算C通常表示为：

C ∝ B × S

4.2 批大小与学习率的协调缩放

在大规模训练中，批大小B和学习率η需要协调调整。经验法则是：

η ∝ √B

这意味着当批大小增加k倍时，学习率应增加√k倍。这种缩放关系源于梯度噪声与批大小的统计关系。

注意事项：这个规则适用于批大小在合理范围内(通常不超过1M tokens)。当批大小极大时，需要更精细的调整策略。

5. 实操建议与经验分享

5.1 批大小设置的实用指南

基于论文发现和实际经验，建议采用以下策略：

1. 初期阶段：使用较小批大小(如0.5-2M tokens)，配合充分warmup
2. 中期阶段：根据验证损失动态调整批大小，目标保持在Bopt附近
3. 后期阶段：逐步增大批大小，但需监控梯度方差

5.2 常见问题排查

1. 训练不稳定：

   • 检查批大小增幅是否过大
   • 确保学习率与批大小协调调整
   • 验证梯度裁剪阈值是否适当

2. 收敛速度慢：

   • 确认当前批大小不低于Bmin
   • 检查学习率是否与批大小匹配
   • 评估是否需要调整WSD调度器参数

3. 显存不足：

   • 考虑使用梯度累积模拟更大批大小
   • 尝试ZeRO优化器或混合精度训练
   • 在数据并行和模型并行间取得平衡

5.3 高级技巧

1. 渐进式调整：批大小调整采用"增加-稳定"的交替模式，每次增幅不超过20%
2. 动态监控：实时跟踪梯度噪声尺度(Gradient Noise Scale)指导调整
3. 混合策略：在不同模型层使用差异化批大小，重点关注关键层

在实际训练Qwen系列模型时，我发现当批大小超过4M tokens后，需要特别注意以下细节：

• 学习率warmup阶段应延长30-50%
• 梯度裁剪阈值应放宽1.5-2倍
• 每步训练时间监控更为关键，避免因同步开销导致效率下降