分布式机器学习网络通信优化技术

张汉钢，邓鑫源，宋 晔，薛旭伟，郭秉礼，黄善国（北京邮电大学，北京 100876）

1 概述

近些年随着互联网的蓬勃发展及大数据技术的不断演进，在图像识别、语言翻译、自动驾驶等相关领域，机器学习取得了众多突破并被广泛成功应用。在进行机器学习训练时，通常需要大量的数据用于计算，以保证模型的准确性。在海量数据的背景下，训练所使用的数据集的规模越来越大，训练的模型越来越复杂，单个机器无法满足存储与算力的需求，分布式机器学习将是未来发展的必然趋势。分布式机器学习需要快速、高效的连接网络作为支撑，该网络结构有2 种主要的并行方式：数据并行与模型并行。数据并行是最为常见的并行模式，广泛应用于大规模分布式系统中。在采用数据并行的方式进行训练时，各个机器之间需要对大量的数据进行同步，例如GPT-3就拥有超过1 750 亿个机器学习参数。通信时间在训练时间中占比较大，如何减少通信开销成为了热门课题。数据并行使用all-reduce操作进行梯度同步，其中Ring all-reduce算法在高性能计算领域得到广泛应用。

Ring all-reduce基于同步通信，在计算节点的性能存在差异时，整体的计算性能会被较慢的计算节点拖慢。为解决同步通信的这一弊端，异步并行（ASP）与延迟同步并行（SSP）被相继提出，但Ring all-reduce 架构难以支持ASP和SSP等模型的一致性协议。本文进一步分析Ring all-reduce 算法特性，提出一种对其同步参数的过程进行加速的优化方案，降低慢节点带来的影响。

2 优化策略及理论分析

2.1 Ring all-reduce基本分析

百度在2017年提出将Ring all-reduce 架构引入到分布式机器学习中来，使这个架构受到广泛关注。Ring all-reduce 算法分为2 个部分：Reduce_Scatter 和Allgather。Reduce_Scatter 阶段按照n个训练节点数将数据划分为n个chunk，将从上一相邻节点接收到的chunk 同本地chunk 进行reduce 操作并发，该阶段包含n-1 个步骤，如图1 所示。Allgather 阶段是将汇总后的梯度段依次传递，用接收到的内容替换本地内容，该阶段包含n-1个步骤（见图2）。

图1 Reduce_Scatter阶段示意

图2 Allgather阶段示意

在Github 上百度公开给出名为baidu-allreduce 的小型c++库，具体演示了环形的all-reduce 与Allgather过程。Ring all-reduce 的同步过程体现在进行reduce时，调用了同步CUDA 流（cudaStreamSynchronize）来等待所有GPU都完成操作。reduce操作执行完毕也就意味着完成了一轮同步过程，随后向后一个节点进行发送操作。

当某个节点的reduce 阶段所用时间显著变长，即该节点为慢节点时，若按照Ring all-reduce 过程进行同步，则需要等慢节点多产生的时间，而这多出的时间将直接加在Ring all-reduce 原有的时间开销上，影响到Ring all-reduce 算法的执行效率。为了减轻慢节点带来的负面影响，如何利用慢节点较正常节点多出来的时间便成为了关键，下文给出一种优化策略思路。

2.2 优化策略

以4 节点为例对Ring all-reduce 的Reduce_Scatter阶段进行分析。设置TRS为节点接收数据、进行reduce操作和发送数据的时间，Ttrans为传输数据的时间。A节点将a0 发送至B 节点，B 节点将收到的a0 与本地的b0进行汇总，并在下一次通信时将汇总的内容发送至C节点（见图3）。在这一时期，B节点因reduce过程缓慢而成为慢节点，TRS变长，C 节点需要等待的TRS+Ttrans时间也随之变长，将导致训练效率下降。

图3 Ring all-reduce的通信步骤

