Repolode
投稿する
facebookincubator

facebookincubator/AITemplate

深層ニューラルネットワークをfp16のCUDA/HIP C++コードにコンパイルし、NVIDIAおよびAMD GPUでの高速推論を実現するPythonフレームワーク。

代替できる製品:TensorRT

GitHub で見る
スター
4,735
フォーク
386
未解決の Issue
163
最終コミット
先月

これは何？

AITemplate（AIT）は、深層ニューラルネットワークを推論サービング向けのCUDA（NVIDIA）またはHIP（AMD）C++コードに変換するPythonフレームワークです。TensorCoreおよびMatrixCoreハードウェアでのfp16ワークロードを対象とし、対応モデルとしてResNet、MaskRCNN、BERT、Vision Transformer、Stable Diffusionなどを挙げています。

コンパイラは、水平融合（並列のGEMMとLayerNorm演算を1つのカーネルにまとめる）、垂直融合（要素単位の演算、reduction、レイアウト変更をTensorCore/MatrixCoreの処理に組み込む）、メモリ融合（連結、分割、スライスを直前の演算と統合する）を適用します。コンパイルされたモデルはそれぞれ単独で動作するバイナリとなり、cuBLAS、cuDNN、TensorRTなどのランタイムに依存しません。

Pythonランタイムは余分なコピーを行わずにPyTorchテンソルを受け取ります。また、付属ツールのFX2AITはPyTorchモデルを変換でき、未対応の演算子を含むモデルも部分的に高速化できます。テスト済みの構成で使用するには、SM80以降のNVIDIA GPUまたはCDNA2 AMD GPUが必要です。

主な機能

  • 01TensorCoreおよびMatrixCoreハードウェアを対象に、fp16ニューラルネットワークをCUDA/HIP C++にコンパイルします。
  • 02水平・垂直・メモリの各演算融合を適用し、カーネル起動回数を減らします。
  • 03cuBLAS、cuDNN、TensorRTに依存しない、単独で動作するモデルバイナリを生成します。
  • 04Pythonランタイムは余分なコピーなしでPyTorchテンソルを直接受け取ります。
  • 05FX2AITはPyTorchモデルを変換し、未対応の演算子を部分的に高速化します。

ユースケース

  • NVIDIA A100クラスのGPUで、ResNet、BERT、Vision Transformerモデルをfp16でサービングする。
  • AMD MIシリーズGPUでStable Diffusionの推論を実行する。
  • AITが対応していない演算子を一部含むPyTorchモデルを高速化する。

クイックスタート

  1. 1

    サブモジュールも含めてリポジトリをクローンする

    git clone --recursive https://github.com/facebookincubator/AITemplate
  2. 2

    python ディレクトリに移動する

    cd python
  3. 3

    Python wheel をビルドする

    python setup.py bdist_wheel
  4. 4

    ビルドした wheel をインストールする

    pip install dist/*.whl --force-reinstall

使う前に知っておきたいこと

  • NVIDIA のサポートは Ampere などの SM80+ GPU でのみテストされています。古い T4 または V100 GPU では、一部のカーネルが動作しない場合があります。
  • AMD のサポートは CDNA2 GPU（MI-210/250）でのみテストされています。古い CDNA1 GPU では、コンパイラに問題が発生する場合があります。
  • ソースからビルドするには、適切な CUDA または ROCm コンパイラが必要です。誤ったコンパイラを使うと、パフォーマンスが低下します。

収録されている特集

GPU で AI を学習・実行するツールコンパイラ、トランスパイラ、プログラミング言語

トピック

#inference#gpu#cuda#fp16#compiler#pytorch#deep-learning

よくある質問

AITemplate とは？
深層ニューラルネットワークをfp16のCUDA/HIP C++コードにコンパイルし、NVIDIAおよびAMD GPUでの高速推論を実現するPythonフレームワーク。 AITemplate（AIT）は、深層ニューラルネットワークを推論サービング向けのCUDA（NVIDIA）またはHIP（AMD）C++コードに変換するPythonフレームワークです。TensorCoreおよびMatrixCoreハードウェアでのfp16ワークロードを対象とし、対応モデルとしてResNet、MaskRCNN、BERT、Vision Transformer、Stable Diffusionなどを挙げています。
AITemplate は無料で使えますか？ライセンスは？
AITemplate は Apache-2.0 ライセンスのオープンソースです。商用製品で使う前にライセンス条件を確認してください。
AITemplate は今もメンテナンスされていますか？
半年以内に更新あり。AITemplate の最終コミットは 2026年10月7日 です。
AITemplate は何に使えますか？
NVIDIA A100クラスのGPUで、ResNet、BERT、Vision Transformerモデルをfp16でサービングする。 AMD MIシリーズGPUでStable Diffusionの推論を実行する。 AITが対応していない演算子を一部含むPyTorchモデルを高速化する。
AITemplate のインストール方法は？
サブモジュールも含めてリポジトリをクローンする: git clone --recursive https://github.com/facebookincubator/AITemplate python ディレクトリに移動する: cd python Python wheel をビルドする: python setup.py bdist_wheel ビルドした wheel をインストールする: pip install dist/*.whl --force-reinstall
AITemplate はどこで入手できますか？
ソースコードとインストール手順は GitHub にあります：https://github.com/facebookincubator/AITemplate

概要はプロジェクトの README をもとに AI が生成しています。GitHub の数値は毎日更新されます。

関連プロジェクト