快速打造分布式深度学习训练平台
如何在开源技术的帮助下,快速搭建起一个分布式的深度学习训练平台,加速深度神经网络的训练,提高公司的竞争力。
接下来我们从四个角度来谈一下这个深度学习训练平台:
- 为什么需要分布式训练平台
- 搭建这个平台需要使用什么关键技术和框架
- 当前的系统架构和实现
- 当前架构的不足和优化
为什么需要分布式训练平台
深度学习框架的多样性和趋同性:
当前的深度学习框架非常之多,耳熟能详的就有Caffe/Caffe2、MxNet、TensorFlow、Torch、CNTK、Theano等等。这些框架开发语言不一,接口设计各异,这对于AI公司,尤其是中小型创业团队的框架选择、技术积累和快速研发带来很多困难。不过,虽然深度学习框架众多,框架本身提供的功能、实现思路和核心技术有一定的趋同性。因此我们也看到有一些项目在这些框架之上做一层封装从而实现逻辑和接口上的统一,比如Kera和DIGIST。不过想短时间内“一统江山”难度会比较大。
因此,我们可以考虑从另一个层面来“统一”多个深度学习的框架,即通过工程的手段开发一个训练平台,将深度网络的训练、测试过程抽象道统一的界面中。同时,配合数据管理(网络存储和对象存储),环境管理(容器),CI和自动部署等,形成一个深度学习的“循环”,这能极高的提升迭代的速度。
深度学习的迭代性和周期性:
深度学习训练的一个特点是具有很强的迭代性,即网络结构确定后,可以周期性地通过增加训练数据而提高模型的泛化能力。这个迭代周期越短,模型更新的速度越快,就能获得更好的效果和更快的产品更新。这个特点使得训练过程非常适合自动化。
网络和算法的复杂度快速增长,带来计算资源的稀缺和浪费:
2012年AlexNet开启这一轮深度学习热潮的时候,网络只有8层,但就算只有8层也需要大量的计算。更何况2016年153层的ResNet,最近甚至出现了上千层的网络结构,这让计算资源变的非常稀缺。另一个方面,随着训练数据的快速增长,分布式训练变得很迫切和必需,这进一步增加了计算资源的稀缺性。
目前大多数的训练过程使用Nvidia GPU,而Nvidia GPU虚拟化的欠缺,使得GPU资源变的不容易管理。当数据中心有成百上千块GPU的时候,很容易出现一个矛盾的现象:GPU计算资源稀缺,同时又有部分GPU资源被闲置浪费,整体的利用率比较低。
其实这个现象在云计算和虚拟化出现之前就存在,只是那个时候浪费的是CPU、内存、IO和存储。现在同样的问题出现在GPU上而且显得更严重和迫切。当然,无论是国内还是国外的公有云、私有云厂商都在积极的解决这个问题。在没有成熟和稳定的方案出来之前,我们需要先自力更生,研发一个私有“GPU云”来应对这个问题。
国内外的互联网和IT公司都在非常积极的解决上诉问题,有更深层次和更广泛的平台级解决方案,比如Google TensorFlow、FaceBook FBLearner、阿里云PAI、腾讯DX-I等等。你可能已经发现,目前能提供这些平台服务的厂商也基本上是公有云的提供商。其实这是自然而然的事情,在Iaas,Saas之后,Aaas(AI as a service)已经是业界的共识和趋势。但是这些公有云平台的现状还在快速发展且不太完善,比较难以满足我们实际的需要。例如,这种通用平台难以快速响应最新的、实验性的深度学习技术;在细节实现和灵活性上难以满足中小客户的需求;另外一个很重要的原因是,目前公有云上的GPU或HPC实例的价格比较高,从成本角度,中小型公司还没有足够的动力把深度学习的业务全部迁移到公有云。
基于上面谈到的几点原因,我们决定研发一个深度学习的训练平台,这个平台要具备的功能包括:
- 管理多台训练服务器,尤其是带有GPU的高性能计算服务器,可把训练任务分到到分布式的计算节点上执行计算
- 集成多种训练框架,抽象训练过程,提供Web界面,上传和指定相关数据和参数,即可启动训练任务并监控和分析训练过程;
- 池化计算资源,尤其是GPU资源,做成“GPU云”。启动训练任务时,平台会自动把训练任务分配到合适的GPU上;
- 打通数据中心,可以直接把数据存储平台中的数据导入到训练节点上中;
- 隔离计算节点中的资源和环境,兼容不同型号的GPU、不同版本的CUDA/CuDNN和不同的深度学习框架
Docker
Docker本身无需多言,这里从深度学习的角度谈一下为什么要使用Docker。
环境隔离
绝大多数的深度学习依赖Nvidia GPU和相应的运算库CUDA、CuDNN等。Nvidia GPU的硬件型号和驱动繁烦复杂且存在一定的兼容性问题。CUDA等基础库众多并更新频繁。同时不同的深度学习框架对底层库的依赖也有所不同。如果直接在裸机上运行深度学习任务,尤其是新手,会花费大量的时间在环境的安装、配置和解决各种不兼容的问题,令人抓狂。我们把一系列的基础库、运算库和深度学习库打包成Docker镜像,充分利用Docker的分层机制在不同层面进行共享和组合。Docker隔离了系统环境和执行环境,即隔离同一台服务器上不同训练任务的环境。我们可以把同一个任务分发到带有不同型号GPU卡的服务器上,也可以在同一台服务器上同时运行不同CUDA版本、不同深度学习框架的多个任务。比较完美的解决了上述问题。
资源隔离
深度学习是典型的计算密集、IO密集和数据密集三者都满足的任务。这里的计算密集除了CPU密集外还增加了GPU密集。这对单机和集群的硬件资源管理和隔离提出了更高的要求。Mesos完成了对集群硬件资源的管理和有效调度,但是我们还需要进一步隔离硬件资源,在满足任务要求的同时,还要避免多任务对资源的恶性和无序竞争。Docker就是解决这个问题的最优选择。Docker在资源隔离方面的优异表现已经被业界充分验证且成功运用在大大小小的云里面。
但是,Docker或者说Linux cgroup对GPU资源的管理没有像CPU那么完善和成熟。不过我们仍然可以通过设备映射的方式,在GPU卡这个粒度上对GPU资源进行隔离,虽然粒度比较粗,但是考虑到深度学习任务一般会把GPU资源吃满,这样的粗粒度尚可接受。其实,不使用Docker,Mesos本身就提供资源隔离的特性。不过术业有专攻,Mesos跟Docker还有一定的距离。
代码共享
虽然目前有很多的深度学习框架和经典的神经网络,为了提高模型的性能和其他的工程因素,我们需要对不同的深度学习框架和计算库作不同层面的修改,作为公司内部的私有版本来使用。这些代码通常是满足特定网络和用途因此在代码管理、共享和使用上非常麻烦。通过Github + CI + Docker,我们把不同repo和branch的代码打包成完成不同任务的Docker镜像从而实现了更灵活和细粒度的共享。
Caffe/MxNet
Caffe和MxNet是在计算机视觉领域比较常用的两个深度学习框架,也是我们第一步支持的深度学习框架。对二者的支持没有本质的区别,主要是把框架的训练过程抽象和统一到训练平台层面,提供统一的接口启动训练任务。