PyTorch (LibTorch) Backend

The Triton backend for PyTorch. You can learn more about Triton backends in the backend repo. Ask questions or report problems on the issues page. This backend is designed to run TorchScript models using the PyTorch C++ API. All models created in PyTorch using the python API must be traced/scripted to produce a TorchScript model.

Where can I ask general questions about Triton and Triton backends? Be sure to read all the information below as well as the general Triton documentation available in the main server repo. If you don't find your answer there you can ask questions on the main Triton issues page.

Build the PyTorch Backend

Use a recent cmake to build. First install the required dependencies.

$ apt-get install patchelf rapidjson-dev python3-dev

An appropriate PyTorch container from NGC must be used. For example, to build a backend that uses the 21.02 version of the PyTorch container from NGC:

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

The following required Triton repositories will be pulled and used in the build. By default the "main" branch/tag will be used for each repo but the listed CMake argument can be used to override.

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]

Build the PyTorch Backend With Custom PyTorch

Currently, Triton requires that a specially patched version of PyTorch be used with the PyTorch backend. The full source for these PyTorch versions are available as Docker images from NGC. For example, the PyTorch version compatible with the 21.02 release of Triton is available as nvcr.io/nvidia/pytorch:21.02-py3.

Copy over the LibTorch and Torchvision headers and libraries from the PyTorch NGC container into local directories. You can see which headers and libraries are needed/copied from the 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

Using the PyTorch Backend

Parameters

Disabling the optimized execution of the PyTorch models is done through the Parameters section of the model's 'config.pbtxt' file.

DISABLE_OPTIMIZED_EXECUTION: Boolean flag to disable the optimized execution of TorchScript models. By default the optimized execuiton is always enabled.

The initial calls to a loaded TorchScript model take extremely long. Due to this longer model warmup issue, Triton also allows execution of models without these optimizations. In some models, optimized execution does not benefit performance as seen here and in other cases impacts performance negatively, as seen here.

The section of model config file specifying these parameters will look like:

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

Important Note