Fix Llama's attention map handling for left padding which causes numerical instability and performance drops
Author: Randolph-zeng
Pull request: yes

Hi this PR is trying to address the performance drop and potential numerical instability caused by vanilla left padding in Llama.
Here is the explanation:
1. If we initialize the tokenizer with left padding and call model.generate without passing in corresponding attention_mask, the code will run, but for the instances who are left padded, its unpadded tokens will "see" the padded tokens. This will cause performance drop a lot ! At least in my case, my performance of llama2 in socialQA drops from 55% to around 20% if I use left padded batch inference instead of one by one generate.
2. If instead, I passed in the attention_map generated by the left_padding tokenizer to model.generate function, the model will throw an error when doing sampling because some values in the hidden states are inf or nan. This numerical instability suddenly appeared because train-test mismatch: **By examining the locations of these infs/nans, I found them only shows up in the position of those padded token and are caused by the attention_map.** 
3. The reason why attention map are causing the numerical instability is because the current way of generating attention mask did not considered the left padded situation and it will cause the left padded tokens to have a fully masked attention tensor ! While the model was never trained with any token that can not see any(including itself) token, the model thus generates anomaly values and creates nan/inf. 

So this PR is trying to fix two bugs I observed:
1. The attention_mask created for left_padded values will contain -inf value due to the operation "expanded_attn_mask + combined_attention_mask".  Consider the attention_map that looks like this ([[1, 1, 1, 1, 1], [0, 0, 0, 1, 1]]). The combined_attention_mask created by line 585 will look like this (under float16)
```
tensor([[[[     0., -65504., -65504., -65504., -65504.],
          [     0.,      0., -65504., -65504., -65504.],
          [     0.,      0.,      0., -65504., -65504.],
          [     0.,      0.,      0.,      0., -65504.],
          [     0.,      0.,      0.,      0.,      0.]]],


        [[[     0., -65504., -65504., -65504., -65504.],
          [     0.,      0., -65504., -65504., -65504.],
          [     0.,      0.,      0., -65504., -65504.],
          [     0.,      0.,      0.,      0., -65504.],
          [     0.,      0.,      0.,      0.,      0.]]]], device='cuda:0',
       dtype=torch.float16)
```
and the expanded_attn_mask created will look like this 
```
tensor([[[[     0.,      0.,      0.,      0.,      0.],
          [     0.,      0.,      0.,      0.,      0.],
          [     0.,      0.,      0.,      0.,      0.],
          [     0.,      0.,      0.,      0.,      0.],
          [     0.,      0.,      0.,      0.,      0.]]],


        [[[-65504., -65504., -65504.,      0.,      0.],
          [-65504., -65504., -65504.,      0.,      0.],
          [-65504., -65504., -65504.,      0.,      0.],
          [-65504., -65504., -65504.,      0.,      0.],
          [-65504., -65504., -65504.,      0.,      0.]]]], device='cuda:0',
       dtype=torch.float16)
```
And in line 598 these two variables are added together. I believe it will be now clear why left padding causes the attention_map itself contains -inf values and why some tokens has a fully masked attn tensor.
3.  My solution then is straightforward, I clamped the variables so it does not overflow, and I forces the left padded values to at least attend to itself. Though the hidden states of the left padded values will not be used by the unpadded tokens due to the attention map, making it cleaned of inf/nan will not break the generation process. 
4. I tested in my local cases and I did not observe any performance drop or nan errors during sampling. Though I am not sure if my patches will break any other use cases. 




Comments:
cc @ArthurZucker 

Hi Arthur, thanks for the comment. Yes this PR should affect the inference because otherwise it causes performance drops or error. I think it applies both to training issue you mentioned and inference. I will do the checks you mentioned above locally and push again today after some meetings. Will keep you posted ! 

Hi Arthur, I tried the make style command and committed the changes. However, for the make fix-copies command, it simply cut my changes out from file ` src/transformers/models/llama/modeling_llama.py` and paste those lines to `src/transformers/models/deprecated/open_llama/modeling_open_llama.py`. I am not sure why script `python utils/check_copies.py --fix_and_overwrite` do this, so I did not include that in the commit. Please let me know if there is any other concern. Thanks!

You can probably just remove the `# Copied from` line above this function in the open llama modeling file! 😉 


Yes that works! I removed the copy statement from both functions in llama and openllama so that this copy script won't be triggered. 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_25284). All of your documentation changes will be reflected on that endpoint.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi @ArthurZucker , it seems this issue still needs to be fixed, at least for inference. Running the batched MMLU evaluation script [here ](https://github.com/allenai/open-instruct/blob/503407ee89553b4e4c09f349052aeeac296eb91e/scripts/eval/mmlu.sh#L5-L13) on llama2 7B checkpoint results in very bad performance with the latest transformers version `4.34.0`, but @Randolph-zeng 's PR gives me a good performance, matching the numbers reported in the llama2 paper. 

Will this PR be merged soon?

Hey pinging @Rocketknight1 as well, I think this is of interest yes, we had a hard time reproducing the nan but it seams that it is mostly related to beam search. We'll investigate and keep you posted! 

@Randolph-zeng do you mind explaining how inf/nan attention values on pad tokens breaks sampling? The values of those tokens should never be read by "real" non-pad tokens anyway, so I don't understand why this would make a difference.