Hugging Face Optimum 與 ONNX Runtime 訓練整合

Hugging Face 與 Microsoft 已將 ONNX Runtime 整合至 Optimum 函式庫，提供一個開放的解決方案，使許多流行的 Hugging Face 模型的訓練時間提升 35% 以上。此合作旨在降低微調大規模語言、語音與視覺模型所需的計算資源與時間。

訓練效能提升

將 ONNX Runtime 與 Optimum 函式庫結合，可為 Hugging Face 模型帶來顯著加速。在單一 Nvidia A100 節點（8 顆 GPU）上進行的基準測試顯示，將 ONNX Runtime 與 DeepSpeed ZeRO Stage 1 結合時，吞吐量提升介於 39% 至 130% 之間。

這些基準測試的關鍵配置細節包括：

Baseline：使用 AdamW 優化器的 PyTorch。
Accelerated：使用融合 Adam 優化器的 ONNX Runtime。
Maximum Gain：ONNX Runtime + DeepSpeed ZeRO Stage 1。

環境規格：

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 函式庫生態系統

Optimum 是 Transformers 函式庫的擴充套件，旨在透過最大化目標硬體的效能來加速模型的訓練與推論。Accelerate 函式庫專注於分散式訓練，而 Optimum 則整合了像 ONNX Runtime 這樣的機器學習加速器，以及 Intel 的 Habana Gaudi 等專用硬體，以降低延遲並減少計算資源需求。

ONNX Runtime 訓練的技術最佳化

ONNX Runtime（ORT）在單獨使用時可提升吞吐量最高 40%，與 DeepSpeed 結合時則可提升最高 130%。這些效能提升源自多項記憶體與計算的最佳化：

Memory Optimizations：高效的記憶體規劃可最大化批次大小並更佳利用可用記憶體。
Compute Optimizations：
- Kernel Optimizations：一般性的執行速度提升。
- Multi Tensor Apply for Adam Optimizer：將所有模型參數的逐元素更新批次化，僅使用少量 kernel 呼叫。
- FP16 Optimizer：減少裝置至主機的記憶體拷貝。
- Mixed Precision Training：提升速度並降低記憶體佔用。
- Graph Optimizations：包括節點融合與節點消除。

ONNX Runtime Training 支援 NVIDIA 與 AMD GPU，並允許使用自訂運算子。

在 Optimum 中實作 ORTTrainer

Optimum 引入 ORTTrainer API，將 Transformers 的 Trainer 擴充以使用 ONNX Runtime 作為後端。此 API 提供完整的訓練與評估迴圈，支援超參數搜尋、混合精度訓練與多 GPU 分散式訓練。

ORTTrainer 允許開發者將 ONNX Runtime 與其他加速技術結合，例如 Distributed Data Parallel（DDP）與 DeepSpeed ZeRO-1（將優化器狀態分割以節省記憶體）。訓練完成後，模型可儲存為 PyTorch 模型或轉換為 ONNX 格式，以進行最佳化的推論部署。

要從標準的 Trainer 移轉至 ORTTrainer，開發者需要進行兩項主要變更：

將 Trainer 替換為 ORTTrainer。
將 TrainingArguments 替換為 ORTTrainingArguments，以取得 ORT 專屬功能，例如 adamw_ort_fused 優化器。

未來路線圖

Hugging Face 與 Microsoft 正合作將這些訓練最佳化擴展至更大型的模型架構，包括 Stable Diffusion 與 Whisper。此外，Microsoft 推出了 Azure Container for PyTorch，這是一個精心打造的環境，內含 DeepSpeed 與 ONNX Runtime，以提升 PyTorch 開發者的生產力。亦正積極開發「邊緣學習」解決方案，專注於在記憶體與電力受限的裝置上進行訓練。