NVIDIA 于 2023 年 10 月 19 日公開發布 TensorRT-LLM，可在 NVIDIA GPU 上加速和優化最新的大語言模型（Large Language Models）的推理性能。該開源程序庫現已作為 NVIDIA NeMo 框架的一部分，在 /NVIDIA/TensorRT-LLM GitHub 資源庫中免費提供。

大語言模型徹底改變了人工智能領域，并創造了與數字世界交互的全新方式。但是，隨著世界各地的企業和應用開發者想要在自己的工作中采用大語言模型，運行這些模型的相關問題開始凸顯。

簡而言之，大語言模型規模龐大。如果不采用正確的技術，那么運行大語言模型的成本不僅會很高，速度也會很慢。

為解決這一問題，從模型優化（如內核融合和量化）到運行時優化（如 C++ 實現、KV 緩存、連續動態批處理 continuous in-flight batching 和分頁注意力 paged attention），眾多優化技術應運而生。但很難確定哪種技術適合您的用例，也很難在這些實施中有時并不兼容的技術間進行交互。

因此，NVIDIA 推出了 TensorRT-LLM，它是一個專門用于編譯和優化大語言模型推理的綜合程序庫。TensorRT-LLM 整合了所有這些優化功能，同時還提供了一個直觀的 Python API 來定義和構建新模型。

TensorRT-LLM 開源程序庫可加快 NVIDIA GPU 上最新大語言模型的推理性能。它是 NVIDIA NeMo 中優化大語言模型推理的骨干力量。NeMo 是一個用于構建和定制生成式 AI 應用并將其部署到生產中的端到端框架，為生成式 AI 的部署提供了完整的容器，如 TensorRT-LLM 和 NVIDIA Triton 等。

TensorRT-LLM 現在也可作為測試版本用于原生 Windows 系統。應用開發者和 AI 愛好者現可在本地運行由 NVIDIA RTX 和 NVIDIA GeForce RTX GPU 驅動的個人電腦和工作站上加速大語言模型。

TensorRT-LLM 包含 TensorRT 的深度學習編譯器，并采用了最新的優化內核，這些內核專為在前沿研究中實施用于大語言模型執行的 FlashAttention 和帶掩碼的多頭注意力 (masked multi-head attention) 而打造。

TensorRT-LLM 還將預處理和后處理步驟以及多 GPU /多節點通信基元整合到一個簡單的開源 Python API 中，可在 GPU 上實現突破性的大語言模型推理性能。

TensorRT-LLM 的特性如下：

• 支持 Llama 1 和 2、Bloom、ChatGLM、Falcon、MPT、Baichuan 及 Starcoder 等大語言模型

• 動態批處理和分頁注意力

• 多 GPU 多節點（MGMN）推理

• FP8 精度的 NVIDIA Hopper Transformer 引擎

• 支持 NVIDIA Ampere 架構、NVIDIA Ada Lovelace 架構和 NVIDIA Hopper GPU

• 原生 Windows 支持（測試版）

過去兩年，NVIDIA 一直與領先的大語言模型企業與團隊開展密切合作，以加速和優化大語言模型推理，包括 Anyscale、百度搜索、百川智能、BOSS直聘、Cohere、Deci、Grammarly、吉利研究院、金山辦公、京東探索研究院、螞蟻集團、美團、Meta、Mistral AI、MosaicML（現已并入 Databricks）、OctoML、Perplexity AI、PPIO 邊緣云、Tabnine、Together.ai、Unisound、微博、蔚來汽車、小鵬汽車語音開發團隊、智譜 AI 、智源、中科創達等（公司名按首字母英文排序）。

為幫助您了解該程序庫及其使用方式，可通過一個示例來了解如何在 Linux 上通過 TensorRT-LLM 和 NVIDIA Triton 使用并部署 Llama 2（一個常用且公開的大語言模型）。如要使用 Windows 測試版，參見用于原生 Windows 的 TensorRT-LLM GitHub 資源庫。

如需了解不同模型、不同優化以及多 GPU 執行等更多信息，請參見完整的 TensorRT-LLM 示例列表：

https://github.com/NVIDIA/TensorRT-LLM/tree/release/0.5.0/examples

開始安裝

首先克隆并構建 TensorRT-LLM 程序庫。構建 TensorRT-LLM 并檢索其所有依賴項的最簡單方法是使用隨附的 Dockerfile：

git clone -b release/0.5.0 https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM
git submodule update --init --recursive
make -C docker release_build

這些命令會提取一個基礎容器，并在容器內安裝 TensorRT-LLM 所需的所有依賴項，然后在容器中構建并安裝 TensorRT-LLM。

檢索模型權重

TensorRT-LLM 是一個用于大語言模型推理的程序庫，因此要使用它，就需要提供一組訓練過的權重。您可以使用自己在 NVIDIA NeMo 等框架中訓練的模型權重，也可以從 HuggingFace Hub 等資源庫中提取一組預訓練權重。

本文中的命令會自動從 HuggingFace Hub 中提取 70 億參數的 Llama 2 模型聊天調優變體的權重和分詞器文件。您還可以使用以下命令，自行下載權重以供離線使用。您只需更新后續命令中的路徑，使其指向此目錄即可：

