Fix patches for Qwen2_5 (#294)

* Fix patch for Qwen2_5

* doc

* fix

* fix

* int2

* attention

* fix ut

* fix

* fix

* fix

* fix aggregation

* mypy

* fix

* lint

* fix

* revert

* stop

* lint

Author: xadupre

.github/workflows/ci.yml

@@ -149,10 +149,12 @@ jobs:
       - name: run tests
         run: |
           pip install pytest
-          PYTHONPATH=. UNITTEST_GOING=1 pytest --durations=10 _unittests --ignore _unittests/ut_reference/test_backend_extended_reference_evaluator.py --ignore _unittests/ut_reference/test_backend_onnxruntime_evaluator.py --ignore _unittests/ut_torch_models/test_tiny_llms_bypassed.py --ignore _unittests/ut_tasks/test_tasks_zero_shot_image_classification.py --ignore _unittests/ut_tasks/test_tasks_image_classification.py --ignore _unittests/ut_torch_models/test_validate_whole_models*.py
+          PYTHONPATH=. UNITTEST_GOING=1 pytest --durations=10 _unittests --ignore _unittests/ut_reference/test_backend_extended_reference_evaluator.py --ignore _unittests/ut_reference/test_backend_onnxruntime_evaluator.py --ignore _unittests/ut_torch_models/test_tiny_llms_bypassed.py --ignore _unittests/ut_tasks/test_tasks_zero_shot_image_classification.py --ignore _unittests/ut_tasks/test_tasks_image_classification.py --ignore-glob _unittests/ut_torch_models/test_validate_whole_models*.py --ignore-glob _unittests/ut_torch_models/test_validate_models*.py
 
       - name: test models
         run: |
+          echo "----"
+          PYTHONPATH=. UNITTEST_GOING=1 python _unittests/ut_torch_models/test_validate_models.py
           echo "----"
           PYTHONPATH=. UNITTEST_GOING=1 python _unittests/ut_torch_models/test_validate_whole_models1.py
           echo "----"

CHANGELOGS.rst

@@ -4,8 +4,7 @@ Change Logs
 0.8.2
 +++++
 
-* :pr:`293`: second series of patches
-* :pr:`292`: new patches for Qwen models
+* :pr:`292`, :pr:`293`, :pr:`294`: new patches for Qwen models
 
 0.8.1
 +++++

_unittests/ut_export/test_api.py

@@ -1,6 +1,12 @@
+import os
 import unittest
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    hide_stdout,
+    has_transformers,
+    ignore_warnings,
+)
 from onnx_diagnostic.helpers import max_diff
 from onnx_diagnostic.helpers.torch_helper import torch_deepcopy
 from onnx_diagnostic.helpers.rt_helper import make_feeds
@@ -36,6 +42,7 @@ def forward(self, x, y):
         )
 
     @hide_stdout()
+    @ignore_warnings(FutureWarning)
     def test_tiny_llm_to_onnx(self):
         import onnxruntime
 
@@ -68,6 +75,8 @@ def test_tiny_llm_to_onnx(self):
                         filename=filename,
                     )
         for exporter, filename in filenames.items():
