LLM 推理 - Nvidia TensorRT-LLM 與 Triton Inference Server

1. LLM 推理 - TensorRT-LLM 與 Triton Inference Server

隨著LLM越來越熱門，LLM的推理服務也得到越來越多的關注與探索。在推理框架方面，tensorrt-llm是非常主流的開源框架，在Nvidia GPU上提供了多種最佳化，加速大語言模型的推理。但是，tensorrt-llm僅是一個推理框架，可以幫助我們完成模型的載入與推理。若是要應用在生產上，仍需要與推理服務例如Triton Inference Server進行配合。

本文會介紹TensorRT-LLM與Triton的基本原理，並展示使用這兩個元件部署LLM推理服務的流程。同時也會介紹其中使用到的推理最佳化技術。

2. 基礎環境

硬體：1 x A10 GPU（24GB視訊記憶體）

作業系統：Ubuntu 22.04

3. TensorRT-LLM

TensorRT-LLM是一個易於使用的Python API,用於定義大型語言模型(LLM),並構建包含最先進最佳化的TensorRT引擎,以在NVIDIA GPU上高效執行推理。TensorRT-LLM包含用於建立執行這些TensorRT引擎的Python和C++執行時的元件。它還包括與NVIDIA Triton推理伺服器整合的後端。使用TensorRT-LLM構建的模型可以在從單個GPU到多個節點(使用張量並行或/和管道並行)的多個GPU的廣泛配置上執行。

為了最大化效能並減小記憶體佔用,TensorRT-LLM允許使用多種量化模式執行模型。TensorRT-LLM支援INT4或INT8量化(以及FP16啟用；也稱INT4/INT8 weight-only)，以及SmoothQuant技術的完整實現。

下面我們先使用TensorRT-LLM來進行模型的推理，然後介紹TensorRT-LLM在提升推理效能上做的部分最佳化。

3.1. 設定TensorRT-LLM環境

下面我們參考TensorRT-LLM的官網[1]進行設定。

# 安裝docker

sudo apt-get install docker

# 部署nvidia ubuntu容器

docker run --runtime=nvidia --gpus all -v /home/ubuntu/data:/data -p 8000:8000 --entrypoint /bin/bash -itd nvidia/cuda:12.4.0-devel-ubuntu22.04

# 進入容器後安裝TensorRT-LLM

# 安裝0.9.0版本，還需要使用低版本的numpy

# 雖然目前最新版是0.10.0，但是必須安裝0.9.0版本

# 因為後續使用triton映象時，裡面的tensorrt_llm最新版本只到0.9.0

pip3 install tensorrt_llm==0.9.0 -U --extra-index-url https://pypi.nvidia.com

pip3 install numpy==1.26.0

# 檢查是否安裝成功

> python3 -c "import tensorrt_llm"

[TensorRT-LLM] TensorRT-LLM version: 0.9.0

3.2. 模型推理

在設定好TensorRT-LLM的環境後，下面對llama2模型進行推理測試。

（這裡為什麼沒有用最新的Llama3是因為在嘗試做部署與推理Llama3-8B-Chinese-Chat模型的過程中遇到了一個暫時未解決的問題，具體報錯為：RuntimeError: 【TensorRT-LLM】【ERROR】 Assertion failed: mpiSize == tp * pp (/home/jenkins/agent/workspace/LLM/release-0.10/L0_PostMerge/tensorrt_llm/cpp/tensorrt_llm/runtime/worldConfig.cpp:99)）。）

下面是具體過程：

# 安裝hugging face 環境

pip install -U huggingface_hub

# 進入與本機磁碟mapping的data目錄

cd data

# 下載模型

huggingface-cli login --token ****

huggingface-cli download --resume-download meta-llama/Llama-2-7b-chat-hf --local-dir llama-2-7b-ckpt

# 下載TensorRT-LLM原始碼

git clone -b v0.9.0 https://github.com/NVIDIA/TensorRT-LLM.git

cd TensorRT-LLM

