Error exporting MPNet model to ONNX and load in ONNXRuntime despite supported

[zhangyilun]

I'm trying to export https://huggingface.co/sentence-transformers/all-mpnet-base-v2 to onnx and run inference on CPU. From https://huggingface.co/docs/optimum/exporters/onnx/overview, MPNet should be supported.

I remember I was able to make it work earlier but suddenly it stopped working and I want to get some help.

Environment:

onnx==1.15.0
onnxruntime==1.16.3
optimum==1.16.1
torch==2.1.2 (also with 2.0.1)

I tried on Mac M chip, ubuntu (within docker container).

The relevant code is:

from optimum.onnxruntime import ORTModelForFeatureExtraction
model = ORTModelForFeatureExtraction.from_pretrained("sentence-transformers/all-mpnet-base-v2", export=True)

Optimum cli also doesn't work:

optimum-cli export onnx --model sentence-transformers/all-mpnet-base-v2 models/

Error:

/Users/tzhang/miniconda3/envs/review-clf-server/lib/python3.10/site-packages/torch/onnx/utils.py:1636: 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 /Users/runner/work/pytorch/pytorch/pytorch/torch/csrc/jit/serialization/export.cpp:1413.)
  _C._check_onnx_proto(proto)

---------------------------------------------------------------------------
InvalidGraph                              Traceback (most recent call last)
Cell In[2], line 1
----> 1 model = ORTModelForFeatureExtraction.from_pretrained("sentence-transformers/all-mpnet-base-v2", export=True)

File ~/miniconda3/envs/review-clf-server/lib/python3.10/site-packages/optimum/onnxruntime/modeling_ort.py:626, in ORTModel.from_pretrained(cls, model_id, export, force_download, use_auth_token, cache_dir, subfolder, config, local_files_only, provider, session_options, provider_options, use_io_binding, **kwargs)
    577 @classmethod
    578 @add_start_docstrings(FROM_PRETRAINED_START_DOCSTRING)
    579 def from_pretrained(
   (...)
    593     **kwargs,
    594 ):
    595     """
    596     provider (`str`, defaults to `"CPUExecutionProvider"`):
    597         ONNX Runtime provider to use for loading the model. See https://onnxruntime.ai/docs/execution-providers/ for
   (...)
    624         `ORTModel`: The loaded ORTModel model.
    625     """
--> 626     return super().from_pretrained(
    627         model_id,
    628         export=export,
    629         force_download=force_download,
    630         use_auth_token=use_auth_token,
    631         cache_dir=cache_dir,
    632         subfolder=subfolder,
    633         config=config,
    634         local_files_only=local_files_only,
    635         provider=provider,
    636         session_options=session_options,
    637         provider_options=provider_options,
    638         use_io_binding=use_io_binding,
    639         **kwargs,
    640     )

File ~/miniconda3/envs/review-clf-server/lib/python3.10/site-packages/optimum/modeling_base.py:372, in OptimizedModel.from_pretrained(cls, model_id, export, force_download, use_auth_token, cache_dir, subfolder, config, local_files_only, trust_remote_code, revision, **kwargs)
    369     trust_remote_code = False
    371 from_pretrained_method = cls._from_transformers if export else cls._from_pretrained
--> 372 return from_pretrained_method(
    373     model_id=model_id,
    374     config=config,
    375     revision=revision,
    376     cache_dir=cache_dir,
    377     force_download=force_download,
    378     use_auth_token=use_auth_token,
    379     subfolder=subfolder,
    380     local_files_only=local_files_only,
    381     trust_remote_code=trust_remote_code,
    382     **kwargs,
    383 )

File ~/miniconda3/envs/review-clf-server/lib/python3.10/site-packages/optimum/onnxruntime/modeling_ort.py:549, in ORTModel._from_transformers(cls, model_id, config, use_auth_token, revision, force_download, cache_dir, subfolder, local_files_only, trust_remote_code, provider, session_options, provider_options, use_io_binding, task)
    546 save_dir = TemporaryDirectory()
    547 save_dir_path = Path(save_dir.name)
--> 549 main_export(
    550     model_name_or_path=model_id,
    551     output=save_dir_path,
    552     task=task,
    553     do_validation=False,
    554     no_post_process=True,
    555     subfolder=subfolder,
    556     revision=revision,
    557     cache_dir=cache_dir,
    558     use_auth_token=use_auth_token,
    559     local_files_only=local_files_only,
    560     force_download=force_download,
    561     trust_remote_code=trust_remote_code,
    562 )
    564 config.save_pretrained(save_dir_path)
    565 maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)

