Benchmarking NLP Model with TensorRT NVIDIA Triton Inference Server and Inference Recommender on Amazon SageMaker

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

This notebook demonstrates the use of Amazon SageMaker Inference recommender to perform custom load testing in order to performance fine tune the NLP BERT Model serving using NVIDIA Triton Serving on SageMaker.

Amazon SageMaker is a fully managed service for data science and machine learning workflows. It helps data scientists and developers to prepare, build, train, and deploy high-quality ML models quickly by bringing together a broad set of capabilities purpose-built for ML.

Now, NVIDIA Triton Inference Server can be used to serve models for inference in Amazon SageMaker. Thanks to the new NVIDIA Triton container image, you can easily serve ML models and benefit from the performance optimizations, dynamic batching, and multi-framework support provided by NVIDIA Triton. Triton helps maximize the utilization of GPU and CPU, further lowering the cost of inference.

SageMaker Inference Recommender is a new capability of SageMaker that reduces the time required to get machine learning (ML) models in production by automating performance benchmarking and load testing models across SageMaker ML instances. You can use Inference Recommender to deploy your model to a real-time inference endpoint that delivers the best performance at the lowest cost.

This notebook was tested with the conda_python3 kernel on an Amazon SageMaker notebook instance of type ml.g4dn.8xlarge with 50GB EBS volume.

Introduction to NVIDIA Triton Server

NVIDIA Triton Inference Server was developed specifically to enable scalable, cost-effective, and easy deployment of models in production. NVIDIA Triton Inference Server is open-source inference serving software that simplifies the inference serving process and provides high inference performance.

Some key features of Triton are:

• Support for Multiple frameworks: Triton can be used to deploy models from all major frameworks. Triton supports TensorFlow GraphDef, TensorFlow SavedModel, ONNX, PyTorch TorchScript, TensorRT, RAPIDS FIL for tree based models, and OpenVINO model formats.
• Model pipelines: Triton model ensemble represents a pipeline of one or more models or pre/post processing logic and the connection of input and output tensors between them. A single inference request to an ensemble will trigger the execution of the entire pipeline.
• Concurrent model execution: Multiple models (or multiple instances of the same model) can run simultaneously on the same GPU or on multiple GPUs for different model management needs.
• Dynamic batching: For models that support batching, Triton has multiple built-in scheduling and batching algorithms that combine individual inference requests together to improve inference throughput. These scheduling and batching decisions are transparent to the client requesting inference.
• Diverse CPUs and GPUs: The models can be executed on CPUs or GPUs for maximum flexibility and to support heterogeneous computing requirements.

Note: This initial release of NVIDIA Triton on SageMaker will only support a single model. Future releases will have multi-model support. A minimal config.pbtxt configuration file is required in the model artifacts. This release doesn't support inferring the model config automatically.

Install packages

Installs the dependencies required to package the model and run inferences using Triton server.

Imports

Set Variables

We set SageMaker variables and other variables below, also define the IAM role that will give Amazon SageMaker access to the model artifacts and the NVIDIA Triton ECR image.

Amazon SageMaker Triton Inference Server Deep Learning Container Image

Let's retrieve Amazon SageMaker NVIDIA Triton Inference server container image based on the account ID you are running this notebook.

Set triton_image_uri based on the account_id and region information

NLP Use case

Deploying and scaling NLP models in a production set up can be quite challenging. NLP models are often very large in size, containing millions of model parameters. Optimal model configurations are required to satisfy stringent performance and scalability of production grade NLP applications

In this notebook, we will benchmark a NLP use case using SageMaker Triton inference server and recommend performance tuning optimizations for the below NLP profile. We will use a large pre-trained transformer based bert-large-uncased model which has about 336 million model parameters. The input sentence used for the binary classification model will be padded and truncated to a maximum input sequence length 512 tokens. The inference load test will simulate to achieve 500 TPS (30000 maximum invocations per minute) and model latency of < 0.5 seconds (500 milliseconds)

NVIDIA Triton Setup with Amazon SageMaker

1. We will use this script generate_models.sh to generate the TensorRT plan to be used with NVIDIA Triton inference server.

2. The script load the pre-trained bert_large_uncased model and saving it ONNX format can be found in this onnx_exporter.py

3. Pre-trained model is loaded in torchscript format and model artifacts are saved used in onnx exporter(model.onnx)

4. trtexec is a tool to quickly utilize TensorRT without having to develop your own application. The trtexec tool has three main purposes:

   • benchmarking networks on random or user-provided input data.
   • generating serialized engines from models.
   • generating a serialized timing cache from the builder.

5. After trtexec execution is complete, model plan file(model_bs16.plan) is generated. This file will be used as model artifact by Triton

6. We used the pre-configured config.pbtxt file provided with this repo to specify model configuration which Triton uses to load the model.

7. We tar the model directory and upload it to s3 to later create a SageMaker Model.

Note: Amazon SageMaker expects the model tarball file to have a top level directory with the same name as the model defined in the config.pbtxt. Below is the sample model directory structure

bert
├── 1
│   └── model.plan
└── config.pbtxt

This step is expected to run for ~60 minutes.

The script saves the model in this workspace directory

Model Configuration

• Each model in a model repository must include a model configuration that provides required and optional information about the model. Typically, this configuration is provided in a config.pbtxt file specified as ModelConfig protobuf.
• The model configuration name property is optional. If the name of the model is not specified in the configuration it is assumed to be the same as the model repository directory containing the model. If name is specified it must match the name of the model repository directory containing the model
• The max_batch_size property indicates the maximum batch size that the model supports for the types of batching that can be exploited by Triton. If the model's batch dimension is the first dimension, and all inputs and outputs to the model have this batch dimension, then Triton can use its dynamic batcher or sequence batcher to automatically use batching with the model. In this case max_batch_size should be set to a value greater-or-equal-to 1 that indicates the maximum batch size that Triton should use with the model
• Each model input and output must specify a name, datatype, and shape. The name specified for an input or output tensor must match the name expected by the model.

The below is the baseline configuration for PyTorch model

You can find other model configurations that are used in the benchmarking exercise in this workspace directory.

1. config-dg.pbtxt (Dynamic batching enabled)
2. config-ig.pbtxt (Multiple instance group)
3. config-db-ig.pbtxt (Dynamic batching, Multiple instance group enabled)

To execute benchmark with different model configurations, copy above files with config.pbtxt name in the model repository and reupload the model tar files.

We will copy the model.pt file and create model tar file to be uploaded to S3. The model archive tar file will be used by Triton Inference server

Create Amazon SageMaker Real Time Endpoint

We start off by creating a sagemaker model from the model files we uploaded to s3 in the previous step.

In this step we also provide an additional Environment Variable i.e. SAGEMAKER_TRITON_DEFAULT_MODEL_NAME which specifies the name of the model to be loaded by Triton. The value of this key should match the folder name in the model package uploaded to s3. This variable is optional in case of a single model. In case of ensemble models, this key has to be specified for Triton to startup in SageMaker.

Additionally, customers can set SAGEMAKER_TRITON_BUFFER_MANAGER_THREAD_COUNT and SAGEMAKER_TRITON_THREAD_COUNT for optimizing the thread counts.

Note: The current release of Triton (21.08-py3) on SageMaker doesn't support running instances of different models on the same server, except in case of ensembles. Only multiple model instances of the same model are supported, which can be specified under the instance-groups section of the config.pbtxt file.

Create payload

Create payload in JSON format and upload it on S3. This will be used by Inference Recommender to run the custom load test.

If you want to change the payload (Token Length), below are the changes -

1. Change the JSON with shape reflecting the right token length below
2. Change the tokenize_text method to reflect the token length
3. Change the config.pbtxt the triton* folder to reflect the input id and attention mask length.

Amazon SageMaker Inference Recommender set up

Set the Domain, Task, Framework, version and Model for Inference Recommender Job.

Create the Triton Container Dictionary object and Model Package group for Inference recommender Job

Amazon SageMaker model registry model package with domain, task and Inference container specification information. Specify the list of supported inference instance types in SupportedRealtimeInferenceInstanceTypes parameter. Also, define the ContentType and MIME type information

Amazon SageMaker Inference Recommender Custom Load Test

Create Custom Inference Recommender Job for Triton Container serving BERT Model with 512 Token length

Use the create_inference_recommendations_job to create an Inference Recommender load test and specify below parameters

• Specify Advanced for the JobType field and provide:
• A job name for your load test (JobName).
• The Amazon Resource Name (ARN) of an IAM role that enables Inference Recommender to perform tasks on your behalf.
• A traffic pattern of the load test (TrafficPattern)
  • Initial number of users = 2
  • Spawn Rate = 3 (creates 3 new users every 3 minutes for a duration of 15 minutes)
• An endpoint configuration dictionary (InputConfig) where you specify an AWS instance type against which to run benchmarks
• StoppingConditions (Inference recommender would adjust the initial number of instances to satisfy below stopping conditions)
  • MaxInvocations is set to 30000
  • ModelLatencyThresholds p95 threshold for 500 ms

Let's get the inference recommender job details using describe_inference_recommendations_job boto3 API

Visualize CloudWatch Metrics

Use get_endpoint_metrics helper functions, visualize the CloudWatch metrics. This will provide detailed overview of resource usage during the load test. Metrics such GPU Memory utilization, Invocations and Model Latency metrics will allow to tweak NVIDIA Triton model configuration to improve application performance.

Conclusion

This notebook provides an overview of NVIDIA Triton support for TensorRT framework, steps to leverage Amazon SageMaker Inference recommender to execute custom load test and find the right configurations, instance types for large NLP workloads. You can use different Triton configurations, visualize resource consumptions and optimize GPU resource utilization during inference.

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.