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

本文涉及的产品
模型在线服务 PAI-EAS,A10/V100等 500元 1个月
模型训练 PAI-DLC,100CU*H 3个月
交互式建模 PAI-DSW,每月250计算时 3个月
简介: 高效大模型训练框架Whale(EPL)入选USENIX ATC

作者:张杰、贾贤艳


近日,阿里云机器学习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计算能力和显存的差异,静态的并行化策略无法充分利用异构资源实现高效训练。


破局

为了应对当前分布式训练的挑战,我们研发了高效、通用、硬件感知的分布式训练框架Whale。Whale抽象并定义了两个分布式原语(replicate和split) 。用户可以通过在模型上添加几行原语标记,即来表达和转换各种并行化策略及其组合,极大降低了分布式框架的使用门槛。Whale runtime将用户的标记信息融合到计算图中,自动完成模型的并行化优化。同时Whale提供了基于硬件感知的自动化分布式并行策略,优化在异构GPU集群上分布式训练性能。Whale的设计很好地平衡了模型用户的干预和系统优化机会,让每一个算法工程师都能轻松高效训练分布式大模型任务。

技术架构

Whale框架如下图所示,主要分为以下几个模块:

  • 接口层:用户的模型编程接口基于TensorFlow,同时Whale提供了易用的并行化策略表达接口,让用户可以组合使用各种混合并行策略;
  • 中间表达层:将用户模型和并行策略转成化内部表达,通过TaskGraph、VirtualDevices和策略抽象来表达各种并行策略;
  • 并行化引擎层:基于中间表达,Whale会对计算图做策略探索,结合硬件资源进行显存/计算/通信优化,并自动生成分布式计算图。
  • Runtime执行引擎:将分布式执行图转成TFGraph,再调用TF 的Runtime来执行;

并行化策略表达

Whale通过strategy annotation的方式来划分模型为多个TaskGraph,并在此基础上进行并行化。 Whale有两类strategy:replicate 和 split。通过这两种并行化接口,可以表达出各种不同的并行化策略,例如:

  • 数据并行: 下面这个例子是一个数据并行的例子,每个模型副本用一张卡来计算。如果用户申请了8张卡,就是一个并行度为8的数据并行任务。

  • 流水并行:在下面的例子里,模型被切分成2个 TaskGraph, "stage0"和"stage1",用户可以通过配置pipeline.num_micro_batch参数来设定pipeline的micro batch数量。 在这个例子里,"stage_0"和"stage_1"组成一个模型副本,共需要2张GPU卡。如果用户申请了8张卡,Whale会自动在pipeline外嵌套一层并行度为4的数据并行(4个pipeline副本并行执行)。

  • 算子拆分并行:在以下例子中,Whale会对split scope下的模型定义做拆分,并放置在不同的GPU卡上做并行计算。

  • 同时Whale支持对上述并行策略进行组合和嵌套,来组成各种混合并行策略,更多示例可以参考开源代码的文档和示例。

Parallel Planner

Paraller Planner是Whale runtime的重要一环,它的职责是生成一个高效的分布式执行plan。Parallel Planner的流程包含 (a) Paraller Planner的输入包含用户模型、用户标记(可选)、硬件计算资源和其他用户配置配置。 (b) 将物理计算资源映射成VirtualDevice, 用户无需担心如何将算子放置在分布式物理设备上。(c) 模型被分割成TaskGraph子图。因为Whale允许对不同的TaskGraph应用不同的分布式策略,所以在TaskGraph之间可能存在输入/输出shape不匹配。在这种情况下,Paraller Planner将自动在两个TaskGraphs之间插入相应的桥接层。

硬件感知的负载均衡算法

当模型训练资源包含异构硬件(比如混合了V100和T4),硬件感知的负载均衡算法可以提高在异构资源下的训练效率。Whale设计了两种平衡策略:Intra-TaskGraph和Inter-TaskGraph的平衡。

(1)对于Intra-TaskGraph的计算平衡,Whale会profile出模型的算力FLOP,并按照机器的计算能力按比例分配对应的计算负载,以此均衡每个step不同卡型上的模型计算时间。对于数据并行,我们可以通过调整不同副本上的batch大小来实现计算负载的均衡(保持全局batch不变)。对于算子拆分,我们可以通过不均匀的维度拆分来实现不同卡上子模块的计算负载均衡。

