Refactors sbs to save memory (#316)

* Refactors sbs to save memory

* add

* a few changes

* fix a few things

* fix ut

* remove doc

* fix import issues

* disable

* mypy

Author: xadupre

CHANGELOGS.rst

@@ -8,7 +8,7 @@ Change Logs
 * :pr:`311`: use custom and local function to use PackedMultiHeadAttention from onnxruntime
 * :pr:`310`: splits patches into multiple files 
 * :pr:`308`: add option --save_ep to dump the exported program as well as torch input
-* :pr:`304`, :pr:`306`: improves side-by-side comparison, creates command line sbs
+* :pr:`304`, :pr:`306`, :pr:`316`: improves side-by-side comparison, creates command line sbs
 
 0.8.2
 +++++
@@ -112,8 +112,7 @@ Change Logs
 * :pr:`203`: Add option to disable patches for torch in command line validate
 * :pr:`202`: add models DeepseekV3ForCausalLM, Gemma3ForCausalLM, Glm4vMoeForConditionalGeneration
 * :pr:`201`: switch CI to 4.55.4
-* :pr:`200`: fixes patches for 4.55.1+, DynamicCache is no longer registered by default,
-  this code moved to executorch.py in transformers
+* :pr:`200`: fixes patches for 4.55.1+, DynamicCache is no longer registered by default, this code moved to executorch.py in transformers
 * :pr:`199`: delete hidden_size and num_attention_heads modification in a config
 * :pr:`198`: support gpt-oss
 * :pr:`197`: updates CI for torch 2.8
@@ -124,15 +123,13 @@ Change Logs
 
 * :pr:`193`: validates with 4.53.3 
 * :pr:`189`: support for task mask-generation
-* :pr:`192`: add support for Gemma-3, add serialization for HybridCache,
-  changes to support ``transformers>=4.54``
+* :pr:`192`: add support for Gemma-3, add serialization for HybridCache, changes to support ``transformers>=4.54``
 
 0.7.5
 +++++
 
 * :pr:`186`: add parameter --output_names to command line validate to change the output names of the onnx exported model
-* :pr:`185`: remove the use of _seen_tokens in DynamicCache (removed in transformers>4.53),
-  updates dummpy inputs for feature-extraction
+* :pr:`185`: remove the use of _seen_tokens in DynamicCache (removed in ``transformers>4.53``), updates dummpy inputs for feature-extraction
 * :pr:`184`: implements side-by-side
 
 0.7.4
@@ -172,12 +169,8 @@ Change Logs
 * :pr:`147`: simplified log processing
 * :pr:`146`: patch for IdeficsAttention, IdeficsEmbedding
 * :pr:`145`: patch for _compute_dynamic_ntk_parameters (Phi3RotaryEmbedding)
-* :pr:`144`: support for second inputs with different dimension,
-  rename test_helper into validate,
-  support ``interpolate_pos_encoding`` for ``VitModel``,
-  update model builder helpers for this PR
-  `Use ONNX IR for model builder
-  <https://github.com/microsoft/onnxruntime-genai/pull/1416>`_
+* :pr:`144`: support for second inputs with different dimension, rename test_helper into validate, support ``interpolate_pos_encoding`` for ``VitModel``, update model builder helpers for this PR
+  `Use ONNX IR for model builder <https://github.com/microsoft/onnxruntime-genai/pull/1416>`_
 * :pr:`143`: compares intermediate results,
 
 0.6.3
@@ -199,8 +192,7 @@ Change Logs
 * :pr:`123`: add subgraphs to TorchOnnxEvaluator
 * :pr:`122`: add local functions to TorchOnnxEvaluator
 * :pr:`120`: enables TorchOnnxEvaluator in command line ``python -m onnx_diagnostic validate ...``
-* :pr:`115`, :pr:`116`, :pr:`117`, :pr:`118`, :pr:`119`, :pr:`127`:
-  first steps for TorchOnnxEvaluator
+* :pr:`115`, :pr:`116`, :pr:`117`, :pr:`118`, :pr:`119`, :pr:`127`: first steps for TorchOnnxEvaluator
 * :pr:`114`: extends the list of known rewritings
 * :pr:`113`: fixes a couple of issues with ModelBuilder
 
@@ -257,10 +249,7 @@ Change Logs
 * :pr:`65`: support SlidingWindowCache
 * :pr:`63`: support option ``--trained``
 * :pr:`61`: improves dynamic shapes for EncoderDecoderCache
-* :pr:`58`: add function use_dyn_not_str to replace string by ``torch.export.Dim.DYNAMIC``,
-  use string instead of ``torch.export.Dim.DYNAMIC`` when returning the dynamic shapes
-  for a specific models, it is a valid definition for ``torch.onnx.export``
-  which can reuse the names
+* :pr:`58`: add function use_dyn_not_str to replace string by ``torch.export.Dim.DYNAMIC``, use string instead of ``torch.export.Dim.DYNAMIC`` when returning the dynamic shapes for a specific models, it is a valid definition for ``torch.onnx.export`` which can reuse the names
 * :pr:`55`: add support for text-classification
 * :pr:`54`: add support for fill-mask, refactoring
 * :pr:`52`: add support for zero-shot-image-classification
@@ -274,28 +263,18 @@ Change Logs
 * :pr:`43`: uses custom patches
 * :pr:`38`: uses the registered serialization functions when it is available
 * :pr:`30`, :pr:`31`: adds command to test a model id, validate the export
-* :pr:`29`: adds helpers to measure the memory peak and run benchmark
-  on different processes
-* :pr:`28`: adds command line to print out the configuration for a model id,
-  support image-text-to-text
-* :pr:`26`: creates a folder ``helpers`` to gather all the functions
-  used in many places
-* :pr:`25`: improve patches for DynamicCache
-  (issue with register_pytree_flatten_spec being deprecated)
-* :pr:`24`: dummy inputs for ``text2text-generation``, add new function
-  ``convert_dynamic_axes_into_dynamic_shapes`` to convert dynamic axes
-  into dynamic shapes, add support for ``T5ForConditionalGeneration``
+* :pr:`29`: adds helpers to measure the memory peak and run benchmark on different processes
+* :pr:`28`: adds command line to print out the configuration for a model id, support image-text-to-text
+* :pr:`26`: creates a folder ``helpers`` to gather all the functions used in many places
+* :pr:`25`: improve patches for DynamicCache (issue with register_pytree_flatten_spec being deprecated)
+* :pr:`24`: dummy inputs for ``text2text-generation``, add new function ``convert_dynamic_axes_into_dynamic_shapes`` to convert dynamic axes into dynamic shapes, add support for ``T5ForConditionalGeneration``
 * :pr:`23`: dummy inputs for ``image-classification``
-* :pr:`22`, :pr:`27`: api to create untrained model copying the architecture
-  of the trained models and dummy inputs for them,
-  support for ``text-generation``
+* :pr:`22`, :pr:`27`: api to create untrained model copying the architecture of the trained models and dummy inputs for them, support for ``text-generation``
 
 0.2.1
 +++++
 
-* :pr:`16`: refactors patches, add model Phi2, implements
-  a tweak to raise an exception with a dynamic dimension
-  becomes static when exporting a model
+* :pr:`16`: refactors patches, add model Phi2, implements a tweak to raise an exception with a dynamic dimension becomes static when exporting a model
 
 0.2.0
 +++++

_doc/cmds/index.rst

@@ -9,4 +9,5 @@ Command Lines
     :maxdepth: 1
 
     config
+    sbs
     validate

_doc/cmds/sbs.rst

@@ -0,0 +1,22 @@
+-m onnx_diagnostic sbs ... runs a side-by-side torch/onnx
+=========================================================
+
+Description
++++++++++++
+
+It compares the intermediate results between an exported program saved with
+:func:`torch.export.save` and an exported model on saved inputs
+with :func:`torch.save`. It assumes intermediate results share the same
+names.
+
+.. runpython::
+
+    from onnx_diagnostic._command_lines_parser import get_parser_sbs
+
+    get_parser_sbs().print_help()
+
+CPU, CUDA
++++++++++
+
+Inputs are saved :func:`torch.save`. The execution will run on CUDA
+if the device of the inputs is CUDA, same goes on CPU.

_unittests/ut_torch_onnx/test_sbs.py

@@ -29,15 +29,15 @@ def test_run_aligned_record(self):
             onnx_name="B",
             ep_target="C",
             onnx_op_type="D",
-            shape_type="E",
+            ep_shape_type="E",
             err_abs=0.1,
             err_rel=0.2,
             err_dev=0.3,
             err_nan=0.4,
         )
         sr = str(r)
         self.assertIn("RunAlignedRecord(", sr)
