Add parallel attention layers and Multi-Query Attention (MQA) from PaLM to the fast path for transformers
Labels: oncall: transformer/mha

### 🚀 The feature, motivation and pitch

Parallel attention layers and MQA are introduced in PaLM [1]. The “standard” encoder layer, as currently implemented in PyTorch, follows the following structure:

Serial attention (with default `norm_first=False`):

`z = LN(x + Attention(x)); y = LN(z + MLP(z))`

A parallel attention layer from PaLM, on the other hand, is implemented as follows:

Parallel attention (with default `norm_first=False`):
`y = LN(x + Attention(x) + MLP(x))`

Parallel attention (with `norm_first=True`):

`y = x + Attention(LN(x)) + MLP(LN(x))`

As for MQA, the description from [1] is pretty concise and explains the advantage as well: “The standard Transformer formulation uses k attention heads, where the input vector for each timestep is linearly projected into “query”, “key”, and “value” tensors of shape [k, h], where h is the attention head size. Here, the key/value projections are shared for each head, i.e. “key” and “value” are projected to [1, h], but “query” is still projected to shape [k, h]. We have found that this has a neutral effect on model quality and training speed (Shazeer, 2019), but results in a significant cost savings at autoregressive decoding time. This is because standard multi-headed attention has low efficiency on accelerator hardware during auto-regressive decoding, because the key/value tensors are not shared between examples, and only a single token is decoded at a time.”

In our own experiments, we have seen a 23% speedup in training for T5-large when using parallel attention, with the same convergence, due to the ability to fuse the attention and feed-forward projection layers into a single matmul operation. Furthermore, adding MQA on top shaves off an extra 10% training time. See the image below for our training loss over time.

![loss_curve](https://user-images.githubusercontent.com/919977/220369662-0fb39569-bfa2-411d-b03c-bbac8f6737da.png)

Of course, this is not free, and we have found that for larger models (3B and 11B parameters), memory can peak high enough that it can OOM an A100 GPU, even when running with FSDP.

As we build larger models and optimize them for inference, we would like to take advantage of the fast path for Transformers already available for the serial implementation, given the already observed advantage in training and the claimed speed-up on inference from using MQA, as well as explore ways to make this approach more memory efficient in PyTorch.

We envision two ways in which this can happen:

1. Through new parameters in both nn.TransformerEncoderLayer and nn.TransformerDecoderLayer (e.g. booleans `use_parallel_attention` and `use_mqa`), that would route the code through the parallel path and/or MQA.

1. New layers nn.TransformerParallelEncoderLayer and nn.TransformerParallelDecoderLayer

Both paths would eventually call the native `transformer_encoder_layer_forward` (https://github.com/pytorch/pytorch/blob/0dceaf07cd1236859953b6f85a61dc4411d10f87/aten/src/ATen/native/transformers/transformer.cpp#L65) function (or a new parallel version of it that ensures kernel fusing) if all conditions to hit the fast path are met.

cc: @daviswer @supriyogit @raghukiran1224 @mudhakar @mayank31398 @cpuhrsch @HamidShojanazeri

References:

[1] [PaLM: Scaling Language Models with Pathways](https://arxiv.org/pdf/2204.02311.pdf)

### Alternatives

_No response_

### Additional context

_No response_

cc @jbschlosser @bhosmer @cpuhrsch @erichan1