NCCL1.3.5対応コードをNCCL2.0に移植してみた

TL;DR

• GitHubにあるNCCL1.3.5はNVLinkに最適化されてない．NVLinkを活用するにはNCCL2を使えばいい
• NCCL1.3.5とNCCL2には若干互換性がないので，そのままでは使えない
• NCCL1.3.5のtestプログラムをNCCL2.0向けに移植してみた．思ったより簡単かも!?

ことのはじまり

NVLink接続された2つのGPUを搭載するサーバを触る機会に恵まれましたので，マルチGPU向けの素敵な集合通信ライブラリNCCLを試してみました．NVLink向けに最適化されたコードが入っているというNCCL2.0では，NVLink向けに最適化されていないNCCL1.3.5に比べて高い性能がでることが確認できました．
...というのが前回の話

比較してみるために，NCCL1.3.5と一緒に公開されているtestプログラムをベンチマークとして使ってみたのですが，NCCL1.3.5とNCCL2には若干の違いがあって手直しが必要であることもわかりました．

今後NCCL1.3.5向けのコードをNCCL2向けに移植していく上で，違いを忘れないようにとメモしておきます．あくまでtestプログラムの移植に必要な分だけなので他にもあるかも．
...チートシートあるよ，というのであれば教えてください．

NCCL1.3.5とNCCL2.0の違い

NCCL1.3.5のtestプログラムをNCCL2.0向けに移植するために，いくつかの修正が必要でした．いいから修正した差分を見たい，という方は，こちらをどうぞ．

(1) ncclGroupStart()/ncclGroupEnd()の追加

NCCL2.0では，NCCLの集合通信APIを叩く前後に，ncclGroupStart()/ncclGroupEnd()を追加する必要があります．たとえば，NCCL1.3.5の↓のようなコードは，

  for (int i = 0; i < nDev; ++i) {
    CUDACHECK(cudaSetDevice(dList[i]));
    NCCLCHECK(ncclAllReduce((const void*)sendbuff[i], (void*)recvbuff[i], std::min(N, 1024 * 1024), type, op, comms[i], s[i]));
  }

NCCL2では，↓のように書く必要があります

  ncclGroupStart();
  for (int i = 0; i < nDev; ++i) {
    NCCLCHECK(ncclAllReduce((const void*)sendbuff[i], (void*)recvbuff[i], std::min(N, 1024 * 1024), type, op, comms[i], s[i]));
  }
  ncclGroupEnd();

この違いは，コンパイルエラーになるわけではなくて，実行時にAPIが呼ばれた際に固まってしまう現象として発現するので，少々厄介．testプログラムの範囲では，

• ncclAllGather
• ncclAllReduce
• ncclBcast
• ncclReduce
• ncclReduceScatter

は，前後にncclGroupStart()/ncclGroupEnd()が必要でした．

(2) ncclReadOp_tとncclDataType_tの変更

リダクション演算子の定義ncclReadOp_tと型の定義ncclDataType_tが，NCCL1.3.5とNCCL2.0で若干異なっています．
NCCL1.3.5では，↓のように定義されていますが，

/* Reduction opperation selector */
typedef enum { ncclSum        = 0,
               ncclProd       = 1,
               ncclMax        = 2,
               ncclMin        = 3,
               nccl_NUM_OPS   = 4 } ncclRedOp_t;

/* Data types */
typedef enum { ncclChar       = 0,
               ncclInt        = 1,
#ifdef CUDA_HAS_HALF
               ncclHalf       = 2,
#endif
               ncclFloat      = 3,
               ncclDouble     = 4,
               ncclInt64      = 5,
               ncclUint64     = 6,
               nccl_NUM_TYPES = 7 } ncclDataType_t;

(地味にoperationの綴りがopperationと間違っているのが気になります)
一方，NCCL2.0の定義は，↓のよう．

/* Reduction operation selector */
typedef enum { ncclSum        = 0,
               ncclProd       = 1,
               ncclMax        = 2,
               ncclMin        = 3,
               ncclNumOps     = 4 } ncclRedOp_t;