-        self.assertIn("shape_type='E'", sr)
+        self.assertIn("ep_shape_type='E'", sr)
 
     @hide_stdout()
     @unittest.skipIf(to_onnx is None, "to_onnx not installed")
@@ -303,8 +303,8 @@ def forward(self, x):
         )
         self.assertEqual(len(results), 14)
         self.assertEqual(
-            [r.err_dev for r in results],
             [None, None, None, None, None, None, None, None, 0, 0, 0, 0, 0, 0],
+            [r.err_dev for r in results],
         )
 
     @hide_stdout()
@@ -349,29 +349,27 @@ def forward(self, x):
             [
                 "ep_id_node",
                 "ep_name",
+                "ep_shape_type",
                 "ep_target",
                 "ep_time_run",
                 "err_abs",
                 "err_dev",
+                "err_h01",
                 "err_nan",
                 "err_rel",
                 "onnx_id_node",
                 "onnx_id_output",
                 "onnx_name",
                 "onnx_op_type",
+                "onnx_shape_type",
                 "onnx_time_run",
-                "shape_type",
             ],
             sorted(df.columns),
         )
-        self.assertEqual(len(results), 12)
-        self.assertEqual(
-            [r.err_dev for r in results],
-            [None, None, None, None, None, None, None, None, None, 0, 0, 0],
-        )
+        self.assertEqual(len(results), 8)
+        self.assertEqual([0, 0, 0, 0, None, 0, 0, 0], [r.err_dev for r in results])
         self.assertEqual(
-            [-1.0, -1.0, -1.0, -1.0, -10.0, -10.0, -10.0, -10.0, -1.0, 0.0, 1.0, 2.0],
-            df["onnx_id_node"].fillna(-10).tolist(),
+            [-1, -1, -1, -1, -1, 0, 1, 2], df["onnx_id_node"].fillna(-10).tolist()
         )
         self.clean_dump()
 
