How to convert distributed Megatron-LM ckpt to nemo
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

I have tried to convert distributed Megatron-LM ckpt to NeMo, the checkpoint files are below (tensor-parallel-size=2, pipeline-model-parallel-size=2): 
<img width="624" alt="image" src="https://github.com/NVIDIA/NeMo/assets/18243223/b8bb5696-5418-4a6d-8edc-3612ca7b812c">

I used this command to convert ckpt format:
```bash
(llama2) linrongjian@infiniAI-nxdx-B01-08:~/NeMo$ torchrun --nproc_per_node=4 \
./examples/nlp/language_modeling/megatron_lm_ckpt_to_nemo.py \
--checkpoint_folder /home/linrongjian/Megatron-LM/gpt2/checkpoint/iter_0000010 \
--checkpoint_name model_optim_rng.pt \
--nemo_file_path \
/home/linrongjian/NeMo/examples/nlp/language_modeling/gpt_distributed.nemo \
--model_type gpt \
--tensor_model_parallel_size 2 \
--pipeline_model_parallel_size 2 \
--gpus_per_node=1 \
--hparams_file /home/linrongjian/NeMo/hparams_distributed.yaml
```

And the output is:
```bash
[2023-11-20 11:34:40,349] torch.distributed.run: [WARNING] 
[2023-11-20 11:34:40,349] torch.distributed.run: [WARNING] *****************************************
[2023-11-20 11:34:40,349] torch.distributed.run: [WARNING] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
[2023-11-20 11:34:40,349] torch.distributed.run: [WARNING] *****************************************
[NeMo I 2023-11-20 11:34:49 nemo_logging:381] Initializing torch.distributed with local_rank: 0, rank: 0, world_size: 4
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] Rank 0 has data parallel group: [0]
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] All data parallel group ranks: [[0], [1], [2], [3]]
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] Ranks 0 has data parallel rank: 0
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] Rank 0 has model parallel group: [0, 1, 2, 3]
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] All model parallel group ranks: [[0, 1, 2, 3]]
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] Rank 0 has tensor model parallel group: [0, 1]
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] All tensor model parallel group ranks: [[0, 1], [2, 3]]
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] Rank 0 has tensor model parallel rank: 0
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] Rank 0 has pipeline model parallel group: [0, 2]
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] Rank 0 has embedding group: [0, 2]
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] All pipeline model parallel group ranks: [[0, 2], [1, 3]]
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] Rank 0 has pipeline model parallel rank 0
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] All embedding group ranks: [[0, 2], [1, 3]]
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] Rank 0 has embedding rank: 0
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
[NeMo I 2023-11-20 11:34:51 nemo_logging:381] loading checkpoint /home/linrongjian/Megatron-LM/gpt2/checkpoint/iter_0000010/mp_rank_00_000/model_optim_rng.pt
converted 102.48M parameters
converted 101.44M parameters
[NeMo W 2023-11-20 11:34:51 nemo_logging:393] the checkpoint version is 3.0
converted 101.44M parameters
converted 102.48M parameters
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] Rank 0 has data parallel group: [0]
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] All data parallel group ranks: [[0], [1], [2], [3]]
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] Ranks 0 has data parallel rank: 0
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] Rank 0 has model parallel group: [0, 1, 2, 3]
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] All model parallel group ranks: [[0, 1, 2, 3]]
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] Rank 0 has tensor model parallel group: [0, 1]
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] All tensor model parallel group ranks: [[0, 1], [2, 3]]
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] Rank 0 has tensor model parallel rank: 0
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] Rank 0 has pipeline model parallel group: [0, 2]
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] Rank 0 has embedding group: [0, 2]
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] All pipeline model parallel group ranks: [[0, 2], [1, 3]]
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] Rank 0 has pipeline model parallel rank 0
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] All embedding group ranks: [[0, 2], [1, 3]]
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] Rank 0 has embedding rank: 0
[NeMo W 2023-11-20 11:34:52 nemo_logging:393] You tried to register an artifact under config key=tokenizer.vocab_file but an artifact for it has already been registered.
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] Getting Megatron tokenizer for pretrained model name: megatron-gpt-345m, custom vocab file: /home/linrongjian/Megatron-LM/gpt2/gpt2-vocab.json, and merges file: /home/linrongjian/Megatron-LM/gpt2/gpt2-merges.txt
[NeMo I 2023-11-20 11:34:52 nemo_logging:381] Getting HuggingFace AutoTokenizer with pretrained_model_name: gpt2, vocab_file: /home/linrongjian/Megatron-LM/gpt2/gpt2-vocab.json, merges_files: /home/linrongjian/Megatron-LM/gpt2/gpt2-merges.txt, special_tokens_dict: {}, and use_fast: False
Using sep_token, but it is not set yet.
Using cls_token, but it is not set yet.
Using pad_token, but it is not set yet.
Using mask_token, but it is not set yet.
Using sep_token, but it is not set yet.
Using cls_token, but it is not set yet.
Using pad_token, but it is not set yet.
Using mask_token, but it is not set yet.
[NeMo I 2023-11-20 11:35:22 nemo_logging:381] Padded vocab_size: 50432, original vocab_size: 50257, dummy tokens: 175.
hidden_size: 1024, num_attention_heads: 16
[NeMo W 2023-11-20 11:35:22 nemo_logging:393] The model: MegatronGPTModel() does not have field.name: fp8_wgrad in its cfg. Add this key to cfg or config_mapping to make to make it configurable.
Using sep_token, but it is not set yet.
Using cls_token, but it is not set yet.
Using pad_token, but it is not set yet.
Using mask_token, but it is not set yet.
hidden_size: 1024, num_attention_heads: 16
hidden_size: 1024, num_attention_heads: 16
Using sep_token, but it is not set yet.
Using cls_token, but it is not set yet.
Using pad_token, but it is not set yet.
Using mask_token, but it is not set yet.
hidden_size: 1024, num_attention_heads: 16
[NeMo W 2023-11-20 11:35:22 nemo_logging:393] /home/linrongjian/anaconda3/envs/llama2/lib/python3.11/site-packages/apex/transformer/layers/layer_norm.py:93: UserWarning: `apex.contrib.layer_norm.FastLayerNorm` isn't available thus falling back to `apex.normalization.FusedLayerNorm`
      warnings.warn("`apex.contrib.layer_norm.FastLayerNorm` isn't available thus falling back to `apex.normalization.FusedLayerNorm`")
    
[NeMo W 2023-11-20 11:35:22 nemo_logging:393] /home/linrongjian/anaconda3/envs/llama2/lib/python3.11/site-packages/pytorch_lightning/core/saving.py:161: UserWarning: Found keys that are in the model state dict but not in the checkpoint: ['model.language_model.encoder.layers.0.input_layernorm.weight', 'model.language_model.encoder.layers.0.input_layernorm.bias', 'model.language_model.encoder.layers.0.post_attention_layernorm.weight', 'model.language_model.encoder.layers.0.post_attention_layernorm.bias', 'model.language_model.encoder.layers.1.input_layernorm.weight', 'model.language_model.encoder.layers.1.input_layernorm.bias', 'model.language_model.encoder.layers.1.post_attention_layernorm.weight', 'model.language_model.encoder.layers.1.post_attention_layernorm.bias', 'model.language_model.encoder.layers.2.input_layernorm.weight', 'model.language_model.encoder.layers.2.input_layernorm.bias', 'model.language_model.encoder.layers.2.post_attention_layernorm.weight', 'model.language_model.encoder.layers.2.post_attention_layernorm.bias', 'model.language_model.encoder.layers.3.input_layernorm.weight', 'model.language_model.encoder.layers.3.input_layernorm.bias', 'model.language_model.encoder.layers.3.post_attention_layernorm.weight', 'model.language_model.encoder.layers.3.post_attention_layernorm.bias', 'model.language_model.encoder.layers.4.input_layernorm.weight', 'model.language_model.encoder.layers.4.input_layernorm.bias', 'model.language_model.encoder.layers.4.post_attention_layernorm.weight', 'model.language_model.encoder.layers.4.post_attention_layernorm.bias', 'model.language_model.encoder.layers.5.input_layernorm.weight', 'model.language_model.encoder.layers.5.input_layernorm.bias', 'model.language_model.encoder.layers.5.post_attention_layernorm.weight', 'model.language_model.encoder.layers.5.post_attention_layernorm.bias', 'model.language_model.encoder.layers.6.input_layernorm.weight', 'model.language_model.encoder.layers.6.input_layernorm.bias', 'model.language_model.encoder.layers.6.post_attention_layernorm.weight', 'model.language_model.encoder.layers.6.post_attention_layernorm.bias', 'model.language_model.encoder.layers.7.input_layernorm.weight', 'model.language_model.encoder.layers.7.input_layernorm.bias', 'model.language_model.encoder.layers.7.post_attention_layernorm.weight', 'model.language_model.encoder.layers.7.post_attention_layernorm.bias', 'model.language_model.encoder.layers.8.input_layernorm.weight', 'model.language_model.encoder.layers.8.input_layernorm.bias', 'model.language_model.encoder.layers.8.post_attention_layernorm.weight', 'model.language_model.encoder.layers.8.post_attention_layernorm.bias', 'model.language_model.encoder.layers.9.input_layernorm.weight', 'model.language_model.encoder.layers.9.input_layernorm.bias', 'model.language_model.encoder.layers.9.post_attention_layernorm.weight', 'model.language_model.encoder.layers.9.post_attention_layernorm.bias', 'model.language_model.encoder.layers.10.input_layernorm.weight', 'model.language_model.encoder.layers.10.input_layernorm.bias', 'model.language_model.encoder.layers.10.post_attention_layernorm.weight', 'model.language_model.encoder.layers.10.post_attention_layernorm.bias', 'model.language_model.encoder.layers.11.input_layernorm.weight', 'model.language_model.encoder.layers.11.input_layernorm.bias', 'model.language_model.encoder.layers.11.post_attention_layernorm.weight', 'model.language_model.encoder.layers.11.post_attention_layernorm.bias']
      rank_zero_warn(
    
[NeMo W 2023-11-20 11:35:22 nemo_logging:393] /home/linrongjian/anaconda3/envs/llama2/lib/python3.11/site-packages/pytorch_lightning/core/saving.py:165: UserWarning: Found keys that are not in the model state dict but in the checkpoint: ['model.language_model.encoder.layers.0.input_norm.weight', 'model.language_model.encoder.layers.0.input_norm.bias', 'model.language_model.encoder.layers.0.post_attention_norm.weight', 'model.language_model.encoder.layers.0.post_attention_norm.bias', 'model.language_model.encoder.layers.1.input_norm.weight', 'model.language_model.encoder.layers.1.input_norm.bias', 'model.language_model.encoder.layers.1.post_attention_norm.weight', 'model.language_model.encoder.layers.1.post_attention_norm.bias', 'model.language_model.encoder.layers.2.input_norm.weight', 'model.language_model.encoder.layers.2.input_norm.bias', 'model.language_model.encoder.layers.2.post_attention_norm.weight', 'model.language_model.encoder.layers.2.post_attention_norm.bias', 'model.language_model.encoder.layers.3.input_norm.weight', 'model.language_model.encoder.layers.3.input_norm.bias', 'model.language_model.encoder.layers.3.post_attention_norm.weight', 'model.language_model.encoder.layers.3.post_attention_norm.bias', 'model.language_model.encoder.layers.4.input_norm.weight', 'model.language_model.encoder.layers.4.input_norm.bias', 'model.language_model.encoder.layers.4.post_attention_norm.weight', 'model.language_model.encoder.layers.4.post_attention_norm.bias', 'model.language_model.encoder.layers.5.input_norm.weight', 'model.language_model.encoder.layers.5.input_norm.bias', 'model.language_model.encoder.layers.5.post_attention_norm.weight', 'model.language_model.encoder.layers.5.post_attention_norm.bias', 'model.language_model.encoder.layers.6.input_norm.weight', 'model.language_model.encoder.layers.6.input_norm.bias', 'model.language_model.encoder.layers.6.post_attention_norm.weight', 'model.language_model.encoder.layers.6.post_attention_norm.bias', 'model.language_model.encoder.layers.7.input_norm.weight', 'model.language_model.encoder.layers.7.input_norm.bias', 'model.language_model.encoder.layers.7.post_attention_norm.weight', 'model.language_model.encoder.layers.7.post_attention_norm.bias', 'model.language_model.encoder.layers.8.input_norm.weight', 'model.language_model.encoder.layers.8.input_norm.bias', 'model.language_model.encoder.layers.8.post_attention_norm.weight', 'model.language_model.encoder.layers.8.post_attention_norm.bias', 'model.language_model.encoder.layers.9.input_norm.weight', 'model.language_model.encoder.layers.9.input_norm.bias', 'model.language_model.encoder.layers.9.post_attention_norm.weight', 'model.language_model.encoder.layers.9.post_attention_norm.bias', 'model.language_model.encoder.layers.10.input_norm.weight', 'model.language_model.encoder.layers.10.input_norm.bias', 'model.language_model.encoder.layers.10.post_attention_norm.weight', 'model.language_model.encoder.layers.10.post_attention_norm.bias', 'model.language_model.encoder.layers.11.input_norm.weight', 'model.language_model.encoder.layers.11.input_norm.bias', 'model.language_model.encoder.layers.11.post_attention_norm.weight', 'model.language_model.encoder.layers.11.post_attention_norm.bias']
      rank_zero_warn(
    
[NeMo W 2023-11-20 11:35:22 nemo_logging:393] You tried to register an artifact under config key=tokenizer.vocab_file but an artifact for it has already been registered.
[NeMo W 2023-11-20 11:35:22 nemo_logging:393] You tried to register an artifact under config key=tokenizer.merge_file but an artifact for it has already been registered.
[NeMo I 2023-11-20 11:35:25 nemo_logging:381] NeMo model saved to: /home/linrongjian/NeMo/examples/nlp/language_modeling/gpt_distributed.nemo
```

Seems like only one ckpt(mp_rank_00_000) got convertted, while the other three didn't.


**Expected behavior**

Should get all nemo ckpts.

**Environment overview (please complete the following information)**

 - Method of NeMo install: from source (main branch)


**Additional context**

GPU: Nvidia A100 * 8


Hi @MachineGunLin 

Can you please let us know how you resolved the issue?