[NLP] Updating model parallel 
State: closed | Pull request: yes

Updating model parallel to work with new PTL DDP API and adding .nemo support.

# Usage
We support 4 main workflows: 
- Training from megatron-lm checkpoint
- Auto-resume training from experiment manager
- Training from model parallel .nemo file
- Evaluation from model parallel .nemo file

To enable these workflows in a NeMo NLP example training script, the PyTorch Lightning `Trainer` must be instantiated with the `NLPDDPPlugin`. This plugin contains the necessary logic to instantiate model parallel Megatron-LM models and train them.

```python
trainer = Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer)
```
Since model parallel models always require multiple GPUs, we must always have a `Trainer` object.
This means that when we restore a model parallel .nemo file, we must also pass in the `Trainer`.
```python
model = TextClassificationModel.restore_from(cfg.model.nemo_path, trainer=trainer)
```

## Training from Megatron-LM checkpoint
Megatron-LM checkpoints have the following specific format:
```
iter_0080000/
├── mp_rank_00
│   └── model_optim_rng.pt
└── mp_rank_01
    └── model_optim_rng.pt
```
In order to start training from these checkpoints, pass in the parent directory and model configuration file:
```
model.language_model.lm_checkpoint=/raid/megatron/bert/iter_0080000 \
model.language_model.config_file=/raid/data/megatron/bert/config.json \
```
The model configuration can be specified in YAML:
```yaml
model:
  language_model:
    config:
        hidden_size: 1024
        num_attention_heads: 16
        num_layers: 24
        max_position_embeddings: 512
```
Additionally, the vocab file must be specified for megatron-lm models:
```
model.tokenizer.vocab_file=/path/to/vocab.txt
```
or if using the Megatron-LM default tokenizer for training BERT:
```
model.tokenizer.tokenizer_name=megatron-bert-uncased
```

## Auto-resume from experiment manager
Auto resume works exactly the same as non-model parallel models. 
While training with PTL, model parallel checkpoint will be saved and loaded properly.
```
checkpoints/
├── mp_rank_00
│   ├── mp_autoresume-last.ckpt
│   ├── mp_autoresume---val_loss=0.35-epoch=0.ckpt
│   ├── mp_autoresume---val_loss=0.38-epoch=1.ckpt
│   └── mp_autoresume---val_loss=0.39-epoch=2.ckpt
└── mp_rank_01
    ├── mp_autoresume-last.ckpt
    ├── mp_autoresume---val_loss=0.35-epoch=0.ckpt
    ├── mp_autoresume---val_loss=0.38-epoch=1.ckpt
    └── mp_autoresume---val_loss=0.39-epoch=2.ckpt
```
## Training from model parallel .nemo
Model parallel checkpoints are saved in the correct format:
```
text_class_350m
├── megatron-bert-uncased_encoder_config.json
├── megatron_checkpoint_version.json
├── model_config.yaml
├── mp_rank_00
│   └── model_weights.ckpt
├── mp_rank_01
│   └── model_weights.ckpt
├── tokenizer_vocab_dict.json
└── tokenizer.vocab_file
```
Training from the .nemo is the same as other NeMo NLP example scripts:
```python
    trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer)
    model = TokenClassificationModel.restore_from(cfg.pretrained_model, trainer=trainer)
    trainer.fit(model)
```

## Evaluation from model parallel .nemo
Current NeMo NLP example scripts often assume using a single GPU for evaluation, which will not work.
However, evaluation with model parallel .nemo is very simple:
```python
    trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer)
    exp_manager(trainer, cfg.get("exp_manager", None))

    model = TextClassificationModel.restore_from(cfg.model.nemo_path, trainer=trainer)
    model.setup_test_data(test_data_config=cfg.model.test_ds)

    trainer.test(model=model, ckpt_path=None)
```


This pull request **introduces 5 alerts** when merging 0cd9b47b003516235ee980d6f536408b4660227e into 6962f75547e7ac2bb576f4e5a862b38d815ed79d - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-655a61c578559798dab12e49c55fba35de8fde65)

**new alerts:**

* 3 for Unused import
* 1 for Except block handles &#39;BaseException&#39;
* 1 for Module is imported with &#39;import&#39; and &#39;import from&#39;