@@ -417,29 +415,27 @@ def forward(self, x):
             [
                 "ep_id_node",
                 "ep_name",
+                "ep_shape_type",
                 "ep_target",
                 "ep_time_run",
                 "err_abs",
                 "err_dev",
+                "err_h01",
                 "err_nan",
                 "err_rel",
                 "onnx_id_node",
                 "onnx_id_output",
                 "onnx_name",
                 "onnx_op_type",
+                "onnx_shape_type",
                 "onnx_time_run",
-                "shape_type",
             ],
             sorted(df.columns),
         )
-        self.assertEqual(len(results), 12)
-        self.assertEqual(
-            [r.err_dev for r in results],
-            [None, None, None, None, None, None, None, None, None, 0, 0, 0],
-        )
+        self.assertEqual(len(results), 8)
+        self.assertEqual([0, 0, 0, 0, None, 0, 0, 0], [r.err_dev for r in results])
         self.assertEqual(
-            [-1.0, -1.0, -1.0, -1.0, -10.0, -10.0, -10.0, -10.0, -1.0, 0.0, 1.0, 2.0],
-            df["onnx_id_node"].fillna(-10).tolist(),
+            [-1, -1, -1, -1, -1, 0, 1, 2], df["onnx_id_node"].fillna(-10).tolist()
         )
         self.clean_dump()
 

