AITemplate

|
|

AITemplate (AIT) 是一个 Python 框架，可将深度神经网络转换为 CUDA (NVIDIA GPU) / HIP (AMD GPU) C++ 代码，以实现闪电般快速的推理服务。AITemplate 的主要特点包括：

• 高性能：在主流模型（如 ResNet、MaskRCNN、BERT、VisionTransformer、Stable Diffusion 等）上，可达到接近理论峰值的 fp16 TensorCore (NVIDIA GPU) / MatrixCore (AMD GPU) 性能。
• 统一、开放且灵活：为 NVIDIA GPU 或 AMD GPU 提供无缝的 fp16 深度神经网络模型。完全开源，提供乐高式易于扩展的高性能原语以支持新模型。相比现有解决方案，为两种 GPU 平台支持了更广泛的算子融合范围。

关于 AITemplate 的更多信息

出色的向后兼容性

AITemplate 不依赖第三方库或运行时，如 cuBLAS、cuDNN、rocBLAS、MIOpen、TensorRT、MIGraphX 等。每个模型都被编译成一个独立、可移植的二进制文件，可在具有相同硬件的任何软件环境中使用。

水平融合

AITemplate 提供独特的高级水平融合功能。它可以将具有不同输入形状的并行 GEMM、LayerNorm 等算子融合到单个 GPU 内核中。

垂直融合

AITemplate 提供强大的垂直融合功能。它可以将大量操作（如逐元素操作、归约操作、布局置换）融合到 TensorCore/MatrixCore 操作中。AITemplate 还支持背靠背式的 TensorCore / MatrixCore 操作融合。

内存融合

AITemplate 提供创新的内存融合功能。它可以将 GEMM、LayerNorm 等算子与后续的内存操作（如拼接、拆分、切片）融合为单个算子。

支持与 PyTorch 协同或独立工作

AITemplate 生成的 Python 运行时可以直接接收 PyTorch 张量作为输入和输出，无需额外拷贝。对于没有 PyTorch 的环境，AITemplate 的 Python/C++ 运行时是自包含的。

轻松扩展

AITemplate 为代码生成中的扩展提供了一种直接的方法。要向 AITemplate 添加新的算子或新的融合内核，大多数情况下只需添加两个 Python 文件：一个用于图节点定义，另一个用于后端代码生成。文本头文件中的 CUDA/HIP 内核可以直接在代码生成中使用。

FX2AIT

FX2AIT 是一个基于 Python 的工具，可将 PyTorch 模型转换为 AITemplate (AIT) 引擎，以实现闪电般快速的推理服务。利用 FX2AIT 内置的 AITLowerer，可以对 AITemplate 中不支持的算子实现部分 AIT 加速。

FX2AIT 的主要特性包括：

• 易于转换：FX2AIT 仅需要一个 PyTorch 模型和输入即可进行转换，并生成一个用于推理服务的 "AITModule" 输出。
• 扩展支持：AITemplate 并不支持所有 PyTorch 算子。FX2AIT 的 AITLowerer 为包含不支持算子的模型提供了部分 AIT 转换的解决方案。更多信息请查看 fx2ait/fx2ait/example/03_lowering_split。

更多信息请访问：https://github.com/facebookincubator/AITemplate/tree/main/fx2ait。

安装

硬件要求：

• NVIDIA：AIT 仅在 SM80+ GPU（Ampere 等）上经过测试。并非所有内核都适用于旧的 SM75/SM70（T4/V100）GPU。
• AMD：AIT 仅在 CDNA2（MI-210/250）GPU 上经过测试。对于旧的 CDNA1（MI-100）GPU 可能存在编译器问题。

克隆代码

克隆代码时，请使用以下命令同时克隆子模块：

git clone --recursive https://github.com/facebookincubator/AITemplate

Docker 镜像

我们强烈建议使用 Docker 来运行 AITemplate，以避免意外使用错误版本的 NVCC 或 HIPCC。

• CUDA：./docker/build.sh cuda
• ROCM：DOCKER_BUILDKIT=1 ./docker/build.sh rocm

这将构建一个标签为 ait:latest 的 Docker 镜像。

从源码安装

以下命令将为 AITemplate 创建一个 Python wheel 包。请确保已安装正确的 CUDA/ROCm 编译器。

• CUDA：CUDA 11.6
• ROCm：我们在 ROCm 5.2.3 上进行了测试，并使用 docker/Dockerfile.rocm#L87-L96 中的命令定制构建了 HIPCC。

不正确的编译器将导致性能下降。

请在进行下一步之前，确保所有子模块都已正确克隆。

cd python
python setup.py bdist_wheel
pip install dist/*.whl --force-reinstall

快速开始

查看 AITemplate 文档 以获取 API 参考。

这里有几个入门教程：

• 01：如何使用 AIT 推理 PyTorch 模型
• 02：如何向 AIT 代码生成中添加算子
• 03：如何可视化 AIT 的优化过程

示例与性能

AITemplate 提供了以下模型模板以及在 A100/MI-250 上的参考性能数据：

• 结合 PyTorch Image Models (TIMM) 的 01_ResNet-50
• 结合 Detectron2 的 02_MaskRCNN-FPN
• 结合 Hugging Face Transformer 的 03_BERT
• 结合 PyTorch Image Models (TIMM) 的 04_Vision Transformer
• 结合 Hugging Face Diffusers 的 05_Stable Diffusion

发布计划

所有当前的开发更新都可以在 AITemplate 仓库中查看。发布没有固定的时间表，仅会在重大功能发布时标记版本。

中期计划：

• 更好的动态形状支持：专注于 Transformer 中的动态序列。增加符号形状支持。
• 更多自动图优化：减少手动重写模型以获得最佳性能的工作量。
• 量化：支持 fp8/int8/int4。
• Gemm 的稀疏性剪枝。
• PT2 集成：Aten2AIT 正在积极开发中。

长期计划：

• 自动模型转换：支持 ONNX、Open-XLA 及其他格式的模型自动转换。
• 可组合内核 CPU 扩展：在 AMD Epyc CPU 上支持 AVX2/AVX-512。

贡献

查看我们的 贡献指南，了解如何为项目做出贡献。

团队

AITemplate 目前由 Meta 工程师维护：Ying Zhang, Yang Chen, Terry Chen, Mu-Chu Lee, Max Podkorytov, Adnan Akhundov。

AITemplate 由 Meta 工程师共同创建：Bing Xu, Ying Zhang, Hao Lu, Yang Chen, 和 Terry Chen，并得到了其他才华横溢的工程师的重要贡献。非详尽名单包括：Mike Iovine, Mu-Chu Lee, Scott Wolchok, Oleg Khabinov, Shirong Wu, Huamin Li, Hui Guo, Zhijing Li, Max Podkorytov。我们还要感谢 Andrew Tulloch, Yinghai Lu, Lu Fang 的宝贵讨论。

FX2AIT 和 Aten2AIT 由 Meta 工程师共同创建和维护：Wei Wei, Shirong Wu 和 Zhijing Li。

致谢

AITemplate 团队与 NVIDIA CUTLASS 团队（由 Andrew Kerr, Haicheng Wu 领导）和 AMD Composable Kernel 团队（由 Chao Liu, Jing Zhang 领导）紧密合作。我们共同设计了许多针对每个平台的高级 GPU 优化，没有我们紧密的合作，这一切都不可能实现。

许可证

AITemplate 采用 Apache 2.0 许可证。