解鎖NVIDIA TensorRT-LLM的卓越性能

NVIDIA TensorRT-LLM 是一個(gè)專為優(yōu)化大語(yǔ)言模型 (LLM) 推理而設(shè)計(jì)的庫(kù)。它提供了多種先進(jìn)的優(yōu)化技術(shù)，包括自定義 Attention Kernel、Inflight Batching、Paged KV Caching、量化技術(shù) (FP8、INT4 AWQ、INT8 SmoothQuant 等) 以及更多功能，確保您的 NVIDIA GPU 能發(fā)揮出卓越的推理性能。

我們深知您對(duì)易用性的需求，為了讓您更快上手，并迅速實(shí)現(xiàn)流行模型的高性能推理，我們開發(fā)了 LLM API，通過簡(jiǎn)潔的指令，您可輕松體驗(yàn) TensorRT-LLM 帶來的卓越性能！

LLM API 是一個(gè) high-level Python API，專為 LLM 推理工作流而設(shè)計(jì)。以下是一個(gè)展示如何使用 TinyLlama 的簡(jiǎn)單示例：

from tensorrt_llm import LLM, SamplingParams
 
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 
llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 
outputs = llm.generate(prompts, sampling_params)
 
# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
print(f"Prompt:{prompt!r},Generatedtext:{generated_text!r}")

希望以上示例能幫助您快速入門 NVIDIA TensorRT-LLM LLM API。

當(dāng)前 TensorRT-LLM LLM API 可支持的模型

Llama (including variants Mistral, Mixtral, InternLM)

GPT (including variants Starcoder-1/2, Santacoder)

Gemma-1/2

Phi-1/2/3

ChatGLM (including variants glm-10b, chatglm, chatglm2, chatglm3, glm4)

QWen-1/1.5/2

Falcon

Baichuan-1/2

GPT-J

Mamba-1/2

一、詳細(xì)介紹

1.1 模型準(zhǔn)備

LLM class 可支持以下三種模型導(dǎo)入來源：

Hugging Face Hub：直接從 Hugging Face 模型庫(kù)下載模型，例如 TinyLlama/TinyLlama-1.1B-Chat-v1.0。

本地 Hugging Face 模型：使用已下載到本地的 Hugging Face 模型。

本地 TensorRT-LLM 引擎：使用通過 trtllm-build 工具構(gòu)建或由 Python LLM API 保存的 Engine。

您可以使用 LLM(model=) 構(gòu)造函數(shù)來靈活切換這些格式。以下各節(jié)將詳細(xì)介紹具體使用方法。

Hugging Face Hub

使用 Hugging Face Hub 來導(dǎo)入模型非常直觀，只需在 LLM 構(gòu)造函數(shù)中指定模型倉(cāng)庫(kù)名稱即可：

llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")

本地 Hugging Face 模型

由于 Hugging Face 模型庫(kù)的廣泛應(yīng)用，API 將 Hugging Face 格式作為主要輸入來源之一。當(dāng)您要使用 Llama 3.1 模型時(shí)，請(qǐng)通過以下命令從 Meta Llama 3.1 8B 模型頁(yè)面下載模型：

git lfs install
gitclone

下載完成后，您可以通過以下方式加載模型：

llm = LLM(model=)

請(qǐng)注意，使用此模型需要遵守特定許可條款：

https://ai.meta.com/resources/models-and-libraries/llama-downloads/

在開始下載之前，請(qǐng)確保同意這些條款并在 Hugging Face 上完成身份驗(yàn)證：

https://huggingface.co/meta-llama/Meta-Llama-3-8B?clone=true

本地 TensorRT-LLM 引擎

LLM API 支持使用 TensorRT-LLM Engine，您可以通過以下兩種方式構(gòu)建 Engine：

您可使用 trtllm-build 工具從 Hugging Face 模型直接構(gòu)建 TensorRT-LLM Engine，并將其保存到磁盤供后續(xù)使用。詳細(xì)說明請(qǐng)參考 GitHub 上的 README 和 examples/llama 倉(cāng)庫(kù)。構(gòu)建完成后，您可以這樣加載模型：

llm = LLM(model=)

或者，您可以使用 LLM instance 來創(chuàng)建 Engine 并保存到本地磁盤：

llm = LLM()
# Save engine to local disk
llm.save()

可以參考第一種方法使用 model 參數(shù)來加載 Engine。

1.2 使用技巧和故障排除

以下是針對(duì)熟悉 TensorRT-LLM 其他 API 的用戶，在剛開始使用 LLM API 時(shí)可能遇到的常見問題及其解決方案：

RuntimeError: only rank 0 can start multi-node session, got 1

在使用 LLM API 進(jìn)行單節(jié)點(diǎn)多 GPU 推理時(shí)，無需添加 mpirun 前綴。您可以直接運(yùn)行 python llm_inference_distributed.py 來執(zhí)行多 GPU 推理。

Slurm 節(jié)點(diǎn)上的掛起問題

