last changes

Author: xadupre

_unittests/ut_tasks/try_export.py

@@ -149,6 +149,7 @@ def _config_reduction(config, task):
         begin = time.perf_counter()
         if not os.environ.get("STOPAT", ""):
             expected = model_to_export(**inputs)
+            expected_big = model_to_export(**big_inputs)
         else:
             expected = None
         print(f"-- MODEL RUN IN {time.perf_counter() - begin}")
@@ -266,7 +267,7 @@ def _config_reduction(config, task):
                     (f"test_qwen25_vli_visual.{device}.{dtype}.{attention}.{exporter}"),
                     filename,
                     model_to_export,
-                    export_inputs,
+                    big_inputs,  # export_inputs,
                     verbose=1,
                     providers=(
                         ["CUDAExecutionProvider", "CPUExecutionProvider"]
@@ -277,7 +278,9 @@ def _config_reduction(config, task):
                     atol=0.05,
                     rtol=10,
                     # ep=pt2_file,
-                    expected=expected,
+                    expected=expected_big,
+                    log_severity_level=0,
+                    log_verbosity_level=0,
                 )
                 print(f"-- MODEL VERIFIED IN {time.perf_counter() - begin}")
         os.environ["QWEN25ATTENTION"] = qwen25_attention

onnx_diagnostic/ext_test_case.py

@@ -1218,7 +1218,7 @@ def tryCall(
     def assert_onnx_disc(
         self,
         test_name: str,
-        proto: "onnx.ModelProto",  # noqa: F821
+        proto: Union[str, "onnx.ModelProto"],  # noqa: F821
         model: "torch.nn.Module",  # noqa: F821
         inputs: Union[Tuple[Any], Dict[str, Any]],
         verbose: int = 0,
@@ -1264,7 +1264,9 @@ def assert_onnx_disc(
             name = f"{test_name}.onnx"
             if verbose:
                 print(f"[{vname}] save the onnx model into {name!r}")
+            model_file = None
             if isinstance(proto, str):
+                model_file = proto
                 name = proto
                 proto = onnx.load(name)
             elif not self.unit_test_going():
@@ -1287,11 +1289,15 @@ def assert_onnx_disc(
             options = onnxruntime.SessionOptions()
             if ort_optimized_graph:
                 options.optimized_model_filepath = f"{name}.optort.onnx"
+            if "log_severity_level" in kwargs:
+                options.log_severity_level = kwargs["log_severity_level"]
+            if "log_verbosity_level" in kwargs:
+                options.log_verbosity_level = kwargs["log_verbosity_level"]
             providers = kwargs.get("providers", ["CPUExecutionProvider"])
             if verbose:
                 print(f"[{vname}] create onnxruntime.InferenceSession with {providers}")
             sess = onnxruntime.InferenceSession(
-                proto.SerializeToString(), options, providers=providers
+                model_file or proto.SerializeToString(), options, providers=providers
             )
             if verbose:
                 print(f"[{vname}] run ort feeds {string_type(feeds, **kws)}")

onnx_diagnostic/torch_export_patches/patches/_patch_transformers_qwen2_5.py

@@ -538,8 +538,6 @@ def forward(
                     break
 
             hidden_states = self.merger(hidden_states)
-            if STOPAT is not None:
-                return hidden_states
             reverse_indices = torch.argsort(window_index)
             hidden_states = hidden_states[reverse_indices, :]
             return hidden_states