Optimum ONNX Runtime: 5 Proven Steps for Faster Inference

Optimum ONNX Runtime is the Hugging Face integration that exports a Transformers model to ONNX and runs it with Microsoft's ONNX Runtime. The first version of this guide promised 40% faster training with ORTTrainer. That class no longer exists: training support was deprecated in Optimum 1.27 and removed in 2.0, and today the import fails. This rewrite keeps what still works in 2026 (install, export, inference, quantization) and shows every step with real output from a test run.

Lab setup: python:3.12-slim container, CPU only (20 threads, no GPU), PyTorch 2.15.0 CPU wheel, then the install command from the original post. The model is distilbert/distilbert-base-uncased-finetuned-sst-2-english, a small sentiment classifier. GPU timings and the CUDA package were not tested here.

Install Optimum ONNX Runtime with pip

Most visitors arrive here searching for pip install optimum[onnxruntime]. That command still works, but since Optimum 2.0 (October 2025) the ONNX code lives in a separate package, optimum-onnx. The Optimum 2.0 release notes say "ONNX export and ONNX Runtime inference related integrations were moved to Optimum ONNX", and the extra simply pulls it in. Here is what pip installed for Optimum ONNX Runtime on 6 November 2026:

$ pip install -q optimum[onnxruntime]
$ pip list | grep -i -E "^(optimum|onnx|onnxruntime|transformers|torch) "
onnx               1.23.2
onnxruntime        1.31.1
optimum            2.1.0
torch              2.15.0+cpu
transformers       4.57.6
$ pip show optimum-onnx | grep -E "^(Name|Version|Requires)"
Name: optimum-onnx
Version: 0.1.0
Requires: onnx, optimum, transformers

Notice the version: PyPI already has optimum 2.3.0, but pip installed 2.1.0. The latest optimum-onnx (0.1.0) declares optimum~=2.1.0 and transformers<4.58.0 in its metadata, so pip backs off to the newest combination that fits. If another package in your project needs a newer Optimum or Transformers, this is where the conflict will come from.

Quote the brackets in zsh when you install Optimum ONNX Runtime

The unquoted command from the old post works in bash. In zsh (the default shell on macOS) the square brackets are a glob pattern:

$ zsh -c "pip install optimum[onnxruntime]"
zsh:1: no matches found: optimum[onnxruntime]
$ zsh -c "pip install -q \"optimum[onnxruntime]\" && echo quoted install OK"
quoted install OK

Always quote it. The optimum-onnx README now recommends installing the package directly with pip install "optimum-onnx[onnxruntime]", or "optimum-onnx[onnxruntime-gpu]" for NVIDIA GPUs. It also warns that onnxruntime and onnxruntime-gpu conflict, so run pip uninstall onnxruntime before installing the GPU variant.

ORTTrainer is gone from Optimum ONNX Runtime

The training code in the original Optimum ONNX Runtime post imported ORTTrainer and ORTTrainingArguments. With the current packages both imports fail:

$ python -c "from optimum.onnxruntime import ORTTrainer"
ImportError: cannot import name 'ORTTrainer' from 'optimum.onnxruntime' (/usr/local/lib/python3.12/site-packages/optimum/onnxruntime/__init__.py)

$ python -c "from optimum.onnxruntime import ORTTrainingArguments"
ImportError: cannot import name 'ORTTrainingArguments' from 'optimum.onnxruntime' (/usr/local/lib/python3.12/site-packages/optimum/onnxruntime/__init__.py)

The history is in the release notes. Optimum 1.27.0 (July 2025) "deprecated support for TFLite, BetterTransformer, and ONNXRuntime-Training, these integrations will be fully removed in v2". Optimum 2.0.0 says "ONNX Runtime Training officially deprecated", and the first optimum-onnx release lists "Remove ort training". The optim="adamw_ort_fused" and feature= arguments from the old example went with it. For fine-tuning, use the normal Transformers Trainer on PyTorch, then export the result. The 40% and "4 hours to 2 hours 45 minutes" figures in the old post could not be checked and are removed.

Export a model for Optimum ONNX Runtime with optimum-cli

Export is where Optimum ONNX Runtime still saves real work. One command downloads the checkpoint, traces it and writes an ONNX graph plus the tokenizer files:

$ optimum-cli export onnx --model distilbert/distilbert-base-uncased-finetuned-sst-2-english distilbert_onnx/
`torch_dtype` is deprecated! Use `dtype` instead!
Weight deduplication check in the ONNX export requires accelerate. Please install accelerate to run it.

$ ls -la distilbert_onnx/
-rw-r--r--. 1 root root       687 Nov  6 06:59 config.json
-rw-r--r--. 1 root root 267932172 Nov  6 06:59 model.onnx
-rw-r--r--. 1 root root       125 Nov  6 06:59 special_tokens_map.json
-rw-r--r--. 1 root root    711396 Nov  6 06:59 tokenizer.json
-rw-r--r--. 1 root root      1278 Nov  6 06:59 tokenizer_config.json
-rw-r--r--. 1 root root    231508 Nov  6 06:59 vocab.txt

The two messages are warnings, not errors; the export finished and the model works (next section). The task is detected from the model card. If you set it yourself, use the exporter's names: --task text-classification worked, while an invented name fails and lists the valid ones:

