Enabling pipeline parallelism will result in nan loss
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

Hi team, when I follow the NeMo Megatron [tutorial](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/nlp/nemo_megatron/gpt/gpt_training.html) to test megatron_gpt_pretraining script on 8 A100 devices, I encountered unexpected result when pipeline parallelism is enabled. The result is that the loss will become nan within few steps when I kick off the training.
```
Epoch 0:   0%|          | 1/351750 [00:03<372:04:57,  3.81s/it]
Epoch 0:   0%|          | 1/351750 [00:03<372:08:48,  3.81s/it, loss=11, v_num=9-02, reduced_train_loss=11.00, global_step=0.000, consumed_samples=0.000]
Epoch 0:   0%|          | 2/351750 [00:04<204:24:55,  2.09s/it, loss=11, v_num=9-02, reduced_train_loss=11.00, global_step=0.000, consumed_samples=0.000]
Epoch 0:   0%|          | 2/351750 [00:04<204:26:30,  2.09s/it, loss=nan, v_num=9-02, reduced_train_loss=nan.0, global_step=1.000, consumed_samples=192.0]
Epoch 0:   0%|          | 3/351750 [00:04<148:27:47,  1.52s/it, loss=nan, v_num=9-02, reduced_train_loss=nan.0, global_step=1.000, consumed_samples=192.0]
Epoch 0:   0%|          | 3/351750 [00:04<148:28:47,  1.52s/it, loss=nan, v_num=9-02, reduced_train_loss=nan.0, global_step=2.000, consumed_samples=384.0]
Epoch 0:   0%|          | 4/351750 [00:04<120:28:47,  1.23s/it, loss=nan, v_num=9-02, reduced_train_loss=nan.0, global_step=2.000, consumed_samples=384.0]
Epoch 0:   0%|          | 4/351750 [00:04<120:29:37,  1.23s/it, loss=nan, v_num=9-02, reduced_train_loss=nan.0, global_step=3.000, consumed_samples=576.0]
```

**Steps/Code to reproduce bug**

Please list *minimal* steps or code snippet for us to be able to reproduce the bug.

A  helpful guide on on how to craft a minimal bug report  http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports. 


Steps: 

1. Follow the [tutorial](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/nlp/nemo_megatron/gpt/gpt_training.html) until `Train GPT-style Model Option 1`
2. Change `trainer.devices=1` to `trainer.devices=8`, `model.pipeline_model_parallel_size=1` to `model.pipeline_model_parallel_size=2`

**Expected behavior**

A clear and concise description of what you expected to happen. 

I can see the loss drop and training to converge. 

**Environment overview (please complete the following information)**

 - Environment location: single AWS EC2 p4d.24xlarge machine
 - Method of NeMo install: [from source]
     - git clone https://github.com/ericharper/apex.git
     - cd apex
     - git checkout nm_v1.15.0
     - python3 -m pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_layer_norm" --global-option="--distributed_adam" --global-option="--deprecated_fused_adam" ./
     - cd ..
     - git clone https://github.com/NVIDIA/NeMo.git
     - git checkout e1b3f5e2fe57174863110d3fa47b6c24238ab58e
     - python3 -m pip install .[nlp]

**Environment details**

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version Ubuntu 20.04.4
- PyTorch version 1.12.0
- Python version 3.8.10

**Additional context**

Add any other context about the problem here.
GPU model: A100