「分布式机器学习系统网络性能优化」研究进展

2022 年 10 月 1 日 专知

以机器学习为代表的人工智能技术需要对海量数据进行处理，对底层算力要求极高。分布式机器学习通过将计算任务分布式地部署到多个计算节点来加快模型的训练速度，从而将训练任务完成时间降低到可接受范围。由于通信开销对分布式机器学习系统的扩展性具有重要影响，因此，分布式机器学习系统网络性能优化受到各界研究者的广泛关注。本文首先分析了分布式机器学习系统扩展性不足的主要原因，并提出了改善其扩展性的关键思路，然后系统地综述了分布式机器学习系统网络性能优化相关的研究工作，并对这些研究工作从多个角度进行了对比分析。最后，对分布式机器学习系统网络性能优化研究的未来发展趋势进行了展望。

引言

近年来，以机器学习，尤其是深度学习，为代表的人工智能技术在图像识别[1]、语音识别[2]、机器翻译[3]和自动驾驶[4]等应用领域都取得了突破性进展。其原因可以分为三个方面：首先，互联网、大数据技术的发展，积累了海量的训练数据，在这些数据中蕴含着丰富的信息；其次，机器学习理论的发展使得机器学习算法和模型不断完善，为从训练数据中挖掘有价值的信息创造可能；最后，GPU 等加速器以及云计算等技术大幅提升计算性能，大大加快了从数据中获取信息的速度，使得机器学习技术的应用成为现实。总之，作为第三次人工智能浪潮的“催化剂”，算力的大幅提升直接将人工智能再次推向新的繁荣期。高性能的机器学习算法往往具有更高的计算需求。据 OpenAI 统计，人工智能训练所需要的算力呈指数级增长，每 3.5 个月翻一倍①。相比之下，近年来计算引擎的发展速度则远远落后于模型计算需求的增长。以 Nvidia GPU 发展为例，表 1 展示了 2012 年以来 Nvidia 的多代 GPU 在训练 ResNet 模型时的性能表现。可以看到近 8 年来， GPU 的计算性能只提高了 16 倍左右，远低于同期模型计算需求的增长。在“后摩尔定律”时代，单个计算引擎的性能提升逐渐进入了瓶颈期。面对日益复杂的计算任务，分布式机器学习被认为是必然的发展趋势，逐渐成为业界的研究热点[5]。

在分布式机器学习训练任务的迭代计算过程中，不同计算节点间需要频繁同步机器学习模型参数，以使得该模型能够遍历完整的数据集，从而保证最终得到的模型与使用单机训练的模型一致。然而，随着计算节点数量的增多，一方面，不同节点间进行参数同步的流量逐渐增加；另一方面，为了避免单轮训练过多数据（即批尺寸过大）带来的模型泛化能力下降问题[6]，每个节点所分配的计算任务会逐渐减少。因此，对于分布式机器学习系统，通信开销和计算开销的比值会随着系统规模的增大而呈现幂增长趋势。这导致通信成为限制大规模分布式机器学习系统扩展效率的主要瓶颈，甚至出现随着节点数量增加，模型训练速度反而下降的情况[7-11]。并且，过多的通信时间会导致 GPU 等昂贵的计算设备大部分时间处于等待参数同步的状态，造成计算资源的浪费。因此，研究如何对分布式机器学习系统的网络性能进行优化，降低通信操作对分布式机器学习系统扩展效率的影响，从而提高机器学习模型的训练速度，具有重要的研究意义和实用价值。

本文将首先介绍分布式机器学习系统的通信特点，并分析网络通信成为分布式机器学习系统扩展性瓶颈的原因，然后提出三种优化网络性能的关键思路，并以这些思路为指导，从通信模式、通信数据量、通信效率以及网络拓扑等方面具体地介绍分布式机器学习系统网络性能优化研究的最新进展，并从加速效果、优化机制、扩展性、对模型收敛性的影响以及是否需要升级硬件设备或更新互联方式等多个角度对这些研究工作进行对比分析，最后讨论分布式机器学习系统中网络性能优化研究的未来发展趋势。

