论文阅读-《Accelerating Distributed Reinforcement Learning with In-Switch Computing》

这是一篇关于在网计算的论文，它讲的是通过借助交换机的计算功能来加速分布式强化学习的训练。

---

作者：Youjie Li，Iou-Jen Liu，Yifan Yuan，Deming Chen，Alexander Schwing，Jian Huang
会议名称：the 46th International Symposium
会议时间：2019/06/22

---

摘要

  随着基于AI的新应用程序不断涌现，强化学习（RL）引起了人们的广泛关注，这些应用程序要求具备对环境变化做出智能反应的能力。与分布式深度神经网络（DNN）训练不同，分布式强化学习训练具有其独特的工作负载特征——它会以较小的大小但更频繁的梯度聚合产生更多数量级的迭代。更具体地说，我们对典型RL算法的研究表明，它们的分布式训练是主要的延迟，并且用于梯度聚合的网络通信最多会占到每次训练迭代执行时间的83.2％。
  在本文中，我们介绍了iSwitch，这是一种交换机内加速解决方案，可将梯度聚合从服务器节点移至网络交换机，从而可以减少梯度聚合的网络跳数。这不仅减少了同步训练的端到端网络等待时间，而且还通过更快的权重更新来提高异步训练的收敛性。在交换机内加速器上，我们通过在网络数据包而不是梯度矢量的粒度上进行动态梯度聚合来进一步减少同步开销。此外，我们重新考虑了分布式RL训练算法，并提出了一种分层聚合机制，以进一步提高机架级分布式RL训练的并行性和可扩展性。
  我们使用真实的可编程交换机NetFPGA板来实现iSwitch，并扩展了可编程交换机的控制和数据平面以支持iSwitch，而不会影响其常规网络功能。与最新的分布式训练方法相比，iSwitch在同步分布式训练中的系统级加速高达3.66倍，对于异步分布式训练的系统级加速高达3.71倍，同时实现了更好的可扩展性。

1 引言

  我们已经看到了一种颠覆性的趋势，即新兴的AI应用程序越来越多地在动态环境中运行，同时会采取相应的行动来应对环境变化。强化学习（RL）自然满足了新兴AI应用程序的这些要求。与其他流行的机器学习技术（例如深层神经网络（DNN））相似，RL还要求分布式训练来提高他们的性能和训练结果，这是基于对分析大量数据和训练更复杂模型的日益增长的需求。
  与分布式DNN训练不同，分布式RL训练通过更小的梯度聚合产生更多数量级的迭代。根据我们对流行的RL算法的研究（见表1），一个典型的RL算法将产生数百万次的迭代，而它的模型大小远远小于一个典型的DNN模型的大小。因此，每次迭代中梯度通信的延迟是影响分布式RL训练性能的关键因素。

  为了支持分布式RL训练，最先进的系统通常使用两种方法之一。他们要么采用集中式参数服务器，将每个worker的局部梯度聚合到中心服务器上进行权重更新，要么采用基于AllReduce的训练，以分散的方式进行梯度聚合。
  在前一种方法（参见图1a）中，集中式参数服务器是限制分布式训练可扩展性的瓶颈，因为所有worker必须与中央服务器交互，以便在每次迭代中传输梯度或接收更新的权重。考虑到RL训练涉及数百万次迭代，这个瓶颈将显著影响训练性能。
  后一种方法（参见图1b）是通过以循环方式执行梯度聚合来解决可扩展性问题的。然而，它需要更多的网络跳数通过交换机来完成集群中所有worker的梯度聚合。当我们使用更多的计算节点来扩展训练时，梯度聚合所需的网络跳数将线性增加。

  为了进一步了解这些方法的性能特点，我们用DQN、A2C、PPO和DDPG等算法量化了分布式训练中关键组件的开销。我们的研究结果表明，梯度聚合的网络通信占每次迭代执行时间的49.9%–83.2%（详见图4）。

  为此，我们提出了iSwitch，一种用于分布式RL训练的交换机内加速方法（见图1c）。iSwitch为可编程交换机内的梯度聚合开发硬件加速器。基于三个观察结果，提出了一种实用有效的解决方案。首先，如前所述，梯度聚合是分布式RL训练中的主要瓶颈，它会带来巨大的网络通信开销。将梯度聚合从服务器节点移动到网络交换机可以显著减少所需的网络跳数。第二，可编程交换机如今已广泛应用于数据中心。它们为开发人员编程硬件提供了灵活性和基本计算能力，从而简化了iSwitch的实现。第三，交换技术已经发展了几十年，目的是扩展集群，在交换计算中，可以利用现有的分层机架级网络结构来扩展分布式RL训练。
  iSwitch是一种通用的方法，有利于同步和异步分布式RL训练。在同步训练中，在每次迭代的梯度聚合过程中，所有worker都被阻塞。iSwitch使用in-switch加速器，减少了端到端网络通信开销，从而减少了阻塞时间。此外，由于iSwitch以网络包的粒度而不是整个梯度向量（由多个网络包组成）进行交换机内聚合，因此iSwitch进一步降低了聚合引起的同步开销。
  对于异步分布式RL训练，每个worker独立运行而不被阻塞。然而，由于训练过程的异步性，在训练worker中，以局部权重和梯度的陈旧性来换来的阻塞开销得去除，这损害了训练的收敛性并增加了训练迭代的次数。由于更快的网络通信使得worker能够提交更新的梯度，因此iSwitch改善了收敛性。因此，训练可以在较少的迭代次数下收敛。为了进一步提高异步分布式RL训练的并行性，我们重新考虑了RL训练算法，并将局部梯度计算、聚合和权值更新的执行完全流水线化。
  此外，iSwitch可以按机架级扩展分布式RL训练。它利用现有的机架级网络层次结构并将交换机内加速器集成到交换机的不同层中，以进行层次结构聚合。通过扩展可编程交换机的网络协议和控制/数据平面，iSwitch将所需的硬件成本降到最低。作为可编程交换机的扩展，iSwitch不会影响其常规网络功能。据我们所知，这是针对分布式RL训练的交换机内加速的第一项工作。总体而言，我们在本文中做出了以下贡献：