git lfs install
git clone https://huggingface.co/meta-llama/Llama-2-7b-chat-hf

該模型的使用受特定許可的限制。

編譯模型

下一步是將模型編譯到 TensorRT 引擎中。為此，像定義模型權重那樣，您也需要在 TensorRT-LLM Python API 中編寫模型定義。

TensorRT-LLM 資源庫包含多種預定義模型架構。在本文中，您可以使用其中包含的 Llama 模型定義，而無需自行編寫。下面是 TensorRT-LLM 中一些可用優化的最簡示例。

有關可用插件和量化的更多信息，參見完整的 Llama 示例和數值精度。

# Launch the Tensorrt-LLM container
make -C docker release_run LOCAL_USER=1


# Log in to huggingface-cli
# You can get your token from huggingface.co/settings/token
huggingface-cli login --token *****


# Compile model
python3 examples/llama/build.py 
    --model_dir meta-llama/Llama-2-7b-chat-hf 
    --dtype float16 
    --use_gpt_attention_plugin float16 
    --use_gemm_plugin float16 
    --remove_input_padding 
    --use_inflight_batching 
    --paged_kv_cache 
    --output_dir examples/llama/out

使用 TensorRT-LLM API 創建模型定義時，可以使用構成神經網絡層的 NVIDIA TensorRT 基元來構建操作圖。這些操作會映射到特定的內核，即為 GPU 預寫的程序。

TensorRT 編譯器可以對圖進行掃描，為每個操作和可用的 GPU 選擇最佳內核。最重要的是，它還能在圖中識別出可以將多個操作“融合”到一個內核中的模式。這就減少了啟動多個 GPU 內核所需的顯存移動量和開銷。

TensorRT 還能將操作圖編譯成單個 CUDA Graph，其可以一次性啟動，從而進一步減少內核啟動開銷。

TensorRT 編譯器在融合多個層和提高執行速度方面非常強大，但有些復雜的層融合（如 FlashAttention 等）會將許多操作交叉在一起，而且無法被自動發現。對此，可以在編譯時使用插件來對圖中的部分內容進行顯式替換。

在本例中，您加入了 gpt_attention 插件（實現類似 FlashAttention 的融合注意力內核）和 gemm 插件（使用 FP32 累加執行矩陣乘法）。您還可以將完整模型的期望精度設為 FP16，以便與您從 HuggingFace 下載的權重的默認精度相匹配。

下面是該腳本運行后所生成的結果。現在，在 /examples/llama/out 文件夾中有以下文件：

• Llama_float16_tp1_rank0.engine：構建腳本的主要輸出，包含嵌入模型權重的可執行操作圖。

• config.json：包含模型的詳細信息，如模型的一般結構和精度以及關于引擎中包含哪些插件的信息。

• model.cache：緩存模型編譯過程中的部分時序和優化信息，使后續構建更快。

運行模型

現在您已經有了模型引擎，接下來該用它做什么呢？

引擎文件包含執行模型所需的信息，但在實際使用大語言模型時，需要的遠不止是一次通過模型的前向傳播。TensorRT-LLM 包含一個高度優化的 C++ 運行時，以執行已構建的大語言模型引擎和管理若干流程，如從模型輸出結果中采樣分詞、管理 KV 緩存和批量處理請求等。

您可以直接使用該運行時在本地執行模型，也可以使用 NVIDIA Triton 推理服務器的 TensorRT-LLM 運行時后端為多個用戶提供模型。

如要在本地運行模型，請執行以下命令：

python3examples/llama/run.py--engine_dir=examples/llama/out--max_output_len100--tokenizer_dirmeta-llama/Llama-2-7b-chat-hf--input_text"HowdoIcounttonineinFrench?"

使用 Triton 推理服務器部署

除了在本地執行之外，您還可以使用 NVIDIA Triton 推理服務器為您的大語言模型創建生產就緒部署。

NVIDIA 為 TensorRT-LLM 發布了一個新的 Triton 推理服務器后端。該后端利用 TensorRT-LLM C++ 運行時實現快速推理執行，并加入動態批處理和分頁 KV 緩存等技術。

首先，可使用 Dockerfile 在容器中為 Triton 推理服務器構建 TensorRT-LLM 后端。

cd ..
git clone -b release/0.5.0 
git@github.com:triton-inference-server/tensorrtllm_backend.git
cd tensorrtllm_backend
git submodule update --init --recursive
git lfs install
git lfs pull
DOCKER_BUILDKIT=1 docker build -t triton_trt_llm -f dockerfile/Dockerfile.trt_llm_backend .

接下來，創建一個模型資源庫，以便 Triton 推理服務器能夠讀取模型和任何相關元數據。tensorrtllm_backend 資源庫包含合適模型資源庫框架，該框架位于 all_models/inflight_batcher_llm/ 目錄下以供使用?，F在該目錄下有四個子文件夾，分別存放模型執行流程中各個部分的工件：

