Optimum + ONNX Runtime: 更容易、更快地訓練你的 Hugging Face 模型

HuggingFace發表於2023-03-03

介紹

基於語言、視覺和語音的 Transformer 模型越來越大,以支援終端使用者複雜的多模態用例。增加模型大小直接影響訓練這些模型所需的資源,並隨著模型大小的增加而擴充套件它們。Hugging Face 和微軟的 ONNX Runtime 團隊正在一起努力,在微調大型語言、語音和視覺模型方面取得進步。Hugging Face 的 ? Optimum 庫,透過和 ONNX Runtime 的整合進行訓練,為許多流行的 Hugging Face 模型提供了一個開放的解決方案,可以將訓練時間縮短 35% 或更多。我們展現了 Hugging Face Optimum 和 ONNX Runtime Training 生態系統的細節,效能資料突出了使用 Optimum 庫的好處。

效能測試結果

下面的圖表表明,當 使用 ONNX Runtime 和 DeepSpeed ZeRO Stage 1 進行訓練時,用 Optimum 的 Hugging Face 模型的加速 從 39% 提高到 130%。效能測試的基準執行是在選定的 Hugging Face PyTorch 模型上進行的,第二次執行是隻用 ONNX Runtime 訓練,最後一次執行是 ONNX Runtime + DeepSpeed ZeRO Stage 1,圖中顯示了最大的收益。基線 PyTorch 執行所用的最佳化器是 AdamW Optimizer,ORT 訓練用的最佳化器是 Fused Adam Optimizer。這些執行是在帶有 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 萬個免費和可訪問的模型 checkpoints 用於各種機器學習任務,1.8 萬個資料集和 2 萬個機器學習演示應用。然而,將 Transformer 模型擴充套件到生產中仍然是工業界的一個挑戰。儘管準確性很高,但基於 Transformer 的模型的訓練和推理可能耗時且昂貴。

為了滿足這些需求,Hugging Face 構建了兩個開源庫: ? Accelerate? Optimum? Accelerate 專注於開箱即用的分散式訓練,而 ? Optimum 作為 Transformer 的擴充套件,透過利用使用者目標硬體的最大效率來加速模型訓練和推理。Optimum 整合了機器學習加速器如 ONNX Runtime,和專業的硬體如 英特爾的 Habana Gaudi,因此使用者可以從訓練和推理的顯著加速中受益。此外,? Optimum 無縫整合了其他 Hugging Face 的工具,同時繼承了 Transformer 的易用性。開發人員可以輕鬆地調整他們的工作,以更少的計算能力實現更低的延遲。

ONNX Runtime 訓練

ONNX Runtime 加速 大型模型訓練,單獨使用時將吞吐量提高 40%,與 DeepSpeed 組合後將吞吐量提高 130%,用於流行的基於 Hugging Face Transformer 的模型。ONNX Runtime 已經整合為 ? Optimum 的一部分,並透過 Hugging Face 的 ? Optimum 訓練框架實現更快的訓練。

ONNX Runtime Training 透過一些記憶體和計算最佳化實現了這樣的吞吐量改進。記憶體最佳化使 ONNX Runtime 能夠最大化批大小並有效利用可用的記憶體,而計算最佳化則加快了訓練時間。這些最佳化包括但不限於,高效的記憶體規劃,核心最佳化,適用於 Adam 最佳化器的多張量應用 (將應用於所有模型引數的按元素更新分批到一個或幾個核心啟動中),FP16 最佳化器 (消除了大量用於主機記憶體複製的裝置),混合精度訓練和圖最佳化,如節點融合和節點消除。ONNX Runtime Training 支援 NVIDIAAMD GPU,並提供自定義操作的可擴充套件性。

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

Optimum 中的 ONNX Runtime Training

Optimum 提供了一個 ORTTrainer API,它擴充套件了 Transformer 中的 Trainer,以使用 ONNX Runtime 作為後端進行加速。ORTTrainer 是一個易於使用的 API,包含完整的訓練迴圈和評估迴圈。它支援像超引數搜尋、混合精度訓練和多 GPU 分散式訓練等功能。ORTTrainer 使 AI 開發人員在訓練 Transformer 模型時能夠組合 ONNX Runtime 和其他第三方加速技術,這有助於進一步加速訓練,並充分發揮硬體的作用。例如,開發人員可以將 ONNX Runtime Training 與 Transformer 訓練器中整合的分散式資料並行和混合精度訓練相結合。此外,ORTTrainer 使你可以輕鬆地將 DeepSpeed ZeRO-1 和 ONNX Runtime Training 組合,透過對最佳化器狀態進行分割槽來節省記憶體。在完成預訓練或微調後,開發人員可以儲存已訓練的 PyTorch 模型,或使用 ? Optimum 實現的 API 將其轉為 ONNX 格式,以簡化推理的部署。和 Trainer 一樣,ORTTrainer 與 Hugging Face Hub 完全整合: 訓練結束後,使用者可以將他們的模型 checkpoints 上傳到 Hugging Face Hub 賬戶。

因此具體來說,使用者應該如何利用 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 Training 團隊合作,為更新和更大的模型架構帶來更多的訓練最佳化,包括 Whisper 和 Stable Diffusion。微軟還將其最先進的訓練加速技術打包在 PyTorch 的 Azure 容器 中。這是一個輕量級的精心營造的環境,包括 DeepSpeed 和 ONNX Runtime,以提高 AI 開發者使用 PyTorch 訓練的生產力。除了大型模型訓練外,ONNX Runtime Training 團隊還在為邊緣學習構建新的解決方案——在記憶體和電源受限的裝置上進行訓練。

準備開始

我們邀請你檢視下面的連結,以瞭解更多關於 Hugging Face 模型的 Optimum ONNX Runtime Training,並開始使用。

感謝閱讀!如果你有任何問題,請透過 Github論壇 隨時聯絡我們。你也可以在 TwitterLinkedIn 上聯絡我。

原文: https://hf.co/blog/optimum-onnxruntime-training

作者: Jingya、Kshama Pawar (guest)、Vincent Wang (guest)、Zhijiang Xu (guest)

譯者: AIboy1993 (李旭東)

審校、排版: zhongdongy (阿東)

相關文章