[Bloom] Remove unused position_ids, improve modeling code (lm head, alibi, attention multiplication)
Author: thomasw21
Pull request: yes

Notable changes:
 - Remove `attention_mask` sum trick, and instead use `torch.masked_fill`
 - Simplify the causal attention creation
 - Move back to `baddbmm` instead of `bmm`. It was unclear why the change was necessary.
 - ~Remove~ Deprecate `position_ids` as they don't make sense in BLOOM.
 - Introduce a fp32 cast for lm_head (and consequently the word embeddings in order to respect the sharing). The intuition is as follows
`
One of the thing we're wondering about if something we'd like to call "max-collapse". Given that 16bit allows to generate  at most 65536 different values, this means that with a vocabulary of 255k+ values are going to collapse, ie multiple values are going to be equal. So if that happens to the max value, this means that greedy decoding can change between fp32 and fp16/bf16.
 `
 - move back test to test generation on 16bit precision

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

First of all, thanks for the quick review!

Concerning the change in default behaviour, I'd advocate that this fixes a bug where we see poor generation specifically due to this issue. The argument is that at some level 255k+ vocab is just too big compared to the precision and we observe some sort of collapse when the top 2 values (2 being arbitrary) that are distinct in fp32 in fact collapse where you run them in fp16 or bf16, and so argmax takes the first value given that matches the max given the [documentation](https://pytorch.org/docs/stable/generated/torch.argmax.html). This issue with collapse actually breaks greedy generation really fast (we've seen cases where the first token generated don't match, and then it just gets a lot worse).

Also maybe we can rename `force_word_embeddings_in_fp32` to `force_lm_head_in_fp32`? The realy reason why I cast the word embeddings in fp32 is because `word_embeddings` and `lm_head` are tied, and I need `lm_head` to be fp32.

**Edit**: Actually after thinking about it, this might be a more generic issue where whenever you have too high a lm_head size you need to upcast before running the final linear layer. Typically models like mt5 might struggle due to this?

This PR tries to group way too many things together which is very bad practice as when we realize after merging it that everything is broken, we won't find the cause easily. Please break it down in at lest three parts:
- clean up code without any change
- removing/deprecating position Ids
- the float32 upcasting (which is where all the heated discussion is, so **really** should be its own PR)

I am veto-ing any merge of everything altogether as we've already had quite enough of "Oh the last BLOOM PR broke Xxx." 🙃 

Here are some comments so far:

- >Move back to baddbmm instead of bmm. It was unclear why the change was necessary.
    - Here: https://github.com/huggingface/transformers/pull/17866#discussion_r921825354
    - But if this PR solves the issue for FP16, OK for me to use `baddbmm`.
    - (**just wondering what are the difference, and if there is any reason you prefer to use `baddbmm`?**)
- There are 4 `force_lm_head_in_fp32` in the test file. Other than the one in `test_force_lm_head_in_fp32_is_close_to_fp16`, I don't know why we set it to `False`.
  - Is it to keep the same behavior as before (the whole model in FP16)?
  - But `prepare_config_and_inputs` has default `force_lm_head_in_fp32=True`, so most tests now use `True`. It is a bit confusing to me we keep them `False` in a few places.
- I agree with @sgugger that the default value for `force_lm_head_in_fp32` should be `False`.
  - Although `True` here is good for generation, this is kind special (casting parts of model weights to different dtype)
  - Also it's good to keep the previous behavior by default -> do not introduce surprise things to users


Also, I really appreciate your finding on "max-collapse" (especially being able to demonstrate it!), and glad that it improves the generation here.

But I personally would not expect FP16 generations will **always** match FP32 generations (even with `force_lm_head_in_fp32=True`), and we don't need to have tests that compare results across FP16/FP32. (I don't remember if we have a common test doing so though).

> Here: https://github.com/huggingface/transformers/pull/17866#discussion_r921825354
But if this PR solves the issue for FP16, OK for me to use baddbmm.
(just wondering what are the difference, and if there is any reason you prefer to use baddbmm?)

I think @younesbelkada and @NouamaneTazi changed the original behaviour, it was unclear what it actually fixed. The reason why I want to use `baddbmm` is because the training codebase used `baddbmm` and so there's no reason to use `bmm`.

> There are 4 force_lm_head_in_fp32 in the test file. Other than the one in test_force_lm_head_in_fp32_is_close_to_fp16, I don't know why we set it to False.
Is it to keep the same behavior as before (the whole model in FP16)?
But prepare_config_and_inputs has default force_lm_head_in_fp32=True, so most tests now use True. It is a bit confusing to me we keep them False in a few places.

Yeah so I initially thought that upcasting would have much better inference (at least in greedy style). turns out that's not true at least for 176b (it was true on the small models in test), so as @sgugger and @patrickvonplaten I'll try to figure out more if that feature is actually necessary at all.

Woops forgot to answer some question:
I agree that default should be `False` now :D

> But I personally would not expect FP16 generations will always match FP32 generations (even with force_lm_head_in_fp32=True), and we don't need to have tests that compare results across FP16/FP32. (I don't remember if we have a common test doing so though).

Well technically given checkpoints are in float16 of bfloat16, there should be little reason that generation don't match. I mean it's the promise of pretraining on those half precision: "use twice less compute/time to get more or less the same model". I would not be surprised that it doesn't match perfectly, but at the same time, now that they do, it's a great signal that the model is robust to numerical inacurracies. Consequently, I think the test matching fp16 (with fp32 lm_head) output with full fp32 output makes sense.

As https://github.com/huggingface/transformers/pull/18344#event-7125942979 has been merged, can you merge main into this branch?

Actually going to close this PR, any reason why you want this branch to still be alive? What should be missing if the `fp32` upcasting that I've done in another branch.

Good to be closed for me