日前,腾讯工程师以2分31秒的成绩,成功打破128卡训练ImageNet的世界纪录。与之前保持的记录相比,快了整整7秒。“我们的实力还没有完全发挥,如果改用RoCE,这个成绩可以进一步提升到2分2秒”参与其中的腾讯工程师这样说。


Tips:ImageNet在图像处理领域非常有名,它是一个已完成标注的海量数据集,也是公认的图像处理算法试金石:谁的算法用了更少的训练资源、在更短的时间内训练完ImageNet,还能得到更高的准确率,谁的算法就能排名靠前。


640?wx_fmt=png&tp=webp&wxfrom=5&wx_lazy=1&wx_co=1

具体点说,腾讯工程师们在25Gbps的VPC网络环境下、使用了128块型号为V100的GPU(业内称为128卡),并借助最新研制的Light大规模分布式多机多卡训练框架,仅耗时2分31秒训练28个epoch就能识别出ImageNet内的128万张图片内容,且TOP5精度达到93%。


640?wx_fmt=png&tp=webp&wxfrom=5&wx_lazy=1&wx_co=1

那么,腾讯工程师们为什么要去攻破训练ImageNet的世界纪录呢?


一个明显的感知:AI模型越来越复杂了

随着AI的广泛应用,AI模型越来越复杂了:


 数据量大:号称史上最大的AI模型GPT-3光训练就使用了高达45TB的数据量,这让它在多轮训练时的数据读取成为非常耗时的环节。

 计算模型复杂:深度网络的层次究竟要多深?这要看AI想要表达的特征有多丰富。当前应用广泛的CNN特征提取网络ResNet-1000,其深度已达1202层。


 参数量大:深度神经网络由于层次很多,参数量也往往大得惊人,GPT-3的参数量高达1750亿个,这势必会给调参工作带来难度。


 超参数范围广泛:随着模型复杂度的提升,模型中可供调节的超参数数量及数值范围也在增多。当多数超参数的取值为连续域的情况下,即使只有少量的超参数仍然可能造成组合爆炸。

 训练时间长:模型越复杂对算力的需求势必越高。从2012年到2018年,业界对算力的需求增长超过2000倍 ,算力紧缺就会导致训练时间拉长。


为了缩短训练时间,各大厂商不惜堆设备增加算力,随之而来的是高额的训练成本。根据估测,训练一个GPT-3大概需要1300万美元,以至于其研究者在论文中说“我们发现了一个bug,但没钱再去重新训练,所以就先算了吧。”


AI模型训练难道只是“氪金游戏”? 


腾讯想做的,

是突破AI模型训练框架性能极限

正是在这样的背景下,腾讯云联合腾讯机智团队、腾讯云智能钛团队、腾讯优图实验室、腾讯大数据团队和香港浸会大学计算机科学系褚晓文教授团队,以ImageNet为训练的标杆,研发出了Light大规模分布式多机多卡训练框架,这套全新的训练解决方案,在单机训练速度、多机多卡通讯优化、batch收敛等方面的多个细节上都进行了优化,可以让AI模型训练更加高效:


单机训练速度


在单机训练速度方面,腾讯工程师主要解决了1.分布式系统中访问远端存储的数据较慢 2.大量线程相互抢占资源导致CPU效率低下 3.JPEG小图片解码制约性能等问题,提升了整体训练的速度。经过优化后,单机训练每卡速度有了明显提升,以96*96*3图片大小为例,训练速度对比如下图所示:

640?wx_fmt=png&tp=webp&wxfrom=5&wx_lazy=1&wx_co=1


具体技术细节如下:


①分布式系统中访问远端存储的数据较慢 :AI训练数据一般以分布式形态进行存储,而分布式存储机器和训练机器并不在一个集群,这会导致访问远端存储的数据较慢。为此,腾讯工程师利用GPU母机的SSD盘/内存,在训练过程中为训练程序提供数据预取和缓存。


②大量线程相互抢占资源导致CPU效率低下:数据预处理时,每个进程分配的线程数是惊人的(比如,单机8卡在数据预处理时就会分配数百个线程)。大量的线程相互抢占资源会导致CPU运行效率低下等问题。于是腾讯工程师根据运行的实时信息和过往经验,自动设置最优数据预处理线程数,降低CPU切换负担,同时也让数据预处理与GPU计算并行。


