create T5MultiheadAttention module

[pmabbo13]

Description

Add T5 architecture to torchtext

Process

The T5 architecture is very similar to the architecture of a traditional transformer. The main differences are that rather than using positional embeddings, it computes a relative attention bias that encodes the relative position of a token within a sequence. This position bias is then passed into each layer and used to compute the attention scores. T5 also uses a simplified layer normalization (root mean square normalization) which occurs at the start of every attention and feed-forward block.

Incorporating relative attention bias requires under the hood changes to the MultiHeadAttention module. We can use HF's implementation for computing relative attention bias and modify the source code for torch.nn.MultiHeadAttention to incorporate relative attention bias. We can also create our own layer normalization, similarly to HF.

Given the above components, we can then define our own T5Layer, T5Stack, and T5Model.

• The T5Layer can be used either as an encoder layer or decoder layer based on an input boolean parameter. The only difference between the decoder layer versus the encoder layer is that the decoder layer also performs cross-attention with the encoder output.
• T5Stack can also be used as either an encoder or decoder based on an input boolean parameter. This dictates which type of layer composes the stack.
• T5Model can be used either as an encoder-only or encoder-decoder model based on an input boolean parameter. If it is an encoder-decoder model, a causal mask is generated for the decoder input tokens.

Testing

To test our implementation of the T5 model, we compared our outputs to the outputs of HuggingFace's T5 encoder-only and T5 encoder-decoder implementations. Testing was done in this notebook. We will update this PR once formal unit and integration tests have been added.

Stack

Stack from ghstack (oldest at bottom):

• add t5 model that can function as both encodery-only or encoder-decoder model #1829
• add t5 stack that can function as either the encoder or decoder of a t5 model #1828
• add t5 layer module that can be used for both encoder or decoder stack #1827
• add layer norm module for t5 model #1826
• adding forward method for multihead attention #1833
• computing attention scores using relative attention bias #1832
• computing relative attention bias #1831
• compute relative position buckets for relative attention bias #1830
• -> create T5MultiheadAttention module #1825

WIP PR where implementation details were discussed: #1812

[Nayef211]

Overall LGTM

[parmeet]

LGTM!

[Nayef211]

Let's rebase the PR stack on the latest main. A lot of the test failures you are seeing should have been resolved by PRs that have already been merged.