如何将深度学习训练速度提升一百倍PAISoar-来了

阿里妹导读： 得力于数据规模增长、神经网络结构的演进和计算能力的增强，深度学习的图像处理、语音识别等领域取得了飞速发展。随着训练数据规模和模型复杂度的不断增大，如何充分利用分布式集群的计算资源加快训练速度，提升业务支持能力成为用户非常关注的问题。今天，我们就来分享阿里工程师的实践成果：将深度学习模型的大规模分布式训练框架 PAISoar 应用于绿网模型（多层 CNN 网络）后，绿网模型在 128 GPU 卡上取得 101 倍的计算加速比，效果显著。

1. 概述

近几年来深度学习发展迅速，图像处理、语音识别等领域都取得了飞速发展。例如在图片识别上，神经网络结构快速演进，分类的网络结构从 AlexNet、VGG、Inception V1 发展到了 Inception V4、Inception-ResNet、SENet。随着模型层次越来越深，参数越来越多，模型能力也越来越强，ImageNet 数据集 Top-5 的错误率越来越低，目前降到了 2.25%（人眼 5.1%）。

随着模型复杂度不断增长、训练样本的指数级增长，分布式进行高效并行的神经网络训练已经变得越发重要。在社区版 Tensorflow 中，分布式训练是基于 Parameter Server 模式进行多机训练的。在这种训练方式下训练任务通常会遇到以下挑战：

1. Variable placement 策略，常用的 replica_device_setter 的策略是 round-robin over all ps tasks，这种策略并没有考虑 Variable 大小，会导致参数分配不均衡，某些 ps 上分配的 Variable size 大就会成为通信瓶颈；
2. 多个 Worker 访问同一个 PS 节点时，受 PS 节点带宽限制和 TCP 的拥塞窗口控制，会导致通信效率大幅降低，并且规模越大，效率越差；
3. 分布式扩展后，模型需要精细调参才能收敛，需要用户有丰富的调参经验。

对此，我们基于 PAI Tensorflow 研发了针对于深度学习模型的高速分布式训练框架 PAISoar，从硬件到软件打造一套分布式训练场景 E2E 的解决方案：

1. 在硬件上，我们和 AIS 网络团队一起搭建了集团内部第一套基于 RoCE 的大规模 RDMA 集群，并针对于深度学习场景进行了相应的参数调优，支持低延迟、高吞吐的无损传输网络；
2. 在软件上，我们基于 Ring AllReduce 算法在 RDMA 网络上实现了高度优化的 PAISoar 分布式训练框架，通过软硬件一体的深度优化大大提高了多机的计算加速比；
3. 在 API 层面，我们提供了简化用户构建分布式 TF 模型训练的 ReplicatedVarsOptimizer，极大地方便了将单机模型扩展成分布式模型，降低了用户构建分布式 TF 训练代码的难度，同时提供支持 warm up 的 learning rate 动态调节方法，帮助模型训练更容易的收敛。

PAISoar 在 Tensorflow 官方 benchmarks 模型上取得了非常不错的加速效果。同时我们还和安全部团队合作，将该研发成果成功的在安全部图像模型业务上落地。安全部的绿网模型训练样本 280 万张图片，单机两卡训练需要 12 天才能收敛，因此有着非常强的分布式训练需求，希望能达到一天内完成训练，提升模型迭代能力。借助于 PAISoar，绿网模型分布式训练取得非常明显的加速效果，模型收敛时间从原先的 12 天降低到目前的一天以内，大大加速了业务的快速迭代。

我们用 images/sec (每秒处理的图片数或样本数) 来统计分布式模型的计算能力。

假设单 GPU 卡下计算能力为 1，绿网模型在 PAISoar 中计算加速比如下图所示：

在 PAISoar 上，绿网模型在 128 GPU 卡上取得了 101 倍的计算加速比，效果非常明显。

2. PAISoar：基于 PAI Tensorflow 的分布式训练框架

2.1 PAISoar 简介

PAISoar 是基于 PAI Tensorflow 实现的分布式训练框架。通过 PAISoar，我们提供了一个从硬件到软件、易用的分布式训练的性能优异框架。

2.1.1 网络层

我们和 AIS 网络团队、RDMA 项目团队一起搭建了集团内部第一套基于 RoCE 的大规模 RDMA 集群，并针对于深度学习场景进行了相应的参数调优，包括：

