使用Triton部署Hugging Face模型

关键要点

部署Hugging Face的embedding模型到Triton需要导出为ONNX格式，并设置模型仓库。
使用Optimum工具可以简化模型导出过程，Triton支持ONNX后端。
启动Triton服务器并使用客户端发送请求是标准流程，可能需要手动配置输入输出。

部署Hugging Face Embedding模型到Triton的步骤

概述

要将Hugging Face上的embedding模型（如SentenceTransformer模型）部署到Triton Inference Server上，需要经过几个步骤，包括模型导出、仓库设置、服务器启动和推理请求发送。以下是简洁的指导，适合初学者。

步骤分解

1. 导出模型到ONNX格式

首先，使用Hugging Face的Optimum工具将模型导出为ONNX格式。这一步会生成一个包含模型文件和分词器配置的目录。例如：
- 安装Optimum：pip install optimum[exporters]。
- 导出命令：optimum-cli export onnx --model sentence-transformers/all-MiniLM-L6-v2 onnx_model。
意外的细节是，SentenceTransformer模型的池化层（pooling layer）会自动包含在ONNX模型中，简化了后续步骤。

2. 设置Triton模型仓库

创建一个目录结构，例如model_repository/embedding_model/1/model.onnx，并添加配置文件config.pbtxt。
config.pbtxt可以手动编写，指定输入（如input_ids和attention_mask）和输出（如sentence_embedding），也可以让Triton自动生成。

3. 启动Triton服务器

使用Docker运行Triton服务器，挂载模型仓库目录。例如：
- 命令：docker run --gpus all --rm -p 8000:8000 -v ${PWD}/model_repository:/models nvcr.io/nvidia/tritonserver:23.04-py3 tritonserver --model-repository=/models。
确保服务器启动成功，可以通过访问http://localhost:8000/v2/health/ready验证。

4. 发送推理请求

使用Triton的Python客户端发送请求，需要先分词输入文本，准备input_ids和attention_mask，然后获取embedding输出。例如：
- 安装客户端：pip install tritonclient[all]。
- 示例代码见下文，输出为numpy数组，形状为(batch_size, embedding_dim)。

详细技术报告

引言

部署Hugging Face的embedding模型到NVIDIA的Triton Inference Server是一种高效的方式，适合生产环境中的语义搜索和相似性匹配任务。本报告详细介绍从模型导出到推理请求的完整流程，涵盖了Optimum工具的使用、Triton配置和客户端交互。

模型导出到ONNX格式

Hugging Face的embedding模型（如sentence-transformers/all-MiniLM-L6-v2）通常基于Transformer架构，适合导出为ONNX格式以在Triton上运行。研究显示，Optimum库提供了便捷的导出工具，支持PyTorch和TensorFlow模型。

安装Optimum：
- 命令：pip install optimum[exporters]。
- 确保环境支持ONNX Runtime。
导出过程：
- 使用optimum-cli export onnx命令导出模型。例如：
  bash 复制代码
```
optimum-cli export onnx --model sentence-transformers/all-MiniLM-L6-v2 onnx_model
```
- 输出目录onnx_model包含：
  - model.onnx：ONNX格式的模型文件。
  - tokenizer.json：分词器配置文件。
- SentenceTransformer模型的池化层（如均值池化）会自动包含在ONNX模型中，简化了后续推理逻辑。证据支持这一点，特别是在Export Sentence Transformer models to ONNX中提到优化和量化选项。
注意事项：
- 如果模型不是SentenceTransformer（如直接的Transformer模型），也适用相同命令，但需确认支持的架构。
- 使用Netron工具(Netron)查看ONNX模型结构，确认输入输出名称和维度。

设置Triton模型仓库

Triton Inference Server使用模型仓库管理模型，目录结构必须符合规范。

目录结构：

示例：

arduino 复制代码

model_repository/
└── embedding_model
    ├── 1
    │   └── model.onnx
    └── config.pbtxt

将步骤1导出的model.onnx复制到1/目录下。

配置config.pbtxt：
- 可以手动编写，也可让Triton自动生成。以下是手动配置的示例：
  yaml 复制代码
```
name: "embedding_model"
platform: "onnxruntime_onnx"
max_batch_size: 4
input [
  {
    name: "input_ids"
    data_type: TYPE_INT64
    dims: [ -1, -1 ]
  },
  {
    name: "attention_mask"
    data_type: TYPE_INT64
    dims: [ -1, -1 ]
  }
]
output [
  {
    name: "sentence_embedding"
    data_type: TYPE_FP32
    dims: [ -1, 384 ]
  }
]
```
- 字段解释 ：
  - name：模型名称，与目录名一致。
  - platform：指定为onnxruntime_onnx，表示使用ONNX Runtime后端。
  - input：输入张量名称和形状，通常为input_ids和attention_mask，dims: [ -1, -1 ]表示批处理大小和序列长度可变。
  - output：输出张量名称和形状，对于all-MiniLM-L6-v2，输出维度为384。
- 如果不确定输入输出名称，查看ONNX模型结构后调整。

启动Triton服务器

使用Docker容器启动Triton服务器，确保挂载模型仓库。

拉取镜像：
- 命令：docker pull nvcr.io/nvidia/tritonserver:23.04-py3。
- 版本号可根据需求调整，当前推荐使用最新稳定版。
启动命令：
- 示例：
  bash 复制代码
```
docker run --gpus all --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v ${PWD}/model_repository:/models nvcr.io/nvidia/tritonserver:23.04-py3 tritonserver --model-repository=/models
```
- 参数解释：
  - --gpus all：启用所有GPU（如果需要CPU运行，可省略）。
  - -p 8000:8000：映射HTTP/REST端口。
  - -v ${PWD}/model_repository:/models：挂载本地模型仓库到容器。
  - --model-repository=/models：指定模型仓库路径。
验证：
- 访问http://localhost:8000/v2/health/ready，返回200表示服务器就绪。
- 检查日志以确认模型加载成功。

发送推理请求

使用Triton的Python客户端发送推理请求，需先准备输入数据。

安装客户端：
- 命令：pip install tritonclient[all]。

准备输入：

embedding模型的输入为tokenized文本，通常包括input_ids和attention_mask。

示例代码：

python 复制代码

from transformers import AutoTokenizer
import numpy as np

# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

# 示例文本
texts = ["Hello world", "This is a test"]

# 分词
inputs = tokenizer(texts, return_tensors="np", padding=True)
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]

发送请求：