GitLab

Small bugfix in bert_model.py: make sure word_embeddings is initialized before instantiating lm_head

Refactor code according to Jared's comments: move pipelining and non-pipelining training loops into separate methods

Also, use mpu.get_*_model_parallel_size() instead of args.*_model_parallel_size

Allocate tensor in `communicate()` method directly on GPU (instead of allocating on CPU and then moving to GPU)

Prevents data_loader from running out of training examples

Make sure all forward and backward operations are accounted for

Also includes following changes for inter-layer model-parallel implementation:
- Refactoring of model implementations
- Training loop changes to support inter-layer communication using `ring_exchange`
- New groups for inter-layer communication
- Checkpoint changes
- Command line arguments