fix many issues

Author: xadupre

_unittests/ut_export/test_dynamic_shapes.py

@@ -3,7 +3,7 @@
 import transformers
 from onnx_diagnostic.ext_test_case import ExtTestCase, requires_transformers
 from onnx_diagnostic.helpers import string_type
-from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache, CacheKeyValue
 from onnx_diagnostic.export import ModelInputs, CoupleInputsDynamicShapes
 from onnx_diagnostic.torch_export_patches import torch_export_patches
 from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
@@ -408,6 +408,7 @@ def forward(self, xy, z):
     def test_guess_dynamic_shapes_cache(self):
         class Model(torch.nn.Module):
             def forward(self, cache, z):
+                cache = CacheKeyValue(cache)
                 return (
                     z
                     + cache.key_cache[0]
@@ -475,6 +476,7 @@ def forward(self, cache, z):
     def test_guess_dynamic_shapes_cache_str(self):
         class Model(torch.nn.Module):
             def forward(self, cache, z):
+                cache = CacheKeyValue(cache)
                 return (
                     z
                     + cache.key_cache[0]
@@ -812,8 +814,9 @@ def test_couple_input_ds_change_dynamic_dimensions_dynamic_cache(self):
         with torch_export_patches(patch_transformers=True):
             new_inputs = inst.change_dynamic_dimensions()
         self.assertIsInstance(new_inputs["A"], transformers.cache_utils.DynamicCache)
-        self.assertEqual((3, 2, 3, 2), new_inputs["A"].key_cache[0].shape)
-        self.assertEqual((3, 2, 3, 2), new_inputs["A"].value_cache[0].shape)
+        new_inputs_A = CacheKeyValue(new_inputs["A"])
+        self.assertEqual((3, 2, 3, 2), new_inputs_A.key_cache[0].shape)
+        self.assertEqual((3, 2, 3, 2), new_inputs_A.value_cache[0].shape)
 
     @requires_transformers("4.51")
     def test_dynamic_cache_replace_by_string(self):

_unittests/ut_export/test_serialization.py

@@ -5,6 +5,7 @@
 from onnx_diagnostic.helpers.cache_helper import (
     make_dynamic_cache,
     flatten_unflatten_for_dynamic_shapes,
+    CacheKeyValue,
 )
 from onnx_diagnostic.export import ModelInputs
 from onnx_diagnostic.torch_export_patches import torch_export_patches
@@ -23,13 +24,15 @@ def _get_cache(self, n_layers=2, bsize=2, nheads=4, slen=1, dim=7):
     def test_dynamic_cache(self):
         class Model(torch.nn.Module):
             def forward(self, cache):
+                cache = CacheKeyValue(cache)
                 return cache.key_cache[0]
 
         cache = self._get_cache()
         DYN = torch.export.Dim.DYNAMIC
         ds = {0: DYN, 1: DYN, 3: DYN}
         dynamic_shapes = ([[ds, ds], [ds, ds]],)
-        exp = torch.export.export(Model(), (cache,), dynamic_shapes=dynamic_shapes)
+        with torch_export_patches(patch_transformers=True):
+            exp = torch.export.export(Model(), (cache,), dynamic_shapes=dynamic_shapes)
         self.assertNotEmpty(exp)
 
     @requires_transformers("4.50")

_unittests/ut_helpers/test_mini_onnx_builder.py

@@ -8,7 +8,7 @@
     create_input_tensors_from_onnx_model,
     MiniOnnxBuilder,
 )
-from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache, CacheKeyValue
 from onnx_diagnostic.helpers import string_type
 
 
@@ -127,8 +127,9 @@ def test_mini_onnx_builder(self):
 
     def test_mini_onnx_builder_transformers(self):
         cache = make_dynamic_cache([(torch.ones((3, 3)), torch.ones((3, 3)) * 2)])
-        self.assertEqual(len(cache.key_cache), 1)
-        self.assertEqual(len(cache.value_cache), 1)
+        dc = CacheKeyValue(cache)
+        self.assertEqual(len(dc.key_cache), 1)
+        self.assertEqual(len(dc.value_cache), 1)
 
         data = [(cache,), cache]
 
@@ -140,8 +141,9 @@ def test_mini_onnx_builder_transformers(self):
 
     def test_mini_onnx_builder_transformers_sep(self):
         cache = make_dynamic_cache([(torch.ones((3, 3)), torch.ones((3, 3)) * 2)])
-        self.assertEqual(len(cache.key_cache), 1)
-        self.assertEqual(len(cache.value_cache), 1)
+        dc = CacheKeyValue(cache)
+        self.assertEqual(len(dc.key_cache), 1)
+        self.assertEqual(len(dc.value_cache), 1)
 
         data = [(cache,), cache]
 

onnx_diagnostic/helpers/cache_helper.py

@@ -41,9 +41,14 @@ def __init__(self, cache=None):
                     f"or value_cache={string_type(self.value_cache)}, "
                     f"cache.layers={string_type(cache.layers)}"
                 )
-        elif cache is not None:
+        elif cache is not None and hasattr(cache, "key_cache"):
             self.key_cache = cache.key_cache
             self.value_cache = cache.value_cache
+        elif cache is None:
+            self.key_cache = None
+            self.value_cache = None
+        else:
+            raise NotImplementedError(f"type(cache)={type(cache)}")
 
     def make_dynamic_cache(self):
         """Do the reverse operation."""
@@ -401,6 +406,13 @@ def __init__(self):
         dtype=key_value_pairs[0][0].dtype,
     )
     ca = CacheKeyValue(cache)
+    if hasattr(cache, "layers") and len(ca.key_cache) == 0:
+        # transformers>= 4.55.2, layers are empty
+        cache_position = torch.arange(key_value_pairs[0][0].shape[2], dtype=torch.int64)
+        for i, (key, value) in enumerate(key_value_pairs):
+            cache.update(key, value, i, cache_kwargs={"cache_position": cache_position})
+        return cache
+
     for i in range(len(key_value_pairs)):
         assert ca.key_cache[i].shape == key_value_pairs[i][0].shape, (
             f"Shape mismatch, expected {cache.key_cache[i].shape}, "

onnx_diagnostic/torch_export_patches/serialization/transformers_impl.py

@@ -261,7 +261,9 @@ def unflatten_encoder_decoder_cache(
 ) -> EncoderDecoderCache:
     """Restores a :class:`transformers.cache_utils.EncoderDecoderCache` from python objects."""
     dictionary = torch.utils._pytree._dict_unflatten(values, context)
-    return EncoderDecoderCache(**dictionary)
+    return EncoderDecoderCache(
+        dictionary["self_attention_cache"], dictionary["cross_attention_cache"]
+    )
 
 
 #############