GTC 2019 - NVIDIAのDGX-2をプログラムするのに知っておくべきこと(4) 複数のGPUで畳み込みの処理を計算する方法

畳み込みを並列処理する例

次の例は3×3の畳み込みを行うケースで、ニューラルネットワークの計算でも良く出てくる計算である。この計算を2つのGPUで行うとなると、下の図のように畳み込む計算を行う3×3の領域が2つのGPUのメモリに分かれて存在するというケースが出てくる。この場合、相手側のメモリにあるデータを持ってきて計算する必要があり、通信が重要になる。

畳み込みを2つのGPUで分担して行う場合、下の図のように対象とするメモリ領域が2つのGPUのメモリに分かれてしまうケースがでる

このようなケースでは、それぞれのGPUに格納するデータを一回り大きくして、それぞれのGPUが自分のメモリにあるデータだけで計算が進められるようにするのが一般的である。この拡大したメモリ領域をハロー(日本語では袖領域)という。

対象領域が2つのGPUのメモリに分かれる場合、袖領域を相手のGPUのメモリにコピーして、計算時には相手のメモリを読まなくても計算ができるようにするのが一般的である

2分割の場合はハローに必要なメモリ容量やコピーの手間はそれほど大きくはないが、GPUの数が増えてくると、メモリ容量もコピーの手間も増えてくる。

4つのGPUに分割すると、袖領域の比率が大きくなり、メモリのムダが増え、転送の手間も大きくなる

しかし、通信をNVLinkに任せて、隣のGPUのメモリをローカルメモリのようにアクセスしてしまえば、ハローのメモリ領域は不要で、コピーなしに直接、隣接するGPUメモリから読むことができる。しかし、このアプローチはどこまで通用するのであろうか？

通信をNVLinkに任せて、他のGPUのメモリもローカルメモリのように読み書きすればハロー領域の確保やハロー領域へのコピーも不要である。これでうまく行くのか？

GPUの数が増えると、ハローとコアの比率が増加する。つまり、オフチップのメモリをアクセスする比率が増加し、オンチップのメモリをアクセスする比率が減少する。これをどんどん進めると、NVLinkのバンド幅がスケーリングの限界になる。しかし、通信は多数のGPUからそれらのメモリとの間のMany-to-Manyの通信になるので、1リンクのバンド幅ではなく通信網全体のバンド幅の制約になる。

GPU数が増えると、ハローの比率が増え、NVLinkを使う通信が多くなり、NVLinkのバンド幅が性能の制約になり得る。しかし、DGX-2の場合、16個のGPUと16個のメモリ間の通信バンド幅が問題になる

オンチップアクセスの場合はHBM2の読み出しバンド幅で決まり、オフチップアクセスの場合はNVLink2のバンド幅で性能が決まると考えると、オフチップアクセスが0%の場合はバンド幅は880GB/s、オフチップアクセスが100%の場合は120GB/sのバンド幅で、中間は、この2点を直線で結んだ形になる筈である。この線の傾きは18.4%である(第1世代のNVLinkのバンド幅は20GB/sであったが、NVLink2ではこれを25GB/sに引き上げている。これを6リンク使うと、NVLinkでは120GB/s、NVLink2では150GB/sとなる。DGX-2はNVLink2を使っているので150GB/sのはずであるが、この発表では120GB/sと書かれている部分もある。また、HBM2のバンド幅も最近では900GB/sになっているはずであるが、ここでは880GB/sとなっている)。

NVLinkのバンド幅は120GB/s、HBM2メモリのバンド幅は880GB/sであるので、リモートアクセスの比率が18.4%以上になるとNVLink経由のアクセスが制約になる

16GPUで3次元のステンシル処理を行わせた場合の性能を次の図に示す。横軸はステンシルのサイズである。このコードを実行したときの性能の傾きは25%であり、理想的な18.4%に対して72%の性能を達成しており、一般的にはこれ以上の高い効率の実現は難しいと言えるレベルである。このレベルの性能がいろいろな工夫をしていないナイーブなプログラムで実現できるのは素晴らしいことである。

3次元のステンシル計算のメモリバンド幅。横軸はステンシルのサイズ。ナイーブなコードであるが、傾きは25%であり理想的なスケーリングの72%の性能となっている

次の図は、流体コードのLULESHステンシルのケースで、4GPUと1GPUの場合の性能を比較している。NVLinkのバンド幅はHBM2のバンド幅の18.4%であるので、リモートアクセスの比率が18.4%を超えると性能が低下すると予想される。このプログラムは非常に良くチューニングされたものであるが、やはり18.4%のところで急激に性能が低下している。