onnx-diagnostic 0.6.2__py3-none-any.whl → 0.7.0__py3-none-any.whl

onnx_diagnostic/helpers/model_builder_helper.py

@@ -3,9 +3,9 @@ import os
 import requests
 import sys
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Optional, Union
 from urllib.parse import urlparse
-from onnx import helper, save_model, external_data_helper, ModelProto
+from onnx import ModelProto, TensorProto
 
 CACHE_SUBDIR = "onnx-diagnostic"
 
@@ -114,87 +114,58 @@ def _make_model(self, model, verbose: int = 0):
                 self.make_lm_head(module)
 
 
-def save_model_builder(self, out_dir: Optional[str] = "", verbose: int = 0) -> ModelProto:
+def save_model_builder(
+    self, out_dir: Optional[str] = "", verbose: int = 0
+) -> Union[str, ModelProto]:
     """
     Saves a model created by function :func:`create_model_builder`.
     If out_dir is empty or not specified, the function still returns the
     generated model.
     """
-    if verbose:
-        print(f"[save_model_builder] Saving ONNX model in {out_dir}")
-
-    # Create ONNX model
-    model = helper.make_model(
-        opset_imports=[
-            self.clear_field(
-                helper.make_operatorsetid("", 21 if self.quant_attrs["use_qdq"] else 14),
-                "domain",
-            ),
-            helper.make_operatorsetid("com.microsoft", 1),
-        ],
-        ir_version=7,
-        producer_name="onnxruntime-genai",
-        producer_version="0.0.0",
-        graph=self.make_graph(
-            name="main_graph",
-            inputs=self.inputs,
-            outputs=self.outputs,
-            initializer=self.initializers,
-            value_info=self.value_infos,
-            nodes=self.nodes,
-        ),
-    )
-
-    # Load external data into ONNX model
-    external_data_helper.load_external_data_for_model(model, self.cache_dir)
-
-    # Delete external data files on disk before re-saving
-    for path in os.listdir(self.cache_dir):
-        if path.endswith(".bin"):
-            os.remove(os.path.join(self.cache_dir, path))
+    import onnx_ir
 
-    # Delete temporary cache dir if empty
-    # if len(os.listdir(self.cache_dir)) == 0:
-    #    os.rmdir(self.cache_dir)
+    if verbose:
+        print(f"[save_model_builder] Saving ONNX model in {out_dir!r}")
 
-    # Quantize ONNX model to desired precision
+    # Skip quantizing `MatMul` in `DequantizeLinear --> Transpose --> MatMul` path
     already_quantized_in_qdq_format = (
         self.quant_type is not None and self.quant_attrs["use_qdq"]
-    )  # Skip quantizing `MatMul` in `DequantizeLinear --> Transpose --> MatMul` path
-    if self.onnx_dtype == "int4" and not already_quantized_in_qdq_format:
-        model = self.to_int4(model)
+    )
+    model = (
+        self.to_int4()
+        if self.onnx_dtype in {onnx_ir.DataType.INT4, onnx_ir.DataType.UINT4}
+        and not already_quantized_in_qdq_format
+        else self.model
+    )
+    model.graph.sort()
+    if not out_dir:
+        return onnx_ir.to_proto(model)
 
-    # Save ONNX model with only one external data file and delete any existing duplicate copies
-    if out_dir:
-        out_path = os.path.join(out_dir, self.filename)
-        data_path = os.path.join(out_dir, os.path.basename(out_path) + ".data")
-        if os.path.exists(out_path):
-            if verbose:
-                print(f"[save_model_builder] Overwriting {out_path!r}")
-            os.remove(out_path)
-        if os.path.exists(data_path):
-            if verbose:
-                print(f"[save_model_builder] Overwriting {data_path!r}")
-            os.remove(data_path)
+    out_path = os.path.join(out_dir, self.filename)
+    data_path = os.path.join(out_dir, os.path.basename(out_path) + ".data")
 
-    if out_dir:
-        location = os.path.basename(data_path)
-        if os.path.exists(location):
-            os.remove(location)
+    # Save ONNX model with only one external data file and delete any existing duplicate copies
+    out_path = os.path.join(out_dir, self.filename)
+    data_path = os.path.join(out_dir, os.path.basename(out_path) + ".data")
+    if os.path.exists(out_path):
         if verbose:
