License

PyTorch (LibTorch) 后端#

Triton 后端用于 PyTorch。您可以在 后端仓库 中了解有关 Triton 后端的更多信息。在 问题页面 上提问或报告问题。此后端旨在运行使用 PyTorch C++ API 的 TorchScript 模型。在 PyTorch 中使用 python API 创建的所有模型都必须经过跟踪/脚本化才能生成 TorchScript 模型。

在哪里可以询问有关 Triton 和 Triton 后端的一般问题？请务必阅读以下所有信息以及主 服务器 仓库中提供的 通用 Triton 文档。如果您在那里找不到答案，可以在主 Triton 问题页面 上提问。

构建 PyTorch 后端#

使用最新的 cmake 进行构建。首先安装所需的依赖项。

$ apt-get install rapidjson-dev python3-dev python3-pip
$ pip3 install patchelf==0.17.2

必须使用来自 NGC 的适当 PyTorch 容器。例如，要构建使用来自 NGC 的 PyTorch 容器 23.04 版本的后端

$ mkdir build
$ cd build
$ cmake -DCMAKE_INSTALL_PREFIX:PATH=`pwd`/install -DTRITON_PYTORCH_DOCKER_IMAGE="nvcr.io/nvidia/pytorch:23.04-py3" ..
$ make install

以下所需的 Triton 仓库将被拉取并在构建中使用。默认情况下，“main”分支/标签将用于每个仓库，但可以使用列出的 CMake 参数来覆盖。

  • triton-inference-server/backend: -DTRITON_BACKEND_REPO_TAG=[tag]

  • triton-inference-server/core: -DTRITON_CORE_REPO_TAG=[tag]

  • triton-inference-server/common: -DTRITON_COMMON_REPO_TAG=[tag]

使用自定义 PyTorch 构建 PyTorch 后端#

目前，Triton 要求 PyTorch 后端使用特别修补版本的 PyTorch。这些 PyTorch 版本的完整源代码可作为 Docker 镜像从 NGC 获取。例如，与 Triton 22.12 版本兼容的 PyTorch 版本可作为 nvcr.io/nvidia/pytorch:22.12-py3 提供。

从 PyTorch NGC 容器 将 LibTorch 和 Torchvision 头文件和库复制到本地目录中。您可以查看 docker 中需要/复制哪些头文件和库。

$ mkdir build
$ cd build
$ cmake -DCMAKE_INSTALL_PREFIX:PATH=`pwd`/install -DTRITON_PYTORCH_INCLUDE_PATHS="<PATH_PREFIX>/torch;<PATH_PREFIX>/torch/torch/csrc/api/include;<PATH_PREFIX>/torchvision" -DTRITON_PYTORCH_LIB_PATHS="<LIB_PATH_PREFIX>" ..
$ make install

使用 PyTorch 后端#

参数#

Triton 通过模型的 config.pbtxt 文件的 Parameters 部分公开了一些标志来控制 TorchScript 模型的执行模式。

  • DISABLE_OPTIMIZED_EXECUTION：布尔标志，用于禁用 TorchScript 模型的优化执行。默认情况下，始终启用优化执行。

首次调用加载的 TorchScript 模型需要很长时间。由于模型预热时间较长 问题，Triton 也允许在没有这些优化的情况下执行模型。在某些模型中，优化执行不会提高性能，如 此处 所示，而在其他情况下，优化执行会对性能产生负面影响，如 此处 所示。

模型配置文件中指定此参数的部分将如下所示

parameters: {
key: "DISABLE_OPTIMIZED_EXECUTION"
    value: {
    string_value: "true"
    }
}

  • INFERENCE_MODE：布尔标志，用于启用 TorchScript 模型的推理模式执行。默认情况下，推理模式已启用。

InferenceMode 是一种新的 RAII 保护机制，类似于 NoGradMode，当您确定您的操作不会与 autograd 交互时使用。与 NoGradMode 相比，在此模式下运行的代码通过禁用 autograd 获得更好的性能。

请注意，在某些模型中，InferenceMode 可能不会提高性能，而在极少数情况下可能会对性能产生负面影响。

模型配置文件中指定此参数的部分将如下所示

parameters: {
key: "INFERENCE_MODE"
    value: {
    string_value: "true"
    }
}

  • DISABLE_CUDNN：布尔标志，用于禁用 cuDNN 库。默认情况下，cuDNN 已启用。

