9n-triton 部署 bert 模型实战经验

2024-04-01
北京
本文字数：2558 字
阅读完需：约 8 分钟

一、背景

对于算法工程师来说，通常采用 python 语言来作为工作语言，但是直接用 python 部署线上服务性能很差。这个问题困扰了我很久，为了缓解深度学习模型工程落地性能问题，探索了 Nvidia 提供的 triton 部署框架，并在九数中台上完成线上部署，发现性能提升近 337%！！ （原服务单次访问模型推理时间 175ms 左右，同模型框架单次访问推理时间缩减至 40ms 左右）。本文仅以本次部署服务经验分享给诸多受制于 python 性能问题的伙伴，希望有所帮助。

二、问题发现

组内原算法服务一直采用 python-backend 的镜像部署方式，将算法模型包装成接口形式，再通过 Flask 外露，打入 docker 中启动服务，但是发现推到线上接口响应时间过长，非常影响用户体验，于是想做出改进。python 后端部署一般存在以下问题：

1.性能问题：

◦由于 python 是一种解释语言，因此对比于其他编译语言（如 C，C++或 go）要慢很多，这对于要求高性能或者低延迟快速响应的线上服务很不友好，用户体验很差。

◦GIL 全局锁限制，多线程困难，限制了高并发能力，限制程序性能。

◦python 内存占用过多，资源浪费。

2.部署复杂：

◦自训练/自搭建算法模型往往需要一系列的相关依赖和系统库，如果利用 python 环境直接部署至线上，则必须解决将所有依赖库同步打包至对应容器中，并很可能需要额外解决不同库版本冲突问题，十分麻烦。

为了解决上述困扰，我们经过调研行业内其他实践经验，决定摒弃传统镜像部署方式，摆脱 python 部署环境，换用 triton 部署算法模型。

三、部署实例

本试验通过 fine-tune 的 Bert 模型进行文本分类任务，并通过九数算法中台包装的 triton 部署框架部署至公网可访问，以下介绍了我们团队部署的全流程和途中踩到的坑，以供大家快速入手（以下均以 Bert 模型为例）：

step 1: 将训练好的模型保存为 onnx 或 torchscript（即.pt）文件格式。据调研行业经验，onnx 一般用于 cpu 环境下推理任务效果较好，.pt 文件在 gpu 环境下推理效果较好，我们最终选择了将模型转化为.pt 文件。代码例子如下：

import torchfrom transformers import BertTokenizer, BertForSequenceClassification
tokenizer = BertTokenizer.from_pretrained('./pretrained-weights/bert-base-chinese')model = BertForSequenceClassification.from_pretrained('./pretrained-weights/bert-base-chinese', num_labels=10)# Trace the wrapped model# input_ids和attention_mask来自tokenizer, 具体可看一下torch.jit.trace用法traced_model = torch.jit.trace(model, (input_ids, attention_mask))traced_model.save("./saved_model_torchscript_v3.pt")

复制代码

本段代码借用了 torch.jit.trace 的存储方式，输入是文本经过 tokenizer 之后的 input_ids 和 attention_mask，输出是未经 softmax 处理的线性层。我们尝试将 softmax 部分包装进模型中一起打包成 pt 文件，使模型能够直接吐出分类类别，方案是外层再包装一层 forward，代码例子如下：

class BertSequenceClassifierWrapper(torch.nn.Module):    def __init__(self, model):        super().__init__()        self.model = model.cuda()
    def forward(self, input_ids, attention_mask):        outputs = self.model(input_ids=input_ids.cuda(), attention_mask=attention_mask.cuda())        # Apply softmax to the logits         probs = softmax(outputs.logits, dim=-1)        # Get the predicted class         predicted_class = torch.argmax(probs, dim=-1)        return predicted_class
wrapped_model = BertSequenceClassifierWrapper(model)traced_model = torch.jit.trace(wrapped_model, (input_ids, attention_mask))traced_model.save("./saved_model_torchscript_v3.pt")

复制代码

然而，参考于行业内实践，我们没有选择将 tokenizer 部分包装至模型内部，这是因为 torchscript 仅接受数组类型输入而非文本。最终，我们封装的模型输入为 input_ids 和 attention_mask，输出为分类结果（标量输出）。

【注】：坑点：

◦转 onnx 需要额外安装 transformers-onnx 包，需要 sudo 权限，九数上不能执行，建议本地转换。

◦对于.pt 类型文件，当前九数 triton 部署仅支持 gpu 方式推理方式，不支持 cpu 推理，且 torch 版本有强要求限制 torch ==1.10.0,cuda11.2（更高版本似乎也可），故我们在保存时必须要在 torch ==1.10.0,cuda11.2 环境下生成 pt 文件，不然会部署失败。

◦转 pt 文件时建议使用 torch.jit.trace 而非 torch.jit.script，后者经常出现引擎不兼容问题。

step 2: 将.pt 文件存放在九数磁盘目录下指定位置，注意，目录格式需要严格按照要求，具体可参考九数帮助文档。

目录格式如下（其中 model.pt 即为我们训练好的模型）：

pytorch-backend-online|-- pytorch-model|   |-- 1|   |   |-- model.pt|   `-- config.pbtxt

复制代码

config.pbtxt 配置如下：

name: "pytorch-model"platform: "pytorch_libtorch"
input [  {    name: "INPUT__0"    data_type: TYPE_INT64    dims: [1, 512]  },  {    name: "INPUT__1"    data_type: TYPE_INT64    dims: [1, 512]  }]output [  {    name: "OUTPUT__0"    data_type: TYPE_INT64    dims: [1, 1]   }]instance_group [  {      count: 1  }]