fix remaining issues

Author: xadupre

_doc/examples/plot_export_with_dynamic_cache.py

@@ -23,18 +23,19 @@
 import pprint
 import torch
 from onnx_diagnostic import doc
-from onnx_diagnostic.ext_test_case import has_transformers
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.helpers.cache_helper import (
     flatten_unflatten_for_dynamic_shapes,
     make_dynamic_cache,
+    CacheKeyValue,
 )
 from onnx_diagnostic.export import ModelInputs
 from onnx_diagnostic.torch_export_patches import torch_export_patches
 
 
 class Model(torch.nn.Module):
     def forward(self, cache, z):
+        cache = CacheKeyValue(cache)
         return (
             z
             + cache.key_cache[0]
@@ -105,13 +106,8 @@ def forward(self, cache, z):
 # registers functions to serialize ``DynamicCache``. This one is modified to make
 # the shape inference implemented in :epkg:`torch` happy.
 
-if has_transformers("4.50"):
+with torch_export_patches(patch_transformers=True):
     ep = torch.export.export(model, inputs[0], dynamic_shapes=ds[0], strict=False)
-else:
-    with torch_export_patches(patch_transformers=True) as modificator:
-        ep = torch.export.export(
-            model, modificator(inputs[0]), dynamic_shapes=ds[0], strict=False
-        )
 print(ep)
 
 # %%

_unittests/ut_torch_models/test_validate_whole_models.py

@@ -258,7 +258,7 @@ def test_validate_model_vit_model(self):
     @requires_torch("2.7")
     @hide_stdout()
     @ignore_warnings(FutureWarning)
-    @requires_transformers("4.53")
+    @requires_transformers("4.55")
     def test_validate_phi35_mini_instruct(self):
         mid = "microsoft/Phi-3.5-mini-instruct"
         summary, data = validate_model(
@@ -282,7 +282,7 @@ def test_validate_phi35_mini_instruct(self):
     @requires_torch("2.7")
     @hide_stdout()
     @ignore_warnings(FutureWarning)
-    @requires_transformers("4.53")
+    @requires_transformers("4.55")
     def test_validate_phi35_4k_mini_instruct(self):
         mid = "microsoft/Phi-3-mini-4k-instruct"
         summary, data = validate_model(

onnx_diagnostic/export/dynamic_shapes.py

@@ -887,27 +887,41 @@ def guess_dynamic_shape_object(
 
         # In case DynamicCache is not registered.
         if obj.__class__.__name__ == "DynamicCache":
-            kc = set(len(o.key_cache) for o in objs)
-            assert (
-                len(kc) == 1
-            ), f"All attribute 'key_cache' should have the same length but found {kc}"
-            vc = set(len(o.value_cache) for o in objs)
-            assert (
-                len(vc) == 1
-            ), f"All attribute 'value_cache' should have the same length but found {vc}"
+            if hasattr(obj, "layers"):
+                kc = set(len(o.layers) for o in objs)
+                assert (
+                    len(kc) == 1
+                ), f"All attribute 'key_cache' should have the same length but found {kc}"
+                vc = kc.copy()
+            else:
+                kc = set(len(o.key_cache) for o in objs)
+                assert (
+                    len(kc) == 1
+                ), f"All attribute 'key_cache' should have the same length but found {kc}"
+                vc = set(len(o.value_cache) for o in objs)
+                assert (
+                    len(vc) == 1
+                ), f"All attribute 'value_cache' should have the same length but found {vc}"
+
             key_cache = []
             for i in range(kc.pop()):
                 key_cache.append(
                     self.guess_dynamic_dimensions(
-                        *[o.key_cache[i] for o in objs],
+                        *[
+                            o.layers[i].keys if hasattr(o, "layers") else o.key_cache[i]
+                            for o in objs
+                        ],
                         auto=auto if isinstance(auto, bool) else f"{auto}_{i}kdc",
                     )
                 )
             value_cache = []
             for i in range(vc.pop()):
                 value_cache.append(
                     self.guess_dynamic_dimensions(
-                        *[o.value_cache[i] for o in objs],
+                        *[
+                            o.layers[i].values if hasattr(o, "layers") else o.value_cache[i]
+                            for o in objs
+                        ],
                         auto=auto if isinstance(auto, bool) else f"{auto}_{i}vdc",
                     )
                 )

onnx_diagnostic/export/shape_helper.py

@@ -9,6 +9,8 @@ def all_dynamic_shape_from_inputs(inputs: Any, dim_prefix: Any = "d") -> Any:
     All dimensions are considered as dynamic.
     ``dim_prefix`` can be a string (the function uses it as a prefix),
     or ``torch.export.Dim.AUTO`` or ``torch.export.Dim.DYNAMIC``.
+    Depending on the version of transformers, serializations function
+    of DynamicCache class is automatically serialized or not (>= 4.51, < 4.55).
 
     .. runpython::
         :showcode:
@@ -17,6 +19,7 @@ def all_dynamic_shape_from_inputs(inputs: Any, dim_prefix: Any = "d") -> Any:
         import torch
         from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
         from onnx_diagnostic.export.shape_helper import all_dynamic_shape_from_inputs
+        from onnx_diagnostic.torch_export_patches import torch_export_patches
 
         bsize, nheads, slen, dim = 2, 1, 30, 96
         inputs = dict(
@@ -25,10 +28,11 @@ def all_dynamic_shape_from_inputs(inputs: Any, dim_prefix: Any = "d") -> Any:
             position_ids=torch.arange(3, dtype=torch.int64),
             past_key_values=make_dynamic_cache(
                 [(torch.randn(bsize, nheads, slen, dim),
-                  torch.randn(bsize, nheads, slen, dim))]
+                    torch.randn(bsize, nheads, slen, dim))]
             ),
         )
-        ds = all_dynamic_shape_from_inputs(inputs)
+        with torch_export_patches(patch_transformers=True):
+            ds = all_dynamic_shape_from_inputs(inputs)
         pprint.pprint(ds)
 
     For this function to work, patches must be enabled if :epkg:`transformers`

onnx_diagnostic/helpers/cache_helper.py

@@ -131,6 +131,8 @@ def is_cache_dynamic_registered(fast: bool = False) -> bool:
     )
     values, spec = torch.utils._pytree.tree_flatten(cache)
     cache2 = torch.utils._pytree.tree_unflatten(values, spec)
+    if hasattr(cache2, "layers") and hasattr(cache, "layers"):
+        return len(cache2.layers) == len(cache.layers)
     return len(cache2.key_cache) == len(cache.value_cache)
 
 

onnx_diagnostic/tasks/feature_extraction.py

@@ -160,5 +160,4 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
         if hasattr(config, att):
             kwargs[att] = getattr(config, att)
     kwargs["decoder_ffn_dim"] = kwargs["encoder_ffn_dim"] = 64
-    print(kwargs)
     return kwargs, get_inputs