This pull request **introduces 5 alerts** when merging 84cfb9b56b2e0fc4bbb656c1ea42dac5929bd9d8 into 3aaaa0e4f53a92fbbe0166b3201113b7fd9006e5 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-3c5a00b698f0385f69c091383646417bbb5eeab9)

**new alerts:**

* 3 for Unused import
* 1 for Except block handles &#39;BaseException&#39;
* 1 for Module is imported with &#39;import&#39; and &#39;import from&#39;

This pull request **introduces 8 alerts** when merging 85200af9decf7f7272e5b235d0e9ac418bc39d0e into 52121810b428e011666e656d313b42c954def26b - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-f29770e79c53e4fbbc3e707f07940b3e826c860c)

**new alerts:**

* 5 for Unused import
* 1 for Unused local variable
* 1 for Except block handles &#39;BaseException&#39;
* 1 for Module is imported with &#39;import&#39; and &#39;import from&#39;

This pull request **introduces 7 alerts** when merging 93dd74e48f5da59f1e367d2b0942b5439bc612a2 into 52121810b428e011666e656d313b42c954def26b - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-b218b3b475b29c5c1dad8d6f283a47eeb287b829)

**new alerts:**

* 5 for Unused import
* 1 for Except block handles &#39;BaseException&#39;
* 1 for Module is imported with &#39;import&#39; and &#39;import from&#39;

This pull request **introduces 7 alerts** when merging 6f87de6fd864d88a3cd4cd111872097d3928a2a3 into 53f33acaabd7ea0f809492d9ab16d432cfb87568 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-44917134f613ac2dfecd77760700b8361b353e19)

**new alerts:**

* 5 for Unused import
* 1 for Except block handles &#39;BaseException&#39;
* 1 for Module is imported with &#39;import&#39; and &#39;import from&#39;

This pull request **introduces 7 alerts** when merging 0600d24afa1acc4d5ea976fdb065602bf3f3f097 into 1b11678b2fed258715df09486ebf71a837d18464 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-4b3e7f69a37d7d59d2ba06a770526cfe7e170332)

**new alerts:**

* 5 for Unused import
* 1 for Except block handles &#39;BaseException&#39;
* 1 for Module is imported with &#39;import&#39; and &#39;import from&#39;

This pull request **introduces 1 alert** and **fixes 1** when merging 528fa33572b8c193e185e6ad2ae5e12c2e2e8d1d into 6a3753b3013dc92a3587853d60c5086e2e64d98f - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-0795bf804015d3643467f8a5ddcd37a401280957)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 1 for Unused import

This pull request **introduces 1 alert** and **fixes 1** when merging 5c183697e304f798d850b6421af361ff0c4c6dfe into 6a3753b3013dc92a3587853d60c5086e2e64d98f - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-459e2bbb725cb00fe7522be1e3802e3bd7729274)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 1 for Unused import

This pull request **introduces 1 alert** and **fixes 1** when merging 4f7af1c76915821b116dc7447b9f9684b4f088d3 into 6a3753b3013dc92a3587853d60c5086e2e64d98f - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-949efdbd5ab8388d1c5a7d53ba9466efaef701a2)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 1 for Unused import

This pull request **introduces 1 alert** and **fixes 1** when merging d85d9e9d8522a7deef8deca694e5f2eb0f3324e6 into 21a17b267fac68d4cdd20f3969a580a0a40dbdb4 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-a5bd4a0b168923cfaa6733d2dd5948899f826cf7)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 1 for Unused import

This pull request **introduces 1 alert** and **fixes 1** when merging f049c57d4a6aa2be43afc533e4a2cc9ea4629f9b into 0b36ea97277db6eb928f84f2c67d7ba499e10b6a - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-d9fbab213e4099c2991f90d6638a53fb73bf16e6)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 1 for Unused import

This pull request **introduces 1 alert** and **fixes 1** when merging 36ea57dc923adcbbb33f46f65884f871ff89b094 into 0b36ea97277db6eb928f84f2c67d7ba499e10b6a - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-c2eac8d6f286d5f1cb41236782b72773d0feac42)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 1 for Unused import

