批次大小对训练的影响

说到Transformer训练，批次大小（batch size）绝对是个绕不开的关键参数。我早期调参时，总觉得这东西无非就是调大调小，影响不大。直到有一次我在一个翻译任务上折腾了整整两周，模型就是无法收敛……后来发现，罪魁祸首就是批次大小设错了。嗯，今天我们就来把这东西彻底捋清楚。

基本原理与数学直觉

批次大小，说白了就是每次参数更新时，你给模型看了多少样本。它直接影响三件事：梯度噪声、收敛速度、内存开销。

我们来简单算笔账。假设总样本数是N，批次大小是B。一个完整的数据集遍历（epoch）会有N/B次参数更新。你想想看，B越大，每次更新看到的样本越多，梯度估计越准，但更新次数变少。B越小，梯度噪声大但更新次数多。这两头的取舍，就是调参的艺术所在。

我在项目中遇到过一件事：同样的模型结构，批次大小从128改成32，验证集上的BLEU值直接涨了1.5个点。但代价是训练时间从3天变成了将近5天。所以，这里面没有绝对的好坏，只有权衡。

收敛稳定性与梯度方差

为什么批次大小会影响收敛？这得从随机梯度下降（SGD）的本质说起。

全批次梯度下降（batch size = 全部样本）给出的梯度方向最准确，但计算太慢，而且容易收敛到离初始点最近的局部极小值——这通常是个尖锐的极小值，泛化能力很差。

小批次随机梯度下降呢？每次梯度方向都有噪声。这个噪声反而成了好事，它能让模型跳出尖锐的极小值，落入平坦的极小值区域。平坦极小值的泛化能力往往更好。

你可能会问：“那是不是批次越小越好？”

也不尽然。批次太小，梯度噪声过大，模型可能一直在晃就是无法收敛。我见过一些初学者把批次大小设为2甚至1，结果loss曲线比心电图还夸张。

内存效率与吞吐量

实际工程中，批次大小受显存的直接约束。Transformer的显存消耗主要由三部分构成：模型参数、激活值（activation）、优化器状态。

我们来算算一笔具体的账。以BERT-base为例：

模型参数：~110M 个参数，约 440MB（FP32）
激活值（batch=32, seq_len=512）：~8GB
优化器状态（Adam）：~880MB（动量和二阶矩）
总计：约 9.3GB

这还只是单卡。如果你把批次大小提升到64，激活值几乎翻倍，显存直接飙到17GB以上。很多24GB的显卡就在这个边界上。

我个人习惯的做法是：先用小批次跑通流程，确认代码和预处理没问题，然后逐步增大批次直到显存溢出。这个过程的调试脚本可以这样写：

import torch

def find_max_batch_size(model, sample_input, start=2, max_memory=0.9):
    device = next(model.parameters()).device
    total_memory = torch.cuda.get_device_properties(device).total_memory
    max_bytes = total_memory * max_memory
    
    batch = start
    while True:
        try:
            # 模拟一次前向+反向传播
            output = model(sample_input)
            loss = output.sum()
            loss.backward()
            optimizer.zero_grad()
            
            # 检查当前显存使用
            current = torch.cuda.memory_allocated(device)
            if current > max_bytes:
                return batch // 2
            
            batch *= 2
            sample_input = torch.randn(batch, *sample_input.shape[1:], device=device)
        except RuntimeError:
            return batch // 2

实践策略：梯度累积与混合精度

如果显存实在不够，但又想用大批次的效果，梯度累积是个好办法。说白了，就是拆成多个小批次，把梯度攒起来一次性更新。

# 梯度累积示例
accumulation_steps = 4  # 等效批次大小 = 物理批次大小 * 累积步数
effective_batch = 32
physical_batch = 8  # 32 / 4 = 8

optimizer.zero_grad()
for i, batch in enumerate(dataloader):
    loss = model(batch)
    loss = loss / accumulation_steps  # 梯度归一化
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

我曾经在这个技巧上吃过亏。有一次没做梯度归一化（就是没除以accumulation_steps），结果等效批次虽然大了，但每次更新的有效梯度被放大了accumulation_steps倍，模型直接炸了。那次debug花了我半天时间。

混合精度训练（AMP）也是一个缓解显存压力的利器。FP16的激活值占用显存只有FP32的一半。但注意，Transformer的某些操作（如LayerNorm、Softmax）在FP16下容易溢出，需要保持FP32精度。

# PyTorch AMP 示例
scaler = torch.cuda.amp.GradScaler()

for batch in dataloader:
    with torch.cuda.amp.autocast():
        loss = model(batch)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

批次大小与学习率的线性缩放

这是大家最容易忽视的点。批次大小变了，学习率必须跟着变。

直觉上很好理解：批次越大，每次梯度的方差越小，更新方向的置信度越高，所以可以迈更大的步子。反之，批次小，噪声大，步子得迈小点。

常用的线性缩放法则：

• 基准配置：batch_size = 256, learning_rate = 1e-4
• 缩放公式：lr_new = lr_base * (batch_size_new / batch_size_base)
• 安全修正：如果用Adam等自适应优化器，缩放系数可以减半

举个例子，你从256换到512：

lr_new = 1e-4 * (512 / 256) = 2e-4

但直接翻倍有时会太大。我建议给个warmup阶段，让学习率从很小逐渐升到目标值。尤其是用大批次时，warmup步数也得相应增加。

真实案例：一次失败的调优经历

之前做一个文本摘要任务，我用T5-small，初始批次设了128，学习率2e-4。跑了3个epoch，loss降到2.1就下不去了。验证集ROUGE-L只有24.5。

我试着把批次降到32（没调学习率）。第一波loss居然降到1.8，但训练曲线抖得厉害，一个epoch就要多花3倍时间。

后来我做了三件事：

1. 把批次设为64
2. 学习率降到5e-5
3. warmup从1000步增加到3000步

结果ROUGE-L涨到了27.3。说白了，就是找到了一个平衡点——梯度噪声适中，更新次数够多，学习率匹配。

这张图把批次大小的影响路径理清楚了。从中心出发，左边是梯度相关（稳定性、泛化），中间是吞吐量，右边是内存相关。底部的五个解决方案基本覆盖了工程实践中95%的场景。

最后给个小技巧：如果你刚开始一个新任务，我建议先用32或64跑一个快速实验，观察loss下降趋势和梯度范数。等确定模型能学会东西之后，再根据显存硬件去调整批次大小和学习率。千万别一上来就追求大批次——我见过好多人把时间浪费在调批次上，结果模型结构本身就有bug。嗯，先跑通，再调优。