feat: Implement support for exporting Torch-TensorRT compiled graphs using torch.export serde APIs #2249

peri044 · 2023-08-21T20:08:36Z

Description

This PR does the following

Implement support for exporting Torch-TensorRT compiled graphs using torch.export serde APIs
Converts TRT call_module submodule nodes into call_function nodes and inserts them accordingly
Generates metadata for TRT nodes and inserts them followed by getitem nodes based on the number of outputs.
These getitem nodes are required by torch._export.save/load to work.
Inlines pytorch submodules into nodes as export serde cannot handle call_module nodes
Adds 5 test cases with different scenarios and tests serde along with save/load functionality.

Type of change

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

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
This change requires a documentation update

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

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to Python style guidelines

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

gs-olive

All changes to the fx/ directory can be removed and ported to Dynamo if not already present

docsrc/user_guide/saving_models.rst

.circleci/config.yml

py/torch_tensorrt/dynamo/utils.py

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

py/torch_tensorrt/dynamo/aten_tracer.py

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

gs-olive

It looks good to me - just added a comment with an idea for a future improvement

gs-olive · 2023-10-03T23:18:23Z

py/torch_tensorrt/dynamo/partitioning/common.py

+def run_shape_analysis(
+    parent_module: torch.fx.GraphModule, inputs: Sequence[Input]
+) -> Tuple[Dict[Any, Sequence[Any]], Dict[Any, Sequence[Any]]]:


I think this is okay to leave for now - in the future we may way to try using Torch's ShapeProp utility to handle this for us

peri044 and others added 28 commits August 2, 2023 11:47

feat: Express TRT engines as nodes instead of modules

cc42ca3

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: Fix input nodes to TRT graph

afcd5ec

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: prototype

58dcc4f

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: minor change

a57f3c0

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

feat: Move tracing to use aot export apis

f1f202e

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: minor changes

abaf047

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: minor change

370099f

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: minor changes

bb1f3cf

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: Rebase with main

3d05b4d

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: rebase

8d99be5

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: minor logging updates

0aad214

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

feat: Add preliminary support for freezing tensors in Dynamo

8899735

fix: Refactor tensor freezing in Dynamo

8af2627

Key op fixes for failing tests

f6969be

fix: Add constant folding utility to freezing

bad1594

chore: Move to new export APIs

db56dd6

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: rebase with dynamo_tensor_freeze branch

bf961f5

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

feat: Add preliminary support for freezing tensors in Dynamo

b13aa82

fix: Refactor tensor freezing in Dynamo

dd95620

Key op fixes for failing tests

6bd3c64

fix: Add constant folding utility to freezing

248073f

feat: Add preliminary support for freezing tensors in Dynamo

3e5f434

fix: Refactor tensor freezing in Dynamo

6bf6945

Key op fixes for failing tests

3b6e1e7

fix: Add constant folding utility to freezing

2107d8e

chore: add BERT test case

fd5a41e

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: remove pdb

f047651

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: rebase with main

4862c68

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

facebook-github-bot added the cla signed label Aug 21, 2023

github-actions bot added the component: api [Python] Issues re: Python API label Aug 21, 2023

github-actions bot approved these changes Oct 2, 2023

View reviewed changes

chore: revert harness.py changes

b2b6373

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions bot approved these changes Oct 2, 2023

View reviewed changes

chore: fix tests

4d82e17

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions bot approved these changes Oct 2, 2023

View reviewed changes

gs-olive reviewed Oct 2, 2023

View reviewed changes

docsrc/user_guide/saving_models.rst Outdated Show resolved Hide resolved

.circleci/config.yml Outdated Show resolved Hide resolved

py/torch_tensorrt/dynamo/utils.py Outdated Show resolved Hide resolved

peri044 added 2 commits October 2, 2023 14:10

chore: address review comments

200b03f

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: updates

52017d2

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions bot approved these changes Oct 2, 2023

View reviewed changes

chore: updates

29073c3

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions bot approved these changes Oct 2, 2023

View reviewed changes

chore: rebase with main

c4c6e5c

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions bot approved these changes Oct 3, 2023

View reviewed changes

chore: fix tests

fbe929f

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions bot approved these changes Oct 3, 2023

View reviewed changes

gs-olive reviewed Oct 3, 2023

View reviewed changes

py/torch_tensorrt/dynamo/aten_tracer.py Show resolved Hide resolved

chore: address review comments

9ea829d

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions bot approved these changes Oct 3, 2023

View reviewed changes

chore: revert fx changes

f24a646

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions bot approved these changes Oct 3, 2023

View reviewed changes

gs-olive approved these changes Oct 3, 2023

View reviewed changes

peri044 merged commit 7b21322 into main Oct 4, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Implement support for exporting Torch-TensorRT compiled graphs using torch.export serde APIs #2249

feat: Implement support for exporting Torch-TensorRT compiled graphs using torch.export serde APIs #2249

peri044 commented Aug 21, 2023 •

edited

Loading

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

gs-olive left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

gs-olive left a comment

gs-olive Oct 3, 2023

feat: Implement support for exporting Torch-TensorRT compiled graphs using torch.export serde APIs #2249

feat: Implement support for exporting Torch-TensorRT compiled graphs using torch.export serde APIs #2249

Conversation

peri044 commented Aug 21, 2023 • edited Loading

Description

Type of change

Checklist:

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

gs-olive left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

gs-olive left a comment

Choose a reason for hiding this comment

gs-olive Oct 3, 2023

Choose a reason for hiding this comment

peri044 commented Aug 21, 2023 •

edited

Loading