/* Data types */
typedef enum { ncclInt8       = 0, ncclChar       = 0,
               ncclUint8      = 1,
               ncclInt32      = 2, ncclInt        = 2,
               ncclUint32     = 3,
               ncclInt64      = 4,
               ncclUint64     = 5,
               ncclFloat16    = 6, ncclHalf       = 6,
               ncclFloat32    = 7, ncclFloat      = 7,
               ncclFloat64    = 8, ncclDouble     = 8,
               ncclNumTypes   = 9 } ncclDataType_t;

operationの綴りが修正されています...じゃなくて，nccl_NUM_OPSがncclNumOpsに，nccl_NUM_TYPESがncclNumTypesに変更されているのが地味に面倒ですね．また，定義名と値にも違いがあるので，定義名を使わずに値を数値として使っている場合には注意が必要です．

(3) APIの引数の順番の変更

引数の順番が変わっているAPIもあります．今回の移植の範囲でコンパイラのエラーで気づいたものは，ncclAllGather．
NCCL1.3.5の定義では，

ncclResult_t  ncclAllGather(const void* sendbuff, int count, ncclDataType_t datatype,
    void* recvbuff, ncclComm_t comm, cudaStream_t stream);

となっているのですが，NCCL2.0では，↓のように，

ncclResult_t  ncclAllGather(const void* sendbuff, void* recvbuff, size_t sendcount,
    ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream);

受信バッファのポインタが引数の2番目に移動されています．all_gather_test.cuのようにncclAllGatherを使うプログラムでは引数の順番の変更が必要です．この手の違いは，コンパイラではじいてくれますので，すぐに気づくことができますね．

簡単に試してみる

簡単に試してみるには，NCCL1.3.5のtest環境をNCCL2.0に入れ変えてみるのが環境を汚さずに試すことができて便利です．
私は以下のようにしてみました．

まず，作業ディレクトリの下に，NCCL1.3.5のプロジェクト一式をクローンして，testプログラム以外を削除します．

$ cd ${作業ディレクトリ}
$ git clone https://github.com/NVIDIA/nccl.git
$ cd nccl
$ rm -rf debian fortran src build
$ mkdir build

NCCL2はdebianパッケージとしてダウンロードすることができます．私がダウンロードしたのは，nccl-repo-ubuntu1604-2.0.4-ga_2.0.4-1_amd64.deb でした．これをとりあえず手で展開して，必要な場所に置くことにします．

$ mkdir ${作業ディレクトリ}/tmp
$ cd ${作業ディレクトリ}/tmp
$ dpkg-deb -x nccl-repo-ubuntu1604-2.0.4-ga_2.0.4-1_amd64.deb .
$ cd var/nccl-repo-2.0.4-ga/
$ dpkg-deb -x libnccl2_2.0.4-1+cuda8.0_amd64.deb .
$ dpkg-deb -x libnccl-dev_2.0.4-1+cuda8.0_amd64.deb .
$ mv usr/* ${作業ディレクトリ}/build
$ cd ${作業ディレクトリ}/build/lib
$ ln -s x86_64-linux-gnu/libnccl* .

これで，NCCL1.3.5のライブラリをビルドした状態と同じようにNCCL2を配置することができました．
最後にMakefileを修正します．修正内容はこんな感じ．diffファイル

• NCCLのバージョンを指定している箇所のコメントアウト
• ライブラリをビルドするルールの削除
• NCCL2関連のコードをcleanで削除しちゃわないように，削除対象をbuild/testに変更
• debiaパッケージ作成関連は削除

これで，make testでNCCL1.3.5のtestプログラムをNCCL2.0を使ってコンパイル・実行することができます．なお，testプログラムの実行時には，LD_LIBRARY_PATHをbuild/libに通しておく必要があります．

まとめ

上記に気をつけてソースコードを書き変えることで，NCCL1.3.5向けに書かれたプログラム(GitHubに公開されているNCCL1.3.5のtestプログラム)をNCCL2.0向けに移植することができました．
少なくとも，このレベルであればNCCL2.0への移植は，そうおそれるものでもないのかなあ...とか．