[BYOC][TensorRT] Reuse TRT engines based on max_batch_size for dynamic batching, improve device buffer allocation #8172
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
added 2 commits
June 1, 2021 20:21
…ve how device buffers are allocated
|
@codeislife99 @comaniac @anijain2305 Could you please review? |
comaniac
requested changes
Jun 1, 2021
added 3 commits
June 2, 2021 19:09
|
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
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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)).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.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.