improve ci

Author: xadupre

.github/workflows/ci.yml

@@ -110,9 +110,6 @@ jobs:
             pip install torch==${{ matrix.torch }} torchvision torchaudio
           fi
 
-      - name: Uninstall triton
-        run: python -m pip uninstall -y triton
-
       - name: Cache pip
         uses: actions/cache@v4
         with:

onnx_diagnostic/_command_lines_parser.py

@@ -553,6 +553,12 @@ def get_parser_validate() -> ArgumentParser:
         action=BooleanOptionalAction,
         help="Enables onnxruntime logging when the session is created",
     )
+    parser.add_argument(
+        "--quiet-input-sets",
+        default="",
+        help="Avoids raising an exception when an input sets does not work with "
+        "the exported model, example: --quiet-input-sets=inputs,inputs22",
+    )
     return parser
 
 
@@ -614,6 +620,7 @@ def _cmd_validate(argv: List[Any]):
             warmup=args.warmup,
             inputs2=args.inputs2,
             ort_logs=args.ort_logs,
+            quiet_input_sets=set(args.quiet_input_sets.split(",")),
             output_names=(
                 None if len(args.outnames.strip()) < 2 else args.outnames.strip().split(",")
             ),

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -12,17 +12,26 @@
     StaticCache,
 )
 
-try:
-    from transformers.models.mamba.modeling_mamba import MambaCache
-except ImportError:
-    from transformers.cache_utils import MambaCache
-
 from ..helpers import string_type
 from .serialization import _lower_name_with_
 
 PATCH_OF_PATCHES: Set[Any] = set()
 
 
+def get_mamba_cache_cls() -> type:
+    try:
+        from transformers.models.mamba.modeling_mamba import MambaCache
+
+        return MambaCache
+    except ImportError:
+        try:
+            from transformers.cache_utils import MambaCache
+
+            return MambaCache
+        except ImportError:
+            return None
+
+
 def register_class_serialization(
     cls,
     f_flatten: Callable,
@@ -203,13 +212,6 @@ def serialization_functions(
                 # f_check=make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]),
                 verbose=verbose,
             ),
-            MambaCache: lambda verbose=verbose: register_class_serialization(
-                MambaCache,
-                flatten_mamba_cache,
-                unflatten_mamba_cache,
-                flatten_with_keys_mamba_cache,
-                verbose=verbose,
-            ),
             EncoderDecoderCache: lambda verbose=verbose: register_class_serialization(
                 EncoderDecoderCache,
                 flatten_encoder_decoder_cache,
@@ -232,6 +234,17 @@ def serialization_functions(
                 verbose=verbose,
             ),
         }
+        MambaCache = get_mamba_cache_cls()
+        if MambaCache:
+            transformers_classes[MambaCache] = (
+                lambda verbose=verbose: register_class_serialization(
+                    MambaCache,
+                    flatten_mamba_cache,
+                    unflatten_mamba_cache,
+                    flatten_with_keys_mamba_cache,
+                    verbose=verbose,
+                )
+            )
         classes.update(transformers_classes)
 
     if patch_diffusers:
@@ -287,7 +300,12 @@ def unregister_class_serialization(cls: type, verbose: int = 0):
 
 def unregister_cache_serialization(undo: Dict[str, bool], verbose: int = 0):
     """Undo all registrations."""
-    cls_ensemble = {MambaCache, DynamicCache, EncoderDecoderCache} | set(undo)
+    MambaCache = get_mamba_cache_cls()
+    cls_ensemble = (
+        {DynamicCache, EncoderDecoderCache}
+        | set(undo)
+        | ({MambaCache} if MambaCache else set())
+    )
     for cls in cls_ensemble:
         if undo.get(cls.__name__, False):
             unregister_class_serialization(cls, verbose)

onnx_diagnostic/torch_models/validate.py

@@ -1438,6 +1438,8 @@ def _mk(key, flavour=flavour):
     keys = [("inputs", "run_expected", "")]
     if second_input_keys:
         keys.extend([(k, f"run_expected2{k[6:]}", f"2{k[6:]}") for k in second_input_keys])
+    if verbose:
+        print(f"[validate_onnx_model] -- keys={keys}")
     for k_input, k_expected, suffix in keys:
         # make_feeds
         assert k_input in data, f"Unable to find {k_input!r} in {sorted(data)}"
@@ -1463,7 +1465,7 @@ def _mk(key, flavour=flavour):
         # run ort
         if verbose:
             print(f"[validate_onnx_model] run session on inputs 'inputs{suffix}'...")
-            if quiet_input_sets:
+            if quiet_input_sets and f"inputs{suffix}" in quiet_input_sets:
                 print(f"[validate_onnx_model] quiet_input_sets={quiet_input_sets}")
 
         got = _quiet_or_not_quiet(