Add option to disable patches for torch in command line validate

.github/workflows/ci.yml

@@ -101,28 +101,28 @@ jobs:
         run: python -m pip freeze
 
       - name: tiny-llm torch.export.export
-        run: PYTHONPATH=. python _unittests/ut_torch_models/test_tiny_llms.py
+        run: PYTHONPATH=. UNITTEST_GOING=1 python _unittests/ut_torch_models/test_tiny_llms.py
 
       - name: tiny-llm onnx
-        run: PYTHONPATH=. python _unittests/ut_torch_models/test_tiny_llms_onnx.py
+        run: PYTHONPATH=. UNITTEST_GOING=1 python _unittests/ut_torch_models/test_tiny_llms_onnx.py
         continue-on-error: true  # connectivity issues
 
       - name: tiny-llm example
-        run: PYTHONPATH=. python _doc/examples/plot_export_tiny_llm.py
+        run: PYTHONPATH=. UNITTEST_GOING=1 python _doc/examples/plot_export_tiny_llm.py
         continue-on-error: true  # connectivity issues
 
       - name: tiny-llm bypass
-        run: PYTHONPATH=. python _doc/examples/plot_export_tiny_llm_patched.py
+        run: PYTHONPATH=. UNITTEST_GOING=1 python _doc/examples/plot_export_tiny_llm_patched.py
         continue-on-error: true  # connectivity issues
 
       - name: run tests bypassed
-        run: PYTHONPATH=. python _unittests/ut_torch_models/test_tiny_llms_bypassed.py
+        run: PYTHONPATH=. UNITTEST_GOING=1 python _unittests/ut_torch_models/test_tiny_llms_bypassed.py
 
       - name: test image_classification
-        run: PYTHONPATH=. python _unittests/ut_tasks/test_tasks_image_classification.py
+        run: PYTHONPATH=. UNITTEST_GOING=1 python _unittests/ut_tasks/test_tasks_image_classification.py
 
       - name: test zero_shot_image_classification
-        run: PYTHONPATH=. python _unittests/ut_tasks/test_tasks_zero_shot_image_classification.py
+        run: PYTHONPATH=. UNITTEST_GOING=1 python _unittests/ut_tasks/test_tasks_zero_shot_image_classification.py
 
       - name: run tests
         run: |

CHANGELOGS.rst

@@ -4,7 +4,11 @@ Change Logs
 0.7.7
 +++++
 
-
+* :pr:`205`: add in_channels in image_text_to_text
+* :pr:`204`: switch default num_hidden_layers to 4
+* :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:`199`: delete hidden_size and num_attention_heads modification in a config

_unittests/ut_tasks/test_tasks.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 import torch
 from onnx_diagnostic.ext_test_case import (
@@ -13,6 +14,11 @@
 
 
 class TestTasks(ExtTestCase):
+    def test_unittest_going(self):
+        assert (
+            os.environ.get("UNITTEST_GOING", "0") == "1"
+        ), "UNITTEST_GOING=1 must be defined for these tests"
+
     @hide_stdout()
     def test_text2text_generation(self):
         mid = "sshleifer/tiny-marian-en-de"

onnx_diagnostic/_command_lines_parser.py

@@ -306,7 +306,7 @@ def __call__(self, parser, namespace, values, option_string=None):
                 value = split_items[1]
 
                 if value in ("True", "true", "False", "false"):
-                    d[key] = bool(value)
+                    d[key] = value in ("True", "true")
                     continue
                 try:
                     d[key] = int(value)
@@ -323,6 +323,54 @@ def __call__(self, parser, namespace, values, option_string=None):
         setattr(namespace, self.dest, d)
 
 
+class _BoolOrParseDictPatch(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None):
+
+        if not values:
+            return
+        if len(values) == 1 and values[0] in (
+            "True",
+            "False",
+            "true",
+            "false",
+            "0",
+            "1",
+            0,
+            1,
+        ):
+            setattr(namespace, self.dest, values[0] in ("True", "true", 1, "1"))
+            return
+        d = getattr(namespace, self.dest) or {}
+        if not isinstance(d, dict):
+            d = {
+                "patch_sympy": d,
+                "patch_torch": d,
+                "patch_transformers": d,
+                "patch_diffusers": d,
+            }
+        for item in values:
+            split_items = item.split("=", 1)
+            key = split_items[0].strip()  # we remove blanks around keys, as is logical
+            value = split_items[1]
+
+            if value in ("True", "true", "False", "false"):
+                d[key] = value in ("True", "true")
+                continue
+            try:
+                d[key] = int(value)
+                continue
+            except (TypeError, ValueError):
+                pass
+            try:
+                d[key] = float(value)
+                continue
+            except (TypeError, ValueError):
+                pass
+            d[key] = _parse_json(value)
+
+        setattr(namespace, self.dest, d)
+
+
 def get_parser_validate() -> ArgumentParser:
     parser = ArgumentParser(
         prog="validate",
@@ -383,8 +431,13 @@ def get_parser_validate() -> ArgumentParser:
     parser.add_argument(
         "--patch",
         default=True,
-        action=BooleanOptionalAction,
-        help="Applies patches before exporting.",
+        action=_BoolOrParseDictPatch,
+        nargs="*",
+        help="Applies patches before exporting, it can be a boolean "
+        "to enable to disable the patches or be more finetuned. It is possible to "
+        "disable patch for torch by adding "
+        '--patch "patch_sympy=False" --patch "patch_torch=False", '
+        "default is True.",
     )
     parser.add_argument(
         "--rewrite",

onnx_diagnostic/helpers/config_helper.py

@@ -1,6 +1,7 @@
 import functools
 import importlib
 import inspect
+import os
 import re
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 import transformers
@@ -110,3 +111,12 @@ def config_class_from_architecture(arch: str, exc: bool = False) -> Optional[typ
     )
     cls_name = unique.pop()
     return getattr(transformers, cls_name)
+
+
+def default_num_hidden_layers():
+    """
+    Returns the default number of layers.
+    It is lower when the unit tests are running
+    when ``UNITTEST_GOING=1``.
+    """
+    return 2 if os.environ.get("UNITTEST_GOING", "0") == "1" else 4

onnx_diagnostic/tasks/automatic_speech_recognition.py

@@ -2,7 +2,11 @@
 import torch
 import transformers
 from ..helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "automatic-speech-recognition"
 
@@ -15,7 +19,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     if hasattr(config, "decoder_layers"):
         config.decoder_layers = min(config.decoder_layers, 2)
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 4)
+        config.num_hidden_layers = min(config.num_hidden_layers, nhl())
     update_config(config, kwargs)
     return kwargs
 

onnx_diagnostic/tasks/feature_extraction.py

@@ -1,15 +1,20 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 from ..helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache
 
+
 __TASK__ = "feature-extraction"
 
 
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_hidden_layers")
-    kwargs = dict(num_hidden_layers=min(config.num_hidden_layers, 4))
+    kwargs = dict(num_hidden_layers=min(config.num_hidden_layers, nhl()))
     update_config(config, kwargs)
     return kwargs
 

onnx_diagnostic/tasks/fill_mask.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "fill-mask"
 
@@ -9,7 +13,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
-        num_hidden_layers=min(config.num_hidden_layers, 4),
+        num_hidden_layers=min(config.num_hidden_layers, nhl()),
         num_attention_heads=min(config.num_attention_heads, 4),
     )
     update_config(config, kwargs)

onnx_diagnostic/tasks/image_classification.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "image-classification"
 
@@ -17,7 +21,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
         num_hidden_layers=(
-            min(config.num_hidden_layers, 4)
+            min(config.num_hidden_layers, nhl())
             if hasattr(config, "num_hidden_layers")
             else len(config.hidden_sizes)
         )

onnx_diagnostic/tasks/image_text_to_text.py

@@ -1,7 +1,12 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
 from ..helpers.cache_helper import make_dynamic_cache, make_hybrid_cache
-from ..helpers.config_helper import update_config, check_hasattr, _pick
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    _pick,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "image-text-to-text"
 
@@ -10,7 +15,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 4)
+        config.num_hidden_layers = min(config.num_hidden_layers, nhl())
     if hasattr(config, "mm_tokens_per_image"):
         config.mm_tokens_per_image = min(config.mm_tokens_per_image, 2)
     if hasattr(config, "vision_config"):

