算法研发-Horovod分布式训练

背景

        训练现代深度学习模型需要大量计算，通常由多个GPU提供。这会遇到如下问题，第一，必须支持GPU间通信，取决硬件支持，这种通信会产生重大开销。其次，用户必须修改训练代码，以利用GPU间通信。所需的修改可能是重要的或很小的。

        在TensorFlow库下，启用多GPU训练需要不可忽视的通信开销，并要求用户大量修改他们的代码。在测试中我们发现，每个使用分布式 TensorFlow 的案例都需要指定初始工作线程和参数服务器。此外，用户必须保证所有的操作都正确地使用 tf.train.device_replica_setter()，并使用 towers 让代码符合服务器中多 GPU 的设置，这通常导致陡峭的学习曲线和大量的代码重构，压缩了实际建模的时间，大规模的计算性能方面，标准的分布式TensorFlow机制无法满足需求，例如，在使用 128个GPU进行训练时，因为低效率损失了一半的计算资源。

Horovod

        通过环形拓扑结构来实现高效的GPU间通信，并且只需对用户代码进行几行修改，即可实现更快、更轻松的分布式训练，吸取了Facebook "Training ImageNet In 1 Hour" 与百度 "Ring Allreduce" 的优点。horovod的特点就是简单易用，深度学习框架很多，但是使用者只需要为Horovod进行配置就可以实现分布式训练，而不是对每个架构有不同的配置方法。

原理

        Horovod的本质就是通过数据并行实现分布式训练，假设每个机器里都有一个完整模型，然后把数据切分成n块，把n块分发给每个计算单元，每个计算单元独自计算出自己的梯度。同时每个计算单元的梯度会进行平均、同步，同步后的梯度可以在每个节点独立去让它修正模型，整个过程结束后每个节点会得到同样的模型。这个方法可以让能够处理的数据量增加，变成了原来的n倍。

 

        实现“数据并行”有两种方法，即构建参数服务器和Ring-AllReduce

        方法一：参数服务器（Parameter Server）。

        在计算单元以外加设新的服务器叫做参数服务器。每次训练的时候每个计算单元把梯度发送给参数服务器，服务器把他们进行汇总计算平均值，把平均值返回到每个计算单元，这样每个计算单元就同步了。

        方法二：Ring-AllReduce。

        该方法基于从高性能计算集合通信，原理是把每个计算单元构建成形成一个环，做梯度平均时每个计算单元先把梯度切分成N部分，然后发送到相邻的下一个模块。假设有N个节点，那么N-1次发送后所有节点就可以获得所有其他节点的数据。优点：如果在大规模分布式训练时候网络占用是最优的；缺点：容错性较差，很多实现都是用MPI实现（MP本身并不是为容错设计的，它更偏向于照顾高性能的计算）。

核心

        Horovod核心基于MPI的size、rank、localrank、allreduce、allgather和broadcast。

        假设我们在4台服务器上启动了一个培训脚本，每台服务器都有4个GPU。如果我们为每个GPU启动一个脚本：

        size是进程数，比如16。

        Rank将是从0到15（size-1）的唯一进程ID。

        Local rank则是服务器中从0到3的唯一进程ID。

        Allreduce是一种在多个进程之间聚合数据并将结果分发回它们的操作。

        Broadcast 是将数据从一个进程传播到其他每个进程的操作

过程

1. 在代码里引入Horovod库，做初始化；
2. 定义优化器的时候，优化器重要参数是学习率；
3. 用一个对象把原来优化器包起来，这就是抽象了Horovod需要进行的梯度平均的逻辑，会比较容易使用；
4. 让Horovod把每次训练的初始状态广播到每个节点，这样保证每个节点从同一个地方开始；
5. 训练的长度，由于把训练分布在N个机器上了，所以可以极大减小，除以N就可以了。

安装

pip install horovod

        或者

pip install horovod[tensorflow,keras,pytorch,mxnet,spark]

使用

        把单个计算单元训练变成多机分布式训练用Horovod是非常简单的，只需要做三步。第一步程序引入Horovod稍微做修改，调整学习率、时间；第二步处理训练的数据，进行分布化；第三步用Horovodrun程序进行启动，就可以进行分布式训练了。

import horovod.torch as hvd

hvd.init()

if torch.cuda.is_available():

    # Horovod: pin GPU to local rank. Pin each GPU to a single process.

    torch.cuda.set_device(hvd.local_rank())



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=batch_size, sampler=train_sampler, **kwargs)



model = ...

model.cuda()



# scale learning rate by the number of GPUs.

optimizer = tf.train.AdagradOptimizer(0.01 * hvd.size())



# Horovod: wrap optimizer with DistributedOptimizer.

optimizer = hvd.DistributedOptimizer(optimizer,named_parameters=model.named_parameters())



# broadcast parameters & optimizer initial variable states from rank 0 to all other processes:

hvd.broadcast_parameters(model.state_dict(), root_rank=0)

hvd.broadcast_optimizer_state(optimizer, 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()))

启动训练

horovodrun -np 4 -H localhost:4 python train.py

horovodrun -np 16 -H server1:4,server2:4,server3:4,server4:4 python train.py



#specify host nodes in a host file. For example:

cat myhostfile

aa slots=2

bb slots=2

cc slots=2

horovodrun -np 6 -hostfile myhostfile python train.py

官方的demo地址：

https://github.com/horovod/horovod/blob/master/examples/pytorch/pytorch_mnist.py

补充   

        做一些小规模的，比如单机八卡这种，horovod的速度比torch.distributed的速度慢，但是要是大规模集群，horovod还是很有优势的。