Add phi

.github/workflows/ci.yml

@@ -10,13 +10,14 @@ on:
 
 jobs:
   run:
-    name: tr-${{ matrix.transformers }}-ci ${{ matrix.os }}-${{ matrix.python }}
+    name: to-${{ matrix.torch }}-tr-${{ matrix.transformers }}-ci ${{ matrix.os }}-${{ matrix.python }}
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
         os: [ubuntu-latest]
         python: ['3.11', '3.12']
         transformers: ['4.48', '4.50', 'main']
+        torch: ['main']
 
     steps:
       - uses: actions/checkout@v3
@@ -26,7 +27,13 @@ jobs:
           python-version: ${{ matrix.python }}
 
       - name: Install pytorch
-        run: python -m pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu
+        run: |
+          if [[ "${{ matrix.torch }}" == "main" ]]; then
+            python -m pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu
+          else
+            echo "install torch==${{ matrix.torch }}"
+            pip install torch==${{ matrix.torch }}
+          fi
 
       - name: Install transformers ${{ matrix.transformers }}
         run: |

README.rst

@@ -50,6 +50,8 @@ Enlightening Examples
 
 * `Use DYNAMIC or AUTO when exporting if dynamic shapes has constraints
   <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_with_dynamic_shapes_auto.html>`_
+* `Find and fix an export issue due to dynamic shapes
+  <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_locate_issue.html>`_
 * `Export with DynamicCache and dynamic shapes
   <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_with_dynamic_cache.html>`_
 * `Steel method forward to guess the dynamic shapes (with Tiny-LLM)

_doc/api/torch_models/llms.rst

@@ -3,5 +3,4 @@ onnx_diagnostic.torch_models.llms
 =================================
 
 .. automodule:: onnx_diagnostic.torch_models.llms
-    :members:
-    :no-undoc-members:
+    :members: get_phi2, get_tiny_llm

_doc/conf.py

@@ -158,6 +158,7 @@
     "ignore_repr_types": "matplotlib\\.(text|axes)",
     # robubstness
     "reset_modules_order": "both",
+    "reset_modules": ("matplotlib", "onnx_diagnostic.doc.reset_torch_transformers"),
 }
 
 if int(os.environ.get("UNITTEST_GOING", "0")):

_doc/examples/plot_export_cond.py

@@ -13,6 +13,7 @@
 """
 
 import torch
+from onnx_diagnostic import doc
 
 
 # %%
@@ -84,3 +85,8 @@ def neg(x):
 
 ep = torch.export.export(model, (x,))
 print(ep.graph)
+
+
+# %%
+
+doc.plot_legend("If -> torch.cond", "torch.export.export", "tomato")

_doc/examples/plot_export_locate_issue.py

@@ -0,0 +1,105 @@
+"""
+.. _l-plot-export-locale-issue:
+
+==================================================
+Find and fix an export issue due to dynamic shapes
+==================================================
+
+LLMs must be exported with dynamic shapes and it is common that
+a static dimension turns into a static ones. The error message from
+:epkg:`pytorch` tells the user to define ``TORCH_LOGS="+dynamic"``
+but it shows a very long list of messages where we need
+to find the string ``range_refined_to_singleton`` and that
+does not really indicates where it comes from. The example
+shows how to tweak pytorch to get that information until
+it gets better.
+
+A model with an export issue
+============================
+
+The following model implies the first dimension of x is equal to 1
+or equal to the number of element in the list ``ys``.
+It is not really dynamic. It looks obvious here but
+it is difficult to find deep inside a big model.
+"""
+
+import traceback
+import torch
+from onnx_diagnostic import doc
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+
+
+class ModelWithIssue(torch.nn.Module):
+    def forward(self, x: torch.Tensor, ys: list[torch.Tensor]):
+        caty = torch.cat([y.unsqueeze(0) for y in ys], axis=0)
+        z = x * caty
+        return z
+
+
+inputs = (torch.rand(2, 3, 1), [torch.rand(3, 4), torch.rand(3, 4)])
+model = ModelWithIssue()
+model(*inputs)
+
+
+# %%
+# Let's export.
+
+DYN = torch.export.Dim.DYNAMIC
+dyn_shapes = ({0: DYN, 1: DYN}, [{0: DYN, 1: DYN}, {0: DYN, 1: DYN}])
+try:
+    ep = torch.export.export(model, inputs, dynamic_shapes=dyn_shapes)
+    print(ep)
+except Exception as e:
+    print("-- ERROR:")
+    print(e)
+
+# %%
+# The error shows:
+#
+# .. code-block::
+#
+#       Constraints violated (L['args'][0][0].size()[0])!
+#           For more information, run with TORCH_LOGS="+dynamic".
+#       - Not all values of RelaxedUnspecConstraint(L['args'][0][0].size()[0])
+#           are valid because L['args'][0][0].size()[0] was inferred to be a constant (2).
+#
+# Where does it happens? That's a tricky question we need to answer.
+# The message is raised from
+# `torch.fx.experimental.symbolic_shapes.ShapeEnv._set_replacement
+# <https://github.com/pytorch/pytorch/blob/main/torch/fx/experimental/symbolic_shapes.py#L6239>`_.
+# One way to find the exact location is to retrieve a stack trace
+# by inserting an assert such as the following:
+#
+# .. code-block::
+#
+#   assert msg != "range_refined_to_singleton", (
+#       f"A dynamic dimension becomes static! "
+#       f"a={a!r}, tgt={tgt!r}, msg={msg!r}, tgt_bound={tgt_bound}"
+#   )
+#
+# Stop when a dynamic dimension turns static
+# ==========================================
+#
+# We use :func:`bypass_export_some_errors
+# <onnx_diagnostic.torch_export_patches.bypass_export_some_errors>`
+# to replace torch implementation by a new one raising the exception
+# mentioned in previous section.
+
+with bypass_export_some_errors(stop_if_static=True, verbose=1):
+    try:
+        torch.export.export(model, inputs, dynamic_shapes=dyn_shapes)
+    except AssertionError:
+        print("-- It failed as excepted. Let's print the stack trace.")
+        print(traceback.format_exc())
+
+# The stack trace is quite long but the first line referring to this example
+# is the following one. It points out the line turing a dynamic dimension into
+# static.
+#
+# .. code-block::
+#
+#   File "onnx-diagnostic/_doc/examples/plot_export_locate_issue.py", line 25, in forward
+#       z = x * caty
+
+
+doc.plot_legend("was inferred to be a constant", "torch.export.export", "tomato")

