Refactor test_trained_tiny_llm.py to match test_random_tiny_llm.py pattern and add CI job (#87)

Copilot · xadupre · web-flow · commit 0af5387562e1 · 2026-04-08T22:46:32.000+02:00
* Initial plan * Replicate refactoring from test_random_tiny_llm.py into test_trained_tiny_llm.py Agent-Logs-Url: https://github.com/xadupre/mbext/sessions/2d48c936-414f-4d71-a1c1-8c6ac81c053e Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com> * Add CI job for trained tiny LLM tests Agent-Logs-Url: https://github.com/xadupre/mbext/sessions/074a6487-04d5-41b1-92f9-8bd78cacc5bf Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com> * fix * fix * Pass HF_TOKEN secret to trained tiny LLM test step Agent-Logs-Url: https://github.com/xadupre/mbext/sessions/ac629472-feba-4e00-9e8b-21d0c8e8d1a2 Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com> * Rename HF_TOKEN to HUGGING_FACE_HUB_TOKEN and add continue-on-error for connectivity tolerance Agent-Logs-Url: https://github.com/xadupre/mbext/sessions/f159bfd1-b592-44be-9bac-3e275644e1b2 Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com> Co-authored-by: Xavier Dupré <xadupre@microsoft.com>
diff --git a/.github/workflows/trained_tiny_llm_tests.yml b/.github/workflows/trained_tiny_llm_tests.yml
@@ -0,0 +1,56 @@
+name: Trained Tiny LLM Tests
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+    branches: ["main"]
+
+permissions:
+  contents: read
+
+jobs:
+  tests:
+    name: ci (${{ matrix.os }}, py${{ matrix.python-version }}, torch-${{ matrix.torch-version || 'stable' }}, transformers-${{ matrix.transformers-version || 'latest' }})
+    runs-on: ${{ matrix.os }}
+    permissions:
+      contents: read
+      id-token: write
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["ubuntu-latest"]
+        python-version: ["3.13"]
+        transformers-version: ["5.5"]
+        torch-version: ["2.10"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+
+      # we install torch first to avoid downloading any CUDA dependency
+
+      - name: Install nightly pytorch
+        if: matrix.torch-version == 'nightly'
+        run: pip install --pre --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu
+
+      - name: Install pytorch 2.10 (Linux/Windows)
+        if: matrix.torch-version == '2.10' && matrix.os != 'macos-latest'
+        run: pip install torch==2.10.0+cpu torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+
+      - name: Install dependencies
+        run: pip install -e ".[dev]" torch transformers tokenizers
+
+      - name: Run trained tiny LLM tests
+        continue-on-error: true
+        env:
+          HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
+        run: LONGTEST=1 pytest tests/trained/test_trained_tiny_llm.py -v --cov=modelbuilder --cov-report=xml --cov-report=term-missing
+
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v5
+        with:
+          files: coverage.xml
+          flags: trained-tiny-llm-tests
+          fail_ci_if_error: false
diff --git a/tests/trained/test_trained_tiny_llm.py b/tests/trained/test_trained_tiny_llm.py
@@ -6,22 +6,21 @@
 import os
 import unittest
 
-import numpy as np
-
-from modelbuilder.ext_test_case import ExtTestCase
+from modelbuilder.ext_test_case import ExtTestCase, hide_stdout, long_test
 
 MODEL_NAME = "arnir0/Tiny-LLM"
 
 
 class TestTrainedTinyLLM(ExtTestCase):
-    def _common_part(self, precision, dtype, int4=False):
+    def _common_part(self, precision, dtype, provider="cpu", int4=False):
         from transformers import AutoModelForCausalLM
+
         from modelbuilder.builder import create_model
 
         # Use 4 layers so that both rope (layers 0-2) and no-rope (layer 3)
         # code paths are exercised with the default no_rope_layer_interval=4.
         output_dir, cache_dir = self.get_dirs(
-            f"test_trained_tiny_llm_{'int4' if int4 else precision}_cpu", clean=False
+            f"test_trained_tiny_llm_{'int4' if int4 else precision}_{provider}", clean=False
         )
         onnx_path = os.path.join(output_dir, "model.onnx")
         if not os.path.exists(onnx_path):
@@ -30,7 +29,7 @@ def _common_part(self, precision, dtype, int4=False):
                 model_name=MODEL_NAME,
                 input_path="",
                 precision="int4" if int4 else precision,
-                execution_provider="cpu",
+                execution_provider=provider,
                 output_dir=output_dir,
                 cache_dir=cache_dir,
             )
@@ -44,22 +43,20 @@ def _common_part(self, precision, dtype, int4=False):
         model.eval()
         return onnx_path, model
 
-    def test_trained_tiny_llm_fp32_discrepancies_cpu(self):
-        """
-        Convert arnir0/Tiny-LLM to an fp32 ONNX model targeting the CPU execution
-        provider.  Only one hidden layer is materialised (via the
-        ``num_hidden_layers`` extra option) to keep the test fast.
-        The test verifies that:
-        * ``create_model`` completes without error.
-        * The expected ``model.onnx`` file is written to the output directory.
-        * The produced ONNX file passes ``onnx.checker.check_model``.
-        """
+    def _dtype_from_precision(self, precision):
         import torch
 
-        precision, dtype, np_dtype = "fp32", torch.float32, np.float32
+        dtype_map = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}
+        return dtype_map.get(precision, torch.float32)
 
