Horovod 通信策略

因為最近的工作要和Horovod打交道，所以分析了Horovod的源碼。在這裡記一筆。

Horovod有幾個亮點，第一，它不依託於某個框架，自己通過MPI建立了一套分散式系統，完成了allreduce, allgather等collective operations通信工作. 第二，發現了Tensor fusion, 梯度傳遞的時候可以將小的tensor合併成一個大的tensor再進行傳遞，從而減小每一次操作的額外開銷(overhead). 第二點比較好理解，本文講一下horovod的通信策略。

為了簡潔，以下稱Horovod為hvd。

hvd將通信和計算框架分離之後，計算框架只需要直接調用hvd介面，如HorovodAllreduceOp來進行梯度求平均即可。但是，因為計算框架往往採用多線程執行訓練的計算圖，所以在多節點情況下，拿allreduce操作來舉例，我們不能保證每個節點上的allreduce請求是有序的。因此MPI_Allreduce並不能直接用。

為了解決這一個問題，hvd 設計了一個主從模式，rank0為master節點，rank 1-n為worker節點。除此之外，我們需要了解hvd在每個節點上定義的數據結構，每個worker節點上都有一個消息隊列，而在master節點上除了一個消息隊列，還有一個消息map。

每當計算框架發來通信請求時，hvd並不直接執行MPI，而是封裝了這個消息並推入自己的消息隊列。我們通常在python文件中使用hvd.init()來初始化hvd，實際上是開了一個後台線程和一個MPI線程。後台線程採用定時輪詢的方式訪問自己的消息隊列，如果非空，worker會將自己收到的所有tensor通信請求都發給master。因為是同步MPI，所以每個節點會阻塞等待MPI完成。master收到worker的消息後，會記錄到自己的消息map中。如果一個tensor的通信請求出現了n次，也就意味著，所有的節點都已經發出了對該tensor的通信請求，那這個tensor就需要且能夠進行通信。master節點會挑選出所有符合要求的tensor進行MPI通信。不符合要求的tensor繼續留在消息map中，等待條件符合。

決定了tensor以後，master又會將可以進行通信的tensor 名字和順序發還給各個節點。至此，所有的節點都得到了即將進行的MPI的tensor和順序，MPI通信得以進行。

以上所有操作。除了MPI通信外，均在後台進程中執行。

hvd的核心在文件https://github.com/uber/horovod/blob/master/horovod/common/operations.cc中，也是本文介紹的內容。而其對Tensorflow, Pytorch的介面定義分別在horovod/tensorflow, horovod/pytorch文件夾中。