アマゾン ウェブ サービス が BERT および Mask R-CNN における最速トレーニングタイムを達成

今日最も多く使用されている機械学習モデルには 2 つあります。自然言語処理 (NLP) 用の BERT と、画像認識用の Mask R-CNN です。AWS では、この数か月にわたり、これら 2 つの良く使われている先進的なモデルにおいて最短のトレーニング時間を達成するため、基盤となるインフラストラクチャ、ネットワーク、機械学習 (ML) フレームワーク、モデルコーディングなどに、大幅な改良を加えてきました。TensorFlow、MXNet、PyTorch に関しクラウド上でこれまで記録された最短のトレーニング時間を、本日、皆様と共有できることを、心から喜んでおります。お客様は、ご自身の TensorFlow、MXNet、PyTorch のモデルでのトレーニングに、これらのハードウェアとソフトウェアに関する最適化手法を、当社と同じ効率とスピードでご利用になれます。

モデルに対するトレーニング時間は、そのモデルの精度への改良を、素早く繰り返すときの作業性に直接影響します。トレーニング時間を削減しようとするときに最初に考えられる手法とは、GPU インスタンスの大規模なクラスターを通じて、トレイニングジョブを供給するということです。しかしこれでは、効率を高めることは困難です。大量のワーカーを通じてトレーニングジョブを供給しても、しばしば、急速にその効果が薄れてしまうことがあります。インスタンス間の通信におけるオーバーヘッドが、GPU を追加した分のコンピューティングパワーを帳消しにしてしまうからです。

BERT

現在、普及している NLP モデルである BERT、つまり Bidirectional Encoder Representations from Transformers は、当初は、いくつかの一般的な NLP タスクを処理するための最新手法として公開されたものです。

NVIDIA V100 GPU を 8 個使用する単一の Amazon EC2 P3dn.24xlarge インスタンスにおいて、TensorFlow と PyTorch を使いながらこの BERT をゼロからトレーニングするには、およそ 3 日間を要します。当社では、Elastic Fabric Adapter (EFA)を使用しつつ、このモデルを大規模クラスター上で集中させる方法を最適化しながら、P3dn.24xlarge インスタンスへの効率的なスケールアウトを実行し、3 日間のトレーニング時間を 60 分を少し超える程度まで削減しました。この記事を執筆している時点で、これは、先進的な (BooksCorpus と English Wikipedia でトレーニングした後の Squad v2 タスクにおける F1 スコアとなる 90.4 以上の) ターゲット精度を達成しながらの、BERT に対するクラウド上のトレーニング時間としては最短となっています。

TensorFlow を使う場合では、256 個の P3dn.24xlarge インスタンスに 2,048 GPU を置くという前例のない規模を実現し、BERT のトレーニングを 62 分間で完了しました。PyTorch を使う場合は、1,536 個の GPU を配置した 192 個の P3dn.24xlarge インスタンスにスケールアウトすることで、トレーニング時間を 69 分間にまで短縮しています。今回、BERT のトレーニングのために行った、ハードウェアとソフトウェアスタックに対するすべての最適化作業において、85% のスケーリング効率を達成しています。つまり、P3dn.24xlarge ノード以上にスケーリングしたときでも、GPU による追加コンピューティングパワーのほとんどを、フレームワークが利用できることが確認できました。これらの向上点を次の表にまとめます。

Mask R-CNN

Mask R-CNN は、インスタンスセグメンテーションのモデルとして広く使用されているものです。自律運転、モーションキャプチャー、およびその他の先進的な物体検知とセグメンテーション機能が求められる用途に使われています。

MXNet、PyTorch、TensorFlow を使用しながら単一の P3dn.24xlarge (8 個の NVIDIA V100 GPU を使用) インスタンス上で Mask R-CNN をトレーニングするには、およそ 80 時間を必要とします。当社は、この 80 時間というトレーニング時間を、MXNet、PyTorch、TensorFlow を使用しつつ、約 25 分にまで短縮しました。3 つの ML フレームワークすべてにある Mask R-CNN トレーニングを、 24 個の P3dn.24xlarge インスタンスにスケーリングし、192 個の GPU が使えるようにしました。これからは、結果を数日間も待つ必要はなくなり、毎日、複数の実験を素早く繰り返すことができるようになります。この記事の執筆時点において今回の記録は、最先端のターゲット精度 (COCO2017 データセット上で、0.377 Box min AP および 0.339 Mask min AP) を達成しつつ行った、クラウドでの Mask R-CNN トレーニング時間としては最速となっています。これらの向上点を次の表にまとめます。

テクノロジースタック

これらの時間を達成するためには、基盤ハードウェア、ネットワーク、そしてソフトウェアスタックの最適化が必須でした。BERT のような大規模モデルのトレーニング時には、稼働している多数の GPU 間で行われる通信がボトルネックとなります。

