disable query_length diff on graph model
Author: jiqing-feng
Pull request: yes

Hi @fxmarty 

In generation tasks, the model will not use `AttentionMaskConverter._unmask_unattended` on the 1st token because no `past_key_values`, but will use it from the 2nd token. It will cause a different implementation while tracing, so we need to disable check `query_length` when using `jit.trace`.

cc @younesbelkada @ArthurZucker 

Comments:
Hi, I don't think this change should be done. jit traced model still uses SDPA that may dispatch to memory-efficient backend, which requires to have no fully masked rows (see https://github.com/pytorch/pytorch/issues/110213), which is what `_unmask_unattended` avoids. I'd recommend to trace models with `query_length > 1`.

@jiqing-feng Could you clarify the motivation and what is not working as intended?

> Hi, I don't think this change should be done. jit traced model still uses SDPA that may dispatch to memory-efficient backend, which requires to have no fully masked rows (see [pytorch/pytorch#110213](https://github.com/pytorch/pytorch/issues/110213)), which is what `_unmask_unattended` avoids. I'd recommend to trace models with `query_length > 1`.
> 
> @jiqing-feng Could you clarify the motivation and what is not working as intended?

Sorry for that I didn't make it clear. Yes, we trace the model with `query_length > 1`, but when we use the trace model to do generation tasks, the inputs do not contain `past_key_values` when inference the 1st token. We used a tricky way that input a `past_key_values` with 0 length, see [here](https://github.com/huggingface/optimum-intel/blob/main/optimum/intel/generation/modeling.py#L311-L313). In this way, we could avoid input `None` to the traced model (None input will cause traced model error), but it does not satisfy `query_length > 0` so the traced model forward failed in the `_unmask_unattended`. 

In conclusion, [TSModelForCausalLM](https://github.com/huggingface/optimum-intel/blob/main/optimum/intel/generation/modeling.py#L331) in `optimum-intel` will fail if I don't make this change. 
 
Do you have any good ideas to avoid it? Thx!

Please ask me if there is anything I didn't make clear.


@jiqing-feng `query_length` is taken from `input_ids` (https://github.com/huggingface/transformers/blob/ee2482b6f8eca320f1eff112673adae215dbc927/src/transformers/models/llama/modeling_llama.py#L1037), not from the sequence length in the KV cache so I am not sure to understand, how is PKV impacting the controlflow `if query_length > 1:`? Feel free to share a repro of the unexpected behavior/bug if that can help.

Though I agree that jit.trace / dynamo / symbolic_trace are likely not able to properly trace https://github.com/huggingface/transformers/blob/6c78bbcb8320d316434262ef003251ca997db0d1/src/transformers/modeling_attn_mask_utils.py#L244-L245

Hi @fxmarty Thanks for your advice, I get your point now. The problem should be fixed in `optimum-intel` instead of transformers because it is related to our traced model inputs.

@jiqing-feng I actually apply the same fix as you in https://github.com/huggingface/transformers/pull/28447 as `_unmask_unattended` can not be properly traced with torch.jit.trace or symbolic_trace. As this method was only implemented as a workaround to a limitation in pytorch with a specific backend on a specific hardware, I think it is fine.

We may want to have a tracable `_unmask_unattended`.