Store the arguments used to trace the exported program in itself (#115)

Summary:
X-link: pytorch/pytorch#107877

Pull Request resolved: #115

Continuation of pytorch/pytorch#107704

Differential Revision: D48637348

fbshipit-source-id: 16ac9543d2d99c79fec58a055e46d8c0599c1821

Author: angelayi

exir/backend/backend_api.py

@@ -302,6 +302,7 @@ def to_backend(
         copy.deepcopy(edge_program.range_constraints),
         copy.deepcopy(edge_program.equality_constraints),
         copy.deepcopy(edge_program.module_call_graph),
+        (),
     )
 
 

exir/capture/_capture.py

@@ -213,6 +213,7 @@ def convert_to_fake(x):
         {},
         [],
         [],
+        (),
     )
     return ExirExportedProgram(ep, False)
 

exir/lowered_backend_module.py

@@ -253,6 +253,7 @@ def create_exported_program_from_submodule(
         range_constraints=copy.deepcopy(owning_program.range_constraints),
         equality_constraints=[],
         module_call_graph=[],
+        original_traced_arguments=(),
     )
 
 

exir/serde/serialize.py

@@ -291,6 +291,7 @@ def serialize_bytes(b: bytes) -> str:
         (
             serialized_original_module,
             serialized_original_state_dict,
+            _,
         ) = ExportedProgramSerializer().serialize(lowered_module.original_module)
 
         serialized_processed_bytes = serialize_bytes(lowered_module.processed_bytes)
@@ -312,8 +313,8 @@ def serialize_bytes(b: bytes) -> str:
 
 class ExportedProgramSerializer(export_serialize.ExportedProgramSerializer):
     def serialize(
-        self, exported_program: torch._export.ExportedProgram
-    ) -> Tuple[schema.ExportedProgram, bytes]:
+        self, exported_program: ep.ExportedProgram
+    ) -> Tuple[schema.ExportedProgram, bytes, bytes]:
         assert isinstance(exported_program, torch._export.ExportedProgram)
         gm_serializer = GraphModuleSerializer(
             exported_program.graph_signature, exported_program.call_spec
@@ -337,7 +338,8 @@ def serialize(
                 equality_constraints=serialized_equality_constraints,
                 schema_version=schema.SCHEMA_VERSION,
             ),
-            export_serialize.serialize_state_dict(gm_serializer.state_dict),
+            export_serialize.serialize_torch_artifact(gm_serializer.state_dict),
+            b"",
         )
 
 
@@ -595,6 +597,7 @@ def deserialize_lowered_module(
         original_module = ExportedProgramDeserializer().deserialize(
             serialized_lowered_module.original_module,
             base64.b64decode(serialized_lowered_module.original_state_dict),
+            b"",
         )
 
         lowered_module = ExirLoweredBackendModule(
@@ -612,6 +615,7 @@ def deserialize(
         self,
         serialized_exported_program: schema.ExportedProgram,
         serialized_state_dict: bytes,
+        serialized_original_traced_args: bytes,
     ) -> exir.ExportedProgram:
         symbol_name_to_range = {
             k: symbolic_shapes.ValueRanges(
@@ -620,7 +624,8 @@ def deserialize(
             )
             for k, v in serialized_exported_program.range_constraints.items()
         }
-        state_dict = export_serialize.deserialize_state_dict(serialized_state_dict)
+        state_dict = export_serialize.deserialize_torch_artifact(serialized_state_dict)
+        assert isinstance(state_dict, dict)
         (
             graph_module,
             sig,
@@ -656,27 +661,29 @@ def deserialize(
             range_constraints,
             equality_constraints,
             [],
+            (),
         )
 
 
 def serialize(
     exported_program: torch._export.ExportedProgram,
     opset_version: Optional[Dict[str, int]] = None,
-) -> Tuple[bytes, bytes]:
-    serialized_exported_program, serialized_state_dict = ExportedProgramSerializer(
+) -> Tuple[bytes, bytes, bytes]:
+    serialized_exported_program, *_ = ExportedProgramSerializer(
         opset_version
     ).serialize(exported_program)
     json_program = json.dumps(
         dataclasses.asdict(serialized_exported_program),
         cls=export_serialize.EnumEncoder,
     )
     json_bytes = json_program.encode("utf-8")
-    return json_bytes, serialized_state_dict
+    return json_bytes, *_
 
 
 def deserialize(
     exported_program_bytes: bytes,
     state_dict: bytes,
+    original_traced_args: bytes,
     expected_opset_version: Optional[Dict[str, int]] = None,
 ) -> exir.ExportedProgram:
     exported_program_str = exported_program_bytes.decode("utf-8")
@@ -689,5 +696,7 @@ def deserialize(
         schema.ExportedProgram, exported_program_dict
     )
     return ExportedProgramDeserializer(expected_opset_version).deserialize(
-        serialized_exported_program, state_dict
+        serialized_exported_program,
+        state_dict,
+        original_traced_args,
     )

sdk/etrecord/_etrecord.py

@@ -38,7 +38,7 @@ def _handle_exported_program(
     etrecord_zip: ZipFile, module_name: str, method_name: str, ep: ExportedProgram
 ) -> None:
     assert isinstance(ep, ExportedProgram)
-    serialized_ep, serialized_state_dict = serialize(ep)
+    serialized_ep, serialized_state_dict, _ = serialize(ep)
     etrecord_zip.writestr(f"{module_name}/{method_name}", serialized_ep)
     etrecord_zip.writestr(
         f"{module_name}/{method_name}_state_dict", serialized_state_dict
@@ -217,6 +217,7 @@ def parse_etrecord(etrecord_path: str) -> ETRecord:
         graph_map[serialized_file] = deserialize(
             etrecord_zip.read(serialized_file),
             etrecord_zip.read(serialized_state_dict_file),
+            b"",
         )
 
     return ETRecord(