• 我们用各种典型的RL算法量化了分布式RL训练的性能特征，并表明用于梯度聚合的网络通信是分布式RL训练的主要瓶颈。
• 我们提出了交换内计算范例，并开发了交换机内聚合加速器，该加速器显着减少了每次训练迭代的梯度聚合所需的网络跳数。
• 我们将交换机内加速应用于同步和异步分布式RL训练。在新的分布式计算范例上，我们重新考虑了RL训练算法，以进一步提高其并行性。
• 我们提出了一种利用机架级集群中现有网络架构来扩展分布式RL训练的分层聚合机制。

  我们用真实的NetFPGA板实现了iSwitch。为了证明iSwitch的功效，我们训练了各种流行的RL算法，包括DQN，A2C，PPO和DDPG。我们的实验结果表明，与最新的分布式训练方法相比较，iSwitch对于同步分布式训练提供了1.72-3.66倍的系统级加速，对于异步分布式训练提供了1.56-3.71倍的系统级加速。我们的评估结果还表明，iSwitch在机架级集群中的同步和异步分布式训练上都实现了更好的可扩展性。

2 背景与动机

2.1 分布式强化学习训练

  如图2所示，标准RL设置假设一个元素通过大量的步骤与给定的环境反复交互。首先，agent从环境接收初始状态，然后根据其策略（由amodel参数化）执行操作，将当前状态映射到可能的操作集中的操作（即action←policy（state））。选定的动作在环境中生效后，将生成下一个状态，并将其与奖励一起返回给agent。这种agent与环境的交互作用一直持续到agent遇到终止状态，并且初始状态和终止状态之间的交互序列形成了一个episode。 之后，交互重新开始以生成新的episode。在产生大量episode的过程中，收集这些状态/动作/奖励以形成trajectory，然后将其用于通过基于计算的更新其模型来改进策略。agent的目标是学习一种使奖励目标最大化的政策，即episode reward，即整个episode中累积的奖励。

  众所周知，DNN训练是非常耗时的。RL训练也是如此。与DNN训练不同，RL训练需要大量的迭代，例如，用DQN算法学习Atari游戏需要2亿次迭代（见表1），而流行的DNN ResNet只需要600K次迭代，因此需要大量的训练时间。例如，在一个GPU上进行8天的DQN训练。
  为了克服这一挑战，分布式RL训练最近开始流行起来。它的关键思想依赖于多个agent，即worker，并行地探索环境，以获得模型改进的局部轨迹，即梯度。中心节点或分散的worker可以“聚合”（即梯度聚合）来自worker的那些计算出的局部梯度，以获得用于更新策略模型的完全求和的梯度。一旦政策得到改进，worker就可以为下一次训练做好准备。

2.2 同步训练和异步训练

  分布式训练中的worker可以同步运行，也可以异步运行。在同步中，在每次迭代的梯度聚合（以及权重更新和转移）期间，所有的worker都被阻塞。在异步中，允许所有的worker独立运行而不阻塞。
  我们在图3中演示了一个使用参数服务器方法的异步训练示例，其中服务器维护最新的权重，而worker独立地提取最新的权重进行本地计算。在本地计算梯度（尽管已经过时）之后，它将被推送到参数服务器以更新当前权重。通过集中式服务器，所有的worker虽然异步运行，但在一定程度上始终保持最新的权重。注意，异步训练不适用于AllReduce方法（参见图1b），因为AllReduce中的循环聚合是一个全局同步的过程。

  由于同步和异步方法提供了不同的折衷，因此它们作为分布式训练的两种主流方法共存。同步分布式训练要求worker之间的同步以进行梯度聚合，并且为每个训练迭代设置全局屏障。这种阻塞聚合停留在同步训练系统的关键路径上，显著影响每次迭代的执行时间，特别是在大规模分布式系统中。
  异步训练打破了worker之间的同步障碍，从而将阻塞开销降至最低。但是，异步的缺点是使用过时的梯度进行模型更新，这会减慢训练收敛速度，即需要更多的训练迭代次数。相反，同步训练没有陈旧性问题，因此享有更快的收敛性，即需要最少的迭代次数。
  理想情况下，我们希望对同步和异步训练都具有快速梯度聚合，以便同步训练将为聚合付出较少的阻塞开销，而异步训练将获得较新的梯度以加快收敛。在本文中，我们提出了iSwitch，它有利于同步RL训练和异步RL训练。

2.3 梯度聚合方法

  如前所述，分布式RL训练中的梯度聚合有两种主流方法：基于集中式参数服务器的方法（PS）和基于分散式AllReduce的方法（AR）。
  我们在图1中比较了这两种方法。我们在图1a中展示了PS方法，其中每个worker中的本地梯度被发送到中央服务器以执行求和，然后是权重更新。然后将更新的权重发送回所有worker以覆盖它们的本地副本，这样下一次迭代就可以开始了。图1b说明了Ring AllReduce方法，其中每个worker将其局部梯度发送到下一个邻居，以循环方式执行部分求和，直到梯度完全聚合。之后，每项工作都使用聚合的梯度来更新局部权重。

  为了便于我们的讨论，我们假设有多个worker和一个与网络交换机连接的中央参数服务器。对于PS方法，每个worker必须经历四个网络跳来完成梯度聚合。而中央服务器是瓶颈。AR方法避免了这个中心瓶颈，但需要更多的网络跳数。对于N个worker连接到交换机的情况，聚合的网络跳数为（4N−4），这与worker数量成线性关系。
  为了进一步了解它们的性能特点，我们在一个连接10Gb以太网的GPU集群中使用PS和AR两种方法运行同步分布式RL训练（参见§5.3中的详细实验设置）。我们将每个迭代的训练过程分解为多个部分：局部梯度计算（包括代理操作、环境反应、轨迹缓冲区采样、内存分配、前向传播、后向传播和GPU内存复制）、梯度聚合、权重更新和其他。我们在图4中量化它们的性能开销。可以看出，对于PS和AR方法，梯度聚合占据了每次迭代执行时间的很大一部分（49.9%–83.2%）。由于梯度聚合只涉及简单的算术运算（如求和），所以其开销主要来自网络通信。

2.4 为什么在交换机中计算

  为此，我们提出了iSwitch，这是一种交换机内计算方法，它利用可编程交换机的计算能力来减少梯度聚合开销。如图1c所示，iSwitch只需要两个网络跳（即从工作节点到交换机，以及从交换机到工作节点）来完成梯度聚合。iSwitch将网络跳数至少减少一半，从而为分布式RL训练的每次迭代提供更低的端到端通信时间。

  我们利用可编程交换机来实现交换机内计算方法，以加速分布式RL训练，这有以下三个原因。首先，可编程交换机如今已普及，在现代数据中心或机架级的集群中，可编程交换机已成为骨干技术，使开发人员可以定义自己的功能来进行网络数据包处理。其次，它们为开发人员提供了对硬件进行编程的灵活性，从而简化了iSwitch的实现。可编程交换机具有控制平面和数据平面。控制平面负责网络管理，而数据平面负责数据传输（即数据包转发）。我们通过扩展控制平面和数据平面而不影响常规网络功能来开发iSwitch。第三，该交换机固有地实现了可扩展性。例如，交换机已被广泛用于扩展数据中心中的群集大小。在此项目中，我们利用典型数据中心的现有网络体系结构来扩展机架级集群中的分布式RL培训。

3 设计与实现

  iSwitch的目标是通过减少网络通信开销、提高并行性和可扩展性来减少分布式RL训练的端到端执行时间。

3.1 挑战

  如前所述，利用可编程交换机进行梯度聚合有利于分布式RL训练。然而，我们必须克服三个挑战：

• 首先，可编程交换机最初是为包转发而设计的。我们的交换机内计算需要在不影响常规网络功能的情况下，实现交换机和工作节点之间的点到点通信以进行梯度聚合（§3.2）。
• 第二，可编程交换机的计算逻辑和片上存储器对于我们的加速是有限的。因此，我们的设计应该简单高效，以满足性能要求（§3.3）。
• 第三，随着机架级集群中工作节点和交换机数量的增加，我们提出的交换机计算应该能够扩展到分布式RL训练（§3.4）。

3.2 网络协议扩展

  为了支持用于分布式RL训练的交换机计算，我们建议基于常规网络协议构建自己的协议和数据包格式。图5演示了iSwitch中控制和数据包的格式。我们使用IP报头中的服务类型（ToS）字段来识别具有我们专用协议的数据包。ToS是IP协议标头中的一个1字节字段，用于区分不同的IP流优先级。我们使用保留的ToS值标记属于交换机内RL训练的数据包。为了区分iSwitch中的控制数据包和数据包，我们使用不同的ToS值。

  控制消息：如图5a所示，用保留的ToS值标记，控制消息包在UDP头之后有一个1字节的强制操作和一个可选的Value有效载荷。 在Action字段中，我们为分布式RL训练的基本操作定义了多个唯一的Action代码（请参见表2）。

  对于某些操作，我们将使用Value字段。具体而言，对于Join消息，Value字段可用于关于训练模型的元数据。此外，对于SetH消息，Value字段用于指定在广播结果之前需要聚合多少个梯度向量（即聚合阈值H）。默认情况下，H等于worker数量。
  数据信息： 图5b描绘了数据分组的格式。类似于控制分组，数据分组也被标记有保留的ToS值。它的UDP有效负载以8字节的Seg字段开头，以指示已传输数据包的索引。每个Seg号对应于梯度向量中的空间偏移，并且来自具有相同Seg号的数据包的梯度数据将被汇总。除了Seg字段外，其余有效负载空间（受以太网帧大小限制，即通常为1,522字节）填充有梯度数据。此外，为了提高数据处理的效率，所有梯度数据都以原始浮点格式在iSwitch中进行传输和计算。

