Make T5 model torchscriptable

[pmabbo13]

Description

Make minor design changes that allow for the T5 model to be scripted

Process

A few design changes had to be made for the model to become scriptable:

1. relative_attention_bias is no longer instantiated when a T5 layer is being instantiated. This is because relative_attention_bias then needs to get passed down to T5MultiheadAttention as an input argument, where the embedding is actually used, but Torchscript does not support nn.Embedding input args. So instead, relative_attention_bias gets instantiated with T5MultiheadAttention. This completely avoids having to pass it in as an input argument.

2. There seemed to be an issue with all_outputs, all_sa_scores, and all_ca_scores (in T5Stack.forward) being tuples (probably because after every layer we "append" tensors to this object, and tuples aren't the best data structures for what we are effectively using as mutable objects). Changing their types to List cleared the torchscript errors.

3. After correcting the above and a few other minor errors, we were then getting a vague error message of torchscript RuntimeError: Unsupported value kind: Tensor. The error message did not provide any pointers to where in the code this issue was arising. It is unclear why exactly this worked, but breaking out T5Layer and T5Stack to be T5EncoderLayer, T5DecoderLayer, T5Encoder, T5Decoder seemed to resolve this issue.

Testing

Updated pre-trained weights saved in S3/Manifold and modified integrations tests to test scripted version of models.

pytest test/prototype/integration_tests/test_models.py

[parmeet]

Overall LGTM! Thanks @pmabbo13 for making the model torchscriptable.

[Nayef211]

LGTM! Great work on this!