test(decoder): add granite unit tests

huggingface · Dec 23, 2024 · 429aeb5 · 429aeb5
1 parent fc34974
commit 429aeb5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/tests/decoder/conftest.py b/tests/decoder/conftest.py
@@ -37,6 +37,10 @@
         "model_id": "Qwen/Qwen2.5-0.5B",
         "export_kwargs": {"batch_size": 4, "sequence_length": 4096, "num_cores": 2, "auto_cast_type": "fp16"},
     },
+    "granite": {
+        "model_id": "ibm-granite/granite-3.1-2b-instruct",
+        "export_kwargs": {"batch_size": 4, "sequence_length": 4096, "num_cores": 2, "auto_cast_type": "bf16"},
+    },
     "mistral": {
         "model_id": "optimum/mistral-1.1b-testing",
         "export_kwargs": {"batch_size": 4, "sequence_length": 4096, "num_cores": 2, "auto_cast_type": "bf16"},

diff --git a/tests/decoder/test_decoder_export.py b/tests/decoder/test_decoder_export.py
@@ -31,6 +31,7 @@
     "mixtral": "dacorvo/Mixtral-tiny",
     "opt": "hf-internal-testing/tiny-random-OPTForCausalLM",
     "qwen2": "yujiepan/qwen2.5-128k-tiny-random",
+    "granite": "hf-internal-testing/tiny-random-GraniteForCausalLM",
 }