[SOT] Add internal API paddle.jit.marker.force_dynamic to ensure function or layer run under dynamic mode (#73059)

SigureMo · web-flow · commit 0d4566d18e87 · 2025-06-04T23:52:30.000+08:00
diff --git a/python/paddle/jit/marker.py b/python/paddle/jit/marker.py
@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import inspect
 from typing import (
     Callable,
     Protocol,
@@ -25,6 +26,8 @@
     ParamSpec,
 )
 
+import paddle
+
 from .dy2static.utils import (
     TransformOptions,
 )
@@ -123,3 +126,24 @@ def _mark_as_unified(fn, *, for_sot: bool, for_ast: bool):
     if fn is None:
         return lambda fn: _mark_as_unified(fn, for_sot=for_sot, for_ast=for_ast)
     return _mark_as_unified(fn, for_sot=for_sot, for_ast=for_ast)
+
+
+def force_dynamic(
+    fn: Callable[_InputT, _RetT] | type[paddle.nn.Layer] | None = None,
+) -> Callable[_InputT, _RetT]:
+    """
+    Mark a function or paddle.nn.Layer to be executed in dynamic mode, it will
+    break the graph and prevent it from being converted to static mode.
+    """
+    from paddle.jit import sot
+
+    if inspect.isclass(fn) and issubclass(fn, paddle.nn.Layer):
+        sot.utils.paddle_api_config.add_break_graph_layer_class(fn)
+        return fn
+    if inspect.isfunction(fn):
+        sot.utils.paddle_api_config.add_break_graph_function(fn)
+        return fn
+
+    raise TypeError(
+        f"Expected a callable or paddle.nn.Layer, but got {type(fn).__name__}."
+    )
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -44,7 +44,6 @@
     get_obj_stable_repr,
     get_static_function,
     hashable,
-    is_break_graph_api,
     is_break_graph_tensor_methods,
     is_builtin_fn,
     is_directly_run_api,
@@ -62,6 +61,7 @@
     DataDependencyOperationBreak,
     FallbackError,
     FallbackInlineCallBreak,
+    ForceBreak,
     InnerError,
     OtherInlineCallBreak,
     PsdbBreakReason,
@@ -70,9 +70,12 @@
     SotErrorBase,
     UnsupportedNumPyAPIBreak,
     UnsupportedOperationBreak,
-    UnsupportedPaddleAPIBreak,
     UnsupportedRandomAPIBreak,
 )
+from ....utils.paddle_api_config import (
+    break_graph_functions,
+    break_graph_layer_classes,
+)
 from ..dispatcher import Dispatcher
 from ..guard import (
     FasterStringifiedExpression,
@@ -146,6 +149,33 @@ def call_function(self, /, *args, **kwargs):
         raise NotImplementedError("call_function is not implemented.")
 
 
+class ForceBreakCallableVariable(CallableVariable):
+    def __init__(self, name: str, graph: FunctionGraph, tracker: Tracker):
+        super().__init__(graph, tracker)
+        self.name = name
+
+    def call_function(self, /, *args, **kwargs) -> VariableBase:
+        raise BreakGraphError(ForceBreak(reason_str=f"Force run {self.name}"))
+
+    def get_py_value(self, allow_tensor=False):
+        return self.value
+
+    @VariableFactory.register_from_value()
+    def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
+        if (
+            isinstance(value, paddle.nn.Layer)
+            and value.__class__ in break_graph_layer_classes
+        ):
+            return ForceBreakCallableVariable(
+                f"Layer({value.__class__.__name__})", graph, tracker
+            )
+        elif hashable(value) and value in break_graph_functions:
+            return ForceBreakCallableVariable(
+                get_obj_stable_repr(value), graph, tracker
+            )
+        return None
+
+
 class FunctionVariable(CallableVariable):
     """
     FunctionVariable is a subclass of CallableVariable used to wrap a function variable.
@@ -364,10 +394,6 @@ def __init__(
         super().__init__(fn, graph, tracker)
 
     def call_function(self, /, *args, **kwargs):
-        if is_break_graph_api(self.value):
-            raise BreakGraphError(
-                UnsupportedPaddleAPIBreak(fn_name=self.value.__name__)
-            )
         return self.graph.call_paddle_api(self.value, *args, **kwargs)
 
     @VariableFactory.register_from_value(
diff --git a/python/paddle/jit/sot/utils/exceptions.py b/python/paddle/jit/sot/utils/exceptions.py
@@ -144,6 +144,24 @@ def __init__(
         )
 
 
+class ForceBreak(UnsupportedOperationBreak):
+    def __init__(
+        self,
+        *,
+        reason_str=None,
+        file_path="",
+        line_number=-1,
+    ):
+        if reason_str is None:
+            reason_str = "Force break graph execution"
+
+        super().__init__(
+            reason_str=reason_str,
+            file_path=file_path,
+            line_number=line_number,
+        )
+
+
 class BuiltinFunctionBreak(UnsupportedOperationBreak):
     """Break reason for unsupported built-in function calls.
 
diff --git a/python/paddle/jit/sot/utils/paddle_api_config.py b/python/paddle/jit/sot/utils/paddle_api_config.py
@@ -116,9 +116,8 @@ def _get_tensor_methods():
     "paddle.nn.functional",
 }
 
