onnx-diagnostic 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -1,6 +1,6 @@
 import contextlib
 import pprint
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, List, Optional, Set
 from .onnx_export_serialization import (
     flatten_with_keys_dynamic_cache,
     flatten_dynamic_cache,
@@ -12,27 +12,36 @@ from .onnx_export_serialization import (
 from .patches import patch_transformers as patch_transformers_list
 
 
-def patch_module(mod, verbose: int = 0) -> Dict[type, Dict[type, Callable]]:
+def patch_module_or_classes(mod, verbose: int = 0) -> Dict[type, Dict[type, Callable]]:
     """
     Applies all patches defined in classes prefixed by ``patched_``
     ``cls._PATCHED_CLASS_`` defines the class to patch,
     ``cls._PATCHES_`` defines the method to patch.
-    The returns information needs to be sent to :func:`unpatch_module`
+    The returns information needs to be sent to :func:`unpatch_module_or_classes`
     to revert the changes.
+
+    :param mod: module of list of clsses to patch
+    :param verbose: verbosity
+    :return: patch info
     """
-    to_patch = []
-    for k in dir(mod):
-        if k.startswith("patched_"):
-            v = getattr(mod, k)
-            if hasattr(v, "_PATCHED_CLASS_") and hasattr(v, "_PATCHES_"):
-                to_patch.append(v)
+    if isinstance(mod, list):
+        to_patch = mod
+        name = "list"
+    else:
+        to_patch = []
+        for k in dir(mod):
+            if k.startswith("patched_"):
+                v = getattr(mod, k)
+                if hasattr(v, "_PATCHED_CLASS_") and hasattr(v, "_PATCHES_"):
+                    to_patch.append(v)
+        name = mod.__name__
 
     res = {}
     for cls in to_patch:
         original = cls._PATCHED_CLASS_
         methods = cls._PATCHES_
         if verbose:
-            print(f"[patch_module] {mod.__name__} - {cls.__name__}: {', '.join(methods)}")
+            print(f"[patch_module_or_classes] {name} - {cls.__name__}: {', '.join(methods)}")
 
         keep = {n: getattr(original, n, None) for n in methods}
         for n in methods:
@@ -42,20 +51,30 @@ def patch_module(mod, verbose: int = 0) -> Dict[type, Dict[type, Callable]]:
     return res
 
 
-def unpatch_module(mod, info: Dict[type, Dict[type, Callable]], verbose: int = 0):
-    """Reverts modification made by :func:`patch_module`."""
-    to_patch = []
-    for k in dir(mod):
-        if k.startswith("patched_"):
-            v = getattr(mod, k)
-            if hasattr(v, "_PATCHED_CLASS_") and hasattr(v, "_PATCHES_"):
-                to_patch.append(v)
+def unpatch_module_or_classes(mod, info: Dict[type, Dict[type, Callable]], verbose: int = 0):
+    """
+    Reverts modification made by :func:`patch_module_or_classes`.
+
+    :param mod: module of list of clsses to patch
+    :param verbose: verbosity
+    """
+    if isinstance(mod, list):
+        to_patch = mod
+        name = "list"
+    else:
+        to_patch = []
+        for k in dir(mod):
+            if k.startswith("patched_"):
+                v = getattr(mod, k)
+                if hasattr(v, "_PATCHED_CLASS_") and hasattr(v, "_PATCHES_"):
+                    to_patch.append(v)
+        name = mod.__name__
     set_patch = set(to_patch)
 
     for cls, methods in info.items():
         assert cls in set_patch, f"No patch registered for {cls} in {mod} (found {set_patch})"
         if verbose:
-            print(f"[unpatch_module] {mod.__name__} - {cls.__name__}: {', '.join(methods)}")
+            print(f"[unpatch_module_or_classes] {name} - {cls.__name__}: {', '.join(methods)}")
         original = cls._PATCHED_CLASS_
         for n, v in methods.items():
             if v is None:
@@ -65,9 +84,14 @@ def unpatch_module(mod, info: Dict[type, Dict[type, Callable]], verbose: int = 0
                 setattr(original, n, v)
 
 
+PATCH_OF_PATCHES: Set[Any] = set()
+
+
 def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
     # Cache serialization: to be moved into appropriate packages
     import torch
+    import transformers
+    import packaging.version as pv
 
     try:
         from transformers.cache_utils import DynamicCache
@@ -100,7 +124,40 @@ def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
             flatten_with_keys_fn=flatten_with_keys_mamba_cache,
         )
 
-    # DynamicCache
+    # DynamicCache serialization is different in transformers and does not
+    # play way with torch.export.export.
+    # see test test_export_dynamic_cache_cat with NOBYPASS=1
+    # :: NOBYBASS=1 python _unittests/ut_torch_export_patches/test_dynamic_class.py -k e_c
+    # This is caused by this line:
+    # torch.fx._pytree.register_pytree_flatten_spec(
+    #           DynamicCache, _flatten_dynamic_cache_for_fx)
+    # so we remove it anyway
+    if (
+        DynamicCache in torch.fx._pytree.SUPPORTED_NODES
+        and not PATCH_OF_PATCHES
+        # and pv.Version(torch.__version__) < pv.Version("2.7")
+        and pv.Version(transformers.__version__) >= pv.Version("4.50")
+    ):
+        if verbose:
+            print(
+                "[_register_cache_serialization] DynamicCache "
+                "is unregistered and registered first."
+            )
+        _unregister(DynamicCache)
+        torch.utils._pytree.register_pytree_node(
+            DynamicCache,
+            flatten_dynamic_cache,
+            unflatten_dynamic_cache,
+            serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
+            flatten_with_keys_fn=flatten_with_keys_dynamic_cache,
+        )
+        if pv.Version(torch.__version__) < pv.Version("2.7"):
+            torch.fx._pytree.register_pytree_flatten_spec(
+                DynamicCache, lambda x, _: [x.key_cache, x.value_cache]
+            )
+        # To avoid doing it multiple times.
+        PATCH_OF_PATCHES.add(DynamicCache)
+
     unregistered_dynamic_cache = True
     if DynamicCache is not None and DynamicCache in torch.utils._pytree.SUPPORTED_NODES:
         if verbose > 1:
@@ -116,12 +173,13 @@ def _register_cache_serialization(verbose: int = 0) -> Dict[str, bool]:
             serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
             flatten_with_keys_fn=flatten_with_keys_dynamic_cache,
         )
-        torch.fx._pytree.register_pytree_flatten_spec(
-            DynamicCache, lambda x, _: [x.key_cache, x.value_cache]
-        )
+        if pv.Version(torch.__version__) < pv.Version("2.7"):
+            torch.fx._pytree.register_pytree_flatten_spec(
+                DynamicCache, lambda x, _: [x.key_cache, x.value_cache]
+            )
 
         # check
-        from ..cache_helpers import make_dynamic_cache
+        from ..helpers.cache_helper import make_dynamic_cache
 
         cache = make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))])
         values, spec = torch.utils._pytree.tree_flatten(cache)
