diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index e39725c7b..f86da3d6c 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -82,7 +82,7 @@ jobs: - name: Test with pytest env: HF_TOKEN: ${{ secrets.HF_TOKEN }} - run: pytest -x --showlocals -s -vv -n=auto --ignore=tests/models/test_openvino.py --ignore=tests/models/test_hf_steered.py --ignore=tests/scripts/test_zeno_visualize.py + run: pytest -x --showlocals -s -vv -n=auto --ignore=tests/models/test_openvino.py --ignore=tests/models/test_hf_steered.py --ignore=tests/models/test_onnxruntime_genai.py --ignore=tests/models/test_winml.py --ignore=tests/scripts/test_zeno_visualize.py # Save test artifacts - name: Archive test artifacts diff --git a/README.md b/README.md index d2711b3ce..50e32f265 100644 --- a/README.md +++ b/README.md @@ -485,6 +485,34 @@ lm_eval --model sglang \ > 2. Lower KV cache pool memory usage by adjusting `mem_fraction_static` - Add to your model arguments for example `--model_args pretrained=...,mem_fraction_static=0.7`. > 3. Increase tensor parallel size `tp_size` (if using multiple GPUs). +### ONNX Runtime GenAI + +We support **ONNX Runtime GenAI** for cross-platform evaluation of ONNX LLMs produced by the [ONNX Runtime GenAI *Model Builder*](https://onnxruntime.ai/docs/genai/howto/build-model.html). Unlike the [Windows ML](#windows-ml) backend, this backend runs on Linux, macOS, and Windows and selects execution providers with the cross-platform `og.Config` API (CPU, CUDA, DirectML, WebGPU, and AMD NPU via VitisAI/RyzenAI). + +Install the backend along with the execution-provider wheel that matches your hardware: + +```bash +# CPU +pip install "lm_eval[onnxruntime-genai]" +# CUDA / DirectML instead of the CPU wheel (mutually exclusive): +# pip install onnxruntime-genai-cuda +# pip install onnxruntime-genai-directml +``` + +Evaluate a Model Builder ONNX model on a chosen execution provider: + +```bash +lm_eval --model onnxruntime-genai \ + --model_args pretrained=/path/to/model_builder_output,execution_provider=cuda \ + --tasks hellaswag \ + --batch_size 1 +``` + +The `pretrained` path is a Model Builder output directory (containing `genai_config.json`, the ONNX graph(s), and an HF tokenizer) or a `.onnx` file inside one. `execution_provider` defaults to `cpu` (which runs with no provider appended); pass `cuda`, `dml`, `VitisAI`, etc. for accelerators. Extra provider settings can be passed via `provider_options`. + +> [!Note] +> Supported architectures are whatever the Model Builder supports (Llama, Phi, Qwen, Gemma, Mistral, Granite, ChatGLM, …). Inference is batch-size 1 and runs one execution provider per run. + ### Windows ML We support **Windows ML** for hardware-accelerated inference on Windows platforms. This enables evaluation on CPU, GPU, and **NPU (Neural Processing Unit)** devices. @@ -556,6 +584,7 @@ Note that for externally hosted models, configs such as `--device` which relate | NVIDIA NeMo | :heavy_check_mark: | `nemo_lm` | [All supported models](https://docs.nvidia.com/nemo-framework/user-guide/24.09/nemotoolkit/core/core.html#nemo-models) | `generate_until`, `loglikelihood`, `loglikelihood_rolling` | | NVIDIA Megatron-LM | :heavy_check_mark: | `megatron_lm` | [Megatron-LM GPT models](https://github.com/NVIDIA/Megatron-LM) (standard and distributed checkpoints) | `generate_until`, `loglikelihood`, `loglikelihood_rolling` | | Watsonx.ai | :heavy_check_mark: | `watsonx_llm` | [Supported Watsonx.ai Engines](https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models.html?context=wx) | `generate_until` `loglikelihood` | +| ONNX Runtime GenAI | :heavy_check_mark: | `onnxruntime-genai` | [ONNX models in GenAI format](https://onnxruntime.ai/docs/genai/howto/build-model.html) (cross-platform: CPU/CUDA/DirectML/NPU) | `generate_until`, `loglikelihood`, `loglikelihood_rolling` | | Windows ML | :heavy_check_mark: | `winml` | [ONNX models in GenAI format](https://code.visualstudio.com/docs/intelligentapps/modelconversion) | `generate_until`, `loglikelihood`, `loglikelihood_rolling` | | [Your local inference server!](docs/API_guide.md) | :heavy_check_mark: | `local-completions` or `local-chat-completions` | Support for OpenAI API-compatible servers, with easy customization for other APIs. | `generate_until`, `loglikelihood`, `loglikelihood_rolling` | @@ -801,6 +830,7 @@ These extras install dependencies required to run specific model backends: | habana | Intel Gaudi backend | | optimum | Intel OpenVINO models | | neuronx | AWS Inferentia2 instances | +| onnxruntime-genai | ONNX Runtime GenAI (cross-platform) - CPU/CUDA/DirectML/NPU | | winml | Windows ML (ONNX Runtime GenAI) - CPU/GPU/NPU | | sparsify | Sparsify model steering | | sae_lens | SAELens model steering | diff --git a/docs/model_guide.md b/docs/model_guide.md index 3c122ffec..8f308db9a 100644 --- a/docs/model_guide.md +++ b/docs/model_guide.md @@ -187,6 +187,8 @@ If not implemented for a given model type, the flags `--apply_chat_template` , ` **Pro tip**: In order to make the Evaluation Harness overestimate total runtimes rather than underestimate it, HuggingFace models come in-built with the ability to provide responses on data points in *descending order by total input length* via `lm_eval.utils.Reorderer`. Take a look at `lm_eval.models.hf_causal.HFLM` to see how this is done, and see if you can implement it in your own model! +**Subclassing an existing backend**: If your backend shares most of its logic with an existing one and differs only in a narrow concern (e.g. how the underlying model is created or which device it targets), prefer subclassing over duplication. For example, `lm_eval.models.optimum_lm.OptimumLM` subclasses `HFLM` and overrides only `_create_model`. + ## Conclusion After reading this guide, you should be able to add new model APIs or implementations to the Eval Harness library! diff --git a/lm_eval/models/__init__.py b/lm_eval/models/__init__.py index e88d61292..616cda39b 100644 --- a/lm_eval/models/__init__.py +++ b/lm_eval/models/__init__.py @@ -45,6 +45,7 @@ MODEL_MAPPING = { "megatron_lm": "lm_eval.models.megatron_lm:MegatronLMEval", "nemo_lm": "lm_eval.models.nemo_lm:NeMoLM", "neuronx": "lm_eval.models.neuron_optimum:NeuronModelForCausalLM", + "onnxruntime-genai": "lm_eval.models.onnxruntime_genai:ONNXRuntimeGenAILM", "openai-chat-completions": "lm_eval.models.openai_completions:OpenAIChatCompletion", "openai-completions": "lm_eval.models.openai_completions:OpenAICompletionsAPI", "openvino": "lm_eval.models.optimum_lm:OptimumLM", @@ -75,4 +76,4 @@ def _register_all_models(): # Call registration on module import _register_all_models() -__all__ = ["MODEL_MAPPING"] \ No newline at end of file +__all__ = ["MODEL_MAPPING"] diff --git a/lm_eval/models/onnxruntime_genai.py b/lm_eval/models/onnxruntime_genai.py new file mode 100644 index 000000000..0a23d4cad --- /dev/null +++ b/lm_eval/models/onnxruntime_genai.py @@ -0,0 +1,405 @@ +"""Cross-platform ONNX Runtime GenAI backend for lm-eval-harness. + +Runs models exported by the ONNX Runtime GenAI *Model Builder* through the +``onnxruntime-genai`` runtime (``og.Model`` / ``og.Generator``). Unlike the +``winml`` backend -- which selects execution providers via the Windows ML +catalog and therefore only runs on Windows -- this backend selects providers +with the cross-platform ``og.Config`` API and runs on CPU, CUDA, DirectML, +WebGPU, and AMD NPU (VitisAI/RyzenAI). + +Example usage: + lm_eval --model onnxruntime-genai --model_args pretrained=path/to/model_builder_output,execution_provider=cuda --tasks hellaswag --limit 10 + +The ``pretrained`` path is a Model Builder output directory (containing +``genai_config.json``, the ONNX graph(s), and an HF tokenizer) or a ``.onnx`` +file inside such a directory. +""" + +import json +import logging +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np +from tqdm import tqdm + +from lm_eval import utils +from lm_eval.api.model import TemplateLM +from lm_eval.api.registry import register_model + + +if TYPE_CHECKING: + from lm_eval.api.instance import Instance + + +eval_logger = logging.getLogger(__name__) + +# ONNX Runtime GenAI treats CPU as "no execution provider appended" rather than +# an appendable provider name, so these aliases are handled specially. +_CPU_EP_ALIASES = {"cpu", "cpuexecutionprovider"} + + +@register_model("onnxruntime-genai") +class ONNXRuntimeGenAILM(TemplateLM): + """Cross-platform backend for Model Builder ONNX models via onnxruntime-genai. + + All lm-eval scoring logic lives here (tokenization, single-pass logits, + log-likelihood, rolling perplexity, and generation). Subclasses need only + override :meth:`_select_ep` to change how execution providers are chosen + (see :class:`lm_eval.models.winml.WindowsML`). + """ + + _DEFAULT_MAX_LENGTH = 2048 + + def __init__( + self, + pretrained: str, + execution_provider: str | None = None, + max_length: int | None = None, + batch_size: int = 1, + max_gen_toks: int = 256, + provider_options: dict[str, Any] | None = None, + **kwargs, + ) -> None: + """Initialize the backend. + + Args: + pretrained: Path to a Model Builder output directory or a ``.onnx`` + file inside one. + execution_provider: EP to run on (e.g. ``cpu``, ``cuda``, ``dml``, + ``VitisAI``). When ``None`` (default), the providers declared in + ``genai_config.json`` are used as-is — matching how a Model + Builder export is meant to run on its target device. Pass + ``cpu`` to force CPU, or a provider name to override the export. + max_length: Maximum sequence length. Defaults to the + ``context_length`` from ``genai_config.json`` when available. + batch_size: Only 1 is supported; other values are coerced to 1. + max_gen_toks: Default maximum number of tokens to generate. + provider_options: Extra key/value options forwarded to the EP. + """ + super().__init__() + + self.og = self._import_genai() + self.pretrained = pretrained + self.model_dir = self._resolve_model_dir(pretrained) + self.execution_provider = execution_provider + self.provider_options = provider_options or {} + self._max_gen_toks = max_gen_toks + + if batch_size != 1: + eval_logger.warning( + f"{type(self).__name__} currently supports batch size 1 only; " + f"overriding requested batch_size={batch_size}." + ) + self.batch_size = 1 + + self._genai_config = self._read_genai_config(self.model_dir) + model_cfg = self._genai_config.get("model", {}) + context_length = model_cfg.get("context_length") + self.max_length = max_length or context_length or self._DEFAULT_MAX_LENGTH + + self._eot_token_id = self._resolve_eot_token_id(model_cfg) + self._bos_token_id = model_cfg.get("bos_token_id") + + self._load() + + @staticmethod + def _import_genai(): + try: + import onnxruntime_genai as og + except ImportError as e: + raise ImportError( + "onnxruntime-genai is required for the onnxruntime-genai backend. " + "Install a matching EP wheel, e.g. `pip install onnxruntime-genai` " + "(CPU) or `pip install onnxruntime-genai-cuda` (CUDA)." + ) from e + eval_logger.info(f"ONNX Runtime GenAI version: {og.__version__}") + return og + + @staticmethod + def _resolve_model_dir(pretrained: str) -> Path: + path = Path(pretrained) + if path.is_file() and path.suffix == ".onnx": + return path.parent + if path.is_dir(): + return path + raise FileNotFoundError( + f"Model path {pretrained} not found or is not a directory / .onnx file" + ) + + @staticmethod + def _read_genai_config(model_dir: Path) -> dict[str, Any]: + config_path = model_dir / "genai_config.json" + if not config_path.exists(): + eval_logger.warning( + f"No genai_config.json found in {model_dir}; falling back to defaults." + ) + return {} + with open(config_path, encoding="utf-8") as f: + return json.load(f) + + def _resolve_eot_token_id(self, model_cfg: dict[str, Any]) -> int: + eos = model_cfg.get("eos_token_id") + if isinstance(eos, list): + eos = eos[0] if eos else None + if eos is not None: + return int(eos) + # Fall back to the tokenizer once it is loaded (see _load). + return None # type: ignore[return-value] + + # ------------------------------------------------------------------ # + # Engine hooks (override in subclasses to change engine / EP selection) + # ------------------------------------------------------------------ # + def _select_ep(self, config) -> None: + """Configure execution providers on an ``og.Config``. + + ``og.Config(model_dir)`` seeds its provider list from + ``genai_config.json``, so we only touch it when the user explicitly + asks for a provider: + + * ``execution_provider is None`` — leave the config untouched, honoring + the providers the export was built with (a CUDA/DML/NPU export runs + on its target device without extra flags). + * ``execution_provider`` is a CPU alias — clear providers so it runs on + CPU (an empty provider list means CPU). + * otherwise — clear and append the named provider with any + ``provider_options``. + """ + ep = self.execution_provider + if ep is None: + eval_logger.info( + "No execution_provider given; using the providers declared in " + "genai_config.json." + ) + return + config.clear_providers() + if ep.lower() in _CPU_EP_ALIASES: + eval_logger.info("Using CPU execution provider (no provider appended).") + return + config.append_provider(ep) + for key, value in self.provider_options.items(): + config.set_provider_option(ep, str(key), str(value)) + eval_logger.info(f"Using execution provider: {ep}") + + def _load(self) -> None: + """Load the tokenizer and the onnxruntime-genai model.""" + from transformers import AutoTokenizer + + self.tokenizer = AutoTokenizer.from_pretrained(str(self.model_dir)) + if self._eot_token_id is None: + self._eot_token_id = ( + self.tokenizer.eos_token_id + if self.tokenizer.eos_token_id is not None + else 0 + ) + if self._bos_token_id is None: + self._bos_token_id = self.tokenizer.bos_token_id + + config = self.og.Config(str(self.model_dir)) + self._select_ep(config) + self.model = self.og.Model(config) + eval_logger.info(f"Loaded onnxruntime-genai model from {self.model_dir}") + + def _forward_logits(self, tokens: list[int]) -> np.ndarray: + """Run one forward pass and return per-position logits. + + Args: + tokens: The full input token sequence. + + Returns: + A ``[seq_len, vocab]`` array where row ``i`` is the distribution over + the token at position ``i + 1`` given ``tokens[: i + 1]``. + """ + params = self.og.GeneratorParams(self.model) + generator = self.og.Generator(self.model, params) + generator.append_tokens(np.asarray(tokens, dtype=np.int32)) + logits = np.asarray(generator.get_output("logits"), dtype=np.float32) + if logits.ndim == 3: # (batch, seq, vocab) + logits = logits[0] + elif logits.ndim != 2: + raise ValueError(f"Unexpected logits shape from GenAI: {logits.shape}") + return logits + + def _generate(self, tokens: list[int], gen_kwargs: dict[str, Any]) -> list[int]: + """Greedily (or with sampling) generate token ids from a prompt. + + Args: + tokens: The prompt token ids. + gen_kwargs: Generation options (``max_gen_toks``, ``until``, + ``temperature``, ``top_p``, ``top_k``, ``do_sample``). + + Returns: + The newly generated token ids (excluding the prompt). + """ + max_gen_toks = int(gen_kwargs.get("max_gen_toks", self.max_gen_toks)) + temperature = float(gen_kwargs.get("temperature", 0.0)) + top_p = float(gen_kwargs.get("top_p", 1.0)) + top_k = int(gen_kwargs.get("top_k", 50)) + do_sample = bool(gen_kwargs.get("do_sample", False)) and temperature > 0 + + params = self.og.GeneratorParams(self.model) + # max_length is total (prompt + generated) length for the genai search. + total_max = min(len(tokens) + max_gen_toks, self.max_length) + if do_sample: + params.set_search_options( + max_length=total_max, + do_sample=True, + temperature=temperature, + top_p=top_p, + top_k=top_k, + ) + else: + params.set_search_options(max_length=total_max, do_sample=False) + + generator = self.og.Generator(self.model, params) + generator.append_tokens(np.asarray(tokens, dtype=np.int32)) + + generated: list[int] = [] + while not generator.is_done() and len(generated) < max_gen_toks: + generator.generate_next_token() + seq = generator.get_sequence(0) + if len(seq) > len(tokens) + len(generated): + generated.append(int(seq[len(tokens) + len(generated)])) + return generated + + # ------------------------------------------------------------------ # + # TemplateLM interface + # ------------------------------------------------------------------ # + @property + def eot_token_id(self) -> int: + return self._eot_token_id + + @property + def prefix_token_id(self) -> int: + if self._bos_token_id is not None: + return int(self._bos_token_id) + return self._eot_token_id + + @property + def max_gen_toks(self) -> int: + return self._max_gen_toks + + def tok_encode( + self, string: str, add_special_tokens: bool | None = None, **kwargs + ) -> list[int]: + if add_special_tokens is None: + add_special_tokens = False + return self.tokenizer.encode(string, add_special_tokens=add_special_tokens) + + def tok_decode(self, tokens: list[int], **kwargs) -> str: + return self.tokenizer.decode(tokens, skip_special_tokens=True) + + def _loglikelihood_tokens( + self, + requests: list[tuple[tuple[str, str] | None, list[int], list[int]]], + disable_tqdm: bool = False, + override_bs: int | None = None, + ) -> list[tuple[float, bool]]: + results: list[tuple[float, bool]] = [] + + for cache_key, context_enc, continuation_enc in tqdm( + requests, disable=disable_tqdm, desc="Computing log-likelihoods" + ): + if not continuation_enc: + results.append((0.0, True)) + continue + + # Feed the full sequence; keep the last max_length tokens so the + # continuation is always scored, mirroring HFLM's left-truncation. + # At least one preceding position is required to score the first + # continuation token, so cap contlen to leave room for context. + full = (context_enc + continuation_enc)[-self.max_length :] + contlen = min(len(continuation_enc), len(full) - 1) + ctxlen = len(full) - contlen + + logits = self._forward_logits(full) + # Logits at position i predict token i+1, so the continuation is + # scored by rows [ctxlen-1 : ctxlen-1+contlen]. + cont_logits = logits[ctxlen - 1 : ctxlen - 1 + contlen] + log_probs = _log_softmax(cont_logits) + + # Left-truncation may have dropped leading continuation tokens. + targets = np.asarray(continuation_enc[-contlen:], dtype=np.int64) + token_ll = log_probs[np.arange(contlen), targets] + total_ll = float(token_ll.sum()) + is_greedy = bool(np.all(log_probs.argmax(axis=-1) == targets)) + + results.append((total_ll, is_greedy)) + + if cache_key is not None: + self.cache_hook.add_partial( + "loglikelihood", cache_key, (total_ll, is_greedy) + ) + + return results + + def loglikelihood_rolling( + self, requests: list["Instance"], disable_tqdm: bool = False + ) -> list[float]: + loglikelihoods: list[float] = [] + + for (string,) in tqdm( + [req.args for req in requests], + disable=disable_tqdm, + desc="Computing rolling log-likelihoods", + ): + rolling_token_windows = [ + (None,) + window + for window in map( + utils.make_disjoint_window, + utils.get_rolling_token_windows( + token_list=self.tok_encode(string), + prefix_token=self.prefix_token_id, + max_seq_len=self.max_length, + context_len=1, + ), + ) + ] + + window_lls = self._loglikelihood_tokens( + rolling_token_windows, disable_tqdm=True + ) + string_ll = sum(ll for ll, _ in window_lls) + loglikelihoods.append(string_ll) + self.cache_hook.add_partial("loglikelihood_rolling", (string,), string_ll) + + return loglikelihoods + + def generate_until( + self, requests: list["Instance"], disable_tqdm: bool = False + ) -> list[str]: + results: list[str] = [] + + for context, gen_kwargs in tqdm( + [req.args for req in requests], + disable=disable_tqdm, + desc="Generating text", + ): + gen_kwargs = dict(gen_kwargs) + until = gen_kwargs.pop("until", None) + if isinstance(until, str): + until = [until] + elif until is None: + until = [] + + prompt_tokens = self.tok_encode(context) + new_tokens = self._generate(prompt_tokens, gen_kwargs) + text = self.tok_decode(new_tokens) + + for stop in until: + if stop: + idx = text.find(stop) + if idx != -1: + text = text[:idx] + + results.append(text) + self.cache_hook.add_partial("generate_until", (context, gen_kwargs), text) + + return results + + +def _log_softmax(logits: np.ndarray) -> np.ndarray: + """Numerically stable log-softmax over the last axis.""" + shifted = logits - logits.max(axis=-1, keepdims=True) + log_sum_exp = np.log(np.exp(shifted).sum(axis=-1, keepdims=True)) + return shifted - log_sum_exp diff --git a/lm_eval/models/winml.py b/lm_eval/models/winml.py index ca819fd66..7a5127897 100644 --- a/lm_eval/models/winml.py +++ b/lm_eval/models/winml.py @@ -1,172 +1,101 @@ -""" -WinML backend for lm-eval-harness with NPU/GPU/CPU support. +"""WinML backend for lm-eval-harness with NPU/GPU/CPU support. -This backend leverages Windows Machine Learning (WinML) to run models on various -hardware backends including NPUs, GPUs, and CPUs. It's particularly useful for -running inference on Windows devices with dedicated Neural Processing Units. +This backend runs Model Builder ONNX models through the Windows Machine Learning +(WinML) execution-provider catalog. It is a thin specialization of +:class:`~lm_eval.models.onnxruntime_genai.ONNXRuntimeGenAILM`: all scoring and +generation logic is inherited, and only execution-provider selection is +Windows-specific. Example usage: - lm_eval --model winml --model_args pretrained=path/to/onnx/model.onnx --tasks hellaswag + lm_eval --model winml --model_args pretrained=path/to/onnx/model --tasks hellaswag """ import logging from pathlib import Path -from typing import TYPE_CHECKING, Any -import numpy as np -from tqdm import tqdm - -from lm_eval import utils -from lm_eval.api.model import TemplateLM from lm_eval.api.registry import register_model +from lm_eval.models.onnxruntime_genai import ONNXRuntimeGenAILM -if TYPE_CHECKING: - from lm_eval.api.instance import Instance - eval_logger = logging.getLogger(__name__) @register_model("winml") -class WindowsML(TemplateLM): - """ - WindowsML backend for lm-eval-harness with NPU/GPU/CPU support. +class WindowsML(ONNXRuntimeGenAILM): + """WindowsML backend: onnxruntime-genai with Windows ML provider selection. - This model class provides integration with Windows Machine Learning (WindowsML) - to enable evaluation on NPUs and other Windows-optimized hardware. + Inherits all lm-eval logic (tokenization, single-pass log-likelihood, + rolling perplexity, and generation) from + :class:`~lm_eval.models.onnxruntime_genai.ONNXRuntimeGenAILM`, overriding + only :meth:`_select_ep` to register execution providers through the Windows + ML catalog. + + The default ``max_length`` / ``max_gen_toks`` are pinned to the historical + winml values so numeric behavior is unchanged; the cross-platform backend + picks its own (more model-aware) defaults. """ - _DEFAULT_MAX_LENGTH = 2048 - - @classmethod - def create_from_arg_obj( - cls, arg_dict: dict[str, Any], additional_config: dict[str, Any] | None = None - ) -> "WindowsML": - """ - Override to properly merge dictionaries and avoid duplicate keyword arguments. - - Args: - arg_dict: Dictionary containing model arguments - additional_config: Optional dictionary containing additional configuration - - Returns: - Instance of WindowsML class - """ - # Merge the dictionaries, with additional_config taking precedence - merged_config = {**(arg_dict or {})} - if additional_config: - # Filter out None values and merge - filtered_additional = { - k: v for k, v in additional_config.items() if v is not None - } - merged_config.update(filtered_additional) - - return cls(**merged_config) + # Historical winml defaults, kept so existing winml runs are unaffected. + _WINML_MAX_LENGTH = 4096 + _WINML_MAX_GEN_TOKS = 4096 def __init__( self, pretrained: str, - max_length: int | None = 4096, - batch_size: int = 1, - max_batch_size: int = 64, + max_length: int | None = None, + max_gen_toks: int | None = None, **kwargs, ) -> None: + super().__init__( + pretrained, + max_length=max_length or self._WINML_MAX_LENGTH, + max_gen_toks=max_gen_toks or self._WINML_MAX_GEN_TOKS, + **kwargs, + ) + + def _select_ep(self, config) -> None: + """Register Windows ML execution providers on the ``og.Config``. + + Discovers providers via the Windows ML catalog and registers them with + onnxruntime-genai. If discovery fails (e.g. running outside Windows), we + fall back to the cross-platform selection in the base class. """ - Initialize WindowsML model. - - Args: - pretrained: Path to ONNX model file or directory containing model files - max_length: Maximum sequence length - batch_size: Batch size for inference - max_batch_size: Maximum batch size for auto-batching - """ - super().__init__() - - # Validate and import dependencies - self._validate_dependencies() - - # Store configuration - self.pretrained = pretrained - self.max_length = max_length or self._DEFAULT_MAX_LENGTH - self.batch_size = batch_size - self.max_batch_size = max_batch_size - - # Warn about batch size limitations - if batch_size != 1 or max_batch_size != 1: - eval_logger.warning( - f"WindowsML backend currently only supports batch size 1. " - f"Requested batch_size={batch_size}, max_batch_size={max_batch_size}. " - f"Setting both to 1." - ) - self.batch_size = 1 - self.max_batch_size = 1 - self._fix_winrt_runtime() + if self._register_winml_providers_to_genai(): + self._log_winml_devices() + else: + eval_logger.warning( + "Windows ML provider registration failed; falling back to " + "cross-platform execution-provider selection." + ) + super()._select_ep(config) - # Initialize Windows ML execution providers - self._register_winml_providers_to_genai() + def _fix_winrt_runtime(self) -> None: + """Remove the bundled msvcp140.dll from winrt-runtime to avoid clashes. - # Setup device and execution providers using new Windows ML API - self._setup_winml_devices_and_providers() - - # Load and compile ONNX model - self._load_and_compile_model(pretrained) - - eval_logger.info( - f"Available EP devices: {len(self.ep_device_map)} execution providers" - ) - - def _validate_dependencies(self) -> None: - """ - Validate that required dependencies are available. - - Raises: - ImportError: If required dependencies are not installed + This DLL, shipped by the ``winrt-runtime`` package, can conflict with + other libraries loaded in the same process. """ try: - import onnxruntime_genai as og - - self.og = og - eval_logger.info(f"ONNX Runtime GenAI version: {og.__version__}") - except ImportError as e: - raise ImportError( - "ONNX Runtime GenAI is required for WinML backend. " - "Install with: pip install onnxruntime-genai" - ) from e - - # Also import regular ONNX Runtime for EP registration + from importlib import metadata + except ImportError: + return try: - import onnxruntime as ort - - self.ort = ort - eval_logger.info(f"ONNX Runtime version: {ort.__version__}") - except ImportError as e: - raise ImportError( - "ONNX Runtime is also required for execution provider registration. " - "Install with: pip install onnxruntime" - ) from e - - def _fix_winrt_runtime(self): - """ - This function removes the msvcp140.dll from the winrt-runtime package. - So it does not cause issues with other libraries. - """ - from importlib import metadata - - site_packages_path = Path( - str(metadata.distribution("winrt-runtime").locate_file("")) - ) + site_packages_path = Path( + str(metadata.distribution("winrt-runtime").locate_file("")) + ) + except metadata.PackageNotFoundError: + return dll_path = site_packages_path / "winrt" / "msvcp140.dll" if dll_path.exists(): dll_path.unlink() def _register_winml_providers_to_genai(self) -> bool: - """ - Register Windows ML execution providers to ONNX Runtime GenAI. + """Discover Windows ML providers and register them with onnxruntime-genai. Returns: - True if registration was successful, False otherwise + True if registration succeeded, False otherwise. """ try: import winui3.microsoft.windows.ai.machinelearning as winml @@ -180,610 +109,44 @@ class WindowsML(TemplateLM): providers = catalog.find_all_providers() for provider in providers: provider.ensure_ready_async().get() - # Register to GenAI instead of regular ONNX Runtime self.og.register_execution_provider_library( provider.name, provider.library_path ) eval_logger.info( - f"Registered {provider.name} to ONNX Runtime GenAI" + "Registered %s to ONNX Runtime GenAI", provider.name ) - return True except ImportError as e: - eval_logger.warning(f"Windows ML import error: {e}") + eval_logger.warning("Windows ML import error: %s", e) return False - except Exception as e: - eval_logger.warning(f"Error registering providers to GenAI: {e}") + except Exception as e: # noqa: BLE001 + eval_logger.warning("Error registering providers to GenAI: %s", e) return False - def _setup_winml_devices_and_providers(self) -> None: - """Setup execution providers using Windows ML device enumeration API. - - This method queries available devices and builds a mapping of execution providers. - """ + def _log_winml_devices(self) -> None: + """Log available execution-provider devices via the Windows ML API.""" try: - # Get available EP devices using Windows ML API - ep_devices = self.ort.get_ep_devices() - self.ep_device_map = {} + import onnxruntime as ort - # Build device map + ep_devices = ort.get_ep_devices() + ep_device_map: dict[str, list] = {} for device in ep_devices: - ep_name = device.ep_name - if ep_name not in self.ep_device_map: - self.ep_device_map[ep_name] = [] - self.ep_device_map[ep_name].append(device) + ep_device_map.setdefault(device.ep_name, []).append(device) - # Log available devices - eval_logger.info("Available execution provider devices:") - for name, devices in self.ep_device_map.items(): - eval_logger.info(f"Execution Provider: {name}") + eval_logger.info( + "Available EP devices: %s execution providers", len(ep_device_map) + ) + for name, devices in ep_device_map.items(): + eval_logger.info("Execution Provider: %s", name) for device in devices: try: - device_type = self.ort.OrtHardwareDeviceType( - device.device.type - ).name - eval_logger.info( - f" | Vendor: {device.ep_vendor:<16} | Device Type: {device_type:<8}" - ) - except Exception: - eval_logger.info( - f" | Vendor: {device.ep_vendor:<16} | Device Type: Unknown" - ) - - except Exception as e: - eval_logger.warning(f"Windows ML device enumeration failed: {e}") - eval_logger.info("Falling back to legacy provider selection") - self.ep_device_map = {} - - def _load_and_compile_model(self, model_path: str) -> None: - """ - Load and optionally compile ONNX model with ONNX Runtime GenAI. - - Args: - model_path: Path to ONNX model file or directory - - Raises: - FileNotFoundError: If model path is not found or invalid - Exception: If model loading fails - """ - model_path = Path(model_path) - - # Handle different input formats - if model_path.is_file() and model_path.suffix == ".onnx": - input_model_path = model_path.parent # GenAI expects directory - elif model_path.is_dir(): - input_model_path = model_path - else: - raise FileNotFoundError(f"Model path {model_path} not found or invalid") - - # Load model using ONNX Runtime GenAI - try: - eval_logger.info( - f"Loading model with ONNX Runtime GenAI from: {input_model_path}" - ) - - # Load model and tokenizer using GenAI - self.genai_model = self.og.Model(str(input_model_path)) - self.genai_tokenizer = self.og.Tokenizer(self.genai_model) - - eval_logger.info( - "Model and tokenizer loaded successfully with ONNX Runtime GenAI" - ) - - # Store model info - self.model_path = input_model_path - - except Exception as e: - eval_logger.error( - f"Failed to load model with ONNX Runtime GenAI from {input_model_path}: {e}" - ) - raise - - @property - def eot_token_id(self) -> int: - """ - Get the end-of-text token ID. - - Returns: - End-of-text token ID from the GenAI tokenizer - """ - try: - # For GenAI tokenizers, try to access EOS from model config or encoding - if hasattr(self.genai_model, "config"): - config = self.genai_model.config - if hasattr(config, "eos_token_id") and config.eos_token_id: - if isinstance(config.eos_token_id, list): - return config.eos_token_id[0] - else: - return config.eos_token_id - - except Exception as e: - eval_logger.warning( - f"Error getting EOS token ID: {e}, using fallback value 2" - ) - - # Fallback: return a common EOS token ID - eval_logger.warning("Could not determine EOS token ID, using fallback value 2") - return 2 - - @property - def prefix_token_id(self) -> int | None: - """ - Get the prefix token ID (typically BOS token). - - Returns: - BOS token ID if available, otherwise None - """ - try: - # For GenAI tokenizers, try to access BOS from model config or encoding - - # If we have access to the model config - if hasattr(self.genai_model, "config"): - config = self.genai_model.config - if hasattr(config, "bos_token_id") and config.bos_token_id is not None: - return config.bos_token_id - - # Check if encoding empty string gives BOS token - try: - empty_enc = self.genai_tokenizer.encode("") - if len(empty_enc) > 0: - return empty_enc[0] - except Exception: - pass - - return None - - except Exception: - return None - - @property - def max_gen_toks(self) -> int: - """ - Get the maximum number of tokens to generate. - - Returns: - Maximum generation tokens (default: 4096) - """ - return 4096 - - def tok_encode( - self, - string: str, - left_truncate_len: int | None = None, - add_special_tokens: bool = True, - ) -> list[int]: - """ - Tokenize string and return token IDs. - - Args: - string: Input string to tokenize - left_truncate_len: If provided, truncate from the left to this length - add_special_tokens: Whether to add special tokens (note: GenAI tokenizer handles this automatically) - - Returns: - List of token IDs - """ - # Use GenAI tokenizer for consistency with model inference - encoding = self.genai_tokenizer.encode(string) - - # Handle left truncation if requested - if left_truncate_len is not None and len(encoding) > left_truncate_len: - encoding = encoding[-left_truncate_len:] - - return encoding - - def tok_decode(self, tokens: list[int]) -> str: - """ - Decode token IDs back to text. - - Args: - tokens: List of token IDs to decode - - Returns: - Decoded text string - """ - return self.genai_tokenizer.decode(tokens) - - def _run_genai_inference_for_full_logits(self, input_text: str) -> np.ndarray: - """ - Run inference using ONNX Runtime GenAI to get full logits sequence. - - Args: - input_text: Input text string to compute logits for - - Returns: - Logits matrix of shape (seq_len, vocab_size) where logits[i] contains - predictions for the token at position i+1 given tokens[0:i+1] - - Raises: - Exception: If inference fails - """ - try: - # Encode input text to tokens - input_tokens = self.genai_tokenizer.encode(input_text) - - if len(input_tokens) == 0: - eval_logger.warning("No tokens to process; returning empty array") - return np.empty((0, 0), dtype=np.float32) - - # Create generator and get full logits in single pass - params = self.og.GeneratorParams(self.genai_model) - params.set_search_options(max_length=4096, do_sample=False) - generator = self.og.Generator(self.genai_model, params) - - # Append all tokens at once - generator.append_tokens(input_tokens) - - # Get FULL logits using get_output("logits") - this gives all positions! - full_logits_tensor = generator.get_output("logits") - logits_array = np.array(full_logits_tensor, dtype=np.float32) - - # Handle different tensor shapes - if len(logits_array.shape) == 3: # (batch_size, seq_len, vocab_size) - logits_matrix = logits_array[0] # Remove batch dimension - elif len(logits_array.shape) == 2: # (seq_len, vocab_size) - logits_matrix = logits_array - else: - raise ValueError(f"Unexpected logits shape: {logits_array.shape}") - - eval_logger.debug( - f"Full logits shape: {logits_matrix.shape} for {len(input_tokens)} input tokens" - ) - return logits_matrix - - except Exception as e: - eval_logger.error(f"GenAI inference failed: {e}") - raise - - def _loglikelihood_tokens( - self, - requests: list[tuple[tuple[str, str], list[int], list[int]]], - disable_tqdm: bool = False, - override_bs: int | None = None, - ) -> list[tuple[float, bool]]: - """ - Stub implementation - not used since we override loglikelihood directly. - WindowsML uses the GenAI tokenizer and overrides loglikelihood to work - with text inputs directly, avoiding tokenization round-trip issues. - - Args: - requests: List of tokenized requests - disable_tqdm: Whether to disable progress bar - override_bs: Optional batch size override - - Returns: - Empty list (method not used) - """ - raise NotImplementedError( - "WindowsML overrides loglikelihood() directly and does not use _loglikelihood_tokens()" - ) - - def loglikelihood( - self, requests: list["Instance"], disable_tqdm: bool = False - ) -> list[tuple[float, bool]]: - """ - Compute log-likelihood using ONNX Runtime GenAI with teacher forcing. - - Args: - requests: List of instances containing (context, continuation) text pairs - disable_tqdm: Whether to disable progress bar - - Returns: - List of tuples containing (log_likelihood, is_greedy) for each request - """ - results = [] - - for request in tqdm( - requests, disable=disable_tqdm, desc="Computing log-likelihoods" - ): - context, continuation = request.args - - if len(continuation) == 0: - results.append((0.0, True)) - continue - - try: - full_text = context + continuation - - # Tokenize to get continuation tokens - if context: - context_enc = self.tok_encode(context) - if ( - len(context_enc) > 0 - and self.prefix_token_id is not None - and context_enc[0] == self.prefix_token_id - ): - context_enc = context_enc[1:] - context_len = len(context_enc) - else: - context_enc = ( - [self.prefix_token_id] - if self.prefix_token_id is not None - else [] + device_type = ort.OrtHardwareDeviceType(device.device.type).name + except Exception: # noqa: BLE001 + device_type = "Unknown" + eval_logger.info( + " | Vendor: %-16s | Device Type: %-8s", + device.ep_vendor, + device_type, ) - context_len = len(context_enc) - - full_tokens = self.tok_encode(full_text) - continuation_tokens = full_tokens[context_len:] - - if len(continuation_tokens) == 0: - results.append((0.0, True)) - continue - - # Create generator for teacher forcing - params = self.og.GeneratorParams(self.genai_model) - params.set_search_options(max_length=4096, do_sample=False) - gen = self.og.Generator(self.genai_model, params) - - # Initialize context: append context tokens to build the initial state - context_tokens = full_tokens[:context_len] - - if len(context_tokens) > 0: - # Append context tokens as 1D numpy array - gen.append_tokens(np.array(context_tokens, dtype=np.int32)) - else: - # If no context, append BOS if available - if self.prefix_token_id is not None: - gen.append_tokens( - np.array([self.prefix_token_id], dtype=np.int32) - ) - - # Teacher forcing: iterate through continuation tokens - total_ll = 0.0 - greedy_ok = True - - for tok in continuation_tokens: - # Get logits for next token prediction (only last position) - logits = gen.get_logits() - logits_array = np.array(logits, dtype=np.float32) - - # Extract logits vector (get_logits returns last position only) - if len(logits_array.shape) == 3: # (batch, 1, vocab) - next_logits = logits_array[0, -1, :] - elif len(logits_array.shape) == 2: # (1, vocab) - next_logits = logits_array[-1, :] - else: # (vocab,) - next_logits = logits_array - - # Compute log probability using numerically stable log_softmax - max_logit = np.max(next_logits) - exp_logits = np.exp(next_logits - max_logit) - log_sum_exp = max_logit + np.log(np.sum(exp_logits)) - log_prob = next_logits[int(tok)] - log_sum_exp - - total_ll += float(log_prob) - - # Check if greedy (argmax matches actual token) - if int(np.argmax(next_logits)) != int(tok): - greedy_ok = False - - # Teacher forcing: append the actual token to continue generation - gen.append_tokens(np.array([int(tok)], dtype=np.int32)) - - results.append((total_ll, greedy_ok)) - - except Exception as e: - eval_logger.warning(f"Failed to compute loglikelihood: {e}") - import traceback - - eval_logger.debug(traceback.format_exc()) - results.append((0.0, False)) - - return results - - def loglikelihood_rolling( - self, requests: list["Instance"], disable_tqdm: bool = False - ) -> list[float]: - """ - Compute rolling log-likelihood for perplexity using ONNX Runtime GenAI. - Uses sliding windows to handle sequences longer than max_length. - - Args: - requests: List of instances containing text sequences - disable_tqdm: Whether to disable progress bar - - Returns: - List of sum of log-likelihood values for each request - """ - loglikelihoods = [] - - for request in tqdm( - requests, disable=disable_tqdm, desc="Computing rolling log-likelihoods" - ): - string = request.args[0] - - # Use sliding window approach for long sequences - rolling_token_windows = list( - map( - utils.make_disjoint_window, - utils.get_rolling_token_windows( - token_list=self.tok_encode(string), - prefix_token=self.prefix_token_id, - max_seq_len=self.max_length, - context_len=1, - ), - ) - ) - - # Prepare windows for loglikelihood (convert to context/continuation text pairs) - window_requests = [] - for context_tokens, continuation_tokens in rolling_token_windows: - context_text = self.tok_decode(context_tokens) - continuation_text = self.tok_decode(continuation_tokens) - window_requests.append((context_text, continuation_text)) - - # Use loglikelihood to compute log-likelihoods for all windows - # Create Instance objects for the windows - from lm_eval.api.instance import Instance - - window_instances = [ - Instance( - request_type="loglikelihood", - doc={}, - arguments=(ctx, cont), - idx=0, - metadata={}, - ) - for ctx, cont in window_requests - ] - - string_nll = self.loglikelihood( - window_instances, - disable_tqdm=True, - ) - - # Extract log-likelihoods (discard is_greedy boolean) - string_nll = [x[0] for x in string_nll] - - # Sum all window log-likelihoods - total_nll = sum(string_nll) - loglikelihoods.append(total_nll) - - # Cache this loglikelihood_rolling request - self.cache_hook.add_partial("loglikelihood_rolling", (string,), total_nll) - - return loglikelihoods - - def generate_until( - self, requests: list["Instance"], disable_tqdm: bool = False - ) -> list[str]: - """ - Generate text until stopping criteria using ONNX Runtime GenAI. - - Args: - requests: List of generation requests with context and generation kwargs - disable_tqdm: Whether to disable progress bar - - Returns: - List of generated text strings - """ - if not requests: - return [] - - results = [] - - for request in tqdm(requests, disable=disable_tqdm, desc="Generating text"): - context, gen_kwargs = request.args - - max_gen_toks = gen_kwargs.get("max_gen_toks", self.max_gen_toks) - until = gen_kwargs.get("until", []) - temperature = gen_kwargs.get("temperature", 0.0) - top_p = gen_kwargs.get("top_p", 1.0) - top_k = gen_kwargs.get("top_k", 50) - do_sample = gen_kwargs.get("do_sample", False) - - try: - # Use GenAI generation API - generated_text = self._run_genai_generation( - context, - max_gen_toks, - until, - temperature=temperature, - top_p=top_p, - top_k=top_k, - do_sample=do_sample, - ) - results.append(generated_text) - - except Exception as e: - eval_logger.warning(f"Generation failed for request: {e}") - results.append("") - - return results - - def _run_genai_generation( - self, - prompt: str, - max_tokens: int = 4096, - stop_sequences: list[str] | None = None, - temperature: float = 0.0, - top_p: float = 1.0, - top_k: int = 50, - do_sample: bool = False, - ) -> str: - """ - Run text generation using ONNX Runtime GenAI. - - Args: - prompt: Input text to generate from - max_tokens: Maximum number of tokens to generate - stop_sequences: List of sequences that will stop generation - temperature: Sampling temperature (0 = greedy, higher = more random) - top_p: Nucleus sampling threshold - top_k: Top-k sampling parameter - do_sample: Whether to use sampling (if False, uses greedy decoding) - - Returns: - Generated text string - """ - try: - # Create generator parameters - params = self.og.GeneratorParams(self.genai_model) - # Set generation parameters based on mode - if do_sample and temperature > 0: - # Sampling mode: use provided parameters - params.set_search_options( - max_length=int(max_tokens), - top_p=float(top_p), - top_k=int(top_k), - temperature=float(temperature), - repetition_penalty=1.0, - ) - else: - # Greedy mode: force temperature to 0 for deterministic output - params.set_search_options( - max_length=int(max_tokens), - top_p=1.0, - top_k=1, - temperature=0.0, - repetition_penalty=1.0, - ) - # Create generator - generator = self.og.Generator(self.genai_model, params) - # Encode prompt and append tokens to the generator - input_tokens = self.genai_tokenizer.encode(prompt) - generator.append_tokens(input_tokens) - - # Generate tokens - generated_tokens = [] - while not generator.is_done(): - generator.generate_next_token() - - if generator.is_done(): - break - - # Get the latest generated token - output_tokens = generator.get_sequence(0) - base = len(input_tokens) + len(generated_tokens) - if len(output_tokens) > base: - new_token = output_tokens[base] - generated_tokens.append(new_token) - - # Check stopping criteria - if len(generated_tokens) >= max_tokens: - break - - # Decode current generation to check stop sequences - if stop_sequences: - current_text = self.genai_tokenizer.decode(generated_tokens) - if any(stop_seq in current_text for stop_seq in stop_sequences): - break - - # Decode generated tokens - if generated_tokens: - generated_text = self.genai_tokenizer.decode(generated_tokens) - - # Remove stop sequences from the end - if stop_sequences: - for stop_seq in stop_sequences: - if generated_text.endswith(stop_seq): - generated_text = generated_text[: -len(stop_seq)] - break - - return generated_text - else: - return "" - - except Exception as e: - eval_logger.error(f"GenAI generation error: {e}") - return "" + except Exception as e: # noqa: BLE001 + eval_logger.warning("Windows ML device enumeration failed: %s", e) diff --git a/pyproject.toml b/pyproject.toml index 7aa5cb2bb..3a2727b2c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +59,10 @@ gptq = ["auto-gptq[triton]>=0.6.0"] gptqmodel = ["gptqmodel>=1.0.9"] ipex = ["optimum-intel"] ibm_watsonx_ai = ["ibm_watsonx_ai>=1.1.22", "python-dotenv"] +# ONNX Runtime GenAI backend. For non-CPU execution providers install the +# matching wheel instead of `onnxruntime-genai`, e.g. `onnxruntime-genai-cuda` +# or `onnxruntime-genai-directml` (these are mutually exclusive). +onnxruntime-genai = ["onnxruntime-genai", "transformers", "numpy"] litellm = ["litellm>=1.60,<1.85", "aiohttp", "requests", "tenacity", "tqdm"] # mamba requires CUDA (nvcc) - cannot build on macOS/CPU-only systems # mamba = ["mamba_ssm", "causal-conv1d==1.0.2", "torch"] diff --git a/tests/models/test_onnxruntime_genai.py b/tests/models/test_onnxruntime_genai.py new file mode 100644 index 000000000..16a6f6977 --- /dev/null +++ b/tests/models/test_onnxruntime_genai.py @@ -0,0 +1,239 @@ +"""Tests for the cross-platform onnxruntime-genai backend. + +The heavy smoke tests require ``onnxruntime-genai`` plus its Model Builder and +build a tiny random Llama model on the fly (CPU execution provider), mirroring +the export-on-the-fly pattern in ``tests/models/test_openvino.py``. They are +excluded from CI (see the ``--ignore`` in ``.github/workflows/unit_tests.yml``) +because the ONNX runtimes are not installed there; run them locally with: + + python -m pytest tests/models/test_onnxruntime_genai.py -vv +""" + +import tempfile + +import numpy as np +import pytest + +from lm_eval.api.model import CacheHook +from lm_eval.models.onnxruntime_genai import ONNXRuntimeGenAILM, _log_softmax + + +pytest.importorskip("onnxruntime_genai") + + +def _build_tiny_model_dir(): + """Build a tiny random Llama Model Builder export (fp32, CPU). + + head_size must be a multiple of 16 for the genai CPU GroupQueryAttention + kernel, so hidden_size=128 / num_heads=8 -> head_size=16. + """ + builder = pytest.importorskip("onnxruntime_genai.models.builder") + from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM + + src = tempfile.mkdtemp(prefix="ort_genai_src_") + out = tempfile.mkdtemp(prefix="ort_genai_out_") + cache = tempfile.mkdtemp(prefix="ort_genai_cache_") + + tokenizer = AutoTokenizer.from_pretrained( + "hf-internal-testing/tiny-random-LlamaForCausalLM" + ) + config = LlamaConfig( + vocab_size=tokenizer.vocab_size, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=8, + num_key_value_heads=8, + max_position_embeddings=512, + tie_word_embeddings=True, + ) + LlamaForCausalLM(config).save_pretrained(src) + tokenizer.save_pretrained(src) + + try: + builder.create_model(None, src, out, "fp32", "cpu", cache) + except Exception as e: # noqa: BLE001 + pytest.skip(f"Model Builder could not create a CPU fixture: {e}") + return out + + +@pytest.fixture(scope="session") +def tiny_model_dir(): + return _build_tiny_model_dir() + + +# --------------------------------------------------------------------------- # +# Unit tests: scoring math, no model load required. +# --------------------------------------------------------------------------- # +def _make_bare_lm(forward_logits, max_length=2048): + """Build an instance without loading a model, wiring only what the token scoring path needs.""" + lm = ONNXRuntimeGenAILM.__new__(ONNXRuntimeGenAILM) + lm.max_length = max_length + lm.cache_hook = CacheHook(None) + lm._forward_logits = forward_logits + return lm + + +def test_log_softmax_matches_reference(): + rng = np.random.default_rng(0) + logits = rng.normal(size=(4, 7)).astype(np.float32) + got = _log_softmax(logits) + # Reference: log(exp(x) / sum(exp(x))) + ref = logits - np.log(np.exp(logits).sum(axis=-1, keepdims=True)) + np.testing.assert_allclose(got, ref, rtol=1e-5, atol=1e-5) + # Each row is a valid log-prob distribution (sums to 1 in prob space). + np.testing.assert_allclose(np.exp(got).sum(axis=-1), np.ones(4), atol=1e-5) + + +def test_loglikelihood_tokens_math(): + # vocab=5; craft logits so the continuation tokens are the argmax everywhere. + vocab = 5 + context_enc = [1, 2] + continuation_enc = [3, 4] + full_len = len(context_enc) + len(continuation_enc) + + # Deterministic logits: position i strongly predicts continuation target. + logits = np.full((full_len, vocab), -10.0, dtype=np.float32) + # Row (ctxlen-1 + j) should predict continuation_enc[j]. + ctxlen = len(context_enc) + for j, tok in enumerate(continuation_enc): + logits[ctxlen - 1 + j, tok] = 10.0 + + lm = _make_bare_lm(lambda toks: logits) + ((ll, is_greedy),) = lm._loglikelihood_tokens( + [(None, context_enc, continuation_enc)], disable_tqdm=True + ) + + # Manually compute expected loglikelihood from the same logits. + lp = _log_softmax(logits[ctxlen - 1 : ctxlen - 1 + len(continuation_enc)]) + expected = float(lp[np.arange(len(continuation_enc)), continuation_enc].sum()) + assert ll == pytest.approx(expected) + assert is_greedy is True + + +def test_loglikelihood_tokens_not_greedy(): + vocab = 5 + context_enc = [1] + continuation_enc = [2] + # Predict a *different* token as argmax, so is_greedy must be False. + logits = np.full((2, vocab), -10.0, dtype=np.float32) + logits[0, 4] = 10.0 # argmax is 4, but target is 2 + lm = _make_bare_lm(lambda toks: logits) + ((ll, is_greedy),) = lm._loglikelihood_tokens( + [(None, context_enc, continuation_enc)], disable_tqdm=True + ) + assert is_greedy is False + assert ll < 0.0 + + +def test_loglikelihood_tokens_empty_continuation(): + lm = _make_bare_lm(lambda toks: np.zeros((1, 3), dtype=np.float32)) + (result,) = lm._loglikelihood_tokens([(None, [1], [])], disable_tqdm=True) + assert result == (0.0, True) + + +class _FakeConfig: + """Records provider mutations made by _select_ep.""" + + def __init__(self): + self.cleared = False + self.appended = [] + self.options = [] + + def clear_providers(self): + self.cleared = True + + def append_provider(self, ep): + self.appended.append(ep) + + def set_provider_option(self, ep, key, value): + self.options.append((ep, key, value)) + + +def _make_ep_lm(execution_provider, provider_options=None): + lm = ONNXRuntimeGenAILM.__new__(ONNXRuntimeGenAILM) + lm.execution_provider = execution_provider + lm.provider_options = provider_options or {} + return lm + + +def test_select_ep_none_leaves_config_untouched(): + # Regression: with no execution_provider, the providers Model Builder wrote + # into genai_config.json must survive (clear_providers must NOT be called), + # else a CUDA/DML/NPU export would be silently downgraded to CPU. + cfg = _FakeConfig() + _make_ep_lm(None)._select_ep(cfg) + assert cfg.cleared is False + assert cfg.appended == [] + + +def test_select_ep_cpu_clears_without_appending(): + cfg = _FakeConfig() + _make_ep_lm("cpu")._select_ep(cfg) + assert cfg.cleared is True + assert cfg.appended == [] + + +def test_select_ep_named_provider_clears_and_appends(): + cfg = _FakeConfig() + _make_ep_lm("cuda", {"device_id": 0})._select_ep(cfg) + assert cfg.cleared is True + assert cfg.appended == ["cuda"] + assert cfg.options == [("cuda", "device_id", "0")] + + +# --------------------------------------------------------------------------- # +# Smoke tests: build a tiny model and run real tasks end-to-end on CPU. +# --------------------------------------------------------------------------- # +@pytest.fixture +def lm(tiny_model_dir): + from lm_eval.api.registry import get_model + + return get_model("onnxruntime-genai").create_from_arg_string( + f"pretrained={tiny_model_dir},execution_provider=cpu", + {"batch_size": 1}, + ) + + +def test_backend_loads(lm): + assert lm.eot_token_id is not None + assert lm.max_length > 0 + ids = lm.tok_encode("hello world") + assert isinstance(ids, list) and len(ids) > 0 + # Tokenization round-trips (up to special-token stripping). + assert "hello" in lm.tok_decode(ids) + + +@pytest.mark.parametrize("task", ["hellaswag", "wikitext", "gsm8k"]) +def test_simple_evaluate(lm, task): + from lm_eval import evaluator + + results = evaluator.simple_evaluate( + model=lm, + tasks=[task], + limit=2, + num_fewshot=0, + bootstrap_iters=0, + ) + assert task in results["results"] + metrics = results["results"][task] + # Some real metric must be present and finite (no silent 0.0-on-error). + numeric = [v for v in metrics.values() if isinstance(v, (int, float))] + assert numeric, f"no numeric metrics for {task}: {metrics}" + assert all(np.isfinite(v) for v in numeric) + + +def test_loglikelihood_deterministic(lm): + """The same request scored twice yields identical log-probs (greedy path).""" + from lm_eval.api.instance import Instance + + req = Instance( + request_type="loglikelihood", + doc={}, + arguments=("The capital of France is", " Paris"), + idx=0, + ) + (first,) = lm.loglikelihood([req], disable_tqdm=True) + (second,) = lm.loglikelihood([req], disable_tqdm=True) + assert first[0] == pytest.approx(second[0]) + assert first[0] < 0.0 # a real, non-zero log-prob diff --git a/tests/models/test_winml.py b/tests/models/test_winml.py new file mode 100644 index 000000000..0ad7e4e96 --- /dev/null +++ b/tests/models/test_winml.py @@ -0,0 +1,125 @@ +"""Tests for the winml backend. + +``winml`` (:class:`lm_eval.models.winml.WindowsML`) is a thin subclass of +:class:`lm_eval.models.onnxruntime_genai.ONNXRuntimeGenAILM` that overrides only +execution-provider selection. The shared scoring/generation logic is covered by +``test_onnxruntime_genai.py``; this module covers the winml-specific wiring: +registration, the subclass relationship, and that Windows-provider registration +degrades gracefully to the cross-platform CPU path when the Windows ML APIs are +unavailable (e.g. off-Windows), so the inherited scoring core still runs. + +Like the genai tests, these are excluded from CI (see the ``--ignore`` in +``.github/workflows/unit_tests.yml``); run locally with: + + python -m pytest tests/models/test_winml.py -vv +""" + +import sys +import tempfile + +import numpy as np +import pytest + +from lm_eval.models.onnxruntime_genai import ONNXRuntimeGenAILM + + +pytest.importorskip("onnxruntime_genai") + + +def test_winml_registered_and_is_subclass(): + from lm_eval.api.registry import get_model + from lm_eval.models.winml import WindowsML + + assert get_model("winml") is WindowsML + assert issubclass(WindowsML, ONNXRuntimeGenAILM) + # winml should specialize only provider selection + defaults, not the + # scoring core. + assert "_select_ep" in vars(WindowsML) + for inherited in ("_loglikelihood_tokens", "loglikelihood", "generate_until"): + assert inherited not in vars(WindowsML) + + +def test_winml_pins_historical_defaults(tiny_model_dir): + # winml keeps its historical max_length / max_gen_toks so existing runs are + # numerically unchanged, independent of the base backend's modern defaults. + from lm_eval.api.registry import get_model + + lm = get_model("winml").create_from_arg_string( + f"pretrained={tiny_model_dir}", {"batch_size": 1} + ) + assert lm.max_length == 4096 + assert lm.max_gen_toks == 4096 + # explicit overrides still win + lm2 = get_model("winml").create_from_arg_string( + f"pretrained={tiny_model_dir},max_length=1024", {"batch_size": 1} + ) + assert lm2.max_length == 1024 + + +def _build_tiny_model_dir(): + """Build a tiny random Llama Model Builder export (fp32, CPU). + + head_size must be a multiple of 16 for the genai CPU GroupQueryAttention + kernel, so hidden_size=128 / num_heads=8 -> head_size=16. + """ + builder = pytest.importorskip("onnxruntime_genai.models.builder") + from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM + + src = tempfile.mkdtemp(prefix="winml_src_") + out = tempfile.mkdtemp(prefix="winml_out_") + cache = tempfile.mkdtemp(prefix="winml_cache_") + + tokenizer = AutoTokenizer.from_pretrained( + "hf-internal-testing/tiny-random-LlamaForCausalLM" + ) + config = LlamaConfig( + vocab_size=tokenizer.vocab_size, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=8, + num_key_value_heads=8, + max_position_embeddings=512, + tie_word_embeddings=True, + ) + LlamaForCausalLM(config).save_pretrained(src) + tokenizer.save_pretrained(src) + + try: + builder.create_model(None, src, out, "fp32", "cpu", cache) + except Exception as e: # noqa: BLE001 + pytest.skip(f"Model Builder could not create a CPU fixture: {e}") + return out + + +@pytest.fixture(scope="session") +def tiny_model_dir(): + return _build_tiny_model_dir() + + +@pytest.mark.skipif( + sys.platform == "win32", + reason="Exercises the non-Windows fallback path; on Windows the Windows ML " + "catalog drives provider selection and requires a real NPU/GPU device.", +) +def test_winml_falls_back_to_cpu_off_windows(tiny_model_dir): + """Off-Windows, winml should load via the base CPU path and score correctly.""" + from lm_eval.api.registry import get_model + + lm = get_model("winml").create_from_arg_string( + f"pretrained={tiny_model_dir},execution_provider=cpu", + {"batch_size": 1}, + ) + from lm_eval import evaluator + + results = evaluator.simple_evaluate( + model=lm, + tasks=["hellaswag"], + limit=2, + num_fewshot=0, + bootstrap_iters=0, + ) + metrics = results["results"]["hellaswag"] + numeric = [v for v in metrics.values() if isinstance(v, (int, float))] + assert numeric + assert all(np.isfinite(v) for v in numeric)