インテル® oneAPI ツールキットを用いた分散型 AI トレーニングの最適化

インクリメンタル・チューニングで大幅な性能向上を実現

ディープラーニングのワークロードは、急速なペースで増加しています。ディープラーニング・ベースのアルゴリズムは、大量のデータを処理してパターンを見つけ、画像分類、物体検出、時系列予測などを行います。データの可用性が高まるにつれ、ディープラーニング・モデルの複雑さも増してきました。ResNet (英語) や VGG (Visual Geometry Group) (英語) などのモデルは、数百万個のパラメーターがあり、数十億回の浮動小数点演算を行います。GPT-3 (Generative Pretrained Transformer 3) (英語)や BERT (Bidirectional Encoder Representations from Transformers) (英語) のような最近のモデルは、数十億から 1 兆個のパラメーターがあります。そのため、ディープラーニング・モデルのトレーニングは、計算量が多く、時間がかかるプロセスです。解決までの時間を短縮するには、シングルコアやマルチコアの性能を最適化する以外に、複数のノードへのスケールアウト (分散学習など) も検討する必要があります。この作業は、モデルを分割する方法 (モデル並列)、データを分割する方法 (データ並列)、あるいは両方を組み合わせた方法 (ハイブリッド並列) で実現できます。

この記事では、計算ノードのクラスター上でのディープラーニング・ベースのアルゴリズムのチューニングとスケーリングに焦点を当てています。ここでは、半教師あり敵対的生成ネットワーク (S-GAN) を用いて画像を分類する方法を説明します。また、OpenMP、TensorFlow、インテル® MPI ライブラリーでさまざまなチューニングを行いました。シングルノードにおいて、インテルの最適化により 2 倍のスピードアップを実現しました。8 ノードのクラスターに拡張した場合、全体で 16 倍のスピードアップを達成しました。インテル® MPIライブラリーの画像スループット (1 秒あたりの画像数) は、Open MPI ライブラリーと比較して、単一ノードで最大 27%、8 ノードのクラスターで最大 18%、一貫して向上しました。これらの性能向上の多くはコードを変更することなく達成されており、これらの最適化がほかのアプリケーションにも効果的に適用できることを示しています。

S-GAN

教師あり学習では、大量のラベル付きデータが必要です。ラベル付けやアノテーションは専門家が手作業で行わなければならないため、手間とコストがかかります。半教師あり学習は、ラベル付きデータとラベルなしデータの両方を使ってモデルを学習させる手法です。通常、ラベル付きのデータポイントの数は、ラベルなしのデータポイントよりもかなり少なくなります。半教師あり学習では、データのパターンや傾向を利用して分類を行います。

S-GAN は、生成モデルを用いてデータポイントを生成することで、膨大な量の学習データの必要性に取り組んでいます。Generative Adversarial Network (GAN) は、大規模な非ラベルデータを用いて、画像識別モデルを介して画像生成モデルを学習するアーキテクチャーです。GAN は、生成モデルと識別モデルの 2 つのモデルで構成されています。両者はゼロサムゲームで一緒に訓練されます。生成器の役割は、データセットに存在するものと同じようなデータを生成することです。識別器の役割は、生成されたデータの中から実際のデータを識別することです。S-GAN は、GAN のアーキテクチャーを拡張し、分類タスクに教師あり識別器 (分類器) を追加したものです (図1)。これにより、見たことのないデータに対しても十分に一般化できる分類器が得られます。

図1. S-GAN のアーキテクチャー

ソフトウェアとハードウェア

これらの実験では以下の 3 つのインテル® oneAPI ツールキット (v2021.1) を使用しました。

  1. インテル® oneAPI ベース・ツールキット
  2. インテル® oneAPI HPC ツールキット
  3. インテル® oneAPI AI Analytics ツールキット (英語)

インテル® oneAPI マス・カーネル・ライブラリーインテル® oneAPI データ アナリティクス ライブラリーを使用したインテル® ディストリビューションの Python (v3.6) は、Python のコア数値パッケージを高速化するために使用されました。S-GAN モデルは、インテル® Optimization for TensorFlow* (v1.15) (英語) を使用して実装されました。分散学習には Horovod (v0.20.2) を使用しました。Horovod はノード内通信に MPI を使用しているため、2つの MPI ライブラリー (インテル® MPI ライブラリーと Open MPI (v4.0.5)) の性能を比較しました。パフォーマンスの解析にはインテル® VTune™ プロファイラーを使用しました。すべてのテストは、インテル® Xeon® Platinum 8260L プロセッサー (英語) を使用し、100Gbps で動作するインテル® Omni-Path ファブリック (英語) を介して接続されたインテルの Endeavor クラスター上で実行されました(図2)。

