Improves investigate_onnxruntime_issue

CHANGELOGS.rst

@@ -1,6 +1,11 @@
 Change Logs
 ===========
 
+0.2.0
++++++
+
+* :pr:`7`: improves function ``investigate_onnxruntime_issue``
+
 0.1.0
 +++++
 

README.rst

@@ -44,6 +44,15 @@ or
 
     pip install onnx-diagnostic
 
+**Enlightening Examples**
+
+* `Use DYNAMIC or AUTO when dynamic shapes has constraints
+  <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_with_dynamic_shapes_auto.html>`_
+* `Steel method forward to guess the dynamic shapes
+  <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_tiny_llm.html>`_
+* `Find where a model is failing by running submodels
+  <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_failing_model_extract.html>`_
+
 Snapshot of usefuls tools
 +++++++++++++++++++++++++
 

_doc/examples/plot_export_cond.py

@@ -0,0 +1,86 @@
+"""
+.. _l-plot-export-cond:
+
+Export a model with a control flow (If)
+=======================================
+
+Control flow cannot be exported with a change.
+The code of the model can be changed or patched
+to introduce function :func:`torch.cond`.
+
+A model with a test
++++++++++++++++++++
+"""
+
+import torch
+
+
+# %%
+# We define a model with a control flow (-> graph break)
+
+
+class ForwardWithControlFlowTest(torch.nn.Module):
+    def forward(self, x):
+        if x.sum():
+            return x * 2
+        return -x
+
+
+class ModelWithControlFlow(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mlp = torch.nn.Sequential(
+            torch.nn.Linear(3, 2),
+            torch.nn.Linear(2, 1),
+            ForwardWithControlFlowTest(),
+        )
+
+    def forward(self, x):
+        out = self.mlp(x)
+        return out
+
+
+model = ModelWithControlFlow()
+
+# %%
+# Let's check it runs.
+x = torch.randn(1, 3)
+model(x)
+
+# %%
+# As expected, it does not export.
+try:
+    torch.export.export(model, (x,))
+    raise AssertionError("This export should failed unless pytorch now supports this model.")
+except Exception as e:
+    print(e)
+
+
+# %%
+# Suggested Patch
+# +++++++++++++++
+#
+# Let's avoid the graph break by replacing the forward.
+
+
+def new_forward(x):
+    def identity2(x):
+        return x * 2
+
+    def neg(x):
+        return -x
+
+    return torch.cond(x.sum() > 0, identity2, neg, (x,))
+
+
+print("the list of submodules")
+for name, mod in model.named_modules():
+    print(name, type(mod))
+    if isinstance(mod, ForwardWithControlFlowTest):
+        mod.forward = new_forward
+
+# %%
+# Let's see what the fx graph looks like.
+
+ep = torch.export.export(model, (x,))
+print(ep.graph)

_doc/examples/plot_export_tiny_llm.py

@@ -1,23 +1,33 @@
 """
 .. _l-plot-tiny-llm-export:
 
-Export LLM with dynamic shapes
-==============================
+Steel method forward to guess the dynamic shapes
+================================================
+
+Inputs are always dynamic with LLMs that is why dyanmic shapes
+needs to be specified when a LLM is exported with:func:`torch.export.export`.
+Most of the examples on :epkg:`HuggingFace` use method
+:meth:`transformers.GenerationMixin.generate` but we only want to
+export the model and its method ``forward``.
+
+That example shows to guess the inputs of this method even though the model
+is executed through meth ``generate``.
 
 We focus on the model
 `Tiny-LLM <https://huggingface.co/arnir0/Tiny-LLM>`_.
 To avoid downloading any weigths, we write a function creating a
 random model based on the same architecture.
 
-Guess the cache dimension
-+++++++++++++++++++++++++
+Steel the forward method
+++++++++++++++++++++++++
 
 The first step is to guess the dummy inputs.
 Let's use the true model for that.
 We use the dummy example from the model page.
 """
 
 import copy
+import pprint
 import torch
 import transformers
 from onnx_diagnostic.helpers import string_type
@@ -64,8 +74,13 @@ def _forward_(*args, _f=None, **kwargs):
 model.forward = keep_model_forward
 
 # %%
-# The model creation
-# ++++++++++++++++++
+# Untrained model
+# +++++++++++++++
+#
+# This part can skipped if you are only interested in exporting
+# the original model. It is useful to create a unit test to ensure
+# a specific architecture can be exported despite the many changes
+# brought to :epkg:`torch` or :epkg:`transformers`.
 #
 # Let's create an untrained model using the config file provided
 # `config.json <https://huggingface.co/arnir0/Tiny-LLM/blob/main/config.json>`_
@@ -126,6 +141,22 @@ def _forward_(*args, _f=None, **kwargs):
 # ++++++++++++++++++++++++++
 #
 # Let's use the same dummy inputs but we use the downloaded model.
+# Dummy inputs and dynamic shapes are created by function
+# :func:`onnx_diagnostic.torch_models.llms.get_tiny_llm`.
+
+data = get_tiny_llm()
+inputs, dynamic_shapes = data["inputs"], data["dynamic_shapes"]
+
+# %%
+# Let's print the inputs.
+
+print(string_type(inputs, with_shape=True))
+
+# %% Let's print the dynamic shapes
+pprint.pprint(dynamic_shapes)
+
+# %%
+# And Let's finally export.
 
 try:
     ep = torch.export.export(model, (), kwargs=cloned_inputs, dynamic_shapes=dynamic_shapes)

_doc/examples/plot_export_with_dynamic_shapes_auto.py

@@ -1,4 +1,6 @@
 """
+.. _l-plot-sxport-with-dynamio-shapes-auto:
+
 Use DYNAMIC or AUTO when dynamic shapes has constraints
 =======================================================
 

_doc/examples/plot_failing_model_extract.py

@@ -0,0 +1,98 @@
+"""
+.. _l-plot-failing-model-extract:
+
+Find where a model is failing by running submodels
+==================================================
+
+Let's assume :epkg:`onnxruntime` crashes without telling why or where.
+The first thing is do is to locate where. For that, we extract every submodel
+starting from the inputs and running the first *n* nodes of the model.
+The model is likely to fail for some *n*. Then the failing is known.
+
+This method only works if the model only contains operator coming
+from the main domain *ai.onnx* otherwise shape inference stops
+at the first non standard operator and the algorithm fails at
+producing :class:`onnx.ModelProto` including the non standard operators.
+
+A failing model
++++++++++++++++
+
+The issue here is a an operator ``Cast`` trying to convert a result
+into a non-existing type.
+"""
+
+import numpy as np
+import onnx
+import onnx.helper as oh
+import onnxruntime
+from onnx_diagnostic.helpers import from_array_extended
+from onnx_diagnostic.ort_session import investigate_onnxruntime_issue
+
+TFLOAT = onnx.TensorProto.FLOAT
+
+model = oh.make_model(
+    oh.make_graph(
+        [
+            oh.make_node("Mul", ["X", "Y"], ["xy"], name="n0"),
+            oh.make_node("Sigmoid", ["xy"], ["sy"], name="n1"),
+            oh.make_node("Add", ["sy", "one"], ["C"], name="n2"),
+            oh.make_node("Cast", ["C"], ["X999"], to=999, name="failing"),
+            oh.make_node("CastLike", ["X999", "Y"], ["Z"], name="n4"),
+        ],
+        "nd",
+        [
+            oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c"]),
+            oh.make_tensor_value_info("Y", TFLOAT, ["a", "b", "c"]),
+        ],
+        [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "c"])],
+        [from_array_extended(np.array([1], dtype=np.float32), name="one")],
+    ),
+    opset_imports=[oh.make_opsetid("", 18)],
+    ir_version=9,
+)
+
+# %%
+# We check it is failing.
+
+try:
+    onnxruntime.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
+except onnxruntime.capi.onnxruntime_pybind11_state.Fail as e:
+    print(e)
+
+
+# %%
+# Shape Inference
+# +++++++++++++++
+#
+# Building submodels requires to known the output type.
+# We run shape inference on the model.
+shaped_model = onnx.shape_inference.infer_shapes(model)
+
+
+# %%
+# Looping over the nodes
+# ++++++++++++++++++++++
+#
+#
+
+failing = investigate_onnxruntime_issue(shaped_model, providers="cpu", verbose=1, quiet=True)
+
+# %%
+# Let's print the failing node.
+print(failing)
+
+
+# %%
+# Detect an issue with shape Inference
+# ++++++++++++++++++++++++++++++++++++
+#
+# We could have caught the error sooner by asking shape inference
+# to raise an exception if one node could not be processed.
+# It means either the node is a custom node
+# and shape inference has no way to guess the output type and shape
+# for this node or shape inference failed.
+
+try:
+    onnx.shape_inference.infer_shapes(model, strict_mode=True)
+except onnx.onnx_cpp2py_export.shape_inference.InferenceError as e:
+    print(e)

_doc/index.rst

@@ -45,8 +45,14 @@ Source are `sdpython/onnx-diagnostic
     CHANGELOGS
     license
 
+**Enlightening Examples**
 
-**Some usefuls tools**
+* :ref:`l-plot-export-cond`
+* :ref:`l-plot-sxport-with-dynamio-shapes-auto`
+* :ref:`l-plot-tiny-llm-export`
+* :ref:`l-plot-failing-model-extract`
+
+**Some Usefuls Tools**
 
 .. code-block:: python
 
@@ -135,7 +141,6 @@ Size of the package:
     gr = df[["dir", "ext", "lines", "chars"]].groupby(["ext", "dir"]).sum()
     print(gr)
 
-Older versions
-++++++++++++++
+**Older versions**
 
 * `0.1.0 <../v0.1.0/index.html>`_

_unittests/ut_xrun_doc/test_helpers.py

@@ -29,6 +29,7 @@
     rename_dynamic_dimensions,
     rename_dynamic_expression,
 )
+from onnx_diagnostic.cache_helpers import make_dynamic_cache
 
 TFLOAT = onnx.TensorProto.FLOAT
 
@@ -94,6 +95,32 @@ def test_pretty_onnx(self):
         pretty_onnx(proto.graph)
         pretty_onnx(proto.graph.node[0])
 
+    @hide_stdout()
+    def test_print_pretty_onnx(self):
+        proto = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("Sigmoid", ["Y"], ["sy"]),
+                    oh.make_node("Mul", ["Y", "sy"], ["ysy"]),
+                    oh.make_node("Mul", ["X", "ysy"], ["final"]),
+                ],
+                "nd",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, [1, "b", "c"]),
+                    oh.make_tensor_value_info("Y", TFLOAT, ["a", "b", "c"]),
+                ],
+                [oh.make_tensor_value_info("final", TFLOAT, ["a", "b", "c"])],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+            ir_version=9,
+        )
+        self.print_onnx(proto)
+        self.print_model(proto)
+        self.dump_onnx("test_print_pretty_onnx", proto)
+        self.check_ort(proto)
+        self.assertNotEmpty(proto)
+        self.assertEmpty(None)
+
     def test_get_onnx_signature(self):
         proto = oh.make_model(
             oh.make_graph(
@@ -115,16 +142,23 @@ def test_get_onnx_signature(self):
         sig = get_onnx_signature(proto)
         self.assertEqual(sig, (("X", 1, (1, "b", "c")), ("Y", 1, ("a", "b", "c"))))
 
+    @hide_stdout()
     def test_flatten(self):
         inputs = (
             torch.rand((3, 4), dtype=torch.float16),
             [
                 torch.rand((5, 6), dtype=torch.float16),
                 torch.rand((5, 6, 7), dtype=torch.float16),
+                {
+                    "a": torch.rand((2,), dtype=torch.float16),
+                    "cache": make_dynamic_cache(
+                        [(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]
+                    ),
+                },
             ],
         )
-        flat = flatten_object(inputs)
-        diff = max_diff(inputs, flat, flatten=True)
+        flat = flatten_object(inputs, drop_keys=True)
+        diff = max_diff(inputs, flat, flatten=True, verbose=10)
         self.assertEqual(diff["abs"], 0)
         d = string_diff(diff)
         self.assertIsInstance(d, str)