This pull request **introduces 1 alert** and **fixes 2** when merging 166e7a302af8331ebec19783ada27840bc1917a1 into 0b36ea97277db6eb928f84f2c67d7ba499e10b6a - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-da395e4a20acbb4d1015a5e7ecde5ae488c511cf)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 2 for Unused import

This pull request **introduces 1 alert** and **fixes 2** when merging 23d3812159f49a92683a32540eb01ae257e71132 into 0b36ea97277db6eb928f84f2c67d7ba499e10b6a - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-0d01551ac0330288223c2edfe3f15ff75d936626)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 2 for Unused import

This pull request **introduces 1 alert** and **fixes 2** when merging 4ef3f70474ccc3fbbc3d267017d7e5b846a80225 into 0b36ea97277db6eb928f84f2c67d7ba499e10b6a - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-dba4d5645a2fb07785b59751c432e3e2db195a31)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 2 for Unused import

This pull request **introduces 1 alert** and **fixes 2** when merging 88551c88b6e78612504240620b3fd8513b90efe2 into 2d68c8e7b872cb9776b8d7d7ae60d4bbcd3f7793 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-8c929c26eae37fc706a050570c1d8839140e8ceb)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 2 for Unused import

This pull request **introduces 1 alert** and **fixes 2** when merging 36522714c5c7fcc6cf9139e49c5150ebe8f884c0 into 2d68c8e7b872cb9776b8d7d7ae60d4bbcd3f7793 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-3c6fb0b8d5c6b92a3db17c8e0ed87087374aea53)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 2 for Unused import

This pull request **introduces 1 alert** and **fixes 2** when merging 1f1c02d15f72e12eb9b151f9bfc8c2658c320b5c into 2d68c8e7b872cb9776b8d7d7ae60d4bbcd3f7793 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-c6b02951d80ab1ea617892895c86c2331861d29a)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 2 for Unused import

This pull request **introduces 1 alert** and **fixes 2** when merging dfff6783379a1c7cb41b5dc74f1387c1d9f0251c into 654003494c3fdf9b996c5f27e9c665c0f1109db2 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-d62dbf30fffd00e5e59fc9b868fede68899c3cf0)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 2 for Unused import

This pull request **introduces 1 alert** and **fixes 2** when merging 7dcfa178c85eaf6cd2ed402c519f340a7a84ba14 into 654003494c3fdf9b996c5f27e9c665c0f1109db2 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-1834597da2286ab8b3a0e76775d7f89bac771d4a)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;

**fixed alerts:**

* 2 for Unused import

@titu1994 I was able to detect model parallel without unpacking the .nemo file. Thanks for catching that.

This pull request **introduces 2 alerts** and **fixes 2** when merging 28901fcb36971681c2a694434709b50a272ed94b into 654003494c3fdf9b996c5f27e9c665c0f1109db2 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-b4eea9614224903daf669af8d47184db08798e6c)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;
* 1 for Unused import

**fixed alerts:**

* 2 for Unused import

This pull request **introduces 2 alerts** and **fixes 2** when merging 24f190aea8fb0da751445c5ccb1c48ed8eeb4091 into 654003494c3fdf9b996c5f27e9c665c0f1109db2 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-79b45c775a322aabdef4ebc92d5f62b391a84d1e)

**new alerts:**

* 1 for Except block handles &#39;BaseException&#39;
* 1 for Unused import

**fixed alerts:**

* 2 for Unused import

This pull request **introduces 1 alert** and **fixes 2** when merging 721bb336b9fb525106df57d7c9924907acb870df into 654003494c3fdf9b996c5f27e9c665c0f1109db2 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-45a23cb1ebcd9c7ae6d5944b3efb04b66275ff2f)

**new alerts:**

* 1 for Unused import

**fixed alerts:**

* 2 for Unused import

This pull request **introduces 1 alert** and **fixes 2** when merging 4cba1c6ee0219713672eb6a79429cc76e78d474d into 654003494c3fdf9b996c5f27e9c665c0f1109db2 - [view on LGTM.com](https://lgtm.com/projects/g/NVIDIA/NeMo/rev/pr-a342b8f0db89e8ef9290fb92ec8fe9f392875fdb)

**new alerts:**

* 1 for Unused import

**fixed alerts:**

* 2 for Unused import