Transformers 中的梯度累积：用 TrainingArguments 扩大有效批量大小并正确处理损失缩放

Transformers 中的梯度累积：用 TrainingArguments 扩大有效批量大小并正确处理损失缩放

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers

梯度累积（Gradient Accumulation）是 TransformersTrainer用于在不增加显存开销的前提下扩大“有效批量大小”的核心训练技巧：先对多个 mini-batch 分别做前向与反向传播、把梯度累加到参数上，待累积满gradient_accumulation_steps次后才执行一次optimizer.step()更新权重。本篇基于官方文档 docs/source/en/grad_accumulation.md 展开，并结合 src/transformers/trainer.py 与 src/transformers/training_args.py 的源码，讲清楚梯度累积在训练循环中的真实执行流程、TrainingArguments相关参数语义，以及自定义损失函数中通过num_items_in_batch按 token 归一化损失这一最容易被忽略的细节。读完本文，你可以正确配置梯度累积、理解日志/保存频率的“步数”含义，并能写出在梯度累积下仍然数值正确的自定义损失。

一、核心机制：梯度跨 n 个 mini-batch 累加后才更新权重

大批量会产生庞大的激活值，很快耗尽 GPU 显存。梯度累积的思路是：把一个大 batch 的梯度计算摊到多个 mini-batch 上。梯度先在 n 个 mini-batch 之间累加，然后优化器才更新一次权重。例如单设备批量大小为 8、累积 4 步时，有效批量大小就是 32。原文档给出的执行流程如下：

Step 1: mini-batch 1 → forward → backward → grads = G₁ Step 2: mini-batch 2 → forward → backward → grads = G₁ + G₂ Step 3: mini-batch 3 → forward → backward → grads = G₁ + G₂ + G₃ Step 4: mini-batch 4 → forward → backward → grads = G₁ + G₂ + G₃ + G₄ → optimizer.step() ← same update as if batch_size × 4 → zero_grad()

需要明确的适用边界是：只有当更大的 batch 放不进显存时才使用梯度累积。相比“一次性喂入真正的大 batch”，它并不会带来吞吐量上的收益，只是用时间换显存的空间手段。

1.1 有效批量大小公式与配置

TrainingArguments中的参数文档（src/transformers/training_args.py 中gradient_accumulation_steps字段定义）给出了精确公式：

Effective batch size =per_device_train_batch_size × num_devices × gradient_accumulation_steps

配置示例（继承自原文档）：

from transformers import TrainingArguments args = TrainingArguments( ..., per_device_train_batch_size=8, gradient_accumulation_steps=4, )

一个必须记住的“步数”语义：在TrainingArguments文档中明确指出——使用梯度累积时，一次“step”指一次带反向传播的 mini-batch 前反向，因此日志、评估和保存会在每gradient_accumulation_steps × xxx_step个训练样本之后发生一次。也就是说logging_steps=50在gradient_accumulation_steps=4时对应的是 50 个优化器更新（200 次反向传播）。

二、Trainer 源码中的训练循环：外层优化器步 + 内层 mini-batch 循环

从源码结构看，Trainer.train 把 epoch 迭代器“分块”为梯度累积步，形成两层循环，与上文流程图一一对应。

外层循环：每个优化器步预取 n 个 mini-batch。先把一个 epoch 的总 mini-batch 数steps_in_epoch对gradient_accumulation_steps取余，得到最后一个不完整块的 batch 数（余数为 0 时即完整的 n 个）：

# We chunkify the epoch iterator into gradient accumulation steps `n` batches remainder = steps_in_epoch % self.args.gradient_accumulation_steps if remainder == 0: remainder = self.args.gradient_accumulation_steps

每个外层迭代通过get_batch_samples一次预取n个 mini-batch（最后一步取remainder个），并同步算出num_items_in_batch：

for update_step in range(num_update_steps_trained, num_update_steps_per_epoch): num_batches = ( self.args.gradient_accumulation_steps if update_step != (num_update_steps_per_epoch - 1) else remainder ) batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches, self.args.device) # This is used to correctly scale the loss when the last accumulation step has fewer batches. # Not used if `num_items_in_batch` is not None. self.current_gradient_accumulation_steps = len(batch_samples)

内层循环：逐 mini-batch 前向 + 反向，最后一个才同步、裁剪、更新。几个关键实现细节：

• 跳过中间步的分布式同步。在最后一个 mini-batch 之外的迭代中，用accelerator.no_sync包裹training_step，避免每次反向传播都做跨进程梯度 all-reduce，只在最后一步（以及 DeepSpeed、sync_each_batch场景）才进入同步上下文：

# We sync the gradients in the following cases: 1. sync_each_batch set to True # 2. Using deepspeed 3. when we are at the last batch sample if ( self.accelerator.gradient_state.plugin_kwargs.get("sync_each_batch", False) or self.accelerator.distributed_type == DistributedType.DEEPSPEED or i == len(batch_samples) - 1 ): sync_context = contextlib.nullcontext else: sync_context = functools.partial(self.accelerator.no_sync, model=model) with sync_context(): tr_loss_step = self.training_step(model, inputs, num_items_in_batch)

• 只在同步步执行梯度裁剪与优化器更新。满足do_sync_step时才依次做max_grad_norm裁剪（_clip_grad_norm）、optimizer.step()、学习率调度lr_scheduler.step()、model.zero_grad()，并触发_maybe_log_save_evaluate；中间步只回调on_substep_end。这正是流程图末尾optimizer.step() → zero_grad()的落点。

