Allow (temporarily?) non-fake input during ONNX export with fake mode (#105246)
Although input and model are expected to be fake during ONNX export with fake mode enabled, apparently some models can create new parameters during tracing. That makes internal checks on dynamo side to fail when we dont set `allow_non_fake_input=True` for `torch._dynamo.export`.
https://github.com/pytorch/pytorch/issues/105077 tracks this issue and if a proper fix is done, we will set `allow_non_fake_input=False` again
Additionally to that, a possible bug was found at torch.nn.Module.state_dict() in which some registered buffers are not listed.
This is being tracked by https://github.com/pytorch/pytorch/issues/105233 but in the mean time, we are merging `state_dict()` and `named_buffers()` results to create a full `state_dict` for the model
Two more complex/larger tests are added to the ONNX export which are the same for the experimental symbolic tracing: tiny gpt2 and toy mlp (https://github.com/pytorch/pytorch/blob/main/test/onnx/test_fx_to_onnx_with_onnxruntime.py#L766-L825)
ps: https://github.com/pytorch/pytorch/issues/105464 tracks pending tasks/limitations from this PR
Pull Request resolved: https://github.com/pytorch/pytorch/pull/105246
Approved by: https://github.com/BowenBao