TFSEQ Part I: 分布式训练的方案和效率对比
本文作者:追一科技算法工程师 Tony
1. 前言
TFSEQ 这个系列总结了笔者在使用 tensorflow 进行自然语言处理的一些经验和思考。计划写三篇文章:
- 分布式训练的方案和效率对比
- 序列模型的实现细节
- Batch size大小,优化和泛化
此为第一篇。
在增大数据集的同时增大模型参数量(Scaling)是提高准确率的一个有效方案,见百度这篇文章。但这也意味着计算量和训练时间的快速上升。为了缩减训练时间,我们可以使用分布式/并行训练。
搭建一个大型的分布式系统是一个耗时耗力的大工程。在量级不那么大的训练场景下,通常多卡并行是一个简单、经济且高效的方案。单机多卡系统可以认为是分布式系统的一种简单特例:卡间通信走 PCIe (或者更加土豪的 Nvlink),要比走以太网(Ethernet)快很多。分布式系统通过 Infinteband 连接方案和 Nvidia 的 GPUDirect RDMA 技术,可以实现不同 host 上的卡间直连,也可以达到甚至超过单机 PCIe 的通信速度。
本文假设读者已经有深度学习在自然语言处理应用上的基本知识,并用 Tensorflow 实现过一些序列模型。为了避免翻译带来的歧义,部分术语会直接使用英文表述(使用中文的话会在括号里加上英文术语),所以中英混杂的文风难以避免。为了讨论方便,以下先做一些术语的规定。
min-batch SGD 是一种迭代式优化(iterative optimization)的算法,每一次迭代都包括以下三个步骤:
- 读取 mini-batch,使用模型进行前馈计算(feedforward or forward)
- 计算 loss,并利用 loss 的值进行反向传播(backpropagation or backprop),得到各个参数的梯度(gradient)
- 根据算出的梯度,利用选定的优化算法(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)。而集合通信中通常有多个发送者和接受者。在分布式训练中比较常用的是集合通信。以下介绍中会用到以下三种操作:
- broadcast,将参数从一个 node 发到多个 node 上
- reduce,将参数从多个 node 收集到一个 node 上,同时对收集到的参数进行归并(求和,求积)。
- 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 的依赖线图,我们可以根据这些信息找到计算阻塞的原因,并尝试优化它们。常见的计算阻塞原因有:
- 对 input pipleline 的等待,这说明 input pipeline 参数没有调好
- 不必要的变量传输。这可以通过用
tf.device
调整变量放置来优化。 - 对 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 × 1000 64\times 1000 64×1000 的数据矩阵 X X X,一个 1000 × 1000 1000\times 1000 1000×1000 的参数矩阵 A A A,一个 1000 × 500 1000\times 500 1000×500 的参数矩阵 B B B。loss 假设为 ∥ X A B ∥ \|XAB\| ∥XAB∥。另假设我们有两个 node,将数据传输到 node 上的耗时忽略不计(使用第一篇文章描述的input pipeline)。我们需要计算 forward 和 backprop 并更新 A A A 和 B B B。
使用模型并行:我们可以把 A 拆成两个 1000 × 500 1000\times 500 1000×500 的矩阵 A = [ A 1 , A 2 ] A = [A_1, A_2] A=[A1,A2]。B 拆成两个 500 × 500 500\times 500 500×500 的矩阵 B T = [ B 1 T , B 2 T ] B^T = [B^T_1, B^T_2] BT=[B1T,B2T] ,则 X A B = X A 1 B 1 + X A 2 B 2 XAB = XA_1B_1 + XA_2B_2 XAB=XA1B1+XA2B2。forward 时每个 node 需要传输/接收 64 × 500 64\times 500 64×500 个浮点数(将 X A i