[SOT] Copy cached meta info to avoid got wrong symbol (#71971)

SigureMo · web-flow · commit 54fd11b2ee5e · 2025-03-29T23:10:09.000+08:00
diff --git a/python/paddle/jit/sot/infer_meta.py b/python/paddle/jit/sot/infer_meta.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
+import copy
 from functools import cached_property
 from typing import TypeVar
 
@@ -102,6 +103,13 @@ def from_value(value: paddle.pir.Value) -> DistInfo:
             value._local_shape,
         )
 
+    def __deepcopy__(self, memo):
+        return DistInfo(
+            mesh=copy.deepcopy(self.mesh),
+            dims_mapping=copy.deepcopy(self.dims_mapping),
+            local_shape=copy.deepcopy(self.local_shape),
+        )
+
     def __repr__(self) -> str:
         return f"DistInfo(mesh={self.mesh}, dims_mapping={self.dims_mapping}, local_shape={self.local_shape})"
 
@@ -270,6 +278,18 @@ def guard_str(self):
         shape = self.shape_with_special_symbol(SymbolicInt())
         return f"({shape}, {self.dtype}, {self.stop_gradient})"
 
+    def __deepcopy__(self, memo):
+        return MetaInfo(
+            list(self.shape),
+            self.dtype,
+            self.stop_gradient,
+            self.name,
+            self.persistable,
+            self.type,
+            self.place,
+            dist_info=copy.deepcopy(self.dist_info),
+        )
+
     def __repr__(self):
         return meta_str(self.shape, self.dtype, self.stop_gradient)
 
@@ -541,6 +561,9 @@ def infermeta_grad(
 
 
 class InferMetaCache(Cache, metaclass=Singleton):
+    def __init__(self):
+        super().__init__(copy=True)
+
     def key_fn(
         self, func, *args, **kwargs
     ):  # args & kwargs have transformed to MetaInfo
@@ -556,6 +579,9 @@ def value_fn(self, func, *args, **kwargs):
 
 
 class LayerInferMetaCache(Cache, metaclass=Singleton):
+    def __init__(self):
+        super().__init__(copy=True)
+
     def key_fn(self, layer, *args, **kwargs):
         params = [
             MetaInfo.from_value(x)
diff --git a/python/paddle/jit/sot/utils/utils.py b/python/paddle/jit/sot/utils/utils.py
@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import builtins
+import copy
 import inspect
 import sys
 import time
@@ -267,12 +268,13 @@ def is_true(*args):
 
 
 class Cache:
-    def __init__(self, weak=False):
+    def __init__(self, weak=False, copy=False):
         if not weak:
             self.cache = {}
         else:
             self.cache = WeakValueDictionary()
         self.hit_num = 0
+        self.copy = copy
 
     def __call__(self, *args, **kwargs):
         cache_key = self.key_fn(*args, **kwargs)
@@ -281,7 +283,10 @@ def __call__(self, *args, **kwargs):
         if cache_key in self.cache:
             log(5, "cache hit: ", cache_key, "\n")
             self.hit_num += 1
-            return self.cache[cache_key]
+            cache_item = self.cache[cache_key]
+            if self.copy:
+                cache_item = copy.deepcopy(cache_item)
+            return cache_item
         value = self.value_fn(*args, **kwargs)
         self.cache[cache_key] = value
         return value