2. 分布式机器学习

随着信息技术快速发展，全球数据呈现爆发式增长，推动人类社会迈入大数据时代。在大数据时代，机器学习训练任务往往需要对海量的训练数据进行大量的计算，以提高模型的准确度。在单机上执行这样的训练任务，无论是在计算速度还是在数据存储方面都显得十分吃力。例如，使用单块 Nvidia Tesla V100 GPU 训练自然语言处理模型 BERT-large 需要耗时 78 天，这显然是不可接受的。分布式机器学习的目标则是将训练任务分布式地部署到多个计算节点，从而提高模型训练的速度，减少任务耗时。因此，分布式机器学习已经成为机器学习最热门的研究领域之一。

分布式机器学习的并行方式主要包括数据并行（Data Parallelism）和模型并行（Model Parallelism）。如图 3 所示，数据并行是指每个计算节点上均具有同一机器学习模型的副本，但不同计算节点分配到的训练数据是不同的，不同计算节点间需要将各自的模型更新进行同步，以保证机器学习模型的全局一致性。模型并行则是将机器学习模型划分为多个子模型，并分别部署在不同计算节点上，训练数据统一输入，前一节点完成子模型计算后将计算结果传递给后一节点继续对该训练样本进行处理。数据并行由于操作简单，且不同节点的计算负载比较均衡，应用最为广泛。目前，TensorFlow[13]、Pytorch[14]和 MXNet[15]等主流机器学习框架均对数据并行提供了支持，并且具有极好的易用性，但模型并行仍需要用户手动对模型进行划分和分布式部署。

假设某分布式机器学习系统共有 m 个计算节点，批尺寸为 n。当采用数据并行时，每个计算节点每轮处理的样本数量为 n/m。各计算节点基于所分配到的训练数据对模型进行更新，然后将不同计算节点更新后的模型进行汇总。汇总后的模型参数如下：

对比公式(2)和公式(3)可知，如果在每一轮训练结束时，将模型参数在所有计算节点间进行同步，则分布式训练时的模型参数变化和单机训练完全相同，即分布式训练可以在不改变模型的收敛性的前提下提高模型的收敛速度。需要说明的是，在实际应用中，一般是对不同计算节点的梯度进行汇总，然后使用汇总后的梯度来更新模型参数，并将新的参数赋给各计算节点的模型副本。相比于单机训练，分布式训练额外引入了节点之间的数据通信，从而导致分布式训练的速度无法随着计算节点数量的增加而线性提高。衡量分布式机器学习训练加速效果的指标主要包括加速比（speedup）和扩展效率（scaling efficiency）。加速比是指同一机器学习训练任务在单机训练和分布式训练时所需时间的比值。加速比越大，分布式训练的加速效果越显著，也就是说，可以更快地完成训练任务。需要注意的是，加速比有可能出现小于 1 的情况，此时分布式训练速度反而不及单机训练。扩展效率是指加速比和计算节点数量的比值。扩展效率越高，各计算节点的计算资源利用率也就越高。图 4 展示了利用 Nvidia Tesla V100 GPU 训练 BERT 模型时加速比和扩展效率随 GPU 数量的变化，可以看出加速比和扩展效率的变化趋势并不相同。一般来说，随着节点数量的增多，扩展效率呈下降趋势，而加速比则呈先升后降趋势。加速比在节点数量增多时反而下降，是由于参数同步引入的通信开销抵消了新增节点带来的性能收益。

