Skip to content

feat(models): add cross-platform onnxruntime-genai backend + refactor winml - #3960

Open
thiagocrepaldi wants to merge 3 commits into
EleutherAI:mainfrom
thiagocrepaldi:feat/onnxruntime-genai-backend
Open

feat(models): add cross-platform onnxruntime-genai backend + refactor winml#3960
thiagocrepaldi wants to merge 3 commits into
EleutherAI:mainfrom
thiagocrepaldi:feat/onnxruntime-genai-backend

Conversation

@thiagocrepaldi

@thiagocrepaldi thiagocrepaldi commented Jul 29, 2026

Copy link
Copy Markdown

Summary

Makes the existing winml backend cross-platform. The new onnxruntime-genai backend is essentially the current winml code with the Windows-specific execution-provider selection factored out — the onnxruntime-genai runtime core (og.Model/og.Generator, tokenization, logits, generation) was already platform-agnostic; only three methods in winml were Windows-locked. This PR lifts the shared logic into a base backend and reduces winml to just its Windows provider selection.

Concretely:

  • New onnxruntime-genai backend (lm_eval/models/onnxruntime_genai.py): ONNXRuntimeGenAILM(TemplateLM) runs Model Builder ONNX exports via the cross-platform og.Config provider API (CPU, CUDA, DirectML, WebGPU, AMD-NPU via VitisAI/RyzenAI). It carries the same lm-eval surface the winml backend implemented — loglikelihood (+ multiple_choice), loglikelihood_rolling, and generate_until.
  • winml becomes a thin subclass (WindowsML(ONNXRuntimeGenAILM), lm_eval/models/winml.py) overriding only execution-provider selection (_select_ep) — the Windows ML catalog methods are kept verbatim; if they're unavailable (e.g. off-Windows) it falls back to the cross-platform CPU path. This follows the existing OptimumLM(HFLM) single-file subclass precedent.

Because the two share one code path, the winml logic also gets the correctness cleanups that were made while lifting it into the base:

  • per-token teacher-forcing loop → a single get_output("logits") pass;
  • the direct loglikelihood override that bypassed TemplateLM token handling → uses the stock TemplateLM machinery (so the space-convention / empty-context handling is inherited);
  • the silent (0.0, False)-on-exception path that corrupted scores → removed.

