当前位置: 首页 > 科技观察

每秒处理240万帧游戏画面,AI训练成本降低80%,谷歌开源RL并行计算框架

时间:2023-03-17 13:33:45 科技观察

本文经AI新媒体量子位授权转载(公众号ID:QbitAI),转载请注明出处转载请联系出处。土豪们最痛苦的事情莫过于硬件多了却达不到1+1=2的效果。AI训练中的并行计算就是这样。就算你有一千个GPU,也达不到单机训练一千倍的效果。最近,不缺钱的谷歌开源了一个SEEDRL框架,可以更轻松地在数千台机器上运行AI训练,效果比之前的方法提高了近四倍。如果不缺钱,在云端进行大规模并行计算,可以节省80%的训练成本。考虑到一个大规模的AI模型要花费数百万的训练成本,着实令人印象深刻。在训练AI踢足球这件事上,SEEDRL的处理速度可以达到每秒240万帧。如果以60fps计算,相当于每秒处理11个小时的游戏画面。SEEDRL架构上一代分布式强化学习代理IMPALA,其架构包括两部分:Actor和Learner。Actor通常在CPU上运行,并在环境中采取的步骤之间进行迭代,并对模型进行推理以预测下一个动作。Actor经常更新推理模型的参数,在收集到足够数量的观察后,将观察和动作的轨迹发送给Learner,以优化Learner。在此架构中,学习者使用来自数百台机器的分布式推理输入在GPU上训练模型。但是IMPALA有很多缺点:1.使用CPU进行神经网络推理效率低下。而且随着模型越来越大,计算量越来越大,问题会越来越严重。2.Actor和Learner之间的模型参数带宽成为性能瓶颈。3、资源利用效率低下。参与者在环境和推理任务之间交替。但是这两个任务的计算需求不同,很难在同一台机器上充分利用资源。SEEDRL架构解决了这些缺点。Actor可以在GPU、TPU等AI硬件加速器上完成推理,通过保证模型参数和状态保存在本地来加速推理,避免数据传输瓶颈。与IMPALA架构相反,SEEDRL中的参与者仅在环境中执行操作。学习者在硬件加速器上使用来自多个参与者的批量数据来集中执行推理。SEEDRL使用gPRC框架的网络库来保持低延迟,同时在每个环境步骤向学习者发送观察结果。这使SEEDRL能够在单台机器上实现每秒高达一百万次查询。Learner可以扩展到数千个core,actor的数量可以扩展到数千台机器,从而达到每秒百万帧的训练速度。SEEDRL使用两种最先进的算法:V-trace和R2D2。V-trace负责从采样的动作中预测动作的分布,而R2D2负责根据其预测的未来值选择动作。V-trace是一种基于策略梯度的方法,最先被IMPALA采用。由于Actor和Learner是异步执行的,因此V-trace在异步架构中运行良好。第二种算法是R2D2,这是DeepMind用来将强化学习代理应用到Atari游戏中的一种Q-learning方法。在上好4倍,并在52场比赛中超越人类表现。这种方法允许Q-learning算法在大规模硬件上运行,同时仍然使用RNN。实验结果Google对GoogleResearchFootball进行了基准测试,这是DeepMind最近研究的一个开源足球游戏项目。使用64个CloudTPU内核实现了每秒240万帧的数据传输速度,与之前最先进的分布式IMPALA相比提高了80倍。IMPALA需要14000个CPU才能达到相同的速度,而SEEDRL仅使用4160个CPU。同样的速度,IMPALA需要的CPU是SEEDRL的3~4倍。通过优化并行计算的硬件加速器,我们可以自信地增加模型的大小。在上述足球比赛任务中,通过增加模型尺寸和输入分辨率,可以解决一些之前没有解决的难点,大大提高了训练模型的效率。传送门论文地址:https://arxiv.org/abs/1910.06591GitHub地址:https://github.com/google-research/seed_rl