git lfs install

# 在載入模型前，需要先將模型格式轉為TensorRT-LLM的checkpoint格式

cd examples/llama/

python3 convert_checkpoint.py --model_dir /data/llama-2-7b-ckpt --output_dir llama-2-7b-ckpt-f16 --dtype float16

# 然後編譯模型

trtllm-build --checkpoint_dir ./llama-2-7b-ckpt-f16 \

--remove_input_padding enable \

--gpt_attention_plugin float16 \

--context_fmha enable \

--gemm_plugin float16 \

--output_dir ./trt_engines/llama2-f16-inflight-0.9/ \

--paged_kv_cache enable \

--max_batch_size 64

# 而後即可進行推理測試

python3 ../run.py --engine_dir ./trt_engines/llama2-f16 --max_output_len 100 --tokenizer_dir llama-2-7b-ckpt --input_text "tell a story"

4. TensorRT-LLM原理

與其他部分推理引擎（例如vLLM）不一樣的是：TensorRT-LLM並不是使用原始的weights進行模型推理。而是先對模型進行編譯，並最佳化核心，實現在英偉達GPU上的高效執行。執行編譯後模型的效能比直接執行原始模型的效能要高得多，這也是為什麼TensorRT-LLM速度快的原因之一。

圖片一：原始模型被編譯成最佳化過的二進位制檔案[3]

原始模型的權重，與設定的最佳化選項（例如量化級別、tensor並行度、pipeline並行度等），一起傳遞給編譯器。然後編譯器根據這些資訊輸出特定針對GPU最佳化過的模型二進位制檔案。

需要注意的一個點是：整個模型編譯的過程必須在GPU上進行。生成的編譯模型是專門針對執行它的GPU進行最佳化的。例如，如果是在A40 GPU上編譯的模型，就無法在A100上執行。所以在推理時，必須使用和編譯時的GPU同一個型別。

TensorRT-LLM並不直接支援所有的LLM，因為每個模型的架構都不一樣，TensorRT會做深層圖級別最佳化，所以這就需要對不同模型進行適配。不過目前大部分模型例如Mistral、Llama、chatGLM、Baichuan、Qwen等都是支援的[4]。

TensorRT-LLM的python包使得開發者可以在不需要了解C++或CUDA的情況下以最高效的方式執行LLM。除此之外，它還提供了例如token streaming、paged attention以及KV cache等使用功能，下面我們具體進行介紹。

4.1. Paged Attention

Paged Attention無論是在TensorRT-LLM還是vLLM中，都是一個核心的加速功能。透過Paged Attention高效地管理attention中快取的張量，實現了比HuggingFace Transformers高14-24倍的吞吐量[5]。

4.1.1. LLM的記憶體管理限制效能

在自迴歸解碼過程中，LLM的所有輸入token都會生成attention機制的keys和values的張量，並且這些張量被保留在GPU記憶體中，用來生成下一個token。大語言模型需要大量記憶體來儲存每個token的keys和values，並且隨著輸入序列變長時，所需的記憶體也會擴充套件到非常大的規模。

這些快取的key和value的張量通常稱為KV快取。KV快取有2個特點：

1. 記憶體佔用大：在LLaMA-13B中，單個序列的KV快取佔用高達1.7GB的記憶體
2. 動態化：其大小取決於序列的長度，而序列長度高度易變，且不可預測

在常規的attention機制中，keys和values的值必須連續儲存。因此，即使我們能在給序列分配的記憶體的中間部分做了記憶體釋放，也無法使用這部分空間供給其他序列，這會導致記憶體碎片和浪費。因此，有效管理KV快取是一個重大挑戰。

4.1.2. Paged Attention機制

為了解決上述的記憶體管理問題，Paged Attention的解決辦法非常直接：它允許在非連續的記憶體空間中儲存連續的keys和values。

