improve documentation

Author: xadupre

README.rst

@@ -56,6 +56,11 @@ or
 Enlightening Examples
 +++++++++++++++++++++
 
+**Where to start to export a model**
+
+* `Export microsoft/phi-2
+  <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_tiny_phi2.html>`_
+
 **Torch Export**
 
 * `Use DYNAMIC or AUTO when exporting if dynamic shapes has constraints

_doc/examples/plot_export_tiny_llm.py

@@ -1,8 +1,8 @@
 """
 .. _l-plot-tiny-llm-export:
 
-Steel method forward to guess the dynamic shapes (with Tiny-LLM)
-================================================================
+Steel method forward to guess inputs and dynamic shapes (with Tiny-LLM)
+=======================================================================
 
 Inputs are always dynamic with LLMs that is why dynamic shapes
 needs to be specified when a LLM is exported with :func:`torch.export.export`.

_doc/examples/plot_export_tiny_phi2.py

@@ -1,11 +1,16 @@
 """
 .. _l-plot-export_tiny_phi2:
 
-Untrained microsoft/phi-2
-=========================
+======================
+Export microsoft/phi-2
+======================
 
-:epkg:`microsoft/phi-2` is not a big models but still quite big
-when it comes to write unittest. Function
+This function exports an smaller untrained model with the same architecture.
+It is faster than the pretrained model.
+When this works, the untrained model can be replaced by the trained one.
+
+:epkg:`microsoft/phi-2` is not a big model but still quite big
+when it comes to write unittests. Function
 :func:`onnx_diagnostic.torch_models.hghub.get_untrained_model_with_inputs`
 can be used to create a reduced untrained version of a model coming from
 :epkg:`HuggingFace`. It downloads the configuration from the website
@@ -14,7 +19,7 @@
 the export or to compare performance. The relevance does not matter.
 
 Create the dummy model
-++++++++++++++++++++++
+======================
 """
 
 import copy
@@ -48,6 +53,8 @@
 print(f"model {size / 2**20:1.3f} Mb with {n_weights // 1000} mille parameters.")
 # %%
 # The original model has 2.7 billion parameters. It was divided by more than 10.
+# However, it can still be used with
+# ``get_untrained_model_with_inputs("microsoft/phi-2", same_as_pretrained=True)``.
 # Let's see the configuration.
 print(config)
 
@@ -72,13 +79,18 @@
 
 
 # %%
-# Export
-# ++++++
+# Export to fx.Graph
+# ==================
+#
+# :func:`torch.export.export` is the first step before converting
+# a model into ONNX. The inputs are duplicated (with ``copy.deepcopy``)
+# because the model may modify them inline (a cache for example).
+# Shapes may not match on the second call with the modified inputs.
 
 
-with torch_export_patches(patch_transformers=True) as modificator:
+with torch_export_patches(patch_transformers=True):
 
-    # Unnecessary steps but useful in case of an error
+    # Two unnecessary steps but useful in case of an error
     # We check the cache is registered.
     assert is_cache_dynamic_registered()
 
@@ -88,24 +100,26 @@
         d["abs"] < 1e-5
     ), f"The model with patches produces different outputs: {string_diff(d)}"
 
-    # Then we export.
+    # Then we export: the only import line in this section.
     ep = torch.export.export(
         untrained_model,
         (),
-        kwargs=modificator(copy.deepcopy(inputs)),
+        kwargs=copy.deepcopy(inputs),
         dynamic_shapes=use_dyn_not_str(dynamic_shapes),
         strict=False,  # mandatory for torch==2.6
     )
 
     # We check the exported program produces the same results as well.
+    # This step is again unnecessary.
     d = max_diff(expected, ep.module()(**copy.deepcopy(inputs)))
     assert d["abs"] < 1e-5, f"The exported model different outputs: {string_diff(d)}"
 
 # %%
 # Export to ONNX
-# ++++++++++++++
+# ==============
 #
-# The export works. We can export to ONNX now.
+# The export works. We can export to ONNX now
+# :func:`torch.onnx.export`.
 # Patches are still needed because the export
 # applies :meth:`torch.export.ExportedProgram.run_decompositions`
 # may export local pieces of the model again.
@@ -157,4 +171,51 @@
 # It looks good.
 
 # %%
-doc.plot_legend("untrained smaller\nmicrosoft/phi-2", "torch.onnx.export", "orange")
+doc.plot_legend("export\nuntrained smaller\nmicrosoft/phi-2", "torch.onnx.export", "orange")
+
+# %%
+# Possible Issues
+# ===============
+#
+# Unknown task
+# ++++++++++++
+#
+# Function :func:`onnx_diagnostic.torch_models.hghub.get_untrained_model_with_inputs`
+# is unabl to guess a task associated to the model.
+# A different set of dummy inputs is defined for every task.
+# The user needs to explicitly give that information to the function.
+# Tasks are the same as the one defined by
+# `HuggingFace/models <https://huggingface.co/models>`_.
+#
+# Inputs are incorrect
+# ++++++++++++++++++++
+#
+# Example :ref:`l-plot-tiny-llm-export` explains
+# how to retrieve that information. If you cannot guess the dynamic
+# shapes - a cache can be tricky sometimes, follow example
+# :ref:`l-plot-export-with-args-kwargs`.
+#
+# DynamicCache or any other cache cannot be exported
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+#
+# That's the role of :func:`onnx_diagnostic.torch_export_patches./torch_export_patches`.
+# It registers the necessary information into pytorch to make the export
+# work with these. Its need should slowly disappear until :epkg:`transformers`
+# includes the serialization functions.
+#
+# Control Flow
+# ++++++++++++
+#
+# Every mixture of models goes through a control flow (a test).
+# It also happens when a cache is truncated. The code of the model
+# needs to be changed. See example :ref:`l-plot-export-cond`.
+#
+# Issue with dynamic shapes
+# +++++++++++++++++++++++++
+#
+# Example :ref:`l-plot-dynamic-shapes-python-int` gives one reason
+# this process may fail but that's not the only one.
+# Example :ref:`l-plot-export-locale-issue` gives an way to locate
+# the cause but that does not cover all the possible causes.
+# Raising an issue on github would be the recommended option
+# until it is fixed.

_doc/index.rst

@@ -65,6 +65,10 @@ or
 Enlightening Examples
 +++++++++++++++++++++
 
+**Where to start to export a model**
+
+* :ref:`l-plot-export_tiny_phi2`
+
 **Torch Export**
 
 * :ref:`l-plot-export-cond`

_doc/recipes/plot_dynamic_shapes_python_int.py

@@ -1,4 +1,6 @@
 """
+.. _l-plot-dynamic-shapes-python-int:
+
 Do not use python int with dynamic shapes
 =========================================