Skip to content

[BYOC][TensorRT] Reuse TRT engines based on max_batch_size for dynamic batching, improve device buffer allocation#8172

Merged
comaniac merged 5 commits into
apache:mainfrom
trevor-m:trt-dynamic-batching-reuse
Jun 3, 2021
Merged

[BYOC][TensorRT] Reuse TRT engines based on max_batch_size for dynamic batching, improve device buffer allocation#8172
comaniac merged 5 commits into
apache:mainfrom
trevor-m:trt-dynamic-batching-reuse

Conversation

@trevor-m

@trevor-m trevor-m commented Jun 1, 2021

Copy link
Copy Markdown
Contributor

This PR makes two changes which will help reduce GPU memory usage by TensorRT for models that use a dynamic batch dimension (tensors with shape like (relay.Any(), 3, 224, 224)).

  1. TensorRT engines are built with a "Max Batch Size" parameter. This means the engine can be used for inputs with any batch size from 1 to max_batch size. Previously, we built a new TensorRT engine for each unique batch size encountered at runtime. With this PR, when we encounter a new batch size, we will first try to match it to an already built engine with an equal or higher batch size. This will cause only one engine to exist at a given point in time, saving memory.. This is now the default mode. The previous behavior of building unique engine for each batch size is available through env variable TVM_TENSORRT_MULTI_ENGINE=1.

  2. Because of the first change, we have to rethink how the GPU device buffers are allocated because now an engine can be used for multiple batch sizes. This PR decouples the device buffers from the engine, so there is only one set of device buffers for subgraph. They will be allocated only for the largest batch size encountered.This will further reduce memory usage since only one buffer per input is allocated, while previously each engine would have its own set of buffer. This will also fix the issue from Fix Segmentation Fault For Tensorrt BYOC when TVM_TENSORRT_CACHE_DIR is Set #7162.

@trevor-m

trevor-m commented Jun 1, 2021

Copy link
Copy Markdown
Contributor Author

@codeislife99 @comaniac @anijain2305 Could you please review?

Comment thread src/runtime/contrib/tensorrt/tensorrt_runtime.cc Outdated
Comment thread src/runtime/contrib/tensorrt/tensorrt_runtime.cc Outdated
Comment thread src/runtime/contrib/tensorrt/tensorrt_runtime.cc Outdated

@comaniac comaniac left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@comaniac comaniac merged commit cc3d60e into apache:main Jun 3, 2021
@comaniac

comaniac commented Jun 3, 2021

Copy link
Copy Markdown
Contributor

Thanks @trevor-m

trevor-m pushed a commit to neo-ai/tvm that referenced this pull request Jun 4, 2021
…c batching, improve device buffer allocation (apache#8172)

* Reuse TRT engines based on max_batch_size for dynamic batching. Improve how device buffers are allocated

* Fix python formatting

* Allow user to configure engine building mode using TVM_TENSORRT_MULTI_ENGINE

* Update doc

* Typo
trevor-m pushed a commit to trevor-m/tvm that referenced this pull request Jun 17, 2021
…c batching, improve device buffer allocation (apache#8172)

* Reuse TRT engines based on max_batch_size for dynamic batching. Improve how device buffers are allocated

* Fix python formatting

* Allow user to configure engine building mode using TVM_TENSORRT_MULTI_ENGINE

* Update doc

* Typo
trevor-m pushed a commit to neo-ai/tvm that referenced this pull request Jun 17, 2021
…c batching, improve device buffer allocation (apache#8172)

* Reuse TRT engines based on max_batch_size for dynamic batching. Improve how device buffers are allocated

* Fix python formatting

* Allow user to configure engine building mode using TVM_TENSORRT_MULTI_ENGINE

* Update doc

* Typo
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants