From ef48a6c5e4c109990267328c2b15a754e657313a Mon Sep 17 00:00:00 2001 From: xadupre Date: Wed, 11 Jun 2025 09:15:26 +0200 Subject: [PATCH 1/9] update index --- .../plot_dump_intermediate_results.py | 118 ++++++++++++++++++ _doc/index.rst | 1 + onnx_diagnostic/helpers/mini_onnx_builder.py | 2 + onnx_diagnostic/helpers/torch_helper.py | 3 +- 4 files changed, 123 insertions(+), 1 deletion(-) create mode 100644 _doc/examples/plot_dump_intermediate_results.py diff --git a/_doc/examples/plot_dump_intermediate_results.py b/_doc/examples/plot_dump_intermediate_results.py new file mode 100644 index 00000000..dbf31e96 --- /dev/null +++ b/_doc/examples/plot_dump_intermediate_results.py @@ -0,0 +1,118 @@ +""" +.. _l-plot-intermediate-results: + +Dumps intermediate results of a torch model +=========================================== + + +codellama/CodeLlama-7b-Python-hf +++++++++++++++++++++++++++++++++ + +""" + +import onnx +import torch +from onnx_array_api.plotting.graphviz_helper import plot_dot +from onnx_diagnostic import doc +from onnx_diagnostic.helpers import string_type +from onnx_diagnostic.helpers.torch_helper import dummy_llm +from onnx_diagnostic.helpers.mini_onnx_builder import create_input_tensors_from_onnx_model +from onnx_diagnostic.helpers.torch_helper import steal_forward + + +model, inputs, ds = dummy_llm(dynamic_shapes=True) + +print(f"type(model)={type(model)}") +print(f"inputs={string_type(inputs, with_shape=True)}") +print(f"ds={string_type(ds, with_shape=True)}") + +# %% +# It contains the following submodules. + +for name, mod in model.named_modules(): + print(f"- {name}: {type(mod)}") + +# %% +# Steal and dump the output of submodules +# +++++++++++++++++++++++++++++++++++++++ +# +# The following context spies on the intermediate results +# for the following module and submodules. It stores +# in one onnx file all the input/output for those. + +with steal_forward( + [ + ("model", model), + ("model.decoder", model.decoder), + ("model.decoder.attention", model.decoder.attention), + ("model.decoder.feed_forward", model.decoder.feed_forward), + ("model.decoder.norm_1", model.decoder.norm_1), + ("model.decoder.norm_2", model.decoder.norm_2), + ], + dump_file="plot_dump_intermediate_results.inputs.onnx", + verbose=1, + storage_limit=2**28, +): + model(*inputs) + + +# %% +# Restores saved inputs/outputs +# +++++++++++++++++++++++++++++ +# +# All the intermediate tensors were saved in one unique onnx model, +# every tensor is stored in a constant node. +# The model can be run with any runtime to restore the inputs +# and function :func:`onnx_diagnostic.mini_onnx_builder.create_input_tensors_from_onnx_model` +# can restore their names. + +saved_tensors = create_input_tensors_from_onnx_model( + "plot_dump_intermediate_results.inputs.onnx" +) +for k, v in saved_tensors.items(): + print(f"{k} -- {string_type(v, with_shape=True)}") + +# %% +# Let's explained the naming convention. +# +# :: +# ('model.decoder.norm_2', 0, 'I') -- ((T1s2x30x16,),{}) +# | | | +# | | +--> input, the format is args, kwargs +# | | +# | +--> iteration, 0 means the first time the execution +# | went through that module +# | it is possible to call multiple times, +# | the model to store more +# | +# +--> the name given to steal forward +# +# The same goes for output except ``'I'`` is replaced by ``'O'``. +# +# :: +# +# ('model.decoder.norm_2', 0, 'O') -- T1s2x30x16 +# +# This trick can be used to compare intermediate results coming +# from pytorch to any other implementation of the same model +# as long as it is possible to map the stored inputs/outputs. + +# %% +# Conversion to ONNX +# ++++++++++++++++++ +# +# The difficult point is to be able to map the saved intermediate +# results to intermediate results in ONNX. +# Let's create the ONNX model. + +epo = torch.onnx.export(model, inputs, dynamic_shapes=ds, dynamo=True) +epo.optimize() +epo.save("plot_dump_intermediate_results.onnx") + +# %% +# It looks like the following. +onx = onnx.load("plot_dump_intermediate_results.onnx") +plot_dot(onx) + +# %% +doc.plot_legend("steal and dump\nintermediate\nresults", "steal_forward", "blue") diff --git a/_doc/index.rst b/_doc/index.rst index 88dbabf7..fdc53930 100644 --- a/_doc/index.rst +++ b/_doc/index.rst @@ -92,6 +92,7 @@ Enlightening Examples * :ref:`l-plot-failing-reference-evaluator` * :ref:`l-plot-failing-onnxruntime-evaluator` * :ref:`l-plot-failing-model-extract` +* :ref:`l-plot-intermediate-results` Some Usefuls Tools ================== diff --git a/onnx_diagnostic/helpers/mini_onnx_builder.py b/onnx_diagnostic/helpers/mini_onnx_builder.py index df3df7a8..4ce656c9 100644 --- a/onnx_diagnostic/helpers/mini_onnx_builder.py +++ b/onnx_diagnostic/helpers/mini_onnx_builder.py @@ -532,6 +532,8 @@ def create_input_tensors_from_onnx_model( :param engine: runtime to use, onnx, the default value, onnxruntime :param sep: separator :return: restored data + + See example :ref:`l-plot-intermediate-results` for an example. """ if engine == "ExtendedReferenceEvaluator": from ..reference import ExtendedReferenceEvaluator diff --git a/onnx_diagnostic/helpers/torch_helper.py b/onnx_diagnostic/helpers/torch_helper.py index 85570f75..10606293 100644 --- a/onnx_diagnostic/helpers/torch_helper.py +++ b/onnx_diagnostic/helpers/torch_helper.py @@ -288,7 +288,8 @@ def steal_forward( """ The necessary modification to steem forward method and prints out inputs and outputs using :func:`onnx_diagnostic.helpers.string_type`. - See example :ref:`l-plot-tiny-llm-export`. + See example :ref:`l-plot-tiny-llm-export` or + :ref:`l-plot-intermediate-results`. :param model: a model or a list of models to monitor, every model can also be a tuple(name, model), name is displayed well. From 39f6f137261b9de94c43e16fecdcabd89f49bb22 Mon Sep 17 00:00:00 2001 From: xadupre Date: Wed, 11 Jun 2025 09:20:31 +0200 Subject: [PATCH 2/9] doc --- onnx_diagnostic/helpers/mini_onnx_builder.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/onnx_diagnostic/helpers/mini_onnx_builder.py b/onnx_diagnostic/helpers/mini_onnx_builder.py index 4ce656c9..8fe1ef7d 100644 --- a/onnx_diagnostic/helpers/mini_onnx_builder.py +++ b/onnx_diagnostic/helpers/mini_onnx_builder.py @@ -393,7 +393,8 @@ def create_onnx_model_from_input_tensors( Creates a model proto including all the value as initializers. They can be restored by executing the model. We assume these inputs are not bigger than 2Gb, - the limit of protobuf. + the limit of protobuf. Nothing is implemented yet to get around + that limit. :param inputs: anything :param switch_low_high: if None, it is equal to ``switch_low_high=sys.byteorder != "big"`` From c48e8671bba3e9057c5551b805aef697235bfbcb Mon Sep 17 00:00:00 2001 From: xadupre Date: Wed, 11 Jun 2025 09:32:01 +0200 Subject: [PATCH 3/9] dot --- _unittests/ut_xrun_doc/test_documentation_technical.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/_unittests/ut_xrun_doc/test_documentation_technical.py b/_unittests/ut_xrun_doc/test_documentation_technical.py index 5dbfb661..e7fa2661 100644 --- a/_unittests/ut_xrun_doc/test_documentation_technical.py +++ b/_unittests/ut_xrun_doc/test_documentation_technical.py @@ -73,7 +73,12 @@ def add_test_methods(cls): if not name.endswith(".py") or not name.startswith("plot_"): continue reason = None - if not reason and not has_dot and name in {"plot_layer_norm_discrepancies.py"}: + if ( + not reason + and not has_dot + and name + in {"plot_layer_norm_discrepancies.py", "plot_dump_intermediate_results.py"} + ): reason = "dot not installed" if reason: From 275c6c60dd5869d684c31f56db287fdc35d2c6d4 Mon Sep 17 00:00:00 2001 From: xadupre Date: Wed, 11 Jun 2025 09:52:09 +0200 Subject: [PATCH 4/9] fix doc --- _doc/examples/plot_dump_intermediate_results.py | 3 ++- onnx_diagnostic/helpers/torch_helper.py | 5 ++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/_doc/examples/plot_dump_intermediate_results.py b/_doc/examples/plot_dump_intermediate_results.py index dbf31e96..683a9e15 100644 --- a/_doc/examples/plot_dump_intermediate_results.py +++ b/_doc/examples/plot_dump_intermediate_results.py @@ -63,7 +63,7 @@ # All the intermediate tensors were saved in one unique onnx model, # every tensor is stored in a constant node. # The model can be run with any runtime to restore the inputs -# and function :func:`onnx_diagnostic.mini_onnx_builder.create_input_tensors_from_onnx_model` +# and function :func:`onnx_diagnostic.helpers.mini_onnx_builder.create_input_tensors_from_onnx_model` # can restore their names. saved_tensors = create_input_tensors_from_onnx_model( @@ -76,6 +76,7 @@ # Let's explained the naming convention. # # :: +# # ('model.decoder.norm_2', 0, 'I') -- ((T1s2x30x16,),{}) # | | | # | | +--> input, the format is args, kwargs diff --git a/onnx_diagnostic/helpers/torch_helper.py b/onnx_diagnostic/helpers/torch_helper.py index 10606293..83ca95a4 100644 --- a/onnx_diagnostic/helpers/torch_helper.py +++ b/onnx_diagnostic/helpers/torch_helper.py @@ -411,12 +411,15 @@ def forward(self, x, y): proto = create_onnx_model_from_input_tensors(storage) if verbose: print("-- dumps stored objects") + location = f"{os.path.split(dump_file)[-1]}.data" + if os.path.exists(location): + os.remove(location) onnx.save( proto, dump_file, save_as_external_data=True, all_tensors_to_one_file=True, - location=f"{os.path.split(dump_file)[-1]}.data", + location=location, ) if verbose: print("-- done dump stored objects") From c2ab5f687292a9f5fab19322ab1d38e9e869f275 Mon Sep 17 00:00:00 2001 From: xadupre Date: Wed, 11 Jun 2025 09:55:05 +0200 Subject: [PATCH 5/9] black --- _doc/examples/plot_dump_intermediate_results.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/_doc/examples/plot_dump_intermediate_results.py b/_doc/examples/plot_dump_intermediate_results.py index 683a9e15..cd477e59 100644 --- a/_doc/examples/plot_dump_intermediate_results.py +++ b/_doc/examples/plot_dump_intermediate_results.py @@ -63,7 +63,8 @@ # All the intermediate tensors were saved in one unique onnx model, # every tensor is stored in a constant node. # The model can be run with any runtime to restore the inputs -# and function :func:`onnx_diagnostic.helpers.mini_onnx_builder.create_input_tensors_from_onnx_model` +# and function :func:`create_input_tensors_from_onnx_model +# ` # can restore their names. saved_tensors = create_input_tensors_from_onnx_model( From ccfa16c68e541a220188a1797e6eda35835cf751 Mon Sep 17 00:00:00 2001 From: xadupre Date: Wed, 11 Jun 2025 10:05:32 +0200 Subject: [PATCH 6/9] doc --- .../plot_dump_intermediate_results.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/_doc/examples/plot_dump_intermediate_results.py b/_doc/examples/plot_dump_intermediate_results.py index cd477e59..2ad7c8c4 100644 --- a/_doc/examples/plot_dump_intermediate_results.py +++ b/_doc/examples/plot_dump_intermediate_results.py @@ -4,9 +4,21 @@ Dumps intermediate results of a torch model =========================================== - -codellama/CodeLlama-7b-Python-hf -++++++++++++++++++++++++++++++++ +Looking for discrepancies is quickly annoying. Discrepancies +come from two results obtained with the same models +implemented in two different ways, :epkg:`pytorch` and :epkg:`onnx`. +Models are big so where do they come from? That's the +unavoidable question. Unless there is an obious reason, +the only way is to compare intermediate outputs alon the computation. +The first step into that direction is to dump the intermediate results +coming from :epkg:`pytorch`. +We use :func:`onnx_diagnostic.helpers.torch_helper import steal_forward` for that. + +A simple LLM Model +++++++++++++++++++ + +See :func:`onnx_diagnostic.helpers.torch_helper.dummy_llm` +for its definition. It is mostly used for unit test or example. """ @@ -87,7 +99,7 @@ # | it is possible to call multiple times, # | the model to store more # | -# +--> the name given to steal forward +# +--> the name given to function steal_forward # # The same goes for output except ``'I'`` is replaced by ``'O'``. # From 588a30f02c35f47902e3d48f858ea2c499efd0c9 Mon Sep 17 00:00:00 2001 From: xadupre Date: Wed, 11 Jun 2025 10:58:02 +0200 Subject: [PATCH 7/9] fix doc --- _unittests/ut_xrun_doc/test_documentation_recipes.py | 3 +++ _unittests/ut_xrun_doc/test_documentation_technical.py | 7 +------ 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/_unittests/ut_xrun_doc/test_documentation_recipes.py b/_unittests/ut_xrun_doc/test_documentation_recipes.py index cef7e9b4..59f1b682 100644 --- a/_unittests/ut_xrun_doc/test_documentation_recipes.py +++ b/_unittests/ut_xrun_doc/test_documentation_recipes.py @@ -73,6 +73,7 @@ def add_test_methods(cls): this = os.path.abspath(os.path.dirname(__file__)) fold = os.path.normpath(os.path.join(this, "..", "..", "_doc", "recipes")) found = os.listdir(fold) + has_dot = int(os.environ.get("UNITTEST_DOT", "0")) for name in found: if not name.endswith(".py") or not name.startswith("plot_"): continue @@ -80,6 +81,8 @@ def add_test_methods(cls): if not reason and not has_torch("4.7"): reason = "torch<2.7" + if not reason and not has_dot and name in {"plot_dump_intermediate_results.py"}: + reason = "dot not installed" if reason: diff --git a/_unittests/ut_xrun_doc/test_documentation_technical.py b/_unittests/ut_xrun_doc/test_documentation_technical.py index e7fa2661..5dbfb661 100644 --- a/_unittests/ut_xrun_doc/test_documentation_technical.py +++ b/_unittests/ut_xrun_doc/test_documentation_technical.py @@ -73,12 +73,7 @@ def add_test_methods(cls): if not name.endswith(".py") or not name.startswith("plot_"): continue reason = None - if ( - not reason - and not has_dot - and name - in {"plot_layer_norm_discrepancies.py", "plot_dump_intermediate_results.py"} - ): + if not reason and not has_dot and name in {"plot_layer_norm_discrepancies.py"}: reason = "dot not installed" if reason: From 8de38e204913381d884bf95e0d2d7b6461f7ad86 Mon Sep 17 00:00:00 2001 From: xadupre Date: Wed, 11 Jun 2025 11:07:33 +0200 Subject: [PATCH 8/9] spell --- _doc/examples/plot_dump_intermediate_results.py | 2 +- _unittests/ut_xrun_doc/test_documentation_examples.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/_doc/examples/plot_dump_intermediate_results.py b/_doc/examples/plot_dump_intermediate_results.py index 2ad7c8c4..8b17cbb7 100644 --- a/_doc/examples/plot_dump_intermediate_results.py +++ b/_doc/examples/plot_dump_intermediate_results.py @@ -8,7 +8,7 @@ come from two results obtained with the same models implemented in two different ways, :epkg:`pytorch` and :epkg:`onnx`. Models are big so where do they come from? That's the -unavoidable question. Unless there is an obious reason, +unavoidable question. Unless there is an obvious reason, the only way is to compare intermediate outputs alon the computation. The first step into that direction is to dump the intermediate results coming from :epkg:`pytorch`. diff --git a/_unittests/ut_xrun_doc/test_documentation_examples.py b/_unittests/ut_xrun_doc/test_documentation_examples.py index 2382e721..3cfd62a4 100644 --- a/_unittests/ut_xrun_doc/test_documentation_examples.py +++ b/_unittests/ut_xrun_doc/test_documentation_examples.py @@ -74,11 +74,15 @@ def add_test_methods(cls): this = os.path.abspath(os.path.dirname(__file__)) fold = os.path.normpath(os.path.join(this, "..", "..", "_doc", "examples")) found = os.listdir(fold) + has_dot = int(os.environ.get("UNITTEST_DOT", "0")) for name in found: if not name.endswith(".py") or not name.startswith("plot_"): continue reason = None + if not reason and not has_dot and name in {"plot_dump_intermediate_results.py"}: + reason = "dot not installed" + if ( not reason and name in {"plot_export_tiny_llm.py"} From 1d139dfcc5cfca252c5502419a2f1039445535c5 Mon Sep 17 00:00:00 2001 From: xadupre Date: Wed, 11 Jun 2025 11:24:47 +0200 Subject: [PATCH 9/9] fix --- _doc/examples/plot_dump_intermediate_results.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/_doc/examples/plot_dump_intermediate_results.py b/_doc/examples/plot_dump_intermediate_results.py index 8b17cbb7..30d784c2 100644 --- a/_doc/examples/plot_dump_intermediate_results.py +++ b/_doc/examples/plot_dump_intermediate_results.py @@ -12,7 +12,7 @@ the only way is to compare intermediate outputs alon the computation. The first step into that direction is to dump the intermediate results coming from :epkg:`pytorch`. -We use :func:`onnx_diagnostic.helpers.torch_helper import steal_forward` for that. +We use :func:`onnx_diagnostic.helpers.torch_helper.steal_forward` for that. A simple LLM Model ++++++++++++++++++