Remove unused tensor literals from memory when saving in dynamo (#218)

* update loading saved model in torch backend to work with accelerate

* remove unused tensor literals from memory

py/torch_migraphx/dynamo/lower_dynamo.py

@@ -37,7 +37,7 @@
 from torch_migraphx.fx.fx2mgx import MGXInterpreter
 from torch_migraphx.fx.passes.pass_utils import validate_inference
 
-from .passes.pass_manager import pre_partition_pass, post_partition_pass
+from .passes.pass_manager import pre_partition_pass, post_partition_pass, post_lowering_pass
 from .passes.partition import partition, get_partition_inputs
 from .utils import print_graph_info
 
@@ -77,8 +77,9 @@ def lower_aten_to_mgx(gm: torch.fx.GraphModule,
         mgx_mod = lower_subgraph(mod, partition_inputs, name=name, **kwargs)
 
         setattr(optim_gm, name, mgx_mod)
-
-    return optim_gm
+
+    lowered_gm = post_lowering_pass(optim_gm)
+    return lowered_gm
 
 
 # @validate_inference(0.1, 0.1)

py/torch_migraphx/dynamo/passes/pass_manager.py

@@ -34,6 +34,7 @@
 from .promote_types import promote_inputs
 from .remove_empty_slice import remove_empty_slices
 from .fix_tensor_meta import fix_tensor_meta
+from .remove_lowered_constants import remove_lowered_constants
 
 
 class MGXPassManager(PassManager):
@@ -60,3 +61,11 @@ def post_partition_pass(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     ]
     post_partition_pass_mgr = MGXPassManager(passes)
     return post_partition_pass_mgr(gm)
+
+
+def post_lowering_pass(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
+    passes = [
+        remove_lowered_constants,
+    ]
+    post_partition_pass_mgr = MGXPassManager(passes)
+    return post_partition_pass_mgr(gm)

py/torch_migraphx/dynamo/passes/remove_lowered_constants.py

@@ -0,0 +1,19 @@
+import itertools
+import torch
+
+def remove_lowered_constants(gm: torch.fx.GraphModule):
+    used_literals = set()
+    for node in gm.graph.nodes:
+        if node.op == "get_attr":
+            used_literals.add(node.name)
+
+    unused_literals = set()
+    for name, _ in itertools.chain(gm.named_parameters(), gm.named_buffers()):
+        if not name in used_literals:
+            if hasattr(gm, name):
+                unused_literals.add(name)
+
+    for name in unused_literals:
+        delattr(gm, name)
+
+    return gm