Refactor PP conversion + add support for TP only conversion

[titu1994]

What does this PR do ?

Updates the TP PP conversion script to support TP only conversion (back ported from NeMo 1.16)
Refactor PP conversion mechanism.

Collection: [NLP]

Changelog

• Refactor components to functions, split apart TP and PP conversion components
• Refactor PP conversion to model specific classes for PP conversion support
• Backport TP only conversion support (irrespective of whether PP conversion is supported or not)

Usage

python megatron_change_num_partitions.py \
    --model_file="ckpt/megatron_gpt_model.nemo" \
    --target_file="ckpt/megatron_gpt_model_tp2_pp1.nemo" \
    --tensor_model_parallel_size=1 \
    --target_tensor_model_parallel_size=2 \
    --pipeline_model_parallel_size=1 \
    --target_pipeline_model_parallel_size=1 \
    --tp_conversion_only

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

[aklife97]

with this handler way of doing things, how do we handle other models? like say BERT? do we add a handler for each model?

[titu1994]

Yep, we'll need a handler for each model. Every minor change in the model's arch can shift around too many layers or require duplication which complicates Pipeline parallelism merging.

[aklife97]

LGTM! thanks