File ~/miniconda3/envs/review-clf-server/lib/python3.10/site-packages/optimum/exporters/onnx/__main__.py:551, in main_export(model_name_or_path, output, task, opset, device, fp16, optimize, monolith, no_post_process, framework, atol, cache_dir, trust_remote_code, pad_token_id, subfolder, revision, force_download, local_files_only, use_auth_token, for_ort, do_validation, model_kwargs, custom_onnx_configs, fn_get_submodels, use_subprocess, _variant, library_name, legacy, **kwargs_shapes)
    547         tokenizer_2.save_pretrained(output.joinpath("tokenizer_2"))
    549     model.save_config(output)
--> 551 _, onnx_outputs = export_models(
    552     models_and_onnx_configs=models_and_onnx_configs,
    553     opset=opset,
    554     output_dir=output,
    555     output_names=onnx_files_subpaths,
    556     input_shapes=input_shapes,
    557     device=device,
    558     dtype="fp16" if fp16 is True else None,
    559     model_kwargs=model_kwargs,
    560 )
    562 if optimize is not None:
    563     from ...onnxruntime import AutoOptimizationConfig, ORTOptimizer

File ~/miniconda3/envs/review-clf-server/lib/python3.10/site-packages/optimum/exporters/onnx/convert.py:753, in export_models(models_and_onnx_configs, output_dir, opset, output_names, device, input_shapes, disable_dynamic_axes_fix, dtype, model_kwargs)
    749     output_path = output_dir / output_name
    750     output_path.parent.mkdir(parents=True, exist_ok=True)
    752     outputs.append(
--> 753         export(
    754             model=submodel,
    755             config=sub_onnx_config,
    756             output=output_path,
    757             opset=opset,
    758             device=device,
    759             input_shapes=input_shapes,
    760             disable_dynamic_axes_fix=disable_dynamic_axes_fix,
    761             dtype=dtype,
    762             model_kwargs=model_kwargs,
    763         )
    764     )
    766 outputs = list(map(list, zip(*outputs)))
    767 return outputs

File ~/miniconda3/envs/review-clf-server/lib/python3.10/site-packages/optimum/exporters/onnx/convert.py:884, in export(model, config, output, opset, device, input_shapes, disable_dynamic_axes_fix, dtype, model_kwargs)
    879     raise RuntimeError(
    880         "You either provided a PyTorch model with only TensorFlow installed, or a TensorFlow model with only PyTorch installed."
    881     )
    883 if not disable_dynamic_axes_fix:
--> 884     config.fix_dynamic_axes(output, device=device, input_shapes=input_shapes, dtype=dtype)
    885 return export_output

File ~/miniconda3/envs/review-clf-server/lib/python3.10/site-packages/optimum/exporters/onnx/base.py:306, in OnnxConfig.fix_dynamic_axes(self, model_path, device, dtype, input_shapes)
    304 session_options = SessionOptions()
    305 session_options.graph_optimization_level = GraphOptimizationLevel.ORT_DISABLE_ALL  # no need to optimize here
--> 306 session = InferenceSession(model_path.as_posix(), providers=providers, sess_options=session_options)
    308 onnx_input_names = [inp.name for inp in session.get_inputs()]
    310 to_fix = []

File ~/miniconda3/envs/review-clf-server/lib/python3.10/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py:419, in InferenceSession.__init__(self, path_or_bytes, sess_options, providers, provider_options, **kwargs)
    416 disabled_optimizers = kwargs["disabled_optimizers"] if "disabled_optimizers" in kwargs else None
    418 try:
--> 419     self._create_inference_session(providers, provider_options, disabled_optimizers)
    420 except (ValueError, RuntimeError) as e:
    421     if self._enable_fallback:

File ~/miniconda3/envs/review-clf-server/lib/python3.10/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py:452, in InferenceSession._create_inference_session(self, providers, provider_options, disabled_optimizers)
    450 session_options = self._sess_options if self._sess_options else C.get_default_session_options()
    451 if self._model_path:
--> 452     sess = C.InferenceSession(session_options, self._model_path, True, self._read_config_from_model)
    453 else:
    454     sess = C.InferenceSession(session_options, self._model_bytes, False, self._read_config_from_model)

InvalidGraph: [ONNXRuntimeError] : 10 : INVALID_GRAPH : Load model from /var/folders/z7/y4b6xh_n0bd58hgf9hhk92940000gn/T/tmpdir8swq5/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.