Hugging Face Transformers Gradient Accumulation Fix

Hugging Face has corrected a mathematical discrepancy in the transformers Trainer where gradient accumulation did not produce results equivalent to full batch training. This fix ensures that losses are calculated based on the total number of non-padding tokens across all batches in an accumulation step, rather than simply averaging per-batch loss values.

The Root Cause of Gradient Accumulation Errors

Gradient accumulation is intended to be mathematically identical to training with a larger batch size. However, a discrepancy was discovered where losses did not match when gradient accumulation was toggled on or off.

This issue stemmed from the "default" loss functions provided within the modeling code of transformers models. These functions are automatically triggered when labels and input_ids are passed to the model, simplifying the API for users. For token-level tasks such as Causal Language Modeling (Causal LM), the correct loss calculation requires the total loss across all batches in a gradient accumulation step to be divided by the total number of non-padding tokens across those batches. The previous implementation was calculating the average of per-batch loss values, which is mathematically incorrect for this specific use case.

Technical Implementation of the Fix

To resolve the precision and calculation errors, Hugging Face modified the loss computation logic. The core change involves switching from a default mean reduction to a sum reduction, then dividing by the total number of items (num_items):

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

Long-term Architectural Changes

Hugging Face is implementing two primary structural changes to prevent future loss calculation issues and increase flexibility:

1. Automatic Loss Correction

For users utilizing the default loss functions, the library will automatically apply the necessary changes during gradient accumulation to ensure the reported and utilized loss is accurate.

2. Customizable Loss API

To prevent internal library issues from blocking users, Hugging Face is introducing an API that allows users to pass their own loss functions directly to the Trainer.

All models inheriting from PreTrainedModel now include a loss_function property. This property is determined by config.loss_type, allowing users to customize the loss by modifying the 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

Deployment and Availability

The fix is being rolled out via two specific Pull Requests:

• PR #34191: Implements the first change for the most popular models.
• PR #34198: Implements the second change, allowing users to provide their own loss functions and utilize the number of samples seen per-batch for calculations.

Users can access these fixes immediately by installing the transformers library from the main branch:

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