Amazon Web Services ブログ
第 2 世代 EFA : クラウドにおける HPC/ML アプリケーションの性能改善
re:Invent 2018 で Elastic Fabric Adapter (EFA) を発表した時点で、私達は計算流体力学 (CFD) や気象予報アプリケーションを加速させるという目標を、柔軟性や可用性、コスト、インスタンスの自由度といった事柄を犠牲にする事無く、達成しました。当初、EFA は C5n と P3dn のインスタンスタイプで、5 つのリージョンにおいて利用可能でした。
図 1 を見ると分かるように、大規模環境での性能は段階的に向上しています。今日では、EFA は Intel、AMD、AWS Graviton の CPU やアクセラレータを搭載した、様々なメモリサイズの 33 種類のインスタンスで利用可能になっています。そして各リージョンに最低でも 1 種類は EFA 搭載インスタンスが存在します。
EFA の利用例は当初の想定を超えた広がりを見せており、機械学習における大規模分散学習や、高解像度のビデオストリーミングのリアルタイムの非圧縮にも利用される様になりました。
図 1: CFD++を使った2400万セルのケースの計算において EFA と TCP を利用した際の1500+コアまでのスケールアウト性能比較。このケースで私達はクリンゴン人のバード・オブ・プレイという宇宙船 (訳注: 『スタートレック』シリーズに登場する架空の宇宙船) が地球に類似した大気中を飛行する状態をシミュレートした。しかし重要なのはスタートレックではなく、コア数の増加に伴なって、EFA を利用した場合と TCP を利用した場合では、前者の方が明らかに効果的にスケールアウトしている事。2019 年 4 月に EFA が登場した直後、C5n のインスタンスを利用した結果。
これまでの 4 年間に渡り EFA の性能と機能の改善を続けてきました。そして今日は第 2 世代の EFA について紹介します。
この記事で紹介する多くの機能改善は、EFA に対応するインスタンスで個々別々に利用可能です。最近リリースした Trn1 インスタンスは、これらの機能全てを有した最初のインスタンスです。この反復的な改善手法——既存のインスタンスに対して改善を適用するという方法——は、私達がどの様に EFA を開発するか、という点で重要です。何故ならユーザは常に新しい利用方法を試行錯誤しているので、新しい世代のインスタンスに対してのみ機能改善を適用するのでは、ユーザの要求を満せないからです。
分散学習
反復的な開発プロセスの一つの例が P4d インスタンスでの機械学習における分散学習です。2 年前は、モデルの学習と言えば少数のインスタンスを使ったデータ並列が大半で、インスタンス間の通信は数ギガバイト程度の AllReduce 通信でした。しかしながら、機械学習のコミュニティではより大規模で複数レベルの並列処理が注目され、その結果として通信パターンが大きく変化し、EFA の性能の限界に達しました。
過去 1 年間かけて、私達は P4d のハードウェアを変更する事なく、小規模又は中規模のメッセージサイズでの性能を最大 50% まで改善しました (図 2 を参照)。これによって、PyTorch の分散学習ライブラリである Fully Shareded Data Parallel (FSDP) では 18% の、NVIDIA のオープンソース分散学習ライブラリである Megatron-LM では 8% 以上の性能向上が確認されました (図3)。
図 2: 2022 年 1 月と 11 月時点での、P4d.24xlarge インスタンスで最新の EFA ソフトウェアスタックを使った Allreduce 性能。この結果は 16 インスタンス、計 128 GPU のクラスタによるもの。大規模なメッセージサイズでは 10%、小規模なメッセージでは 50% の性能向上が見られる。これらの改善はハードウェアから Libfabric、NCCLに至る EFA スタックの修正による。
図 3: 2022 年初頭と後半で比べた場合の FSDP と Megatro-LM の EFA スタックの改善による性能向上 (P4d インスタンスを利用)。FSDP では 18%、Megatron-LM では 8% 性能が上っている。
第 2 世代の改善
第 2 世代の EFA は、特に機械学習アプリケーション向けに、新しい機能を提供しています。汎用 GPU や AWS Tranium のようなアクセラレータを利用した小規模な集団通信演算 (collective operations) では、P4d で利用できる第 1 世代 EFA と比べて、第 2 世代 EFA は通信時間を 50% 短縮しています。同時に、EFA デバイスをホストしている AWS Nitro System カードのスループットを 2 倍にしました。これにより、大規模な集団通信の性能と平均レイテンシを向上させる事ができました。
以降の節では第 1 世代の EFA が登場した後の、EFA の改善について紹介していきます。EFA が利用可能なインスタンスであっても、これらのどの改善点が当てはまるかは異なります。全ての改善の利益を享受できるのは第 2 世代の EFA を搭載したもののみです。
AWS Nitro Sytem のハードウェア改善
第 2 世代の EFA は、Nitro System の通信性能を改善する為の新しいハードウェアと共にスタートしました。これによりエンドポイントレイテンシ——ネットワークケーブルやスイッチを除く、NIC / ホストのソフトウェアによって生じるレイテンシ——を 30% 低減し、同時に Nitro カード 1 枚当たりの帯域幅は 100Gbps から 200Gbps と大幅に上昇、PCIe 帯域幅は 2 倍になりました。
第 2 世代 EFA では、AWS Trainium や汎用 GPU といったアクセラレータのメモリ間での直接的なデータの移動に関するサポートについても大きな改善があり、これは機械学習における分散学習のアプリケーションの改良にも繋っています。第 1 世代 EFA では、NCCL 通信をサポートするために RDMA の read semantic を加えました。第 2 世代では、より複雑な completion-semantics (Nvidia の NCCL 実装での低レイテンシな LL/LL128 プロトコルのような) を可能とする完全な RDMA インタフェースを追加し通信時間をさらに短縮しました。新しい RDMA インタフェースは MPI を使用する HPC アプリケーションでも利用できます。これにより、インスタンス内にある MPI プロセス数が少ない場合のスループットが向上します。これは OpenMP/MPI のハイブリッドアプリケーションをサポートする為に重要です。
ソフトウェアの改善
どの様なネットワーク構造であれ、HPC/ML アプリケーションとネットワークデバイスの間には多数のソフトウェア層があります。EFA の場合であれば、カーネルモジュールや、はたまたネットワークカードへのポータブルなプログラミングインタフェースを提供する RDMA の様な Libfabric というパッケージ、MPI/NCCL パッケージ等がその例です。アプリケーションの性能を向上させる取り組みの一環として、これらのソフトウェア全てに手を入れました。
NCCL のトレーニングスタックに対する主要な 3 つの変更の結果、一般的な集団通信演算にかかる時間を 75% 削減しました。
まず初めに、通信路をアクティブに保つ為に通信プロトコルを変更して、RDMA 通信の設定中にデータを積極的に送信することで、大規模モデルの学習で一般的になりつつある、NCCL を使った中小規模のメッセージの送受信時間が約 40% 短縮されました。
ふたつめとして、GPU メモリ使用量 1% 未満の僅かな上昇する事を許容する事で、ネットワークバッファを増やしました。結果、同時処理可能なトランザクションも増加する事になります。これにより、連続したメッセージをパイプライン化する事でネットワーク遅延をアプリケーションから効果的に隠蔽できます。
さいごに、AWS Trainium インスタンスで実装された独自の集団通信トポロジに対応し、低レイテンシなノード内ネットワーク (インスタンス内部でのチップ間通信) を活用し小規模な集団通信での性能を向上させました。また——そして再び——これらの教訓を MPI に適用し、アクセラレータ搭載インスタンスで MPI を使うユーザが、アクセラレータのバッファに直接データを移動する事を可能にしました。
図 4: 集団通信トポロジ実装前後の trn1.32xlarge インスタンス 16 台を使った Allreduce 性能比較。低遅延のノード内通信と遅延の大きいノード間通信を分離することで、集合演算ルーチンのスケールアウト性能を 75% 以上向上させることができる。
第 2 世代 EFA の活用
第 2 世代 EFA は、2022 年の現時点で第 6 世代 EC2 の内コンピュート最適化、メモリ最適化インスタンスで利用可能です。Trn1 は RDMA セマンティクスをサポートする最初のインスタンスタイプです。
これらの機能強化を活用する為に忘れてはならないのが、バージョン 1.19.0 以降の EFA インストーラを使うという事です (EFA Getting started のドキュメントにより多くの情報があります)。またバージョン 3.3 以降の AWS ParallelCluster には、この記事で言及した改善を活用するのに必要な全ての EFA ソフトウェアが含まれています。Trn1 や P4d/P4de インスタンスであれば、AWS Deep Learning AMI を使えば最新の EFA ソフトウェアが含まれた状態で利用できます。
おわりに
4 年前、第 1 世代の EFA を発表しました。以来、利用者のアプリケーションの為の継続的な性能改善を、EFA デバイスとソフトウェアの改良という形で行なってきました。今年になり、性能改善の面で重要なマイルストーンに辿り着きました。それが、新世代の Nitro カードと組み合わせた第 2 世代 EFA です。
既に、第 6 世代の EC2 コンピュート最適化・メモリ最適化インスタンス及び Trn1 では、RDMA セマンティクス対応を含め、第 2 世代 EFA を活用する事が可能です。また来年にかけて、第 2 世代 EFA を搭載した更なるインスタンスが登場する予定です。
当然、第 2 世代 EFA が最終形ではありません。まだまだハードウェア、ソフトウェアの両面で EFA の力を増すべく改善を行なっていきます。
この記事は 2022 年 11 月 8 日に投稿された「Second generation EFA: improving HPC and ML application performance in the cloud」をソリューションアーキテクトの小野が翻訳したものです。