onnx-diagnostic 0.8.0__py3-none-any.whl

onnx_diagnostic/helpers/config_helper.py

@@ -0,0 +1,170 @@
+import functools
+import importlib
+import inspect
+import os
+import re
+from typing import Any, Callable, Dict, Optional, Tuple, Union
+import transformers
+
+
+def check_hasattr(config: Any, *args: Union[str, Tuple[Any, ...]]):
+    """
+    Checks the confiugation has all the attributes in ``args``.
+    Raises an exception otherwise.
+    """
+    for a in args:
+        assert isinstance(a, (str, tuple)), f"unexpected type {type(a)} in {args!r}"
+        if isinstance(a, str):
+            assert (isinstance(config, dict) and a in config) or hasattr(
+                config, a
+            ), f"Missing attribute {a!r} in\n{config}"
+        elif isinstance(a, tuple):
+            assert any(
+                (isinstance(name, str) and hasattr(config, name))
+                or all(hasattr(config, _) for _ in name)
+                for name in a
+            ), f"All attributes in {a!r} are missing from\n{config}"
+
+
+def update_config(config: Any, mkwargs: Dict[str, Any]):
+    """Updates a configuration with different values."""
+    for k, v in mkwargs.items():
+        if k == "attn_implementation":
+            config._attn_implementation = v
+            if getattr(config, "_attn_implementation_autoset", False):
+                config._attn_implementation_autoset = False
+            continue
+        if isinstance(v, dict):
+            if not hasattr(config, k) or getattr(config, k) is None:
+                setattr(config, k, v)
+                continue
+            existing = getattr(config, k)
+            if type(existing) is dict:
+                existing.update(v)
+            else:
+                update_config(getattr(config, k), v)
+            continue
+        if type(config) is dict:
+            config[k] = v
+        else:
+            setattr(config, k, v)
+
+
+def _pick(config, *atts, exceptions: Optional[Dict[str, Callable]] = None):
+    """Returns the first value found in the configuration."""
+    if (
+        exceptions
+        and hasattr(config, "architectures")
+        and len(config.architectures) == 1
+        and config.architectures[0] in exceptions
+    ):
+        excs = exceptions[config.architectures[0]]
+        return excs(config)
+    for a in atts:
+        if isinstance(a, str):
+            if hasattr(config, a):
+                return getattr(config, a)
+        elif isinstance(a, tuple):
+            if all(hasattr(config, _) for _ in a[1:]):
+                return a[0]([getattr(config, _) for _ in a[1:]])
+    raise AssertionError(f"Unable to find any of these {atts!r} in {config}")
+
+
+def pick(config, name: str, default_value: Any) -> Any:
+    """
+    Returns the value of a attribute if config has it
+    otherwise the default value.
+    """
+    if not config:
+        return default_value
+    if type(config) is dict:
+        return config.get(name, default_value)
+    return getattr(config, name, default_value)
+
+
+@functools.cache
+def config_class_from_architecture(arch: str, exc: bool = False) -> Optional[type]:
+    """
+    Retrieves the configuration class for a given architecture.
+
+    :param arch: architecture (clas name)
+    :param exc: raise an exception if not found
+    :return: type
+    """
+    cls = getattr(transformers, arch)
+    mod_name = cls.__module__
+    mod = importlib.import_module(mod_name)
+    source = inspect.getsource(mod)
+    # [^O] avoids capturing Optional[Something]
+    reg = re.compile("config: ([^O][A-Za-z0-9]+)")
+    fall = reg.findall(source)
+    if len(fall) == 0:
+        assert not exc, (
+            f"Unable to guess Configuration class name for arch={arch!r}, "
+            f"module={mod_name!r}, no candidate, source is\n{source}"
+        )
+        return None
+    unique = set(fall)
+    assert len(unique) == 1, (
+        f"Unable to guess Configuration class name for arch={arch!r}, "
+        f"module={mod_name!r}, found={unique} (#{len(unique)}), "
+        f"source is\n{source}"
+    )
+    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``.
+    """
+    import torch
+
+    if torch.cuda.is_available():
+        capa = torch.cuda.get_device_capability(0)
+        if capa[0] < 9:
+            return 2
+    return 2 if os.environ.get("UNITTEST_GOING", "0") == "1" else 4
+
+
+def build_diff_config(config0, config1):
+    """
+    Returns all the modified values between two configuration
+    """
+    import torch
+
+    diff = {}
+    for k in config0:
+        assert isinstance(k, str), f"k={k!r}, wrong type in {config0}"
+        if k not in config1:
+            v0 = getattr(config0, k) if hasattr(config0, k) else config0[k]
+            diff[k] = f"-{v0}"
+    for k in config1:
+        assert isinstance(k, str), f"k={k!r}, wrong type in {config1}"
+        if k not in config0:
+            v1 = getattr(config1, k) if hasattr(config1, k) else config1[k]
+            diff[k] = f"+{v1}"
+    for k in config0:
+        if k not in config1:
+            continue
+        v0 = getattr(config0, k) if hasattr(config0, k) else config0[k]
+        v1 = getattr(config1, k) if hasattr(config1, k) else config1[k]
+        if (
+            v0 is None
+            or v1 is None
+            or isinstance(v1, (float, int, bool, str, list, tuple, torch.dtype))
+            or (
+                isinstance(v0, dict)
+                and isinstance(v1, dict)
+                and all(isinstance(k, int) for k in v1)
+            )
+        ):
+            if v1 != v0:
+                diff[k] = f"{v0} -> {v1}"
+        else:
+            d = build_diff_config(v0, v1)
+            if d:
+                diff[k] = d
+    return diff

onnx_diagnostic/helpers/doc_helper.py

@@ -0,0 +1,163 @@
+import os
+from typing import Dict, List, Optional, Tuple
+import onnx
+import onnx.helper as oh
+import torch
+from ..reference.torch_ops import OpRunKernel, OpRunTensor
+from .torch_helper import onnx_dtype_to_torch_dtype, torch_dtype_to_onnx_dtype
+from .ort_session import InferenceSessionForTorch
+
+_SAVED: List[str] = []
+_SAVE_OPTIMIZED_MODEL_ = int(os.environ.get("DUMP_ONNX", "0"))
+
+
+def _get_model_name(op_name: str, provider: str) -> Optional[str]:
+    if _SAVE_OPTIMIZED_MODEL_:
+        name = f"dump_doc_layer_norm_{provider}_{len(_SAVED)}.onnx"
+        _SAVED.append(name)
+        return name
+    return None
+
+
+class LayerNormalizationOrt(OpRunKernel):
+    "LayerNormalization with onnxruntime"
+
+    @classmethod
+    def device_dependent(cls) -> bool:
+        "Needs device."
+        return True
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version=None,
+        device: Optional[torch.device] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
+        self.axis = self.get_attribute_int(node, "axis", -1)
+        self.epsilon = self.get_attribute_float(node, "epsilon", 1e-5)
+        self.device = device
+        self.stash_type = onnx_dtype_to_torch_dtype(
+            self.get_attribute_int(node, "stash_type", onnx.TensorProto.FLOAT)  # type: ignore[arg-type]
+        )
+        self.compute_std = len(node.output) > 1
+        assert not self.compute_std, (
+            f"This kernel implementation only work when only one output "
+            f"is required but {node.output} were."
+        )
+        self._cache: Dict[Tuple[int, int], onnx.ModelProto] = {}
+        self.is_cpu = torch.device("cpu") == self.device
+
+    def _make_model(self, itype: int, rank: int, has_bias: bool) -> onnx.ModelProto:
+        shape = [*["d{i}" for i in range(rank - 1)], "last"]
+        layer_model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node(
+                        "LayerNormalization",
+                        ["X", "W", "B"] if has_bias else ["X", "W"],
+                        ["Z"],
+                        axis=self.axis,
+                        epsilon=self.epsilon,
+                    )
+                ],
+                "dummy",
+                (
+                    [
+                        oh.make_tensor_value_info("X", itype, shape),
+                        oh.make_tensor_value_info("W", itype, ["last"]),
+                        oh.make_tensor_value_info("B", itype, ["last"]),
+                    ]
+                    if has_bias
+                    else [
+                        oh.make_tensor_value_info("X", itype, shape),
+                        oh.make_tensor_value_info("W", itype, ["last"]),
+                    ]
+                ),
+                [oh.make_tensor_value_info("Z", itype, shape)],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        provider = "CPUExecutionProvider" if self.is_cpu else "CUDAExecutionProvider"
+        self._provider = provider
+        return InferenceSessionForTorch(
+            layer_model,
+            optimized_model_filepath=_get_model_name("layer_norm", provider),
+            providers=[provider],
+        )
+
+    def run(self, x, scale, bias=None):
+        itype = torch_dtype_to_onnx_dtype(x.dtype)
+        rank = len(x.shape)
+        key = itype, rank
+        if key not in self._cache:
+            self._cache[key] = self._make_model(itype, rank, bias is not None)
+        sess = self._cache[key]
+        if self.verbose:
+            print(f"[LayerNormalizationOrt] running on {self._provider!r}")
+        feeds = dict(X=x.tensor, W=scale.tensor)
+        if bias is not None:
+            feeds["B"] = bias.tensor
+        got = sess.run(None, feeds)[0]
+        return OpRunTensor(got)
+
+
+class MatMulOrt(OpRunKernel):
+    "MatMul with onnxruntime"
+
+    @classmethod
+    def device_dependent(cls) -> bool:
+        "Needs device."
+        return True
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version=None,
+        device: Optional[torch.device] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
+        self.device = device
+        self._cache: Dict[Tuple[int, int, int], onnx.ModelProto] = {}
+        self.is_cpu = torch.device("cpu") == self.device
+
+    def _make_model(self, itype: int, ranka: int, rankb: int) -> onnx.ModelProto:
+        shapea = ["a{i}" for i in range(ranka)]
+        shapeb = ["b{i}" for i in range(rankb)]
+        shapec = ["c{i}" for i in range(max(ranka, rankb))]
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("MatMul", ["A", "B"], ["C"])],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("A", itype, shapea),
+                    oh.make_tensor_value_info("B", itype, shapeb),
+                ],
+                [oh.make_tensor_value_info("C", itype, shapec)],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        provider = "CPUExecutionProvider" if self.is_cpu else "CUDAExecutionProvider"
+        self._provider = provider
+        return InferenceSessionForTorch(
+            model,
+            optimized_model_filepath=_get_model_name("matmul", provider),
+            providers=[provider],
+        )
+
+    def run(self, a, b):
+        itype = torch_dtype_to_onnx_dtype(a.dtype)
+        ranka, rankb = len(a.shape), len(b.shape)
+        key = itype, ranka, rankb
+        if key not in self._cache:
+            self._cache[key] = self._make_model(itype, ranka, rankb)
+        sess = self._cache[key]
+        if self.verbose:
+            print(f"[MatMulOrt] running on {self._provider!r}")
+        feeds = dict(A=a.tensor, B=b.tensor)
+        got = sess.run(None, feeds)[0]
+        return OpRunTensor(got)

onnx_diagnostic/helpers/fake_tensor_helper.py

@@ -0,0 +1,273 @@
+from typing import Any, Dict, Optional, Set, Tuple
+
+
+class FakeTensorContext:
+    """Stores information used to reused same dimension for the same dimension names."""
+
+    def __init__(self, fake_mode: Optional["FakeTensorMode"] = None):  # noqa: F821
+        if fake_mode is None:
+            from torch.fx.experimental.symbolic_shapes import ShapeEnv
+            from torch._subclasses.fake_tensor import FakeTensorMode
+
+            shape_env = ShapeEnv()
+            self.fake_mode = FakeTensorMode(shape_env=shape_env)
+        else:
+            self.fake_mode = fake_mode
+        self._candidates = self._first_primes()
+        self._unique_: Set[str] = set()
+        self._mapping_int: Dict[int, str] = {}
+        self._mapping_str: Dict[str, int] = {}
+
+    @classmethod
+    def _first_primes(cls, n=1000):
+        sieve = [True] * (n + 1)
+        sieve[0:2] = [False, False]
+
+        for i in range(2, int(n**0.5) + 1):
+            if sieve[i]:
+                # Élimine les multiples de i
+                sieve[i * i : n + 1 : i] = [False] * len(range(i * i, n + 1, i))
+
+        return [i for i, prime in enumerate(sieve) if prime and i >= 13]
+
+    def _unique(self) -> int:
+        i = 0
+        c = self._candidates[i]
+        while c in self._unique_ or c in self._mapping_int:
+            i += 1
+            assert i < len(
+                self._candidates
+            ), f"Two many unique dimensions to generate, requested: {len(self._unique_)}"
+            c = self._candidates[i]
+        self._unique_.add(c)
+        return c
+
+    def from_tensor(self, x, static_shapes=False) -> "FakeTensor":  # noqa: F821
+        """
+        Returns a fake tensor.
+        ``pytorch`` returns the same name for the same dimension.
+        """
+        fake = self.fake_mode.from_tensor(x, static_shapes=static_shapes)
+        for i, s in zip(x.shape, fake.shape):
+            assert i not in self._mapping_int or self._mapping_int[i] == s, (
+                f"Inconsistency between {x.shape} and {fake.shape}, "
+                f"mapping has {self._mapping_int[i]} and s={s}"
+            )
+            self._mapping_int[i] = s
+        return fake
+
+    def fake_reshape(
+        self,
+        true_tensor: "torch.Tensor",  # noqa: F821
+        sh: Dict[int, Any],  # noqa: F821
+        fake_tensor: Optional["FakeTensor"] = None,  # noqa: F821
+    ) -> "FakeTensor":  # noqa: F821
+        """
+        Changes the shape of a true tensor to make it dynamic.
+
+        :param true_tensor: true tensor
+        :param sh: dynamic shape
+        :param fake_tensor: fake tensor, if None, make a fake one
+        :return: fake tensor
+        """
+        import torch
+
+        # deal with 0/1
+        for i in sh:
+            if true_tensor.shape[i] <= 1:
+                expanded_shape = list(true_tensor.shape)
+                expanded_shape[i] = self._unique()
+                true_tensor = torch.empty(
+                    tuple(expanded_shape), dtype=true_tensor.dtype, device=true_tensor.device
+                )
+
+        # deal with equivalent dimension
+        new_shape = list(true_tensor.shape)
+        mapping = {}
+        for i, s in sh.items():
+            d = true_tensor.shape[i]
+            if d not in mapping:
+                mapping[d] = s
+            elif mapping[d] != s:
+                d = self._unique()
+                mapping[d] = s
+                new_shape[i] = d
+        true_tensor = torch.empty(
+            tuple(new_shape), dtype=true_tensor.dtype, device=true_tensor.device
+        )
+
+        # now switch to FakeTensor
+        fake_tensor = self.from_tensor(true_tensor, static_shapes=False)
+        new_shape = list(true_tensor.shape)
+        for i in sh:
+            new_shape[i] = fake_tensor.shape[i]
+
+        reduced_tensor = self.from_tensor(true_tensor, static_shapes=True).sum(
+            axis=tuple(sorted(sh)), keepdim=True
+        )
+        return reduced_tensor.expand(*new_shape)
+
+    def make_fake(self, x: Any) -> Optional["FakeTensor"]:  # noqa: F821
+        """See :func:`onnx_diagnostic.helpers.fake_tensor_helper.make_fake`."""
+        if x is None:
+            return None
+        if isinstance(x, (list, tuple)):
+            return x.__class__([self.make_fake(i) for i in x])
+        if isinstance(x, dict):
+            return {k: self.make_fake(v) for k, v in x.items()}
+        if x.__class__.__name__ in {"DynamicCache", "StaticCache", "HybridCache"}:
+            assert hasattr(x, "layers"), (
+                f"Une more recent version of transformers (>=4.55), "
+                f"'layers' not found in class {type(x)}"
+            )
+            for layer in x.layers:
+                assert hasattr(layer, "keys") and hasattr(layer, "values"), (
+                    f"Une more recent version of transformers (>=4.55), 'layers' "
+                    f"not found in class {type(layer)} ({dir(layer)})"
+                )
+                layer.keys = self.make_fake(layer.keys)
+                layer.values = self.make_fake(layer.values)
+            return x
+        if x.__class__.__name__ == "EncoderDecoderCache":
+            self.make_fake(x.self_attention_cache)
+            self.make_fake(x.cross_attention_cache)
+            return x
+        if hasattr(x, "shape"):
+            return self.from_tensor(x, static_shapes=False)
+        from . import string_type
+
+        raise TypeError(
+            f"Unexpected type {type(x)} for x, content is {string_type(x, with_shape=True)}"
+        )
+
+    def make_fake_with_dynamic_dimensions(self, x: Any, dynamic_shapes: Any) -> Any:
+        """
+        See
+        :func:`onnx_diagnostic.export.shape_helper.make_fake_with_dynamic_dimensions`.
+        """
+        if x is None:
+            return None, None
+        if isinstance(x, (list, tuple)):
+            return x.__class__(
+                [
+                    self.make_fake_with_dynamic_dimensions(i, dynamic_shapes=ds)
+                    for i, ds in zip(x, dynamic_shapes)
+                ]
+            )
+        if isinstance(x, dict):
+            return {
+                k: self.make_fake_with_dynamic_dimensions(v, dynamic_shapes=dynamic_shapes[k])
+                for k, v in x.items()
+            }
+        if x.__class__.__name__ in {"DynamicCache", "StaticCache", "HybridCache"}:
+            assert hasattr(x, "layers"), (
+                f"Une more recent version of transformers (>=4.55), "
+                f"'layers' not found in class {type(x)}"
+            )
+            assert isinstance(dynamic_shapes, list) and (
+                not dynamic_shapes or not isinstance(dynamic_shapes[0], list)
+            ), f"Unexpected dynamic_shapes={dynamic_shapes} for a DynamicCache"
+            for il, layer in enumerate(x.layers):
+                assert hasattr(layer, "keys") and hasattr(layer, "values"), (
+                    f"Une more recent version of transformers (>=4.55), 'layers' "
+                    f"not found in class {type(layer)} ({dir(layer)})"
+                )
+                layer.keys = self.make_fake_with_dynamic_dimensions(
+                    layer.keys, dynamic_shapes=dynamic_shapes[il * 2]
+                )
+                layer.values = self.make_fake_with_dynamic_dimensions(
+                    layer.values, dynamic_shapes=dynamic_shapes[il * 2 + 1]
+                )
+            return x
+        if x.__class__.__name__ == "EncoderDecoderCache":
+            self.make_fake_with_dynamic_dimensions(
+                x.self_attention_cache, dynamic_shapes=dynamic_shapes[0]
+            )
+            self.make_fake_with_dynamic_dimensions(
+                x.cross_attention_cache, dynamic_shapes=dynamic_shapes[1]
+            )
+            return x
+        if hasattr(x, "shape"):
+            assert dynamic_shapes is None or isinstance(dynamic_shapes, dict), (
+                f"dynamic_shapes must be a dictionary at this stage but "
+                f"dynamic_shapes={dynamic_shapes}"
+            )
+            # We need to overwrite the values.
+            new_shape = []
+            for idim, dim in enumerate(x.shape):
+                if dynamic_shapes is not None and idim in dynamic_shapes:
+                    s = dynamic_shapes[idim]
+                    assert isinstance(s, str), (
+                        f"Unexpected type {type(s)} in dynamic_shapes={dynamic_shapes} "
+                        f"at index {idim}"
+                    )
+                    if s in self._mapping_str:
+                        dim = self._mapping_str[s]
+                    else:
+                        i = self._unique()
+                        self._mapping_str[s] = i
+                        dim = i
+                assert isinstance(dim, int), (
+                    f"Unexpected type {type(dim)}, dynamic_shapes={dynamic_shapes} "
+                    f"at index {idim}, dim={dim}"
+                )
+                new_shape.append(dim)
+            if tuple(new_shape) != x.shape:
+                import torch
+
+                x = torch.empty(tuple(new_shape), dtype=x.dtype, device=x.device)
+
+            t = self.fake_reshape(x, dynamic_shapes)  # type: ignore[arg-type]
+            assert t.device == x.device, f"device mismatch {x.device} -> {t.device}"
+            assert t.dtype == x.dtype, f"dtype mismatch {x.dtype} -> {t.dtype}"
+            return t
+        from ..helpers import string_type
+
+        raise TypeError(
+            f"Unexpected type {type(x)} for x, content is {string_type(x, with_shape=True)}"
+        )
+
+
+def make_fake(
+    x: Any, context: Optional[FakeTensorContext] = None
+) -> Tuple[Optional["FakeTensor"], Optional[FakeTensorContext]]:  # noqa: F821
+    """
+    Replaces all tensors by fake tensors.
+    This modification happens inplace for caches.
+    This function is only implemented for cache with
+    ``transformers>=4.55``.
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.helpers.fake_tensor_helper import make_fake
+
+        inputs, _ = make_fake(
+            dict(
+                input_ids=torch.randint(30360, size=(2, 3), dtype=torch.int64),
+                attention_mask=torch.randint(1, size=(2, 33), dtype=torch.int64),
+                position_ids=torch.randint(32, size=(2, 3), dtype=torch.int64),
+                past_key_values=make_dynamic_cache(
+                    [
+                        (
+                            torch.rand((2, 32, 30, 96), dtype=torch.float16),
+                            torch.rand((2, 32, 30, 96), dtype=torch.float16),
+                        ),
+                        (
+                            torch.rand((2, 32, 30, 96), dtype=torch.float16),
+                            torch.rand((2, 32, 30, 96), dtype=torch.float16),
+                        ),
+                    ]
+                ),
+            )
+        )
+        pprint.pprint(inputs)
+    """
+    if x is None:
+        return None, None
+    if context is None:
+        context = FakeTensorContext()
+    return context.make_fake(x), context