Pytorch 多GPU訓練-多計算節點並行-All you need

概述

本篇介紹多計算節點上的pytorch分布式訓練。從環境配置到運行demo的所有步驟，step by step。沒有理論原理，理論原理可以參考這里.

基礎環境

• 多台linux計算節點，通過網絡連接，不同主機之間可以相互ping通。網速越快越好，如果通信速度比較慢，就不用怎么考慮分布式訓練。

• 所有linux計算節點都包含若干GPU，GPU數量可以不一致，但是所有GPU計算速度盡量一致，否則模型的同步時會等待大量時間（短板效應）。

• 所有計算節點都擁有Pytorch運行環境，即都可以單獨的運行訓練代碼。

NFS環境搭建

這里NFS有兩個作用：（1）集群初始化時作為多個進程（不同主機上）UDS（unix domain socket）的通信地址；（2）數據集統一存放在NFS目錄上，所有進程都可以同時訪問。

選擇一台與計算節點同一網段的linux主機作為存儲節點，或者使用計算節點中的一台作為存儲節點。

存儲節點操作如下：

apt-get install nfs-kernel-server
<path to share> *(rw,sync,no_root_squash,no_subtree_check) # insert this line to /etc/exports, save & exit
exportfs -a 

數據集保存到共享目錄中。
所有計算節點將共享目錄掛載到本地：

mkdir /mnt/nfs
mount -t nfs <ip of storage node>:<path to share> /mnt/nfs

分布式訓練

代碼見: github ，需要自己更改數據集的地址。

運行demo

分別在三個主機上運行如下的命令。word_size 為 3，當啟動的進程數少於3時，所有已經啟動的進程會等待，進程數達到3時集群創建成功，訓練開始。每個計算節點都使用該計算節點上的兩個GPU --gpu_devices 0 1。

# node 1
python multigpu_demo_v3.py \
    --init_method file://<absolute path to nfs file> \
    --rank 0 \
    --world_size 3 \
    --gpu_devices 0 1

# node 2
python multigpu_demo_v3.py \
    --init_method file://<absolute path to nfs file> \
    --rank 1 \
    --world_size 3 \
    --gpu_devices 0 1

# node 3
python multigpu_demo_v3.py \
    --init_method file://<absolute path to nfs file> \
    --rank 2 \
    --world_size 3\
    --gpu_devices 0 1

更多測試細節見github。

總結

Pytorch里的分布式訓練只能實現增大batch size的作用，對於速度的提升不明顯，對於Batch Norm沒有幫助，如果要提升BN的效果，需要用Sync Bn。