[NLP] Megatron GPT pretraining
State: closed | Pull request: yes

## This PR is highly under construction with many things subject to change.
This PR depends on libraries that have not been released yet.

## Overview
In this PR we are adding Megatron GPT pretraining to NeMo. This includes adding the model architecture to NeMo (instead of importing from megatron-lm). Currently, we still depend on megatron-lm but we will be removing the dependency over time. Our plan is to initially support tensor model parallelism only (models around 5B parameters). Later we will support pipeline model parallelism as well which will enable models with 100B+ parameters.

## Installation
- Install NeMo from this branch.
- Run make inside `nemo/collections/nlp/data/language_modeling/megatron`
- Install Apex
- Install latest version of megatron-lm from internal repo

## Usage
Data must first be preprocessed using: `examples/nlp/language_modeling/megatron_gpt_pretraining.py`
This results in `.idx` and `.bin` files:
```
pile/
├── my-gpt3_00_text_document.bin
├── my-gpt3_00_text_document.idx
├── my-gpt3_01_text_document.bin
├── my-gpt3_01_text_document.idx
```
These files are then used for training by configuring the `data_prefix` either in the config file `examples/nlp/language_modeling/conf/megatron_gpt_config.yaml`:
``` yaml
    data_prefix: 
      - .5
      - /raid/data/pile/my-gpt3_00_text_document
      - .5
      - /raid/data/pile/my-gpt3_01_text_document
```
Or from the CLI: 
```
model.data.data_prefix=[.5,/preproc_data/my-gpt3_00_text_document,.5,/preproc_data/my-gpt3_01_text_document]
```

Sample Python command:
```
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python /code/examples/nlp/language_modeling/megatron_gpt_pretraining.py \
	--config-path=/code/examples/nlp/language_modeling/conf \
	--config-name=megatron_gpt_config \
	trainer.num_nodes=${SLURM_JOB_NUM_NODES} \
	trainer.gpus=${SLURM_NTASKS_PER_NODE} \
	trainer.max_epochs=null \
	trainer.max_steps=${MAX_STEPS} \
	trainer.val_check_interval=${VAL_CHECK_INTERVAL} \
	model.micro_batch_size=${MICRO_BATCH_SIZE} \
	model.tensor_model_parallel_size=${TENSOR_MODEL_PARALLEL_SIZE} \
	model.max_position_embeddings=${SEQ_LENGTH} \
	model.encoder_seq_length=${SEQ_LENGTH} \
	model.num_layers=${NUM_LAYERS} \
	model.hidden_size=${HIDDEN_SIZE} \
	model.num_attention_heads=${NUM_ATTENTION_HEADS} \
	model.vocab_file=/artifacts/vocab.json \
	model.merge_file=/artifacts/merges.txt \
	model.data.data_prefix=[.5,/preproc_data/my-gpt3_00_text_document,.5,/preproc_data/my-gpt3_01_text_document] \
	exp_manager.wandb_logger_kwargs.name=${EXPNAME} \
	exp_manager.wandb_logger_kwargs.project=${PROJECT} \
	+exp_manager.explicit_log_dir=/results \
	+exp_manager.resume_if_exists=True \
	+exp_manager.resume_ignore_no_checkpoint=True \
	exp_manager.create_checkpoint_callback=True \
	exp_manager.checkpoint_callback_params.monitor=val_loss \
	exp_manager.checkpoint_callback_params.save_top_k=10 \
	exp_manager.checkpoint_callback_params.mode=min
```
## TODOs

- [x] Start using apex.mpu
- [x] Checkpoint saving/loading to resume training after time limit (would allow us for e2e training)
- [ ] Mimic megatron-lm log file style to easier compare performance
- [x] train_samples instead of train_iters
- [x] Gradient Accumulation for specifying any global_batch_size
- [ ] Partial activation checkpointing
- [ ] Data broadcasting for model parallel
- [ ] Model parallel gradient clipping and weight decay
- [ ] Text generation and perplexity during validation
- [ ] GPU initialization
- [ ] Barriers when creating index files
- [ ] Upgrade to 21.08 container
- [ ] Use fused kernels from apex
- [ ] Remove megatron-lm dependency
- [ ] Add lambada validation task