more tests

Author: xadupre

_doc/examples/plot_export_hub_codellama.py

@@ -22,9 +22,7 @@
 from onnx_diagnostic import doc
 from onnx_diagnostic.ext_test_case import unit_test_going
 from onnx_diagnostic.helpers import string_type
-from onnx_diagnostic.torch_models.hghub import (
-    get_untrained_model_with_inputs,
-)
+from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
 from onnx_diagnostic.torch_models.hghub.hub_api import (
     get_model_info,
     get_pretrained_config,

_doc/examples/plot_export_tiny_phi2.py

@@ -33,9 +33,7 @@
 from onnx_diagnostic.helpers.rt_helper import make_feeds
 from onnx_diagnostic.torch_export_patches import torch_export_patches
 from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
-from onnx_diagnostic.torch_models.hghub import (
-    get_untrained_model_with_inputs,
-)
+from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
 
 warnings.simplefilter("ignore")
 

_unittests/ut_export/test_api.py

@@ -1,6 +1,11 @@
 import unittest
 import torch
 from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout
+from onnx_diagnostic.helpers import max_diff
+from onnx_diagnostic.helpers.torch_helper import torch_deepcopy
+from onnx_diagnostic.helpers.rt_helper import make_feeds
+from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
+from onnx_diagnostic.torch_export_patches import torch_export_patches
 from onnx_diagnostic.export.api import to_onnx
 
 
@@ -19,16 +24,66 @@ def forward(self, x, y):
             (x, y),
             dynamic_shapes=ds,
             exporter="custom",
-            filename=self.get_dump_file("custom.onnx"),
+            filename=self.get_dump_file("to_onnx_custom.onnx"),
         )
         to_onnx(
             Model(),
             (x, y),
             dynamic_shapes=ds,
             exporter="onnx-dynamo",
-            filename=self.get_dump_file("onnx-dynamo.onnx"),
+            filename=self.get_dump_file("to_onnx_onnx-dynamo.onnx"),
         )
 
+    @hide_stdout()
+    def test_tiny_llm_to_onnx(self):
+        import onnxruntime
+
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM")
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        b1 = data["inputs_batch1"]
+        filenames = {
+            "custom": self.get_dump_file("test_tiny_llm_to_onnx-custom.onnx"),
+            "onnx-dynamo": self.get_dump_file("test_tiny_llm_to_onnx-dynamo.onnx"),
+            "modelbuilder": self.get_dump_file("model.onnx"),
+        }
+        del inputs["position_ids"]
+        del ds["position_ids"]
+        del b1["position_ids"]
+
+        expected = model(**torch_deepcopy(b1))
+
+        with torch_export_patches(patch_transformers=True):
+            for exporter, filename in filenames.items():
+                with self.subTest(exporter=exporter):
+                    to_onnx(
+                        model,
+                        kwargs=inputs,
+                        dynamic_shapes=ds,
+                        exporter=exporter,
+                        filename=filename,
+                    )
+        for exporter, filename in filenames.items():
+            with self.subTest(exporter=f"validate-{exporter}"):
+                sess = onnxruntime.InferenceSession(
+                    filename, providers=["CPUExecutionProvider"]
+                )
+                feeds = make_feeds(sess, b1, use_numpy=True)
+                got = sess.run(None, feeds)
+                diff = max_diff(expected, got)
+                assert diff["abs"] <= 1e-5, f"diff={diff}"
+
+        b1["attention_mask"][:, :] = 1
+        expected = model(**torch_deepcopy(b1))
+        for exporter, filename in filenames.items():
+            with self.subTest(exporter=f"full-mask-{exporter}"):
+                sess = onnxruntime.InferenceSession(
+                    filename, providers=["CPUExecutionProvider"]
+                )
+                feeds = make_feeds(sess, b1, use_numpy=True)
+                got = sess.run(None, feeds)
+                diff = max_diff(expected, got)
+                assert diff["abs"] <= 1e-5, f"diff={diff}"
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_export/test_dynamic_shapes.py

@@ -916,6 +916,27 @@ def forward(self, cache, z):
             ds,
         )
 
