feat: Add Selective ATen decompositions #2173

gs-olive · 2023-08-03T22:30:51Z

Description

Add sets to selectively enable or disable decompositions in Torch
Add new runtime argument enable_experimental_decompositions to enable all core aten decompositions, or a pre-selected subset thereof
Improve documentation of compilation settings overall

Fixes #2160

Type of change

New feature (non-breaking change which adds functionality)

Checklist:

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

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

py/torch_tensorrt/dynamo/lowering/_decompositions.py

gs-olive · 2023-08-04T22:14:55Z

py/torch_tensorrt/dynamo/lowering/_decompositions.py

+ENABLED_TORCH_DECOMPOSITIONS: Dict[
+    torch._ops.OpOverload, Callable
+] = get_torch_decompositions(enabled_decompositions)


Currently, the decompositions are sourced directly from Torch's main registry (via get_torch_decompositions) and may not exactly match with the _core_aten_decompositions. This is because certain decompositions which we depend on (such as native_layer_norm, may occasionally be removed from the core set).

Whenever Torch versions are upgraded, this list should be updated as well.

gs-olive · 2023-08-04T22:15:42Z

py/torch_tensorrt/dynamo/lowering/_decompositions.py

+ENABLED_TORCH_DECOMPOSITIONS: Dict[
+    torch._ops.OpOverload, Callable
+] = get_torch_decompositions(enabled_decompositions)
+TORCH_TRT_DECOMPOSITIONS: Dict[torch._ops.OpOverload, Callable] = {}


The decompositions are three dictionaries:

ENABLED_TORCH_DECOMPOSITIONS - the enabled decompositions we've pre-selected

CORE_ATEN_DECOMPOSITIONS_FILTERED (defined in get_decompositions below) - the complete set of _core_aten_decompositions Torch provides, minus the set of disabled decompositions. Note that TORCH_DECOMPOSITIONS may not be a subset of this set

TORCH_TRT_DECOMPOSITIONS - the decompositions we've written ourselves

Would TORCH_DECOMPOSITIONS only include decompositions from the get_torch_decompoistions set?

Note that TORCH_DECOMPOSITIONS may not be a subset of this set
From what I understand,

it seems like some decompositions in _core_aten_decompositions have been removed eg: aten.native_layer_norm. ENABLED_TORCH_DECOMPOSITIONS is a more complete set (from previous commit maybe ). Is this correct ?

In that case, what if we move aten.native.layer_norm to TORCH_TRT_DECOMPOSITIONS since it is useful to us and maybe other useful ones instead of maintaining a ENABLED_TORCH_DECOMPOSITIONS which overlaps with _core_aten_decompositions one ?

@narendasan
ENABLED_TORCH_DECOMPOSITIONS would only include decompositions from the _core_aten_decompositions set which are not also in disabled_decompositions

@peri044

The interpretation of ENABLED_TORCH_DECOMPOSITIONS is correct

My initial intent for TORCH_TRT_DECOMPOSITIONS was that it would only store decompositions we specifically (custom) wrote, not ones sourced from Torch, as layer_norm would be.

gs-olive · 2023-08-04T22:16:05Z

py/torch_tensorrt/dynamo/lowering/_decompositions.py