• /preprocessing和/postprocessing：包含 Triton 推理服務器 Python 后端的腳本，用于將文本輸入分詞話，并對模型輸出結果進行反分詞化處理，以便在字符串和模型使用的分詞 ID 之間進行轉換。

• /tensorrt_llm：用于放置之前編譯的模型引擎。

• /ensemble：定義一個模型集合，將前面三個組件連接在一起并告訴 Triton 推理服務器如何通過其傳輸數據。

將編譯好的模型復制到模型資源庫：

rm all_models/inflight_batcher_llm/tensorrt_llm/1/*
cp tensorrt_llm/examples/llama/out/*   all_models/inflight_batcher_llm/tensorrt_llm/1/

現在使用以下信息修改資源庫框架中的一些配置文件：

• 編譯后的模型引擎的位置

• 待使用的分詞器

• 成批執行推理時如何處理 KV 緩存的內存分配

python3 tools/fill_template.py --in_place 
      all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt 
      decoupled_mode:true,engine_dir:/app/all_models/inflight_batcher_llm/tensorrt_llm/1,
max_tokens_in_paged_kv_cache:,batch_scheduler_policy:guaranteed_completion,kv_cache_free_gpu_mem_fraction:0.2,
max_num_sequences:4


python tools/fill_template.py --in_place 
    all_models/inflight_batcher_llm/preprocessing/config.pbtxt 
    tokenizer_type:llama,tokenizer_dir:meta-llama/Llama-2-7b-chat-hf


python tools/fill_template.py --in_place 
    all_models/inflight_batcher_llm/postprocessing/config.pbtxt 
    tokenizer_type:llama,tokenizer_dir:meta-llama/Llama-2-7b-chat-hf

最后，您可以啟動 Docker 容器和 Triton 服務器，指定要部署在多少個 GPU 上（總進程數量），并指向剛剛設置的 model_repo 值。

docker run -it --rm --gpus all --network host --shm-size=1g -v $(pwd)/all_models:/app/all_models triton_trt_llm


# Log in to huggingface-cli to get the tokenizer
huggingface-cli login --token *****


# Launch the server
python3 scripts/launch_triton_server.py --world_size 1 --model_repo /app/all_models/inflight_batcher_llm

發送請求

如要向運行中的服務器發送請求并與之交互，您可以使用一個 Triton 推理服務器客戶端程序庫，這些程序庫適用于多種不同的編程語言和環境。在初始環節，請使用下面的 Python 代碼示例或 GitHub 上功能更全面的端到端客戶端流式傳輸。

import numpy as np
import tritonclient.http as httpclient
from tritonclient.utils import np_to_triton_dtype


def prepare_tensor(name, input):
    t = httpclient.InferInput(name, input.shape,
                               np_to_triton_dtype(input.dtype))
    t.set_data_from_numpy(input)
    return t


TRTION_URL = "localhost:8000"
client = httpclient.InferenceServerClient(TRITON_URL)


prompt = "How do I count to nine in French?"


inputs = [
    prepare_tensor("text_input", np.array([[prompt]], dtype=object)),
    prepare_tensor("max_tokens", np.array([[100]], dtype=np.uint32)),
    prepare_tensor("bad_words", np.array([[""]], dtype=object)),
    prepare_tensor("stop_words", np.array([[""]], dtype=object))
]


result = client.infer("ensemble", inputs)
print(result)

總結

TensorRT-LLM 和 Triton 推理服務器是優化、部署和高效運行大語言模型不可或缺的工具。隨著 TensorRT-LLM 以開源程序庫的形式發布在 GitHub 上，企業和應用開發者能夠比以往更容易地挖掘這些模型的潛力。

如果您想要深入了解大語言模型的世界，現在正是開始使用 TensorRT-LLM 的好時機。您可以探索它的功能，嘗試不同的模型和優化方法，并開啟您的旅程，解鎖 AI 驅動的語言模型所蘊含的驚人力量。

更多有關 TensorRT-LLM 的入門信息，參見以下資源：

• 訪問 /NVIDIA/TensorRT-LLM GitHub 資源庫中的開源程序庫：

  https://github.com/NVIDIA/TensorRT-LLM

• 進一步了解為生成式 AI 的部署提供完整容器（包括 TensorRT-LLM 和 NVIDIA Triton）的 NVIDIA NeMo，請訪問：

  https://www.nvidia.com/en-us/ai-data-science/generative-ai/nemo-framework/

• 在 GitHub 上探索示例代碼、基準和 TensorRT-LLM 文檔：
  https://nvidia.github.io/TensorRT-LLM/index.html

• 購買 NVIDIA AI Enterprise。

  該端到端 AI 軟件平臺包含 TensorRT，并將在不久后加入 TensorRT-LLM，其通過提供企業級的安全性、穩定性、可管理性和支持來滿足任務關鍵型 AI 推理的需求。欲了解更多信息，請聯系銷售人員：

  https://www.nvidia.cn/data-center/products/ai-enterprise/contact-sales/

• 探索 TensorRT 和 TensorRT-LLM 的入門資源，請訪問：

  https://developer.nvidia.com/tensorrt-getting-started