fix modelbuilder inputs

Author: xadupre

_doc/api/helpers/index.rst

@@ -14,6 +14,7 @@ onnx_diagnostic.helpers
     helper
     memory_peak
     mini_onnx_builder
+    model_builder_helper
     onnx_helper
     ort_session
     rt_helper

_doc/api/helpers/model_builder_helper.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.helpers.model_builder_helper
+============================================
+
+.. automodule:: onnx_diagnostic.helpers.model_builder_helper
+    :members:
+    :no-undoc-members:

_doc/conf.py

@@ -211,6 +211,7 @@ def linkcode_resolve(domain, info):
     "huggingface_hub": "https://github.com/huggingface/huggingface_hub",
     "Linux": "https://www.linux.org/",
     "ml_dtypes": "https://github.com/jax-ml/ml_dtypes",
+    "ModelBuilder": "https://onnxruntime.ai/docs/genai/howto/build-model.html",
     "monai": "https://monai.io/",
     "numpy": "https://numpy.org/",
     "onnx": "https://onnx.ai/onnx/",

_unittests/ut_export/test_dynamic_shapes.py

@@ -6,6 +6,7 @@
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
 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
 
 
 class TestDynamicShapes(ExtTestCase):
@@ -512,7 +513,6 @@ def forward(self, cache, z):
         mi = ModelInputs(Model(), inputs)
         self.assertIn("DynamicCache", string_type(mi.inputs, with_shape=True))
         ds = mi.guess_dynamic_shapes(auto="dim")
-        print(ds)
         self.assertEqual(
             ds,
             (
@@ -845,6 +845,22 @@ def test_dynamic_cache_replace_by_string(self):
             as_string,
         )
 
+    def test_unbatch_inputs(self):
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM")
+        cpl = CoupleInputsDynamicShapes(
+            None, data["inputs"], dynamic_shapes=data["dynamic_shapes"]
+        )
+        new_dims = cpl.change_dynamic_dimensions(
+            desired_values=dict(batch=1), only_desired=True
+        )
+        s = self.string_type(new_dims, with_shape=True)
+        self.assertEqual(
+            "dict(input_ids:T7s1x3,attention_mask:T7s1x33,position_ids:T7s1x3,"
+            "past_key_values:DynamicCache("
+            "key_cache=#1[T1s1x1x30x96], value_cache=#1[T1s1x1x30x96]))",
+            s,
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_models/test_test_helpers.py

@@ -184,6 +184,24 @@ def test_filter_inputs(self):
         ni, nd = filter_inputs(inputs, dynamic_shapes=ds, drop_names=["a"], model=["a", "b"])
         self.assertEqual((ni, nd), (((None,), {"b": 4}), {"b": 30}))
 
+    @requires_torch("2.7")
+    @hide_stdout()
+    @ignore_warnings(FutureWarning)
+    def test_validate_model_modelbuilder(self):
+        mid = "arnir0/Tiny-LLM"
+        summary, data = validate_model(
+            mid,
+            do_run=True,
+            verbose=10,
+            exporter="modelbuilder",
+            dump_folder="dump_test_validate_model_onnx_dynamo",
+        )
+        self.assertIsInstance(summary, dict)
+        self.assertIsInstance(data, dict)
+        self.assertLess(summary["disc_onnx_ort_run_abs"], 1e-4)
+        onnx_filename = data["onnx_filename"]
+        self.assertExists(onnx_filename)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/export/dynamic_shapes.py

@@ -379,8 +379,9 @@ def _generic_walker_step(
         return torch.utils._pytree.tree_unflatten(res, spec)
 
     class ChangeDimensionProcessor:
-        def __init__(self, desired_values):
+        def __init__(self, desired_values, only_desired):
             self.mapping = desired_values or {}
+            self.only_desired = only_desired
 
         def _build_new_shape(
             self, shape: Tuple[int, ...], ds: Dict[int, Any]
@@ -397,14 +398,16 @@ def _build_new_shape(
                             torch.export.dynamic_shapes._Dim,
                         ),
                     ):
-                        d = str(ds[i])
+                        d = ds[i].__name__
                     elif not isinstance(ds[i], int):
                         raise NotImplementedError(f"Unable to handle type {ds[i]} in {ds}")
                     if d in self.mapping:
                         new_dim = self.mapping[d]
-                    else:
+                    elif not self.only_desired:
                         new_dim = shape[i] + 1
                         self.mapping[d] = new_dim
+                    else:
+                        new_dim = shape[i]
                     new_shape[i] = new_dim
             return tuple(new_shape)
 
@@ -447,7 +450,10 @@ def __call__(self, inputs, ds):
             return self._build_new_tensor(inputs, new_shape)
 
     def change_dynamic_dimensions(
-        self, desired_values: Optional[Dict[str, int]] = None, args_kwargs: bool = False
+        self,
+        desired_values: Optional[Dict[str, int]] = None,
+        args_kwargs: bool = False,
+        only_desired: bool = False,
     ):
         """
         A model exported with dynamic shapes is not necessarily dynamic
@@ -460,6 +466,8 @@ def change_dynamic_dimensions(
 
         :param desired_values: to fixed named dimension to have the desired value
         :param args_kwargs: return both args, kwargs even if empty
+        :param only_desired: if True, only change the dimension specified in
+            ``desired_values``
         :return: new inputs
 
         Example:
@@ -483,7 +491,8 @@ def change_dynamic_dimensions(
             print("-after:", string_type(new_kwargs, with_shape=True))
         """
         return self._generic_walker(
-            self.ChangeDimensionProcessor(desired_values), args_kwargs=args_kwargs
+            self.ChangeDimensionProcessor(desired_values, only_desired=only_desired),
+            args_kwargs=args_kwargs,
         )
 
 

onnx_diagnostic/torch_models/test_helper.py

@@ -345,6 +345,25 @@ def validate_model(
             )
         ),
     )
+
+    if exporter == "modelbuilder":
+        # Models used with ModelBuilder do not like batch size > 1.
+        # Let's change that.
+        for k in ["inputs", "inputs2"]:
+            if k not in data:
+                continue
+            if verbose:
+                print(f"[validate_model] set batch=1 for data[{k!r}]")
+                print(f"[validate_model] batch=1 === {string_type(data[k], with_shape=True)}")
+            cpl = CoupleInputsDynamicShapes(
+                None, data[k], dynamic_shapes=data["dynamic_shapes"]
+            )
+            data[k] = cpl.change_dynamic_dimensions(
+                desired_values=dict(batch=1), only_desired=True
+            )
+            if verbose:
+                print(f"[validate_model] batch=1 --> {string_type(data[k], with_shape=True)}")
+
     data["input_options"] = iop
     data["model_options"] = mop
     data["model_dump_folder"] = dump_folder