完整教程：使用 NVIDIA Triton 构建端到端 GPU 推理管道
在生产级机器学习（ML）应用中，推理往往不局限于对单个模型执行前向传播，而是需要运行多个 ML 模型构成的管道。例如，一个对话式 AI 管道包含三个模块：自动语音识别（ASR）模块将输入音频波形转为文本，大型语言模型（LLM）模块理解输入并生成响应，文本到语音（TTS）模块将 LLM 输出合成语音。
另一个例子是文本到图像应用，其管道由 LLM 和扩散模型组成，分别用于编码输入文本和合成图像。此外，许多应用在将数据送入 ML 模型前需要预处理（如调整大小、裁剪、解码图像，或对文本进行分词），在模型输出后也需要后处理。近年来，ML 模型参数数量激增，且需服务海量用户，优化推理管道变得空前重要。NVIDIA TensorRT 和 FasterTransformer 等工具可优化单个深度学习模型在 GPU 上的推理延迟与吞吐量。
然而，我们的首要目标是加速整个推理管道，而非仅单个模型。例如，当模型在 GPU 上执行时，若预处理和后处理仍在 CPU 上运行，即使模型步骤很快，整体性能也会被拖累。最有效的方案是将预处理、模型执行和后处理全部放在 GPU 上。这种端到端 GPU 推理管道的效率来源于两个关键因素：
- 管道步骤之间无需在 CPU（主机）与 GPU（设备）间来回拷贝数据。
- GPU 强大的算力被用于整个推理管道。
NVIDIA Triton Inference Server 是一款开源推理服务软件，可在 CPU 和 GPU 上大规模部署和运行模型。其众多功能中，Triton 支持集成模型（ensemble models），允许将推理管道定义为有向无环图（DAG）形式的模型集合。Triton 会自动处理整个管道的执行，集成模型定义了一个模型的输出张量如何作为输入馈送给另一个模型。
通过 Triton 集成模型，您可以在 GPU、CPU 或两者混合上运行整个推理管道。当管道包含预处理 / 后处理步骤，或多个 ML 模型且一个模型的输出需输入另一模型时，此功能尤为有用。对于包含循环、条件或自定义逻辑的管道，Triton 还支持业务逻辑脚本（BLS）。
本文将聚焦于集成模型，引导您使用不同框架后端创建包含多个模型的端到端推理管道。Triton 提供了灵活的多框架后端，允许在 GPU、CPU 或混合上运行模型。我们将探讨以下三种管道运行方式：
- 整个管道在 CPU 上执行。
- 预处理和后处理在 CPU 上运行，模型执行在 GPU 上运行。
- 整个管道在 GPU 上执行。
我们还将强调使用 Triton 在 GPU 上运行完整推理管道的优势。本文以 CommonLit Readability Kaggle 挑战赛为例，预测 3-12 年级文学段落的复杂度，并使用 Triton 进行完整推理。请注意，本文使用 NVIDIA Triton 22.11 版本，您也可使用更高版本，但需确保后端版本匹配，以避免兼容性错误。
模型创建
针对此任务，我们训练了两个独立模型：基于 PyTorch 的 BERT-Large 和基于 cuML 的 随机森林回归器。分别命名为 bert-large 和 cuml。两个模型均以预处理后的文本摘录为输入，输出复杂度分数。
创建 BERT-Large 模型
第一个模型从 Hugging Face 预训练模型 bert-large-uncased（3.4 亿参数）微调得到。通过添加一个线性层，将 BERT 最后一层隐藏状态映射到单个输出值，用于任务微调。使用均方根损失、带权重衰减的 Adam 优化器以及 5 折交叉验证。通过传入示例输入并使用以下命令追踪模型，将其序列化为 TorchScript 文件（名为 model.pt）：
traced_script_module = torch.jit.trace(bert_pytorch_model, (example_input_ids, example_attention_mask))
traced_script_module.sa ve("model.pt")创建随机森林模型（cuML）
第二个模型使用 cuML 随机森林回归器，包含 100 棵树，每棵树最大深度 16。为基于树的模型生成以下特征：
- 每段摘录的单词数
- 每段摘录的不同单词数
- 标点符号数
- 每段摘录的句子数
- 每句平均单词数
- 每段摘录的停用词数
- 每句平均停用词数
- 语料库中前 N 高频词的频率分布
- 语料库中后 N 低频词的频率分布
取 N=100，随机森林共包含 207 个特征。通过以下命令将训练好的 cuML 模型转换为 Treeite 检查点（名为 checkpoint.tl）：
cuml_model.convert_to_treelite_model().to_treelite_checkpoint('checkpoint.tl')