Source-linked AI summary
Scaling Up Models and Data with $\texttt{t5x}$ and $\texttt{seqio}$
Adam Roberts, Hyung Won Chung, Anselm Levskaya, Gaurav Mishra, James Bradbury, Daniel Andor, Sharan Narang, Brian Lester, Colin Gaffney, Afroz Mohiuddin, Curtis Hawthorne, Aitor Lewkowycz, Alex Salcianu, Marc van Zee, Jacob Austin, Sebastian Goodman, Livio Baldini Soares, Haitang Hu, Sasha Tsvyashchenko, Aakanksha Chowdhery, Jasmijn Bastings, Jannis Bulian, Xavier Garcia, Jianmo Ni, Andrew Chen, Kathleen Kenealy, Jonathan H. Clark, Stephan Lee, Dan Garrette, James Lee-Thorp, Colin Raffel, Noam Shazeer, Marvin Ritter, Maarten Bosma, Alexandre Passos, Jeremy Maitin-Shepard, Noah Fiedel, Mark Omernick, Brennan Saeta, Ryan Sepassi, Alexander Spiridonov, Joshua Newlan, Andrea Gesmundo
TL;DR
Large-language-model scaling requires specialized systems for distributed computation, data feeding, reproducibility, and consistent evaluation. The paper presents t5x for scalable JAX/Flax model development and seqio for task-based, deterministic data and evaluation pipelines, reporting broad adoption and use in large-scale model training.
Problem
Scaling large language models requires specialized software for distributed training, efficient data pipelines, reproducibility, and consistent downstream evaluation.
Method
The paper introduces t5x for scalable JAX/Flax Transformer training and seqio for task-based, deterministic data processing and evaluation across modeling frameworks.
Results
t5x and seqio achieved widespread internal and external adoption, including training research models from small scale to the largest Google language models.
Takeaways & Limitations
The libraries provide research-friendly infrastructure for scaling models, managing reproducible data workflows, and supporting configurable T5-like and decoder-only model development.
Takeaways & Limitations
t5x assumes the system’s devices can be decomposed into model-parallel and data-parallel axes or submeshes.
Abstract
from arXiv · showhide
Recent neural network-based language models have benefited greatly from scaling up the size of training datasets and the number of parameters in the models themselves. Scaling can be complicated due to various factors including the need to distribute computation on supercomputer clusters (e.g., TPUs), prevent bottlenecks when infeeding data, and ensure reproducible results. In this work, we present two software libraries that ease these issues: $\texttt{t5x}$ simplifies the process of building and training large language models at scale while maintaining ease of use, and $\texttt{seqio}$ provides a task-based API for simple creation of fast and reproducible training data and evaluation pipelines. These open-source libraries have been used to train models with hundreds of billions of parameters on datasets with multiple terabytes of training data. Along with the libraries, we release configurations and instructions for T5-like encoder-decoder models as well as GPT-like decoder-only architectures. $\texttt{t5x}$ and $\texttt{seqio}$ are open source and available at https://github.com/google-research/t5x and https://github.com/google/seqio, respectively.
1. Introduction
t5x and seqio address the practical challenges of scaling Transformer research by combining distributed model training, configurable implementations, and reproducible data workflows. t5x supports flexible partitioning across model, activation, and data dimensions while integrating modular components for training and evaluation.
- Large Transformer models require specialized software for scalable training, rapid research iteration, and consistent downstream evaluation.
- t5x is a JAX-based open-source library for building Transformer models across a wide range of scales.
- t5x supports pretraining, finetuning, evaluation, and inference for Flax-based JAX models, with optimization for TPUs.
- Its modular design separates datasets, evaluation, checkpointing, configuration, models, and partitioning, while allowing use without seqio.
- Partitioning: t5x uses jax.pjit and XLA GSPMD to simplify sharding computation across data, parameters, and activations.
- Partitioning: The system combines orthogonal model- and data-parallel axes and supports multiple parameter and activation partitioning configurations.
- Model integration: Model implementations expose logical tensor axes and wrap Flax modules in t5x.BaseModel to define loss, evaluation, and inference methods.
3. seqio
seqio provides scalable, framework-compatible data processing through a task-based API and deterministic pipeline features. Tasks unify preprocessing, targets, metrics, and architecture-specific conversion, while deterministic processing supports reproducibility and distributed training.
- seqio creates scalable data pipelines for training, inference, and evaluation while converting datasets to NumPy iterators compatible with JAX and PyTorch.
- Task-based API: A seqio Task associates raw data sources with preprocessing, input-target construction, evaluation metrics, and feature conversion.
- Task-based API: The same task can support encoder-decoder and decoder-only architectures, and multiple tasks can be combined into mixtures with user-provided mixing rates.
- Deterministic Pipelines: Deterministic pipelines preserve example order for reproducibility, benchmarking, dataset inspection, and training debugging.
- Deterministic Pipelines: They support resuming from arbitrary training points, avoiding repeated data after restarts and enabling manual skipping of unstable batches.
- Deterministic Pipelines: Distributed caching preprocesses, shuffles, indexes, and shards examples so data-parallel workers can read exclusive file sets efficiently.
- Deterministic Pipelines: The authors report that these features increase throughput, protect against overfitting, ease debugging, and provide fine-grained control over training examples.
4. Example Models
t5x releases minimal, checkpointed implementations of several T5-family models and a scalable T5 variant, alongside a decoder-only configuration compatible with LaMDA. These examples emphasize simple Flax-based implementations and reduced compilation or activation-memory costs.
- t5x provides well-tested minimal model implementations with checkpoints for T5, T5.1.1, mT5, and ByT5.
- The minimal implementations use limited Flax abstractions and closely follow pedagogical Flax examples.
- Scalable T5 uses jax.scan to reduce compilation time and provide finer-grained control over activation memory.
- The release also includes a checkpoint-free decoder-only model configuration compatible with LaMDA.
5. Related Work
t5x builds on a broad ecosystem of open-source sequence-modeling and large-scale training libraries. Its stated differentiators are JAX and Flax model expression, TPU support including TPU v4, and Gin-based configurability.
- Earlier Google systems included Tensor2Tensor, Lingvo, and the Mesh TensorFlow-based T5.
- Related external projects include fairseq, FairScale, DeepSpeed, and Megatron, spanning model libraries and parallelism systems.
- t5x differentiates itself through JAX and Flax model expression, TPU support including TPU v4, and Gin-based configuration.
6. Project Status and Adoption
t5x and seqio have seen broad internal and external adoption, including use for research projects ranging from small-scale studies to Google’s largest language models. The libraries were also validated by reproducing the T5 models.
- t5x has been launched on TPU hundreds of thousands of times at Google, with more than 1,000 total internal users of t5x and seqio.
- Teams use the libraries for research projects spanning small-scale work to the largest language models trained at Google, as well as user-facing products.
- External adopters include academic and commercial Cloud TPU users, including portions of the Big Science project.
- The models were validated by reproducing T5 models originally implemented in Mesh TensorFlow Transformer.
- Users cite the libraries’ usability and research-friendliness as reasons for adoption.
7. Contributions
The project involved leadership, core library designers and implementers, substantial code contributors, advisors, and many additional contributors. Their roles covered overall direction, t5x and seqio development, deterministic pipelines, technical design, and infrastructure.
- Core contributors designed and implemented much of t5x and seqio, including deterministic pipelines, partitioning, and open sourcing.
- Many named contributors made substantial code contributions to the project.
- Additional collaborators helped design deterministic pipelines and seqio, advised on metrics, TensorStore, and overall technical design, and sponsored the project.
- Project leaders contributed to high-level design, roadmap, internal infrastructure, and paper authorship.
- The acknowledgments thank many other contributors to the project.