Chapter 10
auto parallel
北大河图大模型自动并行训练工具Galvatron:https://zhuanlan.zhihu.com/p/591924340
系统特性 为了解决上述问题,研究者们提出了一些系列工作来探索混合并行的自动搜索:一类工作主要讨论了同时考虑数据并行和模型并行的搜索空间,代表性工作包括FlexFlow,Tofu,另一类工作则产生于流水并行场景,将其与数据并行相结合,代表性工作包括PipeDream,DAPPLE。在此基础上还有一些衍生工作,如Unity、Alpa,进一步扩展了自动并行的探索范围。我们提出的系统「惊破天」Galvatron同样属于自动并行搜索的研究领域,但相比于现有工作,我们主要拥有以下三方面优势:
(1)在多样性方面,现有工作能够支持的并行维度仍然比较有限,我们不仅可以支持更多的并行维度,并且面对更加差异化的Transformer模型结构也能够精准建模,以及在不同的集群硬件条件下验证了其自适应调优的能力。
(2)在复杂性方面,由于我们在多样性上的优势,导致我们面临了前所未有的庞大的搜索空间,为此,我们探究了几条目前大规模分布式训练过程中的一些经过实验性或理论性验证的重要观察,作为我们的搜索空间剪枝准则,从而实现高效的分布式执行计划优化。
(3)在实用性方面,我们结合了理论建模和实验测量两者的优势,实现对内存、通信、计算开销的精确估计,甚至考虑到了计算与通信重叠所导致的GPU执行效率下降问题,从而保证我们能够得到足够准确的自动并行优化结果。
另外,Galvatron底层选择PyTorch作为执行引擎,兼容Huggingface等常见的主流Transformer模型实现,所以完全不会对PyTorch用户带来额外的使用负担;同时也不需要用户付出额外的系统安装或者调试代价,使用时只需要添加几行代码,就可以轻松完成自动并行的整个流程。
关键技术
- 基于决策树的搜索空间分解
Galvatron的设计目标是高效地在复杂而庞大的并行策略空间内自动搜索,并为给定的Transformer模型和分布式环境生成最佳的并行执行计划。在搜索空间上,Galvatron是业界首个考虑4种主流并行方法的自动并行训练系统,包括包括数据并行(data parallelism, DP)、分片数据并行(sharded data parallelism, SDP)、张量并行(tensor parallelism, TP)和流水并行(pipeline parallelism, PP)。由于混合并行策略会包含以上四种并行算法的任意组合,在多GPU的场景下这种组合带来的搜索空间十分庞大。例如,对于双机四卡场景,一种可行的策略是在机内使用2-way TP,机间使用2-way PP,另一种可行的策略是在机内使用2-way PP,机间使用2-way DP。当节点内GPU数量扩展到8卡时,模型每一层的候选策略都多达上百种。随着模型层数的增加,其搜索空间规模指数增长,难以有效探索。
为了高效地搜索如此庞大的搜索空间,我们首先提出了以下观察作为指导:
Takeway#1:PP倾向于被跨设备岛放置。此处“设备岛”指具有高内部带宽的一组设备,在绝大多数Transformer模型中,PP的通信量相比于其它并行方式,显著更少。因此,人们通常优先对模型进行PP切分并放置于设备岛之间。
Takeway#2:在同构设备的前提下,并行策略倾向于将设备均匀切分。例如,对于4卡GPU的2-way DP倾向于将设备切分两组2卡的设备,而不是一组1卡和一组3卡的设备。在这种情况下,一个设备组内的最优混合并行策略与其他组内的最优策略保持一致。
Takeway#3:一般来说,在能够混合使用DP和SDP的情况下,只使用SDP在理论上性能更优。根据分析结果,
基于以上重要观察,我们提出了一种基于决策树的搜索空间构建方法:
- 给定一个Transformer模型,基于Takeway#1和Takeway#2,Galvatron首先用PP将模型切分成多个阶段,同时将设备均匀且连续地切分为多个设备组。例如8卡场景下,模型被切分为1/2/4/8-way PP,分别对应设备组大小为8/4/2/1。
- 每种PP切分对应着一棵决策树及一个子搜索空间,决策树叶结点总数为设备组大小,决策树高度为可用的并行方法数量,即决策树每一层可应用一种并行策略。
- 并行策略在决策树不同层间不可重复使用。
- 非叶结点的度数默认在2的指数次幂{2,4,8,…}中选择。
基于以上决策树构建规则,Galvatron构建的决策树能表示以上并行的任意组合。Takeway#1和Takeway#2帮助Galvatron规避了低效的并行组合,缩小了搜索空间。对于8卡GPU训练一层模型的场景,以上规则将产出34种候选的混合并行策略。进一步,使用Takeway#3将DP和SDP同时出现在一棵决策树的情形剪枝后,8卡候选策略数降至22种。
-
基于动态规划的并行优化算法
-
基于混合建模的执行代价估计方法