具体来说，网络性能之所以会成为分布式机器学习系统扩展性瓶颈的主要原因有以下三点：1）机器学习模型越来越复杂，模型参数量不断增加。机器学习算法理论的快速发展催生出各种各样的机器学习模型。例如，OpenAI 最近提出的自然语言处理模型 GPT-3 [16]具有 1750 亿参数，而 Krizhevsky 等人在 2012 年提出的 AlexNet[17]模型参数量仅为 0.45 亿。这导致在相同节点规模下，任意两个计算节点间需要同步的参数量随之大幅增加，加剧了分布式机器学习系统中网络通信的压力；2）复杂的机器学习模型同时也意味着需要更多的算力，导致分布式系统规模的增大。如前所述，在机器学习模型不变的情况下，分布式系统规模越大，每个节点所承担的计算任务越少，计算耗时越短；与之相反，分布式系统规模越大，每个节点需要通信的对端节点数量越多，通信耗时越长。因此，随着分布式机器学习系统规模的增大，通信开销在整体模型训练开销中的占比越来越高；3） GPU 等计算设备的性能提升速度快于网络设备的升级。虽然计算设备的性能提升速度远不及模型算力需求的增长，但仍比网络设备的升级速度更快。计算性能和通信性能的差距越来越大，即，计算资源在单位时间内处理的数据需要更长的时间才能被网络资源处理完。这意味着分布式机器学习系统中的网络瓶颈问题将会日益严重。

如图 5 所示，为消除（或缓解）网络性能对分布式机器学习系统扩展性的限制，从本质上来说，有以下三种根本性思路：

1）降低通信需求。通信需求，即需要通信的数据量和通信次数，对通信耗时具有根本性的影响。显然，数据量越大，或者通信越频繁，通信耗时越长。因此，为了减少通信耗时，可以从机器学习算法层面构建低网络通信需求的训练模型，或采用知识蒸馏和模型剪枝等方式对原模型进行修改以减小模型尺寸。这些方法会造成训练模型的变化，超出了本文的讨论范围，故后文不再作详细描述。参数量化以及参数稀疏化保持训练模型不变，通过降低被传输的参数量来降低通信需求。并行方式的优化通过权衡参数数据量和激活值数据量的大小来切换不同的并行方式：当参数数据量较少时，采用数据并行；反之，采用模型并行。模型一致性协议通过控制参数同步的频率来调节通信需求。参数同步架构对通信次数和每次通信的数据量均会产生影响，通过选取合适的参数同步架构可以有效降低通信需求；

2）提升通信能力。在通信需求一定时，分布式系统的通信能力越高，通信耗时越短。通信能力的提升主要有两种方式。一种是利用 RDMA、NVLink 等高性能传输协议实现高带宽、低时延的网络传输，或利用 MLT 等新型机器学习专用传输协议降低丢包对传输性能的影响；另一种是采用高带宽的网络互联拓扑。例如，BCube 和 BiGraph 均采用多网卡服务器架构，不但大幅提高每个计算节点对外通信的能力，甚至可将节点内的通信流量导出到服务器外部，从而绕过 PCIe 瓶颈；

3）提高通信效率。在通信需求和通信能力均确定的情况下，还可以通过提高通信效率来加速分布式机器学习训练。例如，在 GPU 节点间进行集合通信时，NCCL 通信库由于针对 GPU 设备采取了定制优化，因此具有比传统集合通信库 MPI 更高的性能。网内聚合通过逐跳汇聚参数，增加了单位数据所蕴含的参数信息，从而提高了通信效率。通信调度是在通信需求和通信能力固定的情况下，高效利用网络资源的一种方式，既包括采用小尺寸梯度聚合来降低启动开销的方案，也包括优先传输紧急参数来增加计算和通信重叠程度的方式。

3 参数同步模式优化

在大规模分布式机器学习训练场景中，计算节点间需要频繁地进行参数同步，因此，参数同步模式对整体训练性能具有重要的影响。本节将从模型一致性协议和参数同步架构这两个方面详细介绍对参数同步模式进行优化的相关工作。

3.1 模型一致性协议

在数据并行模式下，每个计算节点都需要保存一份相同的模型副本，然后使用本地的训练数据对模型副本进行更新。因此，在训练过程中，不同计算节点所维护的模型副本会出现差异。为了使得分布式训练能够取得与单机训练相同的效果，需要保证这些模型副本的一致性。

3.2 参数同步架构

