Optimum + ONNX Runtime: 更輕鬆、更快速地訓練您的 Hugging Face 模型

簡介

語言、視覺和語音領域基於 Transformer 的模型正變得越來越大，以支援面向終端使用者的複雜多模態用例。模型大小的增加直接影響了訓練這些模型以及隨著模型規模增大而擴充套件它們所需的資源。Hugging Face 和微軟的 ONNX Runtime 團隊正在合作，致力於在微調大型語言、語音和視覺模型方面取得進展。Hugging Face 的 Optimum 庫透過與 ONNX Runtime 訓練的整合，為許多流行的 Hugging Face 模型提供了開放的解決方案，能夠將訓練時間縮短 35% 或更多。我們在此詳細介紹 Hugging Face Optimum 和 ONNX Runtime 訓練生態系統，並透過效能資料突顯使用 Optimum 庫的好處。

效能結果

下圖顯示，在訓練時使用 ONNX Runtime 和 DeepSpeed ZeRO Stage 1 的情況下，Hugging Face 模型透過 Optimum 實現了 39% 到 130% 的驚人加速。效能測量是在選定的 Hugging Face 模型上進行的，以 PyTorch 作為基準執行，僅使用 ONNX Runtime 進行訓練作為第二次執行，而 ONNX Runtime + DeepSpeed ZeRO Stage 1 作為最終執行，顯示出最大的增益。基準 PyTorch 執行使用的最佳化器是 AdamW 最佳化器，而 ORT 訓練執行則使用 Fused Adam 最佳化器。這些執行是在一個配備 8 個 GPU 的單個 Nvidia A100 節點上執行的。

有關開啟 Optimum 以進行訓練加速的配置設定的更多詳細資訊，請參見此處。這些執行使用的版本資訊如下：

PyTorch: 1.14.0.dev20221103+cu116; ORT: 1.14.0.dev20221103001+cu116; DeepSpeed: 0.6.6; HuggingFace: 4.24.0.dev0; Optimum: 1.4.1.dev0; Cuda: 11.6.2

Optimum 庫

Hugging Face 是一個快速發展的開放社群和平臺，旨在推動優秀機器學習的普及化。繼 Transformers 庫成功之後，我們為了滿足社群的需求，將模態從 NLP 擴充套件到音訊和視覺，現在已覆蓋機器學習的各種用例。目前在 Hugging Face Hub 上，有超過 12 萬個免費且可訪問的各種機器學習任務的模型檢查點，1.8 萬個資料集，以及 2 萬個機器學習演示應用。然而，將 Transformer 模型規模化地應用於生產環境仍然是業界的一大挑戰。儘管精度很高，但基於 Transformer 的模型的訓練和推理可能耗時且昂貴。

為了滿足這些需求，Hugging Face 構建了兩個開源庫：Accelerate 和 Optimum。其中，🤗 Accelerate 專注於開箱即用的分散式訓練，而 🤗 Optimum 作為 Transformers 的擴充套件，透過最大化利用使用者目標硬體的效率來加速模型訓練和推理。Optimum 集成了 ONNX Runtime 等機器學習加速器以及 Intel 的 Habana Gaudi 等專用硬體，因此使用者可以在訓練和推理中獲得顯著的速度提升。此外，Optimum 在繼承與 Transformers 相同易用性的同時，無縫集成了 Hugging Face 的其他工具。開發者可以輕鬆調整他們的工作，以更少的計算能力實現更低的延遲。

ONNX Runtime 訓練

ONNX Runtime 加速大型模型訓練，對於流行的 HuggingFace Transformer 模型，單獨使用可將吞吐量提高達 40%，與 DeepSpeed 結合使用時可提高達 130%。ONNX Runtime 已經整合到 Optimum 中，透過 Hugging Face 的 Optimum 訓練框架實現更快的訓練。

ONNX Runtime 訓練透過多種記憶體和計算最佳化實現瞭如此高的吞吐量提升。記憶體最佳化使 ONNX Runtime 能夠最大化批處理大小並高效利用可用記憶體，而計算最佳化則加快了訓練時間。這些最佳化包括但不限於高效的記憶體規劃、核心最佳化、用於 Adam 最佳化器的多張量應用 (multi tensor apply) (它將應用於模型所有引數的逐元素更新批次化為一次或幾次核心啟動)、FP16 最佳化器 (它消除了大量裝置到主機的記憶體複製)、混合精度訓練以及圖最佳化，如圖節點融合和節點消除。ONNX Runtime 訓練支援 NVIDIA 和 AMD GPU，並提供自定義運算元的可擴充套件性。

