Bloom Optimize operations
Moved the original PR: #17759 here to check if the tests pass

_The documentation is not available anymore as the PR was closed or merged._<|||||>I won't merge this now since I saw that it broke some slow tests, will investigate that!<|||||>With the two proposed changes, all tests are now passing @younesbelkada :)<|||||>Thanks a lot @NouamaneTazi !! Amazing job 🔥 <|||||>Let's merge this together as some improvements to make the inference faster
- [x] create attn mask only once
- [x] broadcast alibi only once instead of each time on the attention layer
- [x] Remove the contiguous calls and test the model
- [ ] Refactor the reshaping (check how it is done in BLOOM Flax) in the attention layer<|||||>Before merging, let's fix the code quality tests...<|||||>All tests should be passing now except for `BloomModelTest::test_simple_generation` 
It seems the issue with this one comes from the fact the we now use `torch.bmm` instead of `torch.baddbmm` in [this line](https://github.com/younesbelkada/transformers/blob/773d8e780fea41b8a8f77bf2bccbfbeacc91d50d/src/transformers/models/bloom/modeling_bloom.py#L307)

And I don't undestand what's happening here: (this only gives different outputs for fp16)
```python
b = torch.baddbmm(
    torch.zeros_like(sliced_alibi, dtype=torch.float16),
    query_layer.transpose(1, 0),
    key_layer.transpose(1, 0).transpose(1, 2),
    beta=1.0,
    alpha=1.0,
)
c = torch.baddbmm(
    sliced_alibi,
    query_layer.transpose(1, 0),
    key_layer.transpose(1, 0).transpose(1, 2),
    beta=1.0,
    alpha=1.0,
) - sliced_alibi
print(b==c)
```

gives:
```
tensor([[[ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True, False, False, False],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True, False, False,  True, False, False],
         [ True,  True, False, False,  True, False, False],
         [ True,  True,  True, False, False,  True,  True],
         [ True, False, False, False,  True, False,  True]],

        [[ True, False,  True,  True,  True,  True, False],
         [ True,  True,  True,  True,  True, False, False],
         [ True, False,  True, False, False,  True,  True],
         [ True,  True, False,  True, False,  True, False],
         [ True,  True,  True,  True,  True, False,  True],
         [ True,  True, False,  True,  True,  True, False],
         [ True, False,  True,  True,  True,  True,  True]],

        [[ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True,  True, False, False,  True],
         [ True,  True,  True, False,  True, False,  True],
         [ True, False,  True,  True, False, False,  True],
         [ True, False,  True, False,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True, False],
         [ True, False,  True, False, False,  True,  True]],

        [[ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True, False,  True,  True, False],
         [ True,  True,  True,  True, False,  True, False],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True,  True,  True, False]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True, False,  True, False,  True],
         [ True,  True,  True, False, False,  True,  True],
         [ True, False,  True,  True,  True, False,  True],
         [ True,  True, False,  True, False, False,  True],
         [ True,  True,  True,  True,  True, False,  True],
         [ True,  True, False,  True,  True,  True, False]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True, False,  True, False, False,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True, False, False,  True],
         [ True,  True, False, False, False,  True,  True],
         [ True, False,  True, False, False, False,  True],
         [ True,  True,  True, False, False,  True,  True]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True,  True,  True, False, False,  True,  True],
         [ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True, False,  True,  True, False,  True,  True],
         [ True, False,  True,  True,  True,  True,  True],
         [ True,  True, False,  True,  True,  True,  True],
         [ True, False, False,  True, False,  True, False]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True, False,  True,  True,  True, False,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True, False, False,  True,  True,  True,  True],
         [ True, False, False,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True, False,  True],
         [ True,  True,  True,  True, False, False,  True],
         [ True, False,  True,  True,  True,  True,  True]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True, False,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True, False,  True]],

        [[ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True, False,  True, False,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True, False,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True, False, False,  True],
         [ True,  True,  True,  True,  True, False,  True],
         [ True,  True, False, False, False,  True, False]],

        [[ True,  True,  True, False, False, False, False],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True, False, False,  True,  True,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True, False,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True, False]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True, False,  True,  True, False,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True, False,  True,  True, False,  True,  True],
         [ True,  True, False,  True, False, False,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True,  True, False, False]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True, False,  True],
         [ True,  True,  True,  True,  True,  True,  True]]], device='cuda:0')
```

And btw this is what we get when print `old_matmul_result == new_matmul_result`
```
tensor([[[ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True, False,  True,  True,  True,  True],
         [ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True,  True,  True, False,  True,  True,  True],
         [ True, False,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True, False,  True,  True]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True, False,  True,  True,  True,  True, False],
         [ True,  True,  True,  True,  True, False,  True],
         [ True,  True, False,  True,  True,  True, False],
         [ True,  True,  True,  True, False,  True,  True],
         [ True, False, False,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True,  True, False,  True,  True,  True, False],
         [ True,  True, False, False,  True,  True, False],
         [ True,  True,  True,  True,  True,  True, False],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True, False,  True,  True,  True,  True],
         [ True, False,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True, False,  True,  True, False, False,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True, False,  True, False,  True,  True,  True],
         [ True,  True,  True, False,  True,  True, False]],

        [[ True,  True,  True, False,  True, False,  True],
         [ True,  True,  True,  True, False,  True, False],
         [ True, False,  True,  True,  True, False, False],
         [ True,  True,  True,  True,  True,  True, False],
         [ True,  True,  True,  True,  True, False, False],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True, False, False, False,  True]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True, False,  True,  True, False],
         [ True,  True, False,  True, False, False,  True],
         [ True,  True,  True,  True,  True, False,  True],
         [ True, False,  True, False, False,  True,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True, False,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True, False],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True,  True,  True,  True,  True,  True, False],
         [ True,  True, False,  True, False,  True,  True],
         [ True,  True,  True,  True,  True,  True, False],
         [ True,  True, False,  True,  True, False,  True],
         [ True, False,  True,  True,  True, False, False],
         [ True, False,  True,  True, False,  True, False],
         [ True, False,  True,  True, False,  True,  True]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True,  True, False,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True, False],
         [ True, False,  True,  True, False,  True, False],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True, False,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True, False,  True, False]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True, False,  True,  True],
         [ True,  True,  True,  True,  True, False,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True, False,  True,  True, False,  True, False],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True, False, False,  True,  True],
         [ True, False,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True, False,  True,  True, False,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True,  True,  True,  True,  True,  True,  True],
         [ True, False,  True,  True,  True, False,  True],
         [ True, False,  True,  True,  True, False,  True],
         [ True,  True,  True, False,  True,  True, False],
         [ True,  True,  True,  True,  True,  True, False],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True, False,  True, False,  True, False,  True]],

        [[ True,  True, False,  True,  True, False, False],
         [ True, False,  True, False,  True,  True, False],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True, False,  True, False, False,  True,  True],
         [ True,  True,  True,  True, False,  True, False],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True]],

        [[ True,  True,  True,  True,  True, False,  True],
         [ True,  True,  True, False,  True,  True, False],
         [ True,  True,  True,  True,  True, False,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True,  True,  True],
         [ True,  True,  True, False, False, False,  True],
         [ True,  True,  True,  True,  True, False,  True]]], device='cuda:0')
```<|||||>all tests are passing now ! Is it ok if we merge this @stas00 (since you are working on DS inference just to check if this PR does not conflict anything with you work) ?<|||||>I didn't have a chance to read this PR, but let me at least run a quick test with it.

**update:** it looks fine for the 350b model - I'm waiting for the 176 to download and will test with it as well.

if in a rush please go ahead and merge and if anything emerges we can fix it after.<|||||>There were lots of changes since you got approval, so please wait for a re-review of @patrickvonplaten and me.<|||||>Thanks a lot @sgugger @patrickvonplaten 
I think that we will do the tests in fp32 instead, we just need to keep in mind that doing batched generation can be flaky for small models (<=350m) as we have identified it with @NouamaneTazi . We will put a comment on the tests explaining what we have found and I think that we should be good to go!<|||||>All tests are passing now (tested on A100) 🎉
<|||||>Now tests pass on both A100 and Titan RTX 🎉 (because we used `fp32`)
(Note that the test `BloomModelTest::test_batch_generation_padd` is still failing on Titan RTX in `fp16` whether for this PR or the `main` branch, because of the issue mentioned above)
