adding forward method for multihead attention

[pmabbo13]

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

[Nayef211]

This is still a relatively large PR and it's a bit hard to tell which lines were added by you. It might be worthwhile leaving a comment on the PR highlighting the lines that are added here and are different from the original implementation (like you did in #1812).

[pmabbo13]

lines 256-274 encapsulate the changes made from _torch.nn.functional.multi_head_attention_forward to include relative attention bias. position_bias is then also added as a return value in lines 393, 400.

[pmabbo13]

Description

The forward method for T5MultiheadAttention is a modified version of nn.Functional.multi_head_attention_forward meant to perform multihead attention with relative attention bias on the input query, key, and value tensors. The main modifications are as follows:

1. Parameters needed to compute relative attention bias are added as input arguments
2. Deprecated non-core functionalities such as add_zero_attn (adds a new batch of zeros to the key and value sequences at dim=1), and adding bias terms to the key and value projections
3. The nn.Functional.multi_head_attention_forward method reshapes the q, k, v to be 3D. Doing so appears to have led to discrepancies to the decoder outputs of the HF implementation when the input decoder sequence had a batch size larger than 1. The resolution was to shape these tensors to 4D (similarly to HF implementation), and that appeared to resolve the issue.

These changes are best visible via this commit.

[pmabbo13]

We were having issues from the original implementation where the q, k, v tensors were reshaped to (batch_size * num_heads, seq_len, head_dim). This led to outputs that differed from the HF output when the input sequence to the decoder had batch size larger 1. The resolution to this was to shape these into 4D tensors of shape (batch_size, num_heads, seq_len, head_dim) which is the same reshaping done in the HF implementation.

[parmeet]

Overall LGTM!

[Nayef211]

LGTM, thanks for adding a detailed description with the changes wrt to original implementation and resolving all PR comments!