Support for OpenAI's gpt-oss open-weight models

[henrylearn2rock]

Please support https://openai.com/open-models/

[kunal-vaishnavi]

It is already supported through ONNX Runtime GenAI. Please check out the Hugging Face repo here for more information.

[LorenRd]

Hey @kunal-vaishnavi , is it possible to export using model builder to onnx with DirectML and int4 quantization?

[Eugene-M]

Load model from d:\Work\HuggingFace\gpt-oss-20b-onnx\cuda\cuda-int4-kquant-block-32-mixed\model.onnx failed:This is an invalid model. In Node, ("/model/layers.0/attn/GroupQueryAttention", GroupQueryAttention, "com.microsoft", -1) : ("/model/layers.0/attn/qkv_proj/Add/output_0": tensor(float16),"","","past_key_values.0.key": tensor(float16),"past_key_values.0.value": tensor(float16),"/model/attn_mask_reformat/attn_mask_subgraph/Sub/Cast/output_0": tensor(int32),"/model/attn_mask_reformat/attn_mask_subgraph/Gather/Cast/output_0": tensor(int32),"cos_cache": tensor(float16),"sin_cache": tensor(float16),"","","model.layers.0.attn.sinks": tensor(float16),) -> ("/model/layers.0/attn/GroupQueryAttention/output_0": tensor(float16),"present.0.key": tensor(float16),"present.0.value": tensor(float16),) , Error Node(/model/layers.0/attn/GroupQueryAttention) with schema(com.microsoft::GroupQueryAttention:1) has input size 12 not in range [min=7, max=11].

[kunal-vaishnavi]

Hey @kunal-vaishnavi , is it possible to export using model builder to onnx with DirectML and int4 quantization?

Yes, it is possible. The files in the Hugging Face repo were made with the model builder. Here is the PR with the model builder changes that were made for gpt-oss.

Load model from d:\Work\HuggingFace\gpt-oss-20b-onnx\cuda\cuda-int4-kquant-block-32-mixed\model.onnx failed:This is an invalid model. In Node, ("/model/layers.0/attn/GroupQueryAttention", GroupQueryAttention, "com.microsoft", -1) : ("/model/layers.0/attn/qkv_proj/Add/output_0": tensor(float16),"","","past_key_values.0.key": tensor(float16),"past_key_values.0.value": tensor(float16),"/model/attn_mask_reformat/attn_mask_subgraph/Sub/Cast/output_0": tensor(int32),"/model/attn_mask_reformat/attn_mask_subgraph/Gather/Cast/output_0": tensor(int32),"cos_cache": tensor(float16),"sin_cache": tensor(float16),"","","model.layers.0.attn.sinks": tensor(float16),) -> ("/model/layers.0/attn/GroupQueryAttention/output_0": tensor(float16),"present.0.key": tensor(float16),"present.0.value": tensor(float16),) , Error Node(/model/layers.0/attn/GroupQueryAttention) with schema(com.microsoft::GroupQueryAttention:1) has input size 12 not in range [min=7, max=11].

You will need a nightly version of ONNX Runtime to get the latest changes needed to load the model. See here for more details.

[LorenRd]

Thanks @kunal-vaishnavi for opening the PR, I followed the steps building the repo for source from your branch and using the nightly version of onnxruntime and running:

python builder.py -m openai/gpt-oss-20b -o ./gpt_oss -p int4 -e dml -c ./cache_dir --extra_options int4_op_types_to_quantize=MatMul/Gather int4_algo_config=k_quant_mixed

The output model.onnx.data was 40GB isn't it huge?

[kunal-vaishnavi]

Upon further investigation, it appears that Hugging Face loads the MoE weights differently than the original OpenAI code (which is what that branch was based on). The weight loading change was causing the above error message. I have pushed a new commit to fix this and verified that both the CPU and CUDA variants do not produce this error message anymore (the CPU model still has QMoE kernel parity issues but the CUDA model does not). I also updated the prior comments to remove any confusion for subsequent readers. Can you try rebuilding the branch on the latest commit and regenerating the ONNX model?

[LorenRd]

I notice that now tensorrt-llm is needed, I have it installed but it fails, do I need a more recent version?

  File "/home/rolo943/onnxruntime-genai/src/python/py/models/builder.py", line 2393, in make_qmoe_weights
    raise RuntimeError("Failed to quantize MoE weights with TensorRT-LLM. Please ensure TensorRT-LLM installs and runs successfully in your environment.")
RuntimeError: Failed to quantize MoE weights with TensorRT-LLM. Please ensure TensorRT-LLM installs and runs successfully in your environment.

Versions I have:

