Support tp pp conversion

[titu1994]

What does this PR do ?

Adds support for changing pipeline parallel version post construction for GPT

Collection: [Core, NLP]

Changelog

• Add new script for pp conversion to avoid breaking old script (should probably be deprecated eventually when new script fully supports all functionality)
• Add some modifications to allow partially loading model parallel models into memory to extract parameters
• Add PP TP conversion support for both Megatron GPT and Megatron T5 (only when shared embeddings are used between encoder and decoder)
• Add support for forcing CPU construction of T5 and GPT models
• Add support for loading model based on AppState parameters instead of Trainer parameters.

Usage

Usage:

# Megatron GPT
python megatron_change_num_partitions.py \
    --model_file=PATH_TO_SRC_FILE \
    --target_file=PATH_TO_TGT_FILE \
    --tensor_model_parallel_size=1 \
    --target_tensor_model_parallel_size=1 \
    --pipeline_model_parallel_size=1 \
    --target_pipeline_model_parallel_size=1 \
    --precision=bf16

# Megatron T5
python megatron_change_num_partitions.py \
    --model_file=PATH_TO_SRC_FILE \
    --target_file=PATH_TO_TGT_FILE \
    --model_class="nemo.collections.nlp.models.language_modeling.megatron_t5_model.MegatronT5Model" \
    --tensor_model_parallel_size=1 \
    --target_tensor_model_parallel_size=1 \
    --pipeline_model_parallel_size=1 \
    --target_pipeline_model_parallel_size=1 \
    --target_pipeline_model_parallel_split_rank=0 \
    --precision=bf16

# NOTE: When converting large models, always ensure that you pre-extract the nemo model and then only perform conversion

$ mkdir "unpacked_nemo_file"
$ tar -xvf "<path to nemo file>" -C "<absolute path to pwd>/unpacked_nemo_file/"

python megatron_change_num_partitions.py \
    ...
    --model_extracted_dir="<Absolute path to pwd>/unpacked_nemo_file/"

# NOTE: Conversion of other model types. 
# Default model type is MegatronGPTModel, if you want another model you need to pass classpath of the model
# For example - MegatronT5Model - 

python megatron_change_num_partitions.py \
    ...
    --model_class="nemo.collections.nlp.models.language_modeling.megatron_t5_model.MegatronT5Model"

# Additional arguments:

--num_gpu_per_node: Number of GPUs per node. Default is 8.
--megatron_legacy: Whether the model is a legacy Megatron model or not. Default is False. May be unsuported for 
    Pipeline Parallelism change.
--tokenizer_model_path: Path to tokenizer model. Default is None. When not None, overrides the tokenizer model path
    in the model config.
--tokenizer_vocab_file: Path to tokenizer vocab file. Default is None. When not None, overrides the tokenizer vocab
    file in the model config.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

[yidong72]

pretty cool PR. just have one minor comment.

[ericharper]

LGTM. Thanks!

[aklife97]

LGTM! just a couple of minor comments, everything else looks great!

[aklife97]

also, can we add a PP change CI as well? would be helpful to keep testing that since the PR brings in global overrides that may cause issues if someone changes it

[titu1994]

Good point about jenkins test - updated old one from tp reduce and increase to jointly increase pp by even or odd number,
Though this tests only GPT. We need nightly tests that test the whole matrix of TP (inc x dec) x PP (inc x dec) x {GPT, T5} - but that is super expensive wrt time and storage space on this CI. Will need to look into how to set that up.

[aklife97]

LGTM, thanks!