如何使用NVIDIA Triton 推理服务器来运行推理管道

在许多生产级机器学习（ ML ）应用程序中，推理并不局限于在单个 ML 模型上运行前向传递。相反，通常需要执行 ML 模型的管道。例如，一个由三个模块组成的对话式人工智能管道：一个将输入音频波形转换为文本的自动语音识别（ ASR ）模块，一个理解输入并提供相关响应的大型语言模型（ LLM ）模块，以及一个从 LLM 输出产生语音的文本到语音（ TTS ）模块。

或者，考虑一个文本到图像的应用程序，其中管道由 LLM 和扩散模型组成，它们分别用于对输入文本进行编码和从编码的文本合成图像。此外，许多应用程序在将输入数据馈送到 ML 模型之前需要对其进行一些预处理步骤，或者对这些模型的输出进行后处理步骤。例如，输入图像在被馈送到计算机视觉模型之前可能需要被调整大小、裁剪和解码，或者文本输入在被馈送给 LLM 之前需要被标记化。

近年来， ML 模型中的参数数量激增，它们越来越多地被要求为庞大的消费者群体提供服务；因此，优化推理管道变得比以往任何时候都更加重要。 NVIDIA TensorRT 和 FasterTransformer 等工具在 GPU s 上执行推理时，优化单个深度学习模型，以获得更低的延迟和更高的吞吐量。

然而，我们的首要目标不是加快单个 ML 模型的推理，而是加快整个推理管道。例如，当在 GPU 上服务模型时，在 CPU 上具有预处理和后处理步骤会降低整个管道的性能，即使模型执行步骤很快。推理管道最有效的解决方案是使预处理、模型执行和后处理步骤都在 GPU s 上运行。这种端到端推理流水线在 GPU 上的效率来自以下两个关键因素。

在流水线步骤之间，数据不需要在 CPU （主机）和 GPU （设备）之间来回复制。

GPU 强大的计算能力被用于整个推理管道。

NVIDIA Triton Inference Server 是一款开源推理服务软件，用于在 CPU 和 GPU 上大规模部署和运行模型。在许多功能中， NVIDIA Triton 支持 ensemble models ，使您能够将推理管道定义为有向非循环图（ DAG ）形式的模型集合。 NVIDIA Triton 将处理整个管道的执行。集成模型定义了如何将一个模型的输出张量作为输入馈送到另一个模型。

使用 NVIDIA Triton 集成模型，您可以在 GPU 或 CPU 上运行整个推理管道，也可以在两者的混合上运行。当涉及预处理和后处理步骤时，或者当管道中有多个 ML 模型，其中一个模型的输出馈送到另一个模型时，这是有用的。对于管道包含循环、条件或其他自定义逻辑的用例， NVIDIA Triton 支持 Business Logic Scripting (BLS) 。

这篇文章只关注整体模型。它将引导您完成使用不同框架后端创建具有多个模型的端到端推理管道的步骤。 NVIDIA Triton 提供了使用多个框架后端构建模型管道的灵活性，并在 GPU 或 CPU 上运行它们，或两者混合运行。我们将探索以下三种管道运行方式。

整个管道在 CPU 上执行。

预处理和后处理步骤在 CPU 上运行，模型执行在 GPU 上运行。

整个管道在 GPU 上执行。

我们还将强调使用 NVIDIA Triton 推理服务器在 GPU 上运行整个推理管道的优势。我们专注于 CommonLit Readability Kaggle challenge ，用于预测 3-12 年级文学段落的复杂度，使用 NVIDIA Triton 进行整个推理。请注意， NVIDIA Triton 222.11 用于本博客文章的目的。您也可以使用 NVIDIA Triton 的更高版本，前提是您在后端使用匹配的版本（表示为），以避免可能的兼容性错误。

模型创建

对于这项任务，训练两个独立的模型：使用 PyTorch 训练的 BERT-Large 和使用 cuML 训练的随机森林回归器。将这些型号命名为bert-large和cuml。这两个模型都会将经过预处理的摘录作为输入，并输出分数或复杂度。

作为第一个模型，从预训练的 Hugging Face 模型中微调基于 transformer 的bert-large模型 bert-large-uncased 其具有 340M 个参数。通过添加一个线性层来微调任务的模型，该线性层将 BERT 的最后一个隐藏层映射到单个输出值。

