使用 TensorFlow 與 XLA 加速文字生成

TL;DR

Hugging Face 已在 TensorFlow 的 transformers 套件中啟用 XLA（Accelerated Linear Algebra）編譯以進行文字生成。此優化可將生成速度提升至多達 100 倍，且在許多基準測試中，表現優於 PyTorch 的文字生成任務。

使用 XLA 加速 TensorFlow

XLA 是一個旨在加速 TensorFlow 模型的編譯器，同時也是 JAX 以及某些 PyTorch 實作的基礎。在使用 Eager Execution 以提升透明度與除錯便利性的 TensorFlow 2 中，部分圖模式的效能優勢會遺失。為了恢復這些優勢，使用者可以將函式包裝在 tf.function 中，將程式碼轉換為圖形。

透過在 tf.function 或 tf.keras.Model.compile 中加入 jit_compile=True 參數，使用者即可觸發 XLA 編譯。雖然首次呼叫 XLA 編譯的函式會因編譯過程而較慢，但對於相同張量形狀與類型的後續呼叫則會顯著加速。

XLA 文字生成的實作需求

XLA 依賴即時（JIT）編譯與多型性。為避免在文字生成過程中產生昂貴的重新編譯（追蹤），必須滿足以下技術需求：

輸入填充

由於 XLA 在遇到不同的張量形狀、類型或非張量參數時會觸發新的編譯步驟，輸入提示必須填充至一致的長度。Hugging Face 建議在 tokenizer 類別中使用 pad_to_multiple_of 參數，以在保持輸入彈性的同時限制可能的形狀數量。

程式碼庫向量化

自回歸文字生成本質上是動態的，常會擴展張量並使用動態切片，這對 XLA 並不友善。為了支援 XLA，Hugging Face 重新編寫了 TensorFlow 文字生成的程式碼庫，將操作向量化並使用帶填充的固定大小結構。此外，NLP 模型也被修改，以確保位置嵌入在這些填充結構下能正確運作。

Transformers 的文字生成功能

transformers 套件中的 generate 函式支援多種解碼策略：

  • Greedy Decoding（貪婪解碼）: 預設的確定性方法（do_sample=False），在每一步選取最可能的 token。
  • Sampling（抽樣）: 一種隨機方法（do_sample=True），可透過 temperature 設定控制隨機性。較低的值會偏好高機率的 token，而較高的值則會增加熵。
  • Beam Search（束搜索）: 當 num_beams 大於 1 時啟用，此方法探索高機率序列，以提升相較於貪婪解碼的輸出品質。

效能基準測試

比較 TensorFlow 與 PyTorch 在多種 GPU 型號上的基準測試顯示兩個主要結果：

  1. Massive Speedups（巨幅加速）: 在使用 XLA 時，TensorFlow 的文字生成速度顯著提升，某些情況下加速幅度超過 100 倍。
  2. Framework Comparison（框架比較）: 在絕大多數情況下，使用 XLA 的 TensorFlow 是最快的選擇，有時在文字生成任務上比 PyTorch 快上至 9 倍。

Sources