-            print(f"[save_model_builder] out_path={out_path!r}")
-            print(f"[save_model_builder] location={location!r}")
-        save_model(
-            model,
-            out_path,
-            save_as_external_data=True,
-            all_tensors_to_one_file=True,
-            location=location,
-            size_threshold=1024,
-            convert_attribute=False,
-        )
-        return None
-    return model
+            print(f"[save_model_builder] Overwriting {out_path!r}")
+        os.remove(out_path)
+    if os.path.exists(data_path):
+        if verbose:
+            print(f"[save_model_builder] Overwriting {data_path!r}")
+        os.remove(data_path)
+
+    onnx_ir.save(
+        model,
+        out_path,
+        external_data=os.path.basename(data_path),
+        size_threshold_bytes=2**10,
+    )
+    if verbose:
+        print(f"[save_model_builder] saved in {out_dir!r}")
+
+    return out_path
 
 
 def create_model_builder(
@@ -220,6 +191,9 @@ def create_model_builder(
     """
     assert cache_dir, "create_model_builder does not work without cache_dir."
     assert os.path.exists(cache_dir), f"cache_dir={cache_dir!r} does not exists"
+    precision = {"float32": "fp32", "float16": "fp16", "bfloat16": "bfp16"}.get(
+        precision, precision
+    )
     download_model_builder_to_cache()
     builder = import_model_builder()
     io_dtype = builder.set_io_dtype(precision, execution_provider, extra_options)
@@ -332,13 +306,23 @@ def create_model_builder(
     for c in remove:
         delattr(config, c)
 
-    onnx_model = cls(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
+    convert = {
+        "fp32": TensorProto.FLOAT,
+        "fp16": TensorProto.FLOAT16,
+        "bfp16": TensorProto.BFLOAT16,
+    }
+    assert (
+        precision in convert
+    ), f"Unexpected value for precision={precision!r}, should be in {convert}"
+    onnx_model = cls(
+        config, io_dtype, convert[precision], execution_provider, cache_dir, extra_options
+    )
 
     if post:
         post(onnx_model)
     _make_model(onnx_model, model, verbose=verbose)
 
-    assert onnx_model.nodes, (
+    assert onnx_model.model, (
         f"No node in the model, io_dtype={io_dtype!r}, "
         f"precision={precision!r}, execution_provider={execution_provider!r}, "
         f"extra_options={extra_options!r}, cache_dir={cache_dir!r}, "

onnx_diagnostic/helpers/onnx_helper.py

@@ -316,7 +316,7 @@ def check_model_ort(
 
 
 @functools.cache
-def onnx_dtype_name(itype: int) -> str:
+def onnx_dtype_name(itype: int, exc: bool = True) -> str:
     """
     Returns the ONNX name for a specific element type.
 
@@ -335,7 +335,11 @@ def onnx_dtype_name(itype: int) -> str:
             v = getattr(TensorProto, k)
             if v == itype:
                 return k
-    raise ValueError(f"Unexpected value itype: {itype}")
+    if exc:
+        raise ValueError(f"Unexpected value itype: {itype}")
+    if itype == 0:
+        return "UNDEFINED"
+    return "UNEXPECTED"
 
 
 def pretty_onnx(
@@ -365,7 +369,7 @@ def pretty_onnx(
         itype = onx.type.tensor_type.elem_type
         shape = tuple((d.dim_param or d.dim_value) for d in onx.type.tensor_type.shape.dim)
         shape_str = ",".join(map(str, shape))
-        return f"{onnx_dtype_name(itype)}[{shape_str}] {name}"
+        return f"{onnx_dtype_name(itype, exc=False)}[{shape_str}] {name}"
 
     if isinstance(onx, AttributeProto):
         att = onx
@@ -767,7 +771,7 @@ def tensor_dtype_to_np_dtype(tensor_dtype: int) -> np.dtype:
 
 
 def iterator_initializer_constant(
-    model: Union[onnx.FunctionProto, onnx.GraphProto, onnx.ModelProto],
+    model: Union[FunctionProto, GraphProto, ModelProto],
     use_numpy: bool = True,
     prefix: str = "",
 ) -> Iterator[Tuple[str, Union["torch.Tensor", np.ndarray]]]:  # noqa: F821
@@ -779,8 +783,8 @@ def iterator_initializer_constant(
     :param prefix: for subgraph
     :return: iterator
     """
-    if not isinstance(model, onnx.FunctionProto):
-        graph = model if isinstance(model, onnx.GraphProto) else model.graph
+    if not isinstance(model, FunctionProto):
+        graph = model if isinstance(model, GraphProto) else model.graph
         if not use_numpy:
             from .torch_helper import to_tensor
         if prefix:
@@ -791,7 +795,7 @@ def iterator_initializer_constant(
             )
         nodes = graph.node
         name = graph.name
-        if isinstance(model, onnx.ModelProto):
+        if isinstance(model, ModelProto):
             for f in model.functions:
                 yield from iterator_initializer_constant(
                     f, use_numpy=use_numpy, prefix=f"{prefix}{f.name}"
@@ -908,3 +912,283 @@ def tensor_statistics(tensor: Union[np.ndarray, TensorProto]) -> Dict[str, Union
     qu = np.quantile(tensor, ii)
     stat.update({f"q{i}": float(q) for i, q in zip(ii, qu)})
     return stat
+
+
+class NodeCoordinates:
+    """
+    A way to localize a node,
+    path is a tuple of three information, node index, node type, node name.
+    """
+
+    __slots__ = ("node", "path")
+
+    def __init__(
+        self,
+        node: Union[onnx.TensorProto, NodeProto, str],
+        path: Tuple[Tuple[int, str, str], ...],
+    ):
+        assert isinstance(path, tuple), f"Unexpected type {type(path)} for path"
+        assert all(isinstance(t, tuple) for t in path), f"Unexpected type in path={path}"
+        self.node = node
+        self.path = path
+
+    def __str__(self) -> str:
+        "usual"
+        if isinstance(self.node, str):
+            return f"{self.path_to_str()} :: {self.node!r}"
+        return f"{self.path_to_str()} :: {pretty_onnx(self.node)}"
+
+    def path_to_str(self) -> str:
+        "Strings representing coordinates."
+        return "x".join(f"({':'.join(map(str, t))})" for t in self.path)
+
+
+class ResultFound:
+    """
+    Class returned by :func:`enumerate_results`.
+    """
+
+    __slots__ = ("consumer", "name", "producer")
+
+    def __init__(
+        self,
+        name: str,
+        producer: Optional[NodeCoordinates],
+        consumer: Optional[NodeCoordinates],
+    ):
+        assert isinstance(name, str), f"unexpected type {type(name)} for name"
+        self.name = name
+        self.producer = producer
+        self.consumer = consumer
+
+    def __str__(self) -> str:
+        "usuals"
+        return (
+            f"<< {self.name} - {self.consumer}"
+            if self.producer is None
+            else f">> {self.name} - {self.producer}"
+        )
+
+
+def enumerate_results(
+    proto: Union[FunctionProto, GraphProto, ModelProto, Sequence[NodeProto]],
+    name: Union[Set[str], str],
+    verbose: int = 0,
+    coordinates: Optional[List[Tuple[int, str, str]]] = None,
+) -> Iterator[ResultFound]:
+    """
+    Iterates on all nodes, attributes to find where a name is used.
+
+    :param proto: a proto
+    :param name: name or names to find
+    :param verbose: verbosity
+    :param coordinates: coordinates of a node
+    :return: iterator on :class:`ResultFound`
+    """
+    if not isinstance(name, set):
+        name = {name}
+    coordinates = coordinates or []
+    assert all(
+        isinstance(c, tuple) for c in coordinates
+    ), f"Unexpected type in coordinates={coordinates}"
+    indent = "  " * len(coordinates)
+    if isinstance(proto, ModelProto):
+        if verbose:
+            print(f"[enumerate_results] {indent}searching for {name!r} into ModelProto...")
+        yield from enumerate_results(proto.graph, name, verbose=verbose)
+    elif isinstance(proto, FunctionProto):
+        if verbose:
+            print(f"[enumerate_results] {indent}searching for {name!r} into FunctionProto...")
+        for i in proto.input:
+            if i in name:
+                r = ResultFound(
+                    i,
+                    NodeCoordinates(i, tuple([*coordinates, (-1, "INPUT", "")])),  # noqa: C409
+                    None,
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+        yield from enumerate_results(proto.node, name, verbose=verbose)
+        for i in proto.output:
+            if i in name:
+                r = ResultFound(
+                    i,
+                    None,
+                    NodeCoordinates(
+                        i, tuple([*coordinates, (len(proto.node), "OUTPUT", "")])  # noqa: C409
+                    ),
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+    elif isinstance(proto, GraphProto):
+        if verbose:
+            print(f"[enumerate_results] {indent}searching for {name!r} into GraphProto...")
+        for i in proto.initializer:
+            if i.name in name:
+                r = ResultFound(
+                    i.name,
+                    NodeCoordinates(i, tuple([*coordinates, (-1, "INIT", "")])),  # noqa: C409
+                    None,
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+        for i in proto.sparse_initializer:
+            if i.name in name:
+                r = ResultFound(
+                    i.name,
+                    NodeCoordinates(i, tuple([*coordinates, (-1, "INIT", "")])),  # noqa: C409
+                    None,
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+        for i in proto.input:
+            if i.name in name:
+                r = ResultFound(
+                    i.name,
+                    NodeCoordinates(i, tuple([*coordinates, (-1, "INPUT", "")])),  # noqa: C409
+                    None,
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+        yield from enumerate_results(
+            proto.node, name, verbose=verbose, coordinates=coordinates
+        )
+        for i in proto.output:
+            if i.name in name:
+                r = ResultFound(
+                    i.name,
+                    None,
+                    NodeCoordinates(
+                        i, tuple([*coordinates, (len(proto.node), "OUTPUT", "")])  # noqa: C409
+                    ),
+                )
+                if verbose > 1:
+                    print(f"[enumerate_results] {indent}-- {r}")
+                yield r
+    else:
+        if verbose:
+            print(
+                f"[enumerate_results] {indent}searching for {name!r} into List[NodeProto]..."
+            )
+        for node_i, node in enumerate(proto):
+            if set(node.input) & name:
+                for n in node.input:
+                    if n in name:
+                        r = ResultFound(
+                            n,
+                            NodeCoordinates(
+                                node,
+                                tuple(  # noqa: C409
+                                    [*coordinates, (node_i, node.op_type, node.name)]
+                                ),
+                            ),
+                            None,
+                        )
+                        if verbose > 1:
+                            print(f"[enumerate_results] {indent}-- {r}")
+                        yield r
+            if node.op_type in {"If", "Scan", "Loop", "SequenceMap"}:
+                for att in node.attribute:
+                    if att.type == onnx.AttributeProto.GRAPH:
+                        yield from enumerate_results(
+                            att.g,
+                            name,
+                            verbose=verbose,
+                            coordinates=[*coordinates, (node_i, node.op_type, node.name)],
+                        )
+            if set(node.output) & name:
+                for n in node.output:
+                    if n in name:
+                        r = ResultFound(
+                            n,
+                            None,
+                            NodeCoordinates(
+                                node,
+                                tuple(  # noqa: C409
+                                    [*coordinates, (node_i, node.op_type, node.name)]
+                                ),
+                            ),
+                        )
+                        if verbose > 1:
+                            print(f"[enumerate_results] {indent}-- {r}")
+                        yield r
+    if verbose:
+        print(f"[enumerate_results] {indent}done")
+
+
+def shadowing_names(
+    proto: Union[FunctionProto, GraphProto, ModelProto, Sequence[NodeProto]],
+    verbose: int = 0,
+    existing: Optional[Set[str]] = None,
+    shadow_context: Optional[Set[str]] = None,
+    post_shadow_context: Optional[Set[str]] = None,
+) -> Tuple[Set[str], Set[str], Set[str]]:
+    """
+    Returns the shadowing names, the names created in the main graph
+    after they were created in a subgraphs and the names created by the nodes.
+    """
+    if isinstance(proto, ModelProto):
+        return shadowing_names(proto.graph)
+    if isinstance(proto, GraphProto):
+        assert (
+            existing is None and shadow_context is None
+        ), "existing must be None if nodes is None"
+        return shadowing_names(
+            proto.node,
+            verbose=verbose,
+            existing=set(i.name for i in proto.initializer)
+            | set(i.name for i in proto.sparse_initializer)
+            | set(i.name for i in proto.input if i.name),
+            shadow_context=set(),
+            post_shadow_context=set(),
+        )
+    if isinstance(proto, FunctionProto):
+        assert (
+            existing is None and shadow_context is None
+        ), "existing must be None if nodes is None"
+        return shadowing_names(
+            proto.node,
+            verbose=verbose,
+            existing=set(i for i in proto.input if i),
+            shadow_context=set(),
+            post_shadow_context=set(),
+        )
+
+    assert (
+        existing is not None and shadow_context is not None
+    ), "existing must not be None if nodes is not None"
+    shadow = set()
+    shadow_context = shadow_context.copy()
+    existing = existing.copy()
+    created = set()
+    post_shadow = set()
+    for node in proto:
+        not_empty = set(n for n in node.input if n)
+        intersection = not_empty & existing
+        assert len(intersection) == len(not_empty), (
+            f"One input in {not_empty}, node={pretty_onnx(node)} "
+            f"was not found in {existing}"
+        )
+        for att in node.attribute:
+            if att.type == AttributeProto.GRAPH:
+                g = att.g
+                shadow |= set(i.name for i in g.input) & shadow_context
+                shadow |= set(i.name for i in g.initializer) & shadow_context
+                shadow |= set(i.name for i in g.sparse_initializer) & shadow_context
+                s, ps, c = shadowing_names(
+                    g.node, verbose=verbose, existing=existing, shadow_context=existing
+                )
+                shadow |= s
+                created |= c
+
+        not_empty = set(n for n in node.output if n)
+        post_shadow |= not_empty & created
+        shadow |= not_empty & shadow_context
+        existing |= not_empty
+        created |= not_empty
+    return shadow, post_shadow, created

onnx_diagnostic/helpers/torch_helper.py

@@ -16,6 +16,7 @@ from .cache_helper import (
     make_encoder_decoder_cache,
     make_sliding_window_cache,
     make_mamba_cache,
+    make_static_cache,
 )
 from .mini_onnx_builder import create_onnx_model_from_input_tensors
 from .onnx_helper import (
@@ -288,7 +289,8 @@ def steal_forward(
     """
     The necessary modification to steem forward method and prints out inputs
     and outputs using :func:`onnx_diagnostic.helpers.string_type`.
-    See example :ref:`l-plot-tiny-llm-export`.
+    See example :ref:`l-plot-tiny-llm-export` or
+    :ref:`l-plot-intermediate-results`.
 
     :param model: a model or a list of models to monitor,
         every model can also be a tuple(name, model), name is displayed well.
@@ -410,12 +412,15 @@ def steal_forward(
             proto = create_onnx_model_from_input_tensors(storage)
             if verbose:
                 print("-- dumps stored objects")
+            location = f"{os.path.split(dump_file)[-1]}.data"
+            if os.path.exists(location):
+                os.remove(location)
             onnx.save(
                 proto,
                 dump_file,
                 save_as_external_data=True,
                 all_tensors_to_one_file=True,
-                location=f"{os.path.split(dump_file)[-1]}.data",
+                location=location,
             )
             if verbose:
                 print("-- done dump stored objects")
@@ -723,6 +728,15 @@ def to_any(value: Any, to_value: Union[torch.dtype, torch.device, str]) -> Any:
                 )
             )
         )
+    if value.__class__.__name__ == "StaticCache":
+        return make_static_cache(
+            list(
+                zip(
+                    [t.to(to_value) for t in value.key_cache],
+                    [t.to(to_value) for t in value.value_cache],
+                )
+            )
+        )
     if value.__class__.__name__ == "EncoderDecoderCache":
         return make_encoder_decoder_cache(
             to_any(value.self_attention_cache, to_value),
@@ -769,6 +783,8 @@ def torch_deepcopy(value: Any) -> Any:
         return make_dynamic_cache(
             torch_deepcopy(list(zip(value.key_cache, value.value_cache)))
         )
+    if value.__class__.__name__ == "StaticCache":
+        return make_static_cache(torch_deepcopy(list(zip(value.key_cache, value.value_cache))))
     if value.__class__.__name__ == "SlidingWindowCache":
         return make_sliding_window_cache(
             torch_deepcopy(list(zip(value.key_cache, value.value_cache)))

onnx_diagnostic/reference/__init__.py

@@ -1,3 +1,4 @@
 from .evaluator import ExtendedReferenceEvaluator
 from .ort_evaluator import OnnxruntimeEvaluator
 from .torch_evaluator import TorchOnnxEvaluator
+from .report_results_comparison import ReportResultComparison