一行程式碼的威力到底有多大?今天我們要介紹的這個 Kernl 函式庫,使用者只需一行程式碼,在 GPU 上就能以快幾倍的速度運行 Pytorch transformer 模型,從而極大的加快了模型的推理速度。
具體而言,有了 Kernl 的加持,Bert 的推理速度比 Hugging Face 基準快了 12 倍。這項成果主要得益於 Kernl 以新的 OpenAI 程式語言 Triton 和 TorchDynamo 編寫了客製化的 GPU 核心。計畫作者來自 Lefebvre Sarrut。
GitHub 網址:https://github.com/ELS-RD/kernl/
以下是Kernl 與其他推理引擎的比較,橫座標中括號裡的數字分別表示batch size、序列長度,縱座標為推理加速情況。
基準測試在 3090 RTX GPU 運行,以及 12 核心 Intel CPU。
由上述結果可得,在長序列輸入這一塊,Kernl 可以說是最快的推理引擎(上圖中的右半部),在短輸入序列上接近英偉達的TensorRT(上圖的左半部)。除此之外,Kernl 內核程式碼非常簡短,易於理解和修改。該專案甚至添加了 Triton 偵錯器和工具 (基於 Fx) 來簡化核心替換,因此不需要修改 PyTorch 模型原始程式碼。
專案作者Michaël Benesty 對這項研究進行了總結,他們發布的Kernl 是一個用於加速transformer 推理的庫,速度非常快,有時會到達SOTA 性能,可破解以匹配大多數transformer 架構。
他們也在 T5 上做了測試,速度提高 6 倍,Benesty 表示這只是個開始。
在 Lefebvre Sarrut,專案作者在生產中運行幾個 transformers 模型,其中一些對延遲敏感,主要是搜尋和 recsys。他們也正在使用 OnnxRuntime 和 TensorRT,甚至創建了 transformer-deploy OSS 庫來與社區分享知識。
最近,作者在測試生成語言,並努力加速它們。然而事實證明,使用傳統工具要做到這些非常困難。在他們看來,Onnx 是另一種有趣的格式,它是一種針對機器學習所設計的開放式檔案格式,用於儲存訓練好的模型,具有廣泛的硬體支援。
但是,當他們處理新的LLM 架構時,Onnx 生態系統(主要是推理引擎)有以下幾個限制:
一個非常煩人的事實是新模型永遠不會被加速,你需要等著其他人來為此編寫自訂 CUDA 核心。現有解決方案並不是不好,OnnxRuntime 的一大優點是它的多硬體支持,TensorRT 則以非常快速著稱。
所以,專案作者想要在 Python/PyTorch 上有像 TensorRT 一樣快速的最佳化器,這也是他們創建 Kernl 的原因。
記憶體頻寬通常是深度學習的瓶頸,為了加速推理,減少記憶體存取往往是一個很好的策略。在短輸入序列上,瓶頸通常與 CPU 開銷有關,它必須被消除。專案作者主要利用了以下3 項技術:
首先是OpenAI Triton,它是一種編寫CUDA 等GPU 核心的語言,不要將它與Nvidia Triton 推理伺服器混淆,它的效率更高。幾個操作的融合實現了改進,使得他們不在 GPU 記憶體中保留中間結果的情況下連結計算。作者使用它重寫注意力(由 Flash Attention 取代)、線性層和激活以及 Layernorm/Rmsnorm。
其次是 CUDA 圖。在預熱(warmup)步驟中,它將保存每個啟動的核心及它們的參數。然後,專案作者重建了整個推理過程。
最後是 TorchDynamo,這個由 Meta 提出的原型機幫助專案作者應對動態行為。在預熱步驟中,它會追蹤模型並提供一個 Fx 圖(靜態計算圖)。他們使用自己的核心替換了 Fx 圖的一些操作,並在 Python 中重新編譯。
未來,專案路線圖將涵蓋更快的預熱、ragged 推理(padding 中沒有損失計算)、訓練支援(長序列支援)、多GPU 支援(多並行化模式)、量化(PTQ)、新batch 的Cutlass 核心測試以及提升硬體支援等。
更多詳細內容請參閱原始項目。
以上是OpenAI程式語言加速Bert推理12倍,引擎備受關注的詳細內容。更多資訊請關注PHP中文網其他相關文章!