Improve optimum coverage in ET (more models, xnnpack on mac)

Author: jackzhxng

.ci/scripts/test_huggingface_optimum_model.py

@@ -1,7 +1,11 @@
 import argparse
+import gc
+import logging
+import math
 import subprocess
 import tempfile
 from pathlib import Path
+from typing import List
 
 import torch
 from datasets import load_dataset
@@ -15,6 +19,7 @@
 )
 from transformers import (
     AutoConfig,
+    AutoModelForCausalLM,
     AutoModelForImageClassification,
     AutoProcessor,
     AutoTokenizer,
@@ -37,6 +42,56 @@ def cli_export(command, model_dir):
         print(f"Export failed with error: {e}")
 
 
+def check_causal_lm_output_quality(
+    model_id: str, generated_tokens: List[int], max_perplexity_threshold: float = 100.0
+):
+    """
+    Evaluates the quality of text generated by a causal language model by calculating its perplexity.
+
+    Args:
+        model_id: HuggingFace model identifier (e.g., "google/gemma2-2b")
+        generated_tokens: The tokens generated by the exported model to evaluate
+        max_perplexity_threshold: Maximum acceptable perplexity (lower is better)
+
+    Returns:
+        tuple: (is_quality_ok, reason) with boolean result and explanation
+    """
+    logging.info(f"Starting perplexity check with model '{model_id}' ...")
+    # Load model
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        low_cpu_mem_usage=True,
+        use_cache=False,
+        torch_dtype=torch.bfloat16,
+    )
+
+    with torch.no_grad():
+        outputs = model(input_ids=generated_tokens, labels=generated_tokens)
+
+    # Get the loss (negative log-likelihood)
+    loss = outputs.loss.item()
+
+    # Calculate perplexity (exp of the average negative log-likelihood)
+    perplexity = math.exp(loss)
+
+    is_quality_ok = perplexity <= max_perplexity_threshold
+    if is_quality_ok:
+        logging.info(
+            f"✓ Perplexity check passed: {perplexity:.2f} <= {max_perplexity_threshold}"
+        )
+    else:
+        logging.warning(
+            f"✗ Perplexity check failed: {perplexity:.2f} > {max_perplexity_threshold}"
+        )
+
+    # Clean up immediately
+    del model
+    del outputs
+    gc.collect()
+
+    return is_quality_ok
+
+
 def test_text_generation(model_id, model_dir, recipe, *, quantize=True, run_only=False):
     command = [
         "optimum-cli",
@@ -51,7 +106,15 @@ def test_text_generation(model_id, model_dir, recipe, *, quantize=True, run_only
         "--output_dir",
         model_dir,
     ]
-    if "coreml" in recipe:
+    if "xnnpack" in recipe:
+        if quantize:
+            command += [
+                "--qlinear",
+                "8da4w",
+                "--qembedding",
+                "8w",
+            ]
+    elif "coreml" in recipe:
         command += [
             "--disable_dynamic_shapes",
         ]
@@ -63,7 +126,9 @@ def test_text_generation(model_id, model_dir, recipe, *, quantize=True, run_only
                 "8w",
             ]
     else:
-        assert not quantize, "Quantization is not supported for non-CoreML recipes yet"
+        assert (
+            not quantize
+        ), "Quantization is only supported for XnnPack and CoreML recipes at the moment."
 
     if not run_only:
         cli_export(command, model_dir)
@@ -77,6 +142,14 @@ def test_text_generation(model_id, model_dir, recipe, *, quantize=True, run_only
         max_seq_len=64,
     )
     print(f"\nGenerated text:\n\t{generated_text}")
+    generated_tokens = tokenizer(generated_text, return_tensors="pt").input_ids
+
+    # Free memory before loading eager for quality check
+    del model
+    del tokenizer
+    gc.collect()
+
+    assert check_causal_lm_output_quality(model_id, generated_tokens) is True
 
 
 def test_fill_mask(model_id, model_dir, recipe, *, quantize=True, run_only=False):
@@ -278,23 +351,39 @@ def test_vit(model_id, model_dir, recipe, *, quantize=False, run_only=False):
     )
     args = parser.parse_args()
 
-    model_to_model_id_and_test_function = {
-        "smollm": ("HuggingFaceTB/SmolLM2-135M", test_text_generation),  # works
-        "qwen3": ("Qwen/Qwen3-0.6B", test_text_generation),  # works
-        "olmo": ("allenai/OLMo-1B-hf", test_text_generation),  # works
-        "gemma3": ("unsloth/gemma-3-1b-it", test_text_generation),  # does not export
-        "phi4": (
+    _text_generation_mapping = {
+        "llama3.2-1b": ("NousResearch/Llama-3.2-1B", test_text_generation),
+        "qwen3-0.6b": ("Qwen/Qwen3-0.6B", test_text_generation),
+        "qwen3-1.7b": ("Qwen/Qwen3-1.7B", test_text_generation),
+        "gemma3-1b": (
+            "unsloth/gemma-3-1b-it",
+            test_text_generation,
+        ),  # does not export for CoreML
+        "phi4-mini": (
             "microsoft/Phi-4-mini-instruct",
             test_text_generation,
-        ),  # fails to lower
-        "llama3": ("NousResearch/Llama-3.2-1B", test_text_generation),  # works
-        "bert": ("google-bert/bert-base-uncased", test_fill_mask),  # works
-        "roberta": ("FacebookAI/xlmcl-roberta-base", test_fill_mask),  # works
-        "distilbert": ("distilbert/distilbert-base-uncased", test_fill_mask),  # works
-        "whisper": ("openai/whisper-tiny", test_whisper),  # works
+        ),  # fails to lower for CoreML
+        "smollm2-135m": ("HuggingFaceTB/SmolLM2-135M", test_text_generation),
+        "smollm3-3b": ("HuggingFaceTB/SmolLM3-3B", test_text_generation),
+        "olmo": ("allenai/OLMo-1B-hf", test_text_generation),
+    }
+
+    _mask_fill_mapping = {
+        "bert": ("google-bert/bert-base-uncased", test_fill_mask),
+        "roberta": ("FacebookAI/xlmcl-roberta-base", test_fill_mask),
+        "distilbert": ("distilbert/distilbert-base-uncased", test_fill_mask),
+    }
+
+    _misc_model_mapping = {
+        "whisper": ("openai/whisper-tiny", test_whisper),
         "t5": ("google-t5/t5-small", test_t5),  # CoreML runime failure
-        "vit": ("google/vit-base-patch16-224", test_vit),  # works
+        "vit": ("google/vit-base-patch16-224", test_vit),
     }
+
+    model_to_model_id_and_test_function = (
+        _text_generation_mapping + _mask_fill_mapping + _misc_model_mapping
+    )
+
     if args.model not in model_to_model_id_and_test_function:
         raise ValueError(
             f"Unknown model name: {args.model}. Available models: {model_to_model_id_and_test_function.keys()}"