3.3 数据平面和控制平面扩展

  在iSwitch中，我们将in-switch加速器设计并集成到数据平面中，作为“bump-in-the-wire”组件。
  数据平面扩展：我们在图6中展示了以太网交换机的扩展数据平面。传入的网络数据包由PHY收发器（PHY）和以太网媒体访问控制（ETH MAC）接收，然后存储在Rx队列中以进行进一步处理。输入仲裁器始终选择一个非空的Rx队列以按优先级顺序提取数据包，并将选定的数据包送入数据包处理模块。之后，提取、解析数据包的报头信息，并将其与查找表中的不同转发规则进行比较，以进行目的地识别。然后，数据包将被分派到其相应的出口Tx队列，在此处最终将通过以太网MAC和PHY收发器传输数据包。

  为了实现我们的交换机内加速，我们增强了输入仲裁器的功能，以便它可以根据其ToS字段检测已标记的数据包并将其传送到加速器，而不是传送到原始的数据包处理模块。并且，输入仲裁器将交换机内加速器的输出视为来自入口Rx队列的输出，以便可以将梯度聚合的结果作为常规流量发送到工作节点。
  交换机内加速器设计： 为了最大化数据级并行度，我们的交换机内加速器以“burst”的粒度处理每个数据包，该粒度是指内部总线可以在单个时钟周期（例如256位）中传递的数据。因此，每个数据分组被分为多个burst以进行处理和计算。
  我们在图7中描述了交换机内加速器的体系结构。当数据包传入时，分离器会将报头的burst与有效负载的burst分开。报头burst（包括以太网、IP、UDP和iSwitch协议字段）将被被送到Seg解码器。包括梯度矢量的有效载荷burst将被送到累积循环中。Seg解码器将提取段号，并将其传递给段计数器和地址生成器。分段计数器通过为每个分段分配一个聚合计数器来跟踪梯度分段的聚合状态。计数器针对每个聚合的梯度段递增，直到达到指定的“聚合阈值”H。

  在聚合的过程中，每个有效负载burst将被切片器切片成单独的32位浮点元素，并送到加法器。加法器并行计算，并继续将传入的有效载荷burst与存储直接聚合的梯度数据的缓冲区提供的累加段数据相加。为了使总和数据对准相同的Seg编号和burst偏移，采用了Addr Generator来动态生成缓冲区地址。该过程一直持续到段完全聚合为止。然后，(1)计数器重置；(2)使用多路复用器将缓冲区写回零；(3)通过连接结果，在报头前添加并发送包含聚合段的数据包来触发输出模块的输出。
  除了在加速器中对每个数据包进行细粒度处理之外，iSwitch还以网络数据包的粒度进行梯度聚合。与传统方法（请参见图8a）不同，在传统方法中，它们必须在求和操作之前等待整个梯度向量的到达，iSwitch会在收到具有相同Seg编号的数据包后立即开始计算。这种动态聚合方法隐藏了求和操作和数据传输的开销，从而进一步减少了梯度聚合的等待时间。

  控制平面扩展：为了支持具有inswitch加速功能的分布式训练，我们扩展了控制平面，为当前训练工作中涉及的worker和交换机维护了一个轻量级成员表。如图9所示，成员资格表记录每个参与的worker/switch的ID号（每个成员资格条目的唯一编号）、IP地址、UDP端口号、类型以及网络类型中相应的父ID。成员资格表中的条目可以使用控制消息进行更新，例如表2中的Join和Leave。这些信息可以被数据平面用于数据收集、计算、转发和广播。除了维护成员资格表外，控制平面还管理交换机内加速器的初始化、配置和重置。这可以通过表2中的Reset和SetH之类的控制消息来实现。
  控制平面也有助于处理丢失的数据包，尽管它在集群环境中并不常见，开销最小。具体来说，我们将处理有损数据包的大部分任务卸载给了worker，并且只在交换机中留下简单的任务，例如接受/转发控制消息（例如，FBcast和Help）。

3.4 扩展交换机内计算

  我们已经讨论了在工作节点机架内进行分布式RL训练的交换机内加速。现在，我们讨论如何扩展机架式集群或数据中心中的交换机内计算。图10显示了典型集群或数据中心的网络架构。同一机架中的所有服务器均通过具有10Gb以太网的机架顶交换机（ToR）连接。在较高级别上，存在具有更高网络带宽（例如40Gb至100Gb）连接的聚合交换机（AGG）和核心交换机（Core）。

  为了通过机架式集群中的iSwitch扩展分布式RL训练，我们开发了一种“分层聚合”方法。具体来说，如果某个交换机完成了对存储在缓冲区中的梯度矢量中某个段的本地聚合，它将把聚合后的段转发给更高级别的交换机以进行全局聚合。如果较高级别中有多于一台交换机，则该交换机将选择IP地址值最小的一台交换机，以便最终将所有梯度数据汇总到Core交换机中。然后，将全局聚合的梯度广播到较低级别的交换机，以进行进一步分发。这样的设计利用了现有的机架级网络架构，并且不会引入其他硬件或网络拓扑更改。

