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

大語(yǔ)言模型徹底改變了人工智能領(lǐng)域，并創(chuàng)造了與數(shù)字世界交互的全新方式。但是，隨著世界各地的企業(yè)和應(yīng)用開(kāi)發(fā)者想要在自己的工作中采用大語(yǔ)言模型，運(yùn)行這些模型的相關(guān)問(wèn)題開(kāi)始凸顯。

簡(jiǎn)而言之，大語(yǔ)言模型規(guī)模龐大。如果不采用正確的技術(shù)，那么運(yùn)行大語(yǔ)言模型的成本不僅會(huì)很高，速度也會(huì)很慢。

為解決這一問(wèn)題，從模型優(yōu)化（如內(nèi)核融合和量化）到運(yùn)行時(shí)優(yōu)化（如 C++ 實(shí)現(xiàn)、KV 緩存、連續(xù)動(dòng)態(tài)批處理 continuous in-flight batching 和分頁(yè)注意力 paged attention），眾多優(yōu)化技術(shù)應(yīng)運(yùn)而生。但很難確定哪種技術(shù)適合您的用例，也很難在這些實(shí)施中有時(shí)并不兼容的技術(shù)間進(jìn)行交互。

因此，NVIDIA 推出了 TensorRT-LLM，它是一個(gè)專門(mén)用于編譯和優(yōu)化大語(yǔ)言模型推理的綜合程序庫(kù)。TensorRT-LLM 整合了所有這些優(yōu)化功能，同時(shí)還提供了一個(gè)直觀的 Python API 來(lái)定義和構(gòu)建新模型。

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

TensorRT-LLM 現(xiàn)在也可作為測(cè)試版本用于原生 Windows 系統(tǒng)。應(yīng)用開(kāi)發(fā)者和 AI 愛(ài)好者現(xiàn)可在本地運(yùn)行由 NVIDIA RTX 和 NVIDIA GeForce RTX GPU 驅(qū)動(dòng)的個(gè)人電腦和工作站上加速大語(yǔ)言模型。

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

TensorRT-LLM 還將預(yù)處理和后處理步驟以及多 GPU /多節(jié)點(diǎn)通信基元整合到一個(gè)簡(jiǎn)單的開(kāi)源 Python API 中，可在 GPU 上實(shí)現(xiàn)突破性的大語(yǔ)言模型推理性能。

TensorRT-LLM 的特性如下：

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

• 動(dòng)態(tài)批處理和分頁(yè)注意力

• 多 GPU 多節(jié)點(diǎn)（MGMN）推理

• FP8 精度的 NVIDIA Hopper Transformer 引擎

• 支持 NVIDIA Ampere 架構(gòu)、NVIDIA Ada Lovelace 架構(gòu)和 NVIDIA Hopper GPU

• 原生 Windows 支持（測(cè)試版）

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

為幫助您了解該程序庫(kù)及其使用方式，可通過(guò)一個(gè)示例來(lái)了解如何在 Linux 上通過(guò) TensorRT-LLM 和 NVIDIA Triton 使用并部署 Llama 2（一個(gè)常用且公開(kāi)的大語(yǔ)言模型）。如要使用 Windows 測(cè)試版，參見(jiàn)用于原生 Windows 的 TensorRT-LLM GitHub 資源庫(kù)。

如需了解不同模型、不同優(yōu)化以及多 GPU 執(zhí)行等更多信息，請(qǐng)參見(jiàn)完整的 TensorRT-LLM 示例列表：

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

開(kāi)始安裝

首先克隆并構(gòu)建 TensorRT-LLM 程序庫(kù)。構(gòu)建 TensorRT-LLM 并檢索其所有依賴項(xiàng)的最簡(jiǎn)單方法是使用隨附的 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

這些命令會(huì)提取一個(gè)基礎(chǔ)容器，并在容器內(nèi)安裝 TensorRT-LLM 所需的所有依賴項(xiàng)，然后在容器中構(gòu)建并安裝 TensorRT-LLM。

檢索模型權(quán)重

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

本文中的命令會(huì)自動(dòng)從 HuggingFace Hub 中提取 70 億參數(shù)的 Llama 2 模型聊天調(diào)優(yōu)變體的權(quán)重和分詞器文件。您還可以使用以下命令，自行下載權(quán)重以供離線使用。您只需更新后續(xù)命令中的路徑，使其指向此目錄即可：

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

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

編譯模型

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

TensorRT-LLM 資源庫(kù)包含多種預(yù)定義模型架構(gòu)。在本文中，您可以使用其中包含的 Llama 模型定義，而無(wú)需自行編寫(xiě)。下面是 TensorRT-LLM 中一些可用優(yōu)化的最簡(jiǎn)示例。

有關(guān)可用插件和量化的更多信息，參見(jiàn)完整的 Llama 示例和數(shù)值精度。

# 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 創(chuàng)建模型定義時(shí)，可以使用構(gòu)成神經(jīng)網(wǎng)絡(luò)層的 NVIDIA TensorRT 基元來(lái)構(gòu)建操作圖。這些操作會(huì)映射到特定的內(nèi)核，即為 GPU 預(yù)寫(xiě)的程序。

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

TensorRT 還能將操作圖編譯成單個(gè) CUDA Graph，其可以一次性啟動(dòng)，從而進(jìn)一步減少內(nèi)核啟動(dòng)開(kāi)銷。

TensorRT 編譯器在融合多個(gè)層和提高執(zhí)行速度方面非常強(qiáng)大，但有些復(fù)雜的層融合（如 FlashAttention 等）會(huì)將許多操作交叉在一起，而且無(wú)法被自動(dòng)發(fā)現(xiàn)。對(duì)此，可以在編譯時(shí)使用插件來(lái)對(duì)圖中的部分內(nèi)容進(jìn)行顯式替換。

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

