[BUG] Converting Megatron-LM GPT-2 to hf format raise error: KeyError: 'position_embeddings'
Author: lishuai-97
Pull request: no
Labels: name: bug

### System Info

I train a GPT-2 345M model with Megatron-LM and try to convert the model to huggingface transformers format with `transformers/src/transformers/models/megatron_gpt2
/convert_megatron_gpt2_checkpoint.py`. Then it raise error like following:
```bash
Traceback (most recent call last):
  File "tools/checkpoint/convert_megatron_gpt2_checkpoint.py", line 371, in <module>
    main()
  File "tools/checkpoint/convert_megatron_gpt2_checkpoint.py", line 330, in main
    output_state_dict = convert_megatron_checkpoint(args, input_state_dict, config)
  File "tools/checkpoint/convert_megatron_gpt2_checkpoint.py", line 138, in convert_megatron_checkpoint
    pos_embeddings = embeddings["position_embeddings"]["weight"]
KeyError: 'position_embeddings'
```

It seems that the Megatron-LM checkpoint doesn't have the `position_embeddings` state.

```bash
# model                                           
..# language_model                                
....# embedding                                   
......# word_embeddings                           
........# weight                                   : torch.Size([50304, 1024])
....# encoder                                     
......# layers.0.input_norm.weight                 : torch.Size([1024])
......# layers.0.input_norm.bias                   : torch.Size([1024])
......# layers.0.self_attention.query_key_value.weight : torch.Size([3072, 1024])
......# layers.0.self_attention.query_key_value.bias : torch.Size([3072])
......# layers.0.self_attention.dense.weight       : torch.Size([1024, 1024])
......# layers.0.self_attention.dense.bias         : torch.Size([1024])
......# layers.0.post_attention_norm.weight        : torch.Size([1024])
......# layers.0.post_attention_norm.bias          : torch.Size([1024])
......# layers.0.mlp.dense_h_to_4h.weight          : torch.Size([4096, 1024])
......# layers.0.mlp.dense_h_to_4h.bias            : torch.Size([4096])
......# layers.0.mlp.dense_4h_to_h.weight          : torch.Size([1024, 4096])
......# layers.0.mlp.dense_4h_to_h.bias            : torch.Size([1024])
......# layers.1.input_norm.weight                 : torch.Size([1024])
......# layers.1.input_norm.bias                   : torch.Size([1024])
......# layers.1.self_attention.query_key_value.weight : torch.Size([3072, 1024])
......# layers.1.self_attention.query_key_value.bias : torch.Size([3072])
......# layers.1.self_attention.dense.weight       : torch.Size([1024, 1024])
......# layers.1.self_attention.dense.bias         : torch.Size([1024])
......# layers.1.post_attention_norm.weight        : torch.Size([1024])
......# layers.1.post_attention_norm.bias          : torch.Size([1024])
......# layers.1.mlp.dense_h_to_4h.weight          : torch.Size([4096, 1024])
......# layers.1.mlp.dense_h_to_4h.bias            : torch.Size([4096])
......# layers.1.mlp.dense_4h_to_h.weight          : torch.Size([1024, 4096])
......# layers.1.mlp.dense_4h_to_h.bias            : torch.Size([1024])
```

What should I do to correctly convert the megatron-lm checkpoint to Huggingface format? Whether commenting out the position_embeddings line works. I'm not sure if this is right.

### Who can help?

_No response_

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```bash
PYTHONPATH=$megatron_path python tools/checkpoint/convert_megatron_gpt2_checkpoint.py \
    --print-checkpoint-structure \
    --path_to_checkpoint $src_dir
```

### Expected behavior

Convert the Megatron-LM checkpoint to HuggingFace format successfully.

Comments:
My guess would be that its not loading the positional_embeddings.weight correctly. See this issue I raised a couple of hours before yours: https://github.com/huggingface/transformers/issues/39038
Which transformers version are you using? If its the newest (4.52.4) setting it back to a slightly older one could solve this issue. 

