[megatron-bert-uncased-345m] fix conversion

[stas00]

Fixes: #16638

The original conversion script made an assumption that all released megatron-bert-*-345m checkpoints had the same vocab, but https://huggingface.co/nvidia/megatron-bert-cased-345m/blob/main/vocab.txt and https://huggingface.co/nvidia/megatron-bert-uncased-345m/blob/main/vocab.txt are quite different.

This PR sets config.vocab_size to the actual size of one of the params of vocab dimension.

I tested that both checkpoints mentioned above convert and load correctly:

python src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py megatron-bert-cased-345m/checkpoint.zip
python -c 'from transformers import MegatronBertForMaskedLM; MegatronBertForMaskedLM.from_pretrained("megatron-bert-cased-345m")'

python src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py megatron-bert-uncased-345m/checkpoint.zip
python -c 'from transformers import MegatronBertForMaskedLM; MegatronBertForMaskedLM.from_pretrained("megatron-bert-uncased-345m")'

both succeed.

Before this PR only the former worked, and the 2nd failed with:

RuntimeError: Error(s) in loading state_dict for MegatronBertForMaskedLM:
	size mismatch for cls.predictions.bias: copying a param with shape torch.Size([30592]) from checkpoint, the shape in current model is torch.Size([29056])

29056 is the vocab size of megatron-bert-cased-345m

@LysandreJik, @sgugger

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

Thanks for fixing!