図2.ノードの構成 ($ cpuinfo -g)

チューニング方法

アプリケーションを最適化する前に、ベースラインのパフォーマンスを測定することは良い習慣です。プロファイリング・ツールは、スレッディング、ベクトル化、I/O、マルチノード・コミュニケーションなど、最適化の可能性がある領域を特定するのに役立ちます。ここでは、インテル® VTune™ プロファイラーの「アプリケーション・パフォーマンス・スナップショット (英語)」を使用しました。Application Performance Snapshot プロファイルによると、アプリケーションによってあまりにも多くのスレッドが生成されていました。いくつかのスレッドは OpenMP* からのもので、ほかのスレッドは TensorFlow* が呼び出す Eigen ライブラリーからのものでした。OpenMP* と Eigen のスレッド数は、ノードあたりの論理コア数を超えており、リソースのオーバーサブスクリプションが発生して、通常はパフォーマンスが低下します。

最適なスレッド数の選択

最初に行ったのは、最適なスレッド数を見つけることでした。1 つの計算ノード上で環境変数 OMP_NUM_THREADS を設定し、低解像度の画像 (256×256ピクセル) を使用し,時間を節約するためにエポック数 (2回) を制限して、さまざまな数の OpenMP* スレッドをテストしました。その結果、50 スレッドで 1 秒あたり 19.67 枚という最高のパフォーマンスが得られました(図3)。

図3.最適な OpenMP* スレッドの数を見つける

次に、TensorFlow* のスレッディング API を使って、スレッド数を制御しました。inter_op_parallelism_threads の値は、独立したノンブロッキング演算で使用するスレッド数を指定し、intra_op_parallelism_threads の値は、行列の乗算やリダクションなどの個別演算で使用するスレッド数を指定します。図4 は、この 2 つのパラメーターの値を変えたときのパフォーマンスを示しています。濃い緑色のブロックは良好なパフォーマンスを示し、濃い赤色のブロックはパフォーマンスが低いことを示しています。白のブロックは実行に失敗したことを示しています。

図4.inter_op_parallelism_threads と intra_op_parallelism_threads の最適値を求める

ノードあたりの MPI ランクが 1 つの場合、inter_op_parallelism_threads と intra_op_parallelism_threads の最適値はそれぞれ 0 と 45 であり、これは 1 秒あたり 12.34 枚の画像に相当します。これは、OMP_NUM_THREADS を使用して達成した 19.67 画像/秒の値よりも低かったため (図 3)、TensorFlow* のスレッド API を使用する代わりに、この環境変数を使用してスレッド数を制御しました。

ノードあたりの最適な MPI ランク数の選択

このアプリケーションでは MPI と OpenMP* のハイブリッド並列を使用しているため,ノードあたりの OpenMP* スレッドと MPI ランクの最適な組み合わせを見つけることが重要です。図 3 のテストを繰り返し行いましたが (例えば、1≦OMP_NUM_THREADS≦96) ノードあたりの MPI ランク数を変化させました(21≦ppn≦24、ppn はノードあたりのランク数、図 5 参照)。その結果、各ノードに 8 つの MPI ランクと 9 つの OpenMP* スレッドを配置した場合に,最高の性能 (26.3枚/秒) が得られました。

図5.1 つのノードで OpenMP* スレッドと MPI ランクの最適な組み合わせを見つける

メモリーリークの克服