③JPEG小图片解码制约性能:每个小图片虽然计算时间少,但单位时间内需要处理的数量多,也会导致CPU负载过大。经过分析后发现,其中制约性能的环节为JPEG图片解码。因此,腾讯工程师预先将数据集中的JPEG图片解码,并缓存于内存中,计算时直接加载解码数据从而加快计算。


多机多卡通讯优化


在多机扩展训练方面,以往在TCP环境下,跨机通信的数据需要从显存拷到主存,再通过CPU收发,每次计算时间短但通信时间长。为此,腾讯展开了多机多卡通讯优化,充分利用通信时的网络带宽,缩短跨机通信的时间。经过测算,优化后腾讯将训练速度上升至3100样本/秒,业界其它类似的算法速度是2608样本/秒。


具体技术细节如下:


①自适应梯度融合技术优化通信时间:将小数据块融合成大数据块,减少通信次数从而降低通信延迟、提升通信效率的方式可以减少通信时间。但在某些情况下,通信需等待所有编译优化所融合的计算完成。因此,腾讯提出了自适应梯度融合技术,根据编译优化的结果自适应的选择梯度融合的阈值,解决了这个问题。


②2D通信+多流提升带宽利用率:TCP网络下全局规约操作有很大的额外延时。针对这个问题,腾讯使用了2D通信和多流来提升网络带宽利用率。以单机有8块卡的机型为例:


2D通信在TCP网络下,8块卡可以同时做跨机通信,带宽空闲的时间少,而且每块卡做跨机通信的次数比全局规约时少得多。


多流又使得多个梯度的2D通信可以形成流水线,某个梯度在做单机内通信没有利用到网络带宽时,其它梯度的跨机通信可以填补网络带宽的间隙。


③梯度压缩通信技术减少通信量,突破带宽瓶颈:网络带宽得到充分利用之后,为进一步提高扩展性,腾讯引入了梯度压缩通信技术。减少通信量,突破带宽瓶颈。


640?wx_fmt=png&tp=webp&wxfrom=5&wx_lazy=1&wx_co=1

梯度压缩通信技术流程


batch收敛问题


为了在最大程度提升训练速度的同时,最小化对精度的影响,腾讯增加了batch size,并使用了大batch调参策略、梯度压缩精度补偿、AutoML调参等方法,结果证明:腾讯工程师仅需训练 ImageNet 28个epoch,TOP5精度就能达到93%,而业界其它算法需要训练90个epoch才能达到相近的结果。


具体细节如下:


①大batch调参策略:综合考虑性能和收敛性,腾讯工程师使用了变分辨率的多阶训练:训练前期利用多个epoch的低分辨率样本快速收敛,而在训练后期使用少量epoch的高分辨率调高精度。此外,腾讯工程师还在模型、优化器、实际超参选择、损失函数上进行了诸多优化。


②梯度压缩精度补偿:在整个训练过程当中,图片大小为96*96*3时通信的时间占比最大,因此腾讯只在此时采用梯度压缩通信,对训练速度的提升最好且对精度的影响也最小。


③AutoML调参:目前,利用业界开源框架研发自动机器学习算法难、集成自研训练平台工作量大、在大规模长任务训练场景优化效果差。为此,腾讯研发了天风 (TianFeng)自动机器学习框架。它高度抽象出了自动机器学习的通用过程,能最大程度地将人力从繁琐的手工调参工作中解放。在同一模型下,工程师仅需指定搜索的超参数及其范围,即可对超参数空间进行探索并对调参思路进行快速验证。


以上多重优化点相结合,腾讯最终打破了128卡训练ImageNet的世界记录。
不得不提的是,这个世界记录是基于腾讯公有云环境下创造的,目前,相关能力已集成到腾讯云智能钛机器学习平台,并广泛应用在腾讯内外部的业务。


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


640?wx_fmt=png&tp=webp&wxfrom=5&wx_lazy=1&wx_co=1

总结


下一次挑战ImageNet榜首的算法,会有多快?尚未可知。


可以确定的是,腾讯工程师们将进一步提升机器学习平台易用性、训练和推理性能,为广大从业者提供更有力的机器学习工具。