fix

Author: xadupre

_unittests/ut_torch_export_patches/test_patch_serialization.py

@@ -5,6 +5,7 @@
 from onnx_diagnostic.torch_export_patches.onnx_export_errors import (
     bypass_export_some_errors,
 )
+from transformers.modeling_outputs import BaseModelOutput
 
 
 class TestPatchSerialization(ExtTestCase):
@@ -83,6 +84,21 @@ def forward(self, cache):
         with bypass_export_some_errors():
             torch.export.export(model, (cache,), dynamic_shapes=(ds,))
 
+    @ignore_warnings(UserWarning)
+    def test_base_model_output(self):
+        bo = BaseModelOutput(last_hidden_state=torch.rand((4, 4, 4)))
+        with bypass_export_some_errors():
+            flat, _spec = torch.utils._pytree.tree_flatten(bo)
+            self.assertEqual(
+                "#1[T1s4x4x4]",
+                self.string_type(flat, with_shape=True),
+            )
+            bo2 = torch.utils._pytree.tree_unflatten(flat, _spec)
+            self.assertEqual(
+                self.string_type(bo, with_shape=True, with_min_max=True),
+                self.string_type(bo2, with_shape=True, with_min_max=True),
+            )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -1,35 +1,20 @@
 import pprint
 from typing import Any, Dict, List, Set, Tuple
+import packaging.version as pv
 import optree
 import torch
 import transformers
-import packaging.version as pv
+from transformers.cache_utils import DynamicCache, MambaCache, EncoderDecoderCache
+from transformers.modeling_outputs import BaseModelOutput
 
 
 PATCH_OF_PATCHES: Set[Any] = set()
 
 
 def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
-    # Cache serialization: to be moved into appropriate packages
-
-    try:
-        from transformers.cache_utils import DynamicCache
-    except ImportError:
-        DynamicCache = None
-
-    try:
-        from transformers.cache_utils import MambaCache
-    except ImportError:
-        MambaCache = None
-
-    try:
-        from transformers.cache_utils import EncoderDecoderCache
-    except ImportError:
-        EncoderDecoderCache = None
-
     # MambaCache
     unregistered_mamba_cache = True
-    if MambaCache is not None and MambaCache in torch.utils._pytree.SUPPORTED_NODES:
+    if MambaCache in torch.utils._pytree.SUPPORTED_NODES:
         if verbose > 1:
             print(f"[_register_cache_serialization] {MambaCache} already registered")
         # It is already registered because bypass_export_some_errors was called
@@ -82,6 +67,26 @@ def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
         # To avoid doing it multiple times.
         PATCH_OF_PATCHES.add(DynamicCache)
 
+    # BaseModelOutput serialization is incomplete.
+    # It does not include dynamic shapes mapping.
+    if BaseModelOutput in torch.fx._pytree.SUPPORTED_NODES and not PATCH_OF_PATCHES:
+        if verbose:
+            print(
+                "[_register_cache_serialization] BaseModelOutput "
+                "is unregistered and registered first."
+            )
+        _unregister(BaseModelOutput)
+        torch.utils._pytree.register_pytree_node(
+            BaseModelOutput,
+            flatten_base_model_output,
+            unflatten_base_model_output,
+            serialized_type_name=f"{BaseModelOutput.__module__}.{BaseModelOutput.__name__}",
+            flatten_with_keys_fn=flatten_with_keys_base_model_output,
+        )
+
+        # To avoid doing it multiple times.
+        PATCH_OF_PATCHES.add(BaseModelOutput)
+
     unregistered_dynamic_cache = True
     if DynamicCache is not None and DynamicCache in torch.utils._pytree.SUPPORTED_NODES:
         if verbose > 1:
@@ -123,7 +128,7 @@ def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
         # within a section already calling bypass_export_some_errors or transformers
         # has updated its code to do it.
         # No need to register and unregister then.
-        unregistered_mamba_cache = False
+        unregistered_encode_decode_cache = False
     else:
         if verbose:
             print("[_register_cache_serialization] register EncoderDecoderCache")
@@ -135,10 +140,32 @@ def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
             flatten_with_keys_fn=flatten_with_keys_encoder_decoder_cache,
         )
 
