将 Triton 与 Inferentia 1 结合使用#

从 21.11 版本开始,Triton 支持 AWS InferentiaNeuron 运行时

目录#

Inferentia 设置#

将 Triton 与 Inferentia 结合使用的第一步是创建具有深度学习 AMI (已在 Ubuntu 18.04 上测试) 的 AWS Inferentia 实例。 ssh -i <private-key-name>.pem ubuntu@<instance address> 注意:建议将存储空间设置为大于默认值 110 GiB。当前版本的 Triton 已在 500 GiB 的存储空间下进行过测试。

登录到 inf1* 实例后,您需要克隆 当前 Github 仓库。按照 Github 上的步骤设置 ssh 访问权限,或者直接使用 https 克隆。使用 Github 将此仓库克隆到主目录 /home/ubuntu

 chmod 777 /home/ubuntu/python_backend/inferentia/scripts/setup-pre-container.sh
 sudo /home/ubuntu/python_backend/inferentia/scripts/setup-pre-container.sh

然后,使用以下命令启动 Triton 实例

 docker run --device /dev/neuron0 <more neuron devices> -v /home/ubuntu/python_backend:/home/ubuntu/python_backend -v /lib/udev:/mylib/udev --shm-size=1g --ulimit memlock=-1 -p 8000:8000 -p 8001:8001 -p 8002:8002 --ulimit stack=67108864 -ti nvcr.io/nvidia/tritonserver:<xx.yy>-py3

