Hugging Face Transformers 梯度累积修复

Hugging Face 已修正 transformers Trainer 中的数学差异,之前梯度累积的结果并未等同于完整批次训练。此修复确保损失是基于累积步骤中所有批次的非填充 token 总数来计算,而不是仅对每批的损失值进行平均。

梯度累积错误的根本原因

梯度累积的目标是与使用更大批次大小的训练在数学上完全相同。然而,发现当打开或关闭梯度累积时,损失并不匹配。

该问题源于 transformers 模型代码中提供的“默认”损失函数。这些函数会在向模型传入 labelsinput_ids 时自动触发,简化了用户的 API。对于诸如因果语言建模(Causal Language Modeling,Causal LM)等 token 级任务,正确的损失计算方式是将梯度累积步骤中所有批次的总损失除以这些批次中非填充 token 的总数。之前的实现是对每批的损失值取平均,这在该特定用例下是数学错误的。

修复的技术实现

为了解决精度和计算错误,Hugging Face 修改了损失计算逻辑。核心改动是将默认的均值 reduction 改为求和 reduction,然后除以总项目数 (num_items):

# Corrected loss calculation
loss = nn.functional.cross_entropy(shift_logits, shift_labels, ignore_index=-100, reduction="sum")
loss = loss / num_items

长期架构变化

Hugging Face 正在实施两项主要结构性改动,以防止未来的损失计算问题并提升灵活性:

1. 自动损失校正

对于使用默认损失函数的用户,库将在梯度累积期间自动应用必要的更改,确保报告和实际使用的损失准确无误。

2. 可定制的损失 API

为防止内部库问题阻塞用户,Hugging Face 正在引入一个 API,允许用户直接向 Trainer 传入自定义损失函数。

所有继承自 PreTrainedModel 的模型现在都包含一个 loss_function 属性。该属性由 config.loss_type 决定,用户可以通过修改 LOSS_MAPPING 来自定义损失:

def my_super_loss(logits, labels):
    return nn.functional.cross_entropy(logits, labels, ignore_index=-100)

LOSS_MAPPING["my_loss_type"] = my_super_loss

部署与可用性

此修复通过两个特定的 Pull Request 推出:

  • PR #34191:为最流行的模型实现第一项改动。
  • PR #34198:实现第二项改动,允许用户提供自定义损失函数并利用每批次看到的样本数量进行计算。

用户可以通过从 main 分支安装 transformers 库立即获取这些修复:

pip install git+https://github.com/huggingface/transformers

Sources