@@ -180,7 +238,7 @@ def _unregister_cache_serialization(undo: Dict[str, bool], verbose: int = 0):
 def register_additional_serialization_functions(
     patch_transformers: bool = False, verbose: int = 0
 ) -> Callable:
-    """The necessary modification to run the fx Graph."""
+    """The necessary modifications to run the fx Graph."""
     fct_callable = replacement_before_exporting if patch_transformers else (lambda x: x)
     done = _register_cache_serialization(verbose=verbose)
     try:
@@ -195,9 +253,10 @@ def bypass_export_some_errors(
     patch_torch: bool = True,
     patch_transformers: bool = False,
     catch_constraints: bool = True,
-    stop_if_static: bool = False,
+    stop_if_static: int = 0,
     verbose: int = 0,
     patch: bool = True,
+    custom_patches: Optional[List[type["torch.nn.Module"]]] = None,  # noqa: F821
 ) -> Callable:
     """
     Tries to bypass some situations :func:`torch.export.export` does not support.
@@ -211,9 +270,14 @@ def bypass_export_some_errors(
         can be put to stop at that stage.
     :param stop_if_static: see example :ref:`l-plot-export-locale-issue`,
         to stop the export as soon as an issue is detected with dynamic shapes
-        and show a stack trace indicating the exact location of the issue
+        and show a stack trace indicating the exact location of the issue,
+        ``if stop_if_static > 1``, more methods are replace to catch more
+        issues
     :param patch: if False, disable all patches except the registration of
         serialization function
+    :param custom_patches: to apply custom patches,
+        every patched class must define static attributes
+        ``_PATCHES_``, ``_PATCHED_CLASS_``
     :param verbose: to show which patches is applied
 
     The list of available patches.
@@ -301,6 +365,7 @@ def bypass_export_some_errors(
             f_sympy_name = getattr(sympy.core.numbers.IntegerConstant, "name", None)
 
             if verbose:
+                print(f"[bypass_export_some_errors] sympy.__version__={sympy.__version__!r}")
                 print("[bypass_export_some_errors] patch sympy")
 
             sympy.core.numbers.IntegerConstant.name = lambda self: f"IntCst{str(self)}"
@@ -318,6 +383,8 @@ def bypass_export_some_errors(
             )
 
             if verbose:
+                print(f"[bypass_export_some_errors] torch.__version__={torch.__version__!r}")
+                print(f"[bypass_export_some_errors] stop_if_static={stop_if_static!r}")
                 print("[bypass_export_some_errors] patch pytorch")
 
             # torch.jit.isinstance
@@ -359,23 +426,46 @@ def bypass_export_some_errors(
             )
 
         if stop_if_static:
+            from torch.fx.experimental.symbolic_shapes import ShapeEnv
+            from .patches.patch_torch import patched_ShapeEnv
+
             if verbose:
                 print(
                     "[bypass_export_some_errors] assert when a dynamic dimension turns static"
                 )
-
-            from torch.fx.experimental.symbolic_shapes import ShapeEnv
-            from .patches.patch_torch import patched_ShapeEnv
+                print("[bypass_export_some_errors] replaces ShapeEnv._set_replacement")
 
             f_shape_env__set_replacement = ShapeEnv._set_replacement
             ShapeEnv._set_replacement = patched_ShapeEnv._set_replacement
 
+            if stop_if_static > 1:
+                if verbose:
+                    print("[bypass_export_some_errors] replaces ShapeEnv._check_frozen")
+                f_shape_env__check_frozen = ShapeEnv._check_frozen
+                ShapeEnv._check_frozen = patched_ShapeEnv._check_frozen
+
         ####################
         # patch transformers
         ####################
 
         if patch_transformers:
-            revert_patches_info = patch_module(patch_transformers_list, verbose=verbose)
+            if verbose:
+                import transformers
+
+                print(
+                    f"[bypass_export_some_errors] transformers.__version__="
+                    f"{transformers.__version__!r}"
+                )
+            revert_patches_info = patch_module_or_classes(
+                patch_transformers_list, verbose=verbose
+            )
+
+        if custom_patches:
+            if verbose:
+                print("[bypass_export_some_errors] applies custom patches")
+            revert_custom_patches_info = patch_module_or_classes(
+                custom_patches, verbose=verbose
+            )
 
         ########
         # export
@@ -397,7 +487,6 @@ def bypass_export_some_errors(
                 print("[bypass_export_some_errors] remove patches")
 
             if patch_sympy:
-
                 # tracked by https://github.com/pytorch/pytorch/issues/143494
                 if f_sympy_name:
                     sympy.core.numbers.IntegerConstant.name = f_sympy_name
@@ -428,6 +517,10 @@ def bypass_export_some_errors(
                     print("[bypass_export_some_errors] restored ShapeEnv._set_replacement")
 
                 ShapeEnv._set_replacement = f_shape_env__set_replacement
+                if stop_if_static > 1:
+                    if verbose:
+                        print("[bypass_export_some_errors] restored ShapeEnv._check_frozen")
+                    ShapeEnv._check_frozen = f_shape_env__check_frozen
 
             if catch_constraints:
                 # to catch or skip dynamic_shapes issues
@@ -440,12 +533,23 @@ def bypass_export_some_errors(
                 if verbose:
                     print("[bypass_export_some_errors] restored shape constraints")
 
+            if custom_patches:
+                if verbose:
+                    print("[bypass_export_some_errors] unpatch custom patches")
+                unpatch_module_or_classes(
+                    custom_patches, revert_custom_patches_info, verbose=verbose
+                )
+
             ##############
             # transformers
             ##############
 
             if patch_transformers:
-                unpatch_module(patch_transformers_list, revert_patches_info, verbose=verbose)
+                if verbose:
+                    print("[bypass_export_some_errors] unpatch transformers")
+                unpatch_module_or_classes(
+                    patch_transformers_list, revert_patches_info, verbose=verbose
+                )
 
             ########
             # caches

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -97,6 +97,10 @@ def flatten_dynamic_cache(
     dynamic_cache: transformers.cache_utils.DynamicCache,
 ) -> Tuple[List[Any], torch.utils._pytree.Context]:
     """Serializes a :class:`transformers.cache_utils.DynamicCache` with python objects."""
+    import transformers.cache_utils
+
+    if hasattr(transformers.cache_utils, "_flatten_dynamic_cache"):
+        return transformers.cache_utils._flatten_dynamic_cache(dynamic_cache)
     flat = [
         (k, getattr(dynamic_cache, k))
         for k in ["key_cache", "value_cache"]
@@ -111,7 +115,10 @@ def flatten_with_keys_dynamic_cache(d: Dict[Any, Any]) -> Tuple[
 ]:
     """Serializes a :class:`transformers.cache_utils.DynamicCache` with python objects."""
     import torch
+    import transformers.cache_utils
 
+    if hasattr(transformers.cache_utils, "_flatten_with_keys_dynamic_cache"):
+        return transformers.cache_utils._flatten_with_keys_dynamic_cache(d)
     values, context = flatten_dynamic_cache(d)
     return [(torch.utils._pytree.MappingKey(k), v) for k, v in zip(context, values)], context
 
@@ -122,9 +129,13 @@ def unflatten_dynamic_cache(
     output_type=None,
 ) -> transformers.cache_utils.DynamicCache:
     """Restores a :class:`transformers.cache_utils.DynamicCache` from python objects."""
-    from transformers.cache_utils import DynamicCache
+    import transformers.cache_utils
+
+    if hasattr(transformers.cache_utils, "_unflatten_dynamic_cache"):
+        assert output_type is None, f"output_type={output_type} not supported"
+        return transformers.cache_utils._unflatten_dynamic_cache(values, context)
 
-    cache = DynamicCache()
+    cache = transformers.cache_utils.DynamicCache()
     values = dict(zip(context, values))
     for k, v in values.items():
         setattr(cache, k, v)

onnx_diagnostic/torch_export_patches/patch_inputs.py

@@ -0,0 +1,174 @@
+import inspect
+from typing import Any, Dict, Optional, Tuple
+import torch
+import transformers
+from ..helpers import string_type
+from ..helpers.cache_helper import make_dynamic_cache
+
+
+def _process_cache(k: str, v):
+    assert k != "position_ids" or isinstance(
+        k, torch.Tensor
+    ), f"Unexpected type for parameter {k!r} {string_type(v, with_shape=True)}"
+    if (
+        isinstance(v, list)
+        and all(isinstance(i, tuple) for i in v)
+        and set(len(t) for t in v) == {2}
+    ):
+        # A dynamicCache
+        cache = make_dynamic_cache(v)
+        return cache
+    if isinstance(v, torch.Tensor):
+        return v
+    raise NotImplementedError(
+        f"Unable to process parameter {k!r} with v={string_type(v,with_shape=True)}"
+    )
+
+
+def _make_shape(subset: Dict, cls: type, value: Any) -> Any:
+    if cls is transformers.cache_utils.DynamicCache:
+        assert subset, "DynamicCache cannot be empty"
+        values = set(map(str, subset.values()))
+        assert len(values) == 1, (
+            f"Inconsistencies in subset={subset}, found={values}, "
+            f"it cannot be a {cls}, value={string_type(value)}"
+        )
+        cache_length = len(value.key_cache)
+        for v in subset.values():
+            axes = v
+            break
+        new_shape = [[axes for i in range(cache_length)], [axes for i in range(cache_length)]]
+        return new_shape
+    if value.__class__ in torch.utils._pytree.SUPPORTED_NODES:
+        raise NotImplementedError(
+            f"_make_shape not implemented for registered class={cls}, "
+            f"subset={subset}, value={string_type(value)}"
+        )
+    raise NotImplementedError(
+        f"_make_shape not implemented for cls={cls}, "
+        f"subset={subset}, value={string_type(value)}"
+    )
+
+
+def convert_dynamic_axes_into_dynamic_shapes(
+    model: torch.nn.Module,
+    args: Optional[Tuple[Any, ...]] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None,
+    prefix_mapping: Optional[Dict[str, str]] = None,
+    verbose: int = 0,
+) -> Tuple[Tuple[Any, ...], Dict[str, Any], Dict[str, Any]]:
+    """
+    Converts the input from an export to something :func:`torch.export.export` can handle.
+
+    :param model: model to convert (used to extract the signature)
+    :param args: positional arguments
+    :param kwargs: named arguments
+    :param dynamic_axes: dynamic axes
+    :param prefix_mapping: prefix mapping
+    :param verbose: verbosity
+    :return: (args, kwargs, dynamic shapes)
+    """
+    new_kwargs = {}
+    if args:
+        assert hasattr(model, "forward"), f"Missing method 'forward' for {model!r}"
+        plus = 0 if isinstance(model, torch.nn.Module) else 1
+        print(
+            f"[convert_dynamic_axes_into_dynamic_shapes] "
+            f"mapping args to kwargs for model="
+            f"{model if plus else model.__class__.__name__}"
+        )
+        pars = inspect.signature(model.forward).parameters
+        assert len(pars) >= len(
+            args
+        ), f"Length mismatch, len(args)={len(args)}, pars={list(pars)}"
+
+        for i, p in enumerate(pars):
+            if i < plus:
+                continue
+            if i - plus >= len(args):
+                break
+            if verbose:
+                print(
+                    f"[convert_dynamic_axes_into_dynamic_shapes] mapping args[{i-plus}] "
+                    f"to {p!r} ({string_type(args[i-plus])})"
+                )
+            new_kwargs[p] = args[i - plus]
+
+    if kwargs:
+        for k, v in kwargs.items():
+            assert k not in new_kwargs, f"Argument {k!r} from kwargs already present in args."
+            new_kwargs[k] = v
+
+    # process
+    updated_kwargs = {}
+    changes = {}
+    for k, v in new_kwargs.items():
+        if isinstance(v, torch.Tensor):
+            updated_kwargs[k] = v
+            continue
+        if isinstance(v, list):
+            # cache?
+            updated_kwargs[k] = _process_cache(k, v)
+            if type(updated_kwargs[k]) is not type(v):
+                # A cache was introduced.
+                if verbose:
+                    print(
+                        f"[convert_dynamic_axes_into_dynamic_shapes] parameter "
+                        f"{k!r} was changed into {type(updated_kwargs[k])}"
+                    )
+                changes[k] = type(updated_kwargs[k])
+                continue
+        raise NotImplementedError(
+            f"Unexpected type {type(v)} for parameter {k!r} "
+            f"({string_type(v, with_shape=True)})"
+        )
+
+    # process dynamic axes
+    if changes:
+        dynamic_shapes = {}
+        done = set()
+        for k, v in dynamic_axes.items():
+            if k not in changes and k in updated_kwargs and isinstance(v, dict):
+                dynamic_shapes[k] = v
+                continue
+            if "." in k:
+                # something like present.0.key
+                prefix = k.split(".")[0]
+                if prefix in done:
+                    continue
+                args_prefix = (
+                    prefix_mapping[prefix]
+                    if prefix_mapping and prefix in prefix_mapping
+                    else prefix
+                )
+                if args_prefix in updated_kwargs and args_prefix in changes:
+                    # A cache.
+                    cls = changes[args_prefix]
+                    dynamic_shapes[args_prefix] = _make_shape(
+                        {
+                            _: __
+                            for _, __ in dynamic_axes.items()
+                            if _.startswith(f"{prefix}.")
+                        },
+                        cls,
+                        updated_kwargs[args_prefix],
+                    )
+                    done.add(prefix)
+                    continue
+            if k not in updated_kwargs:
+                # dynamic axes not in the given inputs, should be raise an exception?
+                if verbose:
+                    print(
+                        f"[convert_dynamic_axes_into_dynamic_shapes] dropping axes "
+                        f"{k!r}-{v!r}, not found in {set(updated_kwargs)}"
+                    )
+                continue
+            raise NotImplementedError(
+                f"Unable to process dynamic axes {k!r}, axes={v}, "
+                f"value={string_type(updated_kwargs[k], with_shape=True)}, "
+                f"dynamic axes={dynamic_axes}, "
+                f"updated_kwargs={string_type(updated_kwargs, with_shape=True)}"
+            )
+
+    return (), updated_kwargs, dynamic_shapes

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -131,7 +131,7 @@ def patched__broadcast_shapes(*_shapes):
         assert isinstance(shape, Sequence)
 
     # Computes common shape
-    common_shape: List[Union[int, torch.SymInt]] = [
+    common_shape = [  # List[Union[int, torch.SymInt]]
         1,
     ] * reduce(max, (len(shape) for shape in shapes))
     for _arg_idx, shape in enumerate(shapes):
@@ -150,6 +150,16 @@ def patched__broadcast_shapes(*_shapes):
 
 class patched_ShapeEnv:
 
+    def _check_frozen(
+        self, expr: "sympy.Basic", concrete_val: "sympy.Basic"  # noqa: F821
+    ) -> None:
+        if self.frozen:
+            self.counter["ignored_backward_guard"] += 1
+            raise AssertionError(
+                f"[patched_ShapeEnv] Ignored guard {expr} == {concrete_val}, "
+                f"this could result in accuracy problems."
+            )
+
     def _set_replacement(
         self, a: "sympy.Symbol", tgt: "sympy.Expr", msg: str  # noqa: F821
     ) -> None:
@@ -314,7 +324,7 @@ class patched_ShapeEnv:
             #     )
             #     self.log.debug("SPECIALIZATION", stack_info=True)
         assert msg != "range_refined_to_singleton", (
-            f"A dynamic dimension becomes static! "
+            f"patched_ShapeEnv: A dynamic dimension becomes static! "
             f"a={a!r}, tgt={tgt!r}, msg={msg!r}, tgt_bound={tgt_bound}"
         )
         # log.info("set_replacement %s = %s (%s) %s", a, tgt, msg, tgt_bound)

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -4,9 +4,9 @@ from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Tuple
 import torch
 import transformers
-import transformers.modeling_attn_mask_utils
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
 from transformers.cache_utils import StaticCache, Cache, DynamicCache
-from ...torch_test_helper import is_torchdynamo_exporting
+from ...helpers.torch_test_helper import is_torchdynamo_exporting
 
 
 def _patch_make_causal_mask(
@@ -54,7 +54,7 @@ if sys.version_info[:2] <= (3, 11):
         """
 
         _PATCHES_ = ["_make_causal_mask"]
-        _PATCHED_CLASS_ = transformers.modeling_attn_mask_utils.AttentionMaskConverter
+        _PATCHED_CLASS_ = AttentionMaskConverter
 
         @staticmethod
         def _make_causal_mask(
@@ -79,7 +79,7 @@ else:
         """
 
         _PATCHES_ = ["_make_causal_mask"]
-        _PATCHED_CLASS_ = transformers.modeling_attn_mask_utils.AttentionMaskConverter
+        _PATCHED_CLASS_ = AttentionMaskConverter
 
         @staticmethod
         def _make_causal_mask(

onnx_diagnostic/torch_models/hghub/__init__.py

@@ -0,0 +1 @@
+from .model_inputs import get_untrained_model_with_inputs