-        onnx_path, model = self._common_part(precision, dtype)
-        sess = self._check_with_ort(onnx_path, cpu=True)
+    def common_discrepancies(self, precision, provider):
+        import torch
+
+        dtype = self._dtype_from_precision(precision)
+        np_dtype = self.get_input_np_dtype(precision)
+
+        onnx_path, model = self._common_part(precision, dtype, provider=provider)
+        sess = self._check_with_ort(onnx_path, cpu=provider == "cpu")
         config = model.config
 
         batch_size = 1
@@ -70,7 +67,7 @@ def test_trained_tiny_llm_fp32_discrepancies_cpu(self):
             batch_size=batch_size,
             seq_len=seq_len,
             np_dtype=np_dtype,
-            provider="cpu",
+            provider=provider,
             head_size=head_size,
             num_hidden_layers=config.num_hidden_layers,
             num_key_value_heads=config.num_key_value_heads,
@@ -89,31 +86,16 @@ def test_trained_tiny_llm_fp32_discrepancies_cpu(self):
                 precision=precision,
                 model_id=MODEL_NAME,
                 experiment="forward",
-                provider="cpu",
-                test="test_trained_tiny_llm_fp32_discrepancies_cpu",
+                provider=provider,
+                test=f"test_trained_tiny_llm_{precision}_discrepancies_{provider}",
                 input_type="text",
+                kind="trained",
             )
         )
         self.log_results(disc)
         self.assertLess(disc["max_abs_err"], 0.05)
 
-    def test_trained_tiny_llm_genai_generate_cpu(self):
-        """
-        Compare ``transformers.generate`` with ``onnxruntime-genai`` generate
-        on the ``arnir0/Tiny-LLM`` trained model.
-
-        The test converts the model to fp32 ONNX (CPU), then:
-
-        * Runs ``transformers.generate(do_sample=False)`` on a short prompt.
-        * Runs ``onnxruntime_genai`` greedy generation on the same prompt and
-          ONNX model.
-
-        Both backends use greedy (argmax) decoding, so the generated token
-        sequences must be bit-for-bit identical.
-
-        The test is skipped automatically when ``onnxruntime-genai`` is not
-        installed.
-        """
+    def common_genai_generate(self, precision, provider, int4=False):
         try:
             import onnxruntime_genai as og
         except ImportError:
@@ -124,9 +106,9 @@ def test_trained_tiny_llm_genai_generate_cpu(self):
         import torch
         from transformers import AutoTokenizer
 
-        precision, dtype = "fp16", torch.float16
+        dtype = self._dtype_from_precision(precision)
 
-        onnx_path, model = self._common_part(precision, dtype)
+        onnx_path, model = self._common_part(precision, dtype, provider=provider, int4=int4)
 
         genai_config_path = os.path.join(os.path.dirname(onnx_path), "genai_config.json")
         self.assertExists(genai_config_path)
@@ -174,102 +156,65 @@ def test_trained_tiny_llm_genai_generate_cpu(self):
 
         # Greedy decoding is deterministic: both backends must produce the
         # exact same token sequence (prompt + all generated tokens).
+        test_precision = "int4" if int4 else precision
         disc = self.first_token_diff(pt_tokens[start_sequence:], og_tokens)
         disc.update(
             dict(
-                precision=precision,
+                precision=test_precision,
                 model_id=MODEL_NAME,
                 experiment="generate",
-                provider="cpu",
-                test="test_trained_tiny_llm_genai_generate_cpu",
+                provider=provider,
+                test=f"test_trained_tiny_llm_genai_generate_{test_precision}_{provider}",
                 expected_text=tokenizer.decode(
                     pt_tokens[start_sequence:], skip_special_tokens=False
                 ),
                 genai_text=tokenizer.decode(og_tokens, skip_special_tokens=False),
                 input_type="text",
+                kind="trained",
             )
         )
         self.log_results(disc)
         self.assertEqual(pt_tokens[start_sequence:], og_tokens)
 
