Accelerate 怎样配置梯度累积

我在显存比较紧的训练任务里，最常用的办法不是先把模型换小，而是把一个“大 batch”拆成多个 micro-batch，等梯度累积到指定次数后再更新一次参数。Accelerate 已经把这件事封装好了：在 Accelerator 中设置 gradient_accumulation_steps，再用 accelerator.accumulate(model) 包住一次训练单元即可。

官方地址：https://huggingface.co/docs/accelerate/en/usage_guides/gradient_accumulation

最小配置是“累积次数交给 Accelerator，训练循环交给 accumulate”。普通样本的 loss 不需要再手动除以累积次数；只有变长 token 任务，才需要额外按窗口内的非填充 token 总数处理归一化。

先把有效 batch 算明白

假设每张设备一次能放下 4 条样本，设置 gradient_accumulation_steps=8，那么一次参数更新覆盖的本地有效样本数约为 32 条；如果有多张设备，还要结合进程数理解全局 batch。累积次数增加后，optimizer.step() 的频率会降低，训练日志里的“更新步”也不再等于 dataloader 的迭代次数。

这也是我第一次接入时最容易误判的地方：显存下降不代表每个 micro-batch 都会更新参数。真正的更新边界是累积窗口，学习率 scheduler、保存频率和评估频率都应该围绕这个边界规划。

在 Accelerator 中配置累积次数

把累积次数放进 Accelerator 构造函数，通常就足够覆盖普通的分类、回归或固定长度 batch。随后交给 prepare 的模型、优化器、dataloader 和 scheduler 会被统一接管。

from accelerate import Accelerator

# 每 4 个 micro-batch 才形成一次参数更新
gradient_accumulation_steps = 4
accelerator = Accelerator(
    gradient_accumulation_steps=gradient_accumulation_steps
)

# 让 Accelerate 统一处理设备放置和分布式包装
model, optimizer, train_loader, scheduler = accelerator.prepare(
    model, optimizer, train_loader, scheduler
)

这里的数值不是越大越好。它一方面扩大有效 batch，另一方面会降低参数更新的即时性；如果任务对快速反馈敏感，应该先从能稳定放入显存的较小累积次数开始。

用 accumulate 包住一次训练单元

进入训练循环后，每一轮只保留一次 forward 和一次 backward。accumulate(model) 会根据当前累积位置决定是否同步梯度、是否真正执行优化器更新，并按普通梯度累积规则调整 backward 的 loss。

for batch in train_loader:
    # 一个上下文对应一个 micro-batch，不要在里面再套第二次 forward
    with accelerator.accumulate(model):
        outputs = model(**batch)
        loss = outputs.loss

        # 用 Accelerate 的 backward 适配当前设备和混合精度
        accelerator.backward(loss)

        # 这三行每轮都写，Accelerate 会在正确的累积边界执行更新
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

普通固定长度样本不需要再写 loss = loss / gradient_accumulation_steps。如果手动除一次，又让 Accelerate 自动调整一次，梯度会被重复缩小。这个“不要重复缩放”是配置迁移时最值得保留的检查点。

scheduler 和 dataloader 尾批要单独看

Accelerate 默认会把 scheduler 的步数按累积行为调整，前提是 scheduler 没有在外部提前按更新次数做过同样的换算。若你的 scheduler 已经自行按真实 optimizer 更新次数配置，就要重新确认是否需要关闭或改写这层调整。

默认情况下，Accelerate 会把 dataloader 的迭代状态纳入累积判断，并假设到 dataloader 末尾需要同步。如果数据量不能被累积次数整除，最后一个窗口就可能比常规窗口短。对于希望严格控制尾批语义的训练，可以使用插件：

from accelerate import Accelerator
from accelerate.utils import GradientAccumulationPlugin

# 关闭“到 dataloader 末尾自动同步”的默认假设
plugin = GradientAccumulationPlugin(
    num_steps=4,
    sync_with_dataloader=False,
)
accelerator = Accelerator(gradient_accumulation_plugin=plugin)

# num_steps 与 gradient_accumulation_steps 二选一，不能同时传入

gradient_accumulation_plugin 和 gradient_accumulation_steps 不能同时传入；需要更细控制时用插件替代简单整数配置。

变长 token 样本为什么不能直接平均每个 batch 的 loss

做因果语言模型或其他 token-level 任务时，不同样本的有效 token 数可能不同。一个 batch 的平均 loss 代表的是“这个 batch 的平均”，直接再平均多个 batch，会让短样本和长样本拥有相同权重，结果不等于整个累积窗口的 token 平均损失。

更稳妥的思路是让损失函数先求和，统计累积窗口内所有非填充 token，再用这个总数做归一化。分布式训练时，局部 token 数还需要通过 accelerator.gather 汇总。

# 先按累积窗口收集 batch，才能得到完整的 token 分母
num_items_in_batch = sum(
    (batch["labels"] != -100).sum() for batch in batch_samples
)

# 多设备时汇总所有进程的有效 token 数
num_items_in_batch = accelerator.gather(
    num_items_in_batch
).sum().item()

# reduction="sum" 保留 token 总损失，再按全局 token 数归一化
loss = loss_function(outputs, labels, reduction="sum")
loss = (loss * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch
accelerator.backward(loss)

这个写法只适用于确实按 token 计数的 loss。普通样本级任务继续使用 accumulate(model) 的自动处理即可，不要为了套用 token 归一化而改变原本的损失定义。

我会保留的配置清单

最后的判断

如果你的问题只是显存不够，先用整数形式配置累积次数，并把 forward、backward、step 和 zero_grad 放进 accumulate(model)。如果出现 scheduler 步数不对、尾批行为不一致或 token-level loss 漂移，再进入插件和全局 token 统计层面排查。对我来说，Accelerate 最省心的地方不是少写几行代码，而是把分布式梯度同步和更新边界放到了同一个上下文里。

相关问题

• gradient_accumulation_steps 越大，训练一定越稳定吗？不一定，它会改变更新频率和有效 batch，需要结合学习率与任务收敛观察。
• 可以同时设置 gradient_accumulation_steps 和 gradient_accumulation_plugin 吗？不可以，二者是两种配置入口。
• 为什么变长文本不能直接平均 batch loss？因为每个 batch 的有效 token 数不同，平均 batch loss 会改变 token 权重。