具體來說，Paged Attention將每個序列的KV快取分為若干個block，每個block包含固定數量token的key和value張量。在注意力計算過程中，Paged Attention核心能夠高效地識別和提取這些block。由於這些block在記憶體中不需要連續，因此也就可以像作業系統的虛擬記憶體一樣，以更靈活的方式管理key和value張量——將block看作page，token看作bytes，序列看作process。

序列的連續邏輯塊透過塊表對映到非連續的物理塊。隨著生成新的token，物理塊會按需進行分配。這種方式可以防止記憶體碎片化，提高記憶體利用率。在生成輸出序列時，page即可根據需要動態分配和釋放。因此,如果我們中間釋放了一些page，那這些空間就可以被重新用於儲存其他序列的KV。

透過這種機制，可以使得極大減少模型在執行復雜取樣演算法（例如parallel sampling和beam search）的記憶體開銷，提升推理吞吐效能。

4.2. KV Cache

KV Cache是LLM推理最佳化裡的一個常用技術，可以在不影響計算精度的情況下，透過空間換時間的辦法，提高推理效能。KV Cache發生在多個token生成的步驟中，並且只發生在Decoder-only模型中（例如GPT。BERT這樣的encoder模型不是生成式模型，而是判別性模型）。

我們知道在生成式模型中，是基於現有的序列生成下一個token。我們給一個輸入文字，模型會輸出一個回答（長度為 N），其實該過程中執行了 N 次推理過程。模型一次推理只輸出一個token，輸出的 token 會與輸入序列拼接在一起，然後作為下一次推理的輸入，這樣不斷反覆直到遇到終止符。

而由於每個token都需要計算其Key和Value（Attention機制），所會存在一個問題：每次生成新的token時，都需要計算之前token的KV，存在大量冗餘計算。而在整個計算過程中，每個token的K、V計算方式是一樣，所以可以把每個token（也就是生成過程中每個token推理時）的K、V放入記憶體進行快取。這樣在進行下一次token預測時，就可以不需要對之前的token KV再進行一次計算，從而實現推理加速。這便是KV Cache的基本原理，如果各位希望瞭解更多的細節，可以繼續閱讀這篇文章[6]。

5. TensorRT-LLM與vLLM對比

vLLM也是開源社群一個非常熱門的推理引擎，上面提到的Page Attention這一重要就是源自vLLM。在使用兩個引擎的部署過程中，筆者對兩者有如下感受：

1. vLLM非常易於使用：vLLM無論是環境準備（pip即可）、模型部署（不需要編譯模型）還是上手難度（文件全面、社群活躍），均優於TensorRT-LLM。TensorRT-LLM的學習成本以及故障排查的難度都要高於vLLM
2. vLLM支援的硬體部署選項更多：TensorRT-LLM作為英偉達推出的框架，主要是針對自家的GPU進行的最佳化，所以僅適用於Nvidia CUDA。而vLLM目前除了支援Nvidia CUDA,外，還支援AMD ROCm, AWS Neuron和CPU的部署
3. 在效能表現上，參考社群在2024年6月使用 LLAMA3-8B 和 LLAMA3-70B with 4bit量化模型做的benchmark[7]測試，可以看到vLLM在TTFT（Time to First Token，即從傳送請求到生成第一個token的時間，在互動式聊天場景較為重要）上表現不俗，優於TensorRT-LLM。但是在Token生成速率上TensorRT-LLM表現更好（特別是在做了4-bit量化後，兩者表現差異更大）

總的來說，如果是團隊有一定的技術能力、僅使用英偉達GPU，以及使用量化後的模型做推理的場景，可以優先考慮TensorRT-LLM，可以獲得更高效的推理效能。

6. Nvidia Triton Inference Server

在TensorRT-LLM裡我們可以看到它實現了模型的載入以及推理的呼叫。但是在生產上，我們不可能僅有這套API工具就足夠，而還需要對應一套Serving服務，也就是一套Server框架，幫助我們部署各種模型，並以API的方式（例如HTTP/REST等）提供模型的推理服務。同時，還要能提供各種監控指標（例如GPU利用率、伺服器吞吐、延遲等），可以幫助使用者瞭解負載情況，併為後續擴縮容的推理服務的實現提供指標基礎。