在使用 Slurm 管理的節(jié)點(diǎn)上遇到掛起或其他問題時(shí)，請(qǐng)?jiān)趩?dòng)腳本中添加前綴 mpirun -n 1 --oversubscribe --allow-run-as-root。

示例命令：

mpirun-n1--oversubscribe--allow-run-as-rootpythonllm_inference_distributed.py

在通訊器 MPI_COMM_WORLD 中，MPI_ABORT 在 rank 1 上被調(diào)用，錯(cuò)誤代碼為 1。

由于 LLM API 依賴 mpi4py 庫(kù)，為避免 mpi4py 中的遞歸生成進(jìn)程，請(qǐng)將 LLM 類放在函數(shù)中，并在 __main__ 命名空間下保護(hù)程序的主入口點(diǎn)。

注意：此限制僅適用于多 GPU 推理。

二、常見自定義操作

2.1 量化

TensorRT-LLM 可以通過在 LLM 實(shí)例中設(shè)置適當(dāng) Flags，自動(dòng)對(duì) Hugging Face 模型進(jìn)行量化。例如，要執(zhí)行 Int4 AWQ 量化，以下代碼會(huì)觸發(fā)模型量化。請(qǐng)參考完整的支持的標(biāo)志列表和可接受的值。

fromtensorrt_llm.llmapiimportQuantConfig,QuantAlgo
quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_AWQ)
llm=LLM(,quant_config=quant_config)

2.2 采樣

SamplingParams 可以自定義采樣策略以控制 LLM 生成的響應(yīng)，如 Beam Search、Temperature 和其他參數(shù)。

例如，要啟用 Beam Size 為 4 的 Beam Search，請(qǐng)按如下方式設(shè)置 Sampling_Params：

from tensorrt_llm.llmapi import LLM, SamplingParams, BuildConfig


build_config = BuildConfig()
build_config.max_beam_width = 4


llm = LLM(, build_config=build_config)
# Let the LLM object generate text with the default sampling strategy, or
# you can create a SamplingParams object as well with several fields set manually
sampling_params = SamplingParams(beam_width=4) # current limitation: beam_width should be equal to max_beam_width


for output in llm.generate(, sampling_params=sampling_params):
print(output)

SamplingParams 管理并分發(fā)字段到 C++ classes，包括：

SamplingConfig:

https://nvidia.github.io/TensorRT-LLM/_cpp_gen/runtime.html#_CPPv4N12tensorrt_llm7runtime14SamplingConfigE

OutputConfig:

https://nvidia.github.io/TensorRT-LLM/_cpp_gen/executor.html#_CPPv4N12tensorrt_llm8executor12OutputConfigE

更多詳情請(qǐng)參考 class 文檔：

https://nvidia.github.io/TensorRT-LLM/llm-api/reference.html#tensorrt_llm.llmapi.SamplingParams

2.3 Build 配置

除了上述參數(shù)外，您還可以使用 build_config 類和從 trtllm-build CLI 借用的其他參數(shù)來自定義構(gòu)建配置。這些構(gòu)建配置選項(xiàng)為目標(biāo)硬件和用例構(gòu)建 Engine 提供了靈活性。請(qǐng)參考以下示例：

11m = LLM(,
          build_config=Buildconfig(
            max_num_tokens=4096,
            max batch size=128,
            max_beam_width=4))

?

更多詳情請(qǐng)參考 buildconfig 文檔：

https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/builder.py#L476-L509

2.4 自定義 Runtime

類似于 build_config，您也可以使用 runtime_config、peft_cache_config，或其他從 Executor API 借用的參數(shù)來自定義 Runtime 配置。這些 Runtime 配置選項(xiàng)在 KV cache management、GPU memory allocation 等方面提供了額外的靈活性。請(qǐng)參考以下示例：

fromtensorrt_llm.llmapiimportLLM,KvCacheConfig


llm=LLM(,
kv_cache_config=KvCacheConfig(
free_gpu_memory_fraction=0.8))

2.5 自定義 Tokenizer

默認(rèn)情況下，LLM API 使用 transformers 的 AutoTokenizer。您可以在創(chuàng)建 LLM 對(duì)象時(shí)傳入自己的分詞器來覆蓋它。請(qǐng)參考以下示例：

llm=LLM(,tokenizer=)

?

LLM () 工作流將使用您的 tokenizer 。

也可以使用以下代碼直接輸入 token ID，由于未使用 Tokenizers，該代碼生成的是不帶文本的 token ID。

llm=LLM()
foroutputinllm.generate([32,12]):
...

關(guān)于作者

嚴(yán)春偉

NVIDIA 性能架構(gòu)師，目前主要聚焦于大模型推理架構(gòu)和優(yōu)化。

張國(guó)銘

NVIDIA 性能架構(gòu)師，目前主要從事大模型推理架構(gòu)和優(yōu)化 。

Adam Zheng

NVIDIA 產(chǎn)品經(jīng)理，負(fù)責(zé) NVIDIA AI 平臺(tái)軟件產(chǎn)品管理，目前主要聚焦于大模型推理架構(gòu)和優(yōu)化。