機器之心報導
機器之心編輯部
剛剛，Meta 發布了革命性的推理引擎 AITemplate。測試結果顯示，相比 PyTorch Eager，AITemplate 在 NVIDIA GPU 上實現了最高 12 倍的性能提升，在 AMD GPU 上實現了高達 4 倍的性能提升。
眾所周知，GPU 在各種視覺、自然語言和多模態模型推理任務中都占據重要位置。然而，對於高性能 GPU 推理引擎，AI 從業者幾乎沒有選擇權，必須使用一些平台專有的黑盒系統。這意味著如果要切換 GPU 供應商，就必須重新實現一遍部署系統。在生產環境中當涉及複雜的依賴狀況時，這種靈活性的缺失使維護疊代成本變得更加高昂。
在 AI 產品落地過程中，經常需要模型快速疊代。儘管一些閉源系統（如 TensorRT）提供了一些定製化功能，但這些定製化功能完全不能滿足需求。更進一步來說，這些閉源專有的解決方案，會使 debug 更加困難，對開發敏捷性造成影響。
針對這些業界難題，Meta AI 開發了擁有 NVIDIA GPU 和 AMD GPU 後端的統一推理引擎——AITemplate。
目前 AITemplate 已開源，項目地址：https://github.com/facebookincubator/AITemplate
AITemplate 在 CNN、Transformer 和 Diffusion 模型上都能提供接近硬體上限的 TensorCore (NVIDIA GPU) 和 MatrixCore (AMD GPU) 性能。使用 AITemplate 後，在 NVIDIA GPU 上對比 PyTorch Eager 的提速最高能達到 12 倍，在 AMD GPU 上對比 PyTorch Eager 的提速達到 4 倍。
這意味著，當應用於超大規模集群時，AITemplate 能夠節約的成本數額將是驚人的。
具體而言，AITemplate 是一個能把 AI 模型轉換成高性能 C++ GPU 模板代碼的 Python 框架。該框架在設計上專注於性能和簡化系統。AITemplate 系統一共分為兩層：前段部分進行圖優化，後端部分針對目標 GPU 生成 C++ 模板代碼。AITemplate 不依賴任何額外的庫或 Runtime，如 cuBLAS、cudnn、rocBLAS、MIOpen、TensorRT、MIGraphX 等。任何 AITemplate 編譯的模型都是自洽的。
AITemplate 中提供了大量性能提升創新，包括更先進的 GPU Kernel fusion，和一些專門針對 Transformer 的先進優化。這些優化極大提升了 NVIDIA TensorCore 和 AMD MatrixCore 的利用率。
目前，AITemplate 支持 NVIDIA A100 和 MI-200 系列 GPU，兩種 GPU 都被廣泛應用在科技公司、研究實驗室和雲計算提供商的數據中心。
團隊對 AITemplate 進行了一系列測試。下圖的測試展示了 AITemplate 和 PyTorch Eager 在 NVIDIA A100 上的主流模型中的加速比。
在帶有 Cuda 11.6 的 Nvidia A100 上運行 BERT 和 ResNet-50，AITemplate 在 ResNet-50 中提供了 3 到 12 倍的加速，在 BERT 上提供了 2 到 5 倍的加速。
經測試，AITemplate 在 AMD MI250 GPU 上較 PyTorch Eager 也有較大的加速比。
使用 ROCm 5.2 和 MI250 加速器，ResNet-50 和 BERT 的加速在 1.5-2 倍範圍內。
AITemplate 的統一 GPU 後端支持，讓深度學習開發者在最小開銷的情況下，擁有了更多的硬體提供商選擇。下圖直觀展示了 AITemplate 在 NVIDIA A100 GPU 和 AMD MI250 GPU 上的加速對比：
此外，AITemaplte 的部署較其他方案也更為簡潔。由於 AI 模型被編譯成了自洽的二進位文件並且不存在任何第三方庫的依賴，任何被編譯的二進位文件都能在相同硬體、CUDA 11/ ROCm 5 或者更新的軟體環境中運行，也不用擔心任何後向兼容問題。AITemplate 提供了開箱即用的模型樣例，如 Vision Transformer、BERT、Stable Diffusion、ResNet 和 MaskRCNN，使得部署 PyTorch 模型更加簡單。
AITemplate 的優化
AITemplate 提供了目前最先進的 GPU Kernel 融合技術：支持縱向、水平和內存融合為一體的多維融合技術。縱向融合將同一條鏈上的操作進行融合；水平融合將並行無依賴的操作進行融合；內存融合把所有內存移動操作和計算密集算子進行融合。
在水平融合中，AITemplate 目前可以把不同輸入形狀的矩陣乘法 (GEMM)、矩陣乘法和激活函數，以及 LayerNorm、LayerNorm 和激活函數進行融合。
在縱向融合中，AITemplate 支持超過傳統標準的 Elementwise 融合，包括：
- 通過 CUTLASS 和 Composable Kernel 支持了矩陣和 Elementwise 算子融合；
- 為 Transformer 的 Multi-head Attention 提供了矩陣乘法和內存布局轉置融合；
- 通過張量訪問器對內存操作，如 split、slice、concatenate 等進行融合來消除內存搬運。
在標準的 Transformer Multi head attention 模塊，目前 AITemplate 在 CUDA 平台使用了 Flash Attention，在 AMD 平台上使用了 Composable Kernel 提供的通用背靠背矩陣乘法融合。兩種解決方案都能大幅減小內存帶寬需求，在長序列問題中，提升更為明顯。如下圖所示：
AITemplate 與 Composable Kernel 的廣義背靠背融合顯著提高了長序列 Transformer 的推理效率。在 batch size 為 1 時，使用 AITemplate 的兩張 GPU 均比原生框架加速了 80%。
開發 AITemplat
AITemplate 有兩層模版系統：第一層在 Python 中使用 Jinja2 模板，第二層在 GPU TensorCore/MatrixCore 中使用 C++ 模板（NVIDIA GPU 上使用 CUTLASS，AMD GPU 上使用 Composable Kernel）。AITemplate 在 Python 中找到性能最優的 GPU 模板參數，再通過 Jinja2 渲染出最終的 C++ 代碼。
在代碼生成後，就能使用 GPU C++ 編譯器（NVIDIA 平台上的 NVCC 和 AMD 平台上的 HIPCC）編譯出最終的二進位代碼。AITemplate 提供了一套類似於 PyTorch 的前端，方便用戶直接將模型轉換到 AITemplate 而不是通過多層 IR 轉換。
總體來看，AITemplate 對當前一代及下一代 NVIDIA GPU 和 AMD GPU 提供了 SOTA 性能並大幅簡化了系統複雜度。
Meta 表示，這只是創建高性能多平台推理引擎旅程的開始：「我們正在積極擴展 AITemplate 的完全動態輸入支持。我們也有計劃推廣 AITemplate 到其他平台，例如 Apple 的 M 系列 GPU，以及來自其他供應商的 CPU 等等。」
此外，AITemplate 團隊也正在開發自動 PyTorch 模型轉換系統，使其成為開箱即用的 PyTorch 部署方案。「AITemplate 對支持 ONNX 和 Open-XLA 也持開放態度。我們希望能構建一個更為綠色高效的 AI 推理系統，能擁有更高的性能，更強的靈活性和更多的後端選擇。」團隊表示。
參考連結：https://ai.facebook.com/blog/gpu-inference-engine-nvidia-amd-open-source/