Tests

  • tests/models/test_onnxruntime_genai.py — unit tests for the log-softmax/gather scoring math (fake _forward_logits, no model load) + CPU smoke tests over hellaswag/wikitext/gsm8k using a tiny Model Builder fixture built on the fly (mirrors test_openvino.py's export-on-the-fly pattern).
  • tests/models/test_winml.py — winml-specific wiring: registration, the subclass relationship, and graceful fallback to the cross-platform CPU path off-Windows.
  • Both are added to the CI --ignore list (like test_openvino.py) because the ONNX runtimes are not installed in CI.
  • Local runs for both backends for phi 3.5 mini instruct to make sure the new backends in lm-eval numbers match the pytorch's

onnxruntime-genai

Task Docs Metric torch fp32 GPU ONNX fp32 CPU
arc_challenge 1,172 acc 0.5913 0.5913
arc_challenge acc_norm 0.6092 0.6092
arc_easy 2,376 acc 0.8434 0.8434
arc_easy acc_norm 0.8329 0.8329
hellaswag 10,042 acc 0.5888 0.5888
hellaswag acc_norm 0.7693 0.7693
mmlu 14,042 acc 0.6877 0.6878
gsm8k 1,319 exact_match (flexible) 0.7953 0.7953
gsm8k exact_match (strict) 0.7930 0.7930

winml (using --limit 200 for brevity)

Task Metric torch fp32 ONNX fp32 winml reference
hellaswag acc 0.5200 0.5200 0.5200 0.5200
hellaswag acc_norm 0.6600 0.6600 0.6600 0.6600
arc_challenge acc 0.6000 0.6000 0.6000 0.6000
arc_challenge acc_norm 0.5950 0.5950 0.5950 0.5950
arc_easy acc 0.7900 0.7900 0.7900 0.7900
arc_easy acc_norm 0.8250 0.8250 0.8250 0.8250
mmlu acc 0.6895 0.6895 0.6895 0.6895
gsm8k exact_match (strict) 0.9000 0.9000 0.9000 0.9000

Packaging & docs

  • pip install "lm_eval[onnxruntime-genai]" (EP-specific wheels like onnxruntime-genai-cuda/-directml are user-installed and mutually exclusive).
  • README backend section + model table + extras table.

Scope

First stage of a larger effort to standardize ONNX evaluation in lm-eval. Out of scope here (follow-up PRs): a raw onnxruntime.InferenceSession backend (matches the exact runtime a customer deploys, and reaches EPs GenAI can't, e.g. ROCm/MIGraphX), a cross-backend genai-vs-raw numerical parity CI test, and an OpenAI-compatible completions server.

Test plan

  • pip install -e ".[onnxruntime-genai]"
  • Backends resolve: get_model("onnxruntime-genai"), get_model("winml")
  • python -m pytest tests/models/test_onnxruntime_genai.py tests/models/test_winml.py -vv (11 passed on Linux/CPU)
  • Local CPU smoke: hellaswag / wikitext / gsm8k produce sane, finite, non-zero-on-error metrics
  • ruff check / ruff format --check clean
  • winml on Windows 11 (testing CPU with --limit 200 for phi 3.5 mini instruct)
  • onnxruntime-genai on Linux (testing CPU and CUDA evaluation for phi 3.5 mini instruct)
  • winml on Windows for NPU (checking whether NPU testing is actually needed for this PR)

… winml

Add a cross-platform `onnxruntime-genai` backend that runs ONNX LLMs exported
by the ONNX Runtime GenAI Model Builder on CPU/CUDA/DirectML/WebGPU/AMD-NPU via
the `og.Config` provider API. All lm-eval logic (single-pass log-likelihood,
multiple-choice, rolling perplexity, generation) lives in
`ONNXRuntimeGenAILM(TemplateLM)`, following the `OptimumLM(HFLM)` single-file
subclass precedent.

Refactor `winml` into a thin `WindowsML(ONNXRuntimeGenAILM)` subclass that
overrides only execution-provider selection, inheriting the shared scoring
core. This removes the previous winml quality debt: the per-token
teacher-forcing loop (now a single `get_output("logits")` pass), the direct
`loglikelihood` override that bypassed `TemplateLM` token handling, and the
silent `(0.0, False)`-on-exception path that corrupted scores.

Add a shared parametrized test suite (unit tests for the scoring math with a
fake `_forward_logits`, plus CPU smoke tests over hellaswag/wikitext/gsm8k
using a tiny Model Builder fixture), packaging extra, and docs. The test is
added to the CI `--ignore` list since the ONNX runtimes are not installed in CI.

Part of QUARK-606 Phase 1. The raw `onnxruntime` backend, cross-backend parity
CI, OpenAI-compatible server, and Quark migration are follow-up stages.

Co-Authored-By: Claude <[email protected]>
@CLAassistant

CLAassistant commented Jul 29, 2026

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

thiagocrepaldi and others added 2 commits July 29, 2026 00:30
Trim the subclassing note to the pre-existing OptimumLM example instead of
documenting the onnxruntime-genai backend internals, which belong in the
backend's own docstrings/README rather than the general model-authoring guide.

Co-Authored-By: Claude <[email protected]>
Move winml-specific coverage (registration, subclass relationship, non-Windows
CPU fallback) into tests/models/test_winml.py, leaving test_onnxruntime_genai.py
focused on the genai backend and its scoring math. Keeps each backend's tests
self-contained and adds the new file to the CI ignore list.

Co-Authored-By: Claude <[email protected]>
@thiagocrepaldi
thiagocrepaldi marked this pull request as draft July 29, 2026 17:52
@thiagocrepaldi
thiagocrepaldi marked this pull request as ready for review August 4, 2026 19:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants