Optimum CLI export to ONNX fails for MPNET based sentence-transformers

Khreas commented 8 months ago

System Info

optimum version: 1.16.2
sentence-transformers version: 2.3.1
onnx version: 1.15.0
transformers version: 4.37.0

Who can help?

@michaelbenayoun

Information

[X] The official example scripts
[ ] My own modified scripts

Tasks

[X] An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
[ ] My own task or dataset (give details below)

Reproduction (minimal, reproducible, runnable)

Hello!

I recently tried converting and optimizing a setfit model from Torch to ONNX using Optimum and encountered a shape issue during the conversion.

After some digging, I realized that converting the base model sentence-transformers/paraphrase-mpnet-base-v2 without any fine tuning led to the same error, which can be easily reproduced using the following CLI command:

optimum-cli export onnx \
--model sentence-transformers/paraphrase-mpnet-base-v2 \
--task feature-extraction \
--device cpu \
paraphrase-mpnet-base-v2-O2

Which outputs the following:

Framework not specified. Using pt to export to ONNX.
Using the export variant default. Available variants are:
    - default: The default ONNX variant.
Using framework PyTorch: 2.2.0+cu121
path/to/project-GKls2EMF-py3.10/lib/python3.10/site-packages/torch/onnx/utils.py:1703: UserWarning: The exported ONNX model failed ONNX shape inference. The model will not be executable by the ONNX Runtime. If this is unintended and you believe there is a bug, please report an issue at https://github.com/pytorch/pytorch/issues. Error reported by strict ONNX shape inference: [ShapeInferenceError] (op_type:Min, node name: /0/auto_model/encoder/Min): data_0 typestr: T, has unsupported type: tensor(int64) (Triggered internally at ../torch/csrc/jit/serialization/export.cpp:1484.)
  _C._check_onnx_proto(proto)
Traceback (most recent call last):
  File "path/to/project-GKls2EMF-py3.10/bin/optimum-cli", line 8, in <module>
    sys.exit(main())
  File "path/to/project-GKls2EMF-py3.10/lib/python3.10/site-packages/optimum/commands/optimum_cli.py", line 163, in main
    service.run()
  File "path/to/project-GKls2EMF-py3.10/lib/python3.10/site-packages/optimum/commands/export/onnx.py", line 246, in run
    main_export(
  File "path/to/project-GKls2EMF-py3.10/lib/python3.10/site-packages/optimum/exporters/onnx/__main__.py", line 551, in main_export
    _, onnx_outputs = export_models(
  File "path/to/project-GKls2EMF-py3.10/lib/python3.10/site-packages/optimum/exporters/onnx/convert.py", line 747, in export_models
    export(
  File "path/to/project-GKls2EMF-py3.10/lib/python3.10/site-packages/optimum/exporters/onnx/convert.py", line 878, in export
    config.fix_dynamic_axes(output, device=device, input_shapes=input_shapes, dtype=dtype)
  File "path/to/project-GKls2EMF-py3.10/lib/python3.10/site-packages/optimum/exporters/onnx/base.py", line 306, in fix_dynamic_axes
    session = InferenceSession(model_path.as_posix(), providers=providers, sess_options=session_options)
  File "path/to/project-GKls2EMF-py3.10/lib/python3.10/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 419, in __init__
    self._create_inference_session(providers, provider_options, disabled_optimizers)
  File "path/to/project-GKls2EMF-py3.10/lib/python3.10/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 472, in _create_inference_session
    sess = C.InferenceSession(session_options, self._model_path, True, self._read_config_from_model)
onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph: [ONNXRuntimeError] : 10 : INVALID_GRAPH : Load model from paraphrase-mpnet-base-v2-O2/model.onnx failed:This is an invalid model. Type Error: Type 'tensor(int64)' of input parameter (/0/auto_model/encoder/Add_1_output_0) of operator (Min) in node (/0/auto_model/encoder/Min) is invalid.

I've tried converting other MPNET-based sentence-transformers as well which seem to indicate that the issue lies in the handling of the MPNET architecture:

sentence-transformers/multi-qa-mpnet-base-cos-v1: fails with same issue
sentence-transformers/multi-qa-mpnet-base-dot-v1: fails with same issue
sentence-transformers/all-mpnet-base-v2: fails with same issue
sentence-transformers/paraphrase-multilingual-mpnet-base-v2: succeeds without any trouble
- This model uses a XLM-Roberta student model that learnt from the MPNET teacher, which could explain the conversion success

Thanks for your help. :)

Expected behavior

After converting the model from Torch to ONNX, the model should be executable by ONNX runtime without any shape inference issue

fxmarty commented 8 months ago

Hi @Khreas thank you for the report, could you try against with optimum 1.17.0 released today? I believe this is fixed.

Khreas commented 8 months ago

I can confirm that the newest release fixed it. Thanks a bunch!

huggingface / optimum