tensorrt                 10.11.0.33
tensorrt-cu12            10.11.0.33
tensorrt-cu12-bindings   10.11.0.33
tensorrt-cu12-libs       10.11.0.33
tensorrt-cu13            10.13.2.6
tensorrt-cu13-bindings   10.13.2.6
tensorrt-cu13-libs       10.13.2.6
tensorrt-llm             0.21.0

[kunal-vaishnavi]

I use the nvcr.io/nvidia/pytorch:25.03-py3 Docker image for my environment and I set the PIP_CONSTRAINT environment variable to empty (i.e. export PIP_CONSTRAINT='') inside it to allow package version changes. Here is the full pip list.

Package                    Version
-------------------------- -----------------------------
absl-py                    2.1.0
accelerate                 1.10.0
aenum                      3.1.16
aiohappyeyeballs           2.5.0
aiohttp                    3.11.13
aiosignal                  1.3.2
annotated-types            0.7.0
anyio                      4.8.0
apex                       0.1
argon2-cffi                23.1.0
argon2-cffi-bindings       21.2.0
arrow                      1.3.0
asciitree                  0.3.3
asttokens                  3.0.0
astunparse                 1.6.3
async-lru                  2.0.4
attrs                      25.1.0
audioread                  3.0.1
autocommand                2.2.2
babel                      2.17.0
backoff                    2.2.1
backports.tarfile          1.2.0
beautifulsoup4             4.13.3
black                      25.1.0
bleach                     6.2.0
blis                       0.7.11
build                      1.3.0
cachetools                 5.5.2
catalogue                  2.0.10
certifi                    2025.1.31
cffi                       1.17.1
charset-normalizer         3.4.1
click                      8.1.8
click-option-group         0.5.7
cloudpathlib               0.21.0
cloudpickle                3.1.1
cmake                      4.0.3
colored                    2.3.1
coloredlogs                15.0
comm                       0.2.2
confection                 0.1.5
contourpy                  1.3.1
cuda-bindings              12.8.0
cuda-python                12.8.0
cudf                       25.2.0
cudf-polars                25.2.0
cugraph                    25.2.0
cugraph-dgl                25.2.0
cugraph-pyg                25.2.0
cugraph-service-client     25.2.0
cugraph-service-server     25.2.0
cuml                       25.2.0
cupy-cuda12x               13.3.0
cuvs                       25.2.0
cycler                     0.12.1
cymem                      2.0.11
Cython                     3.0.12
dask                       2024.12.1
dask-cuda                  25.2.0
dask-cudf                  25.2.0
dask-expr                  1.1.21
datasets                   3.1.0
debugpy                    1.8.13
decorator                  5.2.1
defusedxml                 0.7.1
diffusers                  0.34.0
dill                       0.3.8
distributed                2024.12.1
distributed-ucxx           0.42.0
distro                     1.9.0
dm-tree                    0.1.9
einops                     0.8.1
evaluate                   0.4.5
execnet                    2.1.1
executing                  2.2.0
expecttest                 0.3.0
fastapi                    0.115.4
fasteners                  0.19
fastjsonschema             2.21.1
fastrlock                  0.8.3
filelock                   3.17.0
flash_attn                 2.7.3
flashinfer-python          0.2.5
flatbuffers                25.2.10
fonttools                  4.56.0
fqdn                       1.5.1
frozenlist                 1.5.0
fsspec                     2024.9.0
gast                       0.6.0
grpcio                     1.62.1
h11                        0.16.0
h5py                       3.12.1
hf-xet                     1.1.7
httpcore                   1.0.9
httpx                      0.28.1
huggingface-hub            0.34.4
humanfriendly              9.2
hypothesis                 6.127.9
idna                       3.10
importlib_metadata         8.0.0
inflect                    7.3.1
iniconfig                  2.0.0
intel-openmp               2021.4.0
ipykernel                  6.29.5
ipython                    9.0.1
ipython_pygments_lexers    1.1.1
isoduration                20.11.0
isort                      6.0.1
jaraco.collections         5.1.0
jaraco.context             5.3.0
jaraco.functools           4.0.1
jaraco.text                3.12.1
jedi                       0.19.2
Jinja2                     3.1.6
jiter                      0.10.0
joblib                     1.4.2
json5                      0.10.0
jsonpointer                3.0.0
jsonschema                 4.23.0
jsonschema-specifications  2024.10.1
jupyter_client             8.6.3
jupyter_core               5.7.2
jupyter-events             0.12.0
jupyter-lsp                2.2.5
jupyter_server             2.15.0
jupyter_server_terminals   0.5.3
jupyterlab                 4.3.5
jupyterlab_code_formatter  3.0.2
jupyterlab_pygments        0.3.0
jupyterlab_server          2.27.3
jupyterlab_tensorboard_pro 4.0.0
jupytext                   1.16.7
kiwisolver                 1.4.8
kvikio                     25.2.0
langcodes                  3.5.0
language_data              1.3.0
lark                       1.2.2
lazy_loader                0.4
libcudf                    25.2.0
libcugraph                 25.2.0
libcuml                    25.2.0
libcuvs                    25.2.0
libkvikio                  25.2.0
libraft                    25.2.0
librmm                     25.2.0
librosa                    0.10.1
libucxx                    0.42.0
lightning-thunder          0.2.2.dev0
lightning-thunder          0.2.2.dev0
lightning-utilities        0.14.0
lightning-utilities        0.14.0
lintrunner                 0.12.7
llvmlite                   0.42.0
locket                     1.0.0
looseversion               1.3.0
looseversion               1.3.0
marisa-trie                1.2.1
Markdown                   3.7
markdown-it-py             3.0.0
MarkupSafe                 3.0.2
matplotlib                 3.10.1
matplotlib-inline          0.1.7
mdit-py-plugins            0.4.2
mdurl                      0.1.2
mistune                    3.1.2
mkl                        2021.1.1
mkl-devel                  2021.1.1
mkl-include                2021.1.1
ml_dtypes                  0.5.3
mock                       5.2.0
more-itertools             10.3.0
mpi4py                     4.1.0
mpmath                     1.3.0
msgpack                    1.1.0
multidict                  6.1.0
multiprocess               0.70.16
murmurhash                 1.0.12
mypy-extensions            1.0.0
nbclient                   0.10.2
nbconvert                  7.16.6
nbformat                   5.10.4
nest-asyncio               1.6.0
networkx                   3.4.2
ninja                      1.11.1.3
notebook                   7.3.2
notebook_shim              0.2.4
numba                      0.59.1
numba-cuda                 0.2.0
numcodecs                  0.13.1
numpy                      1.26.4
nvfuser                    0.2.26a0+c5e1555
nvfuser                    0.2.26a0+c5e1555
nvidia-cublas-cu12         12.8.4.1
nvidia-cuda-cupti-cu12     12.8.90
nvidia-cuda-nvrtc-cu12     12.8.93
nvidia-cuda-runtime-cu12   12.8.90
nvidia-cudnn-cu12          9.10.2.21
nvidia-cudnn-frontend      1.10.0
nvidia-cufft-cu12          11.3.3.83
nvidia-cufile-cu12         1.13.1.3
nvidia-curand-cu12         10.3.9.90
nvidia-cusolver-cu12       11.7.3.90
nvidia-cusparse-cu12       12.5.8.93
nvidia-cusparselt-cu12     0.7.1
nvidia-dali-cuda120        1.47.0
nvidia-ml-py               12.570.86
nvidia-modelopt            0.29.0
nvidia-modelopt-core       0.29.0
nvidia-nccl-cu12           2.27.3
nvidia-nvimgcodec-cu12     0.4.1.21
nvidia-nvjitlink-cu12      12.8.93
nvidia-nvjpeg2k-cu12       0.8.1.40
nvidia-nvtiff-cu12         0.4.0.62
nvidia-nvtx-cu12           12.8.90
nvtx                       0.2.11
nx-cugraph                 25.2.0
onnx                       1.17.0
onnx_graphsurgeon          0.5.8
onnx-ir                    0.1.6
onnxruntime-genai-cuda     0.9.0.dev0
onnxruntime-gpu            1.23.0.dev20250810002
openai                     1.99.5
opencv                     4.10.0
opencv-python-headless     4.11.0.86
opt_einsum                 3.4.0
opt_einsum                 3.4.0
optimum                    1.27.0
optree                     0.14.1
ordered-set                4.1.0
overrides                  7.7.0
packaging                  23.2
pandas                     2.2.3
pandocfilters              1.5.1
parso                      0.8.4
partd                      1.4.2
pathspec                   0.12.1
peft                       0.17.0
pexpect                    4.9.0
pillow                     10.3.0
pip                        25.0.1
platformdirs               4.3.6
pluggy                     1.5.0
ply                        3.11
polars                     1.21.0
polygraphy                 0.49.18
pooch                      1.8.2
preshed                    3.0.9
prometheus_client          0.21.1
prompt_toolkit             3.0.50
propcache                  0.3.0
protobuf                   6.31.1
psutil                     7.0.0
ptyprocess                 0.7.0
PuLP                       3.0.2
pure_eval                  0.2.3
pyarrow                    19.0.1
pybind11                   2.13.6
pybind11_global            2.13.6
pycocotools                2.0+nv0.8.1
pycparser                  2.22
pydantic                   2.10.6
pydantic_core              2.27.2
Pygments                   2.19.1
pylibcudf                  25.2.0
pylibcugraph               25.2.0
pylibcugraphops            25.2.0
pylibraft                  25.2.0
pylibwholegraph            25.2.0
pynvjitlink                0.3.0
pynvml                     12.0.0
pyparsing                  3.2.1
pyproject_hooks            1.2.0
pytest                     8.1.1
pytest-flakefinder         1.1.0
pytest-rerunfailures       15.0
pytest-shard               0.1.2
pytest-xdist               3.6.1
python-dateutil            2.9.0.post0
python-hostlist            2.2.1
python-json-logger         3.3.0
pytorch-triton             3.2.0+gitb2684bf3b.nvinternal
pytz                       2023.4
PyYAML                     6.0.2
pyzmq                      26.2.1
raft-dask                  25.2.0
rapids-dask-dependency     25.2.0a0
referencing                0.36.2
regex                      2024.11.6
requests                   2.32.3
rfc3339-validator          0.1.4
rfc3986-validator          0.1.1
rich                       13.9.4
rmm                        25.2.0
rpds-py                    0.23.1
safetensors                0.5.3
scikit-learn               1.6.1
scipy                      1.15.2
Send2Trash                 1.8.3
sentencepiece              0.2.0
setuptools                 79.0.1
shellingham                1.5.4
six                        1.16.0
smart-open                 7.1.0
sniffio                    1.3.1
sortedcontainers           2.4.0
soundfile                  0.13.1
soupsieve                  2.6
soxr                       0.5.0.post1
spacy                      3.7.5
spacy-legacy               3.0.12
spacy-loggers              1.0.5
srsly                      2.5.1
stack-data                 0.6.3
starlette                  0.41.3
StrEnum                    0.4.15
sympy                      1.14.0
tabulate                   0.9.0
tbb                        2021.13.1
tblib                      3.0.0
tensorboard                2.16.2
tensorboard-data-server    0.7.2
tensorrt                   10.9.0.34
tensorrt_cu12              10.9.0.34
tensorrt_cu12_bindings     10.9.0.34
tensorrt_cu12_libs         10.9.0.34
tensorrt_llm               0.20.0
terminado                  0.18.1
thinc                      8.2.5
threadpoolctl              3.5.0
thriftpy2                  0.5.2
tiktoken                   0.10.0
tinycss2                   1.4.0
tokenizers                 0.21.4
tomli                      2.0.1
toolz                      1.0.0
torch                      2.8.0
torch-geometric            2.6.1
torch_tensorrt             2.7.0
torchaudio                 2.8.0
torchprofile               0.0.4
torchvision                0.23.0
tornado                    6.5.2
tqdm                       4.67.1
traitlets                  5.14.3
transformer_engine         2.1.0+450146a
transformers               4.55.0
treelite                   4.4.1
triton                     3.4.0
typeguard                  4.3.0
typer                      0.15.2
types-dataclasses          0.6.6
types-python-dateutil      2.9.0.20241206
typing_extensions          4.12.2
tzdata                     2025.1
ucx-py                     0.42.0
ucxx                       0.42.0
uri-template               1.3.0
urllib3                    2.0.7
uvicorn                    0.35.0
wasabi                     1.1.3
wcwidth                    0.2.13
weasel                     0.4.1
webcolors                  24.11.1
webencodings               0.5.1
websocket-client           1.8.0
Werkzeug                   3.1.3
wheel                      0.45.1
wrapt                      1.17.2
xdoctest                   1.0.2
xgboost                    2.1.3
xgrammar                   0.1.16
xxhash                     3.5.0
yarl                       1.18.3
zarr                       2.18.4
zict                       3.0.0
zipp                       3.19.2

If needed, you can uninstall and reinstall the following packages to the latest versions to resolve any runtime errors in the Docker environment: onnx-ir, onnxruntime-gpu, transformers, torch, torchaudio, torchvision, torch_tensorrt.

[LorenRd]

Good and bad news @kunal-vaishnavi , thanks for the support!

I've been able to export and run the model after installing tensorrt_llm, the bad part is that the model as you mentioned performs very slow even on powerful workstations (Screenshots after exporting on CPU):

Huge use of memory:

Degraded output:

[apsonawane]

Hello @LorenRd, we are currently working on the performance and accuracy of the CPU kernel for QMoE we should have an update soon.

[kunal-vaishnavi]

The PR to create your own ONNX models for OpenAI's gpt-oss has now been merged. The kernel changes needed to support running the CPU and WebGPU variants are in progress inside ONNX Runtime.

[LorenRd]

Any specific PR I should follow for DML/CPU optimization on inference? @kunal-vaishnavi

[kunal-vaishnavi]

The CPU PR to follow can be found here. The WebGPU work is still in development and a PR has not been opened in ONNX Runtime yet. Once ready, you can use WebGPU to get runtime support for gpt-oss across many hardware options.