主题
NVIDIA Triton Inference Server 提供针对 CPU 和 GPU 优化的云端和边缘推理解决方案。此连接器允许 LlamaIndex 远程与使用 Triton 部署的 TensorRT-LLM 模型交互。
启动 Triton Inference Server
此连接器需要一个运行中的 Triton Inference Server 实例,并加载 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 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 Server 实例,请运行以下命令:
triton remove -m all
triton import -m gpt2 --backend tensorrtllm
triton start &默认情况下,Triton 监听 localhost:8000(HTTP)和 localhost:8001(gRPC)。本示例使用 gRPC 端口。
更多信息请参阅 Triton CLI GitHub 仓库。
安装 tritonclient
由于我们要与 Triton Inference Server 交互,你需要安装 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 Inference Server 的更多信息,请参阅以下资源: