check llms separatly

Author: xadupre

.github/workflows/ci.yml

@@ -52,6 +52,12 @@ jobs:
       - name: pip freeze
         run: python -m pip freeze
 
+      - name: tiny-llm torch.export.export
+        run: python _unittests/ut_torch_models/test_llms.py
+
+      - name: tiny-llm onnx
+        run: python _unittests/ut_torch_models/test_llms_onnx.py
+
       - name: run tests
         run: |
           pip install pytest

.github/workflows/documentation.yml

@@ -39,8 +39,7 @@ jobs:
         run: python -m pip install -r requirements.txt
 
       - name: Install requirements dev
-        run: |
-          python -m pip install -r requirements-dev.txt
+        run: python -m pip install -r requirements-dev.txt
 
       - name: Cache pip
         uses: actions/cache@v4
@@ -54,6 +53,12 @@ jobs:
       - name: pip freeze
         run: python -m pip freeze
 
+      - name: tiny-llm torch.export.export
+        run: python _unittests/ut_torch_models/test_llms.py
+
+      - name: tiny-llm onnx
+        run: python _unittests/ut_torch_models/test_llms_onnx.py
+
       - name: Generate coverage report
         run: |
           pip install pytest

_unittests/ut_torch_models/test_llms.py

@@ -17,16 +17,17 @@ def test_get_tiny_llm(self):
     def test_export_tiny_llm_1(self):
         data = get_tiny_llm()
         model, inputs = data["model"], data["inputs"]
+        self.assertEqual({"attention_mask", "past_key_values", "input_ids"}, set(inputs))
         ep = torch.export.export(
             model, (), kwargs=inputs, dynamic_shapes=data["dynamic_shapes"]
         )
         assert ep
-        print(ep)
 
     @ignore_warnings(UserWarning)
     def test_export_tiny_llm_2_bypassed(self):
         data = get_tiny_llm()
         model, inputs = data["model"], data["inputs"]
+        self.assertEqual({"attention_mask", "past_key_values", "input_ids"}, set(inputs))
         with bypass_export_some_errors():
             ep = torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=data["dynamic_shapes"]

_unittests/ut_xrun_doc/test_torch_test_helper.py

@@ -6,13 +6,21 @@
 import onnx.numpy_helper as onh
 import torch
 from onnx_diagnostic.ext_test_case import ExtTestCase
-from onnx_diagnostic.torch_test_helper import dummy_llm, check_model_ort, to_numpy
+from onnx_diagnostic.torch_test_helper import (
+    dummy_llm,
+    check_model_ort,
+    to_numpy,
+    is_torchdynamo_exporting,
+)
 
 TFLOAT = onnx.TensorProto.FLOAT
 
 
 class TestOrtSession(ExtTestCase):
 
+    def test_is_torchdynamo_exporting(self):
+        self.assertFalse(is_torchdynamo_exporting())
+
     def test_dummy_llm(self):
         for cls_name in ["AttentionBlock", "MultiAttentionBlock", "DecoderLayer", "LLM"]:
             model, inputs = dummy_llm(cls_name)

onnx_diagnostic/torch_test_helper.py

@@ -9,10 +9,23 @@
 from .helpers import pretty_onnx
 
 
+def is_torchdynamo_exporting() -> bool:
+    """Tells if torch is exporting a model."""
+    import torch
+
+    try:
+        return torch.compiler.is_exporting()
+    except Exception:
+        try:
+            import torch._dynamo as dynamo
+
+            return dynamo.is_exporting()
+        except Exception:
+            return False
+
+
 def to_numpy(tensor: "torch.Tensor"):  # noqa: F821
-    """
-    Converts a torch tensor to numy.
-    """
+    """Converts a torch tensor to numy."""
     try:
         return tensor.numpy()
     except TypeError: