刷新世界纪录：2 分 31 秒完成 ImageNet 训练

　　基于腾讯公有云 25Gbps 的 VPC 网络环境，使用 128 块 V100，借助 Light 大规模分布式多机多卡训练框架，在 2 分 31 秒内训练 ImageNet 28 个 epoch，TOP5 精度达到 93%，创造 128 卡训练 imagenet 业界新记录。

　　若跨机网络改为 RoCE，则训练时间进一步减少到 2 分 2 秒。之前业界同样使用 128 块 V100 的最好结果是友商的 2 分 38 秒。

　　28 个 epoch 分别是：

　　模型越来越大，算力需求暴涨

　　在 AlexNet[3]网络模型出现后的过去几年中，ResNet[4]、Transformer[5]、BERT[6]等优秀模型相继出现，深度学习有了长足的发展和进步，尤其是在图像、语音、机器翻译、自然语言处理等领域带来了跨越式提升。

　　在 AlphaGo 使用深度学习方法战胜世界围棋冠军李世石之后，近期拥有 1750 亿海量参数，堪称“通用语言智能的希望”——GPT-3[7]的出世，再一次点燃了大家对人工智能未来的期望。但与之相伴的也有很多问题：

　　数据量大。随着互联网上可使用数据资源愈加丰富，模型训练使用的数据集往往越来越大。GPT-3 模型训练则使用了高达 45TB 的数据量，使得多轮训练时数据读取成为非常耗时的部分。

　　计算模型复杂。深度网络的一个特点就是结构越深、越复杂，所表达的特征就越丰富，在此思想下，最新的网络结构越来越复杂，从 AlexNet 的 8 层，VGG-19 的 19 层，ResNet-50 的 50 层，到 Inception-ResNet-V2 的 467 层和 ResNet-1000 的 1202 层等。

　　参数量大。深度神经网络由于层次很多，参数量往往很大。ResNet-50 有 2500 万参数量，AlexNet 有 6200 万的参数量，而 VGG-16 参数量则达到 1.38 亿，BERT-Large 参数量高达 3 亿，而 GPT-3 语言模型的参数量甚至超过 1750 亿。

　　超参数范围广泛。随着模型复杂度的提升，模型中可供调节的超参数数量及数值范围也在增多。例如，在 CIFAR-10 数据集上训练的 ResNet 模型有 16 个可调的超参数，当多数超参数的取值为连续域的情况下，如此少量的超参数仍然可能造成组合爆炸。

　　训练时间长。随着深度学习模型越来越复杂，对算力的需求也越来越高(2012 年到 2018 年，对算力的需求增长 2000 倍+ )。ResNet-50 训练 ImageNet 128 万张图片 90 个 epoch 可以达到 76.5% 的 Top-1 精度，用一块 V100 的 GPU 需要大概 1 天多的时间，而各大厂都分别施展了各自的大规模分布式训练大法来缩短该标杆的训练时间。腾讯机智团队在 2018 年 7 月在 ImageNet 数据集上，用 4 分钟训练好 AlexNet，6.6 分钟训练好 ResNet50，创造了当时 AI 训练世界新纪录[1]。

　　如何做高性能 AI 训练和计算，关乎到 AI 生产研发效率，对于 AI 产品的迭代效率和成功至关重要。高效训练的一个 baseline 是在更短时间内做 ImageNet 训练。腾讯云希望 AI 技术能成为普惠的技术帮助人们生活更便捷，通过向客户提供高效的 AI 生产工具，AI 高效训练可以触手可及。

　　但是要做好高效的 AI 生产工具，需要做大量的工作，也充满挑战。接下来将从具体挑战，联合团队的解决方案，及最终效果几个方面来陈述，并最终通过 2 分 31 秒训练 ImageNet 这个案例为典型 case 来呈现成果。

　　多机多卡难扩展，调参耗时难收敛

　　数据供给制约计算

　　深度学习训练数据输入管道包括以下三个部分：从硬盘上读取数据-解析与数据预处理-拷贝到 GPU，数据输入管道与计算部分并行执行，为保证 GPU 能专心来进行模型训练，数据输入部分使用 CPU 来进行。但随着 GPU 性能的提高，计算时间越来越短，CPU，内存，网络带宽开始成为训练时的瓶颈，数据供给也成为训练效率的关键一环。

　　TCP 网络下的多机多卡扩展性差

　　随着 GPU 硬件的不断升级，GPU 的计算速度一直在提升，加上 XLA 和混合精度等技术也大大提升了单卡的计算速度。NVLink2 使单机内部的 AllReduce 速度有了大幅提升，超过了 120GB/s。但腾讯云的 VPC 网络只有 25Gbps，与单机内的通信带宽有着明显的差异。再加上 TCP 环境下，跨机通信的数据需要从显存拷到主存，再通过 CPU 去收发数据，除了实际网络收发数据的延迟，还有很大的额外延迟，所以通信时间更长了。计算时间短加上通信时间长，使得多机多卡的扩展性受到了很大的挑战。

　　大 batch 收敛难

　　为了充分利用大规模集群算力以达到提升训练速度的目的，人们不断的提升训练的 batch size，这是因为更大的 batch size 允许在扩展 GPU 数量的同时不降低每个 GPU 的计算负载。

　　然而，只增大 batch size 会对精度带来影响和损失。这是因为在大 batch size（相对于训练样本数）情况下，样本随机性降低，梯度下降方向趋于稳定，训练就由 SGD 向 GD 趋近，这导致模型更容易收敛于初始点附近的某个局部最优解，从而抵消了计算力增加带来的好处。如何既增大 batch size，又不降低精度，是另一个重大挑战。

　　超参数选择多

　　超参数的选择对模型最终的效果有极大的影响，但由于超参较多，而每一个超参分布范围较广，使得超参调优的耗时较长，特别是针对 ImageNet 这种大数据集的情况。如何快速寻找到较优解，是另一个重大挑战。

　　解决方案

　　联合团队研发了 Light 大规模分布式多机多卡训练框架来进行高效训练，并将能力平台化。

　　单机训练速度优化

　　1）分布式缓存与数据预取

　　由于训练数据一般保存在分布式存储，而分布式存储机器和训练机器并不在一个集群。为了加速访问远端存储的数据，团队利用 GPU 母机的 SSD 盘/内存，在训练过程中为训练程序提供数据预取和缓存。

　　2）自动调整最优数据预处理线程数

　　在预设的运行程序中，每个进程会分配大量线程用于数据预处理。仅仅在数据预处理上，单机 8 卡就会分配数百个线程。大量的线程的相互抢占导致了 cpu 运行效率低下，同时也会与 tensorflow 的 OP 调度线程争抢而出现 OP 无法正常并行的问题。

　　因此团队根据经验和采集运行时信息自动设置最优数据预处理线程数，降低 cpu 的切换负担，同时也让数据预处理可以和 gpu 计算并行。

　　3）本地预解码图片缓存

　　对于小图片，由于每个迭代步的计算时间少，单位时间需要处理的图片变多，cpu 负载仍然过大，导致性能较差。经过分析后发现，JPEG 图片解码为制约性能的主要部分。为了解决这个问题，团队预先将数据集中的 JPEG 图片解码，计算时直接加载解码数据，缓存在内存中，从而直接去除这部分计算压力。

　　LightCC 高效扩展多机训练

　　1）自适应梯度融合技术优化通信时间

　　针对小数据块通信问题，通过将小数据块融合成大数据块的方式，减少了通信次数，降低了通信延迟，提升了通信效率。但当应用了 XLA 技术后，由于 TensorFlow 中 XlaOp 的同步性，使得通信需等待编译优化所融合的所有计算完成，导致计算通信无法很好的 overlap，为解决这个问题，团队进一步提出了自适应梯度融合技术，根据编译优化的结果自适应的选择梯度融合的阈值，从而达到更好的 overlap。

　　2）层级通信+多流提升带宽利用率

　　TCP 网络下全局 AllReduce，除了网络收发包的时间，还有很大的额外延时，导致带宽利用率不高。

　　针对这个问题，团队使用了 2D 通信[2]和多流来提升网络收发的并发性，从而提升带宽的利用率。以单机有 8 块卡的机型为例，2D 通信在 TCP 网络下有不错的效果，主要是因为可以 8 块卡同时做跨机通信，从而竞争带宽资源，带宽空闲的时间更少，而且跨机通信时建立的环上的节点数只有总节点数的1/8，所以每块卡做跨机通信的次数比全局 AllReduce 时少得多。多流又使得多个梯度的 2D 通信可以形成流水线，既某个梯度在做单机内通信时没有利用到网络带宽，这时其它梯度的跨机通信可以填补网络带宽的间隙。

