GitHub - google-research/t5x

T5X is a modular, composable, research-friendly framework for high-performance, configurable, self-service training, evaluation, and inference of sequence models (starting with language) at many scales. It is essentially a new and improved implementation of the T5 codebase (based on Mesh TensorFlow) in JAX and Flax. Below is a quick start guide for training models with TPUs on Google Cloud. For additional tutorials and background, see the complete documentation. Vertex AI is a platform for training that creates TPU instances and runs code on the TPUs.