下面是該腳本運(yùn)行后所生成的結(jié)果。現(xiàn)在，在 /examples/llama/out 文件夾中有以下文件：

• Llama_float16_tp1_rank0.engine：構(gòu)建腳本的主要輸出，包含嵌入模型權(quán)重的可執(zhí)行操作圖。

• config.json：包含模型的詳細(xì)信息，如模型的一般結(jié)構(gòu)和精度以及關(guān)于引擎中包含哪些插件的信息。

• model.cache：緩存模型編譯過(guò)程中的部分時(shí)序和優(yōu)化信息，使后續(xù)構(gòu)建更快。

運(yùn)行模型

現(xiàn)在您已經(jīng)有了模型引擎，接下來(lái)該用它做什么呢？

引擎文件包含執(zhí)行模型所需的信息，但在實(shí)際使用大語(yǔ)言模型時(shí)，需要的遠(yuǎn)不止是一次通過(guò)模型的前向傳播。TensorRT-LLM 包含一個(gè)高度優(yōu)化的 C++ 運(yùn)行時(shí)，以執(zhí)行已構(gòu)建的大語(yǔ)言模型引擎和管理若干流程，如從模型輸出結(jié)果中采樣分詞、管理 KV 緩存和批量處理請(qǐng)求等。

您可以直接使用該運(yùn)行時(shí)在本地執(zhí)行模型，也可以使用 NVIDIA Triton 推理服務(wù)器的 TensorRT-LLM 運(yùn)行時(shí)后端為多個(gè)用戶提供模型。

如要在本地運(yùn)行模型，請(qǐng)執(zhí)行以下命令：

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

使用 Triton 推理服務(wù)器部署

除了在本地執(zhí)行之外，您還可以使用 NVIDIA Triton 推理服務(wù)器為您的大語(yǔ)言模型創(chuàng)建生產(chǎn)就緒部署。

NVIDIA 為 TensorRT-LLM 發(fā)布了一個(gè)新的 Triton 推理服務(wù)器后端。該后端利用 TensorRT-LLM C++ 運(yùn)行時(shí)實(shí)現(xiàn)快速推理執(zhí)行，并加入動(dòng)態(tài)批處理和分頁(yè) KV 緩存等技術(shù)。

首先，可使用 Dockerfile 在容器中為 Triton 推理服務(wù)器構(gòu)建 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 .

接下來(lái)，創(chuàng)建一個(gè)模型資源庫(kù)，以便 Triton 推理服務(wù)器能夠讀取模型和任何相關(guān)元數(shù)據(jù)。tensorrtllm_backend 資源庫(kù)包含合適模型資源庫(kù)框架，該框架位于 all_models/inflight_batcher_llm/ 目錄下以供使用?，F(xiàn)在該目錄下有四個(gè)子文件夾，分別存放模型執(zhí)行流程中各個(gè)部分的工件：

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

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

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

將編譯好的模型復(fù)制到模型資源庫(kù)：

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

現(xiàn)在使用以下信息修改資源庫(kù)框架中的一些配置文件：

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

• 待使用的分詞器

• 成批執(zhí)行推理時(shí)如何處理 KV 緩存的內(nèi)存分配

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

最后，您可以啟動(dòng) Docker 容器和 Triton 服務(wù)器，指定要部署在多少個(gè) GPU 上（總進(jìn)程數(shù)量），并指向剛剛設(shè)置的 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

發(fā)送請(qǐng)求

如要向運(yùn)行中的服務(wù)器發(fā)送請(qǐng)求并與之交互，您可以使用一個(gè) Triton 推理服務(wù)器客戶端程序庫(kù)，這些程序庫(kù)適用于多種不同的編程語(yǔ)言和環(huán)境。在初始環(huán)節(jié)，請(qǐng)使用下面的 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)

總結(jié)

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

如果您想要深入了解大語(yǔ)言模型的世界，現(xiàn)在正是開(kāi)始使用 TensorRT-LLM 的好時(shí)機(jī)。您可以探索它的功能，嘗試不同的模型和優(yōu)化方法，并開(kāi)啟您的旅程，解鎖 AI 驅(qū)動(dòng)的語(yǔ)言模型所蘊(yùn)含的驚人力量。

更多有關(guān) TensorRT-LLM 的入門(mén)信息，參見(jiàn)以下資源：

• 訪問(wèn) /NVIDIA/TensorRT-LLM GitHub 資源庫(kù)中的開(kāi)源程序庫(kù)：

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

• 進(jìn)一步了解為生成式 AI 的部署提供完整容器（包括 TensorRT-LLM 和 NVIDIA Triton）的 NVIDIA NeMo，請(qǐng)?jiān)L問(wèn)：

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

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

• 購(gòu)買 NVIDIA AI Enterprise。

  該端到端 AI 軟件平臺(tái)包含 TensorRT，并將在不久后加入 TensorRT-LLM，其通過(guò)提供企業(yè)級(jí)的安全性、穩(wěn)定性、可管理性和支持來(lái)滿足任務(wù)關(guān)鍵型 AI 推理的需求。欲了解更多信息，請(qǐng)聯(lián)系銷售人員：

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

• 探索 TensorRT 和 TensorRT-LLM 的入門(mén)資源，請(qǐng)?jiān)L問(wèn)：

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

GTC 2024 將于 2024 年 3 月 18 至 21 日在美國(guó)加州圣何塞會(huì)議中心舉行，線上大會(huì)也將同期開(kāi)放。點(diǎn)擊“閱讀原文”或掃描下方海報(bào)二維碼，關(guān)注更多會(huì)議及活動(dòng)信息。