+    def test_invalid_dimensions_for_export(self):
+        ags = []
+        kws = dict(
+            input_ids=torch.randint(0, 10, (2, 3)),
+            attention_mask=torch.randint(0, 1, (2, 33)),
+            position_ids=torch.randint(0, 10, (2, 3)),
+            past_key_values=make_dynamic_cache(
+                [torch.rand((2, 1, 30, 96)), torch.rand((2, 1, 30, 96))]
+            ),
+        )
+        ds = dict(
+            input_ids={0: "batch", 1: "seq_length"},
+            attention_mask={0: "batch", 1: "seq_length"},
+            position_ids={0: "batch", 1: "seq_length"},
+            past_key_values=[{0: "batch", 2: "cache_length"}, {0: "batch", 2: "cache_length"}],
+        )
+        with torch_export_patches(patch_transformers=True):
+            cpl = CoupleInputsDynamicShapes(ags, kws, ds)
+            backed_size_oblivious = cpl.invalid_dimensions_for_export()
+            self.assertFalse(backed_size_oblivious)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_helpers/test_model_builder_helper.py

@@ -12,9 +12,7 @@
     create_model_builder,
     save_model_builder,
 )
-from onnx_diagnostic.torch_models.hghub import (
-    get_untrained_model_with_inputs,
-)
+from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
 from onnx_diagnostic.helpers.rt_helper import make_feeds
 
 

_unittests/ut_helpers/test_rt_helper.py

@@ -113,7 +113,7 @@ def test_onnx_generate(self):
                 kwargs=inputs,
                 dynamic_shapes=ds,
                 filename=model_name,
-                exporter="custom",
+                exporter="modelbuilder",
             )
 
             print("-- test_onnx_generate: generate")

_unittests/ut_xrun_doc/test_documentation_examples.py

@@ -102,7 +102,7 @@ def add_test_methods(cls):
 
             if (
                 not reason
-                and name in {"plot_export_with_dynamic_cache.py", "plot_export_tiny_phi2.py"}
+                and name in {"plot_export_tiny_phi2.py", "plot_export_with_dynamic_cache.py"}
                 and not has_transformers("4.55")
             ):
                 reason = "transformers<4.55"
@@ -117,6 +117,7 @@ def add_test_methods(cls):
                     "plot_export_locate_issue.py",
                     "plot_export_with_auto.py",
                     "plot_export_tiny_llm.py",
+                    "plot_export_with_dynamic_cache.py",
                 }
                 and not has_torch("2.8")
             ):

onnx_diagnostic/export/api.py

@@ -42,7 +42,7 @@ def to_onnx(
         ), f"output_dynamic_shapes not supported for exporter={exporter!r}"
         epo = torch.onnx.export(
             mod,
-            args=args,
+            args=args or tuple(),
             kwargs=kwargs,
             input_names=input_names,
             output_names=output_names,
@@ -54,4 +54,40 @@ def to_onnx(
         epo.save(filename)
         return epo
 
+    if exporter == "modelbuilder":
+        import os
+        from ..helpers import flatten_object, string_type
+        from ..helpers.model_builder_helper import create_model_builder, save_model_builder
+
+        assert filename, f"filename must be specified for exporter={exporter!r}"
+        assert (
+            not output_dynamic_shapes
+        ), f"output_dynamic_shapes not supported for exporter={exporter!r}"
+        assert hasattr(mod, "config"), f"configuration is missing in model class {type(mod)}"
+        assert not args, f"only kwargs can be defined with exporter={exporter!r}"
+        assert list(kwargs) == ["input_ids", "attention_mask", "past_key_values"], (
+            f"Only a specified set of inputs is supported for exporter={exporter!r}, "
+            f"but it is {list(kwargs)}"
+        )
+        flat_inputs = flatten_object(kwargs, drop_keys=True)
+        first = flat_inputs[0]
+        first_float = [
+            t
+            for t in flat_inputs
+            if t.dtype in {torch.float32, torch.double, torch.float16, torch.bfloat16}
+        ]
+        assert first_float, (
+            f"Unable to find a float tensor in the inputs "
+            f"{string_type(kwargs, with_shape=True)}"
+        )
+        onx = create_model_builder(
+            mod.config,
+            mod,
+            precision=str(first_float[0].dtype).split(".")[-1],
+            execution_provider="cuda" if first.is_cuda else "cpu",
+            cache_dir=os.path.dirname(filename),
+        )
+        save_model_builder(onx, os.path.dirname(filename))
+        return onx
+
     raise ValueError(f"Unknown exporter={exporter!r}")