本文提出一种优化策略，设置T表示在没有慢节点情况下的TRS的平均值，则Reduce_Scatter 阶段1 个步骤的实际开销TRS=T+△T，其中△T为因外在因素导致的时间波动；αT为节点触发策略的时间阈值，α>1。当出现慢节点时，△T变大；当TRS=T+△T>αT时，则触发策略，具体过程如图4所示。

图4 改变策略后的通信步骤

假设B 节点是慢节点，Tt为C 节点等待时间。当TRS=T+△T>Tt=αT+Ttrans时，改变C 节点的操作步骤。此时C 节点将不再等待接收B 节点发送的a0＋b0，而是直接将本地的c0 发送至d0；B 节点达到TRS时将跳过C节点直接与D节点进行通信。为了更直观地分析时间上的重叠，假设如下。

a）各个结点之间传输数据的时间Ttrans都一致。

b）若节点为慢节点，其TRS=T+△T且△T>0。

c）若节点为非慢节点，其TRS=T。

d）设定B节点为慢节点。如图5 所示，当B 节点是慢节点且C 节点等待至αT+Ttrans时，C 节点直接向D 节点传递数据，D 节点先收到来自C 节点的数据并等待B 节点发送来的数据。在收到2 份数据后，D 节点一并进行汇总。截至D 节点汇总所有数据时，策略带来的时间减少值如式（1）所示。

图5 B节点发送的消息先到达D节点

由式（1）可知，当△T较大时，在T+△T>αT+Ttrans情况下策略带来的时间减少值为固定时间。

如图6 所示，当TRS=T+△T

图6 C节点发送的消息先到达D节点

当△T与α取合适的值时，可以使Tdec取得一个较大值，即在有慢节点的情况下能减轻慢节点带来的负面影响，达到优化效果。

3 仿真实验设置及结果分析

采用OMNet++框架对Ring all-reduce 算法和改进策略进行模拟。因为Allgather 阶段这2 种方式的开销相同，故仅比较2 种方式完成Reduce_Scatter 所用的时间（分别是TRing和TStrategy）。基于第2.2 节做出的基本假设，对相关的仿真参数进行设定：T=3 μs；Ttrans为固定链路时延5 μs 与数据包传递时间之和；数据包大小取8 KB；链路带宽为100 Gbit/s。设置不同的△T与α的值，优化效果即Tdec与TRing的比值。

如图7 所示，Ring all-reduce 算法完成Reduce_Scatter 阶段所用时间随着△T的增加而线性增加，而采用优化策略完成Reduce_Scatter 阶段所用时间均得到减少，说明采取优化策略能减少时间开销。进一步分析，如图8 所示，当α不变时，优化效果大体上随着△T的增大而先增后减。当△T较小时，TRS=T+△T

图7 Ring all-reduce与优化策略时间开销比较

图8 △T与α对优化效果的影响

当TRS=T+△T>Tt=αT+Ttrans，优化效果与α的变化无关（见图9），结果符合式（1）预期。

图9 α变化对优化效果影响

T+△T<αT+Ttrans时，即慢节点跳过后一节点发送的数据先到达，α 增大，说明后一节点最长等待时间将延长，触发机制的机会将相对减少，机制带来的增益受到限制，优化效果随之减弱。

当△T是一个较大偏置时，有T+△T>αT+Ttrans，即后一节点发送的数据先到达而慢节点发送的数据后到达，此时获得的时间减少值Tdec是一个定值。△T的增加会导致TRing的时间线性增加，Tdec与TRing的比值也随之减少，即优化效果减弱（见图10）。

图10 △T变化对优化效果影响

4 总结

在分布式计算中，Ring all-reduce是一种常用的通信模式，然而由于网络中存在慢节点，通信延迟分布不均匀，会影响整体性能。本文聚焦Ring all-reduce算法的Reduce_Scatter阶段，提出一种创新性的优化策略实现对慢节点的感知与通信调整，减小慢节点对整体通信性能的负面影响。仿真结果显示，相较于传统的Ring all-reduce 算法，本文提出的优化策略成功降低了通信的时间开销，训练时长最高节省了25.3%。