作者：苗又山
编者按：深度学习技术已经在图像识别、语音识别、自然语言处理、搜索推荐等多个领域不断展现出巨大的应用价值。然而，随着模型规模的不断增大，深度学习模型的训练变得耗时且昂贵，设计最优的并行策略组合以提高其在多设备上的执行性能是目前该领域的一大挑战。在此背景下，微软亚洲研究院提出 nnScaler 技术，通过一套并行化原语和策略限定搜索的方法来寻求最佳的并行策略组合。这一尝试为寻求深度学习并行策略最优化提供了方案和工具，有效应对了当前的深度学习训练效率难题。相关论文已被计算机系统领域顶级学术会议 OSDI 2024 收录。
目前，深度学习已经透过多种应用形式广泛深入人们的日常生活当中。但随着模型规模的不断增大，其训练变得耗时且昂贵。为此，科研人员不断设计出各种并行策略来提高其在多设备上的执行性能。
在这个过程中，来自微软亚洲研究院的研究员们注意到现有的方法依然会遗漏许多潜在的高效并行策略。如何充分探索并行策略的可能性从而进一步释放性能，成为研究员们关注的焦点。
并行化原语：并行策略基本构成的统一表示
为了解决现存的并行策略问题，研究员们重新审视了并行策略的现有实现方式。目前主流的训练系统（如 Megatron-LM、DeepSpeed、Alpa 等）中通常内置了数种并行策略，例如，数据并行（data-parallelism）、张量并行（tensor-parallelism）和流水线并行（pipeline-parallelism）等，并且允许通过配置实现策略的联合应用。这样的方式对于使用者来说较为简便，但仍然可能遗漏高效的策略。而且当人们需要使用新的高效并行策略时，也会因为只能通过修改大量的系统底层代码来实现而显得捉襟见肘。
为此，研究员们重新思考了并行策略的基本构成。众所周知，深度学习的模型执行通常可以表示为一张数据流图（data-flow graph），其中以表示数据的张量（tensor）作为图边，张量的计算算子（operator）作为顶点。并行化的过程实际上就是将原本针对单一设备执行的计算数据流图重新进行分割变换，从而变成针对多个设备执行的分布化数据流图。因此，研究员们提出了一套基本操作作为并行化原语，包含：
- op-trans 表示如何对算子和张量进行划分；
- op-assign 表示对划分后的算子设定哪台设备来执行；
- op-order 可以设定同一台设备上算子的执行先后顺序。
利用这套原语，数据流图中的每个算子和张量如何划分，以及在空间（设备间）和时间（同一设备）上如何调度，都能够得到准确的描述和解释。
并行原语可以用于表示目前广泛应用的各种并行策略。例如，数据并行就可以表达为将所有前向（forward pass）和后向（backward pass）算子进行数据样本维度的划分，均匀地分配到所有的设备上，而将优化器（optimizer）的算子复制到每台设备上。并且所有设备上的算子保持原图中相同的执行顺序。通过这样一种并行原语的引入，研究员们就能使各种并行策略在同一个框架下进行描述和融合，从而极大地拓展了并行策略空间的表示范围。
这种通用性不仅系统化地表达了现有的并行策略，还为探索新的并行策略提供了可能。例如，对于产生单设备无法容纳中间结果的算子，原有的策略自然是在张量并行的划分后，分配到多个设备上并且采用通信协调，而通过原语仅对该算子进行张量划分后，即可全部调度到同一个设备上顺序执行，成功执行的同时也避免了可能产生的通信开销。
策略搜索限定：提高最佳策略搜索效率
尽管并行原语扩展了策略空间，但这种扩展也带来了新的挑战——巨大的策略空间让搜索在有限的时间内完成变得非常困难。面对海量的可能性组合，如何高效地找到最佳策略成为了一个亟待解决的问题。
研究员们的解决方案是借助“专家”的经验智慧，进行有效的策略搜索限定。这种指导通过并行原语进行描述，可以无缝衔接到研究员们的系统中。例如，下图中研究员们通过 op-trans 的参数限定算子 op 只使用 algo 集合指定的划分方案，并且只切分为 num 限定的份数。
“专家”通过设置特定的约束条件，能够大幅度缩小搜索空间，使得搜索过程更加高效且有针对性。在实验中，研究员们观察到通过约束搜索空间可以获得高达10倍的搜索效率提升，同时不损失策略的效果。
利用这种方法，研究员们不仅发现了现有方法所遗漏的高性能策略，还能够在较短的时间内完成搜索，提升深度学习训练的效率。目前这一系统已经在多个深度学习模型的训练中得到了验证，展示了显著的性能提升。
更多细节可参考论文详情：
nnScaler: Constraint-Guided Parallelization Plan Generation for Deep Learning Training
https://www.usenix.org/conference/osdi24/presentation/lin-zhiqi (opens in new tab)
借助引入并行原语和专家智慧指导的策略搜索，微软亚洲研究院的研究员们有效地解决了深度学习并行策略设计中的诸多问题，并在极大地拓展深度学习并行策略空间的同时，为未来的并行策略研究提供了新的方向和工具。研究员们非常期待这一方法能够在更广泛的应用中展示其潜力，为深度学习的发展带来更多的可能性。