当我们拥有大量的数据后，尤其是大规模文本、多模态、视频序列数据及其使用大型预训练模型等等情况下，训练好一个模型不得不借助分布式策略来提高计算资源的使用效率，进而缩短模型的训练时间。文本总结一下Tensorflow的分布式多卡训练，包括单机多卡训练与多机的分布式训练。

Tensorflow提供API包括tf.distribute.Strategy方便我们实现多卡或多记训练策略。

基本理论

首先介绍若干个概念。

模型并行与数据并行：

  • 模型并行，不同的GPU上各自运行模型的不同部分，如多个层分别运行于不同的GPU
  • 数据并行，不同的GPU上运行同一个模型的副本，但是使用不同的数据来计算梯度

数据并行由于不同GPU使用不同的数据来计算，这就涉及到参数更新时同步的问题。方法有同步更新与异步更新：

  • 同步更新，每个batch内所有GPU都计算完成后，整合好各个GPU上的计算结果统一计算新权重，接着更新所有GPU上的副本模型参数后进行下一轮计算。同步更新需要所有GPU协调，更新速度取决于最慢的GPU，即受到木桶效应约束。
  • 异步更新，每个GPU计算完梯度后，不需要协调其他GPU，马上更新整体权重与同步。这种同步方式无需等待，但过程涉及复杂的更新方法。

实践中，参数同步更新方法常见有Parameter Server与Ring All Reduce：

  • Parameter Server，Reducer服务器负责把batch数据分发到各个GPU上，各个GPU计算梯度后send到Reducer上累加，得到该batch的准确梯度后更新参数，参数再同步到各个GPU上。这种方法由于各个GPU都需要与Reducer进行数据、梯度和模型参数的通信，每个batch的计算通信开销很大，且和GPU数量线性相关。此外，Reduce过程需要等待多有GPU，受到木桶效应约束。Google的第一代分布式机器学习框架DistBelief，其多卡训练也是使用Parameter Server架构。
  • Ring All Reduce，分为两个步骤Scatter Reduce与All Gather。Ring All Reduce的通信成本与GPU数量无关且恒定。

Scatter Reduce与All Gather两个过程可以参考 https://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/，这里简单叙述一下：

  • Scatter Reduce，GPU间构成一个环，假设参数分成N份，相邻的GPU传递与接收不同的参数k（$1 \le k \le N$），在传递N-1次后，可以得到分散到不同GPU上的N份参数的累积。
  • All Gather，由于这些参数的累积是分散到不同GPU上，于是再同步一下就能够让参数的累积到同步到不同GPU上。于是每个GPU都获得最新的完整的参数更新。

环境准备

首先我们需要准备单机多卡环境或多机多卡环境。单主机情况下，可以使用如下方式查看CPU或GPU设备的具体信息。tf.config.list_physical_devices可以查看本地GPU、CPU设备，

1
2
3
>>> gpus = tf.config.list_physical_devices(device_type="GPU")
>>> cpus = tf.config.list_physical_devices(device_type="CPU")
>>> print(gpus, cpus)

如果不具备多卡环境，可以使用单GPU来模拟多个GPU环境，例如一个8G显存的显卡，虚拟化为两个4G显存的GPU，具体如下，

1
2
3
4
5
6
7
# 环境下只有一个GPU
gpus = tf.config.list_physical_devices("GPU")
tf.config.set_logical_device_configuration(
    gpus[0],
    [tf.config.LogicalDeviceConfiguration(memory_limit=4096),
     tf.config.LogicalDeviceConfiguration(memory_limit=4096)]
)

然后我们就可以看到两个逻辑GPU，

1
2
>>> tf.config.list_logical_devices("GPU")
[LogicalDevice(name='/device:GPU:0', device_type='GPU'), LogicalDevice(name='/device:GPU:1', device_type='GPU')]

单机多卡训练

单机多卡训练是指多个GPU在同一台主机上训练，在Tensorflow中，tf.distribute.MirroredStrategy提供该训练策略的实现，

1
2
3
4
5
6
# 创建MirroredStrategy的实例，并指定参与对卡训练的GPU，这里指定0~2GPU参与单机多卡训练
strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1", "/gpu:2"])
# 把要支持单机多卡训练的模型放到strategy的Python context下
with strategy.scope():
    # 实现模型的代码
    # 预加载权重

MirroredStrategy策略，即镜像策略，每个GPU上都有一份模型的副本。

那么MirroredStrategy下的单机多卡训练的原理是怎么样的？

  • 首先是在训练前，context下的模型会被复制完整的模型到MirroredStrategy所有指定的K个计算设备上
  • 对于每个batch的数据，会平分成K份，分别送到K个设备上，并使用模型的本地变量计算当前获得数据的梯度
  • 各个设备互相交换梯度数据（即All-reduce操作，默认是使用NVIDIA NCCL），使得每个设备都具有该batch所有K份数据的梯度的和
  • 根据梯度的和结果更新各个设备的本地变量
  • 接着继续进行下一轮训练

多机训练

多机训练是指多个GPU在不同的主机上，联合多个GPU（或其他计算设备）的分布式训练。它的原理和单机多卡训练中的tf.distribute.MirroredStrategy，只不过这里换成tf.distribute.MultiWorkerMirroredStrategy。MultiWorkerMirroredStrategy其实是单机多卡的镜像策略推广到多节点后的情况。

此外，还需要设置TF_CONFIG环境变量，其包括各个节点的元数据，如IP地址、端口号、角色等，例如下面两个节点的情况，

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 第一台机器设置为0
# 第二台机器设置为1
index = 0
os.environ["TF_CONFIG"] = json.dumps({
    # 说明分布式集群的节点信息，包括IP地址与端口号
    "cluster": {
        "worker": ["192.168.8.115:20000", "192.168.8.116:20001"]
    },
    # 说明第index个机器的角色
    # 0 --> "192.168.8.115:20000"
    # 1 --> "192.168.8.116:20001"
    # index需要每个训练脚本单独设置
    "task": {"type": "worker", "index": index}
})

这些配置直接写到Python代码中即可，以便修改后在多个节点间同步。这些配置好后，使用MultiWorkerMirroredStrategy即可实现分布式训练，

1
2
3
4
5
strategy = tf.distribute.MultiWorkerMirroredStrategy()

with strategy.scope():
    # 实现模型的代码
    # 预加载权重

接着每个脚本单独运行，然后会根据配置的元数据信息建立各个节点互连状态，当所有节点都建立连接后，即可进行多机训练。参数更显策略，多机训练与单机多卡训练是类似的。

实现

由于Tensorflow本身已经对多卡训练做了抽象和实现，因此实现上并不困难。源码地址：https://github.com/allenwind/tensorflow-in-large-dataset 。

总结

本文总结了多卡训练的基本原理，包括参数同步算法Parameter Server与Ring All Reduce，然后讲述Tensorflow对卡训练的实践。

参考

[1] https://tensorflow.google.cn/guide/distributed_training

[2] 《TensorFlow 内核剖析》刘光聪

[3] https://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/

转载请包括本文地址：https://allenwind.github.io/blog/16113
更多文章请参考：https://allenwind.github.io/blog/archives/