Optimize causal mask using torch.where
Author: Akababa
Pull request: yes

Instead of multiplying by 1.0 float mask, use torch.where with a bool mask for increased performance.

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/2715?src=pr&el=h1) Report
> Merging [#2715](https://codecov.io/gh/huggingface/transformers/pull/2715?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/33ef7002e17fe42b276dc6d36c07a3c39b1f09ed?src=pr&el=desc) will **decrease** coverage by `0.01%`.
> The diff coverage is `100%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/2715/graphs/tree.svg?width=650&token=9qOlN6Hb1c&height=150&src=pr)](https://codecov.io/gh/huggingface/transformers/pull/2715?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #2715      +/-   ##
==========================================
- Coverage    77.8%   77.79%   -0.02%     
==========================================
  Files         100      100              
  Lines       17051    17052       +1     
==========================================
- Hits        13267    13266       -1     
- Misses       3784     3786       +2
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/2715?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/modeling\_gpt2.py](https://codecov.io/gh/huggingface/transformers/pull/2715/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19ncHQyLnB5) | `86.2% <100%> (+0.04%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/2715/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl91dGlscy5weQ==) | `88.15% <0%> (-0.18%)` | :arrow_down: |
| [src/transformers/modeling\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/2715/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ191dGlscy5weQ==) | `91.81% <0%> (-0.14%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/2715?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/2715?src=pr&el=footer). Last update [33ef700...a54a418](https://codecov.io/gh/huggingface/transformers/pull/2715?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Thanks – what's the PyTorch compatibility on this?

Not sure about that, where can I find more info on compatibility? I think it only relies on torch.where (introduced <= 1.0.0) and tensors of dtype torch.bool (introduced in 1.2.0). Does the None (newaxis) slicing introduce compatibility issues?

If we want to maintain compatibility with 1.0.0, I think we can use torch.uint8 instead of torch.bool.

Hi, I'd recommend to make the following changes:
1. Keep the original shapes of _bias_ buffer (because otherwise it breaks loading of already trained models) and make dtype equal to torch.uint8, so it'd be compatible with pytorch 1.0.0 as no torch.bool type available.
`self.register_buffer("bias", torch.tril(torch.ones((n_ctx, n_ctx), dtype=torch.uint8)).view(1, 1, n_ctx, n_ctx))`
2. Keep -1e4 constant in a buffer to reduce allocations on each _attn call and make it works automatically with different devices (CPU and CUDA):
`self.register_buffer("masked_bias", torch.tensor(-1e4))`
3. Keep `b = self.bias[:, :, ns - nd : ns, :ns]` line as _bias_ buffer have the original shape now
4. So the _where_ statement should look like `w = torch.where(b, w, self.masked_bias)`

As a result, overall speedup will be at 10-15% here as I measured, and the code should be 100% compatible with pytorch 1.0.0

Hi @Akababa, 

Thanks for the PR. I think this is a great change. I checked and it does lead to a significant speed-up :-) 

Could you fix the tests and I think then we can merge (see https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md)

1) You should fetch the master branch and rebase your branch on top of it.
2) Make sure to run `make style` in the root folder before pushing to pass the "check_code_quality" test.

Great work @Akababa - this looks good to me! 

@LysandreJik @thomwolf - could you check and merge? 

Checked slow hardcoded GPT2 tests and it looks all good!