(2)当多个TaskGraph在多个异构GPU上执行流水并行,我们需要Inter-TaskGraph的计算平衡来提升整体的计算效率。由于流水并行的计算特点,前面的TaskGraph会比后面的TaskGraph多缓存一些前向计算的结果,因此对于Transformer类模型均衡layer切分的情况下,前面的TaskGraph会有更大的显存压力。因此在TaskGraph放置的时候,Whale会优先将前面的TaskGraph放置在显存容量更大的GPU卡上。与此同时,Whale会按照GPU卡的算力去切分模型,使得切分后的TaskGraph计算load和GPU算力成正比。

应用示例

借助Whale框架,我们4行代码实现M6模型数据并行+专家并行的混合并行训练。

如下图所示,MoElayer采用专家并行,其他layer采用数据并行:

并首次在480 V100 上,3天内完成万亿M6模型的预训练。相比此前业界训练同等规模的模型,此次仅使用480张V100 32G GPU就成功训练出万亿模型M6,节省算力资源超80%,且训练效率提升近11倍。进一步使用512 GPU在10天内即训练出具有可用水平的10万亿模型。

结语

Whale 通过统一的抽象支持各种并行化策略,并提出简洁的并行化原语来提高框架的易用性。同时,Whale提供了异构硬件感知的自动分布式计算图优化,实现高效的分布式模型训练。我们希望Whale能够成为一个大规模深度学习训练的基石,进一步推进模型算法创新和系统优化,使大模型训练技术能够广泛应用在实际生产环境中。Whale已经开源(https://github.com/alibaba/EasyParallelLibrary),欢迎大家来试用和共建。


[1] Junyang Lin, An Yang, Jinze Bai, Chang Zhou, Le Jiang, Xianyan Jia, Ang Wang, Jie Zhang, Yong Li, Wei Lin, et al. M6-10t: A sharing-delinking paradigm for effi- cient multi-trillion parameter pretraining. arXiv preprint arXiv:2110.03888, 2021.

[2] Junyang Lin, Rui Men, An Yang, Chang Zhou, Ming Ding, Yichang Zhang, Peng Wang, Ang Wang, Le Jiang, Xianyan Jia, Jie Zhang, Jianwei Zhang, Xu Zou, Zhikang Li, Xiaodong Deng, Jie Liu, Jinbao Xue, Huiling Zhou, Jianxin Ma, Jin Yu, Yong Li, Wei Lin, Jingren Zhou, Jie Tang, and Hongxia Yang. M6: A chinese multimodal pretrainer, 2021.

[3] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021.

[4] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia. Pipedream: gen- eralized pipeline parallelism for dnn training. In Pro- ceedings of the 27th ACM Symposium on Operating Systems Principles, pages 1–15, 2019.


了解更多精彩内容,欢迎关注我们的阿里灵杰公众号,获取大数据+AI最新资讯

qrcode_for_gh_0ea952036c3b_344.jpg

相关实践学习
在云上部署ChatGLM2-6B大模型(GPU版)
ChatGLM2-6B是由智谱AI及清华KEG实验室于2023年6月发布的中英双语对话开源大模型。通过本实验,可以学习如何配置AIGC开发环境,如何部署ChatGLM2-6B大模型。
相关文章
|
16天前
|
传感器 算法 安全
基于分布式模型预测控制DMPC的单向拓扑结构下异构车辆车队研究(Matlab代码实现)
基于分布式模型预测控制DMPC的单向拓扑结构下异构车辆车队研究(Matlab代码实现)
|
21天前
|
算法 调度
【孤岛划分】分布式能源接入弹性配电网模型研究【IEEE33节点】(Matlab代码实现)
【孤岛划分】分布式能源接入弹性配电网模型研究【IEEE33节点】(Matlab代码实现)
136 10
|
2月前
|
边缘计算 运维 算法
含分布式电源的配电网日前两阶段优化调度模型(Matlab代码实现)
含分布式电源的配电网日前两阶段优化调度模型(Matlab代码实现)
|
1月前
|
算法 安全 新能源
基于DistFlow的含分布式电源配电网优化模型【IEEE39节点】(Python代码实现)
基于DistFlow的含分布式电源配电网优化模型【IEEE39节点】(Python代码实现)
|
4月前
|
存储 机器学习/深度学习 自然语言处理
避坑指南:PAI-DLC分布式训练BERT模型的3大性能优化策略
本文基于电商搜索场景下的BERT-Large模型训练优化实践,针对数据供给、通信效率与计算资源利用率三大瓶颈,提出异步IO流水线、梯度压缩+拓扑感知、算子融合+混合精度等策略。实测在128卡V100集群上训练速度提升3.2倍,GPU利用率提升至89.3%,训练成本降低70%。适用于大规模分布式深度学习任务的性能调优。
177 2
|
7月前
|
并行计算 PyTorch 算法框架/工具
融合AMD与NVIDIA GPU集群的MLOps:异构计算环境中的分布式训练架构实践
本文探讨了如何通过技术手段混合使用AMD与NVIDIA GPU集群以支持PyTorch分布式训练。面对CUDA与ROCm框架互操作性不足的问题,文章提出利用UCC和UCX等统一通信框架实现高效数据传输,并在异构Kubernetes集群中部署任务。通过解决轻度与强度异构环境下的挑战,如计算能力不平衡、内存容量差异及通信性能优化,文章展示了如何无需重构代码即可充分利用异构硬件资源。尽管存在RDMA验证不足、通信性能次优等局限性,但该方案为最大化GPU资源利用率、降低供应商锁定提供了可行路径。源代码已公开,供读者参考实践。
499 3
融合AMD与NVIDIA GPU集群的MLOps:异构计算环境中的分布式训练架构实践
|
7月前
|
存储 人工智能 固态存储
轻量级AI革命:无需GPU就能运算的DeepSeek-R1-1.5B模型及其低配部署指南
随着AI技术发展,大语言模型成为产业智能化的关键工具。DeepSeek系列模型以其创新架构和高效性能备受关注,其中R1-1.5B作为参数量最小的版本,适合资源受限场景。其部署仅需4核CPU、8GB RAM及15GB SSD,适用于移动对话、智能助手等任务。相比参数更大的R1-35B与R1-67B+,R1-1.5B成本低、效率高,支持数学计算、代码生成等多领域应用,是个人开发者和初创企业的理想选择。未来,DeepSeek有望推出更多小型化模型,拓展低资源设备的AI生态。
1278 8
|
5月前
|
存储 机器学习/深度学习 数据库
阿里云服务器X86/ARM/GPU/裸金属/超算五大架构技术特点、场景适配参考
在云计算技术飞速发展的当下,云计算已经渗透到各个行业,成为企业数字化转型的关键驱动力。选择合适的云服务器架构对于提升业务效率、降低成本至关重要。阿里云提供了多样化的云服务器架构选择,包括X86计算、ARM计算、GPU/FPGA/ASIC、弹性裸金属服务器以及高性能计算等。本文将深入解析这些架构的特点、优势及适用场景,以供大家了解和选择参考。
935 61
|
8月前
|
存储 机器学习/深度学习 人工智能
2025年阿里云GPU服务器租用价格、选型策略与应用场景详解
随着AI与高性能计算需求的增长,阿里云提供了多种GPU实例,如NVIDIA V100、A10、T4等,适配不同场景。2025年重点实例中,V100实例GN6v单月3830元起,适合大规模训练;A10实例GN7i单月3213.99元起,适用于混合负载。计费模式有按量付费和包年包月,后者成本更低。针对AI训练、图形渲染及轻量级推理等场景,推荐不同配置以优化成本和性能。阿里云还提供抢占式实例、ESSD云盘等资源优化策略,支持eRDMA网络加速和倚天ARM架构,助力企业在2025年实现智能计算的效率与成本最优平衡。 (该简介为原文内容的高度概括,符合要求的字符限制。)
|
6月前
|
存储 机器学习/深度学习 算法
阿里云X86/ARM/GPU/裸金属/超算等五大服务器架构技术特点、场景适配与选型策略
在我们选购阿里云服务器的时候,云服务器架构有X86计算、ARM计算、GPU/FPGA/ASIC、弹性裸金属服务器、高性能计算可选,有的用户并不清楚他们之间有何区别。本文将深入解析这些架构的特点、优势及适用场景,帮助用户更好地根据实际需求做出选择。