Can't train and load TFGPT2LMHeadModel from disc
Author: danielricks
Pull request: no

### System Info

Ubuntu 20.04
Python 3.8.10
tokenizers 0.13.3
transformers  4.31.0

### Who can help?

@ArthurZucker @Rocketknight

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I'm wondering why I can't train and load a TFGPT2LMHeadModel from disc (specifically TF, the torch library doesn't seem to work on my machine and I'd like to work with the TF version unless it's absolutely not possible). I can train a Tokenizer just fine (I know there are pretrained Tokenizers out there, but I need to train my own at the word level.)

The code so far will train both a Tokenizer and a TFGPT2LMHeadModel and save them both, but when I go to load them from disc, only the Tokenizer survives the journey. The Model not only reports the following error,  but is noticeably untrained when I go to generate output (this difference in output is not shown using an English corpus, but the following error still appears, and I assume that once that error goes away so will my issue. If the layers can be loaded, then I'll assume the model can be loaded "trained").

Error on "load_pretrained":
```
Some layers of TFGPT2LMHeadModel were not initialized from the model checkpoint at model_folder and are newly initialized: ['transformer/h_._7/ln_2/beta:0', 'transformer/h_._2/mlp/c_fc/weight:0', 'transformer/h_._9/mlp/c_proj/bias:0', 'transformer/h_._11/attn/c_proj/bias:0', 'transformer/h_._0/ln_2/gamma:0', 'transformer/h_._6/mlp/c_proj/bias:0', 'transformer/h_._9/ln_2/beta:0', 'transformer/h_._6/ln_1/beta:0', 'transformer/h_._5/ln_2/beta:0', 'transformer/h_._8/attn/c_proj/weight:0', 'transformer/h_._8/attn/c_proj/bias:0', 'transformer/h_._1/attn/c_attn/bias:0', 'transformer/h_._1/attn/c_proj/weight:0', 'transformer/h_._6/ln_1/gamma:0', 'transformer/h_._11/attn/c_attn/bias:0', 'transformer/h_._0/attn/c_attn/weight:0', 'transformer/h_._0/mlp/c_proj/weight:0', 'transformer/h_._6/mlp/c_proj/weight:0', 'transformer/h_._7/attn/c_proj/weight:0', 'transformer/ln_f/gamma:0', 'transformer/h_._4/ln_2/beta:0', 'transformer/h_._9/mlp/c_fc/bias:0', 'transformer/h_._8/mlp/c_fc/weight:0', 'transformer/h_._8/mlp/c_proj/weight:0', 'transformer/h_._7/mlp/c_proj/weight:0', 'transformer/h_._0/ln_2/beta:0', 'transformer/h_._9/attn/c_proj/weight:0', 'transformer/h_._1/mlp/c_proj/bias:0', 'transformer/h_._6/mlp/c_fc/bias:0', 'transformer/h_._10/attn/c_proj/weight:0', 'transformer/h_._5/ln_1/gamma:0', 'transformer/h_._6/mlp/c_fc/weight:0', 'transformer/h_._8/attn/c_attn/bias:0', 'transformer/h_._10/mlp/c_fc/bias:0', 'transformer/h_._7/attn/c_proj/bias:0', 'transformer/h_._6/attn/c_proj/weight:0', 'transformer/h_._9/attn/c_proj/bias:0', 'transformer/h_._2/attn/c_proj/bias:0', 'transformer/h_._8/ln_1/beta:0', 'transformer/h_._3/mlp/c_fc/weight:0', 'transformer/h_._5/attn/c_proj/bias:0', 'transformer/h_._0/mlp/c_proj/bias:0', 'transformer/wpe/weight:0', 'transformer/h_._1/ln_1/gamma:0', 'transformer/h_._11/ln_2/gamma:0', 'transformer/h_._6/attn/c_proj/bias:0', 'transformer/h_._0/attn/c_proj/bias:0', 'transformer/h_._4/ln_1/gamma:0', 'transformer/h_._1/attn/c_proj/bias:0', 'transformer/h_._4/ln_2/gamma:0', 'transformer/h_._9/mlp/c_proj/weight:0', 'transformer/h_._11/ln_1/beta:0', 'transformer/h_._10/mlp/c_fc/weight:0', 'transformer/h_._4/attn/c_proj/bias:0', 'transformer/h_._10/attn/c_proj/bias:0', 'transformer/h_._0/attn/c_attn/bias:0', 'transformer/h_._2/ln_1/gamma:0', 'transformer/ln_f/beta:0', 'transformer/h_._7/mlp/c_fc/weight:0', 'transformer/h_._3/attn/c_attn/weight:0', 'transformer/h_._7/mlp/c_proj/bias:0', 'transformer/h_._8/ln_2/gamma:0', 'transformer/h_._2/mlp/c_proj/weight:0', 'transformer/h_._11/ln_2/beta:0', 'transformer/h_._1/ln_2/beta:0', 'transformer/h_._5/mlp/c_fc/weight:0', 'transformer/h_._2/attn/c_attn/bias:0', 'transformer/h_._7/mlp/c_fc/bias:0', 'transformer/h_._9/ln_2/gamma:0', 'transformer/h_._11/mlp/c_fc/bias:0', 'transformer/h_._7/ln_2/gamma:0', 'transformer/h_._3/attn/c_proj/bias:0', 'transformer/h_._6/ln_2/gamma:0', 'transformer/h_._3/mlp/c_proj/weight:0', 'transformer/h_._5/attn/c_proj/weight:0', 'transformer/h_._2/attn/c_attn/weight:0', 'transformer/h_._11/mlp/c_fc/weight:0', 'transformer/h_._5/ln_2/gamma:0', 'transformer/h_._6/ln_2/beta:0', 'transformer/h_._8/attn/c_attn/weight:0', 'transformer/h_._10/mlp/c_proj/bias:0', 'transformer/h_._10/ln_1/beta:0', 'transformer/h_._3/attn/c_proj/weight:0', 'transformer/h_._4/ln_1/beta:0', 'transformer/h_._11/mlp/c_proj/bias:0', 'transformer/h_._4/mlp/c_fc/weight:0', 'transformer/h_._11/ln_1/gamma:0', 'transformer/h_._1/attn/c_attn/weight:0', 'transformer/h_._8/ln_1/gamma:0', 'transformer/h_._0/ln_1/beta:0', 'transformer/h_._10/mlp/c_proj/weight:0', 'transformer/h_._9/attn/c_attn/bias:0', 'transformer/h_._2/ln_1/beta:0', 'transformer/h_._1/mlp/c_proj/weight:0', 'transformer/h_._2/attn/c_proj/weight:0', 'transformer/h_._7/attn/c_attn/bias:0', 'transformer/h_._5/mlp/c_proj/weight:0', 'transformer/h_._4/attn/c_proj/weight:0', 'transformer/h_._10/attn/c_attn/weight:0', 'transformer/h_._8/ln_2/beta:0', 'transformer/h_._9/ln_1/gamma:0', 'transformer/h_._2/ln_2/gamma:0', 'transformer/h_._2/ln_2/beta:0', 'transformer/h_._10/ln_2/beta:0', 'transformer/h_._7/ln_1/gamma:0', 'transformer/h_._7/attn/c_attn/weight:0', 'transformer/h_._6/attn/c_attn/weight:0', 'transformer/h_._5/attn/c_attn/bias:0', 'transformer/h_._0/mlp/c_fc/weight:0', 'transformer/h_._8/mlp/c_fc/bias:0', 'transformer/h_._10/attn/c_attn/bias:0', 'transformer/h_._5/ln_1/beta:0', 'transformer/h_._3/mlp/c_fc/bias:0', 'transformer/h_._10/ln_2/gamma:0', 'transformer/h_._11/attn/c_proj/weight:0', 'transformer/h_._6/attn/c_attn/bias:0', 'transformer/h_._4/mlp/c_proj/bias:0', 'transformer/h_._3/ln_1/gamma:0', 'transformer/h_._0/ln_1/gamma:0', 'transformer/h_._4/attn/c_attn/weight:0', 'transformer/h_._8/mlp/c_proj/bias:0', 'transformer/h_._3/attn/c_attn/bias:0', 'transformer/h_._5/mlp/c_fc/bias:0', 'transformer/h_._5/attn/c_attn/weight:0', 'transformer/h_._3/ln_2/gamma:0', 'transformer/h_._3/ln_1/beta:0', 'transformer/h_._0/attn/c_proj/weight:0', 'transformer/h_._4/mlp/c_proj/weight:0', 'transformer/h_._11/mlp/c_proj/weight:0', 'transformer/h_._11/attn/c_attn/weight:0', 'transformer/h_._2/mlp/c_fc/bias:0', 'transformer/h_._9/mlp/c_fc/weight:0', 'transformer/h_._0/mlp/c_fc/bias:0', 'transformer/h_._3/ln_2/beta:0', 'transformer/h_._1/mlp/c_fc/weight:0', 'transformer/h_._7/ln_1/beta:0', 'transformer/h_._1/ln_2/gamma:0', 'transformer/h_._4/mlp/c_fc/bias:0', 'transformer/h_._10/ln_1/gamma:0', 'transformer/h_._1/mlp/c_fc/bias:0', 'transformer/h_._5/mlp/c_proj/bias:0', 'transformer/h_._4/attn/c_attn/bias:0', 'transformer/h_._1/ln_1/beta:0', 'transformer/h_._3/mlp/c_proj/bias:0', 'transformer/h_._2/mlp/c_proj/bias:0', 'transformer/h_._9/ln_1/beta:0', 'transformer/wte/weight:0', 'transformer/h_._9/attn/c_attn/weight:0']
```

Again, the following code is a toy example in English that shows the same error. My goal is to make that error go away without changing the model type (TFGPT2LMHeadModel) and without changing the WordLevel Tokenizer. The above error is really a problem here. I've seen other answers where it's not actually a problem (IE https://github.com/huggingface/transformers/issues/11192), but the model really appears to be untrained after it's loaded. Am I doing something wrong, or is this particular model broken/doesn't support my operation and it's just not documented? Thanks!

```python
import os, logging, pathlib, time
from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from tokenizers.normalizers import NFKC, Sequence
from tokenizers.pre_tokenizers import WhitespaceSplit
from tokenizers.trainers import WordLevelTrainer

from transformers import PreTrainedTokenizerFast
from transformers import GPT2Config, TFGPT2LMHeadModel
from transformers import CONFIG_NAME

import tensorflow as tf


data_folder = "data_folder"
model_folder = "model_folder"

pathlib.Path(data_folder).mkdir(parents=True, exist_ok=True)
pathlib.Path(model_folder).mkdir(parents=True, exist_ok=True)

training_data_filename = "training_data.txt"
training_data_filepath = data_folder + "/" + training_data_filename
paths = [training_data_filepath]

text = "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Tincidunt praesent semper feugiat nibh sed pulvinar proin gravida hendrerit. Nisi lacus sed viverra tellus in hac habitasse platea dictumst. Convallis convallis tellus id interdum. Sed libero enim sed faucibus. Luctus accumsan tortor posuere ac ut consequat semper viverra nam. Fermentum odio eu feugiat pretium nibh ipsum consequat nisl. Augue mauris augue neque gravida in. Vitae suscipit tellus mauris a diam. Eleifend quam adipiscing vitae proin. Arcu cursus euismod quis viverra nibh cras pulvinar mattis nunc. Amet mauris commodo quis imperdiet massa tincidunt nunc. Pulvinar mattis nunc sed blandit libero. Ultrices tincidunt arcu non sodales neque sodales ut. Mi in nulla posuere sollicitudin. Elit ullamcorper dignissim cras tincidunt. Imperdiet sed euismod nisi porta lorem mollis aliquam. Lectus magna fringilla urna porttitor. Id donec ultrices tincidunt arcu. Tempor nec feugiat nisl pretium fusce id velit. Aliquam etiam erat velit scelerisque in. Risus nec feugiat in fermentum posuere urna. Lacus luctus accumsan tortor posuere ac. Feugiat scelerisque varius morbi enim nunc faucibus a pellentesque. Eget dolor morbi non arcu risus quis varius. Non enim praesent elementum facilisis leo vel fringilla. Placerat duis ultricies lacus sed turpis tincidunt id aliquet risus. Commodo quis imperdiet massa tincidunt nunc. Egestas erat imperdiet sed euismod nisi. Pulvinar elementum integer enim neque volutpat ac tincidunt. Tristique senectus et netus et malesuada fames ac. Dignissim cras tincidunt lobortis feugiat vivamus at augue. Et malesuada fames ac turpis egestas. Diam quam nulla porttitor massa id neque aliquam vestibulum morbi. Vitae congue eu consequat ac felis donec. Enim praesent elementum facilisis leo vel. Eleifend donec pretium vulputate sapien nec. Mauris ultrices eros in cursus. Amet cursus sit amet dictum sit amet justo donec. Sollicitudin nibh sit amet commodo. Mi in nulla posuere sollicitudin aliquam ultrices sagittis orci. Ac felis donec et odio. Tellus id interdum velit laoreet. Nibh tellus molestie nunc non blandit massa enim nec dui. In fermentum posuere urna nec tincidunt praesent semper feugiat nibh. Semper viverra nam libero justo laoreet. Ultricies integer quis auctor elit sed vulputate mi sit amet. Diam maecenas ultricies mi eget mauris pharetra et. Dui nunc mattis enim ut tellus elementum sagittis vitae et. Gravida in fermentum et sollicitudin. Tellus at urna condimentum mattis pellentesque id nibh tortor id. Laoreet id donec ultrices tincidunt arcu non sodales neque sodales. Elit at imperdiet dui accumsan sit amet nulla facilisi. Suspendisse ultrices gravida dictum fusce ut placerat orci nulla. Blandit aliquam etiam erat velit. Sodales ut eu sem integer vitae justo eget. Dolor sit amet consectetur adipiscing elit duis. Purus in mollis nunc sed id. Augue mauris augue neque gravida in fermentum et. Justo nec ultrices dui sapien eget mi. Facilisis mauris sit amet massa. Orci dapibus ultrices in iaculis nunc sed. Sapien faucibus et molestie ac feugiat sed lectus. Consequat mauris nunc congue nisi vitae suscipit tellus mauris a. Augue mauris augue neque gravida. Iaculis nunc sed augue lacus viverra. Ultrices neque ornare aenean euismod elementum nisi quis. Cras tincidunt lobortis feugiat vivamus at augue eget arcu dictum. In hac habitasse platea dictumst quisque. At erat pellentesque adipiscing commodo elit at imperdiet. Vulputate eu scelerisque felis imperdiet proin fermentum leo vel. Elit scelerisque mauris pellentesque pulvinar pellentesque habitant morbi tristique. Nibh praesent tristique magna sit amet purus gravida. Faucibus interdum posuere lorem ipsum dolor sit. Vitae purus faucibus ornare suspendisse sed. Donec adipiscing tristique risus nec feugiat in. Neque volutpat ac tincidunt vitae semper quis. Pellentesque massa placerat duis ultricies lacus sed turpis tincidunt. Justo nec ultrices dui sapien eget mi proin sed libero. Quisque sagittis purus sit amet volutpat consequat mauris nunc congue. Gravida in fermentum et sollicitudin ac orci phasellus. Eget nullam non nisi est. Neque convallis a cras semper. Erat imperdiet sed euismod nisi porta lorem mollis. Ultricies mi quis hendrerit dolor magna. Risus commodo viverra maecenas accumsan lacus vel. Tempor commodo ullamcorper a lacus vestibulum sed. Et magnis dis parturient montes. Est pellentesque elit ullamcorper dignissim cras tincidunt lobortis feugiat. Tincidunt id aliquet risus feugiat in ante metus. Condimentum mattis pellentesque id nibh tortor id. Blandit aliquam etiam erat velit scelerisque in. Laoreet non curabitur gravida arcu ac. Auctor neque vitae tempus quam pellentesque nec. Vitae aliquet nec ullamcorper sit. Convallis convallis tellus id interdum velit laoreet id. Lobortis scelerisque fermentum dui faucibus in ornare. Elementum nibh tellus molestie nunc. Arcu cursus euismod quis viverra nibh. Mi sit amet mauris commodo. Duis ultricies lacus sed turpis tincidunt id aliquet. Interdum varius sit amet mattis. Et molestie ac feugiat sed lectus vestibulum. Risus feugiat in ante metus dictum. Risus feugiat in ante metus dictum at tempor. Est velit egestas dui id. Scelerisque eu ultrices vitae auctor eu augue ut. Aliquam etiam erat velit scelerisque in dictum non. Justo eget magna fermentum iaculis eu non. Platea dictumst quisque sagittis purus sit amet volutpat consequat mauris. Aliquam ut porttitor leo a diam. Ante metus dictum at tempor commodo ullamcorper a lacus vestibulum. Quis ipsum suspendisse ultrices gravida dictum fusce ut placerat. Nunc sed augue lacus viverra vitae congue eu. Arcu ac tortor dignissim convallis aenean et tortor at risus. Pretium quam vulputate dignissim suspendisse in est ante in nibh. A arcu cursus vitae congue mauris. Ut pharetra sit amet aliquam id diam maecenas ultricies mi. Et molestie ac feugiat sed lectus vestibulum mattis ullamcorper velit. Eget mauris pharetra et ultrices neque ornare aenean. Eu tincidunt tortor aliquam nulla facilisi. Nibh cras pulvinar mattis nunc sed blandit libero. Massa eget egestas purus viverra accumsan in nisl. Bibendum enim facilisis gravida neque convallis. Neque vitae tempus quam pellentesque nec nam aliquam sem et. Aliquam malesuada bibendum arcu vitae elementum curabitur. Adipiscing bibendum est ultricies integer quis auctor elit. Est lorem ipsum dolor sit amet. Tellus elementum sagittis vitae et leo duis ut. Mollis nunc sed id semper risus. Sapien faucibus et molestie ac feugiat sed lectus vestibulum. Fusce id velit ut tortor pretium viverra suspendisse potenti nullam. Morbi non arcu risus quis. Posuere urna nec tincidunt praesent semper. Urna et pharetra pharetra massa. Tristique magna sit amet purus gravida quis blandit turpis. Egestas integer eget aliquet nibh. Habitant morbi tristique senectus et netus et malesuada fames. In nisl nisi scelerisque eu ultrices vitae auctor. Sed velit dignissim sodales ut eu sem integer. Vulputate odio ut enim blandit. Enim diam vulputate ut pharetra. Amet luctus venenatis lectus magna fringilla. Etiam sit amet nisl purus in mollis. Arcu cursus euismod quis viverra nibh cras pulvinar mattis nunc. Eget mauris pharetra et ultrices neque ornare aenean. Pellentesque id nibh tortor id aliquet lectus proin nibh. Nunc mi ipsum faucibus vitae aliquet nec ullamcorper. Mi tempus imperdiet nulla malesuada pellentesque elit eget. Ut consequat semper viverra nam. Aliquet eget sit amet tellus cras adipiscing enim. Fames ac turpis egestas sed tempus. Dui vivamus arcu felis bibendum ut. Aliquet porttitor lacus luctus accumsan tortor. Rhoncus dolor purus non enim praesent elementum facilisis leo. Egestas erat imperdiet sed euismod nisi porta lorem. Enim sed faucibus turpis in eu mi. Amet porttitor eget dolor morbi non arcu risus quis varius. Euismod elementum nisi quis eleifend quam adipiscing. Dictumst quisque sagittis purus sit amet volutpat consequat mauris. Faucibus scelerisque eleifend donec pretium vulputate sapien nec sagittis. Maecenas ultricies mi eget mauris pharetra. Nulla facilisi cras fermentum odio eu feugiat pretium nibh. Rhoncus aenean vel elit scelerisque mauris pellentesque pulvinar pellentesque. Vestibulum morbi blandit cursus risus at ultrices mi tempus imperdiet. Ac odio tempor orci dapibus ultrices in iaculis nunc. Gravida quis blandit turpis cursus in hac habitasse platea dictumst. Malesuada fames ac turpis egestas maecenas. Aenean pharetra magna ac placerat vestibulum lectus."

with open(training_data_filepath, "w") as f:
    f.write(text)

symbol_count = len(list(set(text.split(" "))))

tokenizer = Tokenizer(WordLevel())
tokenizer.normalizer = Sequence([NFKC()])
tokenizer.pre_tokenizer = WhitespaceSplit()

trainer = WordLevelTrainer(
    vocab_size=symbol_count,
    show_progress=True
)
tokenizer.train(trainer=trainer, files=paths)

print("text", "consectetur", "labore")
encoded = tokenizer.encode("consectetur", "labore")
print("encoded", encoded.ids)

fast_tokenizer = PreTrainedTokenizerFast(tokenizer_object=tokenizer)

print("text", "consectetur", "labore")
encoded = fast_tokenizer.encode("consectetur", "labore")
print("encoded", encoded)

tokenizer = fast_tokenizer


logging.info("String_tokenized")
training_data_filename = "training_data.txt"
training_data_filepath = data_folder + "/" + training_data_filename
single_string = ""
for filename in [training_data_filepath]:
    with open(filename, "r", encoding="utf-8") as f:
        x = f.read()
        single_string += x
string_tokenized = tokenizer.encode(single_string)

logging.info("Batching dataset")
examples = []
block_size = 100
BATCH_SIZE = 12
BUFFER_SIZE = 1000
for i in range(0, len(string_tokenized) - block_size + 1, block_size):
    examples.append(string_tokenized[i : i + block_size])
inputs, labels = [], []
for ex in examples:
    inputs.append(ex[:-1])
    labels.append(ex[1:])
dataset = tf.data.Dataset.from_tensor_slices((inputs, labels))
dataset = dataset.shuffle(BUFFER_SIZE).batch(
    BATCH_SIZE, drop_remainder=True
)

config = GPT2Config(vocab_size=len(tokenizer.get_vocab()))
model = TFGPT2LMHeadModel(config)
optimizer = tf.keras.optimizers.Adam(
    learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0
)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metric = tf.keras.metrics.SparseCategoricalAccuracy("accuracy")
model.compile(
    optimizer=optimizer,
    loss=[loss, *[None] * model.config.n_layer],
    metrics=[metric],
)


num_epochs = 5
time0 = time.time()
logging.info("Beginning training: epoch {0}".format(time0))
model.fit(dataset, epochs=num_epochs, verbose=0)
logging.info("Training took {0} Seconds".format(str(time.time() - time0)))


output_config_file = os.path.join(model_folder, CONFIG_NAME)
model.config.to_json_file(output_config_file)

model.save_pretrained(model_folder)
tokenizer.save_pretrained(model_folder)


def generate(model, tokenizer, text):
    input_ids = tokenizer.encode(text, return_tensors="tf")

    top_k = 50
    top_p = 0.95
    output = model.generate(
        input_ids,
        max_length=300,
        do_sample=True,
        temperature=0.3,
        no_repeat_ngram_size=2,
        num_return_sequences=5,
        top_k=top_k,
        top_p=top_p,
    )
    print(output[0])
    output = tokenizer.decode(output[0])
    print(output)
    return output

gen_sequences = [generate(model, tokenizer, "Lorem").split(" ")]

# Load model from disc
tokenizer = PreTrainedTokenizerFast.from_pretrained(model_folder)
output_config_file = os.path.join(model_folder, CONFIG_NAME)
model = TFGPT2LMHeadModel.from_pretrained(model_folder, config=output_config_file)

print("text", "consectetur", "labore")
encoded = fast_tokenizer.encode("consectetur", "labore")
print("encoded", encoded)

gen_sequences = [generate(model, tokenizer, "Lorem").split(" ")]
```

### Expected behavior

I expect the model to be loaded without any errors.

Comments:
Hi,

This involves a lot of custom code, and the [forums](https://discuss.huggingface.co/) is the place for questions like this. We keep issues for clear bugs in the library and feature requests only.

But if you find a simple `save_pretrained` then `from_pretrained` doesn't work (without all those custom code), we are happy to take a look!

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.