T5 pre-training is now supported in JAX/FLAX. You can check out the example script here: transformers/examples/flax/language-modeling at master · huggingface/transformers · GitHub 41. It actually includes 2 scripts:


t5_tokenizer_model.py 15, to train a T5 tokenizer (i.e. SentencePiece) from scratch.

run_t5_mlm_flax.py 30, to pre-train T5. It’s suited to run on TPUs (for which you can obtain access for free by applying to Google’s TFRC program 5).

@patrickvonplaten also demonstrates how to run the script in this video 14 (starts around 13:35).
This script was developed for the JAX/FLAX community event 5. It would be really cool if someone contributes the PyTorch version of it. It would mean translating the script from FLAX to PyTorch, which is probably straightforward.