• 最后一个不完整块的损失修正。若 epoch 的 mini-batch 数不能被累积步数整除，最后一个优化器步实际只有remainder个 mini-batch；current_gradient_accumulation_steps = len(batch_samples)记录真实数量，用于按实际数量归一化损失（当没有走num_items_in_batch路径时）。

• DeepSpeed 特例。DeepSpeed 引擎自身管理梯度缩放，因此training_step中对 DeepSpeed 传入scale_wrt_gas=False，关闭 Trainer 侧针对梯度累积的损失缩放，避免双重缩放。

三、损失缩放：用 num_items_in_batch 按 token 归一化

这是原文档中最具实战价值的一节。当自定义损失函数通过compute_loss_func传入Trainer时，应当接收并使用num_items_in_batch，让 [Trainer] 用“所有 mini-batch 中预测目标的总个数”来归一化损失，而不是用固定的gradient_accumulation_steps计数。原文档的示例：

import torch.nn.functional as F def compute_loss(outputs, labels, num_items_in_batch=None): logits = outputs["logits"] loss = F.cross_entropy(logits, labels, reduction="sum") return loss / num_items_in_batch

3.1 谁来决定是否计算 num_items_in_batch

Trainer._get_num_items_in_batch只有在以下条件同时满足时才计数：预取的 batch 非空、batch 中带有labels，且模型接受损失关键字参数（model_accepts_loss_kwargs）或定义了compute_loss_func。计数方式非常直接——统计非-100（labels.ne(-100)）的标签个数之和：

num_items_in_batch = sum(labels.ne(-100).sum() for labels in labels_for_count)

在多设备场景下，average_tokens_across_devices（默认True）为真时会用accelerator.gather(...).sum()把各设备的 token 数汇总求和，从而在全量数据上做精确的按 token 归一化；未开启该选项的多卡 DataParallel 场景则退化为按 GPU 数整除的近似。

3.2 因果语言模型：统计的是“移位后”的 labels

对于 causal LM 模型，num_items_in_batch统计的是移位后的 labels。原因来自损失函数的构造：因果 LM 损失把 labels 右移一位，使位置i的预测目标是i + 1位置的 token，于是每条序列的第 0 位都没有预测目标。为了和损失真正覆盖的目标数保持一致，Trainer统计的是labels[..., 1:]上的有效 token。源码中的选择逻辑：

labels_for_count = [ batch["shift_labels"] if "shift_labels" in batch else batch["labels"][..., 1:] if self._loss_shifts_labels else batch["labels"] for batch in batch_samples ]

三段优先级依次是：

1. 若数据 collator 直接提供shift_labels张量（例如 padding-free collator），Trainer直接对该张量计数；
2. 否则若该模型的损失会移位 labels（self._loss_shifts_labels），对labels[..., 1:]计数；
3. 其他损失类型（masked LM、分类等）统计完整的 labels 张量。

而_loss_shifts_labels的判定从源码结构看相当严谨：在Trainer.__init__中，它通过检查模型实际使用的loss_type是否经LOSS_MAPPING路由到ForCausalLMLoss来确定，并显式排除 encoder-decoder 模型——因为 encoder-decoder 的目标与（右移后的）decoder_input_ids对齐，每个非-100标签都是预测目标，若套用labels[..., 1:]的计数规则会少算、导致损失被错误放大。

3.3 不走 num_items_in_batch 路径时的默认缩放

如果模型不接受损失关键字参数、也没有自定义损失函数（即num_items_in_batch未被计算/使用），Trainer.training_step会在 backward 之前兜底缩放：

# Finally we need to normalize the loss for reporting if GA loss bug is not fixed # during compute loss if (not self.model_accepts_loss_kwargs or num_items_in_batch is None) and self.compute_loss_func is None: # If the model does not accept loss kwargs, we need to normalize the loss # by the number of gradient accumulation steps loss = loss / self.current_gradient_accumulation_steps

这正是原文档所说“否则 [Trainer] 会除以gradient_accumulation_steps”的实现。两种方式的差异在于：按固定步数除法假设每个 mini-batch 的预测目标数相同；而按num_items_in_batch除法对变长序列、带 padding 的批数据在数值上更精确。此外compute_loss的文档字符串也提醒：若你重写了compute_loss却不使用num_items_in_batch，应手动把self.model_accepts_loss_kwargs置为False，否则梯度累积下的损失可能略有偏差。

四、参数速查与使用建议

注意适用前提：以上行为均以当前仓库源码中的Trainer实现为准；使用 DeepSpeed 时损失缩放交由引擎处理（scale_wrt_gas=False），手动在损失中再除以累积步数会造成双重缩放。

五、延伸阅读

围绕训练显存的三个相关主题，官方文档给出了配套指南（均位于 docs/source/en/）：

• GPU memory usage：理解训练时 GPU 显存消耗由什么驱动，判断是否真的需要梯度累积；
• Gradient checkpointing：通过重计算激活值而非缓存来降低激活显存，可与梯度累积组合使用；
• Mixed precision training：用低精度数据类型（bf16/fp16）减少显存并加速训练。

原文档还建议参考社区博客理解“梯度累积如何被计算”（Gradient Accumulation Fix，unsloth.ai 的公开博文），该主题即本文第三节所述的损失缩放问题。

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers