Hugging Face Optimum と ONNX Runtime のトレーニング統合

Hugging Face と Microsoft は ONNX Runtime を Optimum ライブラリに統合し、多くの人気 Hugging Face モデルのトレーニング時間を 35％ 以上短縮するオープンなソリューションを提供しています。このコラボレーションは、大規模な言語、音声、ビジョンモデルのファインチューニングに必要な計算リソースと時間を削減することを目的としています。

トレーニング性能向上

ONNX Runtime を Optimum ライブラリに統合することで、Hugging Face モデルの大幅な加速が実現します。8 GPU を搭載した単一の Nvidia A100 ノードで実施したベンチマークでは、ONNX Runtime と DeepSpeed ZeRO Stage 1 を組み合わせた場合、スループットが 39％ から 130％ 向上することが示されました。

ベンチマークの主要な構成は次のとおりです:

  • Baseline: AdamW オプティマイザを使用した PyTorch。
  • Accelerated: Fused Adam Optimizer を使用した ONNX Runtime。
  • Maximum Gain: ONNX Runtime + DeepSpeed ZeRO Stage 1。

Environment Specifications:

  • 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: すべてのモデルパラメータに対する要素ごとの更新を数回のカーネル起動にバッチ化します。
    • 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 により、開発者は Distributed Data Parallel (DDP) や DeepSpeed ZeRO-1（メモリ節約のためにオプティマイザ状態を分割）といった他の加速手法と ONNX Runtime を組み合わせて利用できます。トレーニング後、モデルは PyTorch モデルとして保存するか、最適化された推論デプロイのために ONNX 形式に変換できます。

標準の Trainer から ORTTrainer へ移行するには、主に次の 2 つの変更が必要です:

  1. Trainer を ORTTrainer に置き換える。
  2. TrainingArguments を ORTTrainingArguments に置き換え、adamw_ort_fused オプティマイザなど ORT 固有の機能にアクセスできるようにする。

今後のロードマップ

Hugging Face と Microsoft は、Stable Diffusion や Whisper などのより大規模なモデルアーキテクチャにこれらのトレーニング最適化を展開するために協力しています。さらに、Microsoft は DeepSpeed と ONNX Runtime を含む、PyTorch 開発者の生産性を向上させる Azure Container for PyTorch を導入しました。また、メモリと電力が制限されたデバイス上でのトレーニングに焦点を当てた「エッジでの学習」ソリューションの開発も進められています。

Sources