TensorRT-LLM 專案詳細介紹

專案概述

TensorRT-LLM 是 NVIDIA 開發的開源函式庫，專門用於最佳化大型語言模型（LLM）在 NVIDIA GPU 上的推論效能。它提供易於使用的 Python API 來定義大型語言模型，並支援最先進的最佳化技術，以在 NVIDIA GPU 上高效執行推論。

核心功能

1. 進階最佳化技術

TensorRT-LLM 提供多種先進的最佳化功能，包括：

• 自訂注意力核心：專門最佳化的注意力機制實作
• 動態批次處理（Inflight Batching）：即時處理不同長度的輸入序列
• 分頁 KV 快取：高效的鍵值快取管理
• 投機解碼（Speculative Decoding）：透過預測多個 token 來加速生成
• 多種量化支援：FP8、FP4、INT4 AWQ、INT8 SmoothQuant 等

2. 量化技術詳解

TensorRT-LLM 提供業界領先的統一量化工具包，顯著加速深度學習/生成式 AI 在 NVIDIA 硬體上的部署，同時保持模型精準度。

主要量化方法：

• FP8：在大批次推論情境中通常提供最佳效能和精準度，適用於批次大小 ≥ 16 的情境
• INT8 SmoothQuant：權重平滑和 INT8 通道量化，啟用範圍張量級標定
• INT4 AWQ：權重重新縮放和區塊級量化到 INT4，推薦用於小批次推論情境（批次大小 ≤ 4）
• W4A8 AWQ：權重量化到 INT4，啟用量化到 INT8

效能提升：

根據基準測試，量化技術可以帶來顯著的效能提升：

• FP8 量化：相較於 FP16 基準，Llama 3 8B 模型可獲得 1.45 倍加速，70B 模型可獲得 1.81 倍加速
• INT4 AWQ：在批次大小為 1 的情境下，70B 模型可獲得高達 2.66 倍的效能提升
• 記憶體最佳化：所有量化版本的 Llama 3 70B 模型都可以在單個 NVIDIA H100 GPU 上執行，而 FP16 精準度需要至少兩個 GPU

3. 多 GPU 和多節點支援

TensorRT-LLM 包含預處理和後處理步驟以及多 GPU 多節點通訊原語，透過簡單開源的模型定義 API 實現突破性的 LLM 推論效能。

4. 廣泛的硬體支援

TensorRT-LLM 支援基於 NVIDIA Hopper、NVIDIA Ada Lovelace 和 NVIDIA Ampere 架構的 GPU。特別是：

• H100 GPU：支援 FP8 格式的自動轉換和最佳化核心
• H200 GPU：在 Llama2-13B 上可達到近 12,000 token/秒的效能
• RTX 系列：支援消費級 GPU 的大型模型推論

安裝與使用

Docker 安裝（推薦）

# 執行預建置的 Docker 容器
docker run --ipc host --gpus all -it nvcr.io/nvidia/tensorrt-llm/release

LLM API 使用範例

from tensorrt_llm import BuildConfig, SamplingParams
from tensorrt_llm._tensorrt_engine import LLM

def main():
    build_config = BuildConfig()
    build_config.max_batch_size = 256
    build_config.max_num_tokens = 1024
    
    # 支援 HuggingFace 模型名稱、本地 HF 模型路徑或 TensorRT 模型最佳化器量化檢查點
    llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", 
              build_config=build_config)
    
    # 範例提示詞
    prompts = [
        "Hello, my name is",
        "The capital of France is",
        "The future of AI is",
    ]
    
    # 建立取樣參數
    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
    
    for output in llm.generate(prompts, sampling_params):
        print(f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}")

線上服務部署

# 啟動 OpenAI 相容的伺服器
trtllm-serve --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --port 8000

量化工作流程

基本量化指令

# FP8 量化
python quantize.py --model_dir $MODEL_PATH --qformat fp8 --kv_cache_dtype fp8 --output_dir $OUTPUT_PATH

# INT4 AWQ 量化
python quantize.py --model_dir $MODEL_PATH --qformat int4_awq --awq_block_size 64 --tp_size 4 --output_dir $OUTPUT_PATH

# INT8 SmoothQuant 量化
python quantize.py --model_dir $MODEL_PATH --qformat int8_sq --kv_cache_dtype int8 --output_dir $OUTPUT_PATH

# 自動量化（多種方法組合）
python quantize.py --model_dir $MODEL_PATH --autoq_format fp8,int4_awq,w4a8_awq --output_dir $OUTPUT_PATH --auto_quantize_bits 5 --tp_size 2

支援的模型

TensorRT-LLM 支援大量流行的 LLM 架構，包括但不限於：

• Llama 系列：Llama 2, Llama 3, Llama 3.1, Llama 3.3
• Falcon 系列：包括 Falcon-180B
• GPT 系列：ChatGPT 相關架構
• Gemma 系列：Google 的開源模型
• Mixtral 系列：混合專家模型
• DeepSeek 系列：包括 DeepSeek R1
• CodeLlama：程式碼生成專用模型

生態系統整合

NVIDIA 生態系統

• NVIDIA NeMo：端到端框架，用於建構、客製化和部署生成式 AI 應用
• Triton Inference Server：生產級推論伺服器
• NVIDIA Dynamo：資料中心規模的分散式推論服務框架

第三方整合

• HuggingFace Hub：提供預量化模型
• LlamaIndex：RAG 應用開發
• SageMaker LMI：AWS 託管推論

效能基準

效能提升範例：

• 相較於 CPU 平台：推論速度提升高達 36 倍
• 相較於 RTX 無最佳化：Windows RTX 平台上大型語言模型速度提升高達 4 倍
• Falcon-180B：在單個 H200 GPU 上使用 INT4 AWQ 實現推論
• Llama-70B：相較於 A100 實現 6.7 倍速度提升

最佳實踐建議

量化方法選擇

根據不同情境選擇合適的量化方法：

1. 小批次推論（批次大小 ≤ 4）：

   • 推薦使用權重量化方法（如 INT4 AWQ）
   • 主要考量記憶體頻寬限制

2. 大批次推論（批次大小 ≥ 16）：

   • 優先選擇 FP8 量化，通常提供最佳效能和精準度
   • 如果結果不符合需求，可嘗試 INT8 SmoothQuant，然後是 AWQ 和/或 GPTQ

3. 特定領域應用：

   • 對於程式碼補齊等高度特化的應用，建議使用領域相關的資料集進行校準

技術優勢

1. 易用性：提供進階 Python API，簡化 LLM 定義和最佳化過程
2. 效能：包含所有主流最佳化技術，如核心融合、量化、執行時最佳化等
3. 擴展性：支援從單一 GPU 到多節點的各種部署情境
4. 相容性：與 PyTorch 深度整合，支援主要的推論生態系統
5. 開源：完全開源，社群驅動的持續發展

未來發展

TensorRT-LLM 透過開源模組化模型定義 API 提高易用性和擴展性，用於定義、最佳化和執行新架構和增強功能，可隨著 LLM 的發展輕鬆客製化。

專案持續發展方向包括：

• 更多模型架構支援
• 更先進的量化技術
• 更好的多節點擴展能力
• 更緊密的生態系統整合