Cannot Convert Megatron GPT checkpoint
Author: seanbenhur
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.11.2
- Platform: Linux-5.4.0-91-generic-x86_64-with-glibc2.10
- Python version: 3.8.10
- PyTorch version (GPU?): 1.10.0a0+3fd9dcf (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

###  Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.-->

@stas00 @LysandreJik @jdemouth-nvidia 

## Information

I am trying to convert trained Megatron GPT2 checkpoint to huggingface GPT model using the provided [script]("https://github.com/huggingface/transformers/blob/master/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py"), I have followed the steps as it is in the provided [model card]("https://huggingface.co/nvidia/megatron-gpt2-345m")

But I get this error

```
root@blr-dgxa100-1:/workspace/home/sean/kannada-gpt# python3 $MYDIR/transformers/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py $MYDIR/nvidia/megatron-gpt2-345m/model_optim_rng.pt
Extracting PyTorch state dictionary from /root/nvidia/megatron-gpt2-345m/model_optim_rng.pt
Traceback (most recent call last):
  File "/root/transformers/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py", line 345, in <module>
    main()
  File "/root/transformers/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py", line 256, in main
    input_state_dict = torch.load(args.path_to_checkpoint, map_location="cpu")
  File "/opt/conda/lib/python3.8/site-packages/torch/serialization.py", line 607, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/opt/conda/lib/python3.8/site-packages/torch/serialization.py", line 882, in _load
    result = unpickler.load()
  File "/opt/conda/lib/python3.8/site-packages/torch/serialization.py", line 875, in find_class
    return super().find_class(mod_name, name)
ModuleNotFoundError: No module named 'megatron.model.enums'
root@blr-dgxa100-1:/workspace/home/sean/kannada-gpt# cp kannada_gpt_checkpoint/iter_0210000/mp_rank_00/model_optim_rng.pt $MYDIR/nvidia/megatron-gpt2-345mroot@blr-dgxa100-1:/workspace/home/sean/kannada-gpt# python3 $MYDIR/transformers/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py $MYDIR/nvidia/megatron-gpt2-345m/model_optim_rng.pt
Extracting PyTorch state dictionary from /root/nvidia/megatron-gpt2-345m/model_optim_rng.pt
Traceback (most recent call last):
  File "/root/transformers/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py", line 345, in <module>
    main()
  File "/root/transformers/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py", line 256, in main
    input_state_dict = torch.load(args.path_to_checkpoint, map_location="cpu")
  File "/opt/conda/lib/python3.8/site-packages/torch/serialization.py", line 607, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/opt/conda/lib/python3.8/site-packages/torch/serialization.py", line 882, in _load
    result = unpickler.load()
  File "/opt/conda/lib/python3.8/site-packages/torch/serialization.py", line 875, in find_class
    return super().find_class(mod_name, name)
ModuleNotFoundError: No module named 'megatron.model.enums'

```

## To reproduce

Steps to reproduce the behavior:

Unfortunately, this happens only on my trained checkpoint, I am able to convert the existing  open sourced checkpoints


<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

Should have converted the megatron checkpoint

<!-- A clear and concise description of what you would expect to happen. -->


Comments:
It's because your custom checkpoint includes a module namespace visible during training (and embedded in the checkpoint) but not at the conversion time.

The official checkpoint converts w/o errors:

```
git clone https://github.com/huggingface/transformers/
cd transformers
mkdir megatron_lm_345m
wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/megatron_lm_345m/versions/v0.0/zip -O megatron_lm_345m/checkpoint.zip
python src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py megatron_lm_345m/checkpoint.zip
Extracting PyTorch state dictionary from megatron_lm_345m/checkpoint.zip
Converting
Saving config
Adding GPT2TokenizerFast tokenizer files
Saving checkpoint to "megatron_lm_345m/pytorch_model.bin"
```
so it was probably created before Meg-LM added 'megatron.model.enums' and thus it works.

One way to solve your problem is to add `Megatron-LM` cloned path to your `PYTHONPATH`, e.g. using your setup:

```
git clone https://github.com/NVIDIA/Megatron-LM
git clone https://github.com/huggingface/transformers/
cd transformers
PYTHONPATH=../Megatron-LM python src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py /root/nvidia/megatron-gpt2-345m/model_optim_rng.pt
```

but basically **python is looking for `megatron/model/enums.py` and can't find it, so you need to tell python where to find it.**

If the relative path approach doesn't work, use the full path instead.

Note sure where to best document this though as it impacts only newer Megatron-LM code and the conversion script can't magically find Megatron-LM clone. i.e. we can't fix it on our side other than documenting this issue.

And since one can't install Megatron-LM as a package we can't make the script require it.


Thank you @stas00 , adding to PYTHONPATH has helped, maybe editing the [model card ](https://huggingface.co/nvidia/megatron-gpt2-345m) would be the right thing, let me know if I want to close this issue

Great, thank you for validating my suggestion.

We need to update 3 model cards:

- https://huggingface.co/nvidia/megatron-gpt2-345m/blob/main/README.md
- https://huggingface.co/nvidia/megatron-bert-uncased-345m/blob/main/README.md
- https://huggingface.co/nvidia/megatron-bert-cased-345m/blob/main/README.md

I need to figure out how to get perms to do so. So let's keep this open until then.

and I will also add a note to both megatron conversion scripts since all of them are impacted.


If I understand correctly, you recommend to add `PYTHONPATH=../Megatron-LM` right before the call to the `python` interpreter. That would go to line 118 of https://huggingface.co/nvidia/megatron-gpt2-345m/blob/main/README.md and similarly on the two other files. Is that right? 

Hi @jdemouth,

I was planning to add the same content as I proposed in this PR: https://github.com/huggingface/transformers/pull/14956

Note that this is not required, but needed in some cases. And for example if the checkpoint was produced with another Megatron-LM clone, like https://github.com/microsoft/Megatron-DeepSpeed/ - it's that clone that will be needed instead, since its `enums` are yet different. So basically the clone of the repo that the checkpoint was created with is needed. I probably should add this additional info to the PR as well.

I see you have access to the nvidia org on the hub and perhaps you can do the honours of editing the cards?

Thanks.

Sorry for the late reply. I have just reviewed the PR #14956. Let me update the cards. 

Thank you for editing the cards, @jdemouth!