AITemplateMeta 开源的 Python 框架

联合创作 · 2023-09-26 06:35

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

浏览 6
点赞
评论
收藏
分享

手机扫一扫分享

编辑
举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

编辑
举报