[Phi] Extend implementation to use GQA/MQA.
Author: gugarosa
Pull request: yes

# What does this PR do?

As we discussed on the repositories and the e-mail thread, these are minor changes that we would like to integrate into HF. 

One thing that we need to discuss is how to leverage the current batch of models (using `transformers>=4.36.0`), since the proposed change will make a shape difference in `qkv` weights and biases.

We could change the conversion script from Phi (transformers=4.36.0) to reflect this new implementation, while we compromise in converting our current repositories weights to this new format and use the HF-based code in the next deploys. With that, we could avoid having two conversions, i.e., phi-msft -> phi and phi (4.36.0) -> new_phi.
 
Please let me know your thoughts!

## Changes

- Adds support for using GQA/MQA with Phi-based models. This is a combined implementation between the old `PhiAttention` and `LlamaAttention`.
- Fixes documentation official Phi-based models paths.
- Adds support for dynamically pad the vocab_size to a multiple of 64 (better use of Ampere/Hopper-based GPUs).

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

## Before submitting
- [X] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [X] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [X] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [X] Did you write any new necessary tests?


## Who can review?

@susnato @LysandreJik @ArthurZucker @philschmid @osanseviero 

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
TODO: address results differences when evaluating models.

> The idea for us is that if the model phi2 fits in phi (modeling_phi) with these changes, then a single conversion is great but we don't mind updating it.

I believe yes `phi2` fits into `phi1/1.5` modeling file... I guess this is to add more functionalities (such as GQA, MQA) to both ph1 and phi2 models.

