在Nvidia Jetson ADX Orin中使用TensorRT-LLM运行llama3-8b

背景：

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

在本指南中，我将以 Llama-3-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-8B-Instruct。因此，请确保您已接受许可条款并生成相应的 HuggingFace 令牌以下载模型。
首先，您需要通过在终端中运行以下命令来设置 Git LFS：

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

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

git lfs install

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

https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct

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

pip install -U "huggingface_hub[cli]"

使用 huggingface-cli login 命令验证您的 Hugging Face 帐户。输入您的 Hugging Face API 密钥。
使用以下命令下载 Llama-3-8B-Instruct 模型：

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

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

wangyh@ubuntu:/media/wangyh/My Passport/llama3-8b-instruct$ ls -l
total 15695249
-rwxrwxrwx 1 wangyh wangyh       9802  1月  6 09:36 added_tokens.yaml
-rwxrwxrwx 1 wangyh wangyh        654  1月  6 09:36 config.json
-rwxrwxrwx 1 wangyh wangyh        187  1月  6 09:36 generation_config.json
-rwxrwxrwx 1 wangyh wangyh    2183982  1月  6 09:36 llama3_8b.tiktoken
-rwxrwxrwx 1 wangyh wangyh 4976698672  1月  6 09:38 model-00001-of-00004.safetensors
-rwxrwxrwx 1 wangyh wangyh 4999802720  1月  6 09:40 model-00002-of-00004.safetensors
-rwxrwxrwx 1 wangyh wangyh 4915916176  1月  6 09:42 model-00003-of-00004.safetensors
-rwxrwxrwx 1 wangyh wangyh 1168138808  1月  6 09:42 model-00004-of-00004.safetensors
-rwxrwxrwx 1 wangyh wangyh      23950  1月  6 09:42 model.safetensors.index.json
drwxrwxrwx 1 wangyh wangyh          0  1月  6 09:42 original
-rwxrwxrwx 1 wangyh wangyh         73  1月  6 09:42 special_tokens_map.json
-rwxrwxrwx 1 wangyh wangyh      50977  1月  6 09:42 tokenizer_config.json
-rwxrwxrwx 1 wangyh wangyh    9085698  1月  6 09:42 tokenizer.json

第 2 步：准备

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

python3 -m venv tensorrt-llm-env

激活新创建的虚拟环境

source tensorrt-llm-env/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 submodule update --init --recursive
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-0.12.0-cp310-cp310-linux_aarch64.whl

预期输出：

Successfully installed tensorrt-llm-0.12.0

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

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

【注】到这一步当在虚拟环境下面使用tensorrt_llm库的时候会报错没有找到“tensorrt_llm”，如下图所示：
在这里插入图片描述
提示缺少Python版本的tensorrt库，但是jetpack环境里面没有留下Python版tensorrt的安装包，这个时候有两种方法：

到网上下载jetson版本的TensorRT-10.3的.wheel文件，然后自己手动安装
将本地/usr路径下的tensorrt拷贝出来

这里我们选择第二种方法

cd ~/tensorrt-llm-env/lib/python3.10/site-packages
cp -r /usr/lib/python3.10/dist-packages/tensorrt* .

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

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

python examples/llama/convert_checkpoint.py --model_dir /media/wangyh/My\ Passport/llama3-8b-instruct --output_dir /media/wangyh/My\ Passport/Llama-3.1-8B-Instruct-convert --dtype float16

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

[TensorRT-LLM] TensorRT-LLM version: 0.12.0
0.12.0
230it [00:04, 50.42it/s]
Total time of converting checkpoints: 00:03:11

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

trtllm-build --checkpoint_dir /media/wangyh/My\ Passport/Llama-3.1-8B-Instruct-convert --gpt_attention_plugin float16 --gemm_plugin float16 --output_dir /media/wangyh/My\ Passport/Llama-3.1-8B

此命令使用转换后的检查点构建 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 开发者套件上运行推理
构建模型引擎后，您可以通过使用以下命令运行模型来测试它：

python examples/run.py --engine_dir /media/wangyh/My\ Passport/Llama-3.1-8B-Instruct-engine --max_output_len 100 --max_attention_window_size 1024 --tokenizer_dir /media/wangyh/My\ Passport/llama3-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 服务器。运行以下命令：

python examples/apps/openai_server.py /media/wangyh/My\ Passport/Llama-3.1-8B-Instruct-engine --tokenizer /media/wangyh/My\ Passport/llama3-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-d2c4ab76cea845fa80374af3133d4ef4","object":"chat.completion","created":1736154766,"model":"meta-llama/Llama-3.1-8B","choices":[{"index":0,"message":{"role":"assistant","content":"The Los Angeles Dodgers won the World Series in 2020! They defeated the","tool_calls":[]},"logprobs":null,"finish_reason":null,"stop_reason":null}],"usage":{"prompt_tokens":32,"total_tokens":48,"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)

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

(tensorrt-llm-env) wangyh@ubuntu:~/TensorRT-LLM$ python tensorrt_llm_orin_openai_request.py 
Chat response: ChatCompletion(id='chatcmpl-2123732b12b84b2f8145272da88bb0f7', choices=[Choice(finish_reason=None, index=0, logprobs=None, message=ChatCompletionMessage(content='Kazakhstan! A vast and fascinating country located in Central Asia, bordering', refusal=None, role='assistant', audio=None, function_call=None, tool_calls=[]), stop_reason=None)], created=1736154822, model='meta-llama/Llama-3.1-8B-Instruct', object='chat.completion', service_tier=None, system_fingerprint=None, usage=CompletionUsage(completion_tokens=16, prompt_tokens=28, total_tokens=44, 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 在生产环境中部署模型的所有内容。