This page looks best with JavaScript enabled

Speeding Up Model Inference with TensorRT

 ·  ☕ 4 min read

1. What Is TensorRT

TensorRT is a C++ library mainly used for high-performance inference acceleration on NVIDIA GPUs. It provides both a C++ API and a Python API for integration.

The mainstream deep learning frameworks TensorRT supports are:

  • Caffe, where TensorRT can read the prototxt format directly
  • TensorFlow, where the TensorFlow pb must be converted to the uff format
  • PyTorch, where the PyTorch pth format must be converted to the onnx format
  • MXNet, where the MxNet params format must be converted to the onnx format

TensorRT is an inference engine built specifically for NVIDIA GPUs and does not apply to other vendors.

Also, TensorRT is not fully open source: the core runtime library libnvinfer.so is closed source, and only the surrounding libraries and APIs are open source.

2. How TensorRT Optimizes

  • Layer fusion

During inference, a great deal of time is wasted on launching CUDA cores and on reading and writing each layer’s inputs and outputs, creating a memory bandwidth bottleneck and wasting GPU resources.

TensorRT can fuse layers horizontally and vertically into a single CBR (Convolution-BatchNorm-ReLU) layer, so the model has fewer layers and higher GPU core utilization, which improves inference performance.

  • Quantization

When a model is trained, the precision of the network’s parameters is usually FP32; inference with 32-bit floats performs poorly and consumes a large amount of GPU memory. But at inference time, since backpropagation is not needed, parameter precision can be lowered appropriately to improve inference performance.

TensorRT provides automated support for this quantization process, reducing model accuracy loss while improving inference performance.

  • Automatic kernel tuning

TensorRT can pick the most appropriate strategy and computation method based on the graphics card architecture, the number of SMs, the core clock, and so on.

  • Dynamic tensor memory

At runtime, TensorRT allocates GPU memory dynamically to improve memory utilization and support larger networks.

  • Multi-stream parallelism

TensorRT can execute multiple streams on the same GPU at the same time, improving GPU utilization.

3. Converting Models to TensorRT

3.1 PyTorch

PyTorch uses a dynamic computation graph. Converting a PyTorch model to ONNX requires calling PyTorch’s torch.onnx.export function. There are two ways to convert a PyTorch model to ONNX:

  • trace, the tracing method

This exports a static graph of the model by actually running the model once, but it cannot recognize control flow in the model, such as loops. The idea is to have the model perform one inference and record the computation graph. What trace exports is a static graph, and the inference engine executes it more efficiently.

  • script, the scripting method

This parses the model to record every computation.

3.2 TensorFlow

TensorFlow uses a static computation graph and already has a complete graph structure.

Because the ckpt format carries a lot of redundant information, and the pb format is smaller, the model is usually converted to pb first. But the computation graph optimization for pb is worse than for uff, so it is less efficient. For this reason the model is converted to uff first, and then to TensorRT.

4. PyTorch to ONNX to TensorRT

  • Download the model
1

docker run --security-opt apparmor=unconfined --security-opt seccomp=unconfined -v $PWD:/runtime shaowenchen/huggingface-cli download --resume-download --local-dir-use-symlinks False THUDM/ChatGLM2-6B --local-dir ChatGLM2-6B
  • Download the conversion script

When converting PyTorch to the ONNX format, all you need to do is run the torch.onnx.export function. But you will often run into problems such as unsupported operators, RuntimeError, or ShapeError, which means you have to debug the model parameters and operators. There are shared, verified scripts online that can be used to convert specific models.

1

git clone https://github.com/luchangli03/export_llama_to_onnx
  • Convert to the ONNX format
1

pip install numpy transformers torch==2.1 onnx -i https://pypi.tuna.tsinghua.edu.cn/simple

1
2
3
4
5
6
7
8
9

python3 export_llama_to_onnx/export_chatglm2.py -m ChatGLM2-6B -o ./ChatGLM2-6B-onnx -p float16 -d cuda

