Nvidia Triton¶

NVIDIA Triton 推理服务器 提供了一个为 CPU 和 GPU 优化的云和边缘推理解决方案。该连接器允许 llama_index 远程与部署在 Triton 上的 TRT-LLM 模型交互。

启动 Triton 推理服务器¶

该连接器需要运行 Triton 推理服务器并带有 TensorRT-LLM 模型。在此示例中，我们将使用 Triton 命令行接口 (Triton CLI) 在 Triton 上部署一个 GPT2 模型。

在主机上使用 Triton 和相关工具（Triton 容器镜像外部）时，各种工作流程可能需要一些额外的依赖项。大多数系统依赖项问题可以通过在最新的相应 tritonserver 容器镜像内部安装和运行 CLI 来解决，该容器镜像应该已经安装了所有必要的系统依赖项。

对于 TRT-LLM，您可以使用 nvcr.io/nvidia/tritonserver:{YY.MM}-trtllm-python-py3 镜像，其中 YY.MM 对应于 tritonserver 的版本，例如在本例中，我们使用的是容器的 24.02 版本。要获取可用版本的列表，请参考 Triton 推理服务器 NGC。

要启动容器，请在 Linux 终端中运行

docker run -ti --gpus all --network=host --shm-size=1g --ulimit memlock=-1 nvcr.io/nvidia/tritonserver:24.02-trtllm-python-py3

接下来，我们需要安装以下依赖项

pip install \
  "psutil" \
  "pynvml>=11.5.0" \
  "torch==2.1.2" \
  "tensorrt_llm==0.8.0" --extra-index-url https://pypi.nvidia.com/

最后，运行以下命令安装 Triton CLI。

pip install git+https://github.com/triton-inference-server/triton_cli.git

生成 GPT2 模型的模型仓库并启动一个 Triton 服务器实例

triton remove -m all
triton import -m gpt2 --backend tensorrtllm
triton start &

请注意，默认情况下 Triton 开始监听 localhost:8000 HTTP 端口和 localhost:8001 GRPC 端口。本示例将使用后者。有关更多操作方法和问题，请访问 Triton 命令行接口 (Triton CLI) 问题列表。

安装 tritonclient¶

由于我们与 Triton 推理服务器交互，我们需要 安装 tritonclient 包。

pip install tritonclient[all]

接下来，我们将安装 llama index 连接器。

pip install llama-index-llms-nvidia-triton

基本用法¶

使用提示调用 complete¶

from llama_index.llms.nvidia_triton import NvidiaTriton

# A Triton server instance must be running. Use the correct URL for your desired Triton server instance.
triton_url = "localhost:8001"
model_name = "gpt2"
resp = NvidiaTriton(server_url=triton_url, model_name=model_name, tokens=32).complete("The tallest mountain in North America is ")
print(resp)

您应该会收到以下响应

the Great Pyramid of Giza, which is about 1,000 feet high. The Great Pyramid of Giza is the tallest mountain in North America.

使用提示调用 stream_complete¶

resp = NvidiaTriton(server_url=triton_url, model_name=model_name, tokens=32).stream_complete("The tallest mountain in North America is ")
for delta in resp:
    print(delta.delta, end=" ")

您应该会收到以下流式响应

the Great Pyramid of Giza, which is about 1,000 feet high. The Great Pyramid of Giza is the tallest mountain in North America.

更多示例¶

有关 Triton 推理服务器的更多信息，请参考 快速入门 指南、NVIDIA 开发者 Triton 页面 和 GitHub Issues 频道。