分散コンピューティング (大規模トレーニングがまさにその一例です) では、異なるワーカー (各 GPU) で配列の数 (今回の場合はニューラルネットワークのパラメーター) を削減し、全ワーカーに対する結果の配列を返すため、AllReduce が 1 つの役割を果たします。各トレーニングを終了すると、各 GPU は供に AllReduce のオペレーションを実行します。各トレーニングには、ネットワークを通じて、前方および後方へのパスがそれぞれ 1 つずつあります。

複数の GPU 上で AllReduce を実行する最も一般的なアプローチは、NVIDIA Collective Communications Library (NCCL) や、OpenMPI もしくは Intel MPI Library といった MPI ライブラリーなどを使うことです。これらのライブラリーは均一なクラスター用に設計されています。AllReduce は、ネットワークをトレーニングするのと同じインスタンス上で動作します。均一なクラスターに置かれた AllReduce アルゴリズムは各ワーカーと関連し、それぞれのオペレーションが対応するモデルに比べ、約 2 倍のサイズのデータを送受信します。たとえば、(パラメーターが 3 億 4 千万個ある) BERT 用の AllReduce オペレーション
では、約 650 MB の半精度データを 2 回送信し、同じ量のデータを 2 回受信します。ほとんどのモデルでは、各トレーニングの後に、こういた通信を行う必要が生じボトルネックとなります。

AllReduce アルゴリズムの選択は、通常、ネットワークアーキテクチャにより変わります。たとえば、各ノードが隣のノードとだけ接続されリングを形成しているようなネットワークには、Ring-AllReduce の選択が適しています。各ノードが周囲の 4 つと接続され 2 次元の長方形格子を形成しているようなネットワークには、Torus AllReduce アルゴリズムが適しています。AWS では、任意のノードが他の任意のノードと全帯域幅で通信できるような、より柔軟性の高い内部接続を採用しています。たとえば、128 P3dn インスタンスにあるクラスターでは、任意のインスタンス同士が 100 Gbps で通信することが可能です。

さらに、 100 Gbps という内部接続が使えるのは、 P3dn インスタンスに限られているわけではありません。C5n のような CPU 最適化タイプのインスタンスも同じクラスターに追加でき、そのまま任意のノードペア間で 100 Gbps の通信ができます。

AWS の内部接続が持つ高い柔軟性は、そのユニークな接続性能をフルに活用できる AllReduce アルゴリズムのためにあるといってもよいでしょう。という訳で、AWS ネットワークに最適化された、カスタムの AllReduce アルゴリズムが開発されました。このカスタム AllReduce アルゴリズムでは、不均一なクラスター内で任意のノードペア間での 100 Gbps の内部接続を利用することで、各ワーカーにより送信および受信されるデータ量を半分にまで削減できました。AllReduce アルゴリズムのコンピューティングフェーズでは、計算に最適化された C5 インスタンスに負荷を手渡すことで、GPU を計算グラジエントから素早く解放しています。GPU インスタンスでは削減のための処理は行わないため、グラジエントの送出と AllReduced グラジエントの受け取りは、同時に処理することができます。AllReduce が必要とするグラジエントの移動回数は、ネットワーク内の移動回数がノード数に応じて増加する均一な AllReduce アルゴリズムと比較しても、2 倍程度に削減されています。また、P3dn ノードを使って行う場合に比べトレーニングがかなり速く完了するので、トータルのコストも削減されました。

まとめ

BERT と Mask R-CNN の単一ノードでの実行では、テストの結果に顕著な改善がみられました。P3dn ノードが 1 から 16、32、64、128、192、256 に増加するにつれ、スループットも概ねリニアに増加しています。これは、コストを増やすことなく P3dn.24xlarge インスタンスにスケーリングするだけで、最終的にモデルのトレーニング時間を短縮させられる要因となっています。これらの最適化手法を通じ、AWS では、最先端のコンピュータビジョンや NLP モデルなどに、クラウド上における最短のトレーニング時間をご提供してまいります。

Amazon SageMaker で、今すぐ TensorFlow、MXNet、PyTorch の使用を開始しましょう。

著者について

Aditya Bindal は、AWS Deep Learning のシニアプロダクトマネージャーです。彼は、顧客のために深層学習エンジンを使いやすくする製品に取り組んでいます。余暇には、テニスをしたり、歴史小説を読んだり、旅行を楽しんだりしています。

Kevin Haas は、AWS Deep Learning チームのエンジニアリングリーダーです。AWS の機械学習を使うお客様に、向上したパフォーマンスと利用性をご提供することに取り組んでいます。彼は、お客様が、ご自身のアプリケーションに機械学習を導入する際の大変さを減らす事に熱心です。仕事以外の時間では、オープンソフトソフトウェアのダビングや、ボーイスカウトでのボランティアなどをしています。

Indu Thangakrishnan は、AWS のソフトウェア開発エンジニアです。彼は、深層ニューラルネットワークのトレーニングを高速化することに取り組んでいます。余暇には、Audible を繰り返し聴いたりピンポンを楽しんでいます。