Intel Technologies による PyTorch 分散ファインチューニングの加速

Hugging Face は、Intel Xeon Scalable CPU サーバーのクラスターに PyTorch のファインチューニング・ジョブを分散させることで、トレーニング時間を大幅に短縮できることを実証しました。これは、転移学習タスクにおいて GPU ベースのトレーニングに代わる実行可能な選択肢となります。Intel の Ice Lake アーキテクチャと最適化されたソフトウェア・ライブラリを活用することで、1 ノードから 4 ノードへスケールアップした際に、最大 3 倍のトレーニング速度向上が達成されました。

Intel Ice Lake によるハードウェア加速

最適なパフォーマンスを得るために、このセットアップでは Ice Lake アーキテクチャに基づいた Intel サーバーを利用しています。このハードウェアは、ディープラーニング操作を加速させる特定の機能をサポートしています。

  • Intel AVX-512: ハイパフォーマンス・コンピューティング用の Advanced Vector Extensions。
  • Intel Vector Neural Network Instructions (VNNI): ニューラルネットワークの推論とトレーニングを高速化するために設計された専用命令。

これらの機能は、Amazon EC2 (M6i および C6i インスタンス)、Azure (Dv5, Dsv5, Ddv5, Ddsv5, Edv5, および Edsv5 シリーズ)、および Google Cloud Platform (N2 Compute Engine VMs) を含む主要なクラウド・プロバイダーで利用可能です。

ソフトウェア最適化スタック

Intel CPU のハードウェア機能を最大限に活用するために、2 つの主要なソフトウェア・コンポーネントが PyTorch のワークフローに統合されています。

Intel Extension for PyTorch

Intel extension for PyTorch は、PyTorch が AVX-512 と VNNI を活用できるようにすることで、トレーニングと推論の両方において即座に利用可能な速度向上を提供します。

Intel oneAPI Collective Communications Library (oneCCL)

分散トレーニングでは、大規模なモデルが状態情報を交換する際に、ネットワークのボトルネックが発生することがよくあります。Intel oneAPI Collective Communications Library (oneCCL) は、torch.distributed の通信バックエンドとして使用され、分散トレーニング中にノードを同期させるために不可欠な all-reduce のような通信パターンを効率的に処理します。

パフォーマンス・ベンチマーク

Hugging Face は、MRPC データセット (GLUE ベンチマークの一部) でファインチューニングされた BERT モデルを使用して、異なるクラスター・サイズにおけるトレーニング時間を測定しました。ベースラインは単一ノード (Amazon EC2 c6i.16xlarge インスタンス) で設定されました。

MRPC データセットの結果

クラスター・サイズ トレーニング時間 速度向上
1 ノード 7m 46s ベースライン
2 ノード 4m 39s 1.7x
4 ノード 2m 36s 3x

QQP データセットの結果

Quora Question Pairs (QQP) タスクでは、400,000 サンプルを超えるはるかに大規模なデータセットが含まれますが、速度向上は一貫していました。

| クラスター・サイズ | トレーニング時間 | 速度向上 | | :--- | :--- | :--- | | | 1 ノード | 11h 22m | ベースライン | | 2 ノード | 6h 38m | 1.71x | | 4 ノード | 3h 51m | 2.95x |

実装要件

分散環境のセットアップには、互換性とパフォーマンスを確保するために、特定の構成ステップが必要です。

  • Infrastructure: マスターとワーカー・ノード間でパスワードなしの SSH が構成されており、かつ、内部 oneCCL 通信のためにすべての TCP ポートが開いている、同一のインスタンス。
  • Dependency Matching: PyTorch と Intel extension for PyTorch のバージョンが一致している必要があります (例: PyTorch 1.9.0 と torch_ipex 1.9.0)。
  • Script Modifications: トレーニング・スクリプトは、torch_ccl をインポートし、環境変数 (例: PMI_RANK) を介してローカル・ランクを処理し、CCL バックエンド用のマスター・ノードのアドレスとポートを構成するように更新する必要があります。
  • Execution: ジョブは mpirun を使用して、プロセス数とノードあたりのプロセス数を指定して起動されます。

Sources

関連