Unable to save t5 model locally after training t5 using run_t5_mlm_flax.py
### System Info

I want to develop a POC to train a t5 model on a domain dataset (txt file where each line is a sentence). I came across the run_t5_mlm_flax.py file and followed the steps mentioned in the README file.

After a lot of trial and error, I finally was able to get it running in colab using GPU (using 8 batchsize). After running successfully, I am unable to find the saved model anywhere locally within colab (checked provided output directory). Can anyone help me overcome this issue?

This is the command I am using to run the file (t5-trained is a folder I created during runtime):

**python run_t5_mlm_flax.py --output_dir="./t5-trained" --model_type="t5-small" --config_name="./t5-trained" --tokenizer_name="./t5-trained" --train_file="Input_Sent.txt" --max_seq_length="512" --per_device_train_batch_size="8" --per_device_eval_batch_size="8" --adafactor --learning_rate="0.005" --weight_decay="0.001" --warmup_steps="2000" --overwrite_output_dir --logging_steps="500" --save_steps="10000" --eval_steps="2500"**

<img width="513" alt="Screen Shot 2022-12-27 at 4 15 47 PM" src="https://user-images.githubusercontent.com/31246787/209727198-44a9b280-3a2b-43e7-893f-48e672095a90.png">


### Who can help?

_No response_

### Information

- [X] The official example scripts
- [ ] 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

1. Create a small dataset of sentences and save it in txt file. Each line in txt file is a sentence.
2. Create a folder named "t5-trained"
3. Run code to generate tokenizer file using t5_tokenizer_model.py and code mentioned in README.
4. Generate config file using code mentioned in README file.
5. Run this command to train the t5 model -> python run_t5_mlm_flax.py --output_dir="./t5-trained" --model_type="t5-small" --config_name="./t5-trained" --tokenizer_name="./t5-trained" --train_file="Input_Sent.txt" --max_seq_length="512" --per_device_train_batch_size="8" --per_device_eval_batch_size="8" --adafactor --learning_rate="0.005" --weight_decay="0.001" --warmup_steps="2000" --overwrite_output_dir --logging_steps="500" --save_steps="10000" --eval_steps="2500" 

### Expected behavior

Once these steps are completed, I am expecting a saved model somewhere locally which I can import and utilize for text generation or embedding generation from the encoder.
cc @sanchit-gandhi <|||||>Hey @patelvishwa112! Sorry for the late reply here! The fine-tuned checkpoint is saved periodically every `save_steps` training steps:
https://github.com/huggingface/transformers/blob/12313838d33373d06d35b48c3c501fa832f16443/examples/flax/language-modeling/run_t5_mlm_flax.py#L950

It looks as though you're setting `save_steps=10000`, but you're only training for 156 train steps. Since your maximum number of train steps is less than your `save_steps`, we never hit the minimum number of steps required to save the model!

If you set:
```
--save_steps="50"
```
You should see that the model is saved every 50 steps. Since this is less than our total number of train steps, we should see that the model is saved during training for a total of 3 times: at 50, 100 and 150 train steps respectively.

It would indeed be nice to update the examples to save the model at the end of training (irrespective of the value for `save_steps`). Feel free to open a PR for this change if you're interested! I'd be more than happy to help guide you through the process and help with the integration!<|||||>Thank you @sanchit-gandhi for the help. I was able to run the code successfully and it generated flax_model.msgpack (~800MB) file amoung others. 

Can you tell me how can I import this model using either transformers or tensorflow to either get embedding from encoder or use it for text generation? 

And for creating PR request, I will create one just to update the example so that everyone can follow it without error and I will reach out to you if I need any assistance :). <|||||>Hey @patelvishwa112!

You should be able to load the Flax model using:

```python
from transformers import FlaxT5ForConditionalGeneration

model = FlaxT5ForConditionalGeneration.from_pretrained(<path to your checkpoint>)
```

Looking at your training args, the model weights should be saved under `"./t5-trained"`, so this is the path to your checkpoint.

Here's an example of how you can get the encoder embeddings: https://huggingface.co/docs/transformers/model_doc/t5#transformers.FlaxT5ForConditionalGeneration.encode.example

And an example of how you can generate a sequence of text outputs using the Flax T5 model: https://huggingface.co/docs/transformers/model_doc/t5#transformers.FlaxT5ForConditionalGeneration.__call__.example

Hope that helps! Let me know if you have any other questions regarding how to use the trained Flax T5 model for inference 🤗

That sounds good regarding the PR - feel free to open one with the changes required to save the model at the end of training. You can tag me in the PR for a review! Feel free to reach out if you have any questions on the PR - I'm more than happy to help if you have any questions regarding the changes!<|||||>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.