fix modelbuilder

Author: xadupre

_unittests/ut_export/test_dynamic_shapes.py

@@ -862,6 +862,60 @@ def test_unbatch_inputs(self):
             s,
         )
 
+    def test_guess_dynamic_cache_without_patches(self):
+        n_layers = 2
+        bsize, nheads, slen, dim = 2, 4, 3, 7
+        cache = make_dynamic_cache(
+            [
+                (torch.randn(bsize, nheads, slen, dim), torch.randn(bsize, nheads, slen, dim))
+                for i in range(n_layers)
+            ]
+        )
+        z = torch.randn((1, 1, 1, 7))
+        cache2 = make_dynamic_cache(
+            [
+                (
+                    torch.randn(bsize + 1, nheads, slen + 1, dim + 1),
+                    torch.randn(bsize + 1, nheads, slen + 1, dim + 1),
+                )
+                for i in range(n_layers)
+            ]
+        )
+        inputs = [
+            (cache, z),
+            (cache2, torch.randn((1, 1, 1, 8))),
+        ]
+
+        class Model(torch.nn.Module):
+            def forward(self, cache, z):
+                cache = CacheKeyValue(cache)
+                return (
+                    z
+                    + cache.key_cache[0]
+                    + cache.key_cache[1]
+                    + cache.value_cache[0]
+                    + cache.value_cache[1]
+                )
+
+        mi = ModelInputs(Model(), inputs)
+        ds = mi.guess_dynamic_shapes()
+        DYN = torch.export.Dim.DYNAMIC
+        self.assertEqual(
+            (
+                (
+                    [
+                        {0: DYN, 2: DYN, 3: DYN},
+                        {0: DYN, 2: DYN, 3: DYN},
+                        {0: DYN, 2: DYN, 3: DYN},
+                        {0: DYN, 2: DYN, 3: DYN},
+                    ],
+                    {3: DYN},
+                ),
+                {},
+            ),
+            ds,
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_export_patches/test_patch_inputs.py

@@ -88,8 +88,8 @@ def test_convert_dynamic_axes_into_dynamic_shapes_2(self):
         )
         self.assertEqual(
             [
-                [{0: "batch_size", 2: "past_sequence_length"}],
-                [{0: "batch_size", 2: "past_sequence_length"}],
+                {0: "batch_size", 2: "past_sequence_length"},
+                {0: "batch_size", 2: "past_sequence_length"},
             ],
             res[2]["past_key_values"],
         )

onnx_diagnostic/export/dynamic_shapes.py

@@ -1,4 +1,5 @@
 import inspect
+import itertools
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import numpy as np
 import torch
@@ -934,7 +935,7 @@ def guess_dynamic_shape_object(
                         auto=auto if isinstance(auto, bool) else f"{auto}_{i}vdc",
                     )
                 )