• 机器上搭载 Mellanox 25G 网卡支持基于 RoCE v2 的 RDMA，打造低延迟高吞吐通信网络；
• ASW（接入层交换机，32 口）和 PSW（聚合层交换机）间采用 8*100Gb 高速传输光纤，支持交换机 1:1 收敛比，搭建无损传输网络；
• 构建 TCP 和 RDMA 多级混合流控策略，解决各种混跑场景的流量干扰；

2.1.2 软件层

• 在 PAI Tensorflow 中接入 RDMA 驱动，使用 verbs 库进行 RDMA 通信，并和思科交换机进行适配调参；
• 对 RDMA 通信的关键路径进行梳理，加速内存拷贝，异步化数据发送，优化通信状态机，提高 RDMA 通信的效率和稳定性；
• 自研了深度优化的 Ring AllReduce 同步算法，通过针对 RDMA 网络的深度适配、多路通信融合等关键点优化，大大提升了多机的分布式训练性能。

2.1.3 API 层

• 我们同时提供了简化用户构建分布式 TF 模型训练的 ReplicatedVarsOptimizer，极大的方便了将单机模型扩展成分布式模型，降低了用户构建分布式 TF 训练代码的难度和学习成本；
• 为方便用户在分布式训练中调参，我们提供 lr(learning rate) 的动态调节算法 smooth_exponential_decay，算法中，lr 经过一定轮数的 warm up，从一个较小的 lr 平滑增加到目标 lr，再采用指数衰退，每隔固定 epoch 衰退一次，动态调节的 lr 可以加速用户调参过程。

2.2 性能指标

基于 PAISoar，深度神经网络模型的分布式训练性能提升明显。

我们使用 Tensorflow 官方 Benchmarks 进行性能测试。Tensorflow Benchmarks 是一套 CNN 模型的基准测试集合，包括 Inception v3、ResNet-50、ResNet-152、VGG16、Alexnet 等经典 CNN 模型。下面是我们分别在 1、2、4、8、16、32、64GPU 卡规模下测试了 Inception v3、ResNet-50、ResNet-152、VGG16 这 4 个模型的性能：

图中参数说明：

• 上面 4 个图分别是 Inception v3、ResNet-50、ResNet-152、VGG16 这四个模型的性能测试数据；
• 横坐标是 GPU 卡数，纵坐标是模型训练总的 Images/sec；
• ps、allreduce 分别表示模型参数同步的模式是 worker+ps 模式或 Ring AllReduce 模式；
• grpc、rdma 分别表示通信底层使用的是 grpc 协议还是 rdma 协议。

Tensorflow 默认的分布式方法是 worker+ps 模式，variable 定义在 ps 上，worker 会训练并更新 ps 上的参数。我们将底层通信协议替换成 RDMA 后，各模型性能都有提升，在 64GPU 卡下，Inception v3、ResNet-50、ResNet-152、VGG16 四个模型性能分别提升：24.94%、44.83%、38.80%、23.38%。

相较于 worker+ps 模式，PAISoar 中提供的 Ring AllReduce 通信模式对网络带宽利用率更高，延迟更小。与 worker+ps (使用 grpc 通信) 相比，各模型性能提升非常明显，在 64 GPU 卡下，Inception v3、ResNet-50、ResNet-152、VGG16 四个模型性能分别提升：84.77%、125.43%、56.40%、40.04%。

同时，我们也和开源的 horovod 进行对比，在 64 GPU 卡下，Inception v3、ResNet-50、ResNet-152、VGG16 四个模型性能分别提升：-6.4%、2%、21%、36%，整体性能要好于 horovod。

接下来将会介绍 PAISoar 中的关键技术：RDMA 和 Ring AllReduce 的技术细节。

2.3 RDMA 技术

RDMA(Remote Direct Memory Access) 是一种硬件 IO 技术，通过将传输协议固化在网卡硬件中，网卡就可以实现内核旁路（kernel bypass）和零拷贝操作，从而大幅降低网络 IO 的延迟时间，大幅提高网络吞吐，同时不影响 CPU 的负载。RDMA 在以太网上有 2 个标准，2002 年，ITEF 制定了 iWARP（Internet Wide Area RDMA Protocol），2010 年，IBTA 制定了 RoCE（RDMA over Converged Ethernet）规范。阿里数据中心网络里采用 RoCE(v2) 方案。在时延测试中，RoCE 可以达到单向为 2 – 3us 左右，而与之对应的 TCP 时延为 10 – 15us 左右。