_doc/examples/plot_export_tiny_llm.py

@@ -30,6 +30,7 @@
 import pprint
 import torch
 import transformers
+from onnx_diagnostic import doc
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.torch_models.llms import get_tiny_llm
 
@@ -44,10 +45,11 @@
 
 def _forward_(*args, _f=None, **kwargs):
     assert _f is not None
-    if not torch.compiler.is_exporting():
+    if not hasattr(torch.compiler, "is_exporting") or not torch.compiler.is_exporting():
+        # torch.compiler.is_exporting requires torch>=2.7
         print("<-", string_type((args, kwargs), with_shape=True, with_min_max=True))
     res = _f(*args, **kwargs)
-    if not torch.compiler.is_exporting():
+    if not hasattr(torch.compiler, "is_exporting") or not torch.compiler.is_exporting():
         print("->", string_type((args, kwargs), with_shape=True, with_min_max=True))
     return res
 
@@ -67,7 +69,8 @@ def _forward_(*args, _f=None, **kwargs):
 )
 
 generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-print(generated_text)
+print("-- prompt", prompt)
+print("-- answer", generated_text)
 
 # %%
 # Let's restore the forward as it was.
@@ -84,7 +87,8 @@ def _forward_(*args, _f=None, **kwargs):
 #
 # Let's create an untrained model using the config file provided
 # `config.json <https://huggingface.co/arnir0/Tiny-LLM/blob/main/config.json>`_
-# to create an untrained model: :func:`....get_tiny_llm`.
+# to create an untrained model:
+# :func:`onnx_diagnostic.torch_models.llms.get_tiny_llm`.
 # Then let's use it.
 
 experiment = get_tiny_llm()
@@ -138,7 +142,7 @@ 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:`....get_tiny_llm`.
+# :func:`onnx_diagnostic.torch_models.llms.get_tiny_llm`.
 
 data = get_tiny_llm()
 inputs, dynamic_shapes = data["inputs"], data["dynamic_shapes"]
@@ -167,3 +171,5 @@ def _forward_(*args, _f=None, **kwargs):
 # %%
 # If you have any error, then look at example
 # :ref:`l-plot-tiny-llm-export-patched`.
+
+doc.plot_legend("Tiny-LLM fails", "torch.export.export", "tomato")

_doc/examples/plot_export_tiny_llm_patched.py

@@ -66,6 +66,7 @@
 import pprint
 import torch
 import transformers
+from onnx_diagnostic import doc
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.torch_export_patches.onnx_export_errors import bypass_export_some_errors
 from onnx_diagnostic.torch_models.llms import get_tiny_llm
@@ -122,3 +123,6 @@
     )
     print("It worked:")
     print(ep)
+
+# %%
+doc.plot_legend("Tiny-LLM patched", "torch.export.export", "green")

_doc/examples/plot_export_with_dynamic_cache.py

@@ -23,6 +23,7 @@
 
 import pprint
 import torch
+from onnx_diagnostic import doc
 from onnx_diagnostic.cache_helpers import make_dynamic_cache
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.export import ModelInputs
@@ -221,3 +222,7 @@ def forward(self, cache, z):
             model, modificator(inputs[0]), dynamic_shapes=ds[0], strict=False
         )
 print(ep)
+
+# %%
+
+doc.plot_legend("dynamic shapes", "torch.export.export", "tomato")

_doc/examples/plot_export_with_dynamic_shapes_auto.py

