onnx-diagnostic 0.8.0__py3-none-any.whl

onnx_diagnostic/export/shape_helper.py

@@ -0,0 +1,296 @@
+from typing import Any, Dict, List, Set, Optional, Tuple, Union
+from ..helpers.cache_helper import flatten_unflatten_for_dynamic_shapes
+from .dynamic_shapes import ModelInputs
+
+
+def all_dynamic_shapes_from_inputs(inputs: Any, dim_prefix: Any = "d") -> Any:
+    """
+    Returns the dynamic shapes for the given inputs.
+    All dimensions are considered as dynamic.
+    ``dim_prefix`` can be a string (the function uses it as a prefix),
+    or ``torch.export.Dim.AUTO`` or ``torch.export.Dim.DYNAMIC``.
+    Depending on the version of transformers, serializations function
+    of DynamicCache class is automatically serialized or not (>= 4.51, < 4.55).
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs
+        from onnx_diagnostic.torch_export_patches import torch_export_patches
+
+        bsize, nheads, slen, dim = 2, 1, 30, 96
+        inputs = dict(
+            input_ids=torch.randint(15, size=(2, 3), dtype=torch.int64),
+            attention_mask=torch.randint(1, size=(2, 33), dtype=torch.int64),
+            position_ids=torch.arange(3, dtype=torch.int64),
+            past_key_values=make_dynamic_cache(
+                [(torch.randn(bsize, nheads, slen, dim),
+                    torch.randn(bsize, nheads, slen, dim))]
+            ),
+        )
+        with torch_export_patches(patch_transformers=True):
+            ds = all_dynamic_shapes_from_inputs(inputs)
+        pprint.pprint(ds)
+
+    For this function to work, patches must be enabled if :epkg:`transformers`
+    does not implement the serialization functions.
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import (
+            make_dynamic_cache,
+            make_encoder_decoder_cache,
+            make_mamba_cache,
+            make_sliding_window_cache,
+            make_static_cache,
+        )
+        from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs
+        from onnx_diagnostic.torch_export_patches import torch_export_patches
+
+        caches = [
+            make_dynamic_cache(
+                [
+                    (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                    (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                    (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                ]
+            ),
+            make_encoder_decoder_cache(
+                make_dynamic_cache(
+                    [
+                        (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                        (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                        (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                    ]
+                ),
+                make_dynamic_cache(
+                    [
+                        (torch.rand((5, 5, 5)), torch.rand((5, 5, 5))),
+                        (torch.rand((5, 5, 5)), torch.rand((5, 5, 5))),
+                        (torch.rand((5, 5, 5)), torch.rand((5, 5, 5))),
+                    ]
+                ),
+            ),
+            make_sliding_window_cache(
+                [
+                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
+                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
+                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
+                ]
+            ),
+            make_static_cache(
+                [
+                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
+                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
+                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
+                ],
+                max_cache_len=15,
+            ),
+            make_mamba_cache(
+                [
+                    (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                    (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                    (torch.rand((4, 4, 4)), torch.rand((4, 4, 4))),
+                ]
+            ),
+        ]
+
+        with torch_export_patches(patch_transformers=True):
+            for cache in caches:
+                print(f"-- {cache.__class__.__name__}")
+                pprint.pprint(all_dynamic_shapes_from_inputs(cache))
+    """
+    if isinstance(dim_prefix, str):
+        prefixes: Set[str] = set()
+
+        def tensor_to_shape(tensor):
+            n = len(prefixes)
+            p = f"{dim_prefix}_{n}"
+            prefixes.add(p)
+            return {i: f"{p}_{i}" for i in range(tensor.ndim)}
+
+    else:
+
+        def tensor_to_shape(tensor):
+            return {i: dim_prefix for i in range(tensor.ndim)}  # noqa: C420
+
+    return flatten_unflatten_for_dynamic_shapes(
+        inputs, change_function=tensor_to_shape, use_dict=True
+    )
+
+
+def guess_dynamic_shapes_from_inputs(
+    inputs: List[Any], auto: Union[bool, str] = False
+) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
+    """
+    Guesses which dimension is dimension from a set of inputs.
+    Every dimension having different values over multiple sets
+    of inputs. Every dimension not changing remains static.
+
+    :param inputs: a list of input sets
+    :param auto: True for ``torch.export.Dim.AUTO``,
+        False for ``torch.export.Dim.DYNAMIC``,
+        a string to get a unique string for every dynamic dimension
+    :return: args and kwargs
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.export.shape_helper import guess_dynamic_shapes_from_inputs
+
+        bsize, nheads, slen, dim = 2, 1, 30, 96
+        inputs1 = dict(
+            input_ids=torch.randint(15, size=(2, 3), dtype=torch.int64),
+            attention_mask=torch.randint(1, size=(2, 33), dtype=torch.int64),
+            position_ids=torch.arange(3, dtype=torch.int64),
+            past_key_values=make_dynamic_cache(
+                [
+                    (
+                        torch.randn(bsize, nheads, slen, dim),
+                        torch.randn(bsize, nheads, slen, dim),
+                    ),
+                ]
+            ),
+        )
+        bsize, nheads, slen, dim = 3, 1, 33, 96
+        inputs2 = dict(
+            input_ids=torch.randint(15, size=(3, 4), dtype=torch.int64),
+            attention_mask=torch.randint(1, size=(3, 34), dtype=torch.int64),
+            position_ids=torch.arange(4, dtype=torch.int64),
+            past_key_values=make_dynamic_cache(
+                [
+                    (
+                        torch.randn(bsize, nheads, slen, dim),
+                        torch.randn(bsize, nheads, slen, dim),
+                    ),
+                ]
+            ),
+        )
+        ds = guess_dynamic_shapes_from_inputs([inputs1, inputs2], auto="d")
+        pprint.pprint(ds)
+
+    This function returns something equivalent to function
+    :class:`torch.export.dynamic_shapes.AdditionalInputs` but this
+    one needs a model.
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.export.shape_helper import guess_dynamic_shapes_from_inputs
+        from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
+
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM", add_second_input=True)
+        ds = torch.export.dynamic_shapes.AdditionalInputs()
+        ds.add((), data["inputs"])
+        ds.add((), data["inputs2"])
+        pprint.pprint(ds.dynamic_shapes(data["model"], (), data["inputs"]))
+    """
+    mi = ModelInputs(None, inputs)
+    return mi.guess_dynamic_shapes(auto=auto)
+
+
+def make_fake_with_dynamic_dimensions(
+    x: Any, dynamic_shapes: Any, context: Optional["FakeTensorContext"] = None  # noqa: F821
+) -> Tuple[Any, "FakeTensorContext"]:  # noqa: F821
+    """
+    Replaces all tensors by fake tensor respecting the same
+    constraints as the following dynamic shapes.
+    This uses function :func:`onnx_diagnostic.helpers.fake_tensor_helper.make_fake`.
+    Parameter ``existing`` is used to reused the same object when the dynamic
+    dimension is given the same name as another one.
+
+    A simple tensor:
+
+    .. runpython::
+        :showcode:
+
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.export.shape_helper import make_fake_with_dynamic_dimensions
+
+        inputs, _ = make_fake_with_dynamic_dimensions(
+            torch.rand((2, 3, 4, 5), dtype=torch.float32),
+            {0: "batch", 2: "cache_length"},
+        )
+        print(inputs)
+
+    Two tensors:
+
+    .. runpython::
+        :showcode:
+
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.export.shape_helper import make_fake_with_dynamic_dimensions
+
+        inputs, _ = make_fake_with_dynamic_dimensions(
+            (
+                torch.rand((2, 3, 4, 5), dtype=torch.float32),
+                torch.rand((2, 3, 4, 5), dtype=torch.float32),
+            ),
+            ({0: "batch", 2: "cache_length"}, {0: "batch", 2: "cache_length"}),
+        )
+        print(inputs)
+
+    With a cache:
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        import torch
+        from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+        from onnx_diagnostic.export.shape_helper import make_fake_with_dynamic_dimensions
+
+        inputs, _ = make_fake_with_dynamic_dimensions(
+            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),
+                        ),
+                    ]
+                ),
+            ),
+            dynamic_shapes={
+                "input_ids": {0: "batch", 1: "seq_length"},
+                "attention_mask": {0: "batch", 1: "cache+seq"},
+                "position_ids": {0: "batch", 1: "seq_length"},
+                "past_key_values": [
+                    {0: "batch", 2: "cache_length"},
+                    {0: "batch", 2: "cache_length"},
+                    {0: "batch", 2: "cache_length"},
+                    {0: "batch", 2: "cache_length"},
+                ],
+            },
+        )
+        pprint.pprint(inputs)
+    """
+    if x is None:
+        return None, None
+    if context is None:
+        from ..helpers.fake_tensor_helper import FakeTensorContext
+
+        context = FakeTensorContext()
+
+    return context.make_fake_with_dynamic_dimensions(x, dynamic_shapes), context

