在 Nvidia Jetson AGX Orin 上使用 TensorRT-LLM 运行 LLM

NVIDIA刚发布了TensorRT-LLM for Jetson:TensorRT-LLM：开启Jetson平台上大语言模型推理的新篇章，这不教程就来了么！

背景：

大型语言模型 （LLM） 推理的关键瓶颈在于 GPU 内存资源短缺。因此，各种加速框架主要强调减少峰值 GPU 内存使用和提高 GPU 利用率。 这就是 Nvidia 的游戏规则改变者 TensorRT-LLM 发挥作用的地方。TensorRT-LLM 是一个开源库，可加速和优化 NVIDIA AI 平台上 LLM 的推理性能。TensorRT 为 NVIDIA GPU 上的推理提供了高度专业化的优化。请参阅 Github 存储库以获取更多示例和文档！

在本指南中，我将以 Llama-3.1-8B-Instruct 模型为例，演示如何在 NVIDIA Jetson AGX Orin 64GB 开发者套件上使用 TensorRT-LLM 部署 LLM 推理引擎。NVIDIA Jetson AGX Orin 开发者套件在基于 Arm 的 CPU 内核和基于 NVIDIA Ampere 架构的 GPU 之间具有统一的内存架构。系统具有 64 GB 的共享内存，由 CPU 和 GPU 共享。

Nvidia JetPack 6.1 是 JetPack 6 的最新生产版本。确保您的 Jetson AGX Orin 开发者套件已使用最新的 JetPack 6.1 进行刷新。

使用 apt show nvidia-jetpack 命令检查当前 jetpack 版本：

Package: nvidia-jetpack
Version: 6.1+b123
Priority: standard
Section: metapackages
Source: nvidia-jetpack (6.1)
Maintainer: NVIDIA Corporation
Installed-Size: 199 kB
Depends: nvidia-jetpack-runtime (= 6.1+b123), nvidia-jetpack-dev (= 6.1+b123)
Homepage: http://developer.nvidia.com/jetson
Download-Size: 29.3 kB
APT-Sources: https://repo.download.nvidia.com/jetson/common r36.4/main arm64 Packages
Description: NVIDIA Jetpack Meta Package

您可以使用预构建的 Docker 映像，因此您可以按照 Nvidia Jetson AI Lab 页面上的文档轻松入门：https://www.jetson-ai-lab.com/tensorrt_llm.html

为了增加 VRAM，我在 Jetson AGX Orin 上禁用了桌面 GUI。

步骤 1.获取模型权重

我们需要下载我们将使用的模型的权重，即 Meta-Llama-3.1-8B-Instruct。因此，请确保您已接受许可条款并生成相应的 HuggingFace 令牌以下载模型。

首先，您需要通过在终端中运行以下命令来设置 Git LFS：

sudo apt-get update && sudo apt-get -y install git-lfs

通过运行以下命令安装 Git LFS：

git lfs install

使用以下命令克隆 Llama-3.1-8B 模型存储库：

git clone https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct

或者，您可以使用 Hugging Face CLI 下载模型。使用以下命令安装 CLI：

pip install -U "huggingface_hub[cli]"

使用 huggingface-cli login 命令验证您的 Hugging Face 帐户。输入您的 Hugging Face API 密钥。

使用以下命令下载 Llama-3.1-8B-Instruct 模型：

huggingface-cli download meta-llama/Llama-3.1-8B-Instruct --local-dir ./Llama-3.1-8B-Instruct

下载完成后，您可以使用以下命令验证下载文件夹的内容：

ls -l ./Llama-3.1-8B-Instruct/
total 15693184
-rw-rw-r-- 1 jetson jetson        826 Nov 16 21:22 config.json
-rw-rw-r-- 1 jetson jetson        185 Nov 16 21:22 generation_config.json
-rw-rw-r-- 1 jetson jetson       7627 Nov 16 21:21 LICENSE
-rw-rw-r-- 1 jetson jetson 4976698672 Nov 16 21:30 model-00001-of-00004.safetensors
-rw-rw-r-- 1 jetson jetson 4999802720 Nov 16 21:30 model-00002-of-00004.safetensors
-rw-rw-r-- 1 jetson jetson 4915916176 Nov 16 21:29 model-00003-of-00004.safetensors
-rw-rw-r-- 1 jetson jetson 1168138808 Nov 16 21:24 model-00004-of-00004.safetensors
-rw-rw-r-- 1 jetson jetson     23950 Nov 16 21:22 model.safetensors.index.json
drwxrwxr-x 2 jetson jetson       4096 Nov 16 21:42 original
-rw-rw-r-- 1 jetson jetson      40883 Nov 16 21:21 README.md
-rw-rw-r-- 1 jetson jetson         73 Nov 16 21:22 special_tokens_map.json
-rw-rw-r-- 1 jetson jetson      50500 Nov 16 21:22 tokenizer_config.json
-rw-rw-r-- 1 jetson jetson    9085658 Nov 16 21:22 tokenizer.json
-rw-rw-r-- 1 jetson jetson       4691 Nov 16 21:22 USE_POLICY.md

