CUDAで複数のGPUを使うプログラムを書く

CPUで動くプログラムはそれぞれのGPU(デバイス)を独立に制御する。CUDAのStreamはデバイスごとの仕事のキューになっている。プログラムの実行はStreamの実行開始と同期で行われている。

そして、Stream間の同期を行うにはeventを使う。eventはカーネルの完了を示すことができる。また、1つのデバイス上の1つのStreamにつながれているカーネルは、他のeventを待ち合わせることができる。

  • CUDA

    プログラムはそれぞれのGPUを独立に制御する。Streamの同期にはeventを使い、eventはカーネルの実行終了を知らせることができる

複数のGPUで処理を起動する一番基本的な方法は、次のようなものである。まず、cudaSetDevice()でGPUに準備を行わせ、CudaStreamCreate()でStreamに中身を入れる。これを最初のループで全部のGPUについて行う。2番目のfor文のループではfirstKernel<<< >>>()で、全GPUでfirstKernelが起動される。

3番目のfor文ではcudaStreamSynchronize()で各デバイスでのカーネルの実行終了を待ち合わせる。そして、全GPUでカーネルの実行が終わると次に進み、4番目のfor文のループに入り、secondKernel<<<>>>()で2番目のカーネルの実行を開始させる。

  • CUDA

    全GPUでfirstKernelを実行させ、全部のGPUでの実行の終了を待ち合わせてsecondKernelの実行を開始するのが一番基本的なやり方である

これでも良いのであるが、最初のカーネルの実行が16GPUすべてで終わらないと2番目のカーネルの実行が開始されないというのがムダである。次の例では、firstKernelの実行を開始したのちに、cudaEventRecordでこのカーネルの完了を示すEventの受け取りを指示する。そして、2番目のfor文のループのcudaStreamWaitEvent()で他のGPUからのevent通知を受け取り他のGPUでの実行の終了を待ち合わせる。

次のループではそれぞれのGPUがGPU0のカーネル実行の完了を待ち合わせる。すべてのGPUがGPU0の実行終了を受け取れば、最初のステップは終わりである。

このようにeventを使ってカーネル実行の終了を通知すれば時間のかかる同期(cudaStreamSynchronize()など)は全体の終わり以外には使う必要がなくなり、プログラムの実行時間を短縮することができる。

  • CUDA

    各Kernelの終了をcudaStreamSynchronize()で待ち合わせるのでなく、eventでカーネルの実行終了を伝える方が効率が良い

CUDAのCooperativeKernel機能を使えば、もっと簡単に記述できる。どのStreamのどのカーネルが協調して処理を行っているかをlaunchParamsに記述してcudaLaunchCooperativeKernelMultiDevice()を呼び出し、カーネルの中でthis_multi_grid().syncを呼び出すとGPUコードの中で必要な同期が行なわれる。後は、最後にcudaStreamSynchronize()呼び出すだけで良い。

  • CUDA

    Cooprative Luanchを使えば協調するカーネルの実行が複数GPUに跨っていても、簡単に同期させることができる

複数GPUでのメモリ管理

一般に複数のGPUがそれぞれメモリを持っている場合のメモリの管理は面倒であるが、DGX-2ではユニファイドメモリですべてのGPUとCPUのメモリがCPUからもGPUからもアクセスできるので、メモリの管理は容易である。

  • CUDA

    従来、GPUではGPUごとのメモリは独立であったが、V100 GPUではユニファイドメモリが使えるようになった

DGX-2では16個のGPUにそれぞれ32GBのHBM2メモリが付いている。しかし、NVSwitchが全対全の高バンド幅のメモリマッピングを提供しており、この接続は通常のアクセスだけでなくAtomicアクセスなどもサポートしている。

  • CUDA

    従来、16個のGPUのメモリは独立であったが、NVLinkではGPU間の高バンド幅のメモリマッピングができるようになった。また、この接続はAtomicアクセスなどもサポートしている

これとユニファイドメモリを組み合わせると全GPUで共用される512GBのメモリが存在するように見える。

  • CUDA

    ユニファイドメモリ機能を使うとすべてのGPUのメモリが1つの共通のメモリのように見える

本当は、全部のメモリがリニアに連続した領域になって欲しいのであるが、cudaMallocでは区分したグローバルなアドレス空間ができてしまい、データとGPUの関連が見えてしまう。

  • CUDA

    cudaMallocでは区分されたグローバルメモリになってしまいGPUとメモリの関連付けが残ってしまう

しかし、これもユニファイドメモリを使うと、(この例では)4つのGPUのメモリが連続になり、ポインタ演算が使えるようになる。

  • CUDA

    しかし、ユニファイドメモリ機能を使うと、4つのメモリが連続アドレスになり、ポインタ演算なども使えるようになる

次の例は、N*N*Nの3次元のアレイを作る場合で、まず、cudaMallocManaged()でN*N*Nのメモリを確保する。次にcudaMemAdvise()で全GPUからこのメモリが見えるようにする。そして、最後にGPUごとに存在するメモリが連続するメモリアドレスになるように配置する。

  • CUDA

    cudaMallocManagedでメモリを確保し、cudaMemAdviseで他のGPUからも見えるようにする。そして、cudaMemAdviseで連続アドレスになるようにして、cudaMemPrefetchAsyncでメモリを読んで使えるようにする

(次回は4月19日に掲載します)