+            if not os.path.exists(filename):
+                continue
             with self.subTest(exporter=f"validate-{exporter}"):
                 sess = onnxruntime.InferenceSession(
                     filename, providers=["CPUExecutionProvider"]
@@ -90,6 +99,8 @@ def test_tiny_llm_to_onnx(self):
 
         expected = model(**torch_deepcopy(problem))
         for exporter, filename in filenames.items():
+            if not os.path.exists(filename):
+                continue
             with self.subTest(exporter=f"full-mask-{exporter}"):
                 sess = onnxruntime.InferenceSession(
                     filename, providers=["CPUExecutionProvider"]

_unittests/ut_helpers/test_helper.py

@@ -342,13 +342,15 @@ def test_size_type_onnx(self):
                     continue
                 if name in {"NAME_FIELD_NUMBER"}:
                     continue
-                if name not in {"STRING", "UINT4", "INT4", "FLOAT4E2M1"}:
+                if name not in {"STRING", "UINT4", "INT4", "FLOAT4E2M1", "INT2", "UINT2"}:
                     size_type(i)
 
                 if name not in {
                     "STRING",
                     "UINT4",
                     "INT4",
+                    "INT2",
+                    "UINT2",
                     "FLOAT4E2M1",
                     "FLOAT8E5M2FNUZ",
                     "FLOAT8E5M2",

_unittests/ut_reference/test_backend_onnxruntime_evaluator.py

@@ -248,6 +248,25 @@ def run_node(cls, node, inputs, device=None, outputs_info=None, **kwargs):
 )
 
 if onnx_opset_version() <= 25:
+    exc = "|".join(
+        [
+            "batchnorm_.*_training",
+            "convinteger_with_padding",
+            "rms_normalization",
+            "rotary_embedding_3d",
+            "rotary_embedding",
+            # cuda,
+            "test_Conv3d_dilated.*_cuda",
+            "test_reduce_.*_empty_set_cuda",
+            "test_reduce_sum_square_.*_expanded_cuda",
+            "test_reduce_l1_.*_expanded_cuda",
+            "test_reduce_l2_.*_expanded_cuda",
+            "test_reduce_log_sum_.*_expanded_cuda",
+        ]
+    )
+    backend_test.exclude(f"({exc})")
+
+if onnx_opset_version() <= 26:
     backend_test.exclude(
         "(deform_conv"
         "|gru"
@@ -261,36 +280,25 @@ def run_node(cls, node, inputs, device=None, outputs_info=None, **kwargs):
         "|layer_normalization.*expanded"
         "|layer_normalization.*expanded"
         "|affine_grid.*expanded"
+        "|test_attention_4d_diff_heads_mask4d_padded_kv.*"
+        "|test_convinteger_with_padding"
         "|test_rnn_seq"
         "|test_roialign_aligned_false"
         "|test_roialign_aligned_true"
         "|test_roialign_mode_max"
+        "|test_rotary_embedding_no_position_ids_rotary_dim.*"
+        "|test_rotary_embedding_with_interleaved_rotary_dim.*"
+        "|test_rotary_embedding_with_rotary_dim*"
         "|test_simple_rnn_batchwise"
         "|test_simple_rnn_defaults"
         "|test_simple_rnn_with_initial_bias"
+        "|test_swish*"
+        "|test_tensorscatter*"
+        "|test_top_k*"
         ")"
     )
 
 
-if onnx_opset_version() <= 25:
-    exc = "|".join(
-        [
-            "batchnorm_.*_training",
-            "convinteger_with_padding",
-            "rms_normalization",
-            "rotary_embedding_3d",
-            "rotary_embedding",
-            # cuda,
-            "test_Conv3d_dilated.*_cuda",
-            "test_reduce_.*_empty_set_cuda",
-            "test_reduce_sum_square_.*_expanded_cuda",
-            "test_reduce_l1_.*_expanded_cuda",
-            "test_reduce_l2_.*_expanded_cuda",
-            "test_reduce_log_sum_.*_expanded_cuda",
-        ]
-    )
-    backend_test.exclude(f"({exc})")
-
 if pv.Version(onnxruntime.__version__) <= pv.Version("1.24"):
     backend_test.exclude("(test_attention_4d_with|test_attention_4d_gqa)")
 

_unittests/ut_tasks/try_tasks.py

@@ -1011,6 +1011,7 @@ def test_imagetext2text_qwen_2_5_vl_instruct(self):
                 return_dict:bool
             )
         """
+        import transformers
         from transformers import AutoModel, AutoProcessor
         from qwen_vl_utils import process_vision_info
 
@@ -1068,6 +1069,20 @@ def config_reduction(config, task):
         print(f"-- processor {type(processor)}")
         print(f"-- inputs={self.string_type(inputs, with_shape=True, with_min_max=True)}")
 
+        f_ = transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.apply_multimodal_rotary_pos_emb
+
+        def _apply_multimodal_rotary_pos_emb(*args, **kwargs):
+            print(
+                "-- apply_multimodal_rotary_pos_emb:",
+                self.string_type(args, with_shape=True),
+                self.string_type(kwargs, with_shape=True),
+            )
+            return f_(*args, **kwargs)
+
+        transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.apply_multimodal_rotary_pos_emb = (
+            _apply_multimodal_rotary_pos_emb
+        )
+
         print()
         with (
             torch_export_patches(
@@ -1077,7 +1092,7 @@ def config_reduction(config, task):
                 verbose=1,
             ),
             steal_forward(
-                [model, model.visual],
+                [model, model.visual, model.visual.blocks[0].attn],
                 dump_file=self.get_dump_file("test_imagetext2text_qwen_2_5_vl_instruct.onnx"),
                 dump_drop={"attention_mask", "past_key_values", "pixel_values"},
                 save_as_external_data=False,

_unittests/ut_torch_export_patches/test_dynamic_class.py

@@ -21,7 +21,6 @@
 )
 from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
 from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
-import onnx_diagnostic.torch_export_patches.patches.patch_transformers as patch_transformers
 
 
 class TestOnnxExportErrors(ExtTestCase):
@@ -344,11 +343,6 @@ def test_phi2_export_interpreter(self):
         )
 
         with torch_export_patches(patch_transformers=True, verbose=1):
-            if masking_utils is not None:
-                self.assertEqual(
-                    masking_utils.ALL_MASK_ATTENTION_FUNCTIONS["sdpa"],
-                    patch_transformers.patched_sdpa_mask_recent_torch,
-                )
             ep = torch.export.export(
                 model,
                 (),