-    def test_trained_tiny_llm_genai_generate_int4_cpu(self):
-        try:
-            import onnxruntime_genai as og
-        except ImportError:
-            raise unittest.SkipTest(
-                "onnxruntime-genai is not installed; skipping genai comparison test."
-            )
-
-        import torch
-        from transformers import AutoTokenizer
-
-        precision, dtype = "fp32", torch.float32
-
-        onnx_path, model = self._common_part(precision, dtype, int4=True)
-
-        genai_config_path = os.path.join(os.path.dirname(onnx_path), "genai_config.json")
-        self.assertExists(genai_config_path)
-
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-        prompt = "Once upon a time"
-        max_new_tokens = 20
+    @long_test()
+    @hide_stdout()
+    def test_trained_tiny_llm_fp32_discrepancies_cpu(self):
+        """
+        Convert arnir0/Tiny-LLM to an fp32 ONNX model targeting the CPU execution
+        provider.  Only one hidden layer is materialised (via the
+        ``num_hidden_layers`` extra option) to keep the test fast.
+        The test verifies that:
+        * ``create_model`` completes without error.
+        * The expected ``model.onnx`` file is written to the output directory.
+        * The produced ONNX file passes ``onnx.checker.check_model``.
+        """
+        self.common_discrepancies("fp32", "cpu")
 
-        # ------------------------------------------------------------------
-        # transformers greedy generation (reference)
-        # ------------------------------------------------------------------
-        inputs = tokenizer(prompt, return_tensors="pt")
-        start_sequence = inputs["input_ids"].shape[1]
-        inputs = inputs.to("cpu")
-        with torch.no_grad():
-            pt_output = model.generate(
-                **inputs,
-                max_new_tokens=max_new_tokens,
-                do_sample=False,
-                pad_token_id=tokenizer.eos_token_id,
-            )
-        pt_tokens = pt_output[0].tolist()
+    @long_test()
+    @hide_stdout()
+    def test_trained_tiny_llm_genai_generate_cpu(self):
+        """
+        Compare ``transformers.generate`` with ``onnxruntime-genai`` generate
+        on the ``arnir0/Tiny-LLM`` trained model.
 
-        # ------------------------------------------------------------------
-        # onnxruntime-genai greedy generation
-        # ------------------------------------------------------------------
-        og_model = og.Model(os.path.dirname(onnx_path))
+        The test converts the model to fp16 ONNX (CPU), then:
 
-        params = og.GeneratorParams(og_model)
-        params.set_search_options(
-            do_sample=False,
-            max_length=inputs["input_ids"].shape[1] + max_new_tokens,
-            temperature=1.0,
-            top_k=1,
-        )
+        * Runs ``transformers.generate(do_sample=False)`` on a short prompt.
+        * Runs ``onnxruntime_genai`` greedy generation on the same prompt and
+          ONNX model.
 
-        generator = og.Generator(og_model, params)
-        generator.append_tokens(inputs["input_ids"])
+        Both backends use greedy (argmax) decoding, so the generated token
+        sequences must be bit-for-bit identical.
 
-        og_tokens = []
-        while not generator.is_done():
-            generator.generate_next_token()
-            og_token = generator.get_next_tokens()[0]
-            og_tokens.append(int(og_token))
+        The test is skipped automatically when ``onnxruntime-genai`` is not
+        installed.
+        """
+        self.common_genai_generate("fp16", "cpu")
 
-        # Greedy decoding is deterministic: both backends must produce the
-        # exact same token sequence (prompt + all generated tokens).
-        disc = self.first_token_diff(pt_tokens[start_sequence:], og_tokens)
-        disc.update(
-            dict(
-                precision="int4",
-                model_id=MODEL_NAME,
-                experiment="generate",
-                provider="cpu",
-                test="test_trained_tiny_llm_genai_generate_int4_cpu",
-                expected_text=tokenizer.decode(
-                    pt_tokens[start_sequence:], skip_special_tokens=False
-                ),
-                genai_text=tokenizer.decode(og_tokens, skip_special_tokens=False),
-                input_type="text",
-            )
-        )
-        self.log_results(disc)
-        self.assertEqual(pt_tokens[start_sequence:], og_tokens)
+    @long_test()
+    @hide_stdout()
+    def test_trained_tiny_llm_genai_generate_int4_cpu(self):
+        self.common_genai_generate("fp32", "cpu", int4=True)
 
 
 if __name__ == "__main__":
