Exporting Llava decoder into ONNX format
Author: EricJi150
Pull request: no

I am working on exporting Llava into ONNX format. I came across this previous issue: https://github.com/huggingface/transformers/issues/33637 which had a notebook that outlined how to export in three separate parts. I noticed there wasn't any actual code on how the decoder was exported unlike the other two components. Does anyone know how they were able to export the decoder in the original notebook?

Notebook: https://colab.research.google.com/drive/1IhC8YOV68cze0XWGfuqSclnVTt_FskUd?usp=sharing

Comments:
Hey, the notebook afaik is self consistent and has code for exporting all 3 components outline here

```python
!onnxslim ./transformers.js/embed_tokens.onnx ./final/embed_tokens.onnx
!onnxslim ./transformers.js/vision_encoder.onnx ./final/vision_encoder.onnx
!onnxslim ./transformers.js/models/language-model/onnx/model.onnx ./final/decoder_model_merged.onnx
```

However it might be a bit outdated, are you seeing issues after running it?

Hi, thanks for the response. I believe those three lines are only used to optimize the models once they are already in ONNX format. However, I can't seem to see the code where decoder_model_merged.onnx is exported. The code is there for embed_tokens and vision_encoder, but seems like the cell for decoder_model_merged isn't there. I haven't been able to run the whole notebook on collab as I run out of memory after the first few cells.

Oh yeah, you're right! It must be the cell right before slimming, smth like

`!python -m scripts.convert --model_id ./language-model --tokenizer_id ./tokenizer \
--task text-generation-with-past --skip_validation --trust_remote_code`

I am quite sure the code saves all 3 parts. I am not sure which model you are testing with, you can try a tiny [0.5B llava](https://huggingface.co/llava-hf/llava-interleave-qwen-0.5b-hf/tree/main) to test it without running OOM

I tried just exporting the language model as is in the notebook, but the code doesn't seem to work anymore. Do you know of any updated versions? I am trying to export LlavaForConditionalGeneration.from_pretrained() with the Pixtral-12b config/weights.

Sorry, we didn't update the notebook in a while and I don't have at hand anything up-to-date. I'll try to update the existing notebook later next week, since I am off on a vacation



Thank you!

@EricJi150, [here](https://colab.research.google.com/drive/1mtcxKWHAR7D9LbjzRcTB_b5CURm2jw-z?usp=sharing) is the updated notebook. It works with `transformers==4.52.0` and the latest `optimum`

LMK if that works for you, I was going out of RAM in Colab so tested it locally only

Thanks! I gave it a try, but got an error with mismatched shapes in the `key_cache` for the language model. I am trying to export Mistral, except my hidden size is 5120 instead of the default 4096. It seems like the issue has to do with ONNX generating a dummy input with size 160 (5120/32) instead of 128 (4096/32) for the head dimension. Is there a workaround for this?

```
File "/app/.venv/lib/python3.12/site-packages/transformers/cache_utils.py", line 446, in update
    self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 160 but got size 128 for tensor number 1 in the list.
```

Ahh seems like an issue with sliding window, iirc Mistral had a sliding window length of 4096. Export of an LM is officially supported in `optimum` and [the documentation](https://huggingface.co/docs/optimum/main/en/exporters/onnx/overview) says Mistral is supported

I suggest to try to export an LM-only Mistral checkpoint and check if it works, and if not open an issue in `optimum`

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.