Repository: pytorch
Issue: [FYI] MultiheadAttention / Transformer
Labels: proposal accepted,module: nn,triaged

### This issue is created to track the progress to refine `nn.MultiheadAttention` and `nn.Transformer`. 

Since the release of both modules in PyTorch v1.2.0, we have received a lot of feedback from users, including feature requests, bug fixes and doc improvement. Here is a quick summary of the work items. Feel free to comment if you have any suggestions.

#### NOTE: Help wanted: Please create an issue and wait on feedback before working on them. This way you won't create a PR we can't accept for some fundamental reasons.

#### Working items

- [ ] Decompose `multi_head_attention_forward` functional in `torch.nn.functional`. For example, an individual operator for the standard scaled dot product function, in/out projection operator, self-attention module https://github.com/pytorch/text/pull/720
- [ ] Add quantization support. #32764 Currently, it is blocked by the indexing function [here](https://github.com/pytorch/pytorch/blob/ef5637f85e43dff55b03276176ba2949bd83a085/torch/nn/functional.py#L3272-L3350)
- [ ] The `multi_head_attention_forward` functional is too long to understand. It's also too difficult to modify some parts in the function. The function maintains too many branches. https://github.com/pytorch/text/pull/720
- [ ] Fix the dimension convention for masks in Transformer (Propose to use (N, S, E), instead of (S, N, E). Requested by https://github.com/pytorch/pytorch/issues/25132,  https://github.com/pytorch/pytorch/issues/25100. PR https://github.com/pytorch/pytorch/pull/37597

#### Resolved issues

- [x] Support ByteTensor for `key_padding_mask` and `attn_mask` in `nn.MultiheadAttention`. Requested by Fairseq 
- [x] 3D attention mask in `nn.MultiheadAttention`. Requested by https://github.com/pytorch/pytorch/pull/25359, https://github.com/pytorch/pytorch/issues/29520, https://github.com/pytorch/pytorch/issues/30678, https://github.com/pytorch/pytorch/issues/24461. Merged in https://github.com/pytorch/pytorch/pull/31996
- [x] Add TorchScript support for `nn.MultiheadAttention` and `nn.Transformer`. In the master branch now. Will be released in v1.5.0
- [x] Add gelu activation func to Transformer for BERT model. Merged in https://github.com/pytorch/pytorch/pull/24181
- [x] MultiheadAttention's compatibility with DDP. Merged in https://github.com/pytorch/pytorch/issues/26698
- [x] `float` and `double` precision in MultiheadAttention module https://github.com/pytorch/pytorch/issues/34060, https://github.com/pytorch/pytorch/issues/33841

#### Open issues

- [ ] Flexibility to change the norm func in TransformerEncoderLayer and TransformerDecoderLayer https://github.com/pytorch/pytorch/issues/26342. The order of layer normalization in Transformer https://github.com/pytorch/pytorch/issues/21876, https://github.com/pytorch/pytorch/issues/24930. Users has the flexibility to pass custom encoder/decoder layers to the transformer modules. Update the doc to explain it https://github.com/pytorch/pytorch/pull/32373
- [ ] Fix the constraint that embed_dim should be equal to the dimension of key and value. Explore adding new arguments in `multi_head_attention_forward` functional - requested by Fairseq
- [ ] Move the func `generate_square_subsequent_mask ` out of the transformer module https://github.com/pytorch/pytorch/pull/25538
- [ ] Doc update https://github.com/pytorch/pytorch/issues/28719, https://github.com/pytorch/pytorch/issues/28657
- [ ] Additional attention layers (e.g. AugmentedConv, StandAloneAttention) - need research https://github.com/pytorch/pytorch/issues/32529
- [ ] Add embedding layer and positional encoding layer in Transformer https://github.com/pytorch/pytorch/issues/24826