ModelBuilder for onnxruntime-genai

The code base comes from https://github.com/microsoft/onnxruntime-genai/tree/main/src/python/py/models. It adds fast unit tests checking discrepancies, end to end test with the trained model. It supports more architectures.

Why mbext? Its main advantage is a short CI and fast tests: every supported architecture is validated with a tiny, randomly-initialised model that runs fully offline (no weights downloaded) in seconds, while still checking the numerical discrepancy between PyTorch and ONNX. See the documentation for details.

Documentation

The documentation is published at xadupre.github.io/docs/mbext. It lives under docs/ and is built with Sphinx:

pip install -e .[doc]
sphinx-build -b html docs docs/_build/html

It covers the design, the command lines, an automatically generated list of the supported architectures, the differences with the onnxruntime-genai model builder and with Mobius, the --private scenario, and a gallery with a runnable create_model example on random weights.

Convert a model

Example converting Qwen/Qwen3-8B to ONNX for CPU with int4 precision:

python -m modelbuilder.builder \
    -m Qwen/Qwen3-8B \
    -o qwen3-8b-cpu-int4 \
    -p int4 \
    -e cpu \
    -c cache_dir

The arguments are:

• -m/--model_name: model name on Hugging Face (use -i/--input instead for a local folder).
• -o/--output: folder where the ONNX model and additional files are written.
• -p/--precision: precision of the model (int4, bf16, fp16 or fp32).
• -e/--execution_provider: execution provider to target (cpu here).
• -c/--cache_dir: cache directory for Hugging Face files and temporary ONNX external data files.
• -r/--reuse-weights: for float models stored as safetensors, make the ONNX initializers reference the downloaded checkpoint bytes directly. The checkpoint is downloaded once under <output>/.weights and remains there as the ONNX model's external data.

For example, this exports an FP16 model without creating a second copy of its downloaded weights:

python -m modelbuilder.builder \
    -m Qwen/Qwen3-8B \
    -o qwen3-8b-cuda-fp16 \
    -p fp16 \
    -e cuda \
    -c cache_dir \
    --reuse-weights

The option is limited to floating-point exports because quantization changes the weight representation. Initializers which the builder derives rather than reads directly from the checkpoint continue to use model.onnx.data. With --input, the local checkpoint must already be inside the output directory; --model_name handles that layout automatically.

When the checkpoint dtype differs from the requested ONNX dtype, the model uses ONNX Cast nodes so the checkpoint bytes can still be reused. This avoids a second on-disk copy; ONNX Runtime may materialize the converted values in memory.

Custom model in a separate file

The --private option supports a custom model whose implementation lives outside this package. Its value is made of up to three ;-separated file paths:

python -m modelbuilder.builder \
    -i my-model \
    -o my-model-cpu-fp32 \
    -p fp32 \
    -e cpu \
    --private "modeling.py;convert.py;test.py"

• modeling-file (optional, may be empty): imported before the Hugging Face config is loaded so a custom architecture can register itself with transformers.
• convert-file: defines the ONNX builder used for the conversion. The builder is the module-level MODEL_BUILDER attribute if present, otherwise the single :class:modelbuilder.builders.Model subclass defined in the file.
• fast-test-file (optional): a fast test file used to validate the conversion.

Running the fast tests

If no model id is given on the command line (both -m/--model_name and -i/--input are omitted), the fast-test-file from the --private option is executed as a script (python fast-test-file) instead of converting a model. This runs the tests that validate the custom builder:

python -m modelbuilder.builder --private "modeling.py;convert.py;test.py"

The --private option must then provide a fast-test-file (the third path); otherwise an error is raised. In this test mode -o/--output, -p/--precision and -e/--execution_provider are not required.

A complete, runnable example lives in examples/private_model: it defines the modeling file (a Mixture-of-Experts model implemented from scratch as PrivateDecoderLayer / PrivateModel / PrivateModelForCausalLM, plus config and tokenizer), a converter that builds the custom MoE decoder layer, and the fast test (discrepancy + genai) on a two-layer dummy model, plus a CI job (.github/workflows/private_example.yml) that checks the --private "modeling.py;convert.py;test.py" command line works.

Style

black . && ruff check .

Development

pip install -e .[dev]

Converting a model downloads a full Hugging Face model into ./cache_dir (the default --cache_dir) and writes large ONNX artifacts (*.onnx / *.onnx.data). These directories are git-ignored, and .vscode/settings.json excludes them from the VS Code file watcher so the editor does not run out of file watches and repeatedly prompt to reload the window. The conversion also drops a .vscode/settings.json inside the output model folder with the same exclusions, so opening that folder directly in VS Code stays stable too.

Fast Unit tests

pytest tests/fast

Long Unit tests

python tests/trained/test_trained_tiny_llm.py

With a better machine:

LONGTEST=1 pytest tests/trained

Llama.cpp tests

tests/fast_llama_cpp compares llama.cpp and onnxruntime-genai on the same model. The conversion script convert_hf_to_gguf.py comes from the llama.cpp repository and its requirements pin a specific torch version (for example torch==2.11.0). Installing them downgrades torch from the version installed for the rest of the test suite.

torchaudio (and torchvision) compiled against the previous torch build is then incompatible with the downgraded torch. Because transformers imports torchaudio lazily, the mismatch surfaces while loading the model as:

ModuleNotFoundError: Could not import module 'LlamaForCausalLM'.

Uninstalling torchaudio (and torchvision) removes the mismatch, since these tests only need torch itself:

pip uninstall -y torchaudio torchvision

You can see the results in stats/end2end_results.json. Example:

{'first_diff': 0, 'delta_length': 4, 'expected_length': 16, 'total_diff': 16, 'precision': 'fp32', 'model_id': 'HuggingFaceTB/SmolLM3-3B', 'experiment': 'generate', 'provider': 'cpu'}
{'max_abs_err': 1.6875, '%_gt_0.1': np.float64(0.5278832959081836), '%_gt_0.01': np.float64(0.962624750499002), 'avg_abs_discrepancy': 0.1749267578125, 'shape': (1, 5, 128256), 'dtype': dtype('float16'), 'precision': 'fp16', 'model_id': 'HuggingFaceTB/SmolLM3-3B', 'experiment': 'forward'}