TensorRT-LLM | 大模型部署專用框架
本文經自動駕駛之心公眾號授權轉載,轉載請聯系出處。
TensorRT-LLM是NVIDIA推出的一款高性能深度學習推理優化庫,專注于提升大型語言模型(LLM)在NVIDIA GPU上的推理速度和效率。如果您繞不開Nvidia的芯片,那么一定要好好了解這款推理庫。
項目鏈接:https://github.com/NVIDIA/TensorRT-LLM
一、TensorRT-LLM的優勢
TensorRT-LLM(TensorRT for Large Language Models)旨在解決大型語言模型在實際應用中面臨的性能瓶頸問題。通過提供一系列專為LLM推理設計的優化工具和技術,TensorRT-LLM能夠顯著提升模型的推理速度,降低延遲,并優化內存使用。
二、TensorRT-LLM的核心功能
1)易于使用的Python API
- TensorRT-LLM提供了一個簡潔易用的Python API,允許用戶定義大型語言模型并構建包含先進優化的TensorRT引擎。
- 該API設計類似于PyTorch,使得具有PyTorch經驗的開發者能夠輕松遷移和集成。
2)模型優化
- TensorRT-LLM支持多種量化選項(如FP16、INT8等),用戶可以根據具體需求選擇合適的配置,實現性能與精度的平衡。
- 通過層級融合、內核選擇和精度調整等優化技術,TensorRT-LLM能夠顯著提升模型的推理速度。
3)內存管理
- TensorRT-LLM通過智能內存分配和分頁注意力機制,優化了內存使用,降低了內存占用。
4)多線程并行與硬件加速
- 支持多線程并行處理,提高處理速度。
- 充分利用NVIDIA GPU的計算能力,加速模型推理。
5)動態批處理
- TensorRT-LLM支持動態批處理,通過同時處理多個請求來優化文本生成,減少了等待時間并提高了GPU利用率。
6)多GPU與多節點推理
- 支持在多個GPU或多個節點上進行分布式推理,提高了吞吐量并減少了總體推理時間。
7)FP8支持
- 配備TensorRT-LLM的NVIDIA H100 GPU能夠輕松地將模型權重轉換為新的FP8格式,并自動編譯模型以利用優化的FP8內核。這得益于NVIDIA Hopper架構,且無需更改任何模型代碼。
8)最新GPU支持
- TensorRT-LLM 支持基于 NVIDIA Hopper、NVIDIA Ada Lovelace、NVIDIA Ampere、NVIDIA Turing 和 NVIDIA Volta 架構的GPU。
三、TensorRT-LLM支持部署的模型
1)LLM系列
2)多模態大模型
四、量化相關
INT8 SmoothQuant (W8A8)
SmoothQuant技術在:https://arxiv.org/abs/2211.10438中被介紹。它是一種使用INT8對激活和權重進行推理的方法,同時保持網絡(在下游任務中)的準確性。如研究論文所述,必須對模型的權重進行預處理。TensorRT-LLM包含用于準備模型以使用SmoothQuant方法運行的腳本。
關于如何為GPT、GPT-J和LLaMA啟用SmoothQuant的示例,可以在版本的examples/quantization
文件夾中找到。
INT4和INT8僅權重量化 (W4A16和W8A16)
INT4和INT8僅權重量化技術包括對模型的權重進行量化,并在線性層(Matmuls)中動態地對這些權重進行反量化。激活使用浮點數(FP16或BF16)進行編碼。要使用INT4/INT8僅權重量化方法,用戶必須確定用于量化和反量化模型權重的縮放因子。
GPTQ和AWQ (W4A16)
GPTQ和AWQ技術分別在https://arxiv.org/abs/2210.17323和https://arxiv.org/abs/2306.00978中介紹。TensorRT-LLM支持在線性層中使用每組縮放因子和零偏移來實現GPTQ和AWQ方法。有關詳細信息,請參閱WeightOnlyGroupwiseQuantMatmulPlugin
插件和相應的weight_only_groupwise_quant_matmul
Python函數。
代碼中包括將GPTQ應用于GPT-NeoX和LLaMA-v2的示例,以及使用AWQ與GPT-J的示例。這些示例是實驗性實現,并可能在未來的版本中有所改進。
FP8 (Hopper)
TensorRT-LLM包含為GPT-NeMo、GPT-J和LLaMA實現的FP8。這些示例可以在examples/quantization
中找到。
五、TensorRT-LLM支持的硬件和軟件
六、TensorRT-LLM的應用場景
TensorRT-LLM在多個領域展現了其強大的應用能力,包括但不限于:
- 在線客服系統:通過實時的對話生成,提供無縫的人工智能輔助服務。
- 搜索引擎:利用模型對查詢進行增強,提供更精準的搜索結果。
- 自動代碼補全:在IDE中集成模型,幫助開發者自動完成代碼編寫。
- 內容創作平臺:自動生成文章摘要或建議,提升創作者的工作效率。