T5X is a framework for training, evaluation, and inference of sequence models (starting with language). It is based on JAX and Flax. To learn more, see the T5X Paper.

Please refer to Rosetta T5X, NVIDIA's project that enables seamless training of LLMs, CV models and multimodal models in JAX, for information about running models and experiments on GPUs in T5X.