AITemplate

| |

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

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

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 参考。

以下是几个入门教程：

示例和性能

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

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

发布

所有当前开发更新都可以在 AITemplate 仓库中看到。发布不遵循固定的时间表，只会在重大功能发布时进行标记。

中期计划：

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

长期计划：

自动 ONNX、Open-XLA 和其他格式模型转换。
基于 AVX2/AVX-512 的 AMD Epyc CPU 的 Composable Kernel CPU 扩展。

贡献

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

团队

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优化方案，这一切都离不开我们的紧密协作。