[Paper Review] Scaling Up Models and Data with $ exttt{t5x}$ and $ exttt{seqio}$

Adam Roberts, Hyung Won Chung+41 more|arXiv (Cornell University)|Mar 31, 2022

Topic ModelingComputer Science48 citations

TL;DR

The paper presents two open-source libraries, t5x and seqio, to simplify scaling of Transformer models and data pipelines, enabling large-scale training and reproducible evaluation across TPU and other hardware.

ABSTRACT

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: $ exttt{t5x}$ simplifies the process of building and training large language models at scale while maintaining ease of use, and $ exttt{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. $ exttt{t5x}$ and $ exttt{seqio}$ are open source and available at https://github.com/google-research/t5x and https://github.com/google/seqio, respectively.

Motivation & Objective

• Motivate the need for scalable, reproducible training of large language models.
• Introduce t5x as a JAX-based library for building, training, evaluating, and inferring with Transformer models at scale.
• Introduce seqio as a task-based data pipeline API for efficient, deterministic, and reproducible data processing.
• Demonstrate how these libraries support encoder-decoder and decoder-only architectures and integrate with existing frameworks.

Proposed method

• Describe the modular architecture of t5x and how it wraps jax.pjit with XLA GSPMD for model, data, and activation partitioning.
• Explain the partitioning options (data vs. model, 1D/2D parameter and activation partitioning) and relate them to known schemes like ZeRO-3 and Megatron.
• Detail how Flax and Gin-config are used to implement models and configurations, including compatibility with legacy T5 and Mesh TensorFlow models.
• Present seqio as a task-based API built on top of tensorflow.data for scalable, deterministic data pipelines and multi-task mixtures.
• Outline the deterministic pipeline features for reproducibility, recoverability, sharding, and global shuffling implemented via Apache Beam.

Experimental results

Research questions

• RQ1 How can we simplify scaling of Transformer models across data and model parallelism axes using high-level abstractions?
• RQ2 How can we ensure reproducible, deterministic data pipelines to enable fair comparisons and efficient debugging at scale?
• RQ3 What are practical configurations and workflows for encoder-decoder vs. decoder-only architectures within t5x and seqio?
• RQ4 How do t5x and seqio integrate with existing model implementations and training ecosystems (Flax, TensorFlow, PyTorch)?

Key findings

• t5x provides a high-level interface to JAX/XLA GSPMD enabling flexible data, parameter, and activation partitioning for large-scale Transformer models.
• seqio enables task-based data pipelines with deterministic pipelines, reproducibility, recoverability, and efficient distributed reads, aiding large-scale training and evaluation.
• The libraries support encoder-decoder and decoder-only model configurations, with compatibility and conversion paths from legacy T5 codebases and Mesh TensorFlow models.
• Open-source release with configurations and guidance for T5-like and GPT-like architectures facilitates rapid experimentation and scaling across TPU environments.
• Adoption within Google and external researchers demonstrates usability and research-friendliness for large-scale language modeling.