Tensorflow上手4: 初探分布式训练

Source: Deep Learning on Medium


Go to the profile of Dong Wang

今天打算简单讲讲Tensorflow分布式训练的方法和遇到的一些问题。感兴趣的朋友也可以进入Google Colab提供的文档中进行更多学习。

一般情况下,分布式训练主要有两种模式,一种是模型并行,一种是数据并行。模型并行指将模型分散到不同的GPU当中,每个GPU只计算和更新自己负责的参数,这种模型并行主要解决模型参数量过大的问题。数据并行是在多个GPU上运行同样的模型,但是他们的输入数据却不一样,实现更大的batch同时计算,加快模型训练速度。接下来的内容针对数据并行展开进行讨论。

Tensorflow分布式训练的支持主要是通过tf.distribute.Strategy来实现。现在的Tensorflow分布式库中主要包含五种不同的分布式策略:

  • MirroredStrategy
  • TPUStrategy
  • MultiWorkerMirroredStrategy
  • CentralStorageStrategy
  • ParameterServerStrategy

在这里面我很少听说有人用CentralStorageStrategy,自己也没有过尝试,所以暂时不好说些什么。TPUStrategy可以让用户更好的在Google的Tensor Processing Unites(TPUs)上面进行训练。他与MirroredStrategy类似,实现了同步的分布式训练(Synchronized distributed training)。TPUs本身有很快的all-reduce操作,而TPUStrategy则实现了一系列TPU ops来利用这一特性。对于一个不在Google的工程师,我常接触的则是多个GPU,多台机器的情况,所以这里也不对TPUStrategy进行太多探讨。

MirroredStrategy

MirroredStrategy是一种支持多张GPU在同一个机器上的同步训练方法。在训练开始时,Mirrored会在每张卡上复制一份模型,个显卡会收到tf.data.Dataset传来的数据,独立计算梯度,然后采用all-reduce的方法进行同步更新。多个显卡在通信时默认使用Nvidia NCCL进行。

我们可以深入MirroredStrategy的实现了解一下。基本上所有的distributed strategy都是通过某些collective ops和cross device ops进行数据通讯。MirroredStrategy也是如此,它是这样选择cross device ops的:

if len(workers) > 1:
if not isinstance(self._cross_device_ops, cross_device_ops_lib.MultiWorkerAllReduce):
raise ValueError(
"In-graph multi-worker training with `MirroredStrategy` is not "
"supported.")
self._inferred_cross_device_ops = self._cross_device_ops
else:
# TODO(yuefengz): make `choose_the_best` work with device strings
# containing job names.
self._inferred_cross_device_ops = cross_device_ops_lib.NcclAllReduce()

这也就印证了MirroredStrategy在单机多卡的情况下默认使用NCCL来进行通信的说明。具体的实现大家可以去查看AllReduceCrossDeviceOps的实现。

同时,上面的程序也说明MirroredStrategy可以运用到多机多卡的情况中去,然而多机多卡的情况下用户需要自己传入cross_device_ops_lib.MultiWorkerAllReduce进行通讯,这里MultiWorkerAllReduce支持若干种通讯方式,比如nccl, nccl/xring, nccl/rechd, nccl/pscpu, xring, pscpu, pscpu/pscpu等等。由于目前最佳的通讯方式需要NCCL2.0加上xring,然而Tensorflow目前使用NCCL 1.1,并且nccl/xring在现有的代码中有bug无法工作,所以这一模式常常被大家诟病。

在Esitmator当中使用MirroredStrategy是非常简单的,我们只需要加入以下代码即可:

mirrored_strategy = tf.contrib.distribute.MirroredStrategy(
num_gpus_per_worker=x)
estimator_config = tf.estimator.RunConfig(
train_distribute=mirrored_strategy)
estimator = tf.estimator.Estimator(
model_fn=...
config=estimator_config,
...)

训练脚本就会自动进行分布式训练。

ParameterServerStrategy

接下来我们先跳过MultiWorkerMirroredStrategy讨论一下ParameterServerStrategy,他是Tensorflow最初的分布式训练方法。其概念如下图所示,它由若干个parameter servers和若干个worker servers构成,parameter servers用于存储参数,workers用于计算。

https://eng.uber.com/horovod/

ParameterServerStrategy是一种严格的异步训练方法,在训练过程中worker servers会和不同的parameter servers沟通获得参数,然后计算,向parameter servers传递参数的梯度。配置一个这样的训练环境非常简单,只需要在程序运行时设置好环境变量TF_CONFIG,需要注意的是需要给分布式集群里每一个机子不同的task。

os.environ["TF_CONFIG"] = json.dumps({
"cluster": {
"worker": ["host1:port", "host2:port", "host3:port"],
"ps": ["host4:port", "host5:port"]
},
"task": {"type": "worker", "index": 1}
})

同时,ParameterServerStrategy还有比较神奇的功能,它可以通过传入num_gpus_per_worker在一个worker上进行多GPU的同步计算,然后不同worker之间进行异步计算。但是由于单一worker上多GPU并没有利用NCCL进行通讯,而是直接将结果发送到CPU,所以效率非常低下。

最后要说的是Parameter Server的配置数量也非常复杂,不同的网络环境,模型大小都会对效率有影响,所以现在官方好像也不怎么推荐这种做法了。

MultiWorkerMirroredStrategy

tf.distribute.experimental.MultiWorkerMirroredStrategy在之前的版本里应该对应的是tf.contrib.distribute.CollectiveAllReduceStrategy,他与MirroredStrategy非常类似,都在每一个device上存储一份模型的备份,进行同步的分布式训练。

该策略采用CollectiveOps作为多个worker之间通讯的操作。所谓的collective op是Tensorflow自己实现的根据当前硬件环境,网络结构,和Tensor大小自动采用最佳算法进行all-reduce的计算操作。一个collective op的实现逻辑十分简单

if (CanProceedWithCompute(c, col_exec, done)) {
col_exec->ExecuteAsync(
c, col_params_, GetCollectiveKey(c), actual_done);
}

c是当前op的计算状态,col_exec是Tensorflow根据系统情况选择的collective executor,所有的all reduce,boardcast和receive操作都有collective executor去执行。

该策略目前也实现了很多优化,比方将很多个小tensor的all reduce操作变成几个大tensor的all reduce操作,以及在开发当中的采用最新NCCL 2.0进行通讯的操作,具体可以参见Issue 24505。可以看出Tensorflow分布式训练在被吐槽很多次后,感受到了来自Pytorch,Horovod的压力,在努力的提升自己。

最后,关于MultiWorkerMirroredStrategy的配置,有两点需要注意。一点是collective ops的策略选择,目前支持CollectiveCommunication.RING,采用与Horovod类似的ring-based通讯策略。另一个是CollectiveCommunication.NCCL,采用Nvidia NCCL进行通讯,在启动策略时可以传入参数指定:

multiworker_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
tf.distribute.experimental.CollectiveCommunication.NCCL)

另一个需要注意的是关于TF_CONFIG的设置,该策略并不需要指定Parameter server,只需要一系列worker即可,其配置如下:

TF_CONFIG = {
'cluster': {
'worker': ['worker1:port1', 'worker2:port2', 'worker3:port3', ...]
},
'task': {'type': 'worker', 'index': 0}
})

结语

今天讨论了这么久关于Tensorflow分布式训练的内容,主要目的是抛砖引玉,Tensorflow分布式训练代码里的bug不少,还有很多需要摸索的地方,欢迎大家多多交流。