Remove deepcopy (#2502)

Summary:
Add fake program to remove deepcopy.
See: D54826270


Test Plan:
Imported from GitHub, without a `Test Plan:` line.

Peak memory usage: 38.7 GiB, down from 45GiB.
https://lookaside.facebook.com/intern/diff/file/data/?number=1470390332&download=1

Differential Revision: D55047794

Pulled By: lucylq

Author: lucylq

exir/backend/backend_api.py

@@ -28,6 +28,7 @@
     LoweredBackendModule,
 )
 from executorch.exir.pass_base import ExportPass
+from executorch.exir.program._fake_program import get_fake_program, get_real_program
 from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
 from torch.export import ExportedProgram
 
@@ -343,8 +344,9 @@ def to_backend(
     Returns:
         ExportedProgram: The input program, with some portions targeted for delegation.
     """
-    copied_edge_program = copy.deepcopy(edge_program)
-    partitioner_result = partitioner_instance(copied_edge_program)
+    # Use fake program, with FakeTensors in the state dict, to avoid copying large constant values.
+    fake_edge_program = get_fake_program(edge_program)
+    partitioner_result = partitioner_instance(fake_edge_program)
     tagged_exported_program = partitioner_result.tagged_exported_program
 
     # Check that the partitioner did not modify the original graph
@@ -360,6 +362,7 @@ def to_backend(
         partitioner_result.partition_tags is not None
     ), f"Partitioner {partitioner_instance} needs a `partition_tags` field containing a mapping of tags to delegate spec"
 
+    get_real_program(tagged_exported_program, edge_program)
     tagged_graph_module = _partition_and_lower(
         tagged_exported_program.graph_module, partitioner_result, edge_program
     )

exir/program/TARGETS

@@ -8,6 +8,7 @@ python_library(
         "__init__.py",
     ],
     deps = [
+        ":fake_program",
         ":program",
     ],
 )
@@ -38,3 +39,13 @@ python_library(
         "//executorch/exir/verification:verifier",
     ],
 )
+
+python_library(
+    name = "fake_program",
+    srcs = [
+        "_fake_program.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)

exir/program/__init__.py

@@ -6,6 +6,7 @@
 
 # pyre-strict
 
+from executorch.exir.program._fake_program import get_fake_program
 from executorch.exir.program._program import (
     _to_edge,
     edge_to_executorch_passes,
@@ -24,4 +25,6 @@
     "edge_to_executorch_passes",
     "EdgeProgramManager",
     "ExecutorchProgramManager",
+    "get_fake_program",
+    "get_real_program",
 ]

exir/program/_fake_program.py

@@ -0,0 +1,46 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import copy
+from typing import Dict, Union
+
+import torch
+
+from torch._guards import detect_fake_mode
+from torch.export import ExportedProgram
+
+
+def get_fake_program(real_exported_program: ExportedProgram) -> ExportedProgram:
+    fake_mode = detect_fake_mode(
+        tuple(
+            node.meta["val"]
+            for node in real_exported_program.graph.nodes
+            if node.op == "placeholder"
+        )
+    )
+    new_state_dict: Dict[str, Union[torch.Tensor, torch.nn.Parameter]] = {}
+    for key, tensor in real_exported_program.state_dict.items():
+        fake = fake_mode.from_tensor(tensor, static_shapes=True)
+        new_state_dict[key] = fake
+
+    gm = copy.deepcopy(real_exported_program.graph_module)
+    fake_exported_program = ExportedProgram(
+        root=gm,
+        graph=gm.graph,
+        graph_signature=copy.deepcopy(real_exported_program.graph_signature),
+        state_dict=new_state_dict,
+        range_constraints=copy.deepcopy(real_exported_program.range_constraints),
+        module_call_graph=copy.deepcopy(real_exported_program.module_call_graph),
+        verifier=copy.deepcopy(real_exported_program.verifier),
+        constants=real_exported_program.constants,
+    )
+    return fake_exported_program
+
+
+def get_real_program(
+    fake_exported_program: ExportedProgram, real_exported_program: ExportedProgram
+) -> None:
+    fake_exported_program._state_dict = real_exported_program.state_dict