Creating an inference server for NeMo GPT-5B leads to a bus error during model conversion
State: closed | Pull request: no
Labels: name: bug

name: stale

**Describe the bug**

Attempting to follow instructions from a post on the Nvidia developer blog: [Deploying a 1.3B GPT-3 Model with NVIDIA NeMo Framework](https://developer.nvidia.com/blog/deploying-a-1-3b-gpt-3-model-with-nvidia-nemo-megatron/). The instructions in the article work perfectly when using the 1.3B param pretrained model. However, using the same instructions for 5B leads to an error. 

**Steps/Code to reproduce bug**
1. Download the 5B parameter pretrained model using: 
```mkdir ~/nemo && cd ~/nemo && wget https://huggingface.co/nvidia/nemo-megatron-gpt-5B/resolve/main/nemo_gpt5B_fp16_tp1.nemo```
2. Run the following command:
```
docker run --rm \
    --gpus all \
    --shm-size=16GB \
    -v /home/$USER/nemo/:/checkpoints \
    -v /home/$USER/nemo/outputs-5B/:/model_repository \
    nvcr.io/ea-bignlp/bignlp-training:22.08.01-py3 \
    bash -c 'export PYTHONPATH=/opt/bignlp/FasterTransformer:${PYTHONPATH} && \
    cd /opt/bignlp && \
    python3 FasterTransformer/examples/pytorch/gpt/utils/nemo_ckpt_convert.py \
        --in-file /checkpoints/nemo_gpt5B_fp16_tp1.nemo \
        --infer-gpu-num 1 \
        --saved-dir /model_repository/gpt3_5b \
        --weight-data-type fp16 \
        --load-checkpoints-to-cpu 0 && \
    python3 /opt/bignlp/bignlp-scripts/bignlp/collections/export_scripts/prepare_triton_model_config.py \
        --model-train-name gpt3_5b \
        --template-path /opt/bignlp/fastertransformer_backend/all_models/gpt/fastertransformer/config.pbtxt \
        --ft-checkpoint /model_repository/gpt3_5b/1-gpu \
        --config-path /model_repository/gpt3_5b/config.pbtxt \
        --max-batch-size 256 \
        --pipeline-model-parallel-size 1 \
        --tensor-model-parallel-size 1 \
        --data-type bf16'
```

This command fails with the following output: 
```
This container image and its contents are governed by the NVIDIA Deep Learning Container License.
By pulling and using the container, you accept the terms and conditions of this license:
https://developer.nvidia.com/ngc/nvidia-deep-learning-container-license

NOTE: CUDA Forward Compatibility mode ENABLED.
  Using CUDA 11.7 driver version 515.48.08 with kernel driver version 470.182.03.
  See https://docs.nvidia.com/deploy/cuda-compatibility/ for details.

2023-06-23 06:31:41,635 __main__ [INFO] Spent 0:01:16.565989 (h:m:s) to unpack NeMo archive

=============== Argument ===============
saved_dir: /model_repository/gpt3_5b
in_file: /checkpoints/nemo_gpt5B_fp16_tp1.nemo
prompt_in_file: None
prompt_saved_dir: None
infer_gpu_num: 1
fused_qkv: 1
processes: 16
weight_data_type: fp16
load_checkpoints_to_cpu: 0
vocab_path: None
merges_path: None
tokenizer_model_path: None
verbose: False
========================================
bash: line 16:   102 Bus error               (core dumped) python3 FasterTransformer/examples/pytorch/gpt/utils/nemo_ckpt_convert.py --in-file /checkpoints/nemo_gpt5B_fp16_tp1.nemo --infer-gpu-num 1 --saved-dir /model_repository/gpt3_5b --weight-data-type fp16 --load-checkpoints-to-cpu 0
```


**Expected behavior**
An inference model should have been saved to `~/nemo/outputs-5B`

**Environment overview (please complete the following information)**

 - Environment location: Docker
 - Docker images used:
```
docker login nvcr.io
docker pull nvcr.io/ea-bignlp/bignlp-inference:22.08-py3
docker pull nvcr.io/ea-bignlp/bignlp-training:22.08.01-py3
```

**Additional context**
Attempting to run this on a machine that has ~200+GB RAM and an A100 GPU.