onnx_diagnostic/tasks/mask_generation.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "mask-generation"
 
@@ -9,7 +13,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 4)
+        config.num_hidden_layers = min(config.num_hidden_layers, nhl())
     if hasattr(config, "vision_config") and hasattr(config.vision_config, "num_hidden_layers"):
         config.vision_config.num_hidden_layers = min(config.vision_config.num_hidden_layers, 2)
     update_config(config, kwargs)

onnx_diagnostic/tasks/mixture_of_expert.py

@@ -2,7 +2,7 @@
 import torch
 
 # from ..helpers.cache_helper import make_dynamic_cache
-from ..helpers.config_helper import update_config  # , check_hasattr, _pick
+from ..helpers.config_helper import update_config, default_num_hidden_layers as nhl
 
 __TASK__ = "MoE"
 
@@ -11,7 +11,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 4)
+        config.num_hidden_layers = min(config.num_hidden_layers, nhl())
     if hasattr(config, "vision_config") and hasattr(config.vision_config, "num_hidden_layers"):
         config.vision_config.num_hidden_layers = min(config.vision_config.num_hidden_layers, 2)
     if hasattr(config, "audio_processor") and hasattr(

onnx_diagnostic/tasks/object_detection.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "object-detection"
 
@@ -10,7 +14,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
         num_hidden_layers=(
-            min(config.num_hidden_layers, 4)
+            min(config.num_hidden_layers, nhl())
             if hasattr(config, "num_hidden_layers")
             else len(config.hidden_sizes)
         )

onnx_diagnostic/tasks/sentence_similarity.py

@@ -1,6 +1,10 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.config_helper import update_config, check_hasattr
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "sentence-similarity"
 
@@ -9,7 +13,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
-        num_hidden_layers=min(config.num_hidden_layers, 4),
+        num_hidden_layers=min(config.num_hidden_layers, nhl()),
         num_attention_heads=min(config.num_attention_heads, 4),
     )
     update_config(config, kwargs)

onnx_diagnostic/tasks/summarization.py

@@ -1,7 +1,12 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
 from ..helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache
-from ..helpers.config_helper import update_config, check_hasattr, _pick
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    _pick,
+    default_num_hidden_layers as nhl,
+)
 
 __TASK__ = "summarization"
 
@@ -12,7 +17,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     if hasattr(config, "num_decoder_layers"):
         config.num_decoder_layers = min(config.num_decoder_layers, 2)
     if hasattr(config, "num_hidden_layers"):
-        config.num_hidden_layers = min(config.num_hidden_layers, 4)
+        config.num_hidden_layers = min(config.num_hidden_layers, nhl())
     update_config(config, kwargs)
     return kwargs