所以，要建立一個生產可用的LLM部署，可使用例如 Triton Inference Server的Serving服務，結合TensorRT-LLM的推理引擎，利用TensorRT-LLM C++ runtime實現快速推理執行。同時它還提供了in-flight batching 和paged KV caching等推理最佳化。

這裡還需要再強調的是：Triton Inference Server，如其名字所示，它只是一個Server服務，本身不帶模型推理功能，而是提供的Serving的服務。模型推理功能的實現，在Triton裡是透過一個backend的抽象來實現的。TensorRT-LLM就是其中一種backend，可以對接到Triton Inference Server裡，提供最終的模型推理功能。所以，Triton不僅僅是隻能和TensorRT-LLM整合使用，還可以和其他推理引擎整合，例如vLLM。

在對Triton Inference Server有了簡單瞭解後，下面我們介紹如何實現部署。

6.1. 部署推理服務

建立docker容器來部署推理服務

# 準備triton docker

docker run -it --gpus all --network host --shm-size=2g --ulimit memlock=-1 --ulimit stack=67108864 -v /home/ubuntu/data:/data nvcr.io/nvidia/tritonserver:24.05-trtllm-python-py3

# 下載tensorrtllm_backend

git clone https://github.com/triton-inference-server/tensorrtllm_backend.git --branch v0.9.0

cd tensorrtllm_backend

apt-get update && apt-get install git-lfs -y --no-install-recommends

git lfs install

git submodule update --init --recursive

# 將上面編譯過的模型複製到其指定目錄

cd tensorrtllm_backend/

