fix patches

Author: xadupre

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -65,9 +65,13 @@ def unpatch_module(mod, info: Dict[type, Dict[type, Callable]], verbose: int = 0
                 setattr(original, n, v)
 
 
+PATCH_OF_PATCHES = set()
+
+
 def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
     # Cache serialization: to be moved into appropriate packages
     import torch
+    import transformers
     import packaging.version as pv
 
     try:
@@ -109,11 +113,17 @@ def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
     # torch.fx._pytree.register_pytree_flatten_spec(
     #           DynamicCache, _flatten_dynamic_cache_for_fx)
     # so we remove it anyway
-    if DynamicCache in torch.fx._pytree.SUPPORTED_NODES and pv.Version(
-        torch.__version__
-    ) >= pv.Version("2.7"):
+    if (
+        DynamicCache in torch.fx._pytree.SUPPORTED_NODES
+        and not PATCH_OF_PATCHES
+        and pv.Version(torch.__version__) >= pv.Version("2.7")
+        and pv.Version(transformers.__version__) >= pv.Version("4.50")
+    ):
         if verbose:
-            print("[_register_cache_serialization] DynamicCache is unregistered first.")
+            print(
+                "[_register_cache_serialization] DynamicCache "
+                "is unregistered and registered first."
+            )
         _unregister(DynamicCache)
         torch.utils._pytree.register_pytree_node(
             DynamicCache,
@@ -122,6 +132,8 @@ def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
             serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
             flatten_with_keys_fn=flatten_with_keys_dynamic_cache,
         )
+        # To avoid doing it multiple times.
+        PATCH_OF_PATCHES.add(DynamicCache)
 
     unregistered_dynamic_cache = True
     if DynamicCache is not None and DynamicCache in torch.utils._pytree.SUPPORTED_NODES:
@@ -138,9 +150,10 @@ def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
             serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
             flatten_with_keys_fn=flatten_with_keys_dynamic_cache,
         )
-        torch.fx._pytree.register_pytree_flatten_spec(
-            DynamicCache, lambda x, _: [x.key_cache, x.value_cache]
-        )
+        if pv.Version(torch.__version__) < pv.Version("2.7"):
+            torch.fx._pytree.register_pytree_flatten_spec(
+                DynamicCache, lambda x, _: [x.key_cache, x.value_cache]
+            )
 
         # check
         from ..helpers.cache_helper import make_dynamic_cache