使用均方根损失、带有权重衰减的 Adam 优化器和 5 倍交叉验证进行微调。通过在模型中传递示例输入并使用以下命令跟踪模型，将模型序列化为 TorchScript 文件（名为model.pt）：

 

traced_script_module = torch.jit.trace(bert_pytorch_model,          
                      (example_input_ids, example_attention_mask))
traced_script_module.save("model.pt")

 

作为第二个模型，使用 cuML 随机森林回归器，其中有 100 棵树，每棵树的最大深度为 16 。为基于树的模型生成以下特征：每个摘录的单词数、每个摘录的不同单词数、标点符号数、每个节选的句子数、每个句子的平均单词数、每次摘录的停止单词数、每句子的平均停止单词数， N 最常出现的单词在语料库中的频率分布，以及 N 最不频繁单词在语料库中的频率分布。

使用 N=100 ，这样随机森林总共包含 207 个特征。通过使用以下命令转换 cuML 模型实例，将经过训练的 cuML 模型序列化为 Treeite 检查点（名为checkpoint.tl）：

 

cuml_model.convert_to_treelite_model().to_treelite_checkpoint('checkpoint.tl')

 

请注意，与模型关联的 Treelite 版本需要与用于推断的 NVIDIA Triton 容器中的 Treeliet 版本相匹配。

在 NVIDIA Triton 上运行 ML 模型

要在 NVIDIA Triton 中部署的每个模型都必须包括一个模型配置。默认情况下，当模型元数据可用时， NVIDIA Triton 将尝试使用模型元数据自动创建配置。在模型元数据不足的情况下，可以手动提供模型配置文件，也可以覆盖推断的设置。有关更多详细信息，请参阅 GitHub 上的 triton-inference-server/server 。

要在 NVIDIA Triton 上运行 PyTorch 格式的 BERT Large ，请使用 PyTorch (LibTorch) backend 。将以下行添加到模型配置文件中以指定此后端：

 

backend: "pytorch"

 

要在 NVIDIA Triton 上运行基于树的随机林模型，请使用 FIL (Forest Inference Library) backend ，在模型配置文件中添加以下内容：

 

backend: "fil"

 

此外，在模型配置文件中添加以下行，以指定所提供的模型为 Treeite 二进制格式：

 

parameters {
    key: "model_type"
    value: { string_value: "treelite_checkpoint" } 
}

 

最后，在每个模型配置文件中，包括instance_group[{kind:KIND_GPU}]或instance_group[{kind:KIND_CPU}]，这取决于模型是在 GPU 还是 CPU 上提供服务。

到目前为止，生成的模型存储库目录结构如下：

 

.
├── bert-large
│   ├── 1
│   │   └── model.pt
│   └── config.pbtxt
└── cuml
    ├── 1
    │   └── checkpoint.tl
    └── config.pbtxt

 

预处理和后处理

预处理和后处理可以在 NVIDIA Triton 服务器之外执行，也可以作为 NVIDIA Triton 中模型集合的一部分进行合并。对于本例，预处理和后处理由 Python 中执行的操作组成。使用 Python backend 作为集成的一部分来运行这些操作。

如果 NVIDIA Triton 服务器容器中提供的默认 Python 版本可以运行 Python model ，则可以忽略以下部分，直接跳到下面标题为“比较推理管道”的部分否则，您将需要创建一个自定义 Python 后端存根和一个自定义执行环境，如下所述。

自定义 Python 后端存根