begin load model from chatglm2-6b
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:19<00:00,  2.78s/it]
convert model to float16
convert model to cuda
finish load model from chatglm2-6b
begin export chat_glm_model
layer_num: 28

A large number of files with the .onnx extension are generated under the ./ChatGLM2-6B-onnx directory.

  • Validate the model
 1
 2
 3
 4
 5
 6
 7
 8
 9
10

import onnx

try:
    onnx.checker.check_model("./ChatGLM2-6B-onnx/chat_glm_model.onnx")
    onnx_model = onnx.load("./ChatGLM2-6B-onnx/chat_glm_model.onnx")
    print("ONNX Version:", onnx_model.ir_version)
except Exception as e:
    print("Model incorrect: {}".format(e))
else:
    print("Model correct")

At this point the output should be

1
2

ONNX Version: 8
Model correct
  • Visualize the ONNX model
1

pip install netron

1

netron --host 0.0.0.0 -p 8080 ./ChatGLM2-6B-onnx/chat_glm_model.onnx

Visit http://localhost:8080/ to inspect the model structure.

  • ONNX to TensorRT

Enter the container build environment

1

docker run --security-opt apparmor=unconfined --security-opt seccomp=unconfined --gpus device=all -v $PWD:/workspace -it --rm nvcr.io/nvidia/tensorrt:23.12-py3 bash

Start the conversion

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

trtexec --onnx=./ChatGLM2-6B-onnx/chat_glm_model.onnx --saveEngine=./ChatGLM2-6B-trt-engines/chat_glm_model.engine

[02/04/2024-07:25:07] [I] === Performance summary ===
[02/04/2024-07:25:07] [I] Throughput: 63.7494 qps
[02/04/2024-07:25:07] [I] Latency: min = 15.8738 ms, max = 16.2778 ms, mean = 15.935 ms, median = 15.9253 ms, percentile(90%) = 15.9812 ms, percentile(95%) = 16.0046 ms, percentile(99%) = 16.1168 ms
[02/04/2024-07:25:07] [I] Enqueue Time: min = 2.09204 ms, max = 6.43628 ms, mean = 3.0315 ms, median = 2.94957 ms, percentile(90%) = 3.47656 ms, percentile(95%) = 4.021 ms, percentile(99%) = 4.52615 ms
[02/04/2024-07:25:07] [I] H2D Latency: min = 0.18042 ms, max = 0.573242 ms, mean = 0.206923 ms, median = 0.194153 ms, percentile(90%) = 0.258423 ms, percentile(95%) = 0.273682 ms, percentile(99%) = 0.375793 ms
[02/04/2024-07:25:07] [I] GPU Compute Time: min = 15.5573 ms, max = 15.6531 ms, mean = 15.6016 ms, median = 15.6018 ms, percentile(90%) = 15.6212 ms, percentile(95%) = 15.6265 ms, percentile(99%) = 15.6391 ms
[02/04/2024-07:25:07] [I] D2H Latency: min = 0.12439 ms, max = 0.132263 ms, mean = 0.12646 ms, median = 0.126465 ms, percentile(90%) = 0.12793 ms, percentile(95%) = 0.128418 ms, percentile(99%) = 0.131042 ms
[02/04/2024-07:25:07] [I] Total Host Walltime: 3.04316 s
[02/04/2024-07:25:07] [I] Total GPU Compute Time: 3.02671 s
[02/04/2024-07:25:07] [I] Explanations of the performance metrics are printed in the verbose logs.
[02/04/2024-07:25:07] [I]
&&&& PASSED TensorRT.trtexec [TensorRT v8601] # trtexec --onnx=./ChatGLM2-6B-onnx/chat_glm_model.onnx --saveEngine=./ChatGLM2-6B-trt-engines/chat_glm_model.engine

There will be some performance test results, such as P95 and P99, for reference.

  • Inspect the output TensorRT model
1
2
3

ls -alh ./ChatGLM2-6B-trt-engines/chat_glm_model.engine