@@ -12,6 +12,7 @@
 """
 
 import torch
+from onnx_diagnostic import doc
 
 
 class Model(torch.nn.Module):
@@ -57,7 +58,7 @@ def forward(self, x, y, z):
         },
     )
     print(ep)
-    raise AssertionError("able to export this moel, please update the tutorial")
+    raise AssertionError("able to export this model, please update the tutorial")
 except torch._dynamo.exc.UserError as e:
     print(f"unable to use Dim('dz') because {type(e)}, {e}")
 
@@ -90,3 +91,7 @@ def forward(self, x, y, z):
         dynamic_shapes=({0: AUTO, 1: AUTO}, {0: AUTO, 1: AUTO}, {0: AUTO, 1: AUTO}),
     )
 )
+
+# %%
+
+doc.plot_legend("dynamic shapes inferred", "torch.export.export", "tomato")

_doc/examples/plot_failing_model_extract.py

@@ -25,6 +25,7 @@
 import onnx
 import onnx.helper as oh
 import onnxruntime
+from onnx_diagnostic import doc
 from onnx_diagnostic.helpers import from_array_extended
 from onnx_diagnostic.ort_session import investigate_onnxruntime_issue
 
@@ -96,3 +97,7 @@
     onnx.shape_inference.infer_shapes(model, strict_mode=True)
 except onnx.onnx_cpp2py_export.shape_inference.InferenceError as e:
     print(e)
+
+# %%
+
+doc.plot_legend("Run until it fails", "onnxruntime.InferenceSession", "lightgrey")

_doc/examples/plot_failing_onnxruntime_evaluator.py

@@ -28,6 +28,7 @@
 import onnx.helper as oh
 import torch
 import onnxruntime
+from onnx_diagnostic import doc
 from onnx_diagnostic.ext_test_case import has_cuda
 from onnx_diagnostic.helpers import from_array_extended
 from onnx_diagnostic.reference import OnnxruntimeEvaluator
@@ -104,3 +105,5 @@
 # This runtime is useful when it fails for a numerical reason.
 # It is possible to insert prints in the python code to print
 # more information or debug if needed.
+
+doc.plot_legend("onnxruntime running step by step", "OnnxruntimeEvaluator", "lightgrey")

_doc/examples/plot_failing_reference_evaluator.py

@@ -19,6 +19,7 @@
 import onnx
 import onnx.helper as oh
 import onnxruntime
+from onnx_diagnostic import doc
 from onnx_diagnostic.helpers import from_array_extended
 from onnx_diagnostic.reference import ExtendedReferenceEvaluator
 
@@ -79,3 +80,5 @@
 # This runtime is useful when it fails for a numerical reason.
 # It is possible to insert prints in the python code to print
 # more information or debug if needed.
+
+doc.plot_legend("Python Runtime for ONNX", "ExtendedReferenceEvalutor", "lightgrey")

_doc/index.rst

@@ -61,6 +61,7 @@ Enlightening Examples
 * :ref:`l-plot-export-cond`
 * :ref:`l-plot-sxport-with-dynamio-shapes-auto`
 * :ref:`l-plot-export-with-dynamic-shape`
+* :ref:`l-plot-export-locale-issue`
 * :ref:`l-plot-tiny-llm-export`
 * :ref:`l-plot-tiny-llm-export-patched`
 
@@ -162,5 +163,7 @@ Size of the package:
 Older versions
 ++++++++++++++
 
+* `0.2.2 <../v0.2.2/index.html>`_
+* `0.2.1 <../v0.2.1/index.html>`_
 * `0.2.0 <../v0.2.0/index.html>`_
 * `0.1.0 <../v0.1.0/index.html>`_

_unittests/ut_torch_models/test_llm_phi2.py

@@ -0,0 +1,41 @@
+import unittest
+import torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings, requires_transformers
+from onnx_diagnostic.torch_models.llms import get_phi2
+from onnx_diagnostic.helpers import string_type
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+
+
+class TestLlmPhi(ExtTestCase):
+    def test_get_phi2(self):
+        data = get_phi2(num_hidden_layers=2)
+        model, inputs = data["model"], data["inputs"]
+        self.assertIn("DynamicCache", string_type(inputs))
+        model(**inputs)
+
+    @ignore_warnings(UserWarning)
+    @requires_transformers("4.52")
+    def test_export_phi2_1(self):
+        data = get_phi2(num_hidden_layers=2)
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        self.assertEqual(
+            {"attention_mask", "past_key_values", "input_ids", "position_ids"}, set(inputs)
+        )
+        ep = torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds)
+        assert ep
+
+    @ignore_warnings(UserWarning)
+    def test_export_phi2_2_bypassed(self):
+        data = get_phi2(num_hidden_layers=2)
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        self.assertEqual(
+            {"attention_mask", "past_key_values", "input_ids", "position_ids"}, set(inputs)
+        )
+        with bypass_export_some_errors(patch_transformers=True) as modificator:
+            inputs = modificator(inputs)
+            ep = torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds, strict=False)
+            assert ep
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)