【SOT】fix sot memory leakage when fallback (#58423)

2742195759 · SigureMo · web-flow · commit 3d6b6b218136 · 2023-11-01T04:04:03.000-07:00
* fix sot memory leakage when fallback

* inplace clean call_stack

* [bugfix] fix layer created in forward function bugs: add check and make weakref to string ref.

---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py b/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import gc
 import traceback
 import types
 from typing import List, Tuple
@@ -228,3 +229,5 @@ def start_translate(frame: types.FrameType, **kwargs) -> GuardedFunction:
         raise InnerError(OpcodeExecutorBase.error_message_summary(e)) from e
     finally:
         simulator.cleanup()
+        del simulator
+        gc.collect()
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -26,7 +26,7 @@
 
 from ...infer_meta import InferMetaCache, LayerInferMetaCache, MetaInfo
 from ...profiler import EventGuard, event_register
-from ...symbolic.statement_ir import Symbol
+from ...symbolic.statement_ir import Reference, Symbol
 from ...symbolic.symbolic_context import SymbolicTraceContext
 from ...utils import (
     ENV_SHOW_TRACKERS,
@@ -426,6 +426,7 @@ def get_opcode_executor_stack():
     def call_layer(
         self,
         layer: PaddleLayerVariable,
+        weak_ref: bool,
         *args: VariableBase,
         **kwargs: VariableBase,
     ):
@@ -442,7 +443,7 @@ def infer_meta_fn(layer, *metas, **kwmetas):
 
         def compute_fn(layer, inputs, outputs, stacks):
             self.sir_ctx.call_LAYER(
-                layer.value,
+                Reference(layer.value, weak_ref),
                 inputs=inputs,
                 outputs=outputs,
                 stacks=stacks,
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -1460,6 +1460,7 @@ def __init__(self, frame: types.FrameType, **kwargs):
     def cleanup(self):
         self._graph.pycode_gen = None
         Dispatcher.graph = None
+        self.call_stack[:] = []
 
     @event_register("OpcodeExecutor: _prepare_virtual_env", event_level=2)
     def _prepare_virtual_env(self):
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -522,7 +522,10 @@ def __init__(
 
     def call_function(self, /, *args, **kwargs):
         self.graph.add_global_guarded_variable(self)
-        return self.graph.call_layer(self, *args, **kwargs)
+        # when layer is created in forward function, we use strong ref because it can't have
+        # weigths and buffers, see PaddleLayerClassVariable for details.
+        weak_ref = not isinstance(self.tracker, CreateLayerTracker)
+        return self.graph.call_layer(self, weak_ref, *args, **kwargs)
 
     def make_stringify_guard(self) -> list[StringifyExpression]:
         if isinstance(self.tracker, CreateLayerTracker):
@@ -740,10 +743,18 @@ class PaddleLayerClassVariable(ClassVariable):
     def __init__(self, class_: type, graph: FunctionGraph, tracker: Tracker):
         super().__init__(class_, graph, tracker)
 
+    def check_no_weight_and_buffers(self, paddle_layer):
+        has_parameters = len(paddle_layer.parameters()) > 0
+        has_buffers = len(paddle_layer.buffers()) > 0
+        return not has_parameters and not has_buffers
+
     def call_function(self, /, *args, **kwargs):
         input_py_args = [var.get_py_value() for var in args]
         input_py_kwargs = {k: v.get_py_value() for k, v in kwargs.items()}
         new_layer = self.value(*input_py_args, **input_py_kwargs)
+        assert self.check_no_weight_and_buffers(
+            new_layer
+        ), "You have created a layer in to_static function which may have Potential bugs. please create it in __init__/main function."
         return PaddleLayerVariable(
             new_layer, self.graph, CreateLayerTracker(self, args, kwargs)
         )
diff --git a/python/paddle/jit/sot/symbolic/statement_ir.py b/python/paddle/jit/sot/symbolic/statement_ir.py
@@ -22,12 +22,26 @@
 import weakref
 from typing import Any, Callable
 
-import paddle
 from paddle.utils import is_sequence, map_structure
 
 from ..utils import NameGenerator, OrderedSet, Singleton, flatten_extend
 
 
+class Reference:  # to unify weak_ref and strong_ref
+    def __init__(self, value, is_weak):
+        self.is_weak = is_weak
+        if is_weak is True:
+            self.ref = weakref.ref(value)
+        else:
+            self.ref = value
+
+    def __call__(self):
+        if self.is_weak is True:
+            return self.ref()
+        else:
+            return self.ref
+
+
 class Symbol:
     """
     Symbol is used to distinguish a string and a `math variable`.
@@ -139,15 +153,15 @@ def __init__(
 class LayerStatement(Statement):
     def __init__(
         self,
-        layer: paddle.nn.Layer,
+        layer: Reference,  # Reference of paddle.nn.Layer
         inputs: list[Symbol],
         outputs: list[Symbol],
         stacks: list[str],
     ):
         super().__init__(
             "layer", layer.__class__.__name__, inputs, outputs, stacks
         )
-        self.layer = weakref.ref(layer)
+        self.layer = layer
 
 
 class StatementIR:
diff --git a/test/sot/test_simulate_initialize.py b/test/sot/test_simulate_initialize.py
@@ -31,6 +31,18 @@ def foo(x, y):
     return out
 
 
+def foo2(x, y):
+    t = nn.Softmax()
+    out1 = t(paddle.to_tensor([x, y], dtype="float32"))
+    out2 = t(paddle.to_tensor([x, y], dtype="float32"))
+    return out1 + out2
+
+
+def error_foo(x):
+    t = nn.Linear(10, 10)
+    return t(x)
+
+
 def bar(x):
     a = A(x)
     t = paddle.to_tensor(x)
@@ -40,12 +52,20 @@ def bar(x):
 class TestInit(TestCaseBase):
     def test_init_paddle_layer(self):
         self.assert_results(foo, 1, 2)
+        self.assert_results(foo2, 1, 2)
 
     def test_init_python_object(self):
         sot_output = symbolic_translate(bar)([1.0, 2.0])
         dyn_output = bar([1.0, 2.0])
         self.assert_nest_match(sot_output, dyn_output)
 
+    def test_error(self):
+        def run():
+            inputs = paddle.randn((10, 10))
+            symbolic_translate(error_foo)(inputs)
+
+        self.assertRaises(paddle.jit.sot.utils.exceptions.InnerError, run)
+
 
 if __name__ == "__main__":
     unittest.main()
