英伟达 Triton¶
NVIDIA Triton 推理服务器 提供针对 CPU 和 GPU 优化的云端及边缘推理解决方案。该连接器支持 llama_index 远程与通过 Triton 部署的 TRT-LLM 模型进行交互。
启动 Triton 推理服务器¶
该连接器需要运行一个搭载 TensorRT-LLM 模型的 Triton 推理服务器实例。
本示例中，我们将使用 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 Inference Server 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) 的 issues 页面。
pip install tritonclient[all]
接下来，我们将安装 llama index 连接器。
pip install llama-index-llms-nvidia-triton
基本用法¶
调用 complete 方法并传入提示词¶
from llama_index.llms.nvidia_triton import NvidiaTriton
# 必须运行一个Triton服务器实例。请根据您所需的Triton服务器实例填写正确的URL。
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 问题反馈频道。