Repository: pytorch
Issue: Implementation of Self Attention vs Encoder Decoder Attention Causing Early Evaluation for PyTorch/XLA
Labels: triaged

## 🐛 Bug

The following code snippet from multihead attention module is using tensor.equal method to compare query, key and value to determine if the attention module is being used as self-attention or encoder-decoder attention:
https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L4783

Use of .equal op to make this determination causes early evaluation when used with pytorch-xla.
In pytorch-xla, pytorch tensor are lazy. i.e executed only if a step barrier is encountered or the value of the tensor is fetched. If the tensor values are fetched too frequently the advantage of lazy tensor approach and compiler optimizations are offset by graph compilation costs and the training becomes slow.


## To Reproduce

Steps to reproduce the behavior:

1.  Use Unit Transformer training, glue_qnli task (facebookresearch/mmf):
on a TPU-VM instance run the following instance:
export USE_TORCH=ON
export XRT_TPU_CONFIG="localservice;0;localhost:51011"
export PT_XLA_DEBUG=1
export XRT_START_LOCAL_SERVER=1

python3 mmf_cli/run.py \
    config=./projects/unit/configs/all_8_datasets/shared_dec_without_task_embedding.yaml \
    dataset=glue_qnli \
    model=unit \
    training.batch_size=32 \
    training.device=xla \
    distributed.world_size=8 \
        training.log_interval=100 \
        training.max_updates=1500 \

You will notice in training log something like:
```
_qnli/loss_0/avg: 0.7835, train/total_loss: 0.7778, train/total_loss/avg: 0.7835, experiment: run, epoch: 1, num_updates: 31, iterations: 31, max_updates: 1500, lr: 0., ups: 0.33, time: 03s 633ms, time_since_start: 02m 13s 793ms, eta: 01h 30m 27s 829ms
```

Another observation

As per the original paper (Attention is All You Need):
"
In "encoder-decoder attention" layers, the queries come from the previous decoder layer,
and the memory keys and values come from the output of the encoder. This allows every
position in the decoder to attend over all positions in the input sequence. This mimics the
typical encoder-decoder attention mechanisms in sequence-to-sequence models such as
[38, 2, 9].
• The encoder contains self-attention layers. In a self-attention layer all of the keys, values
and queries come from the same place, in this case, the output of the previous layer in the
encoder. Each position in the encoder can attend to all positions in the previous layer of the
encoder.
"
It does not prescribe attention module itself to be designed differently.  The same multi-head attention module was indeed used for both the purposes, with key and value input coming from the encoder output instead of the preceding decoder self-attention layer. In contrast for self-attention all, key, value, query come from the preceding self-attention layer.

In addition using key = value = query dimensions is generally true for self-attention but using this to determine if it's self-attention vs encoded-decoder attention is not a generalizable assumption.

The advantage of detecting these conditions was to reduce the number of cuda kernels in the resulting calculations.

There are two alternatives proposed for this issue:

1. Work with no assumption on q, k, v. (This is followed in huggingface/transformers multihead attention implementation.
2. Introduced, is_self_attention, is_encoder_decoder_attention parameters to the multihead attention module (default set to False) for both. (This was followed in the fairseq multihead attention module before pytorch core attention module was integrated.


## Environment
Instance: TPU-VM
PyTorch version: 1.8.0a0+56b43f4
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.10.2

Python version: 3.6 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] pytorch-lightning==1.2.7
[pip3] torch==1.8.0a0+56b43f4
[pip3] torch-xla==1.8
[pip3] torchaudio==0.6.0a0+c3e673a
[pip3] torchmetrics==0.2.0
[pip3] torchtext==0.5.0
[pip3] torchvision==0.9.0

Thanks @apsdehal and @myleott for the discussion on this.


cc @zhangguanheng66