+    # BaseModelOutput
+    unregistered_base_model_output = True
+    if BaseModelOutput is not None and BaseModelOutput in torch.utils._pytree.SUPPORTED_NODES:
+        if verbose > 1:
+            print(f"[_register_cache_serialization] {BaseModelOutput} already registered")
+        # It is already registered because bypass_export_some_errors was called
+        # within a section already calling bypass_export_some_errors or transformers
+        # has updated its code to do it.
+        # No need to register and unregister then.
+        unregistered_base_model_output = False
+    else:
+        if verbose:
+            print("[_register_cache_serialization] register BaseModelOutput")
+        torch.utils._pytree.register_pytree_node(
+            BaseModelOutput,
+            flatten_encoder_decoder_cache,
+            unflatten_encoder_decoder_cache,
+            serialized_type_name=f"{BaseModelOutput.__module__}.{BaseModelOutput.__name__}",
+            flatten_with_keys_fn=flatten_with_keys_base_model_output,
+        )
+
     return dict(
         DynamicCache=unregistered_dynamic_cache,
         MambaCache=unregistered_mamba_cache,
         EncoderDecoderCache=unregistered_encode_decode_cache,
+        BaseModelOutput=unregistered_base_model_output,
     )
 
 
@@ -167,20 +194,11 @@ def _unregister(cls: type, verbose: int = 0):
 
 
 def _unregister_cache_serialization(undo: Dict[str, bool], verbose: int = 0):
-    if undo.get("MambaCache", False):
-        _unregister(transformers.cache_utils.MambaCache, verbose)
-    elif verbose > 1:
-        print("[_unregister_cache_serialization] skip unregister MambaCache")
-
-    if undo.get("DynamicCache", False):
-        _unregister(transformers.cache_utils.DynamicCache, verbose)
-    elif verbose > 1:
-        print("[_unregister_cache_serialization] skip unregister DynamicCache")
-
-    if undo.get("EncoderDecoderCache", False):
-        _unregister(transformers.cache_utils.EncoderDecoderCache, verbose)
-    elif verbose > 1:
-        print("[_unregister_cache_serialization] skip unregister EncoderDecoderCache")
+    for cls in [MambaCache, DynamicCache, EncoderDecoderCache, BaseModelOutput]:
+        if undo.get(cls.__name__, False):
+            _unregister(cls, verbose)
+        elif verbose > 1:
+            print(f"[_unregister_cache_serialization] skip unregister {cls.__name__}")
 
 
 ############