cp /data/TensorRT-LLM/examples/llama/trt_engines/llama2-f16/* all_models/inflight_batcher_llm/tensorrt_llm/1/

# 配置config.pbtxt，指定例如模型的位置、使用的tokenizer、啟用in-flight batching等

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/postprocessing/config.pbtxt tokenizer_dir:meta-llama/Llama-2-7b-chat-hf,tokenizer_type:auto,triton_max_batch_size:64,postprocessing_instance_count:1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/preprocessing/config.pbtxt tokenizer_dir:meta-llama/Llama-2-7b-chat-hf,tokenizer_type:auto,triton_max_batch_size:64,preprocessing_instance_count:1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:64,decoupled_mode:false,bls_instance_count:1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/ensemble/config.pbtxt triton_max_batch_size:64

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt triton_max_batch_size:64,decoupled_mode:true,engine_dir:all_models/inflight_batcher_llm/tensorrt_llm/1,max_queue_delay_microseconds:10000,batching_strategy:inflight_fused_batching

# 啟動serving服務，這裡world_size表示使用多少個GPU用作serving

huggingface-cli login --token xxxx

chmod +r ./all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt

python3 scripts/launch_triton_server.py --world_size 1 --model_repo=./all_models/inflight_batcher_llm/

而後即可正常啟動Triton Server推理服務：

6.2. 使用推理服務

在Triton Inference Server啟動後，即可使用HTTP的方式進行模型推理，例如：

curl -X POST localhost:8000/v2/models/ensemble/generate -d '{"text_input": "What is ML?", "max_tokens": 50, "bad_words": "", "stop_words": "", "pad_id": 2, "end_id": 2}'

{"context_logits":0.0,"cum_log_probs":0.0,"generation_logits":0.0,"model_name":"ensemble","model_version":"1","output_log_probs":[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],"sequence_end":false,"sequence_id":0,"sequence_start":false,"text_output":"What is ML?\n\nMachine learning (ML) is a subfield of artificial intelligence (AI) that involves the use of algorithms and statistical models to enable machines to learn from data, make decisions, and improve their performance on a specific task over time.\n"}

6.3. In-flight Batching

在部署過程中，大家可能有注意到：在指定tensorrt_llm的backend時，我們指定了一個batching_strategy:inflight_fused_batching的配置項。這裡In-flight Batching也是模型推理場景裡的一個最佳化項，一般也稱為Continuous Batching。

首先，Batching（批處理）的基本原理是：將多個獨立的推理請求組合為一個批次，然後一次性提交給模型進行推理，從而提高計算資源的利用率。在GPU上的推理請求可以以4種方式進行Batching：

1. No Batching：每次處理1個請求
2. Static Batching：請求放入batch佇列，在佇列滿時一次性執行這個批次
3. Dynamic Batching：收到請求後放入batch佇列，在佇列滿時或達到一定時間閾值，一次性執行這個批次
4. In-flight Batching/Continuous Batching：請求按照token逐個處理，如果有請求已經完成並釋放了其佔據的GPU資源，則新的請求可以直接進行處理，不需要等待整個批次完成

為什麼需要In-flight Batching？我們可以想象一個例子，在Static Batching的場景下，會有多個請求提交給模型去處理。假設一個批次有2個請求，其中request_a的推理需要生成3個token（假設耗時3s），request_b的推理要生成100個token（假設耗時100s），所以當前這個批次需要耗時100s才能結束。同時又來了request_c的推理請求，此時request_c智慧等待。這樣會造成2個問題：

1. request_a只耗時3s即結束，但是由於它與request_b是同一個批次，所以無法在完成後立即返回結果給客戶端，而是需要等待request_b也完成之後才能一起返回結果
2. 新的請求request_c也必須等待前一個批次完成後（例如當前這個批次耗時100s），才能被模型進行推理

在這種方式下，可以看到，會有GPU資源空閒，得不到完全利用。並且會導致整體推理延遲上升。

而In-flight Batching的方式是：它可以動態修改構成當前批次的請求，即使是這個批次正在執行當中。還是以上面的場景舉例，假設當前過去了3s，request_a已經完成，request_b仍需97s完成。這時候request_a由於已經完成，所以可以直接返回結果給客戶端並結束。而由於request_a釋放了其佔據的資源，request_c的推理請求可以立即被處理。這種方式就提升了整個系統的GPU使用率，並降低了整體推理延遲。對應過程如下圖所示：

圖片二：Iteration Batching[8]

總結

從TensorRT-LLM+Triton這套方法的部署過程來看，還是較為複雜的，需要先對模型進行編譯，還要特別注意環境、引數的設定，稍不注意就無法正常部署。除此之外，英偉達的文件也不夠清晰。總體來說，學習成本比較高，但也帶來了相對應的推理效果提升的收益（相對於vLLM在量化模型推理、以及token生成速度的場景）。但是，在模型推理場景下，推理效果也並不是唯一選擇因素，其他例如框架易用性、不同底層硬體的支援，也需要根據實際情況進行考量，選擇最合適推理引擎。

References

[1] TensorRT-LLM Installing on Linux：https://nvidia.github.io/TensorRT-LLM/installation/linux.html

[2] Llama2-7b-chat： https://huggingface.co/meta-llama/Llama-2-7b-chat-hf

[3] Deploying LLMs Into Production Using TensorRT LLM: https://towardsdatascience.com/deploying-llms-into-production-using-tensorrt-llm-ed36e620dac4

[4] TensorRT-LLM 支援的模型： https://github.com/NVIDIA/TensorRT-LLM/tree/main/tensorrt_llm

[5] 比HuggingFace快24倍！伯克利LLM推理系統開源碾壓SOTA，GPU砍半：https://baijiahao.baidu.com/s?id=1769290162439240324&wfr=spider&for=pc

[6] 大模型推理最佳化技術-KV Cache： https://zhuanlan.zhihu.com/p/700197845

[7] Benchmarking LLM Inference Backends: vLLM, LMDeploy, MLC-LLM, TensorRT-LLM, and TGI https://bentoml.com/blog/benchmarking-llm-inference-backends

[8] Iteration batching (a.k.a. continuous batching) to increase LLM inference serving throughput：https://friendli.ai/blog/llm-iteration-batching/