-break_graph_set = set()
-
-
+break_graph_functions = set()
+break_graph_layer_classes = set()
 break_graph_tensor_method = {
     'register_hook',
     'numpy',
@@ -139,8 +138,12 @@ def is_break_graph_tensor_methods(method_name):
     return method_name in break_graph_tensor_method
 
 
-def add_break_graph_apis(apis: list):
-    break_graph_set.update(apis)
+def add_break_graph_function(fn):
+    break_graph_functions.add(fn)
+
+
+def add_break_graph_layer_class(layer_class: type[paddle.nn.Layer]):
+    break_graph_layer_classes.add(layer_class)
 
 
 def is_directly_run_api(api):
diff --git a/python/paddle/jit/sot/utils/utils.py b/python/paddle/jit/sot/utils/utils.py
@@ -38,7 +38,7 @@
     ENV_STRICT_MODE,
 )
 from .paddle_api_config import (
-    break_graph_set,
+    break_graph_functions,
     paddle_api_list,
     paddle_api_module_prefix,
 )
@@ -237,7 +237,7 @@ def in_paddle_module(func):
 
 
 def is_break_graph_api(func):
-    return func in break_graph_set
+    return func in break_graph_functions
 
 
 def is_namedtuple_class(cls):
diff --git a/test/sot/test_break_graph.py b/test/sot/test_break_graph.py
@@ -18,7 +18,7 @@
 from test_case_base import TestCaseBase
 
 import paddle
-from paddle.jit.sot.utils.paddle_api_config import add_break_graph_apis
+from paddle.jit.sot.utils.paddle_api_config import add_break_graph_function
 
 
 def ifelse_func(x, y):
@@ -74,7 +74,7 @@ def to_tensor_break_graph(x, y):
 
 class TestToTensor(TestCaseBase):
     def test_simple(self):
-        add_break_graph_apis([paddle.to_tensor])
+        add_break_graph_function(paddle.to_tensor)
         x = paddle.to_tensor(2)
         y = paddle.to_tensor(3)
         self.assert_results(to_tensor_break_graph, x, y)
diff --git a/test/sot/test_force_dynamic.py b/test/sot/test_force_dynamic.py
@@ -0,0 +1,77 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+from test_case_base import (
+    TestCaseBase,
+    test_instruction_translator_cache_context,
+)
+
+import paddle
+from paddle.jit import sot
+
+
+class EmbeddingLayer(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.embedding = paddle.nn.Embedding(10, 10)
+
+    def forward(self, x):
+        x = x + 1 - 1
+        x = self.embedding(x)
+        return x + 1
+
+
+def call_embedding_layer(x: paddle.Tensor, layer: paddle.nn.Layer):
+    return layer(x)
+
+
+def call_functional_embedding(x: paddle.Tensor, weight: paddle.Tensor):
+    x = x + 1 - 1
+    x = paddle.nn.functional.embedding(x, weight)
+    return x + 1
+
+
+class TestForceDynamic(TestCaseBase):
+    def test_embedding_layer(self):
+        paddle.jit.marker.force_dynamic(paddle.nn.Embedding)
+
+        layer = EmbeddingLayer()
+        with test_instruction_translator_cache_context() as ctx:
+            self.assertEqual(ctx.translate_count, 0)
+            self.assert_results(
+                call_embedding_layer,
+                paddle.randint(0, 10, [1, 3, 224, 224], dtype='int64'),
+                layer,
+            )
+            self.assertGreater(ctx.translate_count, 1)
+
+        sot.utils.paddle_api_config.break_graph_layer_classes.clear()
+
+    def test_functional_embedding(self):
+        paddle.jit.marker.force_dynamic(paddle.nn.functional.embedding)
+
+        weight = paddle.randn([10, 10])
+        x = paddle.randint(0, 10, [1, 3, 224, 224], dtype='int64')
+        with test_instruction_translator_cache_context() as ctx:
+            self.assertEqual(ctx.translate_count, 0)
+            self.assert_results(call_functional_embedding, x, weight)
+            self.assertGreater(ctx.translate_count, 1)
+
+        sot.utils.paddle_api_config.break_graph_functions.clear()
+
+
+if __name__ == "__main__":
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -38,7 +38,7 @@`
`38`	`38`	`ENV_STRICT_MODE,`
`39`	`39`	`)`
`40`	`40`	`from .paddle_api_config import (`
`41`		`- break_graph_set,`
	`41`	`+ break_graph_functions,`
`42`	`42`	`paddle_api_list,`
`43`	`43`	`paddle_api_module_prefix,`
`44`	`44`	`)`
`@@ -237,7 +237,7 @@ def in_paddle_module(func):`
`237`	`237`
`238`	`238`
`239`	`239`	`def is_break_graph_api(func):`
`240`		`- return func in break_graph_set`
	`240`	`+ return func in break_graph_functions`
`241`	`241`
`242`	`242`
`243`	`243`	`def is_namedtuple_class(cls):`