AITemplateMeta 开源的 Python 框架

AITemplate（AIT）是一个 Python 框架，它将深度神经网络转化为 CUDA（NVIDIA GPU）/ HIP（AMD GPU）C++ 代码，以实现快速的推理服务。AITemplate 的亮点包括：

• 高性能：在主要模型上接近 roofline fp16 TensorCore（NVIDIA GPU）/MatrixCore（AMD GPU）性能，包括 ResNet、MaskRCNN、BERT、VisionTransformer、Stable Diffusion 等。
• 统一、开放、灵活：用于 NVIDIA GPU 或 AMD GPU 的 Seamless fp16 深度神经网络模型。完全开放源代码，乐高式的易扩展高性能基元，支持新的模型。

安装

硬件要求：

• NVIDIA ：AIT 仅在 SM80+ GPU 上进行测试，并非所有内核都适用于旧的 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

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