microsoft / microsoft/onnxruntime-genai

Running custom encoder-decoder models in onnxruntime-genai

Open
#875 6 comments 0 reactions 0 assignees View on GitHub
model:transformer
Dominant language
C++
Stars
1.1k
Forks
354
Avg merge
2d 16h
Merged PRs (30d)
85

Description

## background

My question is about executing encoder-decoder models with onnx genai runtime. My goal is to convert the DONUT transformer https://arxiv.org/abs/2111.15664, a sequence-to-sequence transformer for document understanding with swin encoder and mbart decoder to onnx and run it using onnxruntime-genai.

I managed to convert the individual components to onnx. Now I'm stuck at writing a `genai_config.json` suitable for encoder-decoder models.

## steps

I started with the huggingface implementation of DONUT (https://huggingface.co/docs/transformers/v4.42.0/en/model_doc/donut#overview) and converted the encoder and merged decoder (with kv-cache) to onnx using optimum https://huggingface.co/docs/optimum/index. I converted the DONUT processor, which consists of an image processor and sentencepiece-tokenizer to onnx using onnx runtime extensions (https://github.com/microsoft/onnxruntime-extensions/blob/main/onnxruntime_extensions/tools/). I merged the image processor and swin-encoder into a single graph. I can provide the conversion scripts if needed.

My components have the following inputs/outputs shapes:

**tokenizer:**
![image](https://github.com/user-attachments/assets/10a7cd82-de94-4fbd-9543-36af8f2040c9)

**encoder with image processor:**
![image](https://github.com/user-attachments/assets/1d3458ad-6da1-4ca5-b175-bb9982f4875e)

**decoder:**
![image](https://github.com/user-attachments/assets/6db3d42b-4a2a-469e-ad57-75528d1f9b4e)
![image](https://github.com/user-attachments/assets/4bf1c668-ff0a-47c0-94f1-594cecb4ae15)
(more past keys + values)

## question

However, I am stuck at manually writing/loading a suitable `genai_config.json`, with the hidden states from the encoder for use in the decoder's attention mechanism. I'm aware of https://onnxruntime.ai/docs/genai/reference/config.html to write configs, but it seemingly focuses on decoder-only models. I'm also aware of `make_genai_config` https://github.com/microsoft/onnxruntime-genai/blob/c7eba3c63a454edd6662eb007ff397d1146cc081/src/python/py/models/builder.py for auto-config generation of supported models.

I tried the following config:
```json
{
"model": {
"bos_token_id": 1,
"context_length": 1536,
"decoder": {
"session_options": {
"log_id": "onnxruntime-genai",
"provider_options": []
},
"filename": "decoder_model_merged.onnx",
"head_size": 96,
"hidden_size": 1024,
"inputs": {
"input_ids": "input_ids",
"encoder_hidden_states": "encoder_hidden_states",
"past_key_names": "past_key_values.%d.decoder.key",
"past_value_names": "past_key_values.%d.decoder.value"
},
"outputs": {
"logits": "logits",
"present_key_names": "present.%d.decoder.key",
"present_value_names": "present.%d.decoder.value"
},
"num_attention_heads": 16,
"num_hidden_layers": 12,
"num_key_value_heads": 32
},
"encoder": {
"filename": "encoder_with_img_processor.onnx",
"inputs": {
"pixel_values": "pixel_values"
},
"outputs": {
"encoder_hidden_states": "last_hidden_state"
}
},
"eos_token_id": 2,
"pad_token_id": 1,
"type": "sentencepiece",
"vocab_size": 57654
},
"search": {
"diversity_penalty": 0.0,
"do_sample": false,
"early_stopping": true,
"length_penalty": 1.0,
"max_length": 1536,
"min_length": 0,
"no_repeat_ngram_size": 0,
"num_beams": 1,
"num_return_sequences": 1,
"past_present_share_buffer": true,
"repetition_penalty": 1.0,
"temperature": 1.0,
"top_k": 50,
"top_p": 1.0
}
}
```

When loading the config I receive the error `RuntimeError: Error encountered while parsing 'output/genai_config.json' JSON Error: Unknown value: encoder_hidden_states at line 15 index 64`. As far as I can tell from https://github.com/microsoft/onnxruntime-genai/blob/c7eba3c63a454edd6662eb007ff397d1146cc081/src/python/py/models/builder.py#L71 it's currently not possible to pass encoder hidden states to the model as inputs.

Do you have any plans of extending onnx-genai runtime for encoder-decoder models? Could you please give me a hint/advice how to work around this?

Thank you for your assistance.

Contributor guide

No contributing guide indexed for this repository

Research direction

Read the linked config reference and src/python/py/models/builder.py around line 71, then inspect how encoder and decoder filenames, inputs, and outputs are parsed. Reproduce the supplied genai_config.json load error involving encoder_hidden_states; done would be a documented or implemented path for passing encoder hidden states to the decoder, if encoder-decoder support is pursued.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
ai, backend, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
20/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.