　　ResNet-50 v1.5 的参数量大概是 25M 个 FP32，如果通信时转成 FP16，则通信的梯度大小大概是 50MB，但因为 AllReduce 的通信量大概是梯度大小的 2 倍，所以通信量大概是 100MB。通信时间如下图所示。

　　3）层级 topk 压缩通信算法减少通信量，突破带宽瓶颈

　　因为网络带宽得到充分利用之后，通信还是瓶颈，为进一步提高扩展性，团队引入了梯度压缩通信技术[8]。梯度压缩通信的做法是每张卡对要通信的梯度做 Topk 压缩，只剩下绝对值最大的k个 value 和对应的k个 index，然后所有卡的 value 和 index 汇总到一起，再计算每个 index 的总和。通过这种方法，使通信量大大减少。

　　在此基础上，团队进一步提出了以下几点优化：

　　1) 首次提出层级 Topk 压缩通信算法。

　　还是以单机 8 卡的机型为例，首先第一步先用 ReduceScatter 在单机内通信，利用 NVLink2 的高速带宽，使每块卡快速得到单机内1/8 梯度的总和。

　　第二步的流程与普通的 Topk 压缩通信一样（如上图），不同的是每块卡只对原梯度1/8 大小的数据量做 Topk 压缩，从而减少 Topk 压缩的计算时间，而且每块卡在 AllGather 时的节点数减少到1/8，因为 AllGather 的通信量与节点数成正比，所以 8 卡的总通信量也减少到普通 Topk 压缩通信的1/8。而且单机的 8 张卡是并行跨机通信的，也提升了带宽利用率.。

　　第三步再通过单机内的 AllGather 使每张卡得到整个梯度的通信结果。

　　2) Topk 压缩性能优化。

　　如何高效进行 Topk 压缩也是提高整体通信效率至关重要的一环，尤其是在 GPU 上，若使用 Tensorflow 自身的 Topk 操作，在 8MB 数据下选取最大的1% 个 value 需要花费 42ms，而图片大小为 96*96*3 时的前、反向计算的总时间也才几十 ms，因此严重影响了整个梯度压缩通信的性能。

　　为了解决 Topk 压缩慢的问题，团队调研并借鉴了多种压缩算法，使用了性能较好的阈值二分搜索选择算法[8]，并在平衡了效率和精度后做了一些改进，二分搜索过程中记录最大的k'(k'<k且最接近k)个 value 并确保在最终结果中，从而使结果的k个 value 最接近真正的 topk value。<=”” font=””>

　　团队也为算法实现了高性能的 CUDA Kernel，首先为了最大化 GPU 利用率，合理配置 Kernel 使用的 block 和 thread 个数，并保证每个 SM 上拥有至少 16 个活动 warp，以有效的隐藏访存延迟。其次为了提高内存吞吐量，将多次传输合并为一次传输来减少 CPU 和 GPU 之间的内存传输；另外使用带宽更高的共享内存，并通过动态分配的方式，提高共享内存的使用效率；保证 block 内的各个线程访问连续对齐的内存块，提高缓存命中率。

　　最后为了提高指令吞吐量，使用效率更高的 CUDA 接口，并对循环进行展开，减少条件分支。通过以上对 CUDA Kernel 的深度优化后，在 8MB 数据下选取前1% 的 value 只需要 2ms，与 Tensorflow 本身的 Topk 实现相比，加速比为 21。

　　3) 针对 allgather 通信进行优化。

　　当所有卡要通信的数据大小一致时，使用 ncclAllGather 来通信，不一致时使用 ncclBroadCast 来模拟 allgatherv，这个优化也已经提交到 Horovod 社区。因为在压缩通信这个场景下，每块卡压缩后的数据大小是一致的，所以使用 ncclAllGather 通信，比 mpi 的 allgatherv 更快。

　　ResNet-50 v1.5 的参数量大概是 25M 个 FP32，以压缩率 0.01 为例，通信的 value 用 FP16 表示的话有 0.5MB，而 index 是 int32，所以有 1MB。使用层级梯度压缩通信时，跨机 AllGather 的节点数是 16，所以总的跨机通信量大概是 24MB。

　　4) LARS 计算并行化

　　为保证大 batch 下收敛，团队引入了 LARS 优化器，但 LARS 也引入了更多的计算量，每块卡会串行地计算每一层梯度的本地学习率。团队发现，由于该部分时间的消耗，计算时间无法与通信时间完全覆盖，导致整体训练时间的延长。由于 LARS 的计算是梯度通信之后，每张卡的针对模型计算出的本地学习率都是一致的，也就是每张卡都在做同样的计算，且计算出来的本地学习率对每层梯度为单值。

