Bloom Optimize operations
Author: younesbelkada
Pull request: yes

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


Comments:
_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)