3.5 实现

  我们用真实的NetFPGA-SUME板来实现iSwitch。NetFPGA-SUME有一个x8 Gen3 PCIe适配卡，其中包含Xilinx Virtex-7 FPGA和四个10Gbp以太网端口。我们使用NetFPGA社区提供的参考交换机设计进行进一步开发。为了充分利用其内部AXI4-Streambus的位宽（即256位/周期），我们采用八个32位浮点加法器进行并行梯度聚合。我们的交换机内加速器已集成到此参考交换机设计中，并使用标准256位AXI4-Stream总线（频率为200MHz）与其他组件进行交互。在片上资源利用率方面，与之相比，iSwitch加速器消耗了额外的18.6％的查找表（LUT）、17.3％的触发器（FF）、44.5％的Block RAM（BRAM）和17个DSP Slice。 请注意，iSwitch硬件和网络协议的实现对于同步和异步分布式训练都是通用的。

4 对强化学习的影响

  在本节中，我们讨论了如何利用交换中的计算范例来促进硬件/算法的协同设计，并进一步提高同步和异步分布式RL训练的性能。
  如第2.2节中所述，对于同步训练，我们可以直接使用iSwitch通过用交换机内聚合替换聚合操作（例如AllReduce操作）来减少梯度聚合的端到端执行时间。讨论细节（有关性能优势，请参见第6.1节）。对于异步训练，iSwitch提供了新的优化空间，以利用交换机内计算范例改善训练并行性，这展示了iSwitch对分布式RL训练的影响的一个有用案例。我们将在本节中详细讨论。

4.1 重新思考异步强化学习训练

  异步分布式训练的常规方法（请参见图3）依赖于中央参数服务器来维护最新的权重，每个worker与中央服务器进行交互以获取最新的权重，以使训练可以收敛。为了从iSwitch获得收益，一种直接的方法是将参数服务器的功能转移到网络交换机。但是，这将大大增加硬件成本，因为在参数服务器上运行的任务不仅需要大量的计算资源，而且还需要很大的存储空间来存储权重和历史更新。通过交换机内聚合，我们重新考虑了异步分布式训练，并提出了两种优化技术来进一步分散训练并提高其并行性。

  分散权重储存。我们没有将梯度推送到中央服务器，而是通过交换机聚合来自异步worker的梯度，然后将求和的梯度广播给每个worker，以便在每次迭代中更新权重。由于我们在所有worker中初始化相同的模型权重，并且广播相同的聚合梯度，因此尽管异步训练，权重的分散存储总是在迭代过程中达成一致。
  三级流水线。我们将训练迭代中的三个关键阶段解耦：(1)本地梯度计算（LGC），(2)梯度聚合（GA）和(3)本地权重更新（LWU）。 第一阶段发生在workers中，它执行环境交互、轨迹收集和梯度生成以及上传。 第二个是在交换机中，它执行梯度收集、求和和广播。 第三是对worker进行权重更新。
  对于训练迭代中的三个阶段，我们可以将它们流水线化，以提高分布式训练的并行性，如图11所示。在LGC阶段，每个worker独立运行，而不与其他worker或交换机同步，并不断地将计算出的梯度上传到交换机。在GA阶段，交换机以异步方式聚合梯度，并不断聚合传入的梯度。一旦接收到足够的梯度向量，聚合的梯度就被广播回工作者，以便LWU阶段开始。这种方法鼓励速度更快的worker为聚合贡献更多，而速度较慢的worker在不妨碍训练的情况下贡献更少。

  由于训练过程的不同步性，不可避免地会出现权重和梯度的陈旧性，从而影响训练的收敛速度。在这项工作中，我们建议明确地限制梯度的陈旧性。具体地说，我们检查每个worker上的局部梯度的陈旧性，并且只在交换机的边界内提交陈旧的梯度。我们展示了算法1中的详细过程。我们用经验评估和理论推导证明了我们提出的异步训练的收敛性（请参阅§6.2），如下所示。

4.2 收敛性证明

  为了证明异步iSwitch的收敛性，我们将其转换为基于经典参数服务器的异步训练。通过证明前者在数学上等同于后者，我们得出了与《More Effective Distributed ML via A Stale Synchronous Parallel Parameter Server》《Slow Learners are Fast》中相同的结论，但是常数发生了变化。
  具体来说，我们假设异步iSwitch中有一个虚拟参数服务器（请参阅算法1），该服务器存储最新的权重并按照经典设计的方式执行权重更新。这样的参数服务器等效于每个worker上的LWU线程。如所讨论的，所有worker在迭代中执行相同的权重更新，因此，分散的约定权重可以被视为存储在单个中央服务器上。因此，梯度推送、聚合和广播可以减少到与参数服务器的上游通信，而每个worker上LGC线程中的权重复制可以减少到参数服务器的下游通信。所有worker异步并行运行，以将梯度（通过交换机）推送到参数服务器以执行更新，然后更新的权重将用于新的迭代中。我们的方法与《More Effective Distributed ML via A Stale Synchronous Parallel Parameter Server》《Slow Learners are Fast》之间的微小区别在于梯度向量的集合，我们的可以减少使用较大的批次大小进行训练，而不会改变收敛速度。因此，我们提出的异步训练可以简化为传统方法《More Effective Distributed ML via A Stale Synchronous Parallel Parameter Server》《Slow Learners are Fast》，并通过随机梯度下降提供$O(T^{-0.5})$的收敛速度，其中T为训练迭代次数。