これまでの実験では、低解像度の画像を使用しました。高解像度の画像 (1360×1360 ピクセル) では、各ワーカープロセスのメモリー・フットプリントが大幅に増加するため、ノードあたり 8 つの MPI ランクを実行すると、メモリー不足エラーが発生します。1  つのランクのメモリー・フットプリントは 129GB 以上になることがわかりました。各ノードには約 200GB の DRAM しかなく、ノードごとに 1 つのランクしか起動できませんでした。これはデュアルソケットのマシンでは (NUMA (Non-Uniform Memory Access)の問題から) 最適ではありませんでした。このアプリケーションを Python* ユーティリティーのメモリー・プロファイラー (英語) で実行すると、メモリーリークが見つかりました。約 56GB のメモリーがモデルのフォワードパス中に解放されていませんでした (図6、上)。これは、TensorFlow* の Keras Model.predict メソッドの既知のバグ (TensorFlow* issue #33009 (英語) および Keras issue #13118 (英語)) であることが判明しました。TensorFlow* コミュニティーからの提言に基づき、Model.predict を Model.predict_on_batch に置き換えたところ、プロセスごとのメモリー消費量が全体的に減少しました (図6、下)。メモリーリークが修正されても、高解像度画像の場合はノードあたり 2 つのランクに制限されていました。これ以上のメモリー消費量の最適化は行いませんでしたが、可能性はあります。

図6.メモリー・プロファイラーで検出されたメモリーリーク

その他の OpenMP* およびインテル® MPI ライブラリー環境変数の設定

また、ほかの 3 つの環境変数 (KMP_BLOCKTIME、KMP_AFFINITY、I_MPI_PIN_DOMAIN) の最適な設定についても検討しました。これまでの推奨性能 (英語) に基づいて、表 1 に示す 6 つの実験を行いました。なお、セット #5 の環境変数の設定は、ほかの 5 つのセットでも使用しました。セット #3 が最も優れたトレーニング性能を示しました (図7)。

表1. その他の OpenMP* および インテル® MPI ライブラリー の環境変数の実験的設定

図7. 表 1 の 6 つの実験の性能比較

MPI ライブラリーの選び方

次に、これまでのチューニング実験で得られた最適な設定を集め、高解像度の画像(1360×1360 ピクセル) を対象に、ノードあたり 2 つの MPI ランクと 33 の OpenMP* スレッドを用いて、合計 100 のエポックを作成し、インテル® MPI ライブラリーと Open MPI の性能を比較しました。使用したインテル® MPI ライブラリーのコマンドラインは以下の通りです。

これらのほぼ同等のランタイム設定を Open MPI にも採用しました。

インテル® MPI ライブラリーは、シングルノードおよびマルチノードのシナリオで Open MPI を上回りました(図8)。

図8. インテル® MPI ライブラリーと Open MPI の性能比較

1 秒間の画像数に基づく並列スケーラビリティーは線形に見えますが、時間に基づくスケーラビリティーは劣線形であり、アプリケーションの非効率性を示しています。さらに解析を進めると、評価の頻度がノード総数に応じて正しくスケーリングされていないことがわかりました。ノード数の増加に伴い、評価の頻度も増加しました。また、これらの評価は MPI のマスターランクが行っていたため、マスターランクが評価を終えるまでほかのランクは停滞していました。この問題を解決するために、ノード数に関わらず、15 回の学習反復ごとに評価ステップを実行しました (図9、10)。この変更によるモデルの精度への影響はありませんでした。

図9. 最適化前 (左) と最適化後 (右) のトレーニングと評価に費やした総時間の割合

図10. 線形タイム・スケーリングによるマルチノードのパフォーマンス

まとめ

この記事では、S-GAN のトレーニング性能を向上させるために、段階的な最適化アプローチを紹介しました。ランタイムおよびソース・コードベースの最適化を行い、メモリーの問題、スケーリングの問題、およびシングルノードのパフォーマンスの非効率性を解決しました。マルチノード・クラスター上での S-GAN モデルの分散学習を実装するためにHorovod を使用しました。2 種類の MPI ライブラリー (インテル® MPI ライブラリーおよび Open MPI) を比較しました。

図11. インテル® MPI ライブラリーで得られた最終的なスピードアップ

図11 は、今回の最適化による性能向上をまとめたものです。シングルノードの最適化では 2 倍、マルチノードの最適化ではさらに 8 倍のスピードアップを実現し、モデルの精度を落とすことなく、全体の解答までの時間を 16 分の 1 に短縮しました。


当記事でご紹介しているインテル®ベース・ツールキットに関する詳細は こちらからご参照ください。


参照記事:
Optimize Distributed AI Training Using Intel® oneAPI Toolkits