除模型一致性协议外，参数同步架构对大规模分布式机器学习系统的性能也有至关重要的影响。参数同步架构是指不同计算节点上的模型更新进行汇总，并对模型副本进行更新的方式。按每次模型更新是否推送给所有其他计算节点，参数同步架构可以分为中心化架构和去中心化架构；按模型更新是否在单一逻辑节点汇总，可以分为集中式架构和分布式架构。一般来说，中心化架构既可以是集中式架构，也可以是分布式架构；而去中心化架构一定是分布式架构。目前常用的参数同步架构大多为中心化架构，因此，除特别说明外，下文所提到的参数同步架构均为中心化架构。

4. 优化方案

通信效率优化

在实际部署中，通信的性能和效率也会对参数同步过程产生重要的影响。即使采用相同的参数同步模式，不同的通信方式也会对整体的训练性能造成很大的差异。因此，为了提高分布式机器学习训练时的通信效率，研究人员在以下方面进行了深入研究。

并行方式优化

如上文所述，数据并行和模型并行是两种经典的分布式机器学习训练方式。对于数据并行来说，通信开销主要来自不同计算节点间的参数同步；对于模型并行来说，当某个计算节点的输入来自另一个计算节点的输出时，便会产生通信开销。当模型参数量小于中间计算结果的数据量时，数据并行带来的通信开销较小；反之，模型并行的通信开销更小。然而，对模型整体使用某一种并行方式，可能无法达到最优的训练性能。因此，一些工作[68-74]提出使用混合并行、流水并行等方式，通过细粒度的并行优化来提升分布式训练性能。

网络拓扑优化

除以上优化方案外，分布式机器学习系统领域的研究者也对分布式训练集群所使用的底层物理网络拓扑提出了优化方案。

表 4 综合对比了近年来研究人员所提出的分布式机器学习系统网络性能优化研究相关工作。对比的主要指标包括优化机制、训练加速效果、节点扩展性、对模型收敛性的影响以及是否需要更换硬件设备或者互联方式等。这些工作从多个层面对分布式机器学习系统的网络性能进行优化，不同机制之间各有优劣。从训练加速效果来看，ASP 这一模型一致性协议将通信开销从模型训练的核心路径上移除，使得网络通信不会阻塞训练过程，加速效果非常好；模型压缩或并行方式优化等方案，有效地减少了各计算节点通信的数据量，而网内聚合方案则逐跳减少了网络中的流量，因此，这些方案具有非常好的加速效果；传输协议和通信库优化类方案提高了点到点通信性能，网络拓扑优化类方案提高了通信节点之间的互联带宽，这些方案的加速效果也很好；虽然不同参数同步架构的理论参数同步时间之间的差距主要来自于时延开销，但在实际中，负载均衡、多流竞争等都会影响不同参数同步架构的实际参数同步时间，总体来说，参数同步架构类方案的加速效果不如前面几种方案；通信调度类方案的加速效果与训练模型的通信/计算比高度相关，相比其他方案来说，加速效果比较有限。从节点扩展性来看，网内聚合类方案受限于交换机硬件计算能力和存储空间限制，通常应用于单机架规模的训练集群，扩展性较差；随着计算节点数量的增多，通信/计算比越来越高，导致通信调度类方案在节点数量较多时的扩展性较差；由于交换机端口数量、布线难度等因素的限制，底层物理网络拓扑的规模往往不能无限增大，如 BCube 适用于集装箱规模的数据中心，故网络拓扑类方案的扩展性一般；虽然传输协议和通信库类方案可将通信性