Regarding the [issues](https://github.com/huggingface/transformers/issues/28049#issuecomment-1863021572) with the current conversion script, should I create a separate PR to address that or should it be included in this PR? @ArthurZucker 

Thanks @ArthurZucker!

It is exactly what @susnato mentioned, either phi-1, phi-1.5 and phi-2 will fit in the updated files. We are just extending its scope to leverage any potential use of GQA/MQA in newer models.

Regarding the conversion, we can do either way that works for you. For example:

1. Merge susnato changes -> Merge PR and re-upload the new weights / conversion script.
2. Merge PR / conversion script directly and convert phi-msft (current code in repos) to new structure.

Same PR works. 
Understood if the changes are just GQA works great ! (f we have to add new layers or do some if else in the forward, we can't keep it in the same model as per the transformers philosophy but that should be alright 🤗 ) 

Yeah we can:
1. merge the code needed in transformers (using the correct revision for examples) 
2. convert the weights and open PRs to the hub at the same time. With both the checkpoints and updated modeling code
🤗 

> TODO: address results differences when evaluating models.

Hi @gugarosa, could you please give more context in what you mean by that? Do you mean the difference between the logits between the model on the library and the model on the Hub?

Hi @ArthurZucker, I have completed the conversion script and added the  `phi2` test in this [PR](https://github.com/huggingface/transformers/pull/28211) 

Thanks for the review @ArthurZucker! I just got back from vacation and will proceed with the updates.

@susnato Exactly! I still need to re-evaluate some of our tasks due to the logits. There were slight differences on the results of the tasks, but I will double check everything now that I am back.

Just ironed out the differences, it was basically a silly mistake on my end 😄.

Any advice on fixing the `check_repository_consistency` test error?

Yes! Copied from missused somewhere let me review! 

Please disregard the "force-push" messages, I was squashing some of the commits.

No worries and feel free to ping me again whenever! 🤗 

@ArthurZucker could you please do a final review? It should be all good now.

I have updated the `microsoft/phi-1` repository to reflect this implementation as well. As soon as everything is merged, we will rollout the changes to 1.5 and 2.

Thanks for all the attention!

If you can adresse the final comments + the merge we'll be able to merge! 🥳 

> If you can adresse the final comments + the merge we'll be able to merge! 🥳

Thanks @ArthurZucker! I have addressed the latest comments and merge. Everything should be good now! 🎉 

@ArthurZucker 

Is there any more time for a small addition? There is a "bug" that we can solve with upcasting queries and keys to fp32 in `PhiAttention`.

Code to reproduce (`main` branch):
```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("susnato/phi-2", torch_dtype="auto").to("cuda") # or .half()
tokenizer = AutoTokenizer.from_pretrained("susnato/phi-2")
input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)

logits = model(input_ids.to("cuda")).logits
print(logits)
```

Logits are always NaN due to attention overflow when using the `PhiAttention` class with FP16. We have observed this before, that's why we had this https://huggingface.co/microsoft/phi-2/blob/main/modeling_phi.py#L359. Funny thing is that it does not happen on Phi-1 and Phi-1.5 😆 

I am using @susnato model as reference, but it is extensible to any Phi-2 model. 

That's interesting, thanks for pointing this out @gugarosa!

Now I am curious why does this not happen for phi-1 and phi-1.5. 😅 

Were there any difference in the dtype between the phi-1 and phi-2 during training?

> That's interesting, thanks for pointing this out @gugarosa!
> 
> Now I am curious why does this not happen for phi-1 and phi-1.5. 😅
> 
> Were there any difference in the dtype between the phi-1 and phi-2 during training?

Unfortunately no, it was the same training code. The only difference is that phi-2 has been trained for longer and it is a bigger model. We were baffled when we saw it and so far we don't have a clue.

It could have been this very specific checkpoint and/or a combination of factors that caused the overflow, but still, no clear idea 😭 

Sure let me have a look. 
We also had similar issue with Llama I believed, so having a look

> Sure let me have a look. We also had similar issue with Llama I believed, so having a look

Thanks! The fix is on the last commit, only two more lines (and one changed) were able to solve it.

> 2 nits and we can merge! Thanks a lot for bearing with me @gugarosa

No problems! Just added the changes and thanks for all the help with this PR.

@ArthurZucker could we please merge this?

I am anxious to fix Phi-2 😄 

Was waiting for the CIs and fell asleep! Sorry @gugarosa !
Thanks @susnato and @gugarosa for this awesome join contribution! 🤗 
🚀 


@gugarosa @susnato @ArthurZucker 
guys it would be great to communicate on your main Phi-2 README that the checkpoints have changed.
You had 200K+ downloads and many people finetuned on the old fused QKV ckpt.

Hi @vince62s, I think there has been a misunderstanding here, the checkpoints are not changed, they are the same but in different order.

The checkpoints which are present at microsoft/phi-2 are not in the right order to be used with the library model so I reordered the weights and pushed them to susnato/phi-2.
But they are same values. If you compare the logits of those two different models( in fp32 ) you will find that they are same within a acceptable range (1e-3). 

Please let me know if this explanation helps or not. 🤗 

used to be:
```
  "weight_map": {
    "lm_head.linear.bias": "model-00002-of-00002.safetensors",
    "lm_head.linear.weight": "model-00002-of-00002.safetensors",
    "lm_head.ln.bias": "model-00002-of-00002.safetensors",
    "lm_head.ln.weight": "model-00002-of-00002.safetensors",
    "transformer.embd.wte.weight": "model-00001-of-00002.safetensors",
    "transformer.h.0.ln.bias": "model-00001-of-00002.safetensors",
    "transformer.h.0.ln.weight": "model-00001-of-00002.safetensors",
    "transformer.h.0.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
    "transformer.h.0.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
    "transformer.h.0.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
    "transformer.h.0.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
    "transformer.h.0.mlp.fc1.bias": "model-00001-of-00002.safetensors",
    "transformer.h.0.mlp.fc1.weight": "model-00001-of-00002.safetensors",
    "transformer.h.0.mlp.fc2.bias": "model-00001-of-00002.safetensors",
    "transformer.h.0.mlp.fc2.weight": "model-00001-of-00002.safetensors",
```
now is

```
"lm_head.bias": "model-00002-of-00002.safetensors",
--
  | "lm_head.weight": "model-00002-of-00002.safetensors",
  | "model.embed_tokens.weight": "model-00001-of-00002.safetensors",
  | "model.final_layernorm.bias": "model-00002-of-00002.safetensors",
  | "model.final_layernorm.weight": "model-00002-of-00002.safetensors",
  | "model.layers.0.input_layernorm.bias": "model-00001-of-00002.safetensors",
  | "model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
  | "model.layers.0.mlp.fc1.bias": "model-00001-of-00002.safetensors",
  | "model.layers.0.mlp.fc1.weight": "model-00001-of-00002.safetensors",
  | "model.layers.0.mlp.fc2.bias": "model-00001-of-00002.safetensors",
  | "model.layers.0.mlp.fc2.weight": "model-00001-of-00002.safetensors",
  | "model.layers.0.self_attn.dense.bias": "model-00001-of-00002.safetensors",
  | "model.layers.0.self_attn.dense.weight": "model-00001-of-00002.safetensors",
  | "model.layers.0.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
  | "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
  | "model.layers.0.self_attn.q_proj.bias": "model-00001-of-00002.safetensors",
  | "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
  | "model.layers.0.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
  | "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
```

Am I wrong?

I am not saying the weights have changed but the tensors name have changed and now QKV are unmerged (I presume)

and btw you can see people complaining https://huggingface.co/microsoft/phi-2/discussions/76

Yes @vince62s, there seems to be some recent updates, sorry I wasn't aware of this commit. 

People who finetuned can either convert the checkpoint or use truste_remote_code with a specific revision before the changes, which should be the best solution IMO

I'm fine with this, just thinking that would be great to post some info on the README page about the change.

Yep for sure! 

Sorry all for the mess :( 

> People who finetuned can either convert the checkpoint or use truste_remote_code with a specific revision before the changes, which should be the best solution IMO

Yes, but right now the conversion script won't work because the conversion script generates [a single matrix](https://github.com/huggingface/transformers/blob/64bdbd888c78dcef5aeaeabc842e12981c8aae7a/src/transformers/models/phi/convert_phi_weights_to_hf.py#L54) for queries, keys and values whereas it was changed to seperate q,k,v layers in this PR. 😕 



Let's open a PR for the conversion script to support both scheme @susnato would you be down for that? 

At worse we can add a loading state dict hook to make sure this is BC 

Sure @ArthurZucker ! I would love to do that! I will open a PR ASAP.



Thanks! Ping me for a review! I think it's the simplest way of dealing with this

It's motivated by the fact that the format of the model was changed, but not much appart from that. It's not super transformers like but will help us a lot IMO. WDYT @amyeroberts 

Yes, I think having a conversion script is a good idea. I'd even suggest having a separate conversion script that converts between the old and new format e.g. `convert_old_remote_phi.py`

I'd be less in favour of a hook on the loading of the state dict: it's doing something more "magic" and as you said isn't in-line with transformers in general.