onnx-diagnostic 0.7.5__py3-none-any.whl → 0.7.7__py3-none-any.whl

onnx_diagnostic/tasks/text2text_generation.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__ = "text2text-generation"
 
@@ -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, 2)
+        config.num_hidden_layers = min(config.num_hidden_layers, nhl())
     update_config(config, kwargs)
     return kwargs
 

onnx_diagnostic/tasks/text_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__ = "text-classification"
 
@@ -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, 2),
+        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/text_generation.py

@@ -1,13 +1,17 @@
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 import torch
-import transformers
 from ..helpers.cache_helper import (
     make_dynamic_cache,
     make_mamba_cache,
     make_sliding_window_cache,
     make_static_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__ = "text-generation"
 
@@ -26,7 +30,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
     if config.__class__.__name__ == "FalconMambaConfig":
         check_hasattr(config, "conv_kernel", "state_size", "intermediate_size")  # 4 and 8
         kwargs = dict(
-            num_hidden_layers=min(config.num_hidden_layers, 2),
+            num_hidden_layers=min(config.num_hidden_layers, nhl()),
             intermediate_size=256 if config is None else min(512, config.intermediate_size),
             hidden_size=512 if config is None else min(512, config.hidden_size),
             cls_cache="MambaCache",
@@ -38,24 +42,13 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
             head_dim=getattr(
                 config, "head_dim", config.hidden_size // config.num_attention_heads
             ),
-            num_hidden_layers=min(config.num_hidden_layers, 2),
+            num_hidden_layers=min(config.num_hidden_layers, nhl()),
             num_key_value_heads=(
                 config.num_key_value_heads
                 if hasattr(config, "num_key_value_heads")
                 else config.num_attention_heads
             ),
-            hidden_size=(
-                min(config.hidden_size, 4096 // 4)
-                if config.hidden_size % 64 == 0
-                else config.hidden_size
-            ),
         )
-        if config is None or hasattr(config, "intermediate_size"):
-            kwargs["intermediate_size"] = (
-                min(config.intermediate_size, 24576 // 4)
-                if config.intermediate_size % 4 == 0
-                else config.intermediate_size
-            )
     update_config(config, kwargs)
     return kwargs
 
@@ -95,9 +88,14 @@ def get_inputs(
     cache_length = "cache_length"  # torch.export.Dim("cache_length", min=1, max=4096)
 
     if config is not None and config.__class__.__name__ == "FalconMambaConfig":
+        try:
+            from transformers.models.mamba.modeling_mamba import MambaCache
+        except ImportError:
+            from transformers.cache_utils import MambaCache
+
         assert cls_cache in (
             "MambaCache",
-            transformers.cache_utils.MambaCache,
+            MambaCache,
         ), f"Unexpected value for cls_cache={cls_cache} and config={config}"
         seq_length_multiple = 8
         sequence_length = (

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -361,7 +361,7 @@ def torch_export_patches(
             torch._meta_registrations._broadcast_shapes = patched__broadcast_shapes
 
         # torch._export.non_strict_utils.produce_guards_and_solve_constraints
-        if catch_constraints:
+        if patch_torch and catch_constraints:
             if verbose:
                 print("[torch_export_patches] modifies shape constraints")
             f_produce_guards_and_solve_constraints = (
@@ -513,7 +513,7 @@ def torch_export_patches(
                 if verbose:
                     print("[torch_export_patches] restored pytorch functions")
 
-            if stop_if_static:
+            if patch_torch and stop_if_static:
                 if verbose:
                     print("[torch_export_patches] restored ShapeEnv._set_replacement")
 
@@ -529,7 +529,7 @@ def torch_export_patches(
                         print("[torch_export_patches] restored ShapeEnv._check_frozen")
                     ShapeEnv._check_frozen = f_shape_env__check_frozen
 
-            if catch_constraints:
+            if patch_torch and catch_constraints:
                 # to catch or skip dynamic_shapes issues
                 torch._export.non_strict_utils.produce_guards_and_solve_constraints = (
                     f_produce_guards_and_solve_constraints

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -6,12 +6,17 @@ import torch
 import transformers
 from transformers.cache_utils import (
     DynamicCache,
-    MambaCache,
     EncoderDecoderCache,
+    HybridCache,
     SlidingWindowCache,
     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_
 
@@ -161,6 +166,9 @@ def serialization_functions(
             flatten_dynamic_cache,
             unflatten_dynamic_cache,
             flatten_with_keys_dynamic_cache,
+            flatten_hybrid_cache,
+            unflatten_hybrid_cache,
+            flatten_with_keys_hybrid_cache,
             flatten_mamba_cache,
             unflatten_mamba_cache,
             flatten_with_keys_mamba_cache,
@@ -187,6 +195,14 @@ def serialization_functions(
                 # f_check=make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]),
                 verbose=verbose,
             ),
+            HybridCache: lambda verbose=verbose: register_class_serialization(
+                HybridCache,
+                flatten_hybrid_cache,
+                unflatten_hybrid_cache,
+                flatten_with_keys_hybrid_cache,
+                # 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,

onnx_diagnostic/torch_export_patches/patch_inputs.py

@@ -34,7 +34,7 @@ def _make_shape(subset: Dict, cls: type, value: Any) -> Any:
             f"Inconsistencies in subset={subset}, found={values}, "
             f"it cannot be a {cls}, value={string_type(value)}"
         )
-        cache_length = len(value.key_cache)
+        cache_length = len(value.layers if hasattr(value, "layers") else value.key_cache)
         for v in subset.values():
             axes = v
             break
@@ -70,6 +70,8 @@ def convert_dynamic_axes_into_dynamic_shapes(
     :param verbose: verbosity
     :return: (args, kwargs, dynamic shapes)
     """
+    from ..helpers.cache_helper import CacheKeyValue
+
     new_kwargs = {}
     if args:
         assert hasattr(model, "forward"), f"Missing method 'forward' for {model!r}"
@@ -121,7 +123,8 @@ def convert_dynamic_axes_into_dynamic_shapes(
                 changes[k] = type(updated_kwargs[k])
                 continue
         if isinstance(v, transformers.cache_utils.DynamicCache):
-            updated_kwargs[k] = [v.key_cache, v.value_cache]
+            ca = CacheKeyValue(v)
+            updated_kwargs[k] = [ca.key_cache, ca.value_cache]
             changes[k] = type(v)
             continue
         raise NotImplementedError(

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -27,8 +27,8 @@ def _catch_produce_guards_and_solve_constraints(
     dynamic_shapes: Union[Dict[str, Any], Tuple[Any], List[Any], None],
     equalities_inputs: "EqualityConstraint",  # noqa: F821
     original_signature: inspect.Signature,
-    _is_torch_jit_trace: bool = False,
     verbose: int = 0,
+    **kwargs,
 ):
     try:
         return previous_function(
@@ -37,7 +37,7 @@ def _catch_produce_guards_and_solve_constraints(
             dynamic_shapes=dynamic_shapes,
             equalities_inputs=equalities_inputs,
             original_signature=original_signature,
-            _is_torch_jit_trace=_is_torch_jit_trace,
+            **kwargs,
         )
     except Exception as e:
         if not int(os.environ.get("SKIP_SOLVE_CONSTRAINTS", "1")):
@@ -51,7 +51,7 @@ def _catch_produce_guards_and_solve_constraints(
                 f"dynamic_shapes={dynamic_shapes}\n"
                 f"equalities_inputs={equalities_inputs}\n"
                 f"original_signature={original_signature}\n"
-                f"_is_torch_jit_trace={_is_torch_jit_trace}\n"
+                f"kwargs={kwargs}\n"
                 f"exc={e}\ngm={gm}"
             )
         torch._dynamo.reset()
@@ -174,7 +174,7 @@ class patched_ShapeEnv:
             self.counter["ignored_backward_guard"] += 1
             raise AssertionError(
                 f"[patched_ShapeEnv] Ignored guard {expr} == {concrete_val}, "
-                f"this could result in accuracy problems."
+                f"this could result in accuracy problems"
             )
 
     def _set_replacement(