Convert megatron lm ckpt to nemo
**Describe the bug**


I have tried converting megatron lm ckpt to nemo (pytorch .pt to .nemo file ) with model_optim_rng.pt from nvidia/megatron_bert_345m by using [https://github.com/NVIDIA/NeMo/tree/main/examples/nlp/language_modeling)/megatron_ckpt_to_nemo.py](url) script

> wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/megatron_bert_345m/versions/v0.1_uncased/zip -O megatron_bert_345m_v0.1_uncased.zip
> unzip megatron_bert_345m_v0.1_uncased.zip

> python -m torch.distributed.launch --nproc_per_node=1 megatron_lm_ckpt_to_nemo.py \
>      --checkpoint_folder release/mp_rank_00/ \
>      --checkpoint_name model_optim_rng.pt \
>      --nemo_file_path all-thai-lm.nemo \
>      --model_type bert \
>      --tensor_model_parallel_size 2 \
>      --pipeline_model_parallel_size 2

seem like torch has  some issue with GPU


```
 [NeMo W 2022-11-29 02:35:59 optimizers:77] Could not import distributed_fused_adam optimizer from Apex
OMP: Info #276: omp_set_nested routine deprecated, please use omp_set_max_active_levels instead.
[NeMo W 2022-11-29 02:36:11 experimental:27] Module <class 'nemo.collections.nlp.data.language_modeling.megatron.megatron_batch_samplers.MegatronPretrainingRandomBatchSampler'> is experimental, not ready for production and is not fully supported. Use at your own risk.
[NeMo W 2022-11-29 02:36:11 experimental:27] Module <class 'nemo.collections.nlp.models.text_normalization_as_tagging.thutmose_tagger.ThutmoseTaggerModel'> is experimental, not ready for production and is not fully supported. Use at your own risk.
[NeMo I 2022-11-29 02:36:11 distributed:31] Initializing torch.distributed with local_rank: 0, rank: 0, world_size: 1
Traceback (most recent call last):
  File "megatron_lm_ckpt_to_nemo.py", line 476, in <module>
    assert world_size % args.tensor_model_parallel_size == 0
AssertionError
/opt/conda/lib/python3.8/site-packages/torch/distributed/launch.py:178: FutureWarning: The module torch.distributed.launch is deprecated
and will be removed in future. Use torchrun.
Note that --use_env is set by default in torchrun.
If your script expects `--local_rank` argument to be set, please
change it to read from `os.environ['LOCAL_RANK']` instead. See
https://pytorch.org/docs/stable/distributed.html#launch-utility for
further instructions
   warnings.warn(
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 9735) of binary: /opt/conda/bin/python
Traceback (most recent call last):
  File "/opt/conda/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/launch.py", line 193, in <module>
    main()
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/launch.py", line 189, in main
    launch(args)
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/launch.py", line 174, in launch
    run(args)
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/run.py", line 752, in run
    elastic_launch(
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 132, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 246, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
============================================================
megatron_lm_ckpt_to_nemo.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2022-11-29_02:36:16
  host      : e792987d9454
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 9735)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
```


inside docker can detect all GPUs (nvidia-smi)
![image](https://user-images.githubusercontent.com/78535791/204431001-cae6140e-ac86-4eba-b51a-e25c6a96bd7c.png)
    

**Expected behavior**

Should get nemo LM model

**Environment overview (please complete the following information)**

 - Environment location: nvcr.io/nvidia/nemo:22.08 docker 
 - Method of NeMo install: inside image already had nemo
 

> sudo docker pull nvcr.io/nvidia/nemo:22.08 && sudo nvidia-docker run -it -v --shm-size=16g -p 8888:8888 -p 6006:6006 --ulimit memlock=-1 --ulimit stack=67108864 nvcr.io/nvidia/nemo:22.08


**Additional context**

GPU: Tesla V100
server: nvidia DGX1