128 Commits

Author SHA1 Message Date
b1aca36f4c [export] Allow legacy IR to be unflattened with weaker submodule ordering. (#123192)
Summary: In some cases we don't have information from the old IR about submodule ordering, in this case unflattener should still work in best effort mode.

Differential Revision: D55642005

Pull Request resolved: https://github.com/pytorch/pytorch/pull/123192
Approved by: https://github.com/angelayi
2024-04-02 23:08:55 +00:00
0465a90b00 [export][reland] Fix unflattened submodule ordering. (#122341) (#122507)
Summary:

Make sure the order of submodules is the same as the original eager module.

bypass-github-export-checks

Test Plan: buck test mode/opt caffe2/test:test_export -- -r test_unflatten_submodule_ordering

Differential Revision: D55251277

Pull Request resolved: https://github.com/pytorch/pytorch/pull/122507
Approved by: https://github.com/tugsbayasgalan
2024-03-25 15:22:01 +00:00
4b18ab869f [torch.export] Support is_compiling() flag for non-strict mode (#119602)
Summary: In non-strict mode of torch.export() we didn't set those `is_compiling()` to `True` which is needed by some models.

Test Plan: Unit tests and manual testing.

Differential Revision: D53624452

Pull Request resolved: https://github.com/pytorch/pytorch/pull/119602
Approved by: https://github.com/suo
2024-02-29 05:52:51 +00:00
bf4e171539 [export] support non-persistent buffers (#118969)
Summary:
X-link: https://github.com/pytorch/executorch/pull/1817

Basic support for non-persistent buffers, which are buffers that do not show up in the state dict.

One weird twist is that most of our other systems (FX, aot_export, dynamo) have completely buggy handling of non-persistent buffers. I tried to go on a wild goose chase to fix them all, but it got to be too much. So I introduced some sad rewrite passes in `_export` make the final state dict correctly align with the original module's state dict.

This exposed some bugs/ambiguous handling of parameters/buffers in existing test code. For example, `TestSaveLoad.test_save_buffer` traced over a module that was not in the root module hierarchy and caused some weird behavior. I think we should error explicitly on use cases like this: https://github.com/pytorch/pytorch/issues/118410. For now I just rewrote the tests or skipped them.

As a side effect, this diff tightened up quite a few sloppy  behaviors around state dict handling:
- Tensor attributes were getting promoted to be buffers—bad!
- Tracing through a module not in the children of the root module would add its parameters/buffers to the state dict—bad!

This behavior is unlikely to show up in user code since the model would be totally broken, but did show up in a bunch of tests.

#buildmore

Test Plan:
unit tests
sandcastle

Differential Revision: D53340041

Pull Request resolved: https://github.com/pytorch/pytorch/pull/118969
Approved by: https://github.com/guangy10, https://github.com/huydhn, https://github.com/titaiwangms
2024-02-02 19:16:08 +00:00
221747507d Revert "[export] support non-persistent buffers (#118612) (#118722)"
This reverts commit a43c28368c184ba1bf964f4fb99bec300917e2f4.

Reverted https://github.com/pytorch/pytorch/pull/118722 on behalf of https://github.com/atalman due to broke linux-jammy-py3-clang12-executorch ([comment](https://github.com/pytorch/pytorch/pull/118722#issuecomment-1921484565))
2024-02-01 14:39:29 +00:00
a43c28368c [export] support non-persistent buffers (#118612) (#118722)
Summary:
X-link: https://github.com/pytorch/executorch/pull/1769

Basic support for non-persistent buffers, which are buffers that do not show up in the state dict.

One weird twist is that most of our other systems (FX, aot_export, dynamo) have completely buggy handling of non-persistent buffers. I tried to go on a wild goose chase to fix them all, but it got to be too much. So I introduced some sad rewrite passes in `_export` make the final state dict correctly align with the original module's state dict.

This exposed some bugs/ambiguous handling of parameters/buffers in existing test code. For example, `TestSaveLoad.test_save_buffer` traced over a module that was not in the root module hierarchy and caused some weird behavior. I think we should error explicitly on use cases like this: https://github.com/pytorch/pytorch/issues/118410. For now I just rewrote the tests or skipped them.

Test Plan: added a unit test

Differential Revision: D53253905

Pull Request resolved: https://github.com/pytorch/pytorch/pull/118722
Approved by: https://github.com/SherlockNoMad, https://github.com/angelayi
2024-02-01 00:36:09 +00:00
suo
d0627cc2af [export] do not rewrite state dict when unlifting (#118611)
This is Very Bad; changing state dict keys violates one of the key contracts we have, which is "do not mess with the state dict".

Change unlift to use a similar `_assign_attr` approach that fx.GraphModule and unflatten do.

Also took the opportunity to improve the interface of `_assign_attr` to be more general.

Differential Revision: [D53139277](https://our.internmc.facebook.com/intern/diff/D53139277/)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/118611
Approved by: https://github.com/zhxchen17
ghstack dependencies: #118607, #118608, #118609, #118610
2024-01-30 19:14:19 +00:00
suo
4ee8aa6028 [export] adopt KeyPath API in nonstrict mode (#118609)
This PR rewrites two paths to use the newly-added keypaths API in pytree:
First: we were hand-rolling a tree_map during fakification because we wanted to track sources. This PR uses keypaths instead, which can do the same thing without needing custom code.

Second: our constraint error formatting was referencing placeholder names in error messages. These placeholder names are not otherwise user-visible, so they are super confusing to users (e.g. "which input does arg1_3 correspond to?"). This diff uses the `keystr` API to format the error message.

This necessitated some small refactors—generating the keystr is expensive so doing it in an f-string was very bad.

It can also be further improved—we can inspect the signature so that instead of `*args[0]` we can give people the actual argument name, which would be the ideal UX. But leaving that for later.

Differential Revision: [D53139358](https://our.internmc.facebook.com/intern/diff/D53139358/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/118609
Approved by: https://github.com/zhxchen17
ghstack dependencies: #118607, #118608
2024-01-30 19:14:11 +00:00
suo
ca090b2c77 [export] do not use tree_flatten_spec (#118608)
tree_flatten_spec is bad; it isn't synced up with `register_pytree_node` so it will not handle arbitrary custom pytrees. It's also not really maintained.

We only use it for two purposes:
- To retain kwarg ordering stability, so that if the user passes in kwargs in a different order things will still work.
- To do "structural" checks that ignore types.

In both cases, tree_flatten_spec is probably *not* the ideal way to implement the desired behavior.

## kwargs ordering
- tree_flatten_spec overwrites the behavior of ALL dictionaries, not just kwargs. This is not correct, dictionary ordering is meaningful in Python, and it's pretty trivial to write a program that relies on dict ordering.
- For kwargs, we do sort of expect that the order in which arguments are passed shouldn't matter. BUT there is one exception: `**kwargs`. In fact, [PEP 468](https://peps.python.org/pep-0468/) was introduced specifically to clarify that ordering does matter when the function being called uses `**kwargs`.

In this diff I introduce a utility function that *only* reorders kwargs. This gets us most of the way to correct—dicts are no longer reordered, but kwargs can be passed in any order.

A "fully correct" solution would need fix the corner case from PEP468. We could detect whether the top-level fn being traced uses `**kwargs` (via `inspect`), then serialize a flag for it. In ExportedProgram, we would check that flag and only re-order if `**kwargs` was unused; otherwise error if the key order doesn't match. This is a super corner case though, so I'll file it as a followup task.

## structural equivalence checking

This is another use case, where again `tree_flatten_spec` is too broad. Generally we want to treat a precise two types as the same, not override the behavior of comparison generally. So I introduce an `is_equivalent` util for this purpose.

Differential Revision: [D53168420](https://our.internmc.facebook.com/intern/diff/D53168420/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/118608
Approved by: https://github.com/zhxchen17
ghstack dependencies: #118607
2024-01-30 19:14:04 +00:00
suo
bb6eba189f [export][ez] remove unused argument from InterpreterModule (#118364)
small thing I noticed

Differential Revision: [D53113926](https://our.internmc.facebook.com/intern/diff/D53113926/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/118364
Approved by: https://github.com/angelayi
2024-01-27 06:46:01 +00:00
b5c9623835 [export] Add node meta into UnflattenedModule (#118138)
Summary: Reland of #117686

Test Plan: CI

Differential Revision: D53012028

Pull Request resolved: https://github.com/pytorch/pytorch/pull/118138
Approved by: https://github.com/zhxchen17
2024-01-25 23:37:41 +00:00
suo
7b0979ef8e [export] fixes to unflatten + custom obj composition (#117978)
The test I added for this didn't actually enable torchbind tracing, oops. Fix that and fix the issues that cropped up.

Differential Revision: [D52962205](https://our.internmc.facebook.com/intern/diff/D52962205/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/117978
Approved by: https://github.com/avikchaudhuri
ghstack dependencies: #115222
2024-01-23 05:50:41 +00:00
suo
f612e96180 [export] set proper fqn in lift constant tensor pass (#115222)
See comments: previously we were populating the lifted constant in the buffer list without an FQN, which messed up unflattening.

Differential Revision: [D50568062](https://our.internmc.facebook.com/intern/diff/D50568062/)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/115222
Approved by: https://github.com/tugsbayasgalan
2024-01-22 18:13:49 +00:00
7451dd0585 Revert "Add node meta value into UnflattenedModule (#117686)"
This reverts commit cbf24ba962f72175ec1c71a25f3379f7d9149ec1.

Reverted https://github.com/pytorch/pytorch/pull/117686 on behalf of https://github.com/PaliC due to breaks internal modeling tests ([comment](https://github.com/pytorch/pytorch/pull/117686#issuecomment-1898939899))
2024-01-18 17:46:38 +00:00
suo
ccc8440609 [export] introduce WrapperModule (#117571)
Simple module to wrap a callable. This is a useful utility for when we start requiring that torch.export take an nn.Module.

Differential Revision: [D52791310](https://our.internmc.facebook.com/intern/diff/D52791310/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/117571
Approved by: https://github.com/tugsbayasgalan, https://github.com/avikchaudhuri
ghstack dependencies: #117570
2024-01-18 03:40:34 +00:00
cbf24ba962 Add node meta value into UnflattenedModule (#117686)
Fixes #116670
Following the lead of #116720, added node.meta['val'] back to newly created subgraphs.

node.meta['val'] is essential to ONNX in terms of the shape and type information.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/117686
Approved by: https://github.com/angelayi
2024-01-18 02:37:15 +00:00
26a63907ba Ordering placeholder and get_attr nodes in unflattened module (#116910)
Previous to this PR, the generated unflattened module could mix up the order of `placeholder` and newly created `get_attr`. As `placeholder` is the input of a function, it should be placed ahead of `get_attr` nodes.

Before:
```bash
test/export/test_unflatten.py::TestUnflatten::test_placeholder_and_get_attr_ordering_after_unflattened opcode         name         target                    args                                                            kwargs
-------------  -----------  ------------------------  --------------------------------------------------------------  --------
get_attr       bias         bias                      ()                                                              {}
get_attr       weight       weight                    ()                                                              {}
placeholder    l_x_         l_x_                      ()                                                              {}
call_function  convolution  aten.convolution.default  (l_x_, weight, bias, [2, 2], [0, 0], [1, 1], False, [0, 0], 1)  {}
output         output       output                    (convolution,)                                                  {}
```

After:
```bash
test/export/test_unflatten.py::TestUnflatten::test_placeholder_and_get_attr_ordering_after_unflattened opcode         name         target                    args                                                            kwargs
-------------  -----------  ------------------------  --------------------------------------------------------------  --------
placeholder    l_x_         l_x_                      ()                                                              {}
get_attr       weight       weight                    ()                                                              {}
get_attr       bias         bias                      ()                                                              {}
call_function  convolution  aten.convolution.default  (l_x_, weight, bias, [2, 2], [0, 0], [1, 1], False, [0, 0], 1)  {}
output         output       output                    (convolution,)                                                  {}
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/116910
Approved by: https://github.com/tugsbayasgalan
ghstack dependencies: #117409, #116667, #117591, #117500
2024-01-17 23:03:15 +00:00
cb0bfcf590 Revert "Ordering placeholder and get_attr nodes in unflattened module (#116910)"
This reverts commit 12561bb5fed08283baf7a31e6678341a04e83adb.

Reverted https://github.com/pytorch/pytorch/pull/116910 on behalf of https://github.com/PaliC due to breaking internal discussed with author offline ([comment](https://github.com/pytorch/pytorch/pull/117500#issuecomment-1896516512))
2024-01-17 19:34:26 +00:00
12561bb5fe Ordering placeholder and get_attr nodes in unflattened module (#116910)
Previous to this PR, the generated unflattened module could mix up the order of `placeholder` and newly created `get_attr`. As `placeholder` is the input of a function, it should be placed ahead of `get_attr` nodes.

Before:
```bash
test/export/test_unflatten.py::TestUnflatten::test_placeholder_and_get_attr_ordering_after_unflattened opcode         name         target                    args                                                            kwargs
-------------  -----------  ------------------------  --------------------------------------------------------------  --------
get_attr       bias         bias                      ()                                                              {}
get_attr       weight       weight                    ()                                                              {}
placeholder    l_x_         l_x_                      ()                                                              {}
call_function  convolution  aten.convolution.default  (l_x_, weight, bias, [2, 2], [0, 0], [1, 1], False, [0, 0], 1)  {}
output         output       output                    (convolution,)                                                  {}
```

After:
```bash
test/export/test_unflatten.py::TestUnflatten::test_placeholder_and_get_attr_ordering_after_unflattened opcode         name         target                    args                                                            kwargs
-------------  -----------  ------------------------  --------------------------------------------------------------  --------
placeholder    l_x_         l_x_                      ()                                                              {}
get_attr       weight       weight                    ()                                                              {}
get_attr       bias         bias                      ()                                                              {}
call_function  convolution  aten.convolution.default  (l_x_, weight, bias, [2, 2], [0, 0], [1, 1], False, [0, 0], 1)  {}
output         output       output                    (convolution,)                                                  {}
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/116910
Approved by: https://github.com/tugsbayasgalan
ghstack dependencies: #117409, #116667, #117591, #117500
2024-01-17 19:12:33 +00:00
8c7e3a18ff Revert "Ordering placeholder and get_attr nodes in unflattened module (#116910)"
This reverts commit 5e0e78585d9f662ecb957c327c8d3fa31bff4f9a.

Reverted https://github.com/pytorch/pytorch/pull/116910 on behalf of https://github.com/PaliC due to breaking internal discussed with author offline ([comment](https://github.com/pytorch/pytorch/pull/117500#issuecomment-1896426304))
2024-01-17 18:42:39 +00:00
5e0e78585d Ordering placeholder and get_attr nodes in unflattened module (#116910)
Previous to this PR, the generated unflattened module could mix up the order of `placeholder` and newly created `get_attr`. As `placeholder` is the input of a function, it should be placed ahead of `get_attr` nodes.

Before:
```bash
test/export/test_unflatten.py::TestUnflatten::test_placeholder_and_get_attr_ordering_after_unflattened opcode         name         target                    args                                                            kwargs
-------------  -----------  ------------------------  --------------------------------------------------------------  --------
get_attr       bias         bias                      ()                                                              {}
get_attr       weight       weight                    ()                                                              {}
placeholder    l_x_         l_x_                      ()                                                              {}
call_function  convolution  aten.convolution.default  (l_x_, weight, bias, [2, 2], [0, 0], [1, 1], False, [0, 0], 1)  {}
output         output       output                    (convolution,)                                                  {}
```

After:
```bash
test/export/test_unflatten.py::TestUnflatten::test_placeholder_and_get_attr_ordering_after_unflattened opcode         name         target                    args                                                            kwargs
-------------  -----------  ------------------------  --------------------------------------------------------------  --------
placeholder    l_x_         l_x_                      ()                                                              {}
get_attr       weight       weight                    ()                                                              {}
get_attr       bias         bias                      ()                                                              {}
call_function  convolution  aten.convolution.default  (l_x_, weight, bias, [2, 2], [0, 0], [1, 1], False, [0, 0], 1)  {}
output         output       output                    (convolution,)                                                  {}
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/116910
Approved by: https://github.com/tugsbayasgalan
2024-01-16 22:58:37 +00:00
6413511713 [export][refactor][4/n] Make equality_constraints optional (#116233)
Summary: needed to remove equality_contraints eventually :P

Test Plan: CI

Differential Revision: D52351709

Pull Request resolved: https://github.com/pytorch/pytorch/pull/116233
Approved by: https://github.com/tugsbayasgalan
2024-01-05 00:50:52 +00:00
eb958d7552 Fix bug in unflatten pytree (#116750)
Summary: Title

Test Plan: CI

Differential Revision: D52529088

Pull Request resolved: https://github.com/pytorch/pytorch/pull/116750
Approved by: https://github.com/zhxchen17
2024-01-04 14:23:40 +00:00
199e07f108 [pytree][BE] update treespec num_children access (#116370)
Change `len(treespec.children_spes) -> treespec.num_children`.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/116370
Approved by: https://github.com/Skylion007
2023-12-24 20:54:32 +00:00
suo
d2d129de65 [sigmoid] replace unflatten with upstream version (#115468)
as title

Differential Revision: [D52000213](https://our.internmc.facebook.com/intern/diff/D52000213/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/115468
Approved by: https://github.com/zhxchen17
2023-12-22 00:56:19 +00:00
suo
b5c866db13 [export] Add FlatArgsAdapter to unflatten (#115467)
This is the final divergence between our internal/external unflatteners.

Differential Revision: [D52001135](https://our.internmc.facebook.com/intern/diff/D52001135/)

@diff-train-skip-merge
Pull Request resolved: https://github.com/pytorch/pytorch/pull/115467
Approved by: https://github.com/zhxchen17
ghstack dependencies: #115466, #115795
2023-12-21 20:52:36 +00:00
suo
01ec3d1113 [export] upstream some final fixes to OSS unflatten (#115795)
as title

Differential Revision: [D52141387](https://our.internmc.facebook.com/intern/diff/D52141387/)

@diff-train-skip-merge
Pull Request resolved: https://github.com/pytorch/pytorch/pull/115795
Approved by: https://github.com/zhxchen17
ghstack dependencies: #115466
2023-12-21 20:52:36 +00:00
suo
bc3ef1684e [export] refactor unflatten.py to be a top-level API (#115466)
This is in preparation for the merging of the internal and external versions of
the unflattener. Unflatten needs to be its own API because we are adding more
options to it in forthcoming diffs.

Differential Revision: [D52001133](https://our.internmc.facebook.com/intern/diff/D52001133/)

@diff-train-skip-merge
Pull Request resolved: https://github.com/pytorch/pytorch/pull/115466
Approved by: https://github.com/zhxchen17
2023-12-21 20:52:29 +00:00