Issue when converting Whisper using --collect_cross_qk on CPU

[axelman03]

Describe the issue

I am currently using the nightly build of the ONNX runtime to convert Whisper to ONNX. I am specifically interested in getting the cross QK of the model, to be used eventually for timestamps. I am trying to convert the model to run on CPU. My issue is that, when I convert it, a runtime exception occurs:

An error occurred while trying to verify parity between PyTorch and ONNX Runtime: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running WhisperBeamSearch node. Name:'BeamSearch_zcode' Status Message: C:\a\_work\1\s\onnxruntime\contrib_ops/cpu/transformers/beam_search_impl_whisper.h:300 onnxruntime::contrib::transformers::BeamSearchWhisper<float>::Execute decoder_subgraph_.has_decoder_masked_attention_ was false. decoder subgraph: output_cross_qk could only work with has_decoder_masked_attention

Traceback (most recent call last):
  File "C:\git\onnx-rt\1.17.0-pre\onnxruntime\onnxruntime\python\tools\transformers\models\whisper\convert_to_onnx.py", line 481, in main
    max_diff = WhisperHelper.verify_onnx(args.model_name_or_path, ort_session, device)
  File "C:\git\onnx-rt\1.17.0-pre\onnxruntime\onnxruntime\python\tools\transformers\models\whisper\whisper_helper.py", line 338, in verify_onnx
    ort_outputs = ort_session.run(None, inputs)[0][0]
  File "C:\Users\alexander.bolejack\Anaconda3\envs\S2T\lib\site-packages\onnxruntime\capi\onnxruntime_inference_collection.py", line 220, in run
    return self._sess.run(output_names, input_feed, run_options)
onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running WhisperBeamSearch node. Name:'BeamSearch_zcode' Status Message: 
C:\a\_work\1\s\onnxruntime\contrib_ops/cpu/transformers/beam_search_impl_whisper.h:300 onnxruntime::contrib::transformers::BeamSearchWhisper<float>::Execute decoder_subgraph_.has_decoder_masked_attention_ was false. decoder subgraph: output_cross_qk could only work with has_decoder_masked_attention

When looking into it, it seems that DecoderMaskedMultiHeadAttention is only used when the --use_gpu flag is enabled, as is cross QK.

Is there any way I can build and run this model on the CPU?

To reproduce

Here is the exact command I ran. Use the latest version of ONNX Runtime.

python -m models.whisper.convert_to_onnx -m openai/whisper-base --output whisperbase-timestamps --use_external_data_format --precision int8 --quantize_embedding_layer --extra_decoding_ids --output_sequence_scores --overwrite --output_no_speech_probs --output_cross_qk --collect_cross_qk --use_whisper_beamsearch --provider cpu

Urgency

No response

Platform

Windows

OS Version

Windows 10

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

#18206

ONNX Runtime API

Other / Unknown

Architecture

X64

Execution Provider

Default CPU

Execution Provider Library Version

No response

[github-actions]

This issue has been automatically marked as stale due to inactivity and will be closed in 7 days if no further activity occurs. If further support is needed, please provide an update and/or more details.

[thiagocrepaldi]

@shubhambhokare1 @kunal-vaishnavi maybe this is something you are familiar with

[axelman03]

I was trying to look a little further into it. It seems that Cross QK support requires the model to be compiled with DecoderMaskedMultiHeadAttention, which is implemented only for CUDA.
I am not familiar with it enough to say if DecoderMaskedMultiHeadAttention needs to be used for cross QK, or if it can be brought to CPU/other execution providers though.

[kunal-vaishnavi]

It seems that Cross QK support requires the model to be compiled with DecoderMaskedMultiHeadAttention, which is implemented only for CUDA.

Cross QK support is added in this PR as part of the WhisperBeamSearch op and DecoderMaskedMultiHeadAttention op. The feature is currently supported on CUDA and will be supported on CPU in the future.

I am not familiar with it enough to say if DecoderMaskedMultiHeadAttention needs to be used for cross QK, or if it can be brought to CPU/other execution providers though.

DecoderMaskedMultiHeadAttention is used specifically for CUDA. Cross QK support can be added to other attention ops that run on CPU.

[github-actions]

This issue has been automatically marked as stale due to inactivity and will be closed in 30 days if no further activity occurs. If further support is needed, please provide an update and/or more details.