test

Author: xadupre

_unittests/ut_helpers/test_torch_test_helper.py

@@ -3,18 +3,22 @@
 import onnx
 import torch
 from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout
+from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.helpers.torch_test_helper import (
     dummy_llm,
     to_numpy,
     is_torchdynamo_exporting,
     steel_forward,
     replace_string_by_dynamic,
+    to_any,
+    torch_deepcopy,
 )
+from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache
 
 TFLOAT = onnx.TensorProto.FLOAT
 
 
-class TestOrtSession(ExtTestCase):
+class TestTorchTestHelper(ExtTestCase):
 
     def test_is_torchdynamo_exporting(self):
         self.assertFalse(is_torchdynamo_exporting())
@@ -67,6 +71,29 @@ def test_replace_string_by_dynamic(self):
             sproc,
         )
 
+    def test_to_any(self):
+        c1 = make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))])
+        c2 = make_encoder_decoder_cache(
+            make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]),
+            make_dynamic_cache([(torch.rand((5, 5, 5)), torch.rand((5, 5, 5)))]),
+        )
+        a = {"t": [(torch.tensor([1, 2]), c1, c2), {4, 5}]}
+        at = to_any(a, torch.float16)
+        self.assertIn("T10r", string_type(at))
+
+    def test_torch_deepcopy(self):
+        c1 = make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))])
+        c2 = make_encoder_decoder_cache(
+            make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]),
+            make_dynamic_cache([(torch.rand((5, 5, 5)), torch.rand((5, 5, 5)))]),
+        )
+        a = {"t": [(torch.tensor([1, 2]), c1, c2), {4, 5}]}
+        at = torch_deepcopy(a)
+        hash1 = string_type(at, with_shape=True, with_min_max=True)
+        c1.key_cache[0] += 1000
+        hash2 = string_type(at, with_shape=True, with_min_max=True)
+        self.assertEqual(hash1, hash2)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_models/test_test_helpers.py

@@ -1,6 +1,6 @@
 import copy
 import unittest
-from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout
 from onnx_diagnostic.torch_models.test_helper import get_inputs_for_task, validate_model
 from onnx_diagnostic.torch_models.hghub.model_inputs import get_get_inputs_function_for_tasks
 
@@ -15,13 +15,39 @@ def test_get_inputs_for_task(self):
             self.assertIn("dynamic_shapes", data)
             copy.deepcopy(data["inputs"])
 
+    @hide_stdout()
     def test_validate_model(self):
         mid = "arnir0/Tiny-LLM"
         summary, data = validate_model(mid, do_run=True, verbose=2)
         self.assertIsInstance(summary, dict)
         self.assertIsInstance(data, dict)
         validate_model(mid, do_run=True, verbose=2, quiet=True)
 
+    @hide_stdout()
+    def test_validate_model_dtype(self):
+        mid = "arnir0/Tiny-LLM"
+        summary, data = validate_model(
+            mid, do_run=True, verbose=2, dtype="float32", device="cpu"
+        )
+        self.assertIsInstance(summary, dict)
+        self.assertIsInstance(data, dict)
+        validate_model(mid, do_run=True, verbose=2, quiet=True)
+
+    @hide_stdout()
+    def test_validate_model_export(self):
+        mid = "arnir0/Tiny-LLM"
+        summary, data = validate_model(
+            mid,
+            do_run=True,
+            verbose=2,
+            dtype="float32",
+            device="cpu",
+            exporter="export-nostrict",
+        )
+        self.assertIsInstance(summary, dict)
+        self.assertIsInstance(data, dict)
+        validate_model(mid, do_run=True, verbose=2, quiet=False)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/_command_lines_parser.py

@@ -254,9 +254,12 @@ def get_parser_validate() -> ArgumentParser:
     parser.add_argument(
         "-e",
         "--export",
-        default=False,
-        action=BooleanOptionalAction,
-        help="runs the model to check it runs",
+        help="export the model with this exporter",
+    )
+    parser.add_argument(
+        "-o",
+        "--opt",
+        help="optimization to apply after the export",
     )
     parser.add_argument(
         "-r",
@@ -272,18 +275,22 @@ def get_parser_validate() -> ArgumentParser:
         action=BooleanOptionalAction,
         help="catches exception, report them in the summary",
     )