$ optimum-cli export onnx --model distilbert/distilbert-base-uncased-finetuned-sst-2-english --task sentiment x/
ValueError: Asked to export a distilbert model for the task sentiment, but the Optimum ONNX exporter only supports the tasks feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification for distilbert. ...

Run Optimum ONNX Runtime inference with ORTModel classes

The ORTModelForXXX classes replace the AutoModelForXXX classes and plug into the normal pipeline(). This script (infer.py) loads the exported folder, classifies two sentences and compares the raw logits with the original PyTorch model:

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
from optimum.onnxruntime import ORTModelForSequenceClassification
m = "distilbert/distilbert-base-uncased-finetuned-sst-2-english"
tok = AutoTokenizer.from_pretrained(m)
ort = ORTModelForSequenceClassification.from_pretrained("distilbert_onnx")
pt = AutoModelForSequenceClassification.from_pretrained(m).eval()
texts = ["This guide is clear and the export worked perfectly.", "The training loop crashed again after three hours."]
clf = pipeline("text-classification", model=ort, tokenizer=tok)
for t, r in zip(texts, clf(texts)):
    print(f"{r['label']:8} {r['score']:.4f}  {t}")
enc = tok(texts, return_tensors="pt", padding=True)
with torch.no_grad():
    a = pt(**enc).logits
b = ort(**enc).logits
print("providers:", ort.providers)
print("max abs logit difference PyTorch vs ONNX Runtime:", float((a - b).abs().max()))
$ python infer.py
Device set to use cpu
POSITIVE 0.9998  This guide is clear and the export worked perfectly.
NEGATIVE 0.9989  The training loop crashed again after three hours.
providers: ['CPUExecutionProvider']
max abs logit difference PyTorch vs ONNX Runtime: 1.430511474609375e-06

A difference of about 1e-6 in the logits is floating-point noise, so the exported model gives the same answers as PyTorch, which is the main promise of Optimum ONNX Runtime.

Asking for CUDA on the CPU package

The CPU wheel has no CUDA provider. Requesting it does not fall back to CPU, it stops with a clear error, which is the quickest way to find out which package you really installed:

available providers: ['AzureExecutionProvider', 'CPUExecutionProvider']
ValueError: Asked to use CUDAExecutionProvider as an ONNX Runtime execution provider, but the available execution providers are ['AzureExecutionProvider', 'CPUExecutionProvider'].

One trap I hit while writing the test: a helper script named cuda.py in the working directory broke import torch with "partially initialized module 'torch' has no attribute 'nn'", because PyTorch imports a package called cuda and Python found my file first. Do not name your own scripts after packages.

Quantize with Optimum ONNX Runtime and measure the speed

Dynamic int8 quantization is one command. You must choose a target instruction set; the help lists --arm64, --avx2, --avx512, --avx512_vnni, --ppc64le and --tensorrt:

$ optimum-cli onnxruntime quantize --onnx_model distilbert_onnx/ --avx2 -o distilbert_q/
$ ls -la distilbert_q/*.onnx distilbert_onnx/*.onnx
-rw-r--r--. 1 root root 267932172 Nov  6 06:59 distilbert_onnx/model.onnx
-rw-r--r--. 1 root root  67369391 Nov  6 07:00 distilbert_q/model_quantized.onnx

The file is four times smaller. Is it faster? The benchmark script runs one 128-token request 50 times after 5 warm-up runs, on the same CPU, for PyTorch, the exported fp32 model and the int8 model (load it with file_name="model_quantized.onnx"):

input tokens: 128 | torch threads: 14
PyTorch eager                   91.5 ms per request
ONNX Runtime fp32               33.3 ms per request
ONNX Runtime int8 (dynamic)     36.6 ms per request
fp32 labels: [1, 0] | int8 labels: [1, 0]
Optimum ONNX Runtime on CPU: ORTTrainer import error and latency of PyTorch, ONNX fp32 and int8
The removed ORTTrainer and the CPU latency test from the lab.

On this machine Optimum ONNX Runtime fp32 answered almost three times faster than PyTorch eager, and the int8 model kept the same labels but was slightly slower than fp32. Treat these as one measurement on one CPU, not a benchmark: quantization mainly saves memory and disk, and whether it speeds things up depends on the CPU and the --avx2/--avx512_vnni choice. Measure Optimum ONNX Runtime with your own model and input length before you switch.

When Optimum ONNX Runtime is worth it

Use it when you serve a fine-tuned encoder model (classification, embeddings, NER) on CPU and want lower latency with no change in output, or when the target only has ONNX Runtime. Keep training in PyTorch; there is no supported ONNX Runtime training path in Optimum any more. Check accuracy after quantizing, and pin the Optimum ONNX Runtime package optimum-onnx together with transformers in your requirements, because the package fixes both versions.

If you are still setting up the Python stack itself, see Linux AI development with PyTorch. For the export options of other model types, the Optimum ONNX export guide lists the supported tasks and architectures.

Comments

Popular posts from this blog

How to Play Minecraft Bedrock Edition on Linux: A Comprehensive Guide for Tech Professionals

Linux desktop frozen? My 5 go-to tricks to try - before forcing a hard reboot

The Ultimate Guide: How to Set Up DXVK in Wine on Linux for Enhanced Gaming Performance