-def get_decompositions():
-    return DECOMPOSITIONS
+def get_decompositions(


The dictionary returned by get_decompositions is either ENABLED_TORCH_DECOMPOSITIONS or CORE_ATEN_DECOMPOSITIONS_FILTERED concatenated with our TORCH_TRT_DECOMPOSITIONS

py/torch_tensorrt/dynamo/lowering/_decompositions.py

narendasan · 2023-08-05T01:39:31Z

Have we thought about what this might look like if its user accessible?

narendasan · 2023-08-05T01:41:52Z

Can we add a tool to monitor these decomposition sets similar to the opset coverage tool?

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

py/torch_tensorrt/dynamo/lowering/_decompositions.py

gs-olive · 2023-08-08T00:53:10Z

The existing opset coverage tool is compatible with this PR, meaning that get_decompositions will reflect the degree of coverage of the lowering passes we currently apply via decompositions. It will not show the coverage percentages of using the core_aten_decompositions, but this is also not the default usage as of now. We can add in coverage for the enable_experimental_decompositions set to the tool as a command line argument, if preferred.

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

py/torch_tensorrt/dynamo/lowering/_decompositions.py

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

- Add sets to selectively enable or disable decompositions in Torch - Add new runtime argument `enable_experimental_decompositions` to enable all core aten decompositions, or a pre-selected subset thereof - Improve documentation of compilation settings overall

- Add decorator-wrapper to perform import-time checks on decompositions and alert the user if any custom decompositions conflict with existing registered or specified operators - Simplify code logic for dictionary merging in `get_decompositions` function - Add safety logic to ensure invariants about the decompositions are not violated

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

peri044

LGTM

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

gs-olive added component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths Story: Export/Compile Unification Issues relating to unification of Dynamo compile/export paths labels Aug 3, 2023

gs-olive requested a review from peri044 August 3, 2023 22:30

gs-olive self-assigned this Aug 3, 2023

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

github-actions bot added component: api [Python] Issues re: Python API component: lowering Issues re: The lowering / preprocessing passes component: torch_compile labels Aug 3, 2023

gs-olive changed the title ~~fix: Add enabled/disabled sets for decompositions~~ feat: Add Selective ATen decompositions Aug 3, 2023

gs-olive requested a review from narendasan August 3, 2023 22:31

github-actions bot approved these changes Aug 3, 2023

View reviewed changes

gs-olive force-pushed the selective_decompositions branch from 34a190e to bdb06d8 Compare August 4, 2023 21:56

github-actions bot approved these changes Aug 4, 2023

View reviewed changes

gs-olive force-pushed the selective_decompositions branch from bdb06d8 to 368a20e Compare August 4, 2023 22:12

github-actions bot approved these changes Aug 4, 2023

View reviewed changes

gs-olive commented Aug 4, 2023

View reviewed changes

py/torch_tensorrt/dynamo/lowering/_decompositions.py Outdated Show resolved Hide resolved

py/torch_tensorrt/dynamo/lowering/_decompositions.py Outdated Show resolved Hide resolved

py/torch_tensorrt/dynamo/lowering/_decompositions.py Show resolved Hide resolved

gs-olive commented Aug 4, 2023

View reviewed changes

github-actions bot approved these changes Aug 8, 2023

View reviewed changes

gs-olive force-pushed the selective_decompositions branch from 8b12de5 to 806b348 Compare August 8, 2023 00:47

github-actions bot approved these changes Aug 8, 2023

View reviewed changes

gs-olive commented Aug 8, 2023

View reviewed changes

py/torch_tensorrt/dynamo/lowering/_decompositions.py Show resolved Hide resolved

gs-olive mentioned this pull request Aug 8, 2023

Explore stable diffusion support in Torch-TRT #1545

Closed

gs-olive force-pushed the selective_decompositions branch from 806b348 to 7fa036f Compare August 8, 2023 20:48

github-actions bot approved these changes Aug 8, 2023

View reviewed changes

gs-olive commented Aug 8, 2023

View reviewed changes

py/torch_tensorrt/dynamo/lowering/_decompositions.py Show resolved Hide resolved

gs-olive force-pushed the selective_decompositions branch from 7fa036f to d49cadb Compare August 9, 2023 05:32

github-actions bot approved these changes Aug 9, 2023

View reviewed changes

gs-olive added 4 commits August 15, 2023 14:59

fix: Add additional decompositions for SD

7834907

mypy: Reformatting

25191a9

gs-olive force-pushed the selective_decompositions branch from d49cadb to 1e3d12e Compare August 15, 2023 22:01

github-actions bot approved these changes Aug 15, 2023

View reviewed changes

peri044 approved these changes Aug 17, 2023

View reviewed changes

feat: Move decomposition default groups to new file

2064f4f

gs-olive force-pushed the selective_decompositions branch from 1e3d12e to 2064f4f Compare August 17, 2023 20:13

github-actions bot approved these changes Aug 17, 2023

View reviewed changes

gs-olive merged commit 91fcea4 into main Aug 17, 2023

gs-olive deleted the selective_decompositions branch August 17, 2023 23:37

feat: Add Selective ATen decompositions #2173

feat: Add Selective ATen decompositions #2173

Uh oh!

Conversation

gs-olive commented Aug 3, 2023

Description

Type of change

Checklist:

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gs-olive Aug 4, 2023

Choose a reason for hiding this comment

Uh oh!

gs-olive Aug 4, 2023

Choose a reason for hiding this comment

Uh oh!

narendasan Aug 5, 2023

Choose a reason for hiding this comment

Uh oh!

peri044 Aug 5, 2023

Choose a reason for hiding this comment

Uh oh!

gs-olive Aug 7, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gs-olive Aug 4, 2023

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

narendasan commented Aug 5, 2023

Uh oh!

narendasan commented Aug 5, 2023

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gs-olive Aug 7, 2023 •

edited

Loading