NVIDIA TensorRT助力腾讯加速“开悟”AI开放研究平台

作者 英伟达中国

案例简介

•    本案例中通过NVIDIA V100 GPU和TensorRT,腾讯“开悟”AI开放研究平台将强化学习模型训练中的自对弈推理部分,从CPU迁移到GPU上,实现了10倍的加速,并使推理成本下降90%,整体训练成本下降 67%。

•    本案例主要应用到 NVIDIA V100 GPU 和 TensorRT。

Case Introduction

•    With NVIDIA V100 GPU and TensorRT, the self-play part of reinforcement learning model training in Tencent AI Arena open research platform is migrated from CPU to GPU and achieve 10x speedup. The inference cost is reduced by 90% and the whole training cost is reduced by 67%.

•    The major products utilized in the case are NVIDIA V100 GPU and TensorRT.

客户简介及应用背景

腾讯一直积极地推动强化学习在游戏AI领域的发展,并在2019年推出了“开悟”AI开放研究平台,提供不同游戏的训练场景、支撑AI进行强化训练的大规模算力、统一的强化学习框架以加速研发速度、通用的训练与推理服务,加快AI训练速度。

客户挑战

和图像以及语音的训练方式不同,目前在游戏AI训练上表现最好的方式是强化学习。强化学习除了需要大量的算力来训练深度学习网络之外,还需要一个自对弈的模块用来产生训练数据。在自对弈模块当中,会模拟游戏环境,并且选择不同时期的模型来操控游戏内的角色,AI对游戏内角色的每一操控都需要对模型进行一次前向推理。由于更新模型的训练数据来自于自对弈模块,因此自对弈的推理速度会对整个模型的训练速度和效果造成非常大的影响。

而自对弈推理每一次前向推理对延时的要求比实际的线上服务小,因此常见的方式是通过CPU来进行自对弈中的推理,但CPU成本太高。为了提高自对弈的速度与吞吐,减少推理成本,腾讯希望在“开悟”AI开放研究平台里面充分利用GPU去进行自对弈中的模型前向推理。

应用方案

为了解决上述在自对弈推理当中GPU利用率不高的问题, 腾讯“开悟”AI开放研究平台选择使用NVIDIA V100 GPU和NVIDIA TensorRT推理引擎来加速推理。

  1. 为了自动化地将模型从TensorFlow转换到TensorRT, 腾讯“开悟”AI开放研究平台一开始通过自行开发parser,将TensorFlow的算子映射到TensorRT的算子。同时,为了支持更广泛的模型与算子,减少维护和开发的成本, 腾讯“开悟”AI开放研究平台也积极地与NVIDIA合作,推动从TensorFlow转换成ONNX模型,再通过TensorRT ONNX parser转换到TensorRT的流程。
  2. 在自对弈的过程中,需要频繁地更新模型的权重,让自对弈模型始终能保持在较新的状态。这个更新的频率大概几分钟一次,每次必须限制在几百个毫秒。如果通过重新build engine 的方式来更新模型的话,无法满足上述要求。因此 腾讯“开悟”AI开放研究平台采用 TensorRT refit engine的功能来更新权重。同时,为了对更新权重有更好的支持、以及支持更多的算子, 腾讯“开悟”AI开放研究平台从原本的TensorRT 5 升级到TensorRT 7。
  3. TensorRT 7虽然在部分算子上支持权重更新,但并不支持LSTM这个在游戏AI当中很重要的算子。为了解决这个问题, 腾讯“开悟”AI开放研究平台通过开发TensorRT插件的方式封装LSTM算子,并在插件当中更新权重。
  4. 为了充分利用NVIDIA V100 GPU的Tensor core, 腾讯“开悟”AI开放研究平台希望能够使用TensorRT FP16精度来加速推理。由于TensorRT对FP16的支持非常成熟和简便,整个推理流程很快被切换到FP16,并取得了2倍左右的加速。
  5. 寻找模型推理时的性能瓶颈,通过开发TensorRT插件进行算子融合,进一步地提升推理的速度。

在完成以上的工作之后,对比TensorFlow的基础版本,TensorRT 7 能提供5倍以上的加速效果。

使用效果及影响

通过NVIDIA V100 GPU以及TensorRT推理引擎加速自对弈训练的推理部分,腾讯“开悟”AI开放研究平台极大地提升了自对弈推理的吞吐量与速度,进一步地提升了整个模型训练的速度与降低训练成本,加快模型迭代的周期。