_unittests/ut_xrun_doc/test_command_lines_exe.py

@@ -2,10 +2,12 @@
 import unittest
 from contextlib import redirect_stdout
 from io import StringIO
+import pandas
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings
+from onnx_diagnostic.ext_test_case import ExtTestCase, ignore_warnings, requires_transformers
 from onnx_diagnostic._command_lines_parser import main
 from onnx_diagnostic.helpers.log_helper import enumerate_csv_files
+from onnx_diagnostic.export.api import to_onnx
 
 
 class TestCommandLines(ExtTestCase):
@@ -88,6 +90,73 @@ def test_g_parser_agg(self):
         self.assertIn("[CubeLogs.to_excel] plots 1 plots", text)
         self.assertExists(output)
 
+    @ignore_warnings(UserWarning)
+    @requires_transformers("4.53")
+    def test_h_parser_sbs(self):
+        import torch
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super(Model, self).__init__()
+                self.fc1 = torch.nn.Linear(10, 32)  # input size 10 → hidden size 32
+                self.relu = torch.nn.ReLU()
+                self.fc2 = torch.nn.Linear(32, 1)  # hidden → output
+
+            def forward(self, x):
+                x = self.relu(self.fc1(x))
+                x = self.fc2(x)
+                return x
+
+        inputs = dict(x=torch.randn((5, 10)))
+        ds = dict(x={0: "batch"})
+        input_file = self.get_dump_file("test_h_parser_sbs.inputs.pt")
+        ep_file = self.get_dump_file("test_h_parser_sbs.ep")
+        onnx_file = self.get_dump_file("test_h_parser_sbs.model.onnx")
+        torch.save(inputs, input_file)
+        to_onnx(
+            Model(),
+            kwargs=inputs,
+            dynamic_shapes=ds,
+            exporter="custom",
+            save_ep=(ep_file, 2**30),
+            filename=onnx_file,
+        )
+
+        output = self.get_dump_file("test_h_parser_sbs.xlsx")
+        st = StringIO()
+        with redirect_stdout(st):
+            main(
+                [
+                    "sbs",
+                    "-v",
+                    "2",
+                    "--first",
+                    "-i",
+                    input_file,
+                    "-e",
+                    f"{ep_file}.ep.pt2",
+                    "-o",
+                    output,
+                    "-m",
+                    onnx_file,
+                ]
+            )
+        text = st.getvalue()
+        self.assertIn("[run_aligned", text)
+        self.assertExists(output)
+        df = pandas.read_excel(output).apply(
+            lambda col: col.fillna("") if col.dtype == "object" else col
+        )
+        self.assertLess(df["err_abs"].max(), 1e-5)
+        self.assertEqual(df["err_h01"].max(), 0)
+        self.assertIn("p_fc1_weight", set(df["ep_name"]))
+        self.assertIn("fc1.bias", set(df["onnx_name"]))
+        self.assertNotIn("NaN", set(df["ep_name"]))
+        # print(f"{df}\n{st.getvalue()}")
+        self.assertIn("[run_aligned] done", st.getvalue())
+        sdf = df[(df.ep_target == "placeholder") & (df.onnx_op_type == "initializer")]
+        self.assertEqual(sdf.shape[0], 4)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

clean_onnx.sh

@@ -30,7 +30,8 @@ rm _plot_torch_sklearn_201_knnpy.py
 
 rm _doc/sg_execution_times.rst
 
-rm _doc/examples/plot*.onnx
+rm _doc/examples/_debug*
+rm _doc/examples/plot*.onnx*
 rm _doc/examples/plot*.txt
 rm _doc/examples/ort*.onnx
 rm _doc/examples/*.sarif
@@ -83,6 +84,7 @@ rm _doc/technical/*.dynamo.onnx
 rm _doc/technical/*.script.onnx
 rm _doc/technical/dump_models -rf
 rm _doc/technical/dump_onx_*
+rm _doc/technical/model_*.onnx* -rf
 
 rm _tools/bin -rf
 rm _tools/mambaroot -rf