NVIDIA TensorRT 是性能最优、应用最广的 GPU 推理框架,但用户常常因为不支持的算子而面临手写插件的痛点。为此,腾讯 TEG 与 NVIDIA 联合开发了一款能自动生成插件的工具:TensorRT Plugin Autogen Tool,简称 TPAT。它能支持几乎所有的 ONNX 算子,在节省人力成本的同时,性能对比手写插件也毫不逊色。
腾讯 TEG(技术工程团队)为腾讯各个团队提供技术与工程服务,日常任务是为各种产品提供工程设计与实现,以满足其他团队的业务需要。
随着 AI 应用的普及,腾讯 TEG 在深度学习模型的推理上也开始了技术积累。腾讯 TEG 选择了 TensorRT 作为 NVIDIA GPU 上推理应用的基础软件,开发了大量基于 TensorRT 的 AI 模型,并与 NVIDIA 展开了深入合作。
TensorRT 是一个由 NVIDIA 开发的用于推理的高性能 SDK,包含深度学习推理优化器和运行时环境,可以为深度学习推理应用提供低延迟和高吞吐量,这也是 TensorRT 相比其他框架最大的优势。
不过,TensorRT 也存在一定局限性,当现有模型的算子不能被 TensorRT 直接支持时,工程师不得不手写插件,比较耗时耗力。
腾讯 TEG 表示,手写插件往往是 TensorRT 模型开发中最耗时的一部分,其难点在于:
手写插件需要熟悉 TensorRT 插件的编程接口,需要学习 GPU 和 CUDA 相关知识。
新入门的工程师通常需要 1~2 周时间来编写一个算子实现,模型中如果包含多个不支持算子,就需要更多时间来逐个编写和调试。
如果有一种自动生成插件的工具,将极大地提高 TensorRT 模型开发的效率,于是腾讯 TEG 与 NVIDIA 开展了这方面的合作,历时数月,共同开发了 TPAT(TensorRT Plugin Autogen Tool)。
TPAT 本质上是对 TensorRT 的功能扩展。它要解决的关键问题在于如下方面:
为了生成通用的插件,需要开发出 TensorRT 插件的代码模板。与此同时,需要提供部分插件的手工参考实现,用于比对计算结果与性能。这一部分由 NVIDIA 负责。
对于给定的具体算子,其 GPU 上的 CUDA kernel 需要一种自动生成机制。经腾讯 TEG 与 NVIDIA 协商,选择 ONNX 作为候选算子集合,在编写中间代码之后由 TVM 生成 CUDA kernel。这一部分由腾讯 TEG 负责。
在完成上述开发之后,腾讯 TEG 与 NVIDIA 进行联调,并后续同步进行版本迭代。
历经腾讯 TEG 与 NVIDIA 工程师的共同努力,TPAT 已正式发布 1.0 版本。它基于 TVM 和 ONNX,实现了 TensorRT 插件的全自动生成,使得 TensorRT 的部署和上线能基本流程化,不再需要人工参与。
另外,TPAT 全自动生成一个算子插件耗时仅需要 30-60 分钟的时间(该时间用于搜索算子的高性能 CUDA kernel),TensorRT 会因此成为一个真正端到端的推理框架。
TPAT 亮点:
覆盖度:几乎支持 ONNX 所有的算子
全自动:端到端全自动生成用户指定的 TensorRT Plugin
高性能:大部分算子的性能与手写插件相比并不逊色
值得一提的是,最新发布的 TPAT 也支持了动态 batch,从而成为真正具有业界实用性的工具。
TPAT 已经用在了腾讯游戏中心的多个游戏 AI 模型上,体现了优秀的好用性与性能。此外,很多其他腾讯团队也对 TPAT 感兴趣,在评估与尝试部署 TPAT 生成的 TensorRT 模型。
“对内部业务模型里的部分算子进行了测试,TPAT 的性能几乎全面达到甚至超越 CUDA 工程师手写的插件,并且端到端的设计能够大幅减少人力投入。”来自腾讯 TEG 的工程师邱骞表示。