为了实现 RoCE 的性能优势，以太网需要提供一个无损网络做为基础，否则对性能的负面影响非常大。以太网的无损机制主要通过端到端的 QoS 来实现，即所有的 QoS 机制需要在服务器和交换机上做统一的设置，包括 DSCP、队列、DCQCN、ECN、PFC 等。通过使用这些技术，能够保证在 burst，incast 等 IO 场景下 RDMA 也有稳定的表现。下图是在一个 32 个节点的集群中做长时间压测时采集到的发送带宽，可以看到各个节点的带宽非常稳定。

RoCE 技术虽然有很多技术优势，但是它对无损网络的要求也导致网络配置非常复杂，除了新增很多网络设备的管控配置，同时也把之前相对独立的网卡设备纳入了网络运营的范围，大大增加了运营的内容和难度。与此相适应，在主机上需要安装专门的 RDMA 软件平台包，包括网卡驱动、用户库和管控软件等部件。这些软件可以做到一键部署、自动配置和定期巡检，并且与相应的监控系统都有数据对接。通过安装这些软件包，就可以做到及时发现物理网络的配置错误，硬件故障，让应用远离复杂的网络管理，顺利享受 RDMA 带来的技术红利。

2.3 Ring AllReduce 技术

对于许多在大型数据集上训练的现代深度学习模型来说，基于数据并行（Data Parallelism）的同步分布式训练是最合适的训练方法。数据并行的同步分布式训练在每个 worker 上都有一份完整的模型，每个 worker 读取训练集的不同部分计算出不同的梯度，然后同步各个 worker 上的梯度并求平均值，再用平均梯度来更新每个 worker 上的模型。在社区版 Tensorflow 中，梯度平均通过将 Variable 分配到公共的 ps（parameter server）节点上实现，通信代价受 worker 数量和 Variable 分配策略影响很大。在 PAISoar 中我们使用 Ring AllReduce 技术实现梯度平均，通信代价的上限与 worker 数量无关，且不再需要额外的 ps 节点。

Ring Allreduce 算法的原理与核心功能如下：

Ring AllReduce 算法将 device 放置在一个逻辑环路（logical ring）中。每个 device 从上行的 device 接收数据，并向下行的 deivce 发送数据，因此可以充分利用每个 device 的上下行带宽。

使用 Ring Allreduce 算法进行某个稠密梯度的平均值的基本过程如下：

将每个设备上的梯度 tensor 切分成长度大致相等的 num_devices 个分片；

ScatterReduce 阶段：通过 num_devices – 1 轮通信和相加，在每个 device 上都计算出一个 tensor 分片的和；

AllGather 阶段：通过 num_devices – 1 轮通信和覆盖，将上个阶段计算出的每个 tensor 分片的和广播到其他 device；

在每个设备上合并分片，得到梯度和，然后除以 num_devices，得到平均梯度；

以 4 个 device 上的梯度求和过程为例：

ScatterReduce 阶段：

经过 num_devices – 1 轮后，每个 device 上都有一个 tensor 分片进得到了这个分片各个 device 上的和。

AllGather 阶段：

经过 num_devices – 1 轮后，每个 device 上都每个 tensor 分片都得到了这个分片各个 device 上的和；由上例可以看出，一次 Ring Allreduce 中单个节点总的通信数据量是

每个节点上通信数据量的上限不会随分布式规模变大而变大。

在 PaiSoar 中，我们基于 Tensorflow 原生的通信原语 Rendezvous 实现 Ring AllReduce 算法，并提供简化用户构建分布式 TF 模型训练的

ReplicatedVarsOptimizer 接口（用户文档），核心功能如下：

• 简单的用户接口：只需开启 PAISoar 之后将原有的 Optimizer 包装为 ReplicatedVarsOptimizer，即可将单机模型改成分布式模型，不需要修改太多的代码；
• 支持多种通信协议：支持 gRPC、RDMA 等多种通信协议；
• 优化的通信策略：根据计算图的信息，平衡计算通信重叠度和通信效率，合并计算时间相近的梯度一起传输。

3. 安全部绿网模型

3.1 模型介绍

绿网模型开始于 2013 年，在安全部成立初期就已经存在，用于电商场景中的色情内容的识别。绿网模型最初基于 BOW（Bag of Word）类型的图像分类模型，这是在深度学习（Deep Learning）出现之前学术界以及工业界普遍使用的建模方式，在图像检索以及分类上都得到了广范的使用。