cuDNN 是一个 GPU 加速的深度神经网络原语库。cuDNN 为标准例程提供高度优化的实现。

通常，启用 cuDNN 运行的模型速度更快。但是，在某些例外情况下，使用 cuDNN 可能会更慢、导致更高的内存使用率或导致错误。

模型配置文件中指定此参数的部分将如下所示

parameters: {
key: "DISABLE_CUDNN"
    value: {
    string_value: "true"
    }
}

  • ENABLE_WEIGHT_SHARING：布尔标志，用于使同一设备上的模型实例能够共享权重。此优化不应与有状态模型一起使用。如果未指定，则禁用权重共享。

模型配置文件中指定此参数的部分将如下所示

parameters: {
key: "ENABLE_WEIGHT_SHARING"
    value: {
    string_value: "true"
    }
}

  • ENABLE_CACHE_CLEANING：布尔标志，用于在每次模型执行后启用 CUDA 缓存清理。如果未指定，则禁用缓存清理。如果模型位于 CPU 上，则此标志无效。由于每次模型执行后都会进行额外的 CUDA 缓存清理操作，因此将此标志设置为 true 将对性能产生负面影响。因此，只有当您使用 Triton 服务多个模型并在模型执行期间遇到 CUDA 内存不足问题时，才应使用此标志。

模型配置文件中指定此参数的部分将如下所示

parameters: {
key: "ENABLE_CACHE_CLEANING"
    value: {
    string_value:"true"
    }
}

  • INTER_OP_THREAD_COUNT:

PyTorch 允许在 TorchScript 模型推理期间使用多个 CPU 线程。一个或多个推理线程在给定输入上执行模型的前向传递。每个推理线程调用一个 JIT 解释器，该解释器逐个内联执行模型的操作。此参数设置此线程池的大小。此设置的默认值是 CPU 核心数。请参阅 此 文档，了解如何正确设置此参数。

模型配置文件中指定此参数的部分将如下所示

parameters: {
key: "INTER_OP_THREAD_COUNT"
    value: {
    string_value:"1"
    }
}

  • INTRA_OP_THREAD_COUNT:

除了操作间并行性之外，PyTorch 还可以在操作内利用多个线程（操作内并行性）。这在许多情况下都很有用，包括大型张量上的逐元素操作、卷积、GEMM、嵌入查找等。此设置的默认值是 CPU 核心数。请参阅 此 文档，了解如何正确设置此参数。

模型配置文件中指定此参数的部分将如下所示

parameters: {
key: "INTRA_OP_THREAD_COUNT"
    value: {
    string_value:"1"
    }
}

  • 其他优化：还有三个额外的布尔参数可用于禁用某些 Torch 优化，这些优化有时会导致具有复杂执行模式和动态形状的模型的延迟回归。如果未指定，则默认情况下全部启用。

    ENABLE_JIT_EXECUTOR

    ENABLE_JIT_PROFILING

支持#

模型实例组类型#

PyTorch 后端支持以下类型的 模型实例组，其中输入张量按如下方式放置

  • KIND_GPU：输入在与模型实例关联的 GPU 设备上准备。

  • KIND_CPU：输入在 CPU 上准备。

  • KIND_MODEL：输入在 CPU 上准备。加载模型时，后端不会为模型选择 GPU 设备；相反，它会遵循模型中指定的设备，并在推理期间按原样使用它们。当模型内部使用多个 GPU 时，这很有用，如 此示例模型 中所示。如果模型中未指定设备，则后端使用第一个可用的 GPU 设备。此功能从 23.06 版本开始提供。