第 2 步：准备

使用 venv 模块创建新的虚拟环境，以隔离您的项目依赖项。

python3 -m venv tensorrt-llm

激活新创建的虚拟环境

source tensorrt-llm/bin/activate

更新包并安装所需的包

sudo apt-get update
sudo apt-get install -y python3-pip libopenblas-dev ccache

下载 CuSparseLT 安装脚本：

wget https://raw.githubusercontent.com/pytorch/pytorch/9b424aac1d70f360479dd919d6b7933b5a9181ac/.ci/docker/common/install_cusparselt.sh

设置 CUDA 版本：通过运行以下命令将 CUDA 版本设置为 12.6：

export CUDA_VERSION=12.6

通过运行以下命令安装 CuSparseLT：

sudo -E bash ./install_cusparselt.sh

此过程可能需要一些时间才能完成。确保安装成功完成。

最后，使用 pip 安装 NumPy 版本 1.26.1：bash 复制代码

python3 -m pip install numpy=='1.26.1'

此步骤完成准备阶段，设置后续步骤所需的环境和依赖项。

第 3 步：构建 TensorRT-LLM 引擎

使用以下命令从 GitHub 克隆 TensorRT-LLM 存储库：

git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM
git checkout v0.12.0-jetson
git lfs pull

然后，执行以下命令，为 TensorRT-LLM 构建一个 wheel 文件：

 sudo python3 scripts/build_wheel.py --clean --cuda_architectures 87 -DENABLE_MULTI_DEVICE=0 --build_type Release --benchmarks --use_ccache -DCMAKE_CUDA_COMPILER=/usr/local/cuda-12.6/bin/nvcc

此命令为 TensorRT-LLM 构建一个 wheel 文件。可能需要一些时间才能完成。

您将看到如下所示的输出：

Successfully built tensorrt_llm-0.12.0-cp310-cp310-linux_aarch64.whl

使用 pip 安装构建的 wheel 文件：

pip3 install build/tensorrt_llm-*.whl

预期输出：

Successfully installed tensorrt-llm-0.12.0

通过导入库并打印其版本来验证安装：

python3 -c "import tensorrt_llm; print(tensorrt_llm.__version__)"

要在 GPU 上高效运行这些模型，我们必须将它们转换为 TensorRT-LLM 格式。然后，我们将使用 trtllm-build 命令行工具从 Hugging Face 检查点构建优化的 TensorRT 引擎。

HuggingFace 模型的转换可以通过一个命令完成：

sudo python /path/to/TensorRT-LLM/examples/llama/convert_checkpoint.py \
--model_dir /path/to/Llama-3.1-8B-Instruct \
--output_dir /path/to/Llama-3.1-8B-Instruct-convert \
--dtype float16

您将看到如下所示的输出：

[TensorRT-LLM] TensorRT-LLM version: 0.12.0
0.12.0
230it [00:01, 124.05it/s]
Total time of converting checkpoints: 00:00:30

这应该会产生两个文件：模型配置 （config.json） 和权重 （rank0.safetensors）。接下来，我们构建模型引擎：

sudo trtllm-build \
--checkpoint_dir /path/to/Llama-3.1-8B-Instruct-convert \
--gpt_attention_plugin float16 \
--gemm_plugin float16 \
--output_dir /path/to/Llama-3.1-8B-Instruct-engine

此命令使用转换后的检查点构建 TensorRT 引擎，并将结果存储在指定的目录中。

如果构建成功，您应该会看到以下输出：

[11/17/2024-16:26:23] [TRT-LLM] [I] Build phase peak memory: 32793.48 MB, children: 0.00 MB
[11/17/2024-16:26:23] [TRT-LLM] [I] Serializing engine to /home/jetson/Projects/tensorrtllm/Llama-3.1-8B-final/rank0.engine...
[11/17/2024-16:26:44] [TRT-LLM] [I] Engine serialized. Total time: 00:00:20
[11/17/2024-16:26:45] [TRT-LLM] [I] Total time of building all engines: 00:01:06

