SIGN IN SIGN UP

support non-persistent buffers (#1817)

Summary:
X-link: https://github.com/pytorch/pytorch/pull/118969

Pull Request resolved: 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

bypass-github-pytorch-ci-checks

Reviewed By: zhxchen17

Differential Revision: D53340041

fbshipit-source-id: 53ae6069daf13589f6981d673535e83213d2aded
M
Michael Suo committed
c49e1eef73c8d9e689593e0ad3e9004391fff8bd
Parent: 687af0b
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> on 2/2/2024, 6:55:02 PM