@@ -205,7 +223,7 @@ def _unregister_cache_serialization(undo: Dict[str, bool], verbose: int = 0):
 #     dtype=dtype,
 # )
 def flatten_mamba_cache(
-    mamba_cache: transformers.cache_utils.MambaCache,
+    mamba_cache: MambaCache,
 ) -> Tuple[List[Any], torch.utils._pytree.Context]:
     """Serializes a :class:`transformers.cache_utils.MambaCache` with python objects."""
     flat = [
@@ -224,10 +242,8 @@ def flatten_mamba_cache(
 
 
 def unflatten_mamba_cache(
-    values: List[Any],
-    context: torch.utils._pytree.Context,
-    output_type=None,
-) -> transformers.cache_utils.MambaCache:
+    values: List[Any], context: torch.utils._pytree.Context, output_type=None
+) -> MambaCache:
     """Restores a :class:`transformers.cache_utils.MambaCache` from python objects."""
     conv_states, ssm_states = values
 
@@ -258,12 +274,12 @@ def __init__(self):
     return cache
 
 
-def flatten_with_keys_mamba_cache(d: Dict[Any, Any]) -> Tuple[
+def flatten_with_keys_mamba_cache(cache: MambaCache) -> Tuple[
     List[Tuple[torch.utils._pytree.KeyEntry, Any]],
     torch.utils._pytree.Context,
 ]:
     """Serializes a :class:`transformers.cache_utils.MambaCache` with python objects."""
-    values, context = flatten_mamba_cache(d)
+    values, context = flatten_mamba_cache(cache)
     return [(torch.utils._pytree.MappingKey(k), v) for k, v in zip(context, values)], context
 
 
@@ -273,7 +289,7 @@ def flatten_with_keys_mamba_cache(d: Dict[Any, Any]) -> Tuple[
 
 
 def flatten_dynamic_cache(
-    dynamic_cache: transformers.cache_utils.DynamicCache,
+    dynamic_cache: DynamicCache,
 ) -> Tuple[List[Any], torch.utils._pytree.Context]:
     """Serializes a :class:`transformers.cache_utils.DynamicCache` with python objects."""
     if hasattr(transformers.cache_utils, "_flatten_dynamic_cache"):
@@ -287,11 +303,8 @@ def flatten_dynamic_cache(
 
 
 def flatten_with_keys_dynamic_cache(
-    dynamic_cache: transformers.cache_utils.DynamicCache,
-) -> Tuple[
-    List[Tuple[torch.utils._pytree.KeyEntry, Any]],
-    torch.utils._pytree.Context,
-]:
+    dynamic_cache: DynamicCache,
+) -> Tuple[List[Tuple[torch.utils._pytree.KeyEntry, Any]], torch.utils._pytree.Context]:
     """Serializes a :class:`transformers.cache_utils.DynamicCache` with python objects."""
     if hasattr(transformers.cache_utils, "_flatten_with_keys_dynamic_cache"):
         return transformers.cache_utils._flatten_with_keys_dynamic_cache(dynamic_cache)
@@ -300,10 +313,8 @@ def flatten_with_keys_dynamic_cache(
 
 
 def unflatten_dynamic_cache(
-    values: List[Any],
-    context: torch.utils._pytree.Context,
-    output_type=None,
-) -> transformers.cache_utils.DynamicCache:
+    values: List[Any], context: torch.utils._pytree.Context, output_type=None
+) -> DynamicCache:
     """Restores a :class:`transformers.cache_utils.DynamicCache` from python objects."""
     if hasattr(transformers.cache_utils, "_unflatten_dynamic_cache"):
         assert output_type is None, f"output_type={output_type} not supported"
@@ -322,7 +333,7 @@ def unflatten_dynamic_cache(
 
 
 def flatten_encoder_decoder_cache(
-    ec_cache: transformers.cache_utils.DynamicCache,
+    ec_cache: EncoderDecoderCache,
 ) -> Tuple[List[Any], torch.utils._pytree.Context]:
     """
     Serializes a :class:`transformers.cache_utils.EncoderDecoderCache`
@@ -335,9 +346,7 @@ def flatten_encoder_decoder_cache(
     return torch.utils._pytree._dict_flatten(dictionary)
 
 
-def flatten_with_keys_encoder_decoder_cache(
-    ec_cache: transformers.cache_utils.DynamicCache,
-) -> Tuple[
+def flatten_with_keys_encoder_decoder_cache(ec_cache: EncoderDecoderCache) -> Tuple[
     List[Tuple[torch.utils._pytree.KeyEntry, Any]],
     torch.utils._pytree.Context,
 ]:
@@ -353,10 +362,46 @@ def flatten_with_keys_encoder_decoder_cache(
 
 
 def unflatten_encoder_decoder_cache(
-    values: List[Any],
-    context: torch.utils._pytree.Context,
-    output_type=None,
-) -> transformers.cache_utils.EncoderDecoderCache:
+    values: List[Any], context: torch.utils._pytree.Context, output_type=None
+) -> EncoderDecoderCache:
     """Restores a :class:`transformers.cache_utils.EncoderDecoderCache` from python objects."""
     dictionary = torch.utils._pytree._dict_unflatten(values, context)
     return transformers.cache_utils.EncoderDecoderCache(**dictionary)
+
+
+#################
+# BaseModelOutput
+#################
+
+
+def flatten_base_model_output(
+    bo: BaseModelOutput,
+) -> Tuple[List[Any], torch.utils._pytree.Context]:
+    """
+    Serializes a :class:`transformers.modeling_outputs.BaseModelOutput`
+    with python objects.
+    """
+    return list(bo.values()), list(bo.keys())
+
+
+def flatten_with_keys_base_model_output(
+    bo: BaseModelOutput,
+) -> Tuple[List[Tuple[torch.utils._pytree.KeyEntry, Any]], torch.utils._pytree.Context]:
+    """
+    Serializes a :class:`transformers.modeling_outputs.BaseModelOutput`
+    with python objects.
+    """
+    values, context = flatten_dynamic_cache(bo)
+    return [(torch.utils._pytree.MappingKey(k), v) for k, v in zip(context, values)], context
+
+
+def unflatten_base_model_output(
+    values: List[Any],
+    context: torch.utils._pytree.Context,
+    output_type=None,
+) -> BaseModelOutput:
+    """
+    Restores a :class:`transformers.modeling_outputs.BaseModelOutput`
+    from python objects.
+    """
+    return BaseModelOutput(**dict(zip(context, values)))