修复梯度累积
我们的朋友 Unsloth 昨天分享了一个关于梯度累积的问题,该问题影响了 transformers Trainer。最初的报告来自 @bnjmn_marie(向他致敬!)。
梯度累积应该在数学上等同于全批次训练;然而,在启用和禁用该设置的训练运行之间,损失不匹配。
问题根源何在?
在每个模型的建模代码中,transformers
提供了一个“默认”损失函数,这是模型任务最常用的损失函数。它由建模类应 SONY 的任务决定:问答、token 分类、因果语言模型、掩码语言模型。
这是默认的损失函数,它不应该被定制:只有当 labels
和 input_ids
作为输入传递给模型时,才会计算它,这样用户就不必计算损失。默认损失函数很有用,但设计上是有限制的:对于任何不同的操作,我们期望标签不直接传递,并且用户从模型中获取 logits 并使用它们在模型外部计算损失。
然而,transformers Trainer 以及许多 Trainer,都严重依赖这些方法,因为它提供了简单性:这是一把双刃剑。提供一个随着用例不同而变得不同的简单 API 并不是一个经过深思熟虑的 API,我们自己也感到惊讶。
准确地说,对于像因果语言模型训练这样的 token 级任务的梯度累积,正确的损失应该通过梯度累积步骤中所有批次的总损失除以这些批次中所有非填充 token 的总数来计算。这与每批次损失值的平均值不同。修复方法很简单,请参阅以下内容
def ForCausalLMLoss(logits, labels, vocab_size, **kwargs):
# Upcast to float if we need to compute the loss to avoid potential precision issues
logits = logits.float()
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
shift_logits = shift_logits.view(-1, vocab_size)
shift_labels = shift_labels.view(-1)
# Enable model parallelism
shift_labels = shift_labels.to(shift_logits.device)
num_items = kwargs.pop("num_items", None)
+ loss = nn.functional.cross_entropy(shift_logits, shift_labels, ignore_index=-100, reduction="sum")
+ loss = loss / num_items
- loss = nn.functional.cross_entropy(shift_logits, shift_labels, ignore_index=-100)
return loss
我们如何修复它
为了解决这个问题,我们正在以两种方式改变模型的训练方式
- 如果用户使用“默认”损失函数,我们将在使用梯度累积时自动考虑所需的更改,以确保报告和使用了正确的损失,从而解决核心问题。
- 为了确保将来计算损失的任何问题都不会阻碍用户,我们将公开一个 API,让用户可以直接将自己的损失函数传递给
Trainer
,这样他们就可以轻松使用自己的修复,直到我们内部修复任何问题并发布新的 transformers 版本。
所有继承自 PreTrainedModel
的模型现在都有一个 loss_function
属性,它由以下任一方式确定:
config.loss_type
:这是为了确保任何人都可以使用他们的自定义损失。您可以通过修改LOSS_MAPPING
来实现这一点
def my_super_loss(logits, labels):
return loss = nn.functional.cross_entropy(logits, labels, ignore_index=-100)
LOSS_MAPPING["my_loss_type"] = my_super_loss
我们正在努力在本次 PR 中为最流行的模型提供第一个更改:https://github.com/huggingface/transformers/pull/34191#pullrequestreview-2372725010。在此之后,将呼吁贡献者帮助将此更改传播到其余模型,以便大多数模型在下一次发布时得到支持。
我们还在积极努力在此 PR 中提供第二个更改:https://github.com/huggingface/transformers/pull/34198,这将允许用户使用自己的损失函数,并利用每批次看到的样本数量来帮助计算其损失(并且随着更多模型得到先前更改的支持,将在梯度累积期间执行正确的损失计算)
—
到明天,您应该会看到 Trainer 在梯度累积方面表现正常。届时请从 main
安装以享受此修复
pip install git+https://github.com/huggingface/transformers
通常,我们对提交到我们的问题跟踪器中的错误报告非常及时响应:https://github.com/huggingface/transformers/issues
这个问题已经在 Transformers 中存在了一段时间,因为它主要是一个应该由最终用户更新的默认设置;然而,当默认设置变得不直观时,它们就注定要被更改。在这种情况下,我们在不到 24 小时内更新了代码并发布了修复,这是我们针对 Transformers 中的此类问题所追求的目标。请您如果有任何问题,随时提交;这是我们改进 Transformers 并使其更好地适应您的不同用例的唯一途径。
Transformers 团队 🤗