TFSEQ Part I: 分布式训练的方案和效率对比

TFSEQ Part I: 分布式训练的方案和效率对比

本文作者：追一科技算法工程师 Tony

1. 前言

TFSEQ 这个系列总结了笔者在使用 tensorflow 进行自然语言处理的一些经验和思考。计划写三篇文章：

1. 分布式训练的方案和效率对比
2. 序列模型的实现细节
3. Batch size大小，优化和泛化

此为第一篇。

在增大数据集的同时增大模型参数量(Scaling)是提高准确率的一个有效方案，见百度这篇文章。但这也意味着计算量和训练时间的快速上升。为了缩减训练时间，我们可以使用分布式/并行训练。

搭建一个大型的分布式系统是一个耗时耗力的大工程。在量级不那么大的训练场景下，通常多卡并行是一个简单、经济且高效的方案。单机多卡系统可以认为是分布式系统的一种简单特例：卡间通信走 PCIe (或者更加土豪的 Nvlink)，要比走以太网(Ethernet)快很多。分布式系统通过 Infinteband 连接方案和 Nvidia 的 GPUDirect RDMA 技术，可以实现不同 host 上的卡间直连，也可以达到甚至超过单机 PCIe 的通信速度。

本文假设读者已经有深度学习在自然语言处理应用上的基本知识，并用 Tensorflow 实现过一些序列模型。为了避免翻译带来的歧义，部分术语会直接使用英文表述(使用中文的话会在括号里加上英文术语)，所以中英混杂的文风难以避免。为了讨论方便，以下先做一些术语的规定。

min-batch SGD 是一种迭代式优化(iterative optimization)的算法，每一次迭代都包括以下三个步骤：

1. 读取 mini-batch，使用模型进行前馈计算(feedforward or forward)
2. 计算 loss，并利用 loss 的值进行反向传播(backpropagation or backprop)，得到各个参数的梯度(gradient)
3. 根据算出的梯度，利用选定的优化算法(tensorflow 中称为 Optimizer)，如标准 SGD 或者更加流行的 Adam 对参数进行更新。

每个主机(host)上都有多个设备(device，可以是 GPU 或 CPU)。每个 device 都有对应的内存(memory)。host 和 device 通常抽象成计算节点(node)，可以进行运算以及和其他节点通信。本文用 node 做一般性的讨论，但在实现上还是以单机多卡的方案为主，此时 node 即为 device。

在计算时 node 间通常需要相互通信(communication)。Message Passing Model 是常用并行计算的通信模型，其通信操作在 MPI (message passing interface) 中定义。MPI 中的通信方式分两种：点对点通信(Point-to-point communication)和集合通信(collective communication)。点对点通信中只有一个发送者(sender) 和接受者(receiver)。而集合通信中通常有多个发送者和接受者。在分布式训练中比较常用的是集合通信。以下介绍中会用到以下三种操作：

1. broadcast，将参数从一个 node 发到多个 node 上
2. reduce，将参数从多个 node 收集到一个 node 上，同时对收集到的参数进行归并(求和，求积)。
3. allreduce，每个 node 都从其他 node 上面收集参数，同时对收集到的参数进行归并。

这里有集合通信操作的详细图解。

分布式/并行训练的优化目标之一是减少通信对计算的阻塞。

2. Tensorflow 的简单 profiling

为了了解训练过程中 GPU/CPU 的使用情况，我们可以使用 Cuda 提供的 cupti 来对 GPU 运行情况进行 profiling。我们可以从 tensorflow 中的函数调用这个库并打印出 profile 的 log。以下代码可以用来记录一个迭代中 GPU 的占用情况，文档主要来源于这个issue。由于文档的稀缺使得部分内容只能靠猜。google 同时也在开发新的 profiler。

import tensorflow as tf
from tensorflow.python.client import timeline
def profile(fetch_keys, sess, step, output_dir):
    run_metadata = tf.RunMetadata()
    options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    sess.run(fetch_keys,
             options=options,
             run_metadata=run_metadata)
    trace = timeline.Timeline(step_stats=run_metadata.step_stats)
    trace_file = open(os.path.join(output_dir, 'timeline_%d.json' % step), 'w')
    trace_file.write(trace.generate_chrome_trace_format())

拿到保存的 timeline.json 后，在 Chrome 的地址栏中输入 chrome://tracing/，点击 load。便可以得到这个图：

用 chrome 的 tracing 功能读取单个 step 的计算情况。

其中最底下的 CPU:0, GPU:0,1,2,… 即为对应的 tensorflow 对 CPU/GPU 的占用情况，上面的是 CUDA 内部线程的运行记录。点击图中每个 block 可以看到详细的情况，包括 block 对应的 op(operation) 名字，block 的前后依赖，运行时间等。

每个 block 都是根据一个 tensorflow 后端 kernel(GPU/CPU实现)划分的。点击右上方的 view option 可以看到不同 block 的依赖线图，我们可以根据这些信息找到计算阻塞的原因，并尝试优化它们。常见的计算阻塞原因有：

1. 对 input pipleline 的等待，这说明 input pipeline 参数没有调好
2. 不必要的变量传输。这可以通过用 tf.device 调整变量放置来优化。
3. 对 cpu 计算的依赖。常见的需要 cpu 计算的 kernal 是各种 sampler。我们可以尝试把这一部分计算放到 input pipeline 里。

在图中可以看到在使用 tensorflow 提供的 op 去实现计算图的时候，有大量的时间花在 op 的发起和切换上了。这会极大地降低 GPU 的使用率。所以进一步优化的方案便是自己用 cuda 和 C++ 写 tensorflow kernel，或者使用现成的 kernel(如 tf.contrib.cudnn_rnn)

3. 数据并行和模型并行

深度学习模型的并行有两种方案：模型并行(model parallel)和数据并行(data parallel)。

假设我们有多个 node：

• 模型并行：不同 node 输入相同数据，运行模型的不同部分
• 数据并行：不同 node 输入不同数据，运行相同的完整的模型。

为了完成一次更新，不同 node 间需要交换 forward 和 backprop 的信息，所以通信数据量是选择这两种并行方案的一个考量因素。另一个考量因素是由数据依赖(data dependency)带来的计算的阻塞(blocking)。最后一个考量因素是内存限制。当模型参数以及计算产生的中间变量无法放入一个 node 的内存时，我们只能使用模型并行。

当参数量巨大的时候，数据通信量会成为模型运行的瓶颈。模型并行的数据通信量可以比数据并行更少。可以看从这里化用的例子。

假设有一个 $64\times 1000$ 的数据矩阵 $X$，一个 $1000\times 1000$ 的参数矩阵 $A$，一个 $1000\times 500$ 的参数矩阵 $B$。loss 假设为 $\Vert XAB\Vert$。另假设我们有两个 node，将数据传输到 node 上的耗时忽略不计(使用第一篇文章描述的input pipeline)。我们需要计算 forward 和 backprop 并更新 $A$ 和 $B$。

使用模型并行：我们可以把 A 拆成两个 $1000\times 500$ 的矩阵 $A=[A_1,A_2]$。B 拆成两个$500\times 500$ 的矩阵 $B^T=[B_1^T,B_2^T]$ ，则 $XAB=X{A}_1{B}_1+X{A}_2{B}_2$。forward 时每个 node 需要传输/接收 $64\times 500$ 个浮点数(将 X A i