AI 模板

| |

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 等算子融合，然后进行 concatenation、split、slice 等内存操作。

与 PyTorch 一起工作

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

无痛苦的扩展

AITemplate 提供了一种在 codegen 中进行扩展的简单方法。要将新运算符或新的融合内核添加到 AITemplate 中，大多数情况下只需要添加两个 Python 文件：一个用于图形节点定义，另一个用于后端 codegen。文本头文件中的 CUDA/HIP 内核可以直接在 codegen 中使用。

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 镜像

我们强烈建议将 AITemplate 与 Docker 一起使用，以避免意外使用错误版本的 NVCC 或 HIPCC。

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

这将构建一个带有标签的 docker 镜像。ait:latest

从源

以下命令将为 AITemplate 创建一个 Python 轮子。请确保您安装了正确的 CUDA/ROCm 编译器。

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

不正确的编译器将导致性能回归。

在进行下一步之前，请检查所有子模块是否已正确克隆。

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

开始

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

有一些入门教程：

示例和性能

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

01_ResNet-50 与 PyTorch 图像模型（TIMM）
带 Detectron2 的 02_MaskRCNN-FPN
03_BERT 与 Hugging Face Transformer
带有 PyTorch 图像模型的 04_Vision Transformer （TIMM）
05_Stable Diffusion with Hugging Face Differs

释放

所有当前的开发更新都可以在 AITemplate 存储库中看到。版本不按固定的计划进行，只会针对重要功能版本进行标记。

中期计划：

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

长期计划：

自动 ONNX、Open-XLA 和其他格式模型转换。
AVX2/AVX-512 上适用于 AMD Epyc CPU 的可组合内核 CPU 扩展。

贡献

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

团队

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

AITemplate 由 Meta 工程师共同创建：徐兵、张英、璐浩、陈阳和 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 可组合内核团队（由 Chao Liu 和 Jing Zhang 领导）密切合作。我们共同设计了许多专门用于每个平台的高级 GPU 优化，没有我们的密切合作，一切都不可能。

许可证

AITemplate 根据 Apache 2.0 许可证获得许可。

yuanzhongqiao/AITemplate