绿网模型分为互联网场景与电商场景两个模型。本文中主要集中介绍互联网场景的模型。互联网场景是绝大多数网站 / App 的需求，特点是只检出有明确证据判断为色情内容的图片 / 视频，标准较宽松，可以容忍一些大尺度性感照、不雅姿势、猥琐动作等。

目前绿网模型互联网场景主要是在阿里云的阿里绿网（内容安全），绿网的付费客户数量随着互联网内容管控的需求加大呈现了指数级增长。目前已经为集团内和集团外的大量客户提供了安全稳定的内容保障。

目前绿网模型互联网场景每天调用量已经达到亿级别，平均的 RT 为 80ms 左右。

由于图像计算量复杂以及图像数据量大，GPU 单机训练已经无法满足目前模型迭代的速度，单机 2 卡训练一个模型需要长达 12 天之久，因此分布式训练势在必行：

3.2 分布式调参

如上文，使用 PAISoar 进行分布式训练后，绿网模型取得了非常明显的计算加速效果，在 128 GPU 卡上，计算加速比能达到 101 倍。

分布式训练另一个重要的工作就是调参，分布式规模扩大相当于增加 batch size，如果沿用以前的学习策略和参数（learning Rate 等）会导致模型收敛慢或不收敛，达不到单机训练的精度。我们需要进行调参来让模型训练收敛。

我们调参所用的方法：

1. 训练数据分片，在分布式训练时，需要对数据进行分片，确保每个 worker 读到的数据不一样，最好是每几个 epoch 后整体数据 shuffle 一次，避免模型对输入数据顺序的依赖。

2. 使用 PaiSoar 中的 learning rate 动态调节方法 smooth_exponential_decay，在训练初期采用 lr(learning rate) 平滑上升的策略，训练刚开始时，使用一个较小的 lr，尽量避免一开始由于 lr 太大导致收敛不稳定。后面再平滑增大 lr，尽量让 weights 更新到离初始的 init weight 比较远的地方。lr 增长到目标值后再采用指数衰退的方式，每个几个 epoch 衰退一次，逐步减少 lr，避免训练后期过大的 lr 导致训练波动，不收敛，learning rate 变化趋势如下图所示：

3. 分布式 N 个 worker 训练，相当于单机 batch size 增大 N 倍，此时需要调大 learning rate，一般调大 sqrt(N)~N 倍，可以结合 Tensorboard 来观察 loss 和 lr 的变化趋势来决定调大还是调小 lr.

经过调参在不同 GPU 卡下绿网模型都调试收敛，不同 GPU 卡下训练轮数如下表所示：

在 16 机（32 GPU 卡）下绿网模型在 20 个小时收敛，比单机（2 GPU 卡）快了 14.4 倍，完成一天内训练收敛的目标。同时对所有分布式训练的模型使用了安全部的的火眼平台进行 ROC 测试，测试集有 4 万多条数据。经测试，分布式模型效果和单机 2 GPU 卡效果一致：

3.3 模型迭代

我们通过分布式训练加快了绿网模型的收敛，这时就可以尝试采用更复杂的神经网络结构来提高 AUC，之前绿网模型内使用的是 Inception v3 结构，我们升级到 Inception v4 进行测试，发现 ROC 曲线提升明显：

PAISoar 的分布式加速能力让我们可以尝试使用更复杂的模型结构来提高模型效果。

4. 总结和展望

通过与 AIS 网络团队、RDMA 项目团队的协同合作，PAI Tensorflow 的分布式训练框架 PAISoar 成功上线，并和安全部在绿网模型上进行合作落地，取得非常好的效果，128 GPU 卡上计算加速比达到 101 倍。期望后续能服务更多的用户，并为用户提供简单、高效、可靠的分布式服务。

接下来我们会继续优化分布式训练框架 PAISoar：

• 尝试不同的网络拓扑结构和通信架构，进一步充分利用网络带宽；
• 参数稀疏化通信，尝试发现高频变化特征从而达到进一步提高通信效率并且将模型精度衰减在一定范围之内；
• 提供便捷的使用方法，方便用户快速接入，降低学习成本。

本文作者：王林

阅读原文

本文来自云栖社区合作伙伴“阿里技术”，如需转载请联系原作者。