computing attention scores using relative attention bias

[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

WIP PR to workshop implementation: #1812

[pmabbo13]

Description

Having computed the relative attention bias term, this method computes the attention scores. The implementation is very similar to the nn.Functional._scaled_dot_product_attention, expect that we pass in position_bias as an input argument so that relative attention bias can be incorporated in the computation of the attention scores.

Since the input tensors to this function are 4-dimensional, we replace the torch.baddbmm and torch.bmm with torch.matmul. Since we are no longer using torch.baddbmm where attn_mask was passed as an input argument, we instead add the attn_mask directly to position_bias to ensure the mask is still applied before the softmax is taken to get the final attention scores.

[parmeet]

LGTM!