簡而言之，它使 AI 開發者能夠充分利用他們熟悉的生態系統，如 PyTorch 和 Hugging Face，並在他們選擇的目標裝置上使用 ONNX Runtime 的加速功能，從而節省時間和資源。

Optimum 中的 ONNX Runtime 訓練

Optimum 提供了一個 ORTTrainer API，它擴充套件了 Transformers 中的 Trainer，使用 ONNX Runtime 作為加速後端。ORTTrainer 是一個易於使用的 API，包含功能完備的訓練迴圈和評估迴圈。它支援超引數搜尋、混合精度訓練和多 GPU 分散式訓練等功能。ORTTrainer 使 AI 開發者能夠在訓練 Transformers 模型時組合 ONNX Runtime 和其他第三方加速技術，這有助於進一步加速訓練並充分利用硬體效能。例如，開發者可以將 ONNX Runtime 訓練與 Transformers 的 Trainer 中整合的分散式資料並行和混合精度訓練結合起來。此外，ORTTrainer 使得將 ONNX Runtime 訓練與 DeepSpeed ZeRO-1 結合變得容易，後者透過分割槽最佳化器狀態來節省記憶體。預訓練或微調完成後，開發者既可以儲存訓練好的 PyTorch 模型，也可以使用 Optimum 為 ONNX Runtime 實現的 API 將其轉換為 ONNX 格式，以便於推理部署。就像 Trainer 一樣，ORTTrainer 與 Hugging Face Hub 完全整合：訓練結束後，使用者可以將其模型檢查點上傳到他們的 Hugging Face Hub 帳戶。

那麼具體來說，使用者應該如何使用 Optimum 來利用 ONNX Runtime 的訓練加速功能呢？如果您已經在使用 Trainer，您只需修改幾行程式碼即可享受到上述所有改進。主要需要進行兩處替換。首先，將 Trainer 替換為 ORTTrainer，然後將 TrainingArguments 替換為 ORTTrainingArguments，後者包含了訓練器用於訓練和評估的所有超引數。ORTTrainingArguments 擴充套件了 TrainingArguments，以應用一些由 ONNX Runtime 賦能的額外引數。例如，使用者可以應用 Fused Adam 最佳化器以獲得額外的效能提升。以下是一個示例：

-from transformers import Trainer, TrainingArguments
+from optimum.onnxruntime import ORTTrainer, ORTTrainingArguments

# Step 1: Define training arguments
-training_args = TrainingArguments(
+training_args = ORTTrainingArguments(
    output_dir="path/to/save/folder/",
-   optim = "adamw_hf",
+   optim = "adamw_ort_fused",
    ...
)

# Step 2: Create your ONNX Runtime Trainer
-trainer = Trainer(
+trainer = ORTTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
+   feature="sequence-classification",
    ...
)

# Step 3: Use ONNX Runtime for training!🤗
trainer.train()

展望未來

Hugging Face 團隊正在努力開源更多大型模型，並透過訓練和推理的加速工具降低使用者從中受益的門檻。我們正與 ONNX Runtime 訓練團隊合作，為更新、更大的模型架構 (包括 Whisper 和 Stable Diffusion) 帶來更多的訓練最佳化。微軟還將其最先進的訓練加速技術打包到了 Azure Container for PyTorch 中。這是一個輕量級的精選環境，包括 DeepSpeed 和 ONNX Runtime，旨在提高 AI 開發者使用 PyTorch 進行訓練的生產力。除了大型模型訓練，ONNX Runtime 訓練團隊還在為邊緣學習構建新的解決方案——在記憶體和功耗受限的裝置上進行訓練。

開始使用

我們邀請您檢視以下連結，以瞭解更多關於 Optimum ONNX Runtime 訓練的資訊，並開始為您的 Hugging Face 模型使用它。

• Optimum ONNX Runtime 訓練文件
• Optimum ONNX Runtime 訓練示例
• Optimum Github 倉庫
• ONNX Runtime 訓練示例
• ONNX Runtime 訓練 Github 倉庫
• ONNX Runtime
• DeepSpeed 和 ZeRO 教程
• Azure Container for PyTorch

---

🏎感謝閱讀！如果您有任何問題，歡迎透過 Github 或論壇與我們聯絡。您也可以在 Twitter 或 LinkedIn 上與我聯絡。