onnx_diagnostic/export/validate.py

@@ -0,0 +1,173 @@
+import inspect
+import itertools
+import time
+from typing import Any, Dict, List, Optional, Tuple, Union
+import torch
+from ..helpers import string_type, max_diff, string_diff
+from ..helpers.torch_helper import torch_deepcopy
+from .dynamic_shapes import CoupleInputsDynamicShapes
+
+
+def compare_modules(
+    modep: torch.nn.Module,
+    mod: Optional[torch.nn.Module] = None,
+    args: Optional[Tuple[Any, ...]] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    copy: bool = False,
+    exc: bool = True,
+    verbose: int = 0,
+    atol: float = 1e-2,
+    rtol: float = 1e-1,
+) -> Dict[str, Any]:
+    """
+    Compares two torch modules, usually one coming from an exported program,
+    the other being the origin model.
+
+    :param model: first module
+    :param mod: second module (it produces the expected values)
+    :param args: positional arguments
+    :param kwargs: named arguments
+    :param copy: copy the inputs before executing the model (they may modify them inplace)
+    :param exc: raise exception if discrepancies are too high
+    :param verbose: verbosity level
+    :param atol: absolute tolerance
+    :param rtol: relative tolerance
+    :return: dictionary with inputs, outputs and tolerance
+
+    Example:
+
+    .. runpython::
+        :showcode:
+
+        import torch
+        from onnx_diagnostic.export import validate_ep, CoupleInputsDynamicShapes
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                return x + y
+
+        model = Model()
+        x = torch.randn((5, 6))
+        y = torch.randn((1, 6))
+        model(x, y)  # to make it is running
+
+        ds = ({0: "a", 1: "b"}, {1: "b"})
+        cpl = CoupleInputsDynamicShapes((x, y), {}, ds)
+        ep = torch.export.export(model, (x, y), dynamic_shapes=cpl.replace_string_by())
+        validate_ep(
+            ep,
+            model,
+            args=(x, y),
+            verbose=2,
+            copy=True,
+            dynamic_shapes=ds,
+            values_to_try={"a": [5, 10], "b": [10, 20]},
+        )
+
+    """
+    args = args or ()
+    kwargs = kwargs or {}
+
+    def _get(a):
+        return torch_deepcopy(a) if copy else a
+
+    if verbose:
+        begin = time.perf_counter()
+        print(
+            f"[compare_modules] check ep with "
+            f"args={string_type(args, with_shape=True, with_device=True)}, "
+            f"kwargs={string_type(kwargs, with_shape=True, with_device=True)}..."
+        )
+    got = modep(*_get(args), **_get(kwargs))
+    if verbose:
+        d = time.perf_counter() - begin
+        print(f"[compare_modules] done in {d} with output={string_type(got, with_shape=True)}")
+    if mod:
+        if verbose:
+            begin = time.perf_counter()
+            print("[compare_modules] run torch module...")
+        expected = mod(*_get(args), **_get(kwargs))
+        diff = max_diff(expected, got)
+        if verbose:
+            d = time.perf_counter() - begin
+            print(
+                f"[compare_modules] done in {d} with "
+                f"output={string_type(expected, with_shape=True)}"
+            )
+            print(f"[compare_modules] discrepancies={string_diff(diff)}")
+        assert not exc or (
+            isinstance(diff["abs"], float)
+            and isinstance(diff["rel"], float)
+            and diff["abs"] <= atol
+            and diff["rel"] <= rtol
+        ), f"Discrepancies={string_diff(diff)} higher than expected."
+        return dict(args=args, kwargs=kwargs, expected=expected, got=got, diff=diff)
+    return dict(args=args, kwargs=kwargs, got=got)
+
+
+def validate_ep(
+    ep: Union[torch.nn.Module, torch.export.ExportedProgram],
+    mod: Optional[torch.nn.Module] = None,
+    args: Optional[Tuple[Any, ...]] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    copy: bool = False,
+    dynamic_shapes: Optional[Any] = None,
+    values_to_try: Optional[Dict[str, List[int]]] = None,
+    exc: bool = True,
+    verbose: int = 0,
+    atol: float = 1e-2,
+    rtol: float = 1e-1,
+) -> List[Dict[str, Any]]:
+    """
+    Validates an exported program.
+
+    :param model: first module
+    :param mod: second module (it produces the expected values)
+    :param args: positional arguments
+    :param kwargs: named arguments
+    :param copy: copy the inputs before executing the model (they may modify them inplace)
+    :param dynamic_shapes: dynamic shapes, string should be used not ``torch.export.Dim``
+    :param values_to_try: dictionary with the values to try for every dynamic dimension
+    :param exc: raise exception if discrepancies are too high
+    :param verbose: verbosity level
+    :param atol: absolute tolerance
+    :param rtol: relative tolerance
+    :return: dictionary with inputs, outputs and tolerance
+    """
+    modep = ep.module() if isinstance(ep, torch.export.ExportedProgram) else ep
+
+    results = [
+        compare_modules(
+            modep, mod, args, kwargs, copy=copy, verbose=verbose, atol=atol, rtol=rtol
+        )
+    ]
+
+    assert (dynamic_shapes and values_to_try) or (
+        not dynamic_shapes and not values_to_try
+    ), "Either both dynamic_shapes and values_to_try are specified, either none."
+    if not dynamic_shapes or not values_to_try:
+        return results
+
+    items = list(values_to_try.items())
+    keys = [_[0] for _ in items]
+    values = [_[1] for _ in items]
+    all_vals = list(itertools.product(*values))
+    cpl = CoupleInputsDynamicShapes(
+        args or (),
+        kwargs or {},
+        dynamic_shapes,
+        args_names=(
+            list(inspect.signature(modep.forward).parameters) if args and kwargs else None
+        ),
+    )
+    for i, vals in enumerate(all_vals):
+        change_dims = dict(zip(keys, vals))
+        if verbose:
+            print(f"[validate_ep] try {i}/{len(all_vals)}: {change_dims}")
+        new_params = cpl.change_dynamic_dimensions(change_dims, args_kwargs=True)
+        na, nkw = new_params
+        c = compare_modules(
+            modep, mod, na, nkw, copy=copy, verbose=max(verbose - 1, 0), atol=atol, rtol=rtol
+        )
+        results.append(c)
+    return results