Megatron encoder-decoder refactor

[michalivne]

This PR refactor Megatron T5 to a modular base encoder-decoder class hierarchy, and a T5 class which customizes the base class. The goal of the refactor is to support multiple encoder-decoder models with minimal code redundancy.

Refactor High Level Description

Hierarchy tree is expressed below via list alignment (i.e., right alignment for child class).

• MegatronModule sub-classes:
  • MegatronTransformerEncoderModule/MegatronTransformerDecoderModule - various architectures for manipulation of hidden --> hidden
  • MegatronTransformerEncoderDecoderModule - contains above classes and generates required joint encoder/decoder mask given independent masks, and reparametrization (e.g., vae)
  • TokensEncoderDecoderModule - modality-specific (e.g., tokens-tokens) model which map between pre-processed data (e.g., token ids) and encoder/decoder continuous input (i.e., hidden); handles embedders (i.e., embedding + positional embedding); handles projection from hidden to data distribution (e.g., project from hidden dimensions to the number of tokens)
• NLPModel sub-classes:
  • MegatronBaseModel - common initialization of Megatron
    • MegatronLMEncoderDecoderModel - modality-specific (e.g., text-text) model which pre-process raw data (e.g., tokenizer pre-process text to token ids); generic modality-specific training/testing/inference/loss methods
      • MegatronT5Model - a specific model (i.e., T5) which prepares datasets (i.e., with corresponding data augmentation), and model-specific features (e.g., sentinels in tokenizer)

Notable Changes

1. Encoder/decoder mask is 2D with 1 for active element, and 0 for masked element (i.e., as opposed to 3D Megatron mask with False for active and True for masked). This allows easy manipulation of masks. Conversion to 3D Megatron masks is done on-the-fly.
2. Joint cross-attention mask is not precomputed, but is rather done on-the-fly.

[lgtm-com]

This pull request introduces 6 alerts and fixes 1 when merging 059c4d3 into 101977e - view on LGTM.com

new alerts:

• 2 for Unused import
• 2 for Wrong name for an argument in a class instantiation
• 1 for Unused local variable
• 1 for Syntax error

fixed alerts:

• 1 for Unused import

[lgtm-com]

This pull request introduces 6 alerts and fixes 1 when merging a7127a9 into 1b83bec - view on LGTM.com

new alerts:

• 2 for Unused import
• 2 for Wrong name for an argument in a class instantiation
• 1 for Unused local variable
• 1 for Syntax error

fixed alerts:

• 1 for Unused import

[lgtm-com]

This pull request introduces 9 alerts and fixes 1 when merging d4cbc7c into 1b83bec - view on LGTM.com

new alerts:

• 2 for Wrong name for an argument in a call
• 2 for Unused import
• 2 for Wrong name for an argument in a class instantiation
• 1 for Unused local variable
• 1 for Syntax error
• 1 for Variable defined multiple times

fixed alerts:

• 1 for Unused import

[lgtm-com]

This pull request introduces 9 alerts and fixes 1 when merging 802eae7 into 1b83bec - view on LGTM.com

new alerts:

• 3 for Unused import
• 2 for Wrong name for an argument in a call
• 2 for Wrong name for an argument in a class instantiation
• 1 for Unused local variable
• 1 for Syntax error

fixed alerts:

• 1 for Unused import

[lgtm-com]

This pull request introduces 12 alerts and fixes 1 when merging e4008e0 into 9a1cc36 - view on LGTM.com

new alerts:

• 3 for Unused import
• 2 for Wrong name for an argument in a call
• 2 for Unused local variable
• 2 for Wrong name for an argument in a class instantiation
• 1 for First parameter of a method is not named 'self'
• 1 for Superclass attribute shadows subclass method
• 1 for Syntax error

fixed alerts:

• 1 for Unused import

[lgtm-com]

This pull request introduces 9 alerts and fixes 1 when merging 5cee721 into ddcc2a6 - view on LGTM.com

new alerts:

• 2 for Unused local variable
• 2 for Unused import
• 2 for Wrong name for an argument in a class instantiation
• 1 for First parameter of a method is not named 'self'
• 1 for Superclass attribute shadows subclass method
• 1 for Syntax error

fixed alerts:

• 1 for Unused import

[lgtm-com]

This pull request introduces 2 alerts and fixes 4 when merging 7e66395 into 00df30e - view on LGTM.com

new alerts:

• 2 for Unused import

fixed alerts:

• 3 for Unused import
• 1 for First parameter of a method is not named 'self'

[lgtm-com]

This pull request introduces 2 alerts and fixes 4 when merging eda6288 into 00df30e - view on LGTM.com

new alerts:

• 2 for Unused import

fixed alerts:

• 3 for Unused import
• 1 for First parameter of a method is not named 'self'

[ericharper]

LGTM. Thanks!

Had some minor comments, but will approve. Please change your copyrights to 2022.

[lgtm-com]

This pull request introduces 2 alerts and fixes 4 when merging a2c145b into 00df30e - view on LGTM.com

new alerts:

• 2 for Unused import

fixed alerts:

• 3 for Unused import
• 1 for First parameter of a method is not named 'self'

[lgtm-com]

This pull request introduces 2 alerts and fixes 4 when merging d6bc1ec into 00df30e - view on LGTM.com

new alerts:

• 2 for Unused import

fixed alerts:

• 3 for Unused import
• 1 for First parameter of a method is not named 'self'

[lgtm-com]

This pull request introduces 2 alerts and fixes 4 when merging be1c9ec into 00df30e - view on LGTM.com

new alerts:

• 2 for Unused import

fixed alerts:

• 3 for Unused import
• 1 for First parameter of a method is not named 'self'

[lgtm-com]

This pull request introduces 1 alert and fixes 4 when merging 67e56da into 00df30e - view on LGTM.com

new alerts:

• 1 for Unused import

fixed alerts:

• 3 for Unused import
• 1 for First parameter of a method is not named 'self'

[lgtm-com]

This pull request introduces 1 alert and fixes 4 when merging 17f00d3 into 00df30e - view on LGTM.com

new alerts:

• 1 for Unused import

fixed alerts:

• 3 for Unused import
• 1 for First parameter of a method is not named 'self'

[ericharper]

LGTM. Thanks!