Python 后端使用存根进程将model.py文件连接到 NVIDIA Triton C++核心。 Python 后端可以使用安装在当前[ZGK8环境（虚拟环境或Conda环境）或全局 Python 环境中的库。

请注意，这假设用于编译后端存根的 Python 版本与用于安装依赖项的版本相同。写入时使用的 NVIDIA Triton 容器中的默认 Python 版本为 3.8 。如果需要运行预处理和后处理的 Python 版本与 NVIDIA Triton 容器中的版本不匹配，则需要编译 custom Python backend stub 。

要创建自定义 Python 后端存根，请在 NVIDIA conda容器内安装conda、cmake、rapidjson和 Triton 。接下来，创建一个 Conda 虚拟环境（请参见 documentation ），并使用以下命令激活它：

 

conda create -n custom_env python=
conda init bash
bash
conda activate custom_env

 

将替换为感兴趣的版本，例如 3.9 。然后克隆 Python 后端回购，并使用以下代码编译 Python 后端存根：

 

git clone https://github.com/triton-inference-server/python_backend -b r
cd python_backend
mkdir build && cd build
cmake -DTRITON_ENABLE_GPU=ON -DCMAKE_INSTALL_PREFIX:PATH=`pwd`/install ..
make triton-python-backend-stub

 

请注意，必须替换为 NVIDIA Triton 容器版本。运行上面的命令将创建名为triton-python-backend-stub的存根文件。这个 Python 后端存根现在可以用来加载安装有匹配版本[ZGK8的库。

自定义执行环境

如果您想为不同的 Python 型号使用不同的 Python 环境，您需要创建一个 custom execution environment 。要为 Python 模型创建自定义执行环境，第一步是在已经激活的 Conda 环境中安装任何必要的依赖项。然后运行conda-pack将我们的自定义执行环境打包为 tar 文件。这将创建一个名为custom_env.tar.gz的文件。

在撰写本文时， NVIDIA Triton 仅支持conda-pack用于捕获执行环境。请注意，当在 NVIDIA Triton Docker 容器中工作时，容器中包含的包也会在conda-pack创建的执行环境中捕获。

使用 Python 后端存根和自定义执行环境

在创建 Python 后端存根和自定义执行环境后，将两个生成的文件从容器复制到模型存储库所在的本地系统。在本地系统中，将存根文件复制到每个需要使用存根的 Python 模型（即预处理和后处理模型）的模型目录中。对于这些模型，目录结构如下：

 

model_repository
 ├── postprocess
 │   ├── 1
 │   │   └── model.py
 │   ├── config.pbtxt
 │   └── triton_python_backend_stub
 └── preprocess
     ├── 1
     │   └── model.py
     ├── config.pbtxt
     └── triton_python_backend_stub

 

对于预处理和后处理模型，还需要在配置文件中提供自定义执行环境的 tar 文件的路径。例如，预处理模型的配置文件将包括以下代码：

 

name: "preprocess"
backend: "python"

...
parameters: {
  key: "EXECUTION_ENV_PATH",
  value: {string_value: "path/to/custom_env.tar.gz"}
}

 

要使此步骤生效，请将custom_env.tar.gz存储在 NVIDIA Triton 推理服务器容器可以访问的路径中。

用于预处理和后处理的 model.py 文件的结构

每个 Python 模型都需要遵循 documentation 中描述的特定结构。在model.py文件中定义以下三个功能：

1. initialize（可选，在加载模型时运行）：用于在推理之前加载必要的组件，以减少每个客户端请求的开销。特别是对于预处理，加载 cuDF tokenizer ，它将用于标记基于 BERT 的模型的摘录。还加载用于随机森林特征生成的停止词列表，作为对基于树的模型的输入的一部分。

2 .execute（必需，请求推理时运行）：接收推理请求，修改输入，并返回推理响应。由于preprocess是推理的入口点，如果要对 GPU 执行推理，则需要将输入移动到[Z1K1’。

通过创建cudf.Series的实例将摘录输入张量移动到 GPU ，然后利用initialize中加载的 cuDF 标记器在[Z1K1’上标记整批摘录。

类似地，使用字符串操作生成基于树的特征，然后使用在 GPU 上操作的CuPY数组对特征进行归一化。要在 GPU 上输出张量，请使用toDlpack和from_dlpack（参见 documentation ）将张量封装到推理响应中。

最后，为了保持张量在 GPU 上，并避免在集合中的步骤之间复制到 CPU ，请将以下内容添加到每个模型的配置文件中：

 

parameters: {
 key: "FORCE_CPU_ONLY_INPUT_TENSORS"
 value: {
   string_value:"no"
 }
}

 

postprocess的输入分数已经在 GPU 上了，所以只需将分数与CuPY数组再次组合，并输出最终的管道分数。对于在 CPU 上进行后处理的情况，在预处理步骤中将ML模型的输出移动到[Z1K2’。

3. finalize（可选，在卸载模型时运行）：使您能够在从 NVIDIA Triton 服务器卸载模型之前完成任何必要的清理。

推理管道比较

本节介绍了不同的推理管道，并从延迟和吞吐量方面对它们进行了比较。

CPU 上的预处理和后处理， GPU 上的 ML 模型推理

此设置使用 NVIDIA Triton 对 ML 模型执行推理，同时在客户端所在的本地机器上使用 CPU 执行预处理和后处理（图 1 ）。在预处理模型中，对于给定的一批文本摘录，使用 BERT 标记器对摘录进行标记，并为 cuML 模型生成基于树的特征。

然后将预处理模型的输出作为推理请求发送给 NVIDIA Triton 。 NVIDIA Triton 然后对 GPU 上的 ML 模型执行推理并返回响应。在 CPU 上本地对该响应进行后处理，以生成输出分数。

图 1 。在 NVIDIA Triton 推理服务器上运行 ML 模型的管道设置，其中预处理和后处理步骤在 CPU 上本地执行

在 GPU 上执行 NVIDIA Triton 中的整个管道

在此设置中，使用 NVIDIA Triton 在 GPU 上执行整个推理管道。对于这个例子， NVIDIA Triton 中的数据管道和流可以在图 2 中看到。

图 2 :管道设置，包括预处理、模型执行和后处理，所有这些都作为 NVIDIA Triton 集成在 GPU 上执行

管道从一个预处理模型开始，该模型将文本摘录作为输入，为 BERT 标记摘录，并为随机森林模型提取特征。接下来，在预处理模型的输出上同时运行两个 ML 模型，每个模型生成指示输入文本的复杂度的分数。最后，在后处理步骤中对得到的分数进行组合。

要让 NVIDIA Triton 运行上面的执行管道，请创建一个名为ensemble_all的 ensemble model 。该模型与任何其他模型具有相同的模型目录结构，只是它不存储任何模型，并且只由一个配置文件组成。集成模型的目录如下所示：

 

├── ensemble_all
 │   ├── 1
 │   └── config.pbtxt

 

在配置文件中，首先使用以下脚本指定集成模型名称和后端：

 

name: "ensemble_all"
backend: "ensemble"

 

接下来，定义集合的端点，即集合模型的输入和输出：

 

input [
  {
    name: "excerpt"
    data_type: TYPE_STRING
    dims: [ -1 ]
  },
  {
    name: "BERT_WEIGHT"
    data_type: TYPE_FP32
    dims: [ -1 ]
  }
]
output {
    name: "SCORE"
    data_type: TYPE_FP32
    dims: [ 1 ]
  }

 

管道的输入是可变长度的，因此使用 -1 作为尺寸参数。输出是一个单浮点数。

要创建通过不同模型的管道和数据流，请包括ensemble_scheduling部分。第一个模型被称为preprocess，它将摘录作为输入，并输出 BERT 令牌标识符和注意力掩码以及树特征。调度的第一步显示在模型配置的以下部分中：

 

ensemble_scheduling {
  step [
    {
      model_name: "preprocess"
      model_version: 1
      input_map {
        key: "INPUT0"
        value: "excerpt"
      }
      output_map {
        key: "BERT_IDS"
        value: "bert_input_ids",
      }
      output_map {
        key: "BERT_AM"
        value: "bert_attention_masks",
      }
      output_map {
        key: "TREE_FEATS"
        value: "tree_feats",
      }
    },

 

step部分中的每个元素都指定了要使用的模型，以及如何将模型的输入和输出映射到集合调度器识别的张量名称。然后使用这些张量名称来识别各个输入和输出。

例如，step中的第一个元素指定应使用preprocess模型的版本一，其输入"INPUT0"的内容由"excerpt"张量提供，其输出"BERT_IDS"的内容将映射到"bert_input_ids"张量以供以后使用。类似的推理适用于preprocess的其他两个输出。

继续在配置文件中添加步骤以指定整个管道，将preprocess的输出传递到bert-large和cuml的输入：

 

{
      model_name: "bert-large"
      model_version: 1
      input_map {
        key: "INPUT__0"
        value: "bert_input_ids"
      }
      input_map {
        key: "INPUT__1"
        value: "bert_attention_masks"
      }
      output_map {
        key: "OUTPUT__0"
        value: "bert_large_score"
      }
    },

 

最后，通过在配置文件中添加以下行，将这些分数中的每一个传递到后处理模型，以计算分数的平均值并提供单个输出分数，如下所示：

 

{
      model_name: "postprocess"
      model_version: 1
      input_map {
        key: "BERT_WEIGHT_INPUT"
        value: "BERT_WEIGHT"
      }
      input_map {
        key: "BERT_LARGE_SCORE"
        value: "bert_large_score"
      }
      input_map {
        key: "CUML_SCORE"
        value: "cuml_score"
      }
      output_map {
        key: "OUTPUT0"
        value: "SCORE"
      }
    }
  }
]

 

在集成模型的配置文件中调度整个管道的简单性证明了使用 NVIDIA Triton 进行端到端推理的灵活性。要添加另一个模型或添加另一数据处理步骤，请编辑集成模型的配置文件并更新相应的模型目录。

请注意，集成配置文件中定义的max_batch_size必须小于或等于每个模型中定义的max_batch_size。整个模型目录，包括集成模型，如下所示：

 

├── bert-large
│   ├── 1
│   │   └── model.pt
│   └── config.pbtxt
├── cuml
│   ├── 1
│   │   └── checkpoint.tl
│   └── config.pbtxt
├── ensemble_all
│   ├── 1
│   │   └── empty
│   └── config.pbtxt
├── postprocess
│   ├── 1
│   │   ├── model.py
│   └── config.pbtxt
└── preprocess
    ├── 1
    │   ├── model.py
    └── config.pbtxt

 

要告诉 NVIDIA Triton 执行 GPU 上的所有模型，请在每个模型的配置文件中包括以下行（集成模型的配置文档中除外）：

 

instance_group[{kind:KIND_GPU}]

 

在 CPU 上执行 NVIDIA Triton 中的整个管道若要让 NVIDIA Triton 在 CPU 上执行整个管道，请重复在 GPU 上运行管道的所有步骤。在每个配置文件中将instance_group[{kind:KIND_GPU}]替换为以下行：

 

instance_group[{kind:KIND_CPU}]

 

后果

我们使用 GCP a2 高 GPU -1g VM 在延迟和吞吐量方面比较了以下三个推理管道：

由 NVIDIA Triton 在英特尔至强 CPU 上以 2.20 GHz 执行的完整管道

NVIDIA Triton 在 NVIDIA A100 40 GB GPU 上执行 ML 模型，在英特尔至强 CPU 上以 2.20 GHz 本地执行预处理和后处理

由 NVIDIA Triton 在 NVIDIA A100 40 GB GPU 上执行的完整管道

从表 1 中的结果可以明显看出，使用 NVIDIA Triton 在 GPU 上运行整个管道的优势。对于较大的批处理大小和张量大小，吞吐量的提高更为明显。 NVIDIA A100 40 GB 型号的执行管道比在 2.20 GHz 下在 Intel Xeon CPU 上运行的整个管道效率高得多。当将预处理和后处理从 CPU 移动到 GPU 时，有进一步的改进。

	Full pipeline on CPU	Pre/postprocess on CPU; ML models on GPU	Full pipeline on GPU
Latency (ms)	523	192	31
Throughput (samples/second) for batch size 512	242	7707	8308

表 1 。不同管道的延迟和吞吐量

如图 3 所示， CPU 在非常适中的批量下受到瓶颈限制，在 GPU 上运行整个管道可以显著提高吞吐量。

图 3 。随着批处理大小从 1 到 512 的变化，不同推理管道的吞吐量

结论

这篇文章解释了如何使用 NVIDIA Triton 推理服务器来运行由预处理和后处理组成的推理管道，以及基于 transformer 的语言模型和基于树的模型来解决 Kaggle 挑战。 NVIDIA Triton 为同一管道的模型以及预处理和后处理逻辑提供了使用多个框架/后端的灵活性。这些管道可以在 CPU s 和/或 GPU s 上运行。

我们表明，与在 CPU 上运行预处理和后处理步骤以及在 GPU 上执行模型相比，在模型执行的同时利用 GPU s 进行预处理和后期处理，将端到端延迟减少了 6 倍。我们还表明，使用 NVIDIA Triton 使我们能够同时对多个 ML 模型执行推理，并且从一种部署类型（所有 CPU ）到另一种（所有 GPU ）只需要在配置文件中更改一行即可。