Error converting p-tuned model for inferencing
State: closed | Pull request: no
Labels: name: bug

name: stale

**Describe the bug**

I am trying to p-tune a megatron_gpt_345m model using the steps and the notebook mentioned in the NeMo repo [here](https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/Multitask_Prompt_and_PTuning.ipynb).

I am able to p-tune the model for my tasks, however when I try to convert it to [FasterTransformer](https://github.com/NVIDIA/FasterTransformer) format following the steps mentioned [here](https://developer.nvidia.com/blog/deploying-a-1-3b-gpt-3-model-with-nvidia-nemo-megatron/), I get a `KeyError`.

```
WARNING: Your kernel does not support swap limit capabilities or the cgroup is not mounted. Memory limited without swap.

============
== BigNLP ==
============

NVIDIA Release 22.08.01 (build 44509868)

Container image Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copyright (c) 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

Copyright (c) 2014-2022 Facebook Inc.
Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
Copyright (c) 2011-2013 NYU                      (Clement Farabet)
Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
Copyright (c) 2015      Google Inc.
Copyright (c) 2015      Yangqing Jia
Copyright (c) 2013-2016 The Caffe contributors
All rights reserved.

Various files include modifications (c) NVIDIA CORPORATION & AFFILIATES.  All rights reserved.

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: Mellanox network driver detected, but NVIDIA peer memory driver not
      detected.  Multi-node communication performance may be reduced.

2023-01-03 09:32:00,492 __main__ [INFO] Spent 0:00:00.002772 (h:m:s) to unpack NeMo archive

=============== Argument ===============
saved_dir: /model_repository/p_tuned_gpt
in_file: /checkpoints/multitask_p_tuned_gpt.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
========================================
Traceback (most recent call last):
  File "FasterTransformer/examples/pytorch/gpt/utils/nemo_ckpt_convert.py", line 656, in <module>
    main()
  File "FasterTransformer/examples/pytorch/gpt/utils/nemo_ckpt_convert.py", line 629, in main
    convert_checkpoint(unpacked_checkpoint_dir, args)
  File "FasterTransformer/examples/pytorch/gpt/utils/nemo_ckpt_convert.py", line 280, in convert_checkpoint
    val = model_00.get("state_dict", model_00)["model.language_model.embedding.position_embeddings.weight"]
KeyError: 'model.language_model.embedding.position_embeddings.weight'
```

**Steps/Code to reproduce bug**
Here's the complete docker command for converting the  p-tuned `.nemo` model 

```
docker run --rm \
    --gpus '"device=0"' \
    --shm-size=16GB \
    --cpus 16 --memory 32g \
    -v /opt/NeMo/additional_data:/checkpoints \
    -v /opt/NeMo/additional_data/output:/model_repository \
    --network dbr0 \
    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/multitask_p_tuned_gpt.nemo \
        --infer-gpu-num 1 \
        --saved-dir /model_repository/p_tuned_gpt \
        --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 p_tuned_gpt \
        --template-path /opt/bignlp/fastertransformer_backend/all_models/gpt/fastertransformer/config.pbtxt \
        --ft-checkpoint /model_repository/p_tuned_gpt/1-gpu \
        --config-path /model_repository/p_tuned_gpt/config.pbtxt \
        --max-batch-size 256 \
        --pipeline-model-parallel-size 1 \
        --tensor-model-parallel-size 1 \
        --data-type bf16'
```


**Expected behavior**

The converted is generated and stored in the output directory mounted

**Environment overview (please complete the following information)**

Here is the docker pull command
`docker pull nvcr.io/ea-bignlp/bignlp-training:22.08.01-py3`