　　除了 LARS，LAMB 也是这种梯度规约后的操作，属于每块做重复计算的操作。针对这类操作，根据操作的计算代价和通信代价进行评估，如果是计算代价比较大的，就会自动将其转换成多卡并行计算再进行汇总，从而使总时间最少。

　　大 batch 收敛

　　1）大 batch 调参策略

　　训练样本上，团队在性能和收敛性上做了权衡，使用了变分辨率的多阶训练：训练前期利用多个 epoch 的低分辨率样本快速收敛，并在训练后期使用少量 epoch 的高分辨率调高精度。

　　模型上，使用了业界常用的 ResNet50 v1.5，同时在去除了 bias 及 BN 层参数正则化项的情况下，对 BN 层的 beta 做了一定的初始化调整，在较少迭代步数上取得了较好的收敛效果。

　　优化器上，选择了业界普遍采用的 MomentumSGD 结合 LARS 层级自适应优化器，为了快速收敛，团队对该优化器做了一些改动。对 LARS 优化器采用了随迭代步变化的置信比约束策略：在 Warm up 阶段，为了避免由于噪声过大导致发散的问题，将 lars 优化器的置信比裁剪约束到较小的范围，稳定的收敛；Warm up 阶段之后放松对 lars 优化器置信比的裁剪约束，加速收敛。

　　对于优化器中的动量值采用了全周期一次完整循环的方式，前期为了收敛的稳定性，动量值会随着训练的迭代不断增大，防止训练发散，后期为了训练的模型适配图片变大的分辨率，动量值随着训练的迭代不断减小，放大后期梯度对模型的贡献，加速收敛。

　　由于考虑到训练中使用的多分辨率样本，对学习率采用了分段线性变斜率策略，在 warmup 前后采用不同斜率的线性增长和线性减小，在不同分辨率阶段使用不同的学习率斜率，由 AutoML 进行端点斜率调节确定最终策略。

　　实际超参选择中，在确定学习率策略等超参数方案后，团队利用平台大算力结合自研天风 AutoML 框架，以最后一个 epoch 的精度作为评估指标，对于各阶段不同大小的样本进行分段分块有针对性的超参微调，并结合 AutoML 结果对超参方案进行迭代更新。

　　损失函数上，使用了 label smoothing 这一正则化手段对训练标签做了平滑处理，以防止过拟合帮助收敛。

　　2）梯度压缩精度

　　在整个训练过程当中，只在图片大小为 96*96*3 时采用梯度压缩通信，因为这种情况下通信的时间占比最大，梯度压缩通信对训练速度的提升最好，而且对精度的影响也最小。

　　精度补偿采用的是梯度本地累积的方法，也就是没有被选中通信的 index 的值，会在本地累加，参与下一轮通信。

　　3）AutoML 调参

　　针对业界开源框架研发 AutoML 算法难、集成自研训练平台工作量大和在大规模长任务训练场景优化效果差的问题，研发了天风 (TianFeng)， 是聚焦于大规模训练场景的、轻量的、易扩展的、易使用的 AutoML 框架，通过高度抽象 AutoML 的通用过程，来解决这些业界难题。

　　为了方便算法工程师做 AutoML 算法研发和新平台接入，天风内部分为三层，算法层、逻辑层和接入层。

　　（1）算法层，AutoML 算法都在这一层，通过抽象算法逻辑，方便算法工程师实现和接入 AutoML 算法。

　　（2）逻辑层，负责对接 Client，算法逻辑下发、异常处理和持久化。

　　（3）接入层，负责训练平台的对接，当前可以很简单的接入自研训练平台，结合 AutoML 能力和自研训练平台的调度和训练能力。

　　在 ImageNet 训练中，存在训练规模大，超参数范围广，人工调参效率低且效果不好的问题。因而引入天风用于超参数空间的探索，将人力较大程度地从繁琐的手工调参中解放出来。

　　依托智能钛平台的海量弹性算力，天风 AutoML 利用自研的并行贝叶斯优化算法进行超参数搜索，充分利用智能钛平台资源池中的空闲 GPU 算力资源并行进行多组弹性任务 AutoML 实验，做到了高优先级任务占用资源时实验自动终止，资源空闲时自动加载实验终止前的先验数据恢复。

