[分布式訓(xùn)練] 單機(jī)多卡的正確打開方式：Horovod

轉(zhuǎn)自：https://fyubang.com/2019/07/26/distributed-training4/

講完了單機(jī)多卡的分布式訓(xùn)練的理論、TensorFlow和PyTorch分別的實(shí)現(xiàn)后，今天瓦礫講一個(gè)強(qiáng)大的第三方插件：Horovod。

Horovod是Uber開源的跨平臺(tái)的分布式訓(xùn)練工具，名字來自于俄國(guó)傳統(tǒng)民間舞蹈，舞者手牽手圍成一個(gè)圈跳舞，與Horovod設(shè)備之間的通信模式很像，有以下幾個(gè)特點(diǎn)：

兼容TensorFlow、Keras和PyTorch機(jī)器學(xué)習(xí)框架。

使用Ring-AllReduce算法，對(duì)比Parameter Server算法，有著無需等待，負(fù)載均衡的優(yōu)點(diǎn)。

實(shí)現(xiàn)簡(jiǎn)單，五分鐘包教包會(huì)。（劃重點(diǎn)）

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

TensorFlow

以TF的Custom Training Loop API為例：

import tensorflow as tf import horovod.tensorflow as hvd# 1. 初始化horovod hvd.init() # 2. 給當(dāng)前進(jìn)程分配對(duì)應(yīng)的gpu，local_rank()返回的是當(dāng)前是第幾個(gè)進(jìn)程 config = tf.ConfigProto() config.gpu_options.visible_device_list = str(hvd.local_rank()) # 3. Scale學(xué)習(xí)率，封裝優(yōu)化器 opt = tf.train.AdagradOptimizer(0.01 * hvd.size()) opt = hvd.DistributedOptimizer(opt) # 4. 定義初始化的時(shí)候廣播參數(shù)的hook，這個(gè)是為了在一開始的時(shí)候同步各個(gè)gpu之間的參數(shù) 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實(shí)現(xiàn)初始化，讀寫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

單機(jī)雙卡訓(xùn)練輸入以下命令：

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

這里 -np指的是進(jìn)程的數(shù)量。

執(zhí)行之后可以看到如下的結(jié)果，因?yàn)槎嗑€程，每個(gè)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本身單機(jī)多卡訓(xùn)練已經(jīng)夠簡(jiǎn)單了，API也穩(wěn)定，所以筆者一般做的時(shí)候就是直接用Torch自己的DP和DDP了。

import torch import horovod.torch as hvd# 1. 初始化horovod hvd.init() # 2. 給當(dāng)前進(jìn)程分配對(duì)應(yīng)的gpu，local_rank()返回的是當(dāng)前是第幾個(gè)進(jìn)程 torch.cuda.set_device(hvd.local_rank()) # Define dataset... train_dataset = ... # 3. 用DistributedSampler給各個(gè)worker分?jǐn)?shù)據(jù) 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. 封裝優(yōu)化器 optimizer = optim.SGD(model.parameters()) optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters()) # 5. 初始化的時(shí)候廣播參數(shù)，這個(gè)是為了在一開始的時(shí)候同步各個(gè)gpu之間的參數(shù) hvd.broadcast_parameters(model.state_dict(), root_rank=0) # 訓(xùn)練 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()))

速度

瓦礫還沒有來得及做一個(gè)全面的Horovod、tf.distribute和 Torch的單機(jī)多卡訓(xùn)練速度的橫向?qū)Ρ?#xff0c;不過大家可以參考這兩篇：

Horovod: fast and easy distributed deep learning in TensorFlow

Goodbye Horovod, Hello CollectiveAllReduce

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

總結(jié)

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

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

Reference

Horovod的官方給的一些例子。

Uber：如何用Horovod實(shí)現(xiàn)bert的單機(jī)多卡訓(xùn)練

Goodbye Horovod, Hello CollectiveAllReduce

Horovod: fast and easy distributed deep learning in TensorFlow

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。