能提高数倍，暂缓网络瓶颈出现的时间，但随着节点数量的增多，通信操作又将成为系统瓶颈，故该类方案的扩展性也一般；并行方式优化类方案的模型并行粒度不能无限切分，因此在节点规模很大时，仍会出现大量节点使用数据并行的情况，并且对大量节点求解最优并行方式的算法复杂度也非常高，如 PipeDream 的求解时间与计算节点数的二次方和模型层数的三次方成正比，以上因素导致并行方式优化类方案的扩展性也一般；节点规模很大时，环规约架构的通信时间被时延开销所主导，参数服务器架构的连接数量也会大大增加，从而导致传输性能的降低，故参数同步架构类方案的扩展性一般；模型一致性协议从核心路径上移除了通信，模型压缩类方案可将通信量降低数十乃至上百倍，因此这两类方案的扩展性较高。从对模型收敛性的影响来看，模型压缩会导致参数同步时信息量的丢失，从而影响模型收敛性；SSP和ASP以及去中心化参数同步架构引入了陈旧参数对全局模型的更新，也会对模型收敛性产生一定影响；其他类方案不涉及通信内容的改变，故不影响模型收敛性；从对硬件的依赖性来看，RDMA 和 NVLink 需要专用的硬件设备，故依赖于硬件设备的更新升级；网内聚合类方案依赖于可编程交换机或专用交换机来实现在网络内部对参数进行聚合的目的，故该类方案也需要底层硬件设备的支持；网络拓扑类方案涉及对整个集群互联方式的修改，比较适用于新训练集群的搭建，在现有集群上的部署难度较大；其他类方案均为软件层方案，对底层硬件环境无特殊要求，因此部署难度较低，具有非常好的通用性。2016 年，Google 提出了一种分布式机器学习新形式—联邦学习[85]。本质上，联邦学习是一种加密的分布式机器学习框架，允许各参与方在不共享本地数据的条件下与其他各方共建模型。不同于传统分布式机器学习，联邦学习面临四个新的问题：客户端中数据非独立同分布问题、差分隐私问题、通信开销问题和客户端无状态问题。本文仅关注联邦学习中的通信开销问题。在联邦学习中，通信开销远大于计算开销，这主要是由于客户端与中央服务器之间的网络带宽有限，且连接质量较差，同时不同客户端的连接质量参差不齐造成的。虽然联邦学习是一种分布式机器学习框架，但有些针对传统分布式机器学习的网络性能优化方案却不适用于联邦学习。例如，客户端可能通过无线方式接入网络，故无法对这些客户端之间的互联方式进行改善；一般来说，联邦学习的网络瓶颈点基本在客户端侧，因此，网内聚合的方式不能解决联邦学习场景下的网络传输痛点；联邦学习中各客户端均要使用本地数据进行训练，并且不会将本地数据传输给其他客户端，因此联邦学习只能使用数据并行方式，无法通过并行方式优化的方式来提高训练速度。联邦学习场景下的网络性能优化主要依赖对通信内容的压缩来实现。一般来说，客户端上行链路的带宽比下行链路带宽更小，因此，一些工作 [85-87]最早尝试通过多种梯度压缩方式，如量化、稀疏化、下采样、矩阵分解、计数草图（count sketch）和周期平均等，来减小客户端的上行通信压力。

随后，一些工作[88]通过压缩参数的方式降低下行通信成本。文献[88]采用 Federated Dropout 的方式对神经元进行随机丢弃，这样客户端可以只训练一个更小的子模型，从而既减小了中央服务器到客户端的通信数据量，又能更加高效地完成本地计算。虽然模型压缩会减缓模型的收敛速度，但受限于网络连接质量，联邦学习不得不通过压缩通信内容的方式来降低通信成本，提高训练速度。相比之下，传统分布式机器学习训练集群的互联带宽非常高，并且连接可靠性极高，因此，模型压缩在传统分布式机器学习训练中往往作为可选方案，需要充分权衡收敛性和训练速度来决定是否需要对模型进行压缩，以及使用何种压缩方式和压缩比例。

5 研究趋势展望

分布式机器学习系统性能优化作为分布式机器学习领域最为热门的研究方向之一，正在吸引越来越多学术界和工业界研究人员的关注。由于分布式机器学习系统网络性能优化研究与工业界结合紧密，具有重要的实践价值，可以预计在未来数年内相关研究还将持续成为焦点。当前，国内学术界和工业界关于分布式机器学习系统网络性能优化的研究基本与国际水平处于并跑状态。因此，在国家大力发展新基建的背景下，加强分布式机器学习系统网络性能优化研究，不但能够为人工智能的发展提供内生动力，并且可以为依托人工智能实现外部赋能创造条件，对于推动传统行业信息化、数字化、智能化转型升级具有非常重要的意义。从网络通信的角度看，我们认为未来的分布式机器学习系统性能优化研究主要包括以下四个方向：

