`run_decompositions` in `ExportedProgram` doesn't keep the same model state_dict

[titaiwangms]

🐛 Describe the bug

After running run_decompositions, the exported program is not the same.

pytorch/torch/export/exported_program.py

Lines 351 to 361 in e0d2a24

	def run_decompositions(
	self, decomp_table: Optional[Dict[torch._ops.OperatorBase, Callable]] = None
	) -> "ExportedProgram":
	"""
	Run a set of decompositions on the exported program and returns a new
	exported program. By default we will run the Core ATen decompositions to
	get operators in the
	`Core ATen Operator Set <https://pytorch.org/docs/stable/torch.compiler_ir.html>`_.
	For now, we do not decompose joint graphs.
	"""

However, the state_dict is also updated.

pytorch/torch/export/exported_program.py

Lines 458 to 472 in e0d2a24

	state_dict = self.state_dict.copy()
	lift_constant_tensor_pass(gm, new_graph_signature, state_dict)
	_replace_sym_size_ops_pass(gm)
	exported_program = ExportedProgram(
	gm,
	gm.graph,
	new_graph_signature,
	state_dict,
	new_range_constraints,
	new_equality_constraints,
	copy.deepcopy(self.module_call_graph),
	self.example_inputs,
	self.verifier,
	self.tensor_constants,
	)

Should decomposition happen in-place, instead of creating a whole new one?

Versions

main branch

cc @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @thiagocrepaldi

[angelayi]

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?

[thiagocrepaldi]

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

            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?

[tugsbayasgalan]

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