重要提示#

  • PyTorch 模型在 GPU 上的执行本质上是异步的。有关更多详细信息，请参阅 此处。因此，PyTorch 模型执行中的错误可能会在接下来的几个推理请求中引发到服务器。在启动服务器时设置环境变量 CUDA_LAUNCH_BLOCKING=1 将有助于通过强制同步执行来正确调试失败的案例。

    • 在这种情况下，PyTorch 模型可能会或可能不会从失败状态中恢复，并且可能需要重启服务器才能继续成功提供服务。

  • PyTorch 不支持字符串张量，但它支持接受字符串列表作为输入/生成字符串列表作为输出的模型。对于这些模型，Triton 允许用户使用字符串数据类型传递字符串输入/接收字符串输出。由于使用列表而不是张量进行字符串 I/O 的限制，仅支持 1 维输入/输出的字符串类型 I/O。

  • 在多 GPU 环境中，当使用 Tracing 生成 TorchScript 模型时，可能会发生潜在的运行时问题。此问题是由于模型实例和张量之间的设备不匹配而引起的。默认情况下，Triton 为每个可用的 GPU 创建模型的单个执行实例。当请求发送到模型实例时，运行时错误会发生，该模型实例的 GPU 设备与 TorchScript 生成过程中使用的 GPU 设备不同。为了解决此问题，强烈建议在多 GPU 环境中使用 Scripting 而不是 Tracing 进行模型生成。Scripting 避免了设备不匹配问题，并确保与 Triton 一起使用时与不同 GPU 的兼容性。但是，如果 Tracing 是不可避免的，则有一种可用的解决方法。您可以在 模型配置 中显式指定模型实例的 GPU 设备，以确保模型实例和用于推理的张量分配给与模型跟踪时相同的 GPU 设备。

PyTorch 2.0 后端 [实验性]#

[!WARNING] 此功能可能会更改和移除。

从 24.01 开始，PyTorch 模型可以直接通过 Python 运行时 提供服务。默认情况下，Triton 将为 PyTorch 模型使用 LibTorch 运行时。要使用 Python 运行时，请在模型配置中提供以下 运行时设置

runtime: "model.py"

依赖项#

Python 后端依赖项#

此功能依赖于 Python 后端，有关更多详细信息，请参阅 基于 Python 的后端。

PyTorch 依赖项#

此功能将利用 torch.compile 优化，请确保 PyTorch 2.0+ pip 包 在同一 Python 环境中可用。

或者，可以使用具有 PyTorch 依赖项的 Python 执行环境。可以使用 提供的脚本 创建它。生成的 pb_exec_env_model.py.tar.gz 文件应放置在与 Python 运行时 相同的 后端共享库 目录中。

模型布局#

PyTorch 2.0 模型#

模型仓库应如下所示

model_repository/
`-- model_directory
    |-- 1
    |   |-- model.py
    |   `-- [model.pt]
    `-- config.pbtxt

model.py 包含 PyTorch 模型的类定义。该类应扩展 torch.nn.Module。model.pt 可以选择性地提供，其中包含模型的已保存 state_dict。

TorchScript 模型#

模型仓库应如下所示

model_repository/
`-- model_directory
    |-- 1
    |   `-- model.pt
    `-- config.pbtxt

model.pt 是 TorchScript 模型文件。

自定义#

可以通过在 config.pbtxt 上设置参数来自定义以下 PyTorch 设置。

torch.set_num_threads(int)

  • 键：NUM_THREADS

  • 值：用于 CPU 上操作内并行性的线程数。

torch.set_num_interop_threads(int)

  • 键：NUM_INTEROP_THREADS

  • 值：用于 CPU 上操作间并行性（例如，在 JIT 解释器中）的线程数。

torch.compile() 参数

  • 键：TORCH_COMPILE_OPTIONAL_PARAMETERS

  • 值：以下任何参数，编码为 JSON 对象。

    • fullgraph (bool)：是否可以将模型分解为多个子图。

    • dynamic (bool)：使用动态形状跟踪。

    • backend (str)：要使用的后端。

    • mode (str)：可以是“default”、“reduce-overhead”或“max-autotune”。

    • options (dict)：要传递给后端的选项字典。

    • disable (bool)：将 torch.compile() 转换为用于测试的空操作。

例如

parameters: {
    key: "NUM_THREADS"
    value: { string_value: "4" }
}
parameters: {
    key: "TORCH_COMPILE_OPTIONAL_PARAMETERS"
    value: { string_value: "{\"disable\": true}" }
}

限制#

以下是此功能的一些已知限制

  • 可通过 torch.compile 优化的 Python 函数可能无法在 model.py 文件中直接提供服务，它们需要由扩展 torch.nn.Module 的类封装。

  • 模型权重无法在同一 GPU 设备上的多个实例之间共享。

  • 当使用 KIND_MODEL 作为模型实例类型时，将使用模型上第一个参数的默认设备。

本页内容