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_size8, gradient_accumulation_steps4, )参数默认值说明gradient_accumulation_steps1执行一次参数更新前累加梯度的 mini-batch 次数默认 1 即普通训练per_device_train_batch_size视参数而定单卡 mini-batch 大小与累积步数相乘再乘设备数得到有效批量average_tokens_across_devicesTrue是否跨设备用 all_reduce 汇总 token 数以获得精确的按 token 归一化损失一个必须记住的“步数”语义在TrainingArguments文档中明确指出——使用梯度累积时一次“step”指一次带反向传播的 mini-batch 前反向因此日志、评估和保存会在每gradient_accumulation_steps × xxx_step个训练样本之后发生一次。也就是说logging_steps50在gradient_accumulation_steps4时对应的是 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_batchfor 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, modelmodel) 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-batchcurrent_gradient_accumulation_steps len(batch_samples)记录真实数量用于按实际数量归一化损失当没有走num_items_in_batch路径时。DeepSpeed 特例。DeepSpeed 引擎自身管理梯度缩放因此training_step中对 DeepSpeed 传入scale_wrt_gasFalse关闭 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_batchNone): logits outputs[logits] loss F.cross_entropy(logits, labels, reductionsum) return loss / num_items_in_batch3.1 谁来决定是否计算 num_items_in_batchTrainer._get_num_items_in_batch只有在以下条件同时满足时才计数预取的 batch 非空、batch 中带有labels且模型接受损失关键字参数model_accepts_loss_kwargs或定义了compute_loss_func。计数方式非常直接——统计非-100labels.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 ]三段优先级依次是若数据 collator 直接提供shift_labels张量例如 padding-free collatorTrainer直接对该张量计数否则若该模型的损失会移位 labelsself._loss_shifts_labels对labels[..., 1:]计数其他损失类型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否则梯度累积下的损失可能略有偏差。四、参数速查与使用建议场景建议大 batch 放不进显存想扩大有效批量调大gradient_accumulation_steps同时控制per_device_train_batch_size变长序列 自定义损失在compute_loss_func中接收num_items_in_batch用reductionsum后除以该值多设备训练希望损失精确按 token 归一保持average_tokens_across_devicesTrue默认自定义compute_loss不使用num_items_in_batch显式设置self.model_accepts_loss_kwargs False期望吞吐量提升不要指望梯度累积它是显存手段而非加速手段注意适用前提以上行为均以当前仓库源码中的Trainer实现为准使用 DeepSpeed 时损失缩放交由引擎处理scale_wrt_gasFalse手动在损失中再除以累积步数会造成双重缩放。五、延伸阅读围绕训练显存的三个相关主题官方文档给出了配套指南均位于 docs/source/en/GPU memory usage理解训练时 GPU 显存消耗由什么驱动判断是否真的需要梯度累积Gradient checkpointing通过重计算激活值而非缓存来降低激活显存可与梯度累积组合使用Mixed precision training用低精度数据类型bf16/fp16减少显存并加速训练。原文档还建议参考社区博客理解“梯度累积如何被计算”Gradient Accumulation Fixunsloth.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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考