5 评估方法

5.1 分布式强化学习训练基准（Benchmarks）

  为了评估iSwitch的训练性能，我们使用了四种流行的RL算法作为我们的基准。基于他们的单节点训练代码，我们通过遵循最先进的分布式训练方法，为每个基准开发了三种参考设计：基于同步和异步参数服务器的训练（Sync/Async PS）和基于AllReduce的训练（AR）。我们的参考设计是高度优化的，与使用MPI的OpenAI基线相比（社区中使用的一种流行基线），表现出约10%的性能提升和更高的训练回报。我们将这些RL算法列示如下：

• DQN是街机游戏中最流行的RL算法之一。当应用于玩Atari游戏集的任务时，它的模型大小是6.4 MB，我们从中选择了经典游戏“Pong”。
• A2C是另一种流行的RL游戏算法。当应用到Atari游戏集时，它的模型大小是3.3MB，我们从中选择了一个不同但经典的游戏“Qbert”。
• PPO是一种较新的算法，主要用于模拟机器人运动。将其应用于仿真环境set MuJoCo中的机器人控制时，其模型大小为40KB，从中我们选择了一个经典环境“Hopper”。
• DDPG是另一种连续控制算法。当应用于MuJoCo的机器人控制任务时，DDPG的双模总共是157.5kb，我们从中选择了另一个经典环境“HalfCheetah”。

  我们使用最先进的库实现所有参考设计：PyTorch 1.0、CUDA9.2、CuDNN7.2.1、GYM和OpenMPI 3.1.4。对于iSwitch设计，我们使用与参考设计相同的代码和库，但使用不同的梯度聚合方法，即在交换机聚合中，以及在异步iSwitch中使用双线程训练（参见算法1）。

5.2 指标及其定义

  我们对每个基准测试使用多种训练方法：同步参数服务器（PS）、AllReduce（AR）、iSwitch（iSW）以及异步参数服务器（Async PS）、iSwitch（Async iSW）。我们使用以下指标评估所有方法：

• 最终平均奖励：最后10个episode的平均episode reward（见§2.1），这是RL训练评估中使用的标准指标。
• 迭代次数：完成端到端训练所需的训练迭代次数。对于同步训练方法，它可以在任何worker上测量。对于异步训练方法，可以在PS的参数服务器或iSW的LWU线程上通过计算权值更新的次数来精确测量。
• 每次迭代时间：两次连续迭代之间的平均时间间隔。对于同步方法，它是一个训练迭代的延迟。对于异步方法，它可以通过PS的参数服务器或iSW的LWU线程上两个连续的权值更新操作之间的时间间隔来精确测量。
• 端到端训练时间：针对不同方法的每个基准达到相同水平的“最终平均奖励”所需的总训练时间。

5.3 实验设置

  主群集设置。为了测量实际时钟时间下的训练性能，我们设置了一个由四个节点组成的主集群。每个节点都有一个Vidia Titan RTX GPU和一个IntelXeon CPUE5-2687W@3GHz，我们使用这个四节点集群来评估AR和iSW方法。为了支持PS方法，我们使用了一个额外的节点作为参数服务器。所有节点都通过Intel X540T2 10Gb以太网NIC连接到Netgear 10Gb以太网交换机。考虑到RL模型传输梯度的小尺寸，例如PPO为40KB，我们在实验中不考虑支持更大的网络连接（即40∼100Gbps）。至于iSW方法，我们用NetFPGASUME板替换网络交换机，并充分利用NetFPGASUME板上的四个以太网端口来连接工作节点。
  可扩展性实验设置。对于可扩展实验，我们在由两层常规交换机组成的集群中使用更多工作节点来模拟所有方法的训练性能，如图10所示。具体来说，集群有一个根交换机连接到多个“机架”，每个机架包含三个工作节点（由于NetFPGA板的端口限制）。我们对集群中iSwitch的分层聚合进行仿真，开发了三个目标的仿真：仿真聚合必须具有(1)准确的网络跳数，(2)网络链路中尽可能多的通信量，以及(3)准确的加速器开销。我们通过将合成的梯度数据从每个工作节点传输到其第三个相邻的工作节点来实现这些目标，使得每个梯度消息总是穿过交换机的层次结构。之后，在worker之间设置屏障以捕获最慢的梯度转移，从而可以认为聚合完成。这种仿真方法与(1)和(2)的实际聚合相匹配，尽管对交换机之间的网络通信量有轻微的放大。为了实现目标(3)，我们测量硬件加速器开销并将其添加到聚合时间中。对于局部计算的仿真，我们使用了来自PS/AR方法的相同轨迹，并将其应用于iSwitch以进行公平比较。此外，我们还得到了iSwitch所需的“迭代次数”。对于同步训练，iSwitch与PS/AR具有相同的迭代次数，这是因为它们在分布式训练中的数学等价性（见表4）。对于异步训练，iSwitch所需的迭代可以通过控制同步训练方法中过时梯度的使用来仿真，其中iSwitch的陈旧性可以通过每个训练迭代中三个阶段（见图11）的测量时间比来计算。因此，我们相信该仿真平台能够反映支持交换计算的真实机架级集群的可扩展性。

