来源丨机器之心
刚刚，Meta 发布了革命性的推理引擎 AITemplate。测试结果显示，相比 PyTorch Eager，AITemplate 在 NVIDIA GPU 上实现了最高 12 倍的性能提升，在 AMD GPU 上实现了高达 4 倍的性能提升。
众所周知，GPU 在各种视觉、自然语言和多模态模型推理任务中都占据重要位置。然而，对于高性能 GPU 推理引擎，AI 从业者几乎没有选择权，必须使用一些平台专有的黑盒系统。这意味着如果要切换 GPU 供应商，就必须重新实现一遍部署系统。在生产环境中当涉及复杂的依赖状况时，这种灵活性的缺失使维护迭代成本变得更加高昂。
在 AI 产品落地过程中，经常需要模型快速迭代。尽管一些闭源系统（如 TensorRT）提供了一些定制化功能，但这些定制化功能完全不能满足需求。更进一步来说，这些闭源专有的解决方案，会使 debug 更加困难，对开发敏捷性造成影响。
针对这些业界难题，Meta AI 开发了拥有 NVIDIA GPU 和 AMD GPU 后端的统一推理引擎——AITemplate。
目前 AITemplate 已开源，项目地址：https://github.com/facebookincubator/AITemplate
AITemplate 在 CNN、Transformer 和 Diffusion 模型上都能提供接近硬件上限的 TensorCore (NVIDIA GPU) 和 MatrixCore (AMD GPU) 性能。使用 AITemplate 后，在 NVIDIA GPU 上对比 PyTorch Eager 的提速最高能达到 12 倍，在 AMD GPU 上对比 PyTorch Eager 的提速达到 4 倍。
这意味着，当应用于超大规模集群时，AITemplate 能够节约的成本数额将是惊人的。
具体而言，AITemplate 是一个能把 AI 模型转换成高性能 C++ GPU 模板代码的 Python 框架。该框架在设计上专注于性能和简化系统。AITemplate 系统一共分为两层：前段部分进行图优化，后端部分针对目标 GPU 生成 C++ 模板代码。AITemplate 不依赖任何额外的库或 Runtime，如 cuBLAS、cudnn、rocBLAS、MIOpen、TensorRT、MIGraphX 等。任何 AITemplate 编译的模型都是自洽的。
AITemplate 中提供了大量性能提升创新，包括更先进的 GPU Kernel fusion，和一些专门针对 Transformer 的先进优化。这些优化极大提升了 NVIDIA TensorCore 和 AMD MatrixCore 的利用率。
目前，AITemplate 支持 NVIDIA A100 和 MI-200 系列 GPU，两种 GPU 都被广泛应用在科技公司、研究实验室和云计算提供商的数据中心。
团队对 AITemplate 进行了一系列测试。下图的测试展示了 AITemplate 和 PyTorch Eager 在 NVIDIA A100 上的主流模型中的加速比。
在带有 Cuda 11.6 的 Nvidia A100 上运行 BERT 和 ResNet-50，AITemplate 在 ResNet-50 中提供了 3 到 12 倍的加速，在 BERT 上提供了 2 到 5 倍的加速。
经测试，AITemplate 在 AMD MI250 GPU 上较 PyTorch Eager 也有较大的加速比。
使用 ROCm 5.2 和 MI250 加速器，ResNet-50 和 BERT 的加速在 1.5-2 倍范围内。
AITemplate 的统一 GPU 后端支持，让深度学习开发者在最小开销的情况下，拥有了更多的硬件提供商选择。下图直观展示了 AITemplate 在 NVIDIA A100 GPU 和 AMD MI250 GPU 上的加速对比：
此外，AITemaplte 的部署较其他方案也更为简洁。由于 AI 模型被编译成了自洽的二进制文件并且不存在任何第三方库的依赖，任何被编译的二进制文件都能在相同硬件、CUDA 11/ ROCm 5 或者更新的软件环境中运行，也不用担心任何后向兼容问题。AITemplate 提供了开箱即用的模型样例，如 Vision Transformer、BERT、Stable Diffusion、ResNet 和 MaskRCNN，使得部署 PyTorch 模型更加简单。
AITemplate 的优化
AITemplate 提供了目前最先进的 GPU Kernel 融合技术：支持纵向、水平和内存融合为一体的多维融合技术。纵向融合将同一条链上的操作进行融合；水平融合将并行无依赖的操作进行融合；内存融合把所有内存移动操作和计算密集算子进行融合。
在水平融合中，AITemplate 目前可以把不同输入形状的矩阵乘法 (GEMM)、矩阵乘法和激活函数，以及 LayerNorm、LayerNorm 和激活函数进行融合。
在纵向融合中，AITemplate 支持超过传统标准的 Elementwise 融合，包括：
-
通过 CUTLASS 和 Composable Kernel 支持了矩阵和 Elementwise 算子融合；
-
为 Transformer 的 Multi-head Attention 提供了矩阵乘法和内存布局转置融合；
-
通过张量访问器对内存操作，如 split、slice、concatenate 等进行融合来消除内存搬运。
在标准的 Transformer Multi head attention 模块，目前 AITemplate 在 CUDA 平台使用了 Flash Attention，在 AMD 平台上使用了 Composable Kernel 提供的通用背靠背矩阵乘法融合。两种解决方案都能大幅减小内存带宽需求，在长序列问题中，提升更为明显。如下图所示：
AITemplate 与 Composable Kernel 的广义背靠背融合显着提高了长序列 Transformer 的推理效率。在 batch size 为 1 时，使用 AITemplate 的两张 GPU 均比原生框架加速了 80%。
开发 AITemplate
AITemplate 有两层模版系统：第一层在 Python 中使用 Jinja2 模板，第二层在 GPU TensorCore/MatrixCore 中使用 C++ 模板（NVIDIA GPU 上使用 CUTLASS，AMD GPU 上使用 Composable Kernel）。AITemplate 在 Python 中找到性能最优的 GPU 模板参数，再通过 Jinja2 渲染出最终的 C++ 代码。
在代码生成后，就能使用 GPU C++ 编译器（NVIDIA 平台上的 NVCC 和 AMD 平台上的 HIPCC）编译出最终的二进制代码。AITemplate 提供了一套类似于 PyTorch 的前端，方便用户直接将模型转换到 AITemplate 而不是通过多层 IR 转换。
总体来看，AITemplate 对当前一代及下一代 NVIDIA GPU 和 AMD GPU 提供了 SOTA 性能并大幅简化了系统复杂度。
Meta 表示，这只是创建高性能多平台推理引擎旅程的开始：「我们正在积极扩展 AITemplate 的完全动态输入支持。我们也有计划推广 AITemplate 到其他平台，例如 Apple 的 M 系列 GPU，以及来自其他供应商的 CPU 等等。」
此外，AITemplate 团队也正在开发自动 PyTorch 模型转换系统，使其成为开箱即用的 PyTorch 部署方案。「AITemplate 对支持 ONNX 和 Open-XLA 也持开放态度。我们希望能构建一个更为绿色高效的 AI 推理系统，能拥有更高的性能，更强的灵活性和更多的后端选择。」团队表示。
参考链接：https://ai.facebook.com/blog/gpu-inference-engine-nvidia-amd-open-source/
本文仅做学术分享，如有侵权，请联系删文。
干货下载与学习
后台回复：巴塞罗那自治大学课件，即可下载国外大学沉淀数年3D Vison精品课件
后台回复：计算机视觉书籍，即可下载3D视觉领域经典书籍pdf
后台回复：3D视觉课程，即可学习3D视觉领域精品课程
3D视觉工坊精品课程官网：3dcver.com
1.面向自动驾驶领域的多传感器数据融合技术
2.面向自动驾驶领域的3D点云目标检测全栈学习路线！(单模态+多模态/数据+代码) 3.彻底搞透视觉三维重建：原理剖析、代码讲解、及优化改进 4.国内首个面向工业级实战的点云处理课程 5.激光-视觉-IMU-GPS融合SLAM算法梳理和代码讲解 6.彻底搞懂视觉-惯性SLAM：基于VINS-Fusion正式开课啦 7.彻底搞懂基于LOAM框架的3D激光SLAM: 源码剖析到算法优化 8.彻底剖析室内、室外激光SLAM关键算法原理、代码和实战(cartographer+LOAM +LIO-SAM)
9.从零搭建一套结构光3D重建系统[理论+源码+实践]
10.单目深度估计方法：算法梳理与代码实现
11.自动驾驶中的深度学习模型部署实战
12.相机模型与标定(单目+双目+鱼眼）
13.重磅！四旋翼飞行器：算法与实战
14.ROS2从入门到精通：理论与实战
15.国内首个3D缺陷检测教程：理论、源码与实战
16.基于Open3D的点云处理入门与实战教程
重磅！3DCVer-学术论文写作投稿 交流群已成立
扫码添加小助手微信，可申请加入3D视觉工坊-学术论文写作与投稿 微信交流群，旨在交流顶会、顶刊、SCI、EI等写作与投稿事宜。
同时也可申请加入我们的细分方向交流群，目前主要有3D视觉、CV&深度学习、SLAM、三维重建、点云后处理、自动驾驶、多传感器融合、CV入门、三维测量、VR/AR、3D人脸识别、医疗影像、缺陷检测、行人重识别、目标跟踪、视觉产品落地、视觉竞赛、车牌识别、硬件选型、学术交流、求职交流、ORB-SLAM系列源码交流、深度估计等微信群。
一定要备注：研究方向+学校/公司+昵称，例如：”3D视觉 + 上海交大 + 静静“。请按照格式备注，可快速被通过且邀请进群。原创投稿也请联系。
▲长按加微信群或投稿，加微信：dddvision
▲长按关注公众号
3D视觉从入门到精通知识星球：针对3D视觉领域的视频课程（三维重建系列、三维点云系列、结构光系列、手眼标定、相机标定、激光/视觉SLAM、自动驾驶等）、知识点汇总、入门进阶学习路线、最新paper分享、疑问解答五个方面进行深耕，更有各类大厂的算法工程人员进行技术指导。与此同时，星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息，打造成集技术与就业为一体的铁杆粉丝聚集区，近6000星球成员为创造更好的AI世界共同进步，知识星球入口：
学习3D视觉核心技术，扫描查看介绍，3天内无条件退款
圈里有高质量教程资料、答疑解惑、助你高效解决问题
觉得有用，麻烦给个赞和在看~
