阿里云机器学习平台 PAI 与香港大学合作论文入选 INFOCOM 2022，有效减少大规模神经网络训练时间

近日，阿里云机器学习平台 PAI 与香港大学吴川教授团队合作的论文」Efficient Pipeline Planning for Expedited Distributed DNN Training」入选 INFOCOM（IEEE International Conference on Computer Communications）2022，论文提出了一个支持任意网络拓扑的同步流水线并行训练算法，有效减少大规模神经网络的训练时间。

作为分布式机器学习的一种主流训练方式，流水线并行通过同时进行神经网络计算与中间数据通信，减少训练时间。一个典型的同步流水线并行方案包含模型切分设备部署与微批量（micro-batch）执行调度两个部分。

以下的两个图给出了一个 6 层神经网络模型在 4 块 GPU 上进行同步流水线并行训练的示例。由图表 1 所示，模型被切分成三个片段，其中第二个片段由于其计算量较大，被复制到两个 GPU 上通过数据并行的方式训练。图表 2 表示模型的三个微批量的具体训练过程，其中，由于第二个片段以数据并行方法在 GPU2 和 GPU3 上训练，在全部微批量训练完成后通过 AllReduce 算子同步模型片段参数。

图表 1 模型切分设备部署

图表 2 微批量执行调度

然而，设计高效的流水线并行算法方案仍然存在诸多挑战，例如深度学习模型各异，每层的训练时间也不相同，因此难以找到最优的模型切分部署方案；当前的流水线并行算法局限于同质化的 GPU 间网络拓扑，而现实机器学习集群具有复杂的混合 GPU 间网络拓扑（例如，单个机器上的 GPU 可以通过 PCIe 或者 NVLink 连接，跨机通信可以基于 TCP 或者 RDMA），导致现有方案无法使用等，以上问题导致实际训练中的 GPU 使用效率低。

针对以上难点，团队提出了一个近似最优的同步流水线并行训练算法。算法由三个主要模块构成：

1）一个基于递归最小割的 GPU 排序算法，通过分析 GPU 间网络拓扑确定 GPU 的模型部署顺序，保证最大化利用 GPU 间带宽；

2）一个基于动态规划的模型切分部署算法，高效率找到最优的模型分割与部署方案，平衡模型在每个 GPU 上的运算时间与模型切片间的通信时间；

3）一个近似最优的列表排序算法，决策每个微批量在各个 GPU 上的执行顺序，最小化模型的训练时间。

从理论上对算法做出详尽分析，给出了算法的最坏情况保证。同时，在测试集群中实验证明团队的算法相对 PipeDream，可以取得最高 157% 的训练加速比。

INFOCOM 是计算机网络三大顶级国际会议之一，涉及计算机网络领域的各个方面，在国际上享有盛誉且有广泛的学术影响力。此次入选意味着阿里云机器学习平台 PAI 在分布式深度学习模型训练优化领域的工作获得国际学界的广泛认可，进一步彰显了中国在分布式机器学习系统领域有着核心竞争力。

阿里云机器学习 PAI 是面向企业及开发者，提供轻量化、高性价比的云原生机器学习平台，一站式的机器学习解决方案，全面提升机器学习工程效率。

来源：互联网