Hugging Face Transformers 梯度累積修正

Hugging Face 已修正 transformers Trainer 中的數學差異,先前梯度累積未能產生等同於完整批次訓練的結果。此修正確保損失是根據累積步驟中所有批次的非填充 token 總數來計算,而不是僅僅對每批次的損失值取平均。

梯度累積錯誤的根本原因

梯度累積的目標是與使用更大批次大小的訓練在數學上完全相同。然而,發現了一個差異:在開啟或關閉梯度累積時,損失並不相符。

此問題源於 transformers 模型的建模程式碼中提供的「預設」損失函式。當向模型傳入 labelsinput_ids 時,這些函式會自動觸發,以簡化使用者的 API。對於像因果語言模型(Causal Language Modeling,Causal LM)這類 token 級別的任務,正確的損失計算需要將梯度累積步驟中所有批次的總損失除以這些批次的非填充 token 總數。先前的實作是對每批次的損失值取平均,對此特定情況而言在數學上是不正確的。

修正的技術實作

為了解決精度與計算錯誤,Hugging Face 修改了損失計算邏輯。核心變更是將預設的平均(mean)縮減改為總和(sum)縮減,然後除以項目總數(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