facebookincubator/AITemplate

何を解決するか

AITemplateは、NVIDIAおよびAMD GPU上の深層ニューラルネットワークの推論速度を最大化することを目的としています。汎用ランタイムの性能ボトルネックを克服するために、モデルを高度に最適化された自己完結型C++コードに変換し、ハードウェアの理論的最大性能（ルーフライン）に近い推論サービングを実現します。

動作方法

フレームワークはニューラルネットワークグラフをCUDAまたはHIP C++コードに変換します。メモリオーバーヘッドを削減し、GPUの利用効率を向上させるために、いくつかの高度な融合技術を採用しています：

• 水平融合：異なる入力形状を持つ並列演算子（GEMMやLayerNormなど）を1つのカーネルに統合します。
• 垂直融合：要素演算や還元演算などの範囲の演算を、TensorCoreまたはMatrixCore操作に直接統合します。
• メモリ融合：連結、分割、スライスなどのメモリ操作を、それらに先行する演算子に統合します。

FX2AIT と呼ばれるツールを含んでおり、PyTorchモデルをAITemplateエンジンに変換できます。モデルにAITemplateでまだサポートされていない演算子が含まれる場合、FX2AITはモデルを分割して部分的な高速化を実現できます。

対象ユーザー

Ampere世代のNVIDIA GPU（SM80+）またはCDNA2 AMD GPU（MI-210/250）上で極めて高い推論性能を必要とする開発者やエンジニア。特にBERT、Stable Diffusion、Vision Transformersなどのモデルに適しています。

特徴

• ハードウェア非依存：NVIDIAおよびAMD GPUプラットフォームの両方で動作します。
• 依存なし：コンパイル済みモデルは自己完結型バイナリであり、実行時にcuBLAS、cuDNN、TensorRTなどのサードパーティライブラリを必要としません。
• PyTorch統合：PyTorchテンソルを入出力として使用でき、追加のメモリコピーが不要ですが、PyTorchが完全に存在しない環境でも実行可能です。
• 拡張可能：Pythonでグラフノードとバックエンドコード生成を定義することで、新しい演算子や融合カーネルを追加できます。