Menger的大规模强化学习

Menger的大规模强化学习
2022年11月16日 17:20 科技实记

在过去的十年中,强化学习(RL)已成为机器学习中最有前途的研究领域之一,并且在解决复杂的现实世界问题(例如芯片放置和资源管理)以及解决具有挑战性的游戏(例如,围棋,Dota的2,和隐藏和寻求)。用最简单的术语来说,RL基础结构是数据收集和培训的循环,参与者在其中探索环境并收集样本,然后将样本发送给学习者训练和更新模型。当前大多数RL技术都需要对环境中的数百万个样本进行多次迭代以学习目标任务(例如,Dota 2 每2秒从200万帧的频次中学习)。这样,RL基础设施不仅应有效地扩展规模(例如,增加演员人数)并收集大量样本,而且还应能够在训练过程中快速迭代这些大量样本。

演员在其中向学习者发送轨迹(例如多个样本)的RL系统的概述

学习者使用采样数据训练模型,并将更新后的模型推回给参与者(例如TF-Agents,IMPALA)。

今天,Google介绍Menger,这是一个大规模的大规模分布式RL基础结构,具有本地化推断,可在多个处理集群(例如Borg单元)中扩展至数千个参与者,从而减少了芯片放置任务的总体培训时间。在这篇文章中,Google描述了Google如何使用Google TPU加速器实现Menger进行快速训练迭代,并针对挑战性的芯片放置任务展示了其性能和可扩展性。门格尔减少达8.6训练时间X与基线相比实现。

Menger系统设计

有多种分布式RL系统,例如Acme和SEED RL,每个方法都专注于优化分布式强化学习系统空间中的单个特定设计点。例如,尽管Acme在每个参与者上使用本地推理并从学习者那里频繁获取模型,但SEED RL通过分配一部分TPU内核来执行批量调用而得益于集中推理设计。这些设计点之间的权衡是(1)支付向/从集中式推理服务器发送/接收观察和动作的通信成本,或支付从学习者那里获取模型的通信成本,以及(2)对参与者的推理成本(例如CPU)与加速器(例如TPU / GPU)相比。由于Google目标应用程序的要求(例如观察值,操作和模型大小),Menger以类似于Acme的方式使用局部推断,但实际上将参与者的可扩展性推向了无限的极限。

实现大规模可扩展性和加速器快速培训的主要挑战包括:

1.服务从参与者到学习者的大量读取请求以进行模型检索可以轻松限制学习者,并随着参与者数量的增加而迅速成为主要瓶颈(例如,显着增加收敛时间)。

2.在将训练数据馈送到TPU计算核心时,TPU性能通常受到输入管道效率的限制。随着TPU计算核心数量的增加(例如TPU Pod),输入管道的性能对于整个培训运行时变得更加关键。

高效的模型检索

为了解决第一个挑战,Google在学习者和Actor之间引入了透明且分布式的缓存组件,这些组件在TensorFlow中进行了优化,并由Reverb(在Dota中使用的类似方法)支持。缓存组件的主要职责是在参与者的大量请求和学习者的工作之间取得平衡。添加这些缓存组件不仅显着减轻了学习者为读取请求提供服务的压力,而且还以少量的通信开销在多个Borg单元之间分配了参与者。在Google的研究中,Google表明,对于具有512个参与者的16 MB模型,引入的缓存组件将平均读取延迟降低了约4.0倍。x导致更快的训练迭代,尤其是对于基于策略的算法(例如PPO)。

具有分布在不同Borg单元中的多个参与者的分布式RL系统的概述。为来自不同Borg单元的大量参与者的频繁模型更新请求提供服务会限制学习者以及学习者与参与者之间的通信网络,从而导致总体收敛时间显着增加。虚线表示不同机器之间的gRPC通信。

引入的透明和分布式缓存服务的分布式RL系统概述,其中多个参与者放置在不同的Borg单元中。学习者仅将更新的模型发送到分布式缓存服务。每个缓存服务都处理来自附近参与者(即,位于相同Borg单元上的参与者)和缓存服务的模型请求更新。缓存服务不仅减轻了学习者为模型更新请求提供服务的负担,而且减少了参与者的平均读取延迟。

高吞吐量输入管道

为了提供高吞吐量输入数据管道,Menger使用Reverb,这是一种最近开放源代码的数据存储系统,专门为机器学习应用程序而设计,它提供了一个高效灵活的平台,可以在各种策略/关闭状态下实现体验重放策略算法。但是,使用单个混响重放缓冲区服务当前无法在具有成千上万个参与者的分布式RL设置中很好地扩展,并且就参与者的写入吞吐量而言,效率变得很低下。

具有单个重放缓冲区的分布式RL系统。处理来自actor的大量写请求会限制重播缓冲区并降低其总体吞吐量。另外,随着我们将学习者扩展到具有多个计算引擎(例如TPU Pod)的设置,从单个重放缓冲区服务向这些引擎提供数据的效率变得很低,这对总体收敛时间产生了负面影响。 为了更好地了解分布式设置中重播缓冲区的效率,Google评估了16 MB至512 MB各种有效负载大小和16到2048个参与者之间的平均写入延迟。演员被放置在同一个Borg单元中。随着参与者数量的增加,平均写入延迟也显着增加。扩大行动者的数目为16〜2048,由〜6.2倍的平均写入延迟增加X和〜18.9 X为有效载荷大小16 MB和512 MB,分别。写等待时间的这种增加会对数据收集时间产生负面影响,并导致整体训练时间效率低下。

当actor和重放缓冲区放在相同的Borg单元上时,对于各种有效负载大小(16 MB-512 MB)和各种actor数目(16至2048),单个混响重放缓冲区的平均写入延迟。 为了减轻这种情况,Google使用Reverb提供的分片功能来增加参与者,学习者和重播缓冲区服务之间的吞吐量。分片可在多个重播缓冲区服务器之间平衡来自大量角色的写入负载,而不是限制单个重播缓冲区服务器,并且还可以使每个重播缓冲区服务器的平均写延迟最小化(因为更少的角色共享同一服务器)。这使Menger可以有效地扩展到多个Borg单元中的数千个演员。

分布式RL系统分片重播缓冲区。每个重播缓冲区服务都是专用的数据存储区,用于存储演员,通常位于同一Borg单元上。此外,分片重播缓冲区配置为加速器内核提供了更高吞吐量的输入管道。

案例研究:芯片放置

Google研究了Menger在大型网表的复杂芯片放置任务中的优势。使用512个TPU芯,门格尔实现了在训练时间显著改进(达〜8.6 X,减少从〜训练时间8.6小时下来仅仅一小时最快配置)相比,较强的基线。尽管Menger针对TPU进行了优化,但性能提升的关键因素是架构,Google希望在量身定制GPU时看到类似的提升。

使用门格尔与TPU的可变数目在训练时间的改善内核相比,在任务的基准芯片放置。 Google认为,Menger基础设施及其在复杂的芯片放置任务中所取得的令人鼓舞的结果表明了一条创新的道路,可以进一步缩短芯片设计周期,并且不仅有可能在芯片设计过程中实现进一步的创新,而且还具有其他挑战性的现实世界任务。

财经自媒体联盟更多自媒体作者

新浪首页 语音播报 相关新闻 返回顶部