Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: Runtime output buffer optimization #3276

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open

feat: Runtime output buffer optimization #3276

wants to merge 8 commits into from
+591 −166

Conversation

keehyuna
Copy link
Collaborator

@keehyuna keehyuna commented Nov 4, 2024

Description

Latency hiding by creating the output tensor for next output buffer

Fixes #3275

Type of change

Please delete options that are not relevant and/or add your own.

  • New feature (non-breaking change which adds functionality)

Checklist:

  • My code follows the style guidelines of this project (You can use the linters)
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas and hacks
  • I have made corresponding changes to the documentation
  • I have added tests to verify my fix or my feature
  • New and existing unit tests pass locally with my changes
  • I have added the relevant labels to my PR in so that relevant reviewers are notified

@github-actions github-actions bot added component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: runtime component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Nov 4, 2024
@github-actions github-actions bot requested a review from peri044 November 4, 2024 13:14
@keehyuna keehyuna mentioned this pull request Nov 4, 2024
Open
@keehyuna keehyuna self-assigned this Nov 4, 2024
@keehyuna keehyuna marked this pull request as ready for review November 14, 2024 14:27
@keehyuna keehyuna requested a review from narendasan November 14, 2024 14:28
@narendasan
Copy link
Collaborator

@keehyuna I believe that @peri044 has some shape inference for TRT engines utility in one of his PRs can we make sure we use that or are compatible with it?

@keehyuna
Copy link
Collaborator Author

@keehyuna I believe that @peri044 has some shape inference for TRT engines utility in one of his PRs can we make sure we use that or are compatible with it?

I think this PR doesn't have to do with fake tensor as output shape was inferred from trt function(getTensorShape/get_tensor_shape) for both c++ and python runtime.

peri044
peri044 reviewed Nov 18, 2024
View reviewed changes
core/runtime/execute_engine.cpp Outdated
@@ -263,19 +284,15 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
output_profiler_guard =
std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->output_profile_path);
}
if ((false == compiled_engine->use_pre_allocated_outputs) || shape_changed) {
Copy link
Collaborator

Choose a reason for hiding this comment

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

!compiled_engine->use_pre_allocated_outputs ?

core/runtime/execute_engine.cpp Show resolved Hide resolved
narendasan
narendasan reviewed Nov 22, 2024
View reviewed changes
Copy link
Collaborator

@narendasan narendasan left a comment

Choose a reason for hiding this comment

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

Create a context manager to enable this across subgraphs

@github-actions github-actions bot added the component: tests Issues re: Tests label Nov 25, 2024
keehyuna
keehyuna commented Nov 26, 2024
View reviewed changes
core/runtime/TRTEngine.h Outdated
struct RuntimeStates {
bool need_cudagraphs_record;
bool can_use_pre_allocated_outputs;
};
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

If weight streaming budget is changed in cuda graph mode, new capture is required.
weight streaming state will be added

tests/py/dynamo/runtime/test_002_cudagraphs_cpp.py
@@ -17,6 +17,9 @@
"Torch-TensorRT runtime is not available",
)
class TestCudagraphsCPP(TestCase):
def tearDown(self):
# Reset to default cuda graph mode after each test
torch_tensorrt.runtime.set_cudagraphs_mode(False)
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

if multiple tests are running by pytest, other test can run on cuda graph mode. Ensure to turn off cuda graph mode after test.

@github-actions github-actions bot added the documentation Improvements or additions to documentation label Nov 28, 2024
@keehyuna keehyuna mentioned this pull request Dec 3, 2024
Draft
7 tasks
peri044
peri044 reviewed Dec 9, 2024
View reviewed changes
core/runtime/TRTEngine.h Outdated Show resolved Hide resolved
core/runtime/execute_engine.cpp
"Error while setting the tensor address for shape inputs");

if (CUDAGRAPHS_MODE) {
// @peri044 I dont know if this makes sense since they are supposed to be GPU buffers
Copy link
Collaborator

Choose a reason for hiding this comment

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

For shape tensor input, TRT requires them to be on CPU. I'm not sure though if this holds true in the CUDA graphs case.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Thanks, I'm trying to find the use case that is using isShapeInferenceIO

core/runtime/execute_engine.cpp Show resolved Hide resolved
core/runtime/TRTEngine.h Outdated Show resolved Hide resolved
core/runtime/TRTEngine.h Outdated Show resolved Hide resolved
core/runtime/TRTEngine.h Outdated Show resolved Hide resolved
narendasan
narendasan approved these changes Dec 17, 2024
View reviewed changes
Copy link
Collaborator

@narendasan narendasan left a comment

Choose a reason for hiding this comment

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

LGTM, thanks for this

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

@peri044 peri044 peri044 left review comments

@narendasan narendasan narendasan approved these changes

Assignees

@keehyuna keehyuna

Labels
cla signed component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: runtime component: tests Issues re: Tests documentation Improvements or additions to documentation
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

✨[Feature] Output buffer optimization in runtime module
4 participants
@keehyuna @narendasan @peri044 @facebook-github-bot