UCXを試す日記(9)：allreduceを実装する上での予備調査

とりあえず、UCXで実用（？）コードを書いてみようということで、Allreduceを書いてみようと思っています。Allreduceは、ディープラーニングにおいては重要な通信パターンで、業務においても研究したことがあるので経験があります（これについては、近々会社のブログの方で発表できると思います）。

Allreduceを実装するにはいくつかのポイントがあることがわかっていますので、これを先に調査しておきます。

send/recv通信として, bcopy と zcopy のどちらを使うか

いわゆる普通のsend/recvを実現する手段としては、bcopyとzcopyがあります。bcopyとzcopyの違いはわかりにくいですが、このページによると、

• bcopy：転送を直ちに開始するか、失敗かの２通りの結果しかない。失敗とは、sendするローカル側のリソースが不足している場合で、UCS_ERR_WOULD_BLOCKが返される。この場合、Endpointの作成時にUCT_FLAG_PENDINGとコールバックを設定しておくと、リソースが利用可能になった時点でこのコールバックが呼ばれる
• zcopy：bcopyの動作に加えて、 UCS_INPROGRESS を返すことがある。この場合、通信は非同期に行われ、完了時に uct_ep_am_zcopyの引数に指定したコールバックが呼ばれる。

ということで、非同期通信モデルである zcopy が良さそうです。気になるのは、bcopyのメモリコピーに関する動作がはっきり書いていないところです（bが何を意味するのかよくわからない）。zcopyのzはZeroだろうと思われるので、逆に考えるとbcopyはゼロコピー性は保証されていないのかな…？ソースを読む必要がありそう

通信の順序が保証されるか？

MPIおよびInfinibandでは、通信の到着順は保証されています。これが入れ替わってしまうと、Ring-Allreduceが成り立たないので注意が必要です。下のレイヤがInfinibandであれば自動的に（事実上）保証されることになるでしょうが、APIとしての保証がない場合、例えば下のレイヤをTCPに変えたら動作が変わって動かなくなる、みたいなことになりそうなので注意が必要ですね。

UCTとUCPのどちらを使うか？

UCXには、通信を実現するAPIのレイヤとして UCPとUCTがあります。UCTは、固有のAPIに近い、薄いレイヤを提供するので、特定の機能がサポートされているかどうか、などのフラグによる分岐をユーザーが書く必要があります。一方でUCPは上位の機能を実現するレイヤでタグマッチングなどが実装されています。インターフェースがサポートしない機能はエミュレーションするように実装されているようなので、便利そうではあります。

その場合、 * UCPのオーバーヘッドはどれくらいか？ * UCPのエミュレーション機能は、現段階でどれくらい実装されているのか * UCPのみで実現されている機能が必要か？

ということだけど…

以上の点を調査する必要がありそうですが、とりあえずInfinibandを前提にした高速な実装を目指そうと思うので、これらの調査は先送りしようと思います。