Fix forward pass for merged model

[kunal-vaishnavi]

What does this PR do?

After exporting a merged model with this PR, the forward pass has to support both the prompt and token generation steps. When running the prompt step, input_ids should be of shape (batch_size, sequence_length) and when running the token generation step, input_ids should be of shape (batch_size, 1). However, currently input_ids only has shape (batch_size, 1).

This PR fixes the shape of input_ids during the forward pass by first checking the size of past_sequence_length in one of the past_key_values to automatically determine which step is occurring. Then, the shape of input_ids is modified accordingly.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[fxmarty]

@kunal-vaishnavi Thank you for the PR. Isn't past_key_values expected to be None in the first forward pass? And prepared later at use_cache_branch, past_key_values, known_output_shapes = self.prepare_past_key_values(input_ids, past_key_values, use_torch)?

cc @echarlaix

[fxmarty]

Or did you mean decoder_model_merged.onnx?

[echarlaix]

@kunal-vaishnavi Thank you for the PR. Isn't past_key_values expected to be None in the first forward pass? And prepared later at use_cache_branch, past_key_values, known_output_shapes = self.prepare_past_key_values(input_ids, past_key_values, use_torch)?

cc @echarlaix

Yes

[echarlaix]

Not sure when this could happen, could you share an example so that we can reproduce @kunal-vaishnavi ?

[kunal-vaishnavi]

Yes, this is with decoder_model_merged.onnx. When running the forward pass, each past_key and past_value in past_key_values has shape (batch_size, num_heads, past_seq_len, head_size). For the prompt step, past_seq_len = 0 so each past_key and past_value input shape is (batch_size, num_heads, 0, head_size).

Since self.use_cache is true for ORTModelForCausalLM models and the merged model always has past_key_values, both if statements are true. Then the input_ids currently get modified for both prompt and token generation.

[fxmarty]

@kunal-vaishnavi I think it is not true. While decoder_model_merged.onnx, this is only true at the session.run call, not at the forward call, that is typically called from transformers GenerationMixin.generate. Do you have a different use case?

@echarlaix loading decoder_model_merged.onnx following #1257 prints

The ONNX file decoder_model_merged.onnx is not a regular name used in optimum.onnxruntime that are ['model.onnx', 'model_quantized.onnx', 'model_optimized.onnx', 'decoder_with_past_model.onnx', 'decoder_with_past_model_quantized.onnx', 'decoder_with_past_model_optimized.onnx'], the ORTModelForCausalLM might not behave as expected.

is this expected?

@kunal-vaishnavi Apart from backward compatibility reasons, do you have limiting factors on your side that prevent to simply use the new version of the exported model (with an empty KV cache at the first iteration)?

[kunal-vaishnavi]

In the PR linked below, we want to benchmark exported models using Optimum. To have an equal comparison, we currently evaluate with model(...) and not model.generate(...). We pass both input_ids and past_key_values to the forward call in order for the merged ONNX model to receive both inputs. When calling the forward pass, the input_ids and past_key_values are passed with the right shapes. However, the input_ids shape is changed in Optimum before they are passed with the wrong shape to the ONNX model.

[fxmarty]

Thank you, makes sense indeed.