Repository: pytorch/pytorch

Issue:
`run_decompositions` in `ExportedProgram` doesn't keep the same model state_dict
### 🐛 Describe the bug

After running `run_decompositions`, the exported program is not the same.
https://github.com/pytorch/pytorch/blob/e0d2a24967218d7c39e24f66bb6c4836c9d1d427/torch/export/exported_program.py#L351-L361

However, `the state_dict` is also updated.
https://github.com/pytorch/pytorch/blob/e0d2a24967218d7c39e24f66bb6c4836c9d1d427/torch/export/exported_program.py#L458-L472

Should decomposition happen in-place, instead of creating a whole new one?

### Versions

main branch

cc @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @thiagocrepaldi 


Discussion:
I think running decompositions should create a new exported program. What do you mean by the state dict is updated? Do you happen to have an example code?
> I think running decompositions should create a new exported program. What do you mean by the state dict is updated? Do you happen to have an example code?

I will write a full repro soon, but the issue is that we save the original `torch.export.ExportedProgram` during ONNX export (through `torch.onnx.dynamo_export`) for future use. This is convenient as we don't need to add `ExportedProgram` as input to our APIs for later use when we want to run the ONNX model. We can fetch all initializers and buffers through `ONNXProgram.exported_program.state_dict`, similar to how `ExportedProgram.__call__` does to feed all lifted params/buffers to the model.forward as input

During the ONNX export process, we call `ExportedProgram.run_decompositions` which returns a new `ExportedProgram` with a *copy of the original `state_dict`* (instead of an in-place-updated `state_dict`). The reference to the user `ExportedProgram` (and its `state_dict`) previously saved mismatches (the "pointer" to) the one that is used to convert the model to ONNX.

In the code below

```python
            class CustomModule(torch.nn.Module):
                def __init__(self):
                    super().__init__()
                    self.register_buffer(
                        "my_buffer", torch.tensor(4.0), persistent=persistent
                    )

                def forward(self, x, b):
                    output = x + b
                    (
                        self.my_buffer.add_(1.0) + 3.0
                    )  # Mutate buffer through in-place addition
                    return output
            inputs = (torch.randn(3, 3), torch.randn(3, 3))
            ep = torch.export.export(CustomModule(), args=inputs)
            onnx_program = torch.onnx.dynamo_export(ep, *inputs)
```

ideally, any state_dict change made to `ep` would be immediately perceived by `onnx_program._exported_program` because we share the reference to that object. but as `run_decompositions` create a new `state_dict`, the `onnx_program._exported_program.state_dict` differs from what is returned by `run_decompositions`

Assuming `torch.export.ExportedProgram` will be used as based for other backends to lower to other IRs, having the guarantee we will always have the same instance of `ExportedProgram` (and its `state_dict`) can ease backends life when evaluating the state_dict of the new IR against the ExportedProgram. Makes sense?
Solution: After transformation on the ExportedProgram, we check if the state_dict has been updated and if  not, we will just link the old state_dict. In this case, it shouldn't update the state_dict

Resolving patch:
diff --git a/torch/export/exported_program.py b/torch/export/exported_program.py
--- a/torch/export/exported_program.py
+++ b/torch/export/exported_program.py
@@ -466,14 +466,13 @@ def update_arg(old_arg, new_ph):
             for inp_dim1, inp_dim2 in self.equality_constraints
         ]
 
-        state_dict = self.state_dict.copy()
         lift_constant_tensor_pass(gm, new_graph_signature)
         _replace_sym_size_ops_pass(gm)
         exported_program = ExportedProgram(
             gm,
             gm.graph,
             new_graph_signature,
-            state_dict,
+            self.state_dict,
             new_range_constraints,
             new_equality_constraints,
             copy.deepcopy(self.module_call_graph),


Test patch:
diff --git a/test/export/test_export.py b/test/export/test_export.py
--- a/test/export/test_export.py
+++ b/test/export/test_export.py
@@ -1454,7 +1454,7 @@ def forward(self, x):
         m = M()
         with unittest.mock.patch("torch._export.DECOMP_TABLE", None):
             ep = export(m, inp)
-
+        state_dict = ep.state_dict
 
         FileCheck().check_count(
             "torch.ops.aten.t.default", 1, exactly=True
@@ -1469,6 +1469,7 @@ def forward(self, x):
             "torch.ops.aten.t.default", 0, exactly=True
         ).run(core_aten_ep.graph_module.code)
         self.assertTrue(torch.allclose(core_aten_ep(*inp), m(*inp)))
+        self.assertEqual(id(state_dict), id(ep.state_dict))
 
     def test_export_decomps_dynamic(self):
         class M(torch.nn.Module):