返回首頁

GitHub Repo

Transformers 提出語音模型編譯修補,移除舊梯度檢查點相容邏輯

Transformers 社群於 11 月 10 日提出修補,移除 Wav2Vec2 等十種語音模型中阻擋全圖編譯的梯度設定。提案尚未合併,且會改變重入式梯度檢查點的行為,採用前須確認首層卷積仍能取得梯度。

Marcus Qwertyus · Public domain · Image source
zh-Hant

Transformers 社群於 11 月 9 日回報語音模型訓練的編譯障礙,並於 10 日提交修補:Wav2Vec2 特徵編碼器在未凍結、處於訓練模式時,每次前向運算都會設定中間張量的 requires_grad=True,即使沒有啟用梯度檢查點也會執行,導致 torch.compile(fullgraph=True) 拒絕編譯。問題回報

回報者使用小型隨機初始化模型,只編譯特徵編碼器,再執行反向傳播即可觸發錯誤;環境為 Transformers 5.23.0.dev0 與 PyTorch 2.15.0。這項重現將問題定位在卷積特徵抽取路徑,工程師排查時可先獨立檢查此模組。重現程式

這段邏輯原本用來配合重入式梯度檢查點。檢查點透過反向傳播時重新計算中間結果來節省記憶體,但兩種實作有不同限制:PyTorch 官方建議使用 use_reentrant=False;重入式版本則要求至少一個輸入與輸出需要梯度,否則該區段無法取得梯度。PyTorch 文件

新提案依儲存庫協作者的討論方向,直接刪除相容邏輯與不再使用的旗標,並傳播至 HuBERT、WavLM、SpeechT5、Data2Vec-Audio 等十種模型。代價也寫得明確:若仍使用 use_reentrant=True,首層卷積將不再取得梯度。這使修補同時涉及編譯相容性與訓練設定遷移。修補提案

作者回報重現程式已通過,十組模型測試共有 1,955 項通過、2,438 項略過,另兩項 SpeechT5 整合測試失敗在原始主分支也存在。截至查核,提案仍待審查,CI 亦等待維護者核准。接下來應追蹤合併版本、實際 GPU 訓練效能,以及既有檢查點設定下的梯度完整性;目前尚無證據可量化吞吐量收益。測試與審查狀態

來源

  1. Wav2Vec2-family feature encoder breaks torch.compile(fullgraph=True) — Issue #50527
  2. Removing legacy requires_grad workaround — Pull Request #50544
  3. torch.utils.checkpoint — PyTorch 2.15 documentation