Update T5 Onnx Export and Optimization

[tianleiwu]

Description

Previously, the encoder onnx model adds extra initialization for decoder to generate kv cache from prompt. It is not necessary. Here we redesign onnx export for T5 model to output two separate models for encode and decoder.

Move Linear that generates cross features based on encoder_hidden_states to encoder onnx model. In this way, the encoder does not need output encoder_hidden_states, and only need output the features for cross attention used in decoder.

Major changes:
-[x] update t5 onnx export script
-[x] update convert_generation script
-[x] update beam search to support changes of inputs and outputs (detail can be found below).
-[x] add a tiny t5 model, and enable the generation test for T5 in Linux CI pipelines.

Example change in inputs and outputs for one layer model:
Encoder Inputs:

• encoder_input_ids: int32 (B, encode_sequence_length)
• encoder_attention_mask: int32 (B, encode_sequence_length)
• decoder_input_ids: int32 (B, 1)

Encoder Outputs:

• logits: (B, 1, vocab_size)
• encoder_hidden_states: (B, encode_sequence_length, encoder_hidden_size)
• present_key_self_0: (B, num_heads, 1, head_size)
• present_value_self_0: (B, num_heads, 1, head_size)
• present_key_cross_0: (B, num_heads, encode_sequence_length, head_size)
• present_value_cross_0: (B, num_heads, encode_sequence_length, head_size)

Decoder Inputs:

• input_ids: int32 (B, 1)
• encoder_input_ids: int32 (B, encode_sequence_length) (optional for old format; removed in new format)
• encoder_attention_mask: int32 (B, encode_sequence_length)
• encoder_hidden_states: (B, encode_sequence_length, encoder_hidden_size) (optional for old format; removed in new format)
• past_key_self_0: (B, num_heads, past_decode_sequence_length, head_size)
• past_value_self_0: (B, num_heads, past_decode_sequence_length, head_size)
• past_key_cross_0: (B, num_heads, encode_sequence_length, head_size)
• past_value_cross_0: (B, num_heads, encode_sequence_length, head_size)

Decoder Outputs:

• logits: (B, 1, vocab_size)
• present_key_self_0: (B, num_heads, past_decode_sequence_length + 1, head_size)
• present_value_self_0: (B, num_heads, past_decode_sequence_length + 1, head_size)

Known issues:

• Some postprocessing (like converting to use decoder masked MHA, past and present buffer sharing) is not done. Could be a future work item to integrate with onnxruntime-genai.

Motivation and Context

Make the encoder onnx model simpler and more efficient in inference (no need to output encoder_hidden_states).

To fix the problem, we need to add an explicit return statement at the end of the test_gpt_model function to ensure that it always returns a value explicitly. This will make it clear to other developers that the function can return None and that this is intentional.

• Add an explicit return None statement at the end of the test_gpt_model function.
• This change should be made in the file onnxruntime/python/tools/transformers/convert_generation.py.

[github-actions]

You can commit the suggested changes from lintrunner.

[kunal-vaishnavi]

Adding a note here to check that the changes to the encoder and decoder subgraphs for T5 do not affect Whisper. There are CI tests here that were added to ensure backwards compatibility with the WhisperBeamSearch op.

[kunal-vaishnavi]

Can we add T5 to the Big Models CI Pipeline?

[tianleiwu]

Can we add T5 to the Big Models CI Pipeline?

It is not needed. The T5 tests in test_generation.py is enabled in CI pipeline so there is end to end tests there.