6 评估

  我们使用主集群评估四个基准的训练性能（请参阅第5.1节）。我们测量“端到端训练时间”，并总结表3中的性能提升。在同步训练设置中，使用iSwitch方法（iSW） 与其他方法相比，它具有很大的余量，并且与基准设计（PS）相比，性能提高了1.72–3.66倍。尽管AR方法也提供了DQN和A2C的改进，但PPO和DDPG的性能实际上比PS差一些。至于异步训练设置，与基础PS相比，iSwitch的优势仍然存在，并且性能提升为1.56-3.71倍。请注意，我们分别评估同步（§6.1）和异步（§6.2）分布式训练方法的性能，因为这项工作的主要目标是加速和支持两种方法，而不是像§2.2中讨论的那样对它们进行比较。

6.1 同步iSwitch的优点

  为了理解iSwitch在同步训练设置下的性能改进，我们将iSwitch的“每次迭代时间”与PS和AR在图12中的四个基准上进行了比较。我们还提供了不同方法的“每次迭代时间”的详细时间细分。这一结果表明，与PS相比，iSW缩短了41.9%-72.7%的“每次迭代时间”，这是因为四个基准的梯度聚合时间缩短了81.6%-85.8%。

  iSwitch在梯度聚合方面提供了显着的加速，这有以下三个原因。 首先，与PS设计（四跳网络）相比，iSwitch中的聚合过程仅需要一半的网络跳（两跳网络），从而实现了端到端通信延迟降低到一半。 其次，iSwitch具有动态聚合的独特功能（如图8所示），可将聚合粒度从梯度矢量大小（即基线设计中的模型大小）减小到网络数据包大小，在相同索引的数据包到达后，iSwitch无需等待所有梯度向量的到来才开始计算，而是立即进行聚合（请参见图8），这减少了由梯度聚合引起的同步开销。 第三，iSwitch通过为每个工作节点分配专用的网络链路来提供均衡的通信，从而消除了PS设计中由中央链路引起的瓶颈。

  除了与基线设计（PS）进行比较外，我们还将iSwitch与另一种主流方法进行比较——基于AllReduce的训练（AR），它提供了均衡的通信。图12中的结果表明，在所有四个基准测试中，iSwitch的性能仍胜过AR，即“每次迭代时间”减少了36.7％–48.9％。这些改进仍然归因于iSwitch的加速梯度聚合，即与AR相比，iSW的聚合时间减少了63.4％–87.9％。如第2.3节中所述，PS和AR之间存在性能折衷。与PS相比，AR方法遭受的网络跳数更多，但它消除了由中央参数服务器引起的瓶颈。同时，基准测试因其模型大小而要求不同的通信/计算负载（请参阅第5.1节）。结果，与PS相比，AR在DQN和A2C中的性能更好，但在PPO和DDPG中的性能更差。iSwitch的运行速度比PS和AR都要快，这是因为它减少了端到端的网络延迟以及动态聚合。

  此外，我们在表4中显示了详细的结果，包括迭代次数、绝对训练时间和已获得的训练奖励。我们观察到，所有同步方法都训练相同的“迭代次数”，以达到每个基准相同的“最终平均奖励”。

  为了证明所有同步方法的加速和训练奖励的协同作用，我们评估了所有基准的实际训练曲线，并在图13中展示了DQN的案例研究。

6.2 异步iSwitch的优点

  现在，我们将所有基准测试的iSwitch与异步基准（Async PS）进行比较。为了进行公平的比较，我们为两种方法给出了相同的陈旧度界限（S = 3）（尽管传统的Async PS方法不涉及陈旧性控制机制），使得梯度的陈旧性范围介于0到3次迭代。
  我们在表5中总结了这两种方法的训练性能。我们观察到iSwitch（Async-iSW）提供了更快的收敛速度，即与基线（Async-PS）相比，“迭代次数”减少了44.4%–77.8%。这是由于iSwitch中梯度的平均过时性较小，尽管两种方法的最大过时性相同。iSwitch中的加速梯度聚合的优点可以缓解梯度的陈旧性，因为更快的梯度聚合导致更早/及时的权值更新，从而为下一次迭代提供更新鲜的权值和梯度。另一方面，Async PS遭受双倍的端到端通信延迟（如§6.1所述）以及繁重的中央网络链路，因此增加了梯度/权重通信时间。结果，梯度的陈旧性变大，导致训练迭代次数增加。

  从表5中，我们还观察到iSwitch显示，与基线相比，DQN和A2C的“每次迭代时间”缩短了4.6%–51.5%。这是因为asynchronous iSwitch不仅享受了梯度聚合加速的好处，而且利用流水线训练来隐藏部分执行时间（见图11），特别是加速梯度聚合和权值更新。相比之下，Async PS仍然为长的通信延迟付费，因此增加了两个连续权重更新之间的时间间隔，即更大的“每次迭代时间”。

  请注意，对于PPO和DDPG，iSwitch在“每迭代时间”中未显示改进。这主要是由于PPO和DDPG中梯度聚集时间的比例相对较小。 因此，即使使用流水线，梯度聚合的隐藏时间也只能稍微减少“每次迭代时间”，其有限的优势无法胜过异步PS。 但是，iSwitch的加速梯度聚合可减少梯度的陈旧性，并减少训练迭代的次数。
  为了结合iSwitch方法在减少“迭代次数”和改善“每次迭代时间”方面的有效性，我们在表5中显示了“端到端训练时间”。异步iSwitch比基线异步PS减少了35.7％–73.0％ “端到端训练时间”。

  此外，为了展示两种异步方法在加速和训练奖励方面的协同作用，我们评估了所有基准测试的实际时钟训练时间，并在图14中演示了DQN示例。