（1）模型的高质量压缩。分布式机器学习训练的通信数据量对通信耗时具有决定性影响。从机器学习算法的发展趋势来看，越来越大的机器学习模型已经成为必然[16][89]。因此，如何对训练超大模型时的通信数据进行高质量的压缩，既能大幅降低通信数据量，又不会造成训练信息的大量丢失，是未来缓解甚至彻底消除网络瓶颈的重要方向。当前，模型压缩程度仍然受到相关理论发展的限制，通信数据的压缩是以更多的通信次数为代价的。除相关压缩理论的突破外，未来可能的发展方向还包括细粒度的模型压缩方式，如不同层乃至不同算子采用不同的压缩方式、压缩比例，不同训练轮数采用不同的压缩方式、压缩比例，从而避免最差压缩比例限制整体的压缩效果。另一个可能的方向是综合考虑时空相关性的模型压缩方式，当前的压缩算法大多将每个参数值作为单独个体来处理，部分算法引入时间序列相关性以将相邻两轮训练间的结果相关联，从而降低随时间累积的压缩误差。然而，参数张量的空间相关性尚未得到充分重视。视频压缩领域中，基于时空相关性的视频帧间压缩方法已得到广泛应用。因此，模型压缩可以借鉴视频压缩领域的相关经验，综合考虑参数张量的时空相关性对模型参数采取进一步的有效压缩。

（2）并行方式优化。除压缩通信内容外，改进多节点之间的并行训练方式也是降低通信开销的重要途径。分布式机器学习训练通过将训练数据、训练模型分布到多个计算节点来达到并行训练的目的。即使对同一训练模型和相同训练数据而言，不同的并行方式也会产生完全不同的流量模式和通信数据量。现有方案大多在训练数据、模型不同层等维度对训练任务进行并行化分解，最近一些工作又引入了流水并行来提高计算资源利用率。但这些方案仍远未成熟。对流水并行来说，由于层间计算依赖关系的存在，这些方案或者无法完全消除 “气泡”，或者需要占用大量显存来存储多个模型版本。如何提高流水并行的效率，同时最小化硬件资源占用，对流水并行的应用前景至关重要。对模型并行来说，更细粒度的操作符拆分，使其能够并行化计算，从而提高单个操作符的执行速度上限，也是未来值得探索的重要方向。另外，在大规模分布式机器学习训练场景下，如何快速求解最优并行化方式，将大量的计算节点合理地进行编排，也是该类方案将来能否得到广泛应用的重要基础。

（3）多任务场景下的网络资源复用。现有的网络性能优化方案仍主要针对单任务场景而设计，对多任务之间的联合优化方案仍有待研究。但在实际训练场景中，计算设备往往被单一训练任务所独占，但网络设备却被很多训练任务共享，导致不同训练任务由于彼此竞争网络资源造成性能的互相影响。对于分布式机器学习训练任务来说，流量具有明显的周期特征，即，从宏观结构来看，平均流量并不高，但缩放到毫秒粒度，则会出现链路利用率在满载和空载之间频繁切换的情况。当多个训练任务同时使用网络资源时，所有训练任务的通信时间都会被拉长，导致训练速度的下降。因此，未来一个可能的发展方向便是通过使不同训练任务分时复用网络资源，尽量减小每个训练任务所花费的通信时间，从而提升整体的训练速度。

（4）专用网络设备和架构。网络硬件技术的提升，对于分布式机器学习系统性能的提升具有显著的效果。当前，分布式机器学习训练任务和其他业务一样运行在通用网络硬件设备之上。但是分布式机器学习训练任务具有自己的特点，如流量矩阵的确定性以及数据传输的周期性等。因此，针对分布式机器学习业务设计专用的网络设备和架构，如超低转发时延交换机、GPU 与网卡的一体化设计等，也将成为未来的研究热点。此外，光电互联技术的出现也使得数据中心网络的带宽和容量大幅提高。可以预料，未来在如何合理地利用这些新型网络设备来提升网络传输性能方面也将会产生更多的研究成果。

专知便捷查看