> My guess would be that its not loading the positional_embeddings.weight correctly. See this issue I raised a couple of hours before yours: [#39038](https://github.com/huggingface/transformers/issues/39038) Which transformers version are you using? If its the newest (4.52.4) setting it back to a slightly older one could solve this issue.

The version of `transformers` I used is `4.46.3`. I loaded the Megatron checkpoint with `torch.load()` and inspected the resulting state dict, but it contains no `position_embeddings` key.

The states are following

```bash
{'word_embeddings': OrderedDict([('weight', tensor([[ 0.0128, -0.0070,  0.0398,  ..., -0.0018, -0.0305, -0.0233],
        [-0.0135, -0.0151, -0.0077,  ..., -0.0079, -0.0281,  0.0376],
        [-0.0095, -0.0092,  0.0378,  ..., -0.0005, -0.0190, -0.0471],
        ...,
        [ 0.0050, -0.0028,  0.0073,  ...,  0.0310,  0.0093,  0.0075],
        [ 0.0034, -0.0126, -0.0190,  ...,  0.0281, -0.0049, -0.0193],
        [ 0.0178,  0.0126, -0.0210,  ...,  0.0216, -0.0118, -0.0182]],
       dtype=torch.bfloat16))])}, 'encoder': OrderedDict([('layers.0.input_norm.weight', tensor([0.9961, 0.9961, 0.9961,  ..., 0.9961, 0.9922, 0.9883],
       dtype=torch.bfloat16)), ('layers.0.input_norm.bias', tensor([-0.0007,  0.0010, -0.0002,  ..., -0.0002, -0.0006, -0.0018],
       dtype=torch.bfloat16)), ('layers.0.self_attention.query_key_value.weight', tensor([[-0.0124, -0.0190,  0.0173,  ..., -0.0232, -0.0222,  0.0354],
        [ 0.0136,  0.0117,  0.0098,  ..., -0.0071, -0.0125, -0.0151],
        [ 0.0054, -0.0030, -0.0137,  ...,  0.0200,  0.0245,  0.0332],
        ...,
        [ 0.0009,  0.0090,  0.0098,  ..., -0.0074, -0.0225,  0.0073],
        [-0.0415,  0.0121, -0.0283,  ...,  0.0522, -0.0112, -0.0008],
        [-0.0165, -0.0415,  0.0062,  ...,  0.0060,  0.0173,  0.0146]],
       dtype=torch.bfloat16)), ('layers.0.self_attention.query_key_value.bias', tensor([-0.0055, -0.0002, -0.0010,  ..., -0.0002,  0.0002, -0.0009],
       dtype=torch.bfloat16)), ('layers.0.self_attention.dense.weight', tensor([[ 0.0018,  0.0021, -0.0017,  ...,  0.0011,  0.0029,  0.0011],
        [ 0.0040, -0.0037, -0.0010,  ..., -0.0073, -0.0007,  0.0001],
        [-0.0001,  0.0017,  0.0007,  ...,  0.0057,  0.0017,  0.0030],
        ...,
        [ 0.0014,  0.0027, -0.0007,  ...,  0.0007,  0.0008, -0.0024],
        [ 0.0023,  0.0008,  0.0057,  ...,  0.0003, -0.0012,  0.0005],
        [-0.0007,  0.0015, -0.0019,  ..., -0.0036,  0.0018, -0.0001]],
       dtype=torch.bfloat16)), ('layers.0.self_attention.dense.bias', tensor([-5.0068e-05, -8.4686e-04,  4.1199e-04,  ...,  5.2643e-04,
        -4.1199e-04, -7.8583e-04], dtype=torch.bfloat16)), ('layers.0.post_attention_norm.weight', tensor([0.9961, 1.0000, 1.0000,  ..., 0.9961, 1.0000, 0.9961],
       dtype=torch.bfloat16)), ('layers.0.post_attention_norm.bias', tensor([-2.3556e-04, -2.3499e-03,  3.9339e-05,  ...,  1.9150e-03,
        -9.1553e-04, -1.8921e-03], dtype=torch.bfloat16)), ('layers.0.mlp.dense_h_to_4h.weight', tensor([[-0.0282, -0.0154, -0.0157,  ..., -0.0079,  0.0151, -0.0110],
        [ 0.0238,  0.0204, -0.0201,  ...,  0.0038,  0.0050, -0.0199],
        [-0.0117, -0.0095,  0.0021,  ...,  0.0081, -0.0117,  0.0014],
        ...,
        [ 0.0195, -0.0007, -0.0155,  ..., -0.0028,  0.0227, -0.0112],
        [-0.0019,  0.0464,  0.0038,  ..., -0.0175,  0.0086,  0.0259],
        [ 0.0232, -0.0045,  0.0032,  ...,  0.0210, -0.0225,  0.0176]],
       dtype=torch.bfloat16)), ('layers.0.mlp.dense_h_to_4h.bias', tensor([-0.0017, -0.0014, -0.0012,  ..., -0.0006, -0.0020, -0.0017],
       dtype=torch.bfloat16)), ('layers.0.mlp.dense_4h_to_h.weight', tensor([[ 3.4637e-03, -3.1586e-03, -3.3379e-04,  ..., -2.3956e-03,
         -8.7738e-04,  1.5831e-04],
        [-3.6926e-03, -1.4954e-03,  9.7275e-05,  ..., -5.2185e-03,
          7.0572e-04,  4.7607e-03],
        [-2.0447e-03,  1.4496e-03, -3.9673e-03,  ...,  2.5330e-03,
          5.8594e-03, -4.3640e-03],
        ...,
        [ 5.2185e-03, -5.0964e-03, -3.9673e-03,  ..., -3.8452e-03,
         -4.1199e-03, -4.4861e-03],
        [-4.1504e-03, -2.8229e-03, -2.7008e-03,  ...,  1.4038e-03,
          2.5940e-03, -2.4109e-03],
        [ 2.2173e-05, -3.6011e-03,  5.9509e-03,  ..., -3.6011e-03,
         -5.1498e-04, -4.0283e-03]], dtype=torch.bfloat16)), ('layers.0.mlp.dense_4h_to_h.bias', tensor([-9.4175e-06,  2.3246e-05,  1.2302e-04,  ..., -1.6093e-05,
         1.9836e-04, -1.7166e-05], dtype=torch.bfloat16)), ('layers.1.input_norm.weight', tensor([0.9922, 0.9883, 0.9922,  ..., 0.9844, 0.9844, 0.9883],

....
```

There is an open PR here to update the conversion script, you may want to check it out! If conversion is still failing, please post there with the issues: https://github.com/huggingface/transformers/pull/38969

> There is an open PR here to update the conversion script, you may want to check it out! If conversion is still failing, please post there with the issues: [#38969](https://github.com/huggingface/transformers/pull/38969)

Thank you! It worked for me.

> My guess would be that its not loading the positional_embeddings.weight correctly. See this issue I raised a couple of hours before yours: [#39038](https://github.com/huggingface/transformers/issues/39038) Which transformers version are you using? If its the newest (4.52.4) setting it back to a slightly older one could solve this issue.

I got it! The reason is that I was using `RoPE positional encoding` instead of  `learned absolute position embedding`, so no `position_embeddings` tensor is saved in the checkpoint. I’ve therefore switched the positional encoding type from `rope` to the GPT-2's default `learned_absolute`.