How to resume run_t5_mlm_flax.py script?
Author: jbmaxwell
Pull request: no

I'm trying to resume training when using the `run_t5_mlm_flax.py` script, but I've noticed that I don't see any checkpoint folders in my `output_dir`—there are a few `events.out.tfevents[...].v2` files, and a `flax_model.msgpack` file. The latter sounds like maybe it's the weights file, but I'm not sure. When I try to run training with `--overwrite_output_dir False` it just seems to start again from scratch. I'm obviously missing something here...

Any help appreciated.

Comments:
Ping @patil-suraj 

Hi @jbmaxwell !

checkpoint loading and saving is currently not implemented in flax examples.

> The latter sounds like maybe it's the weights file, but I'm not sure

Yes, you are right.

the script only saves the model weights and not the optimizer states. You can load the saved weights by passing it's path as
`--model_name_or_path` argument. 

To save and load the optimizer state

```python

# save the opt state
opt_state = jax.device_get(state.opt_state)
with (Path(output_dir) / "opt_state.msgpack").open("wb") as f:
    f.write(to_bytes(opt_state))

with Path(Path(output_dir) / "opt_state.msgpack")).open("rb") as f:
    opt_state =  from_bytes(f.read())
```

Ah, okay. Thanks for the note!

I'm finally getting back to this... I've trained the model on my data which seemed to converge pretty well. 

There are two tasks I'm interested in: 1) sentence infilling, and 2) paraphrasing. Paraphrasing seems like it should be pretty simple since it kind of amounts to "translating" a sentence in its own language, optionally varying `top_p`, `temperature`, etc. 

But I'm really not sure about infilling. I've tried the code posted here— https://github.com/huggingface/transformers/issues/3972 —but it seems very hit-and-miss. Mind you, I do get this warning when the model loads: 
`Some weights of T5ForConditionalGeneration were not initialized from the Flax model and are newly initialized: ['decoder.embed_tokens.weight', 'encoder.embed_tokens.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.`, which makes me think that either it hasn't loaded my trained weights, or there are fine tuning layers that are as yet untrained (I've seen this after pretraining other models in the past). 

Any guidance on how to get sentence infilling working? Any blogs, posts, or threads I should check out? Any help much appreciated. 

Thanks!

> 

Hi @jbmaxwell,
When I use `model = T5ForConditionalGeneration.from_pretrained(pretrained_model_name_or_path_by_flax, from_flax=True)`  I also meet this warning:
`Some weights of T5ForConditionalGeneration were not initialized from the Flax model and are newly initialized: ['decoder.embed_tokens.weight', 'encoder.embed_tokens.weight'] You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.`
but I use `model_flax = FlaxT5ForConditionalGeneration.from_pretrained(pretrained_model_name_or_path_by_flax)` don't appear warning, so I check the `model_flax`'s weights. There is no embed_tokens weight in `model_flax`.
Is this normal?