这应该会产生两个文件：模型配置 （config.json） 和权重 （rank0.engine）。


第 4 步：在 NVIDIA Jetson AGX Orin 64GB 开发者套件上运行推理

构建模型引擎后，您可以通过使用以下命令运行模型来测试它：

sudo python3 /path/to/TensorRT-LLM/examples/run.py \
--engine_dir /path/to/Llama-3.1-8B-Instruct-engine \
--max_output_len 100 \
--max_attention_window_size 1024 \
--tokenizer_dir /path/to/Llama-3.1-8B-Instruct \
--input_text "Kazakhstan is" \
--gpu_weights_percent 70 \
--kv_cache_free_gpu_memory_fraction 0.1 \
--num_beams 1

如果模型成功运行，您应该会看到以下输出：

当令牌流式传输时，可以明显地注意到 TensorRT LLM 的性能。

然后我们可以在 OpenAI 兼容模式下对 TensorRT LLM 服务器进行 TensorRT LLM 服务器。运行以下命令：

sudo python3 /path/to/TensorRT-LLM/examples/apps/openai_server.py \
/path/to/Llama-3.1-8B-Instruct-engine \
--tokenizer /path/to/Llama-3.1-8B-Instruct \
--host 0.0.0.0 \
--port 5001

使用 --host 选项，您可以允许外部连接。

您可以通过使用 curl 命令发送请求来测试模型的推理。

 curl http://localhost:5001/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "meta-llama/Llama-3.1-8B",
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Who won the world series in 2020?"}
        ]
    }'

如果响应成功，您应该会看到以下输出：

{"id":"chatcmpl-869061ee5db04f8ca9f4d0b870c7de51","object":"chat.completion","created":1732093982,"model":"meta-llama/Llama-3.1-8B","choices":[{"index":0,"message":{"role":"assistant","content":"The Los Angeles Dodgers won the 2020 World Series, defeating the Tampa Bay","tool_calls":[]},"logprobs":null,"finish_reason":null,"stop_reason":null}],"usage":{"prompt_tokens":52,"total_tokens":68,"completion_tokens":16}}

您可以使用类似于 OpenAI 的客户端函数来调用 TensorRT-LLM 服务。

from openai import OpenAI

openai_api_key = "EMPTY"
openai_api_base = "http://localhost:5001/v1"

client = OpenAI(
    api_key=openai_api_key,
    base_url=openai_api_base,
)

chat_response = client.chat.completions.create(
    model="meta-llama/Llama-3.1-8B-Instruct",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Tell me a about Kazakhstan."},
    ]
)
print("Chat response:", chat_response)

如果响应如下，则部署成功。

Chat response: ChatCompletion(id='chatcmpl-b71842ec0407465b9b5ac32130bfd356', choices=[Choice(finish_reason=None, index=0, logprobs=None, message=ChatCompletionMessage(content='Kazakhstan is a country located in Central Asia, bordered by Russia to the', refusal=None, role='assistant', function_call=None, tool_calls=[]), stop_reason=None)], created=1732086771, model='meta-llama/Llama-3.1-8B-Instruct', object='chat.completion', service_tier=None, system_fingerprint=None, usage=CompletionUsage(completion_tokens=16, prompt_tokens=48, total_tokens=64, completion_tokens_details=None, prompt_tokens_details=None))

创建 Streamlit Web 应用程序以与 TensorRT-LLM 服务交互

import streamlit as st
from openai import OpenAI

st.title("TensorRT-LLM Demo on the NVIDIA Jetson AGX Orin Developer Kit ")

client = OpenAI(base_url="http://localhost:5001/v1", api_key="None")

if "messages" not in st.session_state:
    st.session_state["messages"] = []

prompt = st.chat_input("Say something")
if prompt:
    st.session_state["messages"].append({"role": "user", "content": prompt})
    for message in st.session_state["messages"]:
        st.chat_message(message["role"]).write(message["content"])
    container = st.empty()
    chat_completion = client.chat.completions.create(
        stream=True,
        messages=st.session_state["messages"],
        model="ensemble", 
        max_tokens=512
    )
    response = ""
    for event in chat_completion:
        content = event.choices[0].delta.content
        if content:
            response += content
        container.chat_message("assistant").write(response)
    st.session_state["messages"].append({"role": "assistant", "content": response})

在这篇博文中，我的目标是演示如何在 NVIDIA Jetson AGX Orin 64GB 开发者套件上使用 TensorRT LLM 实现最先进的推理。我涵盖了从编译 LLM 到使用 Streamlit 在生产环境中部署模型的所有内容。