10月20日,全球AI領(lǐng)導(dǎo)者英偉達(dá)(NVIDIA)在官網(wǎng)開源了TensorRT-LLM,以加速和優(yōu)化NVIDIA GPU上最新大語言模型的推理性能。
目前,優(yōu)化大語言模型的方法有很多,例如,內(nèi)核融合和量化、知識蒸餾、KV緩存、模型剪枝、連續(xù)運行中批處理和分頁注意力等,但開發(fā)者面對如此多的技術(shù)很難選擇適合自己的用例,同時無法實現(xiàn)更好的兼容,而TensorRT-LLM提供了統(tǒng)一的解決方案。
TensorRT-LLM支持Llama 1/ 2、Baichuan(百川智能)、ChatGLM、Falcon、MPT、和Starcoder等市面上高性能類ChatGPT開源模型。
開源地址:
https://github.com/NVIDIA/TensorRT-LLM/tree/release/0.5.0
TensorRT-LLM簡單介紹
TensorRT-LLM是一個用于編譯和優(yōu)化大語言模型推理的綜合庫。TensorRT-LLM融合了目前主流優(yōu)化方法,同時提供了用于定義和構(gòu)建新模型的直觀Python API。
TensorRT-LLM封裝了TensorRT的深度學(xué)習(xí)編譯器,并包含最新的優(yōu)化內(nèi)核,用于實現(xiàn)FlashAttention和用于大語言模型執(zhí)行的屏蔽多頭注意力 (MHA)。
此外,TensorRT-LLM還原生支持Windows,可以幫助開發(fā)者在英偉RT和GeForce RTX GPU的PC和工作站上更快、更高效地運行大語言模型。
TensorRT-LLM主要技術(shù)亮點如下:
TensorRT-LLM使用示例
下面展示了TensorRT-LLM如何在Linux上使用、部署和優(yōu)化Llama 2。這里需要用到NVIDIA Triton(推理服務(wù)器)。
1)開始安裝
首先克隆并構(gòu)建TensorRT-LLM庫。構(gòu)建TensorRT-LLM并檢索所有依賴項的最簡單方法是使用包含Dockerfile。并在容器內(nèi)安裝TensorRT-LLM 所需的所有依賴項。
2)檢索模型權(quán)重
要使用TensorRT-LLM,需要提供一組經(jīng)過訓(xùn)練的權(quán)重。可以使用在NVIDIA NeMo等框架中訓(xùn)練的自己的模型權(quán)重,也可以從 HuggingFace Hub 等存儲庫中提取一組預(yù)訓(xùn)練的權(quán)重。
我們從HuggingFaceHub中提取 7B 參數(shù) Llama 2 模型的聊天調(diào)整變體的權(quán)重和分詞器文件。
3)編譯模型
需要將模型編譯為TensorRT引擎,所以,要將模型權(quán)重以及用 TensorRT-LLM Python API 編寫的模型進行定義。
而TensorRT-LLM 存儲庫包含各種預(yù)定義的模型架構(gòu),已經(jīng)附帶了Llama 2的模型定義,所以,開發(fā)者無需自己編寫模型定義了。這也突出了TensorRT-LLM的優(yōu)化功能。
此外,TensorRT-LLM編譯器可以掃描圖表,為每個操作和可用 GPU 選擇最佳內(nèi)核。還可以識別圖中的模式,其中多個操作是融合到單個內(nèi)核中的良好候選者。
這極大減少了所需的內(nèi)存移動量以及啟動多個GPU內(nèi)核的消耗,對于優(yōu)化大語言模型推理至關(guān)重要。
4)運行模型
開發(fā)者可以直接使用本地模型,也可以使用 NVIDIA Triton 推理服務(wù)器的 TensorRT-LLM 運行時后端為多個用戶提供的模型。
5)使用 NVIDIA Triton 推理服務(wù)器進行部署
除了本地執(zhí)行之外,開發(fā)者還可以使用NVIDIA Triton 推理服務(wù)器來部署大語言模型的生產(chǎn)環(huán)境。
首先,為Triton推理服務(wù)器構(gòu)建 TensorRT-LLM后端,使用Dockerfile在容器中構(gòu)建后端。
接下來,創(chuàng)建一個模型存儲庫,以便 Triton Inference Server 可以讀取模型和任何關(guān)聯(lián)的元數(shù)據(jù)。將編譯的模型復(fù)制到模型存儲庫。
使用下面信息修改存儲庫框架中的一些配置文件:
最后,啟動 Docker容器并啟動 Triton 服務(wù)器,指定要部署的 GPU 數(shù)量(全局大?。┎⒅赶騧odel_repo剛剛設(shè)置的值。
總體來說,TensorRT-LLM在大語言模型的模型編譯、資源消耗、流程簡化等方面表現(xiàn)非常優(yōu)秀,同時可以幫助開發(fā)人員快速嘗試不同類型的開源類ChatGPT模型。
本文素材來源英偉達(dá)官網(wǎng),如有侵權(quán)請聯(lián)系刪除
END
聯(lián)系客服