MPIによる分散通信を試す

はじめに

OpenMPやOpenACC asyncを使って1台の計算機上でのマルチGPU制御は紹介しましたが、複数の計算機上のマルチGPUをまとめて扱うことはできません。
現在のところ1台の計算機につなぐことができるGPUの数は、コストや電力が関係すると思いますが8台程度です。

例えばNVIDIA A100が8台あれば、約77.6 TFLOPSの倍精度演算性能と合計640 GBのメモリがあるので必要十分ということが多いかと思います。
ですが様々な要因で計算機あたり1, 2台しかGPUが繋げない、640 GBでもメモリが足りない、GPU 8台でも計算速度が足りない、そもそもどのぐらいの規模まで計算するか推測できない、といろいろな要因でより多くのGPUで並列計算をしたいというケースはありそうです。
このような場合のデファクトスタンダードとなる選択肢がMPIです。

今回は、まずMPIの簡単な使い方を紹介いたします。

MPI

MPI (Message Passing Interface) は、平たく言うと並列計算用向け通信APIの標準化規格です。
HTTPやgRPCといったWeb系における通信技術と比べると、シミュレーションのような分散処理をするための通信パターンやそのための機能に特化しています。
MPIは実際に存在するシステムのハードウェアやミドルウェア間の違いを吸収するため、ユーザーは細かい差を気にせずに分散並列処理を記述できる、としています。
規格上はC, C++, Fortranに対し策定されていますが、Pythonなどの言語でもC, C++実装をベースにライブラリ化されています。

MPIの実装をここでは処理系と呼びますが、代表的なのはMPICH, Open MPI, MVAPICHなどです。
API自体は規格化されていますが、実際にMPIプログラムを実行する方法や環境設定等は各処理系によってまちまちです。
NVIDIA HPC SDKではバージョン23.11から、Open MPIをベースにしたHPC-XがデフォルトのMPI処理系となっています。
HPC-XではInfiniBandスイッチに通信処理の一部をオフロードすることで、CPUの処理負荷を削減する機能が追加されています。実行時オプション等の使い方はOpen MPIから変更はないようです。

総和計算

まずは全プロセスでスカラー値の総和をMPIで記述します。
MPIはプロセス(ランクと呼びます)間で通信を行うため、各プロセスの間はメモリが完全に分離されています。
例えば1本のベクトルを分散並列計算の対象とする場合、各プロセスは部分ベクトルを持っていて、通信によって部分ベクトルを共有・交換します。

#include <stdio.h> #include <assert.h> #include <mpi.h> int sum(int n) { int r = 0; for (int i = 1 ; i <= n ; ++i) r += i; return r; } int main(int argc, char *argv[]) { MPI_Init(&argc, &argv); int rank, nprocs; MPI_Comm_rank(MPI_COMM_WORLD, &rank); MPI_Comm_size(MPI_COMM_WORLD, &nprocs); int v = rank + 1; int result; MPI_Allreduce( &v, // input buffer &result, // output buffer 1, // number of data MPI_INT, // data type MPI_SUM, // operation type MPI_COMM_WORLD ); assert(result == sum(nprocs)); printf("rank %d of %d: my data = %d, summation = %d\n", rank, nprocs, v, result); MPI_Finalize(); }

MPIはプロセスグループのことをCommunicatorと呼んでいますが、MPI_COMM_WORLDはMPIによって立ち上げたプロセス全てを指しています。
MPI_Comm_rankおよびMPI_Comm_sizeで、自分のランク番号と何個のプロセスが立ち上げられたかを確認でき、各プロセスは自分が担当するべき計算領域をこれで求め分散並列化を行います。
MPI_Allreduceはreduction演算を行うためのCommunicator全体通信 (collectiveと呼びます) で、計算結果はすべてのプロセスに渡されます。今回は総和なのでMPI_SUMを使っていますが、OpenMPやOpenACCのreductionと同様に最大値や積などを計算することも可能です。
詳しくはOpen MPIのマニュアル等を確認してください。
collective通信には他に、broadcast(全プロセスへの値の共有)やscatter(データの分散)などがあります。

基本的には上記の様に、MPIを初期化しプロセスがいくつ立ち上がっているかを確認、自分の担当領域の計算と他プロセスとの通信を交互に行いながら計算を進める、という流れです。
このコードのコンパイルと実行ですが、MPIのライブラリリンク等を簡略化するためのコンパイラのwrapperが用意されています。
コンパイラ名は処理系により異なりますが、Open MPIの場合は以下のようになっています。

  • C: mpicc
  • C++: mpicxx, mpic++ or mpiCC
  • Fortran: mpifort, mpif90 or mpif77

それぞれが実際にどの言語規格に沿ったコンパイラなのかは、ベースにしているコンパイラ(NVIDIA HPC SDKではNVIDIA HPCコンパイラ)に依存します。
コンパイルオプション自体はベースにしているコンパイラのものをそのまま使うので、今回はNVIDIA HPCコンパイラ記載の通りです。

