Hierarchical Parameter Server Plugin for TensorRT
Introduction to the HPS Plugin for TensorRT
Hierarchical Parameter Server (HPS) is a distributed inference framework that is dedicated to deploying large embedding tables and realizing the low-latency retrieval of embeddings. The framework combines a high-performance GPU embedding cache with a hierarchical storage architecture that encompasses different types of database backends. The HPS plugin for TensorRT can be integrated into the TensorRT network as a custom layer to build the engine. The HPS-integrated engine can perform low-latency embedding lookup for large tables and TensorRT-accelerated forward propagation for dense network at the same time.
Workflow
The workflow to leverage the HPS plugin for TensorRT can be summarized as:
Convert trained models to ONNX: The models trained with different frameworks are converted to ONNX using the popular tools tf2onnx, torch.onnx, hugectr2onnx, and so on.
Perform ONNX graph surgery: The node for embedding lookup in the ONNX graph is replaced by the placeholder of HPS plugin for TensorRT using the tool ONNX GraphSurgeon, as shown in Fig. 1.
Build the HPS-integrated TensorRT engine: We can build the TensorRT engine based on the modified ONNX graph where the HPS can leveraged as a custom plugin layer.
Deploy the engine with Triton Inference Server TensorRT backend: The HPS-integrated TensorRT engine is deployed with Triton Inference Server TensorRT backend. Set the
LD_PRELOAD=/usr/local/hps_trt/lib/libhps_plugin.so
environment variable to load the plugin shared library when you start Triton Inference Server.
Installation
Compute Capability
The plugin supports the following compute capabilities:
Compute Capability |
GPU |
SM |
---|---|---|
7.0 |
NVIDIA V100 (Volta) |
70 |
7.5 |
NVIDIA T4 (Turing) |
75 |
8.0 |
NVIDIA A100 (Ampere) |
80 |
Installing HPS Using NGC Containers
All NVIDIA Merlin components are available as open source projects. However, a more convenient way to use these components is by using our Merlin NGC containers. These containers allow you to package your software application, libraries, dependencies, and runtime compilers in a self-contained environment. When installing HPS using NGC containers, the application environment remains portable, consistent, reproducible, and agnostic to the underlying host system’s software configuration.
HPS is included in the Merlin Docker containers that are available from the NVIDIA GPU Cloud (NGC) catalog. Access the catalog of containers at https://catalog.ngc.nvidia.com/containers. To use these Docker containers, you must install the NVIDIA Container Toolkit to provide GPU support for Docker.
The following sample commands pull and start the Merlin TensorFlow container, Merlin PyTorch container, or Merlin HugeCTR container:
Merlin TensorFlow
# Run the container in interactive mode
$ docker run --gpus=all --rm -it --cap-add SYS_NICE nvcr.io/nvidia/merlin/merlin-tensorflow:23.02
Merlin PyTorch
# Run the container in interactive mode
$ docker run --gpus=all --rm -it --cap-add SYS_NICE nvcr.io/nvidia/merlin/merlin-pytorch:23.02
Merlin HugeCTR
# Run the container in interactive mode
$ docker run --gpus=all --rm -it --cap-add SYS_NICE nvcr.io/nvidia/merlin/merlin-hugectr:23.02
You can check the existence of the HPS plugin for TensorRT after launching the container by running the following Python statements:
import ctypes
handle = ctypes.CDLL("/usr/local/hps_trt/lib/libhps_plugin.so", mode=ctypes.RTLD_GLOBAL)
Example Notebooks
We provide a collection of examples as Jupyter Notebooks that demonstrate how to build the HPS-integrated TensorRT engine for models trained with TensorFlow, PyTorch, or HugeCTR.
Benchmark
We benchmark the DLRM TRT engine with HPS plugin in hps_dlrm_benchmark.md.