[bloom] convert script tweaks
Author: stas00
Pull request: yes

@younesbelkada, could you please have a look

1. when creating a small model for testing the assert for unexpected keys breaks the conversion - I think we should either not assert or perhaps warn instead?

2. also don't try to set up dtype if it's `None`

actually for the latter if `torch_dtype` is not defined, shouldn't we try to derive the dtype from the Meg-DS checkpoint dtype? since currently it'll just create fp32 weights, ignoring the actual dtype of the weights.

Thank you

Comments:
_The documentation is not available anymore as the PR was closed or merged._

sorry, forgot to follow up here.

It was just a small model that I created on the fly.

that assert is not user-friendly as it just fails w/o telling which keys are unexpected. if it were to tell which keys are unexpected I would be able to answer your question

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

@younesbelkada, this is still an issue - we are training a few variations of bloom for m4 https://github.com/huggingface/m4/blob/text_pretraining/experiments/pretraining/text_pretraining/narrow_gpt.slurm and the conversion fails in 2 asserts:

```
  File "src/transformers/models/bloom/convert_bloom_original_checkpoint_to_pytorch.py", line 194, in convert_bloom_checkpoint_to_pytorch
    assert not other_keys.unexpected_keys
AssertionError
```
and if the above removed, then next in:
```
  File "src/transformers/models/bloom/convert_bloom_original_checkpoint_to_pytorch.py", line 200, in convert_bloom_checkpoint_to_pytorch
    assert not missing_keys
AssertionError
```

it's fine then at converting. so besides my PR a 2nd assert is an issue as well.

Thank you!

the failing command line is:
```
python src/transformers/models/bloom/convert_bloom_original_checkpoint_to_pytorch.py \
--bloom_checkpoint_path    $ajs_ALL_CCFRSCRATCH/m4_text_pretraining/narrow_gpt/checkpoints/main/global_step66000 \
--pytorch_dump_folder_path $ajs_ALL_CCFRSCRATCH/m4_text_pretraining/narrow_gpt/checkpoints/hf/narrow-gpt-66000 \
--pretraining_tp 1 \
--bloom_config_file /gpfsdswork/projects/rech/cnw/commun/experiments/stas/m4/experiments/pretraining/text_pretraining/narrow_config.json
```

cc for awareness: @TevenLeScao 

grr, my apologies, @younesbelkada - it proved to be a misinformation - the pretrained model was actually meg-ds gpt2 model and not bloom :( sorry about that. 

but yes, your suggestion of printing out the unexpected keys rather than how it was before sounds great to me

Let's do that for both asserts?



for posterity since we don't have an official script converting from Megatron-Deepspeed's gpt2 code I ended up using this [Megatron-Deepspeed conversion script](https://github.com/bigscience-workshop/bigscience/tree/aa872e754106f6678e8a9dac8c6962404ba39a6d/train/tr1-13B-base#checkpoint-conversion-and-upload) we wrote when developing [pre-bloom tr1-13b-en model](https://github.com/bigscience-workshop/bigscience/tree/aa872e754106f6678e8a9dac8c6962404ba39a6d/train/tr1-13B-base) and then using HF's GPT2 modeling code to generate with it. It's not perfect as GPT2 != gpt2 in Meg-DS - 3 differences are https://github.com/bigscience-workshop/Megatron-DeepSpeed/issues/138 but it more or less works.

conversion:
```
cd Megatron-DeepSpeed
PYTHONPATH=. $six_ALL_CCFRWORK/code/Megatron-DeepSpeed/tools/convert_checkpoint/deepspeed_to_transformers.py \
--input_folder checkpoints/main/global_step112000 \
--output_folder checkpoints/hf/shallow-gpt-112000
```
validate the conversion produced a usable model:
```
python -c '\
import sys; \
mname = sys.argv[1]; \
from transformers import AutoTokenizer, AutoModelForCausalLM; \
tokenizer = AutoTokenizer.from_pretrained(mname); \
tokenizer.add_special_tokens({"pad_token": tokenizer.eos_token}); \
model = AutoModelForCausalLM.from_pretrained(mname); \
inputs = ["Hello, my dog is cute"]; \
input_tokens = tokenizer.batch_encode_plus(inputs, return_tensors="pt", padding=True)
outputs = model.generate(**input_tokens, do_sample=False); \
outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True); \
print(outputs); \
' $ajs_ALL_CCFRSCRATCH/m4_text_pretraining/shallow_gpt/checkpoints/hf/shallow-gpt-112000

['Hello, my dog is cute." "I\'m sorry, I\'m not allowed to say that."']
```
so we know it works.

Ahh I see thanks a lot for the clarification! 
Yes I think we should update the asserts ;) ! Also I think it might be useful to have a script to convert meg-ds gpt2 models to HF format, where do you think we should put this file? Or maybe just adding an arg `convert-gpt2` on the current file would work too, your call! 

I updated the PR to improve the 2nd assert, so I think it's good to merge now.

> Also I think it might be useful to have a script to convert meg-ds gpt2 models to HF format, where do you think we should put this file?

I'm not sure since it depends on `Megatron-Deepsped`'s internal files: https://github.com/bigscience-workshop/Megatron-DeepSpeed/tree/main/tools/convert_checkpoint and it's specific to this fork of Meg-DS (or rather nobody maintains the parent fork that is under MSFT)

Perhaps we add CONVERSION.md under https://github.com/huggingface/transformers/tree/main/src/transformers/models/gpt2 and show how to convert meg-ds models?

There is also https://github.com/huggingface/transformers/tree/main/src/transformers/models/megatron_gpt2 but that one can only handle Megatron-LM generated checkpoints (not Megatron-Deepspeed ones).


> There is also https://github.com/huggingface/transformers/tree/main/src/transformers/models/megatron_gpt2 but that one can only handle Megatron-LM generated checkpoints (not Megatron-Deepspeed ones).

That one folder would fit best I think, happy to open a PR there, let me know! 

I pushed the new doc here. 

Do you think we want a separate PR for that? can move it there if preferable. in a way it summarizes all the discussions of this PR.

whatever works.

this is perfectly fine for me thanks a lot @stas00 !
Gently pinging @sgugger for a final review / approval 
Thank you!