　　整体上 AutoML 搜索超参数为大 batch size 调参指明了方向，在同一模型下，AutoML 仅需由工程师指定所需要搜索的超参数及其范围，即可对超参数空间进行探索从而对调参思路进行快速验证。

　　4. 4 高性能机器学习平台——智能钛

　　智能钛机器学习平台是为 AI 工程师打造的一站式机器学习服务平台，为用户提供从数据预处理、模型构建、模型训练、模型评估到模型服务的全流程开发支持。智能钛机器学习平台内置丰富的算法组件，支持多种算法框架，满足多种 AI 应用场景的需求。自动化建模（AutoML）的支持与拖拽式任务流设计让 AI 初学者也能轻松上手。 

　　取得的成效

　　通过以上 Light 大规模分布式多机多卡训练框架及平台等一系列完整的解决方案，ImageNet 的训练结果取得了新突破。同时，这些方案也广泛应用在腾讯公有云以及腾讯内部的业务，并形成了腾讯云帆（深度学习框架和加速）Oteam、算力 Oteam 等协同组织，在更广泛的范围内协同共建。以 ImageNet 的训练结果为例说明下。

　　单机优化

　　在单机的速度上，4 种不同的图片大小时都超过了友商。例如图片大小为 96*96*3 时，优化前后的速度对比如下图所示。

　　多机扩展

　　在 25G 的 TCP 网络上，图片大小为 96*96*3 时，由于计算时间相对较少，通信时间占比大，所以多机扩展性是最差的。

　　如果使用 Horovod 的全局 AllReduce，训练速度只有 885 样本/秒，而使用 LightCC 多机多卡通信库的话，不使用压缩通信时的训练速度是 2450 样本/秒（LARS 优化前是 2360 样本/秒），而使用压缩通信的训练速度是 3100 样本/秒。而友商 128 卡的速度是 2608 样本/秒，即不压缩通信时由于带宽限制导致略差于友商，而压缩通信之后超过了友商。

　　在图片大小为 128*128*3 时，由于计算时间变长，非压缩通信和压缩通信都只有最后一个梯度的通信没有被计算隐藏，所以压缩通信并没有速度上的优势，不过非压缩通信的速度也已经超过了友商。

　　收敛

　　在收敛精度方面，通过手动设置超参与 AutoML 调参相结合，在 28 个 epoch 将 top5 精度训练到 93%。

　　2 分 31 秒训练 ImageNet

　　在 2 分 36 秒内训练 ImageNet 28 个 epoch，TOP5 精度达到 93%；使用压缩通信后，在精度达到 93% 情况下，时间进一步减少到 2 分 31 秒，刷新了业界最好成绩，创造了业界新记录。若跨机网络改为 RoCE，则训练时间再进一步减少到 2 分 2 秒。

　　构建稳定、易用、好用、高效的平台和服务，将成为算法工程师的重要生产力工具，也会助力游戏 AI、计算机视觉 AI，广告推荐 AI、翻译 AI、语音 ASR AI 等典型 AI 业务从一个成功走向另一个更大的成功。

　　本次破纪录的 ImageNet 训练，由腾讯机智团队、腾讯云智能钛团队、腾讯优图实验室、腾讯大数据团队和香港浸会大学计算机科学系褚晓文教授团队协同优化完成。

　　参考链接：

　　1. Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes

　　https://arxiv.org/abs/1807.11205

　　2. ImageNet/ResNet-50 Training in 224 Seconds

　　https://nnabla.org/paper/imagenet_in_224sec.pdf

　　Krizhevsky， Alex， Ilya Sutskever, and Geoffrey E. Hinton. “Imagenet 3.classification with deep convolutional neural networks.” Advances in neural information processing systems. 2012.

　　4. He, Kaiming, et al. “Deep residual learning for image recognition.” Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.

　　5. Vaswani, Ashish, et al. “Attention is all you need.” Advances in neural information processing systems. 2017.

　　6. Devlin, Jacob, et al. “Bert: Pre-training of deep bidirectional transformers for language understanding.” arXiv preprint arXiv:1810.04805 (2018).

　　7. Brown, Tom B.， et al. “Language models are few-shot learners.” arXiv preprint arXiv:2005.14165 (2020).

　　8. RedSync : Reducing Synchronization Traffic for Distributed Deep Learning

　　https://arxiv.org/abs/1808.04357