引言
PyTorch distributed currently only supports Linux.
這句話是來自 pytorch 官網 的 torch.distributed 部分，說明 pytorch 支持分佈式訓練，而且只在linux 上支持。
torch.distributed supports three backends, each with different capabilities.
同樣來自上述 頁面，pytorch 在分佈式訓練中，支持三種後端（backend）進行集羣管理或者通信。
那麼，在什麼情況下選用什麼樣的backend？
- 經驗法則
- 使用NCCL後端進行分佈式GPU培訓
- 使用Gloo後端進行分佈式CPU培訓。
- 具有InfiniBand互連的GPU主機
- 使用NCCL，因爲它是目前唯一支持InfiniBand和GPUDirect的後端。
- GPU主機與以太網互連
- 使用NCCL，因爲它目前提供最佳的分佈式GPU訓練性能，特別是對於多進程單節點或多節點分佈式訓練。如果您遇到NCCL的任何問題，請使用Gloo作爲後備選項。（請注意，Gloo目前運行速度比GPU的NCCL慢。）
- 具有InfiniBand互連的CPU主機
- 如果您的InfiniBand已啓用IP over IB，請使用Gloo，否則請使用MPI。我們計劃在即將發佈的版本中爲Gloo添加InfiniBand支持。
- 具有以太網互連的CPU主機
- 除非您有特殊原因要使用MPI，否則請使用Gloo。
實踐部分
- 首先，假設，我們要起一個分佈式訓練，起 2 個 rank。
- 把他們起來兩臺機器上，機器 1 的 ip 爲 192.168.61.55，機器 2 的 ip 爲 192.168.61.56
- 機器 1 上起 1 個 rank，機器 2 上起 1 個 rank
我這裏使用的示例代碼來自於 github 上的一段開源代碼， 地址是 https://github.com/pytorch/examples/tree/master/mnist 。我對這段代碼做了些更改，使其支持起 rank 的操作，測試的時候一臺機器上起一個 rank 。你可以在 網上搜索到類似的修改。
但是，這不是一個最佳的測試例子，原因是因爲這段代碼缺少 類似 all_reduce 這樣的操作，缺少運算過程中所有 rank 的權重進行修改、因而引起的 PCI 通信或者 InfiniBand 網絡通信（非系統通信）等等。我不明白的是，網上有很多人把這段代碼作爲一個 非常有效的例子 用來 測試分佈式運算，估計都是沒有深究，互相抄抄改改。
我找到了一段最佳實驗代碼， 地址是 https://github.com/seba-1511/dist_tuto.pth/blob/gh-pages/train_dist.py 。奈何不是做 深度學習 出身，修改了幾次也沒能成功修改成需要符合需求的測試代碼（有高手的話，麻煩能修改好，聯繫我測試測試）。這段代碼是自帶 創建 rank 的，並且是單機運行的。我想要的是，自己創建每一個 rank，並且支持多機分佈式。
多機分佈式訓練
- 需要安裝 python3 環境，並且安裝 torch 和 torchvision 。兩臺機器的 python 環境應當保持一致。
- /mnt/pytorch/ 是一個分佈式存儲掛載目錄。
- 執行程序需要下載數據，如果你需要數據，可以在 https://download.csdn.net/download/github_37320188/11670737 查找 （由於沒有積分了，小收幾個積分用用，不能老做公益呀）。當然也可以自己下載，有合適的網絡環境的話，相對來說就是花點時間。
# 機器 1 上輸入命令
python3 /mnt/pytorch/mnist.py --rank 0 --world-size 2 --init-method tcp://192.168.61.55:9090 --backend gloo
# 機器 2 上輸入命令
python3 /mnt/pytorch/mnist.py --rank 1 --world-size 2 --init-method tcp://192.168.61.55:9090 --backend gloo
其實這已經是一個多機訓練的例子了。但是我建議在操作之前，先單節點 跑 2 個 rank 試試。
並且 run 一下 官方 示例 的一些代碼，玩一下。
這個操作中，以 gloo 作 backend ，以 rank0 作爲 master 通信，9090 機器 1 的未被佔用通信端口。相關概念可以仔細閱讀官方說明。
docker 分佈式訓練
- 首先你需要一個 支持 python 、torch、torchvision 的 docker image。如果選擇省事，我建議去docker hub 下載一個官方的 tag 爲 pytorch/pytorch:latest 的 image 。
- 然後，這個 image 自帶了 conda 、python3、torch，但是你需要自己安裝 torchvision ，打包一個新的 image。建議通過 pip 安裝，並且受限於 國內環境問題，你可能一次成功不了，可能需要多嘗試幾次。
- 最後，如果實在 操作 docker 容器的耐心都沒有了，你可以從我的 docker hub 裏下載一個。tag 是 docker.io/loongc/pytorch:pytorch
操作一波
# 機器 1
docker run -it --rm -v /mnt/pytorch/:/pytorch -w /pytorch --net=host loongc/pytorch:pytorch python mnist.py --rank 0 --world-size 2 --init-method tcp://192.168.61.55:9090 --backend gloo
# 機器 2
docker run -it --rm -v /mnt/pytorch/:/pytorch -w /pytorch --net=host loongc/pytorch:pytorch python mnist.py --rank 1 --world-size 2 --init-method tcp://192.168.61.55:9090 --backend gloo
討論
以上部分，測試都是在 多節點CPU 上運行的。如果你需要在GPU 進行運算：
- 首先，你需要節點本身 支持gpu 運算，安裝 cuda 環境
- 如果你希望在 docker 中進行 gpu 運算，那麼 docker image 需要支持gpu 運算。幸運的是，前文提到的 image 支持gpu 運算
如果，環境已經準備好了，就可以做一些 GPU 相關的測試了。
使用 基於 Docker 的 GPU 分佈式運算，如果不需要 諸如 運算性能 特殊的需求，你可以直接使用 上一節 的 操作方法即可。如果 需要 諸如使用 分佈式文件系統通信、環境變量通信、InfiniBand通信，使用 nccl 作爲 backend 等等，將在下一節做一個總結。
backend 與通信方式 的一些總結
引言中，已經提到 3 種 backend，分別是：MPI、GLOO、NCCL。這 3 種 backend 支持了 3 種 通信 方式：tcp、環境變量、分佈式文件系統（當前使用的 pytorch 版本如此，後續的版本就不知道了）。
所以，大致上可以認爲不同 rank 之間的通信就是 由 backend 與 通信方式的 組合完成的。當然，在不同的組合方式所帶來的 通信速度（影響運算整體性能）是不同的，能實現的功能也是不太一樣的（詳情見引言中的表格）。
本文前面的內容當中，使用的 backend 都是GLOO，使用的通信方式都是 tcp。
本節內容，主要圍繞 NCCL 的 3 種 通信操作方式展開，以及如何使用 InfiniBand。
MPI 的方式我沒有測過，也不打算測試，主要原因是需要 額外的依賴 。
NCCL 與 tcp
nccl 與 gloo 使用 tcp 的方式是一致的，不需要多餘的操作。在運行 python 腳本的時候，只需要將傳入 backend 的參數 gloo 改爲 nccl 即可。
NCCL 與 環境變量
nccl 使用環境變量，相對於 tcp 要複雜一些。
- 首先，需要將傳入 backend 的參數 gloo 改爲 nccl
- 其次，將傳入 init-method 的參數 由 tcp://ip:port 改爲 env://
- 另外，容器啓動的時候的需要給容器設置 2 個環境變量 MASTER_ADDR ，MASTER_PORT
還是在 192.168.61.55 和 192.168.61.56 兩臺機器上操作，啓動大致如下：
# 機器 1
docker run -it --rm --env-file /mnt/pytorch_env/master -v /mnt/pytorch/:/pytorch -w /pytorch --net=host loongc/pytorch:pytorch python mnist.py --rank 0 --world-size 2 --init-method env:// --backend nccl
# 機器 2
docker run -it --rm --env-file /mnt/pytorch_env/master -v /mnt/pytorch/:/pytorch -w /pytorch --net=host loongc/pytorch:pytorch python mnist.py --rank 1 --world-size 2 --init-method env:// --backend nccl
# 文件 /mnt/pytorch_env/master 內容如下
MASTER_ADDR=192.168.61.55
MASTER_PORT=9090
上述 /mnt/pytorch_env/master 文件的內容中 MASTER_ADDR=192.168.61.55 是 rank0 所在機器的 ip ，MASTER_PORT=9090 是該機器上一個未被徵用的 端口
NCCL 與 分佈式文件系統
nccl 使用分佈式文件系統，那麼就需要把分佈式文件系統 映射進 容器當中，在這裏，我是在 掛載目錄下 /mnt/pytorch/ 創建了一個文件 backendnccl ，因爲本身這個掛載 就是分佈式文件系統。
# 以機器 1 上的操作爲例
docker run -it --rm -v /mnt/pytorch/:/pytorch -w /pytorch --net=host loongc/pytorch:pytorch python mnist.py --rank 0 --world-size 2 --init-method file:///pytorch/backendnccl --backend nccl
使用 InfiniBand
原本，pytorch 是支持 IP over IB 的，所以可操作 InfiniBand 的方式主要有以下幾種：
- 在使用 tcp方式 進行通信 時，可以不使用 IP 地址，而直接使用 IB (InfiniBand) 卡的地址
- 在使用環境變量過程中，設置的環境變量中 將 MASTER_ADDR 設置成 IB 卡的地址
另外，NCCL和Gloo後端都會嘗試找到要使用的正確網絡接口，如果自動檢測到的接口不正確，您可以使用以下環境變量覆蓋它（適用於相應的後端）：
在使用 過程中，如果確定 使用 nccl，並且打算用 InfiniBand ，推薦通過環境變量（寫進 前文提到的 文件 /mnt/pytorch_env/master ） NCCL_SOCKET_IFNAME=ib0
額外說明，如果你的機器上有多張 InfiniBand 卡，你可以通過 NCCL_SOCKET_IFNAME 指定任意一張，如 ib0 或者 ib1 這種
但是，你可以直接使用 NCCL_SOCKET_IFNAME=ib 進行指定，這樣所有的 ib 卡都有機會被調用起來用於通信。你可以通過相關工具，如 ibdump 進行監控 ib 卡的發包 和 數據的發送與接受。
經過分析，我們認爲在 以發送數據包 的方式進行通信的過程中，除了真正用於計算的數據包，還可能包括 機器通信 的數據。
總結
這裏，實現了手動啓動多個rank 進行 基於 docker 的分佈式訓練。那麼，自動化進行分佈式訓練就變得比較容易了。
可以通過 mesos 或者 k8s 啓動多個 容器，進行集羣管理，完成分佈式訓練。通過 設計 api 指定啓動 指定數目的 rank，指定硬件資源，以及是否使用 InfiniBand 。
如果有更好的實現 分佈式訓練 PyTorch ，歡迎留言討論。