近日,阿里云机器学习PAI关于深度学习模型高效的分布式训练框架的论文《 Whale: Efficient Giant Model Training over Heterogeneous GPUs 》被计算机系统领域国际顶级学术会议USENIX ATC'22接收。
Whale是阿里云机器学习PAI平台自研的分布式训练框架,开源后的名称是EPL(Easy Parallel Library),Whale通过对不同并行化策略进行统一抽象、封装,在一套分布式训练框架中支持多种并行策略,并进行显存、计算、通信等全方位的优化,来提供易用、高效的分布式训练框架。Whale提供简洁易用的接口,用户只需添加几行annotation即可组合各种混合并行策略。同时Whale提供了基于硬件感知的自动化分布式并行策略,感知不同硬件的算力、存储等资源,根据这些资源来合理的切分模型,均衡不同硬件上的计算量,最大化计算效率。借助Whale的硬件感知负载均衡算法,Bert-Large、Resnet50和GNMT模型 在异构GPU训练上提速1.2至1.4倍。同时,使用Whale框架, 万亿M6模型使用480 张 V100在3天内完成预训练。相比此前业界训练同等规模的模型,节省算力资源超80%,且训练效率提升近11倍。进一步使用512 GPU在10天内即训练出具有可用水平的10万亿模型。
背景和挑战
随着近些年深度学习的火爆,模型的参数规模也增长迅速,OpenAI数据显示:
- 2012年以前,模型计算耗时每2年增长一倍,和摩尔定律保持一致;
- 2012年后,模型计算耗时每3.4个月翻一倍,远超硬件发展速度;
特别最近两年模型参数规模飞速增长,谷歌、英伟达、阿里、智源研究院都发布了万亿参数模型,有大厂也发布了百亿、千亿参数模型。随着模型参数规模增大,模型效果也在逐步提高,但同时也为训练框架带来更大的挑战。当前已经有一些分布式训练框架,例如:Horovod、Tensorflow Estimator、PyTorch DDP等支持数据并行,Gpipe、PipeDream、PipeMare等支持流水并行,Mesh Tensorflow、FlexFlow、OneFlow、MindSpore等支持算子拆分,但当我们要训练一个超大规模的模型时会面临一些挑战:
- 当前的分布式训练框架只支持少量的并行策略,缺乏一个统一的抽象来支持所有的并行策略及其混合策略。
- 实现复杂的并行策略需要大量的模型代码改造和对底层系统的理解,大大增加了用户的使用难度。
- 由于集群中异构GPU计算能力和显存的差异,静态的并行化策略无法充分利用异构资源实现高效训练。