[分布式訓練] 單機多卡的正確打開方式：Horovod

轉自：https://fyubang.com/2019/07/26/distributed-training4/

講完了單機多卡的分布式訓練的理論、TensorFlow和PyTorch分別的實現后，今天瓦礫講一個強大的第三方插件：Horovod。

Horovod是Uber開源的跨平臺的分布式訓練工具，名字來自于俄國傳統民間舞蹈，舞者手牽手圍成一個圈跳舞，與Horovod設備之間的通信模式很像，有以下幾個特點：

兼容TensorFlow、Keras和PyTorch機器學習框架。

使用Ring-AllReduce算法，對比Parameter Server算法，有著無需等待，負載均衡的優點。

實現簡單，五分鐘包教包會。（劃重點）

Uber官方在git上給了很詳細的例子： https://github.com/horovod/horovod/tree/master/examples，所以這里只簡單講一下大概的使用方法：

TensorFlow

以TF的Custom Training Loop API為例：

import tensorflow as tf import horovod.tensorflow as hvd# 1. 初始化horovod hvd.init() # 2. 給當前進程分配對應的gpu，local_rank()返回的是當前是第幾個進程 config = tf.ConfigProto() config.gpu_options.visible_device_list = str(hvd.local_rank()) # 3. Scale學習率，封裝優化器 opt = tf.train.AdagradOptimizer(0.01 * hvd.size()) opt = hvd.DistributedOptimizer(opt) # 4. 定義初始化的時候廣播參數的hook，這個是為了在一開始的時候同步各個gpu之間的參數 hooks = [hvd.BroadcastGlobalVariablesHook(0)] # 搭建model，定義loss loss = ... train_op = opt.minimize(loss) # 5. 只保存一份ckpt就行 checkpoint_dir = '/tmp/train_logs' if hvd.rank() == 0 else None # 7. 用MonitoredTrainingSession實現初始化，讀寫ckpt with tf.train.MonitoredTrainingSession(checkpoint_dir=checkpoint_dir,config=config,hooks=hooks) as mon_sess:while not mon_sess.should_stop():# Perform synchronous training.mon_sess.run(train_op)

具體的代碼看tensorflow_mnist.py：https://github.com/horovod/horovod/blob/master/examples/tensorflow_mnist.py

單機雙卡訓練輸入以下命令：

CUDA_VISIBLE_DEVICES=6,7 horovodrun -np 2 -H localhost:2 python tensorflow_mnist.py

這里 -np指的是進程的數量。

執行之后可以看到如下的結果，因為多線程，每個step都打印了兩遍。

[1,0]<stderr>:INFO:tensorflow:loss = 0.13126025, step = 300 (0.191 sec) [1,1]<stderr>:INFO:tensorflow:loss = 0.01396352, step = 310 (0.177 sec) [1,0]<stderr>:INFO:tensorflow:loss = 0.063738815, step = 310 (0.182 sec) [1,1]<stderr>:INFO:tensorflow:loss = 0.044452004, step = 320 (0.215 sec) [1,0]<stderr>:INFO:tensorflow:loss = 0.028987963, step = 320 (0.212 sec) [1,0]<stderr>:INFO:tensorflow:loss = 0.09094897, step = 330 (0.206 sec) [1,1]<stderr>:INFO:tensorflow:loss = 0.11366991, step = 330 (0.210 sec) [1,0]<stderr>:INFO:tensorflow:loss = 0.08559138, step = 340 (0.200 sec) [1,1]<stderr>:INFO:tensorflow:loss = 0.037002128, step = 340 (0.201 sec) [1,0]<stderr>:INFO:tensorflow:loss = 0.15422738, step = 350 (0.181 sec) [1,1]<stderr>:INFO:tensorflow:loss = 0.06424393, step = 350 (0.179 sec)

PyTorch

Torch下也是類似的套路，但是由于PyTorch本身單機多卡訓練已經夠簡單了，API也穩定，所以筆者一般做的時候就是直接用Torch自己的DP和DDP了。

import torch import horovod.torch as hvd# 1. 初始化horovod hvd.init() # 2. 給當前進程分配對應的gpu，local_rank()返回的是當前是第幾個進程 torch.cuda.set_device(hvd.local_rank()) # Define dataset... train_dataset = ... # 3. 用DistributedSampler給各個worker分數據 train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, num_replicas=hvd.size(), rank=hvd.rank()) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler) # Build model... model = ... model.cuda() # 4. 封裝優化器 optimizer = optim.SGD(model.parameters()) optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters()) # 5. 初始化的時候廣播參數，這個是為了在一開始的時候同步各個gpu之間的參數 hvd.broadcast_parameters(model.state_dict(), root_rank=0) # 訓練 for epoch in range(100):for batch_idx, (data, target) in enumerate(train_loader):optimizer.zero_grad()output = model(data)loss = F.nll_loss(output, target)loss.backward()optimizer.step()if batch_idx % args.log_interval == 0:print('Train Epoch: {} [{}/{}]\tLoss: {}'.format(epoch, batch_idx * len(data), len(train_sampler), loss.item()))

速度

瓦礫還沒有來得及做一個全面的Horovod、tf.distribute和 Torch的單機多卡訓練速度的橫向對比，不過大家可以參考這兩篇：

Horovod: fast and easy distributed deep learning in TensorFlow

Goodbye Horovod, Hello CollectiveAllReduce

總體而言，用了All-Reduce算法的API，速度應該都差不多，如果你是土豪，擁有NVLINK（卡間通信極快）的話，那忘了我說的這幾篇“廢話”吧朋友。Orz。

總結

終于結束了單機多卡系列的最后一章，由于博客本身的限制，給的例子整體還是比較簡單，以入門為主，大家具體使用的時候肯定還是會遇到一些坑，這里瓦礫把踩過的一些坑和解決辦法列舉在這，以避免大家以后重復踩坑：

• tf.contrib.distributed.MirroredStrategy 需要optimizer支持merge_call（bert實現的optimizer是直接修改apply_gradient的，所以會報錯），這個時候就需要正確地修改optimizer里的_apply_dense、_apply_sparse(參考Issue 23986 和 JayYip)。或者用horovod，就可以避免這個問題。
• Effective batch size，不同的多卡工具對輸入的batch size的操作不一樣，要確定最后進模型的effective batch size才有意義。一般來說，多進程的batch size指的是每張卡的batch size。
• Learning rate scale，學習率要根據effective batch size調整。
• All-Reduce由于是多進程的，數據流各自獨立，為了防止同一個step多gpu的batch重疊，最好的的辦法是在每個進程里根據local_rank設置shard的數據，保證各個gpu采樣的數據不重疊。
• 為了使用horovod，新建docker container時，要加—privileged，否則會瘋狂報warning，雖然沒影響，但是看著難受。
• Pytorch的DP多卡要注意最后一個batch的batch size不能小于gpu的數量，否則會報錯，最保險的做法是drop_last，扔掉最后的batch。
• 并不是所有情況下All-Reduce都比PS好，比如當卡間通信用的是NVLink的時候，在gpu數量不多的情況下，數據傳輸的時間不是瓶頸，All-Reduce的提升就幾乎沒有了。
• DP和DDP有一個區別在于BatchNorm。
• DDP封裝model后不能再改動model。
• 待補充。。。

Reference

Horovod的官方給的一些例子。

Uber：如何用Horovod實現bert的單機多卡訓練

Goodbye Horovod, Hello CollectiveAllReduce

Horovod: fast and easy distributed deep learning in TensorFlow

總結

以上是生活随笔為你收集整理的[分布式训练] 单机多卡的正确打开方式：Horovod的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。