-rw-r--r-- 1 root root 24G Feb  4 07:24 ./ChatGLM2-6B-trt-engines/chat_glm_model.engine

For large models, however, using TensorRT-LLM for the conversion is a better choice — not only for the conversion success rate and efficiency, but also because the result is easier to deploy and integrate with Triton.

5. TensorRT-LLM to TensorRT

  • Enter the build environment
1

docker run --security-opt apparmor=unconfined --security-opt seccomp=unconfined --gpus device=0 -v $PWD:/app/tensorrt_llm/models -it --rm shaowenchen/nvidia-tensorrt-llm:v0.7.1 bash
  • Convert the TensorRT model
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

python examples/chatglm/build.py --model_dir ./models/ChatGLM2-6B \
                 --model_name chatglm2_6b \
                 --dtype float16 \
                 --parallel_build  \
                 --use_inflight_batching \
                 --enable_context_fmha \
                 --use_gemm_plugin float16 \
                 --use_gpt_attention_plugin float16 \
                 --output_dir ./models/ChatGLM2-6B-trt-engines

[02/05/2024-02:50:58] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in engine deserialization: CPU +0, GPU +11908, now: CPU 0, GPU 11908 (MiB)
[02/05/2024-02:50:58] [TRT-LLM] [I] Activation memory size: 210.50 MiB
[02/05/2024-02:50:58] [TRT-LLM] [I] Weights memory size: 11909.66 MiB
[02/05/2024-02:50:58] [TRT-LLM] [I] Max KV Cache memory size: 448.00 MiB
[02/05/2024-02:50:58] [TRT-LLM] [I] Estimated max memory usage on runtime: 12568.16 MiB
[02/05/2024-02:50:58] [TRT-LLM] [I] Serializing engine to models/ChatGLM2-6B-trt-engines/chatglm2_6b_float16_tp1_rank0.engine...
[02/05/2024-02:51:03] [TRT-LLM] [I] Engine serialized. Total time: 00:00:04
[02/05/2024-02:51:03] [TRT] [I] Serialized 59 bytes of code generator cache.
[02/05/2024-02:51:03] [TRT] [I] Serialized 35129 bytes of compilation cache.
[02/05/2024-02:51:03] [TRT] [I] Serialized 315 timing cache entries
[02/05/2024-02:51:03] [TRT-LLM] [I] Timing cache serialized to model.cache
[02/05/2024-02:51:05] [TRT-LLM] [I] Total time of building all 1 engines: 00:00:55
  • Test the model
1
2
3
4
5
6
7
8

python examples/run.py --input_text "世界上第三高的山峰是哪座？" \
                       --max_output_len=200 \
                       --tokenizer_dir ./models/ChatGLM2-6B \
                       --engine_dir=./models/ChatGLM2-6B-trt-engines/

[02/05/2024-02:55:55] [TRT-LLM] [W] Found pynvml==11.4.1. Please use pynvml>=11.5.0 to get accurate memory usage
Input [Text 0]: "世界上第三高的山峰是哪座？"
Output [Text 0 Beam 0]: "世界上第三高的山峰是干城章嘉峰，它位于印度洋上的马达加斯加岛。这座山峰高5895米，被誉为“火山中的活火山”

6. Summary

  1. TensorRT is a C++ library from NVIDIA for high-performance inference acceleration on GPUs. It improves inference performance through techniques such as layer fusion, quantization, and kernel optimization.

  2. TensorRT supports converting models from mainstream deep learning frameworks, for example PyTorch to ONNX and then to TensorRT. During the conversion you may have to deal with operators that are not supported.

  3. PyTorch can export an ONNX model with either the trace or the script method.

  4. TensorFlow models can be converted through the chain ckpt -> pb -> uff -> TensorRT.

  5. TensorRT-LLM is NVIDIA’s TensorRT solution for accelerating large model inference. For large models, using TensorRT-LLM is a better choice.

微信公众号
WRITTEN BY
微信公众号

See Also