$ mpicxx -o ./xtest mpi_allreduce.cc

MPIプログラムの実行には、mpirunまたはmpiexecを使います。Open MPIの場合はどちらも同じものを指しています。
これも処理系によって異なる名前のコマンドが提供されていることがあり、MVAPICHの場合はmpirun_rshというコマンドになっています。
コマンドラインオプションも処理系によって全く異なりますが、多くの場合、-nは立ち上げるプロセスの数を指します。
4プロセスで上記のコードを実行すると、以下のような出力が得られるはずです。(実行環境によってはプロセスのCPUへのbindingに失敗するため、bindingを--bind-to noneで無効化しています)

$ mpirun -n 4 --bind-to none ./xtest rank 2 of 4: my data = 3, summation = 10 rank 0 of 4: my data = 1, summation = 10 rank 1 of 4: my data = 2, summation = 10 rank 3 of 4: my data = 4, summation = 10

各プロセス独立に実行されているため、出力の順序は実行の度に変わると思いますが、全プロセスの計算結果が10になっていればOKです。

P2P通信

次にP2P (1対1) 通信の使い方を示します。
P2Pでは送信者と受信者を明確に区別するため、MPI_Send, MPI_Recv関数とランク番号を使って切り替える必要があります。

#include <stdio.h> #include <assert.h> #include <mpi.h> int main(int argc, char *argv[]) { MPI_Init(&argc, &argv); int rank, nprocs; MPI_Comm_rank(MPI_COMM_WORLD, &rank); MPI_Comm_size(MPI_COMM_WORLD, &nprocs); assert(nprocs % 2 == 0); constexpr int n = 5; for (int i = 0 ; i < n ; ++i) { int v; if (rank % 2 == 0) { // sender v = (n * rank / 2) + i; MPI_Send( &v, // send buffer 1, // number of data MPI_INT, // data type (rank / 2) * 2 + 1, // send to rank 0, // tag MPI_COMM_WORLD ); } else { // receiver MPI_Status st; v = -1; MPI_Recv( &v, // recv buffer 1, // number of data MPI_INT, // data type (rank / 2) * 2, // recv from rank 0, // tag MPI_COMM_WORLD, &st // communication status ); } printf("rank %d of %d: my data = %d\n", rank, nprocs, v); } MPI_Finalize(); }

バッファや個数、データ型の指定はcollective通信と同様です。送信と受信でそれぞれデータの個数が必要なので、送受信者間でデータ個数の合意が必要です。
send/recvは複数のペアを実行できるので、tagを使って送信・受信データを識別可能です。
例えばtag = 0で送信するべきデータの個数を送り、tag = 1で実際にデータのやり取りを行う、というような使い分けが可能です。

MPI_Send/MPI_Recvは同期通信(送受信が完了するまで制御が戻らない)ですが、非同期(正確にはnon-blocking)通信で複数のデータをまとめて送受信することも可能で、tagをうまく使い分ける必要があります。
P2P通信は自由が高いですが、どのプロセスと通信を行うか細かい制御が必要になるため、collective通信を優先利用していくのが良いかと思います。

上記のコードを実行すると、以下のような出力が得られるかと思います。結果は見やすいようソートしていますが、実際にはプロセス単位でバラバラの順序で出力されるかと思います。

$ mpicxx -o xtest p2p.cc $ mpirun -n 4 --bind-to none ./xtest | sort rank 0 of 4: my data = 0 rank 0 of 4: my data = 1 rank 0 of 4: my data = 2 rank 0 of 4: my data = 3 rank 0 of 4: my data = 4 rank 1 of 4: my data = 0 rank 1 of 4: my data = 1 rank 1 of 4: my data = 2 rank 1 of 4: my data = 3 rank 1 of 4: my data = 4 rank 2 of 4: my data = 5 rank 2 of 4: my data = 6 rank 2 of 4: my data = 7 rank 2 of 4: my data = 8 rank 2 of 4: my data = 9 rank 3 of 4: my data = 5 rank 3 of 4: my data = 6 rank 3 of 4: my data = 7 rank 3 of 4: my data = 8 rank 3 of 4: my data = 9

2の倍数のプロセスを起動し、ランク i & i+1 (i = 0, 2, 4, 8…) の間でP2P通信をしました。
ランク0, 1とランク2, 3が同じ値を持っているはずです。

まとめ

今回はマルチノードでのマルチGPU計算に向け、MPIの簡単な書き方をご紹介させていただきました。
GPUの性能向上は目覚ましいですが、メモリ容量の不足、より計算量の大きな計算へのシフトなどによって多くのGPUが必要なケースもあります。

次回以降、MPI上でマルチGPUを行う方法や、cuBLASMpのようなマルチプロセス・マルチGPUのためのライブラリを試してみたいと思います。