6.3 iSwitch的可扩展性

  为了评估可扩展性，我们按照第5.3节中的可扩展性实验设置来测量和比较所有训练方法的端到端训练的加速情况。图15展示了4、6、9和12个工作节点的训练PPO和DDPG的可扩展性案例研究。对于同步分布式训练，如图15a和15c所示，我们观察到AR方法在集群扩展时提供了最差的加速，这是因为它的用于梯度聚合的网络跳数在群集大小上是线性的，如第2.3节所述。 PS方法显示出第二好的可扩展性，但是它在通信和计算方面都存在中心瓶颈，随着我们增加工作节点的数量，这个缺点变得更加严重。 由于以下三个主要原因，iSwitch的性能优于AR和PS：(1)所需的网络跳数最少；(2)分层聚合中均衡且减少的流量负载；(3)iSwitch的交换机内加速器 。

  对于异步分布式训练（请参见图15b和15d），我们观察到异步PS方法不能优于异步iSwitch方法，因为异步PS仍需要更多的网络跃点，尽管异步机制在某种程度上减轻了中央瓶颈。相比之下，Async iSwitch拥有最佳的可扩展性（即，几乎线性的加速），因为它不仅具有交换机内计算所带来的上述优势，而且还具有三级流水线的优势以及加速后的陈旧性聚合。

7 相关工作

  分布式RL训练。机器学习界最近的研究提出了几种分布式RL算法，如A2C和A3C。然而，这些工作主要集中在单节点解决方案上。在本文中，我们的目标是分布式RL训练，涉及多个节点在机架规模。系统社区也对开发分布式RL系统感兴趣，比如Gorila、Ray和Horizon。我们对交换机加速度的贡献与这些工作是正交的。iSwitch可以集成到这些系统中，以进一步提高其性能。
  最近，研究人员提出采用AllReduce技术来降低分布式训练系统的通信开销，因为AllReduce，特别是Ring AllReduce是一种带宽最优的通信方法，适合于带宽要求很高的工作负载，如分布式DNN训练。然而，分布式RL训练对延迟更为敏感。此外，由于AllReduce的全局同步特性，它只支持同步训练。iSwitch支持同步和异步训练。
  在网计算。研究人员提出了几种与高性能计算（HPC）社区中的在网计算相关的解决方案，例如BlueGene Network ，PERCS Interconnect和CM-5。但是，这些工作是使用定制的互连架构为高性能计算应用程序开发的。相比之下，iSwitch设计的商品可编程交换机已在当今的数据中心得到广泛采用。
  最近的研究网络系统已经用可编程交换机（如P4和PISCES）实现了软件定义的网络（SDN）。SDN使平台运营商能够在可编程交换机中实现自己的网络策略，以便于网络管理。然而，以往的工作都没有研究分布式机器学习工作负载中的交换机加速器技术。iSwitch利用可编程交换机的可编程性在交换机内部构建一个加速器，该加速器对分组有效负载进行计算，以促进分布式RL训练的梯度聚合。
  机器学习的硬件加速器。 学术界和工业界的最新进展为机器学习提供了各种专用的硬件加速器。例如，TABLA使用FPGA加速器来提高训练工作量的性能。Google提出了TPU，它是具有脉动阵列架构的ASIC加速器。CosMIC提出了一种使用多个FPGA和ASIC加速器的分布式机器学习训练系统。 然而，大多数这些工作仅专注于DNN的训练或推理，而没有考虑RL。重要的是，这些工作仅加速了机器学习算法的本地计算，而没有触及分布式训练系统中的网络瓶颈。iSwitch致力于通过提出的交换机内计算来降低网络开销。

8 结论

  本文对分布式RL训练中的性能开销进行了初步的量化研究，提出了一种交换机内计算范例iSwitch，通过提供一整套解决方案来消除网络瓶颈：(1)一个交换机内聚合加速器来减少端到端的通信开销；(2)对同步和异步分布式RL训练的加速支持，具有改进的并行性；以及(3)机架级集群的分层设计，以扩展分布式RL训练。在不同RL工作负载下的实验表明，iSwitch与现有的方法相比，在同步分布式训练和异步分布式训练中，系统级的加速比分别高达3.66倍和3.71倍。

这是个开始，慢慢享受，加油！