注意 1:用户需要在容器初始化期间列出要运行的任何 neuron 设备。例如,要在实例上使用 4 个 neuron 设备,用户需要使用以下命令运行

 docker run --device /dev/neuron0 --device /dev/neuron1 --device /dev/neuron2 --device /dev/neuron3 ...`

注意 2:/mylib/udev 用于 Neuron 参数传递。

注意 3:对于 Triton 容器版本 xx.yy,请参阅 Triton 推理服务器容器发行说明。当前的构建脚本已在容器版本 21.10 下进行过测试。

启动 Triton 容器后,进入 python_backend 文件夹并运行 setup 脚本。

 source /home/ubuntu/python_backend/inferentia/scripts/setup.sh

此脚本将

  1. 安装必要的依赖项

  2. 安装 neuron-cc,Neuron 编译器。

  3. 根据您的偏好安装 neuron 框架软件包,例如 pytorch、tensorflow 或两者都安装。

该脚本还提供用户可配置选项。请使用 -h--help 选项了解更多可配置选项。

设置 Inferentia 模型#

目前,我们仅支持 PyTorchTensorFlow 工作流程在 inferentia 上执行。

用户需要创建自己的 *.pt (对于 pytorch) 或 *.savedmodels (对于 tensorflow) 模型。这是一个关键步骤,因为 Inferentia 将需要底层的 .NEFF 图来执行推理请求。请参考

PyTorch#

对于 PyTorch,我们支持通过 PyTorch-Neuron trace python API 跟踪的模型,以便在 Inferentia 上执行。获得支持 Inferentia 的 TorchScript 模型后,使用 gen_triton_model.py 脚本生成 triton python 模型目录。

PyTorch 模型的 gen_triton_model.py 的示例调用可能如下所示

 python3 inferentia/scripts/gen_triton_model.py --model_type pytorch --triton_input INPUT__0,INT64,4x384 INPUT__1,INT64,4x384 INPUT__2,INT64,4x384 --triton_output OUTPUT__0,INT64,4x384 OUTPUT__1,INT64,4x384 --compiled_model /home/ubuntu/bert_large_mlperf_neuron_hack_bs1_dynamic.pt --neuron_core_range 0:3 --triton_model_dir bert-large-mlperf-bs1x4

为了使脚本将编译后的模型视为 TorchScript 模型,需要提供 --model_type pytorch

注意:由于 TorchScript 模型中缺少输入和输出的元数据 - 必须向上述脚本提供输入和输出张量的名称、数据类型和形状,并且名称必须遵循特定的命名约定,即 <name>__<index>。其中 <name> 可以是任何字符串,<index> 指的是相应输入/输出的位置。这意味着如果有两个输入和两个输出,它们必须命名为:“INPUT__0”、“INPUT__1”和“OUTPUT__0”、“OUTPUT__1”,其中 “INPUT__0” 指的是第一个输入,而 INPUT__1 指的是第二个输入,等等。

此外,--neuron_core_range 指定在服务此模型时要使用的 neuron 核心。目前,仅支持 torch.neuron.DataParallel() 模式。有关更多信息,请参阅 数据并行推理。可以使用 --triton_model_instance_count 选项指定 Triton 模型实例计数。neuron 核心将在所有实例之间平均分配。例如,在有两个 triton 模型实例和 4 个 neuron 核心的情况下,第一个实例将加载到核心 0-1 上,第二个实例将加载到核心 2-3 上。为了最好地利用 inferentia 设备,请尝试将 neuron 核心的数量设置为实例计数的适当倍数。

TensorFlow#

对于 TensorFlow,模型必须为 AWS Neuron 编译。请参阅 AWS Neuron TensorFlow 教程,了解如何获取使用 Neuron 核心的编译模型。目前,代码仅在 tensorflow==1.15 上进行过测试。

获得编译后的模型后,使用 gen_triton_model.py 脚本生成 triton python 模型目录。

TensorFlow 模型的 gen_triton_model.py 的示例调用可能如下所示

 python3 gen_triton_model.py --model_type tensorflow --compiled_model /home/ubuntu/inferentia-poc-2.0/scripts-rn50-tf-native/resnet50_mlperf_opt_fp16_compiled_b5_nc1/1 --neuron_core_range 0:3  --triton_model_dir rn50-1neuroncores-bs1x1

注意:与 TorchScript 模型不同,TensorFlow SavedModel 存储了足够的元数据来检测模型的输入和输出张量的名称、数据类型和形状。默认情况下,脚本会将编译后的模型假定为 torchscript。为了使其将编译后的模型视为 TF savedmodel,需要提供 --model_type tensorflow。输入和输出详细信息从模型本身读取。用户必须安装 tensorflow python 模块才能将此脚本用于 tensorflow 模型。

与 PyTorch 类似,可以使用 --neuron_core_range--triton_model_instance_count 来指定 neuron 核心范围和 triton 模型实例的数量。但是,neuron 核心索引不指向芯片中的特定 neuron 核心。对于 TensorFlow,我们使用已弃用的 NEURONCORE_GROUP_SIZES 功能来加载模型。在这种情况下,模型将加载到下一个可用的 Neuron 核心上,而不是特定的核心。有关更多信息,请参阅 使用 NEURONCORE_GROUP_SIZES 的并行执行

另一个注意事项,由于 Neuron-Tensorflow(与 Neuron-Python 不同)没有内置函数来为多个核心运行模型,因此,如果用户启用跨多个核心的处理,建议输入的第一维度为 None

请使用 gen_triton_model.py 中的 -h--help 选项了解更多可配置选项。

在 Triton 中部署 Inferentia 模型#

gen_triton_model.py 应创建一个具有以下结构的 triton 模型目录

bert-large-mlperf-bs1x4
 |
 |- 1
 |  |- model.py
 |
 |- config.pbtxt

查看脚本的使用消息以了解每个选项。

该脚本将生成一个具有用户提供名称的模型目录。将该模型目录移动到 Triton 的模型仓库。确保提供给脚本的编译模型路径指向有效的 torchscript 文件或 tensorflow savedmodel。

现在,可以使用以下命令启动服务器和模型

 tritonserver --model-repository <path_to_model_repository>

注意

  1. config.pbtxtmodel.py 应被视为起点。用户可以根据自己的需要自定义这些文件。

  2. Triton Inferentia 目前已在单个模型上进行过测试。

使用 Triton 的动态批处理#

要启用动态批处理,需要指定 --enable_dynamic_batching 标志。gen_triton_model.py 支持以下三个选项来配置 Triton 的动态批处理

  1. --preferred_batch_size:有关首选批大小的详细信息,请参阅 模型配置文档。为了优化性能,建议将其设置为已使用的 neuron 核心数的倍数。例如,如果每个实例使用 2 个 neuron 核心,则 preferred_batch_size 可以为 2、4 或 6。

  2. --max_queue_delay_microseconds:有关详细信息,请参阅 模型配置文档

  3. --disable_batch_requests_to_neuron:启用 Triton 处理批处理请求的非默认方式。Triton 后端将单独向 neuron 发送每个请求,而不管 Triton 服务器请求是否已批处理。当用户想要优化在没有该标志的情况下批处理性能不佳的模型的性能时,建议使用此标志。

此外,--max_batch_size 将影响最大批处理限制。有关详细信息,请参阅 模型配置文档

测试 Inferentia 设置的准确性#

qa 文件夹包含使用简单的 add_sub 模型设置测试的必要文件。该测试需要一个具有超过 8 个 inferentia 核心的实例才能运行,例如:inf1.6xlarge。要开始测试,请运行

 source <triton path>/python_backend/inferentia/qa/setup_test_enviroment_and_test.sh

其中 <triton path> 通常是 /home/ubuntu/。此脚本将拉取包含 inferentia 测试的 server 仓库。然后,它将构建最新的 Triton Server 和 Triton SDK。

注意:如果您需要更改服务器仓库中的某些测试,则需要运行

 export TRITON_SERVER_REPO_TAG=<your branch name>

在运行脚本之前。

将 Triton 与 Inferentia 2 或 Trn1 结合使用#

pytorch-neuronx 和 tensorflow-neuronx#

  1. 与 inf1 的步骤类似,更改预容器和容器上设置脚本的参数,以包含 -inf2-trn1 标志,例如:

 chmod 777 /home/ubuntu/python_backend/inferentia/scripts/setup-pre-container.sh
 sudo /home/ubuntu/python_backend/inferentia/scripts/setup-pre-container.sh -inf2
  1. 在容器上,在 docker run 命令之后,您可以将类似的参数传递给 setup.sh 脚本,例如对于 Pytorch

source /home/ubuntu/python_backend/inferentia/scripts/setup.sh -inf2 -p

对于 Tensorflow

source /home/ubuntu/python_backend/inferentia/scripts/setup.sh -inf2 -t
  1. 按照上述步骤,当使用 gen_triton_model.py 脚本时,您可以将类似的参数 --inf2 传递给 setup.sh 脚本,例如对于 Pytorch

python3 inferentia/scripts/gen_triton_model.py --inf2 --model_type pytorch --triton_input INPUT__0,INT64,4x384 INPUT__1,INT64,4x384 INPUT__2,INT64,4x384 --triton_output OUTPUT__0,INT64,4x384 OUTPUT__1,INT64,4x384 --compiled_model bert_large_mlperf_neuron_hack_bs1_dynamic.pt --neuron_core_range 0:3 --triton_model_dir bert-large-mlperf-bs1x4
  1. 注意:当使用 --inf2 选项时,--compiled_model 路径应相对于 triton 模型目录提供。model.py 中的 initialize() 函数将通过连接仓库中的模型路径和相对 --compiled_model 路径来派生完整路径。

transformers-neuronx#

要将 inf2/trn1 实例与 transformers-neuronx 软件包一起使用以服务模型,请按照上述说明生成 pytorch 模型。transformers-neuronx 当前支持 此处 列出的模型。

正如 neuronx 文档页面上所述,虽然 neuronx 加载 API 因模型而异,但它们遵循相同的模式。

  1. 要服务 transformers-neuronx 模型,请首先在 inf2 实例上使用 save_pretrained_split() API 跟踪模型(对于大型语言模型,建议使用 inf2.24xl)。之后,在使用 gen_triton_model.py 文件时,将文件夹打包为 ‘–compiled_model’。

  2. 以下树显示了 OPT 模型的示例模型结构

opt/
├── 1
│   └── model.py
├── opt-125m-model
│   └── pytorch_model.bin
└── opt-125m-tp12
    ├── FullyUnrolled.1814.1
    │   ├── penguin-sg0000
    │   └── sg00
    ├── FullyUnrolled.1814.2
    │   ├── penguin-sg0000
    │   └── sg00
    ├── FullyUnrolled.1814.3
    │   ├── penguin-sg0000
    │   └── sg00
    ├── FullyUnrolled.1814.4
    │   ├── penguin-sg0000
    │   └── sg00
    └── FullyUnrolled.1814.5
        ├── penguin-sg0000
        └── sg00
  ├── config.pbtxt
  1. 添加以下导入(例如,对于 OPT 模型)。导入将因您尝试运行的模型而异。

from transformers_neuronx.opt.model import OPTForSampling
  1. initialize() 函数中添加以下行。根据您的要求设置 batch_sizetp_degreen_positionsampunroll 参数。tp_degree 通常应与 inf2 实例上可用的 neuron 核心数匹配。

batch_size = 1
tp_degree = 12
n_positions = 2048
amp = 'bf16'
unroll = None
self.model_neuron = OPTForSampling.from_pretrained(compiled_model, batch_size=batch_size, amp=amp, tp_degree=tp_degree, n_positions=n_positions, unroll=unroll)
self.model_neuron.to_neuron()

self.model_neuron.num_workers = num_threads

您也可以选择将 batch_size 等参数添加到 config.pbtxt 作为参数,并在 initialize() 函数中读取它们,类似于 --compiled-model

  1. 最后,在 execute() 函数中,使用以下 API 运行推理

batched_results = self.model_neuron.sample(batched_tensor, 2048)

上面,2048 是一个足够长的输出 token。如果您想将其指定为有效负载的一部分,也可以将其作为输入之一传入。

  1. 继续加载模型,并提交推理有效负载,类似于任何其他 triton 模型。