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平台。

更多关于AITemplate

卓越的向后兼容性

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

水平融合

AITemplate提供独特的高级水平融合。AITemplate可以将具有不同输入形状的并行GEMM、LayerNorm和其他操作符融合到单个GPU内核中。

垂直融合

AITemplate提供强大的垂直融合。AITemplate可以将大范围的操作融合到TensorCore/MatrixCore操作中，如元素级操作、归约和布局排列。AITemplate还提供背靠背式的TensorCore / MatrixCore操作融合。

内存融合

AITemplate提供创新的内存融合。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 上的参考性能数据：

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

许可证

AITemplate采用Apache 2.0许可证授权。