+    parser.add_argument(
+        "-p",
+        "--patch",
+        default=True,
+        action=BooleanOptionalAction,
+        help="applies patches before exporting",
+    )
     parser.add_argument(
         "--trained",
         default=False,
         action=BooleanOptionalAction,
         help="validate the trained model (requires downloading)",
     )
-    parser.add_argument(
-        "-v",
-        "--verbose",
-        default=0,
-        help="verbosity",
-    )
+    parser.add_argument("-v", "--verbose", default=0, type=int, help="verbosity")
+    parser.add_argument("--dtype", help="changes dtype if necessary")
+    parser.add_argument("--device", help="changes the device if necessary")
     return parser
 
 
@@ -316,9 +323,14 @@ def _cmd_validate(argv: List[Any]):
             verbose=args.verbose,
             quiet=args.quiet,
             trained=args.trained,
+            dtype=args.dtype,
+            device=args.device,
+            patch=args.patch,
+            optimization=args.opt,
+            exporter=args.export,
         )
         print("")
-        print("-- summary")
+        print("-- summary --")
         for k, v in sorted(summary.items()):
             print(f":{k},{v};")
 

onnx_diagnostic/helpers/torch_test_helper.py

@@ -1,7 +1,10 @@
 import contextlib
+from collections.abc import Iterable
 from typing import Any, Optional, Tuple, Union
+import numpy as np
 import torch
 from .helper import string_type
+from .cache_helper import make_dynamic_cache, make_encoder_decoder_cache
 
 
 def _forward_(*args, _f=None, _context=None, **kwargs):
@@ -298,3 +301,66 @@ def forward(self, input_ids):
         return dec, (x,)
 
     raise NotImplementedError(f"cls_name={cls_name}")
+
+
+def to_any(value: Any, to_value: Union[torch.dtype, torch.device]) -> Any:
+    """
+    Applies torch.to is applicables.
+    Goes recursively.
+    """
+    if isinstance(value, (torch.nn.Module, torch.Tensor)):
+        return value.to(to_value)
+    if isinstance(value, list):
+        return [to_any(t, to_value) for t in value]
+    if isinstance(value, tuple):
+        return tuple(to_any(t, to_value) for t in value)
+    if isinstance(value, set):
+        return {to_any(t, to_value) for t in value}
+    if isinstance(value, dict):
+        return {k: to_any(t, to_value) for k, t in value.items()}
+    if hasattr(value, "to"):
+        return value.to(to_value)
+    if value.__class__.__name__ == "DynamicCache":
+        return make_dynamic_cache(
+            list(
+                zip(
+                    [t.to(to_value) for t in value.key_cache],
+                    [t.to(to_value) for t in value.value_cache],
+                )
+            )
+        )
+
+    assert not isinstance(value, Iterable), f"Unsupported type {type(value)}"
+    return value
+
+
+def torch_deepcopy(value: Any) -> Any:
+    """
+    Makes a deepcopy.
+    """
+    if isinstance(value, (int, float, str)):
+        return value
+    if isinstance(value, tuple):
+        return tuple(torch_deepcopy(v) for v in value)
+    if isinstance(value, list):
+        return [torch_deepcopy(v) for v in value]
+    if isinstance(value, set):
+        return {torch_deepcopy(v) for v in value}
+    if isinstance(value, dict):
+        return {k: torch_deepcopy(v) for k, v in value.items()}
+    if isinstance(value, np.ndarray):
+        return value.copy()
+    if hasattr(value, "clone"):
+        return value.clone()
+    if value.__class__.__name__ == "DynamicCache":
+        return make_dynamic_cache(
+            torch_deepcopy(list(zip(value.key_cache, value.value_cache)))
+        )
+    if value.__class__.__name__ == "EncoderDecoderCache":
+        return make_encoder_decoder_cache(
+            torch_deepcopy(value.self_attention_cache),
+            torch_deepcopy(value.cross_attention_cache),
+        )
+    # We should have a code using serialization, deserialization assuming a model
+    # cannot be exported without them.
+    raise NotImplementedError(f"torch_deepcopy not implemented for type {type(value)}")