-            return [key_cache, value_cache]
+            return list(itertools.chain.from_iterable(zip(key_cache, value_cache)))
 
         raise NotImplementedError(
             f"Unable to build dynamic shapes for type {set_types.pop()}: "

onnx_diagnostic/helpers/rt_helper.py

@@ -96,54 +96,9 @@ def make_feeds(
         elif isinstance(i, float):
             i = np.array(i, dtype=np.float32)
         new_flat.append(i)
-
-    # NOTE: model builder has a different order for past_key_values
-    #       we need to reorder them to match the expected order
-    if is_modelbuilder:
-        # We assume that if "past_key_values" is in the names when it's
-        # modelbuilder
-        non_past_kv_input_names = [n for n in names if "past_key_values" not in n]
-        past_kv_names = [n for n in names if "past_key_values" in n]
-        reorder_past_kv_names = reorder_modelbuilder_cache_to_torch(past_kv_names)
-        names = non_past_kv_input_names + reorder_past_kv_names
     return dict(zip(names, new_flat))
 
 
-def reorder_modelbuilder_cache_to_torch(past_kv: List[Any]) -> List[Any]:
-    """
-    Reorders the past_kvs for ModelBuilder to match the expected order
-    by PyTorch exported models.
-
-    .. note::
-        This function can take either the names or the actual tensors
-        as long as they are in a list.
-
-    Conceptually,
-
-    From::
-
-        [past_key_values.0.key, past_key_values.0.value,
-        past_key_values.1.key, past_key_values.1.value, ...]
-
-    To::
-
-        [past_key_values.0.key, past_key_values.1.key,
-        ..., past_key_values.0.value, past_key_values.1.value, ...]
-
-    :param past_kv: list of flattened inputs
-    :return: reordered list of flattened inputs
-    """
-    total_len = len(past_kv)
-    if total_len % 2 != 0:
-        raise ValueError("The length of past_key_values should be even.")
-    keys = []
-    values = []
-    for i in range(0, total_len, 2):
-        keys.append(past_kv[i])
-        values.append(past_kv[i + 1])
-    return keys + values
-
-
 def _get_dim(i: int, s: Union[str, int], batch: int = 1) -> int:
     if isinstance(s, int):
         return s

onnx_diagnostic/tasks/image_text_to_text.py

@@ -1,3 +1,4 @@
+import itertools
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
 from ..helpers.cache_helper import make_dynamic_cache, make_hybrid_cache
@@ -269,10 +270,14 @@ def get_inputs_default(
         "token_type_ids": {0: batch, 1: seq_length},
         "attention_mask": {0: batch, 1: "cache+seq"},
         "position_ids": {0: batch, 1: seq_length},
-        "past_key_values": [
-            [{0: batch} for _ in range(num_hidden_layers)],
-            [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
-        ],
+        "past_key_values": list(
+            itertools.chain.from_iterable(
+                zip(
+                    [{0: batch} for _ in range(num_hidden_layers)],
+                    [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
+                )
+            )
+        ),
         "pixel_values": (
             {0: batch, 1: images}
             if model.__class__.__name__ == "IdeficsForVisionText2Text"

onnx_diagnostic/torch_export_patches/patch_inputs.py

@@ -38,7 +38,7 @@ def _make_shape(subset: Dict, cls: type, value: Any) -> Any:
         for v in subset.values():
             axes = v
             break
-        new_shape = [[axes for i in range(cache_length)], [axes for i in range(cache_length)]]
+        new_shape = [axes for i in range(cache_length * 2)]
         return new_shape
     if value.__class__ in torch.utils._pytree.SUPPORTED_NODES:
         raise NotImplementedError(

onnx_diagnostic/torch_models/validate.py

@@ -12,7 +12,7 @@
 from ..export import CoupleInputsDynamicShapes
 from ..helpers import max_diff, string_type, string_diff
 from ..helpers.helper import flatten_object
-from ..helpers.rt_helper import make_feeds, reorder_modelbuilder_cache_to_torch
+from ..helpers.rt_helper import make_feeds
 from ..helpers.torch_helper import to_any, torch_deepcopy
 from ..helpers.cache_helper import flatten_unflatten_for_dynamic_shapes
 from ..tasks import random_input_kwargs
@@ -1478,7 +1478,7 @@ def _mk(key, flavour=flavour):
             data[k_input],
             use_numpy=True,
             check_flatten=False,
-            is_modelbuilder=data["exporter"] == "modelbuilder",
+            is_modelbuilder=data["exporter"] == "modelbuilder",  # to remove position_ids
         )
         if verbose:
             print(f"[validate_onnx_model] ort inputs={string_type(feeds, with_shape=True)}")
@@ -1501,13 +1501,6 @@ def _mk(key, flavour=flavour):
             repeat=repeat,
             warmup=warmup,
         )
-        # NOTE: modelbuilder has different order on past_kv outputs
-        if data["exporter"] == "modelbuilder":
-            logits = got[:1]
-            past_key_values = got[1:]
-            reorder_past_key_values = reorder_modelbuilder_cache_to_torch(past_key_values)
-            got = logits + reorder_past_key_values
-
         if f"ERR_{_mk(f'time_onnx_ort_run{suffix}')}" in summary:
             return summary, data