支持异构 GPU 集群的超大规模模型的高效的分布式训练框架 Whale

近日，阿里云机器学习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等支持算子拆分，但当我们要训练一个超大规模的模型时会面临一些挑战：

1. 当前的分布式训练框架只支持少量的并行策略，缺乏一个统一的抽象来支持所有的并行策略及其混合策略。
2. 实现复杂的并行策略需要大量的模型代码改造和对底层系统的理解，大大增加了用户的使用难